diff --git a/.agents/skills/how-to-write-component/SKILL.md b/.agents/skills/how-to-write-component/SKILL.md index d38bf53529b..9d013876dd9 100644 --- a/.agents/skills/how-to-write-component/SKILL.md +++ b/.agents/skills/how-to-write-component/SKILL.md @@ -24,6 +24,7 @@ Use this as the component decision guide for Dify web. Existing code is referenc - Search before adding UI, hooks, helpers, query utilities, or styling patterns. Reuse existing base components, feature components, hooks, utilities, and design styles when they fit. - Follow Dify's CSS-first Tailwind v4 contract from `packages/dify-ui/README.md` and `packages/dify-ui/AGENTS.md`. Prefer design-system tokens, utilities, and radius mappings over generic Tailwind choices. +- Preserve visible keyboard focus states on the final focusable element. Prefer styled `@langgenius/dify-ui/*` controls when available, because components such as `Button` and form/control primitives carry the standard Dify UI `focus-visible` styling. Do not assume every Dify UI export provides visual focus styles: headless anatomy parts and direct Base UI re-exports such as dialog/popover/tooltip/drawer triggers usually only provide behavior and semantics. When using native `button` / `a`, custom trigger `render` props, clickable rows, icon buttons, menu-like items, or direct trigger parts, verify the rendered focusable element has a visible focus state. If it does not, add the standard Dify UI focus style: `outline-hidden focus-visible:ring-2 focus-visible:ring-state-accent-solid`. Do not hide outlines without an equivalent visible `focus-visible` indicator. Component-specific focus styles should follow an existing styled primitive pattern or a concrete design constraint, not a new ad hoc style. - Group feature code by workflow, route, or ownership area with route-aligned names: components, hooks, local types, query helpers, atoms, constants, tests, and small utilities should live near the code that changes with them. - For each feature module, keep a module-local `README.md` as a boundary note. Start with the module name, a brief one-sentence description, then split dependencies into `Internal Modules` and `External Modules` sections; keep both sections and write `None.` when one category is empty. `Internal Modules` lists modules inside the same overall feature using paths from that feature root, such as `shared/domain/runtime-status`; `External Modules` lists project modules outside the feature using paths from the web root without a `web/` prefix, such as `app/components/base/skeleton`. Omit npm packages, workspace package dependencies, and whitelisted plumbing modules. Do not copy caller-relative import paths into the README. - Module README whitelist: `@/service/client`, `@/next/*`. @@ -53,6 +54,7 @@ Use this as the component decision guide for Dify web. Existing code is referenc - Treat `useParams`, route args, and `nuqs` query state as framework-owned state. When atom logic needs those values, hydrate primitive atoms at the route or surface boundary, such as with `useHydrateAtoms(..., { dangerouslyForceHydrate: true })`; keep URL updates in the route/query-state APIs instead of write atoms. - Within a route-owned feature, choose one source for route identity. If route params are bridged into feature atoms, use that bridge consistently for route-derived queries and actions instead of also threading the same route id through page, tab, and section props. - For async work tied to atom state, use `atomWithQuery` or `atomWithMutation`; write atoms should update only the inputs that drive those atoms. This applies to pure frontend async work as well as network requests, so do not hand-roll loading/error/in-flight state with `useState` or `useRef` for atom-orchestrated async behavior. For component-owned remote work, use `useQuery` or `useMutation` directly. +- `jotai-tanstack-query` query atoms do not support TanStack Query tracked properties. A component that reads `useAtomValue(queryAtom)` subscribes to the whole query result, even if it only accesses `data`, `isLoading`, or `isError`. Export field-specific derived atoms and have components read the exact fields they render; use `selectAtom(queryAtom, result => result.field)` for query-result fields so unchanged selections do not notify subscribers. Keep direct `useAtomValue(queryAtom)` only when the component or hook genuinely needs the full observer result. - Row-local async state belongs to the row owner unless it participates in a shared Jotai workflow or needs atom-scoped reset semantics. - Leave query and mutation atoms unscoped so they keep shared QueryClient cache and invalidation behavior. Scope resettable primitives and explicit hydration tuples; scope a derived atom only when every dependency should be private to that surface. - For scoped primitives that are always hydrated by `ScopeProvider`, prefer `atomWithLazy(() => { throw new Error(...) })` when consumers should see a non-null type. diff --git a/.github/scripts/check-hotfix-cherry-picks.sh b/.github/scripts/check-hotfix-cherry-picks.sh index 11dc024ccf8..cb97dc20699 100644 --- a/.github/scripts/check-hotfix-cherry-picks.sh +++ b/.github/scripts/check-hotfix-cherry-picks.sh @@ -45,7 +45,7 @@ while IFS= read -r commit_sha; do ) if [[ -z "$source_sha" ]]; then - error "Commit $commit_sha ($subject) is missing cherry-pick provenance. $REMEDIATION_HINT" + error "Commit $commit_sha ($subject) is missing cherry-pick provenance. $REMEDIATION_HINT If version differences prevent using git cherry-pick -x, manually add '(cherry picked from commit )' to the commit message." failed=1 continue fi diff --git a/.github/workflows/main-ci.yml b/.github/workflows/main-ci.yml index 8016ff4db8d..0dad5915d00 100644 --- a/.github/workflows/main-ci.yml +++ b/.github/workflows/main-ci.yml @@ -53,6 +53,10 @@ jobs: filters: | api: - 'api/**' + - 'scripts/check_no_new_getattr.py' + - 'scripts/ast_grep_rules/no_new_getattr.yml' + - '.github/workflows/style.yml' + - '.github/workflows/main-ci.yml' - '.github/workflows/api-tests.yml' - 'docker/.env.example' - 'docker/envs/middleware.env.example' @@ -380,6 +384,8 @@ jobs: needs: pre_job if: needs.pre_job.outputs.should_skip != 'true' uses: ./.github/workflows/style.yml + with: + base-rev: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha }} vdb-tests-run: name: Run VDB Tests diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index ceef6a855b7..eecd7f83e5e 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -2,6 +2,10 @@ name: Style check on: workflow_call: + inputs: + base-rev: + required: true + type: string concurrency: group: style-${{ github.head_ref || github.run_id }} @@ -33,6 +37,7 @@ jobs: scripts/check_no_new_getattr.py scripts/ast_grep_rules/no_new_getattr.yml .github/workflows/style.yml + .github/workflows/main-ci.yml - name: Setup UV and Python if: steps.changed-files.outputs.any_changed == 'true' @@ -54,17 +59,9 @@ jobs: if: steps.changed-files.outputs.any_changed == 'true' run: uv run --project api --dev python api/dev/lint_response_contracts.py --fail-on-mismatch - - name: Fetch merge target ref for getattr guard - if: steps.changed-files.outputs.any_changed == 'true' - run: git fetch --no-tags --depth=1 origin +refs/heads/main:refs/remotes/origin/main - - - name: Bind merge target branch for getattr guard - if: steps.changed-files.outputs.any_changed == 'true' - run: git show-ref --verify --quiet refs/heads/main || git branch main origin/main - - name: Run No New Getattr Guard if: steps.changed-files.outputs.any_changed == 'true' - run: uv run --project api python scripts/check_no_new_getattr.py --mode ci --merge-target main + run: uv run --project api python scripts/check_no_new_getattr.py --base-rev "${{ inputs.base-rev }}" - name: Run Type Checks if: steps.changed-files.outputs.any_changed == 'true' diff --git a/.superpowers/sdd/hitl-timeout-semantics-impl-report.md b/.superpowers/sdd/hitl-timeout-semantics-impl-report.md deleted file mode 100644 index ba36a15d7dc..00000000000 --- a/.superpowers/sdd/hitl-timeout-semantics-impl-report.md +++ /dev/null @@ -1,30 +0,0 @@ -# HITL timeout semantics implementation report - -## What changed - -- Updated `api/core/workflow/nodes/human_input/callback.py` so `DifyHITLCallback` now preserves Dify's timeout split at the boundary: - - `HumanInputFormStatus.TIMEOUT` returns the graphon timeout branch via `Expired(selected_handle="__timeout__", ...)`. - - `HumanInputFormStatus.EXPIRED` is treated as an invalid resume state and raises `AssertionError`. - - `HumanInputFormStatus.WAITING` with a past global deadline is treated as an invalid resume state and raises `AssertionError`. - - `HumanInputFormStatus.WAITING` with only the node-level deadline expired still returns the timeout branch. -- Added `created_at` to `HumanInputFormEntity` and `_HumanInputFormEntityImpl` so the callback can compute the global deadline using Dify's shared `HUMAN_INPUT_GLOBAL_TIMEOUT_SECONDS` invariant. -- Kept the submitted and pause flows unchanged. -- Added focused unit coverage in `api/tests/unit_tests/core/workflow/test_human_input_callback.py` for: - - node timeout branch - - global expiration rejection - - waiting-form past node deadline timeout - - waiting-form past global deadline rejection - -## Verification - -- `uv run --project api pytest -o addopts='' api/tests/unit_tests/core/workflow/test_human_input_callback.py api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py -q` -- `git diff --check` - -## Result - -- The focused test set is expected to pass with the new `created_at` boundary in place. -- No unrelated files were modified. - -## Concerns - -- The callback now fails fast on invalid resume states by design. That is intentional, but any caller that previously relied on `EXPIRED` being mapped to the timeout branch will now see an assertion failure instead. diff --git a/api/clients/agent_backend/__init__.py b/api/clients/agent_backend/__init__.py index 172d63858a5..67175c4795c 100644 --- a/api/clients/agent_backend/__init__.py +++ b/api/clients/agent_backend/__init__.py @@ -18,6 +18,7 @@ from clients.agent_backend.errors import ( AgentBackendValidationError, ) from clients.agent_backend.event_adapter import ( + AgentBackendAgentMessageDeltaInternalEvent, AgentBackendDeferredToolCallInternalEvent, AgentBackendInternalEvent, AgentBackendInternalEventType, @@ -46,6 +47,11 @@ from clients.agent_backend.request_builder import ( AgentBackendWorkflowNodeRunInput, redact_for_agent_backend_log, ) +from clients.agent_backend.session_cleanup import ( + AgentBackendSessionCleanupPayload, + AgentBackendSessionCleanupResult, + cleanup_agent_backend_session, +) __all__ = [ "AGENT_SOUL_PROMPT_LAYER_ID", @@ -57,6 +63,7 @@ __all__ = [ "WORKFLOW_NODE_JOB_PROMPT_LAYER_ID", "WORKFLOW_USER_PROMPT_LAYER_ID", "AgentBackendAgentAppRunInput", + "AgentBackendAgentMessageDeltaInternalEvent", "AgentBackendDeferredToolCallInternalEvent", "AgentBackendError", "AgentBackendHTTPError", @@ -73,6 +80,8 @@ __all__ = [ "AgentBackendRunRequestBuilder", "AgentBackendRunStartedInternalEvent", "AgentBackendRunSucceededInternalEvent", + "AgentBackendSessionCleanupPayload", + "AgentBackendSessionCleanupResult", "AgentBackendStreamError", "AgentBackendStreamInternalEvent", "AgentBackendTransportError", @@ -82,6 +91,7 @@ __all__ = [ "FakeAgentBackendRunClient", "FakeAgentBackendScenario", "RuntimeLayerSpec", + "cleanup_agent_backend_session", "create_agent_backend_run_client", "extract_runtime_layer_specs", "redact_for_agent_backend_log", diff --git a/api/clients/agent_backend/event_adapter.py b/api/clients/agent_backend/event_adapter.py index 54c63f15e46..8fdc165ab3c 100644 --- a/api/clients/agent_backend/event_adapter.py +++ b/api/clients/agent_backend/event_adapter.py @@ -5,6 +5,9 @@ The adapter does not define a new cross-service event contract. It consumes workflow Agent Node maps to Graphon/AppQueue events. Deferred external tool calls remain Dify Agent ``run_succeeded`` payloads on the wire; API code turns them into an internal event so workflow pause/session handling stays local to API. +Agent-message deltas are exposed as annotations on ``PydanticAIStreamRunEvent`` +so API code does not have to parse Pydantic AI stream-event internals to +preserve streaming. The terminal answer remains the ``run_succeeded`` output. """ from __future__ import annotations @@ -32,6 +35,7 @@ class AgentBackendInternalEventType(StrEnum): RUN_STARTED = "run_started" STREAM_EVENT = "stream_event" + AGENT_MESSAGE_DELTA = "agent_message_delta" DEFERRED_TOOL_CALL = "deferred_tool_call" RUN_SUCCEEDED = "run_succeeded" RUN_FAILED = "run_failed" @@ -61,6 +65,13 @@ class AgentBackendStreamInternalEvent(AgentBackendInternalEventBase): data: JsonValue +class AgentBackendAgentMessageDeltaInternalEvent(AgentBackendInternalEventBase): + """API-internal agent-message delta emitted independently from raw stream events.""" + + type: Literal[AgentBackendInternalEventType.AGENT_MESSAGE_DELTA] = AgentBackendInternalEventType.AGENT_MESSAGE_DELTA + delta: str + + class AgentBackendRunSucceededInternalEvent(AgentBackendInternalEventBase): """API-internal terminal success event carrying final output and session state.""" @@ -99,6 +110,7 @@ class AgentBackendRunCancelledInternalEvent(AgentBackendInternalEventBase): type AgentBackendInternalEvent = Annotated[ AgentBackendRunStartedInternalEvent | AgentBackendStreamInternalEvent + | AgentBackendAgentMessageDeltaInternalEvent | AgentBackendDeferredToolCallInternalEvent | AgentBackendRunSucceededInternalEvent | AgentBackendRunFailedInternalEvent @@ -121,6 +133,14 @@ class AgentBackendRunEventAdapter: ) ] case PydanticAIStreamRunEvent(): + if event.agent_message_delta: + return [ + AgentBackendAgentMessageDeltaInternalEvent( + run_id=event.run_id, + source_event_id=event.id, + delta=event.agent_message_delta, + ) + ] data = cast(JsonValue, _EVENT_DATA_ADAPTER.dump_python(event.data, mode="json")) event_kind = data.get("event_kind") if isinstance(data, dict) else None return [ diff --git a/api/clients/agent_backend/request_builder.py b/api/clients/agent_backend/request_builder.py index 3b7b78f299b..1fd7b2cd990 100644 --- a/api/clients/agent_backend/request_builder.py +++ b/api/clients/agent_backend/request_builder.py @@ -11,8 +11,9 @@ composition-driven. from __future__ import annotations +import re from collections.abc import Mapping -from typing import ClassVar +from typing import ClassVar, Literal from agenton.compositor import CompositorSessionSnapshot from agenton.compositor.schemas import LayerSessionSnapshot @@ -46,7 +47,6 @@ from dify_agent.protocol import ( LayerExitSignals, RunComposition, RunLayerSpec, - RunPurpose, RuntimeLayerSpec, ) from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator @@ -63,6 +63,7 @@ DIFY_CORE_TOOLS_LAYER_ID = "core_tools" DIFY_KNOWLEDGE_BASE_LAYER_ID = "knowledge" DIFY_ASK_HUMAN_LAYER_ID = "ask_human" DIFY_SHELL_LAYER_ID = "shell" +type AgentConfigVersionKind = Literal["snapshot", "draft", "build_draft"] def _filter_snapshot_to_specs( @@ -104,6 +105,59 @@ def _shell_config_with_drive_ref( return config.model_copy(update={"agent_stub_drive_ref": drive_config.drive_ref}) +def _markdown_backtick_fence(text: str) -> str: + """Choose a fence that will not terminate inside the prompt body.""" + longest_backtick_run = max((len(match.group(0)) for match in re.finditer(r"`+", text)), default=0) + return "`" * max(3, longest_backtick_run + 1) + + +_BUILD_DRAFT_AGENT_SOUL_PROMPT = """You are running in build mode. + +Objective: +- Improve this agent's working environment, configuration, tools, files, notes, + and context so it can handle the intended task well. + +Guidance: +- Treat the intended task as context for setup work, validation, and configuration decisions. +- Perform concrete investigative or setup steps when they help improve or verify the agent configuration. +- Use the installed `dify-agent` CLI when you need to inspect or persist Agent configuration.""" + + +def _wrap_build_draft_agent_soul_prompt(prompt: str | None) -> str: + """Reframe build-draft Agent Soul prompts as preparation work for a future run.""" + prompt_body = (prompt or "").strip() + if not prompt_body: + return _BUILD_DRAFT_AGENT_SOUL_PROMPT + "\n\nIntended task for later normal runs:\nNo task prompt was provided." + fence = _markdown_backtick_fence(prompt_body) + return ( + _BUILD_DRAFT_AGENT_SOUL_PROMPT + + f"\n\nIntended task for later normal runs:\n{fence}text\n{prompt_body}\n{fence}" + ) + + +def _agent_soul_prompt_for_layer( + prompt: str | None, + *, + config_version_kind: AgentConfigVersionKind, +) -> str | None: + """Preserve normal snapshot/draft prompts and only wrap build-draft prompts. + + The API-side layer adapter is the product boundary where Agent Soul text + becomes the model-facing system-prompt layer. ``snapshot`` and normal + ``draft`` runs pass through the original effective prompt unchanged, while + ``build_draft`` always emits a setup prompt. When an original prompt is + present, it is reframed as future-run context and embedded in a fenced + block; when it is blank, the setup instruction is still kept. + """ + if config_version_kind != "build_draft": + if prompt is None: + return None + if not prompt.strip(): + return None + return prompt + return _wrap_build_draft_agent_soul_prompt(prompt) + + class AgentBackendModelConfig(BaseModel): """API-side model/plugin selection before it is converted to Dify Agent layers.""" @@ -163,7 +217,7 @@ class AgentBackendWorkflowNodeRunInput(BaseModel): workflow_node_job_prompt: str user_prompt: str agent_soul_prompt: str | None = None - purpose: RunPurpose = "workflow_node" + agent_config_version_kind: AgentConfigVersionKind = "snapshot" idempotency_key: str | None = None output: AgentBackendOutputConfig | None = None tools: DifyPluginToolsLayerConfig | None = None @@ -212,7 +266,7 @@ class AgentBackendAgentAppRunInput(BaseModel): execution_context: DifyExecutionContextLayerConfig user_prompt: str agent_soul_prompt: str | None = None - purpose: RunPurpose = "agent_app" + agent_config_version_kind: AgentConfigVersionKind = "snapshot" idempotency_key: str | None = None output: AgentBackendOutputConfig | None = None tools: DifyPluginToolsLayerConfig | None = None @@ -261,13 +315,17 @@ class AgentBackendRunRequestBuilder: prompt. """ layers: list[RunLayerSpec] = [] - if run_input.agent_soul_prompt: + agent_soul_prompt = _agent_soul_prompt_for_layer( + run_input.agent_soul_prompt, + config_version_kind=run_input.agent_config_version_kind, + ) + if agent_soul_prompt: layers.append( RunLayerSpec( name=AGENT_SOUL_PROMPT_LAYER_ID, type=PLAIN_PROMPT_LAYER_TYPE_ID, metadata={**run_input.metadata, "origin": "agent_soul"}, - config=PromptLayerConfig(prefix=run_input.agent_soul_prompt), + config=PromptLayerConfig(prefix=agent_soul_prompt), ) ) @@ -419,7 +477,6 @@ class AgentBackendRunRequestBuilder: return CreateRunRequest( composition=RunComposition(layers=layers), - purpose=run_input.purpose, idempotency_key=run_input.idempotency_key, metadata=run_input.metadata, session_snapshot=run_input.session_snapshot, @@ -467,7 +524,6 @@ class AgentBackendRunRequestBuilder: filtered_snapshot = _filter_snapshot_to_specs(session_snapshot, runtime_layer_specs) return CreateRunRequest( composition=RunComposition(layers=layers), - purpose="workflow_node", idempotency_key=idempotency_key, metadata=request_metadata, session_snapshot=filtered_snapshot, @@ -483,13 +539,17 @@ class AgentBackendRunRequestBuilder: ask_human / structured output. """ layers: list[RunLayerSpec] = [] - if run_input.agent_soul_prompt: + agent_soul_prompt = _agent_soul_prompt_for_layer( + run_input.agent_soul_prompt, + config_version_kind=run_input.agent_config_version_kind, + ) + if agent_soul_prompt: layers.append( RunLayerSpec( name=AGENT_SOUL_PROMPT_LAYER_ID, type=PLAIN_PROMPT_LAYER_TYPE_ID, metadata={**run_input.metadata, "origin": "agent_soul"}, - config=PromptLayerConfig(prefix=run_input.agent_soul_prompt), + config=PromptLayerConfig(prefix=agent_soul_prompt), ) ) @@ -649,7 +709,6 @@ class AgentBackendRunRequestBuilder: return CreateRunRequest( composition=RunComposition(layers=layers), - purpose=run_input.purpose, idempotency_key=run_input.idempotency_key, metadata=run_input.metadata, session_snapshot=run_input.session_snapshot, diff --git a/api/clients/agent_backend/session_cleanup.py b/api/clients/agent_backend/session_cleanup.py new file mode 100644 index 00000000000..370ff544e68 --- /dev/null +++ b/api/clients/agent_backend/session_cleanup.py @@ -0,0 +1,100 @@ +"""Shared API-side helper for Agent backend lifecycle-only session cleanup. + +Product code owns local row retirement and background-task dispatch. This module +only adapts persisted cleanup inputs into the public ``dify-agent`` run +protocol, performs the synchronous ``create_run + wait_run`` loop used by Celery +workers, and reports whether the backend cleanup succeeded, was skipped, or +failed. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar, Literal + +from agenton.compositor import CompositorSessionSnapshot +from dify_agent.protocol import RuntimeLayerSpec +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from clients.agent_backend.client import AgentBackendRunClient +from clients.agent_backend.errors import AgentBackendError +from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder + + +class AgentBackendSessionCleanupPayload(BaseModel): + """Serialized cleanup inputs preserved across API and Celery boundaries.""" + + session_snapshot: CompositorSessionSnapshot | None = None + runtime_layer_specs: list[RuntimeLayerSpec] = Field(default_factory=list) + idempotency_key: str | None = None + metadata: dict[str, JsonValue] = Field(default_factory=dict) + timeout_seconds: float = 30.0 + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +@dataclass(frozen=True, slots=True) +class AgentBackendSessionCleanupResult: + """Terminal outcome of one backend cleanup attempt.""" + + status: Literal["succeeded", "skipped", "failed"] + reason: str | None = None + cleanup_run_id: str | None = None + + @classmethod + def succeeded(cls, cleanup_run_id: str) -> AgentBackendSessionCleanupResult: + return cls(status="succeeded", cleanup_run_id=cleanup_run_id) + + @classmethod + def skipped(cls, reason: str) -> AgentBackendSessionCleanupResult: + return cls(status="skipped", reason=reason) + + @classmethod + def failed(cls, reason: str, cleanup_run_id: str | None = None) -> AgentBackendSessionCleanupResult: + return cls(status="failed", reason=reason, cleanup_run_id=cleanup_run_id) + + +def cleanup_agent_backend_session( + *, + payload: AgentBackendSessionCleanupPayload, + client: AgentBackendRunClient | None, + request_builder: AgentBackendRunRequestBuilder | None = None, +) -> AgentBackendSessionCleanupResult: + """Run lifecycle-only cleanup against the Agent backend and report status.""" + if client is None: + return AgentBackendSessionCleanupResult.skipped("no_agent_backend_client") + if payload.session_snapshot is None: + return AgentBackendSessionCleanupResult.skipped("missing_session_snapshot") + if not payload.runtime_layer_specs: + return AgentBackendSessionCleanupResult.skipped("missing_runtime_layer_specs") + + builder = request_builder or AgentBackendRunRequestBuilder() + request = builder.build_cleanup_request( + session_snapshot=payload.session_snapshot, + runtime_layer_specs=payload.runtime_layer_specs, + idempotency_key=payload.idempotency_key, + metadata=payload.metadata, + ) + + try: + response = client.create_run(request) + except AgentBackendError as exc: + return AgentBackendSessionCleanupResult.failed(str(exc)) + + try: + status_response = client.wait_run(response.run_id, timeout_seconds=payload.timeout_seconds) + except AgentBackendError as exc: + return AgentBackendSessionCleanupResult.failed(str(exc), cleanup_run_id=response.run_id) + + if status_response.status != "succeeded": + reason = status_response.error or f"cleanup run ended with status {status_response.status}" + return AgentBackendSessionCleanupResult.failed(reason, cleanup_run_id=response.run_id) + + return AgentBackendSessionCleanupResult.succeeded(response.run_id) + + +__all__ = [ + "AgentBackendSessionCleanupPayload", + "AgentBackendSessionCleanupResult", + "cleanup_agent_backend_session", +] diff --git a/api/commands/account.py b/api/commands/account.py index dfd57d43142..9ea52dfd248 100644 --- a/api/commands/account.py +++ b/api/commands/account.py @@ -25,7 +25,7 @@ def reset_password(email, new_password, password_confirm): return normalized_email = email.strip().lower() - account = AccountService.get_account_by_email_with_case_fallback(db.session, email.strip()) + account = AccountService.get_account_by_email_with_case_fallback(email.strip(), session=db.session()) if not account: click.echo(click.style(f"Account not found for email: {email}", fg="red")) @@ -67,7 +67,7 @@ def reset_email(email, new_email, email_confirm): return normalized_new_email = new_email.strip().lower() - account = AccountService.get_account_by_email_with_case_fallback(db.session, email.strip()) + account = AccountService.get_account_by_email_with_case_fallback(email.strip(), session=db.session()) if not account: click.echo(click.style(f"Account not found for email: {email}", fg="red")) @@ -133,9 +133,9 @@ def create_tenant(email: str, language: str | None = None, name: str | None = No password=new_password, language=language, create_workspace_required=False, - session=db.session, + session=db.session(), ) - TenantService.create_owner_tenant_if_not_exist(account, name, session=db.session) + TenantService.create_owner_tenant_if_not_exist(account, name, session=db.session()) click.echo( click.style( diff --git a/api/commands/data_migration.py b/api/commands/data_migration.py index bd56c41ea44..8c2627601a6 100644 --- a/api/commands/data_migration.py +++ b/api/commands/data_migration.py @@ -9,6 +9,7 @@ from uuid import UUID import click import sqlalchemy as sa import yaml +from sqlalchemy.orm import Session from core.db.session_factory import session_factory from extensions.ext_database import db @@ -108,7 +109,7 @@ def export_migration_data(input_file: str | None, output_file: str | None, overw raw_config = _load_json_object(input_file, "Export config") selection = ExportConfigParser().parse(raw_config) with session_factory.create_session() as session: - result = MigrationExportService().export(session, selection) + result = MigrationExportService().export(selection, session=session) MigrationPackageService().save_package(result.package, output_file, overwrite=overwrite) click.echo(click.style(f"Output written to {output_file}", fg="green")) _render_report(result.report_items, context=_with_output_path(result.report_context, output_file)) @@ -157,7 +158,6 @@ def import_migration_data( package = MigrationPackageService().load_package(input_file) with session_factory.create_session() as session: result = MigrationImportService().import_package( - session, ImportRequest( package=package, cli_target_tenant=target_tenant, @@ -169,6 +169,7 @@ def import_migration_data( create_app_api_token_on_import=create_app_api_token_on_import, ), ), + session=session, ) _render_report(result.report_items, context=result.report_context) except MigrationDataError as exc: @@ -217,7 +218,9 @@ def migration_data_wizard() -> None: default=True, show_default=False, ) - auto_tools = _discover_auto_tools([app for app in apps if app.id in set(app_ids)], include_referenced_tools) + auto_tools = _discover_auto_tools( + [app for app in apps if app.id in set(app_ids)], include_referenced_tools, session=db.session() + ) auto_tools = _resolve_auto_tool_names(tenant.id, auto_tools) _print_auto_tools(auto_tools) additional_tools = _prompt_additional_tools(tenant.id, auto_tools) @@ -253,7 +256,7 @@ def migration_data_wizard() -> None: output_file=output_file, ) with session_factory.create_session() as session: - result = MigrationExportService().export(session, selection) + result = MigrationExportService().export(selection, session=session) MigrationPackageService().save_package(result.package, output_file, overwrite=overwrite) click.echo(click.style(f"Output written to {output_file}", fg="green")) _print_wizard_step("Report") @@ -394,13 +397,13 @@ def _prompt_import_options() -> tuple[bool, bool, str, str]: return include_secrets, create_tokens, id_strategy, conflict_strategy -def _discover_auto_tools(apps: list[App], include_referenced_tools: bool) -> WizardToolMap: +def _discover_auto_tools(apps: list[App], include_referenced_tools: bool, *, session: Session) -> WizardToolMap: auto_tools: WizardToolMap = {"api_tools": {}, "workflow_tools": {}, "mcp_tools": {}} if not include_referenced_tools: return auto_tools discovery_service = DependencyDiscoveryService() for app in apps: - dsl_content = AppDslService.export_dsl(app_model=app, include_secret=False) + dsl_content = AppDslService.export_dsl(app_model=app, session=session, include_secret=False) raw_dsl = yaml.safe_load(dsl_content) if dsl_content else {} dsl = raw_dsl if isinstance(raw_dsl, dict) else {} for dependency in discovery_service.discover_from_dsl(dsl): diff --git a/api/commands/plugin.py b/api/commands/plugin.py index 71c19f842fb..3695c742921 100644 --- a/api/commands/plugin.py +++ b/api/commands/plugin.py @@ -16,7 +16,7 @@ from core.plugin.plugin_service import PluginService from core.tools.utils.system_encryption import encrypt_system_params from extensions.ext_database import db from models import Tenant -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import TenantPluginAutoUpgradeCategory, TenantPluginAutoUpgradeStrategy from models.oauth import DatasourceOauthParamConfig, DatasourceProvider from models.provider_ids import DatasourceProviderID, ToolProviderID from models.source import DataSourceApiKeyAuthBinding, DataSourceOauthBinding @@ -188,13 +188,13 @@ def transform_datasource_credentials(environment: str): firecrawl_plugin_id = "langgenius/firecrawl_datasource" jina_plugin_id = "langgenius/jina_datasource" if environment == "online": - notion_plugin_unique_identifier = plugin_migration._fetch_plugin_unique_identifier(notion_plugin_id) - firecrawl_plugin_unique_identifier = plugin_migration._fetch_plugin_unique_identifier(firecrawl_plugin_id) - jina_plugin_unique_identifier = plugin_migration._fetch_plugin_unique_identifier(jina_plugin_id) + notion_package_identifier = plugin_migration._fetch_latest_package_identifier(notion_plugin_id) + firecrawl_package_identifier = plugin_migration._fetch_latest_package_identifier(firecrawl_plugin_id) + jina_package_identifier = plugin_migration._fetch_latest_package_identifier(jina_plugin_id) else: - notion_plugin_unique_identifier = None - firecrawl_plugin_unique_identifier = None - jina_plugin_unique_identifier = None + notion_package_identifier = None + firecrawl_package_identifier = None + jina_package_identifier = None oauth_credential_type = CredentialType.OAUTH2 api_key_credential_type = CredentialType.API_KEY @@ -219,9 +219,9 @@ def transform_datasource_credentials(environment: str): installed_plugins = installer_manager.list_plugins(tenant_id) installed_plugins_ids = [plugin.plugin_id for plugin in installed_plugins] if notion_plugin_id not in installed_plugins_ids: - if notion_plugin_unique_identifier: + if notion_package_identifier: # install notion plugin - PluginService.install_from_marketplace_pkg(tenant_id, [notion_plugin_unique_identifier]) + PluginService.install_from_marketplace_pkg(tenant_id, [notion_package_identifier]) auth_count = 0 for notion_tenant_credential in notion_tenant_credentials: auth_count += 1 @@ -279,9 +279,9 @@ def transform_datasource_credentials(environment: str): installed_plugins = installer_manager.list_plugins(tenant_id) installed_plugins_ids = [plugin.plugin_id for plugin in installed_plugins] if firecrawl_plugin_id not in installed_plugins_ids: - if firecrawl_plugin_unique_identifier: + if firecrawl_package_identifier: # install firecrawl plugin - PluginService.install_from_marketplace_pkg(tenant_id, [firecrawl_plugin_unique_identifier]) + PluginService.install_from_marketplace_pkg(tenant_id, [firecrawl_package_identifier]) auth_count = 0 for firecrawl_tenant_credential in firecrawl_tenant_credentials: @@ -343,10 +343,10 @@ def transform_datasource_credentials(environment: str): installed_plugins = installer_manager.list_plugins(tenant_id) installed_plugins_ids = [plugin.plugin_id for plugin in installed_plugins] if jina_plugin_id not in installed_plugins_ids: - if jina_plugin_unique_identifier: + if jina_package_identifier: # install jina plugin - logger.debug("Installing Jina plugin %s", jina_plugin_unique_identifier) - PluginService.install_from_marketplace_pkg(tenant_id, [jina_plugin_unique_identifier]) + logger.debug("Installing Jina plugin %s", jina_package_identifier) + PluginService.install_from_marketplace_pkg(tenant_id, [jina_package_identifier]) auth_count = 0 for jina_tenant_credential in jina_tenant_credentials: @@ -406,7 +406,7 @@ def migrate_data_for_plugin(): def _candidate_auto_upgrade_strategy_tenant_ids_stmt(limit: int | None = None): - category_count = len(TenantPluginAutoUpgradeStrategy.PluginCategory) + category_count = len(TenantPluginAutoUpgradeCategory) stmt = ( select(TenantPluginAutoUpgradeStrategy.tenant_id) .group_by(TenantPluginAutoUpgradeStrategy.tenant_id) @@ -472,6 +472,7 @@ def backfill_plugin_auto_upgrade( try: result = PluginAutoUpgradeService.backfill_strategy_categories( current_tenant_id, + session=db.session(), ) except Exception as e: failed_count += 1 diff --git a/api/commands/rbac.py b/api/commands/rbac.py index 0793d11cbb2..be4993920ad 100644 --- a/api/commands/rbac.py +++ b/api/commands/rbac.py @@ -6,6 +6,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import click from sqlalchemy import select +from sqlalchemy.orm import Session from configs import dify_config from core.db.session_factory import session_factory @@ -131,16 +132,35 @@ def _replace_member_role( operator_account_id: str, member_account_id: str, role_id: str, + *, + session: Session, ) -> str: RBACService.MemberRoles.replace( tenant_id=tenant_id, account_id=operator_account_id, member_account_id=member_account_id, role_ids=[role_id], + session=session, ) return member_account_id +def _replace_member_role_with_new_session( + tenant_id: str, + operator_account_id: str, + member_account_id: str, + role_id: str, +) -> str: + with session_factory.create_session() as session: + return _replace_member_role( + tenant_id=tenant_id, + operator_account_id=operator_account_id, + member_account_id=member_account_id, + role_id=role_id, + session=session, + ) + + @click.command( "rbac-migrate-member-roles", help="Migrate legacy workspace member roles into RBAC member-role bindings." ) @@ -217,14 +237,21 @@ def migrate_member_roles_to_rbac( if replace_jobs: if workers == 1: - for member_account_id, resolved_role_id in replace_jobs: - _replace_member_role(workspace_id, owner_account_id, member_account_id, resolved_role_id) - migrated_count += 1 + with session_factory.create_session() as session: + for member_account_id, resolved_role_id in replace_jobs: + _replace_member_role( + workspace_id, + owner_account_id, + member_account_id, + resolved_role_id, + session=session, + ) + migrated_count += 1 else: with ThreadPoolExecutor(max_workers=workers) as executor: futures = [ executor.submit( - _replace_member_role, + _replace_member_role_with_new_session, workspace_id, owner_account_id, member_account_id, diff --git a/api/configs/packaging/pyproject.py b/api/configs/packaging/pyproject.py index 90b1ecba065..c21c02082c5 100644 --- a/api/configs/packaging/pyproject.py +++ b/api/configs/packaging/pyproject.py @@ -6,6 +6,17 @@ class PyProjectConfig(BaseModel): version: str = Field(description="Dify version", default="") +class DifyToolConfig(BaseModel): + min_difyctl_version: str = Field( + description="Oldest difyctl version served on /openapi/v1", + default="0.0.0", + ) + + +class ToolConfig(BaseModel): + dify: DifyToolConfig = Field(default=DifyToolConfig()) + + class PyProjectTomlConfig(BaseSettings): """ configs in api/pyproject.toml @@ -15,3 +26,8 @@ class PyProjectTomlConfig(BaseSettings): description="configs in the project section of pyproject.toml", default=PyProjectConfig(), ) + + tool: ToolConfig = Field( + description="configs in the [tool.*] section of pyproject.toml", + default=ToolConfig(), + ) diff --git a/api/controllers/common/agent_app_parameters.py b/api/controllers/common/agent_app_parameters.py new file mode 100644 index 00000000000..8c2fbccd513 --- /dev/null +++ b/api/controllers/common/agent_app_parameters.py @@ -0,0 +1,54 @@ +from typing import Any + +from sqlalchemy import select + +from core.app.apps.agent_app.app_feature_projection import merge_agent_app_features +from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError +from extensions.ext_database import db +from models.agent import Agent, AgentConfigSnapshot, AgentStatus +from models.agent_config_entities import AgentSoulConfig +from models.model import App + + +def get_published_agent_app_feature_dict_and_user_input_form( + app_model: App, +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Return public Agent App parameters backed by the published Agent Soul.""" + app_model_config = app_model.app_model_config + + agent_id = app_model.bound_agent_id + if not agent_id: + raise AgentAppGeneratorError("Agent App has no bound Agent") + + agent = db.session.scalar( + select(Agent) + .where( + Agent.tenant_id == app_model.tenant_id, + Agent.id == agent_id, + Agent.status == AgentStatus.ACTIVE, + ) + .limit(1) + ) + if agent is None: + raise AgentAppGeneratorError("Agent App has no bound Agent") + # active_config_is_published means the draft has no unpublished edits; the public app + # can still read parameters from the active snapshot while a newer draft is pending. + if not agent.active_config_snapshot_id: + raise AgentAppNotPublishedError("Agent has not been published") + + snapshot = db.session.scalar( + select(AgentConfigSnapshot) + .where( + AgentConfigSnapshot.tenant_id == app_model.tenant_id, + AgentConfigSnapshot.agent_id == agent.id, + AgentConfigSnapshot.id == agent.active_config_snapshot_id, + ) + .limit(1) + ) + if snapshot is None: + raise AgentAppGeneratorError("Agent published version not found") + + agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) + features_dict = merge_agent_app_features(agent_soul=agent_soul, app_model_config=app_model_config) + return features_dict, agent_app_variables_to_user_input_form(agent_soul.app_variables) diff --git a/api/controllers/common/app_access.py b/api/controllers/common/app_access.py index 863b69d2339..214d2de71b4 100644 --- a/api/controllers/common/app_access.py +++ b/api/controllers/common/app_access.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING +from extensions.ext_database import db from services.enterprise import rbac_service as enterprise_rbac_service if TYPE_CHECKING: @@ -76,7 +77,7 @@ def resolve_app_access_filter( inner-API round trip; otherwise it is fetched here. """ if permissions is None: - permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id) + permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id, session=db.session()) whitelist_scope = enterprise_rbac_service.RBACService.AppAccess.whitelist_resources(tenant_id, account_id) can_manage_own_apps = _MANAGE_OWN_APPS_PERMISSION_KEY in permissions.workspace.permission_keys diff --git a/api/controllers/console/agent/composer.py b/api/controllers/console/agent/composer.py index d089772e3ab..f5d71990ade 100644 --- a/api/controllers/console/agent/composer.py +++ b/api/controllers/console/agent/composer.py @@ -16,6 +16,7 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user_id, ) +from extensions.ext_database import db from fields.agent_fields import ( AgentAppComposerResponse, AgentComposerCandidatesResponse, @@ -69,6 +70,7 @@ class WorkflowAgentComposerApi(Resource): node_id=node_id, account_id=account_id, snapshot_id=query.snapshot_id, + session=db.session(), ), ) @@ -94,6 +96,7 @@ class WorkflowAgentComposerApi(Resource): node_id=node_id, account_id=account_id, payload=payload, + session=db.session(), ), ) @@ -126,6 +129,7 @@ class WorkflowAgentComposerCopyFromRosterApi(Resource): source_agent_id=payload.source_agent_id, source_snapshot_id=payload.source_snapshot_id, idempotency_key=payload.idempotency_key, + session=db.session(), ), ) @@ -149,8 +153,9 @@ class WorkflowAgentComposerValidateApi(Resource): tenant_id=tenant_id, payload=payload, agent_id=AgentComposerService.resolve_workflow_node_agent_id( - tenant_id=tenant_id, app_id=app_model.id, node_id=node_id + tenant_id=tenant_id, app_id=app_model.id, node_id=node_id, session=db.session() ), + session=db.session(), ) return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings}) @@ -174,6 +179,7 @@ class WorkflowAgentComposerCandidatesApi(Resource): app_id=app_model.id, node_id=node_id, user_id=current_user_id, + session=db.session(), ), ) @@ -196,7 +202,9 @@ class WorkflowAgentComposerImpactApi(Resource): ) return dump_response( AgentComposerImpactResponse, - AgentComposerService.calculate_impact(tenant_id=tenant_id, current_snapshot_id=current_snapshot_id), + AgentComposerService.calculate_impact( + tenant_id=tenant_id, current_snapshot_id=current_snapshot_id, session=db.session() + ), ) @@ -224,6 +232,7 @@ class WorkflowAgentComposerSaveToRosterApi(Resource): node_id=node_id, account_id=account_id, payload=payload, + session=db.session(), ), ) @@ -238,7 +247,7 @@ class AgentComposerApi(Resource): def get(self, tenant_id: str, agent_id: UUID): return dump_response( AgentAppComposerResponse, - AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id)), + AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session()), ) @console_ns.expect(console_ns.models[ComposerSavePayload.__name__]) @@ -259,6 +268,7 @@ class AgentComposerApi(Resource): agent_id=str(agent_id), account_id=account_id, payload=payload, + session=db.session(), ), ) @@ -274,7 +284,7 @@ class AgentComposerValidateApi(Resource): @account_initialization_required @with_current_tenant_id def post(self, tenant_id: str, agent_id: UUID): - AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id)) + AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session()) payload = ComposerSavePayload.model_validate(console_ns.payload or {}) ComposerConfigValidator.validate_publish_payload(payload) AgentComposerService.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) @@ -282,6 +292,7 @@ class AgentComposerValidateApi(Resource): tenant_id=tenant_id, payload=payload, agent_id=str(agent_id), + session=db.session(), ) return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings}) @@ -303,5 +314,6 @@ class AgentComposerCandidatesApi(Resource): tenant_id=tenant_id, agent_id=str(agent_id), user_id=current_user_id, + session=db.session(), ), ) diff --git a/api/controllers/console/agent/roster.py b/api/controllers/console/agent/roster.py index 349826e54d7..1467cc0c246 100644 --- a/api/controllers/console/agent/roster.py +++ b/api/controllers/console/agent/roster.py @@ -534,7 +534,7 @@ class AgentAppListApi(Resource): status="normal", ) - app_pagination = AppService().get_paginate_apps(current_user.id, current_tenant_id, params, db.session) + app_pagination = AppService().get_paginate_apps(current_user.id, current_tenant_id, params, db.session()) if app_pagination is None: empty = AgentAppPagination(page=args.page, limit=args.limit, total=0, has_more=False, data=[]) return empty.model_dump(mode="json") @@ -567,7 +567,7 @@ class AgentAppListApi(Resource): icon_background=args.icon_background, ) - app = AppService().create_app(current_tenant_id, params, current_user) + app = AppService().create_app(current_tenant_id, params, current_user, session=db.session()) return _serialize_agent_app_detail(app, current_user=current_user), 201 @@ -607,7 +607,7 @@ class AgentAppApi(Resource): "max_active_requests": args.max_active_requests or 0, "role": args.role, } - updated = AppService().update_app(app_model, args_dict) + updated = AppService().update_app(app_model, args_dict, session=db.session()) return _serialize_agent_app_detail(updated, current_user=current_user) @console_ns.response(204, "Agent app deleted successfully") @@ -619,7 +619,7 @@ class AgentAppApi(Resource): @with_current_tenant_id def delete(self, tenant_id: str, agent_id: UUID): app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id) - AppService().delete_app(app_model) + AppService().delete_app(app_model, session=db.session()) return "", 204 @@ -668,6 +668,7 @@ class AgentPublishApi(Resource): agent_id=str(agent_id), account_id=current_user.id, version_note=args.version_note, + session=db.session(), ) @@ -688,6 +689,7 @@ class AgentBuildDraftCheckoutApi(Resource): agent_id=str(agent_id), account_id=current_user.id, force=args.force, + session=db.session(), ) @@ -705,6 +707,7 @@ class AgentBuildDraftApi(Resource): tenant_id=tenant_id, agent_id=str(agent_id), account_id=current_user.id, + session=db.session(), ) @console_ns.expect(console_ns.models[ComposerSavePayload.__name__]) @@ -722,6 +725,7 @@ class AgentBuildDraftApi(Resource): agent_id=str(agent_id), account_id=current_user.id, payload=payload, + session=db.session(), ) @console_ns.response(200, "Agent build draft discarded", console_ns.models[AgentSimpleResultResponse.__name__]) @@ -736,6 +740,7 @@ class AgentBuildDraftApi(Resource): tenant_id=tenant_id, agent_id=str(agent_id), account_id=current_user.id, + session=db.session(), ) @@ -753,6 +758,7 @@ class AgentBuildDraftApplyApi(Resource): tenant_id=tenant_id, agent_id=str(agent_id), account_id=current_user.id, + session=db.session(), ) @@ -810,7 +816,7 @@ class AgentApiStatusApi(Resource): def post(self, tenant_id: str, agent_id: UUID): app_model = _resolve_agent_app_model(tenant_id=tenant_id, agent_id=agent_id) args = AgentApiStatusPayload.model_validate(console_ns.payload) - app_model = AppService().update_app_api_status(app_model, args.enable_api) + app_model = AppService().update_app_api_status(app_model, args.enable_api, session=db.session()) return _serialize_agent_api_access(app_model) diff --git a/api/controllers/console/app/agent.py b/api/controllers/console/app/agent.py index 99164b4755a..81d17ace37a 100644 --- a/api/controllers/console/app/agent.py +++ b/api/controllers/console/app/agent.py @@ -172,7 +172,7 @@ register_response_schema_models( def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None: if node_id and app_model.mode != AppMode.AGENT: return AgentComposerService.resolve_workflow_node_agent_id( - tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id + tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, session=db.session() ) return app_model.bound_agent_id @@ -202,6 +202,7 @@ def _upload_skill_for_app(*, current_user: Account, app_model: App): tenant_id=app_model.tenant_id, user_id=current_user.id, agent_id=agent_id, + session=db.session(), ) except (SkillPackageError, AgentDriveError) as exc: return {"code": exc.code, "message": exc.message}, exc.status_code @@ -240,6 +241,7 @@ def _commit_drive_file_for_app(*, current_user: Account, app_model: App, allow_n value_owned_by_drive=True, ) ], + session=db.session(), ) except AgentDriveError as exc: return {"code": exc.code, "message": exc.message}, exc.status_code @@ -273,6 +275,7 @@ def _delete_drive_file_for_app(*, current_user: Account, app_model: App, allow_n user_id=current_user.id, agent_id=agent_id, items=[DriveCommitItem(key=key, file_ref=None)], + session=db.session(), ) except AgentDriveError as exc: return {"code": exc.code, "message": exc.message}, exc.status_code @@ -298,6 +301,7 @@ def _delete_skill_for_app(*, current_user: Account, app_model: App, slug: str, a DriveCommitItem(key=f"{slug}/SKILL.md", file_ref=None), DriveCommitItem(key=f"{slug}/.DIFY-SKILL-FULL.zip", file_ref=None), ], + session=db.session(), ) except AgentDriveError as exc: return {"code": exc.code, "message": exc.message}, exc.status_code @@ -313,7 +317,9 @@ def _infer_skill_tools_for_app(*, app_model: App, slug: str): if "/" in slug or not slug.strip(): return {"code": "drive_key_invalid", "message": "skill slug must be a single path segment"}, 400 try: - return SkillToolInferenceService().infer(tenant_id=app_model.tenant_id, agent_id=agent_id, slug=slug) + return SkillToolInferenceService().infer( + tenant_id=app_model.tenant_id, agent_id=agent_id, slug=slug, session=db.session() + ) except SkillToolInferenceError as exc: return {"code": exc.code, "message": exc.message}, exc.status_code @@ -335,7 +341,7 @@ class AgentLogApi(Resource): """Get agent logs""" args = AgentLogQuery.model_validate(request.args.to_dict(flat=True)) - return AgentService.get_agent_logs(app_model, args.conversation_id, args.message_id) + return AgentService.get_agent_logs(app_model, args.conversation_id, args.message_id, db.session()) @console_ns.route("/agent//skills/upload") diff --git a/api/controllers/console/app/agent_app_feature.py b/api/controllers/console/app/agent_app_feature.py index 6990886a511..edd2f31f75f 100644 --- a/api/controllers/console/app/agent_app_feature.py +++ b/api/controllers/console/app/agent_app_feature.py @@ -93,7 +93,7 @@ class AgentAppFeatureConfigResource(Resource): app_model=app_model, account=current_user, config=args.model_dump(exclude_none=True), - session=db.session, + session=db.session(), ) app_model_config_was_updated.send(app_model, app_model_config=new_app_model_config) diff --git a/api/controllers/console/app/agent_app_sandbox.py b/api/controllers/console/app/agent_app_sandbox.py index 50cf4ca6b04..6f3811ccdc8 100644 --- a/api/controllers/console/app/agent_app_sandbox.py +++ b/api/controllers/console/app/agent_app_sandbox.py @@ -25,6 +25,7 @@ from controllers.console import console_ns from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model from controllers.console.app.wraps import get_app_model from controllers.console.wraps import account_initialization_required, setup_required, with_current_tenant_id +from extensions.ext_database import db from fields.base import ResponseModel from libs.login import login_required from models.model import App, AppMode @@ -108,14 +109,8 @@ class SandboxReadResponse(ResponseModel): text: str | None = None -class SandboxToolFileResponse(ResponseModel): - transfer_method: Literal["tool_file"] = "tool_file" - reference: str - - class SandboxUploadResponse(ResponseModel): - path: str - file: SandboxToolFileResponse + url: str register_schema_models( @@ -225,7 +220,7 @@ class AgentAppSandboxReadResource(Resource): @console_ns.route("/agent//sandbox/files/upload") class AgentAppSandboxUploadResource(Resource): @console_ns.doc("upload_agent_app_sandbox_file") - @console_ns.doc(description="Upload one Agent App sandbox file as a Dify ToolFile mapping") + @console_ns.doc(description="Upload one Agent App sandbox file and return a signed download URL") @console_ns.expect(console_ns.models[AgentSandboxUploadPayload.__name__]) @console_ns.response(200, "Uploaded", console_ns.models[SandboxUploadResponse.__name__]) @setup_required @@ -275,6 +270,7 @@ class WorkflowAgentSandboxListResource(Resource): node_id=node_id, node_execution_id=query.node_execution_id, path=query.path, + session=db.session(), ) except Exception as exc: return _handle(exc) @@ -311,6 +307,7 @@ class WorkflowAgentSandboxReadResource(Resource): node_id=node_id, node_execution_id=query.node_execution_id, path=query.path, + session=db.session(), ) except Exception as exc: return _handle(exc) @@ -322,7 +319,7 @@ class WorkflowAgentSandboxReadResource(Resource): ) class WorkflowAgentSandboxUploadResource(Resource): @console_ns.doc("upload_workflow_agent_sandbox_file") - @console_ns.doc(description="Upload one workflow Agent sandbox file as a Dify ToolFile mapping") + @console_ns.doc(description="Upload one workflow Agent sandbox file and return a signed download URL") @console_ns.expect(console_ns.models[WorkflowAgentSandboxUploadPayload.__name__]) @console_ns.response(200, "Uploaded", console_ns.models[SandboxUploadResponse.__name__]) @setup_required @@ -340,6 +337,7 @@ class WorkflowAgentSandboxUploadResource(Resource): node_id=node_id, node_execution_id=payload.node_execution_id, path=payload.path, + session=db.session(), ) except Exception as exc: return _handle(exc) diff --git a/api/controllers/console/app/agent_config_inspector.py b/api/controllers/console/app/agent_config_inspector.py index 83824d6434f..0f7aa80ca78 100644 --- a/api/controllers/console/app/agent_config_inspector.py +++ b/api/controllers/console/app/agent_config_inspector.py @@ -253,6 +253,7 @@ def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None: tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, + session=db.session(), ) return app_model.bound_agent_id @@ -288,13 +289,16 @@ def _resolve_console_version( tenant_id=tenant_id, agent_id=agent_id, account_id=account_id, + session=db.session(), ) draft = state.get("draft") or {} draft_id = draft.get("id") if isinstance(draft_id, str) and draft_id: return draft_id, AgentConfigVersionKind.BUILD_DRAFT else: - state = AgentComposerService.load_agent_composer(tenant_id=tenant_id, agent_id=agent_id) + state = AgentComposerService.load_agent_composer( + tenant_id=tenant_id, agent_id=agent_id, session=db.session() + ) draft = state.get("draft") or {} draft_id = draft.get("id") if isinstance(draft_id, str) and draft_id: diff --git a/api/controllers/console/app/agent_drive_inspector.py b/api/controllers/console/app/agent_drive_inspector.py index 473e7364b3e..5166393b3d9 100644 --- a/api/controllers/console/app/agent_drive_inspector.py +++ b/api/controllers/console/app/agent_drive_inspector.py @@ -28,6 +28,7 @@ from controllers.console import console_ns from controllers.console.agent.app_helpers import resolve_agent_runtime_app_model from controllers.console.app.wraps import get_app_model from controllers.console.wraps import account_initialization_required, setup_required, with_current_tenant_id +from extensions.ext_database import db from fields.base import ResponseModel from libs.login import login_required from models.model import App, AppMode @@ -147,7 +148,7 @@ def _resolve_agent_id(app_model: App, node_id: str | None) -> str | None: """Agent identity for the drive: app-bound agent, or the workflow node binding.""" if node_id: return AgentComposerService.resolve_workflow_node_agent_id( - tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id + tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, session=db.session() ) return app_model.bound_agent_id @@ -184,7 +185,9 @@ class AgentDriveListByAgentApi(Resource): query = query_params_from_request(AgentDriveListByAgentQuery) resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id) try: - items = AgentDriveService().manifest(tenant_id=tenant_id, agent_id=str(agent_id), prefix=query.prefix) + items = AgentDriveService().manifest( + tenant_id=tenant_id, agent_id=str(agent_id), prefix=query.prefix, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) return {"items": [{k: v for k, v in item.items() if k != "file_id"} for item in items]} @@ -203,7 +206,7 @@ class AgentDriveSkillListByAgentApi(Resource): def get(self, tenant_id: str, agent_id: UUID): resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id) try: - items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=str(agent_id)) + items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=str(agent_id), session=db.session()) except AgentDriveError as exc: return _handle(exc) return {"items": items} @@ -227,6 +230,7 @@ class AgentDriveSkillInspectByAgentApi(Resource): tenant_id=tenant_id, agent_id=str(agent_id), skill_path=skill_path, + session=db.session(), ) ) except AgentDriveError as exc: @@ -247,7 +251,9 @@ class AgentDrivePreviewByAgentApi(Resource): query = query_params_from_request(AgentDriveFileByAgentQuery) resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id) try: - return AgentDriveService().preview(tenant_id=tenant_id, agent_id=str(agent_id), key=query.key) + return AgentDriveService().preview( + tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) @@ -266,7 +272,9 @@ class AgentDriveDownloadByAgentApi(Resource): query = query_params_from_request(AgentDriveFileByAgentQuery) resolve_agent_runtime_app_model(tenant_id=tenant_id, agent_id=agent_id) try: - url = AgentDriveService().download_url(tenant_id=tenant_id, agent_id=str(agent_id), key=query.key) + url = AgentDriveService().download_url( + tenant_id=tenant_id, agent_id=str(agent_id), key=query.key, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) return {"url": url} @@ -288,7 +296,9 @@ class AgentDriveListApi(Resource): if not agent_id: return _agent_not_bound() try: - items = AgentDriveService().manifest(tenant_id=app_model.tenant_id, agent_id=agent_id, prefix=query.prefix) + items = AgentDriveService().manifest( + tenant_id=app_model.tenant_id, agent_id=agent_id, prefix=query.prefix, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) # the inner manifest exposes file_id for agent-side pulls; the console @@ -312,7 +322,9 @@ class AgentDriveSkillListApi(Resource): if not agent_id: return _agent_not_bound() try: - items = AgentDriveService().list_skills(tenant_id=app_model.tenant_id, agent_id=agent_id) + items = AgentDriveService().list_skills( + tenant_id=app_model.tenant_id, agent_id=agent_id, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) return {"items": items} @@ -345,6 +357,7 @@ class AgentDriveSkillInspectApi(Resource): tenant_id=app_model.tenant_id, agent_id=agent_id, skill_path=skill_path, + session=db.session(), ) ) except AgentDriveError as exc: @@ -367,7 +380,9 @@ class AgentDrivePreviewApi(Resource): if not agent_id: return _agent_not_bound() try: - return AgentDriveService().preview(tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key) + return AgentDriveService().preview( + tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) @@ -388,7 +403,9 @@ class AgentDriveDownloadApi(Resource): if not agent_id: return _agent_not_bound() try: - url = AgentDriveService().download_url(tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key) + url = AgentDriveService().download_url( + tenant_id=app_model.tenant_id, agent_id=agent_id, key=query.key, session=db.session() + ) except AgentDriveError as exc: return _handle(exc) return {"url": url} diff --git a/api/controllers/console/app/annotation.py b/api/controllers/console/app/annotation.py index d14c7d2a7dc..961f9e2f1d8 100644 --- a/api/controllers/console/app/annotation.py +++ b/api/controllers/console/app/annotation.py @@ -211,7 +211,7 @@ class AppAnnotationSettingDetailApi(Resource): @edit_permission_required @rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT) def get(self, app_id: UUID): - result = AppAnnotationService.get_app_annotation_setting_by_app_id(str(app_id)) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(str(app_id), session=db.session()) return dump_response(AnnotationSettingResponse, result), 200 @@ -235,7 +235,7 @@ class AppAnnotationSettingUpdateApi(Resource): setting_args: UpdateAnnotationSettingArgs = {"score_threshold": args.score_threshold} result = AppAnnotationService.update_app_annotation_setting( - str(app_id), annotation_setting_id_str, setting_args + str(app_id), annotation_setting_id_str, setting_args, session=db.session() ) return dump_response(AnnotationSettingResponse, result), 200 @@ -292,7 +292,9 @@ class AnnotationApi(Resource): limit = args.limit keyword = args.keyword - annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(str(app_id), page, limit, keyword) + annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( + str(app_id), page, limit, keyword, session=db.session() + ) annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True) return AnnotationList( data=annotation_models, has_more=len(annotation_list) == limit, limit=limit, total=total, page=page @@ -321,7 +323,9 @@ class AnnotationApi(Resource): upsert_args["message_id"] = args.message_id if args.question is not None: upsert_args["question"] = args.question - annotation = AppAnnotationService.up_insert_app_annotation_from_message(upsert_args, str(app_id)) + annotation = AppAnnotationService.up_insert_app_annotation_from_message( + upsert_args, str(app_id), session=db.session() + ) return dump_response(Annotation, annotation), 201 @setup_required @@ -345,11 +349,11 @@ class AnnotationApi(Resource): }, 400 app_ref = _get_app_ref(str(app_id)) - AppAnnotationService.delete_app_annotations_in_batch(app_ref, annotation_ids) + AppAnnotationService.delete_app_annotations_in_batch(app_ref, annotation_ids, session=db.session()) return "", 204 # If no annotation_ids are provided, handle clearing all annotations else: - AppAnnotationService.clear_all_annotations(str(app_id)) + AppAnnotationService.clear_all_annotations(str(app_id), session=db.session()) return "", 204 @@ -370,7 +374,7 @@ class AnnotationExportApi(Resource): @edit_permission_required @rbac_permission_required(RBACResourceScope.APP, RBACPermission.APP_VIEW_LAYOUT) def get(self, app_id: UUID): - annotation_list = AppAnnotationService.export_annotation_list_by_app_id(str(app_id)) + annotation_list = AppAnnotationService.export_annotation_list_by_app_id(str(app_id), session=db.session()) annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True) return ( AnnotationExportList(data=annotation_models).model_dump(mode="json"), @@ -406,7 +410,7 @@ class AnnotationUpdateDeleteApi(Resource): update_args["question"] = args.question app_ref = _get_app_ref(str(app_id)) annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id)) - annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session) + annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session()) return Annotation.model_validate(annotation, from_attributes=True).model_dump(mode="json") @setup_required @@ -418,7 +422,7 @@ class AnnotationUpdateDeleteApi(Resource): def delete(self, app_id: UUID, annotation_id: UUID): app_ref = _get_app_ref(str(app_id)) annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id)) - AppAnnotationService.delete_app_annotation(annotation_ref, db.session) + AppAnnotationService.delete_app_annotation(annotation_ref, db.session()) return "", 204 @@ -477,7 +481,7 @@ class AnnotationBatchImportApi(Resource): return dump_response( AnnotationBatchImportResponse, - AppAnnotationService.batch_import_app_annotations(str(app_id), file), + AppAnnotationService.batch_import_app_annotations(str(app_id), file, session=db.session()), ) @@ -538,6 +542,7 @@ class AnnotationHitHistoryListApi(Resource): annotation_ref, page, limit, + session=db.session(), ) history_models = TypeAdapter(list[AnnotationHitHistory]).validate_python( annotation_hit_history_list, from_attributes=True diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index 21b9ae434fa..427e1f50d57 100644 --- a/api/controllers/console/app/app.py +++ b/api/controllers/console/app/app.py @@ -631,6 +631,7 @@ class AppListApi(Resource): permissions = enterprise_rbac_service.RBACService.MyPermissions.get( str(current_tenant_id), current_user_id, + session=db.session(), ) if dify_config.RBAC_ENABLED: access_filter = resolve_app_access_filter( @@ -642,7 +643,7 @@ class AppListApi(Resource): # get app list app_service = AppService() - app_pagination = app_service.get_paginate_apps(current_user_id, current_tenant_id, params, db.session) + app_pagination = app_service.get_paginate_apps(current_user_id, current_tenant_id, params, session) if not app_pagination: response = AppPagination(page=args.page, limit=args.limit, total=0, has_more=False, data=[]) return response.model_dump(mode="json"), 200 @@ -697,6 +698,7 @@ class AppListApi(Resource): str(current_tenant_id), current_user.id, [str(app.id)], + session=db.session(), ) app_detail = AppDetailWithSite.model_validate(app, from_attributes=True).model_copy( update={"permission_keys": permission_keys_map.get(str(app.id), [])} @@ -730,7 +732,7 @@ class StarredAppListApi(Resource): is_created_by_me=args.is_created_by_me, ) - app_pagination = AppService().get_paginate_starred_apps(current_user_id, current_tenant_id, params, db.session) + app_pagination = AppService().get_paginate_starred_apps(current_user_id, current_tenant_id, params, session) if not app_pagination: empty = AppPagination(page=args.page, limit=args.limit, total=0, has_more=False, data=[]) return empty.model_dump(mode="json"), 200 @@ -754,7 +756,7 @@ class AppStarApi(Resource): @with_session @get_app_model(mode=None) def post(self, session: Session, current_user_id: str, app_model: App): - AppService.star_app(session, app=app_model, account_id=current_user_id) + AppService.star_app(app=app_model, account_id=current_user_id, session=session) return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.doc("unstar_app") @@ -770,7 +772,7 @@ class AppStarApi(Resource): @with_session @get_app_model(mode=None) def delete(self, session: Session, current_user_id: str, app_model: App): - AppService.unstar_app(session, app=app_model, account_id=current_user_id) + AppService.unstar_app(app=app_model, account_id=current_user_id, session=session) return SimpleResultResponse(result="success").model_dump(mode="json") @@ -802,6 +804,7 @@ class AppApi(Resource): str(current_tenant_id), current_user.id, app_id=str(app_model.id), + session=db.session(), ) permission_keys_map = permissions.app.permission_keys_by_resource_ids([str(app_model.id)]) @@ -838,7 +841,7 @@ class AppApi(Resource): "use_icon_as_answer_icon": args.use_icon_as_answer_icon or False, "max_active_requests": args.max_active_requests or 0, } - app_model = app_service.update_app(app_model, args_dict) + app_model = app_service.update_app(app_model, args_dict, session=db.session()) return dump_response(AppDetailWithSite, app_model) @console_ns.doc("delete_app") @@ -855,7 +858,7 @@ class AppApi(Resource): def delete(self, app_model: App): """Delete app""" app_service = AppService() - app_service.delete_app(app_model) + app_service.delete_app(app_model, session=db.session()) return "", 204 @@ -884,7 +887,7 @@ class AppCopyApi(Resource): with Session(db.engine, expire_on_commit=False) as session: import_service = AppDslService(session) - yaml_content = import_service.export_dsl(app_model=app_model, include_secret=True) + yaml_content = import_service.export_dsl(app_model=app_model, session=session, include_secret=True) result = import_service.import_app( account=current_user, import_mode=ImportMode.YAML_CONTENT, @@ -926,6 +929,7 @@ class AppCopyApi(Resource): str(current_tenant_id), current_user.id, [str(app.id)], + session=db.session(), ) response_model = AppDetailWithSite.model_validate(app, from_attributes=True).model_copy( update={"permission_keys": permission_keys_map.get(str(app.id), [])} @@ -954,6 +958,7 @@ class AppExportApi(Resource): response = AppExportResponse( data=AppDslService.export_dsl( app_model=app_model, + session=db.session(), include_secret=args.include_secret, workflow_id=args.workflow_id, ) @@ -978,7 +983,7 @@ class AppPublishToCreatorsPlatformApi(Resource): if not dify_config.CREATORS_PLATFORM_FEATURES_ENABLED: return {"error": "Creators Platform features are not enabled"}, 403 - dsl_content = AppDslService.export_dsl(app_model=app_model, include_secret=False) + dsl_content = AppDslService.export_dsl(app_model=app_model, session=db.session(), include_secret=False) dsl_bytes = dsl_content.encode("utf-8") claim_code = upload_dsl(dsl_bytes) @@ -1004,7 +1009,7 @@ class AppNameApi(Resource): args = AppNamePayload.model_validate(console_ns.payload) app_service = AppService() - app_model = app_service.update_app_name(app_model, args.name) + app_model = app_service.update_app_name(app_model, args.name, session=db.session()) return dump_response(AppDetail, app_model) @@ -1031,6 +1036,7 @@ class AppIconApi(Resource): args.icon or "", args.icon_background or "", args.icon_type, + session=db.session(), ) return dump_response(AppDetail, app_model) @@ -1053,7 +1059,7 @@ class AppSiteStatus(Resource): args = AppSiteStatusPayload.model_validate(console_ns.payload) app_service = AppService() - app_model = app_service.update_app_site_status(app_model, args.enable_site) + app_model = app_service.update_app_site_status(app_model, args.enable_site, session=db.session()) return dump_response(AppDetail, app_model) @@ -1075,7 +1081,7 @@ class AppApiStatus(Resource): args = AppApiStatusPayload.model_validate(console_ns.payload) app_service = AppService() - app_model = app_service.update_app_api_status(app_model, args.enable_api) + app_model = app_service.update_app_api_status(app_model, args.enable_api, session=db.session()) return dump_response(AppDetail, app_model) diff --git a/api/controllers/console/app/audio.py b/api/controllers/console/app/audio.py index c6cd71f30f1..0c9ed786a1e 100644 --- a/api/controllers/console/app/audio.py +++ b/api/controllers/console/app/audio.py @@ -161,7 +161,7 @@ class ChatMessageTextApi(Resource): # response-contract:ignore return AudioService.transcript_tts( app_model=app_model, - session=db.session, + session=db.session(), text=payload.text, voice=payload.voice, message_ref=message_ref, diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index 1b54f9f9c65..62cf38b86d8 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -36,7 +36,7 @@ from controllers.console.wraps import ( with_current_user_id, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError -from core.app.entities.app_invoke_entities import InvokeFrom +from core.app.entities.app_invoke_entities import AGENT_RUNTIME_EXIT_INTENT_ARG, InvokeFrom from core.app.features.rate_limiting.rate_limit import RateLimitGenerator from core.errors.error import ( ModelCurrentlyNotSupportError, @@ -124,6 +124,7 @@ Use only the current Build chat message history to identify changes that need to validate old config unless the message history already shows that the old config is invalid. Only update the build-draft config note when the current Build chat contains durable context that later runs need. +Write the config note in the language used by the message history. Do not create, update, delete, inspect, or fill gaps in other Agent config resources, including config files, config skills, config env, tools, models, knowledge, or prompt settings. @@ -415,6 +416,7 @@ def _create_build_chat_finalization_message( "draft_type": "debug_build", "conversation_id": debug_conversation_id, "auto_generate_name": False, + AGENT_RUNTIME_EXIT_INTENT_ARG: "delete", } external_trace_id = get_external_trace_id(request) if external_trace_id: diff --git a/api/controllers/console/app/conversation.py b/api/controllers/console/app/conversation.py index a80935e5e33..b7d422d30b1 100644 --- a/api/controllers/console/app/conversation.py +++ b/api/controllers/console/app/conversation.py @@ -200,7 +200,7 @@ class CompletionConversationDetailApi(Resource): conversation_id_str = str(conversation_id) try: - ConversationService.delete(app_model, conversation_id_str, current_user) + ConversationService.delete(app_model, conversation_id_str, current_user, session=db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -354,7 +354,7 @@ class ChatConversationDetailApi(Resource): conversation_id_str = str(conversation_id) try: - ConversationService.delete(app_model, conversation_id_str, current_user) + ConversationService.delete(app_model, conversation_id_str, current_user, session=db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") diff --git a/api/controllers/console/app/message.py b/api/controllers/console/app/message.py index 6a44ca3db8a..958b356de94 100644 --- a/api/controllers/console/app/message.py +++ b/api/controllers/console/app/message.py @@ -363,6 +363,7 @@ def _list_chat_messages(*, app_model: App, current_user: Account | None = None): app_model=app_model, conversation_id=args.conversation_id, user=current_user, + session=db.session(), ) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -474,7 +475,11 @@ def _get_message_suggested_questions(*, current_user: Account, app_model: App, m try: questions = MessageService.get_suggested_questions_after_answer( - app_model=app_model, message_id=message_id_str, user=current_user, invoke_from=InvokeFrom.DEBUGGER + app_model=app_model, + message_id=message_id_str, + user=current_user, + invoke_from=InvokeFrom.DEBUGGER, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message not found") diff --git a/api/controllers/console/app/ops_trace.py b/api/controllers/console/app/ops_trace.py index 46d5ea56e20..e86f65fc035 100644 --- a/api/controllers/console/app/ops_trace.py +++ b/api/controllers/console/app/ops_trace.py @@ -17,6 +17,7 @@ from controllers.console.wraps import ( rbac_permission_required, setup_required, ) +from extensions.ext_database import db from fields.base import ResponseModel from libs.login import login_required from models import App @@ -78,7 +79,7 @@ class TraceAppConfigApi(Resource): try: trace_config = OpsService.get_tracing_app_config( - app_id=app_model.id, tracing_provider=args.tracing_provider + app_id=app_model.id, tracing_provider=args.tracing_provider, session=db.session() ) if not trace_config: return {"has_not_configured": True} @@ -109,7 +110,10 @@ class TraceAppConfigApi(Resource): try: result = OpsService.create_tracing_app_config( - app_id=app_model.id, tracing_provider=args.tracing_provider, tracing_config=args.tracing_config + app_id=app_model.id, + tracing_provider=args.tracing_provider, + tracing_config=args.tracing_config, + session=db.session(), ) if not result: raise TracingConfigIsExist() @@ -142,7 +146,10 @@ class TraceAppConfigApi(Resource): try: result = OpsService.update_tracing_app_config( - app_id=app_model.id, tracing_provider=args.tracing_provider, tracing_config=args.tracing_config + app_id=app_model.id, + tracing_provider=args.tracing_provider, + tracing_config=args.tracing_config, + session=db.session(), ) if not result: raise TracingConfigNotExist() @@ -168,7 +175,9 @@ class TraceAppConfigApi(Resource): args = TraceProviderQuery.model_validate(request.args.to_dict(flat=True)) try: - result = OpsService.delete_tracing_app_config(app_id=app_model.id, tracing_provider=args.tracing_provider) + result = OpsService.delete_tracing_app_config( + app_id=app_model.id, tracing_provider=args.tracing_provider, session=db.session() + ) if not result: raise TracingConfigNotExist() return "", 204 diff --git a/api/controllers/console/app/permission_keys.py b/api/controllers/console/app/permission_keys.py index 810ea04e377..be10f904021 100644 --- a/api/controllers/console/app/permission_keys.py +++ b/api/controllers/console/app/permission_keys.py @@ -1,6 +1,9 @@ +from extensions.ext_database import db from services.enterprise import rbac_service as enterprise_rbac_service def get_app_permission_keys(tenant_id: str, account_id: str | None, app_id: str) -> list[str]: - permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get(tenant_id, account_id, [app_id]) + permission_keys_map = enterprise_rbac_service.RBACService.AppPermissions.batch_get( + tenant_id, account_id, [app_id], session=db.session() + ) return permission_keys_map.get(app_id, []) diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index 609dbfb82c5..53c7c6ea788 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -2,7 +2,7 @@ import json import logging from collections.abc import Sequence from datetime import datetime -from typing import Any, NotRequired, TypedDict, cast +from typing import Any, NotRequired, TypedDict from flask import abort, request from flask_restx import Resource, fields @@ -522,7 +522,7 @@ class DraftWorkflowApi(Resource): """ # fetch draft workflow by app_model workflow_service = WorkflowService() - workflow = workflow_service.get_draft_workflow(app_model=app_model) + workflow = workflow_service.get_draft_workflow(app_model=app_model, session=db.session()) if not workflow: raise DraftWorkflowNotExist() @@ -533,7 +533,7 @@ class DraftWorkflowApi(Resource): # front-end can treat draft graph node data as the editing source. response = WorkflowResponse.model_validate(workflow, from_attributes=True).model_dump(mode="json") response["graph"] = WorkflowAgentPublishService.project_draft_bindings_to_graph( - session=cast(Session, db.session), + session=db.session(), draft_workflow=workflow, ) return response @@ -602,6 +602,7 @@ class DraftWorkflowApi(Resource): account=current_user, environment_variables=environment_variables, conversation_variables=conversation_variables, + session=db.session(), ) except WorkflowHashNotEqualError: raise DraftWorkflowNotSync() @@ -695,7 +696,12 @@ class AdvancedChatDraftRunIterationNodeApi(Resource): try: response = AppGenerateService.generate_single_iteration( - app_model=app_model, user=current_user, node_id=node_id, args=args, streaming=True + app_model=app_model, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -738,7 +744,12 @@ class WorkflowDraftRunIterationNodeApi(Resource): try: response = AppGenerateService.generate_single_iteration( - app_model=app_model, user=current_user, node_id=node_id, args=args, streaming=True + app_model=app_model, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -777,7 +788,12 @@ class AdvancedChatDraftRunLoopNodeApi(Resource): try: response = AppGenerateService.generate_single_loop( - app_model=app_model, user=current_user, node_id=node_id, args=args, streaming=True + app_model=app_model, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -820,7 +836,12 @@ class WorkflowDraftRunLoopNodeApi(Resource): try: response = AppGenerateService.generate_single_loop( - app_model=app_model, user=current_user, node_id=node_id, args=args, streaming=True + app_model=app_model, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -897,6 +918,7 @@ class AdvancedChatDraftHumanInputFormPreviewApi(Resource): account=current_user, node_id=node_id, inputs=inputs, + session=db.session(), ) return jsonable_encoder(preview) @@ -932,6 +954,7 @@ class AdvancedChatDraftHumanInputFormRunApi(Resource): form_inputs=args.form_inputs, inputs=args.inputs, action=args.action, + session=db.session(), ) return jsonable_encoder(result) @@ -963,6 +986,7 @@ class WorkflowDraftHumanInputFormPreviewApi(Resource): account=current_user, node_id=node_id, inputs=inputs, + session=db.session(), ) return jsonable_encoder(preview) @@ -998,6 +1022,7 @@ class WorkflowDraftHumanInputFormRunApi(Resource): form_inputs=args.form_inputs, inputs=args.inputs, action=args.action, + session=db.session(), ) return jsonable_encoder(result) @@ -1028,6 +1053,7 @@ class WorkflowDraftHumanInputDeliveryTestApi(Resource): node_id=node_id, delivery_method_id=args.delivery_method_id, inputs=args.inputs, + session=db.session(), ) return jsonable_encoder({}) @@ -1138,7 +1164,7 @@ class DraftWorkflowNodeRunApi(Resource): workflow_srv = WorkflowService() # fetch draft workflow by app_model - draft_workflow = workflow_srv.get_draft_workflow(app_model=app_model) + draft_workflow = workflow_srv.get_draft_workflow(app_model=app_model, session=db.session()) if not draft_workflow: raise ValueError("Workflow not initialized") files = _parse_file(draft_workflow, args.get("files")) @@ -1181,7 +1207,7 @@ class PublishedWorkflowApi(Resource): """ # fetch published workflow by app_model workflow_service = WorkflowService() - workflow = workflow_service.get_published_workflow(app_model=app_model) + workflow = workflow_service.get_published_workflow(app_model=app_model, session=db.session()) # return workflow, if not found, return None if workflow is None: @@ -1323,7 +1349,9 @@ class ConvertToWorkflowApi(Resource): # convert to workflow mode workflow_service = WorkflowService() - new_app_model = workflow_service.convert_to_workflow(app_model=app_model, account=current_user, args=args) + new_app_model = workflow_service.convert_to_workflow( + app_model=app_model, account=current_user, args=args, session=db.session() + ) # return app id return { @@ -1358,7 +1386,9 @@ class WorkflowFeaturesApi(Resource): features = args.features.model_dump(mode="json", exclude_unset=True) workflow_service = WorkflowService() - workflow_service.update_draft_workflow_features(app_model=app_model, features=features, account=current_user) + workflow_service.update_draft_workflow_features( + app_model=app_model, features=features, account=current_user, session=db.session() + ) return {"result": "success"} @@ -1439,6 +1469,7 @@ class DraftWorkflowRestoreApi(Resource): app_model=app_model, workflow_id=workflow_id, account=current_user, + session=db.session(), ) except IsDraftWorkflowError as exc: raise BadRequest(RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE) from exc @@ -1553,7 +1584,7 @@ class DraftWorkflowNodeLastRunApi(Resource): @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) def get(self, app_model: App, node_id: str): srv = WorkflowService() - workflow = srv.get_draft_workflow(app_model) + workflow = srv.get_draft_workflow(app_model, session=db.session()) if not workflow: raise NotFound("Workflow not found") node_exec = srv.get_node_last_run( @@ -1606,7 +1637,7 @@ class DraftWorkflowTriggerRunApi(Resource): args = DraftWorkflowTriggerRunPayload.model_validate(console_ns.payload or {}) node_id = args.node_id workflow_service = WorkflowService() - draft_workflow = workflow_service.get_draft_workflow(app_model) + draft_workflow = workflow_service.get_draft_workflow(app_model, session=db.session()) if not draft_workflow: raise ValueError("Workflow not found") @@ -1675,7 +1706,7 @@ class DraftWorkflowTriggerNodeApi(Resource): """ workflow_service = WorkflowService() - draft_workflow = workflow_service.get_draft_workflow(app_model) + draft_workflow = workflow_service.get_draft_workflow(app_model, session=db.session()) if not draft_workflow: raise ValueError("Workflow not found") @@ -1759,7 +1790,7 @@ class DraftWorkflowTriggerRunAllApi(Resource): args = DraftWorkflowTriggerRunAllPayload.model_validate(console_ns.payload or {}) node_ids = args.node_ids workflow_service = WorkflowService() - draft_workflow = workflow_service.get_draft_workflow(app_model) + draft_workflow = workflow_service.get_draft_workflow(app_model, session=db.session()) if not draft_workflow: raise ValueError("Workflow not found") @@ -1828,7 +1859,7 @@ class WorkflowOnlineUsersApi(Resource): return {"data": []} workflow_service = WorkflowService() - accessible_app_ids = workflow_service.get_accessible_app_ids(app_ids, current_tenant_id) + accessible_app_ids = workflow_service.get_accessible_app_ids(app_ids, current_tenant_id, session=db.session()) ordered_accessible_app_ids = [app_id for app_id in app_ids if app_id in accessible_app_ids] users_json_by_app_id: dict[str, Any] = {} diff --git a/api/controllers/console/app/workflow_comment.py b/api/controllers/console/app/workflow_comment.py index 64df78f3748..9de91ac59d5 100644 --- a/api/controllers/console/app/workflow_comment.py +++ b/api/controllers/console/app/workflow_comment.py @@ -490,7 +490,7 @@ class WorkflowCommentMentionUsersApi(Resource): current_tenant = current_user.current_tenant # need the tenant object here if current_tenant is None: raise ValueError("current tenant is required") - members = TenantService.get_tenant_members(current_tenant, session=db.session) + members = TenantService.get_tenant_members(current_tenant, session=db.session()) users = TypeAdapter(list[AccountWithRole]).validate_python(members, from_attributes=True) response = WorkflowCommentMentionUsersPayload(users=users) return response.model_dump(mode="json"), 200 diff --git a/api/controllers/console/app/workflow_draft_variable.py b/api/controllers/console/app/workflow_draft_variable.py index 0ccc67f642d..1ffc01e3cef 100644 --- a/api/controllers/console/app/workflow_draft_variable.py +++ b/api/controllers/console/app/workflow_draft_variable.py @@ -337,7 +337,7 @@ class WorkflowVariableCollectionApi(Resource): # fetch draft workflow by app_model workflow_service = WorkflowService() - workflow_exist = workflow_service.is_workflow_exist(app_model=app_model) + workflow_exist = workflow_service.is_workflow_exist(app_model=app_model, session=db.session()) if not workflow_exist: raise DraftWorkflowNotExist() @@ -553,7 +553,7 @@ class VariableResetApi(Resource): ) workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) if draft_workflow is None: raise NotFoundError( f"Draft workflow not found, app_id={app_model.id}", @@ -606,7 +606,7 @@ class ConversationVariableCollectionApi(Resource): # NOTE(QuantumGhost): Prefill conversation variables into the draft variables table # so their IDs can be returned to the caller. workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) if draft_workflow is None: raise NotFoundError(description=f"draft workflow not found, id={app_model.id}") draft_var_srv = WorkflowDraftVariableService(db.session()) @@ -646,6 +646,7 @@ class ConversationVariableCollectionApi(Resource): app_model=app_model, account=current_user, conversation_variables=conversation_variables, + session=db.session(), ) return {"result": "success"} @@ -683,7 +684,7 @@ class EnvironmentVariableCollectionApi(Resource): """ # fetch draft workflow by app_model workflow_service = WorkflowService() - workflow = workflow_service.get_draft_workflow(app_model=app_model) + workflow = workflow_service.get_draft_workflow(app_model=app_model, session=db.session()) if workflow is None: raise DraftWorkflowNotExist() @@ -740,6 +741,7 @@ class EnvironmentVariableCollectionApi(Resource): app_model=app_model, account=current_user, environment_variables=environment_variables, + session=db.session(), ) return {"result": "success"} diff --git a/api/controllers/console/app/workflow_node_output_inspector.py b/api/controllers/console/app/workflow_node_output_inspector.py index 6ed59d6c566..ea45a718a02 100644 --- a/api/controllers/console/app/workflow_node_output_inspector.py +++ b/api/controllers/console/app/workflow_node_output_inspector.py @@ -41,6 +41,7 @@ from controllers.console.wraps import ( rbac_permission_required, setup_required, ) +from extensions.ext_database import db from libs.exception import BaseHTTPException from libs.login import login_required from models import App, AppMode @@ -92,7 +93,9 @@ def _serve_snapshot(app_model: App, run_id: UUID) -> dict: Flask request context. """ try: - snapshot = _service().snapshot_workflow_run(app_model=app_model, workflow_run_id=str(run_id)) + snapshot = _service().snapshot_workflow_run( + app_model=app_model, workflow_run_id=str(run_id), session=db.session() + ) except NodeOutputInspectorError as error: raise _InspectorNotFound(error) from error return snapshot.model_dump(mode="json") @@ -105,6 +108,7 @@ def _serve_node_detail(app_model: App, run_id: UUID, node_id: str) -> dict: app_model=app_model, workflow_run_id=str(run_id), node_id=node_id, + session=db.session(), ) except NodeOutputInspectorError as error: raise _InspectorNotFound(error) from error @@ -119,6 +123,7 @@ def _serve_output_preview(app_model: App, run_id: UUID, node_id: str, output_nam workflow_run_id=str(run_id), node_id=node_id, output_name=output_name, + session=db.session(), ) except NodeOutputInspectorError as error: raise _InspectorNotFound(error) from error @@ -245,7 +250,7 @@ def _stream_inspector_events(app_model: App, run_id: UUID) -> Iterator[str]: # if the run is gone (raised before yielding any bytes, so Flask turns it # into the normal HTTP 404 path). try: - snapshot = service.snapshot_workflow_run(app_model=app_model, workflow_run_id=run_id_str) + snapshot = service.snapshot_workflow_run(app_model=app_model, workflow_run_id=run_id_str, session=db.session()) except NodeOutputInspectorError as error: raise _InspectorNotFound(error) from error @@ -308,6 +313,7 @@ def _stream_inspector_events(app_model: App, run_id: UUID) -> Iterator[str]: app_model=app_model, workflow_run_id=run_id_str, node_id=message.node_id, + session=db.session(), ) except NodeOutputInspectorError: # Node may not appear in the graph yet (race with persistence); skip. diff --git a/api/controllers/console/auth/activate.py b/api/controllers/console/auth/activate.py index b6045685b55..1f58dbe910f 100644 --- a/api/controllers/console/auth/activate.py +++ b/api/controllers/console/auth/activate.py @@ -90,7 +90,7 @@ class ActivateCheckApi(Resource): token = args.token invitation = RegisterService.get_invitation_with_case_fallback( - workspaceId, args.email, token, session=db.session + workspaceId, args.email, token, session=db.session() ) if invitation: data = invitation.get("data", {}) @@ -140,7 +140,7 @@ class ActivateApi(Resource): normalized_request_email = args.email.lower() if args.email else None invitation = RegisterService.get_invitation_with_case_fallback( - args.workspace_id, args.email, args.token, session=db.session + args.workspace_id, args.email, args.token, session=db.session() ) if invitation is None: raise AlreadyActivateError() @@ -178,7 +178,7 @@ class ActivateApi(Resource): RegisterService.revoke_token(args.workspace_id, normalized_request_email, args.token) if membership_id is None: - TenantService.create_tenant_member(tenant, account, db.session, role=role) + TenantService.create_tenant_member(tenant, account, db.session(), role=role) if setup_fields: account.name = setup_fields[0] @@ -188,6 +188,6 @@ class ActivateApi(Resource): account.status = AccountStatus.ACTIVE account.initialized_at = naive_utc_now() - TenantService.switch_tenant(account, tenant.id, session=db.session) + TenantService.switch_tenant(account, tenant.id, session=db.session()) return {"result": "success"} diff --git a/api/controllers/console/auth/data_source_bearer_auth.py b/api/controllers/console/auth/data_source_bearer_auth.py index 11fab84a831..fac725e8534 100644 --- a/api/controllers/console/auth/data_source_bearer_auth.py +++ b/api/controllers/console/auth/data_source_bearer_auth.py @@ -59,7 +59,7 @@ class ApiKeyAuthDataSource(Resource): @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str): - data_source_api_key_bindings = ApiKeyAuthService.get_provider_auth_list(db.session(), current_tenant_id) + data_source_api_key_bindings = ApiKeyAuthService.get_provider_auth_list(current_tenant_id, session=db.session()) if data_source_api_key_bindings: return { "sources": [ @@ -93,7 +93,7 @@ class ApiKeyAuthDataSourceBinding(Resource): data = payload.model_dump() ApiKeyAuthService.validate_api_key_auth_args(data) try: - ApiKeyAuthService.create_provider_auth(db.session(), current_tenant_id, data) + ApiKeyAuthService.create_provider_auth(current_tenant_id, data, session=db.session()) except Exception as e: raise ApiKeyAuthFailedError(str(e)) return {"result": "success"}, 200 @@ -110,6 +110,6 @@ class ApiKeyAuthDataSourceBindingDelete(Resource): @with_current_tenant_id def delete(self, current_tenant_id: str, binding_id: UUID): # The role of the current user in the table must be admin or owner - ApiKeyAuthService.delete_provider_auth(db.session(), current_tenant_id, str(binding_id)) + ApiKeyAuthService.delete_provider_auth(current_tenant_id, str(binding_id), session=db.session()) return "", 204 diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index ba4fc1275d9..d89caa9224f 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -101,7 +101,7 @@ class EmailRegisterSendEmailApi(Resource): if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(normalized_email): raise AccountInFreezeError() - account = AccountService.get_account_by_email_with_case_fallback(db.session, args.email) + account = AccountService.get_account_by_email_with_case_fallback(args.email, session=db.session()) token = AccountService.send_email_register_email(email=normalized_email, account=account, language=language) return {"result": "success", "data": token} @@ -176,7 +176,7 @@ class EmailRegisterResetApi(Resource): email = register_data.get("email", "") normalized_email = email.lower() - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) if account: raise EmailAlreadyInUseError() @@ -187,7 +187,7 @@ class EmailRegisterResetApi(Resource): timezone=args.timezone, language=args.language, ) - token_pair = AccountService.login(account=account, session=db.session, ip_address=extract_remote_ip(request)) + token_pair = AccountService.login(account=account, session=db.session(), ip_address=extract_remote_ip(request)) AccountService.reset_login_error_rate_limit(normalized_email) return {"result": "success", "data": token_pair.model_dump()} @@ -206,7 +206,7 @@ class EmailRegisterResetApi(Resource): password=password, interface_language=get_valid_language(language), timezone=timezone, - session=db.session, + session=db.session(), ) except AccountRegisterError: raise AccountInFreezeError() diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index 8df9600070c..6456bb480f4 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -82,7 +82,7 @@ class ForgotPasswordSendEmailApi(Resource): else: language = "en-US" - account = AccountService.get_account_by_email_with_case_fallback(db.session, args.email) + account = AccountService.get_account_by_email_with_case_fallback(args.email, session=db.session()) token = AccountService.send_reset_password_email( account=account, @@ -180,7 +180,7 @@ class ForgotPasswordResetApi(Resource): password_hashed = hash_password(args.new_password, salt) email = reset_data.get("email", "") - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) if account: account = db.session.merge(account) @@ -198,10 +198,10 @@ class ForgotPasswordResetApi(Resource): # Create workspace if needed if ( - not TenantService.get_join_tenants(account, session=db.session) + not TenantService.get_join_tenants(account, session=db.session()) and FeatureService.get_system_features().is_allow_create_workspace ): - tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session) - TenantService.create_tenant_member(tenant, account, db.session, role="owner") + tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session()) + TenantService.create_tenant_member(tenant, account, db.session(), role="owner") account.current_tenant = tenant tenant_was_created.send(tenant) diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 81f9ee4bae4..486f79bcae2 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -9,7 +9,12 @@ from werkzeug.exceptions import Unauthorized import services from configs import dify_config from constants.languages import get_valid_language -from controllers.common.fields import SimpleResultDataResponse, SimpleResultOptionalDataResponse, SimpleResultResponse +from controllers.common.fields import ( + SimpleResultDataResponse, + SimpleResultMessageResponse, + SimpleResultOptionalDataResponse, + SimpleResultResponse, +) from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.console import console_ns from controllers.console.auth.error import ( @@ -51,7 +56,7 @@ from models.account import Account from services.account_service import AccountService, InvitationDetailDict, RegisterService, TenantService from services.billing_service import BillingService from services.entities.auth_entities import LoginFailureReason, LoginPayloadBase -from services.errors.account import AccountRegisterError +from services.errors.account import AccountRegisterError, RefreshTokenAccountNotFoundError, RefreshTokenNotFoundError from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError from services.feature_service import FeatureService @@ -87,6 +92,7 @@ register_schema_models(console_ns, LoginPayload, EmailPayload, EmailCodeLoginPay register_response_schema_models( console_ns, SimpleResultDataResponse, + SimpleResultMessageResponse, SimpleResultOptionalDataResponse, SimpleResultResponse, ) @@ -120,7 +126,7 @@ class LoginApi(Resource): invitation_data: InvitationDetailDict | None = None if invite_token: invitation_data = RegisterService.get_invitation_with_case_fallback( - None, request_email, invite_token, session=db.session + None, request_email, invite_token, session=db.session() ) if invitation_data is None: invite_token = None @@ -147,23 +153,26 @@ class LoginApi(Resource): _log_console_login_failure(email=normalized_email, reason=LoginFailureReason.INVALID_CREDENTIALS) raise AuthenticationFailedError() from exc # SELF_HOSTED only have one workspace - tenants = TenantService.get_join_tenants(account, session=db.session) + tenants = TenantService.get_join_tenants(account, session=db.session()) if len(tenants) == 0: system_features = FeatureService.get_system_features() if system_features.is_allow_create_workspace and not system_features.license.workspaces.is_available(): raise WorkspacesLimitExceeded() else: - return { - "result": "fail", - "data": "workspace not found, please contact system admin to invite you to join in a workspace", - } + return SimpleResultOptionalDataResponse( + result="fail", + data="workspace not found, please contact system admin to invite you to join in a workspace", + ).model_dump(mode="json") - token_pair = AccountService.login(account=account, session=db.session, ip_address=extract_remote_ip(request)) + token_pair = AccountService.login(account=account, session=db.session(), ip_address=extract_remote_ip(request)) AccountService.reset_login_error_rate_limit(normalized_email) # Create response with cookies instead of returning tokens in body - response = make_response({"result": "success"}) + # response-contract:ignore cookie-bearing Flask response + response = make_response( + SimpleResultOptionalDataResponse(result="success").model_dump(mode="json", exclude_none=True) + ) set_access_token_to_cookie(request, response, token_pair.access_token) set_refresh_token_to_cookie(request, response, token_pair.refresh_token) @@ -178,12 +187,11 @@ class LogoutApi(Resource): @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @with_current_user def post(self, account: Account): - if isinstance(account, flask_login.AnonymousUserMixin): - response = make_response({"result": "success"}) - else: + # response-contract:ignore cookie-bearing Flask response + response = make_response(SimpleResultResponse(result="success").model_dump(mode="json")) + if not isinstance(account, flask_login.AnonymousUserMixin): AccountService.logout(account=account) flask_login.logout_user() - response = make_response({"result": "success"}) # Clear cookies on logout clear_access_token_from_cookie(response) @@ -219,7 +227,7 @@ class ResetPasswordSendEmailApi(Resource): is_allow_register=FeatureService.get_system_features().is_allow_register, ) - return {"result": "success", "data": token} + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @console_ns.route("/email-code-login") @@ -252,7 +260,7 @@ class EmailCodeLoginSendEmailApi(Resource): else: token = AccountService.send_email_code_login_email(account=account, language=language) - return {"result": "success", "data": token} + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @console_ns.route("/email-code-login/validity") @@ -293,7 +301,7 @@ class EmailCodeLoginApi(Resource): _log_console_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_IN_FREEZE) raise AccountInFreezeError() if account: - tenants = TenantService.get_join_tenants(account, session=db.session) + tenants = TenantService.get_join_tenants(account, session=db.session()) if not tenants: workspaces = FeatureService.get_system_features().license.workspaces if not workspaces.is_available(): @@ -301,8 +309,8 @@ class EmailCodeLoginApi(Resource): if not FeatureService.get_system_features().is_allow_create_workspace: raise NotAllowedCreateWorkspace() else: - new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session) - TenantService.create_tenant_member(new_tenant, account, db.session, role="owner") + new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session()) + TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner") account.current_tenant = new_tenant tenant_was_created.send(new_tenant) @@ -313,7 +321,7 @@ class EmailCodeLoginApi(Resource): name=user_email, interface_language=get_valid_language(language), timezone=args.timezone, - session=db.session, + session=db.session(), ) except WorkSpaceNotAllowedCreateError: raise NotAllowedCreateWorkspace() @@ -322,11 +330,12 @@ class EmailCodeLoginApi(Resource): raise AccountInFreezeError() except WorkspacesLimitExceededError: raise WorkspacesLimitExceeded() - token_pair = AccountService.login(account, session=db.session, ip_address=extract_remote_ip(request)) + token_pair = AccountService.login(account, session=db.session(), ip_address=extract_remote_ip(request)) AccountService.reset_login_error_rate_limit(user_email) # Create response with cookies instead of returning tokens in body - response = make_response({"result": "success"}) + # response-contract:ignore cookie-bearing Flask response + response = make_response(SimpleResultResponse(result="success").model_dump(mode="json")) set_csrf_token_to_cookie(request, response, token_pair.csrf_token) # Set HTTP-only secure cookies for tokens @@ -338,45 +347,53 @@ class EmailCodeLoginApi(Resource): @console_ns.route("/refresh-token") class RefreshTokenApi(Resource): @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(401, "Unauthorized", console_ns.models[SimpleResultMessageResponse.__name__]) def post(self): # Get refresh token from cookie instead of request body refresh_token = extract_refresh_token(request) if not refresh_token: - return {"result": "fail", "message": "No refresh token provided"}, 401 + return SimpleResultMessageResponse(result="fail", message="No refresh token provided").model_dump( + mode="json" + ), 401 try: - new_token_pair = AccountService.refresh_token(refresh_token, session=db.session) + new_token_pair = AccountService.refresh_token(refresh_token, session=db.session()) + except Unauthorized as exc: + return SimpleResultMessageResponse(result="fail", message=exc.description or "Unauthorized.").model_dump( + mode="json" + ), 401 + except (RefreshTokenNotFoundError, RefreshTokenAccountNotFoundError) as exc: + return SimpleResultMessageResponse(result="fail", message=str(exc)).model_dump(mode="json"), 401 - # Create response with new cookies - response = make_response({"result": "success"}) + # Create response with new cookies + # response-contract:ignore cookie-bearing Flask response + response = make_response(SimpleResultResponse(result="success").model_dump(mode="json")) - # Update cookies with new tokens - set_csrf_token_to_cookie(request, response, new_token_pair.csrf_token) - set_access_token_to_cookie(request, response, new_token_pair.access_token) - set_refresh_token_to_cookie(request, response, new_token_pair.refresh_token) - return response - except Exception as e: - return {"result": "fail", "message": str(e)}, 401 + # Update cookies with new tokens + set_csrf_token_to_cookie(request, response, new_token_pair.csrf_token) + set_access_token_to_cookie(request, response, new_token_pair.access_token) + set_refresh_token_to_cookie(request, response, new_token_pair.refresh_token) + return response def _get_account_with_case_fallback(email: str): - account = AccountService.get_user_through_email(email, session=db.session) + account = AccountService.get_user_through_email(email, session=db.session()) if account or email == email.lower(): return account - return AccountService.get_user_through_email(email.lower(), session=db.session) + return AccountService.get_user_through_email(email.lower(), session=db.session()) def _authenticate_account_with_case_fallback( original_email: str, normalized_email: str, password: str, invite_token: str | None ): try: - return AccountService.authenticate(original_email, password, invite_token, session=db.session) + return AccountService.authenticate(original_email, password, invite_token, session=db.session()) except services.errors.account.AccountPasswordError: if original_email == normalized_email: raise - return AccountService.authenticate(normalized_email, password, invite_token, session=db.session) + return AccountService.authenticate(normalized_email, password, invite_token, session=db.session()) def _log_console_login_failure(*, email: str, reason: LoginFailureReason) -> None: diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 65f3a5addde..5afafd43131 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -195,7 +195,7 @@ class OAuthCallback(Resource): db.session.commit() try: - TenantService.create_owner_tenant_if_not_exist(account, session=db.session) + TenantService.create_owner_tenant_if_not_exist(account, session=db.session()) except Unauthorized: return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Workspace not found.") except WorkSpaceNotAllowedCreateError: @@ -206,7 +206,7 @@ class OAuthCallback(Resource): token_pair = AccountService.login( account=account, - session=db.session, + session=db.session(), ip_address=extract_remote_ip(request), ) @@ -225,7 +225,7 @@ def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) -> account: Account | None = Account.get_by_openid(provider, user_info.id) if not account: - account = AccountService.get_account_by_email_with_case_fallback(db.session, user_info.email) + account = AccountService.get_account_by_email_with_case_fallback(user_info.email, session=db.session()) return account @@ -241,13 +241,13 @@ def _generate_account( oauth_new_user = False if account: - tenants = TenantService.get_join_tenants(account, session=db.session) + tenants = TenantService.get_join_tenants(account, session=db.session()) if not tenants: if not FeatureService.get_system_features().is_allow_create_workspace: raise WorkSpaceNotAllowedCreateError() else: - new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session) - TenantService.create_tenant_member(new_tenant, account, db.session, role="owner") + new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace", session=db.session()) + TenantService.create_tenant_member(new_tenant, account, db.session(), role="owner") account.current_tenant = new_tenant tenant_was_created.send(new_tenant) @@ -273,10 +273,10 @@ def _generate_account( provider=provider, language=interface_language, timezone=timezone, - session=db.session, + session=db.session(), ) # Link account - AccountService.link_account_integrate(provider, user_info.id, account, session=db.session) + AccountService.link_account_integrate(provider, user_info.id, account, session=db.session()) return account, oauth_new_user diff --git a/api/controllers/console/auth/oauth_server.py b/api/controllers/console/auth/oauth_server.py index 46e2983c12b..d068fb0785e 100644 --- a/api/controllers/console/auth/oauth_server.py +++ b/api/controllers/console/auth/oauth_server.py @@ -10,6 +10,7 @@ from werkzeug.exceptions import BadRequest, NotFound from controllers.common.schema import register_response_schema_models, register_schema_models from controllers.console.wraps import account_initialization_required, setup_required, with_current_user +from extensions.ext_database import db from graphon.model_runtime.utils.encoders import jsonable_encoder from libs.login import login_required from models import Account @@ -131,7 +132,9 @@ def oauth_server_access_token_required[T, **P, R]( response.headers["WWW-Authenticate"] = "Bearer" return response - account = OAuthServerService.validate_oauth_access_token(oauth_provider_app.client_id, access_token) + account = OAuthServerService.validate_oauth_access_token( + oauth_provider_app.client_id, access_token, db.session() + ) if not account: response = jsonify({"error": "access_token or client_id is invalid"}) response.status_code = 401 diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index d6974fe129c..3a983b50176 100644 --- a/api/controllers/console/billing/billing.py +++ b/api/controllers/console/billing/billing.py @@ -56,7 +56,7 @@ class Subscription(Resource): @with_current_tenant_id def get(self, current_tenant_id: str, current_user: Account): args = SubscriptionQuery.model_validate(request.args.to_dict(flat=True)) - BillingService.is_tenant_owner_or_admin(db.session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=db.session()) return BillingService.get_subscription(args.plan, args.interval, current_user.email, current_tenant_id) @@ -70,7 +70,7 @@ class Invoices(Resource): @with_current_user @with_current_tenant_id def get(self, current_tenant_id: str, current_user: Account): - BillingService.is_tenant_owner_or_admin(db.session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=db.session()) return BillingService.get_invoices(current_user.email, current_tenant_id) diff --git a/api/controllers/console/datasets/data_source.py b/api/controllers/console/datasets/data_source.py index b2c8bda0581..17f027df9b3 100644 --- a/api/controllers/console/datasets/data_source.py +++ b/api/controllers/console/datasets/data_source.py @@ -245,7 +245,7 @@ class DataSourceNotionListApi(Resource): exist_page_ids = [] # import notion in the exist dataset if query.dataset_id: - dataset = DatasetService.get_dataset(query.dataset_id, db.session) + dataset = DatasetService.get_dataset(query.dataset_id, db.session()) if not dataset: raise NotFound("Dataset not found.") if dataset.data_source_type != "notion_import": @@ -400,11 +400,11 @@ class DataSourceNotionDatasetSyncApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT) def get(self, dataset_id: UUID) -> tuple[dict[str, str], int]: dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - documents = DocumentService.get_document_by_dataset_id(dataset_id_str, db.session) + documents = DocumentService.get_document_by_dataset_id(dataset_id_str, db.session()) for document in documents: document_indexing_sync_task.delay(dataset_id_str, document.id) return {"result": "success"}, 200 @@ -420,11 +420,11 @@ class DataSourceNotionDocumentSyncApi(Resource): def get(self, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]: dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if document is None: raise NotFound("Document not found.") document_indexing_sync_task.delay(dataset_id_str, document_id_str) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 0bef535d82b..c8ca1d621c4 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -418,6 +418,7 @@ class DatasetListApi(Resource): permissions = enterprise_rbac_service.RBACService.MyPermissions.get( str(current_tenant_id), current_user.id, + session=db.session(), ) accessible_dataset_ids: list[str] | None = None @@ -461,7 +462,7 @@ class DatasetListApi(Resource): datasets, total = DatasetService.get_datasets( query.page, query.limit, - db.session, + db.session(), current_tenant_id, current_user, query.keyword, @@ -573,6 +574,7 @@ class DatasetListApi(Resource): current_tenant_id, current_user.id, [dataset.id], + session=session, ) item = DatasetDetailWithPartialMembersResponse.model_validate(dataset, from_attributes=True).model_dump( @@ -602,17 +604,18 @@ class DatasetApi(Resource): @with_current_tenant_id def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) permissions = enterprise_rbac_service.RBACService.MyPermissions.get( current_tenant_id, current_user.id, dataset_id=dataset_id_str, + session=db.session(), ) permission_keys_map = permissions.dataset.permission_keys_by_resource_ids([dataset_id_str]) data = dump_response(DatasetDetailResponse, dataset) @@ -622,7 +625,7 @@ class DatasetApi(Resource): provider_id = ModelProviderID(dataset.embedding_model_provider) data["embedding_model_provider"] = str(provider_id) if data.get("permission") == "partial_members": - part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session) + part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session()) data.update({"partial_member_list": part_users_list}) # check embedding setting @@ -666,7 +669,7 @@ class DatasetApi(Resource): @with_session def patch(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") @@ -685,10 +688,10 @@ class DatasetApi(Resource): # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator if not dify_config.RBAC_ENABLED: DatasetPermissionService.check_permission( - session, current_user, dataset, payload.permission, payload.partial_member_list + current_user, dataset, payload.permission, payload.partial_member_list, session=session ) - dataset = DatasetService.update_dataset(session, dataset_id_str, payload_data, current_user) + dataset = DatasetService.update_dataset(dataset_id_str, payload_data, current_user, session=session) if dataset is None: raise NotFound("Dataset not found.") @@ -697,6 +700,7 @@ class DatasetApi(Resource): current_tenant_id, current_user.id, [dataset_id_str], + session=session, ) result_data = dump_response(DatasetDetailResponse, dataset) result_data["permission_keys"] = permission_keys_map.get(dataset_id_str, []) @@ -704,13 +708,13 @@ class DatasetApi(Resource): if payload.partial_member_list is not None and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM: DatasetPermissionService.update_partial_member_list( - tenant_id, dataset_id_str, payload.partial_member_list, db.session + tenant_id, dataset_id_str, payload.partial_member_list, db.session() ) # clear partial member list when permission is only_me or all_team_members elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}: - DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session) + DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session()) - partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session) + partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session()) result_data.update({"partial_member_list": partial_member_list}) return result_data, 200 @@ -729,8 +733,8 @@ class DatasetApi(Resource): raise Forbidden() try: - if DatasetService.delete_dataset(dataset_id_str, current_user, db.session): - DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session) + if DatasetService.delete_dataset(dataset_id_str, current_user, db.session()): + DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session()) return "", 204 else: raise NotFound("Dataset not found.") @@ -755,7 +759,7 @@ class DatasetUseCheckApi(Resource): def get(self, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, db.session) + dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, db.session()) return {"is_using": dataset_is_using}, 200 @@ -776,12 +780,12 @@ class DatasetQueryApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) def get(self, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -917,16 +921,16 @@ class DatasetRelatedAppListApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) def get(self, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - app_dataset_joins = DatasetService.get_related_apps(dataset.id, db.session) + app_dataset_joins = DatasetService.get_related_apps(dataset.id, db.session()) related_apps = [] for app_dataset_join in app_dataset_joins: @@ -1101,7 +1105,7 @@ class DatasetEnableApiApi(Resource): def post(self, dataset_id: UUID, status: str): dataset_id_str = str(dataset_id) - DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", db.session) + DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", db.session()) return {"result": "success"}, 200 @@ -1170,10 +1174,10 @@ class DatasetErrorDocs(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) def get(self, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, db.session) + results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, db.session()) return dump_response(ErrorDocsResponse, {"data": results, "total": len(results)}), 200 @@ -1197,15 +1201,15 @@ class DatasetPermissionUserListApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) def get(self, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session) + partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session()) return dump_response(PartialMemberListResponse, {"data": partial_members_list}), 200 @@ -1227,8 +1231,8 @@ class DatasetAutoDisableLogApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) def get(self, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, db.session) + auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, db.session()) return dump_response(AutoDisableLogsResponse, auto_disable_logs), 200 diff --git a/api/controllers/console/datasets/datasets_document.py b/api/controllers/console/datasets/datasets_document.py index a6263c8e2e3..ee441704b20 100644 --- a/api/controllers/console/datasets/datasets_document.py +++ b/api/controllers/console/datasets/datasets_document.py @@ -183,16 +183,16 @@ class DocumentResource(Resource): def get_document( self, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str ) -> Document: - dataset = DatasetService.get_dataset(dataset_id, db.session) + dataset = DatasetService.get_dataset(dataset_id, db.session()) if not dataset: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - document = DocumentService.get_document(dataset_id, document_id, session=db.session) + document = DocumentService.get_document(dataset_id, document_id, session=db.session()) if not document: raise NotFound("Document not found.") @@ -203,16 +203,16 @@ class DocumentResource(Resource): return document def get_batch_documents(self, dataset_id: str, batch: str, current_user: Account) -> Sequence[Document]: - dataset = DatasetService.get_dataset(dataset_id, db.session) + dataset = DatasetService.get_dataset(dataset_id, db.session()) if not dataset: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - documents = DocumentService.get_batch_documents(dataset_id, batch, db.session) + documents = DocumentService.get_batch_documents(dataset_id, batch, db.session()) if not documents: raise NotFound("Documents not found.") @@ -243,13 +243,13 @@ class GetProcessRuleApi(Resource): # get the latest process rule document = db.get_or_404(Document, document_id) - dataset = DatasetService.get_dataset(document.dataset_id, db.session) + dataset = DatasetService.get_dataset(document.dataset_id, db.session()) if not dataset: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -319,12 +319,12 @@ class DatasetDocumentListApi(Resource): ) except (ArgumentTypeError, ValueError, Exception): fetch = False - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -376,6 +376,7 @@ class DatasetDocumentListApi(Resource): documents=documents, dataset=dataset, tenant_id=current_tenant_id, + session=db.session(), ) if fetch: @@ -423,7 +424,7 @@ class DatasetDocumentListApi(Resource): def post(self, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") @@ -433,7 +434,7 @@ class DatasetDocumentListApi(Resource): raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -447,9 +448,9 @@ class DatasetDocumentListApi(Resource): try: documents, batch = DocumentService.save_document_with_dataset_id( - dataset, knowledge_config, current_user, session=db.session + dataset, knowledge_config, current_user, session=db.session() ) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -468,7 +469,7 @@ class DatasetDocumentListApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) def delete(self, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") # check user's model setting @@ -477,7 +478,7 @@ class DatasetDocumentListApi(Resource): try: document_ids = request.args.getlist("document_id") dataset_ref = DatasetRefService.create_dataset_ref(dataset) - DocumentService.delete_documents(dataset_ref, document_ids, dataset.doc_form, db.session) + DocumentService.delete_documents(dataset_ref, document_ids, dataset.doc_form, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Cannot delete document during indexing.") @@ -536,7 +537,7 @@ class DatasetInitApi(Resource): tenant_id=current_tenant_id, knowledge_config=knowledge_config, account=current_user, - session=db.session, + session=db.session(), ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -873,7 +874,7 @@ class DocumentApi(DocumentResource): if metadata == "only": response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details} elif metadata == "without": - dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session) + dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session()) document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {} response = { "id": document.id, @@ -907,7 +908,7 @@ class DocumentApi(DocumentResource): "need_summary": document.need_summary if document.need_summary is not None else False, } else: - dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session) + dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session()) document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {} response = { "id": document.id, @@ -956,7 +957,7 @@ class DocumentApi(DocumentResource): def delete(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") # check user's model setting @@ -965,7 +966,7 @@ class DocumentApi(DocumentResource): document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id) try: - DocumentService.delete_document(document, db.session) + DocumentService.delete_document(document, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Cannot delete document during indexing.") @@ -989,7 +990,7 @@ class DocumentDownloadApi(DocumentResource): def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID) -> dict[str, Any]: # Reuse the shared permission/tenant checks implemented in DocumentResource. document = self.get_document(str(dataset_id), str(document_id), current_user, current_tenant_id) - return {"url": DocumentService.get_document_download_url(document, db.session)} + return {"url": DocumentService.get_document_download_url(document, db.session())} @console_ns.route("/datasets//documents/download-zip") @@ -1019,7 +1020,7 @@ class DocumentBatchDownloadZipApi(DocumentResource): document_ids=document_ids, tenant_id=current_tenant_id, current_user=current_user, - session=db.session, + session=db.session(), ) # Delegate ZIP packing to FileService, but keep Flask response+cleanup in the route. @@ -1168,7 +1169,7 @@ class DocumentStatusApi(DocumentResource): self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable", "archive", "un_archive"] ): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") @@ -1180,12 +1181,12 @@ class DocumentStatusApi(DocumentResource): DatasetService.check_dataset_model_setting(dataset) # check user's permission - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) document_ids = request.args.getlist("document_id") try: - DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session) + DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session()) except services.errors.document.DocumentIndexingError as e: raise InvalidActionError(str(e)) except ValueError as e: @@ -1209,11 +1210,11 @@ class DocumentPauseApi(DocumentResource): dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) # 404 if document not found if document is None: @@ -1225,7 +1226,7 @@ class DocumentPauseApi(DocumentResource): try: # pause document - DocumentService.pause_document(document, db.session) + DocumentService.pause_document(document, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Cannot pause completed document.") @@ -1244,10 +1245,10 @@ class DocumentRecoverApi(DocumentResource): """recover document.""" dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) # 404 if document not found if document is None: @@ -1258,7 +1259,7 @@ class DocumentRecoverApi(DocumentResource): raise ArchivedDocumentImmutableError() try: # pause document - DocumentService.recover_document(document, db.session) + DocumentService.recover_document(document, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Document is not in paused status.") @@ -1278,13 +1279,13 @@ class DocumentRetryApi(DocumentResource): """retry document.""" payload = DocumentRetryPayload.model_validate(console_ns.payload or {}) dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) retry_documents = [] if not dataset: raise NotFound("Dataset not found.") for document_id in payload.document_ids: try: - document = DocumentService.get_document(dataset.id, document_id, session=db.session) + document = DocumentService.get_document(dataset.id, document_id, session=db.session()) # 404 if document not found if document is None: @@ -1302,7 +1303,7 @@ class DocumentRetryApi(DocumentResource): logger.exception("Failed to retry document, document id: %s", document_id) continue # retry document - DocumentService.retry_document(dataset_id_str, retry_documents, db.session) + DocumentService.retry_document(dataset_id_str, retry_documents, db.session()) return "", 204 @@ -1320,14 +1321,14 @@ class DocumentRenameApi(DocumentResource): # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator if not current_user.is_dataset_editor: raise Forbidden() - dataset = DatasetService.get_dataset(dataset_id, db.session) + dataset = DatasetService.get_dataset(dataset_id, db.session()) if not dataset: raise NotFound("Dataset not found.") - DatasetService.check_dataset_operator_permission(current_user, dataset, session=db.session) + DatasetService.check_dataset_operator_permission(current_user, dataset, session=db.session()) payload = DocumentRenamePayload.model_validate(console_ns.payload or {}) try: - document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name, db.session) + document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Cannot delete document during indexing.") @@ -1345,11 +1346,11 @@ class WebsiteDocumentSyncApi(DocumentResource): def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID): """sync website document.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") document_id_str = str(document_id) - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") if document.tenant_id != current_tenant_id: @@ -1360,7 +1361,7 @@ class WebsiteDocumentSyncApi(DocumentResource): if DocumentService.check_archived(document): raise ArchivedDocumentImmutableError() # sync document - DocumentService.sync_website_document(dataset_id_str, document, db.session) + DocumentService.sync_website_document(dataset_id_str, document, db.session()) return {"result": "success"}, 200 @@ -1380,10 +1381,10 @@ class DocumentPipelineExecutionLogApi(DocumentResource): dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") log = db.session.scalar( @@ -1438,7 +1439,7 @@ class DocumentGenerateSummaryApi(Resource): dataset_id_str = str(dataset_id) # Get dataset - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") @@ -1447,7 +1448,7 @@ class DocumentGenerateSummaryApi(Resource): raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -1472,7 +1473,7 @@ class DocumentGenerateSummaryApi(Resource): raise ValueError("Summary index is not enabled for this dataset. Please enable it in the dataset settings.") # Verify all documents exist and belong to the dataset - documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, db.session) + documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, db.session()) if len(documents) != len(document_list): found_ids = {doc.id for doc in documents} @@ -1488,7 +1489,7 @@ class DocumentGenerateSummaryApi(Resource): DocumentService.update_documents_need_summary( dataset_id=dataset_id_str, document_ids=document_ids_to_update, - session=db.session, + session=db.session(), need_summary=True, ) @@ -1539,13 +1540,13 @@ class DocumentSummaryStatusApi(DocumentResource): document_id_str = str(document_id) # Get dataset - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # Check permissions try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -1555,7 +1556,7 @@ class DocumentSummaryStatusApi(DocumentResource): result = SummaryIndexService.get_document_summary_status_detail( document_id=document_id_str, dataset_id=dataset_id_str, - session=db.session, + session=db.session(), ) return result, 200 diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index 5cccd2453dc..e4f2abeb844 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/api/controllers/console/datasets/datasets_segments.py @@ -173,7 +173,7 @@ def _get_segment_for_document( raise NotFound("Document not found.") segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id) - segment = SegmentService.get_segment_by_ref(segment_ref) + segment = SegmentService.get_segment_by_ref(segment_ref, db.session()) if not segment: raise NotFound("Segment not found.") return segment_ref, segment @@ -193,16 +193,16 @@ class DatasetDocumentSegmentListApi(Resource): def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): dataset_id_str = str(dataset_id) document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -278,7 +278,7 @@ class DatasetDocumentSegmentListApi(Resource): summaries: dict[str, str | None] = {} if segment_ids: summary_records = SummaryIndexService.get_segments_summaries( - segment_ids=segment_ids, dataset_id=dataset_id_str + segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session() ) summaries = {chunk_id: summary.summary_content for chunk_id, summary in summary_records.items()} @@ -303,14 +303,14 @@ class DatasetDocumentSegmentListApi(Resource): def delete(self, current_user: Account, dataset_id: UUID, document_id: UUID): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") segment_ids = request.args.getlist("segment_id") @@ -319,10 +319,10 @@ class DatasetDocumentSegmentListApi(Resource): if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - SegmentService.delete_segments(segment_ids, document, dataset, db.session) + SegmentService.delete_segments(segment_ids, document, dataset, db.session()) return "", 204 @@ -348,11 +348,11 @@ class DatasetDocumentSegmentApi(Resource): action: Literal["enable", "disable"], ): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # check user's model setting @@ -362,7 +362,7 @@ class DatasetDocumentSegmentApi(Resource): raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: @@ -388,7 +388,7 @@ class DatasetDocumentSegmentApi(Resource): if cache_result is not None: raise InvalidActionError("Document is being indexed, please try again later") try: - SegmentService.update_segments_status(segment_ids, action, dataset, document, db.session) + SegmentService.update_segments_status(segment_ids, action, dataset, document, db.session()) except Exception as e: raise InvalidActionError(str(e)) return dump_response(SimpleResultResponse, {"result": "success"}), 200 @@ -411,12 +411,12 @@ class DatasetDocumentSegmentAddApi(Resource): def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") if not current_user.is_dataset_editor: @@ -438,15 +438,20 @@ class DatasetDocumentSegmentAddApi(Resource): except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) # validate args payload = SegmentCreatePayload.model_validate(console_ns.payload or {}) payload_dict = payload.model_dump(exclude_none=True) SegmentService.segment_create_args_validate(payload_dict, document) - segment = type_cast(DocumentSegment, SegmentService.create_segment(payload_dict, document, dataset, db.session)) - summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id_str) + segment = type_cast( + DocumentSegment, + SegmentService.create_segment(payload_dict, document, dataset, db.session()), + ) + summary = SummaryIndexService.get_segment_summary( + segment_id=segment.id, dataset_id=dataset_id_str, session=db.session() + ) response = { "data": segment_response_with_summary(segment, summary.summary_content if summary else None), "doc_form": document.doc_form, @@ -472,21 +477,21 @@ class DatasetDocumentSegmentUpdateApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: @@ -518,9 +523,11 @@ class DatasetDocumentSegmentUpdateApi(Resource): segment, document, dataset, - db.session, + db.session(), + ) + summary = SummaryIndexService.get_segment_summary( + segment_id=segment.id, dataset_id=dataset_id_str, session=db.session() ) - summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id_str) response = { "data": segment_response_with_summary(segment, summary.summary_content if summary else None), "doc_form": document.doc_form, @@ -541,26 +548,26 @@ class DatasetDocumentSegmentUpdateApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) segment_id_str = str(segment_id) _, segment = _get_segment_for_document(dataset, document, segment_id_str) - SegmentService.delete_segment(segment, document, dataset, db.session) + SegmentService.delete_segment(segment, document, dataset, db.session()) return "", 204 @@ -583,12 +590,12 @@ class DatasetDocumentSegmentBatchImportApi(Resource): def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -658,18 +665,18 @@ class ChildChunkAddApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) # check embedding model setting @@ -693,7 +700,7 @@ class ChildChunkAddApi(Resource): # validate args try: payload = ChildChunkCreatePayload.model_validate(console_ns.payload or {}) - child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session) + child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session()) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200 @@ -709,14 +716,14 @@ class ChildChunkAddApi(Resource): def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") segment_id_str = str(segment_id) @@ -759,21 +766,21 @@ class ChildChunkAddApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) segment_id_str = str(segment_id) @@ -781,7 +788,7 @@ class ChildChunkAddApi(Resource): # validate args payload = ChildChunkBatchUpdatePayload.model_validate(console_ns.payload or {}) try: - child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, db.session) + child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, db.session()) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) return dump_response(ChildChunkBatchUpdateResponse, {"data": child_chunks}), 200 @@ -811,31 +818,31 @@ class ChildChunkUpdateApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) segment_id_str = str(segment_id) segment_ref, _ = _get_segment_for_document(dataset, document, segment_id_str) child_chunk_id_str = str(child_chunk_id) - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref) + child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session()) if not child_chunk: raise NotFound("Child chunk not found.") try: - SegmentService.delete_child_chunk(child_chunk, dataset, db.session) + SegmentService.delete_child_chunk(child_chunk, dataset, db.session()) except ChildChunkDeleteIndexServiceError as e: raise ChildChunkDeleteIndexError(str(e)) return "", 204 @@ -862,34 +869,34 @@ class ChildChunkUpdateApi(Resource): ): # check dataset dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if not dataset: raise NotFound("Dataset not found.") # check user's model setting DatasetService.check_dataset_model_setting(dataset) # check document document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor if not current_user.is_dataset_editor: raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) segment_id_str = str(segment_id) segment_ref, segment = _get_segment_for_document(dataset, document, segment_id_str) child_chunk_id_str = str(child_chunk_id) - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref) + child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session()) if not child_chunk: raise NotFound("Child chunk not found.") # validate args try: payload = ChildChunkUpdatePayload.model_validate(console_ns.payload or {}) child_chunk = SegmentService.update_child_chunk( - payload.content, child_chunk, segment, document, dataset, db.session + payload.content, child_chunk, segment, document, dataset, db.session() ) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index 5b036641d4d..9cdca96f69a 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -299,7 +299,9 @@ class ExternalApiTemplateApi(Resource): if not (current_user.has_edit_permission or current_user.is_dataset_operator): raise Forbidden() - ExternalDatasetService.delete_external_knowledge_api(session, current_tenant_id, external_knowledge_api_id_str) + ExternalDatasetService.delete_external_knowledge_api( + current_tenant_id, external_knowledge_api_id_str, session=session + ) return "", 204 @@ -318,9 +320,7 @@ class ExternalApiUseCheckApi(Resource): external_knowledge_api_id_str = str(external_knowledge_api_id) external_knowledge_api_is_using, count = ExternalDatasetService.external_knowledge_api_use_check( - session, - external_knowledge_api_id_str, - current_tenant_id, + external_knowledge_api_id_str, current_tenant_id, session=session ) return {"is_using": external_knowledge_api_is_using, "count": count}, 200 @@ -366,6 +366,7 @@ class ExternalDatasetCreateApi(Resource): str(current_tenant_id), current_user.id, [dataset_id_str], + session=session, ) item["permission_keys"] = permission_keys_map.get(dataset_id_str, []) @@ -393,12 +394,12 @@ class ExternalKnowledgeHitTestingApi(Resource): @with_session def post(self, session: Session, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) diff --git a/api/controllers/console/datasets/hit_testing_base.py b/api/controllers/console/datasets/hit_testing_base.py index cc02a990168..656a426c125 100644 --- a/api/controllers/console/datasets/hit_testing_base.py +++ b/api/controllers/console/datasets/hit_testing_base.py @@ -86,12 +86,12 @@ class DatasetsHitTestingBase: dataset_id: str, current_user: Account | None = None, current_tenant_id: str | None = None ) -> Dataset: current_user, _ = resolve_account_fallback(current_user, current_tenant_id) - dataset = DatasetService.get_dataset(dataset_id, db.session) + dataset = DatasetService.get_dataset(dataset_id, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) diff --git a/api/controllers/console/datasets/metadata.py b/api/controllers/console/datasets/metadata.py index 8802fcf2814..42ae4903673 100644 --- a/api/controllers/console/datasets/metadata.py +++ b/api/controllers/console/datasets/metadata.py @@ -61,13 +61,13 @@ class DatasetMetadataCreateApi(Resource): metadata_args = MetadataArgs.model_validate(console_ns.payload or {}) dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) metadata = MetadataService.create_metadata( - db.session(), dataset_id_str, metadata_args, current_user, current_tenant_id + dataset_id_str, metadata_args, current_user, current_tenant_id, session=db.session() ) return dump_response(DatasetMetadataResponse, metadata), 201 @@ -81,10 +81,10 @@ class DatasetMetadataCreateApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT) def get(self, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - metadata = MetadataService.get_dataset_metadatas(db.session(), dataset) + metadata = MetadataService.get_dataset_metadatas(dataset, session=db.session()) return dump_response(DatasetMetadataListResponse, metadata), 200 @@ -105,13 +105,13 @@ class DatasetMetadataApi(Resource): dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) metadata = MetadataService.update_metadata_name( - db.session(), dataset_id_str, metadata_id_str, name, current_user, current_tenant_id + dataset_id_str, metadata_id_str, name, current_user, current_tenant_id, session=db.session() ) return dump_response(DatasetMetadataResponse, metadata), 200 @@ -125,12 +125,12 @@ class DatasetMetadataApi(Resource): def delete(self, current_user: Account, dataset_id: UUID, metadata_id: UUID): dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) - MetadataService.delete_metadata(db.session(), dataset_id_str, metadata_id_str) + MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session=db.session()) # Frontend callers only await success and invalidate metadata caches; no response body is consumed. return "", 204 @@ -162,16 +162,16 @@ class DatasetMetadataBuiltInFieldActionApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) def post(self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable"]): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) match action: case "enable": - MetadataService.enable_built_in_field(db.session(), dataset) + MetadataService.enable_built_in_field(dataset, session=db.session()) case "disable": - MetadataService.disable_built_in_field(db.session(), dataset) + MetadataService.disable_built_in_field(dataset, session=db.session()) # Frontend callers only await success and invalidate metadata caches; no response body is consumed. return "", 204 @@ -191,14 +191,14 @@ class DocumentMetadataEditApi(Resource): @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) def post(self, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) metadata_args = MetadataOperationData.model_validate(console_ns.payload or {}) - MetadataService.update_documents_metadata(db.session(), dataset, metadata_args, current_user) + MetadataService.update_documents_metadata(dataset, metadata_args, current_user, session=db.session()) # Frontend callers only await success and invalidate caches; no response body is consumed. return "", 204 diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py index a575760ee19..57d6b628d4b 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field from werkzeug.exceptions import Forbidden, NotFound from configs import dify_config -from controllers.common.fields import RedirectResponse, SimpleResultResponse +from controllers.common.fields import SimpleResultResponse from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.console import console_ns from controllers.console.wraps import ( @@ -19,11 +19,14 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user, ) +from core.entities.provider_entities import ProviderConfig from core.plugin.entities.plugin_daemon import PluginOAuthAuthorizationUrlResponse from core.plugin.impl.oauth import OAuthHandler +from core.tools.entities.common_entities import I18nObject +from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.errors.validate import CredentialsValidateFailedError -from graphon.model_runtime.utils.encoders import jsonable_encoder +from libs.helper import dump_response from libs.login import login_required from models import Account from models.provider_ids import DatasourceProviderID @@ -33,7 +36,9 @@ from services.plugin.oauth_service import OAuthProxyService class DatasourceCredentialPayload(BaseModel): name: str | None = Field(default=None, max_length=100) - credentials: dict[str, Any] + credentials: dict[str, Any] = Field( + description="Plugin-defined credential parameters. The schema is declared by the datasource provider." + ) class DatasourceCredentialDeletePayload(BaseModel): @@ -43,11 +48,17 @@ class DatasourceCredentialDeletePayload(BaseModel): class DatasourceCredentialUpdatePayload(BaseModel): credential_id: str name: str | None = Field(default=None, max_length=100) - credentials: dict[str, Any] | None = Field(default=None) + credentials: dict[str, Any] | None = Field( + default=None, + description="Plugin-defined credential parameters. The schema is declared by the datasource provider.", + ) class DatasourceCustomClientPayload(BaseModel): - client_params: dict[str, Any] | None = Field(default=None) + client_params: dict[str, Any] | None = Field( + default=None, + description="Plugin-defined OAuth client parameters. The schema is declared by the datasource provider.", + ) enable_oauth_custom_client: bool | None = None @@ -71,8 +82,48 @@ class DatasourceOAuthCallbackQuery(BaseModel): context_id: str | None = Field(default=None, description="OAuth proxy context ID") -class DatasourceCredentialsResponse(ResponseModel): - result: Any +class DatasourceCredentialResponse(ResponseModel): + credential: dict[str, Any] = Field( + description="Obfuscated plugin-defined credential parameters from the datasource provider." + ) + type: str + name: str + avatar_url: str | None + id: str + is_default: bool + + +class DatasourceCredentialListResponse(ResponseModel): + result: list[DatasourceCredentialResponse] + + +class DatasourceOAuthSchemaResponse(ResponseModel): + client_schema: list[ProviderConfig] + credentials_schema: list[ProviderConfig] + oauth_custom_client_params: dict[str, Any] | None = Field( + description="Masked plugin-defined OAuth client parameters, when configured for the tenant." + ) + is_oauth_custom_client_enabled: bool + is_system_oauth_params_exists: bool + redirect_uri: str + + +class DatasourceProviderAuthResponse(ResponseModel): + author: str + provider: str + plugin_id: str + plugin_unique_identifier: str + icon: str + name: str + label: I18nObject + description: I18nObject + credential_schema: list[ProviderConfig] + oauth_schema: DatasourceOAuthSchemaResponse | None + credentials_list: list[DatasourceCredentialResponse] + + +class DatasourceProviderAuthListResponse(ResponseModel): + result: list[DatasourceProviderAuthResponse] register_schema_models( @@ -88,9 +139,9 @@ register_schema_models( ) register_response_schema_models( console_ns, - DatasourceCredentialsResponse, + DatasourceCredentialListResponse, + DatasourceProviderAuthListResponse, PluginOAuthAuthorizationUrlResponse, - RedirectResponse, SimpleResultResponse, ) @@ -100,7 +151,7 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): @console_ns.doc(params=query_params_from_model(DatasourceOAuthAuthorizationQuery)) @console_ns.response( 200, - "Authorization URL retrieved successfully", + "Datasource OAuth authorization URL generated successfully", console_ns.models[PluginOAuthAuthorizationUrlResponse.__name__], ) @setup_required @@ -140,7 +191,8 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): redirect_uri=redirect_uri, system_credentials=oauth_config, ) - response = make_response(jsonable_encoder(authorization_url_response)) + # response-contract:ignore cookie-bearing Flask response + response = make_response(dump_response(PluginOAuthAuthorizationUrlResponse, authorization_url_response)) response.set_cookie( "context_id", context_id, @@ -154,11 +206,8 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): @console_ns.route("/oauth/plugin//datasource/callback") class DatasourceOAuthCallback(Resource): @console_ns.doc(params=query_params_from_model(DatasourceOAuthCallbackQuery)) - @console_ns.response( - 302, - "Redirect to console OAuth callback page", - console_ns.models[RedirectResponse.__name__], - ) + # response-contract:ignore redirect response + @console_ns.response(302, "Redirect to OAuth callback page") @setup_required def get(self, provider_id: str): context_id = request.cookies.get("context_id") or request.args.get("context_id") @@ -217,7 +266,9 @@ class DatasourceOAuthCallback(Resource): @console_ns.route("/auth/plugin/datasource/") class DatasourceAuth(Resource): @console_ns.expect(console_ns.models[DatasourceCredentialPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response( + 200, "Datasource credential created successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -238,12 +289,16 @@ class DatasourceAuth(Resource): ) except CredentialsValidateFailedError as ex: raise ValueError(str(ex)) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 + @console_ns.response( + 200, + "Datasource credentials retrieved successfully", + console_ns.models[DatasourceCredentialListResponse.__name__], + ) @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__]) @with_current_user @with_current_tenant_id def get(self, current_tenant_id: str, user: Account, provider_id: str): @@ -255,8 +310,9 @@ class DatasourceAuth(Resource): provider=datasource_provider_id.provider_name, plugin_id=datasource_provider_id.plugin_id, user=user, + session=db.session(), ) - return {"result": datasources}, 200 + return dump_response(DatasourceCredentialListResponse, {"result": datasources}), 200 @console_ns.route("/auth/plugin/datasource//delete") @@ -281,14 +337,17 @@ class DatasourceAuthDeleteApi(Resource): auth_id=payload.credential_id, provider=provider_name, plugin_id=plugin_id, + session=db.session(), ) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/auth/plugin/datasource//update") class DatasourceAuthUpdateApi(Resource): @console_ns.expect(console_ns.models[DatasourceCredentialUpdatePayload.__name__]) - @console_ns.response(201, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response( + 201, "Datasource credential updated successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -308,39 +367,53 @@ class DatasourceAuthUpdateApi(Resource): credentials=payload.credentials or {}, name=payload.name, ) - return {"result": "success"}, 201 + return SimpleResultResponse(result="success").model_dump(mode="json"), 201 @console_ns.route("/auth/plugin/datasource/list") class DatasourceAuthListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__]) + @console_ns.response( + 200, + "Datasource credentials retrieved successfully", + console_ns.models[DatasourceProviderAuthListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str): datasource_provider_service = DatasourceProviderService() - datasources = datasource_provider_service.get_all_datasource_credentials(tenant_id=current_tenant_id) - return {"result": jsonable_encoder(datasources)}, 200 + datasources = datasource_provider_service.get_all_datasource_credentials( + tenant_id=current_tenant_id, session=db.session() + ) + return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 @console_ns.route("/auth/plugin/datasource/default-list") class DatasourceHardCodeAuthListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__]) + @console_ns.response( + 200, + "Default datasource credentials retrieved successfully", + console_ns.models[DatasourceProviderAuthListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str): datasource_provider_service = DatasourceProviderService() - datasources = datasource_provider_service.get_hard_code_datasource_credentials(tenant_id=current_tenant_id) - return {"result": jsonable_encoder(datasources)}, 200 + datasources = datasource_provider_service.get_hard_code_datasource_credentials( + tenant_id=current_tenant_id, session=db.session() + ) + return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 @console_ns.route("/auth/plugin/datasource//custom-client") class DatasourceAuthOauthCustomClient(Resource): @console_ns.expect(console_ns.models[DatasourceCustomClientPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response( + 200, "Datasource OAuth custom client saved successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -357,7 +430,7 @@ class DatasourceAuthOauthCustomClient(Resource): client_params=payload.client_params or {}, enabled=payload.enable_oauth_custom_client or False, ) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @setup_required @login_required @@ -371,7 +444,7 @@ class DatasourceAuthOauthCustomClient(Resource): tenant_id=current_tenant_id, datasource_provider_id=datasource_provider_id, ) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/auth/plugin/datasource//default") @@ -393,7 +466,7 @@ class DatasourceAuthDefaultApi(Resource): datasource_provider_id=datasource_provider_id, credential_id=payload.id, ) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/auth/plugin/datasource//update-name") @@ -416,4 +489,4 @@ class DatasourceUpdateProviderNameApi(Resource): name=payload.name, credential_id=payload.credential_id, ) - return {"result": "success"}, 200 + return SimpleResultResponse(result="success").model_dump(mode="json"), 200 diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py b/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py index b0af108444c..873ba130064 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py @@ -3,12 +3,13 @@ from typing import Any from flask_restx import ( # type: ignore Resource, # type: ignore ) -from pydantic import BaseModel, RootModel +from pydantic import BaseModel -from controllers.common.schema import register_response_schema_models, register_schema_models +from controllers.common.schema import register_schema_models from controllers.console import console_ns from controllers.console.datasets.wraps import get_rag_pipeline from controllers.console.wraps import account_initialization_required, setup_required, with_current_user +from extensions.ext_database import db from libs.login import login_required from models import Account from models.dataset import Pipeline @@ -21,18 +22,13 @@ class Parser(BaseModel): credential_id: str | None = None -class DataSourceContentPreviewResponse(RootModel[Any]): - root: Any - - register_schema_models(console_ns, Parser) -register_response_schema_models(console_ns, DataSourceContentPreviewResponse) @console_ns.route("/rag/pipelines//workflows/published/datasource/nodes//preview") class DataSourceContentPreviewApi(Resource): @console_ns.expect(console_ns.models[Parser.__name__]) - @console_ns.response(200, "Success", console_ns.models[DataSourceContentPreviewResponse.__name__]) + @console_ns.response(200, "Success") @setup_required @login_required @account_initialization_required @@ -46,7 +42,7 @@ class DataSourceContentPreviewApi(Resource): inputs = args.inputs datasource_type = args.datasource_type - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) preview_content = rag_pipeline_service.run_datasource_node_preview( pipeline=pipeline, node_id=node_id, diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py index 4027fa487a2..2d824afb6ef 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py @@ -108,7 +108,10 @@ class PipelineTemplateListApi(Resource): query = PipelineTemplateListQuery.model_validate(request.args.to_dict(flat=True)) # get pipeline templates pipeline_templates = RagPipelineService.get_pipeline_templates( - session, query.type, query.language, current_tenant_id + type=query.type, + language=query.language, + current_tenant_id=current_tenant_id, + session=session, ) return dump_response(PipelineTemplateListResponse, pipeline_templates), 200 @@ -124,8 +127,11 @@ class PipelineTemplateDetailApi(Resource): @with_session def get(self, session: Session, template_id: str) -> JsonResponseWithStatus: query = PipelineTemplateDetailQuery.model_validate(request.args.to_dict(flat=True)) - rag_pipeline_service = RagPipelineService() - pipeline_template = rag_pipeline_service.get_pipeline_template_detail(session, template_id, query.type) + pipeline_template = RagPipelineService.get_pipeline_template_detail( + template_id, + type=query.type, + session=session, + ) if pipeline_template is None: raise NotFound("Pipeline template not found from upstream service.") return dump_response(PipelineTemplateDetailResponse, pipeline_template), 200 @@ -145,7 +151,7 @@ class CustomizedPipelineTemplateApi(Resource): payload = CustomizedPipelineTemplatePayload.model_validate(console_ns.payload or {}) pipeline_template_info = PipelineTemplateInfoEntity.model_validate(payload.model_dump()) RagPipelineService.update_customized_pipeline_template( - template_id, pipeline_template_info, current_user, current_tenant_id + template_id, pipeline_template_info, current_user, current_tenant_id, session=db.session() ) return "", 204 @@ -156,7 +162,7 @@ class CustomizedPipelineTemplateApi(Resource): @enterprise_license_required @with_current_tenant_id def delete(self, current_tenant_id: str, template_id: str) -> tuple[str, int]: - RagPipelineService.delete_customized_pipeline_template(template_id, current_tenant_id) + RagPipelineService.delete_customized_pipeline_template(template_id, current_tenant_id, session=db.session()) return "", 204 @setup_required @@ -188,8 +194,8 @@ class PublishCustomizedPipelineTemplateApi(Resource): @with_current_tenant_id def post(self, current_tenant_id: str, current_user: Account, pipeline_id: str) -> tuple[str, int]: payload = CustomizedPipelineTemplatePayload.model_validate(console_ns.payload or {}) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) rag_pipeline_service.publish_customized_pipeline_template( - pipeline_id, payload.model_dump(), current_user, current_tenant_id + pipeline_id, payload.model_dump(), current_user, current_tenant_id, session=db.session() ) return "", 204 diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py index a373c8b1a41..5ad764871e4 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py @@ -65,7 +65,7 @@ class CreateRagPipelineDatasetApi(Resource): yaml_content=payload.yaml_content, ) try: - rag_pipeline_dsl_service = RagPipelineDslService(db.session) + rag_pipeline_dsl_service = RagPipelineDslService(db.session()) import_info = rag_pipeline_dsl_service.create_rag_pipeline_dataset( tenant_id=current_tenant_id, rag_pipeline_dataset_create_entity=rag_pipeline_dataset_create_entity, @@ -75,7 +75,7 @@ class CreateRagPipelineDatasetApi(Resource): current_tenant_id, import_info["dataset_id"], rag_pipeline_dataset_create_entity.partial_member_list, - db.session, + db.session(), ) db.session.commit() except services.errors.dataset.DatasetNameDuplicateError: @@ -110,6 +110,6 @@ class CreateEmptyRagPipelineDatasetApi(Resource): permission=DatasetPermissionEnum.ONLY_ME, partial_member_list=None, ), - session=db.session, + session=db.session(), ) return dump_response(DatasetDetailResponse, dataset), 201 diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py index af417f24dfe..25628a67177 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py @@ -98,7 +98,7 @@ class RagPipelineVariableCollectionApi(Resource): query = PaginationQuery.model_validate(request.args.to_dict()) # fetch draft workflow by app_model - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow_exist = rag_pipeline_service.is_workflow_exist(pipeline=pipeline) if not workflow_exist: raise DraftWorkflowNotExist() @@ -290,7 +290,7 @@ class RagPipelineVariableResetApi(Resource): session=db.session(), ) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) draft_workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if draft_workflow is None: raise NotFoundError( @@ -347,7 +347,7 @@ class RagPipelineEnvironmentVariableCollectionApi(Resource): Get draft workflow """ # fetch draft workflow by app_model - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if workflow is None: raise DraftWorkflowNotExist() diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py index c52385f6cf2..a61fc2639db 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py @@ -197,7 +197,7 @@ class DraftRagPipelineApi(Resource): Get draft rag pipeline's workflow """ # fetch draft workflow by app_model - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if not workflow: @@ -231,7 +231,7 @@ class DraftRagPipelineApi(Resource): return {"message": "Invalid JSON data"}, 400 else: abort(415) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) try: environment_variables_list = Workflow.normalize_environment_variable_mappings( @@ -283,7 +283,7 @@ class RagPipelineDraftRunIterationNodeApi(Resource): try: response = PipelineGenerateService.generate_single_iteration( - pipeline=pipeline, user=current_user, node_id=node_id, args=args, streaming=True + pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True ) return helper.compact_generate_response(response) @@ -318,7 +318,7 @@ class RagPipelineDraftRunLoopNodeApi(Resource): try: response = PipelineGenerateService.generate_single_loop( - pipeline=pipeline, user=current_user, node_id=node_id, args=args, streaming=True + pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True ) return helper.compact_generate_response(response) @@ -343,8 +343,8 @@ class DraftRagPipelineRunApi(Resource): @edit_permission_required @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_current_user - @get_rag_pipeline @with_session + @get_rag_pipeline def post(self, session: Session, current_user: Account, pipeline: Pipeline): """ Run draft workflow @@ -377,8 +377,8 @@ class PublishedRagPipelineRunApi(Resource): @edit_permission_required @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_current_user - @get_rag_pipeline @with_session + @get_rag_pipeline def post(self, session: Session, current_user: Account, pipeline: Pipeline): """ Run published workflow @@ -419,7 +419,7 @@ class RagPipelinePublishedDatasourceNodeRunApi(Resource): """ payload = DatasourceNodeRunPayload.model_validate(console_ns.payload or {}) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) return helper.compact_generate_response( PipelineGenerator.convert_to_event_stream( rag_pipeline_service.run_datasource_workflow_node( @@ -452,7 +452,7 @@ class RagPipelineDraftDatasourceNodeRunApi(Resource): """ payload = DatasourceNodeRunPayload.model_validate(console_ns.payload or {}) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) return helper.compact_generate_response( PipelineGenerator.convert_to_event_stream( rag_pipeline_service.run_datasource_workflow_node( @@ -490,7 +490,7 @@ class RagPipelineDraftNodeRunApi(Resource): payload = NodeRunRequiredPayload.model_validate(console_ns.payload or {}) inputs = payload.inputs - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow_node_execution = rag_pipeline_service.run_draft_workflow_node( pipeline=pipeline, node_id=node_id, user_inputs=inputs, account=current_user ) @@ -543,7 +543,7 @@ class PublishedRagPipelineApi(Resource): if not pipeline.is_published: return None # fetch published workflow by pipeline - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow = rag_pipeline_service.get_published_workflow(pipeline=pipeline) # return workflow, if not found, return None @@ -564,9 +564,9 @@ class PublishedRagPipelineApi(Resource): """ Publish workflow """ - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow = rag_pipeline_service.publish_workflow( - session=db.session, # type: ignore[reportArgumentType,arg-type] + session=db.session(), pipeline=pipeline, account=current_user, ) @@ -599,7 +599,7 @@ class DefaultRagPipelineBlockConfigsApi(Resource): Get default block config """ # Get default block configs - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) return rag_pipeline_service.get_default_block_configs() @@ -631,7 +631,7 @@ class DefaultRagPipelineBlockConfigApi(Resource): raise ValueError("Invalid filters") # Get default block configs - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) return rag_pipeline_service.get_default_block_config(node_type=block_type, filters=filters) @@ -666,7 +666,7 @@ class PublishedAllRagPipelineApi(Resource): if user_id != current_user.id: raise Forbidden() - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) with sessionmaker(db.engine).begin() as session: workflows, has_more = rag_pipeline_service.get_all_published_workflow( session=session, @@ -698,7 +698,7 @@ class RagPipelineDraftWorkflowRestoreApi(Resource): @with_current_user @get_rag_pipeline def post(self, current_user: Account, pipeline: Pipeline, workflow_id: str): - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) try: workflow = rag_pipeline_service.restore_published_workflow_to_draft( @@ -743,7 +743,7 @@ class RagPipelineByIdApi(Resource): if not update_data: return {"message": "No valid fields to update"}, 400 - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow_ref = WorkflowRefService.create_pipeline_workflow_ref(pipeline, workflow_id) # Create a session and manage the transaction @@ -809,7 +809,7 @@ class PublishedRagPipelineSecondStepApi(Resource): """ query = NodeIdQuery.model_validate(request.args.to_dict()) node_id = query.node_id - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) variables = rag_pipeline_service.get_second_step_parameters(pipeline=pipeline, node_id=node_id, is_draft=False) return { "variables": variables, @@ -832,7 +832,7 @@ class PublishedRagPipelineFirstStepApi(Resource): """ query = NodeIdQuery.model_validate(request.args.to_dict()) node_id = query.node_id - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) variables = rag_pipeline_service.get_first_step_parameters(pipeline=pipeline, node_id=node_id, is_draft=False) return { "variables": variables, @@ -855,7 +855,7 @@ class DraftRagPipelineFirstStepApi(Resource): """ query = NodeIdQuery.model_validate(request.args.to_dict()) node_id = query.node_id - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) variables = rag_pipeline_service.get_first_step_parameters(pipeline=pipeline, node_id=node_id, is_draft=True) return { "variables": variables, @@ -879,7 +879,7 @@ class DraftRagPipelineSecondStepApi(Resource): query = NodeIdQuery.model_validate(request.args.to_dict()) node_id = query.node_id - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) variables = rag_pipeline_service.get_second_step_parameters(pipeline=pipeline, node_id=node_id, is_draft=True) return { "variables": variables, @@ -913,7 +913,7 @@ class RagPipelineWorkflowRunListApi(Resource): "limit": query.limit, } - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) result = rag_pipeline_service.get_rag_pipeline_paginate_workflow_runs(pipeline=pipeline, args=args) return WorkflowRunPaginationResponse.model_validate(result, from_attributes=True).model_dump(mode="json") @@ -936,7 +936,7 @@ class RagPipelineWorkflowRunDetailApi(Resource): """ run_id_str = str(run_id) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow_run = rag_pipeline_service.get_rag_pipeline_workflow_run(pipeline=pipeline, run_id=run_id_str) if workflow_run is None: raise NotFound("Workflow run not found") @@ -962,7 +962,7 @@ class RagPipelineWorkflowRunNodeExecutionListApi(Resource): """ run_id_str = str(run_id) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) user = cast("Account | EndUser", current_user) node_executions = rag_pipeline_service.get_rag_pipeline_workflow_run_node_executions( pipeline=pipeline, @@ -998,7 +998,7 @@ class RagPipelineWorkflowLastRunApi(Resource): @account_initialization_required @get_rag_pipeline def get(self, pipeline: Pipeline, node_id: str): - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if not workflow: raise NotFound("Workflow not found") @@ -1051,7 +1051,7 @@ class RagPipelineDatasourceVariableApi(Resource): """ args = DatasourceVariablesPayload.model_validate(console_ns.payload or {}).model_dump() - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) workflow_node_execution = rag_pipeline_service.set_datasource_variables( pipeline=pipeline, args=args, @@ -1074,6 +1074,6 @@ class RagPipelineRecommendedPluginApi(Resource): def get(self, current_tenant_id: str, current_user: Account): query = RagPipelineRecommendedPluginQuery.model_validate(request.args.to_dict()) - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) recommended_plugins = rag_pipeline_service.get_recommended_plugins(query.type, current_user, current_tenant_id) return recommended_plugins diff --git a/api/controllers/console/datasets/wraps.py b/api/controllers/console/datasets/wraps.py index b58a07029c8..b5a9cd753ff 100644 --- a/api/controllers/console/datasets/wraps.py +++ b/api/controllers/console/datasets/wraps.py @@ -2,6 +2,7 @@ from collections.abc import Callable from functools import wraps from sqlalchemy import select +from sqlalchemy.orm import Session from controllers.console.datasets.error import PipelineNotFoundError from extensions.ext_database import db @@ -22,9 +23,10 @@ def get_rag_pipeline[**P, R](view_func: Callable[P, R]) -> Callable[P, R]: del kwargs["pipeline_id"] - pipeline = db.session.scalar( - select(Pipeline).where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id).limit(1) - ) + stmt = select(Pipeline).where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id).limit(1) + # Migrated handlers pass the request Session as args[1]; legacy handlers still use db.session. + session = args[1] if len(args) > 1 and isinstance(args[1], Session) else db.session + pipeline = session.scalar(stmt) if not pipeline: raise PipelineNotFoundError() diff --git a/api/controllers/console/explore/audio.py b/api/controllers/console/explore/audio.py index c0b86c19e43..e5f98f0f655 100644 --- a/api/controllers/console/explore/audio.py +++ b/api/controllers/console/explore/audio.py @@ -113,7 +113,7 @@ class ChatTextApi(InstalledAppResource): response = AudioService.transcript_tts( app_model=app_model, - session=db.session, + session=db.session(), text=text, voice=voice, message_ref=message_ref, diff --git a/api/controllers/console/explore/conversation.py b/api/controllers/console/explore/conversation.py index 2004e648f19..25239203d8d 100644 --- a/api/controllers/console/explore/conversation.py +++ b/api/controllers/console/explore/conversation.py @@ -111,7 +111,7 @@ class ConversationApi(InstalledAppResource): conversation_id = str(c_id) try: - ConversationService.delete(app_model, conversation_id, current_user) + ConversationService.delete(app_model, conversation_id, current_user, session=db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -140,7 +140,7 @@ class ConversationRenameApi(InstalledAppResource): try: conversation = ConversationService.rename( - app_model, conversation_id, current_user, payload.name, payload.auto_generate + app_model, conversation_id, current_user, payload.name, payload.auto_generate, session=db.session() ) return ( TypeAdapter(SimpleConversation) @@ -169,7 +169,7 @@ class ConversationPinApi(InstalledAppResource): conversation_id = str(c_id) try: - WebConversationService.pin(app_model, conversation_id, current_user) + WebConversationService.pin(app_model, conversation_id, current_user, db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -192,6 +192,6 @@ class ConversationUnPinApi(InstalledAppResource): raise NotChatAppError() conversation_id = str(c_id) - WebConversationService.unpin(app_model, conversation_id, current_user) + WebConversationService.unpin(app_model, conversation_id, current_user, db.session()) return ResultResponse(result="success").model_dump(mode="json") diff --git a/api/controllers/console/explore/installed_app.py b/api/controllers/console/explore/installed_app.py index 71cb03ce6a0..1fe1201bab7 100644 --- a/api/controllers/console/explore/installed_app.py +++ b/api/controllers/console/explore/installed_app.py @@ -181,7 +181,7 @@ class InstalledAppsListApi(Resource): if current_user.current_tenant is None: raise ValueError("current_user.current_tenant must not be None") - current_user.role = TenantService.get_user_role(current_user, current_user.current_tenant, session=db.session) + current_user.role = TenantService.get_user_role(current_user, current_user.current_tenant, session=db.session()) installed_app_list: list[dict[str, Any]] = [] for installed_app, app_model in installed_apps: installed_app_list.append( diff --git a/api/controllers/console/explore/message.py b/api/controllers/console/explore/message.py index 0e27e2db25b..7b316b0382d 100644 --- a/api/controllers/console/explore/message.py +++ b/api/controllers/console/explore/message.py @@ -27,6 +27,7 @@ from controllers.console.explore.wraps import InstalledAppResource from controllers.console.wraps import with_current_user from core.app.entities.app_invoke_entities import InvokeFrom from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from extensions.ext_database import db from fields.conversation_fields import ResultResponse from fields.message_fields import ( ExploreMessageInfiniteScrollPagination, @@ -91,6 +92,7 @@ class MessageListApi(InstalledAppResource): args.conversation_id, args.first_id or None, args.limit, + session=db.session(), ) adapter = TypeAdapter(ExploreMessageListItem) items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data] @@ -129,6 +131,7 @@ class MessageFeedbackApi(InstalledAppResource): user=current_user, rating=FeedbackRating(payload.rating) if payload.rating else None, content=payload.content, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -207,7 +210,11 @@ class MessageSuggestedQuestionApi(InstalledAppResource): try: questions = MessageService.get_suggested_questions_after_answer( - app_model=app_model, user=current_user, message_id=message_id_str, invoke_from=InvokeFrom.EXPLORE + app_model=app_model, + user=current_user, + message_id=message_id_str, + invoke_from=InvokeFrom.EXPLORE, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message not found") diff --git a/api/controllers/console/explore/parameter.py b/api/controllers/console/explore/parameter.py index 0bc6e032bf0..680885f9bd3 100644 --- a/api/controllers/console/explore/parameter.py +++ b/api/controllers/console/explore/parameter.py @@ -8,6 +8,7 @@ from controllers.console import console_ns from controllers.console.app.error import AppUnavailableError from controllers.console.explore.wraps import InstalledAppResource from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict +from extensions.ext_database import db from models.model import AppMode, InstalledApp from services.app_service import AppService @@ -64,4 +65,4 @@ class ExploreAppMetaApi(InstalledAppResource): app_model = installed_app.app if not app_model: raise ValueError("App not found") - return AppService().get_app_meta(app_model) + return AppService().get_app_meta(app_model, session=db.session()) diff --git a/api/controllers/console/explore/recommended_app.py b/api/controllers/console/explore/recommended_app.py index abe170bf90a..79eaa305d61 100644 --- a/api/controllers/console/explore/recommended_app.py +++ b/api/controllers/console/explore/recommended_app.py @@ -120,7 +120,7 @@ class RecommendedAppListApi(Resource): language_prefix = _resolve_language(args.language, current_user) return RecommendedAppListResponse.model_validate( - RecommendedAppService.get_recommended_apps_and_categories(db.session, language_prefix), + RecommendedAppService.get_recommended_apps_and_categories(language_prefix, session=db.session()), from_attributes=True, ).model_dump(mode="json") @@ -137,7 +137,7 @@ class LearnDifyAppListApi(Resource): language_prefix = _resolve_language(args.language, current_user) return LearnDifyAppListResponse.model_validate( - RecommendedAppService.get_learn_dify_apps(db.session, language_prefix), + RecommendedAppService.get_learn_dify_apps(language_prefix, session=db.session()), from_attributes=True, ).model_dump(mode="json") @@ -148,4 +148,4 @@ class RecommendedAppApi(Resource): @login_required @account_initialization_required def get(self, app_id: UUID): - return RecommendedAppService.get_recommend_app_detail(db.session, str(app_id)) + return RecommendedAppService.get_recommend_app_detail(str(app_id), session=db.session()) diff --git a/api/controllers/console/explore/saved_message.py b/api/controllers/console/explore/saved_message.py index ce43ff18c93..e3fd730a3cc 100644 --- a/api/controllers/console/explore/saved_message.py +++ b/api/controllers/console/explore/saved_message.py @@ -38,11 +38,7 @@ class SavedMessageListApi(InstalledAppResource): args = SavedMessageListQuery.model_validate(request.args.to_dict()) pagination = SavedMessageService.pagination_by_last_id( - db.session(), - app_model, - current_user, - str(args.last_id) if args.last_id else None, - args.limit, + app_model, current_user, str(args.last_id) if args.last_id else None, args.limit, session=db.session() ) adapter = TypeAdapter(SavedMessageItem) items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data] @@ -65,7 +61,7 @@ class SavedMessageListApi(InstalledAppResource): payload = SavedMessageCreatePayload.model_validate(console_ns.payload or {}) try: - SavedMessageService.save(db.session(), app_model, current_user, str(payload.message_id)) + SavedMessageService.save(app_model, current_user, str(payload.message_id), session=db.session()) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -88,6 +84,6 @@ class SavedMessageApi(InstalledAppResource): if app_model.mode != "completion": raise NotCompletionAppError() - SavedMessageService.delete(db.session(), app_model, current_user, message_id_str) + SavedMessageService.delete(app_model, current_user, message_id_str, session=db.session()) return "", 204 diff --git a/api/controllers/console/explore/trial.py b/api/controllers/console/explore/trial.py index 0d54bf7f473..d01eb9c38b1 100644 --- a/api/controllers/console/explore/trial.py +++ b/api/controllers/console/explore/trial.py @@ -13,7 +13,6 @@ import services from controllers.common.fields import ( AudioBinaryResponse, AudioTranscriptResponse, - GeneratedAppResponse, SimpleResultResponse, ) from controllers.common.fields import Parameters as ParametersResponse @@ -391,7 +390,6 @@ register_response_schema_models( ParametersResponse, AudioBinaryResponse, AudioTranscriptResponse, - GeneratedAppResponse, SimpleResultResponse, SiteResponse, SuggestedQuestionsResponse, @@ -406,7 +404,7 @@ simple_account_model = console_ns.models[TrialSimpleAccount.__name__] class TrialAppWorkflowRunApi(TrialAppResource): @trial_feature_enable @console_ns.expect(console_ns.models[WorkflowRunRequest.__name__]) - @console_ns.response(200, "Success", console_ns.models[GeneratedAppResponse.__name__]) + @console_ns.response(200, "Success") @with_current_user @with_session def post(self, session: Session, current_user: Account, trial_app): @@ -433,7 +431,8 @@ class TrialAppWorkflowRunApi(TrialAppResource): invoke_from=InvokeFrom.EXPLORE, streaming=True, ) - RecommendedAppService.add_trial_app_record(db.session, app_id, user_id) + RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -478,7 +477,7 @@ class TrialAppWorkflowTaskStopApi(TrialAppResource): class TrialChatApi(TrialAppResource): @console_ns.expect(console_ns.models[ChatRequest.__name__]) - @console_ns.response(200, "Success", console_ns.models[GeneratedAppResponse.__name__]) + @console_ns.response(200, "Success") @trial_feature_enable @with_current_user @with_session @@ -512,7 +511,8 @@ class TrialChatApi(TrialAppResource): invoke_from=InvokeFrom.EXPLORE, streaming=True, ) - RecommendedAppService.add_trial_app_record(db.session, app_id, user_id) + RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except services.errors.conversation.ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -551,7 +551,11 @@ class TrialMessageSuggestedQuestionApi(TrialAppResource): try: questions = MessageService.get_suggested_questions_after_answer( - app_model=app_model, user=current_user, message_id=message_id, invoke_from=InvokeFrom.EXPLORE + app_model=app_model, + user=current_user, + message_id=message_id, + invoke_from=InvokeFrom.EXPLORE, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message not found") @@ -589,7 +593,7 @@ class TrialChatAudioApi(TrialAppResource): user_id = current_user.id response = AudioService.transcript_asr(app_model=app_model, file=file, end_user=None) - RecommendedAppService.add_trial_app_record(db.session, app_id, user_id) + RecommendedAppService.add_trial_app_record(app_id, user_id, session=db.session()) return response except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") @@ -645,12 +649,12 @@ class TrialChatTextApi(TrialAppResource): response = AudioService.transcript_tts( app_model=app_model, - session=db.session, + session=db.session(), text=text, voice=voice, message_ref=message_ref, ) - RecommendedAppService.add_trial_app_record(db.session, app_id, user_id) + RecommendedAppService.add_trial_app_record(app_id, user_id, session=db.session()) return response except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") @@ -680,7 +684,7 @@ class TrialChatTextApi(TrialAppResource): class TrialCompletionApi(TrialAppResource): @console_ns.expect(console_ns.models[CompletionRequest.__name__]) - @console_ns.response(200, "Success", console_ns.models[GeneratedAppResponse.__name__]) + @console_ns.response(200, "Success") @trial_feature_enable @with_current_user @with_session @@ -709,7 +713,8 @@ class TrialCompletionApi(TrialAppResource): streaming=streaming, ) - RecommendedAppService.add_trial_app_record(db.session, app_id, user_id) + RecommendedAppService.add_trial_app_record(app_id, user_id, session=session) + # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except services.errors.conversation.ConversationNotExistsError: raise NotFound("Conversation Not Exists.") diff --git a/api/controllers/console/extension.py b/api/controllers/console/extension.py index 4b149b9c08d..cc06204a905 100644 --- a/api/controllers/console/extension.py +++ b/api/controllers/console/extension.py @@ -112,7 +112,7 @@ class APIBasedExtensionAPI(Resource): def get(self, current_tenant_id: str): return dump_response( APIBasedExtensionListResponse, - APIBasedExtensionService.get_all_by_tenant_id(db.session(), current_tenant_id), + APIBasedExtensionService.get_all_by_tenant_id(current_tenant_id, session=db.session()), ) @console_ns.doc("create_api_based_extension") @@ -133,7 +133,7 @@ class APIBasedExtensionAPI(Resource): api_key=payload.api_key, ) - extension = APIBasedExtensionService.save(db.session(), extension_data) + extension = APIBasedExtensionService.save(extension_data, session=db.session()) return APIBasedExtensionResponse( id=extension.id, name=extension.name, @@ -158,7 +158,9 @@ class APIBasedExtensionDetailAPI(Resource): return dump_response( APIBasedExtensionResponse, - APIBasedExtensionService.get_with_tenant_id(db.session(), current_tenant_id, api_based_extension_id), + APIBasedExtensionService.get_with_tenant_id( + current_tenant_id, api_based_extension_id, session=db.session() + ), ) @console_ns.doc("update_api_based_extension") @@ -174,7 +176,7 @@ class APIBasedExtensionDetailAPI(Resource): api_based_extension_id = str(id) extension_data_from_db = APIBasedExtensionService.get_with_tenant_id( - db.session(), current_tenant_id, api_based_extension_id + current_tenant_id, api_based_extension_id, session=db.session() ) payload = APIBasedExtensionPayload.model_validate(console_ns.payload or {}) @@ -187,7 +189,7 @@ class APIBasedExtensionDetailAPI(Resource): extension_data_from_db.api_key = payload.api_key api_key_for_response = payload.api_key - APIBasedExtensionService.save(db.session(), extension_data_from_db) + APIBasedExtensionService.save(extension_data_from_db, session=db.session()) return APIBasedExtensionResponse( id=extension_data_from_db.id, name=extension_data_from_db.name, @@ -208,9 +210,9 @@ class APIBasedExtensionDetailAPI(Resource): api_based_extension_id = str(id) extension_data_from_db = APIBasedExtensionService.get_with_tenant_id( - db.session(), current_tenant_id, api_based_extension_id + current_tenant_id, api_based_extension_id, session=db.session() ) - APIBasedExtensionService.delete(db.session(), extension_data_from_db) + APIBasedExtensionService.delete(extension_data_from_db, session=db.session()) return "", 204 diff --git a/api/controllers/console/files.py b/api/controllers/console/files.py index 24c3a5978e1..5e680c83702 100644 --- a/api/controllers/console/files.py +++ b/api/controllers/console/files.py @@ -27,6 +27,7 @@ from controllers.console.wraps import ( ) from extensions.ext_database import db from fields.file_fields import FileResponse, UploadConfig +from libs.helper import dump_response from libs.login import login_required from models import Account from services.file_service import FileService @@ -117,8 +118,7 @@ class FileApi(Resource): except services.errors.file.BlockedFileExtensionError as blocked_extension_error: raise BlockedFileExtensionError(blocked_extension_error.description) - response = FileResponse.model_validate(upload_file, from_attributes=True) - return response.model_dump(mode="json"), 201 + return dump_response(FileResponse, upload_file), 201 @console_ns.route("/files//preview") @@ -131,7 +131,7 @@ class FilePreviewApi(Resource): def get(self, current_tenant_id: str, file_id: UUID): file_id_str = str(file_id) text = FileService(db.engine).get_file_preview(file_id_str, current_tenant_id) - return {"content": text} + return TextContentResponse(content=text).model_dump(mode="json") @console_ns.route("/files/support-type") @@ -141,4 +141,4 @@ class FileSupportTypeApi(Resource): @account_initialization_required @console_ns.response(200, "Success", console_ns.models[AllowedExtensionsResponse.__name__]) def get(self): - return {"allowed_extensions": list(DOCUMENT_EXTENSIONS)} + return AllowedExtensionsResponse(allowed_extensions=list(DOCUMENT_EXTENSIONS)).model_dump(mode="json") diff --git a/api/controllers/console/init_validate.py b/api/controllers/console/init_validate.py index 27f6bcc36dc..f155f222e19 100644 --- a/api/controllers/console/init_validate.py +++ b/api/controllers/console/init_validate.py @@ -50,7 +50,7 @@ def get_init_status() -> InitStatusResponse: @only_edition_self_hosted def validate_init_password(payload: InitValidatePayload) -> InitValidateResponse: """Validate initialization password.""" - tenant_count = TenantService.get_tenant_count(session=db.session) + tenant_count = TenantService.get_tenant_count(session=db.session()) if tenant_count > 0: raise AlreadySetupError() diff --git a/api/controllers/console/setup.py b/api/controllers/console/setup.py index 3b5c1bbe18f..e0a0fba3329 100644 --- a/api/controllers/console/setup.py +++ b/api/controllers/console/setup.py @@ -13,7 +13,7 @@ from services.account_service import RegisterService, TenantService from .error import AlreadySetupError, NotInitValidateError from .init_validate import get_init_validate_status -from .wraps import only_edition_self_hosted +from .wraps import mark_setup_completed, only_edition_self_hosted class SetupRequestPayload(BaseModel): @@ -79,7 +79,7 @@ def setup_system(payload: SetupRequestPayload) -> SetupResponse: if get_setup_status(): raise AlreadySetupError() - tenant_count = TenantService.get_tenant_count(session=db.session) + tenant_count = TenantService.get_tenant_count(session=db.session()) if tenant_count > 0: raise AlreadySetupError() @@ -94,8 +94,9 @@ def setup_system(payload: SetupRequestPayload) -> SetupResponse: password=payload.password, ip_address=extract_remote_ip(request), language=payload.language, - session=db.session, + session=db.session(), ) + mark_setup_completed() return SetupResponse(result="success") diff --git a/api/controllers/console/socketio/workflow.py b/api/controllers/console/socketio/workflow.py index 99e56df3cb8..db5a4144dd3 100644 --- a/api/controllers/console/socketio/workflow.py +++ b/api/controllers/console/socketio/workflow.py @@ -44,7 +44,7 @@ def socket_connect(sid, environ, auth): return False with sio.app.app_context(): - user = AccountService.load_logged_in_account(account_id=user_id, session=db.session) + user = AccountService.load_logged_in_account(account_id=user_id, session=db.session()) if not user: logging.warning("Socket connect rejected: user not found (user_id=%s, sid=%s)", user_id, sid) return False @@ -69,7 +69,7 @@ def handle_user_connect(sid, data): if not workflow_id: return {"msg": "workflow_id is required"}, 400 - result = collaboration_service.authorize_and_join_workflow_room(workflow_id, sid) + result = collaboration_service.authorize_and_join_workflow_room(workflow_id, sid, session=db.session()) if not result: return {"msg": "unauthorized"}, 401 diff --git a/api/controllers/console/spec.py b/api/controllers/console/spec.py index 27b07b4dd81..70e0d1d14ae 100644 --- a/api/controllers/console/spec.py +++ b/api/controllers/console/spec.py @@ -1,8 +1,9 @@ import logging +from collections.abc import Mapping from typing import Any from flask_restx import Resource -from pydantic import RootModel +from pydantic import Field, RootModel from controllers.common.schema import register_response_schema_models from controllers.console.wraps import ( @@ -10,6 +11,7 @@ from controllers.console.wraps import ( setup_required, ) from core.schemas.schema_manager import SchemaManager +from fields.base import ResponseModel from libs.login import login_required from . import console_ns @@ -17,11 +19,17 @@ from . import console_ns logger = logging.getLogger(__name__) -class SchemaDefinitionsResponse(RootModel[Any]): - root: Any +class SchemaDefinitionItemResponse(ResponseModel): + name: str + label: str + schema_: Mapping[str, Any] = Field(alias="schema") -register_response_schema_models(console_ns, SchemaDefinitionsResponse) +class SchemaDefinitionsResponse(RootModel[list[SchemaDefinitionItemResponse]]): + pass + + +register_response_schema_models(console_ns, SchemaDefinitionItemResponse, SchemaDefinitionsResponse) @console_ns.route("/spec/schema-definitions") diff --git a/api/controllers/console/tag/tags.py b/api/controllers/console/tag/tags.py index c4ec925c9a3..86c1ad9c54c 100644 --- a/api/controllers/console/tag/tags.py +++ b/api/controllers/console/tag/tags.py @@ -137,7 +137,7 @@ class TagListApi(Resource): def get(self, current_tenant_id: str): raw_args = request.args.to_dict() param = TagListQueryParam.model_validate(raw_args) - tags = TagService.get_tags(db.session(), param.type, current_tenant_id, param.keyword) + tags = TagService.get_tags(param.type, current_tenant_id, param.keyword, session=db.session()) return dump_response(TagListResponse, tags), 200 @@ -154,7 +154,7 @@ class TagListApi(Resource): payload = TagBasePayload.model_validate(console_ns.payload or {}) _enforce_snippet_tag_rbac_if_needed(payload.type) - tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=payload.type), db.session) + tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=payload.type), db.session()) return dump_response(TagResponse, {"id": tag.id, "name": tag.name, "type": tag.type, "binding_count": 0}), 200 @@ -175,9 +175,9 @@ class TagUpdateDeleteApi(Resource): payload = TagUpdateRequestPayload.model_validate(console_ns.payload or {}) _enforce_snippet_tag_rbac_by_tag_id(tag_id_str) - tag = TagService.update_tags(UpdateTagPayload(name=payload.name), tag_id_str, db.session) + tag = TagService.update_tags(UpdateTagPayload(name=payload.name), tag_id_str, db.session()) - binding_count = TagService.get_tag_binding_count(tag_id_str, db.session) + binding_count = TagService.get_tag_binding_count(tag_id_str, db.session()) return ( dump_response( @@ -196,7 +196,7 @@ class TagUpdateDeleteApi(Resource): tag_id_str = str(tag_id) _enforce_snippet_tag_rbac_by_tag_id(tag_id_str) - TagService.delete_tag(tag_id_str, db.session) + TagService.delete_tag(tag_id_str, db.session()) return "", 204 @@ -223,7 +223,7 @@ def _create_tag_bindings(current_user: Account) -> tuple[dict[str, str], int]: target_id=payload.target_id, type=payload.type, ), - db.session, + db.session(), ) return {"result": "success"}, 200 @@ -239,7 +239,7 @@ def _remove_tag_bindings(current_user: Account) -> tuple[dict[str, str], int]: target_id=payload.target_id, type=payload.type, ), - db.session, + db.session(), ) return {"result": "success"}, 200 diff --git a/api/controllers/console/workspace/__init__.py b/api/controllers/console/workspace/__init__.py index 13bc98c8047..09bb535e626 100644 --- a/api/controllers/console/workspace/__init__.py +++ b/api/controllers/console/workspace/__init__.py @@ -8,7 +8,7 @@ from werkzeug.exceptions import Forbidden from configs import dify_config from extensions.ext_database import db from libs.login import current_account_with_tenant -from models.account import TenantPluginPermission +from models.account import TenantPluginDebugPermission, TenantPluginInstallPermission, TenantPluginPermission def plugin_permission_required( @@ -40,22 +40,22 @@ def plugin_permission_required( if install_required: match permission.install_permission: - case TenantPluginPermission.InstallPermission.NOBODY: + case TenantPluginInstallPermission.NOBODY: raise Forbidden() - case TenantPluginPermission.InstallPermission.ADMINS: + case TenantPluginInstallPermission.ADMINS: if not user.is_admin_or_owner: raise Forbidden() - case TenantPluginPermission.InstallPermission.EVERYONE: + case TenantPluginInstallPermission.EVERYONE: pass if debug_required: match permission.debug_permission: - case TenantPluginPermission.DebugPermission.NOBODY: + case TenantPluginDebugPermission.NOBODY: raise Forbidden() - case TenantPluginPermission.DebugPermission.ADMINS: + case TenantPluginDebugPermission.ADMINS: if not user.is_admin_or_owner: raise Forbidden() - case TenantPluginPermission.DebugPermission.EVERYONE: + case TenantPluginDebugPermission.EVERYONE: pass return view(*args, **kwargs) diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index c13c8aa162f..6a06ed2d3e6 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -1,12 +1,13 @@ from __future__ import annotations from datetime import datetime -from typing import Any, Literal +from http import HTTPStatus +from typing import Literal import pytz from flask import request from flask_restx import Resource -from pydantic import BaseModel, Field, RootModel, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator, model_validator from sqlalchemy import select from werkzeug.exceptions import NotFound @@ -47,7 +48,7 @@ from controllers.console.wraps import ( ) from extensions.ext_database import db from fields.base import ResponseModel -from fields.member_fields import Account as AccountResponse +from fields.member_fields import AccountResponse from graphon.file import helpers as file_helpers from libs.datetime_utils import naive_utc_now from libs.helper import EmailStr, dump_response, extract_remote_ip, timezone, to_timestamp @@ -194,10 +195,6 @@ register_schema_models( ) -def _serialize_account(account) -> dict[str, Any]: - return AccountResponse.model_validate(account, from_attributes=True).model_dump(mode="json") - - class AccountIntegrateResponse(ResponseModel): provider: str created_at: int | None = None @@ -236,23 +233,15 @@ class EducationAutocompleteResponse(ResponseModel): has_next: bool | None = None -class EducationActivateResponse(RootModel[dict[str, Any]]): - root: dict[str, Any] - - -register_schema_models( - console_ns, - AccountIntegrateResponse, - AccountIntegrateListResponse, - EducationVerifyResponse, - EducationStatusResponse, - EducationAutocompleteResponse, -) register_response_schema_models( console_ns, AccountResponse, + AccountIntegrateResponse, + AccountIntegrateListResponse, AvatarUrlResponse, - EducationActivateResponse, + EducationVerifyResponse, + EducationStatusResponse, + EducationAutocompleteResponse, SimpleResultDataResponse, SimpleResultResponse, VerificationTokenResponse, @@ -262,7 +251,7 @@ register_response_schema_models( @console_ns.route("/account/init") class AccountInitApi(Resource): @console_ns.expect(console_ns.models[AccountInitPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @with_current_user @@ -302,7 +291,7 @@ class AccountInitApi(Resource): account.initialized_at = naive_utc_now() db.session.commit() - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.route("/account/profile") @@ -310,11 +299,11 @@ class AccountProfileApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @enterprise_license_required @with_current_user def get(self, current_user: Account): - return _serialize_account(current_user) + return dump_response(AccountResponse, current_user) @console_ns.route("/account/name") @@ -323,14 +312,14 @@ class AccountNameApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} args = AccountNamePayload.model_validate(payload) - updated_account = AccountService.update_account(current_user, session=db.session, name=args.name) + updated_account = AccountService.update_account(current_user, session=db.session(), name=args.name) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/avatar") @@ -338,7 +327,7 @@ class AccountAvatarApi(Resource): @console_ns.doc("get_account_avatar") @console_ns.doc(description="Get account avatar url") @console_ns.doc(params=query_params_from_model(AccountAvatarQuery)) - @console_ns.response(200, "Success", console_ns.models[AvatarUrlResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AvatarUrlResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -349,7 +338,7 @@ class AccountAvatarApi(Resource): avatar = args.avatar if avatar.startswith(("http://", "https://")): - return dump_response(AvatarUrlResponse, {"avatar_url": avatar}) + return AvatarUrlResponse(avatar_url=avatar).model_dump(mode="json") upload_file = db.session.scalar(select(UploadFile).where(UploadFile.id == avatar).limit(1)) if upload_file is None: @@ -362,21 +351,21 @@ class AccountAvatarApi(Resource): raise NotFound("Avatar file not found") avatar_url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id) - return dump_response(AvatarUrlResponse, {"avatar_url": avatar_url}) + return AvatarUrlResponse(avatar_url=avatar_url).model_dump(mode="json") @console_ns.expect(console_ns.models[AccountAvatarPayload.__name__]) @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} args = AccountAvatarPayload.model_validate(payload) - updated_account = AccountService.update_account(current_user, session=db.session, avatar=args.avatar) + updated_account = AccountService.update_account(current_user, session=db.session(), avatar=args.avatar) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/interface-language") @@ -385,17 +374,17 @@ class AccountInterfaceLanguageApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} args = AccountInterfaceLanguagePayload.model_validate(payload) updated_account = AccountService.update_account( - current_user, session=db.session, interface_language=args.interface_language + current_user, session=db.session(), interface_language=args.interface_language ) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/interface-theme") @@ -404,17 +393,17 @@ class AccountInterfaceThemeApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} args = AccountInterfaceThemePayload.model_validate(payload) updated_account = AccountService.update_account( - current_user, session=db.session, interface_theme=args.interface_theme + current_user, session=db.session(), interface_theme=args.interface_theme ) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/timezone") @@ -423,15 +412,15 @@ class AccountTimezoneApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} args = AccountTimezonePayload.model_validate(payload) - updated_account = AccountService.update_account(current_user, session=db.session, timezone=args.timezone) + updated_account = AccountService.update_account(current_user, session=db.session(), timezone=args.timezone) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/password") @@ -440,7 +429,7 @@ class AccountPasswordApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} @@ -448,11 +437,11 @@ class AccountPasswordApi(Resource): try: assert args.password is not None - AccountService.update_account_password(current_user, args.password, args.new_password, session=db.session) + AccountService.update_account_password(current_user, args.password, args.new_password, session=db.session()) except ServiceCurrentPasswordIncorrectError: raise CurrentPasswordIncorrectError() - return _serialize_account(current_user) + return dump_response(AccountResponse, current_user) @console_ns.route("/account/integrates") @@ -460,7 +449,7 @@ class AccountIntegrateApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountIntegrateListResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountIntegrateListResponse.__name__]) @with_current_user def get(self, account: Account): account_integrates = db.session.scalars( @@ -471,33 +460,29 @@ class AccountIntegrateApi(Resource): oauth_base_path = "/console/api/oauth/login" providers = ["github", "google"] - integrate_data = [] + integrate_data: list[AccountIntegrateResponse] = [] for provider in providers: existing_integrate = next((ai for ai in account_integrates if ai.provider == provider), None) if existing_integrate: integrate_data.append( - { - "id": existing_integrate.id, - "provider": provider, - "created_at": existing_integrate.created_at, - "is_bound": True, - "link": None, - } + AccountIntegrateResponse( + provider=provider, + created_at=to_timestamp(existing_integrate.created_at), + is_bound=True, + link=None, + ) ) else: integrate_data.append( - { - "id": None, - "provider": provider, - "created_at": None, - "is_bound": False, - "link": f"{base_url}{oauth_base_path}/{provider}", - } + AccountIntegrateResponse( + provider=provider, + created_at=None, + is_bound=False, + link=f"{base_url}{oauth_base_path}/{provider}", + ) ) - return AccountIntegrateListResponse( - data=[AccountIntegrateResponse.model_validate(item) for item in integrate_data] - ).model_dump(mode="json") + return AccountIntegrateListResponse(data=integrate_data).model_dump(mode="json") @console_ns.route("/account/delete/verify") @@ -505,19 +490,19 @@ class AccountDeleteVerifyApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultDataResponse.__name__]) @with_current_user def get(self, account: Account): token, code = AccountService.generate_account_deletion_verification_code(account) AccountService.send_account_deletion_verification_email(account, code) - return {"result": "success", "data": token} + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @console_ns.route("/account/delete") class AccountDeleteApi(Resource): @console_ns.expect(console_ns.models[AccountDeletePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -529,15 +514,15 @@ class AccountDeleteApi(Resource): if not AccountService.verify_account_deletion_code(args.token, args.code): raise InvalidAccountDeletionCodeError() - AccountService.delete_account(account) + AccountService.delete_account(account, session=db.session()) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.route("/account/delete/feedback") class AccountDeleteUpdateFeedbackApi(Resource): @console_ns.expect(console_ns.models[AccountDeletionFeedbackPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required def post(self): payload = console_ns.payload or {} @@ -545,7 +530,7 @@ class AccountDeleteUpdateFeedbackApi(Resource): BillingService.update_account_deletion_feedback(args.email, args.feedback) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.route("/account/education/verify") @@ -555,18 +540,19 @@ class EducationVerifyApi(Resource): @account_initialization_required @only_edition_cloud @cloud_edition_billing_enabled - @console_ns.response(200, "Success", console_ns.models[EducationVerifyResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[EducationVerifyResponse.__name__]) @with_current_user def get(self, account: Account): - return EducationVerifyResponse.model_validate( - BillingService.EducationIdentity.verify(account.id, account.email) or {} - ).model_dump(mode="json") + return dump_response( + EducationVerifyResponse, BillingService.EducationIdentity.verify(account.id, account.email) or {} + ) @console_ns.route("/account/education") class EducationApi(Resource): @console_ns.expect(console_ns.models[EducationActivatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[EducationActivateResponse.__name__]) + # response-contract:ignore billing-service activation payload; TODO: model education activation result. + @console_ns.response(HTTPStatus.OK, "Success") @setup_required @login_required @account_initialization_required @@ -577,21 +563,22 @@ class EducationApi(Resource): payload = console_ns.payload or {} args = EducationActivatePayload.model_validate(payload) - return BillingService.EducationIdentity.activate(account, args.token, args.institution, args.role) + result = BillingService.EducationIdentity.activate(account, args.token, args.institution, args.role) + return result @setup_required @login_required @account_initialization_required @only_edition_cloud @cloud_edition_billing_enabled - @console_ns.response(200, "Success", console_ns.models[EducationStatusResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[EducationStatusResponse.__name__]) @with_current_user def get(self, account: Account): res = BillingService.EducationIdentity.status(account.id) or {} # convert expire_at to UTC timestamp from isoformat if res and "expire_at" in res: res["expire_at"] = datetime.fromisoformat(res["expire_at"]).astimezone(pytz.utc) - return EducationStatusResponse.model_validate(res).model_dump(mode="json") + return dump_response(EducationStatusResponse, res) @console_ns.route("/account/education/autocomplete") @@ -602,20 +589,21 @@ class EducationAutoCompleteApi(Resource): @account_initialization_required @only_edition_cloud @cloud_edition_billing_enabled - @console_ns.response(200, "Success", console_ns.models[EducationAutocompleteResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[EducationAutocompleteResponse.__name__]) def get(self): payload = request.args.to_dict(flat=True) args = EducationAutocompleteQuery.model_validate(payload) - return EducationAutocompleteResponse.model_validate( - BillingService.EducationIdentity.autocomplete(args.keywords, args.page, args.limit) or {} - ).model_dump(mode="json") + return dump_response( + EducationAutocompleteResponse, + BillingService.EducationIdentity.autocomplete(args.keywords, args.page, args.limit) or {}, + ) @console_ns.route("/account/change-email") class ChangeEmailSendEmailApi(Resource): @console_ns.expect(console_ns.models[ChangeEmailSendPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultDataResponse.__name__]) @enable_change_email @setup_required @login_required @@ -669,13 +657,13 @@ class ChangeEmailSendEmailApi(Resource): language=language, phase=send_phase, ) - return {"result": "success", "data": token} + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @console_ns.route("/account/change-email/validity") class ChangeEmailCheckApi(Resource): @console_ns.expect(console_ns.models[ChangeEmailValidityPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[VerificationTokenResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[VerificationTokenResponse.__name__]) @enable_change_email @setup_required @login_required @@ -716,7 +704,9 @@ class ChangeEmailCheckApi(Resource): new_token = AccountService.generate_change_email_token(refreshed_token_data, current_user) AccountService.reset_change_email_error_rate_limit(user_email) - return {"is_valid": True, "email": normalized_token_email, "token": new_token} + return VerificationTokenResponse(is_valid=True, email=normalized_token_email, token=new_token).model_dump( + mode="json" + ) @console_ns.route("/account/change-email/reset") @@ -726,7 +716,7 @@ class ChangeEmailResetApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountResponse.__name__]) @with_current_user def post(self, current_user: Account): payload = console_ns.payload or {} @@ -736,7 +726,7 @@ class ChangeEmailResetApi(Resource): if AccountService.is_account_in_freeze(normalized_new_email): raise AccountInFreezeError() - if not AccountService.check_email_unique(normalized_new_email, session=db.session): + if not AccountService.check_email_unique(normalized_new_email, session=db.session()): raise EmailAlreadyInUseError() reset_data = AccountService.get_change_email_data(args.token) @@ -761,20 +751,20 @@ class ChangeEmailResetApi(Resource): AccountService.revoke_change_email_token(args.token) updated_account = AccountService.update_account_email( - current_user, email=normalized_new_email, session=db.session + current_user, email=normalized_new_email, session=db.session() ) AccountService.send_change_email_completed_notify_email( email=normalized_new_email, ) - return _serialize_account(updated_account) + return dump_response(AccountResponse, updated_account) @console_ns.route("/account/change-email/check-email-unique") class CheckEmailUnique(Resource): @console_ns.expect(console_ns.models[CheckEmailUniquePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required def post(self): payload = console_ns.payload or {} @@ -782,6 +772,6 @@ class CheckEmailUnique(Resource): normalized_email = args.email.lower() if AccountService.is_account_in_freeze(normalized_email): raise AccountInFreezeError() - if not AccountService.check_email_unique(normalized_email, session=db.session): + if not AccountService.check_email_unique(normalized_email, session=db.session()): raise EmailAlreadyInUseError() - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") diff --git a/api/controllers/console/workspace/endpoint.py b/api/controllers/console/workspace/endpoint.py index ddb0f7045d9..781b4b6c1a2 100644 --- a/api/controllers/console/workspace/endpoint.py +++ b/api/controllers/console/workspace/endpoint.py @@ -6,13 +6,17 @@ verb-based aliases stay available as deprecated resources so OpenAPI metadata marks only the legacy paths as deprecated. """ +from datetime import datetime +from enum import StrEnum +from http import HTTPStatus from typing import Any from flask import request from flask_restx import Resource from pydantic import BaseModel, Field -from controllers.common.schema import query_params_from_model, register_schema_models +from controllers.common.fields import SuccessResponse +from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.console import console_ns from controllers.console.wraps import ( RBACPermission, @@ -24,8 +28,14 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user_id, ) +from core.entities.parameter_entities import ( + AppSelectorScope, + ModelSelectorScope, + ToolSelectorScope, +) +from core.entities.provider_entities import ProviderConfigType from core.plugin.impl.exc import PluginPermissionDeniedError -from graphon.model_runtime.utils.encoders import jsonable_encoder +from fields.base import ResponseModel from libs.login import login_required from services.plugin.endpoint_service import EndpointService @@ -40,14 +50,17 @@ class EndpointIdPayload(BaseModel): endpoint_id: str -class EndpointUpdatePayload(BaseModel): +class EndpointSettingsPayload(BaseModel): settings: dict[str, Any] name: str = Field(min_length=1) -class LegacyEndpointUpdatePayload(EndpointIdPayload): - settings: dict[str, Any] - name: str = Field(min_length=1) +class EndpointUpdatePayload(EndpointSettingsPayload): + pass + + +class LegacyEndpointUpdatePayload(EndpointIdPayload, EndpointSettingsPayload): + pass class EndpointListQuery(BaseModel): @@ -59,98 +72,158 @@ class EndpointListForPluginQuery(EndpointListQuery): plugin_id: str -class EndpointCreateResponse(BaseModel): - success: bool = Field(description="Operation success") +class EndpointProviderConfigScope(StrEnum): + ALL = AppSelectorScope.ALL.value + CHAT = AppSelectorScope.CHAT.value + WORKFLOW = AppSelectorScope.WORKFLOW.value + COMPLETION = AppSelectorScope.COMPLETION.value + LLM = ModelSelectorScope.LLM.value + TEXT_EMBEDDING = ModelSelectorScope.TEXT_EMBEDDING.value + RERANK = ModelSelectorScope.RERANK.value + TTS = ModelSelectorScope.TTS.value + SPEECH2TEXT = ModelSelectorScope.SPEECH2TEXT.value + MODERATION = ModelSelectorScope.MODERATION.value + VISION = ModelSelectorScope.VISION.value + CUSTOM = ToolSelectorScope.CUSTOM.value + BUILTIN = ToolSelectorScope.BUILTIN.value -class EndpointListResponse(BaseModel): - endpoints: list[dict[str, Any]] = Field( - description="Endpoint information", - ) +class EndpointProviderConfigI18nResponse(ResponseModel): + en_US: str + zh_Hans: str | None = None + pt_BR: str | None = None + ja_JP: str | None = None -class PluginEndpointListResponse(BaseModel): - endpoints: list[dict[str, Any]] = Field( - description="Endpoint information", - ) +class EndpointProviderConfigOptionResponse(ResponseModel): + value: str + label: EndpointProviderConfigI18nResponse -class EndpointDeleteResponse(BaseModel): - success: bool = Field(description="Operation success") +class EndpointProviderConfigResponse(ResponseModel): + type: ProviderConfigType + name: str + scope: EndpointProviderConfigScope | None = None + required: bool = False + default: int | str | float | bool | None = None + options: list[EndpointProviderConfigOptionResponse] | None = None + multiple: bool = False + label: EndpointProviderConfigI18nResponse | None = None + help: EndpointProviderConfigI18nResponse | None = None + url: str | None = None + placeholder: EndpointProviderConfigI18nResponse | None = None -class EndpointUpdateResponse(BaseModel): - success: bool = Field(description="Operation success") +class EndpointDeclarationResponse(ResponseModel): + path: str + method: str + hidden: bool = False -class EndpointEnableResponse(BaseModel): - success: bool = Field(description="Operation success") +class EndpointProviderDeclarationResponse(ResponseModel): + settings: list[EndpointProviderConfigResponse] = Field(default_factory=list) + endpoints: list[EndpointDeclarationResponse] | None = Field(default_factory=list) -class EndpointDisableResponse(BaseModel): - success: bool = Field(description="Operation success") +class EndpointListItemResponse(ResponseModel): + id: str + created_at: datetime + updated_at: datetime + tenant_id: str + plugin_id: str + settings: dict[str, Any] + expired_at: datetime + declaration: EndpointProviderDeclarationResponse = Field(default_factory=EndpointProviderDeclarationResponse) + name: str + enabled: bool + url: str + hook_id: str + + +class EndpointListResponse(ResponseModel): + endpoints: list[EndpointListItemResponse] = Field(description="Endpoint information") register_schema_models( console_ns, EndpointCreatePayload, EndpointIdPayload, + EndpointSettingsPayload, EndpointUpdatePayload, LegacyEndpointUpdatePayload, EndpointListQuery, EndpointListForPluginQuery, - EndpointCreateResponse, +) +register_response_schema_models( + console_ns, + SuccessResponse, + EndpointProviderConfigOptionResponse, + EndpointProviderConfigResponse, + EndpointDeclarationResponse, + EndpointProviderDeclarationResponse, + EndpointListItemResponse, EndpointListResponse, - PluginEndpointListResponse, - EndpointDeleteResponse, - EndpointUpdateResponse, - EndpointEnableResponse, - EndpointDisableResponse, ) -def _create_endpoint(tenant_id: str, user_id: str) -> dict[str, bool]: +def _create_endpoint(tenant_id: str, user_id: str) -> bool: """Create a plugin endpoint for the injected workspace and user.""" args = EndpointCreatePayload.model_validate(console_ns.payload) try: - return { - "success": EndpointService.create_endpoint( - tenant_id=tenant_id, - user_id=user_id, - plugin_unique_identifier=args.plugin_unique_identifier, - name=args.name, - settings=args.settings, - ) - } + return EndpointService.create_endpoint( + tenant_id=tenant_id, + user_id=user_id, + plugin_unique_identifier=args.plugin_unique_identifier, + name=args.name, + settings=args.settings, + ) except PluginPermissionDeniedError as e: raise ValueError(e.description) from e -def _update_endpoint(tenant_id: str, user_id: str, endpoint_id: str) -> dict[str, bool]: +def _update_endpoint(tenant_id: str, user_id: str, endpoint_id: str) -> bool: """Update a plugin endpoint identified by the canonical path parameter.""" args = EndpointUpdatePayload.model_validate(console_ns.payload) - return { - "success": EndpointService.update_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - name=args.name, - settings=args.settings, - ) - } + return EndpointService.update_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + name=args.name, + settings=args.settings, + ) -def _delete_endpoint(tenant_id: str, user_id: str, endpoint_id: str) -> dict[str, bool]: +def _legacy_update_endpoint(tenant_id: str, user_id: str) -> bool: + args = LegacyEndpointUpdatePayload.model_validate(console_ns.payload) + return EndpointService.update_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=args.endpoint_id, + name=args.name, + settings=args.settings, + ) + + +def _delete_endpoint(tenant_id: str, user_id: str, endpoint_id: str) -> bool: """Delete a plugin endpoint identified by the canonical path parameter.""" - return { - "success": EndpointService.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - } + return EndpointService.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + +def _delete_endpoint_from_payload(tenant_id: str, user_id: str) -> bool: + args = EndpointIdPayload.model_validate(console_ns.payload) + return _delete_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id) + + +def _set_endpoint_enabled(tenant_id: str, user_id: str, *, enabled: bool) -> bool: + args = EndpointIdPayload.model_validate(console_ns.payload) + action = EndpointService.enable_endpoint if enabled else EndpointService.disable_endpoint + return action(tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id) @console_ns.route("/workspaces/current/endpoints") @@ -161,11 +234,11 @@ class EndpointCollectionApi(Resource): @console_ns.doc(description="Create a new plugin endpoint") @console_ns.expect(console_ns.models[EndpointCreatePayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint created successfully", - console_ns.models[EndpointCreateResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -174,7 +247,7 @@ class EndpointCollectionApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - return _create_endpoint(tenant_id=tenant_id, user_id=user_id) + return SuccessResponse(success=_create_endpoint(tenant_id=tenant_id, user_id=user_id)).model_dump(mode="json") @console_ns.route("/workspaces/current/endpoints/create") @@ -190,11 +263,11 @@ class DeprecatedEndpointCreateApi(Resource): ) @console_ns.expect(console_ns.models[EndpointCreatePayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint created successfully", - console_ns.models[EndpointCreateResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -203,7 +276,7 @@ class DeprecatedEndpointCreateApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - return _create_endpoint(tenant_id=tenant_id, user_id=user_id) + return SuccessResponse(success=_create_endpoint(tenant_id=tenant_id, user_id=user_id)).model_dump(mode="json") @console_ns.route("/workspaces/current/endpoints/list") @@ -212,7 +285,7 @@ class EndpointListApi(Resource): @console_ns.doc(description="List plugin endpoints with pagination") @console_ns.doc(params=query_params_from_model(EndpointListQuery)) @console_ns.response( - 200, + HTTPStatus.OK, "Success", console_ns.models[EndpointListResponse.__name__], ) @@ -224,20 +297,15 @@ class EndpointListApi(Resource): def get(self, tenant_id: str, user_id: str): args = EndpointListQuery.model_validate(request.args.to_dict(flat=True)) - page = args.page - page_size = args.page_size - - return jsonable_encoder( - { - "endpoints": EndpointService.list_endpoints( - tenant_id=tenant_id, - user_id=user_id, - page=page, - page_size=page_size, - ) - } + endpoints = EndpointService.list_endpoints( + tenant_id=tenant_id, + user_id=user_id, + page=args.page, + page_size=args.page_size, ) + return EndpointListResponse(endpoints=endpoints).model_dump(mode="json") + @console_ns.route("/workspaces/current/endpoints/list/plugin") class EndpointListForSinglePluginApi(Resource): @@ -245,9 +313,9 @@ class EndpointListForSinglePluginApi(Resource): @console_ns.doc(description="List endpoints for a specific plugin") @console_ns.doc(params=query_params_from_model(EndpointListForPluginQuery)) @console_ns.response( - 200, + HTTPStatus.OK, "Success", - console_ns.models[PluginEndpointListResponse.__name__], + console_ns.models[EndpointListResponse.__name__], ) @setup_required @login_required @@ -257,22 +325,16 @@ class EndpointListForSinglePluginApi(Resource): def get(self, tenant_id: str, user_id: str): args = EndpointListForPluginQuery.model_validate(request.args.to_dict(flat=True)) - page = args.page - page_size = args.page_size - plugin_id = args.plugin_id - - return jsonable_encoder( - { - "endpoints": EndpointService.list_endpoints_for_single_plugin( - tenant_id=tenant_id, - user_id=user_id, - plugin_id=plugin_id, - page=page, - page_size=page_size, - ) - } + endpoints = EndpointService.list_endpoints_for_single_plugin( + tenant_id=tenant_id, + user_id=user_id, + plugin_id=args.plugin_id, + page=args.page, + page_size=args.page_size, ) + return EndpointListResponse(endpoints=endpoints).model_dump(mode="json") + @console_ns.route("/workspaces/current/endpoints/") class EndpointItemApi(Resource): @@ -282,11 +344,11 @@ class EndpointItemApi(Resource): @console_ns.doc(description="Delete a plugin endpoint") @console_ns.doc(params={"id": {"description": "Endpoint ID", "type": "string", "required": True}}) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint deleted successfully", - console_ns.models[EndpointDeleteResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -295,18 +357,20 @@ class EndpointItemApi(Resource): @with_current_user_id @with_current_tenant_id def delete(self, tenant_id: str, user_id: str, id: str): - return _delete_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=id) + return SuccessResponse( + success=_delete_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=id) + ).model_dump(mode="json") @console_ns.doc("update_endpoint") @console_ns.doc(description="Update a plugin endpoint") @console_ns.expect(console_ns.models[EndpointUpdatePayload.__name__]) @console_ns.doc(params={"id": {"description": "Endpoint ID", "type": "string", "required": True}}) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint updated successfully", - console_ns.models[EndpointUpdateResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -315,7 +379,9 @@ class EndpointItemApi(Resource): @with_current_user_id @with_current_tenant_id def patch(self, tenant_id: str, user_id: str, id: str): - return _update_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=id) + return SuccessResponse( + success=_update_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=id) + ).model_dump(mode="json") @console_ns.route("/workspaces/current/endpoints/delete") @@ -332,11 +398,11 @@ class DeprecatedEndpointDeleteApi(Resource): ) @console_ns.expect(console_ns.models[EndpointIdPayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint deleted successfully", - console_ns.models[EndpointDeleteResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -345,8 +411,9 @@ class DeprecatedEndpointDeleteApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - args = EndpointIdPayload.model_validate(console_ns.payload) - return _delete_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id) + return SuccessResponse(success=_delete_endpoint_from_payload(tenant_id=tenant_id, user_id=user_id)).model_dump( + mode="json" + ) @console_ns.route("/workspaces/current/endpoints/update") @@ -363,11 +430,11 @@ class DeprecatedEndpointUpdateApi(Resource): ) @console_ns.expect(console_ns.models[LegacyEndpointUpdatePayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint updated successfully", - console_ns.models[EndpointUpdateResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -376,8 +443,9 @@ class DeprecatedEndpointUpdateApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - args = LegacyEndpointUpdatePayload.model_validate(console_ns.payload) - return _update_endpoint(tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id) + return SuccessResponse(success=_legacy_update_endpoint(tenant_id=tenant_id, user_id=user_id)).model_dump( + mode="json" + ) @console_ns.route("/workspaces/current/endpoints/enable") @@ -386,11 +454,11 @@ class EndpointEnableApi(Resource): @console_ns.doc(description="Enable a plugin endpoint") @console_ns.expect(console_ns.models[EndpointIdPayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint enabled successfully", - console_ns.models[EndpointEnableResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -399,13 +467,9 @@ class EndpointEnableApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - args = EndpointIdPayload.model_validate(console_ns.payload) - - return { - "success": EndpointService.enable_endpoint( - tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id - ) - } + return SuccessResponse( + success=_set_endpoint_enabled(tenant_id=tenant_id, user_id=user_id, enabled=True) + ).model_dump(mode="json") @console_ns.route("/workspaces/current/endpoints/disable") @@ -414,11 +478,11 @@ class EndpointDisableApi(Resource): @console_ns.doc(description="Disable a plugin endpoint") @console_ns.expect(console_ns.models[EndpointIdPayload.__name__]) @console_ns.response( - 200, + HTTPStatus.OK, "Endpoint disabled successfully", - console_ns.models[EndpointDisableResponse.__name__], + console_ns.models[SuccessResponse.__name__], ) - @console_ns.response(403, "Admin privileges required") + @console_ns.response(HTTPStatus.FORBIDDEN, "Admin privileges required") @setup_required @login_required @is_admin_or_owner_required @@ -427,10 +491,6 @@ class EndpointDisableApi(Resource): @with_current_user_id @with_current_tenant_id def post(self, tenant_id: str, user_id: str): - args = EndpointIdPayload.model_validate(console_ns.payload) - - return { - "success": EndpointService.disable_endpoint( - tenant_id=tenant_id, user_id=user_id, endpoint_id=args.endpoint_id - ) - } + return SuccessResponse( + success=_set_endpoint_enabled(tenant_id=tenant_id, user_id=user_id, enabled=False) + ).model_dump(mode="json") diff --git a/api/controllers/console/workspace/load_balancing_config.py b/api/controllers/console/workspace/load_balancing_config.py index 5983a4e10be..abeb691be03 100644 --- a/api/controllers/console/workspace/load_balancing_config.py +++ b/api/controllers/console/workspace/load_balancing_config.py @@ -10,6 +10,7 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user, ) +from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.errors.validate import CredentialsValidateFailedError @@ -69,6 +70,7 @@ class LoadBalancingCredentialsValidateApi(Resource): model=payload.model, model_type=payload.model_type, credentials=payload.credentials, + session=db.session(), ) except CredentialsValidateFailedError as ex: result = False @@ -118,6 +120,7 @@ class LoadBalancingConfigCredentialsValidateApi(Resource): model=payload.model, model_type=payload.model_type, credentials=payload.credentials, + session=db.session(), config_id=config_id, ) except CredentialsValidateFailedError as ex: diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index 51a6c9b0f3f..72330aba5f0 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -1,10 +1,12 @@ +from http import HTTPStatus from urllib import parse from uuid import UUID from flask import abort, request from flask_restx import Resource -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, field_validator from sqlalchemy import func, select +from werkzeug.exceptions import NotFound import services from configs import dify_config @@ -30,8 +32,8 @@ from controllers.console.wraps import ( from extensions.ext_database import db from extensions.ext_redis import redis_client from fields.base import ResponseModel -from fields.member_fields import AccountWithRole, AccountWithRoleList -from libs.helper import extract_remote_ip +from fields.member_fields import AccountWithRoleListResponse, AccountWithRoleResponse +from libs.helper import dump_response, extract_remote_ip from libs.login import current_account_with_tenant, login_required from models.account import Account, TenantAccountJoin, TenantAccountRole from services.account_service import AccountService, RegisterService, TenantService @@ -45,6 +47,11 @@ class MemberInvitePayload(BaseModel): role: str language: str | None = None + @field_validator("emails") + @classmethod + def normalize_emails(cls, emails: list[str]) -> list[str]: + return list(dict.fromkeys(email.lower() for email in emails)) + class MemberRoleUpdatePayload(BaseModel): role: str @@ -70,14 +77,14 @@ class MemberInviteResultResponse(ResponseModel): message: str | None = None -class MemberInviteResponse(ResponseModel): +class MemberActionResponse(ResponseModel): result: str - invitation_results: list[MemberInviteResultResponse] tenant_id: str -class MemberActionTenantResponse(ResponseModel): +class MemberInviteResponse(ResponseModel): result: str + invitation_results: list[MemberInviteResultResponse] tenant_id: str @@ -92,13 +99,14 @@ register_schema_models( ) register_response_schema_models( console_ns, - AccountWithRole, - AccountWithRoleList, + AccountWithRoleResponse, + AccountWithRoleListResponse, + MemberActionResponse, + MemberInviteResponse, + MemberInviteResultResponse, SimpleResultDataResponse, SimpleResultResponse, VerificationTokenResponse, - MemberInviteResponse, - MemberActionTenantResponse, ) @@ -124,14 +132,10 @@ def _normalize_enum_value(value: object) -> str: return str(normalized) if normalized is not None else "" -def _normalize_invitee_emails(emails: list[str]) -> list[str]: - return list(dict.fromkeys(email.lower() for email in emails)) - - def _count_new_member_invites(tenant_id: str, emails: list[str]) -> int: new_member_count = 0 for email in emails: - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) if not account: new_member_count += 1 continue @@ -179,14 +183,14 @@ class MemberListApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountWithRoleList.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountWithRoleListResponse.__name__]) @with_current_user def get(self, current_user: Account | None = None): if current_user is None: current_user, _ = current_account_with_tenant() if not current_user.current_tenant: raise ValueError("No current tenant") - members = TenantService.get_tenant_members(current_user.current_tenant, session=db.session) + members = TenantService.get_tenant_members(current_user.current_tenant, session=db.session()) if dify_config.RBAC_ENABLED: member_ids = [member.id for member in members] member_roles = enterprise_rbac_service.RBACService.MemberRoles.batch_get( @@ -216,9 +220,7 @@ class MemberListApi(Resource): } ) - member_models = TypeAdapter(list[AccountWithRole]).validate_python(serialized_members) - response = AccountWithRoleList(accounts=member_models) - return response.model_dump(mode="json"), 200 + return dump_response(AccountWithRoleListResponse, {"accounts": serialized_members}), HTTPStatus.OK @console_ns.route("/workspaces/current/members/invite-email") @@ -226,7 +228,7 @@ class MemberInviteEmailApi(Resource): """Invite a new member by email.""" @console_ns.expect(console_ns.models[MemberInvitePayload.__name__]) - @console_ns.response(201, "Success", console_ns.models[MemberInviteResponse.__name__]) + @console_ns.response(HTTPStatus.CREATED, "Success", console_ns.models[MemberInviteResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -235,26 +237,26 @@ class MemberInviteEmailApi(Resource): payload = console_ns.payload or {} args = MemberInvitePayload.model_validate(payload) - invitee_emails = _normalize_invitee_emails(args.emails) + invitee_emails = args.emails invitee_role = args.role interface_language = args.language if not dify_config.RBAC_ENABLED: if not TenantAccountRole.is_valid_role(invitee_role): - return {"code": "invalid-role", "message": "Invalid role"}, 400 + return {"code": "invalid-role", "message": "Invalid role"}, HTTPStatus.BAD_REQUEST if not TenantAccountRole.is_non_owner_role(TenantAccountRole(invitee_role)): - return {"code": "invalid-role", "message": "Invalid role"}, 400 + return {"code": "invalid-role", "message": "Invalid role"}, HTTPStatus.BAD_REQUEST inviter = current_user if not inviter.current_tenant: raise ValueError("No current tenant") if not _is_role_enabled(invitee_role, inviter.current_tenant.id): - return {"code": "invalid-role", "message": "Invalid role"}, 400 + return {"code": "invalid-role", "message": "Invalid role"}, HTTPStatus.BAD_REQUEST # Check workspace permission for member invitations from libs.workspace_permission import check_workspace_member_invite_permission check_workspace_member_invite_permission(inviter.current_tenant.id) - invitation_results = [] + invitation_results: list[MemberInviteResultResponse] = [] console_web_url = dify_config.CONSOLE_WEB_URL tenant_id = inviter.current_tenant.id @@ -273,67 +275,69 @@ class MemberInviteEmailApi(Resource): language=interface_language, role=invitee_role, inviter=inviter, - session=db.session, + session=db.session(), ) encoded_invitee_email = parse.quote(invitee_email) invitation_results.append( - { - "status": "success", - "email": invitee_email, - "url": f"{console_web_url}/activate?email={encoded_invitee_email}&token={token}", - } + MemberInviteResultResponse( + status="success", + email=invitee_email, + url=f"{console_web_url}/activate?email={encoded_invitee_email}&token={token}", + ) ) except AccountAlreadyInTenantError: invitation_results.append( - { - "status": "already_member", - "email": invitee_email, - "message": "Account already in workspace.", - } + MemberInviteResultResponse( + status="already_member", + email=invitee_email, + message="Account already in workspace.", + ) ) except Exception as e: - invitation_results.append({"status": "failed", "email": invitee_email, "message": str(e)}) + invitation_results.append( + MemberInviteResultResponse(status="failed", email=invitee_email, message=str(e)) + ) - return { - "result": "success", - "invitation_results": invitation_results, - "tenant_id": str(inviter.current_tenant.id) if inviter.current_tenant else "", - }, 201 + return MemberInviteResponse( + result="success", + invitation_results=invitation_results, + tenant_id=inviter.current_tenant.id if inviter.current_tenant else "", + ).model_dump(mode="json"), HTTPStatus.CREATED @console_ns.route("/workspaces/current/members/") class MemberCancelInviteApi(Resource): """Cancel an invitation by member id.""" - @console_ns.response(200, "Success", console_ns.models[MemberActionTenantResponse.__name__]) @setup_required @login_required @account_initialization_required + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[MemberActionResponse.__name__]) @with_current_user def delete(self, current_user: Account, member_id: UUID): if not current_user.current_tenant: raise ValueError("No current tenant") member = db.session.get(Account, str(member_id)) if member is None: - abort(404) + abort(HTTPStatus.NOT_FOUND) else: try: TenantService.remove_member_from_tenant( - current_user.current_tenant, member, current_user, session=db.session + current_user.current_tenant, member, current_user, session=db.session() ) except services.errors.account.CannotOperateSelfError as e: - return {"code": "cannot-operate-self", "message": str(e)}, 400 + return {"code": "cannot-operate-self", "message": str(e)}, HTTPStatus.BAD_REQUEST except services.errors.account.NoPermissionError as e: - return {"code": "forbidden", "message": str(e)}, 403 + return {"code": "forbidden", "message": str(e)}, HTTPStatus.FORBIDDEN except services.errors.account.MemberNotInTenantError as e: - return {"code": "member-not-found", "message": str(e)}, 404 + return {"code": "member-not-found", "message": str(e)}, HTTPStatus.NOT_FOUND except Exception as e: raise ValueError(str(e)) - return { - "result": "success", - "tenant_id": str(current_user.current_tenant.id) if current_user.current_tenant else "", - }, 200 + return MemberActionResponse( + result="success", + tenant_id=current_user.current_tenant.id if current_user.current_tenant else "", + ).model_dump(mode="json"), HTTPStatus.OK @console_ns.route("/workspaces/current/members//update-role") @@ -341,7 +345,7 @@ class MemberUpdateRoleApi(Resource): """Update member role.""" @console_ns.expect(console_ns.models[MemberRoleUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -352,32 +356,32 @@ class MemberUpdateRoleApi(Resource): new_role = args.role if not TenantAccountRole.is_valid_role(new_role): - return {"code": "invalid-role", "message": "Invalid role"}, 400 + return {"code": "invalid-role", "message": "Invalid role"}, HTTPStatus.BAD_REQUEST if not current_user.current_tenant: raise ValueError("No current tenant") if not _is_role_enabled(new_role, current_user.current_tenant.id): - return {"code": "invalid-role", "message": "Invalid role"}, 400 + return {"code": "invalid-role", "message": "Invalid role"}, HTTPStatus.BAD_REQUEST member = db.session.get(Account, str(member_id)) if not member: - abort(404) + abort(HTTPStatus.NOT_FOUND) try: assert member is not None, "Member not found" TenantService.update_member_role( - current_user.current_tenant, member, new_role, current_user, session=db.session + current_user.current_tenant, member, new_role, current_user, session=db.session() ) except services.errors.account.CannotOperateSelfError as e: - return {"code": "cannot-operate-self", "message": str(e)}, 400 + return {"code": "cannot-operate-self", "message": str(e)}, HTTPStatus.BAD_REQUEST except services.errors.account.NoPermissionError as e: - return {"code": "forbidden", "message": str(e)}, 403 + return {"code": "forbidden", "message": str(e)}, HTTPStatus.FORBIDDEN except services.errors.account.MemberNotInTenantError as e: - return {"code": "member-not-found", "message": str(e)}, 404 + return {"code": "member-not-found", "message": str(e)}, HTTPStatus.NOT_FOUND except services.errors.account.RoleAlreadyAssignedError as e: - return {"code": "role-already-assigned", "message": str(e)}, 400 + return {"code": "role-already-assigned", "message": str(e)}, HTTPStatus.BAD_REQUEST except Exception as e: raise ValueError(str(e)) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.route("/workspaces/current/dataset-operators") @@ -387,15 +391,13 @@ class DatasetOperatorMemberListApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[AccountWithRoleList.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AccountWithRoleListResponse.__name__]) @with_current_user def get(self, current_user: Account): if not current_user.current_tenant: raise ValueError("No current tenant") - members = TenantService.get_dataset_operator_members(current_user.current_tenant, session=db.session) - member_models = TypeAdapter(list[AccountWithRole]).validate_python(members, from_attributes=True) - response = AccountWithRoleList(accounts=member_models) - return response.model_dump(mode="json"), 200 + members = TenantService.get_dataset_operator_members(current_user.current_tenant, session=db.session()) + return dump_response(AccountWithRoleListResponse, {"accounts": members}), HTTPStatus.OK @console_ns.route("/workspaces/current/members/send-owner-transfer-confirm-email") @@ -403,7 +405,7 @@ class SendOwnerTransferEmailApi(Resource): """Send owner transfer email.""" @console_ns.expect(console_ns.models[OwnerTransferEmailPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultDataResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultDataResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -418,7 +420,7 @@ class SendOwnerTransferEmailApi(Resource): # check if the current user is the owner of the workspace if not current_user.current_tenant: raise ValueError("No current tenant") - if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session): + if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session()): raise NotOwnerError() if args.language is not None and args.language == "zh-Hans": @@ -435,13 +437,13 @@ class SendOwnerTransferEmailApi(Resource): workspace_name=current_user.current_tenant.name if current_user.current_tenant else "", ) - return {"result": "success", "data": token} + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @console_ns.route("/workspaces/current/members/owner-transfer-check") class OwnerTransferCheckApi(Resource): @console_ns.expect(console_ns.models[OwnerTransferCheckPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[VerificationTokenResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[VerificationTokenResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -453,7 +455,7 @@ class OwnerTransferCheckApi(Resource): # check if the current user is the owner of the workspace if not current_user.current_tenant: raise ValueError("No current tenant") - if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session): + if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session()): raise NotOwnerError() user_email = current_user.email @@ -480,13 +482,13 @@ class OwnerTransferCheckApi(Resource): _, new_token = AccountService.generate_owner_transfer_token(user_email, code=args.code, additional_data={}) AccountService.reset_owner_transfer_error_rate_limit(user_email) - return {"is_valid": True, "email": token_data.get("email"), "token": new_token} + return VerificationTokenResponse(is_valid=True, email=user_email, token=new_token).model_dump(mode="json") @console_ns.route("/workspaces/current/members//owner-transfer") class OwnerTransfer(Resource): @console_ns.expect(console_ns.models[OwnerTransferPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -499,7 +501,7 @@ class OwnerTransfer(Resource): # check if the current user is the owner of the workspace if not current_user.current_tenant: raise ValueError("No current tenant") - if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session): + if not TenantService.is_owner(current_user, current_user.current_tenant, session=db.session()): raise NotOwnerError() if current_user.id == str(member_id): @@ -516,18 +518,17 @@ class OwnerTransfer(Resource): member = db.session.get(Account, str(member_id)) if not member: - abort(404) - return # Never reached, but helps type checker + raise NotFound() if not current_user.current_tenant: raise ValueError("No current tenant") - if not TenantService.is_member(member, current_user.current_tenant, session=db.session): + if not TenantService.is_member(member, current_user.current_tenant, session=db.session()): raise MemberNotInTenantError() try: assert member is not None, "Member not found" TenantService.update_member_role( - current_user.current_tenant, member, "owner", current_user, session=db.session + current_user.current_tenant, member, "owner", current_user, session=db.session() ) AccountService.send_new_owner_transfer_notify_email( @@ -546,4 +547,4 @@ class OwnerTransfer(Resource): except Exception as e: raise ValueError(str(e)) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") diff --git a/api/controllers/console/workspace/model_providers.py b/api/controllers/console/workspace/model_providers.py index 779399f9055..3bafa2ab6a6 100644 --- a/api/controllers/console/workspace/model_providers.py +++ b/api/controllers/console/workspace/model_providers.py @@ -353,7 +353,7 @@ class ModelProviderPaymentCheckoutUrlApi(Resource): def get(self, current_tenant_id: str, current_user: Account, provider: str): if provider != "anthropic": raise ValueError(f"provider name {provider} is invalid") - BillingService.is_tenant_owner_or_admin(db.session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=db.session()) data = BillingService.get_model_provider_payment_link( provider_name=provider, tenant_id=current_tenant_id, diff --git a/api/controllers/console/workspace/models.py b/api/controllers/console/workspace/models.py index 1da72ef4362..0f735a7479e 100644 --- a/api/controllers/console/workspace/models.py +++ b/api/controllers/console/workspace/models.py @@ -24,6 +24,7 @@ from controllers.console.wraps import ( with_current_user, ) from core.entities.provider_entities import CredentialConfiguration +from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.entities.model_entities import ModelType, ParameterRule from graphon.model_runtime.errors.validate import CredentialsValidateFailedError @@ -297,6 +298,7 @@ class ModelProviderModelApi(Resource): model_type=args.model_type, configs=args.load_balancing.configs, config_from=args.config_from or "", + session=db.session(), ) if args.load_balancing.enabled: @@ -356,6 +358,7 @@ class ModelProviderModelCredentialApi(Resource): provider=provider, model=args.model, model_type=args.model_type, + session=db.session(), config_from=args.config_from or "", ) diff --git a/api/controllers/console/workspace/plugin.py b/api/controllers/console/workspace/plugin.py index 2dede15330c..c7644af2b48 100644 --- a/api/controllers/console/workspace/plugin.py +++ b/api/controllers/console/workspace/plugin.py @@ -38,11 +38,20 @@ from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolProviderType from core.tools.tool_manager import ToolManager +from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.utils.encoders import jsonable_encoder from libs.helper import dump_response from libs.login import login_required -from models.account import Account, TenantPluginAutoUpgradeStrategy, TenantPluginPermission +from models.account import ( + Account, + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategy, + TenantPluginAutoUpgradeStrategySetting, + TenantPluginDebugPermission, + TenantPluginInstallPermission, +) from models.provider_ids import ToolProviderID from services.entities.model_provider_entities import ProviderEntityResponse from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService @@ -52,9 +61,9 @@ from services.tools.tools_transform_service import ToolTransformService class AutoUpgradeSettingsResponse(TypedDict): - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting + strategy_setting: TenantPluginAutoUpgradeStrategySetting upgrade_time_of_day: int - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode + upgrade_mode: TenantPluginAutoUpgradeMode exclude_plugins: list[str] include_plugins: list[str] @@ -127,8 +136,8 @@ class ParserUninstall(BaseModel): class ParserPermissionChange(BaseModel): - install_permission: TenantPluginPermission.InstallPermission = TenantPluginPermission.InstallPermission.EVERYONE - debug_permission: TenantPluginPermission.DebugPermission = TenantPluginPermission.DebugPermission.EVERYONE + install_permission: TenantPluginInstallPermission = TenantPluginInstallPermission.EVERYONE + debug_permission: TenantPluginDebugPermission = TenantPluginDebugPermission.EVERYONE class ParserDynamicOptions(BaseModel): @@ -150,16 +159,14 @@ class ParserDynamicOptionsWithCredentials(BaseModel): class PluginPermissionSettingsPayload(BaseModel): - install_permission: TenantPluginPermission.InstallPermission = TenantPluginPermission.InstallPermission.EVERYONE - debug_permission: TenantPluginPermission.DebugPermission = TenantPluginPermission.DebugPermission.EVERYONE + install_permission: TenantPluginInstallPermission = TenantPluginInstallPermission.EVERYONE + debug_permission: TenantPluginDebugPermission = TenantPluginDebugPermission.EVERYONE class PluginAutoUpgradeSettingsPayload(BaseModel): - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting = ( - TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY - ) + strategy_setting: TenantPluginAutoUpgradeStrategySetting = TenantPluginAutoUpgradeStrategySetting.FIX_ONLY upgrade_time_of_day: int = 0 - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode = TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE + upgrade_mode: TenantPluginAutoUpgradeMode = TenantPluginAutoUpgradeMode.EXCLUDE exclude_plugins: list[str] = Field(default_factory=list) include_plugins: list[str] = Field(default_factory=list) @@ -170,15 +177,15 @@ class PluginAutoUpgradeChangeResponse(ResponseModel): class PluginAutoUpgradeSettingsResponseModel(ResponseModel): - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting + strategy_setting: TenantPluginAutoUpgradeStrategySetting upgrade_time_of_day: int - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode + upgrade_mode: TenantPluginAutoUpgradeMode exclude_plugins: list[str] include_plugins: list[str] class PluginAutoUpgradeFetchResponse(ResponseModel): - category: TenantPluginAutoUpgradeStrategy.PluginCategory + category: TenantPluginAutoUpgradeCategory auto_upgrade: PluginAutoUpgradeSettingsResponseModel @@ -209,19 +216,19 @@ class PluginDeclarationResponse(ResponseModel): class ParserAutoUpgradeChange(BaseModel): model_config = ConfigDict(extra="forbid") - category: TenantPluginAutoUpgradeStrategy.PluginCategory + category: TenantPluginAutoUpgradeCategory auto_upgrade: PluginAutoUpgradeSettingsPayload class ParserAutoUpgradeFetch(BaseModel): - category: TenantPluginAutoUpgradeStrategy.PluginCategory + category: TenantPluginAutoUpgradeCategory class ParserExcludePlugin(BaseModel): model_config = ConfigDict(extra="forbid") plugin_id: str - category: TenantPluginAutoUpgradeStrategy.PluginCategory + category: TenantPluginAutoUpgradeCategory class ParserReadme(BaseModel): @@ -339,8 +346,8 @@ class PluginTaskResponse(ResponseModel): class PluginPermissionResponse(ResponseModel): - install_permission: TenantPluginPermission.InstallPermission - debug_permission: TenantPluginPermission.DebugPermission + install_permission: TenantPluginInstallPermission + debug_permission: TenantPluginDebugPermission class PluginDynamicOptionsResponse(ResponseModel): @@ -408,22 +415,22 @@ register_response_schema_models( register_enum_models( console_ns, - TenantPluginPermission.DebugPermission, - TenantPluginAutoUpgradeStrategy.PluginCategory, - TenantPluginAutoUpgradeStrategy.UpgradeMode, - TenantPluginAutoUpgradeStrategy.StrategySetting, - TenantPluginPermission.InstallPermission, + TenantPluginDebugPermission, + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategySetting, + TenantPluginInstallPermission, ) def _default_auto_upgrade_settings( tenant_id: str, - category: TenantPluginAutoUpgradeStrategy.PluginCategory, + category: TenantPluginAutoUpgradeCategory, ) -> AutoUpgradeSettingsResponse: return { "strategy_setting": PluginAutoUpgradeService.default_strategy_setting_for_category(category), "upgrade_time_of_day": PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id), - "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + "upgrade_mode": TenantPluginAutoUpgradeMode.EXCLUDE, "exclude_plugins": [], "include_plugins": [], } @@ -967,7 +974,7 @@ class PluginChangePermissionApi(Resource): args = ParserPermissionChange.model_validate(console_ns.payload) set_permission_result = PluginPermissionService.change_permission( - tenant_id, args.install_permission, args.debug_permission + tenant_id, args.install_permission, args.debug_permission, session=db.session() ) if not set_permission_result: return jsonable_encoder({"success": False, "message": "Failed to set permission"}) @@ -983,12 +990,12 @@ class PluginFetchPermissionApi(Resource): @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str): - permission = PluginPermissionService.get_permission(tenant_id) + permission = PluginPermissionService.get_permission(tenant_id, session=db.session()) if not permission: return jsonable_encoder( { - "install_permission": TenantPluginPermission.InstallPermission.EVERYONE, - "debug_permission": TenantPluginPermission.DebugPermission.EVERYONE, + "install_permission": TenantPluginInstallPermission.EVERYONE, + "debug_permission": TenantPluginDebugPermission.EVERYONE, } ) @@ -1088,6 +1095,7 @@ class PluginChangeAutoUpgradeApi(Resource): auto_upgrade.exclude_plugins, auto_upgrade.include_plugins, category=args.category, + session=db.session(), ) if not set_auto_upgrade_strategy_result: return jsonable_encoder({"success": False, "message": "Failed to set auto upgrade strategy"}) @@ -1105,7 +1113,7 @@ class PluginFetchAutoUpgradeApi(Resource): @with_current_tenant_id def get(self, tenant_id: str): args = ParserAutoUpgradeFetch.model_validate(request.args.to_dict(flat=True)) - auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id, args.category) + auto_upgrade = PluginAutoUpgradeService.get_strategy(tenant_id, args.category, session=db.session()) auto_upgrade_dict = ( _auto_upgrade_settings_to_dict(auto_upgrade) if auto_upgrade @@ -1134,7 +1142,11 @@ class PluginAutoUpgradeExcludePluginApi(Resource): args = ParserExcludePlugin.model_validate(console_ns.payload) return jsonable_encoder( - {"success": PluginAutoUpgradeService.exclude_plugin(tenant_id, args.plugin_id, args.category)} + { + "success": PluginAutoUpgradeService.exclude_plugin( + tenant_id, args.plugin_id, args.category, session=db.session() + ) + } ) diff --git a/api/controllers/console/workspace/rbac.py b/api/controllers/console/workspace/rbac.py index 12bd7074250..9a05e5f0e28 100644 --- a/api/controllers/console/workspace/rbac.py +++ b/api/controllers/console/workspace/rbac.py @@ -14,6 +14,7 @@ from controllers.console import console_ns from controllers.console.wraps import RBACPermission, RBACResourceScope, rbac_permission_required from core.db.session_factory import session_factory from core.rbac import RBACResourceWhitelistScope +from extensions.ext_database import db from libs.login import current_account_with_tenant, login_required from models import Account from services.enterprise import rbac_service as svc @@ -564,6 +565,7 @@ class RBACMyPermissionsApi(Resource): account_id, app_id=request.args.get("app_id") or None, dataset_id=request.args.get("dataset_id") or None, + session=db.session(), ) ) @@ -902,7 +904,7 @@ class RBACMemberRolesApi(Resource): @console_ns.response(200, "Success", console_ns.models[svc.MemberRolesResponse.__name__]) def get(self, member_id): tenant_id, account_id = _current_ids() - return _dump(svc.RBACService.MemberRoles.get(tenant_id, account_id, str(member_id))) + return _dump(svc.RBACService.MemberRoles.get(tenant_id, account_id, str(member_id), session=db.session())) @login_required @console_ns.expect(console_ns.models[_ReplaceMemberRolesRequest.__name__]) @@ -916,6 +918,7 @@ class RBACMemberRolesApi(Resource): account_id, str(member_id), role_ids=list(request.role_ids), + session=db.session(), ) ) diff --git a/api/controllers/console/workspace/snippets.py b/api/controllers/console/workspace/snippets.py index 3a8db2dcaa1..18407b62d52 100644 --- a/api/controllers/console/workspace/snippets.py +++ b/api/controllers/console/workspace/snippets.py @@ -4,7 +4,7 @@ from typing import Any from urllib.parse import quote from flask import Response, request -from flask_restx import Resource, marshal +from flask_restx import Resource from pydantic import Field as PydanticField from pydantic import field_validator from sqlalchemy.orm import Session, sessionmaker @@ -37,8 +37,7 @@ from controllers.console.wraps import ( from core.plugin.entities.plugin import PluginDependency from extensions.ext_database import db from fields.base import ResponseModel -from fields.snippet_fields import snippet_fields, snippet_list_fields -from libs.helper import to_timestamp +from libs.helper import dump_response, to_timestamp from libs.login import login_required from models import Account from models.snippet import SnippetType @@ -189,7 +188,7 @@ class CustomizedSnippetsApi(Resource): snippet_service = _snippet_service() snippets, total, has_more = snippet_service.get_snippets( tenant_id=current_tenant_id, - session=db.session, + session=db.session(), page=query.page, limit=query.limit, keyword=query.keyword, @@ -198,13 +197,16 @@ class CustomizedSnippetsApi(Resource): tag_ids=query.tag_ids, ) - return { - "data": marshal(snippets, snippet_list_fields), - "page": query.page, - "limit": query.limit, - "total": total, - "has_more": has_more, - }, 200 + return dump_response( + SnippetPaginationResponse, + { + "data": snippets, + "page": query.page, + "limit": query.limit, + "total": total, + "has_more": has_more, + }, + ), 200 @console_ns.doc("create_customized_snippet") @console_ns.expect(console_ns.models.get(CreateSnippetPayload.__name__)) @@ -245,7 +247,7 @@ class CustomizedSnippetsApi(Resource): except ValueError as e: return {"message": str(e)}, 400 - return marshal(snippet, snippet_fields), 201 + return dump_response(SnippetResponse, snippet), 201 @console_ns.route("/workspaces/current/customized-snippets/") @@ -268,7 +270,7 @@ class CustomizedSnippetDetailApi(Resource): if not snippet: raise NotFound("Snippet not found") - return marshal(snippet, snippet_fields), 200 + return dump_response(SnippetResponse, snippet), 200 @console_ns.doc("update_customized_snippet") @console_ns.expect(console_ns.models.get(UpdateSnippetPayload.__name__)) @@ -317,7 +319,7 @@ class CustomizedSnippetDetailApi(Resource): except ValueError as e: return {"message": str(e)}, 400 - return marshal(snippet, snippet_fields), 200 + return dump_response(SnippetResponse, snippet), 200 @console_ns.doc("delete_customized_snippet") @console_ns.response(204, "Snippet deleted successfully") @@ -533,4 +535,4 @@ class CustomizedSnippetUseCountIncrementApi(Resource): session.commit() session.refresh(snippet) - return {"result": "success", "use_count": snippet.use_count}, 200 + return SnippetUseCountResponse(result="success", use_count=snippet.use_count).model_dump(mode="json"), 200 diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index 7a3f158b0c3..6e9b4be8ac4 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -1,16 +1,25 @@ import io import logging -from typing import Any, Literal +from collections.abc import Iterable, Mapping +from datetime import datetime +from typing import Any, Literal, cast from urllib.parse import urlparse from flask import make_response, redirect, request, send_file from flask_restx import Resource -from pydantic import BaseModel, Field, HttpUrl, RootModel, field_validator, model_validator +from pydantic import ( + BaseModel, + Field, + HttpUrl, + RootModel, + field_validator, + model_validator, +) from sqlalchemy.orm import sessionmaker from werkzeug.exceptions import Forbidden from configs import dify_config -from controllers.common.fields import BinaryFileResponse, RedirectResponse, SimpleResultResponse +from controllers.common.fields import SimpleResultResponse from controllers.common.schema import ( query_params_from_model, query_params_from_request, @@ -31,22 +40,36 @@ from controllers.console.wraps import ( ) from core.db.session_factory import session_factory from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration +from core.entities.provider_entities import ProviderConfig from core.mcp.auth.auth_flow import auth, handle_callback from core.mcp.error import MCPAuthError, MCPError, MCPRefreshTokenError from core.mcp.mcp_client import MCPClient from core.plugin.entities.plugin_daemon import CredentialType, PluginOAuthAuthorizationUrlResponse from core.plugin.impl.oauth import OAuthHandler -from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration +from core.tools.entities.api_entities import ( + ToolApiEntity, + ToolProviderCredentialApiEntity, + ToolProviderCredentialInfoApiEntity, + ToolProviderTypeApiLiteral, +) +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_bundle import ApiToolBundle +from core.tools.entities.tool_entities import ( + ApiProviderSchemaType, + ToolLabel, + ToolProviderType, + WorkflowToolParameterConfiguration, +) from extensions.ext_database import db -from graphon.model_runtime.utils.encoders import jsonable_encoder -from libs.helper import alphanumeric, uuid_value +from fields.base import ResponseModel +from libs.helper import alphanumeric, dump_response, uuid_value from libs.login import login_required from models import Account from models.provider_ids import ToolProviderID # from models.provider_ids import ToolProviderID from services.plugin.oauth_service import OAuthProxyService -from services.tools.api_tools_manage_service import ApiToolManageService +from services.tools.api_tools_manage_service import ApiToolManageService, ApiToolPreviewResult from services.tools.builtin_tools_manage_service import BuiltinToolManageService from services.tools.mcp_tools_manage_service import MCPToolManageService, OAuthDataType from services.tools.tool_labels_service import ToolLabelsService @@ -85,16 +108,21 @@ class BuiltinToolAddPayload(BaseModel): class BuiltinToolUpdatePayload(BaseModel): credential_id: str - credentials: dict[str, Any] | None = Field(default=None) + credentials: dict[str, Any] | None = None name: str | None = Field(default=None, max_length=30) +class ToolEmojiIcon(BaseModel): + background: str + content: str + + class ApiToolProviderBasePayload(BaseModel): credentials: dict[str, Any] schema_type: ApiProviderSchemaType schema_: str = Field(alias="schema") provider: str - icon: dict[str, Any] + icon: ToolEmojiIcon privacy_policy: str | None = None labels: list[str] | None = None custom_disclaimer: str = "" @@ -144,7 +172,7 @@ class WorkflowToolBasePayload(BaseModel): name: str label: str description: str - icon: dict[str, Any] + icon: ToolEmojiIcon parameters: list[WorkflowToolParameterConfiguration] = Field(default_factory=list) privacy_policy: str | None = "" labels: list[str] | None = None @@ -214,7 +242,7 @@ class BuiltinProviderDefaultCredentialPayload(BaseModel): class ToolOAuthCustomClientPayload(BaseModel): - client_params: dict[str, Any] | None = Field(default=None) + client_params: dict[str, Any] | None = None enable_oauth_custom_client: bool | None = True @@ -225,13 +253,20 @@ class MCPProviderBasePayload(BaseModel): icon_type: str icon_background: str = "" server_identifier: str - configuration: dict[str, Any] | None = Field(default_factory=dict) - headers: dict[str, Any] | None = Field(default_factory=dict) - authentication: dict[str, Any] | None = Field(default_factory=dict) + configuration: MCPConfiguration | None = None + headers: dict[str, str] | None = None + authentication: MCPAuthentication | None = None # None means "leave unchanged" on update; the controller resolves it to a # concrete IdentityMode before calling the service (see _resolve_identity_mode). identity_mode: IdentityMode | None = None + @field_validator("authentication", "configuration", mode="before") + @classmethod + def empty_to_none(cls, value: object) -> object: + if value == {}: + return None + return value + def _resolve_identity_mode(requested: IdentityMode | None, *, current: IdentityMode) -> IdentityMode: """Resolve the effective MCP identity_mode for a create/update request. @@ -276,16 +311,124 @@ class MCPCallbackQuery(BaseModel): state: str -class ToolOAuthCustomClientResponse(RootModel[dict[str, Any]]): - root: dict[str, Any] +class ApiProviderDetailResponse(ResponseModel): + schema_type: ApiProviderSchemaType + schema_: str = Field(alias="schema") + tools: list[ApiToolBundle] + icon: ToolEmojiIcon + description: str | None = None + credentials: Mapping[str, object] = Field(default_factory=dict) + privacy_policy: str | None = None + custom_disclaimer: str | None = None + labels: list[str] = Field(default_factory=list) -class ToolOAuthClientSchemaResponse(RootModel[list[dict[str, Any]]]): - root: list[dict[str, Any]] +class ApiSchemaParseResponse(ResponseModel): + schema_type: ApiProviderSchemaType + parameters_schema: list[ApiToolBundle] + credentials_schema: list[ProviderConfig] + warning: dict[str, str] -class ToolProviderOpaqueResponse(RootModel[Any]): - root: Any +class ApiProviderRemoteSchemaResponse(ResponseModel): + schema_: str = Field(alias="schema") + + +class ApiToolPreviewResponse(RootModel[ApiToolPreviewResult]): + pass + + +class BuiltinProviderOAuthClientSchemaResponse(ResponseModel): + schema_: list[ProviderConfig] = Field(alias="schema") + is_oauth_custom_client_enabled: bool + is_system_oauth_params_exists: bool + client_params: Mapping[str, object] | None = None + redirect_uri: str + + +class MCPAuthResponse(ResponseModel): + result: Literal["success"] | None = None + authorization_url: str | None = None + + +class ToolApiListResponse(RootModel[list[ToolApiEntity]]): + pass + + +# TODO: This duplicates core.tools.entities.api_entities.ToolProviderApiEntity's +# public response projection. Consolidate the core entity and controller response +# shape when the tool-provider API serialization boundary is cleaned up. +class ToolProviderApiEntityResponse(ResponseModel): + id: str + author: str + name: str + description: I18nObject + icon: str | Mapping[str, str] + icon_dark: str | Mapping[str, str] = "" + label: I18nObject + type: ToolProviderType + team_credentials: Mapping[str, object] = Field(default_factory=dict) + is_team_authorization: bool = False + allow_delete: bool = True + plugin_id: str | None = Field(default="", description="The plugin id of the tool") + plugin_unique_identifier: str | None = Field(default="", description="The unique identifier of the tool") + tools: list[ToolApiEntity] = Field(default_factory=list) + labels: list[str] = Field(default_factory=list) + server_url: str | None = Field(default="", description="The server url of the tool") + updated_at: int = Field(default_factory=lambda: int(datetime.now().timestamp())) + server_identifier: str | None = Field(default="", description="The server identifier of the MCP tool") + masked_headers: dict[str, str] | None = Field(default=None, description="The masked headers of the MCP tool") + original_headers: dict[str, str] | None = Field(default=None, description="The original headers of the MCP tool") + authentication: MCPAuthentication | None = Field(default=None, description="The OAuth config of the MCP tool") + is_dynamic_registration: bool = Field(default=True, description="Whether the MCP tool is dynamically registered") + configuration: MCPConfiguration | None = Field( + default=None, description="The timeout and sse_read_timeout of the MCP tool" + ) + identity_mode: str = Field(default="off", description="Identity-forwarding mechanism: 'off' or 'idp_token'") + workflow_app_id: str | None = Field(default=None, description="The app id of the workflow tool") + + @field_validator("tools", mode="before") + @classmethod + def convert_none_to_empty_list(cls, value: list[ToolApiEntity] | None) -> list[ToolApiEntity]: + return value if value is not None else [] + + +class ToolProviderListResponse(RootModel[list[ToolProviderApiEntityResponse]]): + pass + + +def _dump_tool_provider_payload(payload: Mapping[str, Any]) -> dict[str, Any]: + return ToolProviderApiEntityResponse.model_validate(payload).model_dump(mode="json", exclude_unset=True) + + +def _dump_tool_provider_payload_list(payloads: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + return [_dump_tool_provider_payload(payload) for payload in payloads] + + +class ToolProviderCredentialListResponse(RootModel[list[ToolProviderCredentialApiEntity]]): + pass + + +class ProviderConfigListResponse(RootModel[list[ProviderConfig]]): + pass + + +class ToolLabelListResponse(RootModel[list[ToolLabel]]): + pass + + +class WorkflowToolDetailResponse(ResponseModel): + name: str + label: str + workflow_tool_id: str + workflow_app_id: str + icon: ToolEmojiIcon + description: str + parameters: list[WorkflowToolParameterConfiguration] + output_schema: Mapping[str, object] = Field(default_factory=dict) + tool: ToolApiEntity + synced: bool + privacy_policy: str | None = None register_schema_models( @@ -297,6 +440,7 @@ register_schema_models( WorkflowToolGetQuery, WorkflowToolListQuery, MCPCallbackQuery, + ToolEmojiIcon, BuiltinToolCredentialDeletePayload, BuiltinToolAddPayload, BuiltinToolUpdatePayload, @@ -317,63 +461,94 @@ register_schema_models( ) register_response_schema_models( console_ns, - BinaryFileResponse, - PluginOAuthAuthorizationUrlResponse, - RedirectResponse, SimpleResultResponse, - ToolOAuthClientSchemaResponse, - ToolOAuthCustomClientResponse, - ToolProviderOpaqueResponse, + ApiProviderDetailResponse, + ApiSchemaParseResponse, + ApiProviderRemoteSchemaResponse, + ApiToolPreviewResponse, + BuiltinProviderOAuthClientSchemaResponse, + ToolApiListResponse, + ToolProviderApiEntityResponse, + ToolProviderCredentialInfoApiEntity, + ToolProviderCredentialListResponse, + ToolProviderListResponse, + ProviderConfigListResponse, + PluginOAuthAuthorizationUrlResponse, + ToolLabelListResponse, + MCPAuthResponse, + WorkflowToolDetailResponse, ) @console_ns.route("/workspaces/current/tool-providers") class ToolProviderListApi(Resource): @console_ns.doc(params=query_params_from_model(ToolProviderListQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Tool providers retrieved successfully", console_ns.models[ToolProviderListResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = ToolProviderListQuery.model_validate(raw_args) + query = query_params_from_request(ToolProviderListQuery) - return ToolCommonService.list_tool_providers(user.id, tenant_id, query.type) # type: ignore + return _dump_tool_provider_payload_list( + ToolCommonService.list_tool_providers( + user.id, tenant_id, cast(ToolProviderTypeApiLiteral | None, query.type) + ), + ) @console_ns.route("/workspaces/current/tool-provider/builtin//tools") class ToolBuiltinProviderListToolsApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider tools retrieved successfully", + console_ns.models[ToolApiListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str, provider: str): - return jsonable_encoder( + + return dump_response( + ToolApiListResponse, BuiltinToolManageService.list_builtin_tool_provider_tools( tenant_id, provider, - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//info") class ToolBuiltinProviderInfoApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider info retrieved successfully", + console_ns.models[ToolProviderApiEntityResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str, provider: str): - return jsonable_encoder(BuiltinToolManageService.get_builtin_tool_provider_info(tenant_id, provider)) + + return _dump_tool_provider_payload( + BuiltinToolManageService.get_builtin_tool_provider_info(tenant_id, provider).to_dict() + ) @console_ns.route("/workspaces/current/tool-provider/builtin//delete") class ToolBuiltinProviderDeleteApi(Resource): @console_ns.expect(console_ns.models[BuiltinToolCredentialDeletePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider credential deleted successfully", + console_ns.models[SimpleResultResponse.__name__], + ) @setup_required @login_required @is_admin_or_owner_required @@ -381,19 +556,27 @@ class ToolBuiltinProviderDeleteApi(Resource): @account_initialization_required @with_current_tenant_id def post(self, tenant_id: str, provider: str): + payload = BuiltinToolCredentialDeletePayload.model_validate(console_ns.payload or {}) - return BuiltinToolManageService.delete_builtin_tool_provider( - tenant_id, - provider, - payload.credential_id, + return dump_response( + SimpleResultResponse, + BuiltinToolManageService.delete_builtin_tool_provider( + tenant_id, + provider, + payload.credential_id, + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//add") class ToolBuiltinProviderAddApi(Resource): @console_ns.expect(console_ns.models[BuiltinToolAddPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider added successfully", + console_ns.models[SimpleResultResponse.__name__], + ) @setup_required @login_required @account_initialization_required @@ -402,21 +585,28 @@ class ToolBuiltinProviderAddApi(Resource): def post(self, tenant_id: str, user: Account, provider: str): payload = BuiltinToolAddPayload.model_validate(console_ns.payload or {}) - return BuiltinToolManageService.add_builtin_tool_provider( - user_id=user.id, - tenant_id=tenant_id, - provider=provider, - credentials=payload.credentials, - name=payload.name, - api_type=CredentialType.of(payload.type), - visibility=payload.visibility, + return dump_response( + SimpleResultResponse, + BuiltinToolManageService.add_builtin_tool_provider( + user_id=user.id, + tenant_id=tenant_id, + provider=provider, + credentials=payload.credentials, + name=payload.name, + api_type=CredentialType.of(payload.type), + visibility=payload.visibility, + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//update") class ToolBuiltinProviderUpdateApi(Resource): @console_ns.expect(console_ns.models[BuiltinToolUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider updated successfully", + console_ns.models[SimpleResultResponse.__name__], + ) @setup_required @login_required @is_admin_or_owner_required @@ -435,13 +625,17 @@ class ToolBuiltinProviderUpdateApi(Resource): credentials=payload.credentials, name=payload.name or "", ) - return result + return dump_response(SimpleResultResponse, result) @console_ns.route("/workspaces/current/tool-provider/builtin//credentials") class ToolBuiltinProviderGetCredentialsApi(Resource): @console_ns.doc(params=query_params_from_model(BuiltinCredentialListQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider credentials retrieved successfully", + console_ns.models[ToolProviderCredentialListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @@ -455,30 +649,33 @@ class ToolBuiltinProviderGetCredentialsApi(Resource): list_fields=("include_credential_ids",), ) - return jsonable_encoder( + return dump_response( + ToolProviderCredentialListResponse, BuiltinToolManageService.get_builtin_tool_provider_credentials( tenant_id=tenant_id, provider_name=provider, + session=db.session(), user=user, include_credential_ids=query.include_credential_ids or None, - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//icon") class ToolBuiltinProviderIconApi(Resource): - @console_ns.response(200, "Success", console_ns.models[BinaryFileResponse.__name__]) + @console_ns.response(200, "Builtin provider icon") @setup_required def get(self, provider: str): icon_bytes, mimetype = BuiltinToolManageService.get_builtin_tool_provider_icon(provider) icon_cache_max_age = dify_config.TOOL_ICON_CACHE_MAX_AGE + # response-contract:ignore binary send_file response return send_file(io.BytesIO(icon_bytes), mimetype=mimetype, max_age=icon_cache_max_age) @console_ns.route("/workspaces/current/tool-provider/api/add") class ToolApiProviderAddApi(Resource): @console_ns.expect(console_ns.models[ApiToolProviderAddPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "API provider added successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -489,66 +686,77 @@ class ToolApiProviderAddApi(Resource): def post(self, tenant_id: str, user: Account): payload = ApiToolProviderAddPayload.model_validate(console_ns.payload or {}) - return ApiToolManageService.create_api_tool_provider( - user.id, - tenant_id, - payload.provider, - payload.icon, - payload.credentials, - payload.schema_type, - payload.schema_, - payload.privacy_policy or "", - payload.custom_disclaimer or "", - payload.labels or [], + return dump_response( + SimpleResultResponse, + ApiToolManageService.create_api_tool_provider( + user.id, + tenant_id, + payload.provider, + payload.icon.model_dump(mode="json"), + payload.credentials, + payload.schema_type, + payload.schema_, + payload.privacy_policy or "", + payload.custom_disclaimer or "", + payload.labels or [], + ), ) @console_ns.route("/workspaces/current/tool-provider/api/remote") class ToolApiProviderGetRemoteSchemaApi(Resource): @console_ns.doc(params=query_params_from_model(UrlQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Remote API provider schema retrieved successfully", + console_ns.models[ApiProviderRemoteSchemaResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = UrlQuery.model_validate(raw_args) + query = query_params_from_request(UrlQuery) - return ApiToolManageService.get_api_tool_provider_remote_schema( - user.id, - tenant_id, - str(query.url), + return dump_response( + ApiProviderRemoteSchemaResponse, + ApiToolManageService.get_api_tool_provider_remote_schema( + user.id, + tenant_id, + str(query.url), + ), ) @console_ns.route("/workspaces/current/tool-provider/api/tools") class ToolApiProviderListToolsApi(Resource): @console_ns.doc(params=query_params_from_model(ProviderQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "API provider tools retrieved successfully", console_ns.models[ToolApiListResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = ProviderQuery.model_validate(raw_args) + query = query_params_from_request(ProviderQuery) - return jsonable_encoder( + return dump_response( + ToolApiListResponse, ApiToolManageService.list_api_tool_provider_tools( user.id, tenant_id, query.provider, - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/api/update") class ToolApiProviderUpdateApi(Resource): @console_ns.expect(console_ns.models[ApiToolProviderUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "API provider updated successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -559,25 +767,28 @@ class ToolApiProviderUpdateApi(Resource): def post(self, tenant_id: str, user: Account): payload = ApiToolProviderUpdatePayload.model_validate(console_ns.payload or {}) - return ApiToolManageService.update_api_tool_provider( - user.id, - tenant_id, - payload.provider, - payload.original_provider, - payload.icon, - payload.credentials, - payload.schema_type, - payload.schema_, - payload.privacy_policy, - payload.custom_disclaimer, - payload.labels or [], + return dump_response( + SimpleResultResponse, + ApiToolManageService.update_api_tool_provider( + user.id, + tenant_id, + payload.provider, + payload.original_provider, + payload.icon.model_dump(mode="json"), + payload.credentials, + payload.schema_type, + payload.schema_, + payload.privacy_policy, + payload.custom_disclaimer, + payload.labels or [], + ), ) @console_ns.route("/workspaces/current/tool-provider/api/delete") class ToolApiProviderDeleteApi(Resource): @console_ns.expect(console_ns.models[ApiToolProviderDeletePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "API provider deleted successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -588,88 +799,106 @@ class ToolApiProviderDeleteApi(Resource): def post(self, tenant_id: str, user: Account): payload = ApiToolProviderDeletePayload.model_validate(console_ns.payload or {}) - return ApiToolManageService.delete_api_tool_provider( - user.id, - tenant_id, - payload.provider, + return dump_response( + SimpleResultResponse, + ApiToolManageService.delete_api_tool_provider( + user.id, + tenant_id, + payload.provider, + ), ) @console_ns.route("/workspaces/current/tool-provider/api/get") class ToolApiProviderGetApi(Resource): @console_ns.doc(params=query_params_from_model(ProviderQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "API provider retrieved successfully", console_ns.models[ApiProviderDetailResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = ProviderQuery.model_validate(raw_args) + query = query_params_from_request(ProviderQuery) - return ApiToolManageService.get_api_tool_provider( - user.id, - tenant_id, - query.provider, + return dump_response( + ApiProviderDetailResponse, + ApiToolManageService.get_api_tool_provider( + user.id, + tenant_id, + query.provider, + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//credential/schema/") class ToolBuiltinProviderCredentialsSchemaApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider credential schema retrieved successfully", + console_ns.models[ProviderConfigListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str, provider, credential_type): - return jsonable_encoder( + + return dump_response( + ProviderConfigListResponse, BuiltinToolManageService.list_builtin_provider_credentials_schema( provider, CredentialType.of(credential_type), tenant_id - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/api/schema") class ToolApiProviderSchemaApi(Resource): @console_ns.expect(console_ns.models[ApiToolSchemaPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "API schema parsed successfully", console_ns.models[ApiSchemaParseResponse.__name__]) @setup_required @login_required @account_initialization_required def post(self): payload = ApiToolSchemaPayload.model_validate(console_ns.payload or {}) - return ApiToolManageService.parser_api_schema( - schema=payload.schema_, - ) + return dump_response(ApiSchemaParseResponse, ApiToolManageService.parser_api_schema(schema=payload.schema_)) @console_ns.route("/workspaces/current/tool-provider/api/test/pre") class ToolApiProviderPreviousTestApi(Resource): @console_ns.expect(console_ns.models[ApiToolTestPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "API tool test preview completed successfully", + console_ns.models[ApiToolPreviewResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def post(self, current_tenant_id: str): payload = ApiToolTestPayload.model_validate(console_ns.payload or {}) - return ApiToolManageService.test_api_tool_preview( - current_tenant_id, - payload.provider_name or "", - payload.tool_name, - payload.credentials, - payload.parameters, - payload.schema_type, - payload.schema_, + return dump_response( + ApiToolPreviewResponse, + ApiToolManageService.test_api_tool_preview( + current_tenant_id, + payload.provider_name or "", + payload.tool_name, + payload.credentials, + payload.parameters, + payload.schema_type, + payload.schema_, + ), ) @console_ns.route("/workspaces/current/tool-provider/workflow/create") class ToolWorkflowProviderCreateApi(Resource): @console_ns.expect(console_ns.models[WorkflowToolCreatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "Workflow tool created successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -680,24 +909,27 @@ class ToolWorkflowProviderCreateApi(Resource): def post(self, tenant_id: str, user: Account): payload = WorkflowToolCreatePayload.model_validate(console_ns.payload or {}) - return WorkflowToolManageService.create_workflow_tool( - user_id=user.id, - tenant_id=tenant_id, - workflow_app_id=payload.workflow_app_id, - name=payload.name, - label=payload.label, - icon=payload.icon, - description=payload.description, - parameters=payload.parameters, - privacy_policy=payload.privacy_policy or "", - labels=payload.labels or [], + return dump_response( + SimpleResultResponse, + WorkflowToolManageService.create_workflow_tool( + user_id=user.id, + tenant_id=tenant_id, + workflow_app_id=payload.workflow_app_id, + name=payload.name, + label=payload.label, + icon=payload.icon.model_dump(mode="json"), + description=payload.description, + parameters=payload.parameters, + privacy_policy=payload.privacy_policy or "", + labels=payload.labels or [], + ), ) @console_ns.route("/workspaces/current/tool-provider/workflow/update") class ToolWorkflowProviderUpdateApi(Resource): @console_ns.expect(console_ns.models[WorkflowToolUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "Workflow tool updated successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -708,24 +940,27 @@ class ToolWorkflowProviderUpdateApi(Resource): def post(self, tenant_id: str, user: Account): payload = WorkflowToolUpdatePayload.model_validate(console_ns.payload or {}) - return WorkflowToolManageService.update_workflow_tool( - user.id, - tenant_id, - payload.workflow_tool_id, - payload.name, - payload.label, - payload.icon, - payload.description, - payload.parameters, - payload.privacy_policy or "", - payload.labels or [], + return dump_response( + SimpleResultResponse, + WorkflowToolManageService.update_workflow_tool( + user.id, + tenant_id, + payload.workflow_tool_id, + payload.name, + payload.label, + payload.icon.model_dump(mode="json"), + payload.description, + payload.parameters, + payload.privacy_policy or "", + payload.labels or [], + ), ) @console_ns.route("/workspaces/current/tool-provider/workflow/delete") class ToolWorkflowProviderDeleteApi(Resource): @console_ns.expect(console_ns.models[WorkflowToolDeletePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "Workflow tool deleted successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -736,25 +971,29 @@ class ToolWorkflowProviderDeleteApi(Resource): def post(self, tenant_id: str, user: Account): payload = WorkflowToolDeletePayload.model_validate(console_ns.payload or {}) - return WorkflowToolManageService.delete_workflow_tool( - user.id, - tenant_id, - payload.workflow_tool_id, + return dump_response( + SimpleResultResponse, + WorkflowToolManageService.delete_workflow_tool( + user.id, + tenant_id, + payload.workflow_tool_id, + ), ) @console_ns.route("/workspaces/current/tool-provider/workflow/get") class ToolWorkflowProviderGetApi(Resource): @console_ns.doc(params=query_params_from_model(WorkflowToolGetQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Workflow tool retrieved successfully", console_ns.models[WorkflowToolDetailResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = WorkflowToolGetQuery.model_validate(raw_args) + query = query_params_from_request(WorkflowToolGetQuery) if query.workflow_tool_id: tool = WorkflowToolManageService.get_workflow_tool_by_tool_id( @@ -771,105 +1010,112 @@ class ToolWorkflowProviderGetApi(Resource): else: raise ValueError("incorrect workflow_tool_id or workflow_app_id") - return jsonable_encoder(tool) + return dump_response(WorkflowToolDetailResponse, tool) @console_ns.route("/workspaces/current/tool-provider/workflow/tools") class ToolWorkflowProviderListToolApi(Resource): @console_ns.doc(params=query_params_from_model(WorkflowToolListQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Workflow provider tools retrieved successfully", console_ns.models[ToolApiListResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - raw_args = request.args.to_dict() - query = WorkflowToolListQuery.model_validate(raw_args) + query = query_params_from_request(WorkflowToolListQuery) - return jsonable_encoder( + return dump_response( + ToolApiListResponse, WorkflowToolManageService.list_single_workflow_tools( user.id, tenant_id, query.workflow_tool_id, - ) + ), ) @console_ns.route("/workspaces/current/tools/builtin") class ToolBuiltinListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Builtin tools retrieved successfully", console_ns.models[ToolProviderListResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - return jsonable_encoder( + return _dump_tool_provider_payload_list( [ provider.to_dict() for provider in BuiltinToolManageService.list_builtin_tools( user.id, tenant_id, ) - ] + ], ) @console_ns.route("/workspaces/current/tools/api") class ToolApiListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "API tools retrieved successfully", console_ns.models[ToolProviderListResponse.__name__]) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str): - return jsonable_encoder( + + return _dump_tool_provider_payload_list( [ provider.to_dict() for provider in ApiToolManageService.list_api_tools( tenant_id, ) - ] + ], ) @console_ns.route("/workspaces/current/tools/workflow") class ToolWorkflowListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Workflow tools retrieved successfully", console_ns.models[ToolProviderListResponse.__name__] + ) @setup_required @login_required @account_initialization_required @with_current_user @with_current_tenant_id def get(self, tenant_id: str, user: Account): - return jsonable_encoder( + return _dump_tool_provider_payload_list( [ provider.to_dict() for provider in WorkflowToolManageService.list_tenant_workflow_tools( user.id, tenant_id, ) - ] + ], ) @console_ns.route("/workspaces/current/tool-labels") class ToolLabelsApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "Tool labels retrieved successfully", console_ns.models[ToolLabelListResponse.__name__]) @setup_required @login_required @account_initialization_required @enterprise_license_required def get(self): - return jsonable_encoder(ToolLabelsService.list_tool_labels()) + return dump_response(ToolLabelListResponse, ToolLabelsService.list_tool_labels()) @console_ns.route("/oauth/plugin//tool/authorization-url") class ToolPluginOAuthApi(Resource): @console_ns.response( 200, - "Authorization URL retrieved successfully", + "Tool OAuth authorization URL generated successfully", console_ns.models[PluginOAuthAuthorizationUrlResponse.__name__], ) @setup_required @@ -901,7 +1147,8 @@ class ToolPluginOAuthApi(Resource): redirect_uri=redirect_uri, system_credentials=oauth_client_params, ) - response = make_response(jsonable_encoder(authorization_url_response)) + # response-contract:ignore cookie-bearing Flask response + response = make_response(dump_response(PluginOAuthAuthorizationUrlResponse, authorization_url_response)) response.set_cookie( "context_id", context_id, @@ -914,11 +1161,7 @@ class ToolPluginOAuthApi(Resource): @console_ns.route("/oauth/plugin//tool/callback") class ToolOAuthCallback(Resource): - @console_ns.response( - 302, - "Redirect to console OAuth callback page", - console_ns.models[RedirectResponse.__name__], - ) + @console_ns.response(302, "Redirect to OAuth callback page") @setup_required def get(self, provider: str): context_id = request.cookies.get("context_id") @@ -967,13 +1210,14 @@ class ToolOAuthCallback(Resource): api_type=CredentialType.OAUTH2, visibility="only_me", ) + # response-contract:ignore redirect response return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback") @console_ns.route("/workspaces/current/tool-provider/builtin//default-credential") class ToolBuiltinProviderSetDefaultApi(Resource): @console_ns.expect(console_ns.models[BuiltinProviderDefaultCredentialPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "Default credential set successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @is_admin_or_owner_required @@ -982,15 +1226,20 @@ class ToolBuiltinProviderSetDefaultApi(Resource): @with_current_tenant_id def post(self, current_tenant_id: str, provider: str): payload = BuiltinProviderDefaultCredentialPayload.model_validate(console_ns.payload or {}) - return BuiltinToolManageService.set_default_provider( - tenant_id=current_tenant_id, provider=provider, id=payload.id + return dump_response( + SimpleResultResponse, + BuiltinToolManageService.set_default_provider( + tenant_id=current_tenant_id, provider=provider, id=payload.id + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//oauth/custom-client") class ToolOAuthCustomClient(Resource): @console_ns.expect(console_ns.models[ToolOAuthCustomClientPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response( + 200, "Custom OAuth client saved successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @is_admin_or_owner_required @@ -1000,55 +1249,71 @@ class ToolOAuthCustomClient(Resource): def post(self, tenant_id: str, provider: str): payload = ToolOAuthCustomClientPayload.model_validate(console_ns.payload or {}) - return BuiltinToolManageService.save_custom_oauth_client_params( - tenant_id=tenant_id, - provider=provider, - client_params=payload.client_params or {}, - enable_oauth_custom_client=payload.enable_oauth_custom_client - if payload.enable_oauth_custom_client is not None - else True, + return dump_response( + SimpleResultResponse, + BuiltinToolManageService.save_custom_oauth_client_params( + tenant_id=tenant_id, + provider=provider, + client_params=payload.client_params or {}, + enable_oauth_custom_client=payload.enable_oauth_custom_client + if payload.enable_oauth_custom_client is not None + else True, + ), ) + @console_ns.response( + 200, + "Custom OAuth client retrieved successfully", + ) @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[ToolOAuthCustomClientResponse.__name__]) @with_current_tenant_id def get(self, current_tenant_id: str, provider: str): - return jsonable_encoder( - BuiltinToolManageService.get_custom_oauth_client_params(tenant_id=current_tenant_id, provider=provider) - ) + return BuiltinToolManageService.get_custom_oauth_client_params(tenant_id=current_tenant_id, provider=provider) + @console_ns.response( + 200, "Custom OAuth client deleted successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @with_current_tenant_id def delete(self, current_tenant_id: str, provider: str): - return jsonable_encoder( - BuiltinToolManageService.delete_custom_oauth_client_params(tenant_id=current_tenant_id, provider=provider) + return dump_response( + SimpleResultResponse, + BuiltinToolManageService.delete_custom_oauth_client_params(tenant_id=current_tenant_id, provider=provider), ) @console_ns.route("/workspaces/current/tool-provider/builtin//oauth/client-schema") class ToolBuiltinProviderGetOauthClientSchemaApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolOAuthClientSchemaResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider OAuth client schema retrieved successfully", + console_ns.models[BuiltinProviderOAuthClientSchemaResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str, provider: str): - return jsonable_encoder( + return dump_response( + BuiltinProviderOAuthClientSchemaResponse, BuiltinToolManageService.get_builtin_tool_provider_oauth_client_schema( tenant_id=current_tenant_id, provider_name=provider - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/builtin//credential/info") class ToolBuiltinProviderGetCredentialInfoApi(Resource): @console_ns.doc(params=query_params_from_model(BuiltinCredentialListQuery)) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, + "Builtin provider credential info retrieved successfully", + console_ns.models[ToolProviderCredentialInfoApiEntity.__name__], + ) @setup_required @login_required @account_initialization_required @@ -1060,20 +1325,24 @@ class ToolBuiltinProviderGetCredentialInfoApi(Resource): list_fields=("include_credential_ids",), ) - return jsonable_encoder( + return dump_response( + ToolProviderCredentialInfoApiEntity, BuiltinToolManageService.get_builtin_tool_provider_credential_info( tenant_id=tenant_id, provider=provider, + session=db.session(), user=user, include_credential_ids=query.include_credential_ids or None, - ) + ), ) @console_ns.route("/workspaces/current/tool-provider/mcp") class ToolProviderMCPApi(Resource): @console_ns.expect(console_ns.models[MCPProviderCreatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "MCP provider created successfully", console_ns.models[ToolProviderApiEntityResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -1083,9 +1352,8 @@ class ToolProviderMCPApi(Resource): def post(self, tenant_id: str, user: Account): payload = MCPProviderCreatePayload.model_validate(console_ns.payload or {}) - # Parse and validate models - configuration = MCPConfiguration.model_validate(payload.configuration or {}) - authentication = MCPAuthentication.model_validate(payload.authentication) if payload.authentication else None + configuration = payload.configuration or MCPConfiguration() + authentication = payload.authentication # 1) Create provider in a short transaction (no network I/O inside) with session_factory.create_session() as session, session.begin(): @@ -1126,10 +1394,10 @@ class ToolProviderMCPApi(Resource): # Best-effort: if initial fetch fails (e.g., auth required), return created provider as-is logger.warning("Failed to fetch MCP tools after creation", exc_info=True) - return jsonable_encoder(result) + return _dump_tool_provider_payload(result.to_dict()) @console_ns.expect(console_ns.models[MCPProviderUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response(200, "MCP provider updated successfully", console_ns.models[SimpleResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -1137,8 +1405,8 @@ class ToolProviderMCPApi(Resource): @with_current_tenant_id def put(self, current_tenant_id: str): payload = MCPProviderUpdatePayload.model_validate(console_ns.payload or {}) - configuration = MCPConfiguration.model_validate(payload.configuration or {}) - authentication = MCPAuthentication.model_validate(payload.authentication) if payload.authentication else None + configuration = payload.configuration or MCPConfiguration() + authentication = payload.authentication # Step 1: Get provider data for URL validation (short-lived session, no network I/O) validation_data = None @@ -1180,7 +1448,7 @@ class ToolProviderMCPApi(Resource): identity_mode=identity_mode, ) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.expect(console_ns.models[MCPProviderDeletePayload.__name__]) @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @@ -1196,13 +1464,13 @@ class ToolProviderMCPApi(Resource): service = MCPToolManageService(session=session) service.delete_provider(tenant_id=current_tenant_id, provider_id=payload.provider_id) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") @console_ns.route("/workspaces/current/tool-provider/mcp/auth") class ToolMCPAuthApi(Resource): @console_ns.expect(console_ns.models[MCPAuthPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "MCP provider authorized successfully", console_ns.models[MCPAuthResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -1241,7 +1509,7 @@ class ToolMCPAuthApi(Resource): credentials=provider_entity.credentials, authed=True, ) - return {"result": "success"} + return MCPAuthResponse(result="success").model_dump(mode="json") except MCPAuthError as e: try: # Pass the extracted OAuth metadata hints to auth() @@ -1254,7 +1522,7 @@ class ToolMCPAuthApi(Resource): with sessionmaker(db.engine).begin() as session: service = MCPToolManageService(session=session) response = service.execute_auth_actions(auth_result) - return response + return dump_response(MCPAuthResponse, response) except MCPRefreshTokenError as e: with sessionmaker(db.engine).begin() as session: service = MCPToolManageService(session=session) @@ -1277,7 +1545,9 @@ class ToolMCPAuthApi(Resource): @console_ns.route("/workspaces/current/tool-provider/mcp/tools/") class ToolMCPDetailApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "MCP provider retrieved successfully", console_ns.models[ToolProviderApiEntityResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -1286,28 +1556,33 @@ class ToolMCPDetailApi(Resource): with sessionmaker(db.engine).begin() as session: service = MCPToolManageService(session=session) provider = service.get_provider(provider_id=provider_id, tenant_id=tenant_id) - return jsonable_encoder(ToolTransformService.mcp_provider_to_user_provider(provider, for_list=True)) + return _dump_tool_provider_payload( + ToolTransformService.mcp_provider_to_user_provider(provider, for_list=True).to_dict() + ) @console_ns.route("/workspaces/current/tools/mcp") class ToolMCPListAllApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response(200, "MCP tools retrieved successfully", console_ns.models[ToolProviderListResponse.__name__]) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str): + with sessionmaker(db.engine).begin() as session: service = MCPToolManageService(session=session) # Skip sensitive data decryption for list view to improve performance tools = service.list_providers(tenant_id=tenant_id, include_sensitive=False) - return [tool.to_dict() for tool in tools] + return _dump_tool_provider_payload_list([tool.to_dict() for tool in tools]) @console_ns.route("/workspaces/current/tool-provider/mcp/update/") class ToolMCPUpdateApi(Resource): - @console_ns.response(200, "Success", console_ns.models[ToolProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "MCP provider tools refreshed successfully", console_ns.models[ToolProviderApiEntityResponse.__name__] + ) @setup_required @login_required @account_initialization_required @@ -1320,20 +1595,15 @@ class ToolMCPUpdateApi(Resource): tenant_id=tenant_id, provider_id=provider_id, ) - return jsonable_encoder(tools) + return _dump_tool_provider_payload(tools.to_dict()) @console_ns.route("/mcp/oauth/callback") class ToolMCPCallbackApi(Resource): @console_ns.doc(params=query_params_from_model(MCPCallbackQuery)) - @console_ns.response( - 302, - "Redirect to console OAuth callback page", - console_ns.models[RedirectResponse.__name__], - ) + @console_ns.response(302, "Redirect to OAuth callback page") def get(self): - raw_args = request.args.to_dict() - query = MCPCallbackQuery.model_validate(raw_args) + query = query_params_from_request(MCPCallbackQuery) state_key = query.state authorization_code = query.code @@ -1347,4 +1617,5 @@ class ToolMCPCallbackApi(Resource): state_data.provider_id, state_data.tenant_id, tokens.model_dump(), OAuthDataType.TOKENS ) + # response-contract:ignore redirect response return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback") diff --git a/api/controllers/console/workspace/trigger_providers.py b/api/controllers/console/workspace/trigger_providers.py index c30a4453901..aff800a5bc9 100644 --- a/api/controllers/console/workspace/trigger_providers.py +++ b/api/controllers/console/workspace/trigger_providers.py @@ -1,20 +1,19 @@ import logging -from typing import Any, Literal +from typing import Any from flask import make_response, redirect, request from flask_restx import Resource -from pydantic import BaseModel, Field, RootModel, model_validator +from pydantic import BaseModel, RootModel, model_validator from sqlalchemy.orm import sessionmaker from werkzeug.exceptions import BadRequest, Forbidden from configs import dify_config from controllers.common.errors import NotFoundError -from controllers.common.fields import BinaryFileResponse, RedirectResponse, SimpleResultResponse +from controllers.common.fields import SimpleResultResponse from controllers.common.schema import register_response_schema_models, register_schema_models -from core.entities.parameter_entities import AppSelectorScope, ModelSelectorScope, ToolSelectorScope +from core.entities.provider_entities import ProviderConfig from core.plugin.entities.plugin_daemon import CredentialType from core.plugin.impl.oauth import OAuthHandler -from core.tools.entities.common_entities import I18nObject from core.trigger.entities.api_entities import ( SubscriptionBuilderApiEntity, TriggerProviderApiEntity, @@ -24,7 +23,7 @@ from core.trigger.entities.entities import RequestLog, SubscriptionBuilderUpdate from core.trigger.trigger_manager import TriggerManager from extensions.ext_database import db from fields.base import ResponseModel -from graphon.model_runtime.utils.encoders import jsonable_encoder +from libs.helper import dump_response from libs.login import login_required from models.account import Account from models.provider_ids import TriggerProviderID @@ -59,9 +58,9 @@ class TriggerSubscriptionBuilderVerifyPayload(BaseModel): class TriggerSubscriptionBuilderUpdatePayload(BaseModel): name: str | None = None - parameters: dict[str, Any] | None = Field(default=None) - properties: dict[str, Any] | None = Field(default=None) - credentials: dict[str, Any] | None = Field(default=None) + parameters: dict[str, Any] | None = None + properties: dict[str, Any] | None = None + credentials: dict[str, Any] | None = None @model_validator(mode="after") def check_at_least_one_field(self): @@ -71,70 +70,23 @@ class TriggerSubscriptionBuilderUpdatePayload(BaseModel): class TriggerOAuthClientPayload(BaseModel): - client_params: dict[str, Any] | None = Field(default=None) + client_params: dict[str, Any] | None = None enabled: bool | None = None -class TriggerOAuthAuthorizeResponse(BaseModel): - authorization_url: str - subscription_builder_id: str - subscription_builder: SubscriptionBuilderApiEntity - - -class TriggerProviderConfigOptionResponse(BaseModel): - value: str = Field(..., description="The value of the option") - label: I18nObject = Field(..., description="The label of the option") - - -class TriggerProviderConfigResponse(BaseModel): - type: Literal[ - "secret-input", - "text-input", - "select", - "boolean", - "app-selector", - "model-selector", - "array[tools]", - ] = Field(..., description="The type of the credentials") - name: str = Field(..., description="The name of the credentials") - scope: AppSelectorScope | ModelSelectorScope | ToolSelectorScope | None = None - required: bool = False - default: int | str | float | bool | None = None - options: list[TriggerProviderConfigOptionResponse] | None = None - multiple: bool = False - label: I18nObject | None = None - help: I18nObject | None = None - url: str | None = None - placeholder: I18nObject | None = None - - -class TriggerOAuthClientResponse(BaseModel): - configured: bool - system_configured: bool - custom_configured: bool - oauth_client_schema: list[TriggerProviderConfigResponse] - custom_enabled: bool - redirect_uri: str - params: dict[str, Any] - - -class TriggerProviderOpaqueResponse(RootModel[Any]): - root: Any - - class TriggerProviderListResponse(RootModel[list[TriggerProviderApiEntity]]): - root: list[TriggerProviderApiEntity] + pass -class TriggerSubscriptionListResponse(RootModel[list[TriggerProviderSubscriptionApiEntity]]): - root: list[TriggerProviderSubscriptionApiEntity] +class TriggerProviderSubscriptionListResponse(RootModel[list[TriggerProviderSubscriptionApiEntity]]): + pass class TriggerSubscriptionBuilderCreateResponse(ResponseModel): subscription_builder: SubscriptionBuilderApiEntity -class TriggerSubscriptionBuilderVerifyResponse(ResponseModel): +class TriggerVerificationResponse(ResponseModel): verified: bool @@ -142,6 +94,26 @@ class TriggerSubscriptionBuilderLogsResponse(ResponseModel): logs: list[RequestLog] +class TriggerOAuthAuthorizeResponse(ResponseModel): + authorization_url: str + subscription_builder_id: str + subscription_builder: SubscriptionBuilderApiEntity + + +class TriggerOAuthClientResponse(ResponseModel): + configured: bool + system_configured: bool + custom_configured: bool + oauth_client_schema: list[ProviderConfig] + custom_enabled: bool + redirect_uri: str + params: dict[str, Any] + + +class TriggerProviderErrorResponse(ResponseModel): + error: str + + register_schema_models( console_ns, TriggerSubscriptionBuilderCreatePayload, @@ -151,27 +123,24 @@ register_schema_models( ) register_response_schema_models( console_ns, - BinaryFileResponse, - RedirectResponse, SimpleResultResponse, TriggerOAuthAuthorizeResponse, TriggerOAuthClientResponse, - TriggerProviderOpaqueResponse, TriggerProviderApiEntity, + TriggerProviderErrorResponse, TriggerProviderListResponse, - TriggerProviderSubscriptionApiEntity, - TriggerSubscriptionListResponse, - SubscriptionBuilderApiEntity, + TriggerProviderSubscriptionListResponse, TriggerSubscriptionBuilderCreateResponse, - TriggerSubscriptionBuilderVerifyResponse, - RequestLog, TriggerSubscriptionBuilderLogsResponse, + SubscriptionBuilderApiEntity, + TriggerVerificationResponse, ) @console_ns.route("/workspaces/current/trigger-provider//icon") class TriggerProviderIconApi(Resource): - @console_ns.response(200, "Success", console_ns.models[BinaryFileResponse.__name__]) + # response-contract:ignore binary trigger provider icon + @console_ns.response(200, "Trigger provider icon") @setup_required @login_required @account_initialization_required @@ -182,31 +151,45 @@ class TriggerProviderIconApi(Resource): @console_ns.route("/workspaces/current/triggers") class TriggerProviderListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TriggerProviderListResponse.__name__]) + @console_ns.response( + 200, + "Trigger providers retrieved successfully", + console_ns.models[TriggerProviderListResponse.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str): """List all trigger providers for the current tenant""" - return jsonable_encoder(TriggerProviderService.list_trigger_providers(tenant_id)) + return dump_response(TriggerProviderListResponse, TriggerProviderService.list_trigger_providers(tenant_id)) @console_ns.route("/workspaces/current/trigger-provider//info") class TriggerProviderInfoApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TriggerProviderApiEntity.__name__]) + @console_ns.response( + 200, + "Trigger provider retrieved successfully", + console_ns.models[TriggerProviderApiEntity.__name__], + ) @setup_required @login_required @account_initialization_required @with_current_tenant_id def get(self, tenant_id: str, provider: str): """Get info for a trigger provider""" - return jsonable_encoder(TriggerProviderService.get_trigger_provider(tenant_id, TriggerProviderID(provider))) + provider_entity = TriggerProviderService.get_trigger_provider(tenant_id, TriggerProviderID(provider)) + return provider_entity.model_dump(mode="json") @console_ns.route("/workspaces/current/trigger-provider//subscriptions/list") class TriggerSubscriptionListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TriggerSubscriptionListResponse.__name__]) + @console_ns.response( + 200, + "Trigger subscriptions retrieved successfully", + console_ns.models[TriggerProviderSubscriptionListResponse.__name__], + ) + @console_ns.response(404, "Trigger provider not found", console_ns.models[TriggerProviderErrorResponse.__name__]) @setup_required @login_required @edit_permission_required @@ -216,16 +199,18 @@ class TriggerSubscriptionListApi(Resource): @with_current_tenant_id def get(self, tenant_id: str, user: Account, provider: str): """List all trigger subscriptions for the current tenant's provider""" + try: - return jsonable_encoder( + return dump_response( + TriggerProviderSubscriptionListResponse, TriggerProviderService.list_trigger_provider_subscriptions( tenant_id=tenant_id, provider_id=TriggerProviderID(provider), user=user, - ) + ), ) except ValueError as e: - return jsonable_encoder({"error": str(e)}), 404 + return TriggerProviderErrorResponse(error=str(e)).model_dump(mode="json"), 404 except Exception as e: logger.exception("Error listing trigger providers", exc_info=e) raise @@ -236,7 +221,11 @@ class TriggerSubscriptionListApi(Resource): ) class TriggerSubscriptionBuilderCreateApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderCreatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[TriggerSubscriptionBuilderCreateResponse.__name__]) + @console_ns.response( + 200, + "Trigger subscription builder created successfully", + console_ns.models[TriggerSubscriptionBuilderCreateResponse.__name__], + ) @setup_required @login_required @edit_permission_required @@ -246,6 +235,7 @@ class TriggerSubscriptionBuilderCreateApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, user: Account, provider: str): """Add a new subscription instance for a trigger provider""" + payload = TriggerSubscriptionBuilderCreatePayload.model_validate(console_ns.payload or {}) try: @@ -256,7 +246,9 @@ class TriggerSubscriptionBuilderCreateApi(Resource): provider_id=TriggerProviderID(provider), credential_type=credential_type, ) - return jsonable_encoder({"subscription_builder": subscription_builder}) + return TriggerSubscriptionBuilderCreateResponse(subscription_builder=subscription_builder).model_dump( + mode="json" + ) except Exception as e: logger.exception("Error adding provider credential", exc_info=e) raise @@ -266,7 +258,11 @@ class TriggerSubscriptionBuilderCreateApi(Resource): "/workspaces/current/trigger-provider//subscriptions/builder/", ) class TriggerSubscriptionBuilderGetApi(Resource): - @console_ns.response(200, "Success", console_ns.models[SubscriptionBuilderApiEntity.__name__]) + @console_ns.response( + 200, + "Trigger subscription builder retrieved successfully", + console_ns.models[SubscriptionBuilderApiEntity.__name__], + ) @setup_required @login_required @edit_permission_required @@ -274,9 +270,8 @@ class TriggerSubscriptionBuilderGetApi(Resource): @account_initialization_required def get(self, provider: str, subscription_builder_id: str): """Get a subscription instance for a trigger provider""" - return jsonable_encoder( - TriggerSubscriptionBuilderService.get_subscription_builder_by_id(subscription_builder_id) - ) + subscription_builder = TriggerSubscriptionBuilderService.get_subscription_builder_by_id(subscription_builder_id) + return subscription_builder.model_dump(mode="json") @console_ns.route( @@ -284,7 +279,11 @@ class TriggerSubscriptionBuilderGetApi(Resource): ) class TriggerSubscriptionBuilderVerifyApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderVerifyPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[TriggerSubscriptionBuilderVerifyResponse.__name__]) + @console_ns.response( + 200, + "Trigger subscription builder verified successfully", + console_ns.models[TriggerVerificationResponse.__name__], + ) @setup_required @login_required @edit_permission_required @@ -294,11 +293,12 @@ class TriggerSubscriptionBuilderVerifyApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, user: Account, provider: str, subscription_builder_id: str): """Verify and update a subscription instance for a trigger provider""" + payload = TriggerSubscriptionBuilderVerifyPayload.model_validate(console_ns.payload or {}) try: # Use atomic update_and_verify to prevent race conditions - return TriggerSubscriptionBuilderService.update_and_verify_builder( + result = TriggerSubscriptionBuilderService.update_and_verify_builder( tenant_id=tenant_id, user_id=user.id, provider_id=TriggerProviderID(provider), @@ -307,6 +307,7 @@ class TriggerSubscriptionBuilderVerifyApi(Resource): credentials=payload.credentials, ), ) + return dump_response(TriggerVerificationResponse, result) except Exception as e: logger.exception("Error verifying provider credential", exc_info=e) raise ValueError(str(e)) from e @@ -317,7 +318,11 @@ class TriggerSubscriptionBuilderVerifyApi(Resource): ) class TriggerSubscriptionBuilderUpdateApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SubscriptionBuilderApiEntity.__name__]) + @console_ns.response( + 200, + "Trigger subscription builder updated successfully", + console_ns.models[SubscriptionBuilderApiEntity.__name__], + ) @setup_required @login_required @edit_permission_required @@ -326,21 +331,20 @@ class TriggerSubscriptionBuilderUpdateApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, provider: str, subscription_builder_id: str): """Update a subscription instance for a trigger provider""" + payload = TriggerSubscriptionBuilderUpdatePayload.model_validate(console_ns.payload or {}) try: - return jsonable_encoder( - TriggerSubscriptionBuilderService.update_trigger_subscription_builder( - tenant_id=tenant_id, - provider_id=TriggerProviderID(provider), - subscription_builder_id=subscription_builder_id, - subscription_builder_updater=SubscriptionBuilderUpdater( - name=payload.name, - parameters=payload.parameters, - properties=payload.properties, - credentials=payload.credentials, - ), - ) - ) + return TriggerSubscriptionBuilderService.update_trigger_subscription_builder( + tenant_id=tenant_id, + provider_id=TriggerProviderID(provider), + subscription_builder_id=subscription_builder_id, + subscription_builder_updater=SubscriptionBuilderUpdater( + name=payload.name, + parameters=payload.parameters, + properties=payload.properties, + credentials=payload.credentials, + ), + ).model_dump(mode="json") except Exception as e: logger.exception("Error updating provider credential", exc_info=e) raise @@ -350,7 +354,11 @@ class TriggerSubscriptionBuilderUpdateApi(Resource): "/workspaces/current/trigger-provider//subscriptions/builder/logs/", ) class TriggerSubscriptionBuilderLogsApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TriggerSubscriptionBuilderLogsResponse.__name__]) + @console_ns.response( + 200, + "Trigger subscription builder logs retrieved successfully", + console_ns.models[TriggerSubscriptionBuilderLogsResponse.__name__], + ) @setup_required @login_required @edit_permission_required @@ -358,9 +366,10 @@ class TriggerSubscriptionBuilderLogsApi(Resource): @account_initialization_required def get(self, provider: str, subscription_builder_id: str): """Get the request logs for a subscription instance for a trigger provider""" + try: logs = TriggerSubscriptionBuilderService.list_logs(subscription_builder_id) - return jsonable_encoder({"logs": [log.model_dump(mode="json") for log in logs]}) + return dump_response(TriggerSubscriptionBuilderLogsResponse, {"logs": logs}) except Exception as e: logger.exception("Error getting request logs for subscription builder", exc_info=e) raise @@ -371,7 +380,9 @@ class TriggerSubscriptionBuilderLogsApi(Resource): ) class TriggerSubscriptionBuilderBuildApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[TriggerProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Trigger subscription builder built successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @edit_permission_required @@ -395,7 +406,7 @@ class TriggerSubscriptionBuilderBuildApi(Resource): properties=payload.properties, ), ) - return 200 + return SimpleResultResponse(result="success").model_dump(mode="json") except Exception as e: logger.exception("Error building provider credential", exc_info=e) raise ValueError(str(e)) from e @@ -406,7 +417,9 @@ class TriggerSubscriptionBuilderBuildApi(Resource): ) class TriggerSubscriptionUpdateApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderUpdatePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[TriggerProviderOpaqueResponse.__name__]) + @console_ns.response( + 200, "Trigger subscription updated successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @edit_permission_required @@ -415,6 +428,7 @@ class TriggerSubscriptionUpdateApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, subscription_id: str): """Update a subscription instance""" + request = TriggerSubscriptionBuilderUpdatePayload.model_validate(console_ns.payload or {}) subscription = TriggerProviderService.get_subscription_by_id( @@ -440,7 +454,7 @@ class TriggerSubscriptionUpdateApi(Resource): name=request.name, properties=request.properties, ) - return 200 + return SimpleResultResponse(result="success").model_dump(mode="json") # For the rest cases(API_KEY, OAUTH2) # we need to call third party provider(e.g. GitHub) to rebuild the subscription @@ -452,7 +466,7 @@ class TriggerSubscriptionUpdateApi(Resource): credentials=request.credentials or subscription.credentials, parameters=request.parameters or subscription.parameters, ) - return 200 + return SimpleResultResponse(result="success").model_dump(mode="json") except ValueError as e: raise BadRequest(str(e)) except Exception as e: @@ -473,6 +487,7 @@ class TriggerSubscriptionDeleteApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, subscription_id: str): """Delete a subscription instance""" + try: with sessionmaker(db.engine).begin() as session: # Delete trigger provider subscription @@ -487,7 +502,7 @@ class TriggerSubscriptionDeleteApi(Resource): tenant_id=tenant_id, subscription_id=subscription_id, ) - return {"result": "success"} + return SimpleResultResponse(result="success").model_dump(mode="json") except ValueError as e: raise BadRequest(str(e)) except Exception as e: @@ -497,9 +512,10 @@ class TriggerSubscriptionDeleteApi(Resource): @console_ns.route("/workspaces/current/trigger-provider//subscriptions/oauth/authorize") class TriggerOAuthAuthorizeApi(Resource): + # response-contract:ignore cookie-bearing Flask response @console_ns.response( 200, - "Authorization URL retrieved successfully", + "Trigger OAuth authorization URL generated successfully", console_ns.models[TriggerOAuthAuthorizeResponse.__name__], ) @setup_required @@ -509,10 +525,12 @@ class TriggerOAuthAuthorizeApi(Resource): @with_current_tenant_id def get(self, tenant_id: str, user: Account, provider: str): """Initiate OAuth authorization flow for a trigger provider""" + try: provider_id = TriggerProviderID(provider) plugin_id = provider_id.plugin_id provider_name = provider_id.provider_name + tenant_id = tenant_id # Get OAuth client configuration oauth_client_params = TriggerProviderService.get_oauth_client( @@ -556,15 +574,12 @@ class TriggerOAuthAuthorizeApi(Resource): system_credentials=oauth_client_params, ) - # Create response with cookie response = make_response( - jsonable_encoder( - { - "authorization_url": authorization_url_response.authorization_url, - "subscription_builder_id": subscription_builder.id, - "subscription_builder": subscription_builder, - } - ) + TriggerOAuthAuthorizeResponse( + authorization_url=authorization_url_response.authorization_url, + subscription_builder_id=subscription_builder.id, + subscription_builder=subscription_builder, + ).model_dump(mode="json") ) response.set_cookie( "context_id", @@ -583,11 +598,8 @@ class TriggerOAuthAuthorizeApi(Resource): @console_ns.route("/oauth/plugin//trigger/callback") class TriggerOAuthCallbackApi(Resource): - @console_ns.response( - 302, - "Redirect to console OAuth callback page", - console_ns.models[RedirectResponse.__name__], - ) + # response-contract:ignore redirect response + @console_ns.response(302, "Redirect to OAuth callback page") @setup_required def get(self, provider: str): """Handle OAuth callback for trigger provider""" @@ -653,7 +665,11 @@ class TriggerOAuthCallbackApi(Resource): @console_ns.route("/workspaces/current/trigger-provider//oauth/client") class TriggerOAuthClientManageApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TriggerOAuthClientResponse.__name__]) + @console_ns.response( + 200, + "Trigger OAuth client retrieved successfully", + console_ns.models[TriggerOAuthClientResponse.__name__], + ) @setup_required @login_required @is_admin_or_owner_required @@ -662,6 +678,7 @@ class TriggerOAuthClientManageApi(Resource): @with_current_tenant_id def get(self, tenant_id: str, provider: str): """Get OAuth client configuration for a provider""" + try: provider_id = TriggerProviderID(provider) @@ -682,24 +699,24 @@ class TriggerOAuthClientManageApi(Resource): ) provider_controller = TriggerManager.get_trigger_provider(tenant_id, provider_id) redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider}/trigger/callback" - return jsonable_encoder( - { - "configured": bool(custom_params or system_client_exists), - "system_configured": system_client_exists, - "custom_configured": bool(custom_params), - "oauth_client_schema": provider_controller.get_oauth_client_schema(), - "custom_enabled": is_custom_enabled, - "redirect_uri": redirect_uri, - "params": custom_params or {}, - } - ) + return TriggerOAuthClientResponse( + configured=bool(custom_params or system_client_exists), + system_configured=system_client_exists, + custom_configured=bool(custom_params), + oauth_client_schema=provider_controller.get_oauth_client_schema(), + custom_enabled=is_custom_enabled, + redirect_uri=redirect_uri, + params=dict(custom_params or {}), + ).model_dump(mode="json") except Exception as e: logger.exception("Error getting OAuth client", exc_info=e) raise @console_ns.expect(console_ns.models[TriggerOAuthClientPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_ns.response( + 200, "Trigger OAuth client saved successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @is_admin_or_owner_required @@ -708,16 +725,18 @@ class TriggerOAuthClientManageApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, provider: str): """Configure custom OAuth client for a provider""" + payload = TriggerOAuthClientPayload.model_validate(console_ns.payload or {}) try: provider_id = TriggerProviderID(provider) - return TriggerProviderService.save_custom_oauth_client_params( + result = TriggerProviderService.save_custom_oauth_client_params( tenant_id=tenant_id, provider_id=provider_id, client_params=payload.client_params, enabled=payload.enabled, ) + return dump_response(SimpleResultResponse, result) except ValueError as e: raise BadRequest(str(e)) @@ -725,22 +744,26 @@ class TriggerOAuthClientManageApi(Resource): logger.exception("Error configuring OAuth client", exc_info=e) raise + @console_ns.response( + 200, "Trigger OAuth client deleted successfully", console_ns.models[SimpleResultResponse.__name__] + ) @setup_required @login_required @is_admin_or_owner_required @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.PLUGIN_PREFERENCES, resource_required=False) @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) @with_current_tenant_id def delete(self, tenant_id: str, provider: str): """Remove custom OAuth client configuration""" + try: provider_id = TriggerProviderID(provider) - return TriggerProviderService.delete_custom_oauth_client_params( + result = TriggerProviderService.delete_custom_oauth_client_params( tenant_id=tenant_id, provider_id=provider_id, ) + return dump_response(SimpleResultResponse, result) except ValueError as e: raise BadRequest(str(e)) except Exception as e: @@ -753,7 +776,11 @@ class TriggerOAuthClientManageApi(Resource): ) class TriggerSubscriptionVerifyApi(Resource): @console_ns.expect(console_ns.models[TriggerSubscriptionBuilderVerifyPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[TriggerSubscriptionBuilderVerifyResponse.__name__]) + @console_ns.response( + 200, + "Trigger subscription verified successfully", + console_ns.models[TriggerVerificationResponse.__name__], + ) @setup_required @login_required @edit_permission_required @@ -763,6 +790,7 @@ class TriggerSubscriptionVerifyApi(Resource): @with_current_tenant_id def post(self, tenant_id: str, user: Account, provider: str, subscription_id: str): """Verify credentials for an existing subscription (edit mode only)""" + verify_request = TriggerSubscriptionBuilderVerifyPayload.model_validate(console_ns.payload or {}) try: @@ -773,7 +801,7 @@ class TriggerSubscriptionVerifyApi(Resource): subscription_id=subscription_id, credentials=verify_request.credentials, ) - return result + return dump_response(TriggerVerificationResponse, result) except ValueError as e: logger.warning("Credential verification failed", exc_info=e) raise BadRequest(str(e)) from e diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index 418c3eb66e1..23ce116b349 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -1,8 +1,9 @@ import logging from datetime import datetime +from http import HTTPStatus from flask import request -from flask_restx import Resource, fields, marshal +from flask_restx import Resource from pydantic import BaseModel, Field, field_validator from sqlalchemy import select from werkzeug.exceptions import Unauthorized @@ -16,7 +17,12 @@ from controllers.common.errors import ( TooManyFilesError, UnsupportedFileTypeError, ) -from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models +from controllers.common.schema import ( + query_params_from_model, + query_params_from_request, + register_response_schema_models, + register_schema_models, +) from controllers.console import console_ns from controllers.console.admin import admin_required from controllers.console.error import AccountNotLinkTenantError @@ -31,7 +37,7 @@ from controllers.console.wraps import ( from enums.cloud_plan import CloudPlan from extensions.ext_database import db from fields.base import ResponseModel -from libs.helper import OptionalTimestampField, TimestampField, dump_response, to_timestamp +from libs.helper import dump_response, to_timestamp from libs.login import login_required from libs.pagination import paginate_query from models.account import Account, Tenant, TenantAccountJoin, TenantCustomConfigDict, TenantStatus @@ -133,7 +139,7 @@ class WorkspaceListItemResponse(ResponseModel): @field_validator("status", mode="before") @classmethod - def _normalize_status(cls, value): + def _normalize_enum_like(cls, value): if value is None: return None if isinstance(value, str): @@ -146,7 +152,7 @@ class WorkspaceListItemResponse(ResponseModel): return to_timestamp(value) -class WorkspaceListResponse(ResponseModel): +class WorkspacePaginationResponse(ResponseModel): data: list[WorkspaceListItemResponse] has_more: bool limit: int @@ -159,7 +165,7 @@ class SwitchWorkspaceResponse(ResponseModel): new_tenant: TenantInfoResponse -class WorkspaceMutationResponse(ResponseModel): +class WorkspaceTenantResultResponse(ResponseModel): result: str tenant: TenantInfoResponse @@ -174,6 +180,16 @@ class WorkspacePermissionResponse(ResponseModel): allow_owner_transfer: bool +WORKSPACE_LOGO_UPLOAD_PARAMS = { + "file": { + "in": "formData", + "type": "file", + "required": True, + "description": "Workspace web app logo file. Only SVG and PNG files are supported.", + } +} + + register_schema_models( console_ns, WorkspaceListQuery, @@ -184,53 +200,21 @@ register_schema_models( register_response_schema_models( console_ns, TenantInfoResponse, + TenantListItemResponse, TenantListResponse, - WorkspaceListResponse, - SwitchWorkspaceResponse, - WorkspaceMutationResponse, - WorkspaceLogoUploadResponse, WorkspaceCustomConfigResponse, + WorkspaceListItemResponse, + WorkspacePaginationResponse, + SwitchWorkspaceResponse, + WorkspaceTenantResultResponse, + WorkspaceLogoUploadResponse, WorkspacePermissionResponse, ) -provider_fields = { - "provider_name": fields.String, - "provider_type": fields.String, - "is_valid": fields.Boolean, - "token_is_set": fields.Boolean, -} - -tenant_fields = { - "id": fields.String, - "name": fields.String, - "plan": fields.String, - "status": fields.String, - "created_at": TimestampField, - "role": fields.String, - "in_trial": fields.Boolean, - "trial_end_reason": fields.String, - "custom_config": fields.Raw(attribute="custom_config"), - "trial_credits": fields.Integer, - "trial_credits_used": fields.Integer, - "next_credit_reset_date": fields.Integer, -} - -tenants_fields = { - "id": fields.String, - "name": fields.String, - "plan": fields.String, - "status": fields.String, - "created_at": TimestampField, - "last_opened_at": OptionalTimestampField, - "current": fields.Boolean, -} - -workspace_fields = {"id": fields.String, "name": fields.String, "status": fields.String, "created_at": TimestampField} - @console_ns.route("/workspaces") class TenantListApi(Resource): - @console_ns.response(200, "Success", console_ns.models[TenantListResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[TenantListResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -239,7 +223,7 @@ class TenantListApi(Resource): def get(self, current_tenant_id: str, current_user: Account): tenant_rows: list[tuple[Tenant, TenantAccountJoin]] = [ (tenant, membership) - for tenant, membership in TenantService.get_workspaces_for_account(db.session, current_user.id) + for tenant, membership in TenantService.get_workspaces_for_account(current_user.id, session=db.session()) if tenant.status == TenantStatus.NORMAL ] tenants = [tenant for tenant, _ in tenant_rows] @@ -281,18 +265,17 @@ class TenantListApi(Resource): tenant_dicts.append(tenant_dict) - return {"workspaces": marshal(tenant_dicts, tenants_fields)}, 200 + return dump_response(TenantListResponse, {"workspaces": tenant_dicts}), HTTPStatus.OK @console_ns.route("/all-workspaces") class WorkspaceListApi(Resource): @console_ns.doc(params=query_params_from_model(WorkspaceListQuery)) - @console_ns.response(200, "Success", console_ns.models[WorkspaceListResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[WorkspacePaginationResponse.__name__]) @setup_required @admin_required def get(self): - payload = request.args.to_dict(flat=True) - args = WorkspaceListQuery.model_validate(payload) + args = query_params_from_request(WorkspaceListQuery) stmt = select(Tenant).order_by(Tenant.created_at.desc()) tenants = paginate_query(stmt, page=args.page, per_page=args.limit) @@ -301,13 +284,9 @@ class WorkspaceListApi(Resource): if tenants.has_next: has_more = True - return { - "data": marshal(tenants.items, workspace_fields), - "has_more": has_more, - "limit": args.limit, - "page": args.page, - "total": tenants.total, - }, 200 + return WorkspacePaginationResponse( + data=tenants.items, has_more=has_more, limit=args.limit, page=args.page, total=tenants.total or 0 + ).model_dump(mode="json"), HTTPStatus.OK @console_ns.route("/workspaces/current", endpoint="workspaces_current") @@ -316,7 +295,7 @@ class TenantApi(Resource): @setup_required @login_required @account_initialization_required - @console_ns.response(200, "Success", console_ns.models[TenantInfoResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[TenantInfoResponse.__name__]) @with_current_user def post(self, current_user: Account): if request.path == "/info": @@ -327,22 +306,25 @@ class TenantApi(Resource): raise ValueError("No current tenant") if tenant.status == TenantStatus.ARCHIVE: - tenants = TenantService.get_join_tenants(current_user, session=db.session) + tenants = TenantService.get_join_tenants(current_user, session=db.session()) # if there is any tenant, switch to the first one if len(tenants) > 0: - TenantService.switch_tenant(current_user, tenants[0].id, session=db.session) + TenantService.switch_tenant(current_user, tenants[0].id, session=db.session()) tenant = tenants[0] # else, raise Unauthorized else: raise Unauthorized("workspace is archived") - return dump_response(TenantInfoResponse, WorkspaceService.get_tenant_info(tenant)), 200 + return ( + dump_response(TenantInfoResponse, WorkspaceService.get_tenant_info(tenant, session=db.session())), + HTTPStatus.OK, + ) @console_ns.route("/workspaces/switch") class SwitchWorkspaceApi(Resource): @console_ns.expect(console_ns.models[SwitchWorkspacePayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[SwitchWorkspaceResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[SwitchWorkspaceResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -351,9 +333,9 @@ class SwitchWorkspaceApi(Resource): payload = console_ns.payload or {} args = SwitchWorkspacePayload.model_validate(payload) - # check if tenant_id is valid, 403 if not + # Check whether the tenant_id belongs to the current account. try: - TenantService.switch_tenant(current_user, args.tenant_id, session=db.session) + TenantService.switch_tenant(current_user, args.tenant_id, session=db.session()) except Exception: raise AccountNotLinkTenantError("Account not link tenant") @@ -361,13 +343,15 @@ class SwitchWorkspaceApi(Resource): if new_tenant is None: raise ValueError("Tenant not found") - return {"result": "success", "new_tenant": marshal(WorkspaceService.get_tenant_info(new_tenant), tenant_fields)} + return SwitchWorkspaceResponse( + result="success", new_tenant=WorkspaceService.get_tenant_info(new_tenant, session=db.session()) + ).model_dump(mode="json") @console_ns.route("/workspaces/custom-config") class CustomConfigWorkspaceApi(Resource): @console_ns.expect(console_ns.models[WorkspaceCustomConfigPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[WorkspaceMutationResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[WorkspaceTenantResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -390,12 +374,15 @@ class CustomConfigWorkspaceApi(Resource): tenant.custom_config_dict = custom_config_dict db.session.commit() - return {"result": "success", "tenant": marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)} + return WorkspaceTenantResultResponse( + result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=db.session()) + ).model_dump(mode="json") @console_ns.route("/workspaces/custom-config/webapp-logo/upload") class WebappLogoWorkspaceApi(Resource): - @console_ns.response(201, "Logo uploaded", console_ns.models[WorkspaceLogoUploadResponse.__name__]) + @console_ns.doc(consumes=["multipart/form-data"], params=WORKSPACE_LOGO_UPLOAD_PARAMS) + @console_ns.response(HTTPStatus.CREATED, "Logo uploaded", console_ns.models[WorkspaceLogoUploadResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -431,13 +418,13 @@ class WebappLogoWorkspaceApi(Resource): except services.errors.file.UnsupportedFileTypeError: raise UnsupportedFileTypeError() - return {"id": upload_file.id}, 201 + return WorkspaceLogoUploadResponse(id=upload_file.id).model_dump(mode="json"), HTTPStatus.CREATED @console_ns.route("/workspaces/info") class WorkspaceInfoApi(Resource): @console_ns.expect(console_ns.models[WorkspaceInfoPayload.__name__]) - @console_ns.response(200, "Success", console_ns.models[WorkspaceMutationResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[WorkspaceTenantResultResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -453,14 +440,16 @@ class WorkspaceInfoApi(Resource): tenant.name = args.name db.session.commit() - return {"result": "success", "tenant": marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)} + return WorkspaceTenantResultResponse( + result="success", tenant=WorkspaceService.get_tenant_info(tenant, session=db.session()) + ).model_dump(mode="json") @console_ns.route("/workspaces/current/permission") class WorkspacePermissionApi(Resource): """Get workspace permissions for the current workspace.""" - @console_ns.response(200, "Success", console_ns.models[WorkspacePermissionResponse.__name__]) + @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[WorkspacePermissionResponse.__name__]) @setup_required @login_required @account_initialization_required @@ -477,8 +466,8 @@ class WorkspacePermissionApi(Resource): # Get workspace permissions from enterprise service permission = EnterpriseService.WorkspacePermissionService.get_permission(current_tenant_id) - return { - "workspace_id": permission.workspace_id, - "allow_member_invite": permission.allow_member_invite, - "allow_owner_transfer": permission.allow_owner_transfer, - }, 200 + return WorkspacePermissionResponse( + workspace_id=permission.workspace_id, + allow_member_invite=permission.allow_member_invite, + allow_owner_transfer=permission.allow_owner_transfer, + ).model_dump(mode="json"), HTTPStatus.OK diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 017793ffe0b..37d7239170c 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -4,7 +4,7 @@ import os import time from collections.abc import Callable from functools import wraps -from typing import Any, Concatenate, overload +from typing import Any, Concatenate, Protocol, cast, overload from flask import abort, request from pydantic import BaseModel, ValidationError @@ -46,6 +46,60 @@ ERROR_MSG_INVALID_ENCRYPTED_DATA = "Invalid encrypted data" ERROR_MSG_INVALID_ENCRYPTED_CODE = "Invalid encrypted code" +class OnceTrueCallable[**P](Protocol): + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> bool: ... + + def mark_success(self) -> None: ... + + def reset_success(self) -> None: ... + + +def once_true[**P](func: Callable[P, bool]) -> OnceTrueCallable[P]: + """Wrap a predicate so only a strict True result is memoized.""" + has_success = False + + def mark_success() -> None: + nonlocal has_success + + has_success = True + + def reset_success() -> None: + nonlocal has_success + + has_success = False + + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> bool: + nonlocal has_success + + if has_success: + return True + + result = func(*args, **kwargs) + if result is True: + has_success = True + + return result + + wrapper.mark_success = mark_success # type: ignore[attr-defined] + wrapper.reset_success = reset_success # type: ignore[attr-defined] + return cast(OnceTrueCallable[P], wrapper) + + +def mark_setup_completed() -> None: + """Remember in this process that one-time self-hosted setup has completed.""" + _is_setup_completed.mark_success() + + +@once_true +def _is_setup_completed() -> bool: + """Check whether setup exists, caching only successful observations. + + Use `once_true` instead of `@cache` because a pre-setup False result must not be memoized. + """ + return db.session.scalar(select(DifySetup).limit(1)) is not None + + @overload def account_initialization_required[T, **P, R]( view: Callable[Concatenate[T, P], R], @@ -246,7 +300,9 @@ def setup_required[T, **P, R]( @overload -def setup_required[**P, R](view: Callable[P, R]) -> Callable[P, R]: ... +def setup_required[**P, R](view: Callable[P, R]) -> Callable[P, R]: + """Require self-hosted bootstrap setup before serving protected routes.""" + ... def setup_required[R](view: Callable[..., R]) -> Callable[..., R]: @@ -255,7 +311,7 @@ def setup_required[R](view: Callable[..., R]) -> Callable[..., R]: # The overloads keep Resource methods method-aware for pyrefly while # preserving support for plain functions used in tests and utilities. # check setup - if dify_config.EDITION == "SELF_HOSTED" and not db.session.scalar(select(DifySetup).limit(1)): + if dify_config.EDITION == "SELF_HOSTED" and not _is_setup_completed(): if os.environ.get("INIT_PASSWORD"): raise NotInitValidateError() raise NotSetupError() diff --git a/api/controllers/files/agent_drive_archive.py b/api/controllers/files/agent_drive_archive.py index afa6ac79483..8ecec2e9a4c 100644 --- a/api/controllers/files/agent_drive_archive.py +++ b/api/controllers/files/agent_drive_archive.py @@ -8,6 +8,7 @@ from werkzeug.exceptions import Forbidden, NotFound from controllers.common.file_response import enforce_download_for_html from controllers.common.schema import register_schema_models from controllers.files import files_ns +from extensions.ext_database import db from models.agent import AgentDriveFileKind from services.agent_drive_service import AgentDriveError, AgentDriveService @@ -54,6 +55,7 @@ class AgentDriveArchiveMemberApi(Resource): archive_file_kind=args.archive_file_kind, archive_file_id=args.archive_file_id, member_path=args.member_path, + session=db.session(), ) except AgentDriveError as exc: raise NotFound(exc.message) from exc diff --git a/api/controllers/inner_api/app/dsl.py b/api/controllers/inner_api/app/dsl.py index 915a11dcddc..9fd111f86dc 100644 --- a/api/controllers/inner_api/app/dsl.py +++ b/api/controllers/inner_api/app/dsl.py @@ -98,6 +98,7 @@ class EnterpriseAppDSLExport(Resource): data = AppDslService.export_dsl( app_model=app_model, + session=db.session(), include_secret=include_secret, ) diff --git a/api/controllers/inner_api/plugin/agent_drive.py b/api/controllers/inner_api/plugin/agent_drive.py index 0cdb9dab35f..e06720a8e99 100644 --- a/api/controllers/inner_api/plugin/agent_drive.py +++ b/api/controllers/inner_api/plugin/agent_drive.py @@ -17,6 +17,7 @@ from controllers.console.wraps import setup_required from controllers.inner_api import inner_api_ns from controllers.inner_api.plugin.wraps import get_user from controllers.inner_api.wraps import plugin_inner_api_only +from extensions.ext_database import db from services.agent_drive_service import ( AgentDriveError, AgentDriveService, @@ -53,6 +54,7 @@ class AgentDriveManifestApi(Resource): agent_id=agent_id, prefix=request.args.get("prefix", ""), include_download_url=include_download_url, + session=db.session(), ) except AgentDriveError as exc: return _error_response(exc) @@ -71,7 +73,7 @@ class AgentDriveSkillsApi(Resource): tenant_id = (request.args.get("tenant_id") or "").strip() if not tenant_id: raise AgentDriveError("missing_tenant_id", "tenant_id is required", status_code=400) - items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=agent_id) + items = AgentDriveService().list_skills(tenant_id=tenant_id, agent_id=agent_id, session=db.session()) except AgentDriveError as exc: return _error_response(exc) return {"items": items} @@ -96,6 +98,7 @@ class AgentDriveCommitApi(Resource): user_id=user.id, agent_id=agent_id, items=body.items, + session=db.session(), ) except AgentDriveError as exc: return _error_response(exc) diff --git a/api/controllers/inner_api/workspace/workspace.py b/api/controllers/inner_api/workspace/workspace.py index 1f25eb576d3..b3a571112f6 100644 --- a/api/controllers/inner_api/workspace/workspace.py +++ b/api/controllers/inner_api/workspace/workspace.py @@ -47,8 +47,8 @@ class EnterpriseWorkspace(Resource): if account is None: return {"message": "owner account not found."}, 404 - tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session) - TenantService.create_tenant_member(tenant, account, db.session, role="owner") + tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session()) + TenantService.create_tenant_member(tenant, account, db.session(), role="owner") tenant_was_created.send(tenant) @@ -84,7 +84,7 @@ class EnterpriseWorkspaceNoOwnerEmail(Resource): def post(self): args = WorkspaceOwnerlessPayload.model_validate(inner_api_ns.payload or {}) - tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session) + tenant = TenantService.create_tenant(args.name, is_from_dashboard=True, session=db.session()) tenant_was_created.send(tenant) diff --git a/api/controllers/mcp/mcp.py b/api/controllers/mcp/mcp.py index 3830c9585a5..45a5a4e6899 100644 --- a/api/controllers/mcp/mcp.py +++ b/api/controllers/mcp/mcp.py @@ -1,15 +1,15 @@ from typing import Any, Union -from flask import Response +from flask import Response, request from flask_restx import Resource -from pydantic import BaseModel, Field, ValidationError +from pydantic import BaseModel, Field, RootModel, ValidationError from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker -from controllers.common.schema import register_schema_model +from controllers.common.schema import register_response_schema_models, register_schema_model from controllers.mcp import mcp_ns from core.mcp import types as mcp_types -from core.mcp.server.streamable_http import handle_mcp_request +from core.mcp.server.streamable_http import handle_mcp_request, negotiate_protocol_version from extensions.ext_database import db from graphon.variables.input_entities import VariableEntity, VariableEntityType from libs import helper @@ -33,7 +33,12 @@ class MCPRequestPayload(BaseModel): id: int | str | None = Field(default=None, description="Request ID for tracking responses") +class MCPJSONRPCResponse(RootModel[mcp_types.JSONRPCResponse | mcp_types.JSONRPCError]): + pass + + register_schema_model(mcp_ns, MCPRequestPayload) +register_response_schema_models(mcp_ns, MCPJSONRPCResponse) @mcp_ns.route("/server//mcp") @@ -42,13 +47,10 @@ class MCPAppApi(Resource): @mcp_ns.doc("handle_mcp_request") @mcp_ns.doc(description="Handle Model Context Protocol (MCP) requests for a specific server") @mcp_ns.doc(params={"server_code": "Unique identifier for the MCP server"}) - @mcp_ns.doc( - responses={ - 200: "MCP response successfully processed", - 400: "Invalid MCP request or parameters", - 404: "Server or app not found", - } - ) + @mcp_ns.response(200, "MCP JSON-RPC response", mcp_ns.models[MCPJSONRPCResponse.__name__]) + @mcp_ns.response(202, "MCP notification accepted") + @mcp_ns.response(400, "Invalid MCP request or parameters") + @mcp_ns.response(404, "Server or app not found") def post(self, server_code: str): """Handle MCP requests for a specific server. @@ -64,10 +66,22 @@ class MCPAppApi(Resource): Raises: ValidationError: Invalid request format or parameters """ + # response-contract:ignore MCP route returns Flask Response from JSON-RPC handler args = MCPRequestPayload.model_validate(mcp_ns.payload or {}) request_id: Union[int, str] | None = args.id mcp_request = self._parse_mcp_request(args.model_dump(exclude_none=True)) + # Resolve the negotiated protocol version from the MCP-Protocol-Version header. + is_initialize = isinstance(mcp_request.root, mcp_types.InitializeRequest) + header_value = request.headers.get("MCP-Protocol-Version") + protocol_version = negotiate_protocol_version(header_value, is_initialize) + if protocol_version is None: + # A notification never receives a response, even with an unsupported header. + if isinstance(mcp_request, mcp_types.ClientNotification): + protocol_version = mcp_types.DEFAULT_NEGOTIATED_VERSION + else: + return self._protocol_version_error_response(request_id, header_value) + with sessionmaker(db.engine, expire_on_commit=False).begin() as session: # Get MCP server and app mcp_server, app = self._get_mcp_server_and_app(server_code, session) @@ -77,7 +91,28 @@ class MCPAppApi(Resource): user_input_form = self._get_user_input_form(app) # Handle notification vs request differently - return self._process_mcp_message(mcp_request, request_id, app, mcp_server, user_input_form, session) + return self._process_mcp_message( + mcp_request, request_id, app, mcp_server, user_input_form, session, protocol_version + ) + + def _protocol_version_error_response( + self, request_id: Union[int, str] | None, header_value: str | None + ) -> Response: + """Return a JSON-RPC error for an unsupported MCP-Protocol-Version header. + + Per JSON-RPC 2.0, an error whose request id is unknown uses a null id, so we echo the + offending request's id directly (None -> null) instead of fabricating a placeholder. + """ + error_data = mcp_types.ErrorData( + code=mcp_types.INVALID_REQUEST, + message=f"Unsupported MCP-Protocol-Version: {header_value}", + ) + error_response = { + "jsonrpc": "2.0", + "id": request_id, + "error": error_data.model_dump(by_alias=True, mode="json", exclude_none=True), + } + return helper.compact_generate_response(error_response) def _get_mcp_server_and_app(self, server_code: str, session: Session) -> tuple[AppMCPServer, App]: """Get and validate MCP server and app in one query session""" @@ -104,12 +139,15 @@ class MCPAppApi(Resource): mcp_server: AppMCPServer, user_input_form: list[VariableEntity], session: Session, + protocol_version: str, ) -> Response: """Process MCP message (notification or request)""" if isinstance(mcp_request, mcp_types.ClientNotification): return self._handle_notification(mcp_request) else: - return self._handle_request(mcp_request, request_id, app, mcp_server, user_input_form, session) + return self._handle_request( + mcp_request, request_id, app, mcp_server, user_input_form, session, protocol_version + ) def _handle_notification(self, mcp_request: mcp_types.ClientNotification) -> Response: """Handle MCP notification""" @@ -127,12 +165,15 @@ class MCPAppApi(Resource): mcp_server: AppMCPServer, user_input_form: list[VariableEntity], session: Session, + protocol_version: str, ) -> Response: """Handle MCP request""" if request_id is None: raise MCPRequestError(mcp_types.INVALID_REQUEST, "Request ID is required") - result = self._handle_mcp_request(app, mcp_server, mcp_request, user_input_form, session, request_id) + result = self._handle_mcp_request( + app, mcp_server, mcp_request, user_input_form, session, request_id, protocol_version + ) if result is None: # This shouldn't happen for requests, but handle gracefully raise MCPRequestError(mcp_types.INTERNAL_ERROR, "No response generated for request") @@ -229,6 +270,7 @@ class MCPAppApi(Resource): user_input_form: list[VariableEntity], session: Session, request_id: Union[int, str], + protocol_version: str, ) -> mcp_types.JSONRPCResponse | mcp_types.JSONRPCError | None: """Handle MCP request and return response""" end_user = self._retrieve_end_user(mcp_server.tenant_id, mcp_server.id) @@ -238,4 +280,6 @@ class MCPAppApi(Resource): client_name = f"{client_info.name}@{client_info.version}" end_user = self._create_end_user(client_name, app.tenant_id, app.id, mcp_server.id, session) - return handle_mcp_request(session, app, mcp_request, user_input_form, mcp_server, end_user, request_id) + return handle_mcp_request( + session, app, mcp_request, user_input_form, mcp_server, end_user, request_id, protocol_version + ) diff --git a/api/controllers/openapi/__init__.py b/api/controllers/openapi/__init__.py index 81c65ca03be..0260422ec1c 100644 --- a/api/controllers/openapi/__init__.py +++ b/api/controllers/openapi/__init__.py @@ -2,11 +2,13 @@ from flask import Blueprint from flask_restx import Namespace from controllers.openapi._errors import ErrorBody, OpenApiErrorCode, OpenApiErrorFormatter +from controllers.openapi._version_gate import attach_version_gate from libs.device_flow_security import attach_anti_framing from libs.external_api import ExternalApi bp = Blueprint("openapi", __name__, url_prefix="/openapi/v1") attach_anti_framing(bp) +attach_version_gate(bp) api = ExternalApi( bp, diff --git a/api/controllers/openapi/_errors.py b/api/controllers/openapi/_errors.py index 5e82c2614de..92884dfcd50 100644 --- a/api/controllers/openapi/_errors.py +++ b/api/controllers/openapi/_errors.py @@ -45,10 +45,12 @@ class OpenApiErrorCode(StrEnum): TOO_MANY_REQUESTS = "too_many_requests" INTERNAL_ERROR = "internal_server_error" BAD_GATEWAY = "bad_gateway" + UPGRADE_REQUIRED = "upgrade_required" UNKNOWN = "unknown" # domain codes (must match the error_code attribute of the exception # classes raised on the openapi surface) APP_UNAVAILABLE = "app_unavailable" + AGENT_NOT_PUBLISHED = "agent_not_published" CONVERSATION_COMPLETED = "conversation_completed" PROVIDER_NOT_INITIALIZE = "provider_not_initialize" PROVIDER_QUOTA_EXCEEDED = "provider_quota_exceeded" diff --git a/api/controllers/openapi/_models.py b/api/controllers/openapi/_models.py index 6e8a9c9d439..5337612e7b6 100644 --- a/api/controllers/openapi/_models.py +++ b/api/controllers/openapi/_models.py @@ -279,7 +279,7 @@ def _csv_string_query_schema(schema: dict[str, Any]) -> None: class AppDescribeQuery(BaseModel): - """`?fields=` allow-list for GET /apps//describe. + """`?fields=` allow-list for GET /apps/. Empty / omitted → all blocks. Unknown member → ValidationError → 422. """ @@ -441,7 +441,7 @@ class MemberActionResponse(BaseModel): class TaskStopResponse(BaseModel): - """200 body for POST /apps//tasks//stop. The handler always returns + """200 body for POST /apps//tasks/:stop. The handler always returns {"result": "success"}, so `result` is required (no default) — the generated contract types it as a required `'success'` rather than an optional field.""" @@ -473,7 +473,7 @@ class AppDslImportPayload(BaseModel): class AppDslExportQuery(BaseModel): - """Query parameters for GET /apps//export.""" + """Query parameters for GET /apps//dsl.""" include_secret: bool = Field(False, description="Include encrypted secret values in the exported DSL") workflow_id: UUIDStr | None = Field( @@ -488,7 +488,7 @@ class AppDslExportResponse(BaseModel): class FormSubmitResponse(BaseModel): - """Empty 200 body for POST /apps//form/human_input/. `extra='forbid'` + """Empty 200 body for POST /apps//human-input-forms/:submit. `extra='forbid'` pins `additionalProperties: false` so the generated contract is an exact `{}` rather than an under-annotated open object.""" diff --git a/api/controllers/openapi/_version_gate.py b/api/controllers/openapi/_version_gate.py new file mode 100644 index 00000000000..785617fd950 --- /dev/null +++ b/api/controllers/openapi/_version_gate.py @@ -0,0 +1,69 @@ +"""Version gate: reject outdated difyctl clients on /openapi/v1 with HTTP 426. + +difyctl and the ``/openapi/v1`` surface ship in lockstep. A breaking path change +(resource-oriented paths) means an outdated difyctl would call removed paths and +get a bare 404; this gate returns ``426 Upgrade Required`` with an upgrade hint +instead. +""" + +from __future__ import annotations + +import re +from typing import Final + +from flask import Blueprint, Response, request +from packaging.version import InvalidVersion, Version + +from configs import dify_config +from controllers.openapi._errors import ErrorBody, OpenApiErrorCode + +_UPGRADE_HINT: Final = "Upgrade difyctl: https://docs.dify.ai/en/cli/install" + +# difyctl sends `User-Agent: difyctl/ (; ; )`. +_DIFYCTL_UA_RE = re.compile(r"^difyctl/(\d+\.\d+\.\d+(?:-[\w.]+)?)") + +_PREFIX: Final = "/openapi/v1/" + +# Paths a too-old client must still reach to discover that it is outdated. +_ALLOWLIST: Final = frozenset({"/openapi/v1/_version", "/openapi/v1/_health"}) + + +def _upgrade_required_response(client_version: str, min_version: str) -> Response: + body = ErrorBody( + code=OpenApiErrorCode.UPGRADE_REQUIRED, + message=f"difyctl {client_version} is no longer supported; upgrade to >= {min_version}.", + status=426, + hint=_UPGRADE_HINT, + ) + return Response(body.model_dump_json(exclude_none=True), status=426, mimetype="application/json") + + +def attach_version_gate(bp: Blueprint) -> None: + """Reject difyctl clients older than ``[tool.dify] min_difyctl_version`` with 426. + + Registered app-wide (``before_app_request``) rather than blueprint-scoped so it + also fires for requests to *removed* paths — those no longer match an openapi + route and would 404 before a blueprint-scoped ``before_request`` ever runs. The + prefix guard scopes it back to ``/openapi/v1``. Fails open for non-difyctl or + unparseable User-Agents (only a confidently-too-old difyctl is blocked). + """ + + @bp.before_app_request + def _enforce_min_client_version() -> Response | None: # pyright: ignore[reportUnusedFunction] + if not request.path.startswith(_PREFIX): + return None + if request.path in _ALLOWLIST: + return None + match = _DIFYCTL_UA_RE.match(request.headers.get("User-Agent", "")) + if match is None: + return None + try: + client_version = Version(match.group(1)) + except InvalidVersion: + return None + # Compare the numeric core (major.minor.patch) only — a pre-release build + # like 0.2.0-rc.1 must not sort below the 0.2.0 floor. + min_version = dify_config.tool.dify.min_difyctl_version + if client_version.release[:3] < Version(min_version).release[:3]: + return _upgrade_required_response(match.group(1), min_version) + return None diff --git a/api/controllers/openapi/account.py b/api/controllers/openapi/account.py index 8ad0b02f4a0..b4786f2ae25 100644 --- a/api/controllers/openapi/account.py +++ b/api/controllers/openapi/account.py @@ -45,8 +45,10 @@ class AccountApi(Resource): enforce(LIMIT_ME_PER_ACCOUNT, key=f"account:{auth_data.account_id}") account_id_str = str(auth_data.account_id) if auth_data.account_id else None - account = AccountService.get_account_by_id(db.session, account_id_str) if account_id_str else None - memberships = TenantService.get_account_memberships(db.session, account_id_str) if account_id_str else [] + account = AccountService.get_account_by_id(account_id_str, session=db.session()) if account_id_str else None + memberships = ( + TenantService.get_account_memberships(account_id_str, session=db.session()) if account_id_str else [] + ) default_ws_id = _pick_default_workspace(memberships) return AccountResponse( @@ -63,7 +65,7 @@ class AccountSessionsSelfApi(Resource): @auth_router.guard(scope=Scope.FULL, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) @returns(200, RevokeResponse, description="Session revoked") def delete(self, *, auth_data: AuthData): - revoke_oauth_token(db.session, redis_client, str(auth_data.token_id)) + revoke_oauth_token(redis_client, str(auth_data.token_id), session=db.session()) return RevokeResponse(status="revoked") @@ -81,7 +83,7 @@ class AccountSessionsApi(Resource): page = query.page limit = query.limit - all_rows = list_active_sessions(db.session, ctx, now) + all_rows = list_active_sessions(ctx, now, session=db.session()) total = len(all_rows) sliced = all_rows[(page - 1) * limit : page * limit] @@ -117,10 +119,10 @@ class AccountSessionByIdApi(Resource): # 404 (not 403) on cross-subject so the endpoint doesn't leak # token IDs that belong to other subjects. - if not token_belongs_to_subject(db.session, session_id, ctx): + if not token_belongs_to_subject(session_id, ctx, session=db.session()): raise NotFound("session not found") - revoke_oauth_token(db.session, redis_client, session_id) + revoke_oauth_token(redis_client, session_id, session=db.session()) return RevokeResponse(status="revoked") diff --git a/api/controllers/openapi/app_dsl.py b/api/controllers/openapi/app_dsl.py index 9b1abd24bac..d06845dada4 100644 --- a/api/controllers/openapi/app_dsl.py +++ b/api/controllers/openapi/app_dsl.py @@ -30,7 +30,7 @@ class AppDslImportApi(Resource): a new app. Returns 202 when the DSL version requires an explicit confirmation step - (major version mismatch). Callers must then POST to the confirm endpoint. + (major version mismatch). Callers must then POST to the imports :confirm method. Returns 400 when the import failed due to invalid DSL or a business error. """ @@ -79,7 +79,7 @@ class AppDslImportApi(Resource): return result, 200 -@openapi_ns.route("/workspaces//apps/imports//confirm") +@openapi_ns.route("/workspaces//apps/imports/:confirm") class AppDslImportConfirmApi(Resource): """Confirm a pending DSL import identified by ``import_id``. @@ -119,7 +119,7 @@ class AppDslImportConfirmApi(Resource): return result, 200 -@openapi_ns.route("/apps//export") +@openapi_ns.route("/apps//dsl") class AppDslExportApi(Resource): """Export an app's current draft configuration as a DSL YAML string. @@ -145,6 +145,7 @@ class AppDslExportApi(Resource): try: data = AppDslService.export_dsl( app_model=app, + session=db.session(), include_secret=query.include_secret, workflow_id=query.workflow_id, ) @@ -153,7 +154,7 @@ class AppDslExportApi(Resource): return AppDslExportResponse(data=data), 200 -@openapi_ns.route("/apps//check-dependencies") +@openapi_ns.route("/apps//dependencies:check") class AppDslCheckDependenciesApi(Resource): """Check for leaked plugin dependencies after a DSL import. diff --git a/api/controllers/openapi/app_run.py b/api/controllers/openapi/app_run.py index 6074c7c0e02..772513ad417 100644 --- a/api/controllers/openapi/app_run.py +++ b/api/controllers/openapi/app_run.py @@ -1,4 +1,4 @@ -"""POST /openapi/v1/apps//run — mode-agnostic runner.""" +"""POST /openapi/v1/apps/:run — mode-agnostic runner.""" from __future__ import annotations @@ -138,7 +138,7 @@ _DISPATCH: dict[AppMode, Callable[[App, Any, AppRunRequest, Session], Any]] = { } -@openapi_ns.route("/apps//run") +@openapi_ns.route("/apps/:run") class AppRunApi(Resource): @auth_router.guard( scope=Scope.APPS_RUN, @@ -174,7 +174,7 @@ class AppRunApi(Resource): return helper.compact_generate_response(stream_obj) -@openapi_ns.route("/apps//tasks//stop") +@openapi_ns.route("/apps//tasks/:stop") class AppRunTaskStopApi(Resource): @auth_router.guard( scope=Scope.APPS_RUN, diff --git a/api/controllers/openapi/apps.py b/api/controllers/openapi/apps.py index 8d5c9670e77..882b55b7041 100644 --- a/api/controllers/openapi/apps.py +++ b/api/controllers/openapi/apps.py @@ -66,13 +66,13 @@ class AppReadResource(Resource): if is_uuid: # ``str(parsed_uuid)`` normalises to the canonical dashed form. - app = AppService.get_visible_app_by_id(db.session, str(parsed_uuid)) + app = AppService.get_visible_app_by_id(str(parsed_uuid), session=db.session()) if app is None: raise NotFound("app not found") else: if not workspace_id: raise UnprocessableEntity("workspace_id is required for name-based lookup") - matches = AppService.find_visible_apps_by_name(db.session, name=app_id, tenant_id=workspace_id) + matches = AppService.find_visible_apps_by_name(name=app_id, tenant_id=workspace_id, session=db.session()) if len(matches) == 0: raise NotFound("app not found") if len(matches) > 1: @@ -129,7 +129,7 @@ def build_app_describe_response(app: App, fields: set[str] | None) -> AppDescrib return AppDescribeResponse(info=info, parameters=parameters, input_schema=input_schema) -@openapi_ns.route("/apps//describe") +@openapi_ns.route("/apps/") class AppDescribeApi(AppReadResource): @auth_router.guard( scope=Scope.APPS_READ, @@ -177,7 +177,7 @@ class AppListApi(Resource): tenant_name: str | None = None if parsed_uuid is not None: - app: App | None = AppService.get_visible_app_by_id(db.session, str(parsed_uuid)) + app: App | None = AppService.get_visible_app_by_id(str(parsed_uuid), session=db.session()) if app is None or str(app.tenant_id) != workspace_id: return empty if not _is_listable(app): @@ -188,7 +188,7 @@ class AppListApi(Resource): str(app.id), str(app.maintainer) if app.maintainer else None, str(auth_data.account_id) ): return empty - tenant_name = TenantService.get_tenant_name(db.session, workspace_id) + tenant_name = TenantService.get_tenant_name(workspace_id, session=db.session()) item = AppListRow( id=str(app.id), name=app.name, @@ -215,13 +215,13 @@ class AppListApi(Resource): if apply_rbac_filter: access_filter.apply_to_params(params) - pagination = AppService().get_paginate_apps(str(auth_data.account_id), workspace_id, params, db.session) + pagination = AppService().get_paginate_apps(str(auth_data.account_id), workspace_id, params, db.session()) if pagination is None: return empty tenant_name = None if pagination.items: - tenant_name = TenantService.get_tenant_name(db.session, workspace_id) + tenant_name = TenantService.get_tenant_name(workspace_id, session=db.session()) items = [ AppListRow( diff --git a/api/controllers/openapi/apps_permitted_external.py b/api/controllers/openapi/apps_permitted_external.py index 718d3dbd169..353a1ec1cb3 100644 --- a/api/controllers/openapi/apps_permitted_external.py +++ b/api/controllers/openapi/apps_permitted_external.py @@ -55,10 +55,10 @@ class PermittedExternalAppsListApi(Resource): return env apps_by_id: dict[str, App] = { - str(a.id): a for a in AppService.find_visible_apps_by_ids(db.session, page_result.app_ids) + str(a.id): a for a in AppService.find_visible_apps_by_ids(page_result.app_ids, session=db.session()) } tenant_ids = list({str(a.tenant_id) for a in apps_by_id.values()}) - tenants_by_id = {str(t.id): t for t in TenantService.get_tenants_by_ids(db.session, tenant_ids)} + tenants_by_id = {str(t.id): t for t in TenantService.get_tenants_by_ids(tenant_ids, session=db.session())} items: list[AppListRow] = [] for app_id in page_result.app_ids: @@ -87,7 +87,7 @@ class PermittedExternalAppsListApi(Resource): return env -@openapi_ns.route("/permitted-external-apps//describe") +@openapi_ns.route("/permitted-external-apps/") class PermittedExternalAppDescribeApi(Resource): @auth_router.guard( scope=Scope.APPS_READ_PERMITTED_EXTERNAL, diff --git a/api/controllers/openapi/auth/prepare.py b/api/controllers/openapi/auth/prepare.py index 6704b27decc..96cf9a8858f 100644 --- a/api/controllers/openapi/auth/prepare.py +++ b/api/controllers/openapi/auth/prepare.py @@ -23,7 +23,7 @@ def load_app(data: AuthData) -> None: uuid.UUID(app_id) except ValueError: raise NotFound("app not found") - app = AppService.get_app_by_id(db.session, app_id) + app = AppService.get_app_by_id(app_id, session=db.session()) if not app or app.status != AppStatus.NORMAL: raise NotFound("app not found") data.app = app @@ -34,7 +34,7 @@ def load_tenant(data: AuthData) -> None: return if data.app is None: raise InternalServerError("pipeline_invariant_violated: app not loaded before load_tenant") - tenant = TenantService.get_tenant_by_id(db.session, str(data.app.tenant_id)) + tenant = TenantService.get_tenant_by_id(str(data.app.tenant_id), session=db.session()) if tenant is None or tenant.status == TenantStatus.ARCHIVE: raise Forbidden("workspace unavailable") data.tenant = tenant @@ -50,7 +50,7 @@ def load_tenant_from_request(data: AuthData) -> None: uuid.UUID(workspace_id) except ValueError: raise NotFound("workspace not found") - tenant = TenantService.get_tenant_by_id(db.session, workspace_id) + tenant = TenantService.get_tenant_by_id(workspace_id, session=db.session()) if tenant is None or tenant.status == TenantStatus.ARCHIVE: raise NotFound("workspace not found") data.tenant = tenant @@ -59,7 +59,7 @@ def load_tenant_from_request(data: AuthData) -> None: def load_account(data: AuthData) -> None: if data.caller is not None: return - account = AccountService.get_account_by_id(db.session, str(data.account_id)) + account = AccountService.get_account_by_id(str(data.account_id), session=db.session()) if account is None: raise Unauthorized("account not found") if data.tenant: @@ -75,7 +75,7 @@ def load_workspace_role(data: AuthData) -> None: return if data.caller is not None and getattr(data.caller, "status", None) != AccountStatus.ACTIVE: return - role = TenantService.get_account_role_in_tenant(db.session, str(data.account_id), str(data.tenant.id)) + role = TenantService.get_account_role_in_tenant(str(data.account_id), str(data.tenant.id), session=db.session()) if role is None: return data.tenant_role = role diff --git a/api/controllers/openapi/auth/verify.py b/api/controllers/openapi/auth/verify.py index b5f10f66b34..b6ef95e3ea3 100644 --- a/api/controllers/openapi/auth/verify.py +++ b/api/controllers/openapi/auth/verify.py @@ -82,7 +82,7 @@ def check_app_api_enabled(data: AuthData) -> None: def check_app_access(data: AuthData) -> None: if data.tenant is None: return - if not TenantService.account_belongs_to_tenant(db.session, data.account_id, data.tenant.id): + if not TenantService.account_belongs_to_tenant(data.account_id, data.tenant.id, session=db.session()): raise Forbidden("subject_no_app_access") @@ -127,5 +127,5 @@ def _resolve_user_id(data: AuthData) -> str | None: return str(data.account_id) if data.account_id is not None else None if data.external_identity is None: return None - account = AccountService.get_account_by_email(db.session, data.external_identity.email) + account = AccountService.get_account_by_email(data.external_identity.email, session=db.session()) return str(account.id) if account is not None else None diff --git a/api/controllers/openapi/files.py b/api/controllers/openapi/files.py index 7326a4a922e..3b3f68fa36b 100644 --- a/api/controllers/openapi/files.py +++ b/api/controllers/openapi/files.py @@ -1,4 +1,4 @@ -"""POST /openapi/v1/apps//files/upload — upload a file for use in app inputs.""" +"""POST /openapi/v1/apps//files — upload a file for use in app inputs.""" from __future__ import annotations @@ -26,7 +26,7 @@ from libs.oauth_bearer import Scope from services.file_service import FileService -@openapi_ns.route("/apps//files/upload") +@openapi_ns.route("/apps//files") class AppFileUploadApi(Resource): @openapi_ns.doc("upload_file_for_app_input") @openapi_ns.doc(description="Upload a file to use as an input variable when running the app") diff --git a/api/controllers/openapi/human_input_form.py b/api/controllers/openapi/human_input_form.py index 998dd669836..593887ffcfc 100644 --- a/api/controllers/openapi/human_input_form.py +++ b/api/controllers/openapi/human_input_form.py @@ -1,8 +1,8 @@ """ OpenAPI bearer-authed human input form endpoints. -GET /apps//form/human_input/ — fetch paused form definition -POST /apps//form/human_input/ — submit form response +GET /apps//human-input-forms/ — fetch paused form definition +POST /apps//human-input-forms/:submit — submit form response """ from __future__ import annotations @@ -60,7 +60,7 @@ def _ensure_form_is_allowed_for_openapi(form) -> None: raise RecipientSurfaceMismatch() -@openapi_ns.route("/apps//form/human_input/") +@openapi_ns.route("/apps//human-input-forms/") class OpenApiWorkflowHumanInputFormApi(Resource): @openapi_ns.response(200, "Form definition", openapi_ns.models[HumanInputFormDefinitionResponse.__name__]) @auth_router.guard( @@ -79,6 +79,9 @@ class OpenApiWorkflowHumanInputFormApi(Resource): service.ensure_form_active(form) return _jsonify_form_definition(form) + +@openapi_ns.route("/apps//human-input-forms/:submit") +class OpenApiWorkflowHumanInputFormSubmitApi(Resource): @auth_router.guard( scope=Scope.APPS_RUN, rbac=RBACRequirement(resource_type=RBACResourceScope.APP, scene=RBACPermission.APP_TEST_AND_RUN), diff --git a/api/controllers/openapi/oauth_device.py b/api/controllers/openapi/oauth_device.py index cee187daaf3..3ba5f2ee207 100644 --- a/api/controllers/openapi/oauth_device.py +++ b/api/controllers/openapi/oauth_device.py @@ -247,7 +247,6 @@ class DeviceApproveApi(Resource): raise BadRequest(description=str(e)) from None ttl_days = oauth_ttl_days(tenant_id=tenant) mint = mint_oauth_token( - db.session, redis_client, subject_email=account.email, subject_issuer=ACCOUNT_ISSUER_SENTINEL, @@ -256,6 +255,7 @@ class DeviceApproveApi(Resource): device_label=state.device_label, prefix=profile.prefix, ttl_days=ttl_days, + session=db.session(), ) poll_payload = _build_account_poll_payload(account, tenant, mint) @@ -342,7 +342,7 @@ def _audit_cross_ip_if_needed(state) -> None: def _build_account_poll_payload(account, tenant, mint) -> PollPayload: - rows = TenantService.get_workspaces_for_account(db.session, str(account.id)) + rows = TenantService.get_workspaces_for_account(str(account.id), session=db.session()) workspaces = [WorkspacePayload(id=str(t.id), name=t.name, role=getattr(m, "role", "")) for t, m in rows] # Prefer active session tenant → DB-flagged current join → first membership. default_ws_id = None diff --git a/api/controllers/openapi/oauth_device_sso.py b/api/controllers/openapi/oauth_device_sso.py index 79538f48059..fbf7bfa6295 100644 --- a/api/controllers/openapi/oauth_device_sso.py +++ b/api/controllers/openapi/oauth_device_sso.py @@ -194,7 +194,7 @@ def _sso_complete_impl(): if state.status is not DeviceFlowStatus.PENDING: return _device_error_redirect("sso_failed", user_code) - if AccountService.has_active_account_with_email(db.session, claims.email): + if AccountService.has_active_account_with_email(claims.email, session=db.session()): _emit_external_rejection_audit( state, _RejectedClaims(subject_email=claims.email, subject_issuer=claims.issuer), @@ -274,7 +274,7 @@ def approve_external(): if state.status is not DeviceFlowStatus.PENDING: raise Conflict("user_code_not_pending") - if AccountService.has_active_account_with_email(db.session, claims.subject_email): + if AccountService.has_active_account_with_email(claims.subject_email, session=db.session()): _emit_external_rejection_audit(state, claims, reason="email_belongs_to_dify_account") raise Forbidden("email_belongs_to_dify_account") @@ -293,7 +293,6 @@ def approve_external(): ttl_days = oauth_ttl_days(tenant_id=None) mint = mint_oauth_token( - db.session, redis_client, subject_email=claims.subject_email, subject_issuer=claims.subject_issuer, @@ -302,6 +301,7 @@ def approve_external(): device_label=state.device_label, prefix=profile.prefix, ttl_days=ttl_days, + session=db.session(), ) # SSO branch of the shared PollPayload contract: account/workspace diff --git a/api/controllers/openapi/workspaces.py b/api/controllers/openapi/workspaces.py index 49f8fb9656f..7f8eb0f7012 100644 --- a/api/controllers/openapi/workspaces.py +++ b/api/controllers/openapi/workspaces.py @@ -64,14 +64,14 @@ def _member_response(account: Account) -> MemberResponse: def _load_tenant(workspace_id: str) -> Tenant: - tenant = TenantService.get_tenant_by_id(db.session, workspace_id) + tenant = TenantService.get_tenant_by_id(workspace_id, session=db.session()) if tenant is None or tenant.status != TenantStatus.NORMAL: raise NotFound("workspace not found") return tenant def _load_account(account_id: object) -> Account: - account = AccountService.get_account_by_id(db.session, str(account_id)) if account_id else None + account = AccountService.get_account_by_id(str(account_id), session=db.session()) if account_id else None if account is None: raise RuntimeError("authenticated account_id has no Account row") return account @@ -94,7 +94,7 @@ class WorkspacesApi(Resource): @auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) @returns(200, WorkspaceListResponse, description="Workspace list") def get(self, *, auth_data: AuthData): - rows = TenantService.get_workspaces_for_account(db.session, str(auth_data.account_id)) + rows = TenantService.get_workspaces_for_account(str(auth_data.account_id), session=db.session()) return WorkspaceListResponse(workspaces=list(starmap(_workspace_summary, rows))) @@ -104,7 +104,7 @@ class WorkspaceByIdApi(Resource): @auth_router.guard(scope=Scope.WORKSPACE_READ, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT})) @returns(200, WorkspaceDetailResponse, description="Workspace detail") def get(self, workspace_id: str, *, auth_data: AuthData): - row = TenantService.find_workspace_for_account(db.session, str(auth_data.account_id), workspace_id) + row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session()) # 404 (not 403) on non-member so workspace IDs don't leak across tenants. if row is None: raise NotFound("workspace not found") @@ -113,7 +113,7 @@ class WorkspaceByIdApi(Resource): return _workspace_detail(tenant, membership) -@openapi_ns.route("/workspaces//switch") +@openapi_ns.route("/workspaces/:switch") class WorkspaceSwitchApi(Resource): """Server-side switch — equivalent to the console's POST /workspaces/switch. @@ -128,11 +128,11 @@ class WorkspaceSwitchApi(Resource): account = _load_account(auth_data.account_id) try: - TenantService.switch_tenant(account, workspace_id, session=db.session) + TenantService.switch_tenant(account, workspace_id, session=db.session()) except AccountNotLinkTenantError: raise NotFound("workspace not found") - row = TenantService.find_workspace_for_account(db.session, str(auth_data.account_id), workspace_id) + row = TenantService.find_workspace_for_account(str(auth_data.account_id), workspace_id, session=db.session()) if row is None: raise NotFound("workspace not found") tenant, membership = row @@ -152,7 +152,7 @@ class WorkspaceMembersApi(Resource): @accepts(query=MemberListQuery) def get(self, workspace_id: str, *, auth_data: AuthData, query: MemberListQuery): tenant = _load_tenant(workspace_id) - members = TenantService.get_tenant_members(tenant, session=db.session) + members = TenantService.get_tenant_members(tenant, session=db.session()) total = len(members) start = (query.page - 1) * query.limit page_items = members[start : start + query.limit] @@ -184,7 +184,7 @@ class WorkspaceMembersApi(Resource): language=None, role=body.role, inviter=inviter, - session=db.session, + session=db.session(), ) except AccountAlreadyInTenantError as exc: raise BadRequest(str(exc)) @@ -194,7 +194,7 @@ class WorkspaceMembersApi(Resource): raise BadRequest(str(exc)) normalized_email = body.email.lower() - member = AccountService.get_account_by_email_with_case_fallback(db.session, normalized_email) + member = AccountService.get_account_by_email_with_case_fallback(normalized_email, session=db.session()) if member is None: # invite_new_member just created or fetched this account. raise RuntimeError("invited member missing from DB after invite") @@ -212,11 +212,12 @@ class WorkspaceMembersApi(Resource): @openapi_ns.route("/workspaces//members/") class WorkspaceMemberApi(Resource): - """Remove a member. + """Remove a member (DELETE) or change a member's role (PATCH). Self-removal and owner-removal are explicitly rejected by the service layer (CannotOperateSelfError, NoPermissionError) — both surface as - 400 per the spec, with the service's message preserved. + 400 per the spec, with the service's message preserved. Owner can never be + assigned via PATCH (closed enum); admin cannot demote the standing owner. """ @auth_router.guard_workspace( @@ -228,12 +229,12 @@ class WorkspaceMemberApi(Resource): def delete(self, workspace_id: str, member_id: str, *, auth_data: AuthData): operator = _load_account(auth_data.account_id) tenant = _load_tenant(workspace_id) - member = AccountService.get_account_by_id(db.session, member_id) + member = AccountService.get_account_by_id(member_id, session=db.session()) if member is None: raise NotFound("member not found") try: - TenantService.remove_member_from_tenant(tenant, member, operator, session=db.session) + TenantService.remove_member_from_tenant(tenant, member, operator, session=db.session()) except CannotOperateSelfError as exc: raise BadRequest(str(exc)) except NoPermissionError as exc: @@ -243,15 +244,6 @@ class WorkspaceMemberApi(Resource): return MemberActionResponse() - -@openapi_ns.route("/workspaces//members//role") -class WorkspaceMemberRoleApi(Resource): - """Change a member's role. - - Owner cannot be assigned here (closed enum). Admin cannot demote the - standing owner (service NoPermissionError → 400, per spec). - """ - @auth_router.guard_workspace( scope=Scope.WORKSPACE_WRITE, allowed_token_types=frozenset({TokenType.OAUTH_ACCOUNT}), @@ -259,15 +251,15 @@ class WorkspaceMemberRoleApi(Resource): ) @returns(200, MemberActionResponse, description="Role updated") @accepts(body=MemberRoleUpdatePayload) - def put(self, workspace_id: str, member_id: str, *, auth_data: AuthData, body: MemberRoleUpdatePayload): + def patch(self, workspace_id: str, member_id: str, *, auth_data: AuthData, body: MemberRoleUpdatePayload): operator = _load_account(auth_data.account_id) tenant = _load_tenant(workspace_id) - member = AccountService.get_account_by_id(db.session, member_id) + member = AccountService.get_account_by_id(member_id, session=db.session()) if member is None: raise NotFound("member not found") try: - TenantService.update_member_role(tenant, member, body.role, operator, session=db.session) + TenantService.update_member_role(tenant, member, body.role, operator, session=db.session()) except CannotOperateSelfError as exc: raise BadRequest(str(exc)) except NoPermissionError as exc: diff --git a/api/controllers/service_api/app/annotation.py b/api/controllers/service_api/app/annotation.py index 0fbf8125ed9..126c67b5d61 100644 --- a/api/controllers/service_api/app/annotation.py +++ b/api/controllers/service_api/app/annotation.py @@ -201,7 +201,7 @@ class AnnotationListApi(Resource): query = AnnotationListQuery.model_validate(request.args.to_dict(flat=True)) annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app_model.id, query.page, query.limit, query.keyword + app_model.id, query.page, query.limit, query.keyword, session=db.session() ) annotation_models = TypeAdapter(list[Annotation]).validate_python(annotation_list, from_attributes=True) response = AnnotationList( @@ -243,7 +243,9 @@ class AnnotationListApi(Resource): """Create a new annotation.""" payload = AnnotationCreatePayload.model_validate(service_api_ns.payload or {}) insert_args: InsertAnnotationArgs = {"question": payload.question, "answer": payload.answer} - annotation = AppAnnotationService.insert_app_annotation_directly(insert_args, app_model.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + insert_args, app_model.id, session=db.session() + ) response = Annotation.model_validate(annotation, from_attributes=True) return response.model_dump(mode="json"), HTTPStatus.CREATED @@ -285,7 +287,7 @@ class AnnotationUpdateDeleteApi(Resource): update_args: UpdateAnnotationArgs = {"question": payload.question, "answer": payload.answer} app_ref = AppRefService.create_app_ref(app_model) annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id)) - annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session) + annotation = AppAnnotationService.update_app_annotation_directly(update_args, annotation_ref, db.session()) response = Annotation.model_validate(annotation, from_attributes=True) return response.model_dump(mode="json") @@ -316,5 +318,5 @@ class AnnotationUpdateDeleteApi(Resource): """Delete an annotation.""" app_ref = AppRefService.create_app_ref(app_model) annotation_ref = AppRefService.create_annotation_ref(app_ref, str(annotation_id)) - AppAnnotationService.delete_app_annotation(annotation_ref, db.session) + AppAnnotationService.delete_app_annotation(annotation_ref, db.session()) return "", 204 diff --git a/api/controllers/service_api/app/app.py b/api/controllers/service_api/app/app.py index 932ec71c769..60f83d7d070 100644 --- a/api/controllers/service_api/app/app.py +++ b/api/controllers/service_api/app/app.py @@ -2,19 +2,17 @@ from typing import Any, cast from flask_restx import Resource from pydantic import Field -from sqlalchemy import select +from controllers.common.agent_app_parameters import get_published_agent_app_feature_dict_and_user_input_form from controllers.common.fields import Parameters from controllers.common.schema import register_response_schema_models from controllers.service_api import service_api_ns -from controllers.service_api.app.error import AppUnavailableError +from controllers.service_api.app.error import AgentNotPublishedError, AppUnavailableError from controllers.service_api.wraps import validate_app_token from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict -from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError from extensions.ext_database import db from fields.base import ResponseModel -from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus -from models.agent_config_entities import AgentSoulConfig from models.model import App, AppMode from services.app_service import AppService @@ -35,38 +33,13 @@ register_response_schema_models(service_api_ns, Parameters, AppMetaResponse, App def _get_agent_app_feature_dict_and_user_input_form(app_model: App) -> tuple[dict[str, Any], list[dict[str, Any]]]: - app_model_config = app_model.app_model_config - features_dict = cast(dict[str, Any], app_model_config.to_dict()) if app_model_config is not None else {} - - agent = db.session.scalar( - select(Agent) - .where( - Agent.tenant_id == app_model.tenant_id, - Agent.app_id == app_model.id, - Agent.scope == AgentScope.ROSTER, - Agent.source == AgentSource.AGENT_APP, - Agent.status == AgentStatus.ACTIVE, - ) - .limit(1) - ) - if agent is None or not agent.active_config_snapshot_id: + try: + return get_published_agent_app_feature_dict_and_user_input_form(app_model) + except AgentAppNotPublishedError: + raise AgentNotPublishedError() + except AgentAppGeneratorError: raise AppUnavailableError() - snapshot = db.session.scalar( - select(AgentConfigSnapshot) - .where( - AgentConfigSnapshot.tenant_id == app_model.tenant_id, - AgentConfigSnapshot.agent_id == agent.id, - AgentConfigSnapshot.id == agent.active_config_snapshot_id, - ) - .limit(1) - ) - if snapshot is None: - raise AppUnavailableError() - - agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) - return features_dict, agent_app_variables_to_user_input_form(agent_soul.app_variables) - @service_api_ns.route("/parameters") class AppParameterApi(Resource): @@ -150,7 +123,7 @@ class AppMetaApi(Resource): Returns metadata about the application including configuration and settings. """ - return AppService().get_app_meta(app_model) + return AppService().get_app_meta(app_model, session=db.session()) @service_api_ns.route("/info") diff --git a/api/controllers/service_api/app/audio.py b/api/controllers/service_api/app/audio.py index 53b31c8e6c4..68ab5f31ea5 100644 --- a/api/controllers/service_api/app/audio.py +++ b/api/controllers/service_api/app/audio.py @@ -188,7 +188,7 @@ class TextApi(Resource): ) response = AudioService.transcript_tts( app_model=app_model, - session=db.session, + session=db.session(), text=text, voice=voice, end_user=end_user.external_user_id, diff --git a/api/controllers/service_api/app/completion.py b/api/controllers/service_api/app/completion.py index 900d46a0f0f..c240c7d85af 100644 --- a/api/controllers/service_api/app/completion.py +++ b/api/controllers/service_api/app/completion.py @@ -15,6 +15,7 @@ from controllers.common.schema import register_response_schema_models, register_ from controllers.console.app.wraps import with_session from controllers.service_api import service_api_ns from controllers.service_api.app.error import ( + AgentNotPublishedError, AppUnavailableError, CompletionRequestError, ConversationCompletedError, @@ -31,6 +32,7 @@ from controllers.service_api.schema import ( ) from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError +from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.app.entities.app_invoke_entities import InvokeFrom from core.errors.error import ( ModelCurrentlyNotSupportError, @@ -248,6 +250,8 @@ class CompletionApi(Resource): except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") raise AppUnavailableError() + except AgentAppNotPublishedError: + raise AgentNotPublishedError() except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) except QuotaExceededError: @@ -403,6 +407,8 @@ class ChatApi(Resource): except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") raise AppUnavailableError() + except AgentAppNotPublishedError: + raise AgentNotPublishedError() except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) except QuotaExceededError: diff --git a/api/controllers/service_api/app/conversation.py b/api/controllers/service_api/app/conversation.py index 9b5533ea07a..a395dcb93fc 100644 --- a/api/controllers/service_api/app/conversation.py +++ b/api/controllers/service_api/app/conversation.py @@ -249,7 +249,7 @@ class ConversationDetailApi(Resource): conversation_id = str(c_id) try: - ConversationService.delete(app_model, conversation_id, end_user) + ConversationService.delete(app_model, conversation_id, end_user, session=db.session()) except services.errors.conversation.ConversationNotExistsError: raise NotFound("Conversation Not Exists.") return "", 204 @@ -299,7 +299,7 @@ class ConversationRenameApi(Resource): try: conversation = ConversationService.rename( - app_model, conversation_id, end_user, payload.name, payload.auto_generate + app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=db.session() ) return ( TypeAdapter(SimpleConversation) @@ -356,7 +356,13 @@ class ConversationVariablesApi(Resource): try: pagination = ConversationService.get_conversational_variable( - app_model, conversation_id, end_user, query_args.limit, last_id, query_args.variable_name + app_model, + conversation_id, + end_user, + query_args.limit, + last_id, + query_args.variable_name, + session=db.session(), ) return ConversationVariableInfiniteScrollPaginationResponse.model_validate( pagination, from_attributes=True @@ -417,7 +423,7 @@ class ConversationVariableDetailApi(Resource): try: variable = ConversationService.update_conversation_variable( - app_model, conversation_id, variable_id_str, end_user, payload.value + app_model, conversation_id, variable_id_str, end_user, payload.value, session=db.session() ) return ConversationVariableResponse.model_validate(variable, from_attributes=True).model_dump(mode="json") except services.errors.conversation.ConversationNotExistsError: diff --git a/api/controllers/service_api/app/error.py b/api/controllers/service_api/app/error.py index 0e04a04cb24..710fc7878fb 100644 --- a/api/controllers/service_api/app/error.py +++ b/api/controllers/service_api/app/error.py @@ -7,6 +7,12 @@ class AppUnavailableError(BaseHTTPException): code = 400 +class AgentNotPublishedError(BaseHTTPException): + error_code = "agent_not_published" + description = "Agent has not been published. Please publish the Agent before using the API." + code = 400 + + class NotCompletionAppError(BaseHTTPException): error_code = "not_completion_app" description = "Please check if your Completion app mode matches the right API route." diff --git a/api/controllers/service_api/app/message.py b/api/controllers/service_api/app/message.py index 18d1c5d3254..3acb2c74872 100644 --- a/api/controllers/service_api/app/message.py +++ b/api/controllers/service_api/app/message.py @@ -15,6 +15,7 @@ from controllers.service_api.app.error import NotChatAppError from controllers.service_api.schema import expect_with_user from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token from core.app.entities.app_invoke_entities import InvokeFrom +from extensions.ext_database import db from fields.base import ResponseModel from fields.conversation_fields import ResultResponse from fields.message_fields import MessageInfiniteScrollPagination, MessageListItem @@ -109,7 +110,7 @@ class MessageListApi(Resource): try: pagination = MessageService.pagination_by_first_id( - app_model, end_user, conversation_id, first_id, query_args.limit + app_model, end_user, conversation_id, first_id, query_args.limit, session=db.session() ) adapter = TypeAdapter(MessageListItem) items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data] @@ -167,6 +168,7 @@ class MessageFeedbackApi(Resource): user=end_user, rating=FeedbackRating(payload.rating) if payload.rating else None, content=payload.content, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -208,7 +210,9 @@ class AppGetFeedbacksApi(Resource): Returns paginated list of all feedback submitted for messages in this app. """ query_args = FeedbackListQuery.model_validate(request.args.to_dict()) - feedbacks = MessageService.get_all_messages_feedbacks(app_model, page=query_args.page, limit=query_args.limit) + feedbacks = MessageService.get_all_messages_feedbacks( + app_model, page=query_args.page, limit=query_args.limit, session=db.session() + ) return {"data": feedbacks} @@ -258,7 +262,11 @@ class MessageSuggestedApi(Resource): try: questions = MessageService.get_suggested_questions_after_answer( - app_model=app_model, user=end_user, message_id=message_id_str, invoke_from=InvokeFrom.SERVICE_API + app_model=app_model, + user=end_user, + message_id=message_id_str, + invoke_from=InvokeFrom.SERVICE_API, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message Not Exists.") diff --git a/api/controllers/service_api/dataset/dataset.py b/api/controllers/service_api/dataset/dataset.py index 56836f56895..66085ca0642 100644 --- a/api/controllers/service_api/dataset/dataset.py +++ b/api/controllers/service_api/dataset/dataset.py @@ -414,7 +414,7 @@ class DatasetListApi(DatasetApiResource): datasets, total = DatasetService.get_datasets( query.page, query.limit, - db.session, + db.session(), tenant_id, current_user, query.keyword, @@ -565,11 +565,11 @@ class DatasetApi(DatasetApiResource): ) def get(self, _, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) data = _dump_service_dataset_detail(dataset) @@ -601,7 +601,7 @@ class DatasetApi(DatasetApiResource): retrieval_model_dict["search_method"] = "keyword_search" if data.get("permission") == "partial_members": - part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session) + part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session()) data.update({"partial_member_list": part_users_list}) return _dump_service_dataset_with_partial_members(data), 200 @@ -640,7 +640,7 @@ class DatasetApi(DatasetApiResource): @with_session def patch(self, session: Session, _, dataset_id: UUID): dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") @@ -681,10 +681,10 @@ class DatasetApi(DatasetApiResource): dataset, str(payload.permission) if payload.permission else None, payload.partial_member_list, - db.session, + session=db.session(), ) - dataset = DatasetService.update_dataset(session, dataset_id_str, update_data, current_user) + dataset = DatasetService.update_dataset(dataset_id_str, update_data, current_user, session=session) if dataset is None: raise NotFound("Dataset not found.") @@ -695,13 +695,13 @@ class DatasetApi(DatasetApiResource): if payload.partial_member_list and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM: DatasetPermissionService.update_partial_member_list( - tenant_id, dataset_id_str, payload.partial_member_list, db.session + tenant_id, dataset_id_str, payload.partial_member_list, db.session() ) # clear partial member list when permission is only_me or all_team_members elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}: - DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session) + DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session()) - partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session) + partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session()) result_data.update({"partial_member_list": partial_member_list}) return _dump_service_dataset_with_partial_members(result_data), 200 @@ -754,8 +754,8 @@ class DatasetApi(DatasetApiResource): dataset_id_str = str(dataset_id) try: - if DatasetService.delete_dataset(dataset_id_str, current_user, db.session): - DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session) + if DatasetService.delete_dataset(dataset_id_str, current_user, db.session()): + DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session()) return "", 204 else: raise NotFound("Dataset not found.") @@ -820,14 +820,14 @@ class DocumentStatusApi(DatasetApiResource): InvalidActionError: If the action is invalid or cannot be performed. """ dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") # Check user's permission try: - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) @@ -839,7 +839,7 @@ class DocumentStatusApi(DatasetApiResource): document_ids = data.get("document_ids", []) try: - DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session) + DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session()) except services.errors.document.DocumentIndexingError as e: raise InvalidActionError(str(e)) except ValueError as e: @@ -876,7 +876,7 @@ class DatasetTagsApi(DatasetApiResource): assert isinstance(current_user, Account) cid = current_user.current_tenant_id assert cid is not None - tags = TagService.get_tags(db.session(), "knowledge", cid) + tags = TagService.get_tags("knowledge", cid, session=db.session()) return dump_response(KnowledgeTagListResponse, tags), 200 @service_api_ns.doc( @@ -909,7 +909,7 @@ class DatasetTagsApi(DatasetApiResource): raise Forbidden() payload = TagCreatePayload.model_validate(service_api_ns.payload or {}) - tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=TagType.KNOWLEDGE), db.session) + tag = TagService.save_tags(SaveTagPayload(name=payload.name, type=TagType.KNOWLEDGE), db.session()) response = dump_response( KnowledgeTagResponse, @@ -948,10 +948,10 @@ class DatasetTagsApi(DatasetApiResource): payload = TagUpdatePayload.model_validate(service_api_ns.payload or {}) tag_id = payload.tag_id tag = TagService.update_tags( - UpdateTagServicePayload(name=payload.name), tag_id, db.session, tag_type=TagType.KNOWLEDGE + UpdateTagServicePayload(name=payload.name), tag_id, db.session(), tag_type=TagType.KNOWLEDGE ) - binding_count = TagService.get_tag_binding_count(tag_id, db.session, tag_type=TagType.KNOWLEDGE) + binding_count = TagService.get_tag_binding_count(tag_id, db.session(), tag_type=TagType.KNOWLEDGE) response = dump_response( KnowledgeTagResponse, @@ -981,7 +981,7 @@ class DatasetTagsApi(DatasetApiResource): def delete(self, _): """Delete a knowledge type tag.""" payload = TagDeletePayload.model_validate(service_api_ns.payload or {}) - TagService.delete_tag(payload.tag_id, db.session, tag_type=TagType.KNOWLEDGE) + TagService.delete_tag(payload.tag_id, db.session(), tag_type=TagType.KNOWLEDGE) return "", 204 @@ -1015,7 +1015,7 @@ class DatasetTagBindingApi(DatasetApiResource): payload = TagBindingPayload.model_validate(service_api_ns.payload or {}) TagService.save_tag_binding( TagBindingCreatePayload(tag_ids=payload.tag_ids, target_id=payload.target_id, type=TagType.KNOWLEDGE), - db.session, + db.session(), ) return "", 204 @@ -1050,7 +1050,7 @@ class DatasetTagUnbindingApi(DatasetApiResource): payload = TagUnbindingPayload.model_validate(service_api_ns.payload or {}) TagService.delete_tag_binding( TagBindingDeletePayload(tag_ids=payload.tag_ids, target_id=payload.target_id, type=TagType.KNOWLEDGE), - db.session, + db.session(), ) return "", 204 @@ -1086,7 +1086,7 @@ class DatasetTagsBindingStatusApi(DatasetApiResource): assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None tags = TagService.get_tags_by_target_id( - "knowledge", current_user.current_tenant_id, str(dataset_id), db.session + "knowledge", current_user.current_tenant_id, str(dataset_id), db.session() ) tags_list = [{"id": tag.id, "name": tag.name} for tag in tags] return dump_response(DatasetBoundTagListResponse, {"data": tags_list, "total": len(tags)}), 200 diff --git a/api/controllers/service_api/dataset/document.py b/api/controllers/service_api/dataset/document.py index 4c083d3d50f..5e5919a7048 100644 --- a/api/controllers/service_api/dataset/document.py +++ b/api/controllers/service_api/dataset/document.py @@ -401,7 +401,7 @@ def _create_document_by_text(tenant_id: str, dataset_id: UUID) -> tuple[Mapping[ account=current_user, dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None, created_from="api", - session=db.session, + session=db.session(), ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -461,7 +461,7 @@ def _update_document_by_text(tenant_id: str, dataset_id: UUID, document_id: UUID account=current_user, dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None, created_from="api", - session=db.session, + session=db.session(), ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -759,7 +759,7 @@ class DocumentAddByFileApi(DatasetApiResource): account=dataset.created_by_account, dataset_process_rule=dataset_process_rule, created_from="api", - session=db.session, + session=db.session(), ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -836,7 +836,7 @@ def _update_document_by_file(tenant_id: str, dataset_id: UUID, document_id: UUID account=dataset.created_by_account, dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None, created_from="api", - session=db.session, + session=db.session(), ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -955,6 +955,7 @@ class DocumentListApi(DatasetApiResource): documents=documents, dataset=dataset, tenant_id=tenant_id, + session=db.session(), ) response = { @@ -1007,7 +1008,7 @@ class DocumentBatchDownloadZipApi(DatasetApiResource): document_ids=[str(document_id) for document_id in payload.document_ids], tenant_id=str(tenant_id), current_user=current_user, - session=db.session, + session=db.session(), ) with ExitStack() as stack: @@ -1064,7 +1065,7 @@ class DocumentIndexingStatusApi(DatasetApiResource): if not dataset: raise NotFound("Dataset not found.") # get documents - documents = DocumentService.get_batch_documents(dataset_id_str, batch, db.session) + documents = DocumentService.get_batch_documents(dataset_id_str, batch, db.session()) if not documents: raise NotFound("Documents not found.") documents_status = [] @@ -1140,7 +1141,7 @@ class DocumentDownloadApi(DatasetApiResource): @cloud_edition_billing_rate_limit_check("knowledge", "dataset") def get(self, tenant_id, dataset_id: UUID, document_id: UUID): dataset = self.get_dataset(str(dataset_id), str(tenant_id)) - document = DocumentService.get_document(dataset.id, str(document_id), session=db.session) + document = DocumentService.get_document(dataset.id, str(document_id), session=db.session()) if not document: raise NotFound("Document not found.") @@ -1148,7 +1149,7 @@ class DocumentDownloadApi(DatasetApiResource): if document.tenant_id != str(tenant_id): raise Forbidden("No permission.") - return {"url": DocumentService.get_document_download_url(document, db.session)} + return {"url": DocumentService.get_document_download_url(document, db.session())} @service_api_ns.route("/datasets//documents/") @@ -1196,7 +1197,7 @@ class DocumentApi(DatasetApiResource): dataset = self.get_dataset(dataset_id_str, tenant_id) - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -1216,12 +1217,13 @@ class DocumentApi(DatasetApiResource): document_id=document_id_str, dataset_id=dataset_id_str, tenant_id=tenant_id, + session=db.session(), ) if metadata == "only": response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details} elif metadata == "without": - dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session) + dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session()) document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {} data_source_info = document.data_source_detail_dict response = { @@ -1256,7 +1258,7 @@ class DocumentApi(DatasetApiResource): "need_summary": document.need_summary if document.need_summary is not None else False, } else: - dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session) + dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session()) document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {} data_source_info = document.data_source_detail_dict response = { @@ -1351,7 +1353,7 @@ class DocumentApi(DatasetApiResource): if not dataset: raise ValueError("Dataset does not exist.") - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) # 404 if document not found if document is None: @@ -1363,7 +1365,7 @@ class DocumentApi(DatasetApiResource): try: # delete document - DocumentService.delete_document(document, db.session) + DocumentService.delete_document(document, db.session()) except services.errors.document.DocumentIndexingError: raise DocumentIndexingError("Cannot delete document during indexing.") diff --git a/api/controllers/service_api/dataset/metadata.py b/api/controllers/service_api/dataset/metadata.py index aec3b06a91e..1d793583cc2 100644 --- a/api/controllers/service_api/dataset/metadata.py +++ b/api/controllers/service_api/dataset/metadata.py @@ -81,12 +81,12 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource): metadata_args = MetadataArgs.model_validate(service_api_ns.payload or {}) dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) - metadata = MetadataService.create_metadata(db.session(), dataset_id_str, metadata_args) + metadata = MetadataService.create_metadata(dataset_id_str, metadata_args, session=db.session()) return dump_response(DatasetMetadataResponse, metadata), 201 @service_api_ns.doc( @@ -116,10 +116,10 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource): def get(self, tenant_id, dataset_id: UUID): """Get all metadata for a dataset.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - metadata = MetadataService.get_dataset_metadatas(db.session(), dataset) + metadata = MetadataService.get_dataset_metadatas(dataset, session=db.session()) return dump_response(DatasetMetadataListResponse, metadata), 200 @@ -154,12 +154,14 @@ class DatasetMetadataServiceApi(DatasetApiResource): dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) - metadata = MetadataService.update_metadata_name(db.session(), dataset_id_str, metadata_id_str, payload.name) + metadata = MetadataService.update_metadata_name( + dataset_id_str, metadata_id_str, payload.name, session=db.session() + ) return dump_response(DatasetMetadataResponse, metadata), 200 @service_api_ns.doc( @@ -189,12 +191,12 @@ class DatasetMetadataServiceApi(DatasetApiResource): """Delete metadata.""" dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) - MetadataService.delete_metadata(db.session(), dataset_id_str, metadata_id_str) + MetadataService.delete_metadata(dataset_id_str, metadata_id_str, session=db.session()) return "", 204 @@ -257,16 +259,16 @@ class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource): def post(self, tenant_id, dataset_id: UUID, action: Literal["enable", "disable"]): """Enable or disable built-in metadata field.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) match action: case "enable": - MetadataService.enable_built_in_field(db.session(), dataset) + MetadataService.enable_built_in_field(dataset, session=db.session()) case "disable": - MetadataService.disable_built_in_field(db.session(), dataset) + MetadataService.disable_built_in_field(dataset, session=db.session()) return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200 @@ -303,13 +305,13 @@ class DocumentMetadataEditServiceApi(DatasetApiResource): def post(self, tenant_id, dataset_id: UUID): """Update metadata for multiple documents.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, db.session) + dataset = DatasetService.get_dataset(dataset_id_str, db.session()) if dataset is None: raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, db.session) + DatasetService.check_dataset_permission(dataset, current_user, db.session()) metadata_args = MetadataOperationData.model_validate(service_api_ns.payload or {}) - MetadataService.update_documents_metadata(db.session(), dataset, metadata_args) + MetadataService.update_documents_metadata(dataset, metadata_args, session=db.session()) return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200 diff --git a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py index 8c4063398bc..35f3a4c01a0 100644 --- a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py @@ -1,9 +1,10 @@ from collections.abc import Generator +from datetime import datetime from typing import Any from uuid import UUID from flask import request -from pydantic import BaseModel, Field, RootModel +from pydantic import BaseModel, Field, RootModel, field_validator from sqlalchemy import select from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound @@ -13,24 +14,20 @@ from controllers.common.errors import FilenameNotExistsError, NoFileUploadedErro from controllers.common.fields import GeneratedAppResponse from controllers.common.schema import ( query_params_from_model, + query_params_from_request, register_response_schema_models, register_schema_model, - register_schema_models, ) from controllers.console.app.wraps import with_session from controllers.service_api import service_api_ns from controllers.service_api.dataset.error import PipelineRunError -from controllers.service_api.dataset.rag_pipeline.serializers import serialize_upload_file -from controllers.service_api.schema import ( - event_stream_response, - json_or_event_stream_response, - multipart_file_params, -) +from controllers.service_api.schema import event_stream_response, json_or_event_stream_response, multipart_file_params from controllers.service_api.wraps import DatasetApiResource from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from fields.base import ResponseModel from libs import helper +from libs.helper import dump_response from libs.login import current_user from models import Account from models.dataset import Dataset, Pipeline @@ -84,7 +81,7 @@ class DatasourcePluginResponse(ResponseModel): datasource_type: str | None = None title: str | None = None user_input_variables: list[dict[str, Any]] = Field(default_factory=list) - credentials: list[DatasourceCredentialInfoResponse] + credentials: list[DatasourceCredentialInfoResponse] = Field(default_factory=list) class DatasourcePluginListResponse(RootModel[list[DatasourcePluginResponse]]): @@ -100,14 +97,22 @@ class PipelineUploadFileResponse(ResponseModel): created_by: str created_at: str | None = None + @field_validator("created_at", mode="before") + @classmethod + def _normalize_created_at(cls, value: datetime | str | None) -> str | None: + if isinstance(value, datetime): + return value.isoformat() + return value + register_schema_model(service_api_ns, DatasourceNodeRunPayload) +register_schema_model(service_api_ns, DatasourcePluginsQuery) register_schema_model(service_api_ns, PipelineRunApiEntity) -register_schema_models(service_api_ns, DatasourcePluginsQuery) register_response_schema_models( service_api_ns, + DatasourceCredentialInfoResponse, + DatasourcePluginResponse, DatasourcePluginListResponse, - GeneratedAppResponse, PipelineUploadFileResponse, ) @@ -119,8 +124,8 @@ class DatasourcePluginsApi(DatasetApiResource): @service_api_ns.doc( summary="List Datasource Plugins", description=( - "List the datasource nodes configured in the knowledge pipeline. Each node includes the " - "plugin it uses plus the metadata needed to run it." + "List the datasource nodes configured in the knowledge pipeline. Each node includes the plugin it uses " + "plus the metadata needed to run it." ), tags=["Knowledge Pipeline"], responses={ @@ -152,14 +157,13 @@ class DatasourcePluginsApi(DatasetApiResource): if not dataset: raise NotFound("Dataset not found.") - # Get query parameter to determine published or draft - is_published: bool = request.args.get("is_published", default=True, type=bool) + query = query_params_from_request(DatasourcePluginsQuery) - rag_pipeline_service: RagPipelineService = RagPipelineService() + rag_pipeline_service = RagPipelineService(db.session()) datasource_plugins: list[dict[Any, Any]] = rag_pipeline_service.get_datasource_plugins( - tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=is_published + tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=query.is_published ) - return datasource_plugins, 200 + return dump_response(DatasourcePluginListResponse, datasource_plugins), 200 @service_api_ns.route("/datasets//pipeline/datasource/nodes//run") @@ -169,8 +173,8 @@ class DatasourceNodeRunApi(DatasetApiResource): @service_api_ns.doc( summary="Run Datasource Node", description=( - "Execute a single datasource node within the knowledge pipeline. Returns a streaming " - "response with the node execution results." + "Execute a single datasource node within the knowledge pipeline. Returns a streaming response with the " + "node execution results." ), tags=["Knowledge Pipeline"], responses={ @@ -189,11 +193,6 @@ class DatasourceNodeRunApi(DatasetApiResource): } ) @service_api_ns.expect(service_api_ns.models[DatasourceNodeRunPayload.__name__]) - @service_api_ns.response( - 200, - "Datasource node run successfully", - service_api_ns.models[GeneratedAppResponse.__name__], - ) def post(self, tenant_id: str, dataset_id: UUID, node_id: str): """Resource for getting datasource plugins.""" dataset_id_str = str(dataset_id) @@ -205,15 +204,16 @@ class DatasourceNodeRunApi(DatasetApiResource): payload = DatasourceNodeRunPayload.model_validate(service_api_ns.payload or {}) assert isinstance(current_user, Account) - rag_pipeline_service: RagPipelineService = RagPipelineService() + rag_pipeline_service: RagPipelineService = RagPipelineService(db.session()) pipeline: Pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str) datasource_node_run_api_entity = DatasourceNodeRunApiEntity.model_validate( { **payload.model_dump(exclude_none=True), - "pipeline_id": str(pipeline.id), + "pipeline_id": pipeline.id, "node_id": node_id, } ) + # response-contract:ignore compact_generate_response return helper.compact_generate_response( PipelineGenerator.convert_to_event_stream( rag_pipeline_service.run_datasource_workflow_node( @@ -236,8 +236,8 @@ class PipelineRunApi(DatasetApiResource): @service_api_ns.doc( summary="Run Pipeline", description=( - "Execute the full knowledge pipeline for a knowledge base. Supports both streaming and " - "blocking response modes." + "Execute the full knowledge pipeline for a knowledge base. Supports both streaming and blocking response " + "modes." ), tags=["Knowledge Pipeline"], responses={ @@ -272,7 +272,7 @@ class PipelineRunApi(DatasetApiResource): dataset_id_str = str(dataset_id) # Verify dataset ownership stmt = select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id_str) - dataset = db.session.scalar(stmt) + dataset = session.scalar(stmt) if not dataset: raise NotFound("Dataset not found.") @@ -281,8 +281,8 @@ class PipelineRunApi(DatasetApiResource): if not isinstance(current_user, Account): raise Forbidden() - rag_pipeline_service: RagPipelineService = RagPipelineService() - pipeline: Pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str) + rag_pipeline_service = RagPipelineService(session) + pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str) try: response: dict[Any, Any] | Generator[str, Any, None] = PipelineGenerateService.generate( session=session, @@ -293,6 +293,7 @@ class PipelineRunApi(DatasetApiResource): streaming=payload.response_mode == "streaming", ) + # response-contract:ignore compact_generate_response return helper.compact_generate_response(response) except Exception as ex: raise PipelineRunError(description=str(ex)) @@ -368,4 +369,4 @@ class KnowledgebasePipelineFileUploadApi(DatasetApiResource): except services.errors.file.UnsupportedFileTypeError: raise UnsupportedFileTypeError() - return serialize_upload_file(upload_file), 201 + return dump_response(PipelineUploadFileResponse, upload_file), 201 diff --git a/api/controllers/service_api/dataset/rag_pipeline/serializers.py b/api/controllers/service_api/dataset/rag_pipeline/serializers.py deleted file mode 100644 index a5e8484037e..00000000000 --- a/api/controllers/service_api/dataset/rag_pipeline/serializers.py +++ /dev/null @@ -1,32 +0,0 @@ -""" -Serialization helpers for Service API knowledge pipeline endpoints. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, TypedDict - -if TYPE_CHECKING: - from models.model import UploadFile - - -class UploadFileDict(TypedDict): - id: str - name: str - size: int - extension: str - mime_type: str | None - created_by: str - created_at: str | None - - -def serialize_upload_file(upload_file: UploadFile) -> UploadFileDict: - return { - "id": upload_file.id, - "name": upload_file.name, - "size": upload_file.size, - "extension": upload_file.extension, - "mime_type": upload_file.mime_type, - "created_by": upload_file.created_by, - "created_at": upload_file.created_at.isoformat() if upload_file.created_at else None, - } diff --git a/api/controllers/service_api/dataset/segment.py b/api/controllers/service_api/dataset/segment.py index 41fbc709fdd..e911c454c9e 100644 --- a/api/controllers/service_api/dataset/segment.py +++ b/api/controllers/service_api/dataset/segment.py @@ -137,7 +137,7 @@ def _get_segment_for_document( raise NotFound("Document not found.") segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id) - segment = SegmentService.get_segment_by_ref(segment_ref) + segment = SegmentService.get_segment_by_ref(segment_ref, db.session()) if not segment: raise NotFound("Segment not found.") return segment_ref, segment @@ -191,7 +191,7 @@ class SegmentApi(DatasetApiResource): raise NotFound("Dataset not found.") document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") if document.indexing_status != "completed": @@ -227,13 +227,13 @@ class SegmentApi(DatasetApiResource): for args_item in segment_items: SegmentService.segment_create_args_validate(args_item, document) segments = cast( - list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, db.session) + list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, db.session()) ) segment_ids = [segment.id for segment in segments] summaries: dict[str, str | None] = {} if segment_ids: summary_records = SummaryIndexService.get_segments_summaries( - segment_ids=segment_ids, dataset_id=dataset_id_str + segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session() ) summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()} response = { @@ -285,7 +285,7 @@ class SegmentApi(DatasetApiResource): raise NotFound("Dataset not found.") document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") # check embedding model setting @@ -317,7 +317,7 @@ class SegmentApi(DatasetApiResource): summaries: dict[str, str | None] = {} if segment_ids: summary_records = SummaryIndexService.get_segments_summaries( - segment_ids=segment_ids, dataset_id=dataset_id_str + segment_ids=segment_ids, dataset_id=dataset_id_str, session=db.session() ) summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()} @@ -367,12 +367,12 @@ class DatasetSegmentApi(DatasetApiResource): DatasetService.check_dataset_model_setting(dataset) document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") segment_id_str = str(segment_id) _, segment = _get_segment_for_document(dataset, document, segment_id_str) - SegmentService.delete_segment(segment, document, dataset, db.session) + SegmentService.delete_segment(segment, document, dataset, db.session()) return "", 204 @service_api_ns.doc( @@ -410,7 +410,7 @@ class DatasetSegmentApi(DatasetApiResource): DatasetService.check_dataset_model_setting(dataset) document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: @@ -434,8 +434,10 @@ class DatasetSegmentApi(DatasetApiResource): payload = SegmentUpdatePayload.model_validate(service_api_ns.payload or {}) - updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, db.session) - summary = SummaryIndexService.get_segment_summary(segment_id=updated_segment.id, dataset_id=dataset_id_str) + updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, db.session()) + summary = SummaryIndexService.get_segment_summary( + segment_id=updated_segment.id, dataset_id=dataset_id_str, session=db.session() + ) response = { "data": segment_response_with_summary(updated_segment, summary.summary_content if summary else None), "doc_form": document.doc_form, @@ -481,13 +483,15 @@ class DatasetSegmentApi(DatasetApiResource): DatasetService.check_dataset_model_setting(dataset) document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") segment_id_str = str(segment_id) _, segment = _get_segment_for_document(dataset, document, segment_id_str) - summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id_str) + summary = SummaryIndexService.get_segment_summary( + segment_id=segment.id, dataset_id=dataset_id_str, session=db.session() + ) response = { "data": segment_response_with_summary(segment, summary.summary_content if summary else None), "doc_form": document.doc_form, @@ -542,7 +546,7 @@ class ChildChunkApi(DatasetApiResource): document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -570,7 +574,7 @@ class ChildChunkApi(DatasetApiResource): payload = ChildChunkCreatePayload.model_validate(service_api_ns.payload or {}) try: - child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session) + child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session()) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) @@ -613,7 +617,7 @@ class ChildChunkApi(DatasetApiResource): document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -680,7 +684,7 @@ class DatasetChildChunkApi(DatasetApiResource): document_id_str = str(document_id) # check document - document = DocumentService.get_document(dataset.id, document_id_str, session=db.session) + document = DocumentService.get_document(dataset.id, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -689,12 +693,12 @@ class DatasetChildChunkApi(DatasetApiResource): child_chunk_id_str = str(child_chunk_id) # check child chunk - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref) + child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session()) if not child_chunk: raise NotFound("Child chunk not found.") try: - SegmentService.delete_child_chunk(child_chunk, dataset, db.session) + SegmentService.delete_child_chunk(child_chunk, dataset, db.session()) except ChildChunkDeleteIndexServiceError as e: raise ChildChunkDeleteIndexError(str(e)) @@ -741,7 +745,7 @@ class DatasetChildChunkApi(DatasetApiResource): document_id_str = str(document_id) # get document - document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session) + document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session()) if not document: raise NotFound("Document not found.") @@ -750,7 +754,7 @@ class DatasetChildChunkApi(DatasetApiResource): child_chunk_id_str = str(child_chunk_id) # get child chunk - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref) + child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, db.session()) if not child_chunk: raise NotFound("Child chunk not found.") @@ -759,7 +763,7 @@ class DatasetChildChunkApi(DatasetApiResource): try: child_chunk = SegmentService.update_child_chunk( - payload.content, child_chunk, segment, document, dataset, db.session + payload.content, child_chunk, segment, document, dataset, db.session() ) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) diff --git a/api/controllers/web/app.py b/api/controllers/web/app.py index d5722faf00d..6804d072ef0 100644 --- a/api/controllers/web/app.py +++ b/api/controllers/web/app.py @@ -8,8 +8,11 @@ from werkzeug.exceptions import Unauthorized from constants import HEADER_NAME_APP_CODE from controllers.common import fields +from controllers.common.agent_app_parameters import get_published_agent_app_feature_dict_and_user_input_form from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError +from extensions.ext_database import db from libs.passport import PassportService from libs.token import extract_webapp_passport from models.model import App, AppMode, EndUser @@ -19,7 +22,7 @@ from services.feature_service import FeatureService from services.webapp_auth_service import WebAppAuthService from . import web_ns -from .error import AppUnavailableError +from .error import AgentNotPublishedError, AppUnavailableError from .wraps import WebApiResource logger = logging.getLogger(__name__) @@ -74,12 +77,21 @@ class AppParameterApi(WebApiResource): @web_ns.response(200, "Success", web_ns.models[fields.Parameters.__name__]) def get(self, app_model: App, end_user: EndUser): """Retrieve app parameters.""" - if app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: + features_dict: dict[str, Any] + user_input_form: list[dict[str, Any]] + if app_model.mode == AppMode.AGENT: + try: + features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(app_model) + except AgentAppNotPublishedError: + raise AgentNotPublishedError() + except AgentAppGeneratorError: + raise AppUnavailableError() + elif app_model.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: workflow = app_model.workflow if workflow is None: raise AppUnavailableError() - features_dict: dict[str, Any] = workflow.features_dict + features_dict = workflow.features_dict user_input_form = workflow.user_input_form(to_old_structure=True) else: app_model_config = app_model.app_model_config @@ -111,7 +123,7 @@ class AppMeta(WebApiResource): @web_ns.response(200, "Success", web_ns.models[AppMetaResponse.__name__]) def get(self, app_model: App, end_user: EndUser): """Get app meta""" - return AppService().get_app_meta(app_model) + return AppService().get_app_meta(app_model, session=db.session()) @web_ns.route("/webapp/access-mode") @@ -137,7 +149,7 @@ class AppAccessMode(Resource): app_id = args.app_id if args.app_code: - app_id = AppService.get_app_id_by_code(args.app_code) + app_id = AppService.get_app_id_by_code(args.app_code, session=db.session()) if not app_id: raise ValueError("appId or appCode must be provided") @@ -168,7 +180,9 @@ class AppWebAuthPermission(Resource): if not app_id or not app_code: raise ValueError("appId must be provided") - require_permission_check = WebAppAuthService.is_app_require_permission_check(app_id=app_id) + require_permission_check = WebAppAuthService.is_app_require_permission_check( + app_id=app_id, session=db.session() + ) if not require_permission_check: return {"result": True} @@ -189,6 +203,6 @@ class AppWebAuthPermission(Resource): return {"result": True} res = True - if WebAppAuthService.is_app_require_permission_check(app_id=app_id): + if WebAppAuthService.is_app_require_permission_check(app_id=app_id, session=db.session()): res = EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp(str(user_id), app_id) return {"result": res} diff --git a/api/controllers/web/audio.py b/api/controllers/web/audio.py index 47e72ff95a5..b7856f7dd90 100644 --- a/api/controllers/web/audio.py +++ b/api/controllers/web/audio.py @@ -141,7 +141,7 @@ class TextApi(WebApiResource): ) response = AudioService.transcript_tts( app_model=app_model, - session=db.session, + session=db.session(), text=text, voice=voice, end_user=end_user.external_user_id, diff --git a/api/controllers/web/completion.py b/api/controllers/web/completion.py index 2c852e208a5..c1a7d1f8d10 100644 --- a/api/controllers/web/completion.py +++ b/api/controllers/web/completion.py @@ -11,6 +11,7 @@ from controllers.common.schema import register_response_schema_models, register_ from controllers.console.app.wraps import with_session from controllers.web import web_ns from controllers.web.error import ( + AgentNotPublishedError, AppUnavailableError, CompletionRequestError, ConversationCompletedError, @@ -22,12 +23,14 @@ from controllers.web.error import ( ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from controllers.web.wraps import WebApiResource +from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.app.entities.app_invoke_entities import InvokeFrom from core.errors.error import ( ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError, ) +from extensions.ext_database import db from graphon.model_runtime.errors.invoke import InvokeError from libs import helper from libs.helper import uuid_value @@ -138,6 +141,8 @@ class CompletionApi(WebApiResource): except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") raise AppUnavailableError() + except AgentAppNotPublishedError: + raise AgentNotPublishedError() except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) except QuotaExceededError: @@ -215,7 +220,10 @@ class ChatApi(WebApiResource): # Eagerly validate conversation to avoid hanging on invalid conversation_id if payload.conversation_id: ConversationService.get_conversation( - app_model=app_model, conversation_id=payload.conversation_id, user=end_user + app_model=app_model, + conversation_id=payload.conversation_id, + user=end_user, + session=db.session(), ) response = AppGenerateService.generate( @@ -235,6 +243,8 @@ class ChatApi(WebApiResource): except services.errors.app_model_config.AppModelConfigBrokenError: logger.exception("App model config broken.") raise AppUnavailableError() + except AgentAppNotPublishedError: + raise AgentNotPublishedError() except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) except QuotaExceededError: diff --git a/api/controllers/web/conversation.py b/api/controllers/web/conversation.py index 73461b1a294..09a3a508824 100644 --- a/api/controllers/web/conversation.py +++ b/api/controllers/web/conversation.py @@ -112,7 +112,7 @@ class ConversationApi(WebApiResource): conversation_id = str(c_id) try: - ConversationService.delete(app_model, conversation_id, end_user) + ConversationService.delete(app_model, conversation_id, end_user, session=db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") return "", 204 @@ -157,7 +157,7 @@ class ConversationRenameApi(WebApiResource): try: conversation = ConversationService.rename( - app_model, conversation_id, end_user, payload.name, payload.auto_generate + app_model, conversation_id, end_user, payload.name, payload.auto_generate, session=db.session() ) return ( TypeAdapter(SimpleConversation) @@ -192,7 +192,7 @@ class ConversationPinApi(WebApiResource): conversation_id = str(c_id) try: - WebConversationService.pin(app_model, conversation_id, end_user) + WebConversationService.pin(app_model, conversation_id, end_user, db.session()) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -221,6 +221,6 @@ class ConversationUnPinApi(WebApiResource): raise NotChatAppError() conversation_id = str(c_id) - WebConversationService.unpin(app_model, conversation_id, end_user) + WebConversationService.unpin(app_model, conversation_id, end_user, db.session()) return ResultResponse(result="success").model_dump(mode="json") diff --git a/api/controllers/web/error.py b/api/controllers/web/error.py index 789c0fabcc1..077b4726e47 100644 --- a/api/controllers/web/error.py +++ b/api/controllers/web/error.py @@ -7,6 +7,12 @@ class AppUnavailableError(BaseHTTPException): code = 400 +class AgentNotPublishedError(BaseHTTPException): + error_code = "agent_not_published" + description = "Agent has not been published. Please publish the Agent before using the web app." + code = 400 + + class NotCompletionAppError(BaseHTTPException): error_code = "not_completion_app" description = "Please check if your Completion app mode matches the right API route." diff --git a/api/controllers/web/forgot_password.py b/api/controllers/web/forgot_password.py index ecc91113c32..a9374555ed4 100644 --- a/api/controllers/web/forgot_password.py +++ b/api/controllers/web/forgot_password.py @@ -69,7 +69,7 @@ class ForgotPasswordSendEmailApi(Resource): else: language = "en-US" - account = AccountService.get_account_by_email_with_case_fallback(db.session, request_email) + account = AccountService.get_account_by_email_with_case_fallback(request_email, session=db.session()) if account is None: raise AuthenticationFailedError() else: @@ -168,7 +168,7 @@ class ForgotPasswordResetApi(Resource): email = reset_data.get("email", "") - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=db.session()) if account: account = db.session.merge(account) diff --git a/api/controllers/web/login.py b/api/controllers/web/login.py index 2d8c38f5507..0aa42f43687 100644 --- a/api/controllers/web/login.py +++ b/api/controllers/web/login.py @@ -9,6 +9,7 @@ from werkzeug.exceptions import Unauthorized import services from configs import dify_config from controllers.common.fields import ( + AccessTokenData, AccessTokenResultResponse, LoginStatusResponse, SimpleResultDataResponse, @@ -29,6 +30,7 @@ from controllers.console.wraps import ( ) from controllers.web import web_ns from controllers.web.wraps import decode_jwt_token +from extensions.ext_database import db from libs.helper import EmailStr, extract_remote_ip from libs.passport import PassportService from libs.password import valid_password @@ -103,7 +105,7 @@ class LoginApi(Resource): normalized_email = payload.email.lower() try: - account = WebAppAuthService.authenticate(payload.email, payload.password) + account = WebAppAuthService.authenticate(payload.email, payload.password, db.session()) except services.errors.account.AccountLoginError: _log_web_login_failure(email=normalized_email, reason=LoginFailureReason.ACCOUNT_BANNED) raise AccountBannedError() @@ -115,9 +117,10 @@ class LoginApi(Resource): raise AuthenticationFailedError() token = WebAppAuthService.login(account=account) - response = make_response({"result": "success", "data": {"access_token": token}}) # set_access_token_to_cookie(request, response, token, samesite="None", httponly=False) - return response + return AccessTokenResultResponse(result="success", data=AccessTokenData(access_token=token)).model_dump( + mode="json" + ) # this api helps frontend to check whether user is authenticated @@ -136,17 +139,15 @@ class LoginStatusApi(Resource): ) @web_ns.response(200, "Login status", web_ns.models[LoginStatusResponse.__name__]) def get(self): - app_code = request.args.get("app_code") - user_id = request.args.get("user_id") + query = LoginStatusQuery.model_validate(request.args.to_dict(flat=True)) + app_code = query.app_code + user_id = query.user_id token = extract_webapp_access_token(request) if not app_code: - return { - "logged_in": bool(token), - "app_logged_in": False, - } - app_id = AppService.get_app_id_by_code(app_code) + return LoginStatusResponse(logged_in=bool(token), app_logged_in=False).model_dump(mode="json") + app_id = AppService.get_app_id_by_code(app_code, session=db.session()) is_public = not dify_config.ENTERPRISE_ENABLED or not WebAppAuthService.is_app_require_permission_check( - app_id=app_id + app_id=app_id, session=db.session() ) user_logged_in = False @@ -165,10 +166,7 @@ class LoginStatusApi(Resource): except Exception: app_logged_in = False - return { - "logged_in": user_logged_in, - "app_logged_in": app_logged_in, - } + return LoginStatusResponse(logged_in=user_logged_in, app_logged_in=app_logged_in).model_dump(mode="json") @web_ns.route("/logout") @@ -183,7 +181,8 @@ class LogoutApi(Resource): ) @web_ns.response(200, "Logout successful", web_ns.models[SimpleResultResponse.__name__]) def post(self): - response = make_response({"result": "success"}) + # response-contract:ignore hand-crafted response + response = make_response(SimpleResultResponse(result="success").model_dump(mode="json")) # enterprise SSO sets same site to None in https deployment # so we need to logout by calling api clear_webapp_access_token_from_cookie(response, samesite="None") @@ -213,12 +212,11 @@ class EmailCodeLoginSendEmailApi(Resource): else: language = "en-US" - account = WebAppAuthService.get_user_through_email(payload.email) + account = WebAppAuthService.get_user_through_email(payload.email, db.session()) if account is None: raise AuthenticationFailedError() - else: - token = WebAppAuthService.send_email_code_login_email(account=account, language=language) - return {"result": "success", "data": token} + token = WebAppAuthService.send_email_code_login_email(account=account, language=language) + return SimpleResultDataResponse(result="success", data=token).model_dump(mode="json") @web_ns.route("/email-code-login/validity") @@ -267,7 +265,7 @@ class EmailCodeLoginApi(Resource): WebAppAuthService.revoke_email_code_login_token(payload.token) try: - account = WebAppAuthService.get_user_through_email(token_email) + account = WebAppAuthService.get_user_through_email(token_email, db.session()) except Unauthorized as exc: _log_web_login_failure(email=user_email, reason=LoginFailureReason.ACCOUNT_BANNED) raise AccountBannedError() from exc @@ -277,9 +275,10 @@ class EmailCodeLoginApi(Resource): token = WebAppAuthService.login(account=account) AccountService.reset_login_error_rate_limit(user_email) - response = make_response({"result": "success", "data": {"access_token": token}}) # set_access_token_to_cookie(request, response, token, samesite="None", httponly=False) - return response + return AccessTokenResultResponse(result="success", data=AccessTokenData(access_token=token)).model_dump( + mode="json" + ) def _log_web_login_failure(*, email: str, reason: LoginFailureReason) -> None: diff --git a/api/controllers/web/message.py b/api/controllers/web/message.py index 691eba05491..45fea9a328e 100644 --- a/api/controllers/web/message.py +++ b/api/controllers/web/message.py @@ -25,6 +25,7 @@ from controllers.web.error import ( from controllers.web.wraps import WebApiResource from core.app.entities.app_invoke_entities import InvokeFrom from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from extensions.ext_database import db from fields.conversation_fields import ResultResponse from fields.message_fields import SuggestedQuestionsResponse, WebMessageInfiniteScrollPagination, WebMessageListItem from graphon.model_runtime.errors.invoke import InvokeError @@ -86,7 +87,7 @@ class MessageListApi(WebApiResource): try: pagination = MessageService.pagination_by_first_id( - app_model, end_user, query.conversation_id, query.first_id, query.limit + app_model, end_user, query.conversation_id, query.first_id, query.limit, session=db.session() ) adapter = TypeAdapter(WebMessageListItem) items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data] @@ -141,6 +142,7 @@ class MessageFeedbackApi(WebApiResource): user=end_user, rating=FeedbackRating(payload.rating) if payload.rating else None, content=payload.content, + session=db.session(), ) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -231,7 +233,11 @@ class MessageSuggestedQuestionApi(WebApiResource): try: questions = MessageService.get_suggested_questions_after_answer( - app_model=app_model, user=end_user, message_id=message_id_str, invoke_from=InvokeFrom.WEB_APP + app_model=app_model, + user=end_user, + message_id=message_id_str, + invoke_from=InvokeFrom.WEB_APP, + session=db.session(), ) # questions is a list of strings, not a list of Message objects except MessageNotExistsError: diff --git a/api/controllers/web/passport.py b/api/controllers/web/passport.py index 99b75776280..4b0b25fb971 100644 --- a/api/controllers/web/passport.py +++ b/api/controllers/web/passport.py @@ -2,7 +2,7 @@ import uuid from datetime import UTC, datetime, timedelta from typing import Any -from flask import make_response, request +from flask import request from flask_restx import Resource from pydantic import BaseModel, Field from sqlalchemy import func, select @@ -10,11 +10,12 @@ from werkzeug.exceptions import NotFound, Unauthorized from configs import dify_config from constants import HEADER_NAME_APP_CODE -from controllers.common.fields import AccessTokenData from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.web import web_ns from controllers.web.error import WebAppAuthRequiredError from extensions.ext_database import db +from fields.base import ResponseModel +from libs.helper import dump_response from libs.passport import PassportService from libs.token import extract_webapp_access_token from models.enums import EndUserType @@ -28,7 +29,13 @@ class PassportQuery(BaseModel): register_schema_models(web_ns, PassportQuery) -register_response_schema_models(web_ns, AccessTokenData) + + +class PassportAccessTokenResponse(ResponseModel): + access_token: str + + +register_response_schema_models(web_ns, PassportAccessTokenResponse) @web_ns.route("/passport") @@ -45,7 +52,7 @@ class PassportResource(Resource): 404: "Application or user not found", } ) - @web_ns.response(200, "Passport retrieved successfully", web_ns.models[AccessTokenData.__name__]) + @web_ns.response(200, "Passport retrieved successfully", web_ns.models[PassportAccessTokenResponse.__name__]) def get(self): system_features = FeatureService.get_system_features() app_code = request.headers.get(HEADER_NAME_APP_CODE) @@ -55,12 +62,15 @@ class PassportResource(Resource): raise Unauthorized("X-App-Code header is missing.") if system_features.webapp_auth.enabled: enterprise_user_decoded = decode_enterprise_webapp_user_id(access_token) - app_auth_type = WebAppAuthService.get_app_auth_type(app_code=app_code) + app_auth_type = WebAppAuthService.get_app_auth_type(app_code=app_code, session=db.session()) if app_auth_type != WebAppAuthType.PUBLIC: if not enterprise_user_decoded: raise WebAppAuthRequiredError() - return exchange_token_for_existing_web_user( - app_code=app_code, enterprise_user_decoded=enterprise_user_decoded, auth_type=app_auth_type + return dump_response( + PassportAccessTokenResponse, + exchange_token_for_existing_web_user( + app_code=app_code, enterprise_user_decoded=enterprise_user_decoded, auth_type=app_auth_type + ), ) # get site from db and check if it is normal @@ -110,12 +120,7 @@ class PassportResource(Resource): tk = PassportService().issue(payload) - response = make_response( - { - "access_token": tk, - } - ) - return response + return dump_response(PassportAccessTokenResponse, {"access_token": tk}) def decode_enterprise_webapp_user_id(jwt_token: str | None) -> dict[str, Any] | None: @@ -206,12 +211,7 @@ def exchange_token_for_existing_web_user( "exp": exp, } token: str = PassportService().issue(payload) - resp = make_response( - { - "access_token": token, - } - ) - return resp + return {"access_token": token} def _exchange_for_public_app_token(app_model, site, token_decoded): @@ -244,12 +244,7 @@ def _exchange_for_public_app_token(app_model, site, token_decoded): tk = PassportService().issue(payload) - resp = make_response( - { - "access_token": tk, - } - ) - return resp + return {"access_token": tk} def generate_session_id(): diff --git a/api/controllers/web/saved_message.py b/api/controllers/web/saved_message.py index 6e59a85e2b0..d61ffd545c3 100644 --- a/api/controllers/web/saved_message.py +++ b/api/controllers/web/saved_message.py @@ -44,7 +44,7 @@ class SavedMessageListApi(WebApiResource): query = SavedMessageListQuery.model_validate(raw_args) pagination = SavedMessageService.pagination_by_last_id( - db.session(), app_model, end_user, query.last_id, query.limit + app_model, end_user, query.last_id, query.limit, session=db.session() ) adapter = TypeAdapter(SavedMessageItem) items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data] @@ -80,7 +80,7 @@ class SavedMessageListApi(WebApiResource): payload = SavedMessageCreatePayload.model_validate(web_ns.payload or {}) try: - SavedMessageService.save(db.session(), app_model, end_user, payload.message_id) + SavedMessageService.save(app_model, end_user, payload.message_id, session=db.session()) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -108,6 +108,6 @@ class SavedMessageApi(WebApiResource): if app_model.mode != "completion": raise NotCompletionAppError() - SavedMessageService.delete(db.session(), app_model, end_user, message_id_str) + SavedMessageService.delete(app_model, end_user, message_id_str, session=db.session()) return "", 204 diff --git a/api/controllers/web/wraps.py b/api/controllers/web/wraps.py index ccc9c0f8f60..eff4b70ff0f 100644 --- a/api/controllers/web/wraps.py +++ b/api/controllers/web/wraps.py @@ -70,7 +70,7 @@ def decode_jwt_token(app_code: str | None = None, user_id: str | None = None) -> app_web_auth_enabled = False webapp_settings = None if system_features.webapp_auth.enabled: - app_id = AppService.get_app_id_by_code(app_code) + app_id = AppService.get_app_id_by_code(app_code, session=db.session()) webapp_settings = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id) if not webapp_settings: raise NotFound("Web app settings not found.") @@ -86,7 +86,7 @@ def decode_jwt_token(app_code: str | None = None, user_id: str | None = None) -> if system_features.webapp_auth.enabled: if not app_code: raise Unauthorized("Please re-login to access the web app.") - app_id = AppService.get_app_id_by_code(app_code) + app_id = AppService.get_app_id_by_code(app_code, session=db.session()) app_web_auth_enabled = ( EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id=app_id).access_mode != WebAppAccessMode.PUBLIC @@ -129,8 +129,10 @@ def _validate_user_accessibility( if not webapp_settings: raise WebAppAuthRequiredError("Web app settings not found.") - if WebAppAuthService.is_app_require_permission_check(access_mode=webapp_settings.access_mode): - app_id = AppService.get_app_id_by_code(app_code) + if WebAppAuthService.is_app_require_permission_check( + access_mode=webapp_settings.access_mode, session=db.session() + ): + app_id = AppService.get_app_id_by_code(app_code, session=db.session()) if not EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp(user_id, app_id): raise WebAppAuthAccessDeniedError() diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py index 8fcf42ce67d..4d823ca79c0 100644 --- a/api/core/agent/cot_agent_runner.py +++ b/api/core/agent/cot_agent_runner.py @@ -133,6 +133,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): stop=app_generate_entity.model_conf.stop, stream=True, callbacks=[], + request_metadata={"app_id": self.app_config.app_id}, ) usage_dict: dict[str, LLMUsage | None] = {} diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py index 9db5fa08f75..0b92daf93a4 100644 --- a/api/core/agent/fc_agent_runner.py +++ b/api/core/agent/fc_agent_runner.py @@ -101,6 +101,7 @@ class FunctionCallAgentRunner(BaseAgentRunner): stop=app_generate_entity.model_conf.stop, stream=self.stream_tool_call, callbacks=[], + request_metadata={"app_id": self.app_config.app_id}, ) tool_calls: list[tuple[str, str, dict[str, Any]]] = [] diff --git a/api/core/app/app_config/easy_ui_based_app/dataset/manager.py b/api/core/app/app_config/easy_ui_based_app/dataset/manager.py index 140d4e6a2a6..0108e7d7c72 100644 --- a/api/core/app/app_config/easy_ui_based_app/dataset/manager.py +++ b/api/core/app/app_config/easy_ui_based_app/dataset/manager.py @@ -257,7 +257,7 @@ class DatasetConfigManager: @classmethod def is_dataset_exists(cls, tenant_id: str, dataset_id: str) -> bool: # verify if the dataset ID exists - dataset = DatasetService.get_dataset(dataset_id, db.session) + dataset = DatasetService.get_dataset(dataset_id, db.session()) if not dataset: return False diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index f52fd1046f8..23acfb7ea62 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -48,6 +48,7 @@ from core.repositories import DifyCoreRepositoryFactory from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from extensions.ext_database import db from factories import file_factory +from graphon.filters import ResponseStreamFilter from graphon.graph_engine.layers import GraphEngineLayer from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from graphon.runtime import GraphRuntimeState @@ -156,7 +157,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): if conversation_id: try: conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=db.session() ) except ConversationNotExistsError: if invoke_from == InvokeFrom.SERVICE_API: @@ -269,6 +270,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): workflow_node_execution_repository: WorkflowNodeExecutionRepository, graph_runtime_state: GraphRuntimeState, pause_state_config: PauseStateLayerConfig | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ): """ Resume a paused advanced chat execution. @@ -298,6 +300,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): stream=application_generate_entity.stream, pause_state_config=pause_state_config, graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, ) def single_iteration_generate( @@ -492,6 +495,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): pause_state_config: PauseStateLayerConfig | None = None, graph_runtime_state: GraphRuntimeState | None = None, graph_engine_layers: Sequence[GraphEngineLayer] = (), + response_stream_filter: ResponseStreamFilter | None = None, ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]: """ Generate App response. @@ -539,12 +543,14 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): ) graph_layers: list[GraphEngineLayer] = list(graph_engine_layers) + resolved_response_stream_filter = response_stream_filter or ResponseStreamFilter() if pause_state_config is not None: graph_layers.append( PauseStatePersistenceLayer( session_factory=pause_state_config.session_factory, generate_entity=application_generate_entity, state_owner_user_id=pause_state_config.state_owner_user_id, + response_stream_filter=resolved_response_stream_filter, ) ) @@ -565,6 +571,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): "workflow_node_execution_repository": workflow_node_execution_repository, "graph_engine_layers": tuple(graph_layers), "graph_runtime_state": graph_runtime_state, + "response_stream_filter": resolved_response_stream_filter, }, ) @@ -604,6 +611,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): workflow_node_execution_repository: WorkflowNodeExecutionRepository, graph_engine_layers: Sequence[GraphEngineLayer] = (), graph_runtime_state: GraphRuntimeState | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ): """ Generate worker in a new thread. @@ -663,6 +671,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): workflow_node_execution_repository=workflow_node_execution_repository, graph_engine_layers=graph_engine_layers, graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, ) try: diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index b78a3b5b3dc..249cb33a98c 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -44,6 +44,7 @@ from extensions.ext_redis import redis_client from extensions.otel import WorkflowAppRunnerHandler, trace_span from extensions.workflow_warm_shutdown import WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, celery_warm_shutdown_started from graphon.enums import WorkflowType +from graphon.filters import ResponseStreamFilter from graphon.graph_engine.command_channels import RedisChannel from graphon.graph_engine.layers import GraphEngineLayer from graphon.runtime import GraphRuntimeState, VariablePool @@ -78,6 +79,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): workflow_node_execution_repository: WorkflowNodeExecutionRepository, graph_engine_layers: Sequence[GraphEngineLayer] = (), graph_runtime_state: GraphRuntimeState | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ): super().__init__( queue_manager=queue_manager, @@ -95,6 +97,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): self._workflow_execution_repository = workflow_execution_repository self._workflow_node_execution_repository = workflow_node_execution_repository self._resume_graph_runtime_state = graph_runtime_state + self._response_stream_filter = response_stream_filter @trace_span(WorkflowAppRunnerHandler) def run(self): @@ -241,6 +244,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): variable_pool=variable_pool, graph_runtime_state=graph_runtime_state, command_channel=command_channel, + response_stream_filter=self._response_stream_filter, ) self._queue_manager.graph_runtime_state = graph_runtime_state diff --git a/api/core/app/apps/agent_app/app_config_manager.py b/api/core/app/apps/agent_app/app_config_manager.py index 0dc04735cc0..71224534133 100644 --- a/api/core/app/apps/agent_app/app_config_manager.py +++ b/api/core/app/apps/agent_app/app_config_manager.py @@ -3,9 +3,9 @@ An Agent App has no legacy ``app_model_config``: its model / prompt live in the bound Agent Soul snapshot. To ride the existing chat message + SSE pipeline we synthesize an ``app_model_config``-shaped dict from the Soul (model + system -prompt) plus any app-level feature flags (opening statement, follow-up, …) -stored on ``app_model_config`` when present, then reuse the same sub-managers -the chat app type uses. +prompt) plus app-level feature flags from Agent Soul, while preserving any +legacy ``app_model_config`` feature flags when present. Then we reuse the same +sub-managers the chat app type uses. """ from typing import Any, cast @@ -21,6 +21,7 @@ from core.app.app_config.entities import ( EasyUIBasedAppModelConfigFrom, PromptTemplateEntity, ) +from core.app.apps.agent_app.app_feature_projection import merge_agent_app_features from core.app.apps.agent_app.app_variable_projection import agent_app_variables_to_user_input_form from models.agent_config_entities import AgentSoulConfig from models.model import App, AppMode, AppModelConfig, AppModelConfigDict, Conversation @@ -79,12 +80,11 @@ class AgentAppConfigManager(BaseAppConfigManager): ) -> dict[str, Any]: """Shape a Soul + feature flags into an ``app_model_config``-style dict. - Feature flags (opening statement / follow-up / tts / stt / citations / - moderation / annotation) come from ``app_model_config`` when present - (Q3: stored there), otherwise defaults; model + prompt always come from + Feature flags come from Agent Soul and fill gaps in the legacy + ``app_model_config`` when one exists; model + prompt always come from the Agent Soul (the single source of truth for those). """ - base: dict[str, Any] = dict(app_model_config.to_dict()) if app_model_config else {} + base = merge_agent_app_features(agent_soul=agent_soul, app_model_config=app_model_config) model = agent_soul.model if model is not None: diff --git a/api/core/app/apps/agent_app/app_feature_projection.py b/api/core/app/apps/agent_app/app_feature_projection.py new file mode 100644 index 00000000000..cb8efd9f290 --- /dev/null +++ b/api/core/app/apps/agent_app/app_feature_projection.py @@ -0,0 +1,23 @@ +from typing import Any + +from models.agent_config_entities import AgentSoulConfig + + +def merge_agent_app_features( + *, + agent_soul: AgentSoulConfig, + app_model_config: Any | None, +) -> dict[str, Any]: + """Project public Agent App features from legacy config plus Agent Soul. + + The hidden backing app may still carry legacy presentation fields such as + opening statements. Agent Soul is the source of truth for Agent-owned + features like file upload, so Soul fields override same-named legacy keys. + """ + features: dict[str, Any] = dict(app_model_config.to_dict()) if app_model_config else {} + soul_features = agent_soul.app_features.model_dump(mode="json", exclude_none=True) + features.update(soul_features) + return features + + +__all__ = ["merge_agent_app_features"] diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py index 9d45ac71389..4c10004fc7e 100644 --- a/api/core/app/apps/agent_app/app_generator.py +++ b/api/core/app/apps/agent_app/app_generator.py @@ -1,14 +1,11 @@ """Agent App generator: orchestrate Agent App chat and finalize executions. -The primary mode mirrors the agent_chat generator (conversation + message + +Agent App turns mirror the agent_chat generator (conversation + message + queue + streamed response over the EasyUI chat pipeline), but the backing config comes from the bound Agent Soul and the answer is produced by ``AgentAppRunner`` calling the dify-agent backend rather than an in-process -LLM/ReAct loop. - -It also exposes a stateless build-finalize mode that reuses existing runtime -context from the bound debug conversation, triggers the Agent backend side -effect synchronously, and skips Dify-side chat/message persistence. +LLM/ReAct loop. Build-chat finalization uses this same streamed path and only +changes the runtime exit policy carried to the backend. """ from __future__ import annotations @@ -32,6 +29,7 @@ from constants import UUID_NIL from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter from core.app.apps.agent_app.app_config_manager import AgentAppConfigManager from core.app.apps.agent_app.app_runner import AgentAppRunner +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError from core.app.apps.agent_app.generate_response_converter import AgentAppGenerateResponseConverter from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeRequestBuilder from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore @@ -40,13 +38,16 @@ from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import ( + AGENT_RUNTIME_EXIT_INTENT_ARG, AgentAppGenerateEntity, + AgentRuntimeExitIntent, DifyRunContext, InvokeFrom, UserFrom, ) from core.app.llm.model_access import build_dify_model_access from core.ops.ops_trace_manager import TraceQueueManager +from core.workflow.file_reference import build_file_reference, is_canonical_file_reference from extensions.ext_database import db from models import Account, App, EndUser, Message from models.agent import ( @@ -63,16 +64,68 @@ from services.conversation_service import ConversationService logger = logging.getLogger(__name__) - -class AgentAppGeneratorError(ValueError): - """Raised when an Agent App turn cannot be set up.""" +_REFERENCE_FILE_TRANSFER_METHODS = {"local_file", "tool_file", "datasource_file"} def _append_prompt_file_mappings(query: str, prompt_file_mappings: Sequence[JsonValue]) -> str: - """Append raw request file references to the backend user prompt.""" - if not prompt_file_mappings: + """Append labeled, prompt-safe file locators to the backend user prompt.""" + prompt_files = _prompt_file_locators(prompt_file_mappings) + if not prompt_files: return query - return f"{query}\n{json.dumps(list(prompt_file_mappings), ensure_ascii=False)}" + payload = json.dumps(prompt_files, ensure_ascii=False, separators=(",", ":")) + return ( + f"{query}\n" + "User provided files: use dify-agent file download with the listed transfer_method and reference/url " + "to get the files and investigate them\n" + f"{payload}" + ) + + +def _prompt_file_locators(prompt_file_mappings: Sequence[JsonValue]) -> list[dict[str, str]]: + locators: list[dict[str, str]] = [] + for file_mapping in prompt_file_mappings: + if not isinstance(file_mapping, Mapping): + continue + locator = _prompt_file_locator(file_mapping) + if locator is not None: + locators.append(locator) + return locators + + +def _prompt_file_locator(file_mapping: Mapping[str, object]) -> dict[str, str] | None: + transfer_method = _string_value(file_mapping, "transfer_method") + if transfer_method == "remote_url": + url = _string_value(file_mapping, "url") or _string_value(file_mapping, "remote_url") + if url is None: + return None + return {"transfer_method": "remote_url", "url": url} + elif transfer_method in _REFERENCE_FILE_TRANSFER_METHODS: + if transfer_method is None: + return None + reference = _canonical_file_reference( + _string_value(file_mapping, "reference") + or _string_value(file_mapping, "upload_file_id") + or _string_value(file_mapping, "file_id") + or _string_value(file_mapping, "id") + ) + if reference is None: + return None + return {"transfer_method": transfer_method, "reference": reference} + else: + return None + + +def _canonical_file_reference(reference: str | None) -> str | None: + if reference is None: + return None + if reference.startswith("dify-file-ref:"): + return reference if is_canonical_file_reference(reference) else None + return build_file_reference(record_id=reference) + + +def _string_value(mapping: Mapping[str, object], key: str) -> str | None: + value = mapping.get(key) + return value if isinstance(value, str) and value else None class AgentAppGenerator(MessageBasedAppGenerator): @@ -108,7 +161,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): conversation_id = args.get("conversation_id") if conversation_id: conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=db.session() ) # Build the EasyUI-shaped config from the Agent Soul so the chat pipeline @@ -123,6 +176,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): model_conf = ModelConfigConverter.convert(app_config) trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) + agent_runtime_exit_intent = self._resolve_agent_runtime_exit_intent(args) application_generate_entity = AgentAppGenerateEntity( task_id=str(uuid.uuid4()), @@ -152,6 +206,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): agent_config_snapshot_id=agent_config_id, agent_config_version_kind=agent_config_version_kind, agent_runtime_session_snapshot_id=runtime_session_snapshot_id, + agent_runtime_exit_intent=agent_runtime_exit_intent, ) conversation, message = self._init_generate_records(application_generate_entity, conversation) @@ -190,86 +245,6 @@ class AgentAppGenerator(MessageBasedAppGenerator): ) return AgentAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) - def generate_stateless( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - ) -> Mapping[str, Any]: - """Run one Agent App turn without persisting Dify conversation messages.""" - query = self._require_query(args) - conversation_id = args.get("conversation_id") - if not isinstance(conversation_id, str) or not conversation_id: - raise AgentAppGeneratorError("conversation_id is required") - - agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( - app_model, - invoke_from=invoke_from, - draft_type=args.get("draft_type"), - user=user, - ) - runtime_session_snapshot_id = self._runtime_session_snapshot_id( - invoke_from=invoke_from, - snapshot_id=agent_config_id, - ) - - return self._run_stateless( - app_model=app_model, - user=user, - invoke_from=invoke_from, - query=query, - conversation_id=conversation_id, - agent=agent, - agent_config_id=agent_config_id, - agent_config_version_kind=agent_config_version_kind, - agent_soul=agent_soul, - runtime_session_snapshot_id=runtime_session_snapshot_id, - ) - - def _run_stateless( - self, - *, - app_model: App, - user: Account | EndUser, - invoke_from: InvokeFrom, - query: str, - conversation_id: str, - agent: Agent, - agent_config_id: str, - agent_config_version_kind: Literal["snapshot", "draft", "build_draft"], - agent_soul: AgentSoulConfig, - runtime_session_snapshot_id: str | None, - ) -> Mapping[str, Any]: - """Run the Agent backend without creating or updating Dify chat records. - - Build-chat finalization is an action against the Agent backend (for - example, ``dify-agent config push``). It may reuse the active build-chat - runtime snapshot for shell/config context, but the API side must not add - a synthetic user/assistant turn to the debug conversation. - """ - - dify_context = DifyRunContext( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - user_id=user.id, - user_from=UserFrom.ACCOUNT if isinstance(user, Account) else UserFrom.END_USER, - invoke_from=invoke_from, - ) - self._build_runner(dify_context).run_stateless( - dify_context=dify_context, - agent_id=agent.id, - agent_config_snapshot_id=agent_config_id, - agent_config_version_kind=agent_config_version_kind, - agent_soul=agent_soul, - conversation_id=conversation_id, - query=query, - idempotency_key=str(uuid.uuid4()), - session_scope_snapshot_id=runtime_session_snapshot_id, - ) - return {"result": "success"} - def resume_after_form_submission( self, *, @@ -287,7 +262,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): out of scope here — the message is persisted and can be re-fetched. """ conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=db.session() ) agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( app_model, @@ -479,6 +454,7 @@ class AgentAppGenerator(MessageBasedAppGenerator): model_name=application_generate_entity.model_conf.model, queue_manager=queue_manager, session_scope_snapshot_id=application_generate_entity.agent_runtime_session_snapshot_id, + agent_runtime_exit_intent=application_generate_entity.agent_runtime_exit_intent, ) except GenerateTaskStoppedError: pass @@ -495,6 +471,18 @@ class AgentAppGenerator(MessageBasedAppGenerator): raise AgentAppGeneratorError("query is required") return query.replace("\x00", "") + @staticmethod + def _resolve_agent_runtime_exit_intent(args: Mapping[str, Any]) -> AgentRuntimeExitIntent: + """Resolve API-internal runtime exit policy from controller-owned args. + + Only the private controller-injected "delete" value changes behavior. + Normal chat and resume flows default/fallback to "suspend" so public + payloads and invalid internal values preserve existing semantics. + """ + if args.get(AGENT_RUNTIME_EXIT_INTENT_ARG) == "delete": + return "delete" + return "suspend" + @staticmethod def _build_runner(dify_context: DifyRunContext) -> AgentAppRunner: credentials_provider, _ = build_dify_model_access(dify_context) @@ -614,6 +602,10 @@ class AgentAppGenerator(MessageBasedAppGenerator): "build_draft" if draft.draft_type == AgentConfigDraftType.DEBUG_BUILD else "draft" ) return agent, draft.id, config_version_kind, agent_soul + # active_config_is_published tracks whether the editable draft matches the active snapshot. + # Public runtime must keep serving the active snapshot even when unpublished draft edits exist. + if not agent.active_config_snapshot_id: + raise AgentAppNotPublishedError("Agent has not been published") _, snapshot, agent_soul = self._resolve_agent_by_id( tenant_id=app_model.tenant_id, agent_id=agent.id, @@ -709,4 +701,4 @@ class AgentAppGenerator(MessageBasedAppGenerator): return agent, draft, agent_soul -__all__ = ["AgentAppGenerator", "AgentAppGeneratorError"] +__all__ = ["AgentAppGenerator", "AgentAppGeneratorError", "AgentAppNotPublishedError"] diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py index 57dac07f761..c6546cac569 100644 --- a/api/core/app/apps/agent_app/app_runner.py +++ b/api/core/app/apps/agent_app/app_runner.py @@ -1,15 +1,10 @@ -"""Agent App runner: drive Agent backend turns for both chat and finalize flows. +"""Agent App runner: drive Agent backend turns for chat and finalization flows. Unlike the legacy ``AgentChatAppRunner`` (which runs an in-process ReAct loop), -this runner delegates to the Agent backend and supports two execution modes. - -- Normal chat turns build the run request from the Agent Soul + conversation, - consume backend stream events, republish the assistant answer through the - existing EasyUI chat task pipeline, and save the conversation - ``session_snapshot`` on success for multi-turn continuity (S3). -- Stateless build-finalize turns reuse any prior conversation snapshot only to - construct the backend request, wait synchronously for backend completion, and - intentionally do not persist Dify-side chat records or runtime-session state. +this runner delegates to the Agent backend, consumes the streamed event flow, +republishes the assistant answer through the existing EasyUI chat task +pipeline, and then either saves or retires the conversation-owned runtime +session depending on the turn's exit policy. """ from __future__ import annotations @@ -26,6 +21,7 @@ from dify_agent.protocol import DeferredToolResultsPayload from pydantic import JsonValue from clients.agent_backend import ( + AgentBackendAgentMessageDeltaInternalEvent, AgentBackendDeferredToolCallInternalEvent, AgentBackendError, AgentBackendInternalEventType, @@ -35,7 +31,7 @@ from clients.agent_backend import ( AgentBackendStreamInternalEvent, extract_runtime_layer_specs, ) -from configs import dify_config +from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from core.app.apps.agent_app.runtime_request_builder import ( AgentAppRuntimeBuildContext, AgentAppRuntimeRequest, @@ -48,8 +44,13 @@ from core.app.apps.agent_app.session_store import ( ) from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import DifyRunContext -from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueLLMChunkEvent, QueueMessageEndEvent +from core.app.entities.app_invoke_entities import AgentRuntimeExitIntent, DifyRunContext +from core.app.entities.queue_entities import ( + QueueAgentMessageEvent, + QueueAgentThoughtEvent, + QueueLLMChunkEvent, + QueueMessageEndEvent, +) from core.repositories.human_input_repository import HumanInputFormRepository, HumanInputFormRepositoryImpl from core.workflow.nodes.agent_v2.ask_human_hitl import AskHumanFormBuildError, create_ask_human_form from core.workflow.nodes.agent_v2.ask_human_resume import build_deferred_tool_results, resolve_ask_human_form @@ -59,6 +60,7 @@ from graphon.model_runtime.entities.message_entities import AssistantPromptMessa from models.agent_config_entities import AgentSoulConfig from models.enums import CreatorUserRole from models.model import MessageAgentThought +from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session logger = logging.getLogger(__name__) @@ -135,6 +137,25 @@ def publish_text_delta( queue_manager.publish(QueueLLMChunkEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) +def publish_agent_message_delta( + *, + queue_manager: AppQueueManager, + model_name: str, + delta: str, + user_query: str | None = None, +) -> None: + """Publish one agent-process text delta through the EasyUI chat pipeline.""" + if not delta: + return + prompt_messages = _prompt_messages_from_query(user_query) + chunk = LLMResultChunk( + model=model_name, + prompt_messages=prompt_messages, + delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=delta)), + ) + queue_manager.publish(QueueAgentMessageEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) + + def publish_message_end( *, queue_manager: AppQueueManager, @@ -159,7 +180,7 @@ def publish_message_end( class _TextDeltaDebouncer: - """Batch assistant text deltas on stream-event boundaries for final SSE output.""" + """Batch independent model text deltas before agent-message SSE output.""" def __init__(self, *, debounce_seconds: float) -> None: self._debounce_seconds = debounce_seconds @@ -192,7 +213,13 @@ class _TextDeltaDebouncer: class _AgentProcessRecorder: - """Persist Agent v2 thinking/tool process events through the legacy thought model.""" + """Persist Agent v2 process streams through the legacy thought model. + + Thinking and answer rows expose snapshot updates for contiguous model-text + segments. Tool events close currently open text segments so later model text + starts a fresh row instead of replaying content that was already streamed + before the tool. + """ def __init__( self, @@ -206,6 +233,7 @@ class _AgentProcessRecorder: self._queue_manager = queue_manager self._next_position = 1 self._thinking_by_index: dict[int, str] = {} + self._answer_thought_id: str | None = None self._tool_by_index: dict[int, str] = {} self._tool_by_call_id: dict[str, str] = {} self._open_tool_by_name: dict[str, set[str]] = {} @@ -261,6 +289,41 @@ class _AgentProcessRecorder: if part_kind in {"tool-return", "builtin-tool-return"}: self._record_tool_return_part(part) + def append_answer_text(self, content_delta: str) -> None: + if not content_delta: + return + + self._thinking_by_index.clear() + if self._answer_thought_id is None: + self._answer_thought_id = self._create_thought(answer=content_delta) + return + self._update_thought(self._answer_thought_id, answer_delta=content_delta) + + def trim_answer_suffix(self, final_answer: str) -> None: + if not final_answer or self._answer_thought_id is None: + return + + row = db.session.get(MessageAgentThought, self._answer_thought_id) + if row is None: + return + + answer = row.answer or "" + overlap = _suffix_prefix_overlap_length(answer, final_answer) + if overlap == 0: + return + + row.answer = answer[:-overlap] + if _is_empty_answer_only_thought(row): + db.session.delete(row) + self._answer_thought_id = None + db.session.commit() + return + + db.session.commit() + self._queue_manager.publish( + QueueAgentThoughtEvent(agent_thought_id=self._answer_thought_id), PublishFrom.APPLICATION_MANAGER + ) + def _handle_tool_call_event(self, data: dict[str, Any]) -> None: part = data.get("part") if isinstance(part, dict): @@ -281,6 +344,7 @@ class _AgentProcessRecorder: ) def _append_thinking(self, index: int, content_delta: str) -> None: + self._answer_thought_id = None thought_id = self._thinking_by_index.get(index) if thought_id is None: thought_id = self._create_thought(thought=content_delta) @@ -289,6 +353,7 @@ class _AgentProcessRecorder: self._update_thought(thought_id, thought_delta=content_delta) def _record_tool_call_delta(self, index: int, delta: dict[str, Any]) -> None: + self._close_thinking_segments() tool_call_id = _string_or_none(delta.get("tool_call_id")) tool_name = _string_or_none(delta.get("tool_name_delta")) args_delta = delta.get("args_delta") @@ -305,8 +370,10 @@ class _AgentProcessRecorder: tool=tool_name, tool_input_delta=_json_or_text(args_delta), ) + self._remember_tool_thought(index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id) def _record_tool_call_part(self, index: int, part: dict[str, Any]) -> None: + self._close_thinking_segments() tool_call_id = _string_or_none(part.get("tool_call_id")) tool_name = _string_or_none(part.get("tool_name")) thought_id = self._lookup_tool_thought(index=index, tool_call_id=tool_call_id) @@ -322,8 +389,10 @@ class _AgentProcessRecorder: tool=tool_name, tool_input=_json_or_text(part.get("args")), ) + self._remember_tool_thought(index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id) def _record_tool_return_part(self, part: dict[str, Any]) -> None: + self._close_thinking_segments() tool_call_id = _string_or_none(part.get("tool_call_id")) tool_name = _string_or_none(part.get("tool_name")) content = part.get("content") @@ -332,6 +401,7 @@ class _AgentProcessRecorder: self._record_tool_observation(tool_call_id=tool_call_id, tool_name=tool_name, observation=content) def _record_tool_observation(self, *, tool_call_id: str | None, tool_name: str | None, observation: Any) -> None: + self._close_thinking_segments() thought_id = self._lookup_observation_thought(tool_call_id=tool_call_id, tool_name=tool_name) if thought_id is None: thought_id = self._create_thought(tool=tool_name) @@ -366,25 +436,34 @@ class _AgentProcessRecorder: for open_thought_ids in self._open_tool_by_name.values(): open_thought_ids.discard(thought_id) + def _close_thinking_segments(self) -> None: + self._thinking_by_index.clear() + self._answer_thought_id = None + def _create_thought( - self, *, thought: str | None = None, tool: str | None = None, tool_input: str | None = None + self, + *, + thought: str | None = None, + answer: str | None = None, + tool: str | None = None, + tool_input: str | None = None, ) -> str: row = MessageAgentThought( message_id=self._message_id, message_chain_id=None, - thought=thought, - tool=tool, + thought=thought or "", + tool=tool or "", tool_labels_str=_tool_labels(tool), tool_meta_str="{}", - tool_input=tool_input, - observation=None, + tool_input=tool_input or "", + observation="", tool_process_data=None, - message=None, + message="", message_token=0, message_unit_price=Decimal(0), message_price_unit=Decimal("0.001"), message_files="", - answer="", + answer=answer or "", answer_token=0, answer_unit_price=Decimal(0), answer_price_unit=Decimal("0.001"), @@ -414,6 +493,7 @@ class _AgentProcessRecorder: tool_input: str | None = None, tool_input_delta: str | None = None, observation: str | None = None, + answer_delta: str | None = None, ) -> None: row = db.session.get(MessageAgentThought, thought_id) if row is None: @@ -430,6 +510,8 @@ class _AgentProcessRecorder: row.tool_input = f"{row.tool_input or ''}{tool_input_delta}" if observation is not None: row.observation = observation + if answer_delta: + row.answer = f"{row.answer or ''}{answer_delta}" db.session.commit() self._queue_manager.publish( @@ -468,6 +550,18 @@ def _tool_labels(tool: str | None) -> str: return json.dumps({tool: {"en_US": tool, "zh_Hans": tool}}, ensure_ascii=False) +def _suffix_prefix_overlap_length(text: str, prefix_source: str) -> int: + max_length = min(len(text), len(prefix_source)) + for length in range(max_length, 0, -1): + if text.endswith(prefix_source[:length]): + return length + return 0 + + +def _is_empty_answer_only_thought(row: MessageAgentThought) -> bool: + return not any((row.thought, row.answer, row.tool, row.tool_input, row.observation)) + + class AgentAppRunner: """Runs one Agent App conversation turn against the Agent backend.""" @@ -500,7 +594,9 @@ class AgentAppRunner: model_name: str, queue_manager: AppQueueManager, session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId = _DEFAULT_SESSION_SCOPE_SNAPSHOT_ID, + agent_runtime_exit_intent: AgentRuntimeExitIntent = "suspend", ) -> None: + preserve_session = agent_runtime_exit_intent == "suspend" scope = self._build_session_scope( dify_context=dify_context, agent_id=agent_id, @@ -522,10 +618,11 @@ class AgentAppRunner: idempotency_key=message_id, stored=stored, message_id=message_id, + suspend_on_exit=preserve_session, ) create_response = self._agent_backend_client.create_run(runtime.request) - terminal, streamed_answer = self._consume_stream( + terminal, process_recorder = self._consume_stream( create_response.run_id, dify_context=dify_context, message_id=message_id, @@ -535,6 +632,9 @@ class AgentAppRunner: ) if isinstance(terminal, AgentBackendDeferredToolCallInternalEvent): + if not preserve_session: + self._mark_session_cleaned(scope=scope, backend_run_id=terminal.run_id) + raise AgentBackendError("Agent App finalization cannot pause for human input.") # ENG-635: the agent asked a human. End this turn with the question and # a conversation-owned HITL form; a form submission resumes the run. self._pause_for_ask_human( @@ -555,70 +655,49 @@ class AgentAppRunner: error = getattr(terminal, "error", None) or "Agent backend run did not complete successfully." raise AgentBackendError(str(error)) - answer = self._extract_answer(terminal.output) - self._publish_terminal_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - query=query, - streamed_answer=streamed_answer, - usage=_llm_usage_from_agent_backend(terminal.usage), - ) - self._save_session( - scope=scope, - backend_run_id=terminal.run_id, - snapshot=terminal.session_snapshot, - runtime_layer_specs=extract_runtime_layer_specs(runtime.request.composition), - ) - - def run_stateless( - self, - *, - dify_context: DifyRunContext, - agent_id: str, - agent_config_snapshot_id: str, - agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot", - agent_soul: AgentSoulConfig, - conversation_id: str, - query: str, - idempotency_key: str, - session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId = _DEFAULT_SESSION_SCOPE_SNAPSHOT_ID, - ) -> None: - """Run the Agent backend without creating Dify chat message records. - - This path is used by build-chat finalization: the API must trigger the - backend side effects in the existing conversation session, but it must - not persist a synthetic user/assistant turn, update API-side runtime - session rows, or set up HITL state that depends on one. - """ - scope = self._build_session_scope( - dify_context=dify_context, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, - conversation_id=conversation_id, - session_scope_snapshot_id=session_scope_snapshot_id, - ) - runtime = self._build_runtime( - dify_context=dify_context, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, - agent_config_version_kind=agent_config_version_kind, - agent_soul=agent_soul, - conversation_id=conversation_id, - query=query, - idempotency_key=idempotency_key, - stored=self._session_store.load_active_session(scope), - message_id=None, - ) - - create_response = self._agent_backend_client.create_run(runtime.request) - status = self._agent_backend_client.wait_run( - create_response.run_id, - timeout_seconds=dify_config.APP_MAX_EXECUTION_TIME, - ) - if status.status != "succeeded": - error = getattr(status, "error", None) or f"Agent backend run ended with status {status.status}." - raise AgentBackendError(str(error)) + answer = self._terminal_output_to_answer(terminal.output) + try: + process_recorder.trim_answer_suffix(answer) + except Exception: + db.session.rollback() + logger.warning( + "Failed to trim Agent App answer text: run_id=%s message_id=%s", + terminal.run_id, + message_id, + exc_info=True, + ) + if preserve_session: + superseded_sessions = self._load_superseded_sessions(scope=scope) + self._publish_terminal_answer( + queue_manager=queue_manager, + model_name=model_name, + answer=answer, + query=query, + usage=_llm_usage_from_agent_backend(terminal.usage), + ) + session_saved = self._save_session( + scope=scope, + backend_run_id=terminal.run_id, + snapshot=terminal.session_snapshot, + runtime_layer_specs=extract_runtime_layer_specs(runtime.request.composition), + ) + if session_saved: + self._cleanup_superseded_sessions(superseded_sessions) + else: + # The backend has already accepted a terminal success with + # delete-on-exit semantics. Local publish/persistence errors must + # not keep the API-side session row active, and cleanup failures + # must not replace the original publish/error outcome. + try: + self._publish_terminal_answer( + queue_manager=queue_manager, + model_name=model_name, + answer=answer, + query=query, + usage=_llm_usage_from_agent_backend(terminal.usage), + ) + finally: + self._mark_session_cleaned(scope=scope, backend_run_id=terminal.run_id) def _build_session_scope( self, @@ -654,6 +733,7 @@ class AgentAppRunner: idempotency_key: str, stored: StoredAgentAppSession | None, message_id: str | None, + suspend_on_exit: bool, ) -> AgentAppRuntimeRequest: session_snapshot = stored.session_snapshot if stored is not None else None deferred_tool_results = ( @@ -673,6 +753,7 @@ class AgentAppRunner: idempotency_key=idempotency_key, session_snapshot=session_snapshot, deferred_tool_results=deferred_tool_results, + suspend_on_exit=suspend_on_exit, ) ) @@ -775,46 +856,61 @@ class AgentAppRunner: model_name: str, query: str | None, ): - """Consume backend events while preserving raw recorder granularity. - - Process events are recorded immediately for observability. Only the - final assistant text deltas sent through the EasyUI queue are debounced, - with flushes happening on later stream events or terminal boundaries. - """ + """Consume backend events while preserving raw recorder granularity.""" terminal = None - streamed_answer_parts: list[str] = [] - text_delta_debouncer = _TextDeltaDebouncer(debounce_seconds=self._text_delta_debounce_seconds) process_recorder = _AgentProcessRecorder( dify_context=dify_context, message_id=message_id, queue_manager=queue_manager, ) + text_delta_debouncer = _TextDeltaDebouncer(debounce_seconds=self._text_delta_debounce_seconds) - def flush_pending_text() -> None: + def persist_answer_text(content_delta: str) -> None: + try: + process_recorder.append_answer_text(content_delta) + except Exception: + db.session.rollback() + logger.warning( + "Failed to persist Agent App answer text: run_id=%s message_id=%s", + run_id, + message_id, + exc_info=True, + ) + publish_agent_message_delta( + queue_manager=queue_manager, + model_name=model_name, + delta=content_delta, + user_query=query, + ) + + def flush_pending_agent_message_text() -> None: pending_text = text_delta_debouncer.flush() if pending_text: - publish_text_delta( - queue_manager=queue_manager, - model_name=model_name, - delta=pending_text, - user_query=query, - ) + persist_answer_text(pending_text) for public_event in self._agent_backend_client.stream_events(run_id): if queue_manager.is_stopped(): - flush_pending_text() + flush_pending_agent_message_text() self._cancel_run(run_id) raise GenerateTaskStoppedError() for internal_event in self._event_adapter.adapt(public_event): if queue_manager.is_stopped(): - flush_pending_text() + flush_pending_agent_message_text() self._cancel_run(run_id) raise GenerateTaskStoppedError() if internal_event.type in ( AgentBackendInternalEventType.RUN_STARTED, AgentBackendInternalEventType.STREAM_EVENT, + AgentBackendInternalEventType.AGENT_MESSAGE_DELTA, ): + if isinstance(internal_event, AgentBackendAgentMessageDeltaInternalEvent): + debounced_delta = text_delta_debouncer.push(internal_event.delta) + if debounced_delta: + persist_answer_text(debounced_delta) + continue + if isinstance(internal_event, AgentBackendStreamInternalEvent): + flush_pending_agent_message_text() try: process_recorder.handle_stream_event(internal_event) except Exception: @@ -826,26 +922,15 @@ class AgentAppRunner: internal_event.event_kind, exc_info=True, ) - text_delta = self._extract_stream_text_delta(internal_event) - if text_delta: - streamed_answer_parts.append(text_delta) - debounced_delta = text_delta_debouncer.push(text_delta) - if debounced_delta: - publish_text_delta( - queue_manager=queue_manager, - model_name=model_name, - delta=debounced_delta, - user_query=query, - ) continue continue - flush_pending_text() + flush_pending_agent_message_text() terminal = internal_event break if terminal is not None: break - flush_pending_text() - return terminal, "".join(streamed_answer_parts) + flush_pending_agent_message_text() + return terminal, process_recorder def _cancel_run(self, run_id: str) -> None: try: @@ -867,42 +952,15 @@ class AgentAppRunner: model_name: str, answer: str, query: str | None, - streamed_answer: str, usage: LLMUsage | None, ) -> None: - """Finish a successful streamed turn without duplicating the final text.""" - if not answer and streamed_answer: - answer = streamed_answer - - if not streamed_answer: - publish_text_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - user_query=query, - usage=usage, - ) - return - - if answer.startswith(streamed_answer): - publish_text_delta( - queue_manager=queue_manager, - model_name=model_name, - delta=answer[len(streamed_answer) :], - user_query=query, - ) - elif answer != streamed_answer: - logger.warning( - "Agent App streamed answer does not match terminal output; " - "using terminal output for message persistence." - ) - - publish_message_end( + """Finish a successful turn from the backend terminal output.""" + publish_text_answer( queue_manager=queue_manager, model_name=model_name, answer=answer, - user_query=query, usage=usage, + user_query=query, ) def _save_session( @@ -914,7 +972,7 @@ class AgentAppRunner: runtime_layer_specs: Any, pending_form_id: str | None = None, pending_tool_call_id: str | None = None, - ) -> None: + ) -> bool: try: self._session_store.save_active_snapshot( scope=scope, @@ -924,6 +982,7 @@ class AgentAppRunner: pending_form_id=pending_form_id, pending_tool_call_id=pending_tool_call_id, ) + return True except Exception: logger.warning( "Failed to persist Agent App conversation session snapshot: " @@ -934,9 +993,91 @@ class AgentAppRunner: scope.agent_id, exc_info=True, ) + return False + + def _load_superseded_sessions(self, *, scope: AgentAppSessionScope) -> list[StoredAgentAppSession]: + try: + stored_sessions = self._session_store.list_active_sessions_for_conversation( + tenant_id=scope.tenant_id, + app_id=scope.app_id, + conversation_id=scope.conversation_id, + ) + except Exception: + logger.warning( + "Failed to load existing Agent App conversation sessions before snapshot save: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s", + scope.tenant_id, + scope.app_id, + scope.conversation_id, + scope.agent_id, + exc_info=True, + ) + return [] + + return [stored for stored in stored_sessions if stored.scope != scope] + + def _cleanup_superseded_sessions(self, stored_sessions: list[StoredAgentAppSession]) -> None: + for stored_session in stored_sessions: + try: + if stored_session.runtime_layer_specs: + payload = AgentBackendSessionCleanupPayload( + session_snapshot=stored_session.session_snapshot, + runtime_layer_specs=stored_session.runtime_layer_specs, + idempotency_key=( + f"{stored_session.scope.tenant_id}:{stored_session.scope.app_id}:" + f"{stored_session.scope.conversation_id}:{stored_session.scope.agent_id}:" + f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" + f"superseded-session-cleanup:{stored_session.backend_run_id or 'no-run'}" + ), + metadata={ + "tenant_id": stored_session.scope.tenant_id, + "app_id": stored_session.scope.app_id, + "conversation_id": stored_session.scope.conversation_id, + "agent_id": stored_session.scope.agent_id, + "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, + "previous_agent_backend_run_id": stored_session.backend_run_id, + }, + ) + cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) + except Exception: + logger.warning( + "Failed to enqueue Agent backend cleanup for superseded Agent App session: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + stored_session.scope.tenant_id, + stored_session.scope.app_id, + stored_session.scope.conversation_id, + stored_session.scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) + + def _mark_session_cleaned( + self, + *, + scope: AgentAppSessionScope, + backend_run_id: str, + ) -> None: + """Best-effort delete-on-exit cleanup for the API-side session row. + + Once the Agent backend reaches a terminal event, cleanup persistence + must not replace the original publish/error outcome for that turn. + """ + try: + self._session_store.mark_cleaned(scope=scope, backend_run_id=backend_run_id) + except Exception: + logger.warning( + "Failed to retire Agent App conversation session after delete-on-exit: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + scope.tenant_id, + scope.app_id, + scope.conversation_id, + scope.agent_id, + backend_run_id, + exc_info=True, + ) @staticmethod - def _extract_answer(output: JsonValue) -> str: + def _terminal_output_to_answer(output: JsonValue) -> str: """Normalize the backend's terminal output to assistant text. Free-text Agent Apps return a plain string; if a structured output is @@ -954,27 +1095,5 @@ class AgentAppRunner: return json.dumps(output, ensure_ascii=False) return json.dumps(output, ensure_ascii=False) - @staticmethod - def _extract_stream_text_delta(event: AgentBackendStreamInternalEvent) -> str | None: - data = event.data - if not isinstance(data, dict): - return None - - if data.get("event_kind") == "part_delta": - delta = data.get("delta") - if isinstance(delta, dict) and delta.get("part_delta_kind") == "text": - content_delta = delta.get("content_delta") - if isinstance(content_delta, str): - return content_delta - - if data.get("event_kind") == "part_start": - part = data.get("part") - if isinstance(part, dict) and part.get("part_kind") == "text": - content = part.get("content") - if isinstance(content, str): - return content - - return None - __all__ = ["AgentAppRunner", "publish_message_end", "publish_text_answer", "publish_text_delta"] diff --git a/api/core/app/apps/agent_app/errors.py b/api/core/app/apps/agent_app/errors.py new file mode 100644 index 00000000000..51b4e77116a --- /dev/null +++ b/api/core/app/apps/agent_app/errors.py @@ -0,0 +1,6 @@ +class AgentAppGeneratorError(ValueError): + """Raised when an Agent App turn cannot be set up.""" + + +class AgentAppNotPublishedError(AgentAppGeneratorError): + """Raised when a public Agent App runtime is requested before publish.""" diff --git a/api/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py index fb1d054e750..c90ab77d681 100644 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ b/api/core/app/apps/agent_app/runtime_request_builder.py @@ -74,6 +74,7 @@ class AgentAppRuntimeBuildContext: session_snapshot: CompositorSessionSnapshot | None = None # ENG-638: set when resuming a chat turn after a submitted ask_human form. deferred_tool_results: DeferredToolResultsPayload | None = None + suspend_on_exit: bool = True @dataclass(frozen=True, slots=True) @@ -163,6 +164,7 @@ class AgentAppRuntimeRequestBuilder: # no frontend-internal {{#…#}} marker ever reaches the model. agent_soul_prompt=expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip() or None, + agent_config_version_kind=context.agent_config_version_kind, user_prompt=context.user_query, tools=tool_layers.plugin_tools, core_tools=tool_layers.core_tools, @@ -173,6 +175,7 @@ class AgentAppRuntimeRequestBuilder: shell_config=build_shell_layer_config(agent_soul), session_snapshot=context.session_snapshot, deferred_tool_results=context.deferred_tool_results, + suspend_on_exit=context.suspend_on_exit, idempotency_key=context.idempotency_key, metadata=metadata, ) diff --git a/api/core/app/apps/agent_app/session_store.py b/api/core/app/apps/agent_app/session_store.py index 35213114e2f..7696155a1c4 100644 --- a/api/core/app/apps/agent_app/session_store.py +++ b/api/core/app/apps/agent_app/session_store.py @@ -124,6 +124,41 @@ class AgentAppRuntimeSessionStore: runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), ) + def list_active_sessions_for_conversation( + self, *, tenant_id: str, app_id: str, conversation_id: str + ) -> list[StoredAgentAppSession]: + """List all ACTIVE conversation-owned sessions for lifecycle cleanup.""" + stmt = ( + select(AgentRuntimeSession) + .where( + AgentRuntimeSession.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION, + AgentRuntimeSession.tenant_id == tenant_id, + AgentRuntimeSession.app_id == app_id, + AgentRuntimeSession.conversation_id == conversation_id, + AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, + ) + .order_by(AgentRuntimeSession.updated_at.desc()) + ) + with session_factory.create_session() as session: + rows = session.scalars(stmt).all() + return [ + StoredAgentAppSession( + scope=AgentAppSessionScope( + tenant_id=row.tenant_id, + app_id=row.app_id, + conversation_id=row.conversation_id or "", + agent_id=row.agent_id, + agent_config_snapshot_id=row.agent_config_snapshot_id, + ), + session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), + backend_run_id=row.backend_run_id, + runtime_layer_specs=_deserialize_runtime_layer_specs(row.composition_layer_specs), + pending_form_id=row.pending_form_id, + pending_tool_call_id=row.pending_tool_call_id, + ) + for row in rows + ] + def save_active_snapshot( self, *, @@ -134,6 +169,14 @@ class AgentAppRuntimeSessionStore: pending_form_id: str | None = None, pending_tool_call_id: str | None = None, ) -> None: + """Persist the current conversation snapshot and enforce one ACTIVE row. + + Agent App chat treats one conversation as one resumable runtime shell. + Saving the latest snapshot therefore upserts the scoped row back to + ACTIVE and retires any other ACTIVE conversation-owned rows for the + same ``tenant_id + app_id + conversation_id`` so later lookups see a + single active session. + """ if snapshot is None: return snapshot_json = snapshot.model_dump_json() diff --git a/api/core/app/apps/agent_chat/app_generator.py b/api/core/app/apps/agent_chat/app_generator.py index d640bcdc863..a3cc913abf3 100644 --- a/api/core/app/apps/agent_chat/app_generator.py +++ b/api/core/app/apps/agent_chat/app_generator.py @@ -108,7 +108,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator): conversation_id = args.get("conversation_id") if conversation_id: conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=db.session() ) # get app model config app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation) diff --git a/api/core/app/apps/base_app_runner.py b/api/core/app/apps/base_app_runner.py index 7b854fec34a..941ae6b330b 100644 --- a/api/core/app/apps/base_app_runner.py +++ b/api/core/app/apps/base_app_runner.py @@ -5,6 +5,8 @@ from collections.abc import Generator, Mapping, Sequence from mimetypes import guess_extension from typing import TYPE_CHECKING, Any, Union +from sqlalchemy.orm import sessionmaker + from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError @@ -423,7 +425,9 @@ class AppRunner: _logger.exception("Failed to save image file") return - # Create MessageFile record + # Create MessageFile record. + # Use an independent session so this side-effect write does not + # commit or close the caller's request-scoped session. message_file = MessageFile( message_id=message_id, type=FileType.IMAGE, @@ -437,9 +441,8 @@ class AppRunner: created_by=user_id, ) - db.session.add(message_file) - db.session.commit() - db.session.refresh(message_file) + with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: + session.add(message_file) # Publish QueueMessageFileEvent queue_manager.publish( diff --git a/api/core/app/apps/chat/app_generator.py b/api/core/app/apps/chat/app_generator.py index 4873168b885..678525e0f77 100644 --- a/api/core/app/apps/chat/app_generator.py +++ b/api/core/app/apps/chat/app_generator.py @@ -105,7 +105,7 @@ class ChatAppGenerator(MessageBasedAppGenerator): conversation_id = args.get("conversation_id") if conversation_id: conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=db.session() ) # get app model config app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation) diff --git a/api/core/app/apps/chat/app_runner.py b/api/core/app/apps/chat/app_runner.py index 1e3be128296..3ee037c82c9 100644 --- a/api/core/app/apps/chat/app_runner.py +++ b/api/core/app/apps/chat/app_runner.py @@ -232,6 +232,7 @@ class ChatAppRunner(AppRunner): model_parameters=application_generate_entity.model_conf.parameters, stop=stop, stream=application_generate_entity.stream, + request_metadata={"app_id": app_config.app_id}, ) # handle invoke result diff --git a/api/core/app/apps/completion/app_runner.py b/api/core/app/apps/completion/app_runner.py index 3be70c860f9..b9c76569ba8 100644 --- a/api/core/app/apps/completion/app_runner.py +++ b/api/core/app/apps/completion/app_runner.py @@ -193,6 +193,7 @@ class CompletionAppRunner(AppRunner): model_parameters=application_generate_entity.model_conf.parameters, stop=stop, stream=application_generate_entity.stream, + request_metadata={"app_id": app_config.app_id}, ) # handle invoke result diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index ab07454ff5b..168b0e525d6 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -43,6 +43,7 @@ from core.repositories import DifyCoreRepositoryFactory from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from extensions.ext_database import db from factories import file_factory +from graphon.filters import ResponseStreamFilter from graphon.graph_engine.layers import GraphEngineLayer from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from graphon.runtime import GraphRuntimeState @@ -281,6 +282,7 @@ class WorkflowAppGenerator(BaseAppGenerator): graph_engine_layers: Sequence[GraphEngineLayer] = (), pause_state_config: PauseStateLayerConfig | None = None, variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER, + response_stream_filter: ResponseStreamFilter | None = None, ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]: """ Resume a paused workflow execution using the persisted runtime state. @@ -311,6 +313,7 @@ class WorkflowAppGenerator(BaseAppGenerator): graph_engine_layers=graph_engine_layers, graph_runtime_state=graph_runtime_state, pause_state_config=pause_state_config, + response_stream_filter=response_stream_filter, ) def _generate( @@ -329,6 +332,7 @@ class WorkflowAppGenerator(BaseAppGenerator): graph_engine_layers: Sequence[GraphEngineLayer] = (), graph_runtime_state: GraphRuntimeState | None = None, pause_state_config: PauseStateLayerConfig | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]: """ Generate App response. @@ -357,12 +361,14 @@ class WorkflowAppGenerator(BaseAppGenerator): app_mode=app_model.mode, ) + resolved_response_stream_filter = response_stream_filter or ResponseStreamFilter() if pause_state_config is not None: graph_layers.append( PauseStatePersistenceLayer( session_factory=pause_state_config.session_factory, generate_entity=application_generate_entity, state_owner_user_id=pause_state_config.state_owner_user_id, + response_stream_filter=resolved_response_stream_filter, ) ) @@ -385,6 +391,7 @@ class WorkflowAppGenerator(BaseAppGenerator): "workflow_node_execution_repository": workflow_node_execution_repository, "graph_engine_layers": tuple(graph_layers), "graph_runtime_state": graph_runtime_state, + "response_stream_filter": resolved_response_stream_filter, }, ) @@ -591,6 +598,7 @@ class WorkflowAppGenerator(BaseAppGenerator): root_node_id: str | None = None, graph_engine_layers: Sequence[GraphEngineLayer] = (), graph_runtime_state: GraphRuntimeState | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ) -> None: """ Generate worker in a new thread. @@ -639,6 +647,7 @@ class WorkflowAppGenerator(BaseAppGenerator): root_node_id=root_node_id, graph_engine_layers=graph_engine_layers, graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, ) try: diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index 6682a395a8c..95c9d777ebd 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -23,6 +23,7 @@ from extensions.ext_redis import redis_client from extensions.otel import WorkflowAppRunnerHandler, trace_span from extensions.workflow_warm_shutdown import WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, celery_warm_shutdown_started from graphon.enums import WorkflowType +from graphon.filters import ResponseStreamFilter from graphon.graph_engine.command_channels import RedisChannel from graphon.graph_engine.layers import GraphEngineLayer from graphon.runtime import GraphRuntimeState, VariablePool @@ -51,6 +52,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): workflow_node_execution_repository: WorkflowNodeExecutionRepository, graph_engine_layers: Sequence[GraphEngineLayer] = (), graph_runtime_state: GraphRuntimeState | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ): super().__init__( queue_manager=queue_manager, @@ -65,6 +67,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): self._workflow_execution_repository = workflow_execution_repository self._workflow_node_execution_repository = workflow_node_execution_repository self._resume_graph_runtime_state = graph_runtime_state + self._response_stream_filter = response_stream_filter @trace_span(WorkflowAppRunnerHandler) def run(self): @@ -177,6 +180,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): variable_pool=variable_pool, graph_runtime_state=graph_runtime_state, command_channel=command_channel, + response_stream_filter=self._response_stream_filter, ) persistence_layer = WorkflowPersistenceLayer( diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index 690f23c8302..b5f7515cd9f 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -15,6 +15,8 @@ if TYPE_CHECKING: DIFY_RUN_CONTEXT_KEY = "_dify" +AGENT_RUNTIME_EXIT_INTENT_ARG = "_agent_runtime_exit_intent" +type AgentRuntimeExitIntent = Literal["suspend", "delete"] class UserFrom(StrEnum): @@ -228,6 +230,10 @@ class AgentAppGenerateEntity(ChatAppGenerateEntity): ``agent_runtime_session_snapshot_id`` carries the runtime session scope used to resume or suspend within the same editable config surface. + ``agent_runtime_exit_intent`` is API-internal lifecycle policy for the + Agent backend session after this turn finishes. Normal chat/resume turns + suspend on exit; build-chat finalization deletes the backend runtime. + ``prompt_file_mappings`` preserves the raw request ``files`` array for the Agent backend prompt. These references are appended to the backend prompt text while the stored chat message keeps the user's original query. @@ -237,6 +243,7 @@ class AgentAppGenerateEntity(ChatAppGenerateEntity): agent_config_snapshot_id: str agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot" agent_runtime_session_snapshot_id: str | None = None + agent_runtime_exit_intent: AgentRuntimeExitIntent = "suspend" prompt_file_mappings: Sequence[JsonValue] = Field(default_factory=list) diff --git a/api/core/app/entities/task_entities.py b/api/core/app/entities/task_entities.py index eb96063a6e6..ba87ff7cb46 100644 --- a/api/core/app/entities/task_entities.py +++ b/api/core/app/entities/task_entities.py @@ -122,7 +122,7 @@ class MessageStreamResponse(StreamResponse): event: StreamEvent = StreamEvent.MESSAGE id: str answer: str - from_variable_selector: list[str] | None = None + from_variable_selector: list[str] = Field(default_factory=list) class MessageAudioStreamResponse(StreamResponse): @@ -151,7 +151,7 @@ class MessageEndStreamResponse(StreamResponse): event: StreamEvent = StreamEvent.MESSAGE_END id: str metadata: Mapping[str, object] = Field(default_factory=dict) - files: Sequence[Mapping[str, Any]] | None = None + files: Sequence[Mapping[str, Any]] = Field(default_factory=list) class MessageFileStreamResponse(StreamResponse): diff --git a/api/core/app/features/annotation_reply/annotation_reply.py b/api/core/app/features/annotation_reply/annotation_reply.py index 520ba7b85b3..9eff9747764 100644 --- a/api/core/app/features/annotation_reply/annotation_reply.py +++ b/api/core/app/features/annotation_reply/annotation_reply.py @@ -45,7 +45,7 @@ class AnnotationReplyFeature: embedding_model_name = collection_binding_detail.model_name dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding( - embedding_provider_name, embedding_model_name, db.session, CollectionBindingType.ANNOTATION + embedding_provider_name, embedding_model_name, db.session(), CollectionBindingType.ANNOTATION ) dataset = Dataset( @@ -66,7 +66,7 @@ class AnnotationReplyFeature: if documents and documents[0].metadata: annotation_id = documents[0].metadata["annotation_id"] score = documents[0].metadata["score"] - annotation = AppAnnotationService.get_annotation_by_id(annotation_id) + annotation = AppAnnotationService.get_annotation_by_id(annotation_id, session=db.session()) if annotation: if invoke_from in {InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP}: from_source = ConversationFromSource.API @@ -84,6 +84,7 @@ class AnnotationReplyFeature: message.id, from_source, score, + session=db.session(), ) return annotation diff --git a/api/core/app/layers/pause_state_persist_layer.py b/api/core/app/layers/pause_state_persist_layer.py index 2a13c73eccf..5d327adb065 100644 --- a/api/core/app/layers/pause_state_persist_layer.py +++ b/api/core/app/layers/pause_state_persist_layer.py @@ -9,6 +9,7 @@ from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, from core.repositories.human_input_repository import HumanInputFormSubmissionRepository from core.workflow.nodes.human_input.boundary import enrich_graph_pause_reasons from core.workflow.system_variables import SystemVariableKey, get_system_text +from graphon.filters import ResponseStreamFilter from graphon.graph_engine.layers import GraphEngineLayer from graphon.graph_events import GraphEngineEvent, GraphRunPausedEvent from models.model import AppMode @@ -43,6 +44,10 @@ class WorkflowResumptionContext(BaseModel): # Only workflow / chatflow could be paused. generate_entity: _GenerateEntityUnion serialized_graph_runtime_state: str + # Optional so that a workflow run paused before this field existed still + # loads: it just degrades to fresh-filter behavior on resume for that one + # stale run. + serialized_response_stream_filter_state: str | None = None def dumps(self) -> str: return self.model_dump_json() @@ -54,6 +59,12 @@ class WorkflowResumptionContext(BaseModel): def get_generate_entity(self) -> WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity: return self.generate_entity.entity + def get_response_stream_filter(self) -> ResponseStreamFilter: + response_stream_filter = ResponseStreamFilter() + if self.serialized_response_stream_filter_state is not None: + response_stream_filter.loads(self.serialized_response_stream_filter_state) + return response_stream_filter + @dataclass(frozen=True) class PauseStateLayerConfig: @@ -69,11 +80,17 @@ class PauseStatePersistenceLayer(GraphEngineLayer): session_factory: Engine | sessionmaker[Session], generate_entity: WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity, state_owner_user_id: str, + response_stream_filter: ResponseStreamFilter, ): """Create a PauseStatePersistenceLayer. The `state_owner_user_id` is used when creating state file for pause. It generally should id of the creator of workflow. + + `response_stream_filter` must be the exact same instance that + `WorkflowEntry` is using to stream this run's events — this layer + dumps its state on pause, and a different instance would silently + persist the wrong (empty) filter state. """ if isinstance(session_factory, Engine): session_factory = sessionmaker(session_factory) @@ -81,6 +98,7 @@ class PauseStatePersistenceLayer(GraphEngineLayer): self._session_maker = session_factory self._state_owner_user_id = state_owner_user_id self._generate_entity = generate_entity + self._response_stream_filter = response_stream_filter def _get_repo(self) -> APIWorkflowRunRepository: return DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._session_maker) @@ -121,6 +139,7 @@ class PauseStatePersistenceLayer(GraphEngineLayer): state = WorkflowResumptionContext( serialized_graph_runtime_state=self.graph_runtime_state.dumps(), generate_entity=entity_wrapper, + serialized_response_stream_filter_state=self._response_stream_filter.dumps(), ) workflow_run_id = get_system_text( diff --git a/api/core/app/llm/quota.py b/api/core/app/llm/quota.py index 5bf3334a7b2..d26d5d8a998 100644 --- a/api/core/app/llm/quota.py +++ b/api/core/app/llm/quota.py @@ -125,6 +125,7 @@ def _deduct_used_llm_quota(*, tenant_id: str, provider: str, provider_configurat CreditPoolService.deduct_credits_capped( tenant_id=tenant_id, credits_required=used_quota, + session=db.session(), ) case ProviderQuotaType.PAID: from services.credit_pool_service import CreditPoolService @@ -133,6 +134,7 @@ def _deduct_used_llm_quota(*, tenant_id: str, provider: str, provider_configurat tenant_id=tenant_id, credits_required=used_quota, pool_type="paid", + session=db.session(), ) case ProviderQuotaType.FREE: _deduct_free_llm_quota( diff --git a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py index a728069eede..c4224d67442 100644 --- a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py +++ b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py @@ -1,6 +1,6 @@ import logging import time -from collections.abc import Generator +from collections.abc import Generator, Mapping, Sequence from threading import Thread from typing import Any, cast @@ -44,7 +44,7 @@ from core.app.entities.task_entities import ( ) from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline from core.app.task_pipeline.message_cycle_manager import MessageCycleManager -from core.app.task_pipeline.message_file_utils import prepare_file_dict +from core.app.task_pipeline.message_file_utils import MessageFileInfoDict, prepare_file_dict from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk from core.model_manager import ModelInstance from core.ops.entities.trace_entity import TraceTaskName @@ -309,32 +309,20 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat response = self._message_cycle_manager.message_file_to_stream_response(event) if response: yield response - case QueueLLMChunkEvent() | QueueAgentMessageEvent(): + case QueueAgentMessageEvent(): chunk = event.chunk - delta_content = chunk.delta.message.content - if delta_content is None: + delta_text = self._chunk_delta_text(chunk) + if delta_text is None: + continue + yield self._agent_message_to_stream_response( + answer=delta_text, + message_id=self._message_id, + ) + case QueueLLMChunkEvent(): + chunk = event.chunk + delta_text = self._chunk_delta_text(chunk) + if delta_text is None: continue - if isinstance(delta_content, list): - # EasyUI streams text only; structured multimodal chunks contribute their text parts. - delta_text = "" - for content in delta_content: - logger.debug( - "The content type %s in LLM chunk delta message content.: %r", type(content), content - ) - match content: - case TextPromptMessageContent(): - delta_text += content.data - case str(): - delta_text += content # failback to str - case _: - logger.warning( - "Unsupported content type %s in LLM chunk delta message content.: %r", - type(content), - content, - ) - continue - else: - delta_text = delta_content if not self._task_state.llm_result.prompt_messages: self._task_state.llm_result.prompt_messages = chunk.prompt_messages @@ -348,23 +336,16 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat current_content += delta_text self._task_state.llm_result.message.content = current_content - match event: - case QueueLLMChunkEvent(): - # Determine the event type once, on first LLM chunk, and reuse for subsequent chunks - if not hasattr(self, "_precomputed_event_type") or self._precomputed_event_type is None: - self._precomputed_event_type = self._message_cycle_manager.get_message_event_type( - message_id=self._message_id - ) - yield self._message_cycle_manager.message_to_stream_response( - answer=delta_text, - message_id=self._message_id, - event_type=self._precomputed_event_type, - ) - case _: - yield self._agent_message_to_stream_response( - answer=delta_text, - message_id=self._message_id, - ) + # Determine the event type once, on first LLM chunk, and reuse for subsequent chunks + if not hasattr(self, "_precomputed_event_type") or self._precomputed_event_type is None: + self._precomputed_event_type = self._message_cycle_manager.get_message_event_type( + message_id=self._message_id + ) + yield self._message_cycle_manager.message_to_stream_response( + answer=delta_text, + message_id=self._message_id, + event_type=self._precomputed_event_type, + ) case QueueMessageReplaceEvent(): yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text) case QueuePingEvent(): @@ -376,6 +357,32 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat if self._conversation_name_generate_thread: logger.debug("Conversation name generation running as daemon thread") + @staticmethod + def _chunk_delta_text(chunk: LLMResultChunk) -> str | None: + delta_content = chunk.delta.message.content + if delta_content is None: + return None + if not isinstance(delta_content, list): + return delta_content + + delta_text = "" + # EasyUI streams text only; structured multimodal chunks contribute their text parts. + for content in delta_content: + logger.debug("The content type %s in LLM chunk delta message content.: %r", type(content), content) + match content: + case TextPromptMessageContent(): + delta_text += content.data + case str(): + delta_text += content + case _: + logger.warning( + "Unsupported content type %s in LLM chunk delta message content.: %r", + type(content), + content, + ) + continue + return delta_text + def _save_message(self, *, session: Session, trace_manager: TraceQueueManager | None = None): """ Save message. @@ -466,10 +473,10 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat :return: """ self._task_state.metadata.usage = self._task_state.llm_result.usage - metadata_dict = self._task_state.metadata.model_dump() + metadata_dict = self._task_state.metadata.model_dump(exclude_none=True) # Fetch files associated with this message - files = None + files: list[MessageFileInfoDict] = [] with Session(db.engine, expire_on_commit=False) as session: message_files = session.scalars(select(MessageFile).where(MessageFile.message_id == self._message_id)).all() @@ -492,13 +499,13 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat file_dict = prepare_file_dict(message_file, upload_files_map) files_list.append(file_dict) - files = files_list or None + files = files_list return MessageEndStreamResponse( task_id=self._application_generate_entity.task_id, id=self._message_id, metadata=metadata_dict, - files=files, + files=cast(Sequence[Mapping[str, Any]], files), ) def _agent_message_to_stream_response(self, answer: str, message_id: str) -> AgentMessageStreamResponse: @@ -528,11 +535,11 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerat task_id=self._application_generate_entity.task_id, id=agent_thought.id, position=agent_thought.position, - thought=agent_thought.thought, - observation=agent_thought.observation, - tool=agent_thought.tool, + thought=agent_thought.thought or "", + observation=agent_thought.observation or "", + tool=agent_thought.tool or "", tool_labels=agent_thought.tool_labels, - tool_input=agent_thought.tool_input, + tool_input=agent_thought.tool_input or "", message_files=agent_thought.files, ) diff --git a/api/core/app/task_pipeline/message_cycle_manager.py b/api/core/app/task_pipeline/message_cycle_manager.py index 62f27060b4e..5ada7d0ba2d 100644 --- a/api/core/app/task_pipeline/message_cycle_manager.py +++ b/api/core/app/task_pipeline/message_cycle_manager.py @@ -154,7 +154,7 @@ class MessageCycleManager: :param event: event :return: """ - annotation = AppAnnotationService.get_annotation_by_id(event.message_annotation_id) + annotation = AppAnnotationService.get_annotation_by_id(event.message_annotation_id, session=db.session()) if annotation: account = annotation.account self._task_state.metadata.annotation_reply = AnnotationReply( @@ -257,7 +257,7 @@ class MessageCycleManager: task_id=self._application_generate_entity.task_id, id=message_id, answer=answer, - from_variable_selector=from_variable_selector, + from_variable_selector=from_variable_selector or [], event=event_type or StreamEvent.MESSAGE, ) diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py index 5494769082e..d2024454a68 100644 --- a/api/core/callback_handler/index_tool_callback_handler.py +++ b/api/core/callback_handler/index_tool_callback_handler.py @@ -2,7 +2,7 @@ import logging from collections.abc import Sequence from sqlalchemy import select, update -from sqlalchemy.orm import scoped_session +from sqlalchemy.orm import Session, sessionmaker from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.entities.app_invoke_entities import InvokeFrom @@ -10,6 +10,7 @@ from core.app.entities.queue_entities import QueueRetrieverResourcesEvent from core.rag.entities import RetrievalSourceMetadata from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.models.document import Document +from extensions.ext_database import db from models.dataset import ChildChunk, DatasetQuery, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import CreatorUserRole, DatasetQuerySource @@ -29,7 +30,7 @@ class DatasetIndexToolCallbackHandler: self._user_id = user_id self._invoke_from = invoke_from - def on_query(self, query: str, dataset_id: str, session: scoped_session): + def on_query(self, query: str, dataset_id: str, session: Session): """ Handle query. """ @@ -46,47 +47,52 @@ class DatasetIndexToolCallbackHandler: created_by=self._user_id, ) - session.add(dataset_query) - session.commit() + # Use an independent session so this audit-log side effect does + # not commit or close the caller's request-scoped session. + with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as independent_session: + independent_session.add(dataset_query) - def on_tool_end(self, documents: list[Document], session: scoped_session): + def on_tool_end(self, documents: list[Document], session: Session): """Handle tool end.""" - for document in documents: - if document.metadata is not None: - document_id = document.metadata["document_id"] - dataset_document_stmt = select(DatasetDocument).where(DatasetDocument.id == document_id) - dataset_document = session.scalar(dataset_document_stmt) - if not dataset_document: - _logger.warning( - "Expected DatasetDocument record to exist, but none was found, document_id=%s", - document_id, - ) - continue - if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunk_stmt = select(ChildChunk).where( - ChildChunk.index_node_id == document.metadata["doc_id"], - ChildChunk.dataset_id == dataset_document.dataset_id, - ChildChunk.document_id == dataset_document.id, - ) - child_chunk = session.scalar(child_chunk_stmt) - if child_chunk: - session.execute( - update(DocumentSegment) - .where(DocumentSegment.id == child_chunk.segment_id) - .values(hit_count=DocumentSegment.hit_count + 1) + # Use an independent session so hit-count updates do not + # interfere with the caller's request-scoped session. + with Session(db.engine, expire_on_commit=False) as independent_session: + for document in documents: + if document.metadata is not None: + document_id = document.metadata["document_id"] + dataset_document_stmt = select(DatasetDocument).where(DatasetDocument.id == document_id) + dataset_document = independent_session.scalar(dataset_document_stmt) + if not dataset_document: + _logger.warning( + "Expected DatasetDocument record to exist, but none was found, document_id=%s", + document_id, ) - else: - conditions = [DocumentSegment.index_node_id == document.metadata["doc_id"]] + continue + if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: + child_chunk_stmt = select(ChildChunk).where( + ChildChunk.index_node_id == document.metadata["doc_id"], + ChildChunk.dataset_id == dataset_document.dataset_id, + ChildChunk.document_id == dataset_document.id, + ) + child_chunk = independent_session.scalar(child_chunk_stmt) + if child_chunk: + independent_session.execute( + update(DocumentSegment) + .where(DocumentSegment.id == child_chunk.segment_id) + .values(hit_count=DocumentSegment.hit_count + 1) + ) + else: + conditions = [DocumentSegment.index_node_id == document.metadata["doc_id"]] - if "dataset_id" in document.metadata: - conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"]) + if "dataset_id" in document.metadata: + conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"]) - # add hit count to document segment - session.execute( - update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1) - ) + # add hit count to document segment + independent_session.execute( + update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1) + ) - session.commit() + independent_session.commit() # TODO(-LAN-): Improve type check def return_retriever_resource_info(self, resource: Sequence[RetrievalSourceMetadata]): diff --git a/api/core/datasource/__base/datasource_provider.py b/api/core/datasource/__base/datasource_provider.py index 4b47777f0b9..d5633d9f081 100644 --- a/api/core/datasource/__base/datasource_provider.py +++ b/api/core/datasource/__base/datasource_provider.py @@ -3,7 +3,7 @@ from typing import Any from core.datasource.__base.datasource_plugin import DatasourcePlugin from core.datasource.entities.datasource_entities import DatasourceProviderEntityWithPlugin, DatasourceProviderType -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.plugin.impl.tool import PluginToolManager from core.tools.errors import ToolProviderCredentialValidationError @@ -78,11 +78,11 @@ class DatasourcePluginProviderController(ABC): if not credential_schema.required and credentials[credential_name] is None: continue - if credential_schema.type in {ProviderConfig.Type.SECRET_INPUT, ProviderConfig.Type.TEXT_INPUT}: + if credential_schema.type in {ProviderConfigType.SECRET_INPUT, ProviderConfigType.TEXT_INPUT}: if not isinstance(credentials[credential_name], str): raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") - elif credential_schema.type == ProviderConfig.Type.SELECT: + elif credential_schema.type == ProviderConfigType.SELECT: if not isinstance(credentials[credential_name], str): raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") @@ -107,9 +107,9 @@ class DatasourcePluginProviderController(ABC): default_value = credential_schema.default # parse default value into the correct type if credential_schema.type in { - ProviderConfig.Type.SECRET_INPUT, - ProviderConfig.Type.TEXT_INPUT, - ProviderConfig.Type.SELECT, + ProviderConfigType.SECRET_INPUT, + ProviderConfigType.TEXT_INPUT, + ProviderConfigType.SELECT, }: default_value = str(default_value) diff --git a/api/core/entities/mcp_provider.py b/api/core/entities/mcp_provider.py index 03b11c7263f..026022db1fb 100644 --- a/api/core/entities/mcp_provider.py +++ b/api/core/entities/mcp_provider.py @@ -9,7 +9,7 @@ from urllib.parse import urlparse from pydantic import BaseModel from configs import dify_config -from core.entities.provider_entities import BasicProviderConfig +from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType from core.helper import encrypter from core.helper.provider_cache import NoOpProviderCredentialCache from core.mcp.types import OAuthClientInformation, OAuthClientMetadata, OAuthTokens @@ -315,7 +315,7 @@ class MCPProviderEntity(BaseModel): return data # Create dynamic config only for encrypted fields - config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in encrypted_fields] + config = [BasicProviderConfig(type=ProviderConfigType.SECRET_INPUT, name=key) for key in encrypted_fields] encrypter_instance, _ = create_provider_encrypter( tenant_id=self.tenant_id, diff --git a/api/core/entities/provider_configuration.py b/api/core/entities/provider_configuration.py index 0c26a03d07b..95b8b686e05 100644 --- a/api/core/entities/provider_configuration.py +++ b/api/core/entities/provider_configuration.py @@ -1233,16 +1233,24 @@ class ProviderConfiguration(BaseModel): available_credentials_count = session.execute(count_stmt).scalar() or 0 session.delete(credential_record) + model_credentials_cache_identity_id: str | None = None + if provider_model_record and ( + available_credentials_count <= 1 or provider_model_record.credential_id == credential_id + ): + model_credentials_cache_identity_id = provider_model_record.id + if provider_model_record and available_credentials_count <= 1: # If all credentials are deleted, delete the custom model record session.delete(provider_model_record) elif provider_model_record and provider_model_record.credential_id == credential_id: provider_model_record.credential_id = None provider_model_record.updated_at = naive_utc_now() + + if model_credentials_cache_identity_id: provider_model_credentials_cache = ProviderCredentialsCache( tenant_id=self.tenant_id, - identity_id=provider_model_record.id, - cache_type=ProviderCredentialsCacheType.PROVIDER, + identity_id=model_credentials_cache_identity_id, + cache_type=ProviderCredentialsCacheType.MODEL, ) provider_model_credentials_cache.delete() diff --git a/api/core/entities/provider_entities.py b/api/core/entities/provider_entities.py index 4c9cd858e13..4ac7614425c 100644 --- a/api/core/entities/provider_entities.py +++ b/api/core/entities/provider_entities.py @@ -165,34 +165,35 @@ class ModelSettings(BaseModel): model_config = ConfigDict(protected_namespaces=()) +class ProviderConfigType(StrEnum): + SECRET_INPUT = CommonParameterType.SECRET_INPUT + TEXT_INPUT = CommonParameterType.TEXT_INPUT + SELECT = CommonParameterType.SELECT + BOOLEAN = CommonParameterType.BOOLEAN + APP_SELECTOR = CommonParameterType.APP_SELECTOR + MODEL_SELECTOR = CommonParameterType.MODEL_SELECTOR + TOOLS_SELECTOR = CommonParameterType.TOOLS_SELECTOR + + @classmethod + def value_of(cls, value: str) -> ProviderConfigType: + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + class BasicProviderConfig(BaseModel): """ Base model class for common provider settings like credentials """ - class Type(StrEnum): - SECRET_INPUT = CommonParameterType.SECRET_INPUT - TEXT_INPUT = CommonParameterType.TEXT_INPUT - SELECT = CommonParameterType.SELECT - BOOLEAN = CommonParameterType.BOOLEAN - APP_SELECTOR = CommonParameterType.APP_SELECTOR - MODEL_SELECTOR = CommonParameterType.MODEL_SELECTOR - TOOLS_SELECTOR = CommonParameterType.TOOLS_SELECTOR - - @classmethod - def value_of(cls, value: str) -> ProviderConfig.Type: - """ - Get value of given mode. - - :param value: mode value - :return: mode - """ - for mode in cls: - if mode.value == value: - return mode - raise ValueError(f"invalid mode value {value}") - - type: Type = Field(..., description="The type of the credentials") + type: ProviderConfigType = Field(..., description="The type of the credentials") name: str = Field(..., description="The name of the credentials") diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py index 951e065b2cb..c30afc0e745 100644 --- a/api/core/helper/code_executor/code_executor.py +++ b/api/core/helper/code_executor/code_executor.py @@ -13,7 +13,7 @@ from core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTr from core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformer from core.helper.code_executor.template_transformer import TemplateTransformer from core.helper.http_client_pooling import get_pooled_http_client -from graphon.nodes.code.entities import CodeLanguage +from graphon.nodes.code.entities import CodeLanguage as CodeLanguage # noqa: PLC0414 logger = logging.getLogger(__name__) code_execution_endpoint_url = URL(str(dify_config.CODE_EXECUTION_ENDPOINT)) @@ -74,12 +74,16 @@ class CodeExecutor: :param code: code :return: """ + running_language = cls.code_language_to_running_language.get(language) + if running_language is None: + raise CodeExecutionError(f"Unsupported language {language}") + url = code_execution_endpoint_url / "v1" / "sandbox" / "run" headers = {"X-Api-Key": dify_config.CODE_EXECUTION_API_KEY} data = { - "language": cls.code_language_to_running_language.get(language), + "language": running_language, "code": code, "preload": preload, "enable_network": True, @@ -133,7 +137,9 @@ class CodeExecutor: return response_code.data.stdout or "" @classmethod - def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]): + def execute_workflow_code_template( + cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any] + ) -> dict[str, Any]: """ Execute code :param language: code language diff --git a/api/core/helper/code_executor/jinja2/jinja2_transformer.py b/api/core/helper/code_executor/jinja2/jinja2_transformer.py index 9cf5089f7b5..d1c75c981b6 100644 --- a/api/core/helper/code_executor/jinja2/jinja2_transformer.py +++ b/api/core/helper/code_executor/jinja2/jinja2_transformer.py @@ -11,7 +11,7 @@ class Jinja2TemplateTransformer(TemplateTransformer): @classmethod @override - def transform_response(cls, response: str): + def transform_response(cls, response: str) -> dict[str, Any]: """ Transform response to dict :param response: response diff --git a/api/core/helper/code_executor/template_transformer.py b/api/core/helper/code_executor/template_transformer.py index 3a6c314159b..501f460ba95 100644 --- a/api/core/helper/code_executor/template_transformer.py +++ b/api/core/helper/code_executor/template_transformer.py @@ -36,14 +36,14 @@ class TemplateTransformer(ABC): return runner_script, preload_script @classmethod - def extract_result_str_from_response(cls, response: str): + def extract_result_str_from_response(cls, response: str) -> str: result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL) if not result: raise ValueError(f"Failed to parse result: no result tag found in response. Response: {response[:200]}...") return result.group(1) @classmethod - def transform_response(cls, response: str) -> Mapping[str, Any]: + def transform_response(cls, response: str) -> dict[str, Any]: """ Transform response to dict :param response: response @@ -71,7 +71,7 @@ class TemplateTransformer(ABC): return result @classmethod - def _post_process_result(cls, result: dict[Any, Any]) -> dict[Any, Any]: + def _post_process_result(cls, result: dict[str, Any]) -> dict[str, Any]: """ Post-process the result to convert scientific notation strings back to numbers """ @@ -89,7 +89,7 @@ class TemplateTransformer(ABC): return [convert_scientific_notation(v) for v in value] return value - return convert_scientific_notation(result) + return {key: convert_scientific_notation(value) for key, value in result.items()} @classmethod @abstractmethod diff --git a/api/core/helper/creators.py b/api/core/helper/creators.py index b01e16f18a7..4ad61371512 100644 --- a/api/core/helper/creators.py +++ b/api/core/helper/creators.py @@ -24,7 +24,7 @@ def upload_dsl(dsl_file_bytes: bytes, filename: str = "template.yaml") -> str: response.raise_for_status() data = response.json() claim_code = data.get("data", {}).get("claim_code") - if not claim_code: + if not isinstance(claim_code, str) or not claim_code: raise ValueError("Creators Platform did not return a valid claim_code") return claim_code diff --git a/api/core/helper/credential_utils.py b/api/core/helper/credential_utils.py index e8f3ba0a547..a57474a8c12 100644 --- a/api/core/helper/credential_utils.py +++ b/api/core/helper/credential_utils.py @@ -45,7 +45,7 @@ def is_credential_exists(credential_id: str, credential_type: "PluginCredentialT def runtime_check_credential_policy_compliance( credential_id: str, provider: str, credential_type: "PluginCredentialType", check_existence: bool = True -): +) -> None: if dify_config.ENTERPRISE_DISABLE_RUNTIME_CREDENTIAL_CHECK: return check_credential_policy_compliance( diff --git a/api/core/helper/download.py b/api/core/helper/download.py index 364d45b1e9e..e99be256a01 100644 --- a/api/core/helper/download.py +++ b/api/core/helper/download.py @@ -1,4 +1,7 @@ -def download_with_size_limit(url, max_download_size: int, **kwargs): +from typing import Any + + +def download_with_size_limit(url: str, max_download_size: int, **kwargs: Any) -> bytes: from core.file import remote_fetcher response = remote_fetcher.make_request("GET", url, follow_redirects=True, **kwargs) diff --git a/api/core/helper/encrypter.py b/api/core/helper/encrypter.py index 20125ec6b30..f72f6c6be9f 100644 --- a/api/core/helper/encrypter.py +++ b/api/core/helper/encrypter.py @@ -1,5 +1,7 @@ import base64 +from Crypto.PublicKey import RSA + from libs import rsa @@ -11,13 +13,13 @@ def obfuscated_token(token: str) -> str: return token[:6] + "*" * 12 + token[-2:] -def full_mask_token(token_length=20): +def full_mask_token(token_length: int = 20) -> str: return "*" * token_length -def encrypt_token(tenant_id: str, token: str): - from extensions.ext_database import db +def encrypt_token(tenant_id: str, token: str) -> str: from models.account import Tenant + from models.engine import db if not (tenant := db.session.get(Tenant, tenant_id)): raise ValueError(f"Tenant with id {tenant_id} not found") @@ -30,15 +32,15 @@ def decrypt_token(tenant_id: str, token: str) -> str: return rsa.decrypt(base64.b64decode(token), tenant_id) -def batch_decrypt_token(tenant_id: str, tokens: list[str]): +def batch_decrypt_token(tenant_id: str, tokens: list[str]) -> list[str]: rsa_key, cipher_rsa = rsa.get_decrypt_decoding(tenant_id) return [rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) for token in tokens] -def get_decrypt_decoding(tenant_id: str): +def get_decrypt_decoding(tenant_id: str) -> tuple[RSA.RsaKey, object]: return rsa.get_decrypt_decoding(tenant_id) -def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa): +def decrypt_token_with_decoding(token: str, rsa_key: RSA.RsaKey, cipher_rsa: object) -> str: return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) diff --git a/api/core/helper/marketplace.py b/api/core/helper/marketplace.py index 0b77891ce16..e6e1d565769 100644 --- a/api/core/helper/marketplace.py +++ b/api/core/helper/marketplace.py @@ -1,5 +1,6 @@ import logging from collections.abc import Sequence +from typing import Any from urllib.parse import urlencode import httpx @@ -21,7 +22,7 @@ def get_plugin_pkg_url(plugin_unique_identifier: str) -> str: return f"{marketplace_api_url / 'api/v1/plugins/download'}?{query}" -def download_plugin_pkg(plugin_unique_identifier: str): +def download_plugin_pkg(plugin_unique_identifier: str) -> bytes: return download_with_size_limit(get_plugin_pkg_url(plugin_unique_identifier), dify_config.PLUGIN_MAX_PACKAGE_SIZE) @@ -41,7 +42,7 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP return [MarketplacePluginDeclaration.model_validate(plugin) for plugin in response.json()["data"]["plugins"]] -def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]: +def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict[str, Any]]: if not plugin_ids: return [] @@ -55,10 +56,19 @@ def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]: response.raise_for_status() data = response.json() - return data.get("data", {}).get("plugins", []) + plugins = data.get("data", {}).get("plugins", []) + if not isinstance(plugins, list): + raise ValueError("Marketplace did not return a valid plugins list") + + result: list[dict[str, Any]] = [] + for plugin in plugins: + if not isinstance(plugin, dict) or not all(isinstance(key, str) for key in plugin): + raise ValueError("Marketplace did not return a valid plugins list") + result.append(plugin) + return result -def record_install_plugin_event(plugin_unique_identifier: str): +def record_install_plugin_event(plugin_unique_identifier: str) -> None: url = str(marketplace_api_url / "api/v1/stats/plugins/install_count") response = httpx.post(url, json={"unique_identifier": plugin_unique_identifier}, timeout=MARKETPLACE_TIMEOUT) response.raise_for_status() diff --git a/api/core/helper/model_provider_cache.py b/api/core/helper/model_provider_cache.py index 10d79a82392..2b9e6613378 100644 --- a/api/core/helper/model_provider_cache.py +++ b/api/core/helper/model_provider_cache.py @@ -34,7 +34,7 @@ class ProviderCredentialsCache: else: return None - def set(self, credentials: dict[str, Any]): + def set(self, credentials: dict[str, Any]) -> None: """ Cache model provider credentials. @@ -43,7 +43,7 @@ class ProviderCredentialsCache: """ redis_client.setex(self.cache_key, 86400, json.dumps(credentials)) - def delete(self): + def delete(self) -> None: """ Delete cached model provider credentials. diff --git a/api/core/helper/module_import_helper.py b/api/core/helper/module_import_helper.py index 768210d899b..4d37abf4897 100644 --- a/api/core/helper/module_import_helper.py +++ b/api/core/helper/module_import_helper.py @@ -20,17 +20,18 @@ def import_module_from_source[T: (str, bytes)]( raise Exception(f"Failed to load module {module_name} from {py_file_path!r}") else: # Refer to: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly - # FIXME: mypy does not support the type of spec.loader - spec = importlib.util.spec_from_file_location(module_name, py_file_path) # type: ignore[assignment] - if not spec or not spec.loader: + new_spec = importlib.util.spec_from_file_location(module_name, py_file_path) + if not new_spec or not new_spec.loader: raise Exception(f"Failed to load module {module_name} from {py_file_path!r}") if use_lazy_loader: # Refer to: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports - spec.loader = importlib.util.LazyLoader(spec.loader) + new_spec.loader = importlib.util.LazyLoader(new_spec.loader) + spec = new_spec module = importlib.util.module_from_spec(spec) if not existed_spec: sys.modules[module_name] = module - spec.loader.exec_module(module) + if spec.loader is not None: + spec.loader.exec_module(module) return module except Exception as e: logger.exception("Failed to load module %s from script file '%s'", module_name, repr(py_file_path)) diff --git a/api/core/helper/provider_cache.py b/api/core/helper/provider_cache.py index 6ad08dfe178..a3b61887892 100644 --- a/api/core/helper/provider_cache.py +++ b/api/core/helper/provider_cache.py @@ -9,11 +9,11 @@ from extensions.ext_redis import redis_client class ProviderCredentialsCache(ABC): """Base class for provider credentials cache""" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: self.cache_key = self._generate_cache_key(**kwargs) @abstractmethod - def _generate_cache_key(self, **kwargs) -> str: + def _generate_cache_key(self, **kwargs: Any) -> str: """Generate cache key based on subclass implementation""" pass @@ -28,11 +28,11 @@ class ProviderCredentialsCache(ABC): return None return None - def set(self, config: dict[str, Any]): + def set(self, config: dict[str, Any]) -> None: """Cache provider credentials""" redis_client.setex(self.cache_key, 86400, json.dumps(config)) - def delete(self): + def delete(self) -> None: """Delete cached provider credentials""" redis_client.delete(self.cache_key) @@ -48,7 +48,7 @@ class SingletonProviderCredentialsCache(ProviderCredentialsCache): ) @override - def _generate_cache_key(self, **kwargs) -> str: + def _generate_cache_key(self, **kwargs: Any) -> str: tenant_id = kwargs["tenant_id"] provider_type = kwargs["provider_type"] identity_name = kwargs["provider_identity"] @@ -63,7 +63,7 @@ class ToolProviderCredentialsCache(ProviderCredentialsCache): super().__init__(tenant_id=tenant_id, provider=provider, credential_id=credential_id) @override - def _generate_cache_key(self, **kwargs) -> str: + def _generate_cache_key(self, **kwargs: Any) -> str: tenant_id = kwargs["tenant_id"] provider = kwargs["provider"] credential_id = kwargs["credential_id"] @@ -77,10 +77,10 @@ class NoOpProviderCredentialCache: """Get cached provider credentials""" return None - def set(self, config: dict[str, Any]): + def set(self, config: dict[str, Any]) -> None: """Cache provider credentials""" pass - def delete(self): + def delete(self) -> None: """Delete cached provider credentials""" pass diff --git a/api/core/helper/provider_encryption.py b/api/core/helper/provider_encryption.py index 8484a28c05f..de9d5e15be1 100644 --- a/api/core/helper/provider_encryption.py +++ b/api/core/helper/provider_encryption.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from copy import deepcopy from typing import Any, Protocol -from core.entities.provider_entities import BasicProviderConfig +from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType from core.helper import encrypter @@ -60,7 +60,7 @@ class ProviderConfigEncrypter: fields[credential.name] = credential for field_name, field in fields.items(): - if field.type == BasicProviderConfig.Type.SECRET_INPUT: + if field.type == ProviderConfigType.SECRET_INPUT: if field_name in data: encrypted = encrypter.encrypt_token(self.tenant_id, data[field_name] or "") data[field_name] = encrypted @@ -81,7 +81,7 @@ class ProviderConfigEncrypter: fields[credential.name] = credential for field_name, field in fields.items(): - if field.type == BasicProviderConfig.Type.SECRET_INPUT: + if field.type == ProviderConfigType.SECRET_INPUT: if field_name in data: if len(data[field_name]) > 6: data[field_name] = ( @@ -112,7 +112,7 @@ class ProviderConfigEncrypter: fields[credential.name] = credential for field_name, field in fields.items(): - if field.type == BasicProviderConfig.Type.SECRET_INPUT: + if field.type == ProviderConfigType.SECRET_INPUT: if field_name in data: with contextlib.suppress(Exception): # if the value is None or empty string, skip decrypt @@ -125,5 +125,7 @@ class ProviderConfigEncrypter: return data -def create_provider_encrypter(tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache): +def create_provider_encrypter( + tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache +) -> tuple[ProviderConfigEncrypter, ProviderConfigCache]: return ProviderConfigEncrypter(tenant_id=tenant_id, config=config, provider_config_cache=cache), cache diff --git a/api/core/helper/tool_parameter_cache.py b/api/core/helper/tool_parameter_cache.py index bf5bf9af03b..2650eb0c2c2 100644 --- a/api/core/helper/tool_parameter_cache.py +++ b/api/core/helper/tool_parameter_cache.py @@ -37,11 +37,11 @@ class ToolParameterCache: else: return None - def set(self, parameters: dict[str, Any]): + def set(self, parameters: dict[str, Any]) -> None: """Cache model provider credentials.""" redis_client.setex(self.cache_key, 86400, json.dumps(parameters)) - def delete(self): + def delete(self) -> None: """ Delete cached model provider credentials. diff --git a/api/core/helper/trace_id_helper.py b/api/core/helper/trace_id_helper.py index 8b022c1d065..e1ebd45e074 100644 --- a/api/core/helper/trace_id_helper.py +++ b/api/core/helper/trace_id_helper.py @@ -61,7 +61,7 @@ def get_external_trace_id(request: Any) -> str | None: return None -def extract_external_trace_id_from_args(args: Mapping[str, Any]): +def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict[str, Any]: """ Extract 'external_trace_id' from args. diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py index f97f9c38330..29a93fff815 100644 --- a/api/core/llm_generator/llm_generator.py +++ b/api/core/llm_generator/llm_generator.py @@ -6,6 +6,7 @@ from typing import Any, Literal, NotRequired, Protocol, TypedDict, cast import json_repair from sqlalchemy import select +from sqlalchemy.orm import Session from core.app.app_config.entities import ModelConfig from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload @@ -117,7 +118,9 @@ def _parse_string_list(text: str) -> list[str]: class WorkflowServiceInterface(Protocol): - def get_draft_workflow(self, app_model: App, workflow_id: str | None = None) -> Workflow | None: + def get_draft_workflow( + self, app_model: App, workflow_id: str | None = None, *, session: Session + ) -> Workflow | None: pass def get_node_last_run(self, app_model: App, workflow: Workflow, node_id: str) -> WorkflowNodeExecutionModel | None: @@ -758,7 +761,7 @@ class LLMGenerator: app: App | None = session.scalar(select(App).where(App.id == flow_id, App.tenant_id == tenant_id).limit(1)) if not app: raise ValueError("App not found.") - workflow = workflow_service.get_draft_workflow(app_model=app) + workflow = workflow_service.get_draft_workflow(app_model=app, session=session) if not workflow: raise ValueError("Workflow not found for the given app model.") last_run = workflow_service.get_node_last_run(app_model=app, workflow=workflow, node_id=node_id) diff --git a/api/core/mcp/server/streamable_http.py b/api/core/mcp/server/streamable_http.py index 3bb75e485a8..7fd03788c7e 100644 --- a/api/core/mcp/server/streamable_http.py +++ b/api/core/mcp/server/streamable_http.py @@ -15,6 +15,36 @@ from services.app_generate_service import AppGenerateService logger = logging.getLogger(__name__) +# Structured tool output (outputSchema + structuredContent) was introduced in MCP 2025-06-18. +STRUCTURED_OUTPUT_MIN_VERSION = "2025-06-18" + + +def _supports_structured_output(protocol_version: str) -> bool: + """Return True when the negotiated protocol version supports structured tool output. + + MCP protocol versions are YYYY-MM-DD strings, so lexical comparison equals chronological. + """ + return protocol_version >= STRUCTURED_OUTPUT_MIN_VERSION + + +def negotiate_protocol_version(header_value: str | None, is_initialize: bool) -> str | None: + """Resolve the negotiated protocol version for an incoming MCP request. + + The version is taken from the MCP-Protocol-Version header on post-initialize requests. + Returns the version to use for behavior gating, or None when the client sent an explicit + but unsupported header (the caller should reply with a JSON-RPC INVALID_REQUEST error). + Initialize requests negotiate via the request body, so they always receive + DEFAULT_NEGOTIATED_VERSION and their header is never validated or rejected. + """ + if is_initialize: + return mcp_types.DEFAULT_NEGOTIATED_VERSION + # Treat an absent or empty header as "not specified" -> default version. + if not header_value: + return mcp_types.DEFAULT_NEGOTIATED_VERSION + if header_value not in mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS: + return None + return header_value + class ToolParameterSchemaDict(TypedDict): type: str @@ -35,6 +65,7 @@ def handle_mcp_request( mcp_server: AppMCPServer, end_user: EndUser | None = None, request_id: int | str = 1, + protocol_version: str = mcp_types.DEFAULT_NEGOTIATED_VERSION, ) -> mcp_types.JSONRPCResponse | mcp_types.JSONRPCError: """ Handle MCP request and return JSON-RPC response @@ -77,15 +108,24 @@ def handle_mcp_request( # Dispatch request to appropriate handler based on instance type match request_root: case mcp_types.InitializeRequest(): - return create_success_response(handle_initialize(mcp_server.description)) + return create_success_response( + handle_initialize(mcp_server.description, request_root.params.protocolVersion) + ) case mcp_types.ListToolsRequest(): return create_success_response( handle_list_tools( - app.name, app.mode, user_input_form, mcp_server.description, mcp_server.parameters_dict + app.name, + app.mode, + user_input_form, + mcp_server.description, + mcp_server.parameters_dict, + protocol_version, ) ) case mcp_types.CallToolRequest(): - return create_success_response(handle_call_tool(session, app, request, user_input_form, end_user)) + return create_success_response( + handle_call_tool(session, app, request, user_input_form, end_user, protocol_version) + ) case mcp_types.PingRequest(): return create_success_response(handle_ping()) case _: @@ -104,14 +144,22 @@ def handle_ping() -> mcp_types.EmptyResult: return mcp_types.EmptyResult() -def handle_initialize(description: str) -> mcp_types.InitializeResult: - """Handle initialize request""" +def handle_initialize(description: str, requested_version: str | int) -> mcp_types.InitializeResult: + """Handle initialize request, negotiating the protocol version with the client. + + Echoes the client's requested version when the server supports it, otherwise returns the + server's latest supported version (per the MCP lifecycle spec). + """ + negotiated_version: str = mcp_types.SERVER_LATEST_PROTOCOL_VERSION + if isinstance(requested_version, str) and requested_version in mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS: + negotiated_version = requested_version + capabilities = mcp_types.ServerCapabilities( tools=mcp_types.ToolsCapability(listChanged=False), ) return mcp_types.InitializeResult( - protocolVersion=mcp_types.SERVER_LATEST_PROTOCOL_VERSION, + protocolVersion=negotiated_version, capabilities=capabilities, serverInfo=mcp_types.Implementation(name="Dify", version=dify_config.project.version), instructions=description, @@ -124,19 +172,23 @@ def handle_list_tools( user_input_form: list[VariableEntity], description: str, parameters_dict: dict[str, str], + protocol_version: str = mcp_types.DEFAULT_NEGOTIATED_VERSION, ) -> mcp_types.ListToolsResult: """Handle list tools request""" parameter_schema = build_parameter_schema(app_mode, user_input_form, parameters_dict) + supports_structured = _supports_structured_output(protocol_version) - return mcp_types.ListToolsResult( - tools=[ - mcp_types.Tool( - name=app_name, - description=description, - inputSchema=cast(dict[str, Any], parameter_schema), - ) - ], + # For 2025-06-18+ clients, expose an explicit display title and a permissive output + # schema. Both stay None (and are stripped by exclude_none serialization) for older + # clients, so their tool definition is unchanged. + tool = mcp_types.Tool( + name=app_name, + title=app_name if supports_structured else None, + description=description, + inputSchema=cast(dict[str, Any], parameter_schema), + outputSchema={"type": "object"} if supports_structured else None, ) + return mcp_types.ListToolsResult(tools=[tool]) def handle_call_tool( @@ -145,6 +197,7 @@ def handle_call_tool( request: mcp_types.ClientRequest, user_input_form: list[VariableEntity], end_user: EndUser | None, + protocol_version: str = mcp_types.DEFAULT_NEGOTIATED_VERSION, ) -> mcp_types.CallToolResult: """Handle call tool request""" request_obj = cast(mcp_types.CallToolRequest, request.root) @@ -154,16 +207,22 @@ def handle_call_tool( raise ValueError("End user not found") response = AppGenerateService.generate( - session, - app, - end_user, - args, - InvokeFrom.SERVICE_API, + session=session, + app_model=app, + user=end_user, + args=args, + invoke_from=InvokeFrom.SERVICE_API, streaming=app.mode == AppMode.AGENT_CHAT, ) answer = extract_answer_from_response(app, response) - return mcp_types.CallToolResult(content=[mcp_types.TextContent(text=answer, type="text")]) + structured_content = None + if _supports_structured_output(protocol_version): + structured_content = extract_structured_output(app, response, answer) + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(text=answer, type="text")], + structuredContent=structured_content, + ) def build_parameter_schema( @@ -204,6 +263,29 @@ def prepare_tool_arguments(app: App, arguments: dict[str, Any]) -> ToolArguments return {"query": query, "inputs": args_copy} +def extract_structured_output(app: App, response: Any, answer: str) -> dict[str, Any] | None: + """Build MCP structured tool output (2025-06-18) from the app response. + + WORKFLOW mode exposes the raw outputs mapping; chat/agent/completion modes expose the + answer string under an "answer" key. Returns None when no structured output is available. + """ + match app.mode: + case AppMode.WORKFLOW: + if isinstance(response, Mapping): + data = response.get("data") + if isinstance(data, Mapping): + outputs = data.get("outputs") + # All three guards use Mapping for consistency; coerce to a concrete dict + # because structuredContent must be a JSON object (dict[str, Any]). + if isinstance(outputs, Mapping): + return dict(outputs) + return None + case AppMode.ADVANCED_CHAT | AppMode.CHAT | AppMode.AGENT_CHAT | AppMode.COMPLETION: + return {"answer": answer} + case _: + return None + + def extract_answer_from_response(app: App, response: Any) -> str: """Extract answer from app generate response""" answer = "" diff --git a/api/core/mcp/types.py b/api/core/mcp/types.py index 9470d39f414..8d1e6587c5f 100644 --- a/api/core/mcp/types.py +++ b/api/core/mcp/types.py @@ -22,10 +22,13 @@ for reference. * Define additional model classes instead of using dictionaries. Do this even if they're not separate types in the schema. """ -# Client support both version, not support 2025-06-18 yet. +# Latest protocol version the Dify MCP client negotiates with upstream MCP servers. LATEST_PROTOCOL_VERSION = "2025-06-18" -# Server support 2024-11-05 to allow claude to use. -SERVER_LATEST_PROTOCOL_VERSION = "2024-11-05" +# Latest protocol version the Dify MCP server advertises to connecting clients. +SERVER_LATEST_PROTOCOL_VERSION = "2025-06-18" +# Protocol versions the Dify MCP server can negotiate down to (e.g. Claude on 2024-11-05). +SERVER_SUPPORTED_PROTOCOL_VERSIONS: frozenset[str] = frozenset({"2024-11-05", "2025-03-26", "2025-06-18"}) +# Version assumed when a client omits the MCP-Protocol-Version header on post-initialize requests. DEFAULT_NEGOTIATED_VERSION = "2025-03-26" ProgressToken = str | int Cursor = str diff --git a/api/core/model_manager.py b/api/core/model_manager.py index 56a8f3bd98c..29113ac6b2c 100644 --- a/api/core/model_manager.py +++ b/api/core/model_manager.py @@ -124,6 +124,7 @@ class ModelInstance: stop: list[str] | None = None, stream: Literal[True] = True, callbacks: list[Callback] | None = None, + request_metadata: Mapping[str, object] | None = None, ) -> Generator: ... @overload @@ -135,6 +136,7 @@ class ModelInstance: stop: list[str] | None = None, stream: Literal[False] = False, callbacks: list[Callback] | None = None, + request_metadata: Mapping[str, object] | None = None, ) -> LLMResult: ... @overload @@ -146,6 +148,7 @@ class ModelInstance: stop: list[str] | None = None, stream: bool = True, callbacks: list[Callback] | None = None, + request_metadata: Mapping[str, object] | None = None, ) -> Union[LLMResult, Generator]: ... def invoke_llm( @@ -156,6 +159,7 @@ class ModelInstance: stop: Sequence[str] | None = None, stream: bool = True, callbacks: list[Callback] | None = None, + request_metadata: Mapping[str, object] | None = None, ) -> Union[LLMResult, Generator]: """ Invoke large language model @@ -166,6 +170,7 @@ class ModelInstance: :param stop: stop words :param stream: is stream response :param callbacks: callbacks + :param request_metadata: optional request metadata :return: full response or stream response chunk generator result """ if not isinstance(self.model_type_instance, LargeLanguageModel): @@ -182,6 +187,7 @@ class ModelInstance: stop=list(stop) if stop else None, stream=stream, callbacks=callbacks, + request_metadata=request_metadata, ), ) diff --git a/api/core/plugin/backwards_invocation/app.py b/api/core/plugin/backwards_invocation/app.py index 046be355daf..a74be9be2d8 100644 --- a/api/core/plugin/backwards_invocation/app.py +++ b/api/core/plugin/backwards_invocation/app.py @@ -78,7 +78,11 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation): if not user_id: user = EndUserService.get_or_create_end_user(app) else: - user = cls._get_user(user_id, app) + try: + user = cls._get_user(user_id, app) + except ValueError: + # Plugins such as WeCom Bot pass external sender IDs rather than EndUser UUIDs. + user = EndUserService.get_or_create_end_user(app, user_id=user_id) conversation_id = conversation_id or "" @@ -232,6 +236,13 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation): EndUser.app_id == app.id, ) user = session.scalar(stmt) + if not user: + stmt = select(EndUser).where( + EndUser.session_id == user_id, + EndUser.tenant_id == app.tenant_id, + EndUser.app_id == app.id, + ) + user = session.scalar(stmt) if not user: stmt = select(Account).where( Account.id == user_id, diff --git a/api/core/plugin/entities/bundle.py b/api/core/plugin/entities/bundle.py index 83d72d5dfa9..e3b17dc4710 100644 --- a/api/core/plugin/entities/bundle.py +++ b/api/core/plugin/entities/bundle.py @@ -5,12 +5,13 @@ from pydantic import BaseModel from core.plugin.entities.plugin import PluginDeclaration, PluginInstallationSource -class PluginBundleDependency(BaseModel): - class Type(StrEnum): - Github = PluginInstallationSource.Github.value - Marketplace = PluginInstallationSource.Marketplace.value - Package = PluginInstallationSource.Package.value +class PluginBundleDependencyType(StrEnum): + Github = PluginInstallationSource.Github.value + Marketplace = PluginInstallationSource.Marketplace.value + Package = PluginInstallationSource.Package.value + +class PluginBundleDependency(BaseModel): class Github(BaseModel): repo_address: str repo: str @@ -26,5 +27,5 @@ class PluginBundleDependency(BaseModel): unique_identifier: str manifest: PluginDeclaration - type: Type + type: PluginBundleDependencyType value: Github | Marketplace | Package diff --git a/api/core/plugin/entities/endpoint.py b/api/core/plugin/entities/endpoint.py index 64199636687..4aa1da8e918 100644 --- a/api/core/plugin/entities/endpoint.py +++ b/api/core/plugin/entities/endpoint.py @@ -32,6 +32,7 @@ class EndpointEntity(BasePluginEntity): entity of an endpoint """ + # TODO: Confirm daemon masks secret-input settings before endpoint list responses expose them. settings: dict[str, Any] tenant_id: str plugin_id: str diff --git a/api/core/plugin/entities/parameters.py b/api/core/plugin/entities/parameters.py index 14ed8af3ef7..4c0abe4ed4a 100644 --- a/api/core/plugin/entities/parameters.py +++ b/api/core/plugin/entities/parameters.py @@ -57,11 +57,12 @@ class MCPServerParameterType(StrEnum): OBJECT = auto() -class PluginParameterAutoGenerate(BaseModel): - class Type(StrEnum): - PROMPT_INSTRUCTION = auto() +class PluginParameterAutoGenerateType(StrEnum): + PROMPT_INSTRUCTION = auto() - type: Type + +class PluginParameterAutoGenerate(BaseModel): + type: PluginParameterAutoGenerateType class PluginParameterTemplate(BaseModel): diff --git a/api/core/plugin/entities/plugin.py b/api/core/plugin/entities/plugin.py index 89e0e8881cb..9f97de01056 100644 --- a/api/core/plugin/entities/plugin.py +++ b/api/core/plugin/entities/plugin.py @@ -166,12 +166,13 @@ class PluginEntity(PluginInstallation): return self -class PluginDependency(BaseModel): - class Type(StrEnum): - Github = PluginInstallationSource.Github - Marketplace = PluginInstallationSource.Marketplace - Package = PluginInstallationSource.Package +class PluginDependencyType(StrEnum): + Github = PluginInstallationSource.Github + Marketplace = PluginInstallationSource.Marketplace + Package = PluginInstallationSource.Package + +class PluginDependency(BaseModel): class Github(BaseModel): repo: str version: str @@ -194,7 +195,7 @@ class PluginDependency(BaseModel): plugin_unique_identifier: str version: str | None = None - type: Type + type: PluginDependencyType value: Github | Marketplace | Package current_identifier: str | None = None diff --git a/api/core/plugin/entities/plugin_daemon.py b/api/core/plugin/entities/plugin_daemon.py index 507a6ea5cd3..4cf55ef8e3c 100644 --- a/api/core/plugin/entities/plugin_daemon.py +++ b/api/core/plugin/entities/plugin_daemon.py @@ -228,7 +228,7 @@ class CredentialType(enum.StrEnum): OAUTH2 = "oauth2" UNAUTHORIZED = "unauthorized" - def get_name(self): + def get_name(self) -> str: if self == CredentialType.API_KEY: return "API KEY" elif self == CredentialType.OAUTH2: diff --git a/api/core/plugin/impl/model.py b/api/core/plugin/impl/model.py index 80a83fb3f21..c69be8a3933 100644 --- a/api/core/plugin/impl/model.py +++ b/api/core/plugin/impl/model.py @@ -27,10 +27,12 @@ _POLLING_UNSUPPORTED_ERROR_MESSAGE = "does not support polling" class PluginModelClient(BasePluginClient): @staticmethod - def _dispatch_payload(*, user_id: str | None, data: dict[str, Any]) -> dict[str, Any]: + def _dispatch_payload(*, user_id: str | None, data: dict[str, Any], app_id: str | None = None) -> dict[str, Any]: payload: dict[str, Any] = {"data": data} if user_id is not None: payload["user_id"] = user_id + if app_id is not None: + payload["app_id"] = app_id return payload def fetch_model_providers(self, tenant_id: str) -> Sequence[PluginModelProviderEntity]: @@ -166,6 +168,7 @@ class PluginModelClient(BasePluginClient): tools: list[PromptMessageTool] | None = None, stop: list[str] | None = None, stream: bool = True, + app_id: str | None = None, ) -> Generator[LLMResultChunk, None, None]: """ Invoke llm @@ -188,6 +191,7 @@ class PluginModelClient(BasePluginClient): "stop": stop, "stream": stream, }, + app_id=app_id, ) ), headers={ diff --git a/api/core/plugin/impl/model_runtime.py b/api/core/plugin/impl/model_runtime.py index 021d0005e85..454bd38958a 100644 --- a/api/core/plugin/impl/model_runtime.py +++ b/api/core/plugin/impl/model_runtime.py @@ -317,21 +317,39 @@ class PluginModelRuntime(ModelRuntime): stream: bool, request_metadata: Mapping[str, object] | None = None, ) -> LLMResult | Generator[LLMResultChunk, None, None]: - del request_metadata + app_id = request_metadata.get("app_id") if request_metadata else None + if not isinstance(app_id, str): + app_id = None plugin_id, provider_name = self._split_provider(provider) - result = self.client.invoke_llm( - tenant_id=self.tenant_id, - user_id=self.user_id, - plugin_id=plugin_id, - provider=provider_name, - model=model, - credentials=credentials, - model_parameters=model_parameters, - prompt_messages=list(prompt_messages), - tools=tools, - stop=list(stop) if stop else None, - stream=stream, - ) + if app_id is None: + result = self.client.invoke_llm( + tenant_id=self.tenant_id, + user_id=self.user_id, + plugin_id=plugin_id, + provider=provider_name, + model=model, + credentials=credentials, + model_parameters=model_parameters, + prompt_messages=list(prompt_messages), + tools=tools, + stop=list(stop) if stop else None, + stream=stream, + ) + else: + result = self.client.invoke_llm( + tenant_id=self.tenant_id, + user_id=self.user_id, + plugin_id=plugin_id, + provider=provider_name, + model=model, + credentials=credentials, + model_parameters=model_parameters, + prompt_messages=list(prompt_messages), + tools=tools, + stop=list(stop) if stop else None, + stream=stream, + app_id=app_id, + ) if stream: return result diff --git a/api/core/plugin/plugin_service.py b/api/core/plugin/plugin_service.py index 6b306e2df86..89274b635ac 100644 --- a/api/core/plugin/plugin_service.py +++ b/api/core/plugin/plugin_service.py @@ -92,6 +92,7 @@ class PluginService: PLUGIN_MODEL_PROVIDERS_REDIS_KEY_PREFIX = "plugin_model_providers:tenant_id:" PLUGIN_MODEL_PROVIDERS_GENERATION_REDIS_KEY_PREFIX = "plugin_model_providers_generation:tenant_id:" PLUGIN_MODEL_PROVIDERS_LOCK_REDIS_KEY_PREFIX = "plugin_model_providers_refresh_lock:tenant_id:" + PLUGIN_MODEL_PROVIDERS_REMOTE_DEBUG_REDIS_KEY_PREFIX = "plugin_model_providers_remote_debug:tenant_id:" PLUGIN_MODEL_PROVIDERS_LOCK_TTL = 30 PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_TIMEOUT = 2.0 PLUGIN_MODEL_PROVIDERS_LOCK_WAIT_INTERVAL = 0.05 @@ -117,6 +118,10 @@ class PluginService: def _get_plugin_model_providers_lock_key(cls, tenant_id: str, generation: int) -> str: return f"{cls.PLUGIN_MODEL_PROVIDERS_LOCK_REDIS_KEY_PREFIX}{tenant_id}:generation:{generation}" + @classmethod + def _get_plugin_model_providers_remote_debug_cache_key(cls, tenant_id: str) -> str: + return f"{cls.PLUGIN_MODEL_PROVIDERS_REMOTE_DEBUG_REDIS_KEY_PREFIX}{tenant_id}" + @staticmethod def _get_provider_short_name_alias(provider: PluginModelProviderEntity) -> str: """ @@ -259,6 +264,111 @@ class PluginService: except (RedisError, RuntimeError): logger.warning("Failed to cache plugin model providers for tenant %s.", tenant_id, exc_info=True) + @classmethod + def _get_remote_model_plugin_cache_marker(cls, plugins: Sequence[PluginEntity]) -> str | None: + remote_model_plugins = sorted( + f"{plugin.plugin_id}:{plugin.plugin_unique_identifier}" + for plugin in plugins + if plugin.source == PluginInstallationSource.Remote + ) + if not remote_model_plugins: + return None + + return "\n".join(remote_model_plugins) + + @classmethod + def _load_cached_remote_model_plugin_marker(cls, tenant_id: str) -> str | None: + cache_key = cls._get_plugin_model_providers_remote_debug_cache_key(tenant_id) + try: + cached_marker = redis_client.get(cache_key) + except (RedisError, RuntimeError): + logger.warning("Failed to read remote debug model plugin marker for tenant %s.", tenant_id, exc_info=True) + return None + + if cached_marker is None: + return None + if isinstance(cached_marker, bytes): + try: + return cached_marker.decode() + except UnicodeDecodeError: + logger.warning( + "Invalid remote debug model plugin marker for tenant %s; deleting cache marker.", + tenant_id, + exc_info=True, + ) + try: + redis_client.delete(cache_key) + except (RedisError, RuntimeError): + logger.warning( + "Failed to delete invalid remote debug model plugin marker for tenant %s.", + tenant_id, + exc_info=True, + ) + return None + if isinstance(cached_marker, str): + return cached_marker + + logger.warning("Invalid remote debug model plugin marker for tenant %s; deleting cache marker.", tenant_id) + try: + redis_client.delete(cache_key) + except (RedisError, RuntimeError): + logger.warning( + "Failed to delete invalid remote debug model plugin marker for tenant %s.", + tenant_id, + exc_info=True, + ) + return None + + @classmethod + def _store_cached_remote_model_plugin_marker(cls, tenant_id: str, marker: str | None) -> None: + cache_key = cls._get_plugin_model_providers_remote_debug_cache_key(tenant_id) + try: + if marker is None: + redis_client.delete(cache_key) + else: + redis_client.setex(cache_key, dify_config.PLUGIN_MODEL_PROVIDERS_CACHE_TTL, marker) + except (RedisError, RuntimeError): + logger.warning("Failed to cache remote debug model plugin marker for tenant %s.", tenant_id, exc_info=True) + + @classmethod + def _load_cached_plugin_model_provider_plugin_ids(cls, tenant_id: str) -> set[str] | None: + """Return plugin ids represented by the current provider cache, or None when no usable cache exists.""" + generation = cls._load_plugin_model_providers_generation(tenant_id) + cached_providers, _ = cls._load_cached_plugin_model_providers_for_generation(tenant_id, generation) + if cached_providers is None: + return None + + plugin_ids: set[str] = set() + for provider in cached_providers: + last_slash = provider.provider.rfind("/") + if last_slash > 0: + plugin_ids.add(provider.provider[:last_slash]) + + return plugin_ids + + @classmethod + def _should_invalidate_model_provider_cache_for_remote_model_plugins( + cls, + tenant_id: str, + plugins: Sequence[PluginEntity], + ) -> bool: + remote_model_plugin_marker = cls._get_remote_model_plugin_cache_marker(plugins) + cached_remote_model_plugin_marker = cls._load_cached_remote_model_plugin_marker(tenant_id) + if remote_model_plugin_marker is None: + return cached_remote_model_plugin_marker is not None + + if remote_model_plugin_marker != cached_remote_model_plugin_marker: + return True + + remote_model_plugin_ids = { + plugin.plugin_id for plugin in plugins if plugin.source == PluginInstallationSource.Remote + } + cached_plugin_ids = cls._load_cached_plugin_model_provider_plugin_ids(tenant_id) + if cached_plugin_ids is None: + return False + + return not remote_model_plugin_ids.issubset(cached_plugin_ids) + @classmethod @contextmanager def _plugin_model_providers_refresh_lock( @@ -571,7 +681,21 @@ class PluginService: This keeps pagination usable before category is persisted on installation rows. """ manager = PluginInstaller() - return manager.list_plugins_by_category(tenant_id, category, page, page_size) + plugins = manager.list_plugins_by_category(tenant_id, category, page, page_size) + if category == PluginCategory.Model: + should_invalidate_model_provider_cache = ( + PluginService._should_invalidate_model_provider_cache_for_remote_model_plugins( + tenant_id, + plugins.list, + ) + ) + if should_invalidate_model_provider_cache: + PluginService.invalidate_plugin_model_providers_cache(tenant_id) + + remote_model_plugin_marker = PluginService._get_remote_model_plugin_cache_marker(plugins.list) + PluginService._store_cached_remote_model_plugin_marker(tenant_id, remote_model_plugin_marker) + + return plugins @staticmethod def _normalize_endpoint_count(value: object) -> int: diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index e2c710923b5..ebfe77e8f30 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -1544,10 +1544,12 @@ class ProviderManager: trail_pool = CreditPoolService.get_pool( tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL, + session=db.session(), ) paid_pool = CreditPoolService.get_pool( tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID, + session=db.session(), ) else: trail_pool = None diff --git a/api/core/rag/datasource/retrieval_service.py b/api/core/rag/datasource/retrieval_service.py index 50381f5e75c..3b20f8bc530 100644 --- a/api/core/rag/datasource/retrieval_service.py +++ b/api/core/rag/datasource/retrieval_service.py @@ -199,7 +199,7 @@ class RetrievalService: metadata_filtering_conditions: dict[str, Any] | None = None, ): stmt = select(Dataset).where(Dataset.id == dataset_id) - dataset = db.session.scalar(stmt) + dataset = session.scalar(stmt) if not dataset: return [] metadata_condition = ( @@ -208,12 +208,12 @@ class RetrievalService: else None ) all_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( - session, - dataset.tenant_id, - dataset_id, - query, - external_retrieval_model or {}, + tenant_id=dataset.tenant_id, + dataset_id=dataset_id, + query=query, + external_retrieval_parameters=external_retrieval_model or {}, metadata_condition=metadata_condition, + session=session, ) return all_documents diff --git a/api/core/rag/extractor/extract_processor.py b/api/core/rag/extractor/extract_processor.py index 4d11ebe5005..36d879427a5 100644 --- a/api/core/rag/extractor/extract_processor.py +++ b/api/core/rag/extractor/extract_processor.py @@ -1,7 +1,7 @@ import re import tempfile from pathlib import Path -from typing import Union +from typing import Literal, overload from urllib.parse import unquote from configs import dify_config @@ -40,10 +40,22 @@ USER_AGENT = ( class ExtractProcessor: + @overload + @classmethod + def load_from_upload_file( + cls, upload_file: UploadFile, return_text: Literal[True], is_automatic: bool = False + ) -> str: ... + + @overload + @classmethod + def load_from_upload_file( + cls, upload_file: UploadFile, return_text: Literal[False] = False, is_automatic: bool = False + ) -> list[Document]: ... + @classmethod def load_from_upload_file( cls, upload_file: UploadFile, return_text: bool = False, is_automatic: bool = False - ) -> Union[list[Document], str]: + ) -> list[Document] | str: extract_setting = ExtractSetting( datasource_type=DatasourceType.FILE, upload_file=upload_file, document_model="text_model" ) @@ -53,8 +65,16 @@ class ExtractProcessor: else: return cls.extract(extract_setting, is_automatic) + @overload @classmethod - def load_from_url(cls, url: str, return_text: bool = False) -> Union[list[Document], str]: + def load_from_url(cls, url: str, return_text: Literal[True]) -> str: ... + + @overload + @classmethod + def load_from_url(cls, url: str, return_text: Literal[False] = False) -> list[Document]: ... + + @classmethod + def load_from_url(cls, url: str, return_text: bool = False) -> list[Document] | str: response = remote_fetcher.make_request("GET", url, headers={"User-Agent": USER_AGENT}) with tempfile.TemporaryDirectory() as temp_dir: diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py index dd173207b09..b31c1bb634b 100644 --- a/api/core/rag/index_processor/processor/paragraph_index_processor.py +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -5,7 +5,7 @@ import re import uuid from typing import Any, TypedDict, cast, override -from sqlalchemy.orm import scoped_session +from sqlalchemy.orm import Session logger = logging.getLogger(__name__) @@ -162,10 +162,10 @@ class ParagraphIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=segment_ids) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=None) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: vector = Vector(dataset) @@ -226,7 +226,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): all_multimodal_documents.append(file_document) doc.attachments = attachments else: - account = AccountService.load_user(document.created_by, db.session) + account = AccountService.load_user(document.created_by, db.session()) if not account: raise ValueError("Invalid account") doc.attachments = self._get_content_files(doc, current_user=account) @@ -414,12 +414,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor): # First, try to get images from SegmentAttachmentBinding (preferred method) if segment_id: image_files = ParagraphIndexProcessor._extract_images_from_segment_attachments( - tenant_id, segment_id, db.session + tenant_id, segment_id, db.session() ) # If no images from attachments, fall back to extracting from text if not image_files: - image_files = ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, db.session) + image_files = ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, db.session()) # Build prompt messages prompt_messages = [] @@ -473,7 +473,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): return summary_content, usage @staticmethod - def _extract_images_from_text(tenant_id: str, text: str, session: scoped_session) -> list[File]: + def _extract_images_from_text(tenant_id: str, text: str, session: Session) -> list[File]: """ Extract images from markdown text and convert them to File objects. @@ -553,9 +553,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): return file_objects @staticmethod - def _extract_images_from_segment_attachments( - tenant_id: str, segment_id: str, session: scoped_session - ) -> list[File]: + def _extract_images_from_segment_attachments(tenant_id: str, segment_id: str, session: Session) -> list[File]: """ Extract images from SegmentAttachmentBinding table (preferred method). This matches how DatasetRetrieval gets segment attachments. diff --git a/api/core/rag/index_processor/processor/parent_child_index_processor.py b/api/core/rag/index_processor/processor/parent_child_index_processor.py index 78d8b7dcd53..aecb4154d6f 100644 --- a/api/core/rag/index_processor/processor/parent_child_index_processor.py +++ b/api/core/rag/index_processor/processor/parent_child_index_processor.py @@ -169,10 +169,10 @@ class ParentChildIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=segment_ids) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=None) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: delete_child_chunks = kwargs.get("delete_child_chunks") or False @@ -291,7 +291,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor): attachments.append(file_document) doc.attachments = attachments else: - account = AccountService.load_user(document.created_by, db.session) + account = AccountService.load_user(document.created_by, db.session()) if not account: raise ValueError("Invalid account") doc.attachments = self._get_content_files(doc, current_user=account) diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py index 253acebc2c6..7b7443a621f 100644 --- a/api/core/rag/index_processor/processor/qa_index_processor.py +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -173,10 +173,10 @@ class QAIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=segment_ids) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None) + SummaryIndexService.delete_summaries_for_segments(dataset=dataset, segment_ids=None) vector = Vector(dataset) if node_ids: diff --git a/api/core/rag/summary_index/summary_index.py b/api/core/rag/summary_index/summary_index.py index bff5f85decb..d9ce3879890 100644 --- a/api/core/rag/summary_index/summary_index.py +++ b/api/core/rag/summary_index/summary_index.py @@ -74,11 +74,16 @@ class SummaryIndex: def process_segment(segment_id: str) -> None: """Process a single segment in a thread with a fresh DB session.""" with session_factory.create_session() as session: + dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) + if dataset is None: + return segment = session.scalar(select(DocumentSegment).where(DocumentSegment.id == segment_id).limit(1)) if segment is None: return try: - SummaryIndexService.generate_and_vectorize_summary(segment, dataset, summary_index_setting) + SummaryIndexService.generate_and_vectorize_summary( + segment, dataset, summary_index_setting, session=session + ) except Exception: logger.exception( "Failed to generate summary for segment %s", diff --git a/api/core/tools/__base/tool_provider.py b/api/core/tools/__base/tool_provider.py index 70e4fe1ff73..39d61409862 100644 --- a/api/core/tools/__base/tool_provider.py +++ b/api/core/tools/__base/tool_provider.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from copy import deepcopy from typing import Any -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.__base.tool import Tool from core.tools.entities.tool_entities import ( ToolProviderEntity, @@ -71,11 +71,11 @@ class ToolProviderController[ToolProviderEntityT: ToolProviderEntity, ToolProvid if not credential_schema.required and credentials[credential_name] is None: continue - if credential_schema.type in {ProviderConfig.Type.SECRET_INPUT, ProviderConfig.Type.TEXT_INPUT}: + if credential_schema.type in {ProviderConfigType.SECRET_INPUT, ProviderConfigType.TEXT_INPUT}: if not isinstance(credentials[credential_name], str): raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") - elif credential_schema.type == ProviderConfig.Type.SELECT: + elif credential_schema.type == ProviderConfigType.SELECT: if not isinstance(credentials[credential_name], str): raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") @@ -100,9 +100,9 @@ class ToolProviderController[ToolProviderEntityT: ToolProviderEntity, ToolProvid default_value = credential_schema.default # parse default value into the correct type if credential_schema.type in { - ProviderConfig.Type.SECRET_INPUT, - ProviderConfig.Type.TEXT_INPUT, - ProviderConfig.Type.SELECT, + ProviderConfigType.SECRET_INPUT, + ProviderConfigType.TEXT_INPUT, + ProviderConfigType.SELECT, }: default_value = str(default_value) diff --git a/api/core/tools/custom_tool/provider.py b/api/core/tools/custom_tool/provider.py index ade5b894f95..df48f63809d 100644 --- a/api/core/tools/custom_tool/provider.py +++ b/api/core/tools/custom_tool/provider.py @@ -5,7 +5,7 @@ from typing import override from pydantic import Field from sqlalchemy import select -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.__base.tool_provider import ToolProviderController from core.tools.__base.tool_runtime import ToolRuntime from core.tools.custom_tool.tool import ApiTool @@ -41,7 +41,7 @@ class ApiToolProviderController(ToolProviderController[ToolProviderEntity, ApiTo ProviderConfig( name="auth_type", required=True, - type=ProviderConfig.Type.SELECT, + type=ProviderConfigType.SELECT, options=[ ProviderConfig.Option(value="none", label=I18nObject(en_US="None", zh_Hans="无")), ProviderConfig.Option(value="api_key_header", label=I18nObject(en_US="Header", zh_Hans="请求头")), @@ -60,20 +60,20 @@ class ApiToolProviderController(ToolProviderController[ToolProviderEntity, ApiTo name="api_key_header", required=False, default="Authorization", - type=ProviderConfig.Type.TEXT_INPUT, + type=ProviderConfigType.TEXT_INPUT, help=I18nObject(en_US="The header name of the api key", zh_Hans="携带 api key 的 header 名称"), ), ProviderConfig( name="api_key_value", required=True, - type=ProviderConfig.Type.SECRET_INPUT, + type=ProviderConfigType.SECRET_INPUT, help=I18nObject(en_US="The api key", zh_Hans="api key 的值"), ), ProviderConfig( name="api_key_header_prefix", required=False, default="basic", - type=ProviderConfig.Type.SELECT, + type=ProviderConfigType.SELECT, help=I18nObject(en_US="The prefix of the api key header", zh_Hans="api key header 的前缀"), options=[ ProviderConfig.Option(value="basic", label=I18nObject(en_US="Basic", zh_Hans="Basic")), @@ -89,7 +89,7 @@ class ApiToolProviderController(ToolProviderController[ToolProviderEntity, ApiTo name="api_key_query_param", required=False, default="key", - type=ProviderConfig.Type.TEXT_INPUT, + type=ProviderConfigType.TEXT_INPUT, help=I18nObject( en_US="The query parameter name of the api key", zh_Hans="携带 api key 的查询参数名称" ), @@ -97,7 +97,7 @@ class ApiToolProviderController(ToolProviderController[ToolProviderEntity, ApiTo ProviderConfig( name="api_key_value", required=True, - type=ProviderConfig.Type.SECRET_INPUT, + type=ProviderConfigType.SECRET_INPUT, help=I18nObject(en_US="The api key", zh_Hans="api key 的值"), ), ] diff --git a/api/core/tools/entities/api_entities.py b/api/core/tools/entities/api_entities.py index 0217300055f..0bee91ffe14 100644 --- a/api/core/tools/entities/api_entities.py +++ b/api/core/tools/entities/api_entities.py @@ -74,8 +74,6 @@ class ToolProviderApiEntity(BaseModel): for parameter in tool.get("parameters"): if parameter.get("type") == ToolParameter.ToolParameterType.SYSTEM_FILES: parameter["type"] = "files" - if parameter.get("input_schema") is None: - parameter.pop("input_schema", None) # ------------- optional_fields = self.optional_field("server_url", self.server_url) match self.type: diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py index 850571c3f19..fc85f20bdd4 100644 --- a/api/core/tools/tool_manager.py +++ b/api/core/tools/tool_manager.py @@ -16,30 +16,16 @@ from yarl import URL import contexts from configs import dify_config -from core.entities import PluginCredentialType -from core.helper.provider_cache import ToolProviderCredentialsCache -from core.plugin.impl.tool import PluginToolManager -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.mcp_tool.provider import MCPToolProviderController -from core.tools.mcp_tool.tool import MCPTool -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.plugin_tool.tool import PluginTool -from core.tools.utils.uuid_utils import is_valid_uuid -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from extensions.ext_database import db -from graphon.runtime import VariablePool -from models.provider_ids import ToolProviderID -from services.tools.mcp_tools_manage_service import MCPToolManageService - -if TYPE_CHECKING: - pass - from core.agent.entities import AgentToolEntity from core.app.entities.app_invoke_entities import InvokeFrom +from core.entities import PluginCredentialType from core.helper.module_import_helper import load_single_subclass_from_source from core.helper.position_helper import is_filtered +from core.helper.provider_cache import ToolProviderCredentialsCache +from core.plugin.impl.tool import PluginToolManager from core.tools.__base.tool import Tool +from core.tools.__base.tool_provider import ToolProviderController +from core.tools.__base.tool_runtime import ToolRuntime from core.tools.builtin_tool.provider import BuiltinToolProviderController from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort from core.tools.builtin_tool.tool import BuiltinTool @@ -49,6 +35,7 @@ from core.tools.entities.api_entities import ToolProviderApiEntity, ToolProvider from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( ApiProviderAuthType, + ApiProviderSchemaType, EmojiIconDict, ToolInvokeFrom, ToolParameter, @@ -56,12 +43,21 @@ from core.tools.entities.tool_entities import ( emoji_icon_adapter, ) from core.tools.errors import ToolProviderNotFoundError +from core.tools.mcp_tool.provider import MCPToolProviderController +from core.tools.mcp_tool.tool import MCPTool +from core.tools.plugin_tool.provider import PluginToolProviderController +from core.tools.plugin_tool.tool import PluginTool from core.tools.tool_label_manager import ToolLabelManager from core.tools.utils.configuration import ToolParameterConfigurationManager from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter +from core.tools.utils.uuid_utils import is_valid_uuid +from core.tools.workflow_as_tool.provider import WorkflowToolProviderController from core.tools.workflow_as_tool.tool import WorkflowTool -from graphon.model_runtime.utils.encoders import jsonable_encoder +from extensions.ext_database import db +from graphon.runtime import VariablePool +from models.provider_ids import ToolProviderID from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider +from services.tools.mcp_tools_manage_service import MCPToolManageService from services.tools.tools_transform_service import ToolTransformService if TYPE_CHECKING: @@ -921,23 +917,20 @@ class ToolManager: # add tool labels labels = ToolLabelManager.get_tool_labels(controller) + schema_type = provider_obj.schema_type + schema_type_value = schema_type.value if isinstance(schema_type, ApiProviderSchemaType) else schema_type - return cast( - dict, - jsonable_encoder( - { - "schema_type": provider_obj.schema_type, - "schema": provider_obj.schema, - "tools": provider_obj.tools, - "icon": icon, - "description": provider_obj.description, - "credentials": masked_credentials, - "privacy_policy": provider_obj.privacy_policy, - "custom_disclaimer": provider_obj.custom_disclaimer, - "labels": labels, - } - ), - ) + return { + "schema_type": schema_type_value, + "schema": provider_obj.schema, + "tools": [tool.model_dump(mode="json") for tool in provider_obj.tools], + "icon": icon, + "description": provider_obj.description, + "credentials": masked_credentials, + "privacy_policy": provider_obj.privacy_policy, + "custom_disclaimer": provider_obj.custom_disclaimer, + "labels": labels, + } @classmethod def generate_builtin_tool_icon_url(cls, provider_id: str) -> str: diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py index a3afe659563..c26523b9be5 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py @@ -80,7 +80,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): all_documents = rerank_runner.run(query, all_documents, self.score_threshold, self.top_k) for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(all_documents, db.session) + hit_callback.on_tool_end(all_documents, db.session()) document_score_list = {} for item in all_documents: @@ -167,7 +167,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): return [] for hit_callback in hit_callbacks: - hit_callback.on_query(query, dataset.id, db.session) + hit_callback.on_query(query, dataset.id, db.session()) # get retrieval model , if the model is not setting , using default retrieval_model = dataset.retrieval_model or default_retrieval_model diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py index 247bd0705fc..d7e390ca877 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py @@ -65,7 +65,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): if not dataset: return "" for hit_callback in self.hit_callbacks: - hit_callback.on_query(query, dataset.id, db.session) + hit_callback.on_query(query, dataset.id, db.session()) dataset_retrieval = DatasetRetrieval() metadata_filter_document_ids, metadata_condition = dataset_retrieval.get_metadata_filter_condition( session, @@ -162,7 +162,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): else: documents = [] for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(documents, db.session) + hit_callback.on_tool_end(documents, db.session()) document_score_list = {} if dataset.indexing_technique != IndexTechniqueType.ECONOMY: for item in documents: diff --git a/api/core/trigger/utils/encryption.py b/api/core/trigger/utils/encryption.py index 9b958690e59..8409be0813e 100644 --- a/api/core/trigger/utils/encryption.py +++ b/api/core/trigger/utils/encryption.py @@ -1,7 +1,7 @@ from collections.abc import Mapping from typing import Union, override -from core.entities.provider_entities import BasicProviderConfig, ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.helper.provider_cache import ProviderCredentialsCache from core.helper.provider_encryption import ProviderConfigCache, ProviderConfigEncrypter, create_provider_encrypter from core.plugin.entities.plugin_daemon import CredentialType @@ -142,7 +142,7 @@ def masked_credentials( if not config: masked_credentials[key] = value continue - if config.type == BasicProviderConfig.Type.SECRET_INPUT: + if config.type == ProviderConfigType.SECRET_INPUT: if len(value) <= 4: masked_credentials[key] = "*" * len(value) else: diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index 991e79cdc61..5b8db068d0d 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -549,7 +549,11 @@ class DifyNodeFactory(NodeFactory): "credentials_provider": self._llm_credentials_provider, "model_factory": self._llm_model_factory, "model_instance": ( - self._wrap_model_instance_for_node(node_data=validated_node_data, model_instance=model_instance) + self._wrap_model_instance_for_node( + node_data=validated_node_data, + model_instance=model_instance, + request_metadata={"app_id": self._dify_context.app_id}, + ) if wrap_model_instance else model_instance ), @@ -581,13 +585,14 @@ class DifyNodeFactory(NodeFactory): *, node_data: LLMCompatibleNodeData, model_instance: ModelInstance, + request_metadata: Mapping[str, object] | None = None, ) -> DifyPreparedLLM: # Only graphon's LLM node consumes the polling protocol. Keep classifier # and extractor nodes on the existing wrapper even if the same model # advertises polling support. if node_data.type == BuiltinNodeTypes.LLM and DifyNodeFactory._supports_plugin_llm_polling(model_instance): - return DifyPreparedPollingLLM(model_instance) - return DifyPreparedLLM(model_instance) + return DifyPreparedPollingLLM(model_instance, request_metadata=request_metadata) + return DifyPreparedLLM(model_instance, request_metadata=request_metadata) @staticmethod def _supports_plugin_llm_polling(model_instance: ModelInstance) -> bool: diff --git a/api/core/workflow/node_runtime.py b/api/core/workflow/node_runtime.py index 233f404b0a4..d0391b69c73 100644 --- a/api/core/workflow/node_runtime.py +++ b/api/core/workflow/node_runtime.py @@ -150,8 +150,9 @@ class DifyFileReferenceFactory(FileReferenceFactoryProtocol): class DifyPreparedLLM(LLMProtocol): """Workflow-layer adapter that hides the full `ModelInstance` API from `graphon` nodes.""" - def __init__(self, model_instance: ModelInstance) -> None: + def __init__(self, model_instance: ModelInstance, request_metadata: Mapping[str, object] | None = None) -> None: self._model_instance = model_instance + self._request_metadata = request_metadata @property @override @@ -230,6 +231,7 @@ class DifyPreparedLLM(LLMProtocol): tools=list(tools or []), stop=list(stop or []), stream=stream, + request_metadata=self._request_metadata, ) @overload @@ -283,10 +285,10 @@ class DifyPreparedLLM(LLMProtocol): class DifyPreparedPollingLLM(DifyPreparedLLM, LLMPollingCapableProtocol): """Prepared workflow LLM adapter that exposes Graphon's polling protocol.""" - def __init__(self, model_instance: ModelInstance) -> None: + def __init__(self, model_instance: ModelInstance, request_metadata: Mapping[str, object] | None = None) -> None: from core.plugin.impl.model_runtime import PluginModelRuntime - super().__init__(model_instance) + super().__init__(model_instance, request_metadata=request_metadata) model_type_instance = model_instance.model_type_instance if not isinstance(model_type_instance, LargeLanguageModel): raise TypeError("Polling wrapper requires a large-language-model instance.") diff --git a/api/core/workflow/nodes/agent_v2/agent_node.py b/api/core/workflow/nodes/agent_v2/agent_node.py index 41778a0ba50..8e27cf63da4 100644 --- a/api/core/workflow/nodes/agent_v2/agent_node.py +++ b/api/core/workflow/nodes/agent_v2/agent_node.py @@ -16,6 +16,7 @@ from clients.agent_backend import ( AgentBackendRunEventAdapter, AgentBackendRunFailedInternalEvent, AgentBackendRunSucceededInternalEvent, + AgentBackendSessionCleanupPayload, AgentBackendStreamError, AgentBackendStreamInternalEvent, AgentBackendTransportError, @@ -34,6 +35,7 @@ from graphon.node_events import NodeEventBase, NodeRunResult, PauseRequestedEven from graphon.nodes.base.node import Node from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig from services.agent.prompt_mentions import extract_workflow_node_output_selectors +from tasks.agent_backend_session_cleanup_task import cleanup_workflow_agent_runtime_session from .ask_human_hitl import AskHumanFormBuildError, build_ask_human_pause_reason from .ask_human_resume import build_deferred_tool_results, resolve_ask_human_form @@ -609,6 +611,44 @@ class DifyAgentNode(Node[DifyAgentNodeData]): ) -> None: if self._session_store is None: return + stored_session = self._session_store.load_active_session(session_scope) + try: + if stored_session is not None and stored_session.runtime_layer_specs: + payload = AgentBackendSessionCleanupPayload( + session_snapshot=stored_session.session_snapshot, + runtime_layer_specs=stored_session.runtime_layer_specs, + idempotency_key=( + f"{session_scope.tenant_id}:{session_scope.workflow_run_id}:{session_scope.node_id}:" + f"{session_scope.binding_id}:workflow-agent-failure-cleanup:" + f"{stored_session.backend_run_id or 'no-stored-run'}:{backend_run_id}" + ), + metadata={ + "tenant_id": session_scope.tenant_id, + "app_id": session_scope.app_id, + "workflow_id": session_scope.workflow_id, + "workflow_run_id": session_scope.workflow_run_id, + "node_id": session_scope.node_id, + "node_execution_id": session_scope.node_execution_id, + "binding_id": session_scope.binding_id, + "agent_id": session_scope.agent_id, + "agent_config_snapshot_id": session_scope.agent_config_snapshot_id, + "previous_agent_backend_run_id": stored_session.backend_run_id, + "failed_agent_backend_run_id": backend_run_id, + }, + ) + cleanup_workflow_agent_runtime_session.delay(payload.model_dump(mode="json")) + except Exception: + logger.warning( + "Failed to enqueue workflow Agent backend cleanup on agent run failure: " + "tenant_id=%s workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s backend_run_id=%s", + session_scope.tenant_id, + session_scope.workflow_run_id, + session_scope.node_id, + session_scope.binding_id, + session_scope.agent_id, + backend_run_id, + exc_info=True, + ) try: self._session_store.mark_cleaned(scope=session_scope, backend_run_id=backend_run_id) agent_backend = dict(metadata.get("agent_backend") or {}) diff --git a/api/core/workflow/nodes/agent_v2/dify_tools_builder.py b/api/core/workflow/nodes/agent_v2/dify_tools_builder.py index fc2719a6204..0e6fb3d1830 100644 --- a/api/core/workflow/nodes/agent_v2/dify_tools_builder.py +++ b/api/core/workflow/nodes/agent_v2/dify_tools_builder.py @@ -15,7 +15,6 @@ from dify_agent.layers.dify_plugin import ( DifyPluginToolsLayerConfig, ) from sqlalchemy import select -from sqlalchemy.orm import Session from core.agent.entities import AgentToolEntity from core.app.entities.app_invoke_entities import InvokeFrom @@ -132,7 +131,7 @@ def _list_provider_tool_names( def _resolve_mcp_provider_id(*, tenant_id: str, provider_id: str) -> str: """Normalize MCP provider ids to the runtime-facing server identifier.""" - service = MCPToolManageService(session=cast(Session, db.session)) + service = MCPToolManageService(session=db.session()) try: return service.get_provider_entity(provider_id, tenant_id, by_server_id=True).provider_id except ValueError: diff --git a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py index 1a4529b5d97..3afa740e545 100644 --- a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py +++ b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py @@ -48,7 +48,7 @@ from clients.agent_backend import ( ) from configs import dify_config from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom -from core.workflow.system_variables import SystemVariableKey, get_system_text +from core.workflow.system_variables import SystemVariableKey, get_system_text, get_system_value from graphon.file import File, FileTransferMethod from graphon.variables.segments import Segment from models.agent import Agent, AgentConfigSnapshot, WorkflowAgentNodeBinding @@ -354,17 +354,22 @@ class WorkflowAgentRuntimeRequestBuilder: ) -> str: lines: list[str] = [] query = get_system_text(context.variable_pool, SystemVariableKey.QUERY) + uploaded_files = self._summarize_uploaded_workflow_files(context.variable_pool) resolved_outputs = self._resolve_previous_node_outputs( context.variable_pool, node_job.previous_node_output_refs, ) - if not query and not resolved_outputs: + if not query and uploaded_files is None and not resolved_outputs: return "" lines.append("Workflow context loaded for this run:") if query: lines.append(f"- User query: {query}") + if uploaded_files is not None: + lines.append("- Uploaded workflow files:") + lines.append(f" - sys.files: {uploaded_files}") + if resolved_outputs: lines.append("- Previous node outputs:") for item in resolved_outputs: @@ -373,6 +378,14 @@ class WorkflowAgentRuntimeRequestBuilder: lines.append("The above workflow context is run-specific. Do not treat it as Agent Soul or persistent memory.") return "\n".join(lines) + def _summarize_uploaded_workflow_files(self, variable_pool: VariablePoolReader) -> str | None: + files = get_system_value(variable_pool, SystemVariableKey.FILES) + if files is None: + return None + if isinstance(files, list | tuple) and not files: + return None + return self._summarize_value(files) + def _build_workflow_task_prompt( self, context: WorkflowAgentRuntimeBuildContext, @@ -441,7 +454,7 @@ class WorkflowAgentRuntimeRequestBuilder: def _resolve_prompt_payload_value(cls, value: Any) -> tuple[Any, bool]: # File-valued workflow context must surface as Agent Stub download # mappings so the model can materialize those inputs with - # `dify-agent file download --mapping ...` inside the sandbox. + # `dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL` inside the sandbox. download_mapping = cls._agent_stub_download_mapping(value) if download_mapping is not None: return download_mapping, True @@ -486,7 +499,21 @@ class WorkflowAgentRuntimeRequestBuilder: return mapping.model_dump(mode="json", exclude_none=True) if isinstance(value, Mapping): - mapping = AgentStubFileMapping.model_validate(value) + transfer_method = value.get("transfer_method") + if not isinstance(transfer_method, str): + return None + if transfer_method == "remote_url": + mapping = AgentStubFileMapping( + transfer_method="remote_url", + url=value.get("url") or value.get("remote_url"), + ) + elif transfer_method in {"local_file", "tool_file", "datasource_file"}: + mapping = AgentStubFileMapping( + transfer_method=cast(AgentStubFileTransferMethod, transfer_method), + reference=value.get("reference"), + ) + else: + return None return mapping.model_dump(mode="json", exclude_none=True) except ValidationError: return None @@ -587,11 +614,12 @@ class WorkflowAgentRuntimeRequestBuilder: broad generated schema but fails API-side output type checking. For files produced inside an Agent run, the supported persisted shape is - narrower than every downloadable mapping: the sandbox must upload the - local artifact via ``dify-agent file upload ``, which returns a - ``tool_file`` mapping. ``local_file`` and ``datasource_file`` are valid - for existing file references in workflow context, not for newly produced - Agent output files. + narrower than the full CLI upload stdout: the sandbox must upload the + local artifact via ``dify-agent file upload ``, then use the + returned ``reference`` inside the accepted ``tool_file`` mapping shape + for structured ``final_output``. ``local_file`` and ``datasource_file`` + are valid for existing file references in workflow context, not for + newly produced Agent output files. """ return { "title": "AgentStubFileMapping", @@ -646,9 +674,9 @@ class WorkflowAgentRuntimeRequestBuilder: if output.type == DeclaredOutputType.FILE: file_output_lines.append( f"- `{output.name}`: create the file in the sandbox, run `dify-agent file upload `, " - f"and set `final_output.{output.name}` to the returned AgentStubFileMapping JSON object. " - "Do not call `final_output` before the upload command succeeds. Do not use the local path, " - "filename, URL, or a synthesized/base64-encoded value as the `reference`." + f"then set `final_output.{output.name}` to a `tool_file` mapping using the returned " + f"`reference`. Do not call `final_output` before the upload command succeeds. Do not use " + "the local path, filename, URL, or a synthesized/base64-encoded value as the `reference`." ) elif ( output.type == DeclaredOutputType.ARRAY @@ -657,9 +685,9 @@ class WorkflowAgentRuntimeRequestBuilder: ): file_output_lines.append( f"- `{output.name}`: for every produced file, run `dify-agent file upload ` and set " - f"`final_output.{output.name}` to an array of the returned AgentStubFileMapping JSON objects. " - "Do not call `final_output` before all upload commands succeed. Do not use local paths, filenames, " - "URLs, or synthesized/base64-encoded values as `reference` values." + f"`final_output.{output.name}` to an array of `tool_file` mappings using the returned " + f"`reference` values. Do not call `final_output` before all upload commands succeed. Do not use " + "local paths, filenames, URLs, or synthesized/base64-encoded values as `reference` values." ) if not file_output_lines: return None @@ -667,8 +695,12 @@ class WorkflowAgentRuntimeRequestBuilder: return "\n".join( [ "When filling file outputs, do not return a local filesystem path directly.", - "Upload each sandbox-local file through the Agent Stub CLI first. Copy the JSON printed by " - "`dify-agent file upload ` verbatim into the final output; never invent the `reference` value.", + "Upload each sandbox-local file through the Agent Stub CLI first. For structured `final_output`, use " + "only the accepted file-mapping shape and the returned `reference`; never invent the `reference` " + "value.", + "If you are replying to the user in natural language and want them to open or download the produced " + "file, include the returned `download_url` in that reply instead of copying it into structured " + "`final_output` unless the schema explicitly asks for it.", *file_output_lines, ] ) @@ -766,6 +798,26 @@ def build_knowledge_layer_config(agent_soul: AgentSoulConfig) -> DifyKnowledgeBa def _knowledge_retrieval_config(retrieval: AgentKnowledgeRetrievalConfig) -> DifyKnowledgeRetrievalConfig: + weights = None + if retrieval.weights is not None: + # The dify-agent runtime payload only consumes the nested vector/keyword + # settings; ``weight_type`` is an API-side authoring detail and must not + # leak into the inner request shape. + weights = ( + cast( + dict[str, Any], + { + key: value + for key, value in { + "vector_setting": retrieval.weights.vector_setting, + "keyword_setting": retrieval.weights.keyword_setting, + }.items() + if value is not None + }, + ) + or None + ) + return DifyKnowledgeRetrievalConfig( mode=retrieval.mode, top_k=retrieval.top_k, @@ -778,9 +830,7 @@ def _knowledge_retrieval_config(retrieval: AgentKnowledgeRetrievalConfig) -> Dif ) if retrieval.reranking_model is not None else None, - weights=cast(dict[str, Any], retrieval.weights.model_dump(mode="json", exclude_none=True)) - if retrieval.weights is not None - else None, + weights=weights, model=_knowledge_model_config(retrieval.model), ) diff --git a/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py b/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py index 809f63b556f..c3a7c23a56c 100644 --- a/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py +++ b/api/core/workflow/nodes/agent_v2/session_cleanup_layer.py @@ -1,11 +1,11 @@ +"""Workflow terminal layer that retires Agent backend sessions asynchronously.""" + from __future__ import annotations import logging from typing import override -from clients.agent_backend import AgentBackendError, AgentBackendRunClient, AgentBackendRunRequestBuilder -from clients.agent_backend.factory import create_agent_backend_run_client -from configs import dify_config +from clients.agent_backend import AgentBackendSessionCleanupPayload from core.workflow.system_variables import SystemVariableKey, get_system_text from graphon.graph_engine.layers import GraphEngineLayer from graphon.graph_events import ( @@ -15,57 +15,23 @@ from graphon.graph_events import ( GraphRunPartialSucceededEvent, GraphRunSucceededEvent, ) +from tasks.agent_backend_session_cleanup_task import cleanup_workflow_agent_runtime_session from .session_store import StoredWorkflowAgentSession, WorkflowAgentRuntimeSessionStore logger = logging.getLogger(__name__) -# Upper bound on how long a cleanup-only run is allowed to settle before the -# layer gives up and leaves the row ACTIVE so it can be retried later. Cleanup -# work is mostly local agent-backend bookkeeping (no LLM inference), so 30s is -# generous; a hung backend should never block workflow termination beyond this. -_CLEANUP_WAIT_TIMEOUT_SECONDS = 30.0 - - class WorkflowAgentSessionCleanupLayer(GraphEngineLayer): - """Retires workflow Agent session snapshots when a workflow reaches a terminal state. + """Retire workflow-owned Agent runtime sessions when the workflow ends. - Implementation notes — there are two failure modes the cleanup path has to - avoid simultaneously: - - 1. The agenton compositor on the agent-backend side validates the cleanup - request's session snapshot against the replayed composition before - running any lifecycle hook. If the snapshot's layer names diverge from - the composition, the run fails asynchronously with ``run_failed`` — but - the initial ``POST /runs`` already returned 202, so the API side has no - visibility of the failure unless it waits for terminal status. The - ``runtime_layer_specs`` persistence in A.1–A.4 plus the - ``_filter_snapshot_to_specs`` shape in ``build_cleanup_request`` keeps - the two name lists in sync. - - 2. The current agent backend's ``runner.py::_run_agent`` always invokes - ``run.get_layer("llm")`` and the structured-output / history validators - before exiting any slot — there is no ``purpose: "cleanup"`` branch - yet. A truly cleanup-only request (no LLM layer) therefore still - crashes inside the runner with ``Layer 'llm' is not defined in this - compositor run.``. Until the backend grows a cleanup-only purpose, - this layer **does not issue an HTTP cleanup run**: it simply retires - the local snapshot row so stale state cannot be re-resumed, and lets - the agent backend's own retention TTL release the suspended layers. - - The HTTP-cleanup machinery (``build_cleanup_request`` + ``wait_run``) is - intentionally still wired into the request builder + integration tests so - that when the agent backend supports cleanup runs we can flip the switch - here with a one-line change (see ``_HTTP_CLEANUP_SUPPORTED``). + Workflow termination is a product-lifecycle boundary: once the run reaches a + terminal graph event, the local session row must no longer be resumable. The + actual Agent backend cleanup is therefore dispatched asynchronously with the + persisted snapshot/specs payload, while the local row is marked CLEANED + immediately afterwards regardless of enqueue outcome. """ - # Flip to True once dify-agent's runner has a ``purpose=cleanup`` branch - # that skips the LLM/output/user-prompt invariants. Until then we only - # update the local row; the spec list is still persisted so the future - # HTTP cleanup path has everything it needs. - _HTTP_CLEANUP_SUPPORTED: bool = False - _TERMINAL_EVENTS = ( GraphRunSucceededEvent, GraphRunPartialSucceededEvent, @@ -73,19 +39,9 @@ class WorkflowAgentSessionCleanupLayer(GraphEngineLayer): GraphRunAbortedEvent, ) - def __init__( - self, - *, - session_store: WorkflowAgentRuntimeSessionStore, - request_builder: AgentBackendRunRequestBuilder, - agent_backend_client: AgentBackendRunClient | None, - cleanup_wait_timeout_seconds: float = _CLEANUP_WAIT_TIMEOUT_SECONDS, - ) -> None: + def __init__(self, *, session_store: WorkflowAgentRuntimeSessionStore) -> None: super().__init__() self._session_store = session_store - self._request_builder = request_builder - self._agent_backend_client = agent_backend_client - self._cleanup_wait_timeout_seconds = cleanup_wait_timeout_seconds @override def on_graph_start(self) -> None: @@ -112,136 +68,59 @@ class WorkflowAgentSessionCleanupLayer(GraphEngineLayer): def _cleanup_session(self, stored_session: StoredWorkflowAgentSession) -> None: scope = stored_session.scope - if not self._HTTP_CLEANUP_SUPPORTED: - # Agent backend has no cleanup-only run mode yet (see class - # docstring). Retire the local row so future re-entries do not - # resume from stale state, and let the backend's retention TTL - # release the suspended layers on its own schedule. - logger.info( - "Workflow Agent session retired locally; HTTP cleanup is disabled " - "until the agent backend supports a cleanup-only run mode. " - "workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s previous_run_id=%s", + try: + if stored_session.runtime_layer_specs: + payload = AgentBackendSessionCleanupPayload( + session_snapshot=stored_session.session_snapshot, + runtime_layer_specs=stored_session.runtime_layer_specs, + idempotency_key=f"{scope.workflow_run_id}:{scope.node_id}:{scope.binding_id}:agent-session-cleanup", + metadata={ + "tenant_id": scope.tenant_id, + "app_id": scope.app_id, + "workflow_id": scope.workflow_id, + "workflow_run_id": scope.workflow_run_id, + "node_id": scope.node_id, + "node_execution_id": scope.node_execution_id, + "binding_id": scope.binding_id, + "agent_id": scope.agent_id, + "agent_config_snapshot_id": scope.agent_config_snapshot_id, + "previous_agent_backend_run_id": stored_session.backend_run_id, + }, + ) + cleanup_workflow_agent_runtime_session.delay(payload.model_dump(mode="json")) + else: + logger.warning( + "Skipping workflow Agent backend cleanup enqueue: no runtime_layer_specs persisted. " + "workflow_run_id=%s node_id=%s agent_id=%s", + scope.workflow_run_id, + scope.node_id, + scope.agent_id, + ) + except Exception: + logger.warning( + "Failed to enqueue workflow Agent backend cleanup: " + "workflow_run_id=%s node_id=%s agent_id=%s previous_run_id=%s", scope.workflow_run_id, scope.node_id, - scope.binding_id, scope.agent_id, stored_session.backend_run_id, - ) - self._session_store.mark_cleaned(scope=scope, backend_run_id=stored_session.backend_run_id) - return - - if self._agent_backend_client is None: - # HTTP cleanup was enabled by the caller but no client was wired - # in (e.g. the API runs without AGENT_BACKEND_BASE_URL configured). - # Leave the row ACTIVE so an operator restart with proper config - # can drive the cleanup; do not silently retire it. - logger.warning( - "Skipping Agent backend cleanup: HTTP cleanup is enabled but no agent " - "backend client is wired in. workflow_run_id=%s node_id=%s agent_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - ) - return - - if not stored_session.runtime_layer_specs: - # Sessions persisted before A.1 landed do not carry the spec list, - # so we cannot replay a valid cleanup composition. Leave the row - # ACTIVE and warn so the absence shows up in observability rather - # than being silently swallowed by a doomed cleanup run. - logger.warning( - "Skipping Agent backend cleanup: no runtime_layer_specs persisted. " - "workflow_run_id=%s node_id=%s agent_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - ) - return - - request = self._request_builder.build_cleanup_request( - session_snapshot=stored_session.session_snapshot, - runtime_layer_specs=stored_session.runtime_layer_specs, - idempotency_key=f"{scope.workflow_run_id}:{scope.node_id}:{scope.binding_id}:agent-session-cleanup", - metadata={ - "tenant_id": scope.tenant_id, - "app_id": scope.app_id, - "workflow_id": scope.workflow_id, - "workflow_run_id": scope.workflow_run_id, - "node_id": scope.node_id, - "node_execution_id": scope.node_execution_id, - "binding_id": scope.binding_id, - "agent_id": scope.agent_id, - "agent_config_snapshot_id": scope.agent_config_snapshot_id, - "previous_agent_backend_run_id": stored_session.backend_run_id, - }, - ) - try: - response = self._agent_backend_client.create_run(request) - except AgentBackendError: - logger.warning( - "Agent backend session cleanup request failed: workflow_run_id=%s node_id=%s agent_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, exc_info=True, ) - return - - try: - status_response = self._agent_backend_client.wait_run( - response.run_id, timeout_seconds=self._cleanup_wait_timeout_seconds - ) - except AgentBackendError: - logger.warning( - "Agent backend session cleanup wait_run failed: " - "workflow_run_id=%s node_id=%s agent_id=%s cleanup_run_id=%s", - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - response.run_id, - exc_info=True, - ) - return - - if status_response.status != "succeeded": - logger.warning( - "Agent backend session cleanup did not succeed: status=%s error=%s " - "workflow_run_id=%s node_id=%s agent_id=%s cleanup_run_id=%s", - status_response.status, - status_response.error, - scope.workflow_run_id, - scope.node_id, - scope.agent_id, - response.run_id, - ) - return - - self._session_store.mark_cleaned(scope=scope, backend_run_id=response.run_id) + finally: + try: + self._session_store.mark_cleaned(scope=scope, backend_run_id=stored_session.backend_run_id) + except Exception: + logger.warning( + "Failed to retire workflow Agent runtime session after cleanup enqueue: " + "workflow_run_id=%s node_id=%s agent_id=%s previous_run_id=%s", + scope.workflow_run_id, + scope.node_id, + scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) def build_workflow_agent_session_cleanup_layer() -> WorkflowAgentSessionCleanupLayer: - """Wire the cleanup layer with the standard production dependencies. - - The agent backend client is constructed only when ``AGENT_BACKEND_BASE_URL`` - is configured (or the deterministic fake is explicitly enabled). When - neither is set — for example unit tests that bring up the workflow runner - without an Agent node — we pass ``None`` so the layer stays harmless. With - ``_HTTP_CLEANUP_SUPPORTED = False`` the local-retire branch never touches - the client anyway, but keeping it ``None`` avoids importing httpx and lets - test harnesses skip backend configuration. - """ - agent_backend_client: AgentBackendRunClient | None - if dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL: - agent_backend_client = create_agent_backend_run_client( - base_url=dify_config.AGENT_BACKEND_BASE_URL, - use_fake=dify_config.AGENT_BACKEND_USE_FAKE, - fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO, - ) - else: - agent_backend_client = None - - return WorkflowAgentSessionCleanupLayer( - session_store=WorkflowAgentRuntimeSessionStore(), - request_builder=AgentBackendRunRequestBuilder(), - agent_backend_client=agent_backend_client, - ) + """Wire the cleanup layer with the standard workflow-owned session store.""" + return WorkflowAgentSessionCleanupLayer(session_store=WorkflowAgentRuntimeSessionStore()) diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py index 9de26b8214b..fb12922ed7f 100644 --- a/api/core/workflow/workflow_entry.py +++ b/api/core/workflow/workflow_entry.py @@ -46,18 +46,26 @@ logger = logging.getLogger(__name__) _file_access_controller = DatabaseFileAccessController() -def iter_dify_graph_engine_events(engine: GraphEngine) -> Generator[GraphEngineEvent, None, None]: +def iter_dify_graph_engine_events( + engine: GraphEngine, + response_stream_filter: ResponseStreamFilter | None = None, +) -> Generator[GraphEngineEvent, None, None]: """ Apply Dify's response streaming compatibility filter to GraphEngine events. Graphon v0.5.0 emits raw variable stream chunks and requires callers to opt into the legacy response-ordered stream behavior that Dify exposes to its workflow runners and tests. + + ``response_stream_filter``, when supplied, must be the same instance a + caller intends to persist on pause (see ``PauseStatePersistenceLayer``) so + the filter's ``paths_map`` reflects everything the engine has actually + streamed for this run. """ yield from filter_graph_events( engine.run(), context=GraphEventFilterContext.from_engine(engine), - filters=[ResponseStreamFilter()], + filters=[response_stream_filter or ResponseStreamFilter()], ) @@ -167,6 +175,7 @@ class WorkflowEntry: variable_pool: VariablePool, graph_runtime_state: GraphRuntimeState, command_channel: CommandChannel | None = None, + response_stream_filter: ResponseStreamFilter | None = None, ) -> None: """ Init workflow entry @@ -183,6 +192,8 @@ class WorkflowEntry: :param variable_pool: variable pool :param graph_runtime_state: pre-created graph runtime state :param command_channel: command channel for external control (optional, defaults to InMemoryChannel) + :param response_stream_filter: pre-restored filter for resumed runs (optional, defaults to a fresh + ResponseStreamFilter for runs with no prior pause) :param thread_pool_id: thread pool id """ # check call depth @@ -195,6 +206,7 @@ class WorkflowEntry: command_channel = InMemoryChannel() self.command_channel = command_channel + self._response_stream_filter = response_stream_filter or ResponseStreamFilter() execution_context = capture_current_context() graph_runtime_state.execution_context = execution_context self._child_engine_builder = _WorkflowChildEngineBuilder(tenant_id=tenant_id) @@ -240,7 +252,7 @@ class WorkflowEntry: try: # Preserve Dify's response-stream semantics on top of Graphon 0.5.0. - generator = iter_dify_graph_engine_events(graph_engine) + generator = iter_dify_graph_engine_events(graph_engine, self._response_stream_filter) yield from generator except GenerateTaskStoppedError: pass diff --git a/api/events/event_handlers/update_provider_when_message_created.py b/api/events/event_handlers/update_provider_when_message_created.py index 8dec5876a9b..15b40afdbf2 100644 --- a/api/events/event_handlers/update_provider_when_message_created.py +++ b/api/events/event_handlers/update_provider_when_message_created.py @@ -204,6 +204,7 @@ def _deduct_credit_pool_quota_capped(*, tenant_id: str, credits_required: int, p tenant_id=tenant_id, credits_required=credits_required, pool_type=pool_type, + session=db.session(), ) if deducted_credits < credits_required: logger.warning( diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index f6496c70a78..6515b22eb36 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -84,7 +84,7 @@ def load_user_from_request(request_from_flask_login: Request) -> LoginUser | Non if not user_id: raise Unauthorized("Invalid Authorization token.") - logged_in_account = AccountService.load_logged_in_account(account_id=user_id, session=db.session) + logged_in_account = AccountService.load_logged_in_account(account_id=user_id, session=db.session()) return logged_in_account elif request.blueprint == "openapi": # Account-branch device-flow approval routes (approve / deny / @@ -103,7 +103,7 @@ def load_user_from_request(request_from_flask_login: Request) -> LoginUser | Non source = decoded.get("token_source") if source or not user_id: return None - return AccountService.load_logged_in_account(account_id=user_id, session=db.session) + return AccountService.load_logged_in_account(account_id=user_id, session=db.session()) elif request.blueprint == "web": app_code = request.headers.get(HEADER_NAME_APP_CODE) webapp_token = extract_webapp_passport(app_code, request) if app_code else None diff --git a/api/fields/app_fields.py b/api/fields/app_fields.py deleted file mode 100644 index 96d8fbdf34c..00000000000 --- a/api/fields/app_fields.py +++ /dev/null @@ -1,271 +0,0 @@ -import json -from typing import override - -from flask_restx import fields - -from fields.workflow_fields import workflow_partial_fields -from libs.helper import AppIconUrlField, TimestampField - - -class JsonStringField(fields.Raw): - @override - def format(self, value): - if isinstance(value, str): - try: - return json.loads(value) - except (json.JSONDecodeError, TypeError): - return value - return value - - -class OpaqueRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "object"} - - -class StringListRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "array", "items": {"type": "string"}} - - -class ObjectListRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "array", "items": {"type": "object"}} - - -app_detail_kernel_fields = { - "id": fields.String, - "name": fields.String, - "description": fields.String, - "mode": fields.String(attribute="mode_compatible_with_agent"), - "icon_type": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "icon_url": AppIconUrlField, -} - -related_app_list = { - "data": fields.List(fields.Nested(app_detail_kernel_fields)), - "total": fields.Integer, -} - -model_config_fields = { - "opening_statement": fields.String, - "suggested_questions": StringListRawField(attribute="suggested_questions_list"), - "suggested_questions_after_answer": OpaqueRawField(attribute="suggested_questions_after_answer_dict"), - "speech_to_text": OpaqueRawField(attribute="speech_to_text_dict"), - "text_to_speech": OpaqueRawField(attribute="text_to_speech_dict"), - "retriever_resource": OpaqueRawField(attribute="retriever_resource_dict"), - "annotation_reply": OpaqueRawField(attribute="annotation_reply_dict"), - "more_like_this": OpaqueRawField(attribute="more_like_this_dict"), - "sensitive_word_avoidance": OpaqueRawField(attribute="sensitive_word_avoidance_dict"), - "external_data_tools": ObjectListRawField(attribute="external_data_tools_list"), - "model": OpaqueRawField(attribute="model_dict"), - "user_input_form": ObjectListRawField(attribute="user_input_form_list"), - "dataset_query_variable": fields.String, - "pre_prompt": fields.String, - "agent_mode": OpaqueRawField(attribute="agent_mode_dict"), - "prompt_type": fields.String, - "chat_prompt_config": OpaqueRawField(attribute="chat_prompt_config_dict"), - "completion_prompt_config": OpaqueRawField(attribute="completion_prompt_config_dict"), - "dataset_configs": OpaqueRawField(attribute="dataset_configs_dict"), - "file_upload": OpaqueRawField(attribute="file_upload_dict"), - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, -} - -tag_fields = {"id": fields.String, "name": fields.String, "type": fields.String} - -app_detail_fields = { - "id": fields.String, - "name": fields.String, - "description": fields.String, - "mode": fields.String(attribute="mode_compatible_with_agent"), - "icon": fields.String, - "icon_background": fields.String, - "enable_site": fields.Boolean, - "enable_api": fields.Boolean, - "model_config": fields.Nested(model_config_fields, attribute="app_model_config", allow_null=True), - "workflow": fields.Nested(workflow_partial_fields, allow_null=True), - "tracing": OpaqueRawField, - "use_icon_as_answer_icon": fields.Boolean, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, - "access_mode": fields.String, - "tags": fields.List(fields.Nested(tag_fields)), - "permission_keys": fields.List(fields.String()), -} - -prompt_config_fields = { - "prompt_template": fields.String, -} - -model_config_partial_fields = { - "model": OpaqueRawField(attribute="model_dict"), - "pre_prompt": fields.String, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, -} - -app_partial_fields = { - "id": fields.String, - "name": fields.String, - "max_active_requests": OpaqueRawField(), - "description": fields.String(attribute="desc_or_prompt"), - "mode": fields.String(attribute="mode_compatible_with_agent"), - "icon_type": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "icon_url": AppIconUrlField, - "model_config": fields.Nested(model_config_partial_fields, attribute="app_model_config", allow_null=True), - "workflow": fields.Nested(workflow_partial_fields, allow_null=True), - "use_icon_as_answer_icon": fields.Boolean, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, - "tags": fields.List(fields.Nested(tag_fields)), - "access_mode": fields.String, - "create_user_name": fields.String, - "author_name": fields.String, - "has_draft_trigger": fields.Boolean, - "permission_keys": fields.List(fields.String()), -} - - -app_pagination_fields = { - "page": fields.Integer, - "limit": fields.Integer(attribute="per_page"), - "total": fields.Integer, - "has_more": fields.Boolean(attribute="has_next"), - "data": fields.List(fields.Nested(app_partial_fields), attribute="items"), -} - -template_fields = { - "name": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "description": fields.String, - "mode": fields.String, - "model_config": fields.Nested(model_config_fields), -} - -template_list_fields = { - "data": fields.List(fields.Nested(template_fields)), -} - -site_fields = { - "access_token": fields.String(attribute="code"), - "code": fields.String, - "title": fields.String, - "icon_type": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "icon_url": AppIconUrlField, - "description": fields.String, - "default_language": fields.String, - "chat_color_theme": fields.String, - "chat_color_theme_inverted": fields.Boolean, - "customize_domain": fields.String, - "copyright": fields.String, - "privacy_policy": fields.String, - "custom_disclaimer": fields.String, - "customize_token_strategy": fields.String, - "prompt_public": fields.Boolean, - "app_base_url": fields.String, - "show_workflow_steps": fields.Boolean, - "use_icon_as_answer_icon": fields.Boolean, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, -} - -deleted_tool_fields = { - "type": fields.String, - "tool_name": fields.String, - "provider_id": fields.String, -} - -app_detail_fields_with_site = { - "id": fields.String, - "name": fields.String, - "description": fields.String, - "mode": fields.String(attribute="mode_compatible_with_agent"), - "icon_type": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "icon_url": AppIconUrlField, - "enable_site": fields.Boolean, - "enable_api": fields.Boolean, - "model_config": fields.Nested(model_config_fields, attribute="app_model_config", allow_null=True), - "workflow": fields.Nested(workflow_partial_fields, allow_null=True), - "api_base_url": fields.String, - "use_icon_as_answer_icon": fields.Boolean, - "max_active_requests": fields.Integer, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, - "deleted_tools": fields.List(fields.Nested(deleted_tool_fields)), - "access_mode": fields.String, - "tags": fields.List(fields.Nested(tag_fields)), - "permission_keys": fields.List(fields.String()), - "site": fields.Nested(site_fields), -} - - -app_site_fields = { - "app_id": fields.String, - "access_token": fields.String(attribute="code"), - "code": fields.String, - "title": fields.String, - "icon": fields.String, - "icon_background": fields.String, - "description": fields.String, - "default_language": fields.String, - "customize_domain": fields.String, - "copyright": fields.String, - "privacy_policy": fields.String, - "custom_disclaimer": fields.String, - "customize_token_strategy": fields.String, - "prompt_public": fields.Boolean, - "show_workflow_steps": fields.Boolean, - "use_icon_as_answer_icon": fields.Boolean, -} - -leaked_dependency_fields = {"type": fields.String, "value": OpaqueRawField, "current_identifier": fields.String} - -app_import_fields = { - "id": fields.String, - "status": fields.String, - "app_id": fields.String, - "app_mode": fields.String, - "current_dsl_version": fields.String, - "imported_dsl_version": fields.String, - "error": fields.String, -} - -app_import_check_dependencies_fields = { - "leaked_dependencies": fields.List(fields.Nested(leaked_dependency_fields)), -} - -app_server_fields = { - "id": fields.String, - "name": fields.String, - "server_code": fields.String, - "description": fields.String, - "status": fields.String, - "parameters": JsonStringField, - "created_at": TimestampField, - "updated_at": TimestampField, -} diff --git a/api/fields/conversation_fields.py b/api/fields/conversation_fields.py index d256ad96cf0..72eec3f1c7d 100644 --- a/api/fields/conversation_fields.py +++ b/api/fields/conversation_fields.py @@ -110,6 +110,7 @@ class AgentThought(ResponseModel): message_id: str position: int thought: str | None = None + answer: str | None = None tool: str | None = None tool_labels: JSONValue tool_input: str | None = None diff --git a/api/fields/member_fields.py b/api/fields/member_fields.py index 80b93f0a24b..5108522af01 100644 --- a/api/fields/member_fields.py +++ b/api/fields/member_fields.py @@ -2,18 +2,11 @@ from __future__ import annotations from datetime import datetime -from flask_restx import fields from pydantic import Field, computed_field, field_validator from fields.base import ResponseModel from libs.helper import build_avatar_url, to_timestamp -simple_account_fields = { - "id": fields.String, - "name": fields.String, - "email": fields.String, -} - class SimpleAccountResponse(ResponseModel): id: str diff --git a/api/fields/raws.py b/api/fields/raws.py deleted file mode 100644 index c7e047626f1..00000000000 --- a/api/fields/raws.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import override - -from flask_restx import fields - -from graphon.file import File - - -class FilesContainedField(fields.Raw): - @override - def format(self, value): - return self._format_file_object(value) - - def _format_file_object(self, v): - if isinstance(v, File): - return v.model_dump() - if isinstance(v, dict): - return {k: self._format_file_object(vv) for k, vv in v.items()} - if isinstance(v, list): - return [self._format_file_object(vv) for vv in v] - return v diff --git a/api/fields/snippet_fields.py b/api/fields/snippet_fields.py index 699a3687ac1..77aba3fe76e 100644 --- a/api/fields/snippet_fields.py +++ b/api/fields/snippet_fields.py @@ -1,61 +1,74 @@ -from typing import override +from datetime import datetime +from typing import Any -from flask_restx import fields +from pydantic import Field, field_validator -from fields.member_fields import simple_account_fields -from libs.helper import TimestampField +from fields.base import ResponseModel +from fields.member_fields import SimpleAccountResponse +from libs.helper import to_timestamp +from models.snippet import SnippetType -class OpaqueRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "object"} +class SnippetTagResponse(ResponseModel): + id: str + name: str + type: str -tag_fields = {"id": fields.String, "name": fields.String, "type": fields.String} +class SnippetListItemResponse(ResponseModel): + id: str + name: str + description: str | None + type: SnippetType + version: int + use_count: int + is_published: bool + icon_info: dict[str, Any] | None + tags: list[SnippetTagResponse] + created_by: str | None + author_name: str | None + created_at: int + updated_by: str | None + updated_at: int -# Snippet list item fields (lightweight for list display) -snippet_list_fields = { - "id": fields.String, - "name": fields.String, - "description": fields.String, - "type": fields.String, - "version": fields.Integer, - "use_count": fields.Integer, - "is_published": fields.Boolean, - "icon_info": OpaqueRawField, - "tags": fields.List(fields.Nested(tag_fields)), - "created_by": fields.String, - "author_name": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, -} + @field_validator("created_at", "updated_at", mode="before") + @classmethod + def _normalize_timestamp(cls, value: datetime | int | None) -> int: + timestamp = to_timestamp(value) + if timestamp is None: + raise ValueError("timestamp is required") + return timestamp -# Full snippet fields (includes creator info and graph data) -snippet_fields = { - "id": fields.String, - "name": fields.String, - "description": fields.String, - "type": fields.String, - "version": fields.Integer, - "use_count": fields.Integer, - "is_published": fields.Boolean, - "icon_info": OpaqueRawField, - "graph": OpaqueRawField(attribute="graph_dict"), - "input_fields": OpaqueRawField(attribute="input_fields_list"), - "tags": fields.List(fields.Nested(tag_fields)), - "created_by": fields.Nested(simple_account_fields, attribute="created_by_account", allow_null=True), - "created_at": TimestampField, - "updated_by": fields.Nested(simple_account_fields, attribute="updated_by_account", allow_null=True), - "updated_at": TimestampField, -} -# Pagination response fields -snippet_pagination_fields = { - "data": fields.List(fields.Nested(snippet_list_fields)), - "page": fields.Integer, - "limit": fields.Integer, - "total": fields.Integer, - "has_more": fields.Boolean, -} +class SnippetResponse(ResponseModel): + id: str + name: str + description: str | None + type: SnippetType + version: int + use_count: int + is_published: bool + icon_info: dict[str, Any] | None + graph: dict[str, Any] = Field(validation_alias="graph_dict") + input_fields: list[dict[str, Any]] = Field(validation_alias="input_fields_list") + tags: list[SnippetTagResponse] + created_by: SimpleAccountResponse | None = Field(validation_alias="created_by_account") + created_at: int + updated_by: SimpleAccountResponse | None = Field(validation_alias="updated_by_account") + updated_at: int + + @field_validator("created_at", "updated_at", mode="before") + @classmethod + def _normalize_timestamp(cls, value: datetime | int | None) -> int: + timestamp = to_timestamp(value) + if timestamp is None: + raise ValueError("timestamp is required") + return timestamp + + +class SnippetPaginationResponse(ResponseModel): + data: list[SnippetListItemResponse] + page: int + limit: int + total: int + has_more: bool diff --git a/api/fields/workflow_fields.py b/api/fields/workflow_fields.py deleted file mode 100644 index 2d0d8f9f546..00000000000 --- a/api/fields/workflow_fields.py +++ /dev/null @@ -1,129 +0,0 @@ -from typing import override - -from flask_restx import fields - -from core.helper import encrypter -from fields.member_fields import simple_account_fields -from graphon.variables import SecretVariable, SegmentType, VariableBase -from libs.helper import TimestampField - -from ._value_type_serializer import serialize_value_type - -ENVIRONMENT_VARIABLE_SUPPORTED_TYPES = (SegmentType.STRING, SegmentType.NUMBER, SegmentType.SECRET) - - -class OpaqueRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "object"} - - -class JsonValueRawField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return { - "anyOf": [ - {"type": "string"}, - {"type": "integer"}, - {"type": "number"}, - {"type": "boolean"}, - {"type": "object", "additionalProperties": True}, - {"type": "array", "items": {}}, - {"type": "null"}, - ] - } - - -class EnvironmentVariableField(fields.Raw): - @override - def schema(self) -> dict[str, object]: - return {"type": "object"} - - @override - def format(self, value): - # Mask secret variables values in environment_variables - if isinstance(value, SecretVariable): - return { - "id": value.id, - "name": value.name, - "value": encrypter.full_mask_token(), - "value_type": value.value_type.value, - "description": value.description, - } - if isinstance(value, VariableBase): - return { - "id": value.id, - "name": value.name, - "value": value.value, - "value_type": str(value.value_type.exposed_type()), - "description": value.description, - } - if isinstance(value, dict): - value_type_str = value.get("value_type") - if not isinstance(value_type_str, str): - raise TypeError( - f"unexpected type for value_type field, value={value_type_str}, type={type(value_type_str)}" - ) - value_type = SegmentType(value_type_str).exposed_type() - if value_type not in ENVIRONMENT_VARIABLE_SUPPORTED_TYPES: - raise ValueError(f"Unsupported environment variable value type: {value_type}") - return value - - -conversation_variable_fields = { - "id": fields.String, - "name": fields.String, - "value_type": fields.String(attribute=serialize_value_type), - "value": JsonValueRawField, - "description": fields.String, -} - -pipeline_variable_fields = { - "label": fields.String, - "variable": fields.String, - "type": fields.String, - "belong_to_node_id": fields.String, - "max_length": fields.Integer, - "required": fields.Boolean, - "unit": fields.String, - "default_value": JsonValueRawField, - "options": fields.List(fields.String), - "placeholder": fields.String, - "tooltips": fields.String, - "allowed_file_types": fields.List(fields.String), - "allow_file_extension": fields.List(fields.String), - "allow_file_upload_methods": fields.List(fields.String), -} - -workflow_fields = { - "id": fields.String, - "graph": OpaqueRawField(attribute="graph_dict"), - "features": OpaqueRawField(attribute="features_dict"), - "hash": fields.String(attribute="unique_hash"), - "version": fields.String, - "marked_name": fields.String, - "marked_comment": fields.String, - "created_by": fields.Nested(simple_account_fields, attribute="created_by_account"), - "created_at": TimestampField, - "updated_by": fields.Nested(simple_account_fields, attribute="updated_by_account", allow_null=True), - "updated_at": TimestampField, - "tool_published": fields.Boolean, - "environment_variables": fields.List(EnvironmentVariableField()), - "conversation_variables": fields.List(fields.Nested(conversation_variable_fields)), - "rag_pipeline_variables": fields.List(fields.Nested(pipeline_variable_fields)), -} - -workflow_partial_fields = { - "id": fields.String, - "created_by": fields.String, - "created_at": TimestampField, - "updated_by": fields.String, - "updated_at": TimestampField, -} - -workflow_pagination_fields = { - "items": fields.List(fields.Nested(workflow_fields), attribute="items"), - "page": fields.Integer, - "limit": fields.Integer(attribute="limit"), - "has_more": fields.Boolean(attribute="has_more"), -} diff --git a/api/models/account.py b/api/models/account.py index df152e1783c..350b811bd46 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -369,17 +369,19 @@ class InvitationCode(TypeBase): ) +class TenantPluginInstallPermission(enum.StrEnum): + EVERYONE = "everyone" + ADMINS = "admins" + NOBODY = "noone" + + +class TenantPluginDebugPermission(enum.StrEnum): + EVERYONE = "everyone" + ADMINS = "admins" + NOBODY = "noone" + + class TenantPluginPermission(TypeBase): - class InstallPermission(enum.StrEnum): - EVERYONE = "everyone" - ADMINS = "admins" - NOBODY = "noone" - - class DebugPermission(enum.StrEnum): - EVERYONE = "everyone" - ADMINS = "admins" - NOBODY = "noone" - __tablename__ = "account_plugin_permissions" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="account_plugin_permission_pkey"), @@ -390,36 +392,42 @@ class TenantPluginPermission(TypeBase): StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) - install_permission: Mapped[InstallPermission] = mapped_column( - EnumText(InstallPermission, length=16), + install_permission: Mapped[TenantPluginInstallPermission] = mapped_column( + EnumText(TenantPluginInstallPermission, length=16), nullable=False, server_default="everyone", - default=InstallPermission.EVERYONE, + default=TenantPluginInstallPermission.EVERYONE, ) - debug_permission: Mapped[DebugPermission] = mapped_column( - EnumText(DebugPermission, length=16), nullable=False, server_default="noone", default=DebugPermission.NOBODY + debug_permission: Mapped[TenantPluginDebugPermission] = mapped_column( + EnumText(TenantPluginDebugPermission, length=16), + nullable=False, + server_default="noone", + default=TenantPluginDebugPermission.NOBODY, ) +class TenantPluginAutoUpgradeCategory(enum.StrEnum): + TOOL = "tool" + MODEL = "model" + EXTENSION = "extension" + AGENT_STRATEGY = "agent-strategy" + DATASOURCE = "datasource" + TRIGGER = "trigger" + + +class TenantPluginAutoUpgradeStrategySetting(enum.StrEnum): + DISABLED = "disabled" + FIX_ONLY = "fix_only" + LATEST = "latest" + + +class TenantPluginAutoUpgradeMode(enum.StrEnum): + ALL = "all" + PARTIAL = "partial" + EXCLUDE = "exclude" + + class TenantPluginAutoUpgradeStrategy(TypeBase): - class PluginCategory(enum.StrEnum): - TOOL = "tool" - MODEL = "model" - EXTENSION = "extension" - AGENT_STRATEGY = "agent-strategy" - DATASOURCE = "datasource" - TRIGGER = "trigger" - - class StrategySetting(enum.StrEnum): - DISABLED = "disabled" - FIX_ONLY = "fix_only" - LATEST = "latest" - - class UpgradeMode(enum.StrEnum): - ALL = "all" - PARTIAL = "partial" - EXCLUDE = "exclude" - __tablename__ = "tenant_plugin_auto_upgrade_strategies" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="tenant_plugin_auto_upgrade_strategy_pkey"), @@ -431,20 +439,23 @@ class TenantPluginAutoUpgradeStrategy(TypeBase): StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) - category: Mapped[PluginCategory] = mapped_column( - EnumText(PluginCategory, length=32), + category: Mapped[TenantPluginAutoUpgradeCategory] = mapped_column( + EnumText(TenantPluginAutoUpgradeCategory, length=32), nullable=False, server_default="tool", - default=PluginCategory.TOOL, + default=TenantPluginAutoUpgradeCategory.TOOL, ) - strategy_setting: Mapped[StrategySetting] = mapped_column( - EnumText(StrategySetting, length=16), + strategy_setting: Mapped[TenantPluginAutoUpgradeStrategySetting] = mapped_column( + EnumText(TenantPluginAutoUpgradeStrategySetting, length=16), nullable=False, server_default="fix_only", - default=StrategySetting.FIX_ONLY, + default=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, ) - upgrade_mode: Mapped[UpgradeMode] = mapped_column( - EnumText(UpgradeMode, length=16), nullable=False, server_default="exclude", default=UpgradeMode.EXCLUDE + upgrade_mode: Mapped[TenantPluginAutoUpgradeMode] = mapped_column( + EnumText(TenantPluginAutoUpgradeMode, length=16), + nullable=False, + server_default="exclude", + default=TenantPluginAutoUpgradeMode.EXCLUDE, ) exclude_plugins: Mapped[list[str]] = mapped_column(sa.JSON, nullable=False, default_factory=list) include_plugins: Mapped[list[str]] = mapped_column(sa.JSON, nullable=False, default_factory=list) diff --git a/api/models/agent_config_entities.py b/api/models/agent_config_entities.py index 64fe85930c6..a78f7d90dd5 100644 --- a/api/models/agent_config_entities.py +++ b/api/models/agent_config_entities.py @@ -8,7 +8,7 @@ from pydantic import BaseModel, ConfigDict, Field, WithJsonSchema, field_validat from core.rag.entities.metadata_entities import ConditionValue, SupportedComparisonOperator from core.workflow.file_reference import is_canonical_file_reference -from graphon.file import FileTransferMethod +from graphon.file import FileTransferMethod, FileType class AgentKnowledgeQueryMode(StrEnum): @@ -314,8 +314,9 @@ class AgentKnowledgeQueryConfig(BaseModel): Agent v2 stores knowledge as explicit ``knowledge.sets`` rather than the legacy flat ``datasets`` / ``query_mode`` / ``query_config`` shape. Each - set owns its own query policy, so ``user_query`` must carry an explicit - ``value`` while ``generated_query`` leaves that value empty. + set owns its own query policy. Mode-dependent completeness, such as + requiring ``value`` for ``user_query``, is enforced by composer publish + validation so draft saves can persist partially configured knowledge sets. """ model_config = ConfigDict(extra="forbid") @@ -323,12 +324,6 @@ class AgentKnowledgeQueryConfig(BaseModel): mode: AgentKnowledgeQueryMode value: str | None = None - @model_validator(mode="after") - def validate_query(self) -> Self: - if self.mode == AgentKnowledgeQueryMode.USER_QUERY and not (self.value or "").strip(): - raise ValueError("knowledge query.value is required for user_query mode") - return self - class AgentKnowledgeModelConfig(BaseModel): model_config = ConfigDict(extra="forbid") @@ -356,8 +351,9 @@ class AgentKnowledgeRetrievalConfig(BaseModel): """Per-set retrieval policy for Agent v2 knowledge retrieval. Retrieval settings now live on each knowledge set instead of one shared - flat config. A set may use either ``multiple`` retrieval with ``top_k`` or - ``single`` retrieval with a required model config. + flat config. Mode-dependent completeness, such as requiring ``top_k`` for + ``multiple`` or a model for ``single``, is enforced by composer publish + validation so draft saves can persist partially configured knowledge sets. """ model_config = ConfigDict(extra="forbid") @@ -371,14 +367,6 @@ class AgentKnowledgeRetrievalConfig(BaseModel): weights: AgentKnowledgeWeightedScoreConfig | None = None model: AgentKnowledgeModelConfig | None = None - @model_validator(mode="after") - def validate_mode_fields(self) -> Self: - if self.mode == "multiple" and self.top_k is None: - raise ValueError("knowledge retrieval.top_k is required for multiple mode") - if self.mode == "single" and self.model is None: - raise ValueError("knowledge retrieval.model is required for single mode") - return self - class AgentKnowledgeMetadataCondition(BaseModel): model_config = ConfigDict(extra="forbid") @@ -401,6 +389,8 @@ class AgentKnowledgeMetadataFilteringConfig(BaseModel): The Python attribute uses ``metadata_model_config`` for clarity because the model belongs to metadata filtering specifically, while the external API and generated schema keep the historical ``model_config`` field name via alias. + Mode-dependent completeness is enforced by composer publish validation so + draft saves can persist partially configured metadata filters. """ model_config = ConfigDict(extra="forbid", populate_by_name=True) @@ -410,14 +400,6 @@ class AgentKnowledgeMetadataFilteringConfig(BaseModel): metadata_model_config: AgentKnowledgeModelConfig | None = Field(default=None, alias="model_config") conditions: AgentKnowledgeMetadataConditions | None = None - @model_validator(mode="after") - def validate_mode_fields(self) -> Self: - if self.mode == "automatic" and self.metadata_model_config is None: - raise ValueError("metadata_filtering.model_config is required for automatic mode") - if self.mode == "manual" and (self.conditions is None or not self.conditions.conditions): - raise ValueError("metadata_filtering.conditions is required for manual mode") - return self - class AgentKnowledgeSetConfig(BaseModel): """One explicit knowledge set in Agent v2. @@ -547,6 +529,23 @@ class AgentSensitiveWordAvoidanceFeatureConfig(AgentFeatureToggleConfig): config: AgentModerationProviderConfig | None = None +class AgentFileUploadImageFeatureConfig(AgentFeatureToggleConfig): + enabled: bool = True + + +class AgentFileUploadFeatureConfig(AgentFeatureToggleConfig): + enabled: bool = True + allowed_file_extensions: list[str] = Field(default_factory=lambda: ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"]) + allowed_file_types: list[FileType] = Field( + default_factory=lambda: [FileType.DOCUMENT, FileType.IMAGE, FileType.AUDIO, FileType.VIDEO] + ) + allowed_file_upload_methods: list[FileTransferMethod] = Field( + default_factory=lambda: [FileTransferMethod.LOCAL_FILE, FileTransferMethod.REMOTE_URL] + ) + image: AgentFileUploadImageFeatureConfig = Field(default_factory=AgentFileUploadImageFeatureConfig) + number_limits: int = 3 + + class AgentSoulAppFeaturesConfig(AgentFlexibleConfig): opening_statement: str | None = None suggested_questions: list[str] | None = None @@ -555,6 +554,7 @@ class AgentSoulAppFeaturesConfig(AgentFlexibleConfig): text_to_speech: AgentTextToSpeechFeatureConfig | None = None retriever_resource: AgentFeatureToggleConfig | None = None sensitive_word_avoidance: AgentSensitiveWordAvoidanceFeatureConfig | None = None + file_upload: AgentFileUploadFeatureConfig = Field(default_factory=AgentFileUploadFeatureConfig) class WorkflowPreviousNodeOutputRef(AgentFlexibleConfig): diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 5f4ecad431c..7822f330d84 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -141,9 +141,9 @@ Get account avatar url #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [EducationActivateResponse](#educationactivateresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Success | ### [GET] /account/education/autocomplete #### Parameters @@ -1268,7 +1268,7 @@ Read a text/binary preview file in an Agent App conversation sandbox | 200 | Preview returned | **application/json**: [SandboxReadResponse](#sandboxreadresponse)
| ### [POST] /agent/{agent_id}/sandbox/files/upload -Upload one Agent App sandbox file as a Dify ToolFile mapping +Upload one Agent App sandbox file and return a signed download URL #### Parameters @@ -1411,7 +1411,7 @@ Infer CLI tool + ENV suggestions from a standardized Agent App skill | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [WorkspaceListResponse](#workspacelistresponse)
| +| 200 | Success | **application/json**: [WorkspacePaginationResponse](#workspacepaginationresponse)
| ### [GET] /api-based-extension Get all API-based extensions for current tenant @@ -3777,7 +3777,7 @@ Read a text/binary preview file in a workflow Agent node sandbox | 200 | Preview returned | **application/json**: [SandboxReadResponse](#sandboxreadresponse)
| ### [POST] /apps/{app_id}/workflow-runs/{workflow_run_id}/agent-nodes/{node_id}/sandbox/files/upload -Upload one workflow Agent sandbox file as a Dify ToolFile mapping +Upload one workflow Agent sandbox file and return a signed download URL #### Parameters @@ -5106,14 +5106,14 @@ Refresh MCP server configuration and regenerate server code | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [DatasourceCredentialsResponse](#datasourcecredentialsresponse)
| +| 200 | Default datasource credentials retrieved successfully | **application/json**: [DatasourceProviderAuthListResponse](#datasourceproviderauthlistresponse)
| ### [GET] /auth/plugin/datasource/list #### Responses | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [DatasourceCredentialsResponse](#datasourcecredentialsresponse)
| +| 200 | Datasource credentials retrieved successfully | **application/json**: [DatasourceProviderAuthListResponse](#datasourceproviderauthlistresponse)
| ### [GET] /auth/plugin/datasource/{provider_id} #### Parameters @@ -5126,7 +5126,7 @@ Refresh MCP server configuration and regenerate server code | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [DatasourceCredentialsResponse](#datasourcecredentialsresponse)
| +| 200 | Datasource credentials retrieved successfully | **application/json**: [DatasourceCredentialListResponse](#datasourcecredentiallistresponse)
| ### [POST] /auth/plugin/datasource/{provider_id} #### Parameters @@ -5145,7 +5145,7 @@ Refresh MCP server configuration and regenerate server code | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Datasource credential created successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [DELETE] /auth/plugin/datasource/{provider_id}/custom-client #### Parameters @@ -5177,7 +5177,7 @@ Refresh MCP server configuration and regenerate server code | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Datasource OAuth custom client saved successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /auth/plugin/datasource/{provider_id}/default #### Parameters @@ -5234,7 +5234,7 @@ Refresh MCP server configuration and regenerate server code | Code | Description | Schema | | ---- | ----------- | ------ | -| 201 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 201 | Datasource credential updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /auth/plugin/datasource/{provider_id}/update-name #### Parameters @@ -7460,9 +7460,9 @@ Get instruction generation template #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 302 | Redirect to console OAuth callback page | **application/json**: [RedirectResponse](#redirectresponse)
| +| Code | Description | +| ---- | ----------- | +| 302 | Redirect to OAuth callback page | ### [GET] /notification Return the active in-product notification for the current user in their interface language (falls back to English if unavailable). The notification is NOT marked as seen here; call POST /notification/dismiss when the user explicitly closes the modal. @@ -7638,9 +7638,9 @@ Initiate OAuth login process #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 302 | Redirect to console OAuth callback page | **application/json**: [RedirectResponse](#redirectresponse)
| +| Code | Description | +| ---- | ----------- | +| 302 | Redirect to OAuth callback page | ### [GET] /oauth/plugin/{provider_id}/datasource/get-authorization-url #### Parameters @@ -7654,7 +7654,7 @@ Initiate OAuth login process | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Authorization URL retrieved successfully | **application/json**: [PluginOAuthAuthorizationUrlResponse](#pluginoauthauthorizationurlresponse)
| +| 200 | Datasource OAuth authorization URL generated successfully | **application/json**: [PluginOAuthAuthorizationUrlResponse](#pluginoauthauthorizationurlresponse)
| ### [GET] /oauth/plugin/{provider}/tool/authorization-url #### Parameters @@ -7667,7 +7667,7 @@ Initiate OAuth login process | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Authorization URL retrieved successfully | **application/json**: [PluginOAuthAuthorizationUrlResponse](#pluginoauthauthorizationurlresponse)
| +| 200 | Tool OAuth authorization URL generated successfully | **application/json**: [PluginOAuthAuthorizationUrlResponse](#pluginoauthauthorizationurlresponse)
| ### [GET] /oauth/plugin/{provider}/tool/callback #### Parameters @@ -7678,9 +7678,9 @@ Initiate OAuth login process #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 302 | Redirect to console OAuth callback page | **application/json**: [RedirectResponse](#redirectresponse)
| +| Code | Description | +| ---- | ----------- | +| 302 | Redirect to OAuth callback page | ### [GET] /oauth/plugin/{provider}/trigger/callback **Handle OAuth callback for trigger provider** @@ -7693,9 +7693,9 @@ Initiate OAuth login process #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 302 | Redirect to console OAuth callback page | **application/json**: [RedirectResponse](#redirectresponse)
| +| Code | Description | +| ---- | ----------- | +| 302 | Redirect to OAuth callback page | ### [POST] /oauth/provider #### Request Body @@ -8477,9 +8477,9 @@ Initiate OAuth login process #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [DataSourceContentPreviewResponse](#datasourcecontentpreviewresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Success | ### [POST] /rag/pipelines/{pipeline_id}/workflows/published/datasource/nodes/{node_id}/run **Run rag pipeline datasource** @@ -8617,6 +8617,7 @@ Initiate OAuth login process | Code | Description | Schema | | ---- | ----------- | ------ | | 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 401 | Unauthorized | **application/json**: [SimpleResultMessageResponse](#simpleresultmessageresponse)
| ### [POST] /remote-files/upload #### Request Body @@ -9432,9 +9433,9 @@ Bedrock retrieval test (internal use only) #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [GeneratedAppResponse](#generatedappresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Success | ### [POST] /trial-apps/{app_id}/completion-messages #### Parameters @@ -9451,9 +9452,9 @@ Bedrock retrieval test (internal use only) #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [GeneratedAppResponse](#generatedappresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Success | ### [GET] /trial-apps/{app_id}/datasets #### Parameters @@ -9568,9 +9569,9 @@ Returns the site configuration for the application including theme, icons, and t #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [GeneratedAppResponse](#generatedappresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Success | ### [POST] /trial-apps/{app_id}/workflows/tasks/{task_id}/stop **Stop workflow task** @@ -9983,7 +9984,7 @@ Create a new plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint created successfully | **application/json**: [EndpointCreateResponse](#endpointcreateresponse)
| +| 200 | Endpoint created successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### ~~[POST] /workspaces/current/endpoints/create~~ @@ -10002,7 +10003,7 @@ Deprecated legacy alias for creating a plugin endpoint. Use POST /workspaces/cur | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint created successfully | **application/json**: [EndpointCreateResponse](#endpointcreateresponse)
| +| 200 | Endpoint created successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### ~~[POST] /workspaces/current/endpoints/delete~~ @@ -10021,7 +10022,7 @@ Deprecated legacy alias for deleting a plugin endpoint. Use DELETE /workspaces/c | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint deleted successfully | **application/json**: [EndpointDeleteResponse](#endpointdeleteresponse)
| +| 200 | Endpoint deleted successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [POST] /workspaces/current/endpoints/disable @@ -10037,7 +10038,7 @@ Disable a plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint disabled successfully | **application/json**: [EndpointDisableResponse](#endpointdisableresponse)
| +| 200 | Endpoint disabled successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [POST] /workspaces/current/endpoints/enable @@ -10053,7 +10054,7 @@ Enable a plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint enabled successfully | **application/json**: [EndpointEnableResponse](#endpointenableresponse)
| +| 200 | Endpoint enabled successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [GET] /workspaces/current/endpoints/list @@ -10087,7 +10088,7 @@ List endpoints for a specific plugin | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [PluginEndpointListResponse](#pluginendpointlistresponse)
| +| 200 | Success | **application/json**: [EndpointListResponse](#endpointlistresponse)
| ### ~~[POST] /workspaces/current/endpoints/update~~ @@ -10105,7 +10106,7 @@ Deprecated legacy alias for updating a plugin endpoint. Use PATCH /workspaces/cu | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint updated successfully | **application/json**: [EndpointUpdateResponse](#endpointupdateresponse)
| +| 200 | Endpoint updated successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [DELETE] /workspaces/current/endpoints/{id} @@ -10121,7 +10122,7 @@ Delete a plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint deleted successfully | **application/json**: [EndpointDeleteResponse](#endpointdeleteresponse)
| +| 200 | Endpoint deleted successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [PATCH] /workspaces/current/endpoints/{id} @@ -10143,7 +10144,7 @@ Update a plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Endpoint updated successfully | **application/json**: [EndpointUpdateResponse](#endpointupdateresponse)
| +| 200 | Endpoint updated successfully | **application/json**: [SuccessResponse](#successresponse)
| | 403 | Admin privileges required | | ### [GET] /workspaces/current/members @@ -10203,7 +10204,7 @@ Update a plugin endpoint | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [MemberActionTenantResponse](#memberactiontenantresponse)
| +| 200 | Success | **application/json**: [MemberActionResponse](#memberactionresponse)
| ### [POST] /workspaces/current/members/{member_id}/owner-transfer #### Parameters @@ -11659,7 +11660,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Tool labels retrieved successfully | **application/json**: [ToolLabelListResponse](#toollabellistresponse)
| ### [POST] /workspaces/current/tool-provider/api/add #### Request Body @@ -11672,7 +11673,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API provider added successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/tool-provider/api/delete #### Request Body @@ -11685,7 +11686,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API provider deleted successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/api/get #### Parameters @@ -11698,7 +11699,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API provider retrieved successfully | **application/json**: [ApiProviderDetailResponse](#apiproviderdetailresponse)
| ### [GET] /workspaces/current/tool-provider/api/remote #### Parameters @@ -11711,7 +11712,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Remote API provider schema retrieved successfully | **application/json**: [ApiProviderRemoteSchemaResponse](#apiproviderremoteschemaresponse)
| ### [POST] /workspaces/current/tool-provider/api/schema #### Request Body @@ -11724,7 +11725,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API schema parsed successfully | **application/json**: [ApiSchemaParseResponse](#apischemaparseresponse)
| ### [POST] /workspaces/current/tool-provider/api/test/pre #### Request Body @@ -11737,7 +11738,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API tool test preview completed successfully | **application/json**: [ApiToolPreviewResponse](#apitoolpreviewresponse)
| ### [GET] /workspaces/current/tool-provider/api/tools #### Parameters @@ -11750,7 +11751,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API provider tools retrieved successfully | **application/json**: [ToolApiListResponse](#toolapilistresponse)
| ### [POST] /workspaces/current/tool-provider/api/update #### Request Body @@ -11763,7 +11764,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API provider updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/tool-provider/builtin/{provider}/add #### Parameters @@ -11782,7 +11783,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider added successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/credential/info #### Parameters @@ -11796,7 +11797,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider credential info retrieved successfully | **application/json**: [ToolProviderCredentialInfoApiEntity](#toolprovidercredentialinfoapientity)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/credential/schema/{credential_type} #### Parameters @@ -11810,7 +11811,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider credential schema retrieved successfully | **application/json**: [ProviderConfigListResponse](#providerconfiglistresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/credentials #### Parameters @@ -11824,7 +11825,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider credentials retrieved successfully | **application/json**: [ToolProviderCredentialListResponse](#toolprovidercredentiallistresponse)
| ### [POST] /workspaces/current/tool-provider/builtin/{provider}/default-credential #### Parameters @@ -11843,7 +11844,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Default credential set successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/tool-provider/builtin/{provider}/delete #### Parameters @@ -11862,7 +11863,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider credential deleted successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/icon #### Parameters @@ -11873,9 +11874,9 @@ Returns permission flags that control workspace features like member invitations #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [BinaryFileResponse](#binaryfileresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Builtin provider icon | ### [GET] /workspaces/current/tool-provider/builtin/{provider}/info #### Parameters @@ -11888,7 +11889,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider info retrieved successfully | **application/json**: [ToolProviderApiEntityResponse](#toolproviderapientityresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/oauth/client-schema #### Parameters @@ -11901,7 +11902,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolOAuthClientSchemaResponse](#tooloauthclientschemaresponse)
| +| 200 | Builtin provider OAuth client schema retrieved successfully | **application/json**: [BuiltinProviderOAuthClientSchemaResponse](#builtinprovideroauthclientschemaresponse)
| ### [DELETE] /workspaces/current/tool-provider/builtin/{provider}/oauth/custom-client #### Parameters @@ -11914,7 +11915,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Custom OAuth client deleted successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/oauth/custom-client #### Parameters @@ -11925,9 +11926,9 @@ Returns permission flags that control workspace features like member invitations #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolOAuthCustomClientResponse](#tooloauthcustomclientresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Custom OAuth client retrieved successfully | ### [POST] /workspaces/current/tool-provider/builtin/{provider}/oauth/custom-client #### Parameters @@ -11946,7 +11947,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Custom OAuth client saved successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/builtin/{provider}/tools #### Parameters @@ -11959,7 +11960,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider tools retrieved successfully | **application/json**: [ToolApiListResponse](#toolapilistresponse)
| ### [POST] /workspaces/current/tool-provider/builtin/{provider}/update #### Parameters @@ -11978,7 +11979,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin provider updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [DELETE] /workspaces/current/tool-provider/mcp #### Request Body @@ -12004,7 +12005,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | MCP provider created successfully | **application/json**: [ToolProviderApiEntityResponse](#toolproviderapientityresponse)
| ### [PUT] /workspaces/current/tool-provider/mcp #### Request Body @@ -12017,7 +12018,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | MCP provider updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/tool-provider/mcp/auth #### Request Body @@ -12030,7 +12031,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | MCP provider authorized successfully | **application/json**: [MCPAuthResponse](#mcpauthresponse)
| ### [GET] /workspaces/current/tool-provider/mcp/tools/{provider_id} #### Parameters @@ -12043,7 +12044,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | MCP provider retrieved successfully | **application/json**: [ToolProviderApiEntityResponse](#toolproviderapientityresponse)
| ### [GET] /workspaces/current/tool-provider/mcp/update/{provider_id} #### Parameters @@ -12056,7 +12057,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | MCP provider tools refreshed successfully | **application/json**: [ToolProviderApiEntityResponse](#toolproviderapientityresponse)
| ### [POST] /workspaces/current/tool-provider/workflow/create #### Request Body @@ -12069,7 +12070,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow tool created successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/tool-provider/workflow/delete #### Request Body @@ -12082,7 +12083,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow tool deleted successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-provider/workflow/get #### Parameters @@ -12096,7 +12097,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow tool retrieved successfully | **application/json**: [WorkflowToolDetailResponse](#workflowtooldetailresponse)
| ### [GET] /workspaces/current/tool-provider/workflow/tools #### Parameters @@ -12109,7 +12110,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow provider tools retrieved successfully | **application/json**: [ToolApiListResponse](#toolapilistresponse)
| ### [POST] /workspaces/current/tool-provider/workflow/update #### Request Body @@ -12122,7 +12123,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow tool updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/tool-providers #### Parameters @@ -12135,35 +12136,35 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Tool providers retrieved successfully | **application/json**: [ToolProviderListResponse](#toolproviderlistresponse)
| ### [GET] /workspaces/current/tools/api #### Responses | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | API tools retrieved successfully | **application/json**: [ToolProviderListResponse](#toolproviderlistresponse)
| ### [GET] /workspaces/current/tools/builtin #### Responses | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Builtin tools retrieved successfully | **application/json**: [ToolProviderListResponse](#toolproviderlistresponse)
| ### [GET] /workspaces/current/tools/mcp #### Responses | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | MCP tools retrieved successfully | **application/json**: [ToolProviderListResponse](#toolproviderlistresponse)
| ### [GET] /workspaces/current/tools/workflow #### Responses | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [ToolProviderOpaqueResponse](#toolprovideropaqueresponse)
| +| 200 | Workflow tools retrieved successfully | **application/json**: [ToolProviderListResponse](#toolproviderlistresponse)
| ### [GET] /workspaces/current/trigger-provider/{provider}/icon #### Parameters @@ -12174,9 +12175,9 @@ Returns permission flags that control workspace features like member invitations #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Success | **application/json**: [BinaryFileResponse](#binaryfileresponse)
| +| Code | Description | +| ---- | ----------- | +| 200 | Trigger provider icon | ### [GET] /workspaces/current/trigger-provider/{provider}/info **Get info for a trigger provider** @@ -12191,7 +12192,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerProviderApiEntity](#triggerproviderapientity)
| +| 200 | Trigger provider retrieved successfully | **application/json**: [TriggerProviderApiEntity](#triggerproviderapientity)
| ### [DELETE] /workspaces/current/trigger-provider/{provider}/oauth/client **Remove custom OAuth client configuration** @@ -12206,7 +12207,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Trigger OAuth client deleted successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/trigger-provider/{provider}/oauth/client **Get OAuth client configuration for a provider** @@ -12221,7 +12222,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerOAuthClientResponse](#triggeroauthclientresponse)
| +| 200 | Trigger OAuth client retrieved successfully | **application/json**: [TriggerOAuthClientResponse](#triggeroauthclientresponse)
| ### [POST] /workspaces/current/trigger-provider/{provider}/oauth/client **Configure custom OAuth client for a provider** @@ -12242,7 +12243,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| +| 200 | Trigger OAuth client saved successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/build/{subscription_builder_id} **Build a subscription instance for a trigger provider** @@ -12264,7 +12265,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerProviderOpaqueResponse](#triggerprovideropaqueresponse)
| +| 200 | Trigger subscription builder built successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/create **Add a new subscription instance for a trigger provider** @@ -12285,7 +12286,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerSubscriptionBuilderCreateResponse](#triggersubscriptionbuildercreateresponse)
| +| 200 | Trigger subscription builder created successfully | **application/json**: [TriggerSubscriptionBuilderCreateResponse](#triggersubscriptionbuildercreateresponse)
| ### [GET] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/logs/{subscription_builder_id} **Get the request logs for a subscription instance for a trigger provider** @@ -12301,7 +12302,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerSubscriptionBuilderLogsResponse](#triggersubscriptionbuilderlogsresponse)
| +| 200 | Trigger subscription builder logs retrieved successfully | **application/json**: [TriggerSubscriptionBuilderLogsResponse](#triggersubscriptionbuilderlogsresponse)
| ### [POST] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/update/{subscription_builder_id} **Update a subscription instance for a trigger provider** @@ -12323,7 +12324,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SubscriptionBuilderApiEntity](#subscriptionbuilderapientity)
| +| 200 | Trigger subscription builder updated successfully | **application/json**: [SubscriptionBuilderApiEntity](#subscriptionbuilderapientity)
| ### [POST] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/verify-and-update/{subscription_builder_id} **Verify and update a subscription instance for a trigger provider** @@ -12345,7 +12346,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerSubscriptionBuilderVerifyResponse](#triggersubscriptionbuilderverifyresponse)
| +| 200 | Trigger subscription builder verified successfully | **application/json**: [TriggerVerificationResponse](#triggerverificationresponse)
| ### [GET] /workspaces/current/trigger-provider/{provider}/subscriptions/builder/{subscription_builder_id} **Get a subscription instance for a trigger provider** @@ -12361,7 +12362,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [SubscriptionBuilderApiEntity](#subscriptionbuilderapientity)
| +| 200 | Trigger subscription builder retrieved successfully | **application/json**: [SubscriptionBuilderApiEntity](#subscriptionbuilderapientity)
| ### [GET] /workspaces/current/trigger-provider/{provider}/subscriptions/list **List all trigger subscriptions for the current tenant's provider** @@ -12376,7 +12377,8 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerSubscriptionListResponse](#triggersubscriptionlistresponse)
| +| 200 | Trigger subscriptions retrieved successfully | **application/json**: [TriggerProviderSubscriptionListResponse](#triggerprovidersubscriptionlistresponse)
| +| 404 | Trigger provider not found | **application/json**: [TriggerProviderErrorResponse](#triggerprovidererrorresponse)
| ### [GET] /workspaces/current/trigger-provider/{provider}/subscriptions/oauth/authorize **Initiate OAuth authorization flow for a trigger provider** @@ -12391,7 +12393,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Authorization URL retrieved successfully | **application/json**: [TriggerOAuthAuthorizeResponse](#triggeroauthauthorizeresponse)
| +| 200 | Trigger OAuth authorization URL generated successfully | **application/json**: [TriggerOAuthAuthorizeResponse](#triggeroauthauthorizeresponse)
| ### [POST] /workspaces/current/trigger-provider/{provider}/subscriptions/verify/{subscription_id} **Verify credentials for an existing subscription (edit mode only)** @@ -12413,7 +12415,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerSubscriptionBuilderVerifyResponse](#triggersubscriptionbuilderverifyresponse)
| +| 200 | Trigger subscription verified successfully | **application/json**: [TriggerVerificationResponse](#triggerverificationresponse)
| ### [POST] /workspaces/current/trigger-provider/{subscription_id}/subscriptions/delete **Delete a subscription instance** @@ -12449,7 +12451,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerProviderOpaqueResponse](#triggerprovideropaqueresponse)
| +| 200 | Trigger subscription updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /workspaces/current/triggers **List all trigger providers for the current tenant** @@ -12458,7 +12460,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [TriggerProviderListResponse](#triggerproviderlistresponse)
| +| 200 | Trigger providers retrieved successfully | **application/json**: [TriggerProviderListResponse](#triggerproviderlistresponse)
| ### [POST] /workspaces/custom-config #### Request Body @@ -12471,9 +12473,15 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [WorkspaceMutationResponse](#workspacemutationresponse)
| +| 200 | Success | **application/json**: [WorkspaceTenantResultResponse](#workspacetenantresultresponse)
| ### [POST] /workspaces/custom-config/webapp-logo/upload +#### Request Body + +| Required | Schema | +| -------- | ------ | +| Yes | **multipart/form-data**: { **"file"**: binary }
| + #### Responses | Code | Description | Schema | @@ -12491,7 +12499,7 @@ Returns permission flags that control workspace features like member invitations | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Success | **application/json**: [WorkspaceMutationResponse](#workspacemutationresponse)
| +| 200 | Success | **application/json**: [WorkspaceTenantResultResponse](#workspacetenantresultresponse)
| ### [POST] /workspaces/switch #### Request Body @@ -13747,6 +13755,23 @@ Stable Agent Soul reference to one normalized skill archive. | upload_file_id | string | | No | | url | string | | No | +#### AgentFileUploadFeatureConfig + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| allowed_file_extensions | [ string ] | | No | +| allowed_file_types | [ [FileType](#filetype) ] | | No | +| allowed_file_upload_methods | [ [FileTransferMethod](#filetransfermethod) ] | | No | +| enabled | boolean,
**Default:** true | | No | +| image | [AgentFileUploadImageFeatureConfig](#agentfileuploadimagefeatureconfig) | | No | +| number_limits | integer,
**Default:** 3 | | No | + +#### AgentFileUploadImageFeatureConfig + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| enabled | boolean,
**Default:** true | | No | + #### AgentHumanContactConfig | Name | Type | Description | Required | @@ -13890,6 +13915,8 @@ Per-set metadata filtering policy. The Python attribute uses ``metadata_model_config`` for clarity because the model belongs to metadata filtering specifically, while the external API and generated schema keep the historical ``model_config`` field name via alias. +Mode-dependent completeness is enforced by composer publish validation so +draft saves can persist partially configured metadata filters. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -13912,8 +13939,9 @@ Per-set query policy for Agent v2 knowledge retrieval. Agent v2 stores knowledge as explicit ``knowledge.sets`` rather than the legacy flat ``datasets`` / ``query_mode`` / ``query_config`` shape. Each -set owns its own query policy, so ``user_query`` must carry an explicit -``value`` while ``generated_query`` leaves that value empty. +set owns its own query policy. Mode-dependent completeness, such as +requiring ``value`` for ``user_query``, is enforced by composer publish +validation so draft saves can persist partially configured knowledge sets. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -13938,8 +13966,9 @@ set owns its own query policy, so ``user_query`` must carry an explicit Per-set retrieval policy for Agent v2 knowledge retrieval. Retrieval settings now live on each knowledge set instead of one shared -flat config. A set may use either ``multiple`` retrieval with ``top_k`` or -``single`` retrieval with a required model config. +flat config. Mode-dependent completeness, such as requiring ``top_k`` for +``multiple`` or a model for ``single``, is enforced by composer publish +validation so draft saves can persist partially configured knowledge sets. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -14338,6 +14367,7 @@ Visibility and lifecycle scope of an Agent record. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | +| file_upload | [AgentFileUploadFeatureConfig](#agentfileuploadfeatureconfig) | | No | | opening_statement | string | | No | | retriever_resource | [AgentFeatureToggleConfig](#agentfeaturetoggleconfig) | | No | | sensitive_word_avoidance | [AgentSensitiveWordAvoidanceFeatureConfig](#agentsensitivewordavoidancefeatureconfig) | | No | @@ -14596,6 +14626,7 @@ Legacy Chat App model config used only for follow-up question generation. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | +| answer | string | | No | | chain_id | string | | No | | created_at | integer | | No | | files | [ string ] | | Yes | @@ -14843,6 +14874,26 @@ Legacy Chat App model config used only for follow-up question generation. | ---- | ---- | ----------- | -------- | | data | [ [ApiKeyItem](#apikeyitem) ] | | Yes | +#### ApiProviderDetailResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| credentials | object | | No | +| custom_disclaimer | string | | No | +| description | string | | No | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | +| labels | [ string ] | | No | +| privacy_policy | string | | No | +| schema | string | | Yes | +| schema_type | [ApiProviderSchemaType](#apiproviderschematype) | | Yes | +| tools | [ [ApiToolBundle](#apitoolbundle) ] | | Yes | + +#### ApiProviderRemoteSchemaResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| schema | string | | Yes | + #### ApiProviderSchemaType Enum class for api provider schema type. @@ -14851,13 +14902,52 @@ Enum class for api provider schema type. | ---- | ---- | ----------- | -------- | | ApiProviderSchemaType | string | Enum class for api provider schema type. | | +#### ApiSchemaParseResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| credentials_schema | [ [ProviderConfig](#providerconfig) ] | | Yes | +| parameters_schema | [ [ApiToolBundle](#apitoolbundle) ] | | Yes | +| schema_type | [ApiProviderSchemaType](#apiproviderschematype) | | Yes | +| warning | object | | Yes | + +#### ApiToolBundle + +This class is used to store the schema information of an api based tool. + such as the url, the method, the parameters, etc. + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| author | string | | Yes | +| icon | string | | No | +| method | string | | Yes | +| openapi | object | | Yes | +| operation_id | string | | No | +| output_schema | object | | No | +| parameters | [ [ToolParameter](#toolparameter) ] | | No | +| server_url | string | | Yes | +| summary | string | | No | + +#### ApiToolPreviewResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ApiToolPreviewResponse | | | | + +#### ApiToolPreviewResult + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| error | string | | No | +| result | string | | No | + #### ApiToolProviderAddPayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | credentials | object | | Yes | | custom_disclaimer | string | | No | -| icon | object | | Yes | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | | labels | [ string ] | | No | | privacy_policy | string | | No | | provider | string | | Yes | @@ -14876,7 +14966,7 @@ Enum class for api provider schema type. | ---- | ---- | ----------- | -------- | | credentials | object | | Yes | | custom_disclaimer | string | | No | -| icon | object | | Yes | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | | labels | [ string ] | | No | | original_provider | string | | Yes | | privacy_policy | string | | No | @@ -15352,6 +15442,16 @@ Retrieval settings for Amazon Bedrock knowledge base queries. | ---- | ---- | ----------- | -------- | | id | string | | Yes | +#### BuiltinProviderOAuthClientSchemaResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| client_params | object | | No | +| is_oauth_custom_client_enabled | boolean | | Yes | +| is_system_oauth_params_exists | boolean | | Yes | +| redirect_uri | string | | Yes | +| schema | [ [ProviderConfig](#providerconfig) ] | | Yes | + #### BuiltinToolAddPayload | Name | Type | Description | Required | @@ -16103,12 +16203,6 @@ Model class for provider custom model configuration. | ---- | ---- | ----------- | -------- | | info_list | [InfoList](#infolist) | | Yes | -#### DataSourceContentPreviewResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| DataSourceContentPreviewResponse | | | | - #### DataSourceIntegrateIconResponse | Name | Type | Description | Required | @@ -16629,32 +16723,43 @@ Model class for provider custom model configuration. | ---- | ---- | ----------- | -------- | | credential_id | string | | Yes | +#### DatasourceCredentialListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| result | [ [DatasourceCredentialResponse](#datasourcecredentialresponse) ] | | Yes | + #### DatasourceCredentialPayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| credentials | object | | Yes | +| credentials | object | Plugin-defined credential parameters. The schema is declared by the datasource provider. | Yes | | name | string | | No | +#### DatasourceCredentialResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| avatar_url | string | | Yes | +| credential | object | Obfuscated plugin-defined credential parameters from the datasource provider. | Yes | +| id | string | | Yes | +| is_default | boolean | | Yes | +| name | string | | Yes | +| type | string | | Yes | + #### DatasourceCredentialUpdatePayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | credential_id | string | | Yes | -| credentials | object | | No | +| credentials | object | Plugin-defined credential parameters. The schema is declared by the datasource provider. | No | | name | string | | No | -#### DatasourceCredentialsResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| result | | | Yes | - #### DatasourceCustomClientPayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| client_params | object | | No | +| client_params | object | Plugin-defined OAuth client parameters. The schema is declared by the datasource provider. | No | | enable_oauth_custom_client | boolean | | No | #### DatasourceDefaultPayload @@ -16686,6 +16791,39 @@ Model class for provider custom model configuration. | error | string | Error message from OAuth provider | No | | state | string | OAuth state parameter | No | +#### DatasourceOAuthSchemaResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| client_schema | [ [ProviderConfig](#providerconfig) ] | | Yes | +| credentials_schema | [ [ProviderConfig](#providerconfig) ] | | Yes | +| is_oauth_custom_client_enabled | boolean | | Yes | +| is_system_oauth_params_exists | boolean | | Yes | +| oauth_custom_client_params | object | Masked plugin-defined OAuth client parameters, when configured for the tenant. | Yes | +| redirect_uri | string | | Yes | + +#### DatasourceProviderAuthListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| result | [ [DatasourceProviderAuthResponse](#datasourceproviderauthresponse) ] | | Yes | + +#### DatasourceProviderAuthResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| author | string | | Yes | +| credential_schema | [ [ProviderConfig](#providerconfig) ] | | Yes | +| credentials_list | [ [DatasourceCredentialResponse](#datasourcecredentialresponse) ] | | Yes | +| description | [I18nObject](#i18nobject) | | Yes | +| icon | string | | Yes | +| label | [I18nObject](#i18nobject) | | Yes | +| name | string | | Yes | +| oauth_schema | [DatasourceOAuthSchemaResponse](#datasourceoauthschemaresponse) | | Yes | +| plugin_id | string | | Yes | +| plugin_unique_identifier | string | | Yes | +| provider | string | | Yes | + #### DatasourceUpdateNamePayload | Name | Type | Description | Required | @@ -16702,12 +16840,6 @@ Model class for provider custom model configuration. | start_node_id | string | | Yes | | start_node_title | string | | Yes | -#### DebugPermission - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| DebugPermission | string | | | - #### DeclaredArrayItem Per-item shape for an ``array``-typed declared output. @@ -17034,12 +17166,6 @@ Request payload for bulk downloading documents as a zip archive. | role | string | | Yes | | token | string | | Yes | -#### EducationActivateResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| EducationActivateResponse | object | | | - #### EducationAutocompleteQuery | Name | Type | Description | Required | @@ -17149,29 +17275,13 @@ Request payload for bulk downloading documents as a zip archive. | plugin_unique_identifier | string | | Yes | | settings | object | | Yes | -#### EndpointCreateResponse +#### EndpointDeclarationResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| success | boolean | Operation success | Yes | - -#### EndpointDeleteResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| success | boolean | Operation success | Yes | - -#### EndpointDisableResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| success | boolean | Operation success | Yes | - -#### EndpointEnableResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| success | boolean | Operation success | Yes | +| hidden | boolean | | No | +| method | string | | Yes | +| path | string | | Yes | #### EndpointIdPayload @@ -17187,6 +17297,23 @@ Request payload for bulk downloading documents as a zip archive. | page_size | integer | | Yes | | plugin_id | string | | Yes | +#### EndpointListItemResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| created_at | dateTime | | Yes | +| declaration | [EndpointProviderDeclarationResponse](#endpointproviderdeclarationresponse) | | No | +| enabled | boolean | | Yes | +| expired_at | dateTime | | Yes | +| hook_id | string | | Yes | +| id | string | | Yes | +| name | string | | Yes | +| plugin_id | string | | Yes | +| settings | object | | Yes | +| tenant_id | string | | Yes | +| updated_at | dateTime | | Yes | +| url | string | | Yes | + #### EndpointListQuery | Name | Type | Description | Required | @@ -17198,7 +17325,59 @@ Request payload for bulk downloading documents as a zip archive. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| endpoints | [ object ] | Endpoint information | Yes | +| endpoints | [ [EndpointListItemResponse](#endpointlistitemresponse) ] | Endpoint information | Yes | + +#### EndpointProviderConfigI18nResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| en_US | string | | Yes | +| ja_JP | string | | No | +| pt_BR | string | | No | +| zh_Hans | string | | No | + +#### EndpointProviderConfigOptionResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| label | [EndpointProviderConfigI18nResponse](#endpointproviderconfigi18nresponse) | | Yes | +| value | string | | Yes | + +#### EndpointProviderConfigResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| default | integer
string
number
boolean | | No | +| help | [EndpointProviderConfigI18nResponse](#endpointproviderconfigi18nresponse) | | No | +| label | [EndpointProviderConfigI18nResponse](#endpointproviderconfigi18nresponse) | | No | +| multiple | boolean | | No | +| name | string | | Yes | +| options | [ [EndpointProviderConfigOptionResponse](#endpointproviderconfigoptionresponse) ] | | No | +| placeholder | [EndpointProviderConfigI18nResponse](#endpointproviderconfigi18nresponse) | | No | +| required | boolean | | No | +| scope | [EndpointProviderConfigScope](#endpointproviderconfigscope) | | No | +| type | [ProviderConfigType](#providerconfigtype) | | Yes | +| url | string | | No | + +#### EndpointProviderConfigScope + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| EndpointProviderConfigScope | string | | | + +#### EndpointProviderDeclarationResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| endpoints | [ [EndpointDeclarationResponse](#endpointdeclarationresponse) ] | | No | +| settings | [ [EndpointProviderConfigResponse](#endpointproviderconfigresponse) ] | | No | + +#### EndpointSettingsPayload + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| name | string | | Yes | +| settings | object | | Yes | #### EndpointUpdatePayload @@ -17207,12 +17386,6 @@ Request payload for bulk downloading documents as a zip archive. | name | string | | Yes | | settings | object | | Yes | -#### EndpointUpdateResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| success | boolean | Operation success | Yes | - #### EnvSuggestion | Name | Type | Description | Required | @@ -18010,12 +18183,6 @@ Input field definition for snippet parameters. | required | boolean | | No | | type | string | | No | -#### InstallPermission - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| InstallPermission | string | | | - #### InstalledAppCreatePayload | Name | Type | Description | Required | @@ -18254,6 +18421,20 @@ Enum class for large language model mode. | authorization_code | string | | No | | provider_id | string | | Yes | +#### MCPAuthResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| authorization_url | string | | No | +| result | string | | No | + +#### MCPAuthentication + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| client_id | string | | Yes | +| client_secret | string | | No | + #### MCPCallbackQuery | Name | Type | Description | Required | @@ -18261,12 +18442,19 @@ Enum class for large language model mode. | code | string | | Yes | | state | string | | Yes | +#### MCPConfiguration + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| sse_read_timeout | number,
**Default:** 300 | | No | +| timeout | number,
**Default:** 30 | | No | + #### MCPProviderCreatePayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| authentication | object | | No | -| configuration | object | | No | +| authentication | [MCPAuthentication](#mcpauthentication) | | No | +| configuration | [MCPConfiguration](#mcpconfiguration) | | No | | headers | object | | No | | icon | string | | Yes | | icon_background | string | | No | @@ -18286,8 +18474,8 @@ Enum class for large language model mode. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| authentication | object | | No | -| configuration | object | | No | +| authentication | [MCPAuthentication](#mcpauthentication) | | No | +| configuration | [MCPConfiguration](#mcpconfiguration) | | No | | headers | object | | No | | icon | string | | Yes | | icon_background | string | | No | @@ -18321,7 +18509,7 @@ Enum class for large language model mode. | marketplace_plugin_unique_identifier | string | | Yes | | version | string | | No | -#### MemberActionTenantResponse +#### MemberActionResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -19092,13 +19280,13 @@ Enum class for parameter type. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | auto_upgrade | [PluginAutoUpgradeSettingsPayload](#pluginautoupgradesettingspayload) | | Yes | -| category | [PluginCategory](#plugincategory) | | Yes | +| category | [TenantPluginAutoUpgradeCategory](#tenantpluginautoupgradecategory) | | Yes | #### ParserAutoUpgradeFetch | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| category | [PluginCategory](#plugincategory) | | Yes | +| category | [TenantPluginAutoUpgradeCategory](#tenantpluginautoupgradecategory) | | Yes | #### ParserCreateCredential @@ -19196,7 +19384,7 @@ Enum class for parameter type. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| category | [PluginCategory](#plugincategory) | | Yes | +| category | [TenantPluginAutoUpgradeCategory](#tenantpluginautoupgradecategory) | | Yes | | plugin_id | string | | Yes | #### ParserGetCredentials @@ -19284,8 +19472,8 @@ Enum class for parameter type. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| debug_permission | [DebugPermission](#debugpermission) | | No | -| install_permission | [InstallPermission](#installpermission) | | No | +| debug_permission | [TenantPluginDebugPermission](#tenantplugindebugpermission) | | No | +| install_permission | [TenantPluginInstallPermission](#tenantplugininstallpermission) | | No | #### ParserPluginIdentifierQuery @@ -19494,7 +19682,7 @@ Shared permission levels for resources (datasets, credentials, etc.) | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | auto_upgrade | [PluginAutoUpgradeSettingsResponseModel](#pluginautoupgradesettingsresponsemodel) | | Yes | -| category | [PluginCategory](#plugincategory) | | Yes | +| category | [TenantPluginAutoUpgradeCategory](#tenantpluginautoupgradecategory) | | Yes | #### PluginAutoUpgradeSettingsPayload @@ -19502,8 +19690,8 @@ Shared permission levels for resources (datasets, credentials, etc.) | ---- | ---- | ----------- | -------- | | exclude_plugins | [ string ] | | No | | include_plugins | [ string ] | | No | -| strategy_setting | [StrategySetting](#strategysetting) | | No | -| upgrade_mode | [UpgradeMode](#upgrademode) | | No | +| strategy_setting | [TenantPluginAutoUpgradeStrategySetting](#tenantpluginautoupgradestrategysetting) | | No | +| upgrade_mode | [TenantPluginAutoUpgradeMode](#tenantpluginautoupgrademode) | | No | | upgrade_time_of_day | integer | | No | #### PluginAutoUpgradeSettingsResponseModel @@ -19512,8 +19700,8 @@ Shared permission levels for resources (datasets, credentials, etc.) | ---- | ---- | ----------- | -------- | | exclude_plugins | [ string ] | | Yes | | include_plugins | [ string ] | | Yes | -| strategy_setting | [StrategySetting](#strategysetting) | | Yes | -| upgrade_mode | [UpgradeMode](#upgrademode) | | Yes | +| strategy_setting | [TenantPluginAutoUpgradeStrategySetting](#tenantpluginautoupgradestrategysetting) | | Yes | +| upgrade_mode | [TenantPluginAutoUpgradeMode](#tenantpluginautoupgrademode) | | Yes | | upgrade_time_of_day | integer | | Yes | #### PluginCategory @@ -19635,21 +19823,21 @@ Shared permission levels for resources (datasets, credentials, etc.) | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | current_identifier | string | | No | -| type | [Type](#type) | | Yes | +| type | [PluginDependencyType](#plugindependencytype) | | Yes | | value | [Github](#github)
[Marketplace](#marketplace)
[Package](#package) | | Yes | +#### PluginDependencyType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| PluginDependencyType | string | | | + #### PluginDynamicOptionsResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | options | | | Yes | -#### PluginEndpointListResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| endpoints | [ object ] | Endpoint information | Yes | - #### PluginInstallationItemResponse | Name | Type | Description | Required | @@ -19730,7 +19918,13 @@ Shared permission levels for resources (datasets, credentials, etc.) | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| type | [core__plugin__entities__parameters__PluginParameterAutoGenerate__Type](#core__plugin__entities__parameters__pluginparameterautogenerate__type) | | Yes | +| type | [PluginParameterAutoGenerateType](#pluginparameterautogeneratetype) | | Yes | + +#### PluginParameterAutoGenerateType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| PluginParameterAutoGenerateType | string | | | #### PluginParameterOption @@ -19750,15 +19944,15 @@ Shared permission levels for resources (datasets, credentials, etc.) | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| debug_permission | [DebugPermission](#debugpermission) | | Yes | -| install_permission | [InstallPermission](#installpermission) | | Yes | +| debug_permission | [TenantPluginDebugPermission](#tenantplugindebugpermission) | | Yes | +| install_permission | [TenantPluginInstallPermission](#tenantplugininstallpermission) | | Yes | #### PluginPermissionSettingsPayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| debug_permission | [DebugPermission](#debugpermission) | | No | -| install_permission | [InstallPermission](#installpermission) | | No | +| debug_permission | [TenantPluginDebugPermission](#tenantplugindebugpermission) | | No | +| install_permission | [TenantPluginInstallPermission](#tenantplugininstallpermission) | | No | #### PluginReadmeResponse @@ -19840,9 +20034,21 @@ Model class for common provider settings like credentials | placeholder | [I18nObject](#i18nobject) | | No | | required | boolean | | No | | scope | [AppSelectorScope](#appselectorscope)
[ModelSelectorScope](#modelselectorscope)
[ToolSelectorScope](#toolselectorscope) | | No | -| type | [core__entities__provider_entities__BasicProviderConfig__Type](#core__entities__provider_entities__basicproviderconfig__type) | The type of the credentials | Yes | +| type | [ProviderConfigType](#providerconfigtype) | The type of the credentials | Yes | | url | string | | No | +#### ProviderConfigListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ProviderConfigListResponse | array | | | + +#### ProviderConfigType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ProviderConfigType | string | | | + #### ProviderCredentialResponse | Name | Type | Description | Required | @@ -20437,19 +20643,11 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs. | text | string | | No | | truncated | boolean | | Yes | -#### SandboxToolFileResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| reference | string | | Yes | -| transfer_method | string,
**Default:** tool_file | | No | - #### SandboxUploadResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| file | [SandboxToolFileResponse](#sandboxtoolfileresponse) | | Yes | -| path | string | | Yes | +| url | string | | Yes | #### SavedMessageCreatePayload @@ -20484,11 +20682,19 @@ Whitelist scopes accepted by RBAC app and dataset access config APIs. | last_id | string | | No | | limit | integer,
**Default:** 20 | | No | +#### SchemaDefinitionItemResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| label | string | | Yes | +| name | string | | Yes | +| schema | object | | Yes | + #### SchemaDefinitionsResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| SchemaDefinitionsResponse | | | | +| SchemaDefinitionsResponse | array | | | #### SegmentAttachmentResponse @@ -20998,12 +21204,6 @@ Query parameters for listing snippet published workflows. | paused | integer | | Yes | | success | integer | | Yes | -#### StrategySetting - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| StrategySetting | string | | | - #### StringListSource | Name | Type | Description | Required | @@ -21258,6 +21458,36 @@ Tag type | ---- | ---- | ----------- | -------- | | workspaces | [ [TenantListItemResponse](#tenantlistitemresponse) ] | | Yes | +#### TenantPluginAutoUpgradeCategory + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TenantPluginAutoUpgradeCategory | string | | | + +#### TenantPluginAutoUpgradeMode + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TenantPluginAutoUpgradeMode | string | | | + +#### TenantPluginAutoUpgradeStrategySetting + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TenantPluginAutoUpgradeStrategySetting | string | | | + +#### TenantPluginDebugPermission + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TenantPluginDebugPermission | string | | | + +#### TenantPluginInstallPermission + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TenantPluginInstallPermission | string | | | + #### TextContentResponse | Name | Type | Description | Required | @@ -21331,11 +21561,46 @@ Available voices | ---- | ---- | ----------- | -------- | | data | [ [TokensPerSecondStatisticItem](#tokenspersecondstatisticitem) ] | | Yes | -#### ToolOAuthClientSchemaResponse +#### ToolApiEntity | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| ToolOAuthClientSchemaResponse | array | | | +| author | string | | Yes | +| description | [I18nObject](#i18nobject) | | Yes | +| label | [I18nObject](#i18nobject) | | Yes | +| labels | [ string ] | | No | +| name | string | | Yes | +| output_schema | object | | No | +| parameters | [ [ToolParameter](#toolparameter) ] | | No | + +#### ToolApiListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ToolApiListResponse | array | | | + +#### ToolEmojiIcon + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| background | string | | Yes | +| content | string | | Yes | + +#### ToolLabel + +Tool label + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| icon | string | The icon of the tool | Yes | +| label | [I18nObject](#i18nobject) | The label of the tool | Yes | +| name | string | The name of the tool | Yes | + +#### ToolLabelListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ToolLabelListResponse | array | | | #### ToolOAuthCustomClientPayload @@ -21344,11 +21609,29 @@ Available voices | client_params | object | | No | | enable_oauth_custom_client | boolean | | No | -#### ToolOAuthCustomClientResponse +#### ToolParameter + +Overrides type | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| ToolOAuthCustomClientResponse | object | | | +| auto_generate | [PluginParameterAutoGenerate](#pluginparameterautogenerate) | | No | +| default | number
integer
string
boolean
[ object ]
object | | No | +| form | [ToolParameterForm](#toolparameterform) | The form of the parameter, schema/form/llm | Yes | +| human_description | [I18nObject](#i18nobject) | The description presented to the user | No | +| input_schema | object | | No | +| label | [I18nObject](#i18nobject) | The label presented to the user | Yes | +| llm_description | string | | No | +| max | number
integer | | No | +| min | number
integer | | No | +| name | string | The name of the parameter | Yes | +| options | [ [PluginParameterOption](#pluginparameteroption) ] | | No | +| placeholder | [I18nObject](#i18nobject) | The placeholder presented to the user | No | +| precision | integer | | No | +| required | boolean | | No | +| scope | string | | No | +| template | [PluginParameterTemplate](#pluginparametertemplate) | | No | +| type | [ToolParameterType](#toolparametertype) | The type of the parameter | Yes | #### ToolParameterForm @@ -21356,17 +21639,84 @@ Available voices | ---- | ---- | ----------- | -------- | | ToolParameterForm | string | | | +#### ToolParameterType + +removes TOOLS_SELECTOR from PluginParameterType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ToolParameterType | string | removes TOOLS_SELECTOR from PluginParameterType | | + +#### ToolProviderApiEntityResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| allow_delete | boolean,
**Default:** true | | No | +| authentication | [MCPAuthentication](#mcpauthentication) | The OAuth config of the MCP tool | No | +| author | string | | Yes | +| configuration | [MCPConfiguration](#mcpconfiguration) | The timeout and sse_read_timeout of the MCP tool | No | +| description | [I18nObject](#i18nobject) | | Yes | +| icon | string
object | | Yes | +| icon_dark | string
object | | No | +| id | string | | Yes | +| identity_mode | string,
**Default:** off | Identity-forwarding mechanism: 'off' or 'idp_token' | No | +| is_dynamic_registration | boolean,
**Default:** true | Whether the MCP tool is dynamically registered | No | +| is_team_authorization | boolean | | No | +| label | [I18nObject](#i18nobject) | | Yes | +| labels | [ string ] | | No | +| masked_headers | object | The masked headers of the MCP tool | No | +| name | string | | Yes | +| original_headers | object | The original headers of the MCP tool | No | +| plugin_id | string | The plugin id of the tool | No | +| plugin_unique_identifier | string | The unique identifier of the tool | No | +| server_identifier | string | The server identifier of the MCP tool | No | +| server_url | string | The server url of the tool | No | +| team_credentials | object | | No | +| tools | [ [ToolApiEntity](#toolapientity) ] | | No | +| type | [ToolProviderType](#toolprovidertype) | | Yes | +| updated_at | integer | | No | +| workflow_app_id | string | The app id of the workflow tool | No | + +#### ToolProviderCredentialApiEntity + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| created_by | string | User ID of the credential creator | No | +| credential_type | [CredentialType](#credentialtype) | The type of the credential | Yes | +| credentials | object | The credentials of the provider | No | +| from_other_member | boolean | True when this credential is being returned only because a workflow/agent node still references it but it would normally be hidden from this user by the visibility filter (another member's only_me credential). The frontend renders it as 'borrowed' — selectable until the node switches away, but not editable/deletable. | No | +| id | string | The unique id of the credential | Yes | +| is_default | boolean | Whether the credential is the default credential for the provider in the workspace | No | +| name | string | The name of the credential | Yes | +| partial_member_list | [ string ] | List of user IDs allowed when visibility is partial_members | No | +| provider | string | The provider of the credential | Yes | +| visibility | string,
**Default:** all_team_members | Credential visibility: only_me, all_team_members, or partial_members | No | + +#### ToolProviderCredentialInfoApiEntity + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| credentials | [ [ToolProviderCredentialApiEntity](#toolprovidercredentialapientity) ] | The credentials of the provider | Yes | +| is_oauth_custom_client_enabled | boolean | Whether the OAuth custom client is enabled for the provider | No | +| supported_credential_types | [ [CredentialType](#credentialtype) ] | The supported credential types of the provider | Yes | + +#### ToolProviderCredentialListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| ToolProviderCredentialListResponse | array | | | + #### ToolProviderListQuery | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | type | string | | No | -#### ToolProviderOpaqueResponse +#### ToolProviderListResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| ToolProviderOpaqueResponse | | | | +| ToolProviderListResponse | array | | | #### ToolProviderType @@ -21647,7 +21997,7 @@ Enum class for tool provider | configured | boolean | | Yes | | custom_configured | boolean | | Yes | | custom_enabled | boolean | | Yes | -| oauth_client_schema | [ [TriggerProviderConfigResponse](#triggerproviderconfigresponse) ] | | Yes | +| oauth_client_schema | [ [ProviderConfig](#providerconfig) ] | | Yes | | params | object | | Yes | | redirect_uri | string | | Yes | | system_configured | boolean | | Yes | @@ -21670,28 +22020,11 @@ Enum class for tool provider | supported_creation_methods | [ [TriggerCreationMethod](#triggercreationmethod) ] | Supported creation methods for the trigger provider. like 'OAUTH', 'APIKEY', 'MANUAL'. | No | | tags | [ string ] | The tags of the trigger provider | No | -#### TriggerProviderConfigOptionResponse +#### TriggerProviderErrorResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| label | [I18nObject](#i18nobject) | The label of the option | Yes | -| value | string | The value of the option | Yes | - -#### TriggerProviderConfigResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| default | integer
string
number
boolean | | No | -| help | [I18nObject](#i18nobject) | | No | -| label | [I18nObject](#i18nobject) | | No | -| multiple | boolean | | No | -| name | string | The name of the credentials | Yes | -| options | [ [TriggerProviderConfigOptionResponse](#triggerproviderconfigoptionresponse) ] | | No | -| placeholder | [I18nObject](#i18nobject) | | No | -| required | boolean | | No | -| scope | [AppSelectorScope](#appselectorscope)
[ModelSelectorScope](#modelselectorscope)
[ToolSelectorScope](#toolselectorscope) | | No | -| type | string,
**Available values:** "app-selector", "array[tools]", "boolean", "model-selector", "secret-input", "select", "text-input" | The type of the credentials
*Enum:* `"app-selector"`, `"array[tools]"`, `"boolean"`, `"model-selector"`, `"secret-input"`, `"select"`, `"text-input"` | Yes | -| url | string | | No | +| error | string | | Yes | #### TriggerProviderListResponse @@ -21699,12 +22032,6 @@ Enum class for tool provider | ---- | ---- | ----------- | -------- | | TriggerProviderListResponse | array | | | -#### TriggerProviderOpaqueResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| TriggerProviderOpaqueResponse | | | | - #### TriggerProviderSubscriptionApiEntity | Name | Type | Description | Required | @@ -21719,6 +22046,12 @@ Enum class for tool provider | provider | string | The provider id of the subscription | Yes | | workflows_in_use | integer | The number of workflows using this subscription | Yes | +#### TriggerProviderSubscriptionListResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| TriggerProviderSubscriptionListResponse | array | | | + #### TriggerSubscriptionBuilderCreatePayload | Name | Type | Description | Required | @@ -21752,24 +22085,12 @@ Enum class for tool provider | ---- | ---- | ----------- | -------- | | credentials | object | | Yes | -#### TriggerSubscriptionBuilderVerifyResponse +#### TriggerVerificationResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | verified | boolean | | Yes | -#### TriggerSubscriptionListResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| TriggerSubscriptionListResponse | array | | | - -#### Type - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| Type | string | | | - #### UnaddedModelConfiguration Model class for provider unadded model configuration. @@ -21810,12 +22131,6 @@ Payload for updating a snippet. | icon_info | [IconInfo](#iconinfo) | | No | | name | string | | No | -#### UpgradeMode - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| UpgradeMode | string | | | - #### UploadConfig | Name | Type | Description | Required | @@ -22826,7 +23141,7 @@ Query parameters for workflow runs. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | description | string | | Yes | -| icon | object | | Yes | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | | label | string | | Yes | | labels | [ string ] | | No | | name | string | | Yes | @@ -22840,6 +23155,22 @@ Query parameters for workflow runs. | ---- | ---- | ----------- | -------- | | workflow_tool_id | string | | Yes | +#### WorkflowToolDetailResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| description | string | | Yes | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | +| label | string | | Yes | +| name | string | | Yes | +| output_schema | object | | No | +| parameters | [ [WorkflowToolParameterConfiguration](#workflowtoolparameterconfiguration) ] | | Yes | +| privacy_policy | string | | No | +| synced | boolean | | Yes | +| tool | [ToolApiEntity](#toolapientity) | | Yes | +| workflow_app_id | string | | Yes | +| workflow_tool_id | string | | Yes | + #### WorkflowToolGetQuery | Name | Type | Description | Required | @@ -22868,7 +23199,7 @@ Workflow tool configuration | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | description | string | | Yes | -| icon | object | | Yes | +| icon | [ToolEmojiIcon](#toolemojiicon) | | Yes | | label | string | | Yes | | labels | [ string ] | | No | | name | string | | Yes | @@ -22946,7 +23277,13 @@ Workflow tool configuration | limit | integer,
**Default:** 20 | | No | | page | integer,
**Default:** 1 | | No | -#### WorkspaceListResponse +#### WorkspaceLogoUploadResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| id | string | | Yes | + +#### WorkspacePaginationResponse | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -22956,19 +23293,6 @@ Workflow tool configuration | page | integer | | Yes | | total | integer | | Yes | -#### WorkspaceLogoUploadResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| id | string | | Yes | - -#### WorkspaceMutationResponse - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| result | string | | Yes | -| tenant | [TenantInfoResponse](#tenantinforesponse) | | Yes | - #### WorkspacePermissionResponse | Name | Type | Description | Required | @@ -22983,6 +23307,13 @@ Workflow tool configuration | ---- | ---- | ----------- | -------- | | permission_keys | [ string ] | | No | +#### WorkspaceTenantResultResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| result | string | | Yes | +| tenant | [TenantInfoResponse](#tenantinforesponse) | | Yes | + #### _AccessControlLanguageQuery | Name | Type | Description | Required | @@ -23051,18 +23382,6 @@ Workflow tool configuration | ---- | ---- | ----------- | -------- | | scope | [RBACResourceWhitelistScope](#rbacresourcewhitelistscope) | | Yes | -#### core__entities__provider_entities__BasicProviderConfig__Type - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| core__entities__provider_entities__BasicProviderConfig__Type | string | | | - -#### core__plugin__entities__parameters__PluginParameterAutoGenerate__Type - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| core__plugin__entities__parameters__PluginParameterAutoGenerate__Type | string | | | - #### core__tools__entities__common_entities__I18nObject Model class for i18n object. diff --git a/api/openapi/markdown/openapi-openapi.md b/api/openapi/markdown/openapi-openapi.md index 08544f20a9d..c5ec6c3a7db 100644 --- a/api/openapi/markdown/openapi-openapi.md +++ b/api/openapi/markdown/openapi-openapi.md @@ -93,21 +93,7 @@ User-scoped operations | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [GET] /apps/{app_id}/check-dependencies -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| app_id | path | | Yes | string | - -#### Responses - -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Dependencies checked | **application/json**: [CheckDependenciesResult](#checkdependenciesresult)
| -| default | Error | **application/json**: [ErrorBody](#errorbody)
| - -### [GET] /apps/{app_id}/describe +### [GET] /apps/{app_id} #### Parameters | Name | Located in | Description | Required | Schema | @@ -123,7 +109,21 @@ User-scoped operations | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [GET] /apps/{app_id}/export +### [GET] /apps/{app_id}/dependencies:check +#### Parameters + +| Name | Located in | Description | Required | Schema | +| ---- | ---------- | ----------- | -------- | ------ | +| app_id | path | | Yes | string | + +#### Responses + +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Dependencies checked | **application/json**: [CheckDependenciesResult](#checkdependenciesresult)
| +| default | Error | **application/json**: [ErrorBody](#errorbody)
| + +### [GET] /apps/{app_id}/dsl #### Parameters | Name | Located in | Description | Required | Schema | @@ -140,7 +140,7 @@ User-scoped operations | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [POST] /apps/{app_id}/files/upload +### [POST] /apps/{app_id}/files Upload a file to use as an input variable when running the app #### Parameters @@ -160,7 +160,7 @@ Upload a file to use as an input variable when running the app | 415 | Unsupported file type or blocked extension | | | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [GET] /apps/{app_id}/form/human_input/{form_token} +### [GET] /apps/{app_id}/human-input-forms/{form_token} #### Parameters | Name | Located in | Description | Required | Schema | @@ -174,7 +174,7 @@ Upload a file to use as an input variable when running the app | ---- | ----------- | ------ | | 200 | Form definition | **application/json**: [HumanInputFormDefinitionResponse](#humaninputformdefinitionresponse)
| -### [POST] /apps/{app_id}/form/human_input/{form_token} +### [POST] /apps/{app_id}/human-input-forms/{form_token}:submit #### Parameters | Name | Located in | Description | Required | Schema | @@ -196,7 +196,38 @@ Upload a file to use as an input variable when running the app | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [POST] /apps/{app_id}/run +### [GET] /apps/{app_id}/tasks/{task_id}/events +#### Parameters + +| Name | Located in | Description | Required | Schema | +| ---- | ---------- | ----------- | -------- | ------ | +| continue_on_pause | query | Whether to keep the event stream open on pause | No | boolean | +| include_state_snapshot | query | Whether to include workflow state snapshots | No | boolean | +| app_id | path | | Yes | string | +| task_id | path | | Yes | string | + +#### Responses + +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | SSE event stream | **application/json**: [EventStreamResponse](#eventstreamresponse)
| + +### [POST] /apps/{app_id}/tasks/{task_id}:stop +#### Parameters + +| Name | Located in | Description | Required | Schema | +| ---- | ---------- | ----------- | -------- | ------ | +| app_id | path | | Yes | string | +| task_id | path | | Yes | string | + +#### Responses + +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Task stopped | **application/json**: [TaskStopResponse](#taskstopresponse)
| +| default | Error | **application/json**: [ErrorBody](#errorbody)
| + +### [POST] /apps/{app_id}:run #### Parameters | Name | Located in | Description | Required | Schema | @@ -216,37 +247,6 @@ Upload a file to use as an input variable when running the app | 200 | Run result (SSE stream) | **application/json**: [EventStreamResponse](#eventstreamresponse)
| | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| -### [GET] /apps/{app_id}/tasks/{task_id}/events -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| continue_on_pause | query | Whether to keep the event stream open on pause | No | boolean | -| include_state_snapshot | query | Whether to include workflow state snapshots | No | boolean | -| app_id | path | | Yes | string | -| task_id | path | | Yes | string | - -#### Responses - -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | SSE event stream | **application/json**: [EventStreamResponse](#eventstreamresponse)
| - -### [POST] /apps/{app_id}/tasks/{task_id}/stop -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| app_id | path | | Yes | string | -| task_id | path | | Yes | string | - -#### Responses - -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Task stopped | **application/json**: [TaskStopResponse](#taskstopresponse)
| -| default | Error | **application/json**: [ErrorBody](#errorbody)
| - ### [POST] /oauth/device/approve #### Request Body @@ -330,7 +330,7 @@ Upload a file to use as an input variable when running the app | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [GET] /permitted-external-apps/{app_id}/describe +### [GET] /permitted-external-apps/{app_id} #### Parameters | Name | Located in | Description | Required | Schema | @@ -391,7 +391,7 @@ Upload a file to use as an input variable when running the app | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [POST] /workspaces/{workspace_id}/apps/imports/{import_id}/confirm +### [POST] /workspaces/{workspace_id}/apps/imports/{import_id}:confirm #### Parameters | Name | Located in | Description | Required | Schema | @@ -460,7 +460,7 @@ Upload a file to use as an input variable when running the app | 200 | Member removed | **application/json**: [MemberActionResponse](#memberactionresponse)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [PUT] /workspaces/{workspace_id}/members/{member_id}/role +### [PATCH] /workspaces/{workspace_id}/members/{member_id} #### Parameters | Name | Located in | Description | Required | Schema | @@ -482,7 +482,7 @@ Upload a file to use as an input variable when running the app | 422 | Validation error | **application/json**: [ErrorBody](#errorbody)
| | default | Error | **application/json**: [ErrorBody](#errorbody)
| -### [POST] /workspaces/{workspace_id}/switch +### [POST] /workspaces/{workspace_id}:switch #### Parameters | Name | Located in | Description | Required | Schema | @@ -532,7 +532,7 @@ Upload a file to use as an input variable when running the app #### AppDescribeQuery -`?fields=` allow-list for GET /apps//describe. +`?fields=` allow-list for GET /apps/. Empty / omitted → all blocks. Unknown member → ValidationError → 422. @@ -550,7 +550,7 @@ Empty / omitted → all blocks. Unknown member → ValidationError → 422. #### AppDslExportQuery -Query parameters for GET /apps//export. +Query parameters for GET /apps//dsl. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | @@ -762,7 +762,7 @@ future server adds a code. Formatter tests pin emitted values to the enum. #### FormSubmitResponse -Empty 200 body for POST /apps//form/human_input/. `extra='forbid'` +Empty 200 body for POST /apps//human-input-forms/:submit. `extra='forbid'` pins `additionalProperties: false` so the generated contract is an exact `{}` rather than an under-annotated open object. @@ -941,9 +941,15 @@ Strict (extra='forbid'). | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | current_identifier | string | | No | -| type | [Type](#type) | | Yes | +| type | [PluginDependencyType](#plugindependencytype) | | Yes | | value | [Github](#github)
[Marketplace](#marketplace)
[Package](#package) | | Yes | +#### PluginDependencyType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| PluginDependencyType | string | | | + #### RevokeResponse | Name | Type | Description | Required | @@ -1016,7 +1022,7 @@ generated CLI whitelist all derive from it. #### TaskStopResponse -200 body for POST /apps//tasks//stop. The handler always returns +200 body for POST /apps//tasks/:stop. The handler always returns {"result": "success"}, so `result` is required (no default) — the generated contract types it as a required `'success'` rather than an optional field. @@ -1024,12 +1030,6 @@ types it as a required `'success'` rather than an optional field. | ---- | ---- | ----------- | -------- | | result | string | | Yes | -#### Type - -| Name | Type | Description | Required | -| ---- | ---- | ----------- | -------- | -| Type | string | | | - #### UsageInfo | Name | Type | Description | Required | diff --git a/api/openapi/markdown/service-openapi.md b/api/openapi/markdown/service-openapi.md index dc4ec74c9af..7680ac77111 100644 --- a/api/openapi/markdown/service-openapi.md +++ b/api/openapi/markdown/service-openapi.md @@ -1046,12 +1046,12 @@ Execute a single datasource node within the knowledge pipeline. Returns a stream #### Responses -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Streaming response with node execution events. | **text/event-stream**: [GeneratedAppResponse](#generatedappresponse)
| -| 401 | Unauthorized - invalid API token | | -| 403 | Forbidden - dataset API access or workspace access denied | | -| 404 | `not_found` : Dataset not found. | | +| Code | Description | +| ---- | ----------- | +| 200 | Streaming response with node execution events. | +| 401 | Unauthorized - invalid API token | +| 403 | Forbidden - dataset API access or workspace access denied | +| 404 | `not_found` : Dataset not found. | ### [POST] /datasets/{dataset_id}/pipeline/run **Run Pipeline** @@ -2335,6 +2335,7 @@ Retrieve the list of available models by type. Primarily used to query `text-emb | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | +| answer | string | | No | | chain_id | string | | No | | created_at | integer | | No | | files | [ string ] | | Yes | @@ -2960,7 +2961,7 @@ Enum class for custom configuration status. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| credentials | [ [DatasourceCredentialInfoResponse](#datasourcecredentialinforesponse) ] | | Yes | +| credentials | [ [DatasourceCredentialInfoResponse](#datasourcecredentialinforesponse) ] | | No | | datasource_type | string | | No | | node_id | string | | No | | plugin_id | string | | No | diff --git a/api/openapi/markdown/web-openapi.md b/api/openapi/markdown/web-openapi.md index ee50341812a..10cc7a831ad 100644 --- a/api/openapi/markdown/web-openapi.md +++ b/api/openapi/markdown/web-openapi.md @@ -600,7 +600,7 @@ Get authentication passport for web application access | Code | Description | Schema | | ---- | ----------- | ------ | -| 200 | Passport retrieved successfully | **application/json**: [AccessTokenData](#accesstokendata)
| +| 200 | Passport retrieved successfully | **application/json**: [PassportAccessTokenResponse](#passportaccesstokenresponse)
| | 401 | Unauthorized - missing app code or invalid authentication | | | 404 | Application or user not found | | @@ -936,6 +936,7 @@ Returns Server-Sent Events stream. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | +| answer | string | | No | | chain_id | string | | No | | created_at | integer | | No | | files | [ string ] | | Yes | @@ -1430,6 +1431,12 @@ Form input definition. | text_to_speech | [JSONObject](#jsonobject) | | Yes | | user_input_form | [ [JSONObject](#jsonobject) ] | | Yes | +#### PassportAccessTokenResponse + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| access_token | string | | Yes | + #### PassportQuery | Name | Type | Description | Required | diff --git a/api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py b/api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py index 2e6395bacc7..d52a3d15544 100644 --- a/api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py +++ b/api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py @@ -33,7 +33,6 @@ from models.dataset import Dataset, DatasetCollectionBinding if TYPE_CHECKING: from qdrant_client.conversions import common_types - from qdrant_client.http import models as rest type DictFilter = dict[str, str | int | bool | dict | list] type MetadataFilter = DictFilter | common_types.Filter diff --git a/api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py b/api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py index 9e6dc27203d..b352243b92a 100644 --- a/api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py +++ b/api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py @@ -41,7 +41,6 @@ from models.enums import TidbAuthBindingStatus if TYPE_CHECKING: from qdrant_client import grpc # noqa from qdrant_client.conversions import common_types - from qdrant_client.http import models as rest type DictFilter = dict[str, str | int | bool | dict | list] type MetadataFilter = DictFilter | common_types.Filter diff --git a/api/pyproject.toml b/api/pyproject.toml index d3fcb59d694..519e400c15f 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -52,6 +52,11 @@ dependencies = [ # Before adding new dependency, consider place it in # alphabet order (a-z) and suitable group. +[tool.dify] +# Oldest difyctl served on /openapi/v1. Bump in lockstep with breaking /openapi/v1 +# changes (paired with difyctl's own version + its compat.minDify). +min_difyctl_version = "0.2.0" + [tool.setuptools] packages = [] diff --git a/api/schedule/check_upgradable_plugin_task.py b/api/schedule/check_upgradable_plugin_task.py index f19a2bf18d8..8d913f1eb13 100644 --- a/api/schedule/check_upgradable_plugin_task.py +++ b/api/schedule/check_upgradable_plugin_task.py @@ -8,7 +8,7 @@ from sqlalchemy import select import app from core.helper.marketplace import fetch_global_plugin_manifest from extensions.ext_database import db -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginAutoUpgradeStrategySetting from tasks import process_tenant_plugin_autoupgrade_check_task as check_task logger = logging.getLogger(__name__) @@ -34,8 +34,7 @@ def check_upgradable_plugin_task(): TenantPluginAutoUpgradeStrategy.upgrade_time_of_day >= now_seconds_of_day, TenantPluginAutoUpgradeStrategy.upgrade_time_of_day < now_seconds_of_day + AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL, - TenantPluginAutoUpgradeStrategy.strategy_setting - != TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED, + TenantPluginAutoUpgradeStrategy.strategy_setting != TenantPluginAutoUpgradeStrategySetting.DISABLED, ) ).all() diff --git a/api/services/account_service.py b/api/services/account_service.py index 445b1acd3b9..b5439467a23 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -16,7 +16,7 @@ from typing import Any, NotRequired, TypedDict, cast from pydantic import BaseModel, TypeAdapter, ValidationError from sqlalchemy import Row, delete, func, select, update -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from werkzeug.exceptions import Unauthorized from configs import dify_config @@ -38,6 +38,8 @@ from models.account import ( Tenant, TenantAccountJoin, TenantAccountRole, + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, TenantPluginAutoUpgradeStrategy, TenantStatus, ) @@ -63,6 +65,8 @@ from services.errors.account import ( LinkAccountIntegrateError, MemberNotInTenantError, NoPermissionError, + RefreshTokenAccountNotFoundError, + RefreshTokenNotFoundError, RoleAlreadyAssignedError, TenantNotFoundError, ) @@ -184,12 +188,12 @@ class AccountService: raise ValueError(f"Builtin RBAC role not found for {role.value} in tenant {tenant_id}") @staticmethod - def get_workspace_permission_keys(tenant_id: str, account_id: str) -> set[str]: - permissions = RBACService.MyPermissions.get(tenant_id, account_id) + def get_workspace_permission_keys(tenant_id: str, account_id: str, *, session: Session) -> set[str]: + permissions = RBACService.MyPermissions.get(tenant_id, account_id, session=session) return set(getattr(getattr(permissions, "workspace", None), "permission_keys", []) or []) @staticmethod - def get_rbac_workspace_owner_account_id(tenant_id: str, actor_account_id: str) -> str: + def get_rbac_workspace_owner_account_id(tenant_id: str, actor_account_id: str, *, session: Session) -> str: """Return the account id bound to the workspace owner RBAC role.""" owner_role_id = AccountService._resolve_legacy_role_id( tenant_id=tenant_id, @@ -207,11 +211,14 @@ class AccountService: return owner_members[0].account_id @staticmethod - def is_rbac_workspace_owner(tenant_id: str, actor_account_id: str, member_account_id: str) -> bool: + def is_rbac_workspace_owner( + tenant_id: str, actor_account_id: str, member_account_id: str, *, session: Session + ) -> bool: roles = RBACService.MemberRoles.get( tenant_id=tenant_id, account_id=actor_account_id, member_account_id=member_account_id, + session=session, ).roles return any( role.is_builtin and role.category == "global_system_default" and role.role_tag == "owner" for role in roles @@ -242,7 +249,7 @@ class AccountService: ) @staticmethod - def _refresh_account_last_active(account: Account, session: scoped_session | Session) -> None: + def _refresh_account_last_active(account: Account, session: Session) -> None: now = naive_utc_now() refresh_before = now - ACCOUNT_LAST_ACTIVE_REFRESH_INTERVAL @@ -272,7 +279,7 @@ class AccountService: redis_client.delete(AccountService._get_account_refresh_token_key(account_id)) @staticmethod - def get_account_by_email(session: Session | scoped_session, email: str) -> Account | None: + def get_account_by_email(email: str, *, session: Session) -> Account | None: """Plain ``Account`` getter keyed by email. Case-sensitive — use :meth:`has_active_account_with_email` for the case-insensitive existence check that backs the SSO collision rule. @@ -280,7 +287,7 @@ class AccountService: return session.execute(select(Account).where(Account.email == email)).scalar_one_or_none() @staticmethod - def has_active_account_with_email(session: Session | scoped_session, email: str) -> bool: + def has_active_account_with_email(email: str, *, session: Session) -> bool: if not email: return False normalized = email.strip().lower() @@ -295,7 +302,7 @@ class AccountService: return row is not None @staticmethod - def get_account_by_id(session: Session | scoped_session, account_id: str) -> Account | None: + def get_account_by_id(account_id: str, *, session: Session) -> Account | None: """Plain ``Account`` getter — no banned check, no tenant rotation, no ``last_active_at`` write. Use this from read-only identity endpoints (``/openapi/v1/account``) where ``load_user``'s @@ -307,7 +314,7 @@ class AccountService: return session.get(Account, account_id) @staticmethod - def load_user(user_id: str, session: scoped_session | Session) -> None | Account: + def load_user(user_id: str, session: Session) -> None | Account: account = session.get(Account, user_id) if not account: return None @@ -359,9 +366,7 @@ class AccountService: return token @staticmethod - def authenticate( - email: str, password: str, invite_token: str | None = None, *, session: scoped_session | Session - ) -> Account: + def authenticate(email: str, password: str, invite_token: str | None = None, *, session: Session) -> Account: """authenticate account with email and password""" account = session.scalar(select(Account).where(Account.email == email).limit(1)) @@ -392,9 +397,7 @@ class AccountService: return account @staticmethod - def update_account_password( - account: Account, password: str, new_password: str, *, session: scoped_session | Session - ): + def update_account_password(account: Account, password: str, new_password: str, *, session: Session): """update account password""" if account.password and not compare_password(password, account.password, account.password_salt): raise CurrentPasswordIncorrectError("Current password is incorrect.") @@ -425,7 +428,7 @@ class AccountService: is_setup: bool | None = False, timezone: str | None = None, *, - session: scoped_session | Session, + session: Session, ) -> Account: """Create an account, preferring explicit user timezone over language-derived defaults.""" if not FeatureService.get_system_features().is_allow_register and not is_setup: @@ -483,7 +486,7 @@ class AccountService: password: str | None = None, timezone: str | None = None, *, - session: scoped_session | Session, + session: Session, ) -> Account: """Create an account and owner workspace.""" account = AccountService.create_account( @@ -540,12 +543,12 @@ class AccountService: return True @staticmethod - def delete_account(account: Account): + def delete_account(account: Account, *, session: Session): """Delete account. This method only adds a task to the queue for deletion.""" # Queue account deletion sync tasks for all workspaces BEFORE account deletion (enterprise only) from services.enterprise.account_deletion_sync import sync_account_deletion - sync_success = sync_account_deletion(account_id=account.id, source="account_deleted") + sync_success = sync_account_deletion(account_id=account.id, source="account_deleted", session=session) if not sync_success: logger.warning( "Enterprise account deletion sync failed for account %s; proceeding with local deletion.", @@ -556,7 +559,7 @@ class AccountService: delete_account_task.delay(account.id) @staticmethod - def link_account_integrate(provider: str, open_id: str, account: Account, *, session: scoped_session | Session): + def link_account_integrate(provider: str, open_id: str, account: Account, *, session: Session): """Link account integrate""" try: # Query whether there is an existing binding record for the same provider @@ -585,13 +588,13 @@ class AccountService: raise LinkAccountIntegrateError("Failed to link account.") from e @staticmethod - def close_account(account: Account, *, session: scoped_session | Session): + def close_account(account: Account, *, session: Session): """Close account""" account.status = AccountStatus.CLOSED session.commit() @staticmethod - def update_account(account: Account, *, session: scoped_session | Session, **kwargs): + def update_account(account: Account, *, session: Session, **kwargs): """Update account fields""" account = session.merge(account) for field, value in kwargs.items(): @@ -604,7 +607,7 @@ class AccountService: return account @staticmethod - def update_account_email(account: Account, email: str, session: scoped_session | Session) -> Account: + def update_account_email(account: Account, email: str, session: Session) -> Account: """Update account email""" account.email = email account_integrate = session.scalar( @@ -617,7 +620,7 @@ class AccountService: return account @staticmethod - def update_login_info(account: Account, session: scoped_session | Session, *, ip_address: str): + def update_login_info(account: Account, session: Session, *, ip_address: str): """Update last login time and ip""" account.last_login_at = naive_utc_now() account.last_login_ip = ip_address @@ -625,7 +628,7 @@ class AccountService: session.commit() @staticmethod - def login(account: Account, *, session: scoped_session | Session, ip_address: str | None = None) -> TokenPair: + def login(account: Account, *, session: Session, ip_address: str | None = None) -> TokenPair: if ip_address: AccountService.update_login_info(account=account, session=session, ip_address=ip_address) @@ -648,15 +651,15 @@ class AccountService: AccountService._delete_refresh_token(refresh_token.decode("utf-8"), account.id) @staticmethod - def refresh_token(refresh_token: str, *, session: scoped_session | Session) -> TokenPair: + def refresh_token(refresh_token: str, *, session: Session) -> TokenPair: # Verify the refresh token account_id = redis_client.get(AccountService._get_refresh_token_key(refresh_token)) if not account_id: - raise ValueError("Invalid refresh token") + raise RefreshTokenNotFoundError("Invalid refresh token") account = AccountService.load_user(account_id.decode("utf-8"), session) if not account: - raise ValueError("Invalid account") + raise RefreshTokenAccountNotFoundError("Invalid account") # Generate new access token and refresh token new_access_token = AccountService.get_account_jwt_token(account) @@ -669,7 +672,7 @@ class AccountService: return TokenPair(access_token=new_access_token, refresh_token=new_refresh_token, csrf_token=csrf_token) @staticmethod - def load_logged_in_account(*, account_id: str, session: scoped_session | Session): + def load_logged_in_account(*, account_id: str, session: Session): return AccountService.load_user(account_id, session) @classmethod @@ -1000,7 +1003,7 @@ class AccountService: return token @staticmethod - def get_account_by_email_with_case_fallback(session: Session | scoped_session, email: str) -> Account | None: + def get_account_by_email_with_case_fallback(email: str, *, session: Session) -> Account | None: """ Retrieve an account by email and fall back to the lowercase email if the original lookup fails. @@ -1022,7 +1025,7 @@ class AccountService: TokenManager.revoke_token(token, "email_code_login") @classmethod - def get_user_through_email(cls, email: str, *, session: scoped_session | Session): + def get_user_through_email(cls, email: str, *, session: Session): if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email): raise AccountRegisterError( description=( @@ -1230,7 +1233,7 @@ class AccountService: return False @staticmethod - def check_email_unique(email: str, *, session: scoped_session | Session) -> bool: + def check_email_unique(email: str, *, session: Session) -> bool: return session.scalar(select(Account).where(Account.email == email).limit(1)) is None @@ -1241,7 +1244,7 @@ class TenantService: is_setup: bool | None = False, is_from_dashboard: bool | None = False, *, - session: scoped_session | Session, + session: Session, ) -> Tenant: """Create tenant""" if ( @@ -1257,13 +1260,13 @@ class TenantService: session.add(tenant) session.commit() - for category in TenantPluginAutoUpgradeStrategy.PluginCategory: + for category in TenantPluginAutoUpgradeCategory: plugin_upgrade_strategy = TenantPluginAutoUpgradeStrategy( tenant_id=tenant.id, category=category, strategy_setting=PluginAutoUpgradeService.default_strategy_setting_for_category(category), upgrade_time_of_day=PluginAutoUpgradeService.default_upgrade_time_of_day(tenant.id), - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=[], include_plugins=[], ) @@ -1275,13 +1278,13 @@ class TenantService: from services.credit_pool_service import CreditPoolService - CreditPoolService.create_default_pool(tenant.id) + CreditPoolService.create_default_pool(tenant.id, session=session) return tenant @staticmethod def create_owner_tenant_if_not_exist( - account: Account, name: str | None = None, is_setup: bool | None = False, *, session: scoped_session | Session + account: Account, name: str | None = None, is_setup: bool | None = False, *, session: Session ): """Check if user have a workspace or not""" available_ta = session.scalar( @@ -1314,6 +1317,7 @@ class TenantService: account_id=account.id, member_account_id=account.id, role_ids=[owner_role_id], + session=session, ) account.current_tenant = tenant session.commit() @@ -1321,7 +1325,7 @@ class TenantService: @staticmethod def create_tenant_member( - tenant: Tenant, account: Account, session: scoped_session | Session, role: str = "normal" + tenant: Tenant, account: Account, session: Session, role: str = "normal" ) -> TenantAccountJoin: """Create tenant member""" if role == TenantAccountRole.OWNER: @@ -1346,7 +1350,7 @@ class TenantService: return ta @staticmethod - def get_join_tenants(account: Account, *, session: scoped_session | Session) -> list[Tenant]: + def get_join_tenants(account: Account, *, session: Session) -> list[Tenant]: """Get account join tenants""" return list( session.scalars( @@ -1357,10 +1361,7 @@ class TenantService: ) @staticmethod - def get_account_memberships( - session: Session | scoped_session, - account_id: str, - ) -> list[Row[tuple[TenantAccountJoin, Tenant]]]: + def get_account_memberships(account_id: str, *, session: Session) -> list[Row[tuple[TenantAccountJoin, Tenant]]]: """Return ``(TenantAccountJoin, Tenant)`` rows for every workspace the account belongs to. Unlike :meth:`get_join_tenants` this keeps the join row so callers can read ``role``/``current`` alongside the @@ -1381,10 +1382,7 @@ class TenantService: ) @staticmethod - def get_workspaces_for_account( - session: Session | scoped_session, - account_id: str, - ) -> list[Row[tuple[Tenant, TenantAccountJoin]]]: + def get_workspaces_for_account(account_id: str, *, session: Session) -> list[Row[tuple[Tenant, TenantAccountJoin]]]: """``(Tenant, TenantAccountJoin)`` rows for every workspace the account belongs to, ordered by ``Tenant.created_at`` ASC — the canonical ordering for ``/openapi/v1/workspaces``. @@ -1403,11 +1401,7 @@ class TenantService: ) @staticmethod - def account_belongs_to_tenant( - session: Session | scoped_session, - account_id: uuid.UUID | str | None, - tenant_id: str, - ) -> bool: + def account_belongs_to_tenant(account_id: uuid.UUID | str | None, tenant_id: str, *, session: Session) -> bool: """Existence check for ``TenantAccountJoin(account_id, tenant_id)``. Backs the CE-deployment membership fallback in ``controllers.openapi.auth.strategies.MembershipStrategy``. @@ -1427,9 +1421,7 @@ class TenantService: @staticmethod def get_account_role_in_tenant( - session: Session | scoped_session, - account_id: uuid.UUID | str | None, - tenant_id: str, + account_id: uuid.UUID | str | None, tenant_id: str, *, session: Session ) -> TenantAccountRole | None: """Return the caller's role in ``tenant_id``, or ``None`` if not a member. @@ -1455,7 +1447,7 @@ class TenantService: return TenantAccountRole(role) if role is not None else None @staticmethod - def get_tenant_by_id(session: Session | scoped_session, tenant_id: str) -> Tenant | None: + def get_tenant_by_id(tenant_id: str, *, session: Session) -> Tenant | None: """Plain ``session.get(Tenant, tenant_id)`` — no status filter. Callers map ``status == ARCHIVE`` to their own error code (the openapi auth pipeline raises 403 ``workspace unavailable``). @@ -1463,10 +1455,7 @@ class TenantService: return session.get(Tenant, tenant_id) @staticmethod - def get_tenants_by_ids( - session: Session | scoped_session, - tenant_ids: list[str], - ) -> list[Tenant]: + def get_tenants_by_ids(tenant_ids: list[str], *, session: Session) -> list[Tenant]: """Bulk ``Tenant`` fetch by primary-key list. Order is unspecified — callers index by ``tenant.id`` (e.g. for cross-tenant denorm in ``/openapi/v1/permitted-external-apps``). @@ -1479,7 +1468,7 @@ class TenantService: return list(session.execute(select(Tenant).where(Tenant.id.in_(tenant_ids))).scalars().all()) @staticmethod - def get_tenant_name(session: Session | scoped_session, tenant_id: str) -> str | None: + def get_tenant_name(tenant_id: str, *, session: Session) -> str | None: """Single-column tenant name read. Used by openapi list endpoints to denormalize ``workspace_name`` onto each row without dragging the full ``Tenant`` ORM entity through. @@ -1488,9 +1477,7 @@ class TenantService: @staticmethod def find_workspace_for_account( - session: Session | scoped_session, - account_id: str, - workspace_id: str, + account_id: str, workspace_id: str, *, session: Session ) -> Row[tuple[Tenant, TenantAccountJoin]] | None: """Single ``(Tenant, TenantAccountJoin)`` row scoped to the account's membership in ``workspace_id``. ``None`` on non-member @@ -1507,7 +1494,7 @@ class TenantService: ).first() @staticmethod - def get_current_tenant_by_account(account: Account, *, session: scoped_session | Session): + def get_current_tenant_by_account(account: Account, *, session: Session): """Get tenant by account and add the role""" tenant = account.current_tenant if not tenant: @@ -1525,7 +1512,7 @@ class TenantService: return tenant @staticmethod - def switch_tenant(account: Account, tenant_id: str | None = None, *, session: scoped_session | Session): + def switch_tenant(account: Account, tenant_id: str | None = None, *, session: Session): """Switch the current workspace for the account""" # Ensure tenant_id is provided @@ -1558,7 +1545,7 @@ class TenantService: session.commit() @staticmethod - def get_tenant_members(tenant: Tenant, *, session: scoped_session | Session) -> list[Account]: + def get_tenant_members(tenant: Tenant, *, session: Session) -> list[Account]: """Get tenant members""" stmt = ( select(Account, TenantAccountJoin.role) @@ -1577,7 +1564,7 @@ class TenantService: return updated_accounts @staticmethod - def get_dataset_operator_members(tenant: Tenant, *, session: scoped_session | Session) -> list[Account]: + def get_dataset_operator_members(tenant: Tenant, *, session: Session) -> list[Account]: """Get dataset admin members""" stmt = ( select(Account, TenantAccountJoin.role) @@ -1597,7 +1584,7 @@ class TenantService: return updated_accounts @staticmethod - def has_roles(tenant: Tenant, roles: list[TenantAccountRole], *, session: scoped_session | Session) -> bool: + def has_roles(tenant: Tenant, roles: list[TenantAccountRole], *, session: Session) -> bool: """Check if user has any of the given roles for a tenant""" if not all(isinstance(role, TenantAccountRole) for role in roles): raise ValueError("all roles must be TenantAccountRole") @@ -1615,9 +1602,7 @@ class TenantService: ) @staticmethod - def get_user_role( - account: Account, tenant: Tenant, *, session: scoped_session | Session - ) -> TenantAccountRole | None: + def get_user_role(account: Account, tenant: Tenant, *, session: Session) -> TenantAccountRole | None: """Get the role of the current account for a given tenant""" join = session.scalar( select(TenantAccountJoin) @@ -1627,13 +1612,13 @@ class TenantService: return TenantAccountRole(join.role) if join else None @staticmethod - def get_tenant_count(*, session: scoped_session | Session) -> int: + def get_tenant_count(*, session: Session) -> int: """Get tenant count""" return cast(int, session.scalar(select(func.count(Tenant.id)))) @staticmethod def check_member_permission( - tenant: Tenant, operator: Account, member: Account | None, action: str, *, session: scoped_session | Session + tenant: Tenant, operator: Account, member: Account | None, action: str, *, session: Session ): """Check member permission""" if action not in {"add", "remove", "update"}: @@ -1647,6 +1632,7 @@ class TenantService: workspace_permission_keys = AccountService.get_workspace_permission_keys( str(tenant.id), str(operator.id), + session=session, ) required_permission_key = ( "workspace.member.manage" if action in {"add", "remove"} else "workspace.role.manage" @@ -1657,7 +1643,9 @@ class TenantService: if ( action == "remove" and member - and AccountService.is_rbac_workspace_owner(str(tenant.id), str(operator.id), str(member.id)) + and AccountService.is_rbac_workspace_owner( + str(tenant.id), str(operator.id), str(member.id), session=session + ) ): raise NoPermissionError(f"No permission to {action} member.") return @@ -1687,9 +1675,7 @@ class TenantService: raise NoPermissionError(f"No permission to {action} member.") @staticmethod - def remove_member_from_tenant( - tenant: Tenant, account: Account, operator: Account, *, session: scoped_session | Session - ): + def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account, *, session: Session): """Remove member from tenant. Apps and datasets maintained by the removed member are reassigned to @@ -1718,7 +1704,9 @@ class TenantService: owner_id: str | None if dify_config.RBAC_ENABLED: - owner_id = AccountService.get_rbac_workspace_owner_account_id(str(tenant.id), str(operator.id)) + owner_id = AccountService.get_rbac_workspace_owner_account_id( + str(tenant.id), str(operator.id), session=session + ) else: owner_id = session.scalar( select(TenantAccountJoin.account_id) @@ -1792,9 +1780,7 @@ class TenantService: RBACService.MemberRoles.delete_rbac_bindings(tenant_id=tenant.id, account_id=account_id) @staticmethod - def update_member_role( - tenant: Tenant, member: Account, new_role: str, operator: Account, *, session: scoped_session | Session - ): + def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account, *, session: Session): """Update member role""" TenantService.check_member_permission(tenant, operator, member, "update", session=session) new_tenant_role = TenantAccountRole(new_role) @@ -1837,6 +1823,7 @@ class TenantService: account_id=operator.id, member_account_id=str(current_owner_join.account_id), role_ids=[admin_role_id], + session=session, ) # Update the role of the target member @@ -1851,6 +1838,7 @@ class TenantService: account_id=operator.id, member_account_id=member.id, role_ids=[resolved_role_id], + session=session, ) else: target_member_join.role = new_tenant_role @@ -1863,11 +1851,11 @@ class TenantService: return tenant.custom_config_dict @staticmethod - def is_owner(account: Account, tenant: Tenant, *, session: scoped_session | Session) -> bool: + def is_owner(account: Account, tenant: Tenant, *, session: Session) -> bool: return TenantService.get_user_role(account, tenant, session=session) == TenantAccountRole.OWNER @staticmethod - def is_member(account: Account, tenant: Tenant, *, session: scoped_session | Session) -> bool: + def is_member(account: Account, tenant: Tenant, *, session: Session) -> bool: """Check if the account is a member of the tenant""" return TenantService.get_user_role(account, tenant, session=session) is not None @@ -1886,7 +1874,7 @@ class RegisterService: ip_address: str, language: str | None, *, - session: scoped_session | Session, + session: Session, ): """ Setup dify @@ -1939,7 +1927,7 @@ class RegisterService: create_workspace_required: bool | None = True, timezone: str | None = None, *, - session: scoped_session | Session, + session: Session, ) -> Account: """Register account""" session.begin_nested() @@ -2001,7 +1989,7 @@ class RegisterService: role: str = "normal", inviter: Account | None = None, *, - session: scoped_session | Session, + session: Session, ) -> str: if not inviter: raise ValueError("Inviter is required") @@ -2015,7 +2003,7 @@ class RegisterService: check_workspace_member_invite_permission(tenant.id) - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=session) requires_setup = False if not account: @@ -2053,6 +2041,7 @@ class RegisterService: account_id=inviter.id, member_account_id=account.id, role_ids=[role], + session=session, ) if ta or dify_config.RBAC_ENABLED: raise AccountAlreadyInTenantError("Account already in tenant.") @@ -2064,6 +2053,7 @@ class RegisterService: account_id=inviter.id, member_account_id=account.id, role_ids=[role], + session=session, ) token = cls.generate_invite_token(tenant, account, role, requires_setup=requires_setup) @@ -2112,7 +2102,7 @@ class RegisterService: @classmethod def get_invitation_if_token_valid( - cls, workspace_id: str | None, email: str | None, token: str, *, session: scoped_session | Session + cls, workspace_id: str | None, email: str | None, token: str, *, session: Session ) -> InvitationDetailDict | None: invitation_data = cls.get_invitation_by_token(token, workspace_id, email) if not invitation_data: @@ -2165,7 +2155,7 @@ class RegisterService: @classmethod def get_invitation_with_case_fallback( - cls, workspace_id: str | None, email: str | None, token: str, *, session: scoped_session | Session + cls, workspace_id: str | None, email: str | None, token: str, *, session: Session ) -> InvitationDetailDict | None: invitation = cls.get_invitation_if_token_valid(workspace_id, email, token, session=session) if invitation or not email or email == email.lower(): diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index 99817eb6441..4839012dd3a 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -4,6 +4,7 @@ from typing import Any from sqlalchemy import func, or_, select from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session from sqlalchemy.sql.elements import ColumnElement from extensions.ext_database import db @@ -33,6 +34,7 @@ from models.workflow import Workflow from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_validator import ComposerConfigValidator from services.agent.errors import ( + AgentModelNotConfiguredError, AgentNameConflictError, AgentNotFoundError, AgentVersionConflictError, @@ -103,23 +105,35 @@ def _agent_soul_config_json(agent_soul: AgentSoulConfig | dict[str, Any]) -> dic class AgentComposerService: @classmethod def load_workflow_composer( - cls, *, tenant_id: str, app_id: str, node_id: str, account_id: str | None = None, snapshot_id: str | None = None + cls, + *, + tenant_id: str, + app_id: str, + node_id: str, + account_id: str | None = None, + snapshot_id: str | None = None, + session: Session, ) -> dict[str, Any]: - workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) - binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) + workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id, session=session) + binding = cls._get_workflow_binding( + tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id, session=session + ) if not binding: if snapshot_id: raise AgentVersionNotFoundError() return cls._empty_workflow_state(app_id=app_id, workflow_id=workflow.id, node_id=node_id) - agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id) + agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) version = cls._workflow_composer_version( tenant_id=tenant_id, binding=binding, agent=agent, snapshot_id=snapshot_id, + session=session, + ) + return cls._serialize_workflow_state( + binding=binding, agent=agent, version=version, account_id=account_id, session=session ) - return cls._serialize_workflow_state(binding=binding, agent=agent, version=version, account_id=account_id) @classmethod def _workflow_composer_version( @@ -129,6 +143,7 @@ class AgentComposerService: binding: WorkflowAgentNodeBinding, agent: Agent | None, snapshot_id: str | None, + session: Session, ) -> AgentConfigSnapshot | None: if snapshot_id: if agent is None: @@ -146,7 +161,7 @@ class AgentComposerService: raise AgentVersionNotFoundError() else: raise AgentVersionNotFoundError() - return cls._require_version(tenant_id=tenant_id, agent_id=agent.id, version_id=snapshot_id) + return cls._require_version(tenant_id=tenant_id, agent_id=agent.id, version_id=snapshot_id, session=session) version_id = ( agent.active_config_snapshot_id @@ -157,20 +172,31 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent.id if agent else None, version_id=version_id, + session=session, ) @classmethod def save_workflow_composer( - cls, *, tenant_id: str, app_id: str, node_id: str, account_id: str, payload: ComposerSavePayload + cls, + *, + tenant_id: str, + app_id: str, + node_id: str, + account_id: str, + payload: ComposerSavePayload, + session: Session, ) -> dict[str, Any]: if payload.variant != ComposerVariant.WORKFLOW: raise ValueError("Workflow composer endpoint only accepts workflow variant") _backfill_cli_tool_ids(payload.agent_soul) _validate_composer_payload_for_strategy(payload) - cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) - workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) - binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) + if payload.save_strategy in _PUBLISH_SAVE_STRATEGIES: + cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) + workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id, session=session) + binding = cls._get_workflow_binding( + tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id, session=session + ) match payload.save_strategy: case ComposerSaveStrategy.NODE_JOB_ONLY: @@ -182,14 +208,15 @@ class AgentComposerService: account_id=account_id, binding=binding, payload=payload, + session=session, ) case ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION: binding = cls._save_to_current_version( - tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload + tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload, session=session ) case ComposerSaveStrategy.SAVE_AS_NEW_VERSION: binding = cls._save_as_new_version( - tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload + tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload, session=session ) case ComposerSaveStrategy.SAVE_AS_NEW_AGENT: binding = cls._save_as_new_agent( @@ -200,14 +227,15 @@ class AgentComposerService: account_id=account_id, binding=binding, payload=payload, + session=session, ) case ComposerSaveStrategy.SAVE_TO_ROSTER: binding = cls._save_to_roster( - tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload + tenant_id=tenant_id, account_id=account_id, binding=binding, payload=payload, session=session ) - db.session.commit() - agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id) + session.commit() + agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) version_id = ( agent.active_config_snapshot_id if agent and binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT @@ -217,12 +245,16 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent.id if agent else None, version_id=version_id, + session=session, + ) + state = cls._serialize_workflow_state( + binding=binding, agent=agent, version=version, account_id=account_id, session=session ) - state = cls._serialize_workflow_state(binding=binding, agent=agent, version=version, account_id=account_id) state["validation"] = cls.collect_validation_findings( tenant_id=tenant_id, payload=payload, agent_id=binding.agent_id, + session=session, ) return state @@ -237,33 +269,38 @@ class AgentComposerService: source_agent_id: str, source_snapshot_id: str | None = None, idempotency_key: str | None = None, + session: Session, ) -> dict[str, Any]: - workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) + workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id, session=session) binding = cls._require_binding( - cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) + cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id, session=session) ) if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and idempotency_key: - agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id) + agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) version = cls._get_version_if_present( tenant_id=tenant_id, agent_id=agent.id if agent else None, version_id=binding.current_snapshot_id, + session=session, + ) + return cls._serialize_workflow_state( + binding=binding, agent=agent, version=version, account_id=account_id, session=session ) - return cls._serialize_workflow_state(binding=binding, agent=agent, version=version, account_id=account_id) if binding.binding_type != WorkflowAgentBindingType.ROSTER_AGENT: raise InvalidComposerConfigError("Workflow agent node must be bound to a roster agent.") if binding.agent_id != source_agent_id: raise InvalidComposerConfigError("Source agent does not match the current workflow node binding.") - source_agent = cls._require_agent(tenant_id=tenant_id, agent_id=source_agent_id) + source_agent = cls._require_agent(tenant_id=tenant_id, agent_id=source_agent_id, session=session) if source_agent.scope != AgentScope.ROSTER or source_agent.status != AgentStatus.ACTIVE: raise InvalidComposerConfigError("Source agent must be an active roster agent.") source_version = cls._require_version( tenant_id=tenant_id, agent_id=source_agent.id, version_id=source_agent.active_config_snapshot_id, + session=session, ) if source_snapshot_id and source_snapshot_id != source_version.id: raise AgentVersionConflictError() @@ -282,6 +319,7 @@ class AgentComposerService: icon_type=source_agent.icon_type, icon=source_agent.icon, icon_background=source_agent.icon_background, + session=session, ) cls._copy_agent_drive_rows( tenant_id=tenant_id, @@ -290,45 +328,48 @@ class AgentComposerService: account_id=account_id, agent_soul=agent_soul, node_job=WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict), + session=session, ) binding.binding_type = WorkflowAgentBindingType.INLINE_AGENT binding.agent_id = inline_agent.id binding.current_snapshot_id = inline_agent.active_config_snapshot_id binding.updated_by = account_id - db.session.flush() - db.session.commit() + session.flush() + session.commit() version = cls._require_version( tenant_id=tenant_id, agent_id=inline_agent.id, version_id=inline_agent.active_config_snapshot_id, + session=session, ) return cls._serialize_workflow_state( - binding=binding, agent=inline_agent, version=version, account_id=account_id + binding=binding, agent=inline_agent, version=version, account_id=account_id, session=session ) @classmethod - def load_agent_app_composer(cls, *, tenant_id: str, app_id: str) -> dict[str, Any]: - agent = cls._require_agent_app_agent(tenant_id=tenant_id, app_id=app_id) - return cls._load_agent_composer_for_agent(tenant_id=tenant_id, agent=agent) + def load_agent_app_composer(cls, *, tenant_id: str, app_id: str, session: Session) -> dict[str, Any]: + agent = cls._require_agent_app_agent(tenant_id=tenant_id, app_id=app_id, session=session) + return cls._load_agent_composer_for_agent(tenant_id=tenant_id, agent=agent, session=session) @classmethod - def load_agent_composer(cls, *, tenant_id: str, agent_id: str) -> dict[str, Any]: - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) - return cls._load_agent_composer_for_agent(tenant_id=tenant_id, agent=agent) + def load_agent_composer(cls, *, tenant_id: str, agent_id: str, session: Session) -> dict[str, Any]: + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) + return cls._load_agent_composer_for_agent(tenant_id=tenant_id, agent=agent, session=session) @classmethod - def _load_agent_composer_for_agent(cls, *, tenant_id: str, agent: Agent) -> dict[str, Any]: + def _load_agent_composer_for_agent(cls, *, tenant_id: str, agent: Agent, session: Session) -> dict[str, Any]: draft = cls._get_or_create_agent_draft( tenant_id=tenant_id, agent=agent, draft_type=AgentConfigDraftType.DRAFT, account_id=None, created_by=agent.updated_by or agent.created_by, + session=session, ) version = cls._get_version_if_present( - tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id + tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id, session=session ) return { "variant": ComposerVariant.AGENT_APP.value, @@ -345,7 +386,13 @@ class AgentComposerService: @classmethod def save_agent_app_composer( - cls, *, tenant_id: str, app_id: str, account_id: str, payload: ComposerSavePayload + cls, + *, + tenant_id: str, + app_id: str, + account_id: str, + payload: ComposerSavePayload, + session: Session, ) -> dict[str, Any]: if payload.variant != ComposerVariant.AGENT_APP: raise ValueError("Agent App composer endpoint only accepts agent_app variant") @@ -357,9 +404,8 @@ class AgentComposerService: raise ValueError("agent_soul is required") _backfill_cli_tool_ids(payload.agent_soul) _validate_composer_payload_for_strategy(payload) - cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) - agent = cls._get_agent_app_agent(tenant_id=tenant_id, app_id=app_id) + agent = cls._get_agent_app_agent(tenant_id=tenant_id, app_id=app_id, session=session) if not agent: agent = Agent( tenant_id=tenant_id, @@ -374,22 +420,29 @@ class AgentComposerService: created_by=account_id, updated_by=account_id, ) - db.session.add(agent) + session.add(agent) try: - db.session.flush() + session.flush() except IntegrityError as exc: - db.session.rollback() + session.rollback() raise AgentNameConflictError() from exc return cls._save_agent_composer_for_agent( tenant_id=tenant_id, agent=agent, account_id=account_id, payload=payload, + session=session, ) @classmethod def save_agent_composer( - cls, *, tenant_id: str, agent_id: str, account_id: str, payload: ComposerSavePayload + cls, + *, + tenant_id: str, + agent_id: str, + account_id: str, + payload: ComposerSavePayload, + session: Session, ) -> dict[str, Any]: if payload.variant != ComposerVariant.AGENT_APP: raise ValueError("Agent composer endpoint only accepts agent_app variant") @@ -401,18 +454,24 @@ class AgentComposerService: raise ValueError("agent_soul is required") _backfill_cli_tool_ids(payload.agent_soul) _validate_composer_payload_for_strategy(payload) - cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) return cls._save_agent_composer_for_agent( tenant_id=tenant_id, agent=agent, account_id=account_id, payload=payload, + session=session, ) @classmethod def _save_agent_composer_for_agent( - cls, *, tenant_id: str, agent: Agent, account_id: str, payload: ComposerSavePayload + cls, + *, + tenant_id: str, + agent: Agent, + account_id: str, + payload: ComposerSavePayload, + session: Session, ) -> dict[str, Any]: if payload.agent_soul is None: raise ValueError("agent_soul is required") @@ -423,20 +482,23 @@ class AgentComposerService: account_id=None, agent_soul=payload.agent_soul, account_id_for_audit=account_id, + session=session, ) agent.updated_by = account_id agent.active_config_is_published = cls._agent_soul_matches_active_config( tenant_id=tenant_id, agent=agent, agent_soul=payload.agent_soul, + session=session, ) - db.session.commit() - state = cls.load_agent_composer(tenant_id=tenant_id, agent_id=agent.id) + session.commit() + state = cls.load_agent_composer(tenant_id=tenant_id, agent_id=agent.id, session=session) state["validation"] = cls.collect_validation_findings( tenant_id=tenant_id, payload=payload, agent_id=agent.id, + session=session, ) return state @@ -447,6 +509,7 @@ class AgentComposerService: tenant_id: str, agent: Agent, agent_soul: AgentSoulConfig, + session: Session, ) -> bool: if not agent.active_config_snapshot_id: return False @@ -455,6 +518,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id, + session=session, ) if not active_version: return False @@ -490,9 +554,15 @@ class AgentComposerService: @classmethod def publish_agent_app_draft( - cls, *, tenant_id: str, agent_id: str, account_id: str, version_note: str | None = None + cls, + *, + tenant_id: str, + agent_id: str, + account_id: str, + version_note: str | None = None, + session: Session, ) -> dict[str, Any]: - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) if agent.scope != AgentScope.ROSTER or agent.source != AgentSource.AGENT_APP: raise AgentNotFoundError() draft = cls._get_or_create_agent_draft( @@ -501,6 +571,7 @@ class AgentComposerService: draft_type=AgentConfigDraftType.DRAFT, account_id=None, created_by=account_id, + session=session, ) agent_soul = AgentSoulConfig.model_validate(draft.config_snapshot_dict) ComposerConfigValidator.validate_publish_payload( @@ -511,6 +582,8 @@ class AgentComposerService: version_note=version_note, ) ) + if not agent_soul_has_model(agent_soul): + raise AgentModelNotConfiguredError() cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=agent_soul) version = cls._create_config_version( tenant_id=tenant_id, @@ -520,6 +593,7 @@ class AgentComposerService: operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, version_note=version_note, previous_snapshot_id=agent.active_config_snapshot_id, + session=session, ) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(agent_soul) @@ -527,7 +601,7 @@ class AgentComposerService: agent.updated_by = account_id draft.base_snapshot_id = version.id draft.updated_by = account_id - db.session.commit() + session.commit() return { "result": "success", "active_config_snapshot_id": version.id, @@ -537,21 +611,29 @@ class AgentComposerService: @classmethod def checkout_agent_app_build_draft( - cls, *, tenant_id: str, agent_id: str, account_id: str, force: bool = False + cls, + *, + tenant_id: str, + agent_id: str, + account_id: str, + force: bool = False, + session: Session, ) -> dict[str, Any]: - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) normal_draft = cls._get_or_create_agent_draft( tenant_id=tenant_id, agent=agent, draft_type=AgentConfigDraftType.DRAFT, account_id=None, created_by=account_id, + session=session, ) build_draft = cls._get_agent_draft( tenant_id=tenant_id, agent_id=agent.id, draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id=account_id, + session=session, ) if build_draft is not None and not force: return cls._serialize_build_draft_state(build_draft) @@ -564,20 +646,23 @@ class AgentComposerService: draft_owner_key=account_id, created_by=account_id, ) - db.session.add(build_draft) + session.add(build_draft) build_draft.base_snapshot_id = normal_draft.base_snapshot_id build_draft.config_snapshot = AgentSoulConfig.model_validate(normal_draft.config_snapshot_dict) build_draft.updated_by = account_id - db.session.commit() + session.commit() return cls._serialize_build_draft_state(build_draft) @classmethod - def load_agent_app_build_draft(cls, *, tenant_id: str, agent_id: str, account_id: str) -> dict[str, Any]: + def load_agent_app_build_draft( + cls, *, tenant_id: str, agent_id: str, account_id: str, session: Session + ) -> dict[str, Any]: build_draft = cls._get_agent_draft( tenant_id=tenant_id, agent_id=agent_id, draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id=account_id, + session=session, ) if build_draft is None: raise AgentVersionNotFoundError() @@ -585,14 +670,19 @@ class AgentComposerService: @classmethod def save_agent_app_build_draft( - cls, *, tenant_id: str, agent_id: str, account_id: str, payload: ComposerSavePayload + cls, + *, + tenant_id: str, + agent_id: str, + account_id: str, + payload: ComposerSavePayload, + session: Session, ) -> dict[str, Any]: if payload.agent_soul is None: raise ValueError("agent_soul is required") _backfill_cli_tool_ids(payload.agent_soul) ComposerConfigValidator.validate_draft_save_payload(payload) - cls.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul) - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) build_draft = cls._save_agent_draft( tenant_id=tenant_id, agent=agent, @@ -600,18 +690,22 @@ class AgentComposerService: account_id=account_id, agent_soul=payload.agent_soul, account_id_for_audit=account_id, + session=session, ) - db.session.commit() + session.commit() return cls._serialize_build_draft_state(build_draft) @classmethod - def apply_agent_app_build_draft(cls, *, tenant_id: str, agent_id: str, account_id: str) -> dict[str, Any]: - agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id) + def apply_agent_app_build_draft( + cls, *, tenant_id: str, agent_id: str, account_id: str, session: Session + ) -> dict[str, Any]: + agent = cls._require_agent(tenant_id=tenant_id, agent_id=agent_id, session=session) build_draft = cls._get_agent_draft( tenant_id=tenant_id, agent_id=agent.id, draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id=account_id, + session=session, ) if build_draft is None: raise AgentVersionNotFoundError() @@ -624,28 +718,33 @@ class AgentComposerService: agent_soul=applied_agent_soul, account_id_for_audit=account_id, base_snapshot_id=build_draft.base_snapshot_id, + session=session, ) agent.active_config_is_published = cls._agent_soul_matches_active_config( tenant_id=tenant_id, agent=agent, agent_soul=applied_agent_soul, + session=session, ) agent.updated_by = account_id - db.session.delete(build_draft) - db.session.commit() + session.delete(build_draft) + session.commit() return {"result": "success", "draft": cls._serialize_draft(normal_draft)} @classmethod - def discard_agent_app_build_draft(cls, *, tenant_id: str, agent_id: str, account_id: str) -> dict[str, Any]: + def discard_agent_app_build_draft( + cls, *, tenant_id: str, agent_id: str, account_id: str, session: Session + ) -> dict[str, Any]: build_draft = cls._get_agent_draft( tenant_id=tenant_id, agent_id=agent_id, draft_type=AgentConfigDraftType.DEBUG_BUILD, account_id=account_id, + session=session, ) if build_draft is not None: - db.session.delete(build_draft) - db.session.commit() + session.delete(build_draft) + session.commit() return {"result": "success"} @classmethod @@ -655,6 +754,7 @@ class AgentComposerService: tenant_id: str, payload: ComposerSavePayload, agent_id: str | None = None, + session: Session, ) -> dict[str, Any]: """ENG-617 soft findings, with DB-backed dataset and drive mention checks.""" existing_knowledge_set_ids = ( @@ -672,6 +772,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent_id, prompt=payload.agent_soul.prompt.system_prompt, + session=session, ) ) return findings @@ -695,9 +796,9 @@ class AgentComposerService: ) @classmethod - def resolve_bound_agent_id(cls, *, tenant_id: str, app_id: str) -> str | None: + def resolve_bound_agent_id(cls, *, tenant_id: str, app_id: str, session: Session) -> str | None: """The Agent App's bound roster agent id, if any (validate-endpoint context).""" - return db.session.scalar( + return session.scalar( select(Agent.id) .where( Agent.tenant_id == tenant_id, @@ -710,13 +811,17 @@ class AgentComposerService: ) @classmethod - def resolve_workflow_node_agent_id(cls, *, tenant_id: str, app_id: str, node_id: str) -> str | None: + def resolve_workflow_node_agent_id( + cls, *, tenant_id: str, app_id: str, node_id: str, session: Session + ) -> str | None: """The draft workflow node binding's agent id, if any (validate-endpoint context).""" try: - workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) + workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id, session=session) except ValueError: return None - binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) + binding = cls._get_workflow_binding( + tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id, session=session + ) return binding.agent_id if binding else None @classmethod @@ -726,6 +831,7 @@ class AgentComposerService: tenant_id: str, agent_id: str, prompt: str, + session: Session, ) -> list[dict[str, str | None]]: """Soft warnings for missing drive-backed prompt mentions.""" from services.agent.prompt_mentions import MentionKind, parse_prompt_mentions @@ -743,7 +849,7 @@ class AgentComposerService: return [] existing_keys = set( - db.session.scalars( + session.scalars( select(AgentDriveFile.key).where( AgentDriveFile.tenant_id == tenant_id, AgentDriveFile.agent_id == agent_id, @@ -767,22 +873,32 @@ class AgentComposerService: return findings @classmethod - def get_workflow_candidates(cls, *, tenant_id: str, app_id: str, node_id: str, user_id: str) -> dict[str, Any]: + def get_workflow_candidates( + cls, + *, + tenant_id: str, + app_id: str, + node_id: str, + user_id: str, + session: Session, + ) -> dict[str, Any]: """Slash-menu data source for the workflow Agent node composer (ENG-615).""" from services.agent.composer_candidates import previous_node_output_candidates, soul_candidates try: - workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) + workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id, session=session) except ValueError: workflow = None node_job: WorkflowNodeJobConfig | None = None agent_soul: AgentSoulConfig | None = None if workflow is not None: - binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) + binding = cls._get_workflow_binding( + tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id, session=session + ) if binding is not None: node_job = cls._parse_node_job(binding) - agent_soul = cls._load_binding_soul(tenant_id=tenant_id, binding=binding) + agent_soul = cls._load_binding_soul(tenant_id=tenant_id, binding=binding, session=session) truncated = False previous_outputs: list[dict[str, Any]] = [] @@ -793,7 +909,7 @@ class AgentComposerService: graph=workflow.graph_dict, node_id=node_id, declared_outputs_loader=lambda nid: cls._binding_declared_outputs( - tenant_id=tenant_id, workflow_id=workflow.id, node_id=nid + tenant_id=tenant_id, workflow_id=workflow.id, node_id=nid, session=session ), draft_variables_loader=lambda nid: cls._draft_node_variables( session=draft_variable_session, app_id=app_id, node_id=nid, user_id=user_id @@ -828,11 +944,13 @@ class AgentComposerService: return response.model_dump(mode="json") @classmethod - def get_agent_app_candidates(cls, *, tenant_id: str, agent_id: str, user_id: str) -> dict[str, Any]: + def get_agent_app_candidates( + cls, *, tenant_id: str, agent_id: str, user_id: str, session: Session + ) -> dict[str, Any]: """Slash-menu data source for the Agent App (Console) composer (ENG-615).""" from services.agent.composer_candidates import soul_candidates - agent_soul = cls._load_agent_soul(tenant_id=tenant_id, agent_id=agent_id) + agent_soul = cls._load_agent_soul(tenant_id=tenant_id, agent_id=agent_id, session=session) soul_lists, truncated = soul_candidates( agent_soul=agent_soul, dataset_lookup=lambda ids: get_tenant_knowledge_dataset_rows(tenant_id=tenant_id, dataset_ids=ids), @@ -857,18 +975,21 @@ class AgentComposerService: return None @classmethod - def _load_binding_soul(cls, *, tenant_id: str, binding: WorkflowAgentNodeBinding) -> AgentSoulConfig | None: - agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id) + def _load_binding_soul( + cls, *, tenant_id: str, binding: WorkflowAgentNodeBinding, session: Session + ) -> AgentSoulConfig | None: + agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) version = cls._get_version_if_present( tenant_id=tenant_id, agent_id=agent.id if agent else None, version_id=binding.current_snapshot_id, + session=session, ) return cls._parse_soul_snapshot(version) @classmethod - def _load_agent_soul(cls, *, tenant_id: str, agent_id: str) -> AgentSoulConfig | None: - agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=agent_id) + def _load_agent_soul(cls, *, tenant_id: str, agent_id: str, session: Session) -> AgentSoulConfig | None: + agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=agent_id, session=session) if agent is None: return None draft = cls._get_or_create_agent_draft( @@ -877,6 +998,7 @@ class AgentComposerService: draft_type=AgentConfigDraftType.DRAFT, account_id=None, created_by=agent.updated_by or agent.created_by, + session=session, ) return AgentSoulConfig.model_validate(draft.config_snapshot_dict) @@ -892,9 +1014,11 @@ class AgentComposerService: @classmethod def _binding_declared_outputs( - cls, *, tenant_id: str, workflow_id: str, node_id: str + cls, *, tenant_id: str, workflow_id: str, node_id: str, session: Session ) -> list[DeclaredOutputConfig] | None: - binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow_id, node_id=node_id) + binding = cls._get_workflow_binding( + tenant_id=tenant_id, workflow_id=workflow_id, node_id=node_id, session=session + ) if binding is None: return None node_job = cls._parse_node_job(binding) @@ -969,8 +1093,8 @@ class AgentComposerService: return tools @classmethod - def calculate_impact(cls, *, tenant_id: str, current_snapshot_id: str) -> dict[str, Any]: - snapshot = db.session.scalar( + def calculate_impact(cls, *, tenant_id: str, current_snapshot_id: str, session: Session) -> dict[str, Any]: + snapshot = session.scalar( select(AgentConfigSnapshot) .where( AgentConfigSnapshot.tenant_id == tenant_id, @@ -986,7 +1110,7 @@ class AgentComposerService: & (WorkflowAgentNodeBinding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT) ) bindings = list( - db.session.scalars( + session.scalars( select(WorkflowAgentNodeBinding).where( WorkflowAgentNodeBinding.tenant_id == tenant_id, or_(*predicates), @@ -1017,6 +1141,7 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding | None, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: node_job = payload.node_job or WorkflowNodeJobConfig() if binding: @@ -1029,6 +1154,7 @@ class AgentComposerService: account_id=account_id, binding=binding, payload=payload, + session=session, ) binding.node_job_config = node_job if payload.agent_soul is not None and binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT: @@ -1036,6 +1162,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=binding.agent_id, version_id=binding.current_snapshot_id, + session=session, ) version = cls._update_current_version( current_snapshot=current_snapshot, @@ -1043,8 +1170,9 @@ class AgentComposerService: agent_soul=payload.agent_soul, operation=AgentConfigRevisionOperation.SAVE_CURRENT_VERSION, version_note=payload.version_note, + session=session, ) - agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) if agent.scope != AgentScope.WORKFLOW_ONLY: raise ValueError("Inline workflow agent binding must point to a workflow-only agent") agent.active_config_snapshot_id = version.id @@ -1063,6 +1191,7 @@ class AgentComposerService: node_id=node_id, account_id=account_id, agent_soul=agent_soul, + session=session, ) binding = WorkflowAgentNodeBinding( tenant_id=tenant_id, @@ -1077,8 +1206,8 @@ class AgentComposerService: created_by=account_id, updated_by=account_id, ) - db.session.add(binding) - db.session.flush() + session.add(binding) + session.flush() return binding @classmethod @@ -1100,6 +1229,7 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: if payload.binding and (payload.binding.agent_id or payload.binding.current_snapshot_id): raise ValueError("Start from Scratch must not provide an existing inline agent binding.") @@ -1112,13 +1242,14 @@ class AgentComposerService: node_id=node_id, account_id=account_id, agent_soul=agent_soul, + session=session, ) binding.binding_type = WorkflowAgentBindingType.INLINE_AGENT binding.agent_id = agent.id binding.current_snapshot_id = agent.active_config_snapshot_id binding.node_job_config = payload.node_job or binding.node_job_config binding.updated_by = account_id - db.session.flush() + session.flush() return binding @classmethod @@ -1129,6 +1260,7 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding | None, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: binding = cls._require_binding(binding) if payload.agent_soul is None: @@ -1137,6 +1269,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=binding.agent_id, version_id=binding.current_snapshot_id, + session=session, ) version = cls._update_current_version( current_snapshot=current_snapshot, @@ -1144,8 +1277,9 @@ class AgentComposerService: agent_soul=payload.agent_soul, operation=AgentConfigRevisionOperation.SAVE_CURRENT_VERSION, version_note=payload.version_note, + session=session, ) - agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(payload.agent_soul) agent.active_config_is_published = True @@ -1164,6 +1298,7 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding | None, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: binding = cls._require_binding(binding) if not binding.agent_id or payload.agent_soul is None: @@ -1175,8 +1310,9 @@ class AgentComposerService: agent_soul=payload.agent_soul, operation=AgentConfigRevisionOperation.SAVE_NEW_VERSION, version_note=payload.version_note, + session=session, ) - agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id) + agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(payload.agent_soul) agent.active_config_is_published = True @@ -1198,6 +1334,7 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding | None, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: if payload.agent_soul is None: raise ValueError("agent_soul is required") @@ -1214,6 +1351,7 @@ class AgentComposerService: agent_soul=payload.agent_soul, operation=AgentConfigRevisionOperation.SAVE_NEW_AGENT, version_note=payload.version_note, + session=session, ) node_job = payload.node_job or WorkflowNodeJobConfig() if not binding: @@ -1225,13 +1363,13 @@ class AgentComposerService: node_id=node_id, created_by=account_id, ) - db.session.add(binding) + session.add(binding) binding.binding_type = WorkflowAgentBindingType.ROSTER_AGENT binding.agent_id = agent.id binding.current_snapshot_id = agent.active_config_snapshot_id binding.node_job_config = node_job binding.updated_by = account_id - db.session.flush() + session.flush() return binding @classmethod @@ -1242,13 +1380,15 @@ class AgentComposerService: account_id: str, binding: WorkflowAgentNodeBinding | None, payload: ComposerSavePayload, + session: Session, ) -> WorkflowAgentNodeBinding: binding = cls._require_binding(binding) - source_agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id) + source_agent = cls._require_agent(tenant_id=tenant_id, agent_id=binding.agent_id, session=session) source_version = cls._require_version( tenant_id=tenant_id, agent_id=source_agent.id, version_id=binding.current_snapshot_id, + session=session, ) agent_soul = payload.agent_soul or AgentSoulConfig.model_validate(source_version.config_snapshot_dict) agent_name = payload.new_agent_name or source_agent.name @@ -1266,6 +1406,7 @@ class AgentComposerService: agent_soul=agent_soul, operation=AgentConfigRevisionOperation.SAVE_TO_ROSTER, version_note=payload.version_note, + session=session, ) cls._copy_agent_drive_rows( tenant_id=tenant_id, @@ -1274,6 +1415,7 @@ class AgentComposerService: account_id=account_id, agent_soul=agent_soul, node_job=payload.node_job or WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict), + session=session, ) binding.binding_type = WorkflowAgentBindingType.ROSTER_AGENT binding.agent_id = roster_agent.id @@ -1299,8 +1441,9 @@ class AgentComposerService: icon_type: Any | None = None, icon: str | None = None, icon_background: str | None = None, + session: Session, ) -> Agent: - backing_app = AgentRosterService(db.session).create_hidden_backing_app_for_workflow_agent( + backing_app = AgentRosterService(session).create_hidden_backing_app_for_workflow_agent( tenant_id=tenant_id, account_id=account_id, name=name or f"Workflow Agent {node_id}", @@ -1328,8 +1471,8 @@ class AgentComposerService: created_by=account_id, updated_by=account_id, ) - db.session.add(agent) - db.session.flush() + session.add(agent) + session.flush() version = cls._create_config_version( tenant_id=tenant_id, agent_id=agent.id, @@ -1337,6 +1480,7 @@ class AgentComposerService: agent_soul=agent_soul, operation=AgentConfigRevisionOperation.CREATE_VERSION, version_note=None, + session=session, ) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(agent_soul) @@ -1353,6 +1497,7 @@ class AgentComposerService: account_id: str, agent_soul: AgentSoulConfig, node_job: WorkflowNodeJobConfig | None = None, + session: Session, ) -> None: exact_keys, prefixes = cls._drive_copy_scopes_from_agent_configs(agent_soul=agent_soul, node_job=node_job) predicates: list[ColumnElement[bool]] = [] @@ -1363,7 +1508,7 @@ class AgentComposerService: return source_rows = list( - db.session.scalars( + session.scalars( select(AgentDriveFile).where( AgentDriveFile.tenant_id == tenant_id, AgentDriveFile.agent_id == source_agent_id, @@ -1375,7 +1520,7 @@ class AgentComposerService: return existing_target_keys = set( - db.session.scalars( + session.scalars( select(AgentDriveFile.key).where( AgentDriveFile.tenant_id == tenant_id, AgentDriveFile.agent_id == target_agent_id, @@ -1386,7 +1531,7 @@ class AgentComposerService: for row in source_rows: if row.key in existing_target_keys: continue - db.session.add( + session.add( AgentDriveFile( tenant_id=tenant_id, agent_id=target_agent_id, @@ -1450,8 +1595,9 @@ class AgentComposerService: icon_type: AgentIconType | None = None, icon: str | None = None, icon_background: str | None = None, + session: Session, ) -> Agent: - account = cls._require_account(account_id=account_id) + account = cls._require_account(account_id=account_id, session=session) try: app = AppService().create_app( tenant_id, @@ -1465,12 +1611,13 @@ class AgentComposerService: icon_background=icon_background, ), account, + session=session, ) except IntegrityError as exc: - db.session.rollback() + session.rollback() raise AgentNameConflictError() from exc - agent = AgentRosterService(db.session).get_app_backing_agent(tenant_id=tenant_id, app_id=app.id) + agent = AgentRosterService(session).get_app_backing_agent(tenant_id=tenant_id, app_id=app.id) if agent is None: raise AgentNotFoundError() @@ -1478,6 +1625,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id, + session=session, ) version = cls._update_current_version( current_snapshot=current_snapshot, @@ -1485,6 +1633,7 @@ class AgentComposerService: agent_soul=agent_soul, operation=operation, version_note=version_note, + session=session, ) agent.active_config_snapshot_id = version.id agent.active_config_has_model = agent_soul_has_model(agent_soul) @@ -1503,9 +1652,10 @@ class AgentComposerService: operation: AgentConfigRevisionOperation, version_note: str | None, previous_snapshot_id: str | None = None, + session: Session, ) -> AgentConfigSnapshot: next_version = ( - db.session.scalar( + session.scalar( select(func.max(AgentConfigSnapshot.version)).where( AgentConfigSnapshot.tenant_id == tenant_id, AgentConfigSnapshot.agent_id == agent_id, @@ -1521,20 +1671,20 @@ class AgentComposerService: version_note=version_note, created_by=account_id, ) - db.session.add(version) - db.session.flush() + session.add(version) + session.flush() revision = AgentConfigRevision( tenant_id=tenant_id, agent_id=agent_id, previous_snapshot_id=previous_snapshot_id, current_snapshot_id=version.id, - revision=cls._next_revision(tenant_id=tenant_id, agent_id=agent_id), + revision=cls._next_revision(tenant_id=tenant_id, agent_id=agent_id, session=session), operation=operation, version_note=version_note, created_by=account_id, ) - db.session.add(revision) - db.session.flush() + session.add(revision) + session.flush() return version @classmethod @@ -1546,6 +1696,7 @@ class AgentComposerService: agent_soul: AgentSoulConfig, operation: AgentConfigRevisionOperation, version_note: str | None, + session: Session, ) -> AgentConfigSnapshot: return cls._create_config_version( tenant_id=current_snapshot.tenant_id, @@ -1555,12 +1706,13 @@ class AgentComposerService: operation=operation, version_note=version_note, previous_snapshot_id=current_snapshot.id, + session=session, ) @classmethod - def _next_revision(cls, *, tenant_id: str, agent_id: str) -> int: + def _next_revision(cls, *, tenant_id: str, agent_id: str, session: Session) -> int: return ( - db.session.scalar( + session.scalar( select(func.max(AgentConfigRevision.revision)).where( AgentConfigRevision.tenant_id == tenant_id, AgentConfigRevision.agent_id == agent_id, @@ -1570,8 +1722,8 @@ class AgentComposerService: ) + 1 @classmethod - def _get_agent_app_agent(cls, *, tenant_id: str, app_id: str) -> Agent | None: - return db.session.scalar( + def _get_agent_app_agent(cls, *, tenant_id: str, app_id: str, session: Session) -> Agent | None: + return session.scalar( select(Agent) .where( Agent.tenant_id == tenant_id, @@ -1585,8 +1737,8 @@ class AgentComposerService: ) @classmethod - def _require_agent_app_agent(cls, *, tenant_id: str, app_id: str) -> Agent: - agent = cls._get_agent_app_agent(tenant_id=tenant_id, app_id=app_id) + def _require_agent_app_agent(cls, *, tenant_id: str, app_id: str, session: Session) -> Agent: + agent = cls._get_agent_app_agent(tenant_id=tenant_id, app_id=app_id, session=session) if agent is None: raise AgentNotFoundError() return agent @@ -1599,6 +1751,7 @@ class AgentComposerService: agent_id: str, draft_type: AgentConfigDraftType, account_id: str | None, + session: Session, ) -> AgentConfigDraft | None: stmt = select(AgentConfigDraft).where( AgentConfigDraft.tenant_id == tenant_id, @@ -1609,7 +1762,7 @@ class AgentComposerService: stmt = stmt.where(AgentConfigDraft.account_id == account_id) else: stmt = stmt.where(AgentConfigDraft.account_id.is_(None)) - return db.session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1)) + return session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1)) @classmethod def _get_or_create_agent_draft( @@ -1620,12 +1773,14 @@ class AgentComposerService: draft_type: AgentConfigDraftType, account_id: str | None, created_by: str | None, + session: Session, ) -> AgentConfigDraft: draft = cls._get_agent_draft( tenant_id=tenant_id, agent_id=agent.id, draft_type=draft_type, account_id=account_id, + session=session, ) if draft is not None: return draft @@ -1633,6 +1788,7 @@ class AgentComposerService: tenant_id=tenant_id, agent_id=agent.id, version_id=agent.active_config_snapshot_id, + session=session, ) agent_soul = ( AgentSoulConfig.model_validate(base_snapshot.config_snapshot_dict) @@ -1650,8 +1806,8 @@ class AgentComposerService: created_by=created_by, updated_by=created_by, ) - db.session.add(draft) - db.session.flush() + session.add(draft) + session.flush() return draft @classmethod @@ -1665,6 +1821,7 @@ class AgentComposerService: agent_soul: AgentSoulConfig, account_id_for_audit: str, base_snapshot_id: str | None = None, + session: Session, ) -> AgentConfigDraft: draft = cls._get_or_create_agent_draft( tenant_id=tenant_id, @@ -1672,6 +1829,7 @@ class AgentComposerService: draft_type=draft_type, account_id=account_id, created_by=account_id_for_audit, + session=session, ) draft.config_snapshot = agent_soul if base_snapshot_id is not None: @@ -1681,7 +1839,7 @@ class AgentComposerService: draft.updated_by = account_id_for_audit if draft_type == AgentConfigDraftType.DRAFT and account_id is None: agent.active_config_is_published = False - db.session.flush() + session.flush() return draft @classmethod @@ -1709,8 +1867,8 @@ class AgentComposerService: } @classmethod - def _get_draft_workflow(cls, *, tenant_id: str, app_id: str) -> Workflow: - workflow = db.session.scalar( + def _get_draft_workflow(cls, *, tenant_id: str, app_id: str, session: Session) -> Workflow: + workflow = session.scalar( select(Workflow) .where( Workflow.tenant_id == tenant_id, @@ -1725,13 +1883,13 @@ class AgentComposerService: @classmethod def _get_workflow_binding( - cls, *, tenant_id: str, workflow_id: str, node_id: str + cls, *, tenant_id: str, workflow_id: str, node_id: str, session: Session ) -> WorkflowAgentNodeBinding | None: # Composer always operates against the draft workflow row, so this lookup # is scoped to ``workflow_version="draft"``. Published bindings are # materialized by WorkflowAgentPublishService.copy_agent_node_bindings_to_published # and are not edited through the Composer. - return db.session.scalar( + return session.scalar( select(WorkflowAgentNodeBinding) .where( WorkflowAgentNodeBinding.tenant_id == tenant_id, @@ -1749,32 +1907,34 @@ class AgentComposerService: return binding @classmethod - def _require_agent(cls, *, tenant_id: str, agent_id: str | None) -> Agent: + def _require_agent(cls, *, tenant_id: str, agent_id: str | None, session: Session) -> Agent: if not agent_id: raise AgentNotFoundError() - agent = db.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1)) + agent = session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1)) if not agent: raise AgentNotFoundError() return agent @classmethod - def _require_account(cls, *, account_id: str) -> Account: - account = db.session.get(Account, account_id) + def _require_account(cls, *, account_id: str, session: Session) -> Account: + account = session.get(Account, account_id) if not account: raise ValueError("Account not found") return account @classmethod - def _get_agent_if_present(cls, *, tenant_id: str, agent_id: str | None) -> Agent | None: + def _get_agent_if_present(cls, *, tenant_id: str, agent_id: str | None, session: Session) -> Agent | None: if not agent_id: return None - return db.session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1)) + return session.scalar(select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id).limit(1)) @classmethod - def _require_version(cls, *, tenant_id: str, agent_id: str | None, version_id: str | None) -> AgentConfigSnapshot: + def _require_version( + cls, *, tenant_id: str, agent_id: str | None, version_id: str | None, session: Session + ) -> AgentConfigSnapshot: if not agent_id or not version_id: raise AgentVersionNotFoundError() - version = db.session.scalar( + version = session.scalar( select(AgentConfigSnapshot) .where( AgentConfigSnapshot.tenant_id == tenant_id, @@ -1789,11 +1949,11 @@ class AgentComposerService: @classmethod def _get_version_if_present( - cls, *, tenant_id: str, agent_id: str | None, version_id: str | None + cls, *, tenant_id: str, agent_id: str | None, version_id: str | None, session: Session ) -> AgentConfigSnapshot | None: if not agent_id or not version_id: return None - return db.session.scalar( + return session.scalar( select(AgentConfigSnapshot) .where( AgentConfigSnapshot.tenant_id == tenant_id, @@ -1852,6 +2012,7 @@ class AgentComposerService: agent: Agent | None, version: AgentConfigSnapshot | None, account_id: str | None = None, + session: Session, ) -> dict[str, Any]: locked = bool(agent and agent.scope == AgentScope.ROSTER) save_options = [ComposerSaveStrategy.NODE_JOB_ONLY.value] @@ -1870,9 +2031,10 @@ class AgentComposerService: binding=binding, agent=agent, account_id=account_id, + session=session, ) debug_conversation_message_count = ( - AgentRosterService(db.session).count_agent_app_debug_conversation_messages( + AgentRosterService(session).count_agent_app_debug_conversation_messages( conversation_id=debug_conversation_id ) if debug_conversation_id @@ -1905,7 +2067,9 @@ class AgentComposerService: # this is the same list (so callers don't need to special-case). "effective_declared_outputs": cls._serialize_effective_outputs(cls._declared_outputs_from_binding(binding)), "save_options": save_options, - "impact_summary": cls.calculate_impact(tenant_id=binding.tenant_id, current_snapshot_id=version.id) + "impact_summary": cls.calculate_impact( + tenant_id=binding.tenant_id, current_snapshot_id=version.id, session=session + ) if version else None, "app_id": binding.app_id, @@ -1926,6 +2090,7 @@ class AgentComposerService: binding: WorkflowAgentNodeBinding, agent: Agent | None, account_id: str | None, + session: Session, ) -> str | None: if ( not account_id @@ -1937,7 +2102,7 @@ class AgentComposerService: from services.agent.roster_service import AgentRosterService - return AgentRosterService(db.session).get_or_create_agent_app_debug_conversation_id( + return AgentRosterService(session).get_or_create_agent_app_debug_conversation_id( tenant_id=tenant_id, agent_id=agent.id, account_id=account_id, diff --git a/api/services/agent/composer_validator.py b/api/services/agent/composer_validator.py index 4e7d4ff3c63..a227978b764 100644 --- a/api/services/agent/composer_validator.py +++ b/api/services/agent/composer_validator.py @@ -3,6 +3,7 @@ from typing import Any from pydantic import ValidationError +from models.agent_config_entities import AgentKnowledgeQueryMode from services.agent.errors import AgentSoulLockedError, InvalidComposerConfigError, PlaintextSecretNotAllowedError from services.agent.prompt_mentions import ( MAX_MENTIONS_PER_PROMPT, @@ -228,9 +229,40 @@ class ComposerConfigValidator: @classmethod def validate_agent_soul(cls, agent_soul: AgentSoulConfig) -> None: dumped = agent_soul.model_dump(mode="json") + cls._validate_knowledge_runtime_config(agent_soul) cls._reject_plaintext_secrets(dumped, path="agent_soul") cls._validate_shell_config(dumped) + @classmethod + def _validate_knowledge_runtime_config(cls, agent_soul: AgentSoulConfig) -> None: + """Validate knowledge settings that are required only for publish/run. + + Draft composer saves must be able to persist partially configured + knowledge sets while a user is still editing the panel. These checks + stay in the publish validator so invalid runtime configs are still + blocked before a version can be published or executed. + """ + for knowledge_set in agent_soul.knowledge.sets: + if ( + knowledge_set.query.mode == AgentKnowledgeQueryMode.USER_QUERY + and not (knowledge_set.query.value or "").strip() + ): + raise InvalidComposerConfigError("knowledge query.value is required for user_query mode") + + retrieval = knowledge_set.retrieval + if retrieval.mode == "multiple" and retrieval.top_k is None: + raise InvalidComposerConfigError("knowledge retrieval.top_k is required for multiple mode") + if retrieval.mode == "single" and retrieval.model is None: + raise InvalidComposerConfigError("knowledge retrieval.model is required for single mode") + + metadata_filtering = knowledge_set.metadata_filtering + if metadata_filtering.mode == "automatic" and metadata_filtering.metadata_model_config is None: + raise InvalidComposerConfigError("metadata_filtering.model_config is required for automatic mode") + if metadata_filtering.mode == "manual" and ( + metadata_filtering.conditions is None or not metadata_filtering.conditions.conditions + ): + raise InvalidComposerConfigError("metadata_filtering.conditions is required for manual mode") + @classmethod def validate_node_job(cls, node_job: WorkflowNodeJobConfig) -> None: cls._reject_plaintext_secrets(node_job.model_dump(mode="json"), path="node_job") diff --git a/api/services/agent/errors.py b/api/services/agent/errors.py index 6a1dc6fb628..163687815d8 100644 --- a/api/services/agent/errors.py +++ b/api/services/agent/errors.py @@ -1,5 +1,7 @@ from werkzeug.exceptions import BadRequest, Conflict, NotFound +from libs.exception import BaseHTTPException + class AgentNotFoundError(NotFound): description = "Agent not found." @@ -21,6 +23,12 @@ class AgentVersionConflictError(Conflict): description = "Agent config version changed. Please reload and try again." +class AgentModelNotConfiguredError(BaseHTTPException): + error_code = "agent_model_not_configured" + description = "Agent App requires the Agent Soul model to be configured." + code = 400 + + class AgentSoulLockedError(BadRequest): description = "Agent Soul is locked for this workflow node." diff --git a/api/services/agent/prompt_mentions.py b/api/services/agent/prompt_mentions.py index 15d257f474a..5f42bffe3ff 100644 --- a/api/services/agent/prompt_mentions.py +++ b/api/services/agent/prompt_mentions.py @@ -318,9 +318,10 @@ def _format_output_mention(output: DeclaredOutputConfig) -> str: if output.type == DeclaredOutputType.FILE: return ( f"{output.name} (file output; create the file locally, run " - f"`dify-agent file upload `, then copy the returned AgentStubFileMapping JSON " - f"as final_output.{output.name}; do not call final_output before upload succeeds, and do not use " - "the local path, filename, URL, or a synthesized dify-file-ref as the reference)" + f"`dify-agent file upload `, then set final_output.{output.name} to a `tool_file` mapping " + f"using the returned `reference`; if replying to the user in natural language, use the returned " + f"`download_url`; do not call final_output before upload succeeds, and do not use the local path, " + "filename, URL, or a synthesized dify-file-ref as the reference)" ) if ( output.type == DeclaredOutputType.ARRAY @@ -329,9 +330,10 @@ def _format_output_mention(output: DeclaredOutputConfig) -> str: ): return ( f"{output.name} (array[file] output; upload each produced file with " - f"`dify-agent file upload `, then copy the returned AgentStubFileMapping JSON objects " - f"as final_output.{output.name}; do not call final_output before all uploads succeed, and do not use " - "local paths, filenames, URLs, or synthesized dify-file-ref values as references)" + f"`dify-agent file upload `, then set final_output.{output.name} to `tool_file` mappings " + f"using the returned `reference` values; if replying to the user in natural language, use the returned " + f"`download_url`; do not call final_output before all uploads succeed, and do not use local paths, " + "filenames, URLs, or synthesized dify-file-ref values as references)" ) return f"{output.name} ({output.type.value})" diff --git a/api/services/agent/roster_service.py b/api/services/agent/roster_service.py index a34fca67105..33fd372398b 100644 --- a/api/services/agent/roster_service.py +++ b/api/services/agent/roster_service.py @@ -1,9 +1,12 @@ +import logging from typing import Any, TypedDict from sqlalchemy import and_, func, or_, select from sqlalchemy.exc import IntegrityError +from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from constants.model_template import default_app_templates +from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom from libs.datetime_utils import naive_utc_now from libs.helper import to_timestamp @@ -38,6 +41,9 @@ from services.app_service import AppService, CreateAppParams from services.enterprise.enterprise_service import EnterpriseService from services.entities.agent_entities import RosterAgentCreatePayload, RosterAgentUpdatePayload from services.feature_service import FeatureService +from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session + +logger = logging.getLogger(__name__) class AgentReferencingWorkflow(TypedDict): @@ -603,7 +609,15 @@ class AgentRosterService: def refresh_agent_app_debug_conversation_id( self, *, tenant_id: str, agent_id: str, account_id: str, commit: bool = True ) -> str: - """Start a new console debug conversation for the current Agent App editor.""" + """Start a new console debug conversation for the current Agent App editor. + + If this account already has a debug conversation mapping, the previous + conversation is abandoned first: any ACTIVE conversation-owned Agent + runtime sessions for that old conversation are sent through best-effort + backend cleanup and then retired locally even when enqueueing fails. + The debug mapping is then repointed to the freshly created + conversation. + """ agent = self._session.scalar( select(Agent).where( @@ -643,6 +657,16 @@ class AgentRosterService: ) ) else: + previous_app_id = mapping.app_id + previous_conversation_id = mapping.conversation_id + if previous_conversation_id: + self._cleanup_debug_conversation_runtime_sessions( + tenant_id=tenant_id, + agent_id=agent_id, + account_id=account_id, + app_id=previous_app_id or backing_app_id, + conversation_id=previous_conversation_id, + ) mapping.app_id = backing_app_id mapping.conversation_id = conversation_id self._session.flush() @@ -650,6 +674,84 @@ class AgentRosterService: self._session.commit() return conversation_id + def _cleanup_debug_conversation_runtime_sessions( + self, + *, + tenant_id: str, + agent_id: str, + account_id: str, + app_id: str, + conversation_id: str, + ) -> None: + session_store = AgentAppRuntimeSessionStore() + try: + stored_sessions = session_store.list_active_sessions_for_conversation( + tenant_id=tenant_id, + app_id=app_id, + conversation_id=conversation_id, + ) + except Exception: + logger.warning( + "Failed to load Agent App runtime sessions for debug conversation refresh: " + "tenant_id=%s app_id=%s conversation_id=%s", + tenant_id, + app_id, + conversation_id, + exc_info=True, + ) + return + + for stored_session in stored_sessions: + try: + if stored_session.runtime_layer_specs: + payload = AgentBackendSessionCleanupPayload( + session_snapshot=stored_session.session_snapshot, + runtime_layer_specs=stored_session.runtime_layer_specs, + idempotency_key=( + f"{tenant_id}:{agent_id}:{account_id}:{conversation_id}:debug-session-cleanup:" + f"{stored_session.scope.agent_id}:" + f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" + f"{stored_session.backend_run_id or 'no-run'}" + ), + metadata={ + "tenant_id": stored_session.scope.tenant_id, + "app_id": stored_session.scope.app_id, + "conversation_id": stored_session.scope.conversation_id, + "agent_id": stored_session.scope.agent_id, + "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, + "previous_agent_backend_run_id": stored_session.backend_run_id, + }, + ) + cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) + except Exception: + logger.warning( + "Failed to enqueue Agent backend cleanup for debug conversation refresh: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + stored_session.scope.tenant_id, + stored_session.scope.app_id, + stored_session.scope.conversation_id, + stored_session.scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) + finally: + try: + session_store.mark_cleaned( + scope=stored_session.scope, + backend_run_id=stored_session.backend_run_id, + ) + except Exception: + logger.warning( + "Failed to retire Agent App runtime session for debug conversation refresh: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + stored_session.scope.tenant_id, + stored_session.scope.app_id, + stored_session.scope.conversation_id, + stored_session.scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) + def load_or_create_agent_app_debug_conversation_ids_by_agent_id( self, *, tenant_id: str, agents: list[Agent], account_id: str ) -> dict[str, str]: @@ -826,6 +928,7 @@ class AgentRosterService: max_active_requests=source_app.max_active_requests, ), account, + session=self._session, ) target_app.enable_site = source_app.enable_site diff --git a/api/services/agent/skill_standardize_service.py b/api/services/agent/skill_standardize_service.py index cc2ba4b9bdc..2639f7a9a18 100644 --- a/api/services/agent/skill_standardize_service.py +++ b/api/services/agent/skill_standardize_service.py @@ -18,6 +18,8 @@ from __future__ import annotations import re from typing import Any +from sqlalchemy.orm import Session + from core.tools.tool_file_manager import ToolFileManager from services.agent.skill_package_service import SkillPackageService from services.agent_drive_service import AgentDriveService, DriveCommitItem, DriveFileRef, DriveSkillMetadata @@ -59,6 +61,7 @@ class SkillStandardizeService: tenant_id: str, user_id: str, agent_id: str, + session: Session, ) -> dict[str, Any]: """Create two ToolFiles, commit two drive-owned keys, and return skill metadata. @@ -113,6 +116,7 @@ class SkillStandardizeService: value_owned_by_drive=True, ), ], + session=session, ) self.last_committed_items = committed_items diff --git a/api/services/agent/skill_tool_inference_service.py b/api/services/agent/skill_tool_inference_service.py index a6d5e6b2de9..7ce53dd4666 100644 --- a/api/services/agent/skill_tool_inference_service.py +++ b/api/services/agent/skill_tool_inference_service.py @@ -19,6 +19,7 @@ from typing import Any import json_repair from pydantic import BaseModel, Field, ValidationError +from sqlalchemy.orm import Session from core.errors.error import ProviderTokenNotInitError from core.model_manager import ModelManager @@ -91,8 +92,8 @@ class SkillToolInferenceService: def __init__(self, *, drive_service: AgentDriveService | None = None) -> None: self._drive = drive_service or AgentDriveService() - def infer(self, *, tenant_id: str, agent_id: str, slug: str) -> dict[str, Any]: - skill_md = self._load_skill_md(tenant_id=tenant_id, agent_id=agent_id, slug=slug) + def infer(self, *, tenant_id: str, agent_id: str, slug: str, session: Session) -> dict[str, Any]: + skill_md = self._load_skill_md(tenant_id=tenant_id, agent_id=agent_id, slug=slug, session=session) user_prompt = f"SKILL.md of skill '{slug}':\n\n{skill_md}" @@ -115,9 +116,11 @@ class SkillToolInferenceService: tool.inferred_from = slug return result.model_dump(mode="json") - def _load_skill_md(self, *, tenant_id: str, agent_id: str, slug: str) -> str: + def _load_skill_md(self, *, tenant_id: str, agent_id: str, slug: str, session: Session) -> str: try: - preview = self._drive.preview(tenant_id=tenant_id, agent_id=agent_id, key=f"{slug}/SKILL.md") + preview = self._drive.preview( + tenant_id=tenant_id, agent_id=agent_id, key=f"{slug}/SKILL.md", session=session + ) except AgentDriveError as exc: if exc.code == "drive_key_not_found": raise SkillToolInferenceError( diff --git a/api/services/agent_app_feature_service.py b/api/services/agent_app_feature_service.py index 5fd794bb10f..d336cdf29de 100644 --- a/api/services/agent_app_feature_service.py +++ b/api/services/agent_app_feature_service.py @@ -13,7 +13,7 @@ from __future__ import annotations from typing import Any, cast -from sqlalchemy.orm import scoped_session +from sqlalchemy.orm import Session from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager from core.app.app_config.features.opening_statement.manager import OpeningStatementConfigManager @@ -74,7 +74,7 @@ class AgentAppFeatureConfigService: app_model: App, account: Account, config: dict[str, Any], - session: scoped_session, + session: Session, ) -> AppModelConfig: """Persist the presentation features as a new app_model_config version. diff --git a/api/services/agent_app_sandbox_service.py b/api/services/agent_app_sandbox_service.py index be301f5cd14..3f5a0bf41b2 100644 --- a/api/services/agent_app_sandbox_service.py +++ b/api/services/agent_app_sandbox_service.py @@ -3,22 +3,29 @@ These services keep product-facing locators (conversation, workflow run, node) on the API boundary and translate them into the agent backend's ``SandboxLocator`` using persisted non-sensitive runtime layer specs plus the -saved Agenton session snapshot. +saved Agenton session snapshot. Upload responses stay console-facing here: the +agent backend still returns a canonical ToolFile mapping, while this API layer +re-resolves that mapping into a signed browser download URL. """ from __future__ import annotations +import urllib.parse from collections.abc import Callable +from typing import Any from agenton.compositor import CompositorSessionSnapshot from dify_agent.client import Client from dify_agent.protocol import RuntimeLayerSpec, SandboxLocator, build_sandbox_locator_from_layer_specs from pydantic import BaseModel, TypeAdapter from sqlalchemy import select +from sqlalchemy.orm import Session from configs import dify_config from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore -from core.db.session_factory import session_factory +from core.app.file_access import DatabaseFileAccessController +from core.app.workflow.file_runtime import DifyWorkflowFileRuntime +from factories import file_factory from models.agent import AgentRuntimeSessionOwnerType, WorkflowAgentRuntimeSession, WorkflowAgentRuntimeSessionStatus _RUNTIME_LAYER_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) @@ -45,6 +52,12 @@ class AgentSandboxInfo(BaseModel): workspace_cwd: str +class AgentSandboxUploadDownload(BaseModel): + """Signed browser download URL for one sandbox upload result.""" + + url: str + + class AgentAppSandboxService: """Inspect and proxy file access for an Agent App conversation sandbox.""" @@ -77,9 +90,15 @@ class AgentAppSandboxService: locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) return self._client_factory().read_sandbox_file_sync(locator, path) - def upload_file(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str): + def upload_file( + self, *, tenant_id: str, app_id: str, conversation_id: str, path: str + ) -> AgentSandboxUploadDownload: locator = self._resolve_locator(tenant_id=tenant_id, app_id=app_id, conversation_id=conversation_id) - return self._client_factory().upload_sandbox_file_sync(locator, path) + uploaded = self._client_factory().upload_sandbox_file_sync(locator, path) + return _upload_download_response( + tenant_id=tenant_id, + file_mapping=uploaded.file.model_dump(mode="python"), + ) def _resolve_locator(self, *, tenant_id: str, app_id: str, conversation_id: str) -> SandboxLocator: stored = self._session_store.load_active_session_for_conversation( @@ -115,6 +134,7 @@ class WorkflowAgentSandboxService: node_id: str, node_execution_id: str | None, path: str, + session: Session, ): locator = self._resolve_locator( tenant_id=tenant_id, @@ -122,6 +142,7 @@ class WorkflowAgentSandboxService: workflow_run_id=workflow_run_id, node_id=node_id, node_execution_id=node_execution_id, + session=session, ) return self._client_factory().list_sandbox_files_sync(locator, path) @@ -134,6 +155,7 @@ class WorkflowAgentSandboxService: node_id: str, node_execution_id: str | None, path: str, + session: Session, ): locator = self._resolve_locator( tenant_id=tenant_id, @@ -141,6 +163,7 @@ class WorkflowAgentSandboxService: workflow_run_id=workflow_run_id, node_id=node_id, node_execution_id=node_execution_id, + session=session, ) return self._client_factory().read_sandbox_file_sync(locator, path) @@ -153,15 +176,21 @@ class WorkflowAgentSandboxService: node_id: str, node_execution_id: str | None, path: str, - ): + session: Session, + ) -> AgentSandboxUploadDownload: locator = self._resolve_locator( tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, node_id=node_id, node_execution_id=node_execution_id, + session=session, + ) + uploaded = self._client_factory().upload_sandbox_file_sync(locator, path) + return _upload_download_response( + tenant_id=tenant_id, + file_mapping=uploaded.file.model_dump(mode="python"), ) - return self._client_factory().upload_sandbox_file_sync(locator, path) def _resolve_locator( self, @@ -171,6 +200,7 @@ class WorkflowAgentSandboxService: workflow_run_id: str, node_id: str, node_execution_id: str | None, + session: Session, ) -> SandboxLocator: """Resolve one workflow Agent sandbox from product-facing identifiers. @@ -193,8 +223,7 @@ class WorkflowAgentSandboxService: stmt = stmt.where(WorkflowAgentRuntimeSession.node_execution_id == node_execution_id) stmt = stmt.order_by(WorkflowAgentRuntimeSession.updated_at.desc()).limit(1) - with session_factory.create_session() as session: - row = session.scalar(stmt) + row = session.scalar(stmt) if row is None: raise AgentSandboxInspectorError( @@ -246,6 +275,41 @@ def _deserialize_runtime_layer_specs(value: str | None) -> list[RuntimeLayerSpec return _RUNTIME_LAYER_SPECS_ADAPTER.validate_json(value) +def _upload_download_response(*, tenant_id: str, file_mapping: dict[str, Any]) -> AgentSandboxUploadDownload: + """Resolve one uploaded ToolFile mapping into a signed external download URL.""" + + controller = DatabaseFileAccessController() + runtime = DifyWorkflowFileRuntime(file_access_controller=controller) + try: + file = file_factory.build_from_mapping( + mapping=file_mapping, + tenant_id=tenant_id, + access_controller=controller, + ) + url = runtime.resolve_file_url(file=file, for_external=True) + except ValueError as exc: + raise AgentSandboxInspectorError( + "sandbox_upload_download_unavailable", + "uploaded sandbox file could not be converted to a download URL", + status_code=502, + ) from exc + + if not url: + raise AgentSandboxInspectorError( + "sandbox_upload_download_unavailable", + "uploaded sandbox file does not support download URL generation", + status_code=502, + ) + return AgentSandboxUploadDownload(url=_with_as_attachment(url)) + + +def _with_as_attachment(url: str) -> str: + parsed = urllib.parse.urlsplit(url) + query = urllib.parse.parse_qsl(parsed.query, keep_blank_values=True) + query.append(("as_attachment", "true")) + return urllib.parse.urlunsplit(parsed._replace(query=urllib.parse.urlencode(query))) + + def _default_client_factory() -> Client: base_url = dify_config.AGENT_BACKEND_BASE_URL if not base_url: @@ -257,4 +321,10 @@ def _default_client_factory() -> Client: return Client(base_url=base_url) -__all__ = ["AgentAppSandboxService", "AgentSandboxInfo", "AgentSandboxInspectorError", "WorkflowAgentSandboxService"] +__all__ = [ + "AgentAppSandboxService", + "AgentSandboxInfo", + "AgentSandboxInspectorError", + "AgentSandboxUploadDownload", + "WorkflowAgentSandboxService", +] diff --git a/api/services/agent_drive_service.py b/api/services/agent_drive_service.py index d79aa1abe9c..eb375f997d1 100644 --- a/api/services/agent_drive_service.py +++ b/api/services/agent_drive_service.py @@ -41,7 +41,6 @@ from sqlalchemy.orm import Session from configs import dify_config from core.app.file_access.controller import DatabaseFileAccessController -from core.db.session_factory import session_factory from extensions.ext_storage import storage from factories import file_factory from libs.uuid_utils import uuidv7 @@ -195,38 +194,38 @@ class AgentDriveService: *, tenant_id: str, agent_id: str, + session: Session, prefix: str = "", include_download_url: bool = False, ) -> list[dict[str, Any]]: - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - stmt = ( - select(AgentDriveFile) - .where(AgentDriveFile.tenant_id == tenant_id, AgentDriveFile.agent_id == agent_id) - .order_by(AgentDriveFile.key) - ) - if prefix: - stmt = stmt.where(AgentDriveFile.key.startswith(prefix)) - rows = list(session.scalars(stmt)) - items: list[dict[str, Any]] = [] - for row in rows: - item: dict[str, Any] = { - "key": row.key, - "size": row.size, - "hash": row.hash, - "mime_type": row.mime_type, - "file_kind": row.file_kind.value, - "file_id": row.file_id, - "is_skill": row.is_skill, - "skill_metadata": row.skill_metadata, - "created_at": int(row.created_at.timestamp()) if row.created_at else None, - } - if include_download_url: - item["download_url"] = self._resolve_download_url( - tenant_id=tenant_id, file_kind=row.file_kind, file_id=row.file_id - ) - items.append(item) - return items + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + stmt = ( + select(AgentDriveFile) + .where(AgentDriveFile.tenant_id == tenant_id, AgentDriveFile.agent_id == agent_id) + .order_by(AgentDriveFile.key) + ) + if prefix: + stmt = stmt.where(AgentDriveFile.key.startswith(prefix)) + rows = list(session.scalars(stmt)) + items: list[dict[str, Any]] = [] + for row in rows: + item: dict[str, Any] = { + "key": row.key, + "size": row.size, + "hash": row.hash, + "mime_type": row.mime_type, + "file_kind": row.file_kind.value, + "file_id": row.file_id, + "is_skill": row.is_skill, + "skill_metadata": row.skill_metadata, + "created_at": int(row.created_at.timestamp()) if row.created_at else None, + } + if include_download_url: + item["download_url"] = self._resolve_download_url( + tenant_id=tenant_id, file_kind=row.file_kind, file_id=row.file_id + ) + items.append(item) + return items def commit( self, @@ -235,25 +234,25 @@ class AgentDriveService: user_id: str, agent_id: str, items: list[DriveCommitItem], + session: Session, ) -> list[dict[str, Any]]: if not items: raise AgentDriveError("empty_commit", "commit requires at least one item", status_code=400) committed: list[dict[str, Any]] = [] pending_storage_deletes: list[str] = [] - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - for item in items: - committed.append( - self._commit_one( - session, - tenant_id=tenant_id, - user_id=user_id, - agent_id=agent_id, - item=item, - pending_storage_deletes=pending_storage_deletes, - ) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + for item in items: + committed.append( + self._commit_one( + session, + tenant_id=tenant_id, + user_id=user_id, + agent_id=agent_id, + item=item, + pending_storage_deletes=pending_storage_deletes, ) - session.commit() + ) + session.commit() for storage_key in pending_storage_deletes: self._delete_storage(storage_key) return committed @@ -263,6 +262,7 @@ class AgentDriveService: *, tenant_id: str, agent_id: str, + session: Session, prefix: str | None = None, key: str | None = None, ) -> list[str]: @@ -276,59 +276,57 @@ class AgentDriveService: raise AgentDriveError("invalid_delete_scope", "delete requires exactly one of prefix or key") removed_keys: list[str] = [] pending_storage_deletes: list[str] = [] - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - stmt = select(AgentDriveFile).where( - AgentDriveFile.tenant_id == tenant_id, - AgentDriveFile.agent_id == agent_id, - ) - if key is not None: - stmt = stmt.where(AgentDriveFile.key == normalize_drive_key(key)) - else: - stmt = stmt.where(AgentDriveFile.key.startswith(normalize_drive_key(prefix or ""))) - rows = list(session.scalars(stmt)) - for row in rows: - if row.value_owned_by_drive: - self._cleanup_value( - session, - tenant_id=tenant_id, - file_kind=row.file_kind, - file_id=row.file_id, - exclude_row_id=row.id, - pending_storage_deletes=pending_storage_deletes, - ) - removed_keys.append(row.key) - session.delete(row) - session.commit() + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + stmt = select(AgentDriveFile).where( + AgentDriveFile.tenant_id == tenant_id, + AgentDriveFile.agent_id == agent_id, + ) + if key is not None: + stmt = stmt.where(AgentDriveFile.key == normalize_drive_key(key)) + else: + stmt = stmt.where(AgentDriveFile.key.startswith(normalize_drive_key(prefix or ""))) + rows = list(session.scalars(stmt)) + for row in rows: + if row.value_owned_by_drive: + self._cleanup_value( + session, + tenant_id=tenant_id, + file_kind=row.file_kind, + file_id=row.file_id, + exclude_row_id=row.id, + pending_storage_deletes=pending_storage_deletes, + ) + removed_keys.append(row.key) + session.delete(row) + session.commit() for storage_key in pending_storage_deletes: self._delete_storage(storage_key) return removed_keys - def list_skills(self, *, tenant_id: str, agent_id: str) -> list[AgentDriveSkillInfo]: + def list_skills(self, *, tenant_id: str, agent_id: str, session: Session) -> list[AgentDriveSkillInfo]: """Return the drive-backed skill catalog derived from canonical ``SKILL.md`` rows.""" - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - skill_rows = list( - session.scalars( - select(AgentDriveFile) - .where( - AgentDriveFile.tenant_id == tenant_id, - AgentDriveFile.agent_id == agent_id, - AgentDriveFile.is_skill.is_(True), - ) - .order_by(AgentDriveFile.key) - ) - ) - archive_keys = set( - session.scalars( - select(AgentDriveFile.key).where( - AgentDriveFile.tenant_id == tenant_id, - AgentDriveFile.agent_id == agent_id, - AgentDriveFile.key.in_([self._skill_archive_key(row.key) for row in skill_rows]), - ) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + skill_rows = list( + session.scalars( + select(AgentDriveFile) + .where( + AgentDriveFile.tenant_id == tenant_id, + AgentDriveFile.agent_id == agent_id, + AgentDriveFile.is_skill.is_(True), + ) + .order_by(AgentDriveFile.key) + ) + ) + archive_keys = set( + session.scalars( + select(AgentDriveFile.key).where( + AgentDriveFile.tenant_id == tenant_id, + AgentDriveFile.agent_id == agent_id, + AgentDriveFile.key.in_([self._skill_archive_key(row.key) for row in skill_rows]), ) ) + ) skills: list[AgentDriveSkillInfo] = [] for row in skill_rows: @@ -349,14 +347,20 @@ class AgentDriveService: ) return skills - def inspect_skill(self, *, tenant_id: str, agent_id: str, skill_path: str) -> AgentDriveSkillInspectInfo: + def inspect_skill( + self, *, tenant_id: str, agent_id: str, skill_path: str, session: Session + ) -> AgentDriveSkillInspectInfo: """Return the UI-facing skill inspect view for slash-menu hover/detail.""" skill_path = normalize_drive_key(skill_path) skill_md_key = skill_path if skill_path.endswith(_SKILL_MD_SUFFIX) else f"{skill_path}{_SKILL_MD_SUFFIX}" skill_path = self._skill_path_from_key(skill_md_key) catalog = next( - (item for item in self.list_skills(tenant_id=tenant_id, agent_id=agent_id) if item["path"] == skill_path), + ( + item + for item in self.list_skills(tenant_id=tenant_id, agent_id=agent_id, session=session) + if item["path"] == skill_path + ), None, ) if catalog is None: @@ -366,10 +370,11 @@ class AgentDriveService: tenant_id=tenant_id, agent_id=agent_id, skill_md_key=skill_md_key, + session=session, ) - drive_items = self.manifest(tenant_id=tenant_id, agent_id=agent_id, prefix=f"{skill_path}/") + drive_items = self.manifest(tenant_id=tenant_id, agent_id=agent_id, prefix=f"{skill_path}/", session=session) drive_keys = {item["key"] for item in drive_items} - preview = self.preview(tenant_id=tenant_id, agent_id=agent_id, key=skill_md_key) + preview = self.preview(tenant_id=tenant_id, agent_id=agent_id, key=skill_md_key, session=session) files, warnings = self._skill_file_entries( skill_path=skill_path, skill_md_key=skill_md_key, @@ -582,23 +587,24 @@ class AgentDriveService: ) from exc @staticmethod - def _manifest_files_from_skill_metadata(*, tenant_id: str, agent_id: str, skill_md_key: str) -> list[str] | None: - with session_factory.create_session() as session: - row = session.scalar( - select(AgentDriveFile).where( - AgentDriveFile.tenant_id == tenant_id, - AgentDriveFile.agent_id == agent_id, - AgentDriveFile.key == skill_md_key, - AgentDriveFile.is_skill.is_(True), - ) + def _manifest_files_from_skill_metadata( + *, tenant_id: str, agent_id: str, skill_md_key: str, session: Session + ) -> list[str] | None: + row = session.scalar( + select(AgentDriveFile).where( + AgentDriveFile.tenant_id == tenant_id, + AgentDriveFile.agent_id == agent_id, + AgentDriveFile.key == skill_md_key, + AgentDriveFile.is_skill.is_(True), ) - if row is None: - return None - try: - metadata = AgentDriveService._parse_skill_metadata(row.key, row.skill_metadata) - except Exception: - logger.warning("drive skill inspect: malformed skill metadata for %s", skill_md_key, exc_info=True) - return None + ) + if row is None: + return None + try: + metadata = AgentDriveService._parse_skill_metadata(row.key, row.skill_metadata) + except Exception: + logger.warning("drive skill inspect: malformed skill metadata for %s", skill_md_key, exc_info=True) + return None return [str(item) for item in (metadata.manifest_files or []) if str(item).strip()] or None @classmethod @@ -932,15 +938,15 @@ class AgentDriveService: archive_file_kind: AgentDriveFileKind, archive_file_id: str, member_path: str, + session: Session, ) -> bytes: member_path = normalize_drive_key(member_path) - with session_factory.create_session() as session: - storage_key = self._storage_key_for_ref( - session, - tenant_id=tenant_id, - file_kind=archive_file_kind, - file_id=archive_file_id, - ) + storage_key = self._storage_key_for_ref( + session, + tenant_id=tenant_id, + file_kind=archive_file_kind, + file_id=archive_file_id, + ) archive_bytes = b"".join(storage.load_stream(storage_key)) try: with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: @@ -978,26 +984,25 @@ class AgentDriveService: return {"key": key, "size": size, "truncated": truncated, "binary": True, "text": None} return {"key": key, "size": size, "truncated": truncated, "binary": False, "text": text} - def preview(self, *, tenant_id: str, agent_id: str, key: str) -> dict[str, Any]: + def preview(self, *, tenant_id: str, agent_id: str, key: str, session: Session) -> dict[str, Any]: """Truncated text preview of one drive value (binary-safe, never 500s on size).""" - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - try: - row = self._require_row(session, tenant_id=tenant_id, agent_id=agent_id, key=key) - storage_key = self._storage_key_for_row(session, tenant_id=tenant_id, row=row) - size = row.size - response_key = row.key - archive_ref: tuple[AgentDriveFile, str] | None = None - except AgentDriveError: - archive_ref = self._archive_member_for_key( - session, - tenant_id=tenant_id, - agent_id=agent_id, - key=key, - ) - storage_key = None - size = None - response_key = normalize_drive_key(key) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + try: + row = self._require_row(session, tenant_id=tenant_id, agent_id=agent_id, key=key) + storage_key = self._storage_key_for_row(session, tenant_id=tenant_id, row=row) + size = row.size + response_key = row.key + archive_ref: tuple[AgentDriveFile, str] | None = None + except AgentDriveError: + archive_ref = self._archive_member_for_key( + session, + tenant_id=tenant_id, + agent_id=agent_id, + key=key, + ) + storage_key = None + size = None + response_key = normalize_drive_key(key) if archive_ref is not None: archive_row, member_path = archive_ref @@ -1006,6 +1011,7 @@ class AgentDriveService: archive_file_kind=archive_row.file_kind, archive_file_id=archive_row.file_id, member_path=member_path, + session=session, ) return self._preview_bytes(key=response_key, size=len(payload), payload=payload) @@ -1026,47 +1032,47 @@ class AgentDriveService: archive_file_kind: AgentDriveFileKind, archive_file_id: str, member_path: str, + session: Session, ) -> dict[str, Any]: - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) payload = self._load_archive_member_bytes( tenant_id=tenant_id, archive_file_kind=archive_file_kind, archive_file_id=archive_file_id, member_path=member_path, + session=session, ) return self._preview_bytes(key=normalize_drive_key(key), size=len(payload), payload=payload) - def download_url(self, *, tenant_id: str, agent_id: str, key: str) -> str: + def download_url(self, *, tenant_id: str, agent_id: str, key: str, session: Session) -> str: """External signed URL for a browser download of one drive value.""" - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) - try: - row = self._require_row(session, tenant_id=tenant_id, agent_id=agent_id, key=key) - except AgentDriveError: - archive_row, member_path = self._archive_member_for_key( - session, - tenant_id=tenant_id, - agent_id=agent_id, - key=key, - ) - return self.sign_archive_member_url( - tenant_id=tenant_id, - agent_id=agent_id, - key=key, - archive_file_kind=archive_row.file_kind, - archive_file_id=archive_row.file_id, - member_path=member_path, - for_external=True, - as_attachment=True, - ) - url = self._resolve_download_url( + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + try: + row = self._require_row(session, tenant_id=tenant_id, agent_id=agent_id, key=key) + except AgentDriveError: + archive_row, member_path = self._archive_member_for_key( + session, tenant_id=tenant_id, - file_kind=row.file_kind, - file_id=row.file_id, + agent_id=agent_id, + key=key, + ) + return self.sign_archive_member_url( + tenant_id=tenant_id, + agent_id=agent_id, + key=key, + archive_file_kind=archive_row.file_kind, + archive_file_id=archive_row.file_id, + member_path=member_path, for_external=True, as_attachment=True, ) + url = self._resolve_download_url( + tenant_id=tenant_id, + file_kind=row.file_kind, + file_id=row.file_id, + for_external=True, + as_attachment=True, + ) if url is None: raise AgentDriveError("drive_key_not_found", "drive value cannot be resolved", status_code=404) return url @@ -1080,10 +1086,10 @@ class AgentDriveService: archive_file_kind: AgentDriveFileKind, archive_file_id: str, member_path: str, + session: Session, for_external: bool = True, ) -> str: - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) return self.sign_archive_member_url( tenant_id=tenant_id, agent_id=agent_id, @@ -1211,14 +1217,15 @@ class AgentDriveService: archive_file_kind: AgentDriveFileKind, archive_file_id: str, member_path: str, + session: Session, ) -> tuple[bytes, str, str]: - with session_factory.create_session() as session: - self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) + self._assert_agent_belongs_to_tenant(session, tenant_id=tenant_id, agent_id=agent_id) payload = self._load_archive_member_bytes( tenant_id=tenant_id, archive_file_kind=archive_file_kind, archive_file_id=archive_file_id, member_path=member_path, + session=session, ) mime_type = mimetypes.guess_type(member_path)[0] or "application/octet-stream" filename = normalize_drive_key(key).rsplit("/", 1)[-1] diff --git a/api/services/agent_service.py b/api/services/agent_service.py index d8f4e11e758..a201eeb0485 100644 --- a/api/services/agent_service.py +++ b/api/services/agent_service.py @@ -3,13 +3,13 @@ from typing import Any import pytz from sqlalchemy import select +from sqlalchemy.orm import Session import contexts from core.app.app_config.easy_ui_based_app.agent.manager import AgentConfigManager from core.plugin.impl.agent import PluginAgentClient from core.plugin.impl.exc import PluginDaemonClientSideError from core.tools.tool_manager import ToolManager -from extensions.ext_database import db from libs.login import current_user from models import Account from models.model import App, Conversation, EndUser, Message @@ -17,14 +17,14 @@ from models.model import App, Conversation, EndUser, Message class AgentService: @classmethod - def get_agent_logs(cls, app_model: App, conversation_id: str, message_id: str): + def get_agent_logs(cls, app_model: App, conversation_id: str, message_id: str, session: Session): """ Service to get agent logs """ contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) - conversation: Conversation | None = db.session.scalar( + conversation: Conversation | None = session.scalar( select(Conversation) .where( Conversation.id == conversation_id, @@ -36,7 +36,7 @@ class AgentService: if not conversation: raise ValueError(f"Conversation not found: {conversation_id}") - message: Message | None = db.session.scalar( + message: Message | None = session.scalar( select(Message) .where( Message.id == message_id, @@ -52,9 +52,9 @@ class AgentService: if conversation.from_end_user_id: # only select name field - executor_name = db.session.scalar(select(EndUser.name).where(EndUser.id == conversation.from_end_user_id)) + executor_name = session.scalar(select(EndUser.name).where(EndUser.id == conversation.from_end_user_id)) else: - executor_name = db.session.scalar(select(Account.name).where(Account.id == conversation.from_account_id)) + executor_name = session.scalar(select(Account.name).where(Account.id == conversation.from_account_id)) executor = executor_name or "Unknown" assert isinstance(current_user, Account) diff --git a/api/services/agent_tool_inner_service.py b/api/services/agent_tool_inner_service.py index 4420f1b66b0..633ca893007 100644 --- a/api/services/agent_tool_inner_service.py +++ b/api/services/agent_tool_inner_service.py @@ -41,7 +41,7 @@ from services.errors.agent_tool_inner import AgentToolInnerServiceError class AgentToolInnerService: """Invoke one API-owned Agent tool declaration, including explicit plugin-via-core calls.""" - def invoke(self, session: Session, request: AgentToolInvokeRequest) -> AgentToolInvokeResponse: + def invoke(self, request: AgentToolInvokeRequest, *, session: Session) -> AgentToolInvokeResponse: app = session.get(App, request.caller.app_id) if app is None: raise AgentToolInnerServiceError( diff --git a/api/services/annotation_service.py b/api/services/annotation_service.py index 03e445a938b..ccca621aab5 100644 --- a/api/services/annotation_service.py +++ b/api/services/annotation_service.py @@ -4,12 +4,12 @@ from typing import TypedDict import pandas as pd from sqlalchemy import delete, or_, select, update -from sqlalchemy.orm import scoped_session +from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage from werkzeug.exceptions import NotFound from core.helper.csv_sanitizer import CSVSanitizer -from extensions.ext_database import db +from extensions.ext_database import db # noqa: F401 from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from libs.login import current_account_with_tenant @@ -91,7 +91,7 @@ class UpdateAnnotationSettingArgs(TypedDict): class AppAnnotationService: @staticmethod - def _get_annotation_by_ref(annotation_ref: AnnotationRef, session: scoped_session) -> MessageAnnotation | None: + def _get_annotation_by_ref(annotation_ref: AnnotationRef, session: Session) -> MessageAnnotation | None: return session.scalar( select(MessageAnnotation) .where( @@ -102,10 +102,12 @@ class AppAnnotationService: ) @classmethod - def up_insert_app_annotation_from_message(cls, args: UpsertAnnotationArgs, app_id: str) -> MessageAnnotation: + def up_insert_app_annotation_from_message( + cls, args: UpsertAnnotationArgs, app_id: str, *, session: Session + ) -> MessageAnnotation: # get app info current_user, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) @@ -119,9 +121,7 @@ class AppAnnotationService: raw_message_id = args.get("message_id") if raw_message_id: message_id = str(raw_message_id) - message = db.session.scalar( - select(Message).where(Message.id == message_id, Message.app_id == app.id).limit(1) - ) + message = session.scalar(select(Message).where(Message.id == message_id, Message.app_id == app.id).limit(1)) if not message: raise NotFound("Message Not Exists.") @@ -155,10 +155,10 @@ class AppAnnotationService: question=question, account_id=current_user.id, ) - db.session.add(annotation) - db.session.commit() + session.add(annotation) + session.commit() - annotation_setting = db.session.scalar( + annotation_setting = session.scalar( select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).limit(1) ) assert current_tenant_id is not None @@ -213,10 +213,10 @@ class AppAnnotationService: return {"job_id": job_id, "job_status": "waiting"} @classmethod - def get_annotation_list_by_app_id(cls, app_id: str, page: int, limit: int, keyword: str): + def get_annotation_list_by_app_id(cls, app_id: str, page: int, limit: int, keyword: str, *, session: Session): # get app info _, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) @@ -247,7 +247,7 @@ class AppAnnotationService: return annotations.items, annotations.total or 0 @classmethod - def export_annotation_list_by_app_id(cls, app_id: str): + def export_annotation_list_by_app_id(cls, app_id: str, *, session: Session): """ Export all annotations for an app with CSV injection protection. @@ -256,13 +256,13 @@ class AppAnnotationService: """ # get app info _, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) if not app: raise NotFound("App not found") - annotations = db.session.scalars( + annotations = session.scalars( select(MessageAnnotation) .where(MessageAnnotation.app_id == app_id) .order_by(MessageAnnotation.created_at.desc()) @@ -280,10 +280,12 @@ class AppAnnotationService: return annotations @classmethod - def insert_app_annotation_directly(cls, args: InsertAnnotationArgs, app_id: str) -> MessageAnnotation: + def insert_app_annotation_directly( + cls, args: InsertAnnotationArgs, app_id: str, *, session: Session + ) -> MessageAnnotation: # get app info current_user, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) @@ -297,10 +299,10 @@ class AppAnnotationService: annotation = MessageAnnotation( app_id=app.id, content=args["answer"], question=question, account_id=current_user.id ) - db.session.add(annotation) - db.session.commit() + session.add(annotation) + session.commit() # if annotation reply is enabled , add annotation to index - annotation_setting = db.session.scalar( + annotation_setting = session.scalar( select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).limit(1) ) if annotation_setting: @@ -315,7 +317,7 @@ class AppAnnotationService: @classmethod def update_app_annotation_directly( - cls, args: UpdateAnnotationArgs, annotation_ref: AnnotationRef, session: scoped_session + cls, args: UpdateAnnotationArgs, annotation_ref: AnnotationRef, session: Session ): annotation = cls._get_annotation_by_ref(annotation_ref, session) @@ -351,7 +353,7 @@ class AppAnnotationService: return annotation @classmethod - def delete_app_annotation(cls, annotation_ref: AnnotationRef, session: scoped_session): + def delete_app_annotation(cls, annotation_ref: AnnotationRef, session: Session): annotation = cls._get_annotation_by_ref(annotation_ref, session) if not annotation: @@ -384,9 +386,9 @@ class AppAnnotationService: ) @classmethod - def delete_app_annotations_in_batch(cls, app_ref: AppRef, annotation_ids: list[str]): + def delete_app_annotations_in_batch(cls, app_ref: AppRef, annotation_ids: list[str], *, session: Session): # Fetch annotations and their settings in a single query - annotations_to_delete = db.session.execute( + annotations_to_delete = session.execute( select(MessageAnnotation, AppAnnotationSetting) .outerjoin(AppAnnotationSetting, MessageAnnotation.app_id == AppAnnotationSetting.app_id) .where(MessageAnnotation.id.in_(annotation_ids), MessageAnnotation.app_id == app_ref.app_id) @@ -399,7 +401,7 @@ class AppAnnotationService: annotation_ids_to_delete = [annotation.id for annotation, _ in annotations_to_delete] # Step 2: Bulk delete hit histories in a single query - db.session.execute( + session.execute( delete(AppAnnotationHitHistory).where( AppAnnotationHitHistory.app_id == app_ref.app_id, AppAnnotationHitHistory.annotation_id.in_(annotation_ids_to_delete), @@ -414,7 +416,7 @@ class AppAnnotationService: ) # Step 4: Bulk delete annotations in a single query - delete_result = db.session.execute( + delete_result = session.execute( delete(MessageAnnotation).where( MessageAnnotation.id.in_(annotation_ids_to_delete), MessageAnnotation.app_id == app_ref.app_id, @@ -422,11 +424,11 @@ class AppAnnotationService: ) deleted_count = getattr(delete_result, "rowcount", 0) - db.session.commit() + session.commit() return {"deleted_count": deleted_count} @classmethod - def batch_import_app_annotations(cls, app_id: str, file: FileStorage): + def batch_import_app_annotations(cls, app_id: str, file: FileStorage, *, session: Session): """ Batch import annotations from CSV file with enhanced security checks. @@ -441,7 +443,7 @@ class AppAnnotationService: # get app info current_user, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) @@ -560,8 +562,8 @@ class AppAnnotationService: return {"job_id": job_id, "job_status": "waiting", "record_count": len(result)} @classmethod - def get_annotation_hit_histories(cls, annotation_ref: AnnotationRef, page, limit): - annotation = cls._get_annotation_by_ref(annotation_ref, db.session) + def get_annotation_hit_histories(cls, annotation_ref: AnnotationRef, page, limit, *, session: Session): + annotation = cls._get_annotation_by_ref(annotation_ref, session) if not annotation: raise NotFound("Annotation not found") @@ -578,8 +580,8 @@ class AppAnnotationService: return annotation_hit_histories.items, annotation_hit_histories.total or 0 @classmethod - def get_annotation_by_id(cls, annotation_id: str) -> MessageAnnotation | None: - annotation = db.session.get(MessageAnnotation, annotation_id) + def get_annotation_by_id(cls, annotation_id: str, *, session: Session) -> MessageAnnotation | None: + annotation = session.get(MessageAnnotation, annotation_id) if not annotation: return None @@ -597,9 +599,11 @@ class AppAnnotationService: message_id: str, from_source: str, score: float, - ): + *, + session: Session, + ) -> None: # add hit count to annotation - db.session.execute( + session.execute( update(MessageAnnotation) .where(MessageAnnotation.id == annotation_id) .values(hit_count=MessageAnnotation.hit_count + 1) @@ -616,21 +620,23 @@ class AppAnnotationService: annotation_question=annotation_question, annotation_content=annotation_content, ) - db.session.add(annotation_hit_history) - db.session.commit() + session.add(annotation_hit_history) + session.commit() @classmethod - def get_app_annotation_setting_by_app_id(cls, app_id: str) -> AnnotationSettingDict | AnnotationSettingDisabledDict: + def get_app_annotation_setting_by_app_id( + cls, app_id: str, *, session: Session + ) -> AnnotationSettingDict | AnnotationSettingDisabledDict: _, current_tenant_id = current_account_with_tenant() # get app info - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) if not app: raise NotFound("App not found") - annotation_setting = db.session.scalar( + annotation_setting = session.scalar( select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).limit(1) ) if annotation_setting: @@ -656,18 +662,18 @@ class AppAnnotationService: @classmethod def update_app_annotation_setting( - cls, app_id: str, annotation_setting_id: str, args: UpdateAnnotationSettingArgs + cls, app_id: str, annotation_setting_id: str, args: UpdateAnnotationSettingArgs, *, session: Session ) -> AnnotationSettingDict: current_user, current_tenant_id = current_account_with_tenant() # get app info - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) if not app: raise NotFound("App not found") - annotation_setting = db.session.scalar( + annotation_setting = session.scalar( select(AppAnnotationSetting) .where( AppAnnotationSetting.app_id == app_id, @@ -680,8 +686,8 @@ class AppAnnotationService: annotation_setting.score_threshold = args["score_threshold"] annotation_setting.updated_user_id = current_user.id annotation_setting.updated_at = naive_utc_now() - db.session.add(annotation_setting) - db.session.commit() + session.add(annotation_setting) + session.commit() collection_binding_detail = annotation_setting.collection_binding_detail @@ -704,9 +710,9 @@ class AppAnnotationService: } @classmethod - def clear_all_annotations(cls, app_id: str): + def clear_all_annotations(cls, app_id: str, *, session: Session): _, current_tenant_id = current_account_with_tenant() - app = db.session.scalar( + app = session.scalar( select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1) ) @@ -714,19 +720,19 @@ class AppAnnotationService: raise NotFound("App not found") # if annotation reply is enabled, delete annotation index - app_annotation_setting = db.session.scalar( + app_annotation_setting = session.scalar( select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).limit(1) ) - annotations_iter = db.session.scalars( + annotations_iter = session.scalars( select(MessageAnnotation).where(MessageAnnotation.app_id == app_id) ).yield_per(100) for annotation in annotations_iter: - hit_histories_iter = db.session.scalars( + hit_histories_iter = session.scalars( select(AppAnnotationHitHistory).where(AppAnnotationHitHistory.annotation_id == annotation.id) ).yield_per(100) for annotation_hit_history in hit_histories_iter: - db.session.delete(annotation_hit_history) + session.delete(annotation_hit_history) # if annotation reply is enabled, delete annotation index if app_annotation_setting: @@ -734,7 +740,7 @@ class AppAnnotationService: annotation.id, app_id, current_tenant_id, app_annotation_setting.collection_binding_id ) - db.session.delete(annotation) + session.delete(annotation) - db.session.commit() + session.commit() return {"result": "success"} diff --git a/api/services/api_based_extension_service.py b/api/services/api_based_extension_service.py index 25f554b6bdc..e855780d6a1 100644 --- a/api/services/api_based_extension_service.py +++ b/api/services/api_based_extension_service.py @@ -8,7 +8,7 @@ from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint class APIBasedExtensionService: @staticmethod - def get_all_by_tenant_id(session: Session, tenant_id: str) -> list[APIBasedExtension]: + def get_all_by_tenant_id(tenant_id: str, *, session: Session) -> list[APIBasedExtension]: extension_list = list( session.scalars( select(APIBasedExtension) @@ -23,7 +23,7 @@ class APIBasedExtensionService: return extension_list @classmethod - def save(cls, session: Session, extension_data: APIBasedExtension) -> APIBasedExtension: + def save(cls, extension_data: APIBasedExtension, *, session: Session) -> APIBasedExtension: cls._validation(session, extension_data) extension_data.api_key = encrypt_token(extension_data.tenant_id, extension_data.api_key) @@ -33,12 +33,12 @@ class APIBasedExtensionService: return extension_data @staticmethod - def delete(session: Session, extension_data: APIBasedExtension): + def delete(extension_data: APIBasedExtension, *, session: Session): session.delete(extension_data) session.commit() @staticmethod - def get_with_tenant_id(session: Session, tenant_id: str, api_based_extension_id: str) -> APIBasedExtension: + def get_with_tenant_id(tenant_id: str, api_based_extension_id: str, *, session: Session) -> APIBasedExtension: extension = session.scalar( select(APIBasedExtension) .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) diff --git a/api/services/app_dsl_service.py b/api/services/app_dsl_service.py index 52e936bf1ee..e8c12586856 100644 --- a/api/services/app_dsl_service.py +++ b/api/services/app_dsl_service.py @@ -39,6 +39,7 @@ from libs.datetime_utils import naive_utc_now from models import Account, App, AppMode from models.model import AppModelConfig, AppModelConfigDict, IconType from models.workflow import Workflow +from services.dsl_content import DSL_MAX_SIZE, dsl_content_size from services.dsl_version import check_version_compatibility from services.entities.dsl_entities import CheckDependenciesResult, ImportMode, ImportStatus from services.errors.app import WorkflowNotFoundError @@ -51,7 +52,6 @@ logger = logging.getLogger(__name__) IMPORT_INFO_REDIS_KEY_PREFIX = "app_import_info:" CHECK_DEPENDENCIES_REDIS_KEY_PREFIX = "app_check_dependencies:" IMPORT_INFO_REDIS_EXPIRY = 10 * 60 # 10 minutes -DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB CURRENT_DSL_VERSION = CURRENT_APP_DSL_VERSION @@ -131,15 +131,16 @@ class AppDslService: yaml_url = yaml_url.replace("/blob/", "/") response = remote_fetcher.make_request("GET", yaml_url.strip(), follow_redirects=True, timeout=(10, 10)) response.raise_for_status() - content = response.content.decode() + raw_content = response.content - if len(content) > DSL_MAX_SIZE: + if dsl_content_size(raw_content) > DSL_MAX_SIZE: return Import( id=import_id, status=ImportStatus.FAILED, error="File size exceeds the limit of 10MB", ) + content = raw_content.decode("utf-8") if not content: return Import( id=import_id, @@ -160,6 +161,12 @@ class AppDslService: error="yaml_content is required when import_mode is yaml-content", ) content = yaml_content + if dsl_content_size(content) > DSL_MAX_SIZE: + return Import( + id=import_id, + status=ImportStatus.FAILED, + error="File size exceeds the limit of 10MB", + ) # Process YAML content try: @@ -467,7 +474,7 @@ class AppDslService: ] workflow_service = WorkflowService() - current_draft_workflow = workflow_service.get_draft_workflow(app_model=app) + current_draft_workflow = workflow_service.get_draft_workflow(app_model=app, session=self._session) if current_draft_workflow: unique_hash = current_draft_workflow.unique_hash else: @@ -493,6 +500,7 @@ class AppDslService: account=account, environment_variables=environment_variables, conversation_variables=conversation_variables, + session=self._session, ) case AppMode.CHAT | AppMode.AGENT_CHAT | AppMode.COMPLETION: # Initialize model config @@ -514,7 +522,14 @@ class AppDslService: return app @classmethod - def export_dsl(cls, app_model: App, include_secret: bool = False, workflow_id: str | None = None) -> str: + def export_dsl( + cls, + app_model: App, + *, + session: Session, + include_secret: bool = False, + workflow_id: str | None = None, + ) -> str: """ Export app :param app_model: App instance @@ -541,7 +556,11 @@ class AppDslService: if app_mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: cls._append_workflow_export_data( - export_data=export_data, app_model=app_model, include_secret=include_secret, workflow_id=workflow_id + export_data=export_data, + app_model=app_model, + include_secret=include_secret, + workflow_id=workflow_id, + session=session, ) else: cls._append_model_config_export_data(export_data, app_model) @@ -550,7 +569,13 @@ class AppDslService: @classmethod def _append_workflow_export_data( - cls, *, export_data: dict[str, Any], app_model: App, include_secret: bool, workflow_id: str | None = None + cls, + *, + export_data: dict[str, Any], + app_model: App, + include_secret: bool, + session: Session, + workflow_id: str | None = None, ): """ Append workflow export data @@ -558,7 +583,7 @@ class AppDslService: :param app_model: App instance """ workflow_service = WorkflowService() - workflow = workflow_service.get_draft_workflow(app_model, workflow_id) + workflow = workflow_service.get_draft_workflow(app_model, workflow_id, session=session) if not workflow: raise WorkflowNotFoundError("Missing draft workflow configuration, please check.") diff --git a/api/services/app_generate_service.py b/api/services/app_generate_service.py index 3e2c3c96403..7724555b615 100644 --- a/api/services/app_generate_service.py +++ b/api/services/app_generate_service.py @@ -40,35 +40,6 @@ if TYPE_CHECKING: class AppGenerateService: - @classmethod - @trace_span(AppGenerateHandler) - def generate_stateless_agent_app( - cls, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - ): - """Run build-chat finalization as a blocking, non-SSE Agent App action. - - This is the service entry point for the Agent build-chat finalize flow. - It applies the same tracing, quota, and rate-limit guardrails as normal - app generation, but invokes the Agent App generator in stateless mode: - the call waits synchronously for Agent backend completion, triggers only - the backend side effect, and does not create Dify chat/message records. - """ - return cls._run_with_guardrails( - app_model=app_model, - streaming=False, - action=lambda _rate_limit, _request_id: AgentAppGenerator().generate_stateless( - app_model=app_model, - user=user, - args=args, - invoke_from=invoke_from, - ), - ) - @staticmethod def _build_streaming_task_on_subscribe(start_task: Callable[[], None]) -> Callable[[], None]: """ @@ -120,11 +91,12 @@ class AppGenerateService: @trace_span(AppGenerateHandler) def generate( cls, - session: Session, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, + *, + session: Session, streaming: bool = True, root_node_id: str | None = None, ): @@ -141,13 +113,13 @@ class AppGenerateService: app_model=app_model, streaming=streaming, action=lambda rate_limit, request_id: cls._dispatch_generate( - session=session, app_model=app_model, user=user, args=args, invoke_from=invoke_from, streaming=streaming, root_node_id=root_node_id, + session=session, rate_limit=rate_limit, request_id=request_id, ), @@ -189,13 +161,13 @@ class AppGenerateService: def _dispatch_generate( cls, *, - session: Session, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: bool, root_node_id: str | None, + session: Session, rate_limit: RateLimit, request_id: str, ): @@ -251,7 +223,7 @@ class AppGenerateService: ) case AppMode.ADVANCED_CHAT: workflow_id = args.get("workflow_id") - workflow = cls._get_workflow(app_model, invoke_from, workflow_id) + workflow = cls._get_workflow(app_model, invoke_from, workflow_id, session=session) if streaming: # Streaming mode: subscribe to SSE and enqueue the execution on first subscriber @@ -308,7 +280,7 @@ class AppGenerateService: ) case AppMode.WORKFLOW: workflow_id = args.get("workflow_id") - workflow = cls._get_workflow(app_model, invoke_from, workflow_id) + workflow = cls._get_workflow(app_model, invoke_from, workflow_id, session=session) if streaming: with rate_limit_context(rate_limit, request_id): payload = AppExecutionParams.new( @@ -384,12 +356,21 @@ class AppGenerateService: return min(limits) if limits else 0 @classmethod - def generate_single_iteration(cls, app_model: App, user: Account, node_id: str, args: Any, streaming: bool = True): + def generate_single_iteration( + cls, + app_model: App, + user: Account, + node_id: str, + args: Any, + *, + session: Session, + streaming: bool = True, + ): match app_model.mode: case AppMode.COMPLETION | AppMode.CHAT | AppMode.AGENT_CHAT: raise ValueError(f"Invalid app mode {app_model.mode}") case AppMode.ADVANCED_CHAT: - workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER) + workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session) return AdvancedChatAppGenerator.convert_to_event_stream( AdvancedChatAppGenerator().single_iteration_generate( app_model=app_model, @@ -401,7 +382,7 @@ class AppGenerateService: ) ) case AppMode.WORKFLOW: - workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER) + workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session) return AdvancedChatAppGenerator.convert_to_event_stream( WorkflowAppGenerator().single_iteration_generate( app_model=app_model, @@ -419,13 +400,20 @@ class AppGenerateService: @classmethod def generate_single_loop( - cls, app_model: App, user: Account, node_id: str, args: LoopNodeRunPayload, streaming: bool = True + cls, + app_model: App, + user: Account, + node_id: str, + args: LoopNodeRunPayload, + *, + session: Session, + streaming: bool = True, ): match app_model.mode: case AppMode.COMPLETION | AppMode.CHAT | AppMode.AGENT_CHAT: raise ValueError(f"Invalid app mode {app_model.mode}") case AppMode.ADVANCED_CHAT: - workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER) + workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session) return AdvancedChatAppGenerator.convert_to_event_stream( AdvancedChatAppGenerator().single_loop_generate( app_model=app_model, @@ -437,7 +425,7 @@ class AppGenerateService: ) ) case AppMode.WORKFLOW: - workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER) + workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session) return AdvancedChatAppGenerator.convert_to_event_stream( WorkflowAppGenerator().single_loop_generate( app_model=app_model, @@ -456,11 +444,12 @@ class AppGenerateService: @classmethod def generate_more_like_this( cls, - session: Session, app_model: App, user: Account | EndUser, message_id: str, invoke_from: InvokeFrom, + *, + session: Session, streaming: bool = True, ) -> Mapping | Generator: """ @@ -482,7 +471,14 @@ class AppGenerateService: ) @classmethod - def _get_workflow(cls, app_model: App, invoke_from: InvokeFrom, workflow_id: str | None = None) -> Workflow: + def _get_workflow( + cls, + app_model: App, + invoke_from: InvokeFrom, + workflow_id: str | None = None, + *, + session: Session, + ) -> Workflow: """ Get workflow :param app_model: app model @@ -498,20 +494,22 @@ class AppGenerateService: _ = uuid.UUID(workflow_id) except ValueError: raise WorkflowIdFormatError(f"Invalid workflow_id format: '{workflow_id}'. ") - workflow = workflow_service.get_published_workflow_by_id(app_model=app_model, workflow_id=workflow_id) + workflow = workflow_service.get_published_workflow_by_id( + app_model=app_model, workflow_id=workflow_id, session=session + ) if not workflow: raise WorkflowNotFoundError(f"Workflow not found with id: {workflow_id}") return workflow if invoke_from == InvokeFrom.DEBUGGER: # fetch draft workflow by app_model - workflow = workflow_service.get_draft_workflow(app_model=app_model) + workflow = workflow_service.get_draft_workflow(app_model=app_model, session=session) if not workflow: raise ValueError("Workflow not initialized") else: # fetch published workflow by app_model - workflow = workflow_service.get_published_workflow(app_model=app_model) + workflow = workflow_service.get_published_workflow(app_model=app_model, session=session) if not workflow: raise ValueError("Workflow not published") diff --git a/api/services/app_service.py b/api/services/app_service.py index 08cd30974e3..139513e87ee 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -8,7 +8,7 @@ import sqlalchemy as sa from pydantic import BaseModel, Field from sqlalchemy import ColumnElement, select from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from configs import dify_config from constants.model_template import default_app_templates @@ -18,7 +18,7 @@ from core.model_manager import ModelManager from core.tools.tool_manager import ToolManager from core.tools.utils.configuration import ToolParameterConfigurationManager from events.app_event import app_was_created, app_was_deleted, app_was_updated -from extensions.ext_database import db +from extensions.ext_database import db # noqa: F401 from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel from libs.datetime_utils import naive_utc_now @@ -80,7 +80,7 @@ class CreateAppParams(BaseModel): class AppService: @staticmethod def _build_app_list_filters( - user_id: str, tenant_id: str, params: AppListBaseParams, session: scoped_session + user_id: str, tenant_id: str, params: AppListBaseParams, session: Session ) -> list[sa.ColumnElement[bool]]: filters = [App.tenant_id == tenant_id, App.is_universal == False] @@ -153,13 +153,7 @@ class AppService: }[sort_by] @staticmethod - def get_starred_app_ids( - session: Session | scoped_session, - *, - tenant_id: str, - account_id: str, - app_ids: Sequence[str], - ) -> set[str]: + def get_starred_app_ids(*, tenant_id: str, account_id: str, app_ids: Sequence[str], session: Session) -> set[str]: """Return app IDs starred by this account within the tenant.""" if not app_ids: return set() @@ -174,38 +168,24 @@ class AppService: return set(starred_app_ids) @staticmethod - def get_app_by_id( - session: Session | scoped_session, - app_id: str, - ) -> App | None: + def get_app_by_id(app_id: str, *, session: Session) -> App | None: return session.get(App, app_id) @staticmethod - def get_visible_app_by_id( - session: Session | scoped_session, - app_id: str, - ) -> App | None: + def get_visible_app_by_id(app_id: str, *, session: Session) -> App | None: app = session.get(App, app_id) if not app or app.status != "normal" or not is_openapi_visible(app): return None return app @staticmethod - def find_visible_apps_by_ids( - session: Session | scoped_session, - app_ids: Sequence[str], - ) -> list[App]: + def find_visible_apps_by_ids(app_ids: Sequence[str], *, session: Session) -> list[App]: if not app_ids: return [] return list(session.execute(apply_openapi_gate(select(App).where(App.id.in_(list(app_ids))))).scalars().all()) @staticmethod - def find_visible_apps_by_name( - session: Session | scoped_session, - *, - name: str, - tenant_id: str, - ) -> list[App]: + def find_visible_apps_by_name(*, name: str, tenant_id: str, session: Session) -> list[App]: return list( session.execute( apply_openapi_gate( @@ -219,7 +199,7 @@ class AppService: ) def get_paginate_apps( - self, user_id: str, tenant_id: str, params: AppListParams, session: scoped_session + self, user_id: str, tenant_id: str, params: AppListParams, session: Session ) -> PaginatedResult | None: """ Get app list with pagination, filters, and explicit sort order. @@ -238,14 +218,12 @@ class AppService: sa.select(App).where(*filters).order_by(order_by), page=params.page, per_page=params.limit, + session=session, ) app_ids = [str(app.id) for app in app_models.items] starred_app_ids = self.get_starred_app_ids( - db.session, - tenant_id=tenant_id, - account_id=user_id, - app_ids=app_ids, + tenant_id=tenant_id, account_id=user_id, app_ids=app_ids, session=session ) for app in app_models.items: app.is_starred = str(app.id) in starred_app_ids @@ -253,7 +231,7 @@ class AppService: return app_models def get_paginate_starred_apps( - self, user_id: str, tenant_id: str, params: StarredAppListParams, session: scoped_session + self, user_id: str, tenant_id: str, params: StarredAppListParams, session: Session ) -> PaginatedResult | None: """ Get apps starred by the current account with pagination, filters, and explicit sort order. @@ -277,6 +255,7 @@ class AppService: .order_by(order_by), page=params.page, per_page=params.limit, + session=session, ) for app in app_models.items: @@ -285,7 +264,7 @@ class AppService: return app_models @staticmethod - def star_app(session: Session, *, app: App, account_id: str) -> None: + def star_app(*, app: App, account_id: str, session: Session) -> None: """Create the account's app star if it does not already exist.""" existing_star = session.scalar( select(AppStar) @@ -302,7 +281,7 @@ class AppService: session.add(AppStar(tenant_id=app.tenant_id, app_id=app.id, account_id=account_id)) @staticmethod - def unstar_app(session: Session, *, app: App, account_id: str) -> None: + def unstar_app(*, app: App, account_id: str, session: Session) -> None: """Remove the account's app star if present.""" existing_star = session.scalar( select(AppStar) @@ -318,7 +297,7 @@ class AppService: session.delete(existing_star) - def create_app(self, tenant_id: str, params: CreateAppParams, account: Account) -> App: + def create_app(self, tenant_id: str, params: CreateAppParams, account: Account, *, session: Session) -> App: """ Create app :param tenant_id: tenant id @@ -397,15 +376,15 @@ class AppService: app.maintainer = account.id app.updated_by = account.id - db.session.add(app) - db.session.flush() + session.add(app) + session.flush() if default_model_config: app_model_config = AppModelConfig( **default_model_config, app_id=app.id, created_by=account.id, updated_by=account.id ) - db.session.add(app_model_config) - db.session.flush() + session.add(app_model_config) + session.flush() app.app_model_config_id = app_model_config.id elif app_mode == AppMode.AGENT: @@ -418,8 +397,8 @@ class AppService: # left unset so App.is_agent stays False (this is the new Agent App # type, not a legacy function-call/react agent). agent_app_model_config = AppModelConfig(app_id=app.id, created_by=account.id, updated_by=account.id) - db.session.add(agent_app_model_config) - db.session.flush() + session.add(agent_app_model_config) + session.flush() app.app_model_config_id = agent_app_model_config.id @@ -431,7 +410,7 @@ class AppService: from services.agent.roster_service import AgentRosterService icon_type = AgentIconType(params.icon_type) if params.icon_type else None - AgentRosterService(db.session).create_backing_agent_for_app( + AgentRosterService(session).create_backing_agent_for_app( tenant_id=tenant_id, account_id=account.id, app_id=app.id, @@ -443,7 +422,7 @@ class AppService: icon_background=params.icon_background, ) - db.session.commit() + session.commit() app_was_created.send(app, account=account) enterprise_rbac_service.try_sync_creator_access_policy_member_bindings( @@ -542,10 +521,10 @@ class AppService: role: NotRequired[str | None] @staticmethod - def _get_backing_agent_for_update(app: App) -> Agent | None: + def _get_backing_agent_for_update(app: App, *, session: Session) -> Agent | None: if app.mode != AppMode.AGENT: return None - return db.session.scalar( + return session.scalar( select(Agent).where( Agent.tenant_id == app.tenant_id, Agent.app_id == app.id, @@ -574,6 +553,7 @@ class AppService: icon_background: str | None = None, account_id: str | None = None, updated_at: datetime | None = None, + session: Session, ) -> None: """Keep the Roster identity aligned with its Agent App shell. @@ -584,7 +564,7 @@ class AppService: Role omission is intentional: ``role=None`` preserves the backing Agent's current role, while ``role=""`` explicitly clears it. """ - agent = self._get_backing_agent_for_update(app) + agent = self._get_backing_agent_for_update(app, session=session) if agent is None: return @@ -605,16 +585,16 @@ class AppService: agent.updated_at = updated_at @staticmethod - def _commit_app_identity_update(app: App) -> None: + def _commit_app_identity_update(app: App, *, session: Session) -> None: try: - db.session.commit() + session.commit() except IntegrityError as exc: - db.session.rollback() + session.rollback() if app.mode == AppMode.AGENT: raise AgentNameConflictError() from exc raise - def update_app(self, app: App, args: ArgsDict) -> App: + def update_app(self, app: App, args: ArgsDict, *, session: Session) -> App: """ Update app :param app: App instance @@ -649,14 +629,15 @@ class AppService: icon_background=app.icon_background, account_id=current_user.id, updated_at=app.updated_at, + session=session, ) - self._commit_app_identity_update(app) + self._commit_app_identity_update(app, session=session) app_was_updated.send(app) return app - def update_app_name(self, app: App, name: str) -> App: + def update_app_name(self, app: App, name: str, *, session: Session) -> App: """ Update app name :param app: App instance @@ -672,15 +653,22 @@ class AppService: name=app.name, account_id=current_user.id, updated_at=app.updated_at, + session=session, ) - self._commit_app_identity_update(app) + self._commit_app_identity_update(app, session=session) app_was_updated.send(app) return app def update_app_icon( - self, app: App, icon: str, icon_background: str, icon_type: IconType | str | None = None + self, + app: App, + icon: str, + icon_background: str, + icon_type: IconType | str | None = None, + *, + session: Session, ) -> App: """ Update app icon @@ -704,14 +692,15 @@ class AppService: icon_background=app.icon_background, account_id=current_user.id, updated_at=app.updated_at, + session=session, ) - db.session.commit() + session.commit() app_was_updated.send(app) return app - def update_app_site_status(self, app: App, enable_site: bool) -> App: + def update_app_site_status(self, app: App, enable_site: bool, *, session: Session) -> App: """ Update app site status :param app: App instance @@ -724,13 +713,13 @@ class AppService: app.enable_site = enable_site app.updated_by = current_user.id app.updated_at = naive_utc_now() - db.session.commit() + session.commit() app_was_updated.send(app) return app - def update_app_api_status(self, app: App, enable_api: bool) -> App: + def update_app_api_status(self, app: App, enable_api: bool, *, session: Session) -> App: """ Update app api status :param app: App instance @@ -744,20 +733,20 @@ class AppService: app.enable_api = enable_api app.updated_by = current_user.id app.updated_at = naive_utc_now() - db.session.commit() + session.commit() app_was_updated.send(app) return app - def delete_app(self, app: App): + def delete_app(self, app: App, *, session: Session) -> None: """ Delete app :param app: App instance """ app_was_deleted.send(app) - backing_agent = self._get_backing_agent_for_update(app) + backing_agent = self._get_backing_agent_for_update(app, session=session) if backing_agent is not None: now = naive_utc_now() account_id = getattr(current_user, "id", None) @@ -767,8 +756,8 @@ class AppService: backing_agent.updated_by = account_id backing_agent.updated_at = now - db.session.delete(app) - db.session.commit() + session.delete(app) + session.commit() # clean up web app settings if FeatureService.get_system_features().webapp_auth.enabled: @@ -780,7 +769,7 @@ class AppService: # Trigger asynchronous deletion of app and related data remove_app_and_related_data_task.delay(tenant_id=app.tenant_id, app_id=app.id) - def get_app_meta(self, app_model: App): + def get_app_meta(self, app_model: App, *, session: Session): """ Get app meta info :param app_model: app model @@ -833,7 +822,7 @@ class AppService: meta["tool_icons"][tool_name] = url_prefix + provider_id + "/icon" elif provider_type == "api": try: - provider: ApiToolProvider | None = db.session.get(ApiToolProvider, provider_id) + provider: ApiToolProvider | None = session.get(ApiToolProvider, provider_id) if provider is None: raise ValueError(f"provider not found for tool {tool_name}") meta["tool_icons"][tool_name] = json.loads(provider.icon) @@ -843,25 +832,25 @@ class AppService: return meta @staticmethod - def get_app_code_by_id(app_id: str) -> str: + def get_app_code_by_id(app_id: str, *, session: Session) -> str: """ Get app code by app id :param app_id: app id :return: app code """ - site = db.session.scalar(select(Site).where(Site.app_id == app_id).limit(1)) + site = session.scalar(select(Site).where(Site.app_id == app_id).limit(1)) if not site: raise ValueError(f"App with id {app_id} not found") return str(site.code) @staticmethod - def get_app_id_by_code(app_code: str) -> str: + def get_app_id_by_code(app_code: str, *, session: Session) -> str: """ Get app id by app code :param app_code: app code :return: app id """ - site = db.session.scalar(select(Site).where(Site.code == app_code).limit(1)) + site = session.scalar(select(Site).where(Site.code == app_code).limit(1)) if not site: raise ValueError(f"App with code {app_code} not found") return str(site.app_id) diff --git a/api/services/async_workflow_service.py b/api/services/async_workflow_service.py index ceda30e950f..601cad7557a 100644 --- a/api/services/async_workflow_service.py +++ b/api/services/async_workflow_service.py @@ -51,7 +51,7 @@ class AsyncWorkflowService: @classmethod def trigger_workflow_async( - cls, session: Session, user: Account | EndUser, trigger_data: TriggerData + cls, user: Account | EndUser, trigger_data: TriggerData, *, session: Session ) -> AsyncTriggerResponse: """ Universal entry point for async workflow execution - THIS METHOD WILL NOT BLOCK @@ -187,7 +187,7 @@ class AsyncWorkflowService: @classmethod def reinvoke_trigger( - cls, session: Session, user: Account | EndUser, workflow_trigger_log_id: str + cls, user: Account | EndUser, workflow_trigger_log_id: str, *, session: Session ) -> AsyncTriggerResponse: """ Re-invoke a previously failed or rate-limited trigger - THIS METHOD WILL NOT BLOCK @@ -231,7 +231,7 @@ class AsyncWorkflowService: session.commit() # Re-trigger workflow (this will create a new trigger log) - return cls.trigger_workflow_async(session, user, trigger_data) + return cls.trigger_workflow_async(user, trigger_data, session=session) @classmethod def get_trigger_log( @@ -309,7 +309,8 @@ class AsyncWorkflowService: workflow_service: WorkflowService, app_model: App, workflow_id: str | None = None, - session: Session | None = None, + *, + session: Session, ) -> Workflow: """ Get workflow for the app @@ -317,9 +318,7 @@ class AsyncWorkflowService: Args: app_model: App model instance workflow_id: Optional specific workflow ID - session: Reuse this SQLAlchemy session for the lookup when provided, - so the caller's explicit session bears the connection cost - instead of Flask's request-scoped ``db.session``. + session: SQLAlchemy session used for the workflow lookup. Returns: Workflow instance diff --git a/api/services/audio_service.py b/api/services/audio_service.py index 86c56e60a13..52c71edd576 100644 --- a/api/services/audio_service.py +++ b/api/services/audio_service.py @@ -6,7 +6,7 @@ from typing import cast from flask import Response, stream_with_context from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage from constants import AUDIO_EXTENSIONS @@ -32,7 +32,7 @@ logger = logging.getLogger(__name__) class AudioService: @staticmethod - def _get_message_by_ref(session: Session | scoped_session, message_ref: MessageRef) -> Message | None: + def _get_message_by_ref(session: Session, message_ref: MessageRef) -> Message | None: stmt = select(Message).where(Message.id == message_ref.message_id, Message.app_id == message_ref.app_id) if message_ref.end_user_id is not None: stmt = stmt.where(Message.from_end_user_id == message_ref.end_user_id) @@ -89,7 +89,7 @@ class AudioService: cls, app_model: App, *, - session: Session | scoped_session, + session: Session, text: str | None = None, voice: str | None = None, end_user: str | None = None, diff --git a/api/services/auth/api_key_auth_service.py b/api/services/auth/api_key_auth_service.py index 42f1d4d8d40..f9ad7cf27b0 100644 --- a/api/services/auth/api_key_auth_service.py +++ b/api/services/auth/api_key_auth_service.py @@ -11,7 +11,7 @@ from services.auth.api_key_auth_factory import ApiKeyAuthFactory class ApiKeyAuthService: @staticmethod - def get_provider_auth_list(session: Session, tenant_id: str): + def get_provider_auth_list(tenant_id: str, *, session: Session): data_source_api_key_bindings = session.scalars( select(DataSourceApiKeyAuthBinding).where( DataSourceApiKeyAuthBinding.tenant_id == tenant_id, DataSourceApiKeyAuthBinding.disabled.is_(False) @@ -20,7 +20,7 @@ class ApiKeyAuthService: return data_source_api_key_bindings @staticmethod - def create_provider_auth(session: Session, tenant_id: str, args: dict[str, Any]): + def create_provider_auth(tenant_id: str, args: dict[str, Any], *, session: Session): auth_result = ApiKeyAuthFactory(args["provider"], args["credentials"]).validate_credentials() if auth_result: # Encrypt the api key @@ -35,7 +35,7 @@ class ApiKeyAuthService: session.commit() @staticmethod - def get_auth_credentials(session: Session, tenant_id: str, category: str, provider: str): + def get_auth_credentials(tenant_id: str, category: str, provider: str, *, session: Session): data_source_api_key_bindings = session.scalar( select(DataSourceApiKeyAuthBinding).where( DataSourceApiKeyAuthBinding.tenant_id == tenant_id, @@ -52,7 +52,7 @@ class ApiKeyAuthService: return credentials @staticmethod - def delete_provider_auth(session: Session, tenant_id: str, binding_id: str): + def delete_provider_auth(tenant_id: str, binding_id: str, *, session: Session): data_source_api_key_binding = session.scalar( select(DataSourceApiKeyAuthBinding).where( DataSourceApiKeyAuthBinding.tenant_id == tenant_id, diff --git a/api/services/billing_service.py b/api/services/billing_service.py index ec391e51676..4a829590f16 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -7,7 +7,7 @@ from typing import Any, Literal, NotRequired, TypedDict import httpx from pydantic import TypeAdapter from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed from werkzeug.exceptions import InternalServerError @@ -50,9 +50,26 @@ class QuotaReleaseResult(TypedDict): released: int +class QuotaBalanceResult(TypedDict): + available: int + reserved: int + quota: int + usage: int + + +class QuotaConsumeCappedResult(TypedDict): + deducted: int + available: int + reserved: int + quota: int + usage: int + + _quota_reserve_adapter = TypeAdapter(QuotaReserveResult) _quota_commit_adapter = TypeAdapter(QuotaCommitResult) _quota_release_adapter = TypeAdapter(QuotaReleaseResult) +_quota_balance_adapter = TypeAdapter(QuotaBalanceResult) +_quota_consume_capped_adapter = TypeAdapter(QuotaConsumeCappedResult) class _TenantFeatureQuota(TypedDict): @@ -176,6 +193,7 @@ class DismissNotificationDict(TypedDict): class BillingService: base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") + quota_base_url = os.environ.get("BILLING_QUOTA_API_URL") or base_url secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 4, 60) @@ -215,12 +233,18 @@ class BillingService: def get_quota_info(cls, tenant_id: str) -> TenantFeatureQuotaInfo: params = {"tenant_id": tenant_id} return _tenant_feature_quota_info_adapter.validate_python( - cls._send_request("GET", "/quota/info", params=params) + cls._send_quota_request("GET", "/quota/info", params=params) ) @classmethod def quota_reserve( - cls, tenant_id: str, feature_key: str, request_id: str, amount: int = 1, meta: dict | None = None + cls, + tenant_id: str, + feature_key: str, + request_id: str, + amount: int = 1, + meta: dict | None = None, + bucket: str = "", ) -> QuotaReserveResult: """Reserve quota before task execution.""" payload: dict = { @@ -229,13 +253,21 @@ class BillingService: "request_id": request_id, "amount": amount, } + if bucket: + payload["bucket"] = bucket if meta: payload["meta"] = meta - return _quota_reserve_adapter.validate_python(cls._send_request("POST", "/quota/reserve", json=payload)) + return _quota_reserve_adapter.validate_python(cls._send_quota_request("POST", "/quota/reserve", json=payload)) @classmethod def quota_commit( - cls, tenant_id: str, feature_key: str, reservation_id: str, actual_amount: int, meta: dict | None = None + cls, + tenant_id: str, + feature_key: str, + reservation_id: str, + actual_amount: int, + meta: dict | None = None, + bucket: str = "", ) -> QuotaCommitResult: """Commit a reservation with actual consumption.""" payload: dict = { @@ -244,23 +276,57 @@ class BillingService: "reservation_id": reservation_id, "actual_amount": actual_amount, } + if bucket: + payload["bucket"] = bucket if meta: payload["meta"] = meta - return _quota_commit_adapter.validate_python(cls._send_request("POST", "/quota/commit", json=payload)) + return _quota_commit_adapter.validate_python(cls._send_quota_request("POST", "/quota/commit", json=payload)) @classmethod - def quota_release(cls, tenant_id: str, feature_key: str, reservation_id: str) -> QuotaReleaseResult: + def quota_release( + cls, tenant_id: str, feature_key: str, reservation_id: str, bucket: str = "" + ) -> QuotaReleaseResult: """Release a reservation (cancel, return frozen quota).""" - return _quota_release_adapter.validate_python( - cls._send_request( - "POST", - "/quota/release", - json={ - "tenant_id": tenant_id, - "feature_key": feature_key, - "reservation_id": reservation_id, - }, - ) + payload = { + "tenant_id": tenant_id, + "feature_key": feature_key, + "reservation_id": reservation_id, + } + if bucket: + payload["bucket"] = bucket + return _quota_release_adapter.validate_python(cls._send_quota_request("POST", "/quota/release", json=payload)) + + @classmethod + def quota_get_balance(cls, tenant_id: str, feature_key: str, bucket: str = "") -> QuotaBalanceResult: + """Get quota balance for a feature bucket.""" + params = {"tenant_id": tenant_id, "feature_key": feature_key} + if bucket: + params["bucket"] = bucket + return _quota_balance_adapter.validate_python(cls._send_quota_request("GET", "/quota/balance", params=params)) + + @classmethod + def quota_consume_capped( + cls, + tenant_id: str, + feature_key: str, + request_id: str, + amount: int, + meta: dict | None = None, + bucket: str = "", + ) -> QuotaConsumeCappedResult: + """Consume up to the available quota and return the actual deducted amount.""" + payload: dict = { + "tenant_id": tenant_id, + "feature_key": feature_key, + "request_id": request_id, + "amount": amount, + } + if bucket: + payload["bucket"] = bucket + if meta: + payload["meta"] = meta + return _quota_consume_capped_adapter.validate_python( + cls._send_quota_request("POST", "/quota/consume-capped", json=payload) ) @classmethod @@ -334,6 +400,12 @@ class BillingService: params = {"tenant_id": tenant_id, "feature_key": feature_key} return cls._send_request("GET", "/billing/tenant_feature_plan/usage", params=params) + @classmethod + def _send_quota_request( + cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None + ): + return cls._send_request(method, endpoint, json=json, params=params, base_url=cls.quota_base_url) + @classmethod @retry( wait=wait_fixed(2), @@ -341,10 +413,17 @@ class BillingService: retry=retry_if_exception_type(httpx.RequestError), reraise=True, ) - def _send_request(cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None): + def _send_request( + cls, + method: Literal["GET", "POST", "DELETE", "PUT"], + endpoint: str, + json=None, + params=None, + base_url: str | None = None, + ): headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key} - url = f"{cls.base_url}{endpoint}" + url = f"{base_url or cls.base_url}{endpoint}" response = _http_client.request(method, url, json=json, params=params, headers=headers, follow_redirects=True) if method == "GET" and response.status_code != httpx.codes.OK: raise ValueError("Unable to retrieve billing information. Please try again later or contact support.") @@ -363,7 +442,7 @@ class BillingService: return response.json() @staticmethod - def is_tenant_owner_or_admin(session: Session | scoped_session, current_user: Account): + def is_tenant_owner_or_admin(current_user: Account, *, session: Session): tenant_id = current_user.current_tenant_id join: TenantAccountJoin | None = session.scalar( diff --git a/api/services/conversation_service.py b/api/services/conversation_service.py index 557ae8e89f3..6ab9c5a8a2b 100644 --- a/api/services/conversation_service.py +++ b/api/services/conversation_service.py @@ -6,18 +6,17 @@ from typing import Any from sqlalchemy import asc, desc, func, or_, select from sqlalchemy.orm import Session +from clients.agent_backend import AgentBackendSessionCleanupPayload from configs import dify_config +from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom -from core.db.session_factory import session_factory from core.llm_generator.llm_generator import LLMGenerator -from extensions.ext_database import db from factories import variable_factory from graphon.variables.types import SegmentType from libs.datetime_utils import naive_utc_now from libs.infinite_scroll_pagination import InfiniteScrollPagination from models import Account, ConversationVariable from models.model import App, Conversation, EndUser, Message -from services.conversation_variable_updater import ConversationVariableUpdater from services.errors.conversation import ( ConversationNotExistsError, ConversationVariableNotExistsError, @@ -25,6 +24,7 @@ from services.errors.conversation import ( LastConversationNotExistsError, ) from services.errors.message import MessageNotExistsError +from tasks.agent_backend_session_cleanup_task import cleanup_conversation_agent_runtime_session from tasks.delete_conversation_task import delete_conversation_related_data logger = logging.getLogger(__name__) @@ -122,24 +122,26 @@ class ConversationService: user: Account | EndUser | None, name: str | None, auto_generate: bool, + *, + session: Session, ): - conversation = cls.get_conversation(app_model, conversation_id, user) + conversation = cls.get_conversation(app_model, conversation_id, user, session=session) if auto_generate: - return cls.auto_generate_name(app_model, conversation) + return cls.auto_generate_name(app_model, conversation, session=session) else: if name is None: raise ValueError("name is required when auto_generate is false") conversation.name = name conversation.updated_at = naive_utc_now() - db.session.commit() + session.commit() return conversation @classmethod - def auto_generate_name(cls, app_model: App, conversation: Conversation): + def auto_generate_name(cls, app_model: App, conversation: Conversation, *, session: Session): # get conversation first message - message = db.session.scalar( + message = session.scalar( select(Message) .where(Message.app_id == app_model.id, Message.conversation_id == conversation.id) .order_by(Message.created_at.asc()) @@ -156,13 +158,15 @@ class ConversationService: ) conversation.name = name - db.session.commit() + session.commit() return conversation @classmethod - def get_conversation(cls, app_model: App, conversation_id: str, user: Account | EndUser | None): - conversation = db.session.scalar( + def get_conversation( + cls, app_model: App, conversation_id: str, user: Account | EndUser | None, *, session: Session + ): + conversation = session.scalar( select(Conversation) .where( Conversation.id == conversation_id, @@ -181,14 +185,27 @@ class ConversationService: return conversation @classmethod - def delete(cls, app_model: App, conversation_id: str, user: Account | EndUser | None): + def delete(cls, app_model: App, conversation_id: str, user: Account | EndUser | None, *, session: Session): """ Delete a conversation only if it belongs to the given user and app context. + Before removing the conversation row, this best-effort lifecycle path + enumerates any ACTIVE conversation-owned Agent backend runtime sessions, + enqueues asynchronous backend cleanup for rows with persisted runtime + layer specs, and then retires the local session rows even if enqueueing + fails. Conversation deletion and related-data cleanup scheduling still + proceed when that lifecycle bookkeeping only partially succeeds. + Raises: ConversationNotExistsError: When the conversation is not visible to the current user. """ - conversation = cls.get_conversation(app_model, conversation_id, user) + conversation = cls.get_conversation(app_model, conversation_id, user, session=session) + session_store = AgentAppRuntimeSessionStore() + stored_sessions = session_store.list_active_sessions_for_conversation( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + conversation_id=conversation.id, + ) try: logger.info( @@ -196,15 +213,66 @@ class ConversationService: app_model.name, conversation_id, ) + for stored_session in stored_sessions: + try: + if stored_session.runtime_layer_specs: + payload = AgentBackendSessionCleanupPayload( + session_snapshot=stored_session.session_snapshot, + runtime_layer_specs=stored_session.runtime_layer_specs, + idempotency_key=( + f"{stored_session.scope.tenant_id}:{stored_session.scope.app_id}:" + f"{stored_session.scope.conversation_id}:agent-runtime-session-cleanup:" + f"{stored_session.scope.agent_id}:" + f"{stored_session.scope.agent_config_snapshot_id or 'no-config'}:" + f"{stored_session.backend_run_id or 'no-run'}" + ), + metadata={ + "tenant_id": stored_session.scope.tenant_id, + "app_id": stored_session.scope.app_id, + "conversation_id": stored_session.scope.conversation_id, + "agent_id": stored_session.scope.agent_id, + "agent_config_snapshot_id": stored_session.scope.agent_config_snapshot_id, + "previous_agent_backend_run_id": stored_session.backend_run_id, + }, + ) + cleanup_conversation_agent_runtime_session.delay(payload.model_dump(mode="json")) + except Exception: + logger.warning( + "Failed to enqueue Agent backend cleanup for conversation deletion: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + stored_session.scope.tenant_id, + stored_session.scope.app_id, + stored_session.scope.conversation_id, + stored_session.scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) + finally: + try: + session_store.mark_cleaned( + scope=stored_session.scope, + backend_run_id=stored_session.backend_run_id, + ) + except Exception: + logger.warning( + "Failed to retire Agent App runtime session for conversation deletion: " + "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s backend_run_id=%s", + stored_session.scope.tenant_id, + stored_session.scope.app_id, + stored_session.scope.conversation_id, + stored_session.scope.agent_id, + stored_session.backend_run_id, + exc_info=True, + ) - db.session.delete(conversation) - db.session.commit() + session.delete(conversation) + session.commit() delete_conversation_related_data.delay(conversation.id) - except Exception as e: - db.session.rollback() - raise e + except Exception: + session.rollback() + raise @classmethod def get_conversational_variable( @@ -215,8 +283,10 @@ class ConversationService: limit: int, last_id: str | None, variable_name: str | None = None, + *, + session: Session, ) -> InfiniteScrollPagination: - conversation = cls.get_conversation(app_model, conversation_id, user) + conversation = cls.get_conversation(app_model, conversation_id, user, session=session) stmt = ( select(ConversationVariable) @@ -245,18 +315,17 @@ class ConversationService: ) ) - with session_factory.create_session() as session: - if last_id: - last_variable = session.scalar(stmt.where(ConversationVariable.id == last_id)) - if not last_variable: - raise ConversationVariableNotExistsError() + if last_id: + last_variable = session.scalar(stmt.where(ConversationVariable.id == last_id)) + if not last_variable: + raise ConversationVariableNotExistsError() - # Filter for variables created after the last_id - stmt = stmt.where(ConversationVariable.created_at > last_variable.created_at) + # Filter for variables created after the last_id + stmt = stmt.where(ConversationVariable.created_at > last_variable.created_at) - # Apply limit to query: fetch one extra row to determine has_more - query_stmt = stmt.limit(limit + 1) - rows = session.scalars(query_stmt).all() + # Apply limit to query: fetch one extra row to determine has_more + query_stmt = stmt.limit(limit + 1) + rows = session.scalars(query_stmt).all() has_more = False if len(rows) > limit: @@ -282,6 +351,8 @@ class ConversationService: variable_id: str, user: Account | EndUser | None, new_value: Any, + *, + session: Session, ): """ Update a conversation variable's value. @@ -302,7 +373,7 @@ class ConversationService: ConversationVariableTypeMismatchError: If the new value type doesn't match the variable's expected type """ # Verify conversation exists and user has access - conversation = cls.get_conversation(app_model, conversation_id, user) + conversation = cls.get_conversation(app_model, conversation_id, user, session=session) # Get the existing conversation variable stmt = ( @@ -312,48 +383,43 @@ class ConversationService: .where(ConversationVariable.id == variable_id) ) - with session_factory.create_session() as session: - existing_variable = session.scalar(stmt) - if not existing_variable: - raise ConversationVariableNotExistsError() + existing_variable = session.scalar(stmt) + if not existing_variable: + raise ConversationVariableNotExistsError() - # Convert existing variable to Variable object - current_variable = existing_variable.to_variable() + # Convert existing variable to Variable object + current_variable = existing_variable.to_variable() - # Validate that the new value type matches the expected variable type - expected_type = SegmentType(current_variable.value_type) + # Validate that the new value type matches the expected variable type + expected_type = SegmentType(current_variable.value_type) - # There is showing number in web ui but int in db - if expected_type == SegmentType.INTEGER: - expected_type = SegmentType.NUMBER + # There is showing number in web ui but int in db + if expected_type == SegmentType.INTEGER: + expected_type = SegmentType.NUMBER - if not expected_type.is_valid(new_value): - inferred_type = SegmentType.infer_segment_type(new_value) - raise ConversationVariableTypeMismatchError( - f"Type mismatch: variable '{current_variable.name}' expects {expected_type.value}, " - f"but got {inferred_type.value if inferred_type else 'unknown'} type" - ) + if not expected_type.is_valid(new_value): + inferred_type = SegmentType.infer_segment_type(new_value) + raise ConversationVariableTypeMismatchError( + f"Type mismatch: variable '{current_variable.name}' expects {expected_type.value}, " + f"but got {inferred_type.value if inferred_type else 'unknown'} type" + ) - # Create updated variable with new value only, preserving everything else - updated_variable_dict = { - "id": current_variable.id, - "name": current_variable.name, - "description": current_variable.description, - "value_type": current_variable.value_type, - "value": new_value, - "selector": current_variable.selector, - } + # Create updated variable with new value only, preserving everything else + updated_variable_dict = { + "id": current_variable.id, + "name": current_variable.name, + "description": current_variable.description, + "value_type": current_variable.value_type, + "value": new_value, + "selector": current_variable.selector, + } - updated_variable = variable_factory.build_conversation_variable_from_mapping(updated_variable_dict) + updated_variable = variable_factory.build_conversation_variable_from_mapping(updated_variable_dict) + existing_variable.data = updated_variable.model_dump_json() + session.commit() - # Use the conversation variable updater to persist the changes - updater = ConversationVariableUpdater(session_factory.get_session_maker()) - updater.update(conversation_id, updated_variable) - updater.flush() - - # Return the updated variable data - return { - "created_at": existing_variable.created_at, - "updated_at": naive_utc_now(), # Update timestamp - **updated_variable.model_dump(), - } + return { + "created_at": existing_variable.created_at, + "updated_at": naive_utc_now(), # Update timestamp + **updated_variable.model_dump(), + } diff --git a/api/services/credential_permission_service.py b/api/services/credential_permission_service.py index 2b1082d132b..d9ce5e7c502 100644 --- a/api/services/credential_permission_service.py +++ b/api/services/credential_permission_service.py @@ -1,7 +1,7 @@ from collections.abc import Sequence from sqlalchemy import or_, select -from sqlalchemy.orm import InstrumentedAttribute, Session, scoped_session +from sqlalchemy.orm import InstrumentedAttribute, Session from models.account import Account from models.credential_permission import CredentialPermission @@ -16,9 +16,7 @@ class CredentialPermissionService: """ @classmethod - def get_partial_member_list( - cls, session: Session | scoped_session, credential_id: str, credential_type: str - ) -> Sequence[str]: + def get_partial_member_list(cls, credential_id: str, credential_type: str, *, session: Session) -> Sequence[str]: """Return account_ids that have partial-member access to a credential.""" return session.scalars( select(CredentialPermission.account_id).where( diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index 94515309e79..837bf52c082 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -7,25 +7,57 @@ from piling up database transactions while preserving cross-tenant concurrency. import logging from collections.abc import Callable +from dataclasses import dataclass +from uuid import uuid4 from sqlalchemy import select from sqlalchemy.orm import Session from configs import dify_config -from core.db.session_factory import session_factory from core.errors.error import QuotaExceededError -from extensions.ext_database import db from extensions.ext_redis import redis_client from models import TenantCreditPool from models.enums import ProviderQuotaType logger = logging.getLogger(__name__) +FEATURE_KEY_CREDIT_POOL = "credit_pool" CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS = 10 CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS = 5 +@dataclass(frozen=True) +class CreditPoolBalance: + tenant_id: str + pool_type: str + quota_limit: int + quota_used: int + + @property + def remaining_credits(self) -> int: + if self.quota_limit == -1: + return -1 + return max(0, self.quota_limit - self.quota_used) + + def has_sufficient_credits(self, required_credits: int) -> bool: + return self.quota_limit == -1 or self.remaining_credits >= required_credits + + class CreditPoolService: + @staticmethod + def _normalize_pool_type(pool_type: str | ProviderQuotaType) -> str: + return pool_type.value if isinstance(pool_type, ProviderQuotaType) else str(pool_type) + + @staticmethod + def _use_billing_quota() -> bool: + return bool(dify_config.BILLING_ENABLED) + + @staticmethod + def _require_session(session: Session | None) -> Session: + if session is None: + raise ValueError("session is required when billing quota is disabled") + return session + @staticmethod def _get_tenant_lock_key(tenant_id: str) -> str: return f"credit_pool:tenant:{tenant_id}:deduct_lock" @@ -66,7 +98,7 @@ class CreditPoolService: ) @classmethod - def create_default_pool(cls, tenant_id: str) -> TenantCreditPool: + def create_default_pool(cls, tenant_id: str, session: Session) -> TenantCreditPool: """create default credit pool for new tenant""" credit_pool = TenantCreditPool( tenant_id=tenant_id, @@ -74,61 +106,133 @@ class CreditPoolService: quota_used=0, pool_type=ProviderQuotaType.TRIAL, ) - db.session.add(credit_pool) - db.session.commit() + session.add(credit_pool) + session.commit() return credit_pool @classmethod - def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None: + def get_pool( + cls, + tenant_id: str, + pool_type: str | ProviderQuotaType = "trial", + *, + session: Session | None = None, + ) -> TenantCreditPool | CreditPoolBalance | None: """get tenant credit pool""" - with session_factory.get_session_maker().begin() as session: - return session.scalar( - select(TenantCreditPool) - .where( - TenantCreditPool.tenant_id == tenant_id, - TenantCreditPool.pool_type == pool_type, - ) - .limit(1) + normalized_pool_type = cls._normalize_pool_type(pool_type) + if cls._use_billing_quota(): + from services.billing_service import BillingService + + balance = BillingService.quota_get_balance( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, ) + return CreditPoolBalance( + tenant_id=tenant_id, + pool_type=normalized_pool_type, + quota_limit=balance["quota"], + quota_used=balance["usage"], + ) + + session = cls._require_session(session) + return session.scalar( + select(TenantCreditPool) + .where( + TenantCreditPool.tenant_id == tenant_id, + TenantCreditPool.pool_type == normalized_pool_type, + ) + .limit(1) + ) @classmethod def check_credits_available( cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", + *, + session: Session | None = None, ) -> bool: """check if credits are available without deducting""" - pool = cls.get_pool(tenant_id, pool_type) + pool = cls.get_pool(tenant_id, pool_type, session=session) if not pool: return False - return pool.remaining_credits >= credits_required + return pool.has_sufficient_credits(credits_required) @classmethod def check_and_deduct_credits( cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", + *, + session: Session | None = None, ) -> int: """Deduct exactly the requested credits or raise without mutating the pool.""" if credits_required <= 0: return 0 + normalized_pool_type = cls._normalize_pool_type(pool_type) + + if cls._use_billing_quota(): + from services.billing_service import BillingService + + request_id = str(uuid4()) + result = BillingService.quota_reserve( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + request_id=request_id, + amount=credits_required, + meta={"source": "credit_pool.check_and_deduct"}, + ) + reservation_id = result.get("reservation_id", "") + if not reservation_id: + raise QuotaExceededError("Insufficient credits remaining") + try: + BillingService.quota_commit( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + reservation_id=reservation_id, + actual_amount=credits_required, + meta={"source": "credit_pool.check_and_deduct"}, + ) + except Exception: + try: + BillingService.quota_release( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + reservation_id=reservation_id, + ) + except Exception: + logger.warning( + "Failed to release reserved credit pool quota, tenant_id=%s, pool_type=%s, reservation_id=%s", + tenant_id, + normalized_pool_type, + reservation_id, + exc_info=True, + ) + raise + return credits_required + + session = cls._require_session(session) def deduct() -> int: - with session_factory.get_session_maker().begin() as session: - pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type) - if not pool: - raise QuotaExceededError("Credit pool not found") + pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type) + if not pool: + raise QuotaExceededError("Credit pool not found") - remaining_credits = pool.remaining_credits - if remaining_credits <= 0: - raise QuotaExceededError("No credits remaining") - if remaining_credits < credits_required: - raise QuotaExceededError("Insufficient credits remaining") + remaining_credits = pool.remaining_credits + if remaining_credits <= 0: + raise QuotaExceededError("No credits remaining") + if remaining_credits < credits_required: + raise QuotaExceededError("Insufficient credits remaining") - pool.quota_used += credits_required - return credits_required + pool.quota_used += credits_required + session.commit() + return credits_required try: return cls._deduct_with_tenant_lock(tenant_id, deduct) @@ -143,25 +247,43 @@ class CreditPoolService: cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", + *, + session: Session | None = None, ) -> int: """Deduct up to the available balance and return the actual deducted credits.""" if credits_required <= 0: return 0 + normalized_pool_type = cls._normalize_pool_type(pool_type) + + if cls._use_billing_quota(): + from services.billing_service import BillingService + + result = BillingService.quota_consume_capped( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + request_id=str(uuid4()), + amount=credits_required, + meta={"source": "credit_pool.deduct_capped"}, + ) + return result["deducted"] + + session = cls._require_session(session) def deduct() -> int: - with session_factory.get_session_maker().begin() as session: - pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type) - if not pool: - logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, pool_type) - return 0 + pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type) + if not pool: + logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, normalized_pool_type) + return 0 - deducted_credits = min(credits_required, pool.remaining_credits) - if deducted_credits <= 0: - return 0 + deducted_credits = min(credits_required, pool.remaining_credits) + if deducted_credits <= 0: + return 0 - pool.quota_used += deducted_credits - return deducted_credits + pool.quota_used += deducted_credits + session.commit() + return deducted_credits try: return cls._deduct_with_tenant_lock(tenant_id, deduct) diff --git a/api/services/data_migration/export_service.py b/api/services/data_migration/export_service.py index f5d214d230b..c0233006690 100644 --- a/api/services/data_migration/export_service.py +++ b/api/services/data_migration/export_service.py @@ -120,8 +120,8 @@ class MigrationExportService: self.package_service = package_service or MigrationPackageService() self.dependency_discovery_service = dependency_discovery_service or DependencyDiscoveryService() - def export(self, session: Session, selection: ExportSelection) -> ExportResult: - tenant = self._get_tenant(session, selection) + def export(self, selection: ExportSelection, *, session: Session) -> ExportResult: + tenant = self._get_tenant(selection, session=session) package = self.package_service.build_empty_package( source_tenant_id=tenant.id, source_tenant_name=tenant.name, @@ -131,10 +131,12 @@ class MigrationExportService: report_items: list[ResourceReportItem] = [] discovered_dependencies: list[DiscoveredDependency] = [] - apps = self._selected_apps(session, tenant.id, selection) + apps = self._selected_apps(tenant.id, selection, session=session) exported_app_ids = {app.id for app in apps} for app in apps: - dsl_content = AppDslService.export_dsl(app_model=app, include_secret=selection.include_secrets) + dsl_content = AppDslService.export_dsl( + app_model=app, session=session, include_secret=selection.include_secrets + ) package.workflows.append( { "id": app.id, @@ -157,7 +159,6 @@ class MigrationExportService: report_items=report_items, ) self._export_workflow_tools( - session, tenant, self._provider_ids( selection.additional_workflow_tools, discovered_dependencies, DependencyKind.WORKFLOW_TOOL @@ -166,9 +167,9 @@ class MigrationExportService: exported_workflow_tools=package.workflow_tools, dependencies=package.dependencies, report_items=report_items, + session=session, ) self._export_mcp_tools( - session, tenant_id=tenant.id, provider_ids=self._provider_ids( selection.additional_mcp_tools, @@ -179,6 +180,7 @@ class MigrationExportService: exported_mcp_tools=package.mcp_tools, dependencies=package.dependencies, report_items=report_items, + session=session, ) self._record_dependency_metadata( self._dependencies_by_kind(discovered_dependencies, DependencyKind.BUILTIN_OR_PLUGIN_TOOL), @@ -195,7 +197,7 @@ class MigrationExportService: ), ) - def _get_tenant(self, session: Session, selection: ExportSelection) -> Tenant: + def _get_tenant(self, selection: ExportSelection, *, session: Session) -> Tenant: if selection.source_tenant_id: tenant = session.get(Tenant, selection.source_tenant_id) if tenant is None: @@ -214,7 +216,7 @@ class MigrationExportService: ) return tenants[0] - def _selected_apps(self, session: Session, tenant_id: str, selection: ExportSelection) -> list[App]: + def _selected_apps(self, tenant_id: str, selection: ExportSelection, *, session: Session) -> list[App]: query = sa.select(App).where(App.tenant_id == tenant_id, App.mode.in_(SUPPORTED_APP_MODES)) if not selection.export_all_apps: if not selection.app_ids: @@ -267,7 +269,6 @@ class MigrationExportService: def _export_workflow_tools( self, - session: Session, tenant: Tenant, provider_ids: Iterable[str], *, @@ -275,11 +276,12 @@ class MigrationExportService: exported_workflow_tools: list[dict[str, Any]], dependencies: list[dict[str, Any]], report_items: list[ResourceReportItem], + session: Session, ) -> None: provider_ids = self._dedupe(provider_ids) if not provider_ids: return - owner = self._get_tenant_owner(session, tenant.id) + owner = self._get_tenant_owner(tenant.id, session=session) if owner is None: for provider_id in provider_ids: report_items.append( @@ -330,7 +332,7 @@ class MigrationExportService: ResourceReportItem(ResourceType.WORKFLOW_TOOL, provider_id, provider_id, "unresolved", str(exc)) ) - def _get_tenant_owner(self, session: Session, tenant_id: str) -> Account | None: + def _get_tenant_owner(self, tenant_id: str, *, session: Session) -> Account | None: return session.scalar( sa.select(Account) .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) @@ -341,7 +343,6 @@ class MigrationExportService: def _export_mcp_tools( self, - session: Session, *, tenant_id: str, provider_ids: Iterable[str], @@ -349,6 +350,7 @@ class MigrationExportService: exported_mcp_tools: list[dict[str, Any]], dependencies: list[dict[str, Any]], report_items: list[ResourceReportItem], + session: Session, ) -> None: for provider_id in self._dedupe(provider_ids): if not include_secrets: @@ -359,7 +361,7 @@ class MigrationExportService: ) continue try: - provider = self._get_mcp_provider(session, tenant_id, provider_id) + provider = self._get_mcp_provider(tenant_id, provider_id, session=session) exported_mcp_tools.append(self._serialize_mcp_provider(provider)) report_items.append(ResourceReportItem(ResourceType.MCP_TOOL, provider_id, provider.name, "exported")) except Exception as exc: @@ -367,7 +369,7 @@ class MigrationExportService: ResourceReportItem(ResourceType.MCP_TOOL, provider_id, provider_id, "unresolved", str(exc)) ) - def _get_mcp_provider(self, session: Session, tenant_id: str, provider_id: str) -> MCPToolProvider: + def _get_mcp_provider(self, tenant_id: str, provider_id: str, *, session: Session) -> MCPToolProvider: predicates = [MCPToolProvider.server_identifier == provider_id] if self._is_uuid_string(provider_id): predicates.append(MCPToolProvider.id == provider_id) diff --git a/api/services/data_migration/import_service.py b/api/services/data_migration/import_service.py index 3eb251bbaef..b3354413ba1 100644 --- a/api/services/data_migration/import_service.py +++ b/api/services/data_migration/import_service.py @@ -82,11 +82,11 @@ class ImportTargetResolver: "Target tenant must be provided by --target-tenant, import config, or package metadata." ) - def resolve(self, session: Session, request: ImportRequest) -> ImportTarget: + def resolve(self, request: ImportRequest, *, session: Session) -> ImportTarget: target_tenant_name = self.select_target_tenant_name(request) package_target = request.package.metadata.target_tenant or {} if request.cli_target_tenant or request.config_target_tenant: - tenant = self._resolve_tenant_by_id_or_name(session, target_tenant_name) + tenant = self._resolve_tenant_by_id_or_name(target_tenant_name, session=session) elif package_target.get("id") and self._is_uuid(package_target["id"]): tenant = session.get(Tenant, package_target["id"]) if tenant is not None and package_target.get("name") and tenant.name != package_target.get("name"): @@ -94,7 +94,7 @@ class ImportTargetResolver: f"Target tenant id/name mismatch: {package_target['id']} / {package_target['name']}" ) else: - tenant = self._resolve_tenant_by_id_or_name(session, target_tenant_name) + tenant = self._resolve_tenant_by_id_or_name(target_tenant_name, session=session) if tenant is None: raise MigrationDataError(f"Target tenant not found: {target_tenant_name}") @@ -123,7 +123,7 @@ class ImportTargetResolver: operator_email=account.email, ) - def _resolve_tenant_by_id_or_name(self, session: Session, value: str) -> Tenant | None: + def _resolve_tenant_by_id_or_name(self, value: str, *, session: Session) -> Tenant | None: if self._is_uuid(value): tenant = session.get(Tenant, value) if tenant is not None: @@ -149,8 +149,8 @@ class MigrationImportService: def __init__(self, *, target_resolver: ImportTargetResolver | None = None) -> None: self.target_resolver = target_resolver or ImportTargetResolver() - def import_package(self, session: Session, request: ImportRequest) -> ImportResult: - target = self.target_resolver.resolve(session, request) + def import_package(self, request: ImportRequest, *, session: Session) -> ImportResult: + target = self.target_resolver.resolve(request, session=session) options = request.options_override or request.package.metadata.import_options report_items = [ ResourceReportItem( @@ -165,7 +165,6 @@ class MigrationImportService: id_mapping_details: list[ResourceIdMapping] = [] self._import_api_tools( - session, request.package, target, options, @@ -173,14 +172,16 @@ class MigrationImportService: id_mapping, id_mapping_details, self._source_api_provider_ids_by_name(request.package), + session=session, ) - self._import_mcp_tools(session, request.package, target, options, report_items, id_mapping, id_mapping_details) - self._preflight_dependency_only_mcp(session, request.package, target, report_items) + self._import_mcp_tools( + request.package, target, options, report_items, id_mapping, id_mapping_details, session=session + ) + self._preflight_dependency_only_mcp(request.package, target, report_items, session=session) workflow_tool_app_ids = self._workflow_tool_source_app_ids(request.package) imported_workflow_ids: set[str] = set() if workflow_tool_app_ids: self._import_workflows( - session, request.package, target, options, @@ -189,12 +190,12 @@ class MigrationImportService: id_mapping_details=id_mapping_details, imported_workflow_ids=imported_workflow_ids, only_app_ids=workflow_tool_app_ids, + session=session, ) self._import_workflow_tools( - session, request.package, target, options, id_mapping, id_mapping_details, report_items + request.package, target, options, id_mapping, id_mapping_details, report_items, session=session ) self._import_workflows( - session, request.package, target, options, @@ -203,6 +204,7 @@ class MigrationImportService: id_mapping_details=id_mapping_details, imported_workflow_ids=imported_workflow_ids, skip_app_ids=imported_workflow_ids, + session=session, ) return ImportResult( report_items=report_items, @@ -218,7 +220,6 @@ class MigrationImportService: def _import_workflows( self, - session: Session, package: MigrationPackage, target: ImportTarget, options: ImportOptions, @@ -228,6 +229,8 @@ class MigrationImportService: imported_workflow_ids: set[str] | None = None, only_app_ids: set[str] | None = None, skip_app_ids: set[str] | None = None, + *, + session: Session, ) -> None: account = session.get(Account, target.operator_id) tenant = session.get(Tenant, target.tenant_id) @@ -248,7 +251,7 @@ class MigrationImportService: id_mapping, ) existing_app = ( - self._find_existing_app(session, app_id, target.tenant_id) + self._find_existing_app(app_id, target.tenant_id, session=session) if options.id_strategy == IdStrategy.PRESERVE_ID else None ) @@ -270,13 +273,13 @@ class MigrationImportService: continue imported_app_id = self._import_workflow_app( - session=session, account=account, workflow_data=workflow_data, dsl_content=dsl_content, app_id=app_id, existing_app=existing_app, options=options, + session=session, ) if app_id: self._record_id_mappings( @@ -290,7 +293,7 @@ class MigrationImportService: if imported_workflow_ids is not None: imported_workflow_ids.add(app_id) if options.create_app_api_token_on_import: - self._create_or_reuse_app_api_token(session, imported_app_id, target.tenant_id) + self._create_or_reuse_app_api_token(imported_app_id, target.tenant_id, session=session) report_items.append( ResourceReportItem( ResourceType.WORKFLOW, @@ -311,15 +314,15 @@ class MigrationImportService: def _import_workflow_app( self, *, - session: Session, account: Account, workflow_data: dict[str, object], dsl_content: str, app_id: str | None, existing_app: App | None, options: ImportOptions, + session: Session, ) -> str: - import_service = AppDslService(session) + import_service = AppDslService(cast(Session, session)) if existing_app is not None: import_result = import_service.import_app( account=account, @@ -408,12 +411,12 @@ class MigrationImportService: def _should_preserve_source_app_id(self, options: ImportOptions) -> bool: return options.id_strategy == IdStrategy.PRESERVE_ID - def _find_existing_app(self, session: Session, app_id: str | None, tenant_id: str) -> App | None: + def _find_existing_app(self, app_id: str | None, tenant_id: str, *, session: Session) -> App | None: if not self._is_uuid_string(app_id): return None return session.scalar(sa.select(App).where(App.id == app_id, App.tenant_id == tenant_id)) - def _create_or_reuse_app_api_token(self, session: Session, app_id: str, tenant_id: str) -> None: + def _create_or_reuse_app_api_token(self, app_id: str, tenant_id: str, *, session: Session) -> None: existing = session.scalar( sa.select(ApiToken).where( ApiToken.type == ApiTokenType.APP, @@ -433,7 +436,6 @@ class MigrationImportService: def _import_api_tools( self, - session: Session, package: MigrationPackage, target: ImportTarget, options: ImportOptions, @@ -441,6 +443,8 @@ class MigrationImportService: id_mapping: dict[str, str], id_mapping_details: list[ResourceIdMapping], source_provider_ids_by_name: dict[str, set[str]], + *, + session: Session, ) -> None: for tool_data in package.tools: provider_name = self._required_string(tool_data, "provider_name", "api_tool") @@ -510,7 +514,7 @@ class MigrationImportService: icon=icon, ) status = "created" - target_provider = self._find_api_tool_provider(session, target.tenant_id, provider_name) + target_provider = self._find_api_tool_provider(target.tenant_id, provider_name, session=session) if target_provider is not None: self._record_id_mappings( id_mapping, @@ -522,7 +526,9 @@ class MigrationImportService: ) report_items.append(ResourceReportItem(ResourceType.API_TOOL, provider_name, provider_name, status)) - def _find_api_tool_provider(self, session: Session, tenant_id: str, provider_name: str) -> ApiToolProvider | None: + def _find_api_tool_provider( + self, tenant_id: str, provider_name: str, *, session: Session + ) -> ApiToolProvider | None: return session.scalar( sa.select(ApiToolProvider).where( ApiToolProvider.tenant_id == tenant_id, @@ -558,13 +564,14 @@ class MigrationImportService: def _import_workflow_tools( self, - session: Session, package: MigrationPackage, target: ImportTarget, options: ImportOptions, id_mapping: dict[str, str], id_mapping_details: list[ResourceIdMapping], report_items: list[ResourceReportItem], + *, + session: Session, ) -> None: if not package.workflow_tools: return @@ -574,7 +581,10 @@ class MigrationImportService: for workflow_tool_data in package.workflow_tools: app_id = self._optional_string(workflow_tool_data.get("app_id")) resolved_app_id = id_mapping.get(app_id or "", app_id) - if not resolved_app_id or self._find_existing_app(session, resolved_app_id, target.tenant_id) is None: + if ( + not resolved_app_id + or self._find_existing_app(resolved_app_id, target.tenant_id, session=session) is None + ): report_items.append( ResourceReportItem( ResourceType.WORKFLOW_TOOL, @@ -586,7 +596,7 @@ class MigrationImportService: ) continue try: - self._ensure_workflow_app_is_published(session, target, account, resolved_app_id) + self._ensure_workflow_app_is_published(target, account, resolved_app_id, session=session) except Exception as exc: report_items.append( ResourceReportItem( @@ -602,7 +612,7 @@ class MigrationImportService: tool_name = self._required_string(workflow_tool_data, "name", "workflow_tool") lookup_workflow_tool_id = workflow_tool_id if options.id_strategy == IdStrategy.PRESERVE_ID else None existing = self._find_existing_workflow_tool( - session, target.tenant_id, lookup_workflow_tool_id, tool_name, resolved_app_id + target.tenant_id, lookup_workflow_tool_id, tool_name, resolved_app_id, session=session ) if existing is not None and options.conflict_strategy == ConflictStrategy.FAIL: raise MigrationDataError(f"Workflow tool already exists and conflict_strategy=fail: {tool_name}") @@ -669,7 +679,7 @@ class MigrationImportService: ) status = "created" target_provider = self._find_existing_workflow_tool( - session, target.tenant_id, import_id or None, tool_name, resolved_app_id + target.tenant_id, import_id or None, tool_name, resolved_app_id, session=session ) if target_provider is None: raise MigrationDataError(f"Workflow tool was not created: {tool_name}") @@ -686,9 +696,9 @@ class MigrationImportService: report_items.append(ResourceReportItem(ResourceType.WORKFLOW_TOOL, identifier, tool_name, status)) def _ensure_workflow_app_is_published( - self, session: Session, target: ImportTarget, account: Account, app_id: str + self, target: ImportTarget, account: Account, app_id: str, *, session: Session ) -> None: - app = self._find_existing_app(session, app_id, target.tenant_id) + app = self._find_existing_app(app_id, target.tenant_id, session=session) if app is None: raise MigrationDataError(f"Referenced workflow app was not found in target tenant: {app_id}") if app.workflow_id: @@ -714,20 +724,23 @@ class MigrationImportService: def _import_mcp_tools( self, - session: Session, package: MigrationPackage, target: ImportTarget, options: ImportOptions, report_items: list[ResourceReportItem], id_mapping: dict[str, str], id_mapping_details: list[ResourceIdMapping], + *, + session: Session, ) -> None: for mcp_data in package.mcp_tools: name = self._required_string(mcp_data, "name", "mcp_tool") server_identifier = self._required_string(mcp_data, "server_identifier", "mcp_tool") provider_id = self._optional_string(mcp_data.get("id")) lookup_provider_id = provider_id if options.id_strategy == IdStrategy.PRESERVE_ID else None - existing = self._find_existing_mcp_tool(session, target.tenant_id, lookup_provider_id, server_identifier) + existing = self._find_existing_mcp_tool( + target.tenant_id, lookup_provider_id, server_identifier, session=session + ) if existing is not None and options.conflict_strategy == ConflictStrategy.FAIL: raise MigrationDataError(f"MCP tool already exists and conflict_strategy=fail: {name}") if existing is not None and options.conflict_strategy == ConflictStrategy.SKIP: @@ -743,7 +756,7 @@ class MigrationImportService: report_items.append(ResourceReportItem(ResourceType.MCP_TOOL, existing.id, name, "skipped")) continue - service = MCPToolManageService(session=session) + service = MCPToolManageService(session=cast(Session, session)) configuration = MCPConfiguration.model_validate(mcp_data.get("configuration") or {}) authentication = ( MCPAuthentication.model_validate(mcp_data["authentication"]) if mcp_data.get("authentication") else None @@ -784,7 +797,7 @@ class MigrationImportService: authentication=authentication, ) created_provider = self._find_existing_mcp_tool( - session, target.tenant_id, lookup_provider_id, server_identifier + target.tenant_id, lookup_provider_id, server_identifier, session=session ) if created_provider is None: raise MigrationDataError(f"MCP provider was not created: {name}") @@ -812,7 +825,12 @@ class MigrationImportService: provider.authed = True def _find_existing_mcp_tool( - self, session: Session, tenant_id: str, provider_id: str | None, server_identifier: str + self, + tenant_id: str, + provider_id: str | None, + server_identifier: str, + *, + session: Session, ) -> MCPToolProvider | None: predicates = [MCPToolProvider.server_identifier == server_identifier] if self._is_uuid_string(provider_id): @@ -831,7 +849,13 @@ class MigrationImportService: return True def _find_existing_workflow_tool( - self, session: Session, tenant_id: str, workflow_tool_id: str | None, tool_name: str, app_id: str + self, + tenant_id: str, + workflow_tool_id: str | None, + tool_name: str, + app_id: str, + *, + session: Session, ) -> WorkflowToolProvider | None: predicates = [WorkflowToolProvider.name == tool_name, WorkflowToolProvider.app_id == app_id] if self._is_uuid_string(workflow_tool_id): @@ -843,14 +867,21 @@ class MigrationImportService: ) def _preflight_dependency_only_mcp( - self, session: Session, package: MigrationPackage, target: ImportTarget, report_items: list[ResourceReportItem] + self, + package: MigrationPackage, + target: ImportTarget, + report_items: list[ResourceReportItem], + *, + session: Session, ) -> None: for dependency in package.dependencies: if dependency.get("kind") != DependencyKind.MCP_TOOL.value: continue provider_id = str(dependency.get("provider_id", dependency.get("id", ""))) provider_name = self._optional_string(dependency.get("provider_name") or dependency.get("name")) - existing = self._find_dependency_only_mcp_provider(session, target.tenant_id, provider_id, provider_name) + existing = self._find_dependency_only_mcp_provider( + target.tenant_id, provider_id, provider_name, session=session + ) report_name = f"mcp_tool {provider_name or getattr(existing, 'name', None) or provider_id}" if existing is not None: report_items.append( @@ -879,7 +910,12 @@ class MigrationImportService: ) def _find_dependency_only_mcp_provider( - self, session: Session, tenant_id: str, provider_id: str, provider_name: str | None + self, + tenant_id: str, + provider_id: str, + provider_name: str | None, + *, + session: Session, ) -> MCPToolProvider | None: predicates = [MCPToolProvider.server_identifier == provider_id] if self._is_uuid_string(provider_id): diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index b36926a32c5..dda5440f772 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -13,7 +13,7 @@ import sqlalchemy as sa from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator from redis.exceptions import LockNotOwnedError from sqlalchemy import ColumnElement, delete, exists, func, select, update -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound from configs import dify_config @@ -110,13 +110,6 @@ from tasks.sync_website_document_indexing_task import sync_website_document_inde logger = logging.getLogger(__name__) -def _session_for_helpers(session: scoped_session | Session) -> Session: - """Return a concrete SQLAlchemy session for helpers that do not accept scoped_session.""" - if isinstance(session, scoped_session): - return session() - return session - - class ProcessRulesDict(TypedDict): mode: ProcessRuleMode rules: dict[str, Any] @@ -244,11 +237,11 @@ class _EstimateArgs(BaseModel): class DatasetService: @staticmethod - def _can_manage_all_datasets(tenant_id: str, account_id: str) -> bool: + def _can_manage_all_datasets(tenant_id: str, account_id: str, *, session: Session) -> bool: if not dify_config.RBAC_ENABLED: return False - permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id) + permissions = enterprise_rbac_service.RBACService.MyPermissions.get(tenant_id, account_id, session=session) workspace_permission_keys = getattr(getattr(permissions, "workspace", None), "permission_keys", []) or [] return "dataset.create_and_management" in workspace_permission_keys @@ -256,7 +249,7 @@ class DatasetService: def get_datasets( page, per_page, - session: scoped_session | Session, + session: Session, tenant_id=None, user=None, search=None, @@ -291,7 +284,9 @@ class DatasetService: return [], 0 else: if dify_config.RBAC_ENABLED: - can_manage_all_datasets = DatasetService._can_manage_all_datasets(str(tenant_id), str(user.id)) + can_manage_all_datasets = DatasetService._can_manage_all_datasets( + str(tenant_id), str(user.id), session=session + ) should_show_all_datasets = include_all and can_manage_all_datasets else: should_show_all_datasets = user.current_role == TenantAccountRole.OWNER and include_all @@ -361,7 +356,7 @@ class DatasetService: return datasets.items, datasets.total @staticmethod - def get_process_rules(dataset_id, session: scoped_session | Session) -> ProcessRulesDict: + def get_process_rules(dataset_id, session: Session) -> ProcessRulesDict: # get the latest process rule dataset_process_rule = session.execute( select(DatasetProcessRule) @@ -406,7 +401,6 @@ class DatasetService: @staticmethod def create_empty_dataset( - session: Session, tenant_id: str, name: str, description: str | None, @@ -420,6 +414,8 @@ class DatasetService: embedding_model_name: str | None = None, retrieval_model: RetrievalModel | None = None, summary_index_setting: dict[str, Any] | None = None, + *, + session: Session, ): # check if dataset name already exists if session.scalar(select(Dataset).where(Dataset.name == name, Dataset.tenant_id == tenant_id).limit(1)): @@ -473,7 +469,7 @@ class DatasetService: if provider == "external" and external_knowledge_api_id: external_knowledge_api = ExternalDatasetService.get_external_knowledge_api( - session, external_knowledge_api_id, tenant_id + external_knowledge_api_id, tenant_id, session=session ) if not external_knowledge_api: raise ValueError("External API template not found.") @@ -501,7 +497,7 @@ class DatasetService: def create_empty_rag_pipeline_dataset( tenant_id: str, rag_pipeline_dataset_create_entity: RagPipelineDatasetCreateEntity, - session: scoped_session | Session, + session: Session, ): if rag_pipeline_dataset_create_entity.name: # check if dataset name already exists @@ -549,7 +545,7 @@ class DatasetService: return dataset @staticmethod - def get_dataset(dataset_id, session: scoped_session | Session) -> Dataset | None: + def get_dataset(dataset_id, session: Session) -> Dataset | None: dataset: Dataset | None = session.get(Dataset, dataset_id) return dataset @@ -632,7 +628,7 @@ class DatasetService: raise ValueError(ex.description) @staticmethod - def update_dataset(session: Session, dataset_id, data, user): + def update_dataset(dataset_id, data, user, *, session: Session): """ Update dataset configuration and settings. @@ -672,7 +668,7 @@ class DatasetService: return DatasetService._update_internal_dataset(dataset, data, user, session) @staticmethod - def _has_dataset_same_name(tenant_id: str, dataset_id: str, name: str, session: scoped_session | Session): + def _has_dataset_same_name(tenant_id: str, dataset_id: str, name: str, session: Session): dataset = session.scalar( select(Dataset) .where( @@ -725,7 +721,7 @@ class DatasetService: if not external_knowledge_api_id: raise ValueError("External knowledge api id is required.") # Ensure the referenced external API template exists and belongs to the dataset tenant. - ExternalDatasetService.get_external_knowledge_api(session, external_knowledge_api_id, dataset.tenant_id) + ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id, dataset.tenant_id, session=session) # Update metadata fields dataset.updated_by = user.id if user else None dataset.updated_at = naive_utc_now() @@ -743,7 +739,7 @@ class DatasetService: @staticmethod def _update_external_knowledge_binding( - dataset_id, external_knowledge_id, external_knowledge_api_id, session: scoped_session | Session + dataset_id, external_knowledge_id, external_knowledge_api_id, session: Session ): """ Update external knowledge binding configuration. @@ -770,7 +766,7 @@ class DatasetService: session.add(external_knowledge_binding) @staticmethod - def _update_internal_dataset(dataset, data, user, session: scoped_session | Session): + def _update_internal_dataset(dataset, data, user, session: Session): """ Update internal dataset configuration. @@ -836,9 +832,7 @@ class DatasetService: return dataset @staticmethod - def _update_pipeline_knowledge_base_node_data( - dataset: Dataset, updata_user_id: str, session: scoped_session | Session - ): + def _update_pipeline_knowledge_base_node_data(dataset: Dataset, updata_user_id: str, session: Session): """ Update pipeline knowledge base node data. """ @@ -850,7 +844,7 @@ class DatasetService: return try: - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(session) published_workflow = rag_pipeline_service.get_published_workflow(pipeline) draft_workflow = rag_pipeline_service.get_draft_workflow(pipeline) @@ -921,7 +915,7 @@ class DatasetService: raise @staticmethod - def _handle_indexing_technique_change(dataset, data, filtered_data, session: scoped_session | Session): + def _handle_indexing_technique_change(dataset, data, filtered_data, session: Session): """ Handle changes in indexing technique and configure embedding models accordingly. @@ -955,7 +949,7 @@ class DatasetService: return None @staticmethod - def _configure_embedding_model_for_high_quality(data, filtered_data, session: scoped_session | Session): + def _configure_embedding_model_for_high_quality(data, filtered_data, session: Session): """ Configure embedding model settings for high quality indexing. @@ -992,9 +986,7 @@ class DatasetService: raise ValueError(ex.description) @staticmethod - def _handle_embedding_model_update_when_technique_unchanged( - dataset, data, filtered_data, session: scoped_session | Session - ): + def _handle_embedding_model_update_when_technique_unchanged(dataset, data, filtered_data, session: Session): """ Handle embedding model updates when indexing technique remains the same. @@ -1043,7 +1035,7 @@ class DatasetService: del filtered_data["embedding_model"] @staticmethod - def _update_embedding_model_settings(dataset, data, filtered_data, session: scoped_session | Session): + def _update_embedding_model_settings(dataset, data, filtered_data, session: Session): """ Update embedding model settings with new values. @@ -1078,7 +1070,7 @@ class DatasetService: return None @staticmethod - def _apply_new_embedding_settings(dataset, data, filtered_data, session: scoped_session | Session): + def _apply_new_embedding_settings(dataset, data, filtered_data, session: Session): """ Apply new embedding model settings to the dataset. @@ -1176,7 +1168,11 @@ class DatasetService: @staticmethod def update_rag_pipeline_dataset_settings( - session: Session, dataset: Dataset, knowledge_configuration: KnowledgeConfiguration, has_published: bool = False + dataset: Dataset, + knowledge_configuration: KnowledgeConfiguration, + has_published: bool = False, + *, + session: Session, ): if not current_user or not current_user.current_tenant_id: raise ValueError("Current user or current tenant not found") @@ -1335,7 +1331,7 @@ class DatasetService: deal_dataset_index_update_task.delay(dataset.id, action) @staticmethod - def delete_dataset(dataset_id, user, session: scoped_session | Session): + def delete_dataset(dataset_id, user, session: Session): dataset = DatasetService.get_dataset(dataset_id, session) if dataset is None: @@ -1350,12 +1346,12 @@ class DatasetService: return True @staticmethod - def dataset_use_check(dataset_id, session: scoped_session | Session) -> bool: + def dataset_use_check(dataset_id, session: Session) -> bool: stmt = select(exists().where(AppDatasetJoin.dataset_id == dataset_id)) return session.execute(stmt).scalar_one() @staticmethod - def check_dataset_permission(dataset, user, session: scoped_session | Session): + def check_dataset_permission(dataset, user, session: Session): """Validate dataset access for a user, using the injected session for partial-member lookups.""" if dataset.tenant_id != user.current_tenant_id: logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id) @@ -1378,7 +1374,7 @@ class DatasetService: @staticmethod def check_dataset_operator_permission( - user: Account | None = None, dataset: Dataset | None = None, *, session: scoped_session | Session + user: Account | None = None, dataset: Dataset | None = None, *, session: Session ): if not dataset: raise ValueError("Dataset not found") @@ -1409,7 +1405,7 @@ class DatasetService: return dataset_queries.items, dataset_queries.total @staticmethod - def get_related_apps(dataset_id: str, session: scoped_session | Session): + def get_related_apps(dataset_id: str, session: Session): return session.scalars( select(AppDatasetJoin) .where(AppDatasetJoin.dataset_id == dataset_id) @@ -1417,7 +1413,7 @@ class DatasetService: ).all() @staticmethod - def update_dataset_api_status(dataset_id: str, status: bool, session: scoped_session | Session): + def update_dataset_api_status(dataset_id: str, status: bool, session: Session): dataset = DatasetService.get_dataset(dataset_id, session) if dataset is None: raise NotFound("Dataset not found.") @@ -1429,7 +1425,7 @@ class DatasetService: session.commit() @staticmethod - def get_dataset_auto_disable_logs(dataset_id: str, session: scoped_session | Session) -> AutoDisableLogsDict: + def get_dataset_auto_disable_logs(dataset_id: str, session: Session) -> AutoDisableLogsDict: assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) @@ -1628,9 +1624,7 @@ class DocumentService: } @staticmethod - def get_document( - dataset_id: str, document_id: str | None = None, *, session: scoped_session | Session - ) -> Document | None: + def get_document(dataset_id: str, document_id: str | None = None, *, session: Session) -> Document | None: """Fetch a document by id within a dataset using the caller-provided session.""" if document_id: document = session.scalar( @@ -1641,9 +1635,7 @@ class DocumentService: return None @staticmethod - def get_documents_by_ids( - dataset_id: str, document_ids: Sequence[str], session: scoped_session | Session - ) -> Sequence[Document]: + def get_documents_by_ids(dataset_id: str, document_ids: Sequence[str], session: Session) -> Sequence[Document]: """Fetch documents for a dataset in a single batch query.""" if not document_ids: return [] @@ -1661,7 +1653,7 @@ class DocumentService: def update_documents_need_summary( dataset_id: str, document_ids: Sequence[str], - session: scoped_session | Session, + session: Session, need_summary: bool = True, ) -> int: """ @@ -1705,7 +1697,7 @@ class DocumentService: return updated_count @staticmethod - def get_document_download_url(document: Document, session: scoped_session | Session) -> str: + def get_document_download_url(document: Document, session: Session) -> str: """ Return a signed download URL for an upload-file document. """ @@ -1717,6 +1709,7 @@ class DocumentService: documents: Sequence[Document], dataset: Dataset, tenant_id: str, + session: Session, ) -> None: """ Enrich documents with summary_index_status based on dataset summary index settings. @@ -1728,6 +1721,7 @@ class DocumentService: documents: List of Document instances to enrich dataset: Dataset instance containing summary_index_setting tenant_id: Tenant ID for summary status lookup + session: SQLAlchemy session used to read summary status records """ # Check if dataset has summary index enabled has_summary_index = dataset.summary_index_setting and dataset.summary_index_setting.get("enable") is True @@ -1745,6 +1739,7 @@ class DocumentService: document_ids=document_ids_need_summary, dataset_id=dataset.id, tenant_id=tenant_id, + session=session, ) # Add summary_index_status to each document @@ -1763,7 +1758,7 @@ class DocumentService: document_ids: Sequence[str], tenant_id: str, current_user: Account, - session: scoped_session | Session, + session: Session, ) -> tuple[list[UploadFile], str]: """ Resolve upload files for batch ZIP downloads and generate a client-visible filename. @@ -1814,7 +1809,7 @@ class DocumentService: return str(upload_file_id) @staticmethod - def _get_upload_file_for_upload_file_document(document: Document, session: scoped_session | Session) -> UploadFile: + def _get_upload_file_for_upload_file_document(document: Document, session: Session) -> UploadFile: """ Load the `UploadFile` row for an upload-file document. """ @@ -1823,9 +1818,7 @@ class DocumentService: invalid_source_message="Document does not have an uploaded file to download.", missing_file_message="Uploaded file not found.", ) - upload_files_by_id = FileService.get_upload_files_by_ids( - _session_for_helpers(session), document.tenant_id, [upload_file_id] - ) + upload_files_by_id = FileService.get_upload_files_by_ids(document.tenant_id, [upload_file_id], session=session) upload_file = upload_files_by_id.get(upload_file_id) if not upload_file: raise NotFound("Uploaded file not found.") @@ -1837,7 +1830,7 @@ class DocumentService: dataset_id: str, document_ids: Sequence[str], tenant_id: str, - session: scoped_session | Session, + session: Session, ) -> dict[str, UploadFile]: """ Batch load upload files keyed by document id for ZIP downloads. @@ -1865,9 +1858,7 @@ class DocumentService: upload_file_ids.append(upload_file_id) upload_file_ids_by_document_id[document_id] = upload_file_id - upload_files_by_id = FileService.get_upload_files_by_ids( - _session_for_helpers(session), tenant_id, upload_file_ids - ) + upload_files_by_id = FileService.get_upload_files_by_ids(tenant_id, upload_file_ids, session=session) missing_upload_file_ids: set[str] = set(upload_file_ids) - set(upload_files_by_id.keys()) if missing_upload_file_ids: raise NotFound("Only uploaded-file documents can be downloaded as ZIP.") @@ -1878,13 +1869,13 @@ class DocumentService: } @staticmethod - def get_document_by_id(document_id: str, session: scoped_session | Session) -> Document | None: + def get_document_by_id(document_id: str, session: Session) -> Document | None: document = session.get(Document, document_id) return document @staticmethod - def get_document_by_ids(document_ids: list[str], session: scoped_session | Session) -> Sequence[Document]: + def get_document_by_ids(document_ids: list[str], session: Session) -> Sequence[Document]: documents = session.scalars( select(Document).where( Document.id.in_(document_ids), @@ -1896,7 +1887,7 @@ class DocumentService: return documents @staticmethod - def get_document_by_dataset_id(dataset_id: str, session: scoped_session | Session) -> Sequence[Document]: + def get_document_by_dataset_id(dataset_id: str, session: Session) -> Sequence[Document]: documents = session.scalars( select(Document).where( Document.dataset_id == dataset_id, @@ -1907,7 +1898,7 @@ class DocumentService: return documents @staticmethod - def get_working_documents_by_dataset_id(dataset_id: str, session: scoped_session | Session) -> Sequence[Document]: + def get_working_documents_by_dataset_id(dataset_id: str, session: Session) -> Sequence[Document]: documents = session.scalars( select(Document).where( Document.dataset_id == dataset_id, @@ -1920,7 +1911,7 @@ class DocumentService: return documents @staticmethod - def get_error_documents_by_dataset_id(dataset_id: str, session: scoped_session | Session) -> Sequence[Document]: + def get_error_documents_by_dataset_id(dataset_id: str, session: Session) -> Sequence[Document]: documents = session.scalars( select(Document).where( Document.dataset_id == dataset_id, @@ -1930,7 +1921,7 @@ class DocumentService: return documents @staticmethod - def get_batch_documents(dataset_id: str, batch: str, session: scoped_session | Session) -> Sequence[Document]: + def get_batch_documents(dataset_id: str, batch: str, session: Session) -> Sequence[Document]: assert isinstance(current_user, Account) documents = session.scalars( select(Document).where( @@ -1943,7 +1934,7 @@ class DocumentService: return documents @staticmethod - def get_document_file_detail(file_id: str, session: scoped_session | Session): + def get_document_file_detail(file_id: str, session: Session): file_detail = session.get(UploadFile, file_id) return file_detail @@ -1955,7 +1946,7 @@ class DocumentService: return False @staticmethod - def delete_document(document, session: scoped_session | Session): + def delete_document(document, session: Session): # trigger document_was_deleted signal file_id = None if document.data_source_type == DataSourceType.UPLOAD_FILE: @@ -1975,7 +1966,7 @@ class DocumentService: dataset_ref: DatasetRef, document_ids: list[str], doc_form: str | None, - session: scoped_session | Session, + session: Session, ): # Check if document_ids is not empty to avoid WHERE false condition if not document_ids or len(document_ids) == 0: @@ -2006,7 +1997,7 @@ class DocumentService: batch_clean_document_task.delay(deleted_document_ids, dataset_ref.dataset_id, doc_form, file_ids) @staticmethod - def rename_document(dataset_id: str, document_id: str, name: str, session: scoped_session | Session) -> Document: + def rename_document(dataset_id: str, document_id: str, name: str, session: Session) -> Document: assert isinstance(current_user, Account) dataset = DatasetService.get_dataset(dataset_id, session) @@ -2041,7 +2032,7 @@ class DocumentService: return document @staticmethod - def pause_document(document, session: scoped_session | Session): + def pause_document(document, session: Session): if document.indexing_status not in { IndexingStatus.WAITING, IndexingStatus.PARSING, @@ -2063,7 +2054,7 @@ class DocumentService: redis_client.setnx(indexing_cache_key, "True") @staticmethod - def recover_document(document, session: scoped_session | Session): + def recover_document(document, session: Session): if not document.is_paused: raise DocumentIndexingError() # update document to be recover @@ -2080,7 +2071,7 @@ class DocumentService: recover_document_indexing_task.delay(document.dataset_id, document.id) @staticmethod - def retry_document(dataset_id: str, documents: list[Document], session: scoped_session | Session): + def retry_document(dataset_id: str, documents: list[Document], session: Session): for document in documents: # add retry flag retry_indexing_cache_key = f"document_{document.id}_is_retried" @@ -2100,7 +2091,7 @@ class DocumentService: retry_document_indexing_task.delay(dataset_id, document_ids, current_user.id) @staticmethod - def sync_website_document(dataset_id: str, document: Document, session: scoped_session | Session): + def sync_website_document(dataset_id: str, document: Document, session: Session): # add sync flag sync_indexing_cache_key = f"document_{document.id}_is_sync" cache_result = redis_client.get(sync_indexing_cache_key) @@ -2120,7 +2111,7 @@ class DocumentService: sync_website_document_indexing_task.delay(dataset_id, document.id) @staticmethod - def get_documents_position(dataset_id, session: scoped_session | Session): + def get_documents_position(dataset_id, session: Session): document = session.scalar( select(Document).where(Document.dataset_id == dataset_id).order_by(Document.position.desc()).limit(1) ) @@ -2137,7 +2128,7 @@ class DocumentService: dataset_process_rule: DatasetProcessRule | None = None, created_from: str = DocumentCreatedFrom.WEB, *, - session: scoped_session | Session, + session: Session, ) -> tuple[list[Document], str]: # check doc_form DatasetService.check_doc_form(dataset, knowledge_config.doc_form) @@ -2793,7 +2784,7 @@ class DocumentService: return document @staticmethod - def get_tenant_documents_count(session: scoped_session | Session): + def get_tenant_documents_count(*, session: Session): assert isinstance(current_user, Account) documents_count = ( @@ -2817,7 +2808,7 @@ class DocumentService: dataset_process_rule: DatasetProcessRule | None = None, created_from: str = DocumentCreatedFrom.WEB, *, - session: scoped_session | Session, + session: Session, ): assert isinstance(current_user, Account) @@ -2944,7 +2935,7 @@ class DocumentService: @staticmethod def save_document_without_dataset_id( - tenant_id: str, knowledge_config: KnowledgeConfig, account: Account, session: scoped_session | Session + tenant_id: str, knowledge_config: KnowledgeConfig, account: Account, session: Session ): assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None @@ -3133,7 +3124,7 @@ class DocumentService: document_ids: list[str], action: Literal["enable", "disable", "archive", "un_archive"], user, - session: scoped_session | Session, + session: Session, ): """ Batch update document status. @@ -3216,7 +3207,7 @@ class DocumentService: document = update_info["document"] indexing_cache_key = f"document_{document.id}_indexing" redis_client.setex(indexing_cache_key, 600, 1) - except Exception as e: + except Exception: # Log the error but do not rollback the transaction logger.exception("Error setting cache for document %s", update_info["document"].id) # Raise any propagation error after all updates @@ -3340,9 +3331,7 @@ class SegmentService: raise ValueError(f"Exceeded maximum attachment limit of {single_chunk_attachment_limit}") @classmethod - def create_segment( - cls, args: dict[str, Any], document: Document, dataset: Dataset, session: scoped_session | Session - ): + def create_segment(cls, args: dict[str, Any], document: Document, dataset: Dataset, session: Session): assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None @@ -3408,7 +3397,13 @@ class SegmentService: try: keywords = args.get("keywords") keywords_list = [keywords] if keywords is not None else None - VectorService.create_segments_vector(keywords_list, [segment_document], dataset, document.doc_form) + VectorService.create_segments_vector( + keywords_list, + [segment_document], + dataset, + document.doc_form, + session, + ) except Exception as e: logger.exception("create segment index failed") segment_document.enabled = False @@ -3422,9 +3417,7 @@ class SegmentService: pass @classmethod - def multi_create_segment( - cls, segments: list, document: Document, dataset: Dataset, session: scoped_session | Session - ): + def multi_create_segment(cls, segments: list, document: Document, dataset: Dataset, session: Session): assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None @@ -3498,7 +3491,11 @@ class SegmentService: try: # save vector index VectorService.create_segments_vector( - keywords_list, pre_segment_data_list, dataset, document.doc_form + keywords_list, + pre_segment_data_list, + dataset, + document.doc_form, + session, ) except Exception as e: logger.exception("create segment index failed") @@ -3519,7 +3516,7 @@ class SegmentService: segment: DocumentSegment, document: Document, dataset: Dataset, - session: scoped_session | Session, + session: Session, ): assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None @@ -3597,7 +3594,13 @@ class SegmentService: processing_rule = session.get(DatasetProcessRule, document.dataset_process_rule_id) if processing_rule: VectorService.generate_child_chunks( - segment, document, dataset, embedding_model_instance, processing_rule, True + segment, + document, + dataset, + embedding_model_instance, + processing_rule, + session, + True, ) elif document.doc_form in (IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX): if args.enabled or keyword_changed: @@ -3628,7 +3631,12 @@ class SegmentService: from services.summary_index_service import SummaryIndexService try: - SummaryIndexService.update_summary_for_segment(segment, dataset, args.summary) + SummaryIndexService.update_summary_for_segment( + segment, + dataset, + args.summary, + session=session, + ) except Exception: logger.exception("Failed to update summary for segment %s", segment.id) # Don't fail the entire update if summary update fails @@ -3697,7 +3705,13 @@ class SegmentService: processing_rule = session.get(DatasetProcessRule, document.dataset_process_rule_id) if processing_rule: VectorService.generate_child_chunks( - segment, document, dataset, embedding_model_instance, processing_rule, True + segment, + document, + dataset, + embedding_model_instance, + processing_rule, + session, + True, ) elif document.doc_form in (IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX): # update segment vector index @@ -3728,7 +3742,10 @@ class SegmentService: try: SummaryIndexService.generate_and_vectorize_summary( - segment, dataset, dataset.summary_index_setting + segment, + dataset, + dataset.summary_index_setting, + session=session, ) logger.info("Auto-regenerated summary for segment %s after content change", segment.id) except Exception: @@ -3743,7 +3760,12 @@ class SegmentService: from services.summary_index_service import SummaryIndexService try: - SummaryIndexService.update_summary_for_segment(segment, dataset, args.summary) + SummaryIndexService.update_summary_for_segment( + segment, + dataset, + args.summary, + session=session, + ) logger.info("Updated summary for segment %s with user-provided content", segment.id) except Exception: logger.exception("Failed to update summary for segment %s", segment.id) @@ -3760,7 +3782,10 @@ class SegmentService: try: SummaryIndexService.generate_and_vectorize_summary( - segment, dataset, dataset.summary_index_setting + segment, + dataset, + dataset.summary_index_setting, + session=session, ) logger.info( "Regenerated summary for segment %s after content change (summary unchanged)", @@ -3770,7 +3795,7 @@ class SegmentService: logger.exception("Failed to regenerate summary for segment %s", segment.id) # Don't fail the entire update if summary regeneration fails # update multimodel vector index - VectorService.update_multimodel_vector(segment, args.attachment_ids or [], dataset) + VectorService.update_multimodel_vector(segment, args.attachment_ids or [], dataset, session) except Exception as e: logger.exception("update segment index failed") segment.enabled = False @@ -3784,9 +3809,7 @@ class SegmentService: return new_segment @classmethod - def delete_segment( - cls, segment: DocumentSegment, document: Document, dataset: Dataset, session: scoped_session | Session - ): + def delete_segment(cls, segment: DocumentSegment, document: Document, dataset: Dataset, session: Session): indexing_cache_key = f"segment_{segment.id}_delete_indexing" cache_result = redis_client.get(indexing_cache_key) if cache_result is not None: @@ -3821,9 +3844,7 @@ class SegmentService: session.commit() @classmethod - def delete_segments( - cls, segment_ids: list, document: Document, dataset: Dataset, session: scoped_session | Session - ): + def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset, session: Session): assert current_user is not None # Check if segment_ids is not empty to avoid WHERE false condition if not segment_ids or len(segment_ids) == 0: @@ -3882,7 +3903,7 @@ class SegmentService: action: Literal["enable", "disable"], dataset: Dataset, document: Document, - session: scoped_session | Session, + session: Session, ): assert current_user is not None @@ -3948,7 +3969,7 @@ class SegmentService: segment: DocumentSegment, document: Document, dataset: Dataset, - session: scoped_session | Session, + session: Session, ) -> ChildChunk: assert isinstance(current_user, Account) @@ -3997,7 +4018,7 @@ class SegmentService: segment: DocumentSegment, document: Document, dataset: Dataset, - session: scoped_session | Session, + session: Session, ) -> list[ChildChunk]: assert isinstance(current_user, Account) child_chunks = session.scalars( @@ -4072,7 +4093,7 @@ class SegmentService: segment: DocumentSegment, document: Document, dataset: Dataset, - session: scoped_session | Session, + session: Session, ) -> ChildChunk: assert current_user is not None @@ -4092,7 +4113,7 @@ class SegmentService: return child_chunk @classmethod - def delete_child_chunk(cls, child_chunk: ChildChunk, dataset: Dataset, session: scoped_session | Session): + def delete_child_chunk(cls, child_chunk: ChildChunk, dataset: Dataset, session: Session): session.delete(child_chunk) try: VectorService.delete_child_chunk_vector(child_chunk, dataset) @@ -4124,9 +4145,7 @@ class SegmentService: return paginate_query(query, page=page, per_page=limit, max_per_page=100) @classmethod - def get_child_chunk_by_id( - cls, child_chunk_id: str, tenant_id: str, session: scoped_session | Session - ) -> ChildChunk | None: + def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str, session: Session) -> ChildChunk | None: """Get a child chunk by its ID.""" result = session.scalar( select(ChildChunk).where(ChildChunk.id == child_chunk_id, ChildChunk.tenant_id == tenant_id).limit(1) @@ -4134,9 +4153,11 @@ class SegmentService: return result if isinstance(result, ChildChunk) else None @classmethod - def get_child_chunk_by_segment_ref(cls, child_chunk_id: str, segment_ref: SegmentRef) -> ChildChunk | None: + def get_child_chunk_by_segment_ref( + cls, child_chunk_id: str, segment_ref: SegmentRef, session: Session + ) -> ChildChunk | None: """Get a child chunk through the full tenant/dataset/document/segment chain.""" - result = db.session.scalar( + result = session.scalar( select(ChildChunk) .where( ChildChunk.id == child_chunk_id, @@ -4178,9 +4199,7 @@ class SegmentService: return paginated_segments.items, paginated_segments.total @classmethod - def get_segment_by_id( - cls, segment_id: str, tenant_id: str, session: scoped_session | Session - ) -> DocumentSegment | None: + def get_segment_by_id(cls, segment_id: str, tenant_id: str, session: Session) -> DocumentSegment | None: """Get a segment by its ID.""" result = session.scalar( select(DocumentSegment) @@ -4190,9 +4209,9 @@ class SegmentService: return result if isinstance(result, DocumentSegment) else None @classmethod - def get_segment_by_ref(cls, segment_ref: SegmentRef) -> DocumentSegment | None: + def get_segment_by_ref(cls, segment_ref: SegmentRef, session: Session) -> DocumentSegment | None: """Get a segment through the full tenant/dataset/document ownership chain.""" - result = db.session.scalar( + result = session.scalar( select(DocumentSegment) .where( DocumentSegment.id == segment_ref.segment_id, @@ -4209,7 +4228,7 @@ class SegmentService: cls, document_id: str, dataset_id: str, - session: scoped_session | Session, + session: Session, status: str | None = None, enabled: bool | None = None, ) -> Sequence[DocumentSegment]: @@ -4242,7 +4261,7 @@ class SegmentService: class DatasetCollectionBindingService: @classmethod def get_dataset_collection_binding( - cls, provider_name: str, model_name: str, session: scoped_session | Session, collection_type: str = "dataset" + cls, provider_name: str, model_name: str, session: Session, collection_type: str = "dataset" ) -> DatasetCollectionBinding: dataset_collection_binding = session.scalar( select(DatasetCollectionBinding) @@ -4268,7 +4287,7 @@ class DatasetCollectionBindingService: @classmethod def get_dataset_collection_binding_by_id_and_type( - cls, collection_binding_id: str, session: scoped_session | Session, collection_type: str = "dataset" + cls, collection_binding_id: str, session: Session, collection_type: str = "dataset" ) -> DatasetCollectionBinding: dataset_collection_binding = session.scalar( select(DatasetCollectionBinding) @@ -4286,7 +4305,7 @@ class DatasetCollectionBindingService: class DatasetPermissionService: @classmethod - def get_dataset_partial_member_list(cls, dataset_id, session: scoped_session | Session): + def get_dataset_partial_member_list(cls, dataset_id, session: Session): user_list_query = session.scalars( select( DatasetPermission.account_id, @@ -4296,7 +4315,7 @@ class DatasetPermissionService: return user_list_query @classmethod - def update_partial_member_list(cls, tenant_id, dataset_id, user_list, session: scoped_session | Session): + def update_partial_member_list(cls, tenant_id, dataset_id, user_list, session: Session): try: session.execute(delete(DatasetPermission).where(DatasetPermission.dataset_id == dataset_id)) permissions = [] @@ -4315,7 +4334,7 @@ class DatasetPermissionService: raise e @classmethod - def check_permission(cls, session: Session, user, dataset, requested_permission, requested_partial_member_list): + def check_permission(cls, user, dataset, requested_permission, requested_partial_member_list, *, session: Session): if not user.is_dataset_editor: raise NoPermissionError("User does not have permission to edit this dataset.") @@ -4332,7 +4351,7 @@ class DatasetPermissionService: raise ValueError("Dataset operators cannot change the dataset permissions.") @classmethod - def clear_partial_member_list(cls, dataset_id, session: scoped_session | Session): + def clear_partial_member_list(cls, dataset_id, session: Session): try: session.execute(delete(DatasetPermission).where(DatasetPermission.dataset_id == dataset_id)) session.commit() diff --git a/api/services/datasource_provider_service.py b/api/services/datasource_provider_service.py index 12807a41f04..5de194dd262 100644 --- a/api/services/datasource_provider_service.py +++ b/api/services/datasource_provider_service.py @@ -446,12 +446,14 @@ class DatasourceProviderService: is not None ) - def is_tenant_oauth_params_enabled(self, tenant_id: str, datasource_provider_id: DatasourceProviderID) -> bool: + def is_tenant_oauth_params_enabled( + self, tenant_id: str, datasource_provider_id: DatasourceProviderID, *, session: Session + ) -> bool: """ check if tenant oauth params is enabled """ return ( - db.session.scalar( + session.scalar( select(func.count(DatasourceOauthTenantParamConfig.id)).where( DatasourceOauthTenantParamConfig.tenant_id == tenant_id, DatasourceOauthTenantParamConfig.provider == datasource_provider_id.provider_name, @@ -463,12 +465,17 @@ class DatasourceProviderService: ) > 0 def get_tenant_oauth_client( - self, tenant_id: str, datasource_provider_id: DatasourceProviderID, mask: bool = False + self, + tenant_id: str, + datasource_provider_id: DatasourceProviderID, + mask: bool = False, + *, + session: Session, ) -> Mapping[str, Any] | None: """ get tenant oauth client """ - tenant_oauth_client_params = db.session.scalar( + tenant_oauth_client_params = session.scalar( select(DatasourceOauthTenantParamConfig) .where( DatasourceOauthTenantParamConfig.tenant_id == tenant_id, @@ -547,7 +554,7 @@ class DatasourceProviderService: @staticmethod def generate_next_datasource_provider_name( - session: Session, tenant_id: str, provider_id: DatasourceProviderID, credential_type: CredentialType + tenant_id: str, provider_id: DatasourceProviderID, credential_type: CredentialType, *, session: Session ) -> str: db_providers = session.scalars( select(DatasourceProvider).where( @@ -800,6 +807,8 @@ class DatasourceProviderService: provider: str, plugin_id: str, user: "Account | None" = None, + *, + session: Session, ) -> list[dict]: """ list datasource credentials with obfuscated sensitive fields, @@ -829,11 +838,11 @@ class DatasourceProviderService: credential_type=CredPermType.DATASOURCE_PROVIDER, user=user, ) - datasource_providers: list[DatasourceProvider] = list(db.session.scalars(query).all()) + datasource_providers: list[DatasourceProvider] = list(session.scalars(query).all()) if not datasource_providers: return [] copy_credentials_list = [] - default_provider = db.session.execute( + default_provider = session.execute( select(DatasourceProvider.id) .where( DatasourceProvider.tenant_id == tenant_id, @@ -870,7 +879,7 @@ class DatasourceProviderService: return copy_credentials_list - def get_all_datasource_credentials(self, tenant_id: str) -> list[dict]: + def get_all_datasource_credentials(self, tenant_id: str, *, session: Session) -> list[dict]: """ get datasource credentials. @@ -883,7 +892,10 @@ class DatasourceProviderService: for datasource in datasources: datasource_provider_id = DatasourceProviderID(f"{datasource.plugin_id}/{datasource.provider}") credentials = self.list_datasource_credentials( - tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id + tenant_id=tenant_id, + provider=datasource.provider, + plugin_id=datasource.plugin_id, + session=session, ) redirect_uri = ( f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{datasource_provider_id}/datasource/callback" @@ -912,10 +924,10 @@ class DatasourceProviderService: for credential_schema in datasource.declaration.oauth_schema.credentials_schema ], "oauth_custom_client_params": self.get_tenant_oauth_client( - tenant_id, datasource_provider_id, mask=True + tenant_id, datasource_provider_id, mask=True, session=session ), "is_oauth_custom_client_enabled": self.is_tenant_oauth_params_enabled( - tenant_id, datasource_provider_id + tenant_id, datasource_provider_id, session=session ), "is_system_oauth_params_exists": self.is_system_oauth_params_exist(datasource_provider_id), "redirect_uri": redirect_uri, @@ -926,7 +938,7 @@ class DatasourceProviderService: ) return datasource_credentials - def get_hard_code_datasource_credentials(self, tenant_id: str) -> list[dict]: + def get_hard_code_datasource_credentials(self, tenant_id: str, *, session: Session) -> list[dict]: """ get hard code datasource credentials. @@ -945,7 +957,10 @@ class DatasourceProviderService: ]: datasource_provider_id = DatasourceProviderID(f"{datasource.plugin_id}/{datasource.provider}") credentials = self.list_datasource_credentials( - tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id + tenant_id=tenant_id, + provider=datasource.provider, + plugin_id=datasource.plugin_id, + session=session, ) redirect_uri = "{}/console/api/oauth/plugin/{}/datasource/callback".format( dify_config.CONSOLE_API_URL, datasource_provider_id @@ -974,10 +989,10 @@ class DatasourceProviderService: for credential_schema in datasource.declaration.oauth_schema.credentials_schema ], "oauth_custom_client_params": self.get_tenant_oauth_client( - tenant_id, datasource_provider_id, mask=True + tenant_id, datasource_provider_id, mask=True, session=session ), "is_oauth_custom_client_enabled": self.is_tenant_oauth_params_enabled( - tenant_id, datasource_provider_id + tenant_id, datasource_provider_id, session=session ), "is_system_oauth_params_exists": self.is_system_oauth_params_exist(datasource_provider_id), "redirect_uri": redirect_uri, @@ -988,7 +1003,9 @@ class DatasourceProviderService: ) return datasource_credentials - def get_real_datasource_credentials(self, tenant_id: str, provider: str, plugin_id: str) -> list[dict]: + def get_real_datasource_credentials( + self, tenant_id: str, provider: str, plugin_id: str, *, session: Session + ) -> list[dict]: """ get datasource credentials. @@ -998,7 +1015,7 @@ class DatasourceProviderService: """ # Get all provider configurations of the current workspace datasource_providers: list[DatasourceProvider] = list( - db.session.scalars( + session.scalars( select(DatasourceProvider).where( DatasourceProvider.tenant_id == tenant_id, DatasourceProvider.provider == provider, @@ -1110,7 +1127,9 @@ class DatasourceProviderService: datasource_provider.encrypted_credentials = encrypted_credentials - def remove_datasource_credentials(self, tenant_id: str, auth_id: str, provider: str, plugin_id: str) -> None: + def remove_datasource_credentials( + self, tenant_id: str, auth_id: str, provider: str, plugin_id: str, *, session: Session + ) -> None: """ remove datasource credentials. @@ -1119,7 +1138,7 @@ class DatasourceProviderService: :param plugin_id: plugin id :return: """ - datasource_provider = db.session.scalar( + datasource_provider = session.scalar( select(DatasourceProvider) .where( DatasourceProvider.tenant_id == tenant_id, @@ -1130,5 +1149,5 @@ class DatasourceProviderService: .limit(1) ) if datasource_provider: - db.session.delete(datasource_provider) - db.session.commit() + session.delete(datasource_provider) + session.commit() diff --git a/api/services/dsl_content.py b/api/services/dsl_content.py new file mode 100644 index 00000000000..3874d0546db --- /dev/null +++ b/api/services/dsl_content.py @@ -0,0 +1,9 @@ +"""Shared DSL content size and decoding rules.""" + +DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB + + +def dsl_content_size(content: str | bytes) -> int: + if isinstance(content, bytes): + return len(content) + return len(content.encode("utf-8")) diff --git a/api/services/enterprise/account_deletion_sync.py b/api/services/enterprise/account_deletion_sync.py index b5107fb0f66..89c4b80e670 100644 --- a/api/services/enterprise/account_deletion_sync.py +++ b/api/services/enterprise/account_deletion_sync.py @@ -5,9 +5,9 @@ from datetime import UTC, datetime from redis import RedisError from sqlalchemy import select +from sqlalchemy.orm import Session from configs import dify_config -from extensions.ext_database import db from extensions.ext_redis import redis_client from models.account import TenantAccountJoin @@ -87,7 +87,7 @@ def sync_workspace_member_removal(workspace_id: str, member_id: str, *, source: return _queue_task(workspace_id=workspace_id, member_id=member_id, source=source) -def sync_account_deletion(account_id: str, *, source: str) -> bool: +def sync_account_deletion(account_id: str, *, source: str, session: Session) -> bool: """ Sync full account deletion across all workspaces (enterprise only). @@ -97,6 +97,7 @@ def sync_account_deletion(account_id: str, *, source: str) -> bool: Args: account_id: The account ID being deleted source: Source of the sync request (e.g., "account_deleted") + session: SQLAlchemy session used to fetch workspace memberships Returns: bool: True if all tasks were queued (or skipped in community), False if any queueing failed @@ -105,9 +106,7 @@ def sync_account_deletion(account_id: str, *, source: str) -> bool: return True # Fetch all workspaces the account belongs to - workspace_joins = db.session.scalars( - select(TenantAccountJoin).where(TenantAccountJoin.account_id == account_id) - ).all() + workspace_joins = session.scalars(select(TenantAccountJoin).where(TenantAccountJoin.account_id == account_id)).all() # Queue sync task for each workspace success = True diff --git a/api/services/enterprise/rbac_service.py b/api/services/enterprise/rbac_service.py index 787e0fb04ec..c8f6ca68acb 100644 --- a/api/services/enterprise/rbac_service.py +++ b/api/services/enterprise/rbac_service.py @@ -9,9 +9,9 @@ from flask import has_request_context, request from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_validator from sqlalchemy import select from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.orm import Session from configs import dify_config -from core.db.session_factory import session_factory from core.rbac import RBACResourceWhitelistScope from models import TenantAccountJoin, TenantAccountRole from services.enterprise.base import EnterpriseRequest @@ -366,7 +366,6 @@ _LEGACY_WORKSPACE_ADMIN_KEYS: list[str] = [ ] _LEGACY_WORKSPACE_EDITOR_KEYS: list[str] = [ - "workspace.member.manage", "api_extension.manage", "plugin.install", "credential.use", @@ -567,25 +566,24 @@ def _legacy_member_roles_response( ) -def _legacy_my_permissions(tenant_id: str, account_id: str | None) -> MyPermissionsResponse: +def _legacy_my_permissions(tenant_id: str, account_id: str | None, *, session: Session) -> MyPermissionsResponse: if not account_id: return MyPermissionsResponse() try: - with session_factory.create_session() as session: - role = session.scalar( - select(TenantAccountJoin.role).where( - TenantAccountJoin.tenant_id == tenant_id, - TenantAccountJoin.account_id == account_id, - ) + role = session.scalar( + select(TenantAccountJoin.role).where( + TenantAccountJoin.tenant_id == tenant_id, + TenantAccountJoin.account_id == account_id, ) - if not role: - return MyPermissionsResponse() + ) + if not role: + return MyPermissionsResponse() - try: - tenant_role = TenantAccountRole(role) - except ValueError: - return MyPermissionsResponse() + try: + tenant_role = TenantAccountRole(role) + except ValueError: + return MyPermissionsResponse() except SQLAlchemyError: return MyPermissionsResponse() @@ -602,8 +600,10 @@ def _legacy_resource_permission_keys_batch( account_id: str | None, resource_ids: list[str], resource_type: RBACResourceType, + *, + session: Session, ) -> dict[str, list[str]]: - snapshot = _legacy_my_permissions(tenant_id, account_id) + snapshot = _legacy_my_permissions(tenant_id, account_id, session=session) if resource_type == RBACResourceType.APP: permission_keys = snapshot.app.default_permission_keys else: @@ -1600,7 +1600,9 @@ class RBACService: class MemberRoles: @staticmethod - def get(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: + def get( + tenant_id: str, account_id: str | None, member_account_id: str, *, session: Session + ) -> MemberRolesResponse: if dify_config.RBAC_ENABLED: data = _inner_call( "GET", @@ -1612,14 +1614,13 @@ class RBACService: rst = MemberRolesResponse.model_validate(data or {}) return rst else: - with session_factory.create_session() as session: - role = session.scalar( - select(TenantAccountJoin.role).where( - TenantAccountJoin.tenant_id == tenant_id, - TenantAccountJoin.account_id == member_account_id, - ) + role = session.scalar( + select(TenantAccountJoin.role).where( + TenantAccountJoin.tenant_id == tenant_id, + TenantAccountJoin.account_id == member_account_id, ) - return _legacy_member_roles_response(tenant_id, member_account_id, role) + ) + return _legacy_member_roles_response(tenant_id, member_account_id, role) @staticmethod def batch_get( @@ -1649,34 +1650,35 @@ class RBACService: account_id: str | None, member_account_id: str, role_ids: list[str], + *, + session: Session, ) -> MemberRolesResponse: if not dify_config.RBAC_ENABLED: if len(role_ids) != 1: raise ValueError("Legacy workspace member role update requires exactly one role.") tenant_role = TenantAccountRole(role_ids[0]) - with session_factory.create_session() as session: - target_member_join = session.scalar( + target_member_join = session.scalar( + select(TenantAccountJoin).where( + TenantAccountJoin.tenant_id == tenant_id, + TenantAccountJoin.account_id == member_account_id, + ) + ) + if not target_member_join: + raise ValueError("Member not in tenant.") + + if tenant_role == TenantAccountRole.OWNER: + current_owner_join = session.scalar( select(TenantAccountJoin).where( TenantAccountJoin.tenant_id == tenant_id, - TenantAccountJoin.account_id == member_account_id, + TenantAccountJoin.role == TenantAccountRole.OWNER, ) ) - if not target_member_join: - raise ValueError("Member not in tenant.") + if current_owner_join and current_owner_join.account_id != member_account_id: + current_owner_join.role = TenantAccountRole.ADMIN - if tenant_role == TenantAccountRole.OWNER: - current_owner_join = session.scalar( - select(TenantAccountJoin).where( - TenantAccountJoin.tenant_id == tenant_id, - TenantAccountJoin.role == TenantAccountRole.OWNER, - ) - ) - if current_owner_join and current_owner_join.account_id != member_account_id: - current_owner_join.role = TenantAccountRole.ADMIN - - target_member_join.role = tenant_role - session.commit() + target_member_join.role = tenant_role + session.commit() return _legacy_member_roles_response(tenant_id, member_account_id, tenant_role) @@ -1742,11 +1744,15 @@ class RBACService: tenant_id: str, account_id: str | None, app_ids: list[str], + *, + session: Session, ) -> dict[str, list[str]]: if not app_ids: return {} if not dify_config.RBAC_ENABLED: - return _legacy_resource_permission_keys_batch(tenant_id, account_id, app_ids, RBACResourceType.APP) + return _legacy_resource_permission_keys_batch( + tenant_id, account_id, app_ids, RBACResourceType.APP, session=session + ) data = _inner_call( "POST", f"{_INNER_PREFIX}/apps/permission-keys/batch", @@ -1762,12 +1768,14 @@ class RBACService: tenant_id: str, account_id: str | None, dataset_ids: list[str], + *, + session: Session, ) -> dict[str, list[str]]: if not dataset_ids: return {} if not dify_config.RBAC_ENABLED: return _legacy_resource_permission_keys_batch( - tenant_id, account_id, dataset_ids, RBACResourceType.DATASET + tenant_id, account_id, dataset_ids, RBACResourceType.DATASET, session=session ) data = _inner_call( "POST", @@ -1786,9 +1794,10 @@ class RBACService: *, app_id: str | None = None, dataset_id: str | None = None, + session: Session, ) -> MyPermissionsResponse: if not dify_config.RBAC_ENABLED: - return _legacy_my_permissions(tenant_id, account_id) + return _legacy_my_permissions(tenant_id, account_id, session=session) data = _inner_call( "GET", diff --git a/api/services/errors/account.py b/api/services/errors/account.py index 4d3d150e072..700c1dd4aaf 100644 --- a/api/services/errors/account.py +++ b/api/services/errors/account.py @@ -17,6 +17,14 @@ class AccountPasswordError(BaseServiceError): pass +class RefreshTokenNotFoundError(BaseServiceError): + pass + + +class RefreshTokenAccountNotFoundError(BaseServiceError): + pass + + class AccountNotLinkTenantError(BaseServiceError): pass diff --git a/api/services/external_knowledge_service.py b/api/services/external_knowledge_service.py index 42e7eca29d7..cdd6c48342e 100644 --- a/api/services/external_knowledge_service.py +++ b/api/services/external_knowledge_service.py @@ -10,6 +10,7 @@ from sqlalchemy.orm import Session from constants import HIDDEN_VALUE from core.helper import ssrf_proxy from core.rag.entities import MetadataFilteringCondition +from extensions.ext_database import db # noqa: F401 from graphon.nodes.http_request.exc import InvalidHttpMethodError from libs.datetime_utils import naive_utc_now from libs.pagination import paginate_query @@ -57,7 +58,7 @@ class ExternalDatasetService: @staticmethod def create_external_knowledge_api( - tenant_id: str, user_id: str, args: dict[str, Any], session: Session + tenant_id: str, user_id: str, args: dict[str, Any], *, session: Session ) -> ExternalKnowledgeApis: settings = args.get("settings") if settings is None: @@ -105,7 +106,7 @@ class ExternalDatasetService: @staticmethod def get_external_knowledge_api( - session: Session, external_knowledge_api_id: str, tenant_id: str + external_knowledge_api_id: str, tenant_id: str, *, session: Session ) -> ExternalKnowledgeApis: external_knowledge_api: ExternalKnowledgeApis | None = session.scalar( select(ExternalKnowledgeApis) @@ -118,7 +119,12 @@ class ExternalDatasetService: @staticmethod def update_external_knowledge_api( - session: Session, tenant_id: str, user_id: str, external_knowledge_api_id: str, args + tenant_id: str, + user_id: str, + external_knowledge_api_id: str, + args: dict[str, Any], + *, + session: Session, ) -> ExternalKnowledgeApis: external_knowledge_api: ExternalKnowledgeApis | None = session.scalar( select(ExternalKnowledgeApis) @@ -131,9 +137,9 @@ class ExternalDatasetService: if settings and settings.get("api_key") == HIDDEN_VALUE and external_knowledge_api.settings_dict: settings["api_key"] = external_knowledge_api.settings_dict.get("api_key") - external_knowledge_api.name = args.get("name") - external_knowledge_api.description = args.get("description", "") - external_knowledge_api.settings = json.dumps(args.get("settings"), ensure_ascii=False) + external_knowledge_api.name = str(args.get("name")) + external_knowledge_api.description = str(args.get("description", "")) + external_knowledge_api.settings = json.dumps(settings, ensure_ascii=False) external_knowledge_api.updated_by = user_id external_knowledge_api.updated_at = naive_utc_now() session.commit() @@ -141,7 +147,7 @@ class ExternalDatasetService: return external_knowledge_api @staticmethod - def delete_external_knowledge_api(session: Session, tenant_id: str, external_knowledge_api_id: str): + def delete_external_knowledge_api(tenant_id: str, external_knowledge_api_id: str, *, session: Session) -> None: external_knowledge_api = session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) @@ -155,7 +161,7 @@ class ExternalDatasetService: @staticmethod def external_knowledge_api_use_check( - session: Session, external_knowledge_api_id: str, tenant_id: str + external_knowledge_api_id: str, tenant_id: str, *, session: Session ) -> tuple[bool, int]: """ Return usage for an external knowledge API within a single tenant. @@ -176,7 +182,7 @@ class ExternalDatasetService: @staticmethod def get_external_knowledge_binding_with_dataset_id( - session: Session, tenant_id: str, dataset_id: str + tenant_id: str, dataset_id: str, *, session: Session ) -> ExternalKnowledgeBindings: external_knowledge_binding: ExternalKnowledgeBindings | None = session.scalar( select(ExternalKnowledgeBindings) @@ -189,8 +195,12 @@ class ExternalDatasetService: @staticmethod def document_create_args_validate( - session: Session, tenant_id: str, external_knowledge_api_id: str, process_parameter: dict[str, Any] - ): + tenant_id: str, + external_knowledge_api_id: str, + process_parameter: dict[str, Any], + *, + session: Session, + ) -> None: external_knowledge_api = session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) @@ -264,7 +274,7 @@ class ExternalDatasetService: return ExternalKnowledgeApiSetting.model_validate(settings) @staticmethod - def create_external_dataset(tenant_id: str, user_id: str, args: dict[str, Any], session: Session) -> Dataset: + def create_external_dataset(tenant_id: str, user_id: str, args: dict[str, Any], *, session: Session) -> Dataset: # check if dataset name already exists if session.scalar( select(Dataset).where(Dataset.name == args.get("name"), Dataset.tenant_id == tenant_id).limit(1) @@ -314,12 +324,13 @@ class ExternalDatasetService: @staticmethod def fetch_external_knowledge_retrieval( - session: Session, tenant_id: str, dataset_id: str, query: str, external_retrieval_parameters: dict[str, Any], metadata_condition: MetadataFilteringCondition | None = None, + *, + session: Session, ): """Fetch retrieval records from an external knowledge provider. diff --git a/api/services/file_service.py b/api/services/file_service.py index e41d74ad3eb..ec69af4e80c 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -268,7 +268,7 @@ class FileService: @staticmethod def get_upload_files_by_ids( - session: Session, tenant_id: str, upload_file_ids: Sequence[str] + tenant_id: str, upload_file_ids: Sequence[str], *, session: Session ) -> dict[str, UploadFile]: """ Fetch `UploadFile` rows for a tenant in a single batch query. diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index 1b51a2d279b..1bfa4025fa0 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -4,7 +4,7 @@ import time from typing import Any, TypedDict, cast from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from core.app.app_config.entities import ModelConfig from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService @@ -56,9 +56,7 @@ class HitTestingService: } @classmethod - def _dump_retrieval_records( - cls, session: Session | scoped_session, records: list[RetrievalSegments] - ) -> list[dict[str, Any]]: + def _dump_retrieval_records(cls, session: Session, records: list[RetrievalSegments]) -> list[dict[str, Any]]: document_ids = { document_id for record in records @@ -105,7 +103,6 @@ class HitTestingService: @classmethod def retrieve( cls, - session: Session, dataset: Dataset, query: str, account: Account, @@ -113,6 +110,8 @@ class HitTestingService: external_retrieval_model: dict[str, Any], attachment_ids: list | None = None, limit: int = 10, + *, + session: Session, ): start = time.perf_counter() @@ -144,7 +143,7 @@ class HitTestingService: if metadata_filter_document_ids: document_ids_filter = metadata_filter_document_ids.get(dataset.id, []) if metadata_condition and not document_ids_filter: - return cls.compact_retrieve_response(session, query, []) + return cls.compact_retrieve_response(query, [], session=session) all_documents = RetrievalService.retrieve( retrieval_method=RetrievalMethod( resolved_retrieval_model.get("search_method", RetrievalMethod.SEMANTIC_SEARCH) @@ -186,17 +185,18 @@ class HitTestingService: session.add(dataset_query) session.commit() - return cls.compact_retrieve_response(session, query, all_documents) + return cls.compact_retrieve_response(query, all_documents, session=session) @classmethod def external_retrieve( cls, - session: Session, dataset: Dataset, query: str, account: Account, external_retrieval_model: dict[str, Any] | None = None, metadata_filtering_conditions: dict[str, Any] | None = None, + *, + session: Session, ): if dataset.provider != "external": return { @@ -233,7 +233,7 @@ class HitTestingService: @classmethod def compact_retrieve_response( - cls, session: Session | scoped_session, query: str, documents: list[Document] + cls, query: str, documents: list[Document], *, session: Session ) -> RetrieveResponseDict: records = RetrievalService.format_retrieval_documents(documents) diff --git a/api/services/message_service.py b/api/services/message_service.py index e8d1b6232bc..4fbeb61e1f7 100644 --- a/api/services/message_service.py +++ b/api/services/message_service.py @@ -3,7 +3,7 @@ from collections.abc import Sequence from typing import cast from sqlalchemy import select -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager from core.app.entities.app_invoke_entities import InvokeFrom @@ -70,6 +70,8 @@ class MessageService: first_id: str | None, limit: int, order: str = "asc", + *, + session: Session, ) -> InfiniteScrollPagination: if not user: return InfiniteScrollPagination(data=[], limit=limit, has_more=False) @@ -78,20 +80,20 @@ class MessageService: return InfiniteScrollPagination(data=[], limit=limit, has_more=False) conversation = ConversationService.get_conversation( - app_model=app_model, user=user, conversation_id=conversation_id + app_model=app_model, user=user, conversation_id=conversation_id, session=session ) fetch_limit = limit + 1 if first_id: - first_message = db.session.scalar( + first_message = session.scalar( select(Message).where(Message.conversation_id == conversation.id, Message.id == first_id).limit(1) ) if not first_message: raise FirstMessageNotExistsError() - history_messages = db.session.scalars( + history_messages = session.scalars( select(Message) .where( Message.conversation_id == conversation.id, @@ -102,7 +104,7 @@ class MessageService: .limit(fetch_limit) ).all() else: - history_messages = db.session.scalars( + history_messages = session.scalars( select(Message) .where(Message.conversation_id == conversation.id) .order_by(Message.created_at.desc()) @@ -130,6 +132,8 @@ class MessageService: limit: int, conversation_id: str | None = None, include_ids: list | None = None, + *, + session: Session, ) -> InfiniteScrollPagination: if not user: return InfiniteScrollPagination(data=[], limit=limit, has_more=False) @@ -140,7 +144,7 @@ class MessageService: if conversation_id is not None: conversation = ConversationService.get_conversation( - app_model=app_model, user=user, conversation_id=conversation_id + app_model=app_model, user=user, conversation_id=conversation_id, session=session ) stmt = stmt.where(Message.conversation_id == conversation.id) @@ -152,18 +156,18 @@ class MessageService: stmt = stmt.where(Message.id.in_(include_ids)) if last_id: - last_message = db.session.scalar(stmt.where(Message.id == last_id).limit(1)) + last_message = session.scalar(stmt.where(Message.id == last_id).limit(1)) if not last_message: raise LastMessageNotExistsError() - history_messages = db.session.scalars( + history_messages = session.scalars( stmt.where(Message.created_at < last_message.created_at, Message.id != last_message.id) .order_by(Message.created_at.desc()) .limit(fetch_limit) ).all() else: - history_messages = db.session.scalars(stmt.order_by(Message.created_at.desc()).limit(fetch_limit)).all() + history_messages = session.scalars(stmt.order_by(Message.created_at.desc()).limit(fetch_limit)).all() has_more = False if len(history_messages) > limit: @@ -181,16 +185,17 @@ class MessageService: user: Account | EndUser | None, rating: FeedbackRating | None, content: str | None, + session: Session, ): if not user: raise ValueError("user cannot be None") - message = cls.get_message(app_model=app_model, user=user, message_id=message_id) + message = cls.get_message(app_model=app_model, user=user, message_id=message_id, session=session) feedback = message.user_feedback if isinstance(user, EndUser) else message.admin_feedback if not rating and feedback: - db.session.delete(feedback) + session.delete(feedback) elif rating and feedback: feedback.rating = rating feedback.content = content @@ -208,17 +213,17 @@ class MessageService: from_end_user_id=(user.id if isinstance(user, EndUser) else None), from_account_id=(user.id if isinstance(user, Account) else None), ) - db.session.add(feedback) + session.add(feedback) - db.session.commit() + session.commit() return feedback @classmethod - def get_all_messages_feedbacks(cls, app_model: App, page: int, limit: int): + def get_all_messages_feedbacks(cls, app_model: App, page: int, limit: int, *, session: Session): """Get all feedbacks of an app""" offset = (page - 1) * limit - feedbacks = db.session.scalars( + feedbacks = session.scalars( select(MessageFeedback) .where(MessageFeedback.app_id == app_model.id) .order_by(MessageFeedback.created_at.desc(), MessageFeedback.id.desc()) @@ -229,8 +234,8 @@ class MessageService: return [record.to_dict() for record in feedbacks] @classmethod - def get_message(cls, app_model: App, user: Account | EndUser | None, message_id: str): - message = db.session.scalar( + def get_message(cls, app_model: App, user: Account | EndUser | None, message_id: str, *, session: Session): + message = session.scalar( select(Message) .where( Message.id == message_id, @@ -249,15 +254,21 @@ class MessageService: @classmethod def get_suggested_questions_after_answer( - cls, app_model: App, user: Account | EndUser | None, message_id: str, invoke_from: InvokeFrom + cls, + app_model: App, + user: Account | EndUser | None, + message_id: str, + invoke_from: InvokeFrom, + *, + session: Session, ) -> list[str]: if not user: raise ValueError("user cannot be None") - message = cls.get_message(app_model=app_model, user=user, message_id=message_id) + message = cls.get_message(app_model=app_model, user=user, message_id=message_id, session=session) conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=message.conversation_id, user=user + app_model=app_model, conversation_id=message.conversation_id, user=user, session=session ) model_manager = ModelManager.for_tenant(tenant_id=app_model.tenant_id) @@ -266,9 +277,9 @@ class MessageService: if app_model.mode == AppMode.ADVANCED_CHAT: workflow_service = WorkflowService() if invoke_from == InvokeFrom.DEBUGGER: - workflow = workflow_service.get_draft_workflow(app_model=app_model) + workflow = workflow_service.get_draft_workflow(app_model=app_model, session=session) else: - workflow = workflow_service.get_published_workflow(app_model=app_model) + workflow = workflow_service.get_published_workflow(app_model=app_model, session=session) if workflow is None: return [] @@ -288,7 +299,7 @@ class MessageService: ) else: if not conversation.override_model_configs: - app_model_config = db.session.scalar( + app_model_config = session.scalar( select(AppModelConfig) .where(AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id) .limit(1) diff --git a/api/services/metadata_service.py b/api/services/metadata_service.py index 4e83858ea0e..481eb3b2e29 100644 --- a/api/services/metadata_service.py +++ b/api/services/metadata_service.py @@ -23,11 +23,12 @@ logger = logging.getLogger(__name__) class MetadataService: @staticmethod def create_metadata( - session: Session, dataset_id: str, metadata_args: MetadataArgs, current_user: Account | None = None, # TODO: the service_api is not migrated yet current_tenant_id: str | None = None, + *, + session: Session, ) -> DatasetMetadata: # check if metadata name is too long if len(metadata_args.name) > 255: @@ -60,12 +61,13 @@ class MetadataService: @staticmethod def update_metadata_name( - session: Session, dataset_id: str, metadata_id: str, name: str, current_user: Account | None = None, current_tenant_id: str | None = None, # TODO: the service_api is not migrated yet + *, + session: Session, ) -> DatasetMetadata | None: # check if metadata name is too long if len(name) > 255: @@ -126,7 +128,7 @@ class MetadataService: redis_client.delete(lock_key) @staticmethod - def delete_metadata(session: Session, dataset_id: str, metadata_id: str): + def delete_metadata(dataset_id: str, metadata_id: str, *, session: Session): lock_key = f"dataset_metadata_lock_{dataset_id}" try: MetadataService.knowledge_base_metadata_lock_check(dataset_id, None) @@ -172,7 +174,7 @@ class MetadataService: ] @staticmethod - def enable_built_in_field(session: Session, dataset: Dataset): + def enable_built_in_field(dataset: Dataset, *, session: Session): if dataset.built_in_field_enabled: return lock_key = f"dataset_metadata_lock_{dataset.id}" @@ -201,7 +203,7 @@ class MetadataService: redis_client.delete(lock_key) @staticmethod - def disable_built_in_field(session: Session, dataset: Dataset): + def disable_built_in_field(dataset: Dataset, *, session: Session): if not dataset.built_in_field_enabled: return lock_key = f"dataset_metadata_lock_{dataset.id}" @@ -233,11 +235,12 @@ class MetadataService: @staticmethod def update_documents_metadata( - session: Session, dataset: Dataset, metadata_args: MetadataOperationData, current_user: Account | None = None, # TODO: the service_api is not migrated yet current_tenant_id: str | None = None, + *, + session: Session, ): current_user, current_tenant_id = resolve_account_fallback( current_user, current_tenant_id, fallback_tenant_id=dataset.tenant_id @@ -316,7 +319,7 @@ class MetadataService: redis_client.set(lock_key, 1, ex=3600) @staticmethod - def get_dataset_metadatas(session: Session, dataset: Dataset): + def get_dataset_metadatas(dataset: Dataset, *, session: Session): return { "doc_metadata": [ { diff --git a/api/services/model_load_balancing_service.py b/api/services/model_load_balancing_service.py index 2a9094a35f2..6eab1ffbe3e 100644 --- a/api/services/model_load_balancing_service.py +++ b/api/services/model_load_balancing_service.py @@ -3,6 +3,7 @@ import logging from typing import Any, TypedDict, cast from sqlalchemy import or_, select +from sqlalchemy.orm import Session from constants import HIDDEN_VALUE from core.entities.provider_configuration import ProviderConfiguration @@ -14,7 +15,6 @@ from core.helper.model_provider_cache import ( from core.model_manager import LBModelManager from core.plugin.impl.model_runtime_factory import create_plugin_model_assembly, create_plugin_provider_manager from core.provider_manager import ProviderConfigurationCacheSource, ProviderManager -from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.entities.provider_entities import ( ModelCredentialSchema, @@ -93,7 +93,13 @@ class ModelLoadBalancingService: provider_configuration.disable_model_load_balancing(model=model, model_type=ModelType(model_type)) def get_load_balancing_configs( - self, tenant_id: str, provider: str, model: str, model_type: str, config_from: str = "" + self, + tenant_id: str, + provider: str, + model: str, + model_type: str, + session: Session, + config_from: str = "", ) -> tuple[bool, list[LoadBalancingConfigSummaryDict]]: """ Get load balancing configurations. @@ -131,7 +137,7 @@ class ModelLoadBalancingService: # Get load balancing configurations load_balancing_configs = list( - db.session.scalars( + session.scalars( select(LoadBalancingModelConfig) .where( LoadBalancingModelConfig.tenant_id == tenant_id, @@ -158,7 +164,7 @@ class ModelLoadBalancingService: if not inherit_config_exists: # Initialize the inherit configuration - inherit_config = self._init_inherit_config(tenant_id, provider, model, model_type_enum) + inherit_config = self._init_inherit_config(tenant_id, provider, model, model_type_enum, session=session) # prepend the inherit configuration load_balancing_configs.insert(0, inherit_config) @@ -233,7 +239,13 @@ class ModelLoadBalancingService: return is_load_balancing_enabled, datas def get_load_balancing_config( - self, tenant_id: str, provider: str, model: str, model_type: str, config_id: str + self, + tenant_id: str, + provider: str, + model: str, + model_type: str, + config_id: str, + session: Session, ) -> LoadBalancingConfigDetailDict | None: """ Get load balancing configuration. @@ -256,7 +268,7 @@ class ModelLoadBalancingService: model_type_enum = ModelType(model_type) # Get load balancing configurations - load_balancing_model_config = db.session.scalar( + load_balancing_model_config = session.scalar( select(LoadBalancingModelConfig) .where( LoadBalancingModelConfig.tenant_id == tenant_id, @@ -296,7 +308,12 @@ class ModelLoadBalancingService: return result def _init_inherit_config( - self, tenant_id: str, provider: str, model: str, model_type: ModelType + self, + tenant_id: str, + provider: str, + model: str, + model_type: ModelType, + session: Session, ) -> LoadBalancingModelConfig: """ Initialize the inherit configuration. @@ -314,8 +331,8 @@ class ModelLoadBalancingService: model_name=model, name="__inherit__", ) - db.session.add(inherit_config) - db.session.commit() + session.add(inherit_config) + session.commit() ProviderManager.invalidate_configurations_cache( tenant_id, sources=(ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS,), @@ -324,7 +341,14 @@ class ModelLoadBalancingService: return inherit_config def update_load_balancing_configs( - self, tenant_id: str, provider: str, model: str, model_type: str, configs: list[dict], config_from: str + self, + tenant_id: str, + provider: str, + model: str, + model_type: str, + configs: list[dict], + config_from: str, + session: Session, ): """ Update load balancing configurations. @@ -350,7 +374,7 @@ class ModelLoadBalancingService: if not isinstance(configs, list): raise ValueError("Invalid load balancing configs") - current_load_balancing_configs = db.session.scalars( + current_load_balancing_configs = session.scalars( select(LoadBalancingModelConfig).where( LoadBalancingModelConfig.tenant_id == tenant_id, LoadBalancingModelConfig.provider_name == provider_configuration.provider.provider, @@ -377,7 +401,7 @@ class ModelLoadBalancingService: if credential_id: if config_from == "predefined-model": - credential_record = db.session.scalar( + credential_record = session.scalar( select(ProviderCredential) .where( ProviderCredential.id == credential_id, @@ -387,7 +411,7 @@ class ModelLoadBalancingService: .limit(1) ) else: - credential_record = db.session.scalar( + credential_record = session.scalar( select(ProviderModelCredential) .where( ProviderModelCredential.id == credential_id, @@ -440,7 +464,7 @@ class ModelLoadBalancingService: load_balancing_config.name = name load_balancing_config.enabled = enabled load_balancing_config.updated_at = naive_utc_now() - db.session.commit() + session.commit() ProviderManager.invalidate_configurations_cache( tenant_id, sources=(ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS,), @@ -496,8 +520,8 @@ class ModelLoadBalancingService: encrypted_config=json.dumps(credentials), ) - db.session.add(load_balancing_model_config) - db.session.commit() + session.add(load_balancing_model_config) + session.commit() ProviderManager.invalidate_configurations_cache( tenant_id, sources=(ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS,), @@ -506,8 +530,8 @@ class ModelLoadBalancingService: # get deleted config ids deleted_config_ids = set(current_load_balancing_configs_dict.keys()) - updated_config_ids for config_id in deleted_config_ids: - db.session.delete(current_load_balancing_configs_dict[config_id]) - db.session.commit() + session.delete(current_load_balancing_configs_dict[config_id]) + session.commit() ProviderManager.invalidate_configurations_cache( tenant_id, sources=(ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS,), @@ -522,6 +546,7 @@ class ModelLoadBalancingService: model: str, model_type: str, credentials: dict[str, Any], + session: Session, config_id: str | None = None, ): """ @@ -548,7 +573,7 @@ class ModelLoadBalancingService: load_balancing_model_config = None if config_id: # Get load balancing config - load_balancing_model_config = db.session.scalar( + load_balancing_model_config = session.scalar( select(LoadBalancingModelConfig) .where( LoadBalancingModelConfig.tenant_id == tenant_id, diff --git a/api/services/oauth_device_flow.py b/api/services/oauth_device_flow.py index 9ec5711890b..9e59b8c326a 100644 --- a/api/services/oauth_device_flow.py +++ b/api/services/oauth_device_flow.py @@ -13,7 +13,7 @@ from enum import StrEnum from typing import Any, NotRequired, TypedDict from sqlalchemy import and_, func, select, update -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from libs.oauth_bearer import TOKEN_CACHE_KEY_FMT, AuthContext, SubjectType from models.oauth import OAuthAccessToken @@ -335,9 +335,6 @@ def sha256_hex(token: str) -> str: def mint_oauth_token( - # Accept either Session or Flask-SQLAlchemy's request-scoped wrapper — - # the wrapper proxies the same execute/commit surface. - session: Session | scoped_session, redis_client, *, subject_email: str, @@ -347,6 +344,7 @@ def mint_oauth_token( device_label: str, prefix: str, ttl_days: int, + session: Session, ) -> MintResult: """Live row rotates in place via partial unique index ``uq_oauth_active_per_device``; hard-expired rows are excluded by the @@ -390,7 +388,7 @@ def mint_oauth_token( def _upsert( - session: Session | scoped_session, + session: Session, *, subject_email: str, subject_issuer: str | None, @@ -501,11 +499,7 @@ def subject_match_clauses(ctx: AuthContext) -> tuple[Any, ...]: ) -def list_active_sessions( - session: Session | scoped_session, - ctx: AuthContext, - now: datetime, -) -> list[OAuthAccessToken]: +def list_active_sessions(ctx: AuthContext, now: datetime, *, session: Session) -> list[OAuthAccessToken]: return list( session.execute( select(OAuthAccessToken) @@ -524,11 +518,7 @@ def list_active_sessions( ) -def token_belongs_to_subject( - session: Session | scoped_session, - token_id: str, - ctx: AuthContext, -) -> bool: +def token_belongs_to_subject(token_id: str, ctx: AuthContext, *, session: Session) -> bool: row = session.execute( select(OAuthAccessToken.id).where( and_( @@ -540,11 +530,7 @@ def token_belongs_to_subject( return row is not None -def revoke_oauth_token( - session: Session | scoped_session, - redis_client: Any, - token_id: str, -) -> None: +def revoke_oauth_token(redis_client: Any, token_id: str, *, session: Session) -> None: row = ( session.query(OAuthAccessToken.token_hash) .filter( diff --git a/api/services/oauth_server.py b/api/services/oauth_server.py index 5f3277c9525..47aa1bc99bf 100644 --- a/api/services/oauth_server.py +++ b/api/services/oauth_server.py @@ -2,7 +2,7 @@ import enum import uuid from sqlalchemy import select -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import BadRequest from extensions.ext_database import db @@ -83,7 +83,7 @@ class OAuthServerService: return token @staticmethod - def validate_oauth_access_token(client_id: str, token: str) -> Account | None: + def validate_oauth_access_token(client_id: str, token: str, session: Session) -> Account | None: redis_key = OAUTH_ACCESS_TOKEN_REDIS_KEY.format(client_id=client_id, token=token) user_account_id = redis_client.get(redis_key) if not user_account_id: @@ -91,4 +91,4 @@ class OAuthServerService: user_id_str = user_account_id.decode("utf-8") - return AccountService.load_user(user_id_str, db.session) + return AccountService.load_user(user_id_str, session) diff --git a/api/services/ops_service.py b/api/services/ops_service.py index 3ad42faf249..b6f17168b3c 100644 --- a/api/services/ops_service.py +++ b/api/services/ops_service.py @@ -1,23 +1,23 @@ from typing import Any from sqlalchemy import select +from sqlalchemy.orm import Session from core.ops.entities.config_entity import BaseTracingConfig from core.ops.ops_trace_manager import OpsTraceManager, TracingProviderConfigEntry, provider_config_map -from extensions.ext_database import db from models.model import App, TraceAppConfig class OpsService: @classmethod - def get_tracing_app_config(cls, app_id: str, tracing_provider: str): + def get_tracing_app_config(cls, app_id: str, tracing_provider: str, session: Session): """ Get tracing app config :param app_id: app id :param tracing_provider: tracing provider :return: """ - trace_config_data: TraceAppConfig | None = db.session.scalar( + trace_config_data: TraceAppConfig | None = session.scalar( select(TraceAppConfig) .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) .limit(1) @@ -27,7 +27,7 @@ class OpsService: return None # decrypt_token and obfuscated_token - app = db.session.get(App, app_id) + app = session.get(App, app_id) if not app: return None tenant_id = app.tenant_id @@ -137,7 +137,9 @@ class OpsService: return trace_config_data.to_dict() @classmethod - def create_tracing_app_config(cls, app_id: str, tracing_provider: str, tracing_config: dict[str, Any]): + def create_tracing_app_config( + cls, app_id: str, tracing_provider: str, tracing_config: dict[str, Any], session: Session + ): """ Create tracing app config :param app_id: app id @@ -184,7 +186,7 @@ class OpsService: project_url = None # check if trace config already exists - trace_config_data: TraceAppConfig | None = db.session.scalar( + trace_config_data: TraceAppConfig | None = session.scalar( select(TraceAppConfig) .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) .limit(1) @@ -194,7 +196,7 @@ class OpsService: return None # get tenant id - app = db.session.get(App, app_id) + app = session.get(App, app_id) if not app: return None tenant_id = app.tenant_id @@ -206,13 +208,15 @@ class OpsService: tracing_provider=tracing_provider, tracing_config=tracing_config, ) - db.session.add(trace_config_data) - db.session.commit() + session.add(trace_config_data) + session.commit() return {"result": "success"} @classmethod - def update_tracing_app_config(cls, app_id: str, tracing_provider: str, tracing_config: dict[str, Any]): + def update_tracing_app_config( + cls, app_id: str, tracing_provider: str, tracing_config: dict[str, Any], session: Session + ): """ Update tracing app config :param app_id: app id @@ -226,7 +230,7 @@ class OpsService: raise ValueError(f"Invalid tracing provider: {tracing_provider}") # check if trace config already exists - current_trace_config = db.session.scalar( + current_trace_config = session.scalar( select(TraceAppConfig) .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) .limit(1) @@ -236,7 +240,7 @@ class OpsService: return None # get tenant id - app = db.session.get(App, app_id) + app = session.get(App, app_id) if not app: return None tenant_id = app.tenant_id @@ -251,19 +255,19 @@ class OpsService: raise ValueError("Invalid Credentials") current_trace_config.tracing_config = tracing_config - db.session.commit() + session.commit() return current_trace_config.to_dict() @classmethod - def delete_tracing_app_config(cls, app_id: str, tracing_provider: str): + def delete_tracing_app_config(cls, app_id: str, tracing_provider: str, session: Session): """ Delete tracing app config :param app_id: app id :param tracing_provider: tracing provider :return: """ - trace_config = db.session.scalar( + trace_config = session.scalar( select(TraceAppConfig) .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) .limit(1) @@ -272,7 +276,7 @@ class OpsService: if not trace_config: return None - db.session.delete(trace_config) - db.session.commit() + session.delete(trace_config) + session.commit() return True diff --git a/api/services/plugin/dependencies_analysis.py b/api/services/plugin/dependencies_analysis.py index 2f0c5ae3af5..4c83285ebe5 100644 --- a/api/services/plugin/dependencies_analysis.py +++ b/api/services/plugin/dependencies_analysis.py @@ -2,7 +2,7 @@ import re from configs import dify_config from core.helper import marketplace -from core.plugin.entities.plugin import PluginDependency, PluginInstallationSource +from core.plugin.entities.plugin import PluginDependency, PluginDependencyType, PluginInstallationSource from core.plugin.impl.plugin import PluginInstaller from models.provider_ids import ModelProviderID, ToolProviderID @@ -55,7 +55,7 @@ class DependenciesAnalysisService: unique_identifier = dependency.value.plugin_unique_identifier if unique_identifier in missing_plugin_unique_identifiers: # Extract version for Marketplace dependencies - if dependency.type == PluginDependency.Type.Marketplace: + if dependency.type == PluginDependencyType.Marketplace: version_match = _VERSION_REGEX.search(unique_identifier) if version_match: dependency.value.version = version_match.group("version") @@ -84,7 +84,7 @@ class DependenciesAnalysisService: if plugin.source == PluginInstallationSource.Github: result.append( PluginDependency( - type=PluginDependency.Type.Github, + type=PluginDependencyType.Github, value=PluginDependency.Github( repo=plugin.meta["repo"], version=plugin.meta["version"], @@ -96,7 +96,7 @@ class DependenciesAnalysisService: elif plugin.source == PluginInstallationSource.Marketplace: result.append( PluginDependency( - type=PluginDependency.Type.Marketplace, + type=PluginDependencyType.Marketplace, value=PluginDependency.Marketplace( marketplace_plugin_unique_identifier=plugin.plugin_unique_identifier ), @@ -105,7 +105,7 @@ class DependenciesAnalysisService: elif plugin.source == PluginInstallationSource.Package: result.append( PluginDependency( - type=PluginDependency.Type.Package, + type=PluginDependencyType.Package, value=PluginDependency.Package(plugin_unique_identifier=plugin.plugin_unique_identifier), ) ) @@ -130,7 +130,7 @@ class DependenciesAnalysisService: deps = marketplace.batch_fetch_plugin_manifests(dependencies) return [ PluginDependency( - type=PluginDependency.Type.Marketplace, + type=PluginDependencyType.Marketplace, value=PluginDependency.Marketplace(marketplace_plugin_unique_identifier=dep.latest_package_identifier), ) for dep in deps diff --git a/api/services/plugin/plugin_auto_upgrade_service.py b/api/services/plugin/plugin_auto_upgrade_service.py index 0e13214ee77..79770063016 100644 --- a/api/services/plugin/plugin_auto_upgrade_service.py +++ b/api/services/plugin/plugin_auto_upgrade_service.py @@ -12,13 +12,17 @@ from hashlib import sha256 from sqlalchemy import select from sqlalchemy.orm import Session -from core.db.session_factory import session_factory from core.plugin.impl.plugin import PluginInstaller -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import ( + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategy, + TenantPluginAutoUpgradeStrategySetting, +) logger = logging.getLogger(__name__) -PluginCategory = TenantPluginAutoUpgradeStrategy.PluginCategory +PluginCategory = TenantPluginAutoUpgradeCategory PLUGIN_CATEGORIES = tuple(PluginCategory) SECONDS_PER_DAY = 24 * 60 * 60 AUTO_UPGRADE_CHECK_SLOT_SECONDS = 15 * 60 @@ -35,10 +39,10 @@ class PluginAutoUpgradeService: @staticmethod def default_strategy_setting_for_category( category: PluginCategory, - ) -> TenantPluginAutoUpgradeStrategy.StrategySetting: + ) -> TenantPluginAutoUpgradeStrategySetting: if category == PluginCategory.MODEL: - return TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST - return TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY + return TenantPluginAutoUpgradeStrategySetting.LATEST + return TenantPluginAutoUpgradeStrategySetting.FIX_ONLY @staticmethod def default_upgrade_time_of_day(tenant_id: str) -> int: @@ -102,9 +106,9 @@ class PluginAutoUpgradeService: @staticmethod def _has_default_strategy(strategy: TenantPluginAutoUpgradeStrategy) -> bool: return ( - strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY + strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.FIX_ONLY and strategy.upgrade_time_of_day == 0 - and strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE + and strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE and not strategy.exclude_plugins and not strategy.include_plugins ) @@ -114,7 +118,7 @@ class PluginAutoUpgradeService: source_strategy: TenantPluginAutoUpgradeStrategy, category: PluginCategory, source_has_default_strategy: bool, - ) -> TenantPluginAutoUpgradeStrategy.StrategySetting: + ) -> TenantPluginAutoUpgradeStrategySetting: # Only pure legacy defaults adopt the new model=latest default. User-edited # strategies keep their original setting across all categories. if source_has_default_strategy: @@ -136,6 +140,8 @@ class PluginAutoUpgradeService: @staticmethod def backfill_strategy_categories( tenant_id: str, + *, + session: Session, ) -> PluginAutoUpgradeBackfillResult: """Create missing category strategies and split include/exclude lists when needed. @@ -143,89 +149,85 @@ class PluginAutoUpgradeService: New category rows copy it first, then plugin lists are narrowed by real plugin category when the source strategy contains include/exclude IDs. """ - with session_factory.create_session() as session, session.begin(): - strategies = list( - session.scalars( - select(TenantPluginAutoUpgradeStrategy).where( - TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id - ) - ).all() + strategies = list( + session.scalars( + select(TenantPluginAutoUpgradeStrategy).where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id) + ).all() + ) + if not strategies: + return PluginAutoUpgradeBackfillResult(created_count=0, normalized=False) + + # Schema migration marks the historical workspace-level row as tool. + source_strategy = next( + (strategy for strategy in strategies if strategy.category == PluginCategory.TOOL), + strategies[0], + ) + source_has_default_strategy = PluginAutoUpgradeService._has_default_strategy(source_strategy) + strategies_by_category = {strategy.category: strategy for strategy in strategies} + exclude_plugins = source_strategy.exclude_plugins + include_plugins = source_strategy.include_plugins + should_split_plugin_lists = bool(exclude_plugins or include_plugins) + # Query daemon only for tenants that actually customized plugin lists. + plugin_categories = ( + PluginAutoUpgradeService._get_installed_plugin_categories(tenant_id) if should_split_plugin_lists else {} + ) + if should_split_plugin_lists: + PluginAutoUpgradeService._log_unknown_plugin_ids( + tenant_id, + "exclude_plugins", + exclude_plugins, + plugin_categories, ) - if not strategies: - return PluginAutoUpgradeBackfillResult(created_count=0, normalized=False) - - # Schema migration marks the historical workspace-level row as tool. - source_strategy = next( - (strategy for strategy in strategies if strategy.category == PluginCategory.TOOL), - strategies[0], + PluginAutoUpgradeService._log_unknown_plugin_ids( + tenant_id, + "include_plugins", + include_plugins, + plugin_categories, ) - source_has_default_strategy = PluginAutoUpgradeService._has_default_strategy(source_strategy) - strategies_by_category = {strategy.category: strategy for strategy in strategies} - exclude_plugins = source_strategy.exclude_plugins - include_plugins = source_strategy.include_plugins - should_split_plugin_lists = bool(exclude_plugins or include_plugins) - # Query daemon only for tenants that actually customized plugin lists. - plugin_categories = ( - PluginAutoUpgradeService._get_installed_plugin_categories(tenant_id) - if should_split_plugin_lists - else {} + + created_count = 0 + for category in PLUGIN_CATEGORIES: + strategy = strategies_by_category.get(category) + if strategy is None: + # Start from the legacy workspace-level behavior before narrowing lists. + strategy = TenantPluginAutoUpgradeStrategy( + tenant_id=tenant_id, + category=category, + strategy_setting=PluginAutoUpgradeService._strategy_setting_for_category( + source_strategy, category, source_has_default_strategy + ), + upgrade_time_of_day=PluginAutoUpgradeService._upgrade_time_of_day_for_category( + tenant_id, source_strategy, source_has_default_strategy + ), + upgrade_mode=source_strategy.upgrade_mode, + exclude_plugins=source_strategy.exclude_plugins.copy(), + include_plugins=source_strategy.include_plugins.copy(), + ) + session.add(strategy) + created_count += 1 + elif source_has_default_strategy: + strategy.strategy_setting = PluginAutoUpgradeService.default_strategy_setting_for_category( + strategy.category + ) + strategy.upgrade_time_of_day = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id) + + if not should_split_plugin_lists: + continue + + # Narrow include/exclude lists to the current category after all rows exist. + strategy.exclude_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category( + exclude_plugins, + strategy.category, + plugin_categories, + ) + strategy.include_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category( + include_plugins, + strategy.category, + plugin_categories, ) - if should_split_plugin_lists: - PluginAutoUpgradeService._log_unknown_plugin_ids( - tenant_id, - "exclude_plugins", - exclude_plugins, - plugin_categories, - ) - PluginAutoUpgradeService._log_unknown_plugin_ids( - tenant_id, - "include_plugins", - include_plugins, - plugin_categories, - ) - created_count = 0 - for category in PLUGIN_CATEGORIES: - strategy = strategies_by_category.get(category) - if strategy is None: - # Start from the legacy workspace-level behavior before narrowing lists. - strategy = TenantPluginAutoUpgradeStrategy( - tenant_id=tenant_id, - category=category, - strategy_setting=PluginAutoUpgradeService._strategy_setting_for_category( - source_strategy, category, source_has_default_strategy - ), - upgrade_time_of_day=PluginAutoUpgradeService._upgrade_time_of_day_for_category( - tenant_id, source_strategy, source_has_default_strategy - ), - upgrade_mode=source_strategy.upgrade_mode, - exclude_plugins=source_strategy.exclude_plugins.copy(), - include_plugins=source_strategy.include_plugins.copy(), - ) - session.add(strategy) - created_count += 1 - elif source_has_default_strategy: - strategy.strategy_setting = PluginAutoUpgradeService.default_strategy_setting_for_category( - strategy.category - ) - strategy.upgrade_time_of_day = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id) - - if not should_split_plugin_lists: - continue - - # Narrow include/exclude lists to the current category after all rows exist. - strategy.exclude_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category( - exclude_plugins, - strategy.category, - plugin_categories, - ) - strategy.include_plugins = PluginAutoUpgradeService._filter_plugin_ids_for_category( - include_plugins, - strategy.category, - plugin_categories, - ) - - return PluginAutoUpgradeBackfillResult(created_count=created_count, normalized=should_split_plugin_lists) + session.commit() + return PluginAutoUpgradeBackfillResult(created_count=created_count, normalized=should_split_plugin_lists) @staticmethod def _get_strategy( @@ -246,29 +248,27 @@ class PluginAutoUpgradeService: def get_strategy( tenant_id: str, category: PluginCategory, + *, + session: Session, ) -> TenantPluginAutoUpgradeStrategy | None: - with session_factory.create_session() as session: - return PluginAutoUpgradeService._get_strategy(session, tenant_id, category) + return PluginAutoUpgradeService._get_strategy(session, tenant_id, category) @staticmethod - def get_strategies(tenant_id: str) -> list[TenantPluginAutoUpgradeStrategy]: - with session_factory.create_session() as session: - return list( - session.scalars( - select(TenantPluginAutoUpgradeStrategy).where( - TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id - ) - ).all() - ) + def get_strategies(tenant_id: str, *, session: Session) -> list[TenantPluginAutoUpgradeStrategy]: + return list( + session.scalars( + select(TenantPluginAutoUpgradeStrategy).where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id) + ).all() + ) @staticmethod def _change_strategy( session: Session, tenant_id: str, category: PluginCategory, - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting, + strategy_setting: TenantPluginAutoUpgradeStrategySetting, upgrade_time_of_day: int, - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode, + upgrade_mode: TenantPluginAutoUpgradeMode, exclude_plugins: list[str], include_plugins: list[str], ) -> None: @@ -294,26 +294,28 @@ class PluginAutoUpgradeService: @staticmethod def change_strategy( tenant_id: str, - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting, + strategy_setting: TenantPluginAutoUpgradeStrategySetting, upgrade_time_of_day: int, - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode, + upgrade_mode: TenantPluginAutoUpgradeMode, exclude_plugins: list[str], include_plugins: list[str], category: PluginCategory, + *, + session: Session, ) -> bool: - with session_factory.create_session() as session, session.begin(): - PluginAutoUpgradeService._change_strategy( - session, - tenant_id=tenant_id, - category=category, - strategy_setting=strategy_setting, - upgrade_time_of_day=upgrade_time_of_day, - upgrade_mode=upgrade_mode, - exclude_plugins=exclude_plugins, - include_plugins=include_plugins, - ) + PluginAutoUpgradeService._change_strategy( + session, + tenant_id=tenant_id, + category=category, + strategy_setting=strategy_setting, + upgrade_time_of_day=upgrade_time_of_day, + upgrade_mode=upgrade_mode, + exclude_plugins=exclude_plugins, + include_plugins=include_plugins, + ) - return True + session.commit() + return True @staticmethod def _exclude_plugin( @@ -329,28 +331,28 @@ class PluginAutoUpgradeService: session, tenant_id, category, - TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, 0, - TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + TenantPluginAutoUpgradeMode.EXCLUDE, [plugin_id], [], ) else: - if exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE: + if exist_strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE: # In exclude mode, disabling one plugin means adding it to exclude_plugins. if plugin_id not in exist_strategy.exclude_plugins: new_exclude_plugins = exist_strategy.exclude_plugins.copy() new_exclude_plugins.append(plugin_id) exist_strategy.exclude_plugins = new_exclude_plugins - elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL: + elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL: # In partial mode, disabling one plugin means removing it from include_plugins. if plugin_id in exist_strategy.include_plugins: new_include_plugins = exist_strategy.include_plugins.copy() new_include_plugins.remove(plugin_id) exist_strategy.include_plugins = new_include_plugins - elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL: + elif exist_strategy.upgrade_mode == TenantPluginAutoUpgradeMode.ALL: # In all mode, switch to exclude mode so only this plugin is skipped. - exist_strategy.upgrade_mode = TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE + exist_strategy.upgrade_mode = TenantPluginAutoUpgradeMode.EXCLUDE exist_strategy.exclude_plugins = [plugin_id] @staticmethod @@ -358,13 +360,15 @@ class PluginAutoUpgradeService: tenant_id: str, plugin_id: str, category: PluginCategory, + *, + session: Session, ) -> bool: - with session_factory.create_session() as session, session.begin(): - PluginAutoUpgradeService._exclude_plugin( - session, - tenant_id, - category, - plugin_id, - ) + PluginAutoUpgradeService._exclude_plugin( + session, + tenant_id, + category, + plugin_id, + ) - return True + session.commit() + return True diff --git a/api/services/plugin/plugin_migration.py b/api/services/plugin/plugin_migration.py index 82eeb5a7261..5c2ddb77e0f 100644 --- a/api/services/plugin/plugin_migration.py +++ b/api/services/plugin/plugin_migration.py @@ -307,9 +307,9 @@ class PluginMigration: return result @classmethod - def _fetch_plugin_unique_identifier(cls, plugin_id: str) -> str | None: + def _fetch_latest_package_identifier(cls, plugin_id: str) -> str | None: """ - Fetch plugin unique identifier using plugin id. + Fetch the latest marketplace package identifier using a plugin id. """ if not dify_config.MARKETPLACE_ENABLED: return None @@ -328,7 +328,7 @@ class PluginMigration: @classmethod def extract_unique_plugins(cls, extracted_plugins: str) -> ExtractedPluginsDict: - plugins: dict[str, str] = {} + package_identifier_by_plugin_id: dict[str, str] = {} plugin_ids = [] plugin_not_exist = [] logger.info("Extracting unique plugins from %s", extracted_plugins) @@ -341,19 +341,19 @@ class PluginMigration: def fetch_plugin(plugin_id): try: - unique_identifier = cls._fetch_plugin_unique_identifier(plugin_id) - if unique_identifier: - plugins[plugin_id] = unique_identifier + latest_package_identifier = cls._fetch_latest_package_identifier(plugin_id) + if latest_package_identifier: + package_identifier_by_plugin_id[plugin_id] = latest_package_identifier else: plugin_not_exist.append(plugin_id) except Exception: - logger.exception("Failed to fetch plugin unique identifier for %s", plugin_id) + logger.exception("Failed to fetch latest package identifier for %s", plugin_id) plugin_not_exist.append(plugin_id) with ThreadPoolExecutor(max_workers=10) as executor: list(tqdm.tqdm(executor.map(fetch_plugin, plugin_ids), total=len(plugin_ids))) - return {"plugins": plugins, "plugin_not_exist": plugin_not_exist} + return {"plugins": package_identifier_by_plugin_id, "plugin_not_exist": plugin_not_exist} @classmethod def install_plugins(cls, extracted_plugins: str, output_file: str, workers: int = 100): @@ -362,17 +362,22 @@ class PluginMigration: """ manager = PluginInstaller() - plugins = cls.extract_unique_plugins(extracted_plugins) + extracted = cls.extract_unique_plugins(extracted_plugins) + package_identifier_by_plugin_id = extracted["plugins"] not_installed = [] plugin_install_failed = [] # use a fake tenant id to install all the plugins fake_tenant_id = uuid4().hex - logger.info("Installing %s plugin instances for fake tenant %s", len(plugins["plugins"]), fake_tenant_id) + logger.info( + "Installing %s plugin instances for fake tenant %s", + len(package_identifier_by_plugin_id), + fake_tenant_id, + ) thread_pool = ThreadPoolExecutor(max_workers=workers) - response = cls.handle_plugin_instance_install(fake_tenant_id, plugins["plugins"]) + response = cls.handle_plugin_instance_install(fake_tenant_id, package_identifier_by_plugin_id) if response.get("failed"): plugin_install_failed.extend(response.get("failed", [])) @@ -384,21 +389,21 @@ class PluginMigration: # at most 64 plugins one batch for i in range(0, len(plugin_ids), 64): batch_plugin_ids = plugin_ids[i : i + 64] - batch_plugin_identifiers = [ - plugins["plugins"][plugin_id] + batch_package_identifiers = [ + package_identifier_by_plugin_id[plugin_id] for plugin_id in batch_plugin_ids - if plugin_id not in installed_plugins_ids and plugin_id in plugins["plugins"] + if plugin_id not in installed_plugins_ids and plugin_id in package_identifier_by_plugin_id ] - if batch_plugin_identifiers: + if batch_package_identifiers: manager.install_from_identifiers( tenant_id, - batch_plugin_identifiers, + batch_package_identifiers, PluginInstallationSource.Marketplace, metas=[ { - "plugin_unique_identifier": identifier, + "plugin_unique_identifier": package_identifier, } - for identifier in batch_plugin_identifiers + for package_identifier in batch_package_identifiers ], ) PluginService.invalidate_plugin_model_providers_cache(tenant_id) @@ -411,12 +416,9 @@ class PluginMigration: data = _tenant_plugin_adapter.validate_json(line) tenant_id = data["tenant_id"] plugin_ids = data["plugins"] - plugin_not_exist: list[str] = [] - # get plugin unique identifier - for plugin_id in plugin_ids: - unique_identifier = plugins.get(plugin_id) - if unique_identifier: - plugin_not_exist.append(plugin_id) + plugin_not_exist = [ + plugin_id for plugin_id in plugin_ids if plugin_id not in package_identifier_by_plugin_id + ] if plugin_not_exist: not_installed.append( @@ -459,36 +461,44 @@ class PluginMigration: """ manager = PluginInstaller() - plugins = cls.extract_unique_plugins(extracted_plugins) + extracted = cls.extract_unique_plugins(extracted_plugins) + package_identifier_by_plugin_id = extracted["plugins"] plugin_install_failed = [] # use a fake tenant id to install all the plugins fake_tenant_id = uuid4().hex - logger.info("Installing %s plugin instances for fake tenant %s", len(plugins["plugins"]), fake_tenant_id) + logger.info( + "Installing %s plugin instances for fake tenant %s", + len(package_identifier_by_plugin_id), + fake_tenant_id, + ) thread_pool = ThreadPoolExecutor(max_workers=workers) - response = cls.handle_plugin_instance_install(fake_tenant_id, plugins["plugins"]) + response = cls.handle_plugin_instance_install(fake_tenant_id, package_identifier_by_plugin_id) if response.get("failed"): plugin_install_failed.extend(response.get("failed", [])) def install( - tenant_id: str, plugin_ids: dict[str, str], total_success_tenant: int, total_failed_tenant: int + tenant_id: str, + package_identifier_by_plugin_id: dict[str, str], + total_success_tenant: int, + total_failed_tenant: int, ) -> None: - logger.info("Installing %s plugins for tenant %s", len(plugin_ids), tenant_id) + logger.info("Installing %s plugins for tenant %s", len(package_identifier_by_plugin_id), tenant_id) try: # fetch plugin already installed installed_plugins = manager.list_plugins(tenant_id) installed_plugins_ids = [plugin.plugin_id for plugin in installed_plugins] # at most 64 plugins one batch - for i in range(0, len(plugin_ids), 64): - batch_plugin_ids = list(plugin_ids.keys())[i : i + 64] - batch_plugin_identifiers = [ - plugin_ids[plugin_id] + for i in range(0, len(package_identifier_by_plugin_id), 64): + batch_plugin_ids = list(package_identifier_by_plugin_id.keys())[i : i + 64] + batch_package_identifiers = [ + package_identifier_by_plugin_id[plugin_id] for plugin_id in batch_plugin_ids - if plugin_id not in installed_plugins_ids and plugin_id in plugin_ids + if plugin_id not in installed_plugins_ids and plugin_id in package_identifier_by_plugin_id ] - PluginService.install_from_marketplace_pkg(tenant_id, batch_plugin_identifiers) + PluginService.install_from_marketplace_pkg(tenant_id, batch_package_identifiers) total_success_tenant += 1 except Exception: @@ -510,7 +520,7 @@ class PluginMigration: thread_pool.submit( install, tenant_id, - plugins.get("plugins", {}), + package_identifier_by_plugin_id, total_success_tenant, total_failed_tenant, ) @@ -542,12 +552,12 @@ class PluginMigration: @classmethod def handle_plugin_instance_install( - cls, tenant_id: str, plugin_identifiers_map: Mapping[str, str] + cls, tenant_id: str, package_identifier_by_plugin_id: Mapping[str, str] ) -> PluginInstallResultDict: """ Install plugins for a tenant. """ - if plugin_identifiers_map and not dify_config.MARKETPLACE_ENABLED: + if package_identifier_by_plugin_id and not dify_config.MARKETPLACE_ENABLED: raise ValueError( "Marketplace disabled in offline mode; cannot bulk-install plugins. " "Pre-upload plugin packages via Console first." @@ -558,17 +568,17 @@ class PluginMigration: thread_pool = ThreadPoolExecutor(max_workers=10) futures = [] - for plugin_id, plugin_identifier in plugin_identifiers_map.items(): + for plugin_id, package_identifier in package_identifier_by_plugin_id.items(): - def download_and_upload(tenant_id, plugin_id, plugin_identifier): - plugin_package = marketplace.download_plugin_pkg(plugin_identifier) + def download_and_upload(tenant_id, plugin_id, package_identifier): + plugin_package = marketplace.download_plugin_pkg(package_identifier) if not plugin_package: - raise Exception(f"Failed to download plugin {plugin_identifier}") + raise Exception(f"Failed to download plugin {package_identifier}") # upload manager.upload_pkg(tenant_id, plugin_package, verify_signature=True) - futures.append(thread_pool.submit(download_and_upload, tenant_id, plugin_id, plugin_identifier)) + futures.append(thread_pool.submit(download_and_upload, tenant_id, plugin_id, package_identifier)) # Wait for all downloads to complete for future in futures: @@ -578,33 +588,33 @@ class PluginMigration: success = [] failed = [] - reverse_map = {v: k for k, v in plugin_identifiers_map.items()} + plugin_id_by_package_identifier = {v: k for k, v in package_identifier_by_plugin_id.items()} # at most 8 plugins one batch - for i in range(0, len(plugin_identifiers_map), 8): - batch_plugin_ids = list(plugin_identifiers_map.keys())[i : i + 8] - batch_plugin_identifiers = [plugin_identifiers_map[plugin_id] for plugin_id in batch_plugin_ids] + for i in range(0, len(package_identifier_by_plugin_id), 8): + batch_plugin_ids = list(package_identifier_by_plugin_id.keys())[i : i + 8] + batch_package_identifiers = [package_identifier_by_plugin_id[plugin_id] for plugin_id in batch_plugin_ids] try: response = manager.install_from_identifiers( tenant_id=tenant_id, - identifiers=batch_plugin_identifiers, + identifiers=batch_package_identifiers, source=PluginInstallationSource.Marketplace, metas=[ { - "plugin_unique_identifier": identifier, + "plugin_unique_identifier": package_identifier, } - for identifier in batch_plugin_identifiers + for package_identifier in batch_package_identifiers ], ) PluginService.invalidate_plugin_model_providers_cache(tenant_id) except Exception: # add to failed - failed.extend(batch_plugin_identifiers) + failed.extend(batch_plugin_ids) continue if response.all_installed: - success.extend(batch_plugin_identifiers) + success.extend(batch_plugin_ids) continue task_id = response.task_id @@ -614,10 +624,13 @@ class PluginMigration: if status.status in [PluginInstallTaskStatus.Failed, PluginInstallTaskStatus.Success]: PluginService.invalidate_plugin_model_providers_cache(tenant_id) for plugin in status.plugins: + plugin_id = plugin_id_by_package_identifier.get( + plugin.plugin_unique_identifier, plugin.plugin_unique_identifier.split(":", 1)[0] + ) if plugin.status == PluginInstallTaskStatus.Success: - success.append(reverse_map[plugin.plugin_unique_identifier]) + success.append(plugin_id) else: - failed.append(reverse_map[plugin.plugin_unique_identifier]) + failed.append(plugin_id) logger.error( "Failed to install plugin %s, error: %s", plugin.plugin_unique_identifier, diff --git a/api/services/plugin/plugin_permission_service.py b/api/services/plugin/plugin_permission_service.py index 3cca4268d00..339a6ccb89b 100644 --- a/api/services/plugin/plugin_permission_service.py +++ b/api/services/plugin/plugin_permission_service.py @@ -1,35 +1,36 @@ from sqlalchemy import select +from sqlalchemy.orm import Session -from core.db.session_factory import session_factory -from models.account import TenantPluginPermission +from models.account import TenantPluginDebugPermission, TenantPluginInstallPermission, TenantPluginPermission class PluginPermissionService: @staticmethod - def get_permission(tenant_id: str) -> TenantPluginPermission | None: - with session_factory.create_session() as session: - return session.scalar( - select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1) - ) + def get_permission(tenant_id: str, *, session: Session) -> TenantPluginPermission | None: + return session.scalar( + select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1) + ) @staticmethod def change_permission( tenant_id: str, - install_permission: TenantPluginPermission.InstallPermission, - debug_permission: TenantPluginPermission.DebugPermission, - ): - with session_factory.create_session() as session, session.begin(): - permission = session.scalar( - select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1) + install_permission: TenantPluginInstallPermission, + debug_permission: TenantPluginDebugPermission, + *, + session: Session, + ) -> bool: + permission = session.scalar( + select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1) + ) + if not permission: + permission = TenantPluginPermission( + tenant_id=tenant_id, install_permission=install_permission, debug_permission=debug_permission ) - if not permission: - permission = TenantPluginPermission( - tenant_id=tenant_id, install_permission=install_permission, debug_permission=debug_permission - ) - session.add(permission) - else: - permission.install_permission = install_permission - permission.debug_permission = debug_permission + session.add(permission) + else: + permission.install_permission = install_permission + permission.debug_permission = debug_permission - return True + session.commit() + return True diff --git a/api/services/rag_pipeline/pipeline_generate_service.py b/api/services/rag_pipeline/pipeline_generate_service.py index e77ff9687ed..276bfaea158 100644 --- a/api/services/rag_pipeline/pipeline_generate_service.py +++ b/api/services/rag_pipeline/pipeline_generate_service.py @@ -17,12 +17,13 @@ class PipelineGenerateService: @classmethod def generate( cls, - session: Session, pipeline: Pipeline, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: bool = True, + *, + session: Session, ): """ Pipeline Content Generate @@ -34,10 +35,10 @@ class PipelineGenerateService: :return: """ try: - workflow = cls._get_workflow(pipeline, invoke_from) + workflow = cls._get_workflow(pipeline, invoke_from, session) if original_document_id := args.get("original_document_id"): # update document status to waiting - cls.update_document_status(original_document_id, session) + cls.update_document_status(original_document_id, session=session) return PipelineGenerator.convert_to_event_stream( PipelineGenerator().generate( pipeline=pipeline, @@ -64,9 +65,9 @@ class PipelineGenerateService: @classmethod def generate_single_iteration( - cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, streaming: bool = True + cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, session: Session, streaming: bool = True ): - workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER) + workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER, session) return PipelineGenerator.convert_to_event_stream( PipelineGenerator().single_iteration_generate( pipeline=pipeline, workflow=workflow, node_id=node_id, user=user, args=args, streaming=streaming @@ -74,8 +75,10 @@ class PipelineGenerateService: ) @classmethod - def generate_single_loop(cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, streaming: bool = True): - workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER) + def generate_single_loop( + cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, session: Session, streaming: bool = True + ): + workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER, session) return PipelineGenerator.convert_to_event_stream( PipelineGenerator().single_loop_generate( pipeline=pipeline, workflow=workflow, node_id=node_id, user=user, args=args, streaming=streaming @@ -83,14 +86,14 @@ class PipelineGenerateService: ) @classmethod - def _get_workflow(cls, pipeline: Pipeline, invoke_from: InvokeFrom) -> Workflow: + def _get_workflow(cls, pipeline: Pipeline, invoke_from: InvokeFrom, session: Session) -> Workflow: """ Get workflow :param pipeline: pipeline :param invoke_from: invoke from :return: """ - rag_pipeline_service = RagPipelineService() + rag_pipeline_service = RagPipelineService(session) if invoke_from == InvokeFrom.DEBUGGER: # fetch draft workflow by app_model workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) @@ -107,7 +110,7 @@ class PipelineGenerateService: return workflow @classmethod - def update_document_status(cls, document_id: str, session: Session): + def update_document_status(cls, document_id: str, *, session: Session): """ Update document status to waiting :param document_id: document id diff --git a/api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py b/api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py index 6de0be33a4a..d56c239ace2 100644 --- a/api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/built_in/built_in_retrieval.py @@ -23,14 +23,15 @@ class BuiltInPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): @override def get_pipeline_templates( - self, session: Session, language: str, current_tenant_id: str | None = None + self, language: str, current_tenant_id: str | None = None, *, session: Session ) -> dict[str, Any]: - del current_tenant_id + del current_tenant_id, session result = self.fetch_pipeline_templates_from_builtin(language) return result @override - def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: + def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: + del session result = self.fetch_pipeline_template_detail_from_builtin(template_id) return result diff --git a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py index 4faaf342f66..3d6baefcc46 100644 --- a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py @@ -41,16 +41,16 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): @override def get_pipeline_templates( - self, session: Session, language: str, current_tenant_id: str | None = None + self, language: str, current_tenant_id: str | None = None, *, session: Session ) -> dict[str, Any]: current_tenant_id = resolve_tenant_id_fallback(current_tenant_id) return self.fetch_pipeline_templates_from_customized( - session=session, tenant_id=current_tenant_id, language=language + tenant_id=current_tenant_id, language=language, session=session ) @override - def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: - return self.fetch_pipeline_template_detail_from_db(session, template_id) + def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: + return self.fetch_pipeline_template_detail_from_db(template_id, session=session) @override def get_type(self) -> str: @@ -58,7 +58,7 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): @classmethod def fetch_pipeline_templates_from_customized( - cls, session: Session, tenant_id: str, language: str + cls, tenant_id: str, language: str, *, session: Session ) -> dict[str, Any]: """ Fetch pipeline templates from db. @@ -89,7 +89,7 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): return {"pipeline_templates": recommended_pipelines_results} @classmethod - def fetch_pipeline_template_detail_from_db(cls, session: Session, template_id: str) -> dict[str, Any] | None: + def fetch_pipeline_template_detail_from_db(cls, template_id: str, *, session: Session) -> dict[str, Any] | None: """ Fetch pipeline template detail from db. :param template_id: Template ID diff --git a/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py b/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py index f6d2731e21a..d5c31ff74b2 100644 --- a/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py @@ -41,21 +41,21 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase): @override def get_pipeline_templates( - self, session: Session, language: str, current_tenant_id: str | None = None + self, language: str, current_tenant_id: str | None = None, *, session: Session ) -> dict[str, Any]: del current_tenant_id - return self.fetch_pipeline_templates_from_db(session, language) + return self.fetch_pipeline_templates_from_db(language, session=session) @override - def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: - return self.fetch_pipeline_template_detail_from_db(session, template_id) + def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: + return self.fetch_pipeline_template_detail_from_db(template_id, session=session) @override def get_type(self) -> str: return PipelineTemplateType.DATABASE @classmethod - def fetch_pipeline_templates_from_db(cls, session: Session, language: str) -> dict[str, Any]: + def fetch_pipeline_templates_from_db(cls, language: str, *, session: Session) -> dict[str, Any]: """ Fetch pipeline templates from db. :param language: language @@ -83,7 +83,7 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase): return {"pipeline_templates": recommended_pipelines_results} @classmethod - def fetch_pipeline_template_detail_from_db(cls, session: Session, template_id: str) -> dict[str, Any] | None: + def fetch_pipeline_template_detail_from_db(cls, template_id: str, *, session: Session) -> dict[str, Any] | None: """ Fetch pipeline template detail from db. :param pipeline_id: Pipeline ID diff --git a/api/services/rag_pipeline/pipeline_template/pipeline_template_base.py b/api/services/rag_pipeline/pipeline_template/pipeline_template_base.py index ff53dc1f79e..c61ac6d60f2 100644 --- a/api/services/rag_pipeline/pipeline_template/pipeline_template_base.py +++ b/api/services/rag_pipeline/pipeline_template/pipeline_template_base.py @@ -7,9 +7,9 @@ class PipelineTemplateRetrievalBase(Protocol): """Interface for pipeline template retrieval.""" def get_pipeline_templates( - self, session: Session, language: str, current_tenant_id: str | None = None + self, language: str, current_tenant_id: str | None = None, *, session: Session ) -> dict[str, Any]: ... - def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: ... + def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: ... def get_type(self) -> str: ... diff --git a/api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py b/api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py index 7f9fe1b56ea..29acbd198b6 100644 --- a/api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/remote/remote_retrieval.py @@ -18,23 +18,25 @@ class RemotePipelineTemplateRetrieval(PipelineTemplateRetrievalBase): """ @override - def get_pipeline_template_detail(self, session: Session, template_id: str) -> dict[str, Any] | None: + def get_pipeline_template_detail(self, template_id: str, *, session: Session) -> dict[str, Any] | None: try: return self.fetch_pipeline_template_detail_from_dify_official(template_id) except Exception as e: logger.warning("fetch recommended app detail from dify official failed: %r, switch to database.", e) - return DatabasePipelineTemplateRetrieval.fetch_pipeline_template_detail_from_db(session, template_id) + return DatabasePipelineTemplateRetrieval.fetch_pipeline_template_detail_from_db( + template_id, session=session + ) @override def get_pipeline_templates( - self, session: Session, language: str, current_tenant_id: str | None = None + self, language: str, current_tenant_id: str | None = None, *, session: Session ) -> dict[str, Any]: del current_tenant_id try: return self.fetch_pipeline_templates_from_dify_official(language) except Exception as e: logger.warning("fetch pipeline templates from dify official failed: %r, switch to database.", e) - return DatabasePipelineTemplateRetrieval.fetch_pipeline_templates_from_db(session, language) + return DatabasePipelineTemplateRetrieval.fetch_pipeline_templates_from_db(language, session=session) @override def get_type(self) -> str: diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 9e17a05be16..8bd3918eb15 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -27,7 +27,6 @@ from core.datasource.entities.datasource_entities import ( from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin from core.datasource.online_drive.online_drive_plugin import OnlineDriveDatasourcePlugin from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin -from core.db.session_factory import session_factory from core.helper import marketplace from core.rag.entities import DatasourceCompletedEvent, DatasourceErrorEvent, DatasourceProcessingEvent from core.repositories.factory import DifyCoreRepositoryFactory, OrderConfig @@ -96,11 +95,13 @@ def _build_seeded_variable_pool(variables: Sequence[Variable]) -> VariablePool: class RagPipelineService: - def __init__(self, session_maker: sessionmaker | None = None): + _session: Session + + def __init__(self, session: Session, session_maker: sessionmaker | None = None): """Initialize RagPipelineService with repository dependencies.""" + self._session = session if session_maker is None: - session_maker = session_factory.get_session_maker() - self._session_maker = session_maker + session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) self._node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( session_maker ) @@ -109,15 +110,16 @@ class RagPipelineService: @classmethod def get_pipeline_templates( cls, - session: Session, type: str = "built-in", language: str = "en-US", current_tenant_id: str | None = None, + *, + session: Session, ) -> dict[str, Any]: if type == "built-in": mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)() - result = retrieval_instance.get_pipeline_templates(session, language, current_tenant_id) + result = retrieval_instance.get_pipeline_templates(language, current_tenant_id, session=session) if not result.get("pipeline_templates") and language != "en-US": template_retrieval = PipelineTemplateRetrievalFactory.get_built_in_pipeline_template_retrieval() result = template_retrieval.fetch_pipeline_templates_from_builtin("en-US") @@ -125,12 +127,12 @@ class RagPipelineService: else: mode = "customized" retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)() - result = retrieval_instance.get_pipeline_templates(session, language, current_tenant_id) + result = retrieval_instance.get_pipeline_templates(language, current_tenant_id, session=session) return result @classmethod def get_pipeline_template_detail( - cls, session: Session, template_id: str, type: str = "built-in" + cls, template_id: str, type: str = "built-in", *, session: Session ) -> dict[str, Any] | None: """ Get pipeline template detail. @@ -143,7 +145,7 @@ class RagPipelineService: mode = dify_config.HOSTED_FETCH_PIPELINE_TEMPLATES_MODE retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)() built_in_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail( - session, template_id + template_id, session=session ) if built_in_result is None: logger.warning( @@ -156,7 +158,7 @@ class RagPipelineService: mode = "customized" retrieval_instance = PipelineTemplateRetrievalFactory.get_pipeline_template_factory(mode)() customized_result: dict[str, Any] | None = retrieval_instance.get_pipeline_template_detail( - session, template_id + template_id, session=session ) return customized_result @@ -167,7 +169,8 @@ class RagPipelineService: template_info: PipelineTemplateInfoEntity, current_user: Account | None = None, current_tenant_id: str | None = None, - session: Session | None = None, + *, + session: Session, ): """ Update pipeline template. @@ -175,16 +178,6 @@ class RagPipelineService: :param template_info: template info """ current_user, current_tenant_id = resolve_account_fallback(current_user, current_tenant_id) - if session is None: - with session_factory.get_session_maker().begin() as new_session: - return cls.update_customized_pipeline_template( - template_id, - template_info, - current_user, - current_tenant_id, - session=new_session, - ) - customized_template: PipelineCustomizedTemplate | None = session.scalar( select(PipelineCustomizedTemplate) .where( @@ -213,21 +206,17 @@ class RagPipelineService: customized_template.description = template_info.description customized_template.icon = template_info.icon_info.model_dump() customized_template.updated_by = current_user.id + session.commit() return customized_template @classmethod def delete_customized_pipeline_template( - cls, template_id: str, current_tenant_id: str | None = None, session: Session | None = None + cls, template_id: str, current_tenant_id: str | None = None, *, session: Session ): """ Delete customized pipeline template. """ current_tenant_id = resolve_tenant_id_fallback(current_tenant_id) - if session is None: - with session_factory.get_session_maker().begin() as new_session: - cls.delete_customized_pipeline_template(template_id, current_tenant_id, session=new_session) - return - customized_template: PipelineCustomizedTemplate | None = session.scalar( select(PipelineCustomizedTemplate) .where( @@ -239,22 +228,22 @@ class RagPipelineService: if not customized_template: raise ValueError("Customized pipeline template not found.") session.delete(customized_template) + session.commit() def get_draft_workflow(self, pipeline: Pipeline) -> Workflow | None: """ Get draft workflow """ # fetch draft workflow by rag pipeline - with self._session_maker() as session: - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.version == "draft", - ) - .limit(1) + workflow = self._session.scalar( + select(Workflow) + .where( + Workflow.tenant_id == pipeline.tenant_id, + Workflow.app_id == pipeline.id, + Workflow.version == "draft", ) + .limit(1) + ) # return draft workflow return workflow @@ -268,31 +257,29 @@ class RagPipelineService: return None # fetch published workflow by workflow_id - with self._session_maker() as session: - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.id == pipeline.workflow_id, - ) - .limit(1) + workflow = self._session.scalar( + select(Workflow) + .where( + Workflow.tenant_id == pipeline.tenant_id, + Workflow.app_id == pipeline.id, + Workflow.id == pipeline.workflow_id, ) + .limit(1) + ) return workflow def get_published_workflow_by_id(self, pipeline: Pipeline, workflow_id: str) -> Workflow | None: """Fetch a published workflow snapshot by ID for restore operations.""" - with self._session_maker() as session: - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.id == workflow_id, - ) - .limit(1) + workflow = self._session.scalar( + select(Workflow) + .where( + Workflow.tenant_id == pipeline.tenant_id, + Workflow.app_id == pipeline.id, + Workflow.id == workflow_id, ) + .limit(1) + ) if workflow and workflow.version == Workflow.VERSION_DRAFT: raise IsDraftWorkflowError("source workflow must be published") return workflow @@ -350,51 +337,39 @@ class RagPipelineService: Sync draft workflow :raises WorkflowHashNotEqualError """ - with self._session_maker.begin() as session: - managed_pipeline = session.get(Pipeline, pipeline.id) - if not managed_pipeline: - raise ValueError("Pipeline not found") + # fetch draft workflow by app_model + workflow = self.get_draft_workflow(pipeline=pipeline) - # fetch draft workflow by app_model - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == managed_pipeline.tenant_id, - Workflow.app_id == managed_pipeline.id, - Workflow.version == "draft", - ) - .limit(1) + if workflow and workflow.unique_hash != unique_hash: + raise WorkflowHashNotEqualError() + + # create draft workflow if not found + if not workflow: + workflow = Workflow( + tenant_id=pipeline.tenant_id, + app_id=pipeline.id, + features="{}", + type=WorkflowType.RAG_PIPELINE.value, + version="draft", + graph=json.dumps(graph), + created_by=account.id, + environment_variables=environment_variables, + conversation_variables=conversation_variables, + rag_pipeline_variables=rag_pipeline_variables, ) - - if workflow and workflow.unique_hash != unique_hash: - raise WorkflowHashNotEqualError() - - # create draft workflow if not found - if not workflow: - workflow = Workflow( - tenant_id=managed_pipeline.tenant_id, - app_id=managed_pipeline.id, - features="{}", - type=WorkflowType.RAG_PIPELINE.value, - version="draft", - graph=json.dumps(graph), - created_by=account.id, - environment_variables=environment_variables, - conversation_variables=conversation_variables, - rag_pipeline_variables=rag_pipeline_variables, - ) - session.add(workflow) - session.flush() - managed_pipeline.workflow_id = workflow.id - pipeline.workflow_id = workflow.id - # update draft workflow if found - else: - workflow.graph = json.dumps(graph) - workflow.updated_by = account.id - workflow.updated_at = datetime.now(UTC).replace(tzinfo=None) - workflow.environment_variables = environment_variables - workflow.conversation_variables = conversation_variables - workflow.rag_pipeline_variables = rag_pipeline_variables + self._session.add(workflow) + self._session.flush() + pipeline.workflow_id = workflow.id + # update draft workflow if found + else: + workflow.graph = json.dumps(graph) + workflow.updated_by = account.id + workflow.updated_at = datetime.now(UTC).replace(tzinfo=None) + workflow.environment_variables = environment_variables + workflow.conversation_variables = conversation_variables + workflow.rag_pipeline_variables = rag_pipeline_variables + # commit db session changes + self._session.commit() # trigger workflow events TODO # app_draft_workflow_was_synced.send(pipeline, synced_draft_workflow=workflow) @@ -415,48 +390,26 @@ class RagPipelineService: the pipeline-specific flush/link step that wires a newly created draft back onto ``pipeline.workflow_id``. """ - with self._session_maker.begin() as session: - managed_pipeline = session.get(Pipeline, pipeline.id) - if not managed_pipeline: - raise ValueError("Pipeline not found") + source_workflow = self.get_published_workflow_by_id(pipeline=pipeline, workflow_id=workflow_id) + if not source_workflow: + raise WorkflowNotFoundError("Workflow not found.") - source_workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == managed_pipeline.tenant_id, - Workflow.app_id == managed_pipeline.id, - Workflow.id == workflow_id, - ) - .limit(1) - ) - if source_workflow and source_workflow.version == Workflow.VERSION_DRAFT: - raise IsDraftWorkflowError("source workflow must be published") - if not source_workflow: - raise WorkflowNotFoundError("Workflow not found.") + draft_workflow = self.get_draft_workflow(pipeline=pipeline) + draft_workflow, is_new_draft = apply_published_workflow_snapshot_to_draft( + tenant_id=pipeline.tenant_id, + app_id=pipeline.id, + source_workflow=source_workflow, + draft_workflow=draft_workflow, + account=account, + updated_at_factory=lambda: datetime.now(UTC).replace(tzinfo=None), + ) - draft_workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == managed_pipeline.tenant_id, - Workflow.app_id == managed_pipeline.id, - Workflow.version == Workflow.VERSION_DRAFT, - ) - .limit(1) - ) - draft_workflow, is_new_draft = apply_published_workflow_snapshot_to_draft( - tenant_id=managed_pipeline.tenant_id, - app_id=managed_pipeline.id, - source_workflow=source_workflow, - draft_workflow=draft_workflow, - account=account, - updated_at_factory=lambda: datetime.now(UTC).replace(tzinfo=None), - ) + if is_new_draft: + self._session.add(draft_workflow) + self._session.flush() + pipeline.workflow_id = draft_workflow.id - if is_new_draft: - session.add(draft_workflow) - session.flush() - managed_pipeline.workflow_id = draft_workflow.id - pipeline.workflow_id = draft_workflow.id + self._session.commit() return draft_workflow @@ -633,7 +586,7 @@ class RagPipelineService: workflow_node_execution.id ) - with self._session_maker.begin() as session: + with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, app_id=pipeline.id, @@ -1050,22 +1003,23 @@ class RagPipelineService: dataset_id = get_system_segment(variable_pool, SystemVariableKey.DATASET_ID) pipeline_id = get_system_segment(variable_pool, SystemVariableKey.APP_ID) if document_id and dataset_id and pipeline_id: - with self._session_maker.begin() as session: - document = session.scalar( - select(Document) - .join(Dataset, Dataset.id == Document.dataset_id) - .where( - Document.id == document_id.value, - Document.tenant_id == tenant_id, - Document.dataset_id == dataset_id.value, - Dataset.tenant_id == tenant_id, - Dataset.pipeline_id == pipeline_id.value, - ) - .limit(1) + document = self._session.scalar( + select(Document) + .join(Dataset, Dataset.id == Document.dataset_id) + .where( + Document.id == document_id.value, + Document.tenant_id == tenant_id, + Document.dataset_id == dataset_id.value, + Dataset.tenant_id == tenant_id, + Dataset.pipeline_id == pipeline_id.value, ) - if document: - document.indexing_status = IndexingStatus.ERROR - document.error = error + .limit(1) + ) + if document: + document.indexing_status = IndexingStatus.ERROR + document.error = error + self._session.add(document) + self._session.commit() return workflow_node_execution @@ -1276,89 +1230,86 @@ class RagPipelineService: args: dict[str, Any], current_user: Account | None = None, current_tenant_id: str | None = None, + *, + session: Session, ): """ Publish customized pipeline template """ current_user, _ = resolve_account_fallback(current_user, current_tenant_id) - with session_factory.get_session_maker().begin() as session: - pipeline = session.get(Pipeline, pipeline_id) - if not pipeline: - raise ValueError("Pipeline not found") - if not pipeline.workflow_id: - raise ValueError("Pipeline workflow not found") - workflow = session.get(Workflow, pipeline.workflow_id) - if not workflow: - raise ValueError("Workflow not found") - dataset = pipeline.retrieve_dataset(session=session) - if not dataset: - raise ValueError("Dataset not found") + pipeline = session.get(Pipeline, pipeline_id) + if not pipeline: + raise ValueError("Pipeline not found") + if not pipeline.workflow_id: + raise ValueError("Pipeline workflow not found") + workflow = session.get(Workflow, pipeline.workflow_id) + if not workflow: + raise ValueError("Workflow not found") + dataset = pipeline.retrieve_dataset(session=session) + if not dataset: + raise ValueError("Dataset not found") - # check template name is exist - template_name = args.get("name") - if template_name: - template = session.scalar( - select(PipelineCustomizedTemplate) - .where( - PipelineCustomizedTemplate.name == template_name, - PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id, - ) - .limit(1) - ) - if template: - raise ValueError("Template name is already exists") - - max_position = session.scalar( - select(func.max(PipelineCustomizedTemplate.position)).where( - PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id + # check template name is exist + template_name = args.get("name") + if template_name: + template = session.scalar( + select(PipelineCustomizedTemplate) + .where( + PipelineCustomizedTemplate.name == template_name, + PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id, ) + .limit(1) ) + if template: + raise ValueError("Template name is already exists") - from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService - - rag_pipeline_dsl_service = RagPipelineDslService(session) - dsl = rag_pipeline_dsl_service.export_rag_pipeline_dsl(pipeline=pipeline, include_secret=True) - if args.get("icon_info") is None: - args["icon_info"] = {} - if args.get("description") is None: - raise ValueError("Description is required") - if args.get("name") is None: - raise ValueError("Name is required") - pipeline_customized_template = PipelineCustomizedTemplate( - name=args.get("name") or "", - description=args.get("description") or "", - icon=args.get("icon_info") or {}, - tenant_id=pipeline.tenant_id, - yaml_content=dsl, - install_count=0, - position=max_position + 1 if max_position else 1, - chunk_structure=dataset.chunk_structure, - language="en-US", - created_by=current_user.id, + max_position = session.scalar( + select(func.max(PipelineCustomizedTemplate.position)).where( + PipelineCustomizedTemplate.tenant_id == pipeline.tenant_id ) - session.add(pipeline_customized_template) + ) + + from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService + + rag_pipeline_dsl_service = RagPipelineDslService(session) + dsl = rag_pipeline_dsl_service.export_rag_pipeline_dsl(pipeline=pipeline, include_secret=True) + if args.get("icon_info") is None: + args["icon_info"] = {} + if args.get("description") is None: + raise ValueError("Description is required") + if args.get("name") is None: + raise ValueError("Name is required") + pipeline_customized_template = PipelineCustomizedTemplate( + name=args.get("name") or "", + description=args.get("description") or "", + icon=args.get("icon_info") or {}, + tenant_id=pipeline.tenant_id, + yaml_content=dsl, + install_count=0, + position=max_position + 1 if max_position else 1, + chunk_structure=dataset.chunk_structure, + language="en-US", + created_by=current_user.id, + ) + session.add(pipeline_customized_template) + session.commit() def is_workflow_exist(self, pipeline: Pipeline) -> bool: - with self._session_maker() as session: - return ( - session.scalar( - select(func.count(Workflow.id)).where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.version == Workflow.VERSION_DRAFT, - ) + return ( + self._session.scalar( + select(func.count(Workflow.id)).where( + Workflow.tenant_id == pipeline.tenant_id, + Workflow.app_id == pipeline.id, + Workflow.version == Workflow.VERSION_DRAFT, ) - or 0 - ) > 0 + ) + or 0 + ) > 0 def get_node_last_run( self, pipeline: Pipeline, workflow: Workflow, node_id: str ) -> WorkflowNodeExecutionModel | None: - node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( - self._session_maker - ) - - node_exec = node_execution_service_repo.get_node_last_execution( + node_exec = self._node_execution_service_repo.get_node_last_execution( tenant_id=pipeline.tenant_id, app_id=pipeline.id, workflow_id=workflow.id, @@ -1431,7 +1382,7 @@ class RagPipelineService: # Convert node_execution to WorkflowNodeExecution after save workflow_node_execution_db_model = repository._to_db_model(workflow_node_execution) # type: ignore - with self._session_maker.begin() as session: + with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, app_id=pipeline.id, @@ -1465,10 +1416,9 @@ class RagPipelineService: if type and type != "all": stmt = stmt.where(PipelineRecommendedPlugin.type == type) - with self._session_maker() as session: - pipeline_recommended_plugins = session.scalars( - stmt.order_by(PipelineRecommendedPlugin.position.asc()) - ).all() + pipeline_recommended_plugins = self._session.scalars( + stmt.order_by(PipelineRecommendedPlugin.position.asc()) + ).all() if not pipeline_recommended_plugins: return { @@ -1507,173 +1457,41 @@ class RagPipelineService: """ Retry error document """ - with self._session_maker() as session: - document_pipeline_execution_log = session.scalar( - select(DocumentPipelineExecutionLog) - .where(DocumentPipelineExecutionLog.document_id == document.id) - .limit(1) - ) - if not document_pipeline_execution_log: - raise ValueError("Document pipeline execution log not found") - pipeline = session.get(Pipeline, document_pipeline_execution_log.pipeline_id) - if not pipeline: - raise ValueError("Pipeline not found") - # convert to app config - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.id == pipeline.workflow_id, - ) - .limit(1) - ) - if not workflow: - raise ValueError("Workflow not found") - PipelineGenerator().generate( - pipeline=pipeline, - workflow=workflow, - user=user, - args={ - "inputs": document_pipeline_execution_log.input_data, - "start_node_id": document_pipeline_execution_log.datasource_node_id, - "datasource_type": document_pipeline_execution_log.datasource_type, - "datasource_info_list": [json.loads(document_pipeline_execution_log.datasource_info)], - "original_document_id": document.id, - }, - invoke_from=InvokeFrom.PUBLISHED_PIPELINE, - streaming=False, - call_depth=0, - workflow_thread_pool_id=None, - is_retry=True, - ) + document_pipeline_execution_log = self._session.scalar( + select(DocumentPipelineExecutionLog).where(DocumentPipelineExecutionLog.document_id == document.id).limit(1) + ) + if not document_pipeline_execution_log: + raise ValueError("Document pipeline execution log not found") + pipeline = self._session.get(Pipeline, document_pipeline_execution_log.pipeline_id) + if not pipeline: + raise ValueError("Pipeline not found") + # convert to app config + workflow = self.get_published_workflow(pipeline) + if not workflow: + raise ValueError("Workflow not found") + PipelineGenerator().generate( + pipeline=pipeline, + workflow=workflow, + user=user, + args={ + "inputs": document_pipeline_execution_log.input_data, + "start_node_id": document_pipeline_execution_log.datasource_node_id, + "datasource_type": document_pipeline_execution_log.datasource_type, + "datasource_info_list": [json.loads(document_pipeline_execution_log.datasource_info)], + "original_document_id": document.id, + }, + invoke_from=InvokeFrom.PUBLISHED_PIPELINE, + streaming=False, + call_depth=0, + workflow_thread_pool_id=None, + is_retry=True, + ) def get_datasource_plugins(self, tenant_id: str, dataset_id: str, is_published: bool) -> list[dict]: """ Get datasource plugins """ - with self._session_maker() as session: - dataset: Dataset | None = session.scalar( - select(Dataset) - .where( - Dataset.id == dataset_id, - Dataset.tenant_id == tenant_id, - ) - .limit(1) - ) - if not dataset: - raise ValueError("Dataset not found") - pipeline: Pipeline | None = session.scalar( - select(Pipeline) - .where( - Pipeline.id == dataset.pipeline_id, - Pipeline.tenant_id == tenant_id, - ) - .limit(1) - ) - if not pipeline: - raise ValueError("Pipeline not found") - - if is_published: - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.id == pipeline.workflow_id, - ) - .limit(1) - ) - else: - workflow = session.scalar( - select(Workflow) - .where( - Workflow.tenant_id == pipeline.tenant_id, - Workflow.app_id == pipeline.id, - Workflow.version == Workflow.VERSION_DRAFT, - ) - .limit(1) - ) - if not pipeline or not workflow: - raise ValueError("Pipeline or workflow not found") - - datasource_nodes = workflow.graph_dict.get("nodes", []) - datasource_plugins = [] - for datasource_node in datasource_nodes: - if datasource_node.get("data", {}).get("type") == "datasource": - datasource_node_data = datasource_node["data"] - if not datasource_node_data: - continue - - variables = workflow.rag_pipeline_variables - if variables: - variables_map = {item["variable"]: item for item in variables} - else: - variables_map = {} - - datasource_parameters = datasource_node_data.get("datasource_parameters", {}) - user_input_variables_keys = [] - user_input_variables = [] - - for _, value in datasource_parameters.items(): - if value.get("value") and isinstance(value.get("value"), str): - pattern = ( - r"\{\{#([a-zA-Z0-9_]{1,50}" - r"(?:\.[a-zA-Z0-9_][a-zA-Z0-9_]{0,29}){1,10})#\}\}" - ) - match = re.match(pattern, value["value"]) - if match: - full_path = match.group(1) - last_part = full_path.split(".")[-1] - user_input_variables_keys.append(last_part) - elif value.get("value") and isinstance(value.get("value"), list): - last_part = value.get("value")[-1] - user_input_variables_keys.append(last_part) - for key, value in variables_map.items(): - if key in user_input_variables_keys: - user_input_variables.append(value) - - # get credentials - datasource_provider_service: DatasourceProviderService = DatasourceProviderService() - credentials: list[dict[Any, Any]] = datasource_provider_service.list_datasource_credentials( - tenant_id=tenant_id, - provider=datasource_node_data.get("provider_name"), - plugin_id=datasource_node_data.get("plugin_id"), - ) - credential_info_list: list[Any] = [] - for credential in credentials: - credential_info_list.append( - { - "id": credential.get("id"), - "name": credential.get("name"), - "type": credential.get("type"), - "is_default": credential.get("is_default"), - } - ) - - datasource_plugins.append( - { - "node_id": datasource_node.get("id"), - "plugin_id": datasource_node_data.get("plugin_id"), - "provider_name": datasource_node_data.get("provider_name"), - "datasource_type": datasource_node_data.get("provider_type"), - "title": datasource_node_data.get("title"), - "user_input_variables": user_input_variables, - "credentials": credential_info_list, - } - ) - - return datasource_plugins - - def get_pipeline(self, tenant_id: str, dataset_id: str, session: Session | None = None) -> Pipeline: - """ - Get pipeline - """ - if session is None: - with self._session_maker() as new_session: - return self.get_pipeline(tenant_id, dataset_id, session=new_session) - - dataset: Dataset | None = session.scalar( + dataset: Dataset | None = self._session.scalar( select(Dataset) .where( Dataset.id == dataset_id, @@ -1683,7 +1501,106 @@ class RagPipelineService: ) if not dataset: raise ValueError("Dataset not found") - pipeline: Pipeline | None = session.scalar( + pipeline: Pipeline | None = self._session.scalar( + select(Pipeline) + .where( + Pipeline.id == dataset.pipeline_id, + Pipeline.tenant_id == tenant_id, + ) + .limit(1) + ) + if not pipeline: + raise ValueError("Pipeline not found") + + workflow: Workflow | None = None + if is_published: + workflow = self.get_published_workflow(pipeline=pipeline) + else: + workflow = self.get_draft_workflow(pipeline=pipeline) + if not pipeline or not workflow: + raise ValueError("Pipeline or workflow not found") + + datasource_nodes = workflow.graph_dict.get("nodes", []) + datasource_plugins = [] + for datasource_node in datasource_nodes: + if datasource_node.get("data", {}).get("type") == "datasource": + datasource_node_data = datasource_node["data"] + if not datasource_node_data: + continue + + variables = workflow.rag_pipeline_variables + if variables: + variables_map = {item["variable"]: item for item in variables} + else: + variables_map = {} + + datasource_parameters = datasource_node_data.get("datasource_parameters", {}) + user_input_variables_keys = [] + user_input_variables = [] + + for _, value in datasource_parameters.items(): + if value.get("value") and isinstance(value.get("value"), str): + pattern = r"\{\{#([a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z0-9_][a-zA-Z0-9_]{0,29}){1,10})#\}\}" + match = re.match(pattern, value["value"]) + if match: + full_path = match.group(1) + last_part = full_path.split(".")[-1] + user_input_variables_keys.append(last_part) + elif value.get("value") and isinstance(value.get("value"), list): + last_part = value.get("value")[-1] + user_input_variables_keys.append(last_part) + for key, value in variables_map.items(): + if key in user_input_variables_keys: + user_input_variables.append(value) + + # get credentials + datasource_provider_service: DatasourceProviderService = DatasourceProviderService() + credentials: list[dict[Any, Any]] = datasource_provider_service.list_datasource_credentials( + tenant_id=tenant_id, + provider=datasource_node_data.get("provider_name"), + plugin_id=datasource_node_data.get("plugin_id"), + session=self._session, + ) + credential_info_list: list[Any] = [] + for credential in credentials: + credential_info_list.append( + { + "id": credential.get("id"), + "name": credential.get("name"), + "type": credential.get("type"), + "is_default": credential.get("is_default"), + } + ) + + datasource_plugins.append( + { + "node_id": datasource_node.get("id"), + "plugin_id": datasource_node_data.get("plugin_id"), + "provider_name": datasource_node_data.get("provider_name"), + "datasource_type": datasource_node_data.get("provider_type"), + "title": datasource_node_data.get("title"), + "user_input_variables": user_input_variables, + "credentials": credential_info_list, + } + ) + + return datasource_plugins + + def get_pipeline(self, tenant_id: str, dataset_id: str) -> Pipeline: + """ + Get pipeline + """ + dataset: Dataset | None = self._session.scalar( + select(Dataset) + .where( + Dataset.id == dataset_id, + Dataset.tenant_id == tenant_id, + ) + .limit(1) + ) + if not dataset: + raise ValueError("Dataset not found") + pipeline: Pipeline | None = self._session.scalar( select(Pipeline) .where( Pipeline.id == dataset.pipeline_id, diff --git a/api/services/rag_pipeline/rag_pipeline_dsl_service.py b/api/services/rag_pipeline/rag_pipeline_dsl_service.py index 4f9cde37a7b..d562a4b9adf 100644 --- a/api/services/rag_pipeline/rag_pipeline_dsl_service.py +++ b/api/services/rag_pipeline/rag_pipeline_dsl_service.py @@ -15,7 +15,7 @@ from Crypto.Util.Padding import pad, unpad from flask_login import current_user from pydantic import BaseModel from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from core.file import remote_fetcher from core.helper.name_generator import generate_incremental_name @@ -36,6 +36,7 @@ from models import Account from models.dataset import Dataset, DatasetCollectionBinding, Pipeline from models.enums import CollectionBindingType, DatasetRuntimeMode from models.workflow import Workflow, WorkflowType +from services.dsl_content import DSL_MAX_SIZE, dsl_content_size from services.dsl_version import check_version_compatibility from services.entities.dsl_entities import CheckDependenciesResult, ImportMode, ImportStatus from services.entities.knowledge_entities.rag_pipeline_entities import ( @@ -50,7 +51,6 @@ logger = logging.getLogger(__name__) IMPORT_INFO_REDIS_KEY_PREFIX = "app_import_info:" CHECK_DEPENDENCIES_REDIS_KEY_PREFIX = "app_check_dependencies:" IMPORT_INFO_REDIS_EXPIRY = 10 * 60 # 10 minutes -DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB CURRENT_DSL_VERSION = "0.1.0" @@ -83,7 +83,7 @@ class RagPipelineDslService: when generated IDs are needed mid-operation; they never commit or rollback. """ - def __init__(self, session: Session | scoped_session): + def __init__(self, session: Session): self._session = session def import_rag_pipeline( @@ -127,15 +127,16 @@ class RagPipelineDslService: yaml_url = yaml_url.replace("/blob/", "/") response = remote_fetcher.make_request("GET", yaml_url.strip(), follow_redirects=True, timeout=(10, 10)) response.raise_for_status() - content = response.content.decode() + raw_content = response.content - if len(content) > DSL_MAX_SIZE: + if dsl_content_size(raw_content) > DSL_MAX_SIZE: return RagPipelineImportInfo( id=import_id, status=ImportStatus.FAILED, error="File size exceeds the limit of 10MB", ) + content = raw_content.decode("utf-8") if not content: return RagPipelineImportInfo( id=import_id, @@ -156,6 +157,12 @@ class RagPipelineDslService: error="yaml_content is required when import_mode is yaml-content", ) content = yaml_content + if dsl_content_size(content) > DSL_MAX_SIZE: + return RagPipelineImportInfo( + id=import_id, + status=ImportStatus.FAILED, + error="File size exceeds the limit of 10MB", + ) # Process YAML content try: diff --git a/api/services/rag_pipeline/rag_pipeline_transform_service.py b/api/services/rag_pipeline/rag_pipeline_transform_service.py index daefaa9e30e..6a7902c1908 100644 --- a/api/services/rag_pipeline/rag_pipeline_transform_service.py +++ b/api/services/rag_pipeline/rag_pipeline_transform_service.py @@ -96,7 +96,7 @@ class RagPipelineTransformService: # deal document data self._deal_document_data(dataset, session) - session.flush() + session.commit() return { "pipeline_id": pipeline.id, "dataset_id": dataset_id, @@ -194,6 +194,7 @@ class RagPipelineTransformService: def _create_pipeline( self, data: dict[str, Any], + *, session: Session, ) -> Pipeline: """Create a new app or update an existing one.""" @@ -269,11 +270,13 @@ class RagPipelineTransformService: installed_plugins_ids = [plugin.plugin_id for plugin in installed_plugins] dependencies = pipeline_yaml.get("dependencies", []) - need_install_plugin_unique_identifiers = [] + package_identifiers_to_install = [] for dependency in dependencies: if dependency.get("type") == "marketplace": - plugin_unique_identifier = dependency.get("value", {}).get("plugin_unique_identifier") - plugin_id = plugin_unique_identifier.split(":")[0] + package_identifier = dependency.get("value", {}).get("plugin_unique_identifier") + if not package_identifier: + continue + plugin_id = package_identifier.split(":", 1)[0] if plugin_id not in installed_plugins_ids: if not dify_config.MARKETPLACE_ENABLED: logger.warning( @@ -282,14 +285,14 @@ class RagPipelineTransformService: plugin_id, ) continue - plugin_unique_identifier = plugin_migration._fetch_plugin_unique_identifier(plugin_id) # type: ignore - if plugin_unique_identifier: - need_install_plugin_unique_identifiers.append(plugin_unique_identifier) - if need_install_plugin_unique_identifiers: - logger.debug("Installing missing pipeline plugins %s", need_install_plugin_unique_identifiers) - PluginService.install_from_marketplace_pkg(tenant_id, need_install_plugin_unique_identifiers) + latest_package_identifier = plugin_migration._fetch_latest_package_identifier(plugin_id) # type: ignore + if latest_package_identifier: + package_identifiers_to_install.append(latest_package_identifier) + if package_identifiers_to_install: + logger.debug("Installing missing pipeline plugins %s", package_identifiers_to_install) + PluginService.install_from_marketplace_pkg(tenant_id, package_identifiers_to_install) - def _transform_to_empty_pipeline(self, dataset: Dataset, session: Session): + def _transform_to_empty_pipeline(self, dataset: Dataset, *, session: Session): pipeline = Pipeline( tenant_id=dataset.tenant_id, name=dataset.name, @@ -304,7 +307,7 @@ class RagPipelineTransformService: dataset.updated_by = current_user.id dataset.updated_at = datetime.now(UTC).replace(tzinfo=None) session.add(dataset) - session.flush() + session.commit() return { "pipeline_id": pipeline.id, "dataset_id": dataset.id, diff --git a/api/services/recommend_app/buildin/buildin_retrieval.py b/api/services/recommend_app/buildin/buildin_retrieval.py index 03b72a4f57c..d29d754b67e 100644 --- a/api/services/recommend_app/buildin/buildin_retrieval.py +++ b/api/services/recommend_app/buildin/buildin_retrieval.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any, override from flask import current_app +from sqlalchemy.orm import Session from services.recommend_app.database.database_retrieval import DatabaseRecommendAppRetrieval from services.recommend_app.recommend_app_base import RecommendAppRetrievalBase @@ -22,17 +23,19 @@ class BuildInRecommendAppRetrieval(RecommendAppRetrievalBase): return RecommendAppType.BUILDIN @override - def get_recommended_apps_and_categories(self, language: str): + def get_recommended_apps_and_categories(self, language: str, *, session: Session): + del session result = self.fetch_recommended_apps_from_builtin(language) return result @override - def get_learn_dify_apps(self, language: str): - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language) + def get_learn_dify_apps(self, language: str, *, session: Session): + result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language, session=session) return result @override - def get_recommend_app_detail(self, app_id: str): + def get_recommend_app_detail(self, app_id: str, *, session: Session): + del session result = self.fetch_recommended_app_detail_from_builtin(app_id) return result diff --git a/api/services/recommend_app/database/database_retrieval.py b/api/services/recommend_app/database/database_retrieval.py index f6786175896..08d902fdeb5 100644 --- a/api/services/recommend_app/database/database_retrieval.py +++ b/api/services/recommend_app/database/database_retrieval.py @@ -1,9 +1,9 @@ from typing import Any, NotRequired, TypedDict, override from sqlalchemy import select +from sqlalchemy.orm import Session from constants.languages import languages -from extensions.ext_database import db from models.model import App, RecommendedApp from services.app_dsl_service import AppDslService from services.recommend_app.category_order import order_categories @@ -45,18 +45,18 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): """ @override - def get_recommended_apps_and_categories(self, language: str) -> RecommendedAppsResultDict: - result = self.fetch_recommended_apps_from_db(language) + def get_recommended_apps_and_categories(self, language: str, *, session: Session) -> RecommendedAppsResultDict: + result = self.fetch_recommended_apps_from_db(language, session=session) return result @override - def get_learn_dify_apps(self, language: str) -> RecommendedAppsResultDict: - result = self.fetch_learn_dify_apps_from_db(language) + def get_learn_dify_apps(self, language: str, *, session: Session) -> RecommendedAppsResultDict: + result = self.fetch_learn_dify_apps_from_db(language, session=session) return result @override - def get_recommend_app_detail(self, app_id: str) -> RecommendedAppDetailDict | None: - result = self.fetch_recommended_app_detail_from_db(app_id) + def get_recommend_app_detail(self, app_id: str, *, session: Session) -> RecommendedAppDetailDict | None: + result = self.fetch_recommended_app_detail_from_db(app_id, session=session) return result @override @@ -64,42 +64,42 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): return RecommendAppType.DATABASE @classmethod - def fetch_recommended_apps_from_db(cls, language: str) -> RecommendedAppsResultDict: + def fetch_recommended_apps_from_db(cls, language: str, *, session: Session) -> RecommendedAppsResultDict: """ Fetch recommended apps from db. :param language: language :return: """ - recommended_apps = cls._fetch_listed_recommended_apps(language) + recommended_apps = cls._fetch_listed_recommended_apps(language, session=session) if len(recommended_apps) == 0: - recommended_apps = cls._fetch_listed_recommended_apps(languages[0]) + recommended_apps = cls._fetch_listed_recommended_apps(languages[0], session=session) return cls._format_recommended_apps(recommended_apps, language) @classmethod - def fetch_learn_dify_apps_from_db(cls, language: str) -> RecommendedAppsResultDict: + def fetch_learn_dify_apps_from_db(cls, language: str, *, session: Session) -> RecommendedAppsResultDict: """ Fetch listed recommended apps explicitly marked for the Learn Dify section. :param language: language :return: """ - recommended_apps = cls._fetch_listed_recommended_apps(language, is_learn_dify=True) + recommended_apps = cls._fetch_listed_recommended_apps(language, session=session, is_learn_dify=True) if len(recommended_apps) == 0 and language != languages[0]: - recommended_apps = cls._fetch_listed_recommended_apps(languages[0], is_learn_dify=True) + recommended_apps = cls._fetch_listed_recommended_apps(languages[0], session=session, is_learn_dify=True) return cls._format_recommended_apps(recommended_apps, language) @classmethod def _fetch_listed_recommended_apps( - cls, language: str, *, is_learn_dify: bool | None = None + cls, language: str, *, session: Session, is_learn_dify: bool | None = None ) -> list[RecommendedApp]: filters = [RecommendedApp.is_listed.is_(True), RecommendedApp.language == language] if is_learn_dify is not None: filters.append(RecommendedApp.is_learn_dify.is_(is_learn_dify)) - return list(db.session.scalars(select(RecommendedApp).where(*filters)).all()) + return list(session.scalars(select(RecommendedApp).where(*filters)).all()) @classmethod def _format_recommended_apps( @@ -146,14 +146,14 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): ) @classmethod - def fetch_recommended_app_detail_from_db(cls, app_id: str) -> RecommendedAppDetailDict | None: + def fetch_recommended_app_detail_from_db(cls, app_id: str, *, session: Session) -> RecommendedAppDetailDict | None: """ Fetch recommended app detail from db. :param app_id: App ID :return: """ # is in public recommended list - recommended_app = db.session.scalar( + recommended_app = session.scalar( select(RecommendedApp).where(RecommendedApp.is_listed == True, RecommendedApp.app_id == app_id).limit(1) ) @@ -161,7 +161,7 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): return None # get app detail - app_model = db.session.get(App, app_id) + app_model = session.get(App, app_id) if not app_model or not app_model.is_public: return None @@ -171,5 +171,5 @@ class DatabaseRecommendAppRetrieval(RecommendAppRetrievalBase): icon=app_model.icon, icon_background=app_model.icon_background, mode=app_model.mode, - export_data=AppDslService.export_dsl(app_model=app_model), + export_data=AppDslService.export_dsl(app_model=app_model, session=session), ) diff --git a/api/services/recommend_app/recommend_app_base.py b/api/services/recommend_app/recommend_app_base.py index f819cc3a937..821ad476c42 100644 --- a/api/services/recommend_app/recommend_app_base.py +++ b/api/services/recommend_app/recommend_app_base.py @@ -1,13 +1,15 @@ from typing import Any, Protocol +from sqlalchemy.orm import Session + class RecommendAppRetrievalBase(Protocol): """Interface for recommend app retrieval.""" - def get_recommended_apps_and_categories(self, language: str) -> Any: ... + def get_recommended_apps_and_categories(self, language: str, *, session: Session) -> Any: ... - def get_learn_dify_apps(self, language: str) -> Any: ... + def get_learn_dify_apps(self, language: str, *, session: Session) -> Any: ... - def get_recommend_app_detail(self, app_id: str) -> Any: ... + def get_recommend_app_detail(self, app_id: str, *, session: Session) -> Any: ... def get_type(self) -> str: ... diff --git a/api/services/recommend_app/remote/remote_retrieval.py b/api/services/recommend_app/remote/remote_retrieval.py index 2e3222bb978..c676ec907e0 100644 --- a/api/services/recommend_app/remote/remote_retrieval.py +++ b/api/services/recommend_app/remote/remote_retrieval.py @@ -3,6 +3,7 @@ from typing import Any, override import httpx from flask import has_request_context, request +from sqlalchemy.orm import Session from configs import dify_config from services.recommend_app.buildin.buildin_retrieval import BuildInRecommendAppRetrieval @@ -33,7 +34,8 @@ class RemoteRecommendAppRetrieval(RecommendAppRetrievalBase): """ @override - def get_recommend_app_detail(self, app_id: str): + def get_recommend_app_detail(self, app_id: str, *, session: Session): + del session try: result = self.fetch_recommended_app_detail_from_dify_official(app_id) except Exception as e: @@ -42,7 +44,8 @@ class RemoteRecommendAppRetrieval(RecommendAppRetrievalBase): return result @override - def get_recommended_apps_and_categories(self, language: str): + def get_recommended_apps_and_categories(self, language: str, *, session: Session): + del session try: result = self.fetch_recommended_apps_from_dify_official(language) except Exception as e: @@ -51,12 +54,12 @@ class RemoteRecommendAppRetrieval(RecommendAppRetrievalBase): return result @override - def get_learn_dify_apps(self, language: str): + def get_learn_dify_apps(self, language: str, *, session: Session): try: result = self.fetch_learn_dify_apps_from_dify_official(language) except Exception as e: logger.warning("fetch learn dify apps from dify official failed: %s, switch to database.", e) - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language) + result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db(language, session=session) return result @override diff --git a/api/services/recommended_app_service.py b/api/services/recommended_app_service.py index 2d247ba5b71..813aa74754c 100644 --- a/api/services/recommended_app_service.py +++ b/api/services/recommended_app_service.py @@ -1,7 +1,7 @@ from typing import Any from sqlalchemy import select -from sqlalchemy.orm import scoped_session +from sqlalchemy.orm import Session from configs import dify_config from models.model import AccountTrialAppRecord, TrialApp @@ -11,7 +11,7 @@ from services.recommend_app.recommend_app_factory import RecommendAppRetrievalFa class RecommendedAppService: @classmethod - def get_recommended_apps_and_categories(cls, session: scoped_session, language: str): + def get_recommended_apps_and_categories(cls, language: str, *, session: Session): """ Get recommended apps and categories. :param language: language @@ -19,7 +19,7 @@ class RecommendedAppService: """ mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result = retrieval_instance.get_recommended_apps_and_categories(language) + result = retrieval_instance.get_recommended_apps_and_categories(language, session=session) if not result.get("recommended_apps"): result = ( RecommendAppRetrievalFactory.get_buildin_recommend_app_retrieval().fetch_recommended_apps_from_builtin( @@ -35,7 +35,7 @@ class RecommendedAppService: return result @classmethod - def get_learn_dify_apps(cls, session: scoped_session, language: str) -> dict[str, Any]: + def get_learn_dify_apps(cls, language: str, *, session: Session) -> dict[str, Any]: """ Get recommended apps marked for the Learn Dify section. :param language: language @@ -43,7 +43,7 @@ class RecommendedAppService: """ mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result = retrieval_instance.get_learn_dify_apps(language) + result = retrieval_instance.get_learn_dify_apps(language, session=session) if FeatureService.get_system_features().enable_trial_app: for app in result["recommended_apps"]: @@ -52,7 +52,7 @@ class RecommendedAppService: return {"recommended_apps": result["recommended_apps"]} @classmethod - def get_recommend_app_detail(cls, session: scoped_session, app_id: str) -> dict[str, Any] | None: + def get_recommend_app_detail(cls, app_id: str, *, session: Session) -> dict[str, Any] | None: """ Get recommend app detail. :param app_id: app id @@ -60,7 +60,7 @@ class RecommendedAppService: """ mode = dify_config.HOSTED_FETCH_APP_TEMPLATES_MODE retrieval_instance = RecommendAppRetrievalFactory.get_recommend_app_factory(mode)() - result: dict[str, Any] | None = retrieval_instance.get_recommend_app_detail(app_id) + result: dict[str, Any] | None = retrieval_instance.get_recommend_app_detail(app_id, session=session) if result is None: return None if FeatureService.get_system_features().enable_trial_app: @@ -69,7 +69,7 @@ class RecommendedAppService: return result @classmethod - def add_trial_app_record(cls, session: scoped_session, app_id: str, account_id: str): + def add_trial_app_record(cls, app_id: str, account_id: str, *, session: Session): """ Add trial app record. :param app_id: app id @@ -88,6 +88,6 @@ class RecommendedAppService: session.commit() @staticmethod - def _can_trial_app(session: scoped_session, app_id: str) -> bool: + def _can_trial_app(session: Session, app_id: str) -> bool: trial_app_model = session.scalar(select(TrialApp).where(TrialApp.app_id == app_id).limit(1)) return trial_app_model is not None diff --git a/api/services/saved_message_service.py b/api/services/saved_message_service.py index 9a65429748e..6165d74333f 100644 --- a/api/services/saved_message_service.py +++ b/api/services/saved_message_service.py @@ -12,7 +12,7 @@ from services.message_service import MessageService class SavedMessageService: @classmethod def pagination_by_last_id( - cls, session: Session, app_model: App, user: Account | EndUser | None, last_id: str | None, limit: int + cls, app_model: App, user: Account | EndUser | None, last_id: str | None, limit: int, *, session: Session ) -> InfiniteScrollPagination: if not user: raise ValueError("User is required") @@ -28,11 +28,16 @@ class SavedMessageService: message_ids = [sm.message_id for sm in saved_messages] return MessageService.pagination_by_last_id( - app_model=app_model, user=user, last_id=last_id, limit=limit, include_ids=message_ids + app_model=app_model, + user=user, + last_id=last_id, + limit=limit, + include_ids=message_ids, + session=session, ) @classmethod - def save(cls, session: Session, app_model: App, user: Account | EndUser | None, message_id: str): + def save(cls, app_model: App, user: Account | EndUser | None, message_id: str, *, session: Session): if not user: return saved_message = session.scalar( @@ -49,7 +54,7 @@ class SavedMessageService: if saved_message: return - message = MessageService.get_message(app_model=app_model, user=user, message_id=message_id) + message = MessageService.get_message(app_model=app_model, user=user, message_id=message_id, session=session) saved_message = SavedMessage( app_id=app_model.id, @@ -62,7 +67,7 @@ class SavedMessageService: session.commit() @classmethod - def delete(cls, session: Session, app_model: App, user: Account | EndUser | None, message_id: str): + def delete(cls, app_model: App, user: Account | EndUser | None, message_id: str, *, session: Session): if not user: return saved_message = session.scalar( diff --git a/api/services/snippet_dsl_service.py b/api/services/snippet_dsl_service.py index 19e0dabab95..ae1cd7f7e31 100644 --- a/api/services/snippet_dsl_service.py +++ b/api/services/snippet_dsl_service.py @@ -3,12 +3,10 @@ import logging import uuid from collections.abc import Mapping from datetime import UTC, datetime -from enum import StrEnum from urllib.parse import urlparse import yaml -from packaging import version -from pydantic import BaseModel, Field +from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.orm import Session @@ -20,6 +18,9 @@ from graphon.model_runtime.utils.encoders import jsonable_encoder from models import Account from models.snippet import CustomizedSnippet, SnippetType from models.workflow import Workflow +from services.dsl_content import DSL_MAX_SIZE, dsl_content_size +from services.dsl_version import check_version_compatibility +from services.entities.dsl_entities import CheckDependenciesResult, ImportMode, ImportStatus from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.snippet_service import SNIPPET_FORBIDDEN_NODE_TYPES, SnippetService @@ -28,22 +29,9 @@ logger = logging.getLogger(__name__) IMPORT_INFO_REDIS_KEY_PREFIX = "snippet_import_info:" CHECK_DEPENDENCIES_REDIS_KEY_PREFIX = "snippet_check_dependencies:" IMPORT_INFO_REDIS_EXPIRY = 10 * 60 # 10 minutes -DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB CURRENT_DSL_VERSION = "0.1.0" -class ImportMode(StrEnum): - YAML_CONTENT = "yaml-content" - YAML_URL = "yaml-url" - - -class ImportStatus(StrEnum): - COMPLETED = "completed" - COMPLETED_WITH_WARNINGS = "completed-with-warnings" - PENDING = "pending" - FAILED = "failed" - - class SnippetImportInfo(BaseModel): id: str status: ImportStatus @@ -53,32 +41,9 @@ class SnippetImportInfo(BaseModel): error: str = "" -class CheckDependenciesResult(BaseModel): - leaked_dependencies: list[PluginDependency] = Field(default_factory=list) - - def _check_version_compatibility(imported_version: str) -> ImportStatus: - """Determine import status based on version comparison""" - try: - current_ver = version.parse(CURRENT_DSL_VERSION) - imported_ver = version.parse(imported_version) - except version.InvalidVersion: - return ImportStatus.FAILED - - # If imported version is newer than current, always return PENDING - if imported_ver > current_ver: - return ImportStatus.PENDING - - # If imported version is older than current's major, return PENDING - if imported_ver.major < current_ver.major: - return ImportStatus.PENDING - - # If imported version is older than current's minor, return COMPLETED_WITH_WARNINGS - if imported_ver.minor < current_ver.minor: - return ImportStatus.COMPLETED_WITH_WARNINGS - - # If imported version equals or is older than current's micro, return COMPLETED - return ImportStatus.COMPLETED + """Determine import status based on version comparison.""" + return check_version_compatibility(imported_version, CURRENT_DSL_VERSION) class SnippetPendingData(BaseModel): @@ -145,13 +110,14 @@ class SnippetDslService: status=ImportStatus.FAILED, error=f"Failed to fetch YAML from URL: {response.status_code}", ) - content = response.text - if len(content) > DSL_MAX_SIZE: + raw_content = response.content + if dsl_content_size(raw_content) > DSL_MAX_SIZE: return SnippetImportInfo( id=import_id, status=ImportStatus.FAILED, error=f"YAML content size exceeds maximum limit of {DSL_MAX_SIZE} bytes", ) + content = raw_content.decode("utf-8") except Exception as e: logger.exception("Failed to fetch YAML from URL") return SnippetImportInfo( @@ -167,7 +133,7 @@ class SnippetDslService: error="yaml_content is required when import_mode is yaml-content", ) content = yaml_content - if len(content) > DSL_MAX_SIZE: + if dsl_content_size(content) > DSL_MAX_SIZE: return SnippetImportInfo( id=import_id, status=ImportStatus.FAILED, diff --git a/api/services/snippet_service.py b/api/services/snippet_service.py index a54c9f6a069..64c1ec12370 100644 --- a/api/services/snippet_service.py +++ b/api/services/snippet_service.py @@ -6,9 +6,8 @@ from datetime import UTC, datetime from typing import Any from sqlalchemy import delete, func, select -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session, sessionmaker -from core.db import session_factory from core.workflow.node_factory import LATEST_VERSION, NODE_TYPE_CLASSES_MAPPING from graphon.enums import BuiltinNodeTypes, NodeType from libs.infinite_scroll_pagination import InfiniteScrollPagination @@ -59,9 +58,8 @@ class SnippetService: session_maker = None if session is not None: session_maker = sessionmaker(bind=session.get_bind(), expire_on_commit=False) - elif session_maker is None: - session_maker = session_factory.get_session_maker() - assert session_maker is not None + if session_maker is None: + raise ValueError("SnippetService requires a session or session_maker.") self._session = session self._session_maker = session_maker self._node_execution_service_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( @@ -192,7 +190,7 @@ class SnippetService: self, *, tenant_id: str, - session: scoped_session, + session: Session, page: int = 1, limit: int = 20, keyword: str | None = None, diff --git a/api/services/summary_index_service.py b/api/services/summary_index_service.py index 3e065653bdf..a960b219334 100644 --- a/api/services/summary_index_service.py +++ b/api/services/summary_index_service.py @@ -7,7 +7,7 @@ from datetime import UTC, datetime from typing import TypedDict, cast from sqlalchemy import select -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from core.db.session_factory import session_factory from core.model_manager import ModelManager @@ -94,6 +94,8 @@ class SummaryIndexService: dataset: Dataset, summary_content: str, status: SummaryStatus = SummaryStatus.GENERATING, + *, + session: Session, ) -> DocumentSegmentSummary: """ Create or update a DocumentSegmentSummary record. @@ -105,46 +107,48 @@ class SummaryIndexService: summary_content: Generated summary content status: Summary status (default: SummaryStatus.GENERATING) + Keyword Args: + session: SQLAlchemy session used for the summary record. + Returns: Created or updated DocumentSegmentSummary instance """ - with session_factory.create_session() as session: - # Check if summary record already exists - existing_summary = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) + # Check if summary record already exists + existing_summary = session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment.id, + DocumentSegmentSummary.dataset_id == dataset.id, ) + .limit(1) + ) - if existing_summary: - # Update existing record - existing_summary.summary_content = summary_content - existing_summary.status = status - existing_summary.error = None # Clear any previous errors - # Re-enable if it was disabled - if not existing_summary.enabled: - existing_summary.enabled = True - existing_summary.disabled_at = None - existing_summary.disabled_by = None - session.add(existing_summary) - session.flush() - return existing_summary - else: - # Create new record (enabled by default) - summary_record = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=segment.document_id, - chunk_id=segment.id, - summary_content=summary_content, - status=status, - enabled=True, # Explicitly set enabled to True - ) - session.add(summary_record) - session.flush() - return summary_record + if existing_summary: + # Update existing record + existing_summary.summary_content = summary_content + existing_summary.status = status + existing_summary.error = None # Clear any previous errors + # Re-enable if it was disabled + if not existing_summary.enabled: + existing_summary.enabled = True + existing_summary.disabled_at = None + existing_summary.disabled_by = None + session.add(existing_summary) + session.flush() + return existing_summary + else: + # Create new record (enabled by default) + summary_record = DocumentSegmentSummary( + dataset_id=dataset.id, + document_id=segment.document_id, + chunk_id=segment.id, + summary_content=summary_content, + status=status, + enabled=True, # Explicitly set enabled to True + ) + session.add(summary_record) + session.flush() + return summary_record @staticmethod def vectorize_summary( @@ -641,6 +645,8 @@ class SummaryIndexService: segment: DocumentSegment, dataset: Dataset, summary_index_setting: SummaryIndexSettingDict, + *, + session: Session, ) -> DocumentSegmentSummary: """ Generate summary for a segment and vectorize it. @@ -651,106 +657,101 @@ class SummaryIndexService: dataset: Dataset containing the segment summary_index_setting: Summary index configuration + Keyword Args: + session: SQLAlchemy session used for summary record updates. + Returns: Created DocumentSegmentSummary instance Raises: ValueError: If summary generation fails """ - with session_factory.create_session() as session: - try: - # Get or refresh summary record in this session - summary_record_in_session = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) + try: + # Get or refresh summary record in this session + summary_record_in_session = session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment.id, + DocumentSegmentSummary.dataset_id == dataset.id, ) + .limit(1) + ) - if not summary_record_in_session: - # If not found, create one - logger.warning("Summary record not found for segment %s, creating one", segment.id) - summary_record_in_session = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=segment.document_id, - chunk_id=segment.id, - summary_content="", - status=SummaryStatus.GENERATING, - enabled=True, - ) - session.add(summary_record_in_session) - session.flush() - - # Update status to "generating" - summary_record_in_session.status = SummaryStatus.GENERATING - summary_record_in_session.error = None - session.add(summary_record_in_session) - # Don't flush here - wait until after vectorization succeeds - - # Generate summary (returns summary_content and llm_usage) - summary_content, llm_usage = SummaryIndexService.generate_summary_for_segment( - segment, dataset, summary_index_setting + if not summary_record_in_session: + # If not found, create one + logger.warning("Summary record not found for segment %s, creating one", segment.id) + summary_record_in_session = DocumentSegmentSummary( + dataset_id=dataset.id, + document_id=segment.document_id, + chunk_id=segment.id, + summary_content="", + status=SummaryStatus.GENERATING, + enabled=True, ) - - # Update summary content - summary_record_in_session.summary_content = summary_content session.add(summary_record_in_session) - # Flush to ensure summary_content is saved before vectorize_summary queries it session.flush() - # Log LLM usage for summary generation - if llm_usage and llm_usage.total_tokens > 0: - logger.info( - "Summary generation for segment %s used %s tokens (prompt: %s, completion: %s)", - segment.id, - llm_usage.total_tokens, - llm_usage.prompt_tokens, - llm_usage.completion_tokens, - ) + # Update status to "generating" + summary_record_in_session.status = SummaryStatus.GENERATING + summary_record_in_session.error = None + session.add(summary_record_in_session) + # Don't flush here - wait until after vectorization succeeds - # Vectorize summary (will delete old vector if exists before creating new one) - # Pass the session-managed record to vectorize_summary - # vectorize_summary will update status to "completed" and tokens in its own session - # vectorize_summary will also ensure summary_content is preserved - try: - # Pass the session to vectorize_summary to avoid session isolation issues - SummaryIndexService.vectorize_summary(summary_record_in_session, segment, dataset, session=session) - # Refresh the object from database to get the updated status and tokens from vectorize_summary - session.refresh(summary_record_in_session) - # Commit the session - # (summary_record_in_session should have status="completed" and tokens from refresh) - session.commit() - logger.info("Successfully generated and vectorized summary for segment %s", segment.id) - return summary_record_in_session - except Exception as vectorize_error: - # If vectorization fails, update status to error in current session - logger.exception("Failed to vectorize summary for segment %s", segment.id) - summary_record_in_session.status = SummaryStatus.ERROR - summary_record_in_session.error = f"Vectorization failed: {str(vectorize_error)}" - session.add(summary_record_in_session) - session.commit() - raise + # Generate summary (returns summary_content and llm_usage) + summary_content, llm_usage = SummaryIndexService.generate_summary_for_segment( + segment, dataset, summary_index_setting + ) - except Exception as e: - logger.exception("Failed to generate summary for segment %s", segment.id) - # Update summary record with error status - summary_record_in_session = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) + # Update summary content + summary_record_in_session.summary_content = summary_content + session.add(summary_record_in_session) + # Flush to ensure summary_content is saved before vectorize_summary queries it + session.flush() + + # Log LLM usage for summary generation + if llm_usage and llm_usage.total_tokens > 0: + logger.info( + "Summary generation for segment %s used %s tokens (prompt: %s, completion: %s)", + segment.id, + llm_usage.total_tokens, + llm_usage.prompt_tokens, + llm_usage.completion_tokens, ) - if summary_record_in_session: - summary_record_in_session.status = SummaryStatus.ERROR - summary_record_in_session.error = str(e) - session.add(summary_record_in_session) - session.commit() + + try: + SummaryIndexService.vectorize_summary(summary_record_in_session, segment, dataset, session=session) + # vectorize_summary mutates status and token fields; refresh before returning the ORM object. + session.refresh(summary_record_in_session) + session.commit() + logger.info("Successfully generated and vectorized summary for segment %s", segment.id) + return summary_record_in_session + except Exception as vectorize_error: + # If vectorization fails, update status to error in current session + logger.exception("Failed to vectorize summary for segment %s", segment.id) + summary_record_in_session.status = SummaryStatus.ERROR + summary_record_in_session.error = f"Vectorization failed: {str(vectorize_error)}" + session.add(summary_record_in_session) + session.commit() raise + except Exception as e: + logger.exception("Failed to generate summary for segment %s", segment.id) + # Update summary record with error status + summary_record_in_session = session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment.id, + DocumentSegmentSummary.dataset_id == dataset.id, + ) + .limit(1) + ) + if summary_record_in_session: + summary_record_in_session.status = SummaryStatus.ERROR + summary_record_in_session.error = str(e) + session.add(summary_record_in_session) + session.commit() + raise + @staticmethod def generate_summaries_for_document( dataset: Dataset, @@ -840,7 +841,7 @@ class SummaryIndexService: try: summary_record = SummaryIndexService.generate_and_vectorize_summary( - segment, dataset, summary_index_setting + segment, dataset, summary_index_setting, session=session ) summary_records.append(summary_record) except Exception as e: @@ -1048,6 +1049,8 @@ class SummaryIndexService: segment: DocumentSegment, dataset: Dataset, summary_content: str, + *, + session: Session, ) -> DocumentSegmentSummary | None: """ Update summary for a segment and re-vectorize it. @@ -1057,6 +1060,9 @@ class SummaryIndexService: dataset: Dataset containing the segment summary_content: New summary content + Keyword Args: + session: SQLAlchemy session used for summary record updates. + Returns: Updated DocumentSegmentSummary instance, or None if indexing technique is not high_quality """ @@ -1072,67 +1078,22 @@ class SummaryIndexService: if segment.document and segment.document.doc_form == "qa_model": return None - with session_factory.create_session() as session: - try: - # Check if summary_content is empty (whitespace-only strings are considered empty) - if not summary_content or not summary_content.strip(): - # If summary is empty, only delete existing summary vector and record - summary_record = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - - if summary_record: - # Delete old vector if exists - old_summary_node_id = summary_record.summary_index_node_id - if old_summary_node_id: - try: - vector = Vector(dataset) - vector.delete_by_ids([old_summary_node_id]) - except Exception as e: - logger.warning( - "Failed to delete old summary vector for segment %s: %s", - segment.id, - str(e), - ) - - # Delete summary record since summary is empty - session.delete(summary_record) - session.commit() - logger.info("Deleted summary for segment %s (empty content provided)", segment.id) - return None - else: - # No existing summary record, nothing to do - logger.info("No summary record found for segment %s, nothing to delete", segment.id) - return None - - # Find existing summary record - summary_record = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) + try: + summary_record = session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment.id, + DocumentSegmentSummary.dataset_id == dataset.id, ) + .limit(1) + ) + # Check if summary_content is empty (whitespace-only strings are considered empty) + if not summary_content or not summary_content.strip(): + # If summary is empty, only delete existing summary vector and record if summary_record: - # Update existing summary + # Delete old vector if exists old_summary_node_id = summary_record.summary_index_node_id - - # Update summary content - summary_record.summary_content = summary_content - summary_record.status = SummaryStatus.GENERATING - summary_record.error = None # Clear any previous errors - session.add(summary_record) - # Flush to ensure summary_content is saved before vectorize_summary queries it - session.flush() - - # Delete old vector if exists (before vectorization) if old_summary_node_id: try: vector = Vector(dataset) @@ -1144,80 +1105,90 @@ class SummaryIndexService: str(e), ) - # Re-vectorize summary (this will update status to "completed" and tokens in its own session) - # vectorize_summary will also ensure summary_content is preserved - # Note: vectorize_summary may take time due to embedding API calls, but we need to complete it - # to ensure the summary is properly indexed - try: - # Pass the session to vectorize_summary to avoid session isolation issues - SummaryIndexService.vectorize_summary(summary_record, segment, dataset, session=session) - # Refresh the object from database to get the updated status and tokens from vectorize_summary - session.refresh(summary_record) - # Now commit the session (summary_record should have status="completed" and tokens from refresh) - session.commit() - logger.info("Successfully updated and re-vectorized summary for segment %s", segment.id) - return summary_record - except Exception as e: - # If vectorization fails, update status to error in current session - # Don't raise the exception - just log it and return the record with error status - # This allows the segment update to complete even if vectorization fails - summary_record.status = SummaryStatus.ERROR - summary_record.error = f"Vectorization failed: {str(e)}" - session.commit() - logger.exception("Failed to vectorize summary for segment %s", segment.id) - # Return the record with error status instead of raising - # The caller can check the status if needed - return summary_record - else: - # Create new summary record if doesn't exist - summary_record = SummaryIndexService.create_summary_record( - segment, dataset, summary_content, status=SummaryStatus.GENERATING - ) - # Re-vectorize summary (this will update status to "completed" and tokens in its own session) - # Note: summary_record was created in a different session, - # so we need to merge it into current session - try: - # Merge the record into current session first (since it was created in a different session) - summary_record = session.merge(summary_record) - # Pass the session to vectorize_summary - it will update the merged record - SummaryIndexService.vectorize_summary(summary_record, segment, dataset, session=session) - # Refresh to get updated status and tokens from database - session.refresh(summary_record) - # Commit the session to persist the changes - session.commit() - logger.info("Successfully created and vectorized summary for segment %s", segment.id) - return summary_record - except Exception as e: - # If vectorization fails, update status to error in current session - # Merge the record into current session first - error_record = session.merge(summary_record) - error_record.status = SummaryStatus.ERROR - error_record.error = f"Vectorization failed: {str(e)}" - session.commit() - logger.exception("Failed to vectorize summary for segment %s", segment.id) - # Return the record with error status instead of raising - return error_record - - except Exception as e: - logger.exception("Failed to update summary for segment %s", segment.id) - # Update summary record with error status if it exists - summary_record = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - if summary_record: - summary_record.status = SummaryStatus.ERROR - summary_record.error = str(e) - session.add(summary_record) + # Delete summary record since summary is empty + session.delete(summary_record) session.commit() - raise + logger.info("Deleted summary for segment %s (empty content provided)", segment.id) + return None + else: + # No existing summary record, nothing to do + logger.info("No summary record found for segment %s, nothing to delete", segment.id) + return None + + if summary_record: + # Update existing summary + old_summary_node_id = summary_record.summary_index_node_id + + # Update summary content + summary_record.summary_content = summary_content + summary_record.status = SummaryStatus.GENERATING + summary_record.error = None # Clear any previous errors + session.add(summary_record) + # Flush to ensure summary_content is saved before vectorize_summary queries it + session.flush() + + # Delete old vector if exists (before vectorization) + if old_summary_node_id: + try: + vector = Vector(dataset) + vector.delete_by_ids([old_summary_node_id]) + except Exception as e: + logger.warning( + "Failed to delete old summary vector for segment %s: %s", + segment.id, + str(e), + ) + else: + # Create new summary record if doesn't exist + summary_record = SummaryIndexService.create_summary_record( + segment, + dataset, + summary_content, + status=SummaryStatus.GENERATING, + session=session, + ) + + try: + # Vectorization must finish here so the manual summary is searchable immediately. + SummaryIndexService.vectorize_summary(summary_record, segment, dataset, session=session) + session.refresh(summary_record) + session.commit() + logger.info("Successfully updated and re-vectorized summary for segment %s", segment.id) + return summary_record + except Exception as e: + # If vectorization fails, update status to error in current session. + # Return the record with error status so callers can still finish segment updates. + summary_record.status = SummaryStatus.ERROR + summary_record.error = f"Vectorization failed: {str(e)}" + session.commit() + logger.exception("Failed to vectorize summary for segment %s", segment.id) + return summary_record + + except Exception as e: + logger.exception("Failed to update summary for segment %s", segment.id) + # Update summary record with error status if it exists + summary_record = session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment.id, + DocumentSegmentSummary.dataset_id == dataset.id, + ) + .limit(1) + ) + if summary_record: + summary_record.status = SummaryStatus.ERROR + summary_record.error = str(e) + session.add(summary_record) + session.commit() + raise @staticmethod - def get_segment_summary(segment_id: str, dataset_id: str) -> DocumentSegmentSummary | None: + def get_segment_summary( + segment_id: str, + dataset_id: str, + *, + session: Session, + ) -> DocumentSegmentSummary | None: """ Get summary for a single segment. @@ -1225,22 +1196,29 @@ class SummaryIndexService: segment_id: Segment ID (chunk_id) dataset_id: Dataset ID + Keyword Args: + session: SQLAlchemy session used to read summary records. + Returns: DocumentSegmentSummary instance if found, None otherwise """ - with session_factory.create_session() as session: - return session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment_id, - DocumentSegmentSummary.dataset_id == dataset_id, - DocumentSegmentSummary.enabled.is_(True), # Only return enabled summaries - ) - .limit(1) + return session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.chunk_id == segment_id, + DocumentSegmentSummary.dataset_id == dataset_id, + DocumentSegmentSummary.enabled.is_(True), ) + .limit(1) + ) @staticmethod - def get_segments_summaries(segment_ids: list[str], dataset_id: str) -> dict[str, DocumentSegmentSummary]: + def get_segments_summaries( + segment_ids: list[str], + dataset_id: str, + *, + session: Session, + ) -> dict[str, DocumentSegmentSummary]: """ Get summaries for multiple segments. @@ -1248,26 +1226,31 @@ class SummaryIndexService: segment_ids: List of segment IDs (chunk_ids) dataset_id: Dataset ID + Keyword Args: + session: SQLAlchemy session used to read summary records. + Returns: Dictionary mapping segment_id to DocumentSegmentSummary (only enabled summaries) """ if not segment_ids: return {} - with session_factory.create_session() as session: - summary_records = session.scalars( - select(DocumentSegmentSummary).where( - DocumentSegmentSummary.chunk_id.in_(segment_ids), - DocumentSegmentSummary.dataset_id == dataset_id, - DocumentSegmentSummary.enabled.is_(True), # Only return enabled summaries - ) - ).all() - - return {summary.chunk_id: summary for summary in summary_records} + summaries = session.scalars( + select(DocumentSegmentSummary).where( + DocumentSegmentSummary.chunk_id.in_(segment_ids), + DocumentSegmentSummary.dataset_id == dataset_id, + DocumentSegmentSummary.enabled.is_(True), + ) + ).all() + return {summary.chunk_id: summary for summary in summaries} @staticmethod def get_document_summaries( - document_id: str, dataset_id: str, segment_ids: list[str] | None = None + document_id: str, + dataset_id: str, + segment_ids: list[str] | None = None, + *, + session: Session, ) -> list[DocumentSegmentSummary]: """ Get all summary records for a document. @@ -1277,23 +1260,31 @@ class SummaryIndexService: dataset_id: Dataset ID segment_ids: Optional list of segment IDs to filter by + Keyword Args: + session: SQLAlchemy session used to read summary records. + Returns: List of DocumentSegmentSummary instances (only enabled summaries) """ - with session_factory.create_session() as session: - stmt = select(DocumentSegmentSummary).where( - DocumentSegmentSummary.document_id == document_id, - DocumentSegmentSummary.dataset_id == dataset_id, - DocumentSegmentSummary.enabled.is_(True), # Only return enabled summaries - ) + stmt = select(DocumentSegmentSummary).where( + DocumentSegmentSummary.document_id == document_id, + DocumentSegmentSummary.dataset_id == dataset_id, + DocumentSegmentSummary.enabled.is_(True), + ) - if segment_ids: - stmt = stmt.where(DocumentSegmentSummary.chunk_id.in_(segment_ids)) + if segment_ids: + stmt = stmt.where(DocumentSegmentSummary.chunk_id.in_(segment_ids)) - return list(session.scalars(stmt).all()) + return list(session.scalars(stmt).all()) @staticmethod - def get_document_summary_index_status(document_id: str, dataset_id: str, tenant_id: str) -> str | None: + def get_document_summary_index_status( + document_id: str, + dataset_id: str, + tenant_id: str, + *, + session: Session, + ) -> str | None: """ Get summary_index_status for a single document. @@ -1302,26 +1293,28 @@ class SummaryIndexService: dataset_id: Dataset ID tenant_id: Tenant ID + Keyword Args: + session: SQLAlchemy session used to read summary status. + Returns: "SUMMARIZING" if there are pending summaries, None otherwise """ # Get all segments for this document (excluding qa_model and re_segment) - with session_factory.create_session() as session: - segment_ids = list( - session.scalars( - select(DocumentSegment.id).where( - DocumentSegment.document_id == document_id, - DocumentSegment.status != "re_segment", - DocumentSegment.tenant_id == tenant_id, - ) - ).all() - ) + segment_ids = list( + session.scalars( + select(DocumentSegment.id).where( + DocumentSegment.document_id == document_id, + DocumentSegment.status != "re_segment", + DocumentSegment.tenant_id == tenant_id, + ) + ).all() + ) if not segment_ids: return None # Get all summary records for these segments - summaries = SummaryIndexService.get_segments_summaries(segment_ids, dataset_id) + summaries = SummaryIndexService.get_segments_summaries(segment_ids, dataset_id, session=session) summary_status_map = {chunk_id: summary.status for chunk_id, summary in summaries.items()} # Check if there are any "not_started" or "generating" status summaries @@ -1335,7 +1328,11 @@ class SummaryIndexService: @staticmethod def get_documents_summary_index_status( - document_ids: list[str], dataset_id: str, tenant_id: str + document_ids: list[str], + dataset_id: str, + tenant_id: str, + *, + session: Session, ) -> dict[str, str | None]: """ Get summary_index_status for multiple documents. @@ -1345,6 +1342,9 @@ class SummaryIndexService: dataset_id: Dataset ID tenant_id: Tenant ID + Keyword Args: + session: SQLAlchemy session used to read summary status. + Returns: Dictionary mapping document_id to summary_index_status ("SUMMARIZING" or None) """ @@ -1352,14 +1352,13 @@ class SummaryIndexService: return {} # Get all segments for these documents (excluding qa_model and re_segment) - with session_factory.create_session() as session: - segments = session.execute( - select(DocumentSegment.id, DocumentSegment.document_id).where( - DocumentSegment.document_id.in_(document_ids), - DocumentSegment.status != "re_segment", - DocumentSegment.tenant_id == tenant_id, - ) - ).all() + segments = session.execute( + select(DocumentSegment.id, DocumentSegment.document_id).where( + DocumentSegment.document_id.in_(document_ids), + DocumentSegment.status != "re_segment", + DocumentSegment.tenant_id == tenant_id, + ) + ).all() # Group segments by document_id document_segments_map: dict[str, list[str]] = {} @@ -1371,7 +1370,7 @@ class SummaryIndexService: # Get all summary records for these segments all_segment_ids = [seg.id for seg in segments] - summaries = SummaryIndexService.get_segments_summaries(all_segment_ids, dataset_id) + summaries = SummaryIndexService.get_segments_summaries(all_segment_ids, dataset_id, session=session) summary_status_map = {chunk_id: summary.status for chunk_id, summary in summaries.items()} # Calculate summary_index_status for each document @@ -1407,7 +1406,7 @@ class SummaryIndexService: def get_document_summary_status_detail( document_id: str, dataset_id: str, - session: Session | scoped_session, + session: Session, ) -> DocumentSummaryStatusDetailDict: """ Get detailed summary status for a document. @@ -1425,6 +1424,7 @@ class SummaryIndexService: - generating: Number of summaries being generated - error: Number of summaries with errors - not_started: Number of segments without summary records + - timeout: Number of summaries that timed out - summaries: List of summary records with status and content preview """ from services.dataset_service import SegmentService @@ -1448,6 +1448,7 @@ class SummaryIndexService: document_id=document_id, dataset_id=dataset_id, segment_ids=segment_ids, + session=session, ) # Create a mapping of chunk_id to summary diff --git a/api/services/tag_service.py b/api/services/tag_service.py index 2d89bafa920..f404ec0eb37 100644 --- a/api/services/tag_service.py +++ b/api/services/tag_service.py @@ -6,7 +6,7 @@ from flask_login import current_user from pydantic import BaseModel, Field from sqlalchemy import delete, func, select from sqlalchemy.engine import CursorResult -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound from models.dataset import Dataset @@ -14,7 +14,6 @@ from models.enums import TagType from models.model import App, Tag, TagBinding from models.snippet import CustomizedSnippet -type _SessionLike = Session | scoped_session type _TagTypeLike = TagType | str @@ -41,7 +40,7 @@ class TagBindingDeletePayload(BaseModel): class TagService: @staticmethod - def get_tags(session: Session, tag_type: _TagTypeLike, current_tenant_id: str, keyword: str | None = None): + def get_tags(tag_type: _TagTypeLike, current_tenant_id: str, keyword: str | None = None, *, session: Session): stmt = ( select(Tag.id, Tag.type, Tag.name, func.count(TagBinding.id).label("binding_count")) .outerjoin(TagBinding, Tag.id == TagBinding.tag_id) @@ -61,7 +60,7 @@ class TagService: tag_type: _TagTypeLike, current_tenant_id: str, tag_ids: list[str], - session: _SessionLike, + session: Session, *, match_all: bool = False, ): @@ -107,7 +106,7 @@ class TagService: return tag_bindings @staticmethod - def get_tag_by_tag_name(tag_type: _TagTypeLike, current_tenant_id: str, tag_name: str, session: _SessionLike): + def get_tag_by_tag_name(tag_type: _TagTypeLike, current_tenant_id: str, tag_name: str, session: Session): if not tag_type or not tag_name: return [] tags = list( @@ -120,7 +119,7 @@ class TagService: return tags @staticmethod - def get_tags_by_target_id(tag_type: _TagTypeLike, current_tenant_id: str, target_id: str, session: _SessionLike): + def get_tags_by_target_id(tag_type: _TagTypeLike, current_tenant_id: str, target_id: str, session: Session): tags = session.scalars( select(Tag) .join(TagBinding, Tag.id == TagBinding.tag_id) @@ -135,7 +134,7 @@ class TagService: return tags or [] @staticmethod - def save_tags(payload: SaveTagPayload, session: _SessionLike) -> Tag: + def save_tags(payload: SaveTagPayload, session: Session) -> Tag: if TagService.get_tag_by_tag_name(payload.type, current_user.current_tenant_id, payload.name, session): raise ValueError("Tag name already exists") tag = Tag( @@ -151,7 +150,7 @@ class TagService: @staticmethod def update_tags( - payload: UpdateTagPayload, tag_id: str, session: _SessionLike, *, tag_type: TagType | None = None + payload: UpdateTagPayload, tag_id: str, session: Session, *, tag_type: TagType | None = None ) -> Tag: current_tenant_id = current_user.current_tenant_id stmt = select(Tag).where(Tag.id == tag_id, Tag.tenant_id == current_tenant_id) @@ -178,7 +177,7 @@ class TagService: return tag @staticmethod - def get_tag_binding_count(tag_id: str, session: _SessionLike, *, tag_type: TagType | None = None) -> int: + def get_tag_binding_count(tag_id: str, session: Session, *, tag_type: TagType | None = None) -> int: current_tenant_id = current_user.current_tenant_id stmt = ( select(func.count(TagBinding.id)) @@ -191,7 +190,7 @@ class TagService: return count @staticmethod - def delete_tag(tag_id: str, session: _SessionLike, *, tag_type: TagType | None = None): + def delete_tag(tag_id: str, session: Session, *, tag_type: TagType | None = None): current_tenant_id = current_user.current_tenant_id stmt = select(Tag).where(Tag.id == tag_id, Tag.tenant_id == current_tenant_id) if tag_type is not None: @@ -210,7 +209,7 @@ class TagService: session.commit() @staticmethod - def save_tag_binding(payload: TagBindingCreatePayload, session: _SessionLike): + def save_tag_binding(payload: TagBindingCreatePayload, session: Session): TagService.check_target_exists(payload.type, payload.target_id, session) valid_tag_ids = session.scalars( select(Tag.id).where( @@ -237,7 +236,7 @@ class TagService: session.commit() @staticmethod - def delete_tag_binding(payload: TagBindingDeletePayload, session: _SessionLike): + def delete_tag_binding(payload: TagBindingDeletePayload, session: Session): TagService.check_target_exists(payload.type, payload.target_id, session) result = cast( CursorResult, @@ -260,7 +259,7 @@ class TagService: session.commit() @staticmethod - def check_target_exists(type: _TagTypeLike, target_id: str, session: _SessionLike): + def check_target_exists(type: _TagTypeLike, target_id: str, session: Session): if type == "knowledge": dataset = session.scalar( select(Dataset) diff --git a/api/services/tools/api_tools_manage_service.py b/api/services/tools/api_tools_manage_service.py index 5ff2c217492..032e6cba4d1 100644 --- a/api/services/tools/api_tools_manage_service.py +++ b/api/services/tools/api_tools_manage_service.py @@ -1,12 +1,13 @@ import json import logging -from typing import Any, TypedDict, cast +from typing import Any, Literal, TypedDict, cast from httpx import get +from pydantic import TypeAdapter from sqlalchemy import select from sqlalchemy.orm import sessionmaker -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.__base.tool_runtime import ToolRuntime from core.tools.custom_tool.provider import ApiToolProviderController from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity @@ -22,7 +23,6 @@ from core.tools.tool_manager import ToolManager from core.tools.utils.encryption import create_tool_provider_encrypter from core.tools.utils.parser import ApiBasedToolSchemaParser from extensions.ext_database import db -from graphon.model_runtime.utils.encoders import jsonable_encoder from models.tools import ApiToolProvider from services.tools.tools_transform_service import ToolTransformService @@ -36,6 +36,27 @@ class ApiSchemaParseResult(TypedDict): warning: dict[str, str] +class ApiToolPreviewResult(TypedDict, total=False): + result: str + error: str + + +class RemoteSchemaResult(TypedDict): + schema: str + + +class SimpleSuccessResult(TypedDict): + result: Literal["success"] + + +def _dump_api_tool_bundles(tool_bundles: list[ApiToolBundle]) -> list[dict[str, Any]]: + return cast(list[dict[str, Any]], TypeAdapter(list[ApiToolBundle]).dump_python(tool_bundles, mode="json")) + + +def _dump_provider_configs(configs: list[ProviderConfig]) -> list[dict[str, Any]]: + return cast(list[dict[str, Any]], TypeAdapter(list[ProviderConfig]).dump_python(configs, mode="json")) + + class ApiToolManageService: @staticmethod def parser_api_schema(schema: str) -> ApiSchemaParseResult: @@ -52,7 +73,7 @@ class ApiToolManageService: credentials_schema = [ ProviderConfig( name="auth_type", - type=ProviderConfig.Type.SELECT, + type=ProviderConfigType.SELECT, required=True, default="none", options=[ @@ -63,7 +84,7 @@ class ApiToolManageService: ), ProviderConfig( name="api_key_header", - type=ProviderConfig.Type.TEXT_INPUT, + type=ProviderConfigType.TEXT_INPUT, required=False, placeholder=I18nObject(en_US="Enter api key header", zh_Hans="输入 api key header，如：X-API-KEY"), default="api_key", @@ -71,7 +92,7 @@ class ApiToolManageService: ), ProviderConfig( name="api_key_value", - type=ProviderConfig.Type.TEXT_INPUT, + type=ProviderConfigType.TEXT_INPUT, required=False, placeholder=I18nObject(en_US="Enter api key", zh_Hans="输入 api key"), default="", @@ -80,14 +101,12 @@ class ApiToolManageService: return cast( ApiSchemaParseResult, - jsonable_encoder( - { - "schema_type": schema_type, - "parameters_schema": tool_bundles, - "credentials_schema": credentials_schema, - "warning": warnings, - } - ), + { + "schema_type": schema_type.value, + "parameters_schema": _dump_api_tool_bundles(tool_bundles), + "credentials_schema": _dump_provider_configs(credentials_schema), + "warning": warnings, + }, ) except Exception as e: raise ValueError(f"invalid schema: {str(e)}") @@ -118,7 +137,7 @@ class ApiToolManageService: privacy_policy: str, custom_disclaimer: str, labels: list[str], - ) -> dict[str, Any]: + ) -> SimpleSuccessResult: """ Create a new API tool provider. @@ -169,7 +188,7 @@ class ApiToolManageService: schema=schema, description=extra_info.get("description", ""), schema_type_str=schema_type, - tools_str=json.dumps(jsonable_encoder(tool_bundles)), + tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), credentials_str="{}", privacy_policy=privacy_policy, custom_disclaimer=custom_disclaimer, @@ -201,7 +220,7 @@ class ApiToolManageService: return {"result": "success"} @staticmethod - def get_api_tool_provider_remote_schema(user_id: str, tenant_id: str, url: str): + def get_api_tool_provider_remote_schema(user_id: str, tenant_id: str, url: str) -> RemoteSchemaResult: """ get api tool provider remote schema """ @@ -276,7 +295,7 @@ class ApiToolManageService: privacy_policy: str | None, custom_disclaimer: str, labels: list[str], - ) -> dict[str, Any]: + ) -> SimpleSuccessResult: """ Update an existing API tool provider. @@ -322,7 +341,7 @@ class ApiToolManageService: provider.schema = schema provider.description = extra_info.get("description", "") provider.schema_type_str = schema_type - provider.tools_str = json.dumps(jsonable_encoder(tool_bundles)) + provider.tools_str = json.dumps(_dump_api_tool_bundles(tool_bundles)) provider.privacy_policy = privacy_policy provider.custom_disclaimer = custom_disclaimer @@ -365,7 +384,7 @@ class ApiToolManageService: return {"result": "success"} @staticmethod - def delete_api_tool_provider(user_id: str, tenant_id: str, provider_name: str): + def delete_api_tool_provider(user_id: str, tenant_id: str, provider_name: str) -> SimpleSuccessResult: """ Delete an API tool provider. @@ -413,9 +432,9 @@ class ApiToolManageService: tool_name: str, credentials: dict[str, Any], parameters: dict[str, Any], - schema_type: ApiProviderSchemaType, + schema_type: ApiProviderSchemaType | str, schema: str, - ) -> dict[str, Any]: + ) -> ApiToolPreviewResult: """ Test an API tool before adding the API tool provider. @@ -464,7 +483,7 @@ class ApiToolManageService: schema=schema, description="", schema_type_str=ApiProviderSchemaType.OPENAPI, - tools_str=json.dumps(jsonable_encoder(tool_bundles)), + tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), credentials_str=json.dumps(credentials), ) diff --git a/api/services/tools/builtin_tools_manage_service.py b/api/services/tools/builtin_tools_manage_service.py index e49ab8398f1..45480f71d1a 100644 --- a/api/services/tools/builtin_tools_manage_service.py +++ b/api/services/tools/builtin_tools_manage_service.py @@ -327,7 +327,7 @@ class BuiltinToolManageService: @staticmethod def generate_builtin_tool_provider_name( - session: Session, tenant_id: str, provider: str, credential_type: CredentialType + tenant_id: str, provider: str, credential_type: CredentialType, *, session: Session ) -> str: db_providers = session.scalars( select(BuiltinToolProvider) @@ -347,6 +347,7 @@ class BuiltinToolManageService: def get_builtin_tool_provider_credentials( tenant_id: str, provider_name: str, + session: Session, user: Account | None = None, include_credential_ids: list[str] | None = None, ) -> list[ToolProviderCredentialApiEntity]: @@ -367,7 +368,7 @@ class BuiltinToolManageService: from models.credential_permission import CredentialType as CredPermType from services.credential_permission_service import CredentialPermissionService - with db.session.no_autoflush: + with session.no_autoflush: base_filter = ( BuiltinToolProvider.tenant_id == tenant_id, BuiltinToolProvider.provider == provider_name, @@ -383,7 +384,7 @@ class BuiltinToolManageService: credential_type=CredPermType.BUILTIN_TOOL_PROVIDER, user=user, ) - visible_providers = list(db.session.scalars(visible_query).all()) + visible_providers = list(session.scalars(visible_query).all()) # Fetch any explicitly-included IDs that the visibility filter excluded. borrowed_ids: set[str] = set() @@ -397,7 +398,7 @@ class BuiltinToolManageService: .where(*base_filter, BuiltinToolProvider.id.in_(wanted_ids)) .order_by(*order) ) - borrowed_providers = list(db.session.scalars(borrowed_query).all()) + borrowed_providers = list(session.scalars(borrowed_query).all()) borrowed_ids = {p.id for p in borrowed_providers} providers = visible_providers + borrowed_providers @@ -427,7 +428,7 @@ class BuiltinToolManageService: if vis_str == "partial_members": credential_entity.partial_member_list = list( CredentialPermissionService.get_partial_member_list( - db.session, provider.id, CredPermType.BUILTIN_TOOL_PROVIDER + provider.id, CredPermType.BUILTIN_TOOL_PROVIDER, session=session ) ) if provider.id in borrowed_ids: @@ -439,6 +440,7 @@ class BuiltinToolManageService: def get_builtin_tool_provider_credential_info( tenant_id: str, provider: str, + session: Session, user: Account | None = None, include_credential_ids: list[str] | None = None, ) -> ToolProviderCredentialInfoApiEntity: @@ -450,6 +452,7 @@ class BuiltinToolManageService: credentials = BuiltinToolManageService.get_builtin_tool_provider_credentials( tenant_id, provider, + session=session, user=user, include_credential_ids=include_credential_ids, ) diff --git a/api/services/tools/mcp_tools_manage_service.py b/api/services/tools/mcp_tools_manage_service.py index 1654ca87868..ae184ba4561 100644 --- a/api/services/tools/mcp_tools_manage_service.py +++ b/api/services/tools/mcp_tools_manage_service.py @@ -459,13 +459,11 @@ class MCPToolManageService: Returns: JSON string of encrypted data """ - from core.entities.provider_entities import BasicProviderConfig + from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType from core.tools.utils.encryption import create_provider_encrypter # Create config for secret fields - config = [ - BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=field) for field in secret_fields - ] + config = [BasicProviderConfig(type=ProviderConfigType.SECRET_INPUT, name=field) for field in secret_fields] encrypter_instance, _ = create_provider_encrypter( tenant_id=tenant_id, diff --git a/api/services/trigger/schedule_service.py b/api/services/trigger/schedule_service.py index a827222c1dc..495674248b1 100644 --- a/api/services/trigger/schedule_service.py +++ b/api/services/trigger/schedule_service.py @@ -26,10 +26,7 @@ logger = logging.getLogger(__name__) class ScheduleService: @staticmethod def create_schedule( - session: Session, - tenant_id: str, - app_id: str, - config: ScheduleConfig, + tenant_id: str, app_id: str, config: ScheduleConfig, *, session: Session ) -> WorkflowSchedulePlan: """ Create a new schedule with validated configuration. @@ -63,11 +60,7 @@ class ScheduleService: return schedule @staticmethod - def update_schedule( - session: Session, - schedule_id: str, - updates: SchedulePlanUpdate, - ) -> WorkflowSchedulePlan: + def update_schedule(schedule_id: str, updates: SchedulePlanUpdate, *, session: Session) -> WorkflowSchedulePlan: """ Update an existing schedule with validated configuration. @@ -110,10 +103,7 @@ class ScheduleService: return schedule @staticmethod - def delete_schedule( - session: Session, - schedule_id: str, - ) -> None: + def delete_schedule(schedule_id: str, *, session: Session) -> None: """ Delete a schedule plan. @@ -129,7 +119,7 @@ class ScheduleService: session.flush() @staticmethod - def get_tenant_owner(session: Session, tenant_id: str) -> Account: + def get_tenant_owner(tenant_id: str, *, session: Session) -> Account: """ Returns an account to execute scheduled workflows on behalf of the tenant. Prioritizes owner over admin to ensure proper authorization hierarchy. @@ -157,10 +147,7 @@ class ScheduleService: raise AccountNotFoundError(f"Account not found for tenant: {tenant_id}") @staticmethod - def update_next_run_at( - session: Session, - schedule_id: str, - ) -> datetime: + def update_next_run_at(schedule_id: str, *, session: Session) -> datetime: """ Advances the schedule to its next execution time after a successful trigger. Uses current time as base to prevent missing executions during delays. diff --git a/api/services/trigger/trigger_provider_service.py b/api/services/trigger/trigger_provider_service.py index b0a3de1cee8..8506c523a61 100644 --- a/api/services/trigger/trigger_provider_service.py +++ b/api/services/trigger/trigger_provider_service.py @@ -388,7 +388,7 @@ class TriggerProviderService: return subscription @classmethod - def delete_trigger_provider(cls, session: Session, tenant_id: str, subscription_id: str): + def delete_trigger_provider(cls, tenant_id: str, subscription_id: str, *, session: Session): """ Delete a trigger provider subscription within an existing session. diff --git a/api/services/trigger/trigger_subscription_operator_service.py b/api/services/trigger/trigger_subscription_operator_service.py index 5d7785549e6..491723c6ec2 100644 --- a/api/services/trigger/trigger_subscription_operator_service.py +++ b/api/services/trigger/trigger_subscription_operator_service.py @@ -40,12 +40,7 @@ class TriggerSubscriptionOperatorService: return list(subscribers) @classmethod - def delete_plugin_trigger_by_subscription( - cls, - session: Session, - tenant_id: str, - subscription_id: str, - ) -> None: + def delete_plugin_trigger_by_subscription(cls, tenant_id: str, subscription_id: str, *, session: Session) -> None: """Delete a plugin trigger by tenant_id and subscription_id within an existing session Args: diff --git a/api/services/trigger/webhook_service.py b/api/services/trigger/webhook_service.py index 23b3ac55b93..587048e2ccd 100644 --- a/api/services/trigger/webhook_service.py +++ b/api/services/trigger/webhook_service.py @@ -835,11 +835,7 @@ class WebhookService: # NOTE: don not use `with sessionmaker(bind=db.engine, expire_on_commit=False).begin()` # trigger_workflow_async need to handle multipe session commits internally with Session(db.engine, expire_on_commit=False) as session: - AsyncWorkflowService.trigger_workflow_async( - session, - end_user, - trigger_data, - ) + AsyncWorkflowService.trigger_workflow_async(end_user, trigger_data, session=session) quota_charge.commit() except Exception: quota_charge.refund() diff --git a/api/services/vector_service.py b/api/services/vector_service.py index 5b5088ec5a1..faf4fb085d6 100644 --- a/api/services/vector_service.py +++ b/api/services/vector_service.py @@ -1,6 +1,7 @@ import logging from sqlalchemy import delete, select +from sqlalchemy.orm import Session from core.model_manager import ModelInstance, ModelManager from core.rag.datasource.keyword.keyword_factory import Keyword @@ -11,7 +12,6 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind from core.rag.index_processor.index_processor_base import BaseIndexProcessor from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.rag.models.document import AttachmentDocument, Document -from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType from models import UploadFile from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding @@ -24,14 +24,20 @@ logger = logging.getLogger(__name__) class VectorService: @classmethod def create_segments_vector( - cls, keywords_list: list[list[str]] | None, segments: list[DocumentSegment], dataset: Dataset, doc_form: str + cls, + keywords_list: list[list[str]] | None, + segments: list[DocumentSegment], + dataset: Dataset, + doc_form: str, + session: Session, ): + """Create vector records for document segments using the caller's active DB session.""" documents: list[Document] = [] multimodal_documents: list[AttachmentDocument] = [] for segment in segments: if doc_form == IndexStructureType.PARENT_CHILD_INDEX: - dataset_document = db.session.get(DatasetDocument, segment.document_id) + dataset_document = session.get(DatasetDocument, segment.document_id) if not dataset_document: logger.warning( "Expected DatasetDocument record to exist, but none was found, document_id=%s, segment_id=%s", @@ -40,7 +46,7 @@ class VectorService: ) continue # get the process rule - processing_rule = db.session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) + processing_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) if not processing_rule: raise ValueError("No processing rule found.") # get embedding model instance @@ -63,7 +69,13 @@ class VectorService: else: raise ValueError("The knowledge base index technique is not high quality!") cls.generate_child_chunks( - segment, dataset_document, dataset, embedding_model_instance, processing_rule, False + segment, + dataset_document, + dataset, + embedding_model_instance, + processing_rule, + session, + False, ) else: rag_document = Document( @@ -136,8 +148,10 @@ class VectorService: dataset: Dataset, embedding_model_instance: ModelInstance, processing_rule: DatasetProcessRule, + session: Session, regenerate: bool = False, ): + """Generate child chunks and persist them with the caller's active DB session.""" index_processor = IndexProcessorFactory(dataset.doc_form).init_index_processor() assert segment.index_node_id if regenerate: @@ -184,8 +198,8 @@ class VectorService: type=SegmentType.AUTOMATIC, created_by=dataset_document.created_by, ) - db.session.add(child_segment) - db.session.commit() + session.add(child_segment) + session.commit() @classmethod def create_child_chunk_vector(cls, child_segment: ChildChunk, dataset: Dataset): @@ -255,7 +269,10 @@ class VectorService: vector.delete_by_ids([child_chunk.index_node_id]) @classmethod - def update_multimodel_vector(cls, segment: DocumentSegment, attachment_ids: list[str], dataset: Dataset): + def update_multimodel_vector( + cls, segment: DocumentSegment, attachment_ids: list[str], dataset: Dataset, session: Session + ): + """Update multimodal vectors and attachment bindings with the caller's active DB session.""" if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: return @@ -274,19 +291,17 @@ class VectorService: vector.delete_by_ids(old_attachment_ids) # Delete existing segment attachment bindings in one operation - db.session.execute( - delete(SegmentAttachmentBinding).where(SegmentAttachmentBinding.segment_id == segment.id) - ) + session.execute(delete(SegmentAttachmentBinding).where(SegmentAttachmentBinding.segment_id == segment.id)) if not attachment_ids: - db.session.commit() + session.commit() return # Bulk fetch upload files - only fetch needed fields - upload_file_list = db.session.scalars(select(UploadFile).where(UploadFile.id.in_(attachment_ids))).all() + upload_file_list = session.scalars(select(UploadFile).where(UploadFile.id.in_(attachment_ids))).all() if not upload_file_list: - db.session.commit() + session.commit() return # Create a mapping for quick lookup @@ -329,16 +344,16 @@ class VectorService: # Bulk insert all bindings at once if bindings: - db.session.add_all(bindings) + session.add_all(bindings) # Add documents to vector store if any if documents and dataset.is_multimodal: vector.create_multimodal(documents) # Single commit for all operations - db.session.commit() + session.commit() except Exception: logger.exception("Failed to update multimodal vector for segment %s", segment.id) - db.session.rollback() + session.rollback() raise diff --git a/api/services/web_conversation_service.py b/api/services/web_conversation_service.py index 2c8a3be8631..96d95d5f5ac 100644 --- a/api/services/web_conversation_service.py +++ b/api/services/web_conversation_service.py @@ -2,7 +2,6 @@ from sqlalchemy import select from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom -from extensions.ext_database import db from libs.infinite_scroll_pagination import InfiniteScrollPagination from models import Account from models.enums import CreatorUserRole @@ -59,10 +58,10 @@ class WebConversationService: ) @classmethod - def pin(cls, app_model: App, conversation_id: str, user: Account | EndUser | None): + def pin(cls, app_model: App, conversation_id: str, user: Account | EndUser | None, session: Session): if not user: return - pinned_conversation = db.session.scalar( + pinned_conversation = session.scalar( select(PinnedConversation) .where( PinnedConversation.app_id == app_model.id, @@ -77,7 +76,7 @@ class WebConversationService: return conversation = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation_id, user=user + app_model=app_model, conversation_id=conversation_id, user=user, session=session ) pinned_conversation = PinnedConversation( @@ -87,14 +86,14 @@ class WebConversationService: created_by=user.id, ) - db.session.add(pinned_conversation) - db.session.commit() + session.add(pinned_conversation) + session.commit() @classmethod - def unpin(cls, app_model: App, conversation_id: str, user: Account | EndUser | None): + def unpin(cls, app_model: App, conversation_id: str, user: Account | EndUser | None, session: Session): if not user: return - pinned_conversation = db.session.scalar( + pinned_conversation = session.scalar( select(PinnedConversation) .where( PinnedConversation.app_id == app_model.id, @@ -108,5 +107,5 @@ class WebConversationService: if not pinned_conversation: return - db.session.delete(pinned_conversation) - db.session.commit() + session.delete(pinned_conversation) + session.commit() diff --git a/api/services/webapp_auth_service.py b/api/services/webapp_auth_service.py index 6ecc8eb8bc9..33267c53d5c 100644 --- a/api/services/webapp_auth_service.py +++ b/api/services/webapp_auth_service.py @@ -4,10 +4,10 @@ from datetime import UTC, datetime, timedelta from typing import Any from sqlalchemy import select +from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound, Unauthorized from configs import dify_config -from extensions.ext_database import db from libs.helper import TokenManager from libs.passport import PassportService from libs.password import compare_password @@ -33,9 +33,9 @@ class WebAppAuthService: """Service for web app authentication.""" @staticmethod - def authenticate(email: str, password: str) -> Account: + def authenticate(email: str, password: str, session: Session) -> Account: """authenticate account with email and password""" - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + account = AccountService.get_account_by_email_with_case_fallback(email, session=session) if not account: raise AccountNotFoundError() @@ -54,8 +54,8 @@ class WebAppAuthService: return access_token @classmethod - def get_user_through_email(cls, email: str): - account = AccountService.get_account_by_email_with_case_fallback(db.session, email) + def get_user_through_email(cls, email: str, session: Session): + account = AccountService.get_account_by_email_with_case_fallback(email, session=session) if not account: return None @@ -93,11 +93,11 @@ class WebAppAuthService: TokenManager.revoke_token(token, "email_code_login") @classmethod - def create_end_user(cls, app_code, email) -> EndUser: - site = db.session.scalar(select(Site).where(Site.code == app_code).limit(1)) + def create_end_user(cls, app_code, email, session: Session) -> EndUser: + site = session.scalar(select(Site).where(Site.code == app_code).limit(1)) if not site: raise NotFound("Site not found.") - app_model = db.session.get(App, site.app_id) + app_model = session.get(App, site.app_id) if not app_model: raise NotFound("App not found.") end_user = EndUser( @@ -109,8 +109,8 @@ class WebAppAuthService: name="enterpriseuser", external_user_id="enterpriseuser", ) - db.session.add(end_user) - db.session.commit() + session.add(end_user) + session.commit() return end_user @@ -133,7 +133,7 @@ class WebAppAuthService: @classmethod def is_app_require_permission_check( - cls, app_code: str | None = None, app_id: str | None = None, access_mode: str | None = None + cls, app_code: str | None = None, app_id: str | None = None, access_mode: str | None = None, *, session: Session ) -> bool: """ Check if the app requires permission check based on its access mode. @@ -145,7 +145,7 @@ class WebAppAuthService: raise ValueError("Either app_code or app_id must be provided.") if app_code: - app_id = AppService.get_app_id_by_code(app_code) + app_id = AppService.get_app_id_by_code(app_code, session=session) if not app_id: raise ValueError("App ID could not be determined from the provided app_code.") @@ -155,7 +155,9 @@ class WebAppAuthService: return False @classmethod - def get_app_auth_type(cls, app_code: str | None = None, access_mode: str | None = None) -> WebAppAuthType: + def get_app_auth_type( + cls, app_code: str | None = None, access_mode: str | None = None, *, session: Session + ) -> WebAppAuthType: """ Get the authentication type for the app based on its access mode. """ @@ -171,8 +173,8 @@ class WebAppAuthService: return WebAppAuthType.EXTERNAL if app_code: - app_id = AppService.get_app_id_by_code(app_code) + app_id = AppService.get_app_id_by_code(app_code, session=session) webapp_settings = EnterpriseService.WebAppAuth.get_app_access_mode_by_id(app_id=app_id) - return cls.get_app_auth_type(access_mode=webapp_settings.access_mode) + return cls.get_app_auth_type(access_mode=webapp_settings.access_mode, session=session) raise ValueError("Could not determine app authentication type.") diff --git a/api/services/workflow/node_output_inspector_service.py b/api/services/workflow/node_output_inspector_service.py index 66dcfec591f..5d6a8f1c675 100644 --- a/api/services/workflow/node_output_inspector_service.py +++ b/api/services/workflow/node_output_inspector_service.py @@ -52,9 +52,9 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import select +from sqlalchemy.orm import Session from core.app.file_access import DatabaseFileAccessController -from core.db.session_factory import session_factory from core.workflow.nodes.agent_v2.binding_resolver import ( WorkflowAgentBindingError, WorkflowAgentBindingResolver, @@ -410,8 +410,8 @@ class NodeOutputInspectorService: The service is dependency-light: it holds a single :class:`WorkflowAgentBindingResolver` so agent v2 nodes can map to their declared outputs without re-implementing binding lookup. All other I/O - uses the global session factory so workflow runs / executions stay on the - repo-default code path. + receives an explicit SQLAlchemy session from its caller so transaction + ownership stays at the controller/task boundary. Tenancy is enforced via ``app_model.tenant_id`` + ``app_model.id`` on every load — the same scope guard regardless of trigger source. @@ -422,9 +422,13 @@ class NodeOutputInspectorService: # ── public API ──────────────────────────────────────────────────────── - def snapshot_workflow_run(self, *, app_model: App, workflow_run_id: str) -> WorkflowRunSnapshotView: + def snapshot_workflow_run( + self, *, app_model: App, workflow_run_id: str, session: Session + ) -> WorkflowRunSnapshotView: """Build the per-node snapshot for one debug workflow run.""" - workflow_run, executions = self._load_run_and_executions(app_model=app_model, workflow_run_id=workflow_run_id) + workflow_run, executions = self._load_run_and_executions( + app_model=app_model, workflow_run_id=workflow_run_id, session=session + ) executions_by_node = self._index_executions_by_node(executions) graph_nodes = _graph_nodes(workflow_run) @@ -447,9 +451,11 @@ class NodeOutputInspectorService: node_outputs=node_views, ) - def node_detail(self, *, app_model: App, workflow_run_id: str, node_id: str) -> NodeOutputsView: + def node_detail(self, *, app_model: App, workflow_run_id: str, node_id: str, session: Session) -> NodeOutputsView: """Per-node Inspector entry — returns one ``NodeOutputsView``.""" - workflow_run, executions = self._load_run_and_executions(app_model=app_model, workflow_run_id=workflow_run_id) + workflow_run, executions = self._load_run_and_executions( + app_model=app_model, workflow_run_id=workflow_run_id, session=session + ) graph_nodes = _graph_nodes(workflow_run) raw_node = next((n for n in graph_nodes if str(n.get("id")) == node_id), None) if raw_node is None: @@ -474,9 +480,12 @@ class NodeOutputInspectorService: workflow_run_id: str, node_id: str, output_name: str, + session: Session, ) -> OutputPreviewView: """Full payload for one declared output (with signed file URL).""" - workflow_run, executions = self._load_run_and_executions(app_model=app_model, workflow_run_id=workflow_run_id) + workflow_run, executions = self._load_run_and_executions( + app_model=app_model, workflow_run_id=workflow_run_id, session=session + ) graph_nodes = _graph_nodes(workflow_run) raw_node = next((n for n in graph_nodes if str(n.get("id")) == node_id), None) if raw_node is None: @@ -536,7 +545,7 @@ class NodeOutputInspectorService: # ── DB loading ──────────────────────────────────────────────────────── def _load_run_and_executions( - self, *, app_model: App, workflow_run_id: str + self, *, app_model: App, workflow_run_id: str, session: Session ) -> tuple[WorkflowRun, Sequence[WorkflowNodeExecutionModel]]: """Fetch the ``WorkflowRun`` row + every execution that belongs to it. @@ -548,24 +557,23 @@ class NodeOutputInspectorService: deliberately not checked here — D-1 was lifted 2026-05-26 and the Inspector now serves both draft and published runs. """ - with session_factory.create_session() as session: - workflow_run = session.scalar( - select(WorkflowRun).where( - WorkflowRun.id == workflow_run_id, - WorkflowRun.app_id == app_model.id, - WorkflowRun.tenant_id == app_model.tenant_id, - ) + workflow_run = session.scalar( + select(WorkflowRun).where( + WorkflowRun.id == workflow_run_id, + WorkflowRun.app_id == app_model.id, + WorkflowRun.tenant_id == app_model.tenant_id, ) - if workflow_run is None: - raise NodeOutputInspectorError("workflow_run_not_found", "Workflow run not found.") + ) + if workflow_run is None: + raise NodeOutputInspectorError("workflow_run_not_found", "Workflow run not found.") - executions = session.scalars( - select(WorkflowNodeExecutionModel).where( - WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, - WorkflowNodeExecutionModel.tenant_id == app_model.tenant_id, - WorkflowNodeExecutionModel.app_id == app_model.id, - ) - ).all() + executions = session.scalars( + select(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id, + WorkflowNodeExecutionModel.tenant_id == app_model.tenant_id, + WorkflowNodeExecutionModel.app_id == app_model.id, + ) + ).all() return workflow_run, executions diff --git a/api/services/workflow/workflow_converter.py b/api/services/workflow/workflow_converter.py index e279f1daaa3..5f787bb51cd 100644 --- a/api/services/workflow/workflow_converter.py +++ b/api/services/workflow/workflow_converter.py @@ -2,6 +2,7 @@ import json from typing import Any, TypedDict from sqlalchemy import select +from sqlalchemy.orm import Session from core.app.app_config.entities import ( DatasetEntity, @@ -18,7 +19,6 @@ from core.helper import encrypter from core.prompt.simple_prompt_transform import SimplePromptTransform from core.prompt.utils.prompt_template_parser import PromptTemplateParser from events.app_event import app_was_created -from extensions.ext_database import db from graphon.file import FileUploadConfig from graphon.model_runtime.entities.llm_entities import LLMMode from graphon.model_runtime.utils.encoders import jsonable_encoder @@ -53,7 +53,14 @@ class WorkflowConverter: """ def convert_to_workflow( - self, app_model: App, account: Account, name: str, icon_type: str, icon: str, icon_background: str + self, + app_model: App, + account: Account, + name: str, + icon_type: str, + icon: str, + icon_background: str, + session: Session, ): """ Convert app to workflow @@ -77,7 +84,7 @@ class WorkflowConverter: raise ValueError("App model config is required") workflow = self.convert_app_model_config_to_workflow( - app_model=app_model, app_model_config=app_model.app_model_config, account_id=account.id + app_model=app_model, app_model_config=app_model.app_model_config, account_id=account.id, session=session ) # create new app @@ -97,17 +104,19 @@ class WorkflowConverter: new_app.created_by = account.id new_app.maintainer = account.id new_app.updated_by = account.id - db.session.add(new_app) - db.session.flush() + session.add(new_app) + session.flush() workflow.app_id = new_app.id - db.session.commit() + session.commit() app_was_created.send(new_app, account=account) return new_app - def convert_app_model_config_to_workflow(self, app_model: App, app_model_config: AppModelConfig, account_id: str): + def convert_app_model_config_to_workflow( + self, app_model: App, app_model_config: AppModelConfig, account_id: str, session: Session + ): """ Convert app model config to workflow mode :param app_model: App instance @@ -144,6 +153,7 @@ class WorkflowConverter: app_model=app_model, variables=app_config.variables, external_data_variables=app_config.external_data_variables, + session=session, ) for http_request_node in http_request_nodes: @@ -217,8 +227,8 @@ class WorkflowConverter: conversation_variables=[], ) - db.session.add(workflow) - db.session.commit() + session.add(workflow) + session.commit() return workflow @@ -262,7 +272,11 @@ class WorkflowConverter: } def _convert_to_http_request_node( - self, app_model: App, variables: list[VariableEntity], external_data_variables: list[ExternalDataVariableEntity] + self, + app_model: App, + variables: list[VariableEntity], + external_data_variables: list[ExternalDataVariableEntity], + session: Session, ) -> tuple[list[_NodeType], dict[str, str]]: """ Convert API Based Extension to HTTP Request Node @@ -290,7 +304,7 @@ class WorkflowConverter: # get api_based_extension api_based_extension = self._get_api_based_extension( - tenant_id=tenant_id, api_based_extension_id=api_based_extension_id + tenant_id=tenant_id, api_based_extension_id=api_based_extension_id, session=session ) # decrypt api_key @@ -650,14 +664,14 @@ class WorkflowConverter: else: return AppMode.ADVANCED_CHAT - def _get_api_based_extension(self, tenant_id: str, api_based_extension_id: str): + def _get_api_based_extension(self, tenant_id: str, api_based_extension_id: str, session: Session): """ Get API Based Extension :param tenant_id: tenant id :param api_based_extension_id: api based extension id :return: """ - api_based_extension = db.session.scalar( + api_based_extension = session.scalar( select(APIBasedExtension) .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) .limit(1) diff --git a/api/services/workflow_collaboration_service.py b/api/services/workflow_collaboration_service.py index bec61ce666d..5c635d7d66a 100644 --- a/api/services/workflow_collaboration_service.py +++ b/api/services/workflow_collaboration_service.py @@ -9,8 +9,8 @@ from collections.abc import Mapping from typing import Any, override from sqlalchemy import select +from sqlalchemy.orm import Session -from core.db.session_factory import session_factory from models.account import Account from models.model import App from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo @@ -94,20 +94,22 @@ class WorkflowCollaborationService: }, ) - def authorize_and_join_workflow_room(self, workflow_id: str, sid: str) -> tuple[str, bool] | None: + def authorize_and_join_workflow_room( + self, workflow_id: str, sid: str, *, session: Session + ) -> tuple[str, bool] | None: """ Join a collaboration room only after validating the socket session and tenant-scoped app access. The Socket.IO payload still calls the room key `workflow_id`, but the identifier is the workflow app's `App.id`. Returning `None` lets the controller reject the join before any Redis or room state is created. """ - session = self._socketio.get_session(sid) - user_id = session.get("user_id") - tenant_id = session.get("tenant_id") + socket_session = self._socketio.get_session(sid) + user_id = socket_session.get("user_id") + tenant_id = socket_session.get("tenant_id") if not user_id or not tenant_id: return None - if not self._can_access_workflow(workflow_id, str(tenant_id)): + if not self._can_access_workflow(workflow_id, str(tenant_id), session=session): logger.warning( "Workflow collaboration join rejected: workflow_id=%s tenant_id=%s user_id=%s sid=%s", workflow_id, @@ -121,8 +123,8 @@ class WorkflowCollaborationService: session_info: WorkflowSessionInfo = { "user_id": str(user_id), - "username": str(session.get("username", "Unknown")), - "avatar": session.get("avatar"), + "username": str(socket_session.get("username", "Unknown")), + "avatar": socket_session.get("avatar"), "sid": sid, "connected_at": int(time.time()), "server_id": self.server_id, @@ -140,10 +142,9 @@ class WorkflowCollaborationService: return str(user_id), is_leader - def _can_access_workflow(self, workflow_id: str, tenant_id: str) -> bool: + def _can_access_workflow(self, workflow_id: str, tenant_id: str, *, session: Session) -> bool: """Check room access without relying on Flask's app-context-bound scoped session.""" - with session_factory.create_session() as session: - app_id = session.scalar(select(App.id).where(App.id == workflow_id, App.tenant_id == tenant_id).limit(1)) + app_id = session.scalar(select(App.id).where(App.id == workflow_id, App.tenant_id == tenant_id).limit(1)) return app_id is not None def disconnect_session(self, sid: str) -> None: diff --git a/api/services/workflow_run_service.py b/api/services/workflow_run_service.py index 2499e6cc094..1ff0b59514f 100644 --- a/api/services/workflow_run_service.py +++ b/api/services/workflow_run_service.py @@ -2,7 +2,7 @@ import threading from collections.abc import Sequence from typing import TypedDict -from sqlalchemy import Engine +from sqlalchemy import Engine, select from sqlalchemy.orm import sessionmaker import contexts @@ -12,6 +12,7 @@ from models import ( Account, App, EndUser, + Message, WorkflowNodeExecutionModel, WorkflowRun, WorkflowRunTriggeredFrom, @@ -72,9 +73,29 @@ class WorkflowRunService: pagination = self.get_paginate_workflow_runs(app_model, args, triggered_from) + # Batch-load the associated Message for every run in a single query to avoid + # an N+1 pattern: the deprecated WorkflowRun.message property issues one query + # per run. The filter matches that property exactly (app_id + workflow_run_id). + workflow_runs = pagination.data + run_ids = [workflow_run.id for workflow_run in workflow_runs] + messages_by_run_id: dict[str, Message] = {} + if run_ids: + messages = db.session.scalars( + select(Message).where( + Message.app_id == app_model.id, + Message.workflow_run_id.in_(run_ids), + ) + ).all() + for loaded_message in messages: + run_id = loaded_message.workflow_run_id + if run_id is None: + continue + # setdefault mirrors scalar()'s single-row-per-run semantics. + messages_by_run_id.setdefault(run_id, loaded_message) + with_message_workflow_runs = [] - for workflow_run in pagination.data: - message = workflow_run.message + for workflow_run in workflow_runs: + message = messages_by_run_id.get(workflow_run.id) with_message_workflow_run = WorkflowWithMessage(workflow_run=workflow_run) if message: with_message_workflow_run.message_id = message.id diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 048b25c6bf9..95be0f7017a 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from typing import Any, cast from sqlalchemy import exists, select -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager @@ -178,7 +178,7 @@ class WorkflowService: node_id=node_id, ) - def is_workflow_exist(self, app_model: App) -> bool: + def is_workflow_exist(self, app_model: App, *, session: Session) -> bool: stmt = select( exists().where( Workflow.tenant_id == app_model.tenant_id, @@ -186,23 +186,21 @@ class WorkflowService: Workflow.version == Workflow.VERSION_DRAFT, ) ) - return db.session.execute(stmt).scalar_one() + return session.execute(stmt).scalar_one() def get_draft_workflow( - self, app_model: App, workflow_id: str | None = None, session: Session | scoped_session | None = None + self, app_model: App, workflow_id: str | None = None, *, session: Session ) -> Workflow | None: """ Get draft workflow - When ``session`` is provided, reuse it so callers that already hold a - Session avoid checking out an extra request-scoped ``db.session`` - connection. Falls back to ``db.session`` for backward compatibility. + Reuses the caller's active session so workflow reads stay in the same + transaction as the surrounding request or task. """ if workflow_id: return self.get_published_workflow_by_id(app_model, workflow_id, session=session) # fetch draft workflow by app_model - bind = session if session is not None else db.session - workflow = bind.scalar( + workflow = session.scalar( select(Workflow) .where( Workflow.tenant_id == app_model.tenant_id, @@ -215,18 +213,14 @@ class WorkflowService: # return draft workflow return workflow - def get_published_workflow_by_id( - self, app_model: App, workflow_id: str, session: Session | scoped_session | None = None - ) -> Workflow | None: + def get_published_workflow_by_id(self, app_model: App, workflow_id: str, *, session: Session) -> Workflow | None: """ fetch published workflow by workflow_id - When ``session`` is provided, reuse it so callers that already hold a - Session avoid checking out an extra request-scoped ``db.session`` - connection. Falls back to ``db.session`` for backward compatibility. + Reuses the caller's active session so workflow reads stay in the same + transaction as the surrounding request or task. """ - bind = session if session is not None else db.session - workflow = bind.scalar( + workflow = session.scalar( select(Workflow) .where( Workflow.tenant_id == app_model.tenant_id, @@ -244,20 +238,18 @@ class WorkflowService: ) return workflow - def get_published_workflow(self, app_model: App, session: Session | None = None) -> Workflow | None: + def get_published_workflow(self, app_model: App, *, session: Session) -> Workflow | None: """ Get published workflow - When ``session`` is provided, reuse it so callers that already hold a - Session avoid checking out an extra request-scoped ``db.session`` - connection. Falls back to ``db.session`` for backward compatibility. + Reuses the caller's active session so workflow reads stay in the same + transaction as the surrounding request or task. """ if not app_model.workflow_id: return None - bind = session if session is not None else db.session - workflow = bind.scalar( + workflow = session.scalar( select(Workflow) .where( Workflow.tenant_id == app_model.tenant_id, @@ -269,7 +261,7 @@ class WorkflowService: return workflow - def get_accessible_app_ids(self, app_ids: Sequence[str], tenant_id: str) -> set[str]: + def get_accessible_app_ids(self, app_ids: Sequence[str], tenant_id: str, *, session: Session) -> set[str]: """ Return app IDs that belong to the given tenant. """ @@ -277,7 +269,7 @@ class WorkflowService: return set() stmt = select(App.id).where(App.id.in_(app_ids), App.tenant_id == tenant_id) - return {str(app_id) for app_id in db.session.scalars(stmt).all()} + return {str(app_id) for app_id in session.scalars(stmt).all()} def get_all_published_workflow( self, @@ -327,13 +319,14 @@ class WorkflowService: account: Account, environment_variables: Sequence[VariableBase], conversation_variables: Sequence[VariableBase], + session: Session, ) -> Workflow: """ Sync draft workflow :raises WorkflowHashNotEqualError """ # fetch draft workflow by app_model - workflow = self.get_draft_workflow(app_model=app_model) + workflow = self.get_draft_workflow(app_model=app_model, session=session) if workflow and workflow.unique_hash != unique_hash: raise WorkflowHashNotEqualError() @@ -357,7 +350,7 @@ class WorkflowService: environment_variables=environment_variables, conversation_variables=conversation_variables, ) - db.session.add(workflow) + session.add(workflow) # update draft workflow if found else: workflow.graph = json.dumps(graph) @@ -369,19 +362,19 @@ class WorkflowService: from services.agent.workflow_publish_service import WorkflowAgentPublishService - db.session.flush() + session.flush() WorkflowAgentPublishService.sync_agent_bindings_for_draft( - session=cast(Session, db.session), + session=session, draft_workflow=workflow, account_id=account.id, ) WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync( - session=cast(Session, db.session), + session=session, draft_workflow=workflow, ) # commit db session changes - db.session.commit() + session.commit() # trigger app workflow events app_draft_workflow_was_synced.send(app_model, synced_draft_workflow=workflow) @@ -395,12 +388,13 @@ class WorkflowService: app_model: App, environment_variables: Sequence[VariableBase], account: Account, + session: Session, ): """ Update draft workflow environment variables """ # fetch draft workflow by app_model - workflow = self.get_draft_workflow(app_model=app_model) + workflow = self.get_draft_workflow(app_model=app_model, session=session) if not workflow: raise ValueError("No draft workflow found.") @@ -410,7 +404,7 @@ class WorkflowService: workflow.updated_at = naive_utc_now() # commit db session changes - db.session.commit() + session.commit() def update_draft_workflow_conversation_variables( self, @@ -418,12 +412,13 @@ class WorkflowService: app_model: App, conversation_variables: Sequence[VariableBase], account: Account, + session: Session, ): """ Update draft workflow conversation variables """ # fetch draft workflow by app_model - workflow = self.get_draft_workflow(app_model=app_model) + workflow = self.get_draft_workflow(app_model=app_model, session=session) if not workflow: raise ValueError("No draft workflow found.") @@ -433,7 +428,7 @@ class WorkflowService: workflow.updated_at = naive_utc_now() # commit db session changes - db.session.commit() + session.commit() def update_draft_workflow_features( self, @@ -441,12 +436,13 @@ class WorkflowService: app_model: App, features: dict, account: Account, + session: Session, ): """ Update draft workflow features """ # fetch draft workflow by app_model - workflow = self.get_draft_workflow(app_model=app_model) + workflow = self.get_draft_workflow(app_model=app_model, session=session) if not workflow: raise ValueError("No draft workflow found.") @@ -459,7 +455,7 @@ class WorkflowService: workflow.updated_at = naive_utc_now() # commit db session changes - db.session.commit() + session.commit() def restore_published_workflow_to_draft( self, @@ -467,20 +463,23 @@ class WorkflowService: app_model: App, workflow_id: str, account: Account, + session: Session, ) -> Workflow: """Restore a published workflow snapshot into the draft workflow. Secret environment variables are copied server-side from the selected published workflow so the normal draft sync flow stays stateless. """ - source_workflow = self.get_published_workflow_by_id(app_model=app_model, workflow_id=workflow_id) + source_workflow = self.get_published_workflow_by_id( + app_model=app_model, workflow_id=workflow_id, session=session + ) if not source_workflow: raise WorkflowNotFoundError("Workflow not found.") self.validate_features_structure(app_model=app_model, features=source_workflow.normalized_features_dict) self.validate_graph_structure(graph=source_workflow.graph_dict) - draft_workflow = self.get_draft_workflow(app_model=app_model) + draft_workflow = self.get_draft_workflow(app_model=app_model, session=session) draft_workflow, is_new_draft = apply_published_workflow_snapshot_to_draft( tenant_id=app_model.tenant_id, app_id=app_model.id, @@ -491,9 +490,9 @@ class WorkflowService: ) if is_new_draft: - db.session.add(draft_workflow) + session.add(draft_workflow) - db.session.commit() + session.commit() app_draft_workflow_was_synced.send(app_model, synced_draft_workflow=draft_workflow) return draft_workflow @@ -520,7 +519,7 @@ class WorkflowService: from services.feature_service import FeatureService if FeatureService.get_system_features().plugin_manager.enabled: - self._validate_workflow_credentials(draft_workflow) + self._validate_workflow_credentials(draft_workflow, session=session) # validate graph structure self.validate_graph_structure(graph=draft_workflow.graph_dict) @@ -577,7 +576,7 @@ class WorkflowService: # return new workflow return workflow - def _validate_workflow_credentials(self, workflow: Workflow) -> None: + def _validate_workflow_credentials(self, workflow: Workflow, *, session: Session) -> None: """ Validate all credentials in workflow nodes before publishing. @@ -609,7 +608,7 @@ class WorkflowService: ) else: # Check default workspace credential for this provider - self._check_default_tool_credential(workflow.tenant_id, provider) + self._check_default_tool_credential(workflow.tenant_id, provider, session=session) elif node_type == "agent": agent_params = node_data.get("agent_parameters", {}) @@ -622,7 +621,9 @@ class WorkflowService: # Validate load balancing credentials for agent model if load balancing is enabled agent_model_node_data = {"model": model_config} - self._validate_load_balancing_credentials(workflow, agent_model_node_data, node_id) + self._validate_load_balancing_credentials( + workflow, agent_model_node_data, node_id, session=session + ) # Validate agent tools tools = agent_params.get("tools", {}).get("value", []) @@ -636,7 +637,7 @@ class WorkflowService: check_credential_policy_compliance(credential_id, provider, PluginCredentialType.TOOL) else: - self._check_default_tool_credential(workflow.tenant_id, provider) + self._check_default_tool_credential(workflow.tenant_id, provider, session=session) elif node_type in ["llm", "knowledge_retrieval", "parameter_extractor", "question_classifier"]: model_config = node_data.get("model", {}) @@ -647,7 +648,7 @@ class WorkflowService: # Validate that the provider+model combination can fetch valid credentials self._validate_llm_model_config(workflow.tenant_id, provider, model_name) # Validate load balancing credentials if load balancing is enabled - self._validate_load_balancing_credentials(workflow, node_data, node_id) + self._validate_load_balancing_credentials(workflow, node_data, node_id, session=session) else: raise ValueError(f"Node {node_id} ({node_type}): Missing provider or model configuration") @@ -710,7 +711,7 @@ class WorkflowService: f"Failed to validate LLM model configuration (provider: {provider}, model: {model_name}): {str(e)}" ) - def _check_default_tool_credential(self, tenant_id: str, provider: str) -> None: + def _check_default_tool_credential(self, tenant_id: str, provider: str, *, session: Session) -> None: """ Check credential policy compliance for the default workspace credential of a tool provider. @@ -726,7 +727,7 @@ class WorkflowService: # Use the same fallback logic as runtime: get the first available credential # ordered by is_default DESC, created_at ASC (same as tool_manager.py) - default_provider = db.session.scalar( + default_provider = session.scalar( select(BuiltinToolProvider) .where( BuiltinToolProvider.tenant_id == tenant_id, @@ -753,7 +754,9 @@ class WorkflowService: except Exception as e: raise ValueError(f"Failed to validate default credential for tool provider {provider}: {str(e)}") - def _validate_load_balancing_credentials(self, workflow: Workflow, node_data: dict[str, Any], node_id: str) -> None: + def _validate_load_balancing_credentials( + self, workflow: Workflow, node_data: dict[str, Any], node_id: str, *, session: Session + ) -> None: """ Validate load balancing credentials for a workflow node. @@ -773,7 +776,9 @@ class WorkflowService: # Check if this model has load balancing enabled if self._is_load_balancing_enabled(workflow.tenant_id, provider, model_name): # Get all load balancing configurations for this model - load_balancing_configs = self._get_load_balancing_configs(workflow.tenant_id, provider, model_name) + load_balancing_configs = self._get_load_balancing_configs( + workflow.tenant_id, provider, model_name, session=session + ) # Validate each load balancing configuration try: for config in load_balancing_configs: @@ -817,7 +822,9 @@ class WorkflowService: # If we can't determine the status, assume load balancing is not enabled return False - def _get_load_balancing_configs(self, tenant_id: str, provider: str, model_name: str) -> list[dict[str, Any]]: + def _get_load_balancing_configs( + self, tenant_id: str, provider: str, model_name: str, *, session: Session + ) -> list[dict[str, Any]]: """ Get all load balancing configurations for a model. @@ -835,11 +842,17 @@ class WorkflowService: provider=provider, model=model_name, model_type="llm", # Load balancing is primarily used for LLM models + session=session, config_from="predefined-model", # Check both predefined and custom models ) _, custom_configs = model_load_balancing_service.get_load_balancing_configs( - tenant_id=tenant_id, provider=provider, model=model_name, model_type="llm", config_from="custom-model" + tenant_id=tenant_id, + provider=provider, + model=model_name, + model_type="llm", + session=session, + config_from="custom-model", ) all_configs = cast(list[dict[str, Any]], configs) + cast(list[dict[str, Any]], custom_configs) @@ -1047,6 +1060,7 @@ class WorkflowService: account: Account, node_id: str, inputs: Mapping[str, Any] | None = None, + session: Session, ) -> Mapping[str, Any]: """ Build a human input form preview for a draft workflow. @@ -1057,7 +1071,7 @@ class WorkflowService: node_id: Human input node ID. inputs: Values used to fill missing upstream variables referenced in form_content. """ - draft_workflow = self.get_draft_workflow(app_model=app_model) + draft_workflow = self.get_draft_workflow(app_model=app_model, session=session) if not draft_workflow: raise ValueError("Workflow not initialized") @@ -1104,6 +1118,7 @@ class WorkflowService: form_inputs: Mapping[str, Any], inputs: Mapping[str, Any] | None = None, action: str, + session: Session, ) -> Mapping[str, Any]: """ Submit a human input form preview for a draft workflow. @@ -1116,7 +1131,7 @@ class WorkflowService: inputs: Values used to fill missing upstream variables referenced in form_content. action: Selected action ID. """ - draft_workflow = self.get_draft_workflow(app_model=app_model) + draft_workflow = self.get_draft_workflow(app_model=app_model, session=session) if not draft_workflow: raise ValueError("Workflow not initialized") @@ -1189,8 +1204,9 @@ class WorkflowService: node_id: str, delivery_method_id: str, inputs: Mapping[str, Any] | None = None, + session: Session, ) -> None: - draft_workflow = self.get_draft_workflow(app_model=app_model) + draft_workflow = self.get_draft_workflow(app_model=app_model, session=session) if not draft_workflow: raise ValueError("Workflow not initialized") @@ -1529,7 +1545,7 @@ class WorkflowService: node_execution.status = WorkflowNodeExecutionStatus.FAILED node_execution.error = error - def convert_to_workflow(self, app_model: App, account: Account, args: dict[str, Any]) -> App: + def convert_to_workflow(self, app_model: App, account: Account, args: dict[str, Any], *, session: Session) -> App: """ Basic mode of chatbot app(expert mode) to workflow Completion App to Workflow App @@ -1553,6 +1569,7 @@ class WorkflowService: icon_type=args.get("icon_type", "emoji"), icon=args.get("icon", "🤖"), icon_background=args.get("icon_background", "#FFEAD5"), + session=session, ) return new_app diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index 180c077b88a..30853b2cc99 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -1,9 +1,9 @@ from flask_login import current_user from sqlalchemy import select +from sqlalchemy.orm import Session from configs import dify_config from enums.cloud_plan import CloudPlan -from extensions.ext_database import db from models.account import Tenant, TenantAccountJoin, TenantAccountRole from services.account_service import TenantService from services.feature_service import FeatureService @@ -11,7 +11,7 @@ from services.feature_service import FeatureService class WorkspaceService: @classmethod - def get_tenant_info(cls, tenant: Tenant): + def get_tenant_info(cls, tenant: Tenant, session: Session): if not tenant: return None tenant_info: dict[str, object] = { @@ -25,7 +25,7 @@ class WorkspaceService: } # Get role of user - tenant_account_join = db.session.scalar( + tenant_account_join = session.scalar( select(TenantAccountJoin) .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == current_user.id) .limit(1) @@ -37,7 +37,7 @@ class WorkspaceService: can_replace_logo = feature.can_replace_logo if can_replace_logo and TenantService.has_roles( - tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN], session=db.session + tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN], session=session ): base_url = dify_config.FILES_URL replace_webapp_logo = ( @@ -56,7 +56,7 @@ class WorkspaceService: from services.credit_pool_service import CreditPoolService - paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid") + paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid", session=session) # if the tenant is not on the sandbox plan and the paid pool is not full, use the paid pool if ( feature.billing.subscription.plan != CloudPlan.SANDBOX @@ -66,7 +66,7 @@ class WorkspaceService: tenant_info["trial_credits"] = paid_pool.quota_limit tenant_info["trial_credits_used"] = paid_pool.quota_used else: - trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial") + trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial", session=session) if trial_pool: tenant_info["trial_credits"] = trial_pool.quota_limit tenant_info["trial_credits_used"] = trial_pool.quota_used diff --git a/api/tasks/agent_backend_session_cleanup_task.py b/api/tasks/agent_backend_session_cleanup_task.py new file mode 100644 index 00000000000..2c799bda97c --- /dev/null +++ b/api/tasks/agent_backend_session_cleanup_task.py @@ -0,0 +1,68 @@ +"""Celery tasks that execute Agent backend lifecycle-only session cleanup.""" + +from __future__ import annotations + +import logging + +from celery import shared_task + +from clients.agent_backend.factory import create_agent_backend_run_client +from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder +from clients.agent_backend.session_cleanup import ( + AgentBackendSessionCleanupPayload, + cleanup_agent_backend_session, +) +from configs import dify_config + +logger = logging.getLogger(__name__) + + +def _create_agent_backend_client(): + if not (dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL): + return None + return create_agent_backend_run_client( + base_url=dify_config.AGENT_BACKEND_BASE_URL, + use_fake=dify_config.AGENT_BACKEND_USE_FAKE, + fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO, + ) + + +def _run_cleanup_task(payload_dict: dict[str, object]) -> None: + payload = AgentBackendSessionCleanupPayload.model_validate(payload_dict) + result = cleanup_agent_backend_session( + payload=payload, + client=_create_agent_backend_client(), + request_builder=AgentBackendRunRequestBuilder(), + ) + if result.status == "succeeded": + return + + log_fields = { + "tenant_id": payload.metadata.get("tenant_id"), + "app_id": payload.metadata.get("app_id"), + "workflow_run_id": payload.metadata.get("workflow_run_id"), + "node_id": payload.metadata.get("node_id"), + "conversation_id": payload.metadata.get("conversation_id"), + "agent_id": payload.metadata.get("agent_id"), + "previous_agent_backend_run_id": payload.metadata.get("previous_agent_backend_run_id"), + "failed_agent_backend_run_id": payload.metadata.get("failed_agent_backend_run_id"), + "cleanup_run_id": result.cleanup_run_id, + "reason": result.reason, + } + if result.status == "skipped": + logger.info("Agent backend session cleanup skipped: %s", log_fields) + return + + logger.warning("Agent backend session cleanup failed: %s", log_fields) + + +@shared_task(queue="workflow_storage") +def cleanup_workflow_agent_runtime_session(payload_dict: dict[str, object]) -> None: + """Run one workflow-owned Agent backend cleanup payload.""" + _run_cleanup_task(payload_dict) + + +@shared_task(queue="conversation") +def cleanup_conversation_agent_runtime_session(payload_dict: dict[str, object]) -> None: + """Run one conversation-owned Agent backend cleanup payload.""" + _run_cleanup_task(payload_dict) diff --git a/api/tasks/app_generate/workflow_execute_task.py b/api/tasks/app_generate/workflow_execute_task.py index a383839bd05..36bd21e16c1 100644 --- a/api/tasks/app_generate/workflow_execute_task.py +++ b/api/tasks/app_generate/workflow_execute_task.py @@ -25,6 +25,7 @@ from core.repositories import DifyCoreRepositoryFactory from extensions.ext_database import db from graphon.entities import WorkflowStartReason from graphon.enums import WorkflowExecutionStatus +from graphon.filters import ResponseStreamFilter from graphon.runtime import GraphRuntimeState from libs.datetime_utils import naive_utc_now from libs.flask_utils import set_login_user @@ -486,6 +487,7 @@ def _resume_app_execution(payload: dict[str, Any]) -> None: generate_entity = resumption_context.get_generate_entity() graph_runtime_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state) + response_stream_filter = resumption_context.get_response_stream_filter() conversation = None message = None @@ -562,6 +564,7 @@ def _resume_app_execution(payload: dict[str, Any]) -> None: message=message, generate_entity=generate_entity, graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, session_factory=session_factory, pause_state_config=pause_config, workflow_run_id=workflow_run_id, @@ -574,6 +577,7 @@ def _resume_app_execution(payload: dict[str, Any]) -> None: user=user, generate_entity=generate_entity, graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, session_factory=session_factory, pause_state_config=pause_config, workflow_run_id=workflow_run_id, @@ -592,6 +596,7 @@ def _resume_advanced_chat( message: Message, generate_entity: AdvancedChatAppGenerateEntity, graph_runtime_state: GraphRuntimeState, + response_stream_filter: ResponseStreamFilter, session_factory: sessionmaker, pause_state_config: PauseStateLayerConfig, workflow_run_id: str, @@ -631,6 +636,7 @@ def _resume_advanced_chat( workflow_node_execution_repository=workflow_node_execution_repository, graph_runtime_state=graph_runtime_state, pause_state_config=pause_state_config, + response_stream_filter=response_stream_filter, ) except Exception: logger.exception("Failed to resume chatflow execution for workflow run %s", workflow_run_id) @@ -654,6 +660,7 @@ def _resume_workflow( user: Account | EndUser, generate_entity: WorkflowAppGenerateEntity, graph_runtime_state: GraphRuntimeState, + response_stream_filter: ResponseStreamFilter, session_factory: sessionmaker, pause_state_config: PauseStateLayerConfig, workflow_run_id: str, @@ -693,6 +700,7 @@ def _resume_workflow( workflow_execution_repository=workflow_execution_repository, workflow_node_execution_repository=workflow_node_execution_repository, pause_state_config=pause_state_config, + response_stream_filter=response_stream_filter, ) except Exception: logger.exception("Failed to resume workflow execution for workflow run %s", workflow_run_id) diff --git a/api/tasks/async_workflow_tasks.py b/api/tasks/async_workflow_tasks.py index 9f6dfc93f4c..a6cdb0a1f96 100644 --- a/api/tasks/async_workflow_tasks.py +++ b/api/tasks/async_workflow_tasks.py @@ -232,6 +232,7 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None: return graph_runtime_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state) + response_stream_filter = resumption_context.get_response_stream_filter() with session_factory() as session: workflow = session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id)) @@ -294,6 +295,7 @@ def resume_workflow_execution(task_data_dict: dict[str, Any]) -> None: workflow_node_execution_repository=workflow_node_execution_repository, graph_engine_layers=graph_engine_layers, pause_state_config=pause_config, + response_stream_filter=response_stream_filter, ) workflow_run_repo.delete_workflow_pause(pause_entity) diff --git a/api/tasks/batch_create_segment_to_index_task.py b/api/tasks/batch_create_segment_to_index_task.py index 9f19b03544d..0f92af21dc8 100644 --- a/api/tasks/batch_create_segment_to_index_task.py +++ b/api/tasks/batch_create_segment_to_index_task.py @@ -177,7 +177,7 @@ def batch_create_segment_to_index_task( with session_factory.create_session() as session: dataset = session.get(Dataset, dataset_id) if dataset: - VectorService.create_segments_vector(None, document_segments, dataset, document_config["doc_form"]) + VectorService.create_segments_vector(None, document_segments, dataset, document_config["doc_form"], session) redis_client.setex(indexing_cache_key, 600, "completed") end_at = time.perf_counter() diff --git a/api/tasks/process_tenant_plugin_autoupgrade_check_task.py b/api/tasks/process_tenant_plugin_autoupgrade_check_task.py index 0840a595b7d..a796af5f3d3 100644 --- a/api/tasks/process_tenant_plugin_autoupgrade_check_task.py +++ b/api/tasks/process_tenant_plugin_autoupgrade_check_task.py @@ -11,11 +11,15 @@ from core.plugin.entities.plugin import PluginInstallation, PluginInstallationSo from core.plugin.impl.plugin import PluginInstaller from core.plugin.plugin_service import PluginService from extensions.ext_redis import redis_client -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import ( + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategySetting, +) logger = logging.getLogger(__name__) -PluginCategory = TenantPluginAutoUpgradeStrategy.PluginCategory +PluginCategory = TenantPluginAutoUpgradeCategory RETRY_TIMES_OF_ONE_PLUGIN_IN_ONE_TENANT = 3 CACHE_REDIS_KEY_PREFIX = "plugin_autoupgrade_check_task:cached_plugin_snapshot:" CACHE_REDIS_TTL = 60 * 60 # 1 hour @@ -95,9 +99,9 @@ def _plugin_matches_category(plugin: PluginInstallation, category: str | None) - @shared_task(queue="plugin") def process_tenant_plugin_autoupgrade_check_task( tenant_id: str, - strategy_setting: TenantPluginAutoUpgradeStrategy.StrategySetting, + strategy_setting: TenantPluginAutoUpgradeStrategySetting, upgrade_time_of_day: int, - upgrade_mode: TenantPluginAutoUpgradeStrategy.UpgradeMode, + upgrade_mode: TenantPluginAutoUpgradeMode, exclude_plugins: list[str], include_plugins: list[str], category: PluginCategory | str | None = None, @@ -113,14 +117,14 @@ def process_tenant_plugin_autoupgrade_check_task( ) ) - if strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED: + if strategy_setting == TenantPluginAutoUpgradeStrategySetting.DISABLED: return # get plugin_ids to check plugin_ids: list[tuple[str, str, str]] = [] # plugin_id, version, unique_identifier click.echo(click.style(f"Upgrade mode: {upgrade_mode}", fg="green")) - if upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL and include_plugins: + if upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL and include_plugins: all_plugins = manager.list_plugins(tenant_id) for plugin in all_plugins: @@ -137,7 +141,7 @@ def process_tenant_plugin_autoupgrade_check_task( ) ) - elif upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE: + elif upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE: # get all plugins and remove excluded plugins all_plugins = manager.list_plugins(tenant_id) plugin_ids = [ @@ -147,7 +151,7 @@ def process_tenant_plugin_autoupgrade_check_task( and plugin.plugin_id not in exclude_plugins and _plugin_matches_category(plugin, category_value) ] - elif upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL: + elif upgrade_mode == TenantPluginAutoUpgradeMode.ALL: all_plugins = manager.list_plugins(tenant_id) plugin_ids = [ (plugin.plugin_id, plugin.version, plugin.plugin_unique_identifier) @@ -187,8 +191,8 @@ def process_tenant_plugin_autoupgrade_check_task( return False version_checker = { - TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST: operator.ne, - TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY: fix_only_checker, + TenantPluginAutoUpgradeStrategySetting.LATEST: operator.ne, + TenantPluginAutoUpgradeStrategySetting.FIX_ONLY: fix_only_checker, } if version_checker[strategy_setting](latest_version, current_version): diff --git a/api/tasks/regenerate_summary_index_task.py b/api/tasks/regenerate_summary_index_task.py index 16b59fdbba8..5cb8d4281f0 100644 --- a/api/tasks/regenerate_summary_index_task.py +++ b/api/tasks/regenerate_summary_index_task.py @@ -259,9 +259,8 @@ def regenerate_summary_index_task( # Regenerate both summary content and vectors (for summary_model change) SummaryIndexService.generate_and_vectorize_summary( - segment, dataset, summary_index_setting + segment, dataset, summary_index_setting, session=session ) - session.commit() total_segments_processed += 1 except Exception as e: diff --git a/api/tasks/remove_app_and_related_data_task.py b/api/tasks/remove_app_and_related_data_task.py index a4fb6d57207..2010deb4a80 100644 --- a/api/tasks/remove_app_and_related_data_task.py +++ b/api/tasks/remove_app_and_related_data_task.py @@ -5,17 +5,25 @@ from typing import Any, cast import click import sqlalchemy as sa +from agenton.compositor import CompositorSessionSnapshot from celery import shared_task +from dify_agent.protocol import RuntimeLayerSpec +from pydantic import JsonValue, TypeAdapter from sqlalchemy import delete, select from sqlalchemy.engine import CursorResult from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import sessionmaker +from clients.agent_backend.session_cleanup import AgentBackendSessionCleanupPayload from configs import dify_config from core.db.session_factory import session_factory from extensions.ext_database import db from libs.archive_storage import ArchiveStorageNotConfiguredError, get_archive_storage +from libs.datetime_utils import naive_utc_now from models import ( + AgentRuntimeSession, + AgentRuntimeSessionOwnerType, + AgentRuntimeSessionStatus, ApiToken, AppAnnotationHitHistory, AppAnnotationSetting, @@ -50,8 +58,13 @@ from models.workflow import ( ) from repositories.factory import DifyAPIRepositoryFactory from services.api_token_service import ApiTokenCache +from tasks.agent_backend_session_cleanup_task import ( + cleanup_conversation_agent_runtime_session, + cleanup_workflow_agent_runtime_session, +) logger = logging.getLogger(__name__) +_RUNTIME_LAYER_SPECS_ADAPTER: TypeAdapter[list[RuntimeLayerSpec]] = TypeAdapter(list[RuntimeLayerSpec]) @shared_task(queue="app_deletion", bind=True, max_retries=3) @@ -59,6 +72,7 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): logger.info(click.style(f"Start deleting app and related data: {tenant_id}:{app_id}", fg="green")) start_at = time.perf_counter() try: + _cleanup_active_agent_runtime_sessions_for_app(tenant_id, app_id) # Delete related data _delete_app_model_configs(tenant_id, app_id) _delete_app_site(tenant_id, app_id) @@ -99,6 +113,143 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): raise self.retry(exc=e, countdown=60) # Retry after 60 seconds +def _cleanup_active_agent_runtime_sessions_for_app(tenant_id: str, app_id: str, *, batch_size: int = 100) -> None: + """Best-effort fan-out for ACTIVE Agent runtime sessions during app deletion. + + App deletion must not block on synchronous Agent backend lifecycle work, so + this helper scans ACTIVE ``agent_runtime_sessions`` rows in batches, + dispatches owner-specific cleanup tasks only when enough persisted data + exists to replay a lifecycle-only run, and then marks each visited row + ``CLEANED`` locally regardless of enqueue outcome. The local retirement is + the contract that lets the rest of app deletion continue even when backend + cleanup dispatch is skipped or fails. + """ + if batch_size <= 0: + raise ValueError("batch_size must be positive") + + while True: + with session_factory.create_session() as session: + row_ids = session.scalars( + select(AgentRuntimeSession.id) + .where( + AgentRuntimeSession.tenant_id == tenant_id, + AgentRuntimeSession.app_id == app_id, + AgentRuntimeSession.status == AgentRuntimeSessionStatus.ACTIVE, + ) + .order_by(AgentRuntimeSession.updated_at.asc()) + .limit(batch_size) + ).all() + + if not row_ids: + return + + retired_count = 0 + for row_id in row_ids: + with session_factory.create_session() as session: + row = session.get(AgentRuntimeSession, row_id) + if row is None or row.status != AgentRuntimeSessionStatus.ACTIVE: + retired_count += 1 + continue + + try: + payload = _build_agent_runtime_session_cleanup_payload(row) + if payload is not None: + _enqueue_agent_runtime_session_cleanup(row=row, payload=payload) + except Exception: + logger.warning( + "Failed to enqueue Agent backend cleanup during app deletion: " + "tenant_id=%s app_id=%s owner_type=%s conversation_id=%s workflow_run_id=%s " + "node_id=%s agent_id=%s backend_run_id=%s", + row.tenant_id, + row.app_id, + row.owner_type, + row.conversation_id, + row.workflow_run_id, + row.node_id, + row.agent_id, + row.backend_run_id, + exc_info=True, + ) + finally: + try: + row.status = AgentRuntimeSessionStatus.CLEANED + row.cleaned_at = naive_utc_now() + session.commit() + retired_count += 1 + except Exception: + session.rollback() + logger.warning( + "Failed to retire Agent runtime session during app deletion: " + "tenant_id=%s app_id=%s owner_type=%s conversation_id=%s workflow_run_id=%s " + "node_id=%s agent_id=%s backend_run_id=%s", + row.tenant_id, + row.app_id, + row.owner_type, + row.conversation_id, + row.workflow_run_id, + row.node_id, + row.agent_id, + row.backend_run_id, + exc_info=True, + ) + + if retired_count == 0: + logger.warning( + "Failed to retire any active Agent runtime sessions during app deletion: tenant_id=%s app_id=%s", + tenant_id, + app_id, + ) + return + + +def _build_agent_runtime_session_cleanup_payload( + row: AgentRuntimeSession, +) -> AgentBackendSessionCleanupPayload | None: + runtime_layer_specs = _RUNTIME_LAYER_SPECS_ADAPTER.validate_json(row.composition_layer_specs or "[]") + if not runtime_layer_specs: + return None + + metadata: dict[str, JsonValue] = { + "tenant_id": row.tenant_id, + "app_id": row.app_id, + "agent_id": row.agent_id, + "agent_config_snapshot_id": row.agent_config_snapshot_id, + "previous_agent_backend_run_id": row.backend_run_id, + } + if row.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION: + metadata["conversation_id"] = row.conversation_id + idempotency_key = ( + f"{row.tenant_id}:{row.app_id}:{row.conversation_id}:" + f"{row.agent_id}:app-delete-cleanup:{row.id or row.backend_run_id or 'no-session-id'}" + ) + else: + metadata["workflow_run_id"] = row.workflow_run_id + metadata["node_id"] = row.node_id + idempotency_key = ( + f"{row.tenant_id}:{row.app_id}:{row.workflow_run_id}:{row.node_id}:" + f"{row.agent_id}:app-delete-cleanup:{row.id or row.backend_run_id or 'no-session-id'}" + ) + + return AgentBackendSessionCleanupPayload( + session_snapshot=CompositorSessionSnapshot.model_validate_json(row.session_snapshot), + runtime_layer_specs=runtime_layer_specs, + idempotency_key=idempotency_key, + metadata=metadata, + ) + + +def _enqueue_agent_runtime_session_cleanup( + *, + row: AgentRuntimeSession, + payload: AgentBackendSessionCleanupPayload, +) -> None: + payload_dict = payload.model_dump(mode="json") + if row.owner_type == AgentRuntimeSessionOwnerType.CONVERSATION: + cleanup_conversation_agent_runtime_session.delay(payload_dict) + return + cleanup_workflow_agent_runtime_session.delay(payload_dict) + + def _delete_app_model_configs(tenant_id: str, app_id: str): def del_model_config(session, model_config_id: str): session.execute( diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py index fa02afda15f..dddb7715d22 100644 --- a/api/tasks/retry_document_indexing_task.py +++ b/api/tasks/retry_document_indexing_task.py @@ -101,8 +101,9 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_ session.commit() if dataset.runtime_mode == "rag_pipeline": - rag_pipeline_service = RagPipelineService() - rag_pipeline_service.retry_error_document(dataset, document, user) + with session_factory.create_session() as rag_session: + rag_pipeline_service = RagPipelineService(rag_session) + rag_pipeline_service.retry_error_document(dataset, document, user) else: indexing_runner = IndexingRunner() indexing_runner.run([document]) diff --git a/api/tasks/workflow_schedule_tasks.py b/api/tasks/workflow_schedule_tasks.py index 76386520000..38737f96e78 100644 --- a/api/tasks/workflow_schedule_tasks.py +++ b/api/tasks/workflow_schedule_tasks.py @@ -39,7 +39,7 @@ def run_schedule_trigger(schedule_id: str) -> None: if not schedule: raise ScheduleNotFoundError(f"Schedule {schedule_id} not found") - tenant_owner = ScheduleService.get_tenant_owner(session, schedule.tenant_id) + tenant_owner = ScheduleService.get_tenant_owner(schedule.tenant_id, session=session) if not tenant_owner: raise TenantOwnerNotFoundError(f"No owner or admin found for tenant {schedule.tenant_id}") diff --git a/api/tests/integration_tests/conftest.py b/api/tests/integration_tests/conftest.py index ea875e63fe8..25ee1974e92 100644 --- a/api/tests/integration_tests/conftest.py +++ b/api/tests/integration_tests/conftest.py @@ -84,7 +84,7 @@ def setup_account(request) -> Generator[Account, None, None]: password=secrets.token_hex(16), ip_address="localhost", language="en-US", - session=db.session, + session=db.session(), ) with _CACHED_APP.test_request_context(): diff --git a/api/tests/integration_tests/controllers/openapi/test_app_run.py b/api/tests/integration_tests/controllers/openapi/test_app_run.py index b4f383a7cea..fbfedd0f269 100644 --- a/api/tests/integration_tests/controllers/openapi/test_app_run.py +++ b/api/tests/integration_tests/controllers/openapi/test_app_run.py @@ -1,4 +1,4 @@ -"""Integration tests for POST /openapi/v1/apps//run.""" +"""Integration tests for POST /openapi/v1/apps/:run.""" from __future__ import annotations @@ -36,7 +36,7 @@ def test_run_chat_dispatches_to_chat_handler( monkeypatch.setattr("controllers.openapi.app_run.AppGenerateService.generate", staticmethod(_fake_generate)) client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"inputs": {}, "query": "hi", "response_mode": "blocking", "user": "spoof@x.com"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -85,7 +85,7 @@ def test_run_chat_without_query_returns_422( ): client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"inputs": {}, "response_mode": "blocking"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -116,7 +116,7 @@ def test_run_completion_dispatches_to_completion_handler( monkeypatch.setattr("controllers.openapi.app_run.AppGenerateService.generate", staticmethod(_fake_generate)) client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app.id}/run", + f"/openapi/v1/apps/{app.id}:run", json={"inputs": {}, "response_mode": "blocking"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -131,7 +131,7 @@ def test_run_workflow_with_query_returns_422( app = app_with_mode("workflow") client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app.id}/run", + f"/openapi/v1/apps/{app.id}:run", json={"inputs": {}, "query": "hi", "response_mode": "blocking"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -154,7 +154,7 @@ def test_run_workflow_no_query_dispatches_to_workflow_handler( monkeypatch.setattr("controllers.openapi.app_run.AppGenerateService.generate", staticmethod(_fake_generate)) client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app.id}/run", + f"/openapi/v1/apps/{app.id}:run", json={"inputs": {}, "response_mode": "blocking"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -170,7 +170,7 @@ def test_run_unsupported_mode_returns_422( app = app_with_mode("channel") client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app.id}/run", + f"/openapi/v1/apps/{app.id}:run", json={"inputs": {}, "response_mode": "blocking"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -181,7 +181,7 @@ def test_run_unsupported_mode_returns_422( def test_run_without_bearer_returns_401(flask_app: Flask, app_in_workspace): client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"inputs": {}, "query": "hi"}, ) assert res.status_code == 401 @@ -205,7 +205,7 @@ def test_run_with_insufficient_scope_returns_403( client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"inputs": {}, "query": "hi"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -215,7 +215,7 @@ def test_run_with_insufficient_scope_returns_403( def test_run_with_unknown_app_returns_404(flask_app: Flask, account_token): client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{uuid.uuid4()}/run", + f"/openapi/v1/apps/{uuid.uuid4()}:run", json={"inputs": {}, "query": "hi"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -235,7 +235,7 @@ def test_run_streaming_returns_event_stream( client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"inputs": {}, "query": "hi", "response_mode": "streaming"}, headers={"Authorization": f"Bearer {account_token}"}, ) @@ -247,7 +247,7 @@ def test_run_streaming_returns_event_stream( def test_run_without_inputs_returns_422(flask_app: Flask, account_token, app_in_workspace): client = flask_app.test_client() res = client.post( - f"/openapi/v1/apps/{app_in_workspace.id}/run", + f"/openapi/v1/apps/{app_in_workspace.id}:run", json={"query": "hi"}, headers={"Authorization": f"Bearer {account_token}"}, ) diff --git a/api/tests/integration_tests/controllers/openapi/test_apps.py b/api/tests/integration_tests/controllers/openapi/test_apps.py index 20ac46fbbde..ab992d1c5eb 100644 --- a/api/tests/integration_tests/controllers/openapi/test_apps.py +++ b/api/tests/integration_tests/controllers/openapi/test_apps.py @@ -37,7 +37,7 @@ def test_apps_describe_returns_merged_shape( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe", + f"/openapi/v1/apps/{app_in_workspace.id}", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -53,7 +53,7 @@ def test_apps_describe_full_includes_input_schema( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe", + f"/openapi/v1/apps/{app_in_workspace.id}", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -70,7 +70,7 @@ def test_apps_describe_fields_info_only( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=info", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=info", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -86,7 +86,7 @@ def test_apps_describe_fields_parameters_only( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=parameters", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=parameters", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -102,7 +102,7 @@ def test_apps_describe_fields_input_schema_only( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=input_schema", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=input_schema", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -118,7 +118,7 @@ def test_apps_describe_fields_combined( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=info,input_schema", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=info,input_schema", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 200 @@ -134,7 +134,7 @@ def test_apps_describe_fields_unknown_returns_422( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=garbage", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=garbage", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 422 @@ -146,7 +146,7 @@ def test_apps_describe_fields_extra_param_returns_422( account_token: str, ): res = test_client.get( - f"/openapi/v1/apps/{app_in_workspace.id}/describe?fields=info&page=1", + f"/openapi/v1/apps/{app_in_workspace.id}?fields=info&page=1", headers={"Authorization": f"Bearer {account_token}"}, ) assert res.status_code == 422 diff --git a/api/tests/integration_tests/model_runtime/__mock/plugin_model.py b/api/tests/integration_tests/model_runtime/__mock/plugin_model.py index c4146d5ccdd..9b2bb963d2f 100644 --- a/api/tests/integration_tests/model_runtime/__mock/plugin_model.py +++ b/api/tests/integration_tests/model_runtime/__mock/plugin_model.py @@ -246,5 +246,6 @@ class MockModelClass(PluginModelClient): tools: list[PromptMessageTool] | None = None, stop: list[str] | None = None, stream: bool = True, + app_id: str | None = None, ): return MockModelClass.mocked_chat_create_stream(model=model, prompt_messages=prompt_messages, tools=tools) diff --git a/api/tests/integration_tests/services/plugin/test_plugin_lifecycle.py b/api/tests/integration_tests/services/plugin/test_plugin_lifecycle.py index fc76129e3c7..e6e681b426f 100644 --- a/api/tests/integration_tests/services/plugin/test_plugin_lifecycle.py +++ b/api/tests/integration_tests/services/plugin/test_plugin_lifecycle.py @@ -2,12 +2,21 @@ import pytest from sqlalchemy import delete, func, select from core.db.session_factory import session_factory +from extensions.ext_database import db from models import Tenant -from models.account import TenantPluginAutoUpgradeStrategy, TenantPluginPermission +from models.account import ( + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategy, + TenantPluginAutoUpgradeStrategySetting, + TenantPluginDebugPermission, + TenantPluginInstallPermission, + TenantPluginPermission, +) from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService from services.plugin.plugin_permission_service import PluginPermissionService -PLUGIN_CATEGORY = TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL +PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL @pytest.fixture @@ -31,36 +40,39 @@ def tenant(flask_req_ctx): class TestPluginPermissionLifecycle: def test_get_returns_none_for_new_tenant(self, tenant): - assert PluginPermissionService.get_permission(tenant) is None + assert PluginPermissionService.get_permission(tenant, session=db.session()) is None def test_change_creates_row(self, tenant): result = PluginPermissionService.change_permission( tenant, - TenantPluginPermission.InstallPermission.ADMINS, - TenantPluginPermission.DebugPermission.EVERYONE, + TenantPluginInstallPermission.ADMINS, + TenantPluginDebugPermission.EVERYONE, + session=db.session, ) assert result is True - perm = PluginPermissionService.get_permission(tenant) + perm = PluginPermissionService.get_permission(tenant, session=db.session()) assert perm is not None - assert perm.install_permission == TenantPluginPermission.InstallPermission.ADMINS - assert perm.debug_permission == TenantPluginPermission.DebugPermission.EVERYONE + assert perm.install_permission == TenantPluginInstallPermission.ADMINS + assert perm.debug_permission == TenantPluginDebugPermission.EVERYONE def test_change_updates_existing_row(self, tenant): PluginPermissionService.change_permission( tenant, - TenantPluginPermission.InstallPermission.ADMINS, - TenantPluginPermission.DebugPermission.NOBODY, + TenantPluginInstallPermission.ADMINS, + TenantPluginDebugPermission.NOBODY, + session=db.session, ) PluginPermissionService.change_permission( tenant, - TenantPluginPermission.InstallPermission.EVERYONE, - TenantPluginPermission.DebugPermission.ADMINS, + TenantPluginInstallPermission.EVERYONE, + TenantPluginDebugPermission.ADMINS, + session=db.session, ) - perm = PluginPermissionService.get_permission(tenant) + perm = PluginPermissionService.get_permission(tenant, session=db.session()) assert perm is not None - assert perm.install_permission == TenantPluginPermission.InstallPermission.EVERYONE - assert perm.debug_permission == TenantPluginPermission.DebugPermission.ADMINS + assert perm.install_permission == TenantPluginInstallPermission.EVERYONE + assert perm.debug_permission == TenantPluginDebugPermission.ADMINS with session_factory.create_session() as session: count = session.scalar( @@ -73,73 +85,77 @@ class TestPluginPermissionLifecycle: class TestPluginAutoUpgradeLifecycle: def test_get_returns_none_for_new_tenant(self, tenant): - assert PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) is None + assert PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) is None def test_change_creates_row(self, tenant): result = PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, upgrade_time_of_day=3, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, exclude_plugins=[], include_plugins=[], category=PLUGIN_CATEGORY, + session=db.session(), ) assert result is True - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None - assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST + assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST assert strategy.upgrade_time_of_day == 3 def test_change_updates_existing_row(self, tenant): PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, exclude_plugins=[], include_plugins=[], category=PLUGIN_CATEGORY, + session=db.session(), ) PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, upgrade_time_of_day=12, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL, + upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, exclude_plugins=[], include_plugins=["plugin-a"], category=PLUGIN_CATEGORY, + session=db.session(), ) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None - assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST + assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST assert strategy.upgrade_time_of_day == 12 - assert strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL assert strategy.include_plugins == ["plugin-a"] def test_exclude_plugin_creates_strategy_when_none_exists(self, tenant): - PluginAutoUpgradeService.exclude_plugin(tenant, "my-plugin", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin(tenant, "my-plugin", PLUGIN_CATEGORY, session=db.session()) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None - assert strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE assert "my-plugin" in strategy.exclude_plugins def test_exclude_plugin_appends_in_exclude_mode(self, tenant): PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=["existing"], include_plugins=[], category=PLUGIN_CATEGORY, + session=db.session(), ) - PluginAutoUpgradeService.exclude_plugin(tenant, "new-plugin", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin(tenant, "new-plugin", PLUGIN_CATEGORY, session=db.session()) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None assert "existing" in strategy.exclude_plugins assert "new-plugin" in strategy.exclude_plugins @@ -147,32 +163,34 @@ class TestPluginAutoUpgradeLifecycle: def test_exclude_plugin_dedup_in_exclude_mode(self, tenant): PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=["same-plugin"], include_plugins=[], category=PLUGIN_CATEGORY, + session=db.session(), ) - PluginAutoUpgradeService.exclude_plugin(tenant, "same-plugin", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin(tenant, "same-plugin", PLUGIN_CATEGORY, session=db.session()) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None assert strategy.exclude_plugins.count("same-plugin") == 1 def test_exclude_from_partial_mode_removes_from_include(self, tenant): PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL, + upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, exclude_plugins=[], include_plugins=["p1", "p2"], category=PLUGIN_CATEGORY, + session=db.session(), ) - PluginAutoUpgradeService.exclude_plugin(tenant, "p1", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin(tenant, "p1", PLUGIN_CATEGORY, session=db.session()) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None assert "p1" not in strategy.include_plugins assert "p2" in strategy.include_plugins @@ -180,16 +198,17 @@ class TestPluginAutoUpgradeLifecycle: def test_exclude_from_all_mode_switches_to_exclude(self, tenant): PluginAutoUpgradeService.change_strategy( tenant, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, exclude_plugins=[], include_plugins=[], category=PLUGIN_CATEGORY, + session=db.session(), ) - PluginAutoUpgradeService.exclude_plugin(tenant, "excluded-plugin", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin(tenant, "excluded-plugin", PLUGIN_CATEGORY, session=db.session()) - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY) + strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=db.session()) assert strategy is not None - assert strategy.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE assert "excluded-plugin" in strategy.exclude_plugins diff --git a/api/tests/integration_tests/services/test_node_output_inspector_service.py b/api/tests/integration_tests/services/test_node_output_inspector_service.py index 5a8c07e0434..c2253a20c15 100644 --- a/api/tests/integration_tests/services/test_node_output_inspector_service.py +++ b/api/tests/integration_tests/services/test_node_output_inspector_service.py @@ -219,6 +219,36 @@ def _stub_resolver(declared_outputs_payload: list[dict[str, Any]]): return _Resolver() +def _snapshot_workflow_run(service: NodeOutputInspectorService, *, app_model: Any, workflow_run_id: str): + with session_factory.create_session() as session: + return service.snapshot_workflow_run(app_model=app_model, workflow_run_id=workflow_run_id, session=session) + + +def _node_detail(service: NodeOutputInspectorService, *, app_model: Any, workflow_run_id: str, node_id: str): + with session_factory.create_session() as session: + return service.node_detail( + app_model=app_model, workflow_run_id=workflow_run_id, node_id=node_id, session=session + ) + + +def _output_preview( + service: NodeOutputInspectorService, + *, + app_model: Any, + workflow_run_id: str, + node_id: str, + output_name: str, +): + with session_factory.create_session() as session: + return service.output_preview( + app_model=app_model, + workflow_run_id=workflow_run_id, + node_id=node_id, + output_name=output_name, + session=session, + ) + + # ────────────────────────────────────────────────────────────────────────────── # Tests # ────────────────────────────────────────────────────────────────────────────── @@ -229,7 +259,8 @@ def test_snapshot_returns_agent_v2_declared_outputs_with_status_ready(seeded_run real ``WorkflowRun`` + ``WorkflowNodeExecutionModel`` rows.""" app_model, workflow_run, _ = seeded_run service = NodeOutputInspectorService(binding_resolver=_stub_resolver([{"name": "text", "type": "string"}])) - snapshot = service.snapshot_workflow_run( + snapshot = _snapshot_workflow_run( + service, app_model=app_model, workflow_run_id=workflow_run.id, ) @@ -256,7 +287,7 @@ def test_snapshot_404s_for_missing_run(fake_app_model): """Service raises ``workflow_run_not_found`` when the row doesn't exist.""" service = NodeOutputInspectorService(binding_resolver=_stub_resolver([])) with pytest.raises(NodeOutputInspectorError) as exc: - service.snapshot_workflow_run(app_model=fake_app_model, workflow_run_id=str(uuid.uuid4())) + _snapshot_workflow_run(service, app_model=fake_app_model, workflow_run_id=str(uuid.uuid4())) assert exc.value.code == "workflow_run_not_found" @@ -266,7 +297,7 @@ def test_snapshot_404s_for_cross_tenant_access(seeded_run): intruder = SimpleNamespace(id=str(uuid.uuid4()), tenant_id=str(uuid.uuid4())) service = NodeOutputInspectorService(binding_resolver=_stub_resolver([])) with pytest.raises(NodeOutputInspectorError) as exc: - service.snapshot_workflow_run(app_model=intruder, workflow_run_id=workflow_run.id) + _snapshot_workflow_run(service, app_model=intruder, workflow_run_id=workflow_run.id) assert exc.value.code == "workflow_run_not_found" @@ -286,7 +317,7 @@ def test_snapshot_404s_for_published_run_per_decision_d1(flask_req_ctx, fake_app try: service = NodeOutputInspectorService(binding_resolver=_stub_resolver([])) with pytest.raises(NodeOutputInspectorError) as exc: - service.snapshot_workflow_run(app_model=fake_app_model, workflow_run_id=run_id) + _snapshot_workflow_run(service, app_model=fake_app_model, workflow_run_id=run_id) assert exc.value.code == "published_run_inspector_not_implemented" finally: with session_factory.create_session() as session: @@ -328,7 +359,7 @@ def test_snapshot_surfaces_type_check_failure_from_metadata(flask_req_ctx, fake_ try: service = NodeOutputInspectorService(binding_resolver=_stub_resolver([{"name": "summary", "type": "string"}])) - snapshot = service.snapshot_workflow_run(app_model=fake_app_model, workflow_run_id=run_id) + snapshot = _snapshot_workflow_run(service, app_model=fake_app_model, workflow_run_id=run_id) output = snapshot.node_outputs[0].outputs[0] assert output.status == NodeOutputStatus.TYPE_CHECK_FAILED assert output.type_check is not None @@ -375,7 +406,7 @@ def test_snapshot_surfaces_output_check_failure_from_metadata(flask_req_ctx, fak "services.workflow.node_output_inspector_service.file_helpers.get_signed_file_url", return_value="https://signed.example/report", ): - snapshot = service.snapshot_workflow_run(app_model=fake_app_model, workflow_run_id=run_id) + snapshot = _snapshot_workflow_run(service, app_model=fake_app_model, workflow_run_id=run_id) output = snapshot.node_outputs[0].outputs[0] assert output.status == NodeOutputStatus.OUTPUT_CHECK_FAILED assert output.output_check is not None @@ -391,7 +422,8 @@ def test_snapshot_surfaces_output_check_failure_from_metadata(flask_req_ctx, fak def test_node_detail_serves_one_node(seeded_run): app_model, workflow_run, _ = seeded_run service = NodeOutputInspectorService(binding_resolver=_stub_resolver([{"name": "text", "type": "string"}])) - view = service.node_detail( + view = _node_detail( + service, app_model=app_model, workflow_run_id=workflow_run.id, node_id="agent-node-1", @@ -421,7 +453,8 @@ def test_output_preview_for_file_renders_signed_url(seeded_run, fake_app_model): "services.workflow.node_output_inspector_service.file_helpers.get_signed_file_url", return_value="https://signed.example/x.pdf", ): - preview = service.output_preview( + preview = _output_preview( + service, app_model=fake_app_model, workflow_run_id=workflow_run.id, node_id="agent-node-1", @@ -466,7 +499,7 @@ def test_keeps_latest_execution_per_node_by_index(flask_req_ctx, fake_app_model) try: service = NodeOutputInspectorService(binding_resolver=_stub_resolver([{"name": "text", "type": "string"}])) - snapshot = service.snapshot_workflow_run(app_model=fake_app_model, workflow_run_id=run_id) + snapshot = _snapshot_workflow_run(service, app_model=fake_app_model, workflow_run_id=run_id) assert snapshot.node_outputs[0].outputs[0].value_preview == "second attempt" finally: with session_factory.create_session() as session: diff --git a/api/tests/integration_tests/workflow/test_response_stream_filter_pause_resume_integration.py b/api/tests/integration_tests/workflow/test_response_stream_filter_pause_resume_integration.py new file mode 100644 index 00000000000..08ebbfb31a1 --- /dev/null +++ b/api/tests/integration_tests/workflow/test_response_stream_filter_pause_resume_integration.py @@ -0,0 +1,216 @@ +"""Regression test: if-else branch + human_input pause + downstream answer nodes. + +Reproduces https://github.com/langgenius/dify/issues/38525 at the +iter_dify_graph_engine_events layer: without a restored ResponseStreamFilter, +answer nodes downstream of a pre-pause branch never unlock for streaming on +resume, even though the graph executes correctly. +""" + +from datetime import timedelta +from unittest.mock import MagicMock + +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.repositories.human_input_repository import HumanInputFormEntity, HumanInputFormRepository +from core.workflow.nodes.human_input.callback import DifyHITLCallback +from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig +from core.workflow.nodes.human_input.enums import HumanInputFormStatus +from core.workflow.system_variables import build_system_variables +from core.workflow.workflow_entry import iter_dify_graph_engine_events +from graphon.filters import GraphEventFilterContext, ResponseStreamFilter, filter_graph_events +from graphon.graph import Graph +from graphon.graph_engine import GraphEngine, GraphEngineConfig +from graphon.graph_engine.command_channels import InMemoryChannel +from graphon.graph_events import GraphRunPausedEvent, GraphRunSucceededEvent, NodeRunStreamChunkEvent +from graphon.nodes.answer.answer_node import AnswerNode +from graphon.nodes.answer.entities import AnswerNodeData +from graphon.nodes.human_input.human_input_node import HumanInputNode +from graphon.nodes.if_else.entities import IfElseNodeData +from graphon.nodes.if_else.if_else_node import IfElseNode +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.runtime import GraphRuntimeState, VariablePool +from graphon.utils.condition.entities import Condition +from libs.datetime_utils import naive_utc_now +from tests.workflow_test_utils import build_test_graph_init_params + +WORKFLOW_EXECUTION_ID = "wf-exec-38525" + + +def _mock_repo_paused() -> HumanInputFormRepository: + repo = MagicMock(spec=HumanInputFormRepository) + form = MagicMock(spec=HumanInputFormEntity) + form.id = "form-1" + form.submission_token = "token-1" + form.recipients = [] + form.rendered_content = "rendered" + form.submitted = False + repo.create_form.return_value = form + repo.get_form.return_value = None + return repo + + +def _mock_repo_resumed(action_id: str = "continue") -> HumanInputFormRepository: + repo = MagicMock(spec=HumanInputFormRepository) + form = MagicMock(spec=HumanInputFormEntity) + form.id = "form-1" + form.submission_token = "token-1" + form.recipients = [] + form.rendered_content = "rendered" + form.submitted = True + form.selected_action_id = action_id + form.submitted_data = {} + form.status = HumanInputFormStatus.WAITING + form.expiration_time = naive_utc_now() + timedelta(hours=1) + repo.get_form.return_value = form + return repo + + +def _build_graph(runtime_state: GraphRuntimeState, form_repository: HumanInputFormRepository) -> Graph: + params = build_test_graph_init_params( + workflow_id="wf", + graph_config={"nodes": [], "edges": []}, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ) + + start_node = StartNode( + node_id="start", + data=StartNodeData(title="start", variables=[]), + graph_init_params=params, + graph_runtime_state=runtime_state, + ) + + if_else_node = IfElseNode( + node_id="if_else", + data=IfElseNodeData( + title="if-else", + cases=[ + IfElseNodeData.Case( + case_id="true", + logical_operator="and", + conditions=[ + Condition( + variable_selector=["start", "category"], + comparison_operator="is", + value="fruit", + ) + ], + ) + ], + ), + graph_init_params=params, + graph_runtime_state=runtime_state, + ) + + human_data = HumanInputNodeData( + title="human", + form_content="Awaiting human input", + inputs=[], + user_actions=[UserActionConfig(id="continue", title="Continue")], + ) + human_node = HumanInputNode( + node_id="human_input", + data=human_data, + graph_init_params=params, + graph_runtime_state=runtime_state, + hitl_callback=DifyHITLCallback(form_repository=form_repository, node_data=human_data), + ) + + answer_false_node = AnswerNode( + node_id="answer_false", + data=AnswerNodeData(title="answer_false", answer="unreachable branch"), + graph_init_params=params, + graph_runtime_state=runtime_state, + ) + + answer_after_pause = AnswerNode( + node_id="answer_after_pause", + data=AnswerNodeData(title="answer_after_pause", answer="Post-branch answer chunk 1"), + graph_init_params=params, + graph_runtime_state=runtime_state, + ) + + answer_after_pause_2 = AnswerNode( + node_id="answer_after_pause_2", + data=AnswerNodeData(title="answer_after_pause_2", answer="Post-branch answer chunk 2"), + graph_init_params=params, + graph_runtime_state=runtime_state, + ) + + return ( + Graph.new() + .add_root(start_node) + .add_node(if_else_node, from_node_id="start") + .add_node(human_node, from_node_id="if_else", source_handle="true") + .add_node(answer_false_node, from_node_id="if_else", source_handle="false") + .add_node(answer_after_pause, from_node_id="human_input", source_handle="continue") + .add_node(answer_after_pause_2, from_node_id="answer_after_pause") + .build() + ) + + +def _build_runtime_state() -> GraphRuntimeState: + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + workflow_execution_id=WORKFLOW_EXECUTION_ID, + app_id="app", + workflow_id="wf", + user_id="user", + ), + user_inputs={}, + conversation_variables=[], + ) + variable_pool.add(("start", "category"), "fruit") # drives the if-else "true" branch + return GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) + + +def test_if_else_human_input_pause_resume_answer_chunks_survive_resume() -> None: + # ---- Phase 1: run to GraphRunPausedEvent ---- + runtime_state_1 = _build_runtime_state() + graph_1 = _build_graph(runtime_state_1, _mock_repo_paused()) + engine_1 = GraphEngine( + workflow_id="wf", + graph=graph_1, + graph_runtime_state=runtime_state_1, + command_channel=InMemoryChannel(), + config=GraphEngineConfig(), + ) + filter_1 = ResponseStreamFilter() + phase1_events = list( + filter_graph_events( + engine_1.run(), + context=GraphEventFilterContext.from_engine(engine_1), + filters=[filter_1], + ) + ) + + assert any(isinstance(e, GraphRunPausedEvent) for e in phase1_events) + phase1_chunks = [e for e in phase1_events if isinstance(e, NodeRunStreamChunkEvent)] + assert not any(e.node_id in ("answer_after_pause", "answer_after_pause_2") for e in phase1_chunks) + + response_filter_snapshot = filter_1.dumps() + runtime_snapshot = runtime_state_1.dumps() + + # ---- Phase 2: rebuild engine + filter from snapshots, resume to completion ---- + runtime_state_2 = GraphRuntimeState.from_snapshot(runtime_snapshot) + graph_2 = _build_graph(runtime_state_2, _mock_repo_resumed(action_id="continue")) + engine_2 = GraphEngine( + workflow_id="wf", + graph=graph_2, + graph_runtime_state=runtime_state_2, + command_channel=InMemoryChannel(), + config=GraphEngineConfig(), + ) + filter_2 = ResponseStreamFilter() + filter_2.loads(response_filter_snapshot) + + phase2_events = list(iter_dify_graph_engine_events(engine_2, filter_2)) + + assert any(isinstance(e, GraphRunSucceededEvent) for e in phase2_events) + + phase2_chunks = [e for e in phase2_events if isinstance(e, NodeRunStreamChunkEvent)] + answer_1_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause"] + answer_2_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause_2"] + + assert answer_1_chunks, "answer_after_pause produced no stream chunks after resume" + assert answer_2_chunks, "answer_after_pause_2 produced no stream chunks after resume" diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py b/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py index ae37d305670..df9d655fbbc 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py +++ b/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py @@ -500,7 +500,11 @@ class TestWorkflowDraftVariableEndpoints: api = workflow_draft_variable_module.WorkflowVariableCollectionApi() method = unwrap(api.get) - monkeypatch.setattr(workflow_draft_variable_module, "db", SimpleNamespace(engine=MagicMock())) + monkeypatch.setattr( + workflow_draft_variable_module, + "db", + SimpleNamespace(engine=MagicMock(), session=MagicMock()), + ) class DummySessionCtx: def __enter__(self): diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py index e55b46d38bf..ef8c0add709 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py @@ -85,7 +85,7 @@ def test_create_binding_successful( assert response.status_code == 200 assert response.get_json() == {"result": "success"} - create_auth.assert_called_once_with(ANY, tenant_id, payload) + create_auth.assert_called_once_with(tenant_id, payload, session=ANY) def test_create_binding_failure( diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_email_register.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_email_register.py index 109332e16c9..d893e9e6efb 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_email_register.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_email_register.py @@ -270,7 +270,7 @@ def test_get_account_by_email_with_case_fallback_falls_back_to_lowercase(): second_result.scalar_one_or_none.return_value = expected_account mock_session.execute.side_effect = [first_result, second_result] - result = AccountService.get_account_by_email_with_case_fallback(mock_session, "Case@Test.com") + result = AccountService.get_account_by_email_with_case_fallback("Case@Test.com", session=mock_session) assert result is expected_account assert mock_session.execute.call_count == 2 diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_forgot_password.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_forgot_password.py index 812aa299c1b..a7eba9d723c 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_forgot_password.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_forgot_password.py @@ -165,7 +165,7 @@ def test_get_account_by_email_with_case_fallback_falls_back_to_lowercase(): second_result.scalar_one_or_none.return_value = expected_account mock_session.execute.side_effect = [first_result, second_result] - result = AccountService.get_account_by_email_with_case_fallback(mock_session, "Mixed@Test.com") + result = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=mock_session) assert result is expected_account assert mock_session.execute.call_count == 2 diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py index 464e0134a2f..484ca71ca59 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth.py @@ -494,7 +494,7 @@ class TestAccountGeneration: second_result.scalar_one_or_none.return_value = expected_account mock_session.execute.side_effect = [first_result, second_result] - result = AccountService.get_account_by_email_with_case_fallback(mock_session, "Case@Test.com") + result = AccountService.get_account_by_email_with_case_fallback("Case@Test.com", session=mock_session) assert result is expected_account assert mock_session.execute.call_count == 2 diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py index c34810c97d0..6c73b0010ed 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Callable from inspect import unwrap from typing import cast -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest @@ -90,55 +90,51 @@ class TestPipelineTemplateDetailApi: "graph": {"nodes": nodes, "edges": edges, "viewport": viewport}, } - service = MagicMock() - service.get_pipeline_template_detail.return_value = template - with ( app.test_request_context("/?type=built-in"), patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService", - return_value=service, - ), + "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService.get_pipeline_template_detail", + return_value=template, + ) as get_detail_mock, ): response, status = method(api, MagicMock(), "tpl-1") assert status == 200 assert response == {**template, "created_by": None} + get_detail_mock.assert_called_once_with("tpl-1", type="built-in", session=ANY) def test_get_returns_404_when_template_not_found(self, app: Flask) -> None: api = PipelineTemplateDetailApi() method = unwrap(api.get) - service = MagicMock() - service.get_pipeline_template_detail.return_value = None - with ( app.test_request_context("/?type=built-in"), patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService", - return_value=service, - ), + "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService.get_pipeline_template_detail", + return_value=None, + ) as get_detail_mock, ): with pytest.raises(NotFound): method(api, MagicMock(), "non-existent-id") + get_detail_mock.assert_called_once_with("non-existent-id", type="built-in", session=ANY) + def test_get_returns_404_for_customized_type_not_found(self, app: Flask) -> None: api = PipelineTemplateDetailApi() method = unwrap(api.get) - service = MagicMock() - service.get_pipeline_template_detail.return_value = None - with ( app.test_request_context("/?type=customized"), patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService", - return_value=service, - ), + "controllers.console.datasets.rag_pipeline.rag_pipeline.RagPipelineService.get_pipeline_template_detail", + return_value=None, + ) as get_detail_mock, ): with pytest.raises(NotFound): method(api, MagicMock(), "non-existent-id") + get_detail_mock.assert_called_once_with("non-existent-id", type="customized", session=ANY) + class TestCustomizedPipelineTemplateApi: @pytest.fixture @@ -186,7 +182,7 @@ class TestCustomizedPipelineTemplateApi: ): response, status = method(api, tenant_id, "tpl-1") - delete_mock.assert_called_once_with("tpl-1", tenant_id) + delete_mock.assert_called_once_with("tpl-1", tenant_id, session=ANY) assert status == 204 assert response == "" diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py index ff1521800ad..b9a0029d131 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py +++ b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py @@ -14,7 +14,7 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_import import ( RagPipelineImportCheckDependenciesApi, RagPipelineImportConfirmApi, ) -from core.plugin.entities.plugin import PluginDependency +from core.plugin.entities.plugin import PluginDependency, PluginDependencyType from models.dataset import Pipeline from services.entities.dsl_entities import CheckDependenciesResult, ImportStatus from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineImportInfo @@ -237,7 +237,7 @@ class TestRagPipelineImportCheckDependenciesApi: pipeline = MagicMock(spec=Pipeline) dependency = PluginDependency( - type=PluginDependency.Type.Marketplace, + type=PluginDependencyType.Marketplace, value=PluginDependency.Marketplace( marketplace_plugin_unique_identifier="langgenius/example:0.1.0", version="0.1.0", diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py index e1bdff4d23c..4a41aa352d2 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py @@ -607,7 +607,11 @@ class TestMiscApis: method = unwrap(api.get) service = MagicMock() - service.get_recommended_plugins.return_value = [{"id": "p1"}] + recommended_plugins = { + "installed_recommended_plugins": [{"id": "p1"}], + "uninstalled_recommended_plugins": [{"id": "p2"}], + } + service.get_recommended_plugins.return_value = recommended_plugins user = make_account() tenant_id = "tenant-1" @@ -619,7 +623,7 @@ class TestMiscApis: ), ): result = method(api, tenant_id, user) - assert result == [{"id": "p1"}] + assert result == recommended_plugins service.get_recommended_plugins.assert_called_once_with("all", user, tenant_id) diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py b/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py index e60558040a5..4cca4c2170f 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py +++ b/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py @@ -97,13 +97,13 @@ def test_list_scopes_api_based_extensions_to_authenticated_tenant( assert account_create_response.status_code == 201 APIBasedExtensionService.save( - db_session_with_containers, APIBasedExtension( tenant_id=foreign_tenant_id, name="Foreign API", api_endpoint="https://foreign.example.com/hook", api_key="foreign-secret-12345", ), + session=db_session_with_containers, ) response = test_client_with_containers.get( diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py index 8739ca28bd3..37aa68fc00b 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py +++ b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py @@ -44,6 +44,7 @@ from controllers.console.workspace.tool_providers import ( ToolWorkflowProviderUpdateApi, is_valid_url, ) +from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity from models.account import Account, TenantAccountRole from services.tools.mcp_tools_manage_service import ReconnectResult from tests.test_containers_integration_tests.controllers.console.helpers import ( @@ -60,6 +61,148 @@ def empty_list() -> list[object]: return [] +def emoji_icon() -> dict[str, str]: + return {"content": "tool", "background": "#252525"} + + +def i18n(text: str) -> dict[str, str]: + return {"en_US": text} + + +def tool_payload(name: str = "ping") -> dict[str, object]: + return { + "author": "langgenius", + "name": name, + "label": i18n(name.title()), + "description": i18n(f"{name} description"), + "parameters": [], + "labels": ["utilities"], + "output_schema": {}, + } + + +def provider_payload( + *, + provider_id: str = "provider-1", + name: str = "provider", + provider_type: str = "builtin", + tools: list[dict[str, object]] | None = None, +) -> dict[str, object]: + return { + "id": provider_id, + "author": "langgenius", + "name": name, + "description": i18n(f"{name} description"), + "icon": emoji_icon(), + "icon_dark": emoji_icon(), + "label": i18n(name.title()), + "type": provider_type, + "masked_credentials": {"api_key": "[__HIDDEN__]"}, + "original_credentials": {"api_key": "sk-secret"}, + "is_team_authorization": False, + "allow_delete": True, + "plugin_id": "langgenius/provider", + "plugin_unique_identifier": "langgenius/provider:1.0.0", + "tools": tools or [tool_payload()], + "labels": ["utilities"], + "server_url": "", + "updated_at": 1710000000, + "server_identifier": "", + "masked_headers": None, + "original_headers": None, + "authentication": None, + "is_dynamic_registration": True, + "configuration": None, + "identity_mode": "off", + "workflow_app_id": None, + } + + +def provider_entity( + *, + provider_id: str = "provider-1", + name: str = "provider", + provider_type: str = "builtin", + tools: list[dict[str, object]] | None = None, +) -> CoreToolProviderApiEntity: + return CoreToolProviderApiEntity.model_validate( + provider_payload(provider_id=provider_id, name=name, provider_type=provider_type, tools=tools) + ) + + +def credential_payload() -> dict[str, object]: + return { + "id": "credential-1", + "name": "Default credential", + "provider": "provider", + "credential_type": "api-key", + "is_default": True, + "credentials": {"api_key": "masked"}, + "visibility": "all_team_members", + "created_by": "user-1", + "partial_member_list": [], + "from_other_member": False, + } + + +def provider_config_payload() -> dict[str, object]: + return {"type": "secret-input", "name": "api_key", "required": True} + + +def api_tool_bundle_payload() -> dict[str, object]: + return { + "server_url": "https://api.example.com", + "method": "get", + "summary": "Ping", + "operation_id": "ping", + "parameters": [], + "author": "langgenius", + "icon": None, + "openapi": {"operationId": "ping"}, + "output_schema": {}, + } + + +def api_provider_detail_payload() -> dict[str, object]: + return { + "schema_type": "openapi", + "schema": "{}", + "tools": [api_tool_bundle_payload()], + "icon": emoji_icon(), + "description": "API provider", + "credentials": {}, + "privacy_policy": "", + "custom_disclaimer": "", + "labels": ["utilities"], + } + + +def credential_info_payload() -> dict[str, object]: + return { + "supported_credential_types": ["api-key", "oauth2"], + "is_oauth_custom_client_enabled": False, + "credentials": [credential_payload()], + } + + +def oauth_client_schema_payload() -> dict[str, object]: + return { + "schema": [provider_config_payload()], + "is_oauth_custom_client_enabled": False, + "is_system_oauth_params_exists": True, + "client_params": {"client_id": "masked"}, + "redirect_uri": "https://console.example.com/oauth/callback", + } + + +def tool_label_payload() -> dict[str, object]: + return { + "name": "utilities", + "label": i18n("Utilities"), + "icon": "wrench", + } + + @pytest.fixture def _mock_cache() -> None: return @@ -127,7 +270,7 @@ def test_create_mcp_provider_populates_tools( with ( patch( "services.tools.tools_transform_service.ToolTransformService.mcp_provider_to_user_provider", - return_value={"id": "provider-1", "tools": [{"name": "ping"}]}, + return_value=provider_entity(provider_id="provider-1", provider_type="mcp", tools=[tool_payload()]), autospec=True, ), ): @@ -138,13 +281,15 @@ def test_create_mcp_provider_populates_tools( content_type="application/json", ) - # Assert - assert resp.status_code == 200 - body = resp.get_json() - assert body.get("id") == "provider-1" - # 若 transform 后包含 tools 字段，确保非空 - assert isinstance(body.get("tools"), list) - assert body["tools"] + # Assert + assert resp.status_code == 200 + body = resp.get_json() + assert body.get("id") == "provider-1" + assert body["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in body + assert "original_credentials" not in body + assert isinstance(body.get("tools"), list) + assert body["tools"] class TestUtils: @@ -170,10 +315,16 @@ class TestToolProviderListApi: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.ToolCommonService.list_tool_providers", - return_value=["p1"], + return_value=[provider_entity(provider_id="p1").to_dict()], ), ): - assert method(api, "t1", make_account(id="u1")) == ["p1"] + result = method(api, "t1", make_account(id="u1")) + + assert result[0]["id"] == "p1" + assert result[0]["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in result[0] + assert "original_credentials" not in result[0] + assert result[0]["tools"][0]["name"] == "ping" class TestBuiltinProviderApis: @@ -189,10 +340,10 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tool_provider_tools", - return_value=[{"a": 1}], + return_value=[tool_payload()], ), ): - assert method(api, "t1", "provider") == [{"a": 1}] + assert method(api, "t1", "provider")[0]["name"] == "ping" def test_info(self, app: Flask) -> None: api = ToolBuiltinProviderInfoApi() @@ -202,10 +353,15 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_info", - return_value={"x": 1}, + return_value=provider_entity(), ), ): - assert method(api, "t1", "provider") == {"x": 1} + result = method(api, "t1", "provider") + + assert result["id"] == "provider-1" + assert result["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in result + assert "original_credentials" not in result def test_delete(self, app: Flask) -> None: api = ToolBuiltinProviderDeleteApi() @@ -240,10 +396,10 @@ class TestBuiltinProviderApis: app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.add_builtin_tool_provider", - return_value={"id": 1}, + return_value={"result": "success"}, ), ): - assert method(api, "t", make_account(), "provider")["id"] == 1 + assert method(api, "t", make_account(), "provider")["result"] == "success" def test_update(self, app: Flask) -> None: api = ToolBuiltinProviderUpdateApi() @@ -255,10 +411,10 @@ class TestBuiltinProviderApis: app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.update_builtin_tool_provider", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t", make_account(), "provider")["ok"] + assert method(api, "t", make_account(), "provider")["result"] == "success" def test_get_credentials(self, app: Flask) -> None: api = ToolBuiltinProviderGetCredentialsApi() @@ -268,10 +424,10 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credentials", - return_value={"k": "v"}, + return_value=[credential_payload()], ), ): - assert method(api, "t", make_account(id="user-1"), "provider") == {"k": "v"} + assert method(api, "t", make_account(id="user-1"), "provider")[0]["id"] == "credential-1" def test_icon(self, app: Flask) -> None: api = ToolBuiltinProviderIconApi() @@ -295,10 +451,10 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_provider_credentials_schema", - return_value={"schema": {}}, + return_value=[provider_config_payload()], ), ): - assert method(api, "t", "provider", "oauth2") == {"schema": {}} + assert method(api, "t", "provider", "oauth2")[0]["name"] == "api_key" def test_set_default_credential(self, app: Flask) -> None: api = ToolBuiltinProviderSetDefaultApi() @@ -308,10 +464,10 @@ class TestBuiltinProviderApis: app.test_request_context("/", json={"id": "c1"}), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.set_default_provider", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t", "provider")["ok"] + assert method(api, "t", "provider")["result"] == "success" def test_get_credential_info(self, app: Flask) -> None: api = ToolBuiltinProviderGetCredentialInfoApi() @@ -321,10 +477,10 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credential_info", - return_value={"info": "x"}, + return_value=credential_info_payload(), ), ): - assert method(api, "t", make_account(), "provider") == {"info": "x"} + assert method(api, "t", make_account(), "provider")["credentials"][0]["id"] == "credential-1" def test_get_oauth_client_schema(self, app: Flask) -> None: api = ToolBuiltinProviderGetOauthClientSchemaApi() @@ -334,10 +490,10 @@ class TestBuiltinProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_oauth_client_schema", - return_value={"schema": {}}, + return_value=oauth_client_schema_payload(), ), ): - assert method(api, "t", "provider") == {"schema": {}} + assert method(api, "t", "provider")["schema"][0]["name"] == "api_key" class TestApiProviderApis: @@ -354,30 +510,34 @@ class TestApiProviderApis: "schema_type": "openapi", "schema": "{}", "provider": "p", - "icon": empty_mapping(), + "icon": emoji_icon(), } with ( app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.create_api_tool_provider", - return_value={"id": 1}, - ), + return_value={"result": "success"}, + ) as create_api_tool_provider, ): - assert method(api, "t", make_account())["id"] == 1 + assert method(api, "t", make_account()) == {"result": "success"} + + create_api_tool_provider.assert_called_once() + assert create_api_tool_provider.call_args.args[3] == emoji_icon() def test_remote_schema(self, app: Flask) -> None: api = ToolApiProviderGetRemoteSchemaApi() method = unwrap(api.get) + openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' with ( app.test_request_context("/?url=http://x.com"), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider_remote_schema", - return_value={"schema": "x"}, + return_value={"schema": openapi_schema}, ), ): - assert method(api, "t", make_account())["schema"] == "x" + assert method(api, "t", make_account()) == {"schema": openapi_schema} def test_list_tools(self, app: Flask) -> None: api = ToolApiProviderListToolsApi() @@ -387,10 +547,10 @@ class TestApiProviderApis: app.test_request_context("/?provider=p"), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tool_provider_tools", - return_value=[{"tool": 1}], + return_value=[tool_payload("api_ping")], ), ): - assert method(api, "t", make_account()) == [{"tool": 1}] + assert method(api, "t", make_account())[0]["name"] == "api_ping" def test_update(self, app: Flask) -> None: api = ToolApiProviderUpdateApi() @@ -402,7 +562,7 @@ class TestApiProviderApis: "schema": "{}", "provider": "p", "original_provider": "o", - "icon": empty_mapping(), + "icon": emoji_icon(), "privacy_policy": "", "custom_disclaimer": "", } @@ -411,10 +571,13 @@ class TestApiProviderApis: app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.update_api_tool_provider", - return_value={"ok": True}, - ), + return_value={"result": "success"}, + ) as update_api_tool_provider, ): - assert method(api, "t", make_account())["ok"] + assert method(api, "t", make_account()) == {"result": "success"} + + update_api_tool_provider.assert_called_once() + assert update_api_tool_provider.call_args.args[4] == emoji_icon() def test_delete(self, app: Flask) -> None: api = ToolApiProviderDeleteApi() @@ -437,10 +600,10 @@ class TestApiProviderApis: app.test_request_context("/?provider=p"), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider", - return_value={"x": 1}, + return_value=api_provider_detail_payload(), ), ): - assert method(api, "t", make_account()) == {"x": 1} + assert method(api, "t", make_account())["schema"] == "{}" class TestWorkflowApis: @@ -457,7 +620,7 @@ class TestWorkflowApis: "name": "n", "label": "l", "description": "d", - "icon": empty_mapping(), + "icon": emoji_icon(), "parameters": empty_list(), } @@ -465,10 +628,13 @@ class TestWorkflowApis: app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.WorkflowToolManageService.create_workflow_tool", - return_value={"id": 1}, - ), + return_value={"result": "success"}, + ) as create_workflow_tool, ): - assert method(api, "t", make_account())["id"] == 1 + assert method(api, "t", make_account()) == {"result": "success"} + + create_workflow_tool.assert_called_once() + assert create_workflow_tool.call_args.kwargs["icon"] == emoji_icon() def test_update_invalid(self, app: Flask) -> None: api = ToolWorkflowProviderUpdateApi() @@ -479,18 +645,21 @@ class TestWorkflowApis: "name": "Tool", "label": "Tool Label", "description": "A tool", - "icon": empty_mapping(), + "icon": emoji_icon(), } with ( app.test_request_context("/", json=payload), patch( "controllers.console.workspace.tool_providers.WorkflowToolManageService.update_workflow_tool", - return_value={"ok": True}, - ), + return_value={"result": "success"}, + ) as update_workflow_tool, ): result = method(api, "t", make_account()) - assert result["ok"] + assert result == {"result": "success"} + + update_workflow_tool.assert_called_once() + assert update_workflow_tool.call_args.args[5] == emoji_icon() def test_delete(self, app: Flask) -> None: api = ToolWorkflowProviderDeleteApi() @@ -500,10 +669,10 @@ class TestWorkflowApis: app.test_request_context("/", json={"workflow_tool_id": "123e4567-e89b-12d3-a456-426614174000"}), patch( "controllers.console.workspace.tool_providers.WorkflowToolManageService.delete_workflow_tool", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t", make_account())["ok"] + assert method(api, "t", make_account())["result"] == "success" def test_get_error(self, app: Flask) -> None: api = ToolWorkflowProviderGetApi() @@ -525,49 +694,40 @@ class TestLists: api = ToolBuiltinListApi() method = unwrap(api.get) - m = MagicMock() - m.to_dict.return_value = {"x": 1} - with ( app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tools", - return_value=[m], + return_value=[provider_entity(provider_id="builtin-1")], ), ): - assert method(api, "t", make_account()) == [{"x": 1}] + assert method(api, "t", make_account())[0]["id"] == "builtin-1" def test_api_list(self, app: Flask) -> None: api = ToolApiListApi() method = unwrap(api.get) - m = MagicMock() - m.to_dict.return_value = {"x": 1} - with ( app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tools", - return_value=[m], + return_value=[provider_entity(provider_id="api-1", provider_type="api")], ), ): - assert method(api, "t") == [{"x": 1}] + assert method(api, "t")[0]["id"] == "api-1" def test_workflow_list(self, app: Flask) -> None: api = ToolWorkflowListApi() method = unwrap(api.get) - m = MagicMock() - m.to_dict.return_value = {"x": 1} - with ( app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.WorkflowToolManageService.list_tenant_workflow_tools", - return_value=[m], + return_value=[provider_entity(provider_id="workflow-1", provider_type="workflow")], ), ): - assert method(api, "t", make_account()) == [{"x": 1}] + assert method(api, "t", make_account())[0]["id"] == "workflow-1" class TestLabels: @@ -583,10 +743,10 @@ class TestLabels: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.ToolLabelsService.list_tool_labels", - return_value=["l1"], + return_value=[tool_label_payload()], ), ): - assert method(api) == ["l1"] + assert method(api)[0]["name"] == "utilities" class TestOAuth: @@ -630,10 +790,10 @@ class TestOAuthCustomClient: app.test_request_context("/", json={"client_params": {"a": 1}}), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.save_custom_oauth_client_params", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t", "provider")["ok"] + assert method(api, "t", "provider") == {"result": "success"} def test_get_custom_client(self, app: Flask) -> None: api = ToolOAuthCustomClient() @@ -656,7 +816,7 @@ class TestOAuthCustomClient: app.test_request_context("/"), patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.delete_custom_oauth_client_params", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t", "provider")["ok"] + assert method(api, "t", "provider") == {"result": "success"} diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_trigger_providers.py b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_trigger_providers.py index 6684381880c..31d625ac91d 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_trigger_providers.py +++ b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_trigger_providers.py @@ -2,6 +2,7 @@ from __future__ import annotations +from datetime import datetime from inspect import unwrap from unittest.mock import MagicMock, patch @@ -29,6 +30,8 @@ from controllers.console.workspace.trigger_providers import ( TriggerSubscriptionVerifyApi, ) from core.plugin.entities.plugin_daemon import CredentialType +from core.trigger.entities.api_entities import SubscriptionBuilderApiEntity, TriggerProviderApiEntity +from core.trigger.entities.entities import RequestLog from models.account import Account @@ -38,6 +41,47 @@ def mock_user() -> Account: return user +def trigger_provider() -> TriggerProviderApiEntity: + return TriggerProviderApiEntity( + author="Dify", + name="github", + label={"en_US": "GitHub"}, + description={"en_US": "GitHub trigger provider"}, + icon="icon.svg", + icon_dark=None, + tags=["code"], + plugin_id="plugin", + plugin_unique_identifier="plugin:github", + supported_creation_methods=[], + subscription_constructor=None, + subscription_schema=[], + events=[], + ) + + +def subscription_builder() -> SubscriptionBuilderApiEntity: + return SubscriptionBuilderApiEntity( + id="b1", + name="Builder", + provider="github", + endpoint="b1", + parameters={"repo": "dify"}, + properties={"branch": "main"}, + credentials={"token": "secret"}, + credential_type=CredentialType.UNAUTHORIZED, + ) + + +def request_log() -> RequestLog: + return RequestLog( + id="log1", + endpoint="/hooks/b1", + request={"headers": {}, "body": {"event": "push"}}, + response={"status": 200, "body": {"ok": True}}, + created_at=datetime(2024, 1, 1), + ) + + class TestTriggerProviderApis: @pytest.fixture def app(self, flask_app_with_containers: Flask) -> Flask: @@ -77,10 +121,10 @@ class TestTriggerProviderApis: app.test_request_context("/"), patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.get_trigger_provider", - return_value={"id": "p1"}, + return_value=trigger_provider(), ), ): - assert method(api, "t1", "github") == {"id": "p1"} + assert method(api, "t1", "github")["name"] == "github" class TestTriggerSubscriptionListApi: @@ -129,11 +173,11 @@ class TestTriggerSubscriptionBuilderApis: app.test_request_context("/", json={"credential_type": "UNAUTHORIZED"}), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", - return_value={"id": "b1"}, + return_value=subscription_builder(), ), ): result = method(api, "t1", mock_user(), "github") - assert "subscription_builder" in result + assert result["subscription_builder"]["id"] == "b1" def test_get_builder(self, app: Flask) -> None: api = TriggerSubscriptionBuilderGetApi() @@ -143,10 +187,10 @@ class TestTriggerSubscriptionBuilderApis: app.test_request_context("/"), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.get_subscription_builder_by_id", - return_value={"id": "b1"}, + return_value=subscription_builder(), ), ): - assert method(api, "github", "b1") == {"id": "b1"} + assert method(api, "github", "b1")["id"] == "b1" def test_verify_builder(self, app: Flask) -> None: api = TriggerSubscriptionBuilderVerifyApi() @@ -156,10 +200,10 @@ class TestTriggerSubscriptionBuilderApis: app.test_request_context("/", json={"credentials": {"a": 1}}), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_verify_builder", - return_value={"ok": True}, + return_value={"verified": True}, ), ): - assert method(api, "t1", mock_user(), "github", "b1") == {"ok": True} + assert method(api, "t1", mock_user(), "github", "b1") == {"verified": True} def test_verify_builder_error(self, app: Flask) -> None: api = TriggerSubscriptionBuilderVerifyApi() @@ -183,26 +227,24 @@ class TestTriggerSubscriptionBuilderApis: app.test_request_context("/", json={"name": "n"}), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_trigger_subscription_builder", - return_value={"id": "b1"}, + return_value=subscription_builder(), ), ): - assert method(api, "t1", "github", "b1") == {"id": "b1"} + assert method(api, "t1", "github", "b1")["id"] == "b1" def test_logs(self, app: Flask) -> None: api = TriggerSubscriptionBuilderLogsApi() method = unwrap(api.get) - log = MagicMock() - log.model_dump.return_value = {"a": 1} - with ( app.test_request_context("/"), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.list_logs", - return_value=[log], + return_value=[request_log()], ), ): - assert "logs" in method(api, "github", "b1") + result = method(api, "github", "b1") + assert result["logs"][0]["id"] == "log1" def test_build(self, app: Flask) -> None: api = TriggerSubscriptionBuilderBuildApi() @@ -215,7 +257,7 @@ class TestTriggerSubscriptionBuilderApis: return_value=None, ), ): - assert method(api, "t1", mock_user(), "github", "b1") == 200 + assert method(api, "t1", mock_user(), "github", "b1") == {"result": "success"} class TestTriggerSubscriptionCrud: @@ -239,7 +281,7 @@ class TestTriggerSubscriptionCrud: ), patch("controllers.console.workspace.trigger_providers.TriggerProviderService.update_trigger_subscription"), ): - assert method(api, "t1", "s1") == 200 + assert method(api, "t1", "s1") == {"result": "success"} def test_update_not_found(self, app: Flask) -> None: api = TriggerSubscriptionUpdateApi() @@ -275,7 +317,7 @@ class TestTriggerSubscriptionCrud: "controllers.console.workspace.trigger_providers.TriggerProviderService.rebuild_trigger_subscription" ), ): - assert method(api, "t1", "s1") == 200 + assert method(api, "t1", "s1") == {"result": "success"} def test_delete_subscription(self, app: Flask) -> None: api = TriggerSubscriptionDeleteApi() @@ -336,7 +378,7 @@ class TestTriggerOAuthApis: ), patch( "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", - return_value=MagicMock(id="b1"), + return_value=subscription_builder(), ), patch( "controllers.console.workspace.trigger_providers.OAuthProxyService.create_proxy_context", @@ -480,7 +522,7 @@ class TestTriggerOAuthClientManageApi: ), patch( "controllers.console.workspace.trigger_providers.TriggerManager.get_trigger_provider", - return_value=MagicMock(get_oauth_client_schema=lambda: {}), + return_value=MagicMock(get_oauth_client_schema=lambda: []), ), ): result = method(api, "t1", "github") @@ -494,10 +536,10 @@ class TestTriggerOAuthClientManageApi: app.test_request_context("/", json={"enabled": True}), patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.save_custom_oauth_client_params", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t1", "github") == {"ok": True} + assert method(api, "t1", "github") == {"result": "success"} def test_delete_client(self, app: Flask) -> None: api = TriggerOAuthClientManageApi() @@ -507,10 +549,10 @@ class TestTriggerOAuthClientManageApi: app.test_request_context("/"), patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.delete_custom_oauth_client_params", - return_value={"ok": True}, + return_value={"result": "success"}, ), ): - assert method(api, "t1", "github") == {"ok": True} + assert method(api, "t1", "github") == {"result": "success"} def test_oauth_client_post_value_error(self, app: Flask) -> None: api = TriggerOAuthClientManageApi() @@ -540,10 +582,10 @@ class TestTriggerSubscriptionVerifyApi: app.test_request_context("/", json={"credentials": {}}), patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.verify_subscription_credentials", - return_value={"ok": True}, + return_value={"verified": True}, ), ): - assert method(api, "t1", mock_user(), "github", "s1") == {"ok": True} + assert method(api, "t1", mock_user(), "github", "s1") == {"verified": True} @pytest.mark.parametrize("raised_exception", [ValueError("bad"), Exception("boom")]) def test_verify_errors(self, app: Flask, raised_exception: Exception) -> None: diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py index 99cabb6cea5..895f638380b 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py +++ b/api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py @@ -10,7 +10,13 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden from controllers.console.workspace import plugin_permission_required -from models.account import Tenant, TenantPluginPermission, TenantStatus +from models.account import ( + Tenant, + TenantPluginDebugPermission, + TenantPluginInstallPermission, + TenantPluginPermission, + TenantStatus, +) def _create_tenant(db_session: Session) -> Tenant: @@ -24,8 +30,8 @@ def _create_tenant(db_session: Session) -> Tenant: def _create_permission( db_session: Session, tenant_id: str, - install: TenantPluginPermission.InstallPermission = TenantPluginPermission.InstallPermission.EVERYONE, - debug: TenantPluginPermission.DebugPermission = TenantPluginPermission.DebugPermission.EVERYONE, + install: TenantPluginInstallPermission = TenantPluginInstallPermission.EVERYONE, + debug: TenantPluginDebugPermission = TenantPluginDebugPermission.EVERYONE, ) -> TenantPluginPermission: perm = TenantPluginPermission( tenant_id=tenant_id, @@ -59,8 +65,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.NOBODY, - debug=TenantPluginPermission.DebugPermission.EVERYONE, + install=TenantPluginInstallPermission.NOBODY, + debug=TenantPluginDebugPermission.EVERYONE, ) user = SimpleNamespace(is_admin_or_owner=True) @@ -81,8 +87,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.ADMINS, - debug=TenantPluginPermission.DebugPermission.EVERYONE, + install=TenantPluginInstallPermission.ADMINS, + debug=TenantPluginDebugPermission.EVERYONE, ) user = SimpleNamespace(is_admin_or_owner=False) @@ -103,8 +109,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.ADMINS, - debug=TenantPluginPermission.DebugPermission.EVERYONE, + install=TenantPluginInstallPermission.ADMINS, + debug=TenantPluginDebugPermission.EVERYONE, ) user = SimpleNamespace(is_admin_or_owner=True) @@ -124,8 +130,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.EVERYONE, - debug=TenantPluginPermission.DebugPermission.NOBODY, + install=TenantPluginInstallPermission.EVERYONE, + debug=TenantPluginDebugPermission.NOBODY, ) user = SimpleNamespace(is_admin_or_owner=True) @@ -146,8 +152,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.EVERYONE, - debug=TenantPluginPermission.DebugPermission.ADMINS, + install=TenantPluginInstallPermission.EVERYONE, + debug=TenantPluginDebugPermission.ADMINS, ) user = SimpleNamespace(is_admin_or_owner=False) @@ -168,8 +174,8 @@ class TestPluginPermissionRequired: _create_permission( db_session_with_containers, tenant.id, - install=TenantPluginPermission.InstallPermission.EVERYONE, - debug=TenantPluginPermission.DebugPermission.ADMINS, + install=TenantPluginInstallPermission.EVERYONE, + debug=TenantPluginDebugPermission.ADMINS, ) user = SimpleNamespace(is_admin_or_owner=True) diff --git a/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py b/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py index c281f071560..4c9d2e28116 100644 --- a/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py +++ b/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py @@ -8,7 +8,7 @@ from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from flask import Response +from flask import Flask, Response from pydantic import ValidationError import controllers.mcp.mcp as module @@ -37,12 +37,15 @@ class DummyServer: self.app_id = app_id self.tenant_id = tenant_id self.id = server_id + self.description = "Test server" + self.parameters_dict = {} class DummyApp: def __init__(self, mode, workflow=None, app_model_config=None): self.id = _APP_ID self.tenant_id = _TENANT_ID + self.name = "test_app" self.mode = mode self.workflow = workflow self.app_model_config = app_model_config @@ -494,3 +497,220 @@ class TestMCPAppApi: with pytest.raises(module.MCPRequestError) as exc_info: post_fn("server-1") assert "Invalid user_input_form" in str(exc_info.value) + + +_UNSUPPORTED_VERSION = "1999-01-01" + + +def _initialize_payload(protocol_version: str = "2024-11-05") -> dict[str, object]: + return { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": protocol_version, + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + + +def _tools_list_payload(request_id: int | None = 1) -> dict[str, object]: + payload: dict[str, object] = {"jsonrpc": "2.0", "method": "tools/list", "params": {}} + if request_id is not None: + payload["id"] = request_id + return payload + + +def _tools_call_payload() -> dict[str, object]: + return { + "jsonrpc": "2.0", + "method": "tools/call", + "id": 1, + "params": {"name": "test_app", "arguments": {"query": "test question"}}, + } + + +class TestMCPProtocolVersionNegotiationApi: + """MCP protocol version negotiation exercised through the HTTP controller layer. + + Covers the MCP-Protocol-Version header contract (resolution, rejection, threading) + and the serialized JSON responses seen by modern (2025-06-18) vs legacy (2024-11-05) + clients, including the back-compat guarantee that legacy responses carry none of the + structured-output fields. + """ + + def _make_api(self) -> module.MCPAppApi: + server = DummyServer(status=module.AppMCPServerStatus.ACTIVE) + app = DummyApp(mode=module.AppMode.CHAT, app_model_config=DummyConfig()) + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + api._retrieve_end_user = MagicMock(return_value=MagicMock()) + return api + + def _post( + self, flask_app: Flask, api: module.MCPAppApi, payload: dict[str, object], headers: dict[str, str] | None = None + ) -> Response: + fake_payload(payload) + post_fn = unwrap(api.post) + with flask_app.test_request_context(headers=headers): + return post_fn("server-1") + + @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) + def test_initialize_echoes_supported_body_version(self, flask_app_with_containers, version): + """Initialize echoes every supported client-requested version back unchanged.""" + api = self._make_api() + + response = self._post(flask_app_with_containers, api, _initialize_payload(version)) + + body = response.get_json() + assert body["result"]["protocolVersion"] == version + + def test_initialize_falls_back_for_unsupported_body_version(self, flask_app_with_containers): + """An unsupported requested version falls back to the server latest.""" + api = self._make_api() + + response = self._post(flask_app_with_containers, api, _initialize_payload(_UNSUPPORTED_VERSION)) + + body = response.get_json() + assert body["result"]["protocolVersion"] == module.mcp_types.SERVER_LATEST_PROTOCOL_VERSION + + def test_initialize_ignores_unsupported_header(self, flask_app_with_containers): + """Initialize negotiates via the request body, so its header is never rejected.""" + api = self._make_api() + + response = self._post( + flask_app_with_containers, + api, + _initialize_payload("2024-11-05"), + headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, + ) + + body = response.get_json() + assert "error" not in body + assert body["result"]["protocolVersion"] == "2024-11-05" + + @pytest.mark.parametrize("request_id", [5, None]) + def test_unsupported_header_returns_invalid_request_error(self, flask_app_with_containers, request_id): + """An unsupported header gets a JSON-RPC error echoing the request id (missing id -> null).""" + api = self._make_api() + + with patch.object(module, "handle_mcp_request", autospec=True) as mock_handle: + response = self._post( + flask_app_with_containers, + api, + _tools_list_payload(request_id=request_id), + headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, + ) + + body = response.get_json() + assert response.status_code == 200 + assert body["jsonrpc"] == "2.0" + assert body["id"] == request_id + assert body["error"]["code"] == module.mcp_types.INVALID_REQUEST + assert _UNSUPPORTED_VERSION in body["error"]["message"] + mock_handle.assert_not_called() + + def test_notification_with_unsupported_header_is_accepted(self, flask_app_with_containers): + """A notification is accepted (202, no body) even with an unsupported header.""" + api = self._make_api() + + response = self._post( + flask_app_with_containers, + api, + {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}, + headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, + ) + + assert response.status_code == 202 + + @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) + def test_supported_header_is_threaded_to_handler(self, flask_app_with_containers, version): + """Every supported header value is passed through to handle_mcp_request.""" + api = self._make_api() + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: + self._post( + flask_app_with_containers, + api, + _tools_list_payload(), + headers={"MCP-Protocol-Version": version}, + ) + + assert mock_handle.call_args.args[-1] == version + + def test_absent_header_defaults_to_back_compat_version(self, flask_app_with_containers): + """An absent header resolves to the spec's default version (2025-03-26).""" + api = self._make_api() + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: + self._post(flask_app_with_containers, api, _tools_list_payload()) + + assert mock_handle.call_args.args[-1] == module.mcp_types.DEFAULT_NEGOTIATED_VERSION + + def test_tools_list_json_advertises_structured_output_for_modern_client(self, flask_app_with_containers): + """A 2025-06-18 client sees outputSchema and title in the serialized tool JSON.""" + api = self._make_api() + + response = self._post( + flask_app_with_containers, + api, + _tools_list_payload(), + headers={"MCP-Protocol-Version": "2025-06-18"}, + ) + + tool = response.get_json()["result"]["tools"][0] + assert tool["outputSchema"] == {"type": "object"} + assert tool["title"] == "test_app" + + def test_tools_list_json_unchanged_for_legacy_client(self, flask_app_with_containers): + """A 2024-11-05 client sees exactly the pre-upgrade tool JSON keys.""" + api = self._make_api() + + response = self._post( + flask_app_with_containers, + api, + _tools_list_payload(), + headers={"MCP-Protocol-Version": "2024-11-05"}, + ) + + tool = response.get_json()["result"]["tools"][0] + assert set(tool) == {"name", "description", "inputSchema"} + + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_tools_call_json_includes_structured_content_for_modern_client( + self, mock_app_generate, flask_app_with_containers + ): + """A 2025-06-18 client receives structuredContent alongside the text content.""" + api = self._make_api() + mock_app_generate.generate.return_value = {"answer": "test answer"} + + response = self._post( + flask_app_with_containers, + api, + _tools_call_payload(), + headers={"MCP-Protocol-Version": "2025-06-18"}, + ) + + result = response.get_json()["result"] + assert result["structuredContent"] == {"answer": "test answer"} + assert result["content"][0]["text"] == "test answer" + + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_tools_call_json_omits_structured_content_for_legacy_client( + self, mock_app_generate, flask_app_with_containers + ): + """A 2024-11-05 client receives the pre-upgrade tools/call JSON without structuredContent.""" + api = self._make_api() + mock_app_generate.generate.return_value = {"answer": "test answer"} + + response = self._post( + flask_app_with_containers, + api, + _tools_call_payload(), + headers={"MCP-Protocol-Version": "2024-11-05"}, + ) + + result = response.get_json()["result"] + assert "structuredContent" not in result + assert result["content"][0]["text"] == "test answer" diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py index 4cdbec3e30e..4222f49a28a 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py @@ -30,7 +30,6 @@ def _mint_account_token( ) -> MintResult: """Mint a real, persisted ``dfoa_`` access token for ``account``.""" return mint_oauth_token( - db_session, redis_client, subject_email=account.email, subject_issuer=None, @@ -39,6 +38,7 @@ def _mint_account_token( device_label=device_label, prefix=PREFIX_OAUTH_ACCOUNT, ttl_days=14, + session=db_session, ) diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py index 93e8927cfef..4d9bfb5ea17 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py @@ -96,7 +96,7 @@ def _app_and_account(db_session: Session, *, mode: str = "chat") -> tuple[App, A api_rph=100, api_rpm=10, ) - app_model = AppService().create_app(tenant.id, app_args, account) + app_model = AppService().create_app(tenant.id, app_args, account, session=db_session) return app_model, account @@ -167,7 +167,7 @@ class TestDslImportConfirm: api = AppDslImportConfirmApi() with app.test_request_context( - f"/openapi/v1/workspaces/{tenant.id}/apps/imports/{import_id}/confirm", method="POST" + f"/openapi/v1/workspaces/{tenant.id}/apps/imports/{import_id}:confirm", method="POST" ): result, code = unwrap(api.post)( api, workspace_id=tenant.id, import_id=import_id, auth_data=auth_for(account) @@ -198,7 +198,7 @@ class TestDslExport: db_session_with_containers.commit() api = AppDslExportApi() - with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/export"): + with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/dsl"): response, code = unwrap(api.get)( api, app_id=app_model.id, auth_data=auth_for(account, app_model=app_model), query=AppDslExportQuery() ) @@ -216,7 +216,7 @@ class TestDslExport: app_model, account = _app_and_account(db_session_with_containers, mode="workflow") api = AppDslExportApi() - with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/export"): + with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/dsl"): result, code = unwrap(api.get)( api, app_id=app_model.id, auth_data=auth_for(account, app_model=app_model), query=AppDslExportQuery() ) @@ -232,7 +232,7 @@ class TestDslCheckDependencies: app_model, account = _app_and_account(db_session_with_containers, mode="chat") api = AppDslCheckDependenciesApi() - with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/check-dependencies"): + with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/dependencies:check"): result, code = unwrap(api.get)(api, app_id=app_model.id, auth_data=auth_for(account, app_model=app_model)) assert code == 200 diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py index c6fde623677..df4f3873b1e 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py @@ -24,7 +24,7 @@ def _create_app(db_session: Session, account: Account, *, name: str = "Runner") icon="🤖", icon_background="#FF6B6B", ) - app_model = AppService().create_app(tenant.id, params, account) + app_model = AppService().create_app(tenant.id, params, account, session=db_session) db_session.commit() return app_model @@ -38,7 +38,7 @@ class TestAppRunTaskStop: task_id = str(uuid4()) api = AppRunTaskStopApi() - with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/tasks/{task_id}/stop", method="POST"): + with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/tasks/{task_id}:stop", method="POST"): result = unwrap(api.post)( api, app_id=app_model.id, diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py index 22f812e125b..ce1425d9e61 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py @@ -39,7 +39,7 @@ def _create_app( icon="🤖", icon_background="#FF6B6B", ) - app_model = AppService().create_app(tenant.id, params, account) + app_model = AppService().create_app(tenant.id, params, account, session=db_session) # The openapi surface gate keys off ``enable_api``; flip it explicitly so # the test states the visibility precondition rather than relying on the # template default. @@ -120,7 +120,7 @@ class TestAppDescribe: app_model = _create_app(db_session_with_containers, account, name="Describe Me", enable_api=True) api = AppDescribeApi() - with app.test_request_context(f"/openapi/v1/apps/{app_model.id}/describe?fields=info"): + with app.test_request_context(f"/openapi/v1/apps/{app_model.id}?fields=info"): result = unwrap(api.get)( api, app_id=app_model.id, auth_data=auth_for(account), query=AppDescribeQuery(fields="info") ) @@ -138,7 +138,7 @@ class TestAppDescribe: missing_id = str(uuid4()) api = AppDescribeApi() - with app.test_request_context(f"/openapi/v1/apps/{missing_id}/describe"): + with app.test_request_context(f"/openapi/v1/apps/{missing_id}"): with pytest.raises(NotFound): unwrap(api.get)(api, app_id=missing_id, auth_data=auth_for(account), query=AppDescribeQuery()) @@ -151,6 +151,6 @@ class TestAppDescribe: hidden = _create_app(db_session_with_containers, account, name="Hidden", enable_api=False) api = AppDescribeApi() - with app.test_request_context(f"/openapi/v1/apps/{hidden.id}/describe"): + with app.test_request_context(f"/openapi/v1/apps/{hidden.id}"): with pytest.raises(NotFound): unwrap(api.get)(api, app_id=hidden.id, auth_data=auth_for(account), query=AppDescribeQuery()) diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py index b90d5ab907c..31a5485d2b3 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py @@ -25,7 +25,7 @@ def _create_app(db_session: Session, account: Account, *, name: str = "Uploader" icon="🤖", icon_background="#FF6B6B", ) - app_model = AppService().create_app(tenant.id, params, account) + app_model = AppService().create_app(tenant.id, params, account, session=db_session) db_session.commit() return app_model @@ -43,7 +43,7 @@ class TestAppFileUpload: api = AppFileUploadApi() data = {"file": (BytesIO(content), "note.txt", "text/plain")} with app.test_request_context( - f"/openapi/v1/apps/{app_model.id}/files/upload", + f"/openapi/v1/apps/{app_model.id}/files", method="POST", data=data, content_type="multipart/form-data", diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py b/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py index 18075704325..5e794ae1982 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py +++ b/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py @@ -95,7 +95,7 @@ class TestWorkspaceSwitch: ) api = WorkspaceSwitchApi() - with app.test_request_context(f"/openapi/v1/workspaces/{target.id}/switch", method="POST"): + with app.test_request_context(f"/openapi/v1/workspaces/{target.id}:switch", method="POST"): detail = unwrap(api.post)(api, workspace_id=target.id, auth_data=auth_for(account)) # Response reflects the post-switch state. @@ -118,6 +118,6 @@ class TestWorkspaceSwitch: assert outsider_ws is not None api = WorkspaceSwitchApi() - with app.test_request_context(f"/openapi/v1/workspaces/{outsider_ws.id}/switch", method="POST"): + with app.test_request_context(f"/openapi/v1/workspaces/{outsider_ws.id}:switch", method="POST"): with pytest.raises(NotFound): unwrap(api.post)(api, workspace_id=outsider_ws.id, auth_data=auth_for(account)) diff --git a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py b/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py index 372157813cc..d670425be0c 100644 --- a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py +++ b/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py @@ -734,7 +734,8 @@ class TestDatasetApiPatch: assert response["name"] == "Updated Dataset" assert response["partial_member_list"] == ["user-1"] mock_dataset_svc.update_dataset.assert_called_once() - session, _, update_data, _ = mock_dataset_svc.update_dataset.call_args.args + _, update_data, _ = mock_dataset_svc.update_dataset.call_args.args + session = mock_dataset_svc.update_dataset.call_args.kwargs["session"] assert isinstance(session, (Session, scoped_session)) assert update_data["name"] == "Updated Dataset" assert update_data["permission"] == "partial_members" @@ -1013,7 +1014,7 @@ class TestDatasetTagsApiGet: assert status == 200 assert response == [{"id": "tag-1", "name": "Test Tag", "type": "knowledge", "binding_count": "0"}] - mock_tag_svc.get_tags.assert_called_once_with(SessionMatcher(), "knowledge", "tenant-1") + mock_tag_svc.get_tags.assert_called_once_with("knowledge", "tenant-1", session=SessionMatcher()) @patch("controllers.service_api.dataset.dataset.current_user") def test_list_tags_from_db( diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_web_forgot_password.py b/api/tests/test_containers_integration_tests/controllers/web/test_web_forgot_password.py index d568a1c0b04..cd754782df4 100644 --- a/api/tests/test_containers_integration_tests/controllers/web/test_web_forgot_password.py +++ b/api/tests/test_containers_integration_tests/controllers/web/test_web_forgot_password.py @@ -57,7 +57,7 @@ class TestForgotPasswordSendEmailApi: response = ForgotPasswordSendEmailApi().post() assert response == {"result": "success", "data": "token-123"} - mock_get_account.assert_called_once_with(ANY, "User@Example.com") + mock_get_account.assert_called_once_with("User@Example.com", session=ANY) mock_send_mail.assert_called_once_with(account=mock_account, email="user@example.com", language="zh-Hans") mock_extract_ip.assert_called_once() mock_rate_limit.assert_called_once_with("127.0.0.1") @@ -177,7 +177,7 @@ class TestForgotPasswordResetApi: response = ForgotPasswordResetApi().post() assert response == {"result": "success"} - mock_get_account.assert_called_once_with(ANY, "User@Example.com") + mock_get_account.assert_called_once_with("User@Example.com", session=ANY) mock_update_account.assert_called_once() mock_revoke_token.assert_called_once_with("token-123") diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_wraps.py b/api/tests/test_containers_integration_tests/controllers/web/test_wraps.py index 3eab8ccbee5..aa85ac2ca7b 100644 --- a/api/tests/test_containers_integration_tests/controllers/web/test_wraps.py +++ b/api/tests/test_containers_integration_tests/controllers/web/test_wraps.py @@ -19,6 +19,8 @@ from controllers.web.wraps import ( decode_jwt_token, ) +pytestmark = pytest.mark.usefixtures("db_session_with_containers") + class TestValidateWebappToken: def test_enterprise_enabled_and_app_auth_requires_webapp_source(self) -> None: diff --git a/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py b/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py index 66b3392a4b4..84f01ea52ee 100644 --- a/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py +++ b/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py @@ -20,6 +20,7 @@ providing more reliable and realistic test scenarios than mocks. import json import uuid from time import time +from unittest.mock import Mock import pytest from sqlalchemy import Engine, delete, select @@ -35,6 +36,7 @@ from core.workflow.system_variables import build_system_variables from extensions.ext_storage import storage from graphon.entities.pause_reason import SchedulingPause from graphon.enums import WorkflowExecutionStatus +from graphon.filters import GraphEventFilterContext, ResponseStreamFilter from graphon.graph_engine.entities.commands import GraphEngineCommand from graphon.graph_engine.layers.base import GraphEngineLayerNotInitializedError from graphon.graph_events import GraphRunPausedEvent @@ -49,6 +51,22 @@ from services.file_service import FileService from services.workflow_run_service import WorkflowRunService +def _create_initialized_response_stream_filter() -> ResponseStreamFilter: + """Build a `ResponseStreamFilter` that has already run `initialize()`. + + `ResponseStreamFilter.dumps()` raises `RuntimeError` unless the filter has + processed a `GraphEventFilterContext` first. In production this always + happens before any event (including `GraphRunPausedEvent`) reaches + `PauseStatePersistenceLayer.on_event`, so tests that exercise `on_event` + or a subsequent `dumps()` call need a filter in that same state. A + nodeless graph is enough to satisfy the precondition. + """ + response_stream_filter = ResponseStreamFilter() + context = GraphEventFilterContext(graph=Mock(nodes={}), runtime_state=Mock()) + response_stream_filter.initialize(context) + return response_stream_filter + + class _TestCommandChannelImpl: """Real implementation of CommandChannel for testing.""" @@ -295,6 +313,7 @@ class TestPauseStatePersistenceLayerTestContainers: session_factory=self.session.get_bind(), state_owner_user_id=owner_id, generate_entity=entity, + response_stream_filter=_create_initialized_response_stream_filter(), ) def test_complete_pause_flow_with_real_dependencies(self, db_session_with_containers: Session): diff --git a/api/tests/test_containers_integration_tests/services/auth/test_api_key_auth_service.py b/api/tests/test_containers_integration_tests/services/auth/test_api_key_auth_service.py index e2f8c8fc703..e22aa102328 100644 --- a/api/tests/test_containers_integration_tests/services/auth/test_api_key_auth_service.py +++ b/api/tests/test_containers_integration_tests/services/auth/test_api_key_auth_service.py @@ -51,7 +51,7 @@ class TestApiKeyAuthService: self._create_binding(db_session_with_containers, tenant_id=tenant_id, category=category, provider=provider) db_session_with_containers.expire_all() - result = ApiKeyAuthService.get_provider_auth_list(db_session_with_containers, tenant_id) + result = ApiKeyAuthService.get_provider_auth_list(tenant_id, session=db_session_with_containers) assert len(result) >= 1 tenant_results = [r for r in result if r.tenant_id == tenant_id] @@ -61,7 +61,7 @@ class TestApiKeyAuthService: def test_get_provider_auth_list_empty( self, flask_app_with_containers: Flask, db_session_with_containers: Session, tenant_id ): - result = ApiKeyAuthService.get_provider_auth_list(db_session_with_containers, tenant_id) + result = ApiKeyAuthService.get_provider_auth_list(tenant_id, session=db_session_with_containers) tenant_results = [r for r in result if r.tenant_id == tenant_id] assert tenant_results == [] @@ -74,7 +74,7 @@ class TestApiKeyAuthService: ) db_session_with_containers.expire_all() - result = ApiKeyAuthService.get_provider_auth_list(db_session_with_containers, tenant_id) + result = ApiKeyAuthService.get_provider_auth_list(tenant_id, session=db_session_with_containers) tenant_results = [r for r in result if r.tenant_id == tenant_id] assert tenant_results == [] @@ -95,7 +95,7 @@ class TestApiKeyAuthService: mock_factory.return_value = mock_auth_instance mock_encrypter.encrypt_token.return_value = "encrypted_test_key_123" - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=db_session_with_containers) mock_factory.assert_called_once() mock_auth_instance.validate_credentials.assert_called_once() @@ -118,7 +118,7 @@ class TestApiKeyAuthService: mock_auth_instance.validate_credentials.return_value = False mock_factory.return_value = mock_auth_instance - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=db_session_with_containers) db_session_with_containers.expire_all() bindings = db_session_with_containers.query(DataSourceApiKeyAuthBinding).filter_by(tenant_id=tenant_id).all() @@ -142,7 +142,7 @@ class TestApiKeyAuthService: original_key = mock_args["credentials"]["config"]["api_key"] - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=db_session_with_containers) assert mock_args["credentials"]["config"]["api_key"] == "encrypted_test_key_123" assert mock_args["credentials"]["config"]["api_key"] != original_key @@ -166,14 +166,18 @@ class TestApiKeyAuthService: ) db_session_with_containers.expire_all() - result = ApiKeyAuthService.get_auth_credentials(db_session_with_containers, tenant_id, category, provider) + result = ApiKeyAuthService.get_auth_credentials( + tenant_id, category, provider, session=db_session_with_containers + ) assert result == mock_credentials def test_get_auth_credentials_not_found( self, flask_app_with_containers: Flask, db_session_with_containers: Session, tenant_id, category, provider ): - result = ApiKeyAuthService.get_auth_credentials(db_session_with_containers, tenant_id, category, provider) + result = ApiKeyAuthService.get_auth_credentials( + tenant_id, category, provider, session=db_session_with_containers + ) assert result is None @@ -190,7 +194,9 @@ class TestApiKeyAuthService: ) db_session_with_containers.expire_all() - result = ApiKeyAuthService.get_auth_credentials(db_session_with_containers, tenant_id, category, provider) + result = ApiKeyAuthService.get_auth_credentials( + tenant_id, category, provider, session=db_session_with_containers + ) assert result == special_credentials assert result["config"]["api_key"] == "key_with_中文_and_special_chars_!@#$%" @@ -204,7 +210,7 @@ class TestApiKeyAuthService: binding_id = binding.id db_session_with_containers.expire_all() - ApiKeyAuthService.delete_provider_auth(db_session_with_containers, tenant_id, binding_id) + ApiKeyAuthService.delete_provider_auth(tenant_id, binding_id, session=db_session_with_containers) db_session_with_containers.expire_all() remaining = db_session_with_containers.query(DataSourceApiKeyAuthBinding).filter_by(id=binding_id).first() @@ -214,7 +220,7 @@ class TestApiKeyAuthService: self, flask_app_with_containers: Flask, db_session_with_containers: Session, tenant_id ): # Should not raise when binding not found - ApiKeyAuthService.delete_provider_auth(db_session_with_containers, tenant_id, str(uuid4())) + ApiKeyAuthService.delete_provider_auth(tenant_id, str(uuid4()), session=db_session_with_containers) def test_validate_api_key_auth_args_success(self, mock_args): ApiKeyAuthService.validate_api_key_auth_args(mock_args) @@ -291,13 +297,13 @@ class TestApiKeyAuthService: mock_session = MagicMock() mock_session.commit.side_effect = Exception("Database error") with pytest.raises(Exception, match="Database error"): - ApiKeyAuthService.create_provider_auth(mock_session, tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=mock_session) @patch("services.auth.api_key_auth_service.ApiKeyAuthFactory") def test_create_provider_auth_factory_exception(self, mock_factory: MagicMock, tenant_id, mock_args): mock_factory.side_effect = Exception("Factory error") with pytest.raises(Exception, match="Factory error"): - ApiKeyAuthService.create_provider_auth(MagicMock(), tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=MagicMock()) @patch("services.auth.api_key_auth_service.ApiKeyAuthFactory") @patch("services.auth.api_key_auth_service.encrypter") @@ -307,7 +313,7 @@ class TestApiKeyAuthService: mock_factory.return_value = mock_auth_instance mock_encrypter.encrypt_token.side_effect = Exception("Encryption error") with pytest.raises(Exception, match="Encryption error"): - ApiKeyAuthService.create_provider_auth(MagicMock(), tenant_id, mock_args) + ApiKeyAuthService.create_provider_auth(tenant_id, mock_args, session=MagicMock()) def test_validate_api_key_auth_args_none_input(self): with pytest.raises(TypeError): diff --git a/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py b/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py index 9b86ab41f2b..cd3ed01cbfa 100644 --- a/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py +++ b/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py @@ -57,7 +57,7 @@ class TestAuthIntegration: mock_encrypt.return_value = "encrypted_fc_test_key_123" args = {"category": category, "provider": AuthType.FIRECRAWL, "credentials": firecrawl_credentials} - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id_1, args) + ApiKeyAuthService.create_provider_auth(tenant_id_1, args, session=db_session_with_containers) mock_http.assert_called_once() call_args = mock_http.call_args @@ -101,15 +101,15 @@ class TestAuthIntegration: mock_encrypt.return_value = "encrypted_key" args1 = {"category": category, "provider": AuthType.FIRECRAWL, "credentials": firecrawl_credentials} - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id_1, args1) + ApiKeyAuthService.create_provider_auth(tenant_id_1, args1, session=db_session_with_containers) args2 = {"category": category, "provider": AuthType.JINA, "credentials": jina_credentials} - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id_2, args2) + ApiKeyAuthService.create_provider_auth(tenant_id_2, args2, session=db_session_with_containers) db_session_with_containers.expire_all() - result1 = ApiKeyAuthService.get_provider_auth_list(db_session_with_containers, tenant_id_1) - result2 = ApiKeyAuthService.get_provider_auth_list(db_session_with_containers, tenant_id_2) + result1 = ApiKeyAuthService.get_provider_auth_list(tenant_id_1, session=db_session_with_containers) + result2 = ApiKeyAuthService.get_provider_auth_list(tenant_id_2, session=db_session_with_containers) assert len(result1) == 1 assert result1[0].tenant_id == tenant_id_1 @@ -120,7 +120,7 @@ class TestAuthIntegration: self, flask_app_with_containers: Flask, db_session_with_containers: Session, tenant_id_2, category ): result = ApiKeyAuthService.get_auth_credentials( - db_session_with_containers, tenant_id_2, category, AuthType.FIRECRAWL + tenant_id_2, category, AuthType.FIRECRAWL, session=db_session_with_containers ) assert result is None @@ -163,7 +163,7 @@ class TestAuthIntegration: "provider": AuthType.FIRECRAWL, "credentials": {"auth_type": "bearer", "config": {"api_key": "fc_test_key_123"}}, } - ApiKeyAuthService.create_provider_auth(db.session(), tenant_id_1, thread_args) + ApiKeyAuthService.create_provider_auth(tenant_id_1, thread_args, session=db.session()) results.append("success") except Exception as e: exceptions.append(e) @@ -216,7 +216,7 @@ class TestAuthIntegration: args = {"category": category, "provider": AuthType.FIRECRAWL, "credentials": firecrawl_credentials} with pytest.raises(httpx.RequestError): - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id_1, args) + ApiKeyAuthService.create_provider_auth(tenant_id_1, args, session=db_session_with_containers) db_session_with_containers.expire_all() bindings = db_session_with_containers.query(DataSourceApiKeyAuthBinding).filter_by(tenant_id=tenant_id_1).all() @@ -253,12 +253,12 @@ class TestAuthIntegration: mock_encrypt.return_value = "encrypted_key" args = {"category": category, "provider": AuthType.FIRECRAWL, "credentials": firecrawl_credentials} - ApiKeyAuthService.create_provider_auth(db_session_with_containers, tenant_id_1, args) + ApiKeyAuthService.create_provider_auth(tenant_id_1, args, session=db_session_with_containers) db_session_with_containers.expire_all() result = ApiKeyAuthService.get_auth_credentials( - db_session_with_containers, tenant_id_1, category, AuthType.FIRECRAWL + tenant_id_1, category, AuthType.FIRECRAWL, session=db_session_with_containers ) assert result is not None assert result["config"]["api_key"] == "encrypted_key" diff --git a/api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py b/api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py index 646a0592630..0a34733adeb 100644 --- a/api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py +++ b/api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py @@ -6,7 +6,7 @@ Redis queuing, error handling, and community vs enterprise behavior. from __future__ import annotations -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -118,7 +118,7 @@ class TestSyncAccountDeletion: with patch("services.enterprise.account_deletion_sync.dify_config") as mock_config: mock_config.ENTERPRISE_ENABLED = False - result = sync_account_deletion(account_id=str(uuid4()), source="account_deleted") + result = sync_account_deletion(account_id=str(uuid4()), source="account_deleted", session=MagicMock()) assert result is True mock_queue_task.assert_not_called() @@ -137,7 +137,9 @@ class TestSyncAccountDeletion: with patch("services.enterprise.account_deletion_sync.dify_config") as mock_config: mock_config.ENTERPRISE_ENABLED = True - result = sync_account_deletion(account_id=account_id, source="account_deleted") + result = sync_account_deletion( + account_id=account_id, source="account_deleted", session=db_session_with_containers + ) assert result is True assert mock_queue_task.call_count == 3 @@ -151,7 +153,9 @@ class TestSyncAccountDeletion: with patch("services.enterprise.account_deletion_sync.dify_config") as mock_config: mock_config.ENTERPRISE_ENABLED = True - result = sync_account_deletion(account_id=str(uuid4()), source="account_deleted") + result = sync_account_deletion( + account_id=str(uuid4()), source="account_deleted", session=db_session_with_containers + ) assert result is True mock_queue_task.assert_not_called() @@ -176,7 +180,9 @@ class TestSyncAccountDeletion: with patch("services.enterprise.account_deletion_sync.dify_config") as mock_config: mock_config.ENTERPRISE_ENABLED = True - result = sync_account_deletion(account_id=account_id, source="account_deleted") + result = sync_account_deletion( + account_id=account_id, source="account_deleted", session=db_session_with_containers + ) assert result is False assert mock_queue_task.call_count == 3 @@ -196,7 +202,9 @@ class TestSyncAccountDeletion: with patch("services.enterprise.account_deletion_sync.dify_config") as mock_config: mock_config.ENTERPRISE_ENABLED = True - result = sync_account_deletion(account_id=account_id, source="account_deleted") + result = sync_account_deletion( + account_id=account_id, source="account_deleted", session=db_session_with_containers + ) assert result is False mock_queue_task.assert_called_once() diff --git a/api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py b/api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py index e7cf04c0d8b..a52458ac972 100644 --- a/api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py +++ b/api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py @@ -2,11 +2,10 @@ from __future__ import annotations from uuid import uuid4 -import pytest from sqlalchemy import func, select from sqlalchemy.orm import Session -from models.account import TenantPluginPermission +from models.account import TenantPluginDebugPermission, TenantPluginInstallPermission, TenantPluginPermission from services.plugin.plugin_permission_service import PluginPermissionService @@ -32,23 +31,22 @@ class TestGetPermission: tenant_id = _tenant_id() permission = TenantPluginPermission( tenant_id=tenant_id, - install_permission=TenantPluginPermission.InstallPermission.ADMINS, - debug_permission=TenantPluginPermission.DebugPermission.EVERYONE, + install_permission=TenantPluginInstallPermission.ADMINS, + debug_permission=TenantPluginDebugPermission.EVERYONE, ) db_session_with_containers.add(permission) db_session_with_containers.commit() - result = PluginPermissionService.get_permission(tenant_id) + result = PluginPermissionService.get_permission(tenant_id, session=db_session_with_containers) assert result is not None assert result.id == permission.id assert result.tenant_id == tenant_id - assert result.install_permission == TenantPluginPermission.InstallPermission.ADMINS - assert result.debug_permission == TenantPluginPermission.DebugPermission.EVERYONE + assert result.install_permission == TenantPluginInstallPermission.ADMINS + assert result.debug_permission == TenantPluginDebugPermission.EVERYONE - @pytest.mark.usefixtures("flask_app_with_containers") - def test_returns_none_when_not_found(self) -> None: - result = PluginPermissionService.get_permission(_tenant_id()) + def test_returns_none_when_not_found(self, db_session_with_containers: Session) -> None: + result = PluginPermissionService.get_permission(_tenant_id(), session=db_session_with_containers) assert result is None @@ -61,36 +59,38 @@ class TestChangePermission: result = PluginPermissionService.change_permission( tenant_id, - TenantPluginPermission.InstallPermission.EVERYONE, - TenantPluginPermission.DebugPermission.EVERYONE, + TenantPluginInstallPermission.EVERYONE, + TenantPluginDebugPermission.EVERYONE, + session=db_session_with_containers, ) permission = _get_permission(db_session_with_containers, tenant_id) assert result is True assert permission is not None - assert permission.install_permission == TenantPluginPermission.InstallPermission.EVERYONE - assert permission.debug_permission == TenantPluginPermission.DebugPermission.EVERYONE + assert permission.install_permission == TenantPluginInstallPermission.EVERYONE + assert permission.debug_permission == TenantPluginDebugPermission.EVERYONE def test_updates_existing_permission(self, db_session_with_containers: Session) -> None: tenant_id = _tenant_id() existing = TenantPluginPermission( tenant_id=tenant_id, - install_permission=TenantPluginPermission.InstallPermission.EVERYONE, - debug_permission=TenantPluginPermission.DebugPermission.EVERYONE, + install_permission=TenantPluginInstallPermission.EVERYONE, + debug_permission=TenantPluginDebugPermission.EVERYONE, ) db_session_with_containers.add(existing) db_session_with_containers.commit() result = PluginPermissionService.change_permission( tenant_id, - TenantPluginPermission.InstallPermission.ADMINS, - TenantPluginPermission.DebugPermission.ADMINS, + TenantPluginInstallPermission.ADMINS, + TenantPluginDebugPermission.ADMINS, + session=db_session_with_containers, ) permission = _get_permission(db_session_with_containers, tenant_id) assert result is True assert permission is not None assert permission.id == existing.id - assert permission.install_permission == TenantPluginPermission.InstallPermission.ADMINS - assert permission.debug_permission == TenantPluginPermission.DebugPermission.ADMINS + assert permission.install_permission == TenantPluginInstallPermission.ADMINS + assert permission.debug_permission == TenantPluginDebugPermission.ADMINS assert _count_permissions(db_session_with_containers, tenant_id) == 1 diff --git a/api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py b/api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py index 2e7df67d266..75d127ce6b2 100644 --- a/api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py +++ b/api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py @@ -42,7 +42,9 @@ class TestRagPipelineServiceGetPipeline: yield db_session_with_containers.rollback() - def _make_service(self, flask_app_with_containers: Flask) -> RagPipelineService: + def _make_service( + self, flask_app_with_containers: Flask, db_session_with_containers: Session + ) -> RagPipelineService: with ( patch( "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository", @@ -54,7 +56,7 @@ class TestRagPipelineServiceGetPipeline: ), ): session_factory = sessionmaker(bind=flask_app_with_containers.extensions["sqlalchemy"].engine) - return RagPipelineService(session_maker=session_factory) + return RagPipelineService(db_session_with_containers, session_maker=session_factory) def _create_pipeline(self, db_session: Session, tenant_id: str, created_by: str) -> Pipeline: pipeline = Pipeline( @@ -85,7 +87,7 @@ class TestRagPipelineServiceGetPipeline: self, db_session_with_containers: Session, flask_app_with_containers: Flask ) -> None: """get_pipeline raises ValueError when dataset does not exist.""" - service = self._make_service(flask_app_with_containers) + service = self._make_service(flask_app_with_containers, db_session_with_containers) with pytest.raises(ValueError, match="Dataset not found"): service.get_pipeline(tenant_id=str(uuid4()), dataset_id=str(uuid4())) @@ -99,10 +101,10 @@ class TestRagPipelineServiceGetPipeline: dataset = self._create_dataset(db_session_with_containers, tenant_id, created_by, pipeline_id=None) db_session_with_containers.flush() - service = self._make_service(flask_app_with_containers) + service = self._make_service(flask_app_with_containers, db_session_with_containers) with pytest.raises(ValueError, match="Pipeline not found"): - service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id, session=db_session_with_containers) + service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id) def test_get_pipeline_returns_pipeline_when_found( self, db_session_with_containers: Session, flask_app_with_containers: Flask @@ -115,9 +117,9 @@ class TestRagPipelineServiceGetPipeline: dataset = self._create_dataset(db_session_with_containers, tenant_id, created_by, pipeline_id=pipeline.id) db_session_with_containers.flush() - service = self._make_service(flask_app_with_containers) + service = self._make_service(flask_app_with_containers, db_session_with_containers) - result = service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id, session=db_session_with_containers) + result = service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset.id) assert result.id == pipeline.id @@ -185,7 +187,9 @@ class TestUpdateCustomizedPipelineTemplate: icon_info=IconInfo(icon="📄"), ) with pytest.raises(ValueError, match="Customized pipeline template not found"): - RagPipelineService.update_customized_pipeline_template(str(uuid4()), info, account, tenant_id) + RagPipelineService.update_customized_pipeline_template( + str(uuid4()), info, account, tenant_id, session=db_session_with_containers + ) def test_update_template_raises_on_duplicate_name( self, db_session_with_containers: Session, flask_app_with_containers: Flask @@ -264,4 +268,6 @@ class TestDeleteCustomizedPipelineTemplate: tenant_id = str(uuid4()) with pytest.raises(ValueError, match="Customized pipeline template not found"): - RagPipelineService.delete_customized_pipeline_template(str(uuid4()), tenant_id) + RagPipelineService.delete_customized_pipeline_template( + str(uuid4()), tenant_id, session=db_session_with_containers + ) diff --git a/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py b/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py index 0f7c790ba14..1c366d3ee32 100644 --- a/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py +++ b/api/tests/test_containers_integration_tests/services/recommend_app/test_database_retrieval.py @@ -1,6 +1,6 @@ from __future__ import annotations -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 from flask import Flask @@ -82,8 +82,8 @@ class TestDatabaseRecommendAppRetrieval: "fetch_recommended_apps_from_db", return_value={"recommended_apps": [], "categories": []}, ) as mock_fetch: - result = DatabaseRecommendAppRetrieval().get_recommended_apps_and_categories("en-US") - mock_fetch.assert_called_once_with("en-US") + result = DatabaseRecommendAppRetrieval().get_recommended_apps_and_categories("en-US", session=MagicMock()) + mock_fetch.assert_called_once() assert result == {"recommended_apps": [], "categories": []} def test_get_recommend_app_detail_delegates(self): @@ -92,8 +92,8 @@ class TestDatabaseRecommendAppRetrieval: "fetch_recommended_app_detail_from_db", return_value={"id": "app-1"}, ) as mock_fetch: - result = DatabaseRecommendAppRetrieval().get_recommend_app_detail("app-1") - mock_fetch.assert_called_once_with("app-1") + result = DatabaseRecommendAppRetrieval().get_recommend_app_detail("app-1", session=MagicMock()) + mock_fetch.assert_called_once() assert result == {"id": "app-1"} @@ -112,7 +112,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "en-US", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert app1.id in app_ids @@ -135,7 +137,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "en-US", session=db_session_with_containers + ) recommended_app = next(item for item in result["recommended_apps"] if item["app_id"] == created_app.id) assert recommended_app["categories"] == ["writing", "assistant"] @@ -160,7 +164,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "en-US", session=db_session_with_containers + ) recommended_app = next(item for item in result["recommended_apps"] if item["app_id"] == created_app.id) assert "category" not in recommended_app @@ -177,7 +183,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("fr-FR") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "fr-FR", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert app1.id in app_ids @@ -190,7 +198,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "en-US", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert app1.id not in app_ids @@ -202,7 +212,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_recommended_apps_from_db( + "en-US", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert app1.id not in app_ids @@ -235,7 +247,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db("en-US") + result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db( + "en-US", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert learn_dify_app.id in app_ids @@ -261,7 +275,9 @@ class TestFetchRecommendedAppsFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db("fr-FR") + result = DatabaseRecommendAppRetrieval.fetch_learn_dify_apps_from_db( + "fr-FR", session=db_session_with_containers + ) app_ids = {r["app_id"] for r in result["recommended_apps"]} assert learn_dify_app.id in app_ids @@ -269,7 +285,9 @@ class TestFetchRecommendedAppsFromDb: class TestFetchRecommendedAppDetailFromDb: def test_returns_none_when_not_listed(self, flask_app_with_containers: Flask, db_session_with_containers: Session): - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db(str(uuid4())) + result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( + str(uuid4()), session=db_session_with_containers + ) assert result is None @@ -282,7 +300,9 @@ class TestFetchRecommendedAppDetailFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db(app1.id) + result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( + app1.id, session=db_session_with_containers + ) assert result is None @@ -298,7 +318,9 @@ class TestFetchRecommendedAppDetailFromDb: db_session_with_containers.expire_all() - result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db(app1.id) + result = DatabaseRecommendAppRetrieval.fetch_recommended_app_detail_from_db( + app1.id, session=db_session_with_containers + ) assert result is not None assert result["id"] == app1.id diff --git a/api/tests/test_containers_integration_tests/services/test_account_service.py b/api/tests/test_containers_integration_tests/services/test_account_service.py index 65a5b0a96bf..ac8ed39316b 100644 --- a/api/tests/test_containers_integration_tests/services/test_account_service.py +++ b/api/tests/test_containers_integration_tests/services/test_account_service.py @@ -1120,10 +1120,12 @@ class TestAccountService: mock_sync.return_value = True # Delete account - AccountService.delete_account(account) + AccountService.delete_account(account, session=db_session_with_containers) # Verify sync was called - mock_sync.assert_called_once_with(account_id=account.id, source="account_deleted") + mock_sync.assert_called_once_with( + account_id=account.id, source="account_deleted", session=db_session_with_containers + ) # Verify task was added to queue mock_delete_task.delay.assert_called_once_with(account.id) diff --git a/api/tests/test_containers_integration_tests/services/test_agent_service.py b/api/tests/test_containers_integration_tests/services/test_agent_service.py index 0ee0cb84e75..00b4a1563ff 100644 --- a/api/tests/test_containers_integration_tests/services/test_agent_service.py +++ b/api/tests/test_containers_integration_tests/services/test_agent_service.py @@ -132,7 +132,7 @@ class TestAgentService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Update the app model config to set agent_mode for agent-chat mode if app.mode == AppMode.AGENT_CHAT and app.app_model_config: @@ -295,7 +295,7 @@ class TestAgentService: agent_thoughts = self._create_test_agent_thoughts(db_session_with_containers, message) # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result structure assert result is not None @@ -355,7 +355,7 @@ class TestAgentService: # Execute the method under test with non-existent conversation with pytest.raises(ValueError, match="Conversation not found"): - AgentService.get_agent_logs(app, fake.uuid4(), fake.uuid4()) + AgentService.get_agent_logs(app, fake.uuid4(), fake.uuid4(), db_session_with_containers) def test_get_agent_logs_message_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -371,7 +371,7 @@ class TestAgentService: # Execute the method under test with non-existent message with pytest.raises(ValueError, match="Message not found"): - AgentService.get_agent_logs(app, conversation.id, fake.uuid4()) + AgentService.get_agent_logs(app, conversation.id, fake.uuid4(), db_session_with_containers) def test_get_agent_logs_with_end_user( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -452,7 +452,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -524,7 +524,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -569,7 +569,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -593,7 +593,7 @@ class TestAgentService: conversation, message = self._create_test_conversation_and_message(db_session_with_containers, app, account) # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -655,7 +655,7 @@ class TestAgentService: # Execute the method under test with pytest.raises(ValueError, match="App model config not found"): - AgentService.get_agent_logs(app, conversation.id, message.id) + AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) def test_get_agent_logs_agent_config_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -674,7 +674,7 @@ class TestAgentService: # Execute the method under test with pytest.raises(ValueError, match="Agent config not found"): - AgentService.get_agent_logs(app, conversation.id, message.id) + AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) def test_list_agent_providers_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -804,7 +804,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -899,7 +899,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -927,7 +927,7 @@ class TestAgentService: mock_external_service_dependencies["current_user"].timezone = "Asia/Shanghai" # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -968,7 +968,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result assert result is not None @@ -1009,7 +1009,7 @@ class TestAgentService: db_session_with_containers.commit() # Execute the method under test - result = AgentService.get_agent_logs(app, conversation.id, message.id) + result = AgentService.get_agent_logs(app, conversation.id, message.id, db_session_with_containers) # Verify the result - should handle malformed JSON gracefully assert result is not None diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/test_containers_integration_tests/services/test_annotation_service.py index 94d72b19be8..2710df5e56c 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_annotation_service.py @@ -101,7 +101,7 @@ class TestAnnotationService: # Create app app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Setup current_user mock self._mock_current_user(mock_external_service_dependencies, account.id, tenant.id) @@ -207,7 +207,9 @@ class TestAnnotationService: } # Insert annotation directly - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) # Verify annotation was created correctly assert annotation.app_id == app.id @@ -241,7 +243,9 @@ class TestAnnotationService: } with pytest.raises(ValueError): - AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) def test_insert_app_annotation_directly_app_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -263,7 +267,9 @@ class TestAnnotationService: # Try to insert annotation with non-existent app with pytest.raises(NotFound, match="App not found"): - AppAnnotationService.insert_app_annotation_directly(annotation_args, non_existent_app_id) + AppAnnotationService.insert_app_annotation_directly( + annotation_args, non_existent_app_id, session=db_session_with_containers + ) def test_update_app_annotation_directly_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -279,7 +285,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(original_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + original_args, app.id, session=db_session_with_containers + ) # Update the annotation updated_args = { @@ -328,7 +336,9 @@ class TestAnnotationService: } # Insert annotation from message - annotation = AppAnnotationService.up_insert_app_annotation_from_message(annotation_args, app.id) + annotation = AppAnnotationService.up_insert_app_annotation_from_message( + annotation_args, app.id, session=db_session_with_containers + ) # Verify annotation was created correctly assert annotation.app_id == app.id @@ -361,7 +371,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - initial_annotation = AppAnnotationService.up_insert_app_annotation_from_message(initial_args, app.id) + initial_annotation = AppAnnotationService.up_insert_app_annotation_from_message( + initial_args, app.id, session=db_session_with_containers + ) # Update the annotation updated_args = { @@ -369,7 +381,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - updated_annotation = AppAnnotationService.up_insert_app_annotation_from_message(updated_args, app.id) + updated_annotation = AppAnnotationService.up_insert_app_annotation_from_message( + updated_args, app.id, session=db_session_with_containers + ) # Verify annotation was updated correctly (same ID) assert updated_annotation.id == initial_annotation.id @@ -402,7 +416,9 @@ class TestAnnotationService: # Try to insert annotation with non-existent app with pytest.raises(NotFound, match="App not found"): - AppAnnotationService.up_insert_app_annotation_from_message(annotation_args, non_existent_app_id) + AppAnnotationService.up_insert_app_annotation_from_message( + annotation_args, non_existent_app_id, session=db_session_with_containers + ) def test_get_annotation_list_by_app_id_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -420,12 +436,18 @@ class TestAnnotationService: "question": f"Question {i}: {fake.sentence()}", "answer": f"Answer {i}: {fake.text(max_nb_chars=200)}", } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) annotations.append(annotation) # Get annotation list annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword="" + app.id, + page=1, + limit=10, + keyword="", + session=db_session_with_containers, ) # Verify results @@ -452,18 +474,22 @@ class TestAnnotationService: "question": f"Question with {unique_keyword} keyword", "answer": f"Answer with {unique_keyword} keyword", } - AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id, session=db_session_with_containers) # Create another annotation without the keyword other_args = { "question": "Different question without special term", "answer": "Different answer without special content", } - AppAnnotationService.insert_app_annotation_directly(other_args, app.id) + AppAnnotationService.insert_app_annotation_directly(other_args, app.id, session=db_session_with_containers) # Search with keyword annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword=unique_keyword + app.id, + page=1, + limit=10, + keyword=unique_keyword, + session=db_session_with_containers, ) # Verify only matching annotations are returned @@ -490,30 +516,42 @@ class TestAnnotationService: "question": "Question with 50% discount", "answer": "Answer about 50% discount offer", } - AppAnnotationService.insert_app_annotation_directly(annotation_with_percent, app.id) + AppAnnotationService.insert_app_annotation_directly( + annotation_with_percent, app.id, session=db_session_with_containers + ) annotation_with_underscore = { "question": "Question with test_data", "answer": "Answer about test_data value", } - AppAnnotationService.insert_app_annotation_directly(annotation_with_underscore, app.id) + AppAnnotationService.insert_app_annotation_directly( + annotation_with_underscore, app.id, session=db_session_with_containers + ) annotation_with_backslash = { "question": "Question with path\\to\\file", "answer": "Answer about path\\to\\file location", } - AppAnnotationService.insert_app_annotation_directly(annotation_with_backslash, app.id) + AppAnnotationService.insert_app_annotation_directly( + annotation_with_backslash, app.id, session=db_session_with_containers + ) # Create annotation that should NOT match (contains % but as part of different text) annotation_no_match = { "question": "Question with 100% different", "answer": "Answer about 100% different content", } - AppAnnotationService.insert_app_annotation_directly(annotation_no_match, app.id) + AppAnnotationService.insert_app_annotation_directly( + annotation_no_match, app.id, session=db_session_with_containers + ) # Test 1: Search with % character - should find exact match only annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword="50%" + app.id, + page=1, + limit=10, + keyword="50%", + session=db_session_with_containers, ) assert total == 1 assert len(annotation_list) == 1 @@ -521,7 +559,11 @@ class TestAnnotationService: # Test 2: Search with _ character - should find exact match only annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword="test_data" + app.id, + page=1, + limit=10, + keyword="test_data", + session=db_session_with_containers, ) assert total == 1 assert len(annotation_list) == 1 @@ -529,7 +571,11 @@ class TestAnnotationService: # Test 3: Search with \ character - should find exact match only annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword="path\\to\\file" + app.id, + page=1, + limit=10, + keyword="path\\to\\file", + session=db_session_with_containers, ) assert total == 1 assert len(annotation_list) == 1 @@ -537,7 +583,11 @@ class TestAnnotationService: # Test 4: Search with % should NOT match 100% (verifies escaping works) annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id( - app.id, page=1, limit=10, keyword="50%" + app.id, + page=1, + limit=10, + keyword="50%", + session=db_session_with_containers, ) # Should only find the 50% annotation, not the 100% one assert total == 1 @@ -557,7 +607,9 @@ class TestAnnotationService: # Try to get annotation list with non-existent app with pytest.raises(NotFound, match="App not found"): - AppAnnotationService.get_annotation_list_by_app_id(non_existent_app_id, page=1, limit=10, keyword="") + AppAnnotationService.get_annotation_list_by_app_id( + non_existent_app_id, page=1, limit=10, keyword="", session=db_session_with_containers + ) def test_delete_app_annotation_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -573,7 +625,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) annotation_id = annotation.id # Delete the annotation @@ -728,7 +782,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) # Add some hit histories for i in range(3): @@ -742,6 +798,7 @@ class TestAnnotationService: message_id=fake.uuid4(), from_source=ConversationFromSource.CONSOLE, score=0.8 + (i * 0.1), + session=db_session_with_containers, ) # Get hit histories @@ -749,6 +806,7 @@ class TestAnnotationService: self._annotation_ref(app, annotation.id), page=1, limit=10, + session=db_session_with_containers, ) # Verify results @@ -775,7 +833,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) # Get initial hit count initial_hit_count = annotation.hit_count @@ -795,6 +855,7 @@ class TestAnnotationService: message_id=message_id, from_source=ConversationFromSource.CONSOLE, score=score, + session=db_session_with_containers, ) # Verify hit count was incremented @@ -834,10 +895,14 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - created_annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + created_annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) # Get annotation by ID - retrieved_annotation = AppAnnotationService.get_annotation_by_id(created_annotation.id) + retrieved_annotation = AppAnnotationService.get_annotation_by_id( + created_annotation.id, session=db_session_with_containers + ) # Verify annotation was retrieved correctly assert retrieved_annotation is not None @@ -880,7 +945,9 @@ class TestAnnotationService: mock_pd.read_csv.return_value = mock_df # Batch import annotations - result = AppAnnotationService.batch_import_app_annotations(app.id, file_storage) + result = AppAnnotationService.batch_import_app_annotations( + app.id, file_storage, session=db_session_with_containers + ) # Verify result structure assert "job_id" in result @@ -920,7 +987,9 @@ class TestAnnotationService: mock_pd.read_csv.return_value = mock_df # Batch import annotations - result = AppAnnotationService.batch_import_app_annotations(app.id, file_storage) + result = AppAnnotationService.batch_import_app_annotations( + app.id, file_storage, session=db_session_with_containers + ) # Verify error result assert "error_msg" in result @@ -966,7 +1035,9 @@ class TestAnnotationService: ].get_features.return_value.annotation_quota_limit.size = 0 # Batch import annotations - result = AppAnnotationService.batch_import_app_annotations(app.id, file_storage) + result = AppAnnotationService.batch_import_app_annotations( + app.id, file_storage, session=db_session_with_containers + ) # Verify error result assert "error_msg" in result @@ -1008,7 +1079,7 @@ class TestAnnotationService: db_session_with_containers.commit() # Get annotation setting - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, session=db_session_with_containers) # Verify result structure assert result["enabled"] is True @@ -1027,7 +1098,7 @@ class TestAnnotationService: app, account = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies) # Get annotation setting (no setting exists) - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, session=db_session_with_containers) # Verify result structure assert result["enabled"] is False @@ -1072,7 +1143,9 @@ class TestAnnotationService: "score_threshold": 0.9, } - result = AppAnnotationService.update_app_annotation_setting(app.id, annotation_setting.id, update_args) + result = AppAnnotationService.update_app_annotation_setting( + app.id, annotation_setting.id, update_args, session=db_session_with_containers + ) # Verify result structure assert result["enabled"] is True @@ -1101,11 +1174,15 @@ class TestAnnotationService: "question": f"Question {i}: {fake.sentence()}", "answer": f"Answer {i}: {fake.text(max_nb_chars=200)}", } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) annotations.append(annotation) # Export annotation list - exported_annotations = AppAnnotationService.export_annotation_list_by_app_id(app.id) + exported_annotations = AppAnnotationService.export_annotation_list_by_app_id( + app.id, session=db_session_with_containers + ) # Verify results assert len(exported_annotations) == 3 @@ -1132,7 +1209,9 @@ class TestAnnotationService: # Try to export annotation list with non-existent app with pytest.raises(NotFound, match="App not found"): - AppAnnotationService.export_annotation_list_by_app_id(non_existent_app_id) + AppAnnotationService.export_annotation_list_by_app_id( + non_existent_app_id, session=db_session_with_containers + ) def test_insert_app_annotation_directly_with_setting_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -1176,7 +1255,9 @@ class TestAnnotationService: } # Insert annotation directly - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) # Verify annotation was created correctly assert annotation.app_id == app.id @@ -1235,7 +1316,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(original_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + original_args, app.id, session=db_session_with_containers + ) # Reset mock to clear previous calls mock_external_service_dependencies["update_task"].delay.reset_mock() @@ -1312,7 +1395,9 @@ class TestAnnotationService: "question": fake.sentence(), "answer": fake.text(max_nb_chars=200), } - annotation = AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + annotation = AppAnnotationService.insert_app_annotation_directly( + annotation_args, app.id, session=db_session_with_containers + ) annotation_id = annotation.id # Reset mock to clear previous calls @@ -1382,7 +1467,9 @@ class TestAnnotationService: } # Insert annotation from message - annotation = AppAnnotationService.up_insert_app_annotation_from_message(annotation_args, app.id) + annotation = AppAnnotationService.up_insert_app_annotation_from_message( + annotation_args, app.id, session=db_session_with_containers + ) # Verify annotation was created correctly assert annotation.app_id == app.id diff --git a/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py b/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py index 1f88ce90621..de51f5077e6 100644 --- a/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py +++ b/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py @@ -82,7 +82,7 @@ class TestAPIBasedExtensionService: ) # Save extension - saved_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + saved_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Verify extension was saved correctly assert saved_extension.id is not None @@ -120,21 +120,21 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="name must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Test empty api_endpoint extension_data.name = fake.company() extension_data.api_endpoint = "" with pytest.raises(ValueError, match="api_endpoint must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Test empty api_key extension_data.api_endpoint = f"https://{fake.domain_name()}/api" extension_data.api_key = "" with pytest.raises(ValueError, match="api_key must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_get_all_by_tenant_id_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -158,11 +158,11 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - saved_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + saved_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) extensions.append(saved_extension) # Get all extensions for tenant - extension_list = APIBasedExtensionService.get_all_by_tenant_id(db_session_with_containers, tenant.id) + extension_list = APIBasedExtensionService.get_all_by_tenant_id(tenant.id, session=db_session_with_containers) # Verify results assert len(extension_list) == 3 @@ -192,11 +192,11 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - created_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Get extension by ID retrieved_extension = APIBasedExtensionService.get_with_tenant_id( - db_session_with_containers, tenant.id, created_extension.id + tenant.id, created_extension.id, session=db_session_with_containers ) # Verify extension was retrieved correctly @@ -223,7 +223,7 @@ class TestAPIBasedExtensionService: # Try to get non-existent extension with pytest.raises(ValueError, match="API based extension is not found"): APIBasedExtensionService.get_with_tenant_id( - db_session_with_containers, tenant.id, non_existent_extension_id + tenant.id, non_existent_extension_id, session=db_session_with_containers ) def test_delete_extension_success(self, db_session_with_containers: Session, mock_external_service_dependencies): @@ -243,11 +243,11 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - created_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) extension_id = created_extension.id # Delete the extension - APIBasedExtensionService.delete(db_session_with_containers, created_extension) + APIBasedExtensionService.delete(created_extension, session=db_session_with_containers) # Verify extension was deleted @@ -275,7 +275,7 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - APIBasedExtensionService.save(db_session_with_containers, extension_data1) + APIBasedExtensionService.save(extension_data1, session=db_session_with_containers) # Try to create second extension with same name extension_data2 = APIBasedExtension( tenant_id=tenant.id, @@ -285,7 +285,7 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="name must be unique, it is already existed"): - APIBasedExtensionService.save(db_session_with_containers, extension_data2) + APIBasedExtensionService.save(extension_data2, session=db_session_with_containers) def test_save_extension_update_existing( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -306,7 +306,7 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - created_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Save original values for later comparison original_name = created_extension.name @@ -325,7 +325,7 @@ class TestAPIBasedExtensionService: created_extension.api_endpoint = new_endpoint created_extension.api_key = new_api_key - updated_extension = APIBasedExtensionService.save(db_session_with_containers, created_extension) + updated_extension = APIBasedExtensionService.save(created_extension, session=db_session_with_containers) # Verify extension was updated correctly assert updated_extension.id == created_extension.id @@ -342,7 +342,7 @@ class TestAPIBasedExtensionService: # Verify the update by retrieving the extension again retrieved_extension = APIBasedExtensionService.get_with_tenant_id( - db_session_with_containers, tenant.id, created_extension.id + tenant.id, created_extension.id, session=db_session_with_containers ) assert retrieved_extension.name == new_name assert retrieved_extension.api_endpoint == new_endpoint @@ -374,7 +374,7 @@ class TestAPIBasedExtensionService: # Try to save extension with connection error with pytest.raises(ValueError, match="connection error: request timeout"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_invalid_api_key_length( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -397,7 +397,7 @@ class TestAPIBasedExtensionService: # Try to save extension with short API key with pytest.raises(ValueError, match="api_key must be at least 5 characters"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_empty_fields(self, db_session_with_containers: Session, mock_external_service_dependencies): """ @@ -417,21 +417,21 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="name must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Test with None api_endpoint extension_data.name = fake.company() extension_data.api_endpoint = None with pytest.raises(ValueError, match="api_endpoint must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Test with None api_key extension_data.api_endpoint = f"https://{fake.domain_name()}/api" extension_data.api_key = None with pytest.raises(ValueError, match="api_key must not be empty"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_get_all_by_tenant_id_empty_list( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -445,7 +445,7 @@ class TestAPIBasedExtensionService: ) # Get all extensions for tenant (none exist) - extension_list = APIBasedExtensionService.get_all_by_tenant_id(db_session_with_containers, tenant.id) + extension_list = APIBasedExtensionService.get_all_by_tenant_id(tenant.id, session=db_session_with_containers) # Verify empty list is returned assert len(extension_list) == 0 @@ -475,7 +475,7 @@ class TestAPIBasedExtensionService: # Try to save extension with invalid ping response with pytest.raises(ValueError, match="{'result': 'invalid'}"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_missing_ping_result( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -501,7 +501,7 @@ class TestAPIBasedExtensionService: # Try to save extension with missing ping result with pytest.raises(ValueError, match="{'status': 'ok'}"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_get_with_tenant_id_wrong_tenant( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -527,11 +527,13 @@ class TestAPIBasedExtensionService: api_key=fake.password(length=20), ) - created_extension = APIBasedExtensionService.save(db_session_with_containers, extension_data) + created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) # Try to get extension with wrong tenant ID with pytest.raises(ValueError, match="API based extension is not found"): - APIBasedExtensionService.get_with_tenant_id(db_session_with_containers, tenant2.id, created_extension.id) + APIBasedExtensionService.get_with_tenant_id( + tenant2.id, created_extension.id, session=db_session_with_containers + ) def test_save_extension_api_key_exactly_four_chars_rejected( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -551,7 +553,7 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="api_key must be at least 5 characters"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_api_key_exactly_five_chars_accepted( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -570,7 +572,7 @@ class TestAPIBasedExtensionService: api_key="12345", ) - saved = APIBasedExtensionService.save(db_session_with_containers, extension_data) + saved = APIBasedExtensionService.save(extension_data, session=db_session_with_containers) assert saved.id is not None def test_save_extension_requestor_constructor_error( @@ -593,7 +595,7 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="connection error: bad config"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_network_exception( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -617,7 +619,7 @@ class TestAPIBasedExtensionService: ) with pytest.raises(ValueError, match="connection error: network failure"): - APIBasedExtensionService.save(db_session_with_containers, extension_data) + APIBasedExtensionService.save(extension_data, session=db_session_with_containers) def test_save_extension_update_duplicate_name_rejected( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -630,28 +632,28 @@ class TestAPIBasedExtensionService: assert tenant is not None ext1 = APIBasedExtensionService.save( - db_session_with_containers, APIBasedExtension( tenant_id=tenant.id, name="Extension Alpha", api_endpoint=f"https://{fake.domain_name()}/api", api_key=fake.password(length=20), ), + session=db_session_with_containers, ) ext2 = APIBasedExtensionService.save( - db_session_with_containers, APIBasedExtension( tenant_id=tenant.id, name="Extension Beta", api_endpoint=f"https://{fake.domain_name()}/api", api_key=fake.password(length=20), ), + session=db_session_with_containers, ) # Try to rename ext2 to ext1's name ext2.name = "Extension Alpha" with pytest.raises(ValueError, match="name must be unique, it is already existed"): - APIBasedExtensionService.save(db_session_with_containers, ext2) + APIBasedExtensionService.save(ext2, session=db_session_with_containers) def test_get_all_returns_empty_for_different_tenant( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -667,15 +669,15 @@ class TestAPIBasedExtensionService: assert tenant1 is not None APIBasedExtensionService.save( - db_session_with_containers, APIBasedExtension( tenant_id=tenant1.id, name=fake.company(), api_endpoint=f"https://{fake.domain_name()}/api", api_key=fake.password(length=20), ), + session=db_session_with_containers, ) assert tenant2 is not None - result = APIBasedExtensionService.get_all_by_tenant_id(db_session_with_containers, tenant2.id) + result = APIBasedExtensionService.get_all_by_tenant_id(tenant2.id, session=db_session_with_containers) assert result == [] diff --git a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py b/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py index cee08c4c33e..24c14637296 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py @@ -162,7 +162,7 @@ class TestAppDslService: api_rpm=10, ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account def _create_simple_yaml_content(self, app_name: str = "Test App", app_mode: str = "chat") -> str: @@ -841,7 +841,7 @@ class TestAppDslService: # ── Export ───────────────────────────────────────────────────────── - def test_export_dsl_delegates_by_mode(self, monkeypatch: pytest.MonkeyPatch): + def test_export_dsl_delegates_by_mode(self, monkeypatch: pytest.MonkeyPatch, db_session_with_containers: Session): workflow_calls: list[bool] = [] model_calls: list[bool] = [] monkeypatch.setattr( @@ -859,7 +859,7 @@ class TestAppDslService: mode=AppMode.WORKFLOW, icon_type="emoji", ) - AppDslService.export_dsl(workflow_app) + AppDslService.export_dsl(workflow_app, session=db_session_with_containers) assert workflow_calls == [True] chat_app = _app_stub( @@ -867,10 +867,12 @@ class TestAppDslService: icon_type="emoji", app_model_config=SimpleNamespace(to_dict=lambda: {"agent_mode": {"tools": []}}), ) - AppDslService.export_dsl(chat_app) + AppDslService.export_dsl(chat_app, session=db_session_with_containers) assert model_calls == [True] - def test_export_dsl_preserves_icon_and_icon_type(self, monkeypatch: pytest.MonkeyPatch): + def test_export_dsl_preserves_icon_and_icon_type( + self, monkeypatch: pytest.MonkeyPatch, db_session_with_containers: Session + ): monkeypatch.setattr( AppDslService, "_append_workflow_export_data", @@ -886,7 +888,7 @@ class TestAppDslService: description="App with emoji icon", use_icon_as_answer_icon=True, ) - yaml_output = AppDslService.export_dsl(emoji_app) + yaml_output = AppDslService.export_dsl(emoji_app, session=db_session_with_containers) data = yaml.safe_load(yaml_output) assert data["app"]["icon"] == "🎨" assert data["app"]["icon_type"] == "emoji" @@ -901,7 +903,7 @@ class TestAppDslService: description="App with image icon", use_icon_as_answer_icon=False, ) - yaml_output = AppDslService.export_dsl(image_app) + yaml_output = AppDslService.export_dsl(image_app, session=db_session_with_containers) data = yaml.safe_load(yaml_output) assert data["app"]["icon"] == "https://example.com/icon.png" assert data["app"]["icon_type"] == "image" @@ -936,7 +938,7 @@ class TestAppDslService: db_session_with_containers.add(model_config) db_session_with_containers.commit() - exported_dsl = AppDslService.export_dsl(app, include_secret=False) + exported_dsl = AppDslService.export_dsl(app, include_secret=False, session=db_session_with_containers) exported_data = yaml.safe_load(exported_dsl) assert exported_data["kind"] == "app" @@ -972,7 +974,7 @@ class TestAppDslService: "workflow_service" ].return_value.get_draft_workflow.return_value = mock_workflow - exported_dsl = AppDslService.export_dsl(app, include_secret=False) + exported_dsl = AppDslService.export_dsl(app, include_secret=False, session=db_session_with_containers) exported_data = yaml.safe_load(exported_dsl) assert exported_data["kind"] == "app" @@ -1006,7 +1008,7 @@ class TestAppDslService: workflow_id = str(uuid4()) - def mock_get_draft_workflow(app_model, wf_id=None): + def mock_get_draft_workflow(app_model, wf_id=None, **_kwargs): if wf_id == workflow_id: return mock_workflow return None @@ -1015,7 +1017,9 @@ class TestAppDslService: "workflow_service" ].return_value.get_draft_workflow.side_effect = mock_get_draft_workflow - exported_dsl = AppDslService.export_dsl(app, include_secret=False, workflow_id=workflow_id) + exported_dsl = AppDslService.export_dsl( + app, include_secret=False, workflow_id=workflow_id, session=db_session_with_containers + ) exported_data = yaml.safe_load(exported_dsl) assert exported_data["kind"] == "app" @@ -1034,11 +1038,15 @@ class TestAppDslService: WorkflowNotFoundError, match="Missing draft workflow configuration, please check.", ): - AppDslService.export_dsl(app, include_secret=False, workflow_id=str(uuid4())) + AppDslService.export_dsl( + app, include_secret=False, workflow_id=str(uuid4()), session=db_session_with_containers + ) # ── Workflow Export Data ─────────────────────────────────────────── - def test_append_workflow_export_data_filters_and_overrides(self, monkeypatch: pytest.MonkeyPatch): + def test_append_workflow_export_data_filters_and_overrides( + self, monkeypatch: pytest.MonkeyPatch, db_session_with_containers: Session + ): workflow_dict = { "graph": { "nodes": [ @@ -1123,6 +1131,7 @@ class TestAppDslService: app_model=_app_stub(), include_secret=False, workflow_id=None, + session=db_session_with_containers, ) nodes = export_data["workflow"]["graph"]["nodes"] @@ -1138,7 +1147,9 @@ class TestAppDslService: assert nodes[5]["data"]["subscription_id"] == "" assert export_data["dependencies"] == [{"tenant": _DEFAULT_TENANT_ID, "dep": "dep-1"}] - def test_append_workflow_export_data_missing_workflow_raises(self, monkeypatch: pytest.MonkeyPatch): + def test_append_workflow_export_data_missing_workflow_raises( + self, monkeypatch: pytest.MonkeyPatch, db_session_with_containers: Session + ): workflow_service = MagicMock() workflow_service.get_draft_workflow.return_value = None monkeypatch.setattr(app_dsl_service, "WorkflowService", lambda: workflow_service) @@ -1149,6 +1160,7 @@ class TestAppDslService: app_model=_app_stub(), include_secret=False, workflow_id=None, + session=db_session_with_containers, ) # ── Model Config Export Data ────────────────────────────────────── diff --git a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py b/api/tests/test_containers_integration_tests/services/test_app_generate_service.py index 473111f364a..89cc7715d1c 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_generate_service.py @@ -187,7 +187,7 @@ class TestAppGenerateService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account @@ -234,12 +234,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -267,12 +267,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -298,12 +298,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -329,12 +329,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -362,12 +362,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -399,12 +399,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -431,12 +431,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.DEBUGGER, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -461,12 +461,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=db_session_with_containers, ) # Verify the result @@ -503,12 +503,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=end_user, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -535,12 +535,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -574,12 +574,12 @@ class TestAppGenerateService: # StatementError (from EnumText validation during autoflush) with pytest.raises((ValueError, sa.exc.StatementError)): AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) def test_generate_with_workflow_id_format_error( @@ -603,12 +603,12 @@ class TestAppGenerateService: # Execute the method under test and expect WorkflowIdFormatError with pytest.raises(WorkflowIdFormatError) as exc_info: AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -642,12 +642,12 @@ class TestAppGenerateService: # Execute the method under test and expect WorkflowNotFoundError with pytest.raises(WorkflowNotFoundError) as exc_info: AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -673,12 +673,12 @@ class TestAppGenerateService: # Execute the method under test and expect ValueError with pytest.raises(ValueError) as exc_info: AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.DEBUGGER, streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -704,12 +704,12 @@ class TestAppGenerateService: # Execute the method under test and expect ValueError with pytest.raises(ValueError) as exc_info: AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -731,7 +731,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate_single_iteration( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -758,7 +763,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate_single_iteration( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -786,7 +796,12 @@ class TestAppGenerateService: # Execute the method under test and expect ValueError with pytest.raises(ValueError) as exc_info: AppGenerateService.generate_single_iteration( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -808,7 +823,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate_single_loop( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -835,7 +855,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate_single_loop( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -861,7 +886,12 @@ class TestAppGenerateService: # Execute the method under test and expect ValueError with pytest.raises(ValueError) as exc_info: AppGenerateService.generate_single_loop( - app_model=app, user=account, node_id=node_id, args=args, streaming=True + app_model=app, + user=account, + node_id=node_id, + args=args, + streaming=True, + session=db_session_with_containers, ) # Verify error message @@ -1021,12 +1051,12 @@ class TestAppGenerateService: # Execute the method under test and expect exception with pytest.raises(Exception) as exc_info: AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify exception message @@ -1054,12 +1084,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -1094,12 +1124,12 @@ class TestAppGenerateService: # Execute the method under test result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=invoke_from, streaming=True, + session=db_session_with_containers, ) # Verify the result @@ -1137,12 +1167,12 @@ class TestAppGenerateService: mock_exec_params.new.return_value = mock_payload result = AppGenerateService.generate( - session=db_session_with_containers, app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=db_session_with_containers, ) # Verify the result diff --git a/api/tests/test_containers_integration_tests/services/test_app_service.py b/api/tests/test_containers_integration_tests/services/test_app_service.py index f9df99c5594..8deaf6d462d 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_service.py @@ -84,7 +84,7 @@ class TestAppService: # Create app app_service = AppService() - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Verify app was created correctly assert app.name == app_params.name @@ -144,7 +144,7 @@ class TestAppService: icon_background="#4ECDC4", ) - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Verify app mode was set correctly assert app.mode == mode @@ -183,7 +183,7 @@ class TestAppService: ) app_service = AppService() - created_app = app_service.create_app(tenant.id, app_params, account) + created_app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Get app using the service - needs current_user mock mock_current_user = create_autospec(Account, instance=True) @@ -234,7 +234,7 @@ class TestAppService: icon="📱", icon_background="#96CEB4", ) - app_service.create_app(tenant.id, app_params, account) + app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Get paginated apps params = AppListParams(page=1, limit=10, mode="chat") @@ -277,16 +277,19 @@ class TestAppService: tenant.id, CreateAppParams(name="Oldest Created", mode="chat", icon_type="emoji", icon="1"), account, + session=db_session_with_containers, ) newest_modified = app_service.create_app( tenant.id, CreateAppParams(name="Newest Modified", mode="chat", icon_type="emoji", icon="2"), account, + session=db_session_with_containers, ) newest_created = app_service.create_app( tenant.id, CreateAppParams(name="Newest Created", mode="chat", icon_type="emoji", icon="3"), account, + session=db_session_with_containers, ) timestamp_by_app_id = { @@ -362,15 +365,17 @@ class TestAppService: tenant.id, CreateAppParams(name="Starred App", mode="chat", icon_type="emoji", icon="1"), account, + session=db_session_with_containers, ) unstarred_app = app_service.create_app( tenant.id, CreateAppParams(name="Unstarred App", mode="chat", icon_type="emoji", icon="2"), account, + session=db_session_with_containers, ) - app_service.star_app(db_session_with_containers, app=starred_app, account_id=account.id) - app_service.star_app(db_session_with_containers, app=starred_app, account_id=account.id) + app_service.star_app(app=starred_app, account_id=account.id, session=db_session_with_containers) + app_service.star_app(app=starred_app, account_id=account.id, session=db_session_with_containers) db_session_with_containers.commit() star_count = db_session_with_containers.scalar( @@ -386,7 +391,7 @@ class TestAppService: assert starred_by_app_id[starred_app.id] is True assert starred_by_app_id[unstarred_app.id] is False - app_service.unstar_app(db_session_with_containers, app=starred_app, account_id=account.id) + app_service.unstar_app(app=starred_app, account_id=account.id, session=db_session_with_containers) db_session_with_containers.commit() paginated_apps = app_service.get_paginate_apps( @@ -422,26 +427,30 @@ class TestAppService: tenant.id, CreateAppParams(name="Oldest Created Starred App", mode="chat", icon_type="emoji", icon="1"), account, + session=db_session_with_containers, ) newest_modified_app = app_service.create_app( tenant.id, CreateAppParams(name="Newest Modified Starred App", mode="chat", icon_type="emoji", icon="2"), account, + session=db_session_with_containers, ) newest_created_app = app_service.create_app( tenant.id, CreateAppParams(name="Newest Created Starred App", mode="chat", icon_type="emoji", icon="3"), account, + session=db_session_with_containers, ) unstarred_app = app_service.create_app( tenant.id, CreateAppParams(name="Unstarred App", mode="chat", icon_type="emoji", icon="4"), account, + session=db_session_with_containers, ) - app_service.star_app(db_session_with_containers, app=oldest_created_app, account_id=account.id) - app_service.star_app(db_session_with_containers, app=newest_modified_app, account_id=account.id) - app_service.star_app(db_session_with_containers, app=newest_created_app, account_id=account.id) + app_service.star_app(app=oldest_created_app, account_id=account.id, session=db_session_with_containers) + app_service.star_app(app=newest_modified_app, account_id=account.id, session=db_session_with_containers) + app_service.star_app(app=newest_created_app, account_id=account.id, session=db_session_with_containers) timestamp_by_app_id = { oldest_created_app.id: (datetime(2026, 1, 1, 10, 0, 0), datetime(2026, 1, 1, 10, 0, 0)), @@ -535,8 +544,10 @@ class TestAppService: icon_background="#4ECDC4", ) - chat_app = app_service.create_app(tenant.id, chat_app_params, account) - completion_app = app_service.create_app(tenant.id, completion_app_params, account) + chat_app = app_service.create_app(tenant.id, chat_app_params, account, session=db_session_with_containers) + completion_app = app_service.create_app( + tenant.id, completion_app_params, account, session=db_session_with_containers + ) # Test filter by mode chat_apps = app_service.get_paginate_apps( @@ -599,7 +610,7 @@ class TestAppService: icon="💬", icon_background="#FF6B6B", ) - app_service.create_app(tenant.id, app_params, first_account) + app_service.create_app(tenant.id, app_params, first_account, session=db_session_with_containers) other_app_params = CreateAppParams( name="Second Creator App", description="Created by the second account", @@ -608,7 +619,7 @@ class TestAppService: icon="✍️", icon_background="#4ECDC4", ) - app_service.create_app(tenant.id, other_app_params, second_account) + app_service.create_app(tenant.id, other_app_params, second_account, session=db_session_with_containers) filtered_apps = app_service.get_paginate_apps( first_account.id, @@ -654,7 +665,7 @@ class TestAppService: icon="🏷️", icon_background="#FFEAA7", ) - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Mock TagService to return the app ID for tag filtering with patch("services.app_service.TagService.get_target_ids_by_tag_ids") as mock_tag_service: @@ -717,7 +728,7 @@ class TestAppService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Store original values original_name = app.name @@ -741,7 +752,7 @@ class TestAppService: mock_current_user.current_tenant_id = account.current_tenant_id with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app(app, update_args) + updated_app = app_service.update_app(app, update_args, session=db_session_with_containers) # Verify updated fields assert updated_app.name == update_args["name"] @@ -788,6 +799,7 @@ class TestAppService: icon_background="#45B7D1", ), account, + session=db_session_with_containers, ) mock_current_user = create_autospec(Account, instance=True) @@ -805,6 +817,7 @@ class TestAppService: "icon_background": "#FF8C42", "use_icon_as_answer_icon": True, }, + session=db_session_with_containers, ) assert updated_app.icon_type == IconType.EMOJI @@ -841,6 +854,7 @@ class TestAppService: icon_background="#45B7D1", ), account, + session=db_session_with_containers, ) mock_current_user = create_autospec(Account, instance=True) @@ -859,6 +873,7 @@ class TestAppService: "icon_background": "#FF8C42", "use_icon_as_answer_icon": True, }, + session=db_session_with_containers, ) def test_update_app_name_success(self, db_session_with_containers: Session, mock_external_service_dependencies): @@ -892,7 +907,7 @@ class TestAppService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Store original name original_name = app.name @@ -904,7 +919,7 @@ class TestAppService: mock_current_user.current_tenant_id = account.current_tenant_id with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_name(app, new_name) + updated_app = app_service.update_app_name(app, new_name, session=db_session_with_containers) assert updated_app.name == new_name assert updated_app.updated_by == account.id @@ -946,7 +961,7 @@ class TestAppService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_params, account) + app = app_service.create_app(tenant.id, app_params, account, session=db_session_with_containers) # Store original values original_icon = app.icon @@ -961,7 +976,9 @@ class TestAppService: mock_current_user.current_tenant_id = account.current_tenant_id with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_icon(app, new_icon, new_icon_background, new_icon_type) + updated_app = app_service.update_app_icon( + app, new_icon, new_icon_background, new_icon_type, session=db_session_with_containers + ) assert updated_app.icon == new_icon assert updated_app.icon_background == new_icon_background @@ -1007,7 +1024,7 @@ class TestAppService: icon_background="#74B9FF", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Store original site status original_site_status = app.enable_site @@ -1018,13 +1035,13 @@ class TestAppService: mock_current_user.current_tenant_id = account.current_tenant_id with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_site_status(app, False) + updated_app = app_service.update_app_site_status(app, False, session=db_session_with_containers) assert updated_app.enable_site is False assert updated_app.updated_by == account.id # Update site status back to enabled with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_site_status(updated_app, True) + updated_app = app_service.update_app_site_status(updated_app, True, session=db_session_with_containers) assert updated_app.enable_site is True assert updated_app.updated_by == account.id @@ -1067,7 +1084,7 @@ class TestAppService: icon_background="#A29BFE", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Store original API status original_api_status = app.enable_api @@ -1078,13 +1095,13 @@ class TestAppService: mock_current_user.current_tenant_id = account.current_tenant_id with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_api_status(app, False) + updated_app = app_service.update_app_api_status(app, False, session=db_session_with_containers) assert updated_app.enable_api is False assert updated_app.updated_by == account.id # Update API status back to enabled with patch("services.app_service.current_user", mock_current_user): - updated_app = app_service.update_app_api_status(updated_app, True) + updated_app = app_service.update_app_api_status(updated_app, True, session=db_session_with_containers) assert updated_app.enable_api is True assert updated_app.updated_by == account.id @@ -1127,14 +1144,14 @@ class TestAppService: icon_background="#FD79A8", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Store original values original_site_status = app.enable_site original_updated_at = app.updated_at # Update site status to the same value (no change) - updated_app = app_service.update_app_site_status(app, original_site_status) + updated_app = app_service.update_app_site_status(app, original_site_status, session=db_session_with_containers) # Verify app is returned unchanged assert updated_app.id == app.id @@ -1178,7 +1195,7 @@ class TestAppService: icon_background="#E17055", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Store app ID for verification app_id = app.id @@ -1188,7 +1205,7 @@ class TestAppService: mock_delete_task.delay.return_value = None # Delete app - app_service.delete_app(app) + app_service.delete_app(app, session=db_session_with_containers) # Verify async deletion task was called mock_delete_task.delay.assert_called_once_with(tenant_id=tenant.id, app_id=app_id) @@ -1230,7 +1247,7 @@ class TestAppService: icon_background="#00B894", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Store app ID for verification app_id = app.id @@ -1245,7 +1262,7 @@ class TestAppService: mock_delete_task.delay.return_value = None # Delete app - app_service.delete_app(app) + app_service.delete_app(app, session=db_session_with_containers) # Verify webapp auth cleanup was called mock_external_service_dependencies["enterprise_service"].WebAppAuth.cleanup_webapp.assert_called_once_with( @@ -1290,10 +1307,10 @@ class TestAppService: icon_background="#6C5CE7", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Get app metadata - app_meta = app_service.get_app_meta(app) + app_meta = app_service.get_app_meta(app, session=db_session_with_containers) # Verify metadata contains expected fields assert "tool_icons" in app_meta @@ -1329,10 +1346,10 @@ class TestAppService: icon_background="#FDCB6E", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Get app code by ID - app_code = AppService.get_app_code_by_id(app.id) + app_code = AppService.get_app_code_by_id(app.id, session=db_session_with_containers) # Verify app code was retrieved correctly # Note: Site would be created when App is created, site.code is auto-generated @@ -1369,7 +1386,7 @@ class TestAppService: icon_background="#E84393", ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Create a site for the app site = Site() @@ -1384,7 +1401,7 @@ class TestAppService: db_session_with_containers.commit() # Get app ID by code - app_id = AppService.get_app_id_by_code(site.code) + app_id = AppService.get_app_id_by_code(site.code, session=db_session_with_containers) # Verify app ID was retrieved correctly assert app_id == app.id @@ -1462,6 +1479,7 @@ class TestAppService: api_rpm=10, ), account, + session=db_session_with_containers, ) app_with_underscore = app_service.create_app( @@ -1477,6 +1495,7 @@ class TestAppService: api_rpm=10, ), account, + session=db_session_with_containers, ) app_with_backslash = app_service.create_app( @@ -1492,6 +1511,7 @@ class TestAppService: api_rpm=10, ), account, + session=db_session_with_containers, ) # Create app that should NOT match @@ -1508,6 +1528,7 @@ class TestAppService: api_rpm=10, ), account, + session=db_session_with_containers, ) # Test 1: Search with % character @@ -1560,7 +1581,7 @@ class TestAppService: from services.app_service import AppService with pytest.raises(ValueError, match="not found"): - AppService.get_app_code_by_id(str(uuid4())) + AppService.get_app_code_by_id(str(uuid4()), session=db_session_with_containers) def test_get_app_id_by_code_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -1569,7 +1590,7 @@ class TestAppService: from services.app_service import AppService with pytest.raises(ValueError, match="not found"): - AppService.get_app_id_by_code("nonexistent-code") + AppService.get_app_id_by_code("nonexistent-code", session=db_session_with_containers) def test_get_app_meta_returns_empty_when_workflow_missing( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -1582,7 +1603,7 @@ class TestAppService: app_service = AppService() workflow_app = SimpleNamespace(mode="workflow", workflow=None) - meta = app_service.get_app_meta(workflow_app) + meta = app_service.get_app_meta(workflow_app, session=db_session_with_containers) assert meta == {"tool_icons": {}} def test_get_app_meta_returns_empty_when_model_config_missing( @@ -1596,5 +1617,5 @@ class TestAppService: app_service = AppService() chat_app = SimpleNamespace(mode="chat", app_model_config=None) - meta = app_service.get_app_meta(chat_app) + meta = app_service.get_app_meta(chat_app, session=db_session_with_containers) assert meta == {"tool_icons": {}} diff --git a/api/tests/test_containers_integration_tests/services/test_billing_service.py b/api/tests/test_containers_integration_tests/services/test_billing_service.py index a3a4a0e6edd..777fb7721b4 100644 --- a/api/tests/test_containers_integration_tests/services/test_billing_service.py +++ b/api/tests/test_containers_integration_tests/services/test_billing_service.py @@ -417,7 +417,7 @@ class TestBillingServiceIsTenantOwnerOrAdmin: account, _ = self._create_account_with_tenant_role(db_session_with_containers, TenantAccountRole.EDITOR) with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"): - BillingService.is_tenant_owner_or_admin(db_session_with_containers, account) + BillingService.is_tenant_owner_or_admin(account, session=db_session_with_containers) def test_is_tenant_owner_or_admin_dataset_operator_raises_error(self, db_session_with_containers: Session) -> None: """is_tenant_owner_or_admin raises ValueError for DATASET_OPERATOR role.""" @@ -426,4 +426,4 @@ class TestBillingServiceIsTenantOwnerOrAdmin: ) with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"): - BillingService.is_tenant_owner_or_admin(db_session_with_containers, account) + BillingService.is_tenant_owner_or_admin(account, session=db_session_with_containers) diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service.py b/api/tests/test_containers_integration_tests/services/test_conversation_service.py index b19b6b9c984..60858df08e8 100644 --- a/api/tests/test_containers_integration_tests/services/test_conversation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_conversation_service.py @@ -6,11 +6,13 @@ from unittest.mock import patch from uuid import uuid4 import pytest +from agenton.compositor import CompositorSessionSnapshot from sqlalchemy import select from sqlalchemy.orm import Session +from core.app.apps.agent_app.session_store import AgentAppRuntimeSessionStore from core.app.entities.app_invoke_entities import InvokeFrom -from models import TenantAccountRole +from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus, TenantAccountRole from models.account import Account, Tenant, TenantAccountJoin from models.enums import ConversationFromSource, EndUserType from models.model import App, Conversation, EndUser, Message, MessageAnnotation @@ -350,6 +352,7 @@ class TestConversationServiceMessageCreation: conversation_id=conversation.id, first_id=None, # No starting point specified limit=10, + session=db_session_with_containers, ) # Assert - Verify the results @@ -395,6 +398,7 @@ class TestConversationServiceMessageCreation: conversation_id=conversation.id, first_id=first_message.id, limit=10, + session=db_session_with_containers, ) # Assert - Verify the results @@ -426,6 +430,7 @@ class TestConversationServiceMessageCreation: conversation_id=conversation.id, first_id=str(uuid4()), limit=10, + session=db_session_with_containers, ) def test_pagination_with_has_more_flag(self, db_session_with_containers: Session): @@ -461,6 +466,7 @@ class TestConversationServiceMessageCreation: conversation_id=conversation.id, first_id=None, limit=limit, + session=db_session_with_containers, ) # Assert @@ -498,7 +504,8 @@ class TestConversationServiceMessageCreation: conversation_id=conversation.id, first_id=None, limit=10, - order="asc", # Ascending order + order="asc", # Ascending order, + session=db_session_with_containers, ) # Assert @@ -547,7 +554,7 @@ class TestConversationServiceSummarization: mock_llm_generator.return_value = generated_name # Act - result = ConversationService.auto_generate_name(app_model, conversation) + result = ConversationService.auto_generate_name(app_model, conversation, session=db_session_with_containers) # Assert assert conversation.name == generated_name # Name updated on conversation object @@ -572,7 +579,7 @@ class TestConversationServiceSummarization: # Act & Assert with pytest.raises(MessageNotExistsError): - ConversationService.auto_generate_name(app_model, conversation) + ConversationService.auto_generate_name(app_model, conversation, session=db_session_with_containers) @patch("services.conversation_service.LLMGenerator.generate_conversation_name") def test_auto_generate_name_handles_llm_failure_gracefully( @@ -604,7 +611,7 @@ class TestConversationServiceSummarization: mock_llm_generator.side_effect = Exception("LLM service unavailable") # Act - result = ConversationService.auto_generate_name(app_model, conversation) + result = ConversationService.auto_generate_name(app_model, conversation, session=db_session_with_containers) # Assert assert conversation.name == original_name # Name remains unchanged @@ -637,6 +644,7 @@ class TestConversationServiceSummarization: user=user, name=new_name, auto_generate=False, + session=db_session_with_containers, ) # Assert @@ -671,6 +679,7 @@ class TestConversationServiceSummarization: user=user, name=None, auto_generate=True, + session=db_session_with_containers, ) # Assert @@ -719,7 +728,9 @@ class TestConversationServiceMessageAnnotation: args = {"message_id": message.id, "answer": "AI is artificial intelligence"} # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app_model.id) + result = AppAnnotationService.up_insert_app_annotation_from_message( + args, app_model.id, session=db_session_with_containers + ) # Assert assert result.message_id == message.id @@ -753,7 +764,9 @@ class TestConversationServiceMessageAnnotation: } # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app_model.id) + result = AppAnnotationService.up_insert_app_annotation_from_message( + args, app_model.id, session=db_session_with_containers + ) # Assert assert result.message_id is None @@ -802,7 +815,9 @@ class TestConversationServiceMessageAnnotation: args = {"message_id": message.id, "answer": "Updated annotation content"} # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app_model.id) + result = AppAnnotationService.up_insert_app_annotation_from_message( + args, app_model.id, session=db_session_with_containers + ) # Assert assert result.id == existing_annotation.id @@ -838,7 +853,11 @@ class TestConversationServiceMessageAnnotation: # Act result_items, result_total = AppAnnotationService.get_annotation_list_by_app_id( - app_id=app_model.id, page=1, limit=10, keyword="" + app_id=app_model.id, + page=1, + limit=10, + keyword="", + session=db_session_with_containers, ) # Assert @@ -886,7 +905,8 @@ class TestConversationServiceMessageAnnotation: app_id=app_model.id, page=1, limit=10, - keyword="machine", # Search keyword + keyword="machine", # Search keyword, + session=db_session_with_containers, ) # Assert @@ -914,7 +934,9 @@ class TestConversationServiceMessageAnnotation: } # Act - result = AppAnnotationService.insert_app_annotation_directly(args, app_model.id) + result = AppAnnotationService.insert_app_annotation_directly( + args, app_model.id, session=db_session_with_containers + ) # Assert assert result.question == args["question"] @@ -942,7 +964,9 @@ class TestConversationServiceExport: ) # Act - result = ConversationService.get_conversation(app_model=app_model, conversation_id=conversation.id, user=user) + result = ConversationService.get_conversation( + app_model=app_model, conversation_id=conversation.id, user=user, session=db_session_with_containers + ) # Assert assert result == conversation @@ -956,7 +980,12 @@ class TestConversationServiceExport: # Act & Assert with pytest.raises(ConversationNotExistsError): - ConversationService.get_conversation(app_model=app_model, conversation_id=str(uuid4()), user=user) + ConversationService.get_conversation( + app_model=app_model, + conversation_id=str(uuid4()), + user=user, + session=db_session_with_containers, + ) @patch("services.annotation_service.current_account_with_tenant") def test_export_annotation_list(self, mock_current_account, db_session_with_containers: Session): @@ -982,7 +1011,7 @@ class TestConversationServiceExport: mock_current_account.return_value = (account, app_model.tenant_id) # Act - result = AppAnnotationService.export_annotation_list_by_app_id(app_model.id) + result = AppAnnotationService.export_annotation_list_by_app_id(app_model.id, session=db_session_with_containers) # Assert assert len(result) == 10 @@ -1006,7 +1035,9 @@ class TestConversationServiceExport: ) # Act - result = MessageService.get_message(app_model=app_model, user=user, message_id=message.id) + result = MessageService.get_message( + app_model=app_model, user=user, message_id=message.id, session=db_session_with_containers + ) # Assert assert result == message @@ -1020,7 +1051,9 @@ class TestConversationServiceExport: # Act & Assert with pytest.raises(MessageNotExistsError): - MessageService.get_message(app_model=app_model, user=user, message_id=str(uuid4())) + MessageService.get_message( + app_model=app_model, user=user, message_id=str(uuid4()), session=db_session_with_containers + ) def test_get_conversation_for_end_user(self, db_session_with_containers: Session): """ @@ -1041,14 +1074,18 @@ class TestConversationServiceExport: # Act result = ConversationService.get_conversation( - app_model=app_model, conversation_id=conversation.id, user=end_user + app_model=app_model, + conversation_id=conversation.id, + user=end_user, + session=db_session_with_containers, ) # Assert assert result == conversation + @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_conversation(self, mock_delete_task, db_session_with_containers: Session): + def test_delete_conversation(self, mock_delete_task, mock_cleanup_task, db_session_with_containers: Session): """ Test conversation deletion with async cleanup. @@ -1067,9 +1104,25 @@ class TestConversationServiceExport: user, ) conversation_id = conversation.id + runtime_session = AgentRuntimeSession( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id=str(uuid4()), + agent_config_snapshot_id=str(uuid4()), + backend_run_id="backend-run-1", + session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), + composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + conversation_id=conversation.id, + status=AgentRuntimeSessionStatus.ACTIVE, + ) + db_session_with_containers.add(runtime_session) + db_session_with_containers.commit() # Act - Delete the conversation - ConversationService.delete(app_model=app_model, conversation_id=conversation_id, user=user) + ConversationService.delete( + app_model=app_model, conversation_id=conversation_id, user=user, session=db_session_with_containers + ) # Assert - Verify two-step deletion process # Step 1: Immediate database deletion @@ -1079,9 +1132,29 @@ class TestConversationServiceExport: # Step 2: Async cleanup task triggered # The Celery task will handle cleanup of messages, annotations, etc. mock_delete_task.delay.assert_called_once_with(conversation_id) + mock_cleanup_task.delay.assert_called_once() + cleanup_payload = mock_cleanup_task.delay.call_args.args[0] + assert cleanup_payload["metadata"]["conversation_id"] == conversation_id + assert ( + cleanup_payload["idempotency_key"] + == f"{app_model.tenant_id}:{app_model.id}:{conversation_id}:agent-runtime-session-cleanup:" + f"{runtime_session.agent_id}:{runtime_session.agent_config_snapshot_id}:{runtime_session.backend_run_id}" + ) + runtime_session_row = db_session_with_containers.scalar( + select(AgentRuntimeSession).where(AgentRuntimeSession.id == runtime_session.id) + ) + assert runtime_session_row is not None + assert runtime_session_row.status == AgentRuntimeSessionStatus.CLEANED + + @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_conversation_not_owned_by_account(self, mock_delete_task, db_session_with_containers: Session): + def test_delete_conversation_not_owned_by_account( + self, + mock_delete_task, + mock_cleanup_task, + db_session_with_containers: Session, + ): """ Test deletion is denied when conversation belongs to a different account. """ @@ -1104,20 +1177,29 @@ class TestConversationServiceExport: app_model=app_model, conversation_id=conversation.id, user=other_account, + session=db_session_with_containers, ) # Verify no deletion and no async cleanup trigger not_deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation.id)) assert not_deleted is not None mock_delete_task.delay.assert_not_called() + mock_cleanup_task.delay.assert_not_called() + @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") @patch("services.conversation_service.delete_conversation_related_data") - def test_delete_handles_exception_and_rollback(self, mock_delete_task, db_session_with_containers: Session): + def test_delete_handles_exception_and_rollback( + self, + mock_delete_task, + mock_cleanup_task, + db_session_with_containers: Session, + ): """ Test that delete propagates exceptions and does not trigger the cleanup task. - When a DB error occurs during deletion, the service must rollback the - transaction and re-raise the exception without scheduling async cleanup. + When a DB error occurs during deletion, the conversation row stays in + place, but any already-enqueued Agent backend cleanup remains a + best-effort terminal lifecycle action. """ # Arrange app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( @@ -1127,15 +1209,136 @@ class TestConversationServiceExport: db_session_with_containers, app_model, user ) conversation_id = conversation.id + runtime_session = AgentRuntimeSession( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id=str(uuid4()), + agent_config_snapshot_id=str(uuid4()), + backend_run_id="backend-run-rollback", + session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), + composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + conversation_id=conversation.id, + status=AgentRuntimeSessionStatus.ACTIVE, + ) + db_session_with_containers.add(runtime_session) + db_session_with_containers.commit() # Act — force an error during the delete to exercise the rollback path - with patch("services.conversation_service.db.session.delete", side_effect=Exception("DB error")): + with patch.object(db_session_with_containers, "delete", side_effect=Exception("DB error")): with pytest.raises(Exception, match="DB error"): - ConversationService.delete(app_model=app_model, conversation_id=conversation_id, user=user) + ConversationService.delete( + app_model=app_model, + conversation_id=conversation_id, + user=user, + session=db_session_with_containers, + ) - # Assert — async cleanup must NOT have been scheduled + # Assert — related-data deletion is not scheduled, but the backend + # cleanup task was already enqueued before the row delete failed. mock_delete_task.delay.assert_not_called() + mock_cleanup_task.delay.assert_called_once() + cleanup_payload = mock_cleanup_task.delay.call_args.args[0] + assert ( + cleanup_payload["idempotency_key"] + == f"{app_model.tenant_id}:{app_model.id}:{conversation_id}:agent-runtime-session-cleanup:" + f"{runtime_session.agent_id}:{runtime_session.agent_config_snapshot_id}:{runtime_session.backend_run_id}" + ) # Conversation is still present because the deletion was never committed still_there = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation_id)) assert still_there is not None + + @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") + @patch("services.conversation_service.delete_conversation_related_data") + def test_delete_ignores_mark_cleaned_failure( + self, + mock_delete_task, + mock_cleanup_task, + db_session_with_containers: Session, + ): + app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( + db_session_with_containers + ) + conversation = ConversationServiceIntegrationTestDataFactory.create_conversation( + db_session_with_containers, + app_model, + user, + ) + runtime_session = AgentRuntimeSession( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id=str(uuid4()), + agent_config_snapshot_id=str(uuid4()), + backend_run_id="backend-run-cleanup-failure", + session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), + composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + conversation_id=conversation.id, + status=AgentRuntimeSessionStatus.ACTIVE, + ) + db_session_with_containers.add(runtime_session) + db_session_with_containers.commit() + + with patch.object(AgentAppRuntimeSessionStore, "mark_cleaned", side_effect=RuntimeError("cleanup failed")): + ConversationService.delete( + app_model=app_model, + conversation_id=conversation.id, + user=user, + session=db_session_with_containers, + ) + + deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation.id)) + assert deleted is None + mock_delete_task.delay.assert_called_once_with(conversation.id) + mock_cleanup_task.delay.assert_called_once() + + @patch("services.conversation_service.cleanup_conversation_agent_runtime_session") + @patch("services.conversation_service.delete_conversation_related_data") + def test_delete_ignores_cleanup_enqueue_failure_and_still_retires_runtime_session( + self, + mock_delete_task, + mock_cleanup_task, + db_session_with_containers: Session, + ): + app_model, user = ConversationServiceIntegrationTestDataFactory.create_app_and_account( + db_session_with_containers + ) + conversation = ConversationServiceIntegrationTestDataFactory.create_conversation( + db_session_with_containers, + app_model, + user, + ) + conversation_id = conversation.id + runtime_session = AgentRuntimeSession( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id=str(uuid4()), + agent_config_snapshot_id=str(uuid4()), + backend_run_id="backend-run-enqueue-failure", + session_snapshot=CompositorSessionSnapshot(layers=[]).model_dump_json(), + composition_layer_specs='[{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + conversation_id=conversation.id, + status=AgentRuntimeSessionStatus.ACTIVE, + ) + db_session_with_containers.add(runtime_session) + db_session_with_containers.commit() + mock_cleanup_task.delay.side_effect = RuntimeError("queue down") + + ConversationService.delete( + app_model=app_model, + conversation_id=conversation_id, + user=user, + session=db_session_with_containers, + ) + + deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation_id)) + assert deleted is None + mock_delete_task.delay.assert_called_once_with(conversation_id) + mock_cleanup_task.delay.assert_called_once() + runtime_session_row = db_session_with_containers.scalar( + select(AgentRuntimeSession).where(AgentRuntimeSession.id == runtime_session.id) + ) + assert runtime_session_row is not None + assert runtime_session_row.status == AgentRuntimeSessionStatus.CLEANED diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py b/api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py index 33d4563904e..9a725b06b64 100644 --- a/api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py +++ b/api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py @@ -6,10 +6,9 @@ from uuid import uuid4 import pytest from flask import Flask -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom -from extensions.ext_database import db from graphon.variables import FloatVariable, IntegerVariable, StringVariable from models.account import Account, Tenant, TenantAccountJoin from models.enums import ConversationFromSource, EndUserType @@ -152,13 +151,6 @@ class ConversationServiceVariableIntegrationFactory: @pytest.fixture def real_conversation_service_session_factory(flask_app_with_containers: Flask): del flask_app_with_containers - real_session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - - with ( - patch("services.conversation_service.session_factory.create_session", side_effect=lambda: real_session_maker()), - patch("services.conversation_service.session_factory.get_session_maker", return_value=real_session_maker), - ): - yield class TestConversationServiceVariables: @@ -193,6 +185,7 @@ class TestConversationServiceVariables: user=account, limit=10, last_id=None, + session=db_session_with_containers, ) assert [item["id"] for item in result.data] == [first_variable.id, second_variable.id] @@ -237,6 +230,7 @@ class TestConversationServiceVariables: user=account, limit=10, last_id=first_variable.id, + session=db_session_with_containers, ) assert [item["id"] for item in result.data] == [second_variable.id, third_variable.id] @@ -257,6 +251,7 @@ class TestConversationServiceVariables: user=account, limit=10, last_id=str(uuid4()), + session=db_session_with_containers, ) def test_get_conversational_variable_sets_has_more( @@ -282,6 +277,7 @@ class TestConversationServiceVariables: user=account, limit=2, last_id=None, + session=db_session_with_containers, ) assert len(result.data) == 2 @@ -309,6 +305,7 @@ class TestConversationServiceVariables: variable_id=existing.id, user=account, new_value="support", + session=db_session_with_containers, ) db_session_with_containers.expire_all() @@ -335,6 +332,7 @@ class TestConversationServiceVariables: variable_id=str(uuid4()), user=account, new_value="support", + session=db_session_with_containers, ) def test_update_conversation_variable_type_mismatch_raises_error( @@ -358,6 +356,7 @@ class TestConversationServiceVariables: variable_id=existing.id, user=account, new_value="wrong-type", + session=db_session_with_containers, ) def test_update_conversation_variable_integer_number_compatibility( @@ -380,6 +379,7 @@ class TestConversationServiceVariables: variable_id=existing.id, user=account, new_value=42, + session=db_session_with_containers, ) db_session_with_containers.expire_all() diff --git a/api/tests/test_containers_integration_tests/services/test_credit_pool_service.py b/api/tests/test_containers_integration_tests/services/test_credit_pool_service.py index de8e6ba612c..9cbe5252bbf 100644 --- a/api/tests/test_containers_integration_tests/services/test_credit_pool_service.py +++ b/api/tests/test_containers_integration_tests/services/test_credit_pool_service.py @@ -4,10 +4,8 @@ from unittest.mock import patch from uuid import uuid4 import pytest -from flask import has_app_context from sqlalchemy.orm import Session -from core.db.session_factory import session_factory from core.errors.error import QuotaExceededError from models import TenantCreditPool from models.enums import ProviderQuotaType @@ -35,11 +33,10 @@ class TestCreditPoolService: db_session.add(pool) db_session.commit() - @pytest.mark.usefixtures("db_session_with_containers") - def test_create_default_pool(self) -> None: + def test_create_default_pool(self, db_session_with_containers: Session) -> None: tenant_id = self._create_tenant_id() - pool = CreditPoolService.create_default_pool(tenant_id) + pool = CreditPoolService.create_default_pool(tenant_id, session=db_session_with_containers) assert isinstance(pool, TenantCreditPool) assert pool.tenant_id == tenant_id @@ -51,43 +48,46 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=0) - result = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL) + result = CreditPoolService.get_pool( + tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL, session=db_session_with_containers + ) assert result is not None assert result.tenant_id == tenant_id assert result.pool_type == ProviderQuotaType.TRIAL - @pytest.mark.usefixtures("flask_app_with_containers") - def test_get_pool_uses_configured_session_factory_without_flask_app_context(self) -> None: + def test_get_pool_uses_provided_session(self, db_session_with_containers: Session) -> None: tenant_id = self._create_tenant_id() - session_maker = session_factory.get_session_maker() - with session_maker.begin() as session: - session.add( - TenantCreditPool( - tenant_id=tenant_id, - pool_type=ProviderQuotaType.TRIAL, - quota_limit=10, - quota_used=2, - ) + db_session_with_containers.add( + TenantCreditPool( + tenant_id=tenant_id, + pool_type=ProviderQuotaType.TRIAL, + quota_limit=10, + quota_used=2, ) + ) + db_session_with_containers.commit() - assert not has_app_context() - result = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL) + result = CreditPoolService.get_pool( + tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL, session=db_session_with_containers + ) assert result is not None assert result.tenant_id == tenant_id assert result.pool_type == ProviderQuotaType.TRIAL assert result.quota_used == 2 - @pytest.mark.usefixtures("flask_app_with_containers") - def test_get_pool_returns_none_when_not_exists(self) -> None: - result = CreditPoolService.get_pool(tenant_id=self._create_tenant_id(), pool_type=ProviderQuotaType.TRIAL) + def test_get_pool_returns_none_when_not_exists(self, db_session_with_containers: Session) -> None: + result = CreditPoolService.get_pool( + tenant_id=self._create_tenant_id(), pool_type=ProviderQuotaType.TRIAL, session=db_session_with_containers + ) assert result is None - @pytest.mark.usefixtures("flask_app_with_containers") - def test_check_credits_available_returns_false_when_no_pool(self) -> None: - result = CreditPoolService.check_credits_available(tenant_id=self._create_tenant_id(), credits_required=10) + def test_check_credits_available_returns_false_when_no_pool(self, db_session_with_containers: Session) -> None: + result = CreditPoolService.check_credits_available( + tenant_id=self._create_tenant_id(), credits_required=10, session=db_session_with_containers + ) assert result is False @@ -95,7 +95,9 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=0) - result = CreditPoolService.check_credits_available(tenant_id=tenant_id, credits_required=10) + result = CreditPoolService.check_credits_available( + tenant_id=tenant_id, credits_required=10, session=db_session_with_containers + ) assert result is True @@ -103,14 +105,17 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=10) - result = CreditPoolService.check_credits_available(tenant_id=tenant_id, credits_required=1) + result = CreditPoolService.check_credits_available( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) assert result is False - @pytest.mark.usefixtures("flask_app_with_containers") - def test_check_and_deduct_credits_raises_when_no_pool(self) -> None: + def test_check_and_deduct_credits_raises_when_no_pool(self, db_session_with_containers: Session) -> None: with pytest.raises(QuotaExceededError, match="Credit pool not found"): - CreditPoolService.check_and_deduct_credits(tenant_id=self._create_tenant_id(), credits_required=1) + CreditPoolService.check_and_deduct_credits( + tenant_id=self._create_tenant_id(), credits_required=1, session=db_session_with_containers + ) def test_check_and_deduct_credits_returns_zero_for_non_positive_request( self, db_session_with_containers: Session @@ -118,10 +123,12 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=2) - result = CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=0) + result = CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=0, session=db_session_with_containers + ) assert result == 0 - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 2 @@ -130,9 +137,11 @@ class TestCreditPoolService: self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=10) with pytest.raises(QuotaExceededError, match="No credits remaining"): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1) + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 10 @@ -141,10 +150,12 @@ class TestCreditPoolService: self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=2) credits_required = 3 - result = CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=credits_required) + result = CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=credits_required, session=db_session_with_containers + ) assert result == credits_required - pool = CreditPoolService.get_pool(tenant_id=tenant_id) + pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert pool is not None assert pool.quota_used == 5 @@ -155,9 +166,11 @@ class TestCreditPoolService: self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=9) with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=3) + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=3, session=db_session_with_containers + ) - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 9 @@ -171,9 +184,11 @@ class TestCreditPoolService: patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1) + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 2 @@ -181,10 +196,12 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=9) - result = CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=3) + result = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=3, session=db_session_with_containers + ) assert result == 1 - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 10 @@ -194,16 +211,19 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=2) - result = CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=0) + result = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=0, session=db_session_with_containers + ) assert result == 0 - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 2 - @pytest.mark.usefixtures("flask_app_with_containers") - def test_deduct_credits_capped_returns_zero_when_no_pool(self) -> None: - result = CreditPoolService.deduct_credits_capped(tenant_id=self._create_tenant_id(), credits_required=1) + def test_deduct_credits_capped_returns_zero_when_no_pool(self, db_session_with_containers: Session) -> None: + result = CreditPoolService.deduct_credits_capped( + tenant_id=self._create_tenant_id(), credits_required=1, session=db_session_with_containers + ) assert result == 0 @@ -211,10 +231,12 @@ class TestCreditPoolService: tenant_id = self._create_tenant_id() self._create_pool(db_session_with_containers, tenant_id=tenant_id, quota_limit=10, quota_used=10) - result = CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + result = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) assert result == 0 - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 10 @@ -226,9 +248,11 @@ class TestCreditPoolService: patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 2 @@ -240,8 +264,10 @@ class TestCreditPoolService: patch.object(CreditPoolService, "_get_locked_pool", side_effect=QuotaExceededError("quota unavailable")), pytest.raises(QuotaExceededError, match="quota unavailable"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=1, session=db_session_with_containers + ) - updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id) + updated_pool = CreditPoolService.get_pool(tenant_id=tenant_id, session=db_session_with_containers) assert updated_pool is not None assert updated_pool.quota_used == 2 diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service.py b/api/tests/test_containers_integration_tests/services/test_dataset_service.py index 40c00267043..912e00b0b7d 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service.py @@ -602,10 +602,7 @@ class TestDatasetServiceUpdateAndDeleteDataset: # Act / Assert with pytest.raises(ValueError, match="Dataset name already exists"): DatasetService.update_dataset( - db_session_with_containers, - source_dataset.id, - {"name": "Existing Dataset"}, - account, + source_dataset.id, {"name": "Existing Dataset"}, account, session=db_session_with_containers ) def test_delete_dataset_with_documents_success(self, db_session_with_containers: Session): @@ -728,7 +725,7 @@ class TestDatasetServiceRetrievalConfiguration: } # Act - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, account) + result = DatasetService.update_dataset(dataset.id, update_data, account, session=db_session_with_containers) # Assert db_session_with_containers.refresh(dataset) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py index ced144e8d6e..6b32273624b 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py @@ -574,11 +574,7 @@ class TestDatasetPermissionServiceIntegration: with pytest.raises(NoPermissionError, match="does not have permission"): DatasetPermissionService.check_permission( - db_session_with_containers, - user, - dataset, - DatasetPermissionEnum.ALL_TEAM, - [], + user, dataset, DatasetPermissionEnum.ALL_TEAM, [], session=db_session_with_containers ) def test_check_permission_prevents_dataset_operator_from_changing_permission_mode( @@ -589,11 +585,7 @@ class TestDatasetPermissionServiceIntegration: with pytest.raises(NoPermissionError, match="cannot change the dataset permissions"): DatasetPermissionService.check_permission( - db_session_with_containers, - user, - dataset, - DatasetPermissionEnum.ONLY_ME, - [], + user, dataset, DatasetPermissionEnum.ONLY_ME, [], session=db_session_with_containers ) def test_check_permission_requires_partial_member_list_for_partial_members_mode( @@ -604,11 +596,7 @@ class TestDatasetPermissionServiceIntegration: with pytest.raises(ValueError, match="Partial member list is required"): DatasetPermissionService.check_permission( - db_session_with_containers, - user, - dataset, - DatasetPermissionEnum.PARTIAL_TEAM, - [], + user, dataset, DatasetPermissionEnum.PARTIAL_TEAM, [], session=db_session_with_containers ) def test_check_permission_rejects_dataset_operator_member_list_changes(self, db_session_with_containers: Session): @@ -618,11 +606,11 @@ class TestDatasetPermissionServiceIntegration: with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]): with pytest.raises(ValueError, match="cannot change the dataset permissions"): DatasetPermissionService.check_permission( - db_session_with_containers, user, dataset, DatasetPermissionEnum.PARTIAL_TEAM, [{"user_id": "user-2"}], + session=db_session_with_containers, ) def test_check_permission_allows_dataset_operator_when_member_list_is_unchanged( @@ -633,11 +621,11 @@ class TestDatasetPermissionServiceIntegration: with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]): DatasetPermissionService.check_permission( - db_session_with_containers, user, dataset, DatasetPermissionEnum.PARTIAL_TEAM, [{"user_id": "user-1"}], + session=db_session_with_containers, ) def test_clear_partial_member_list_deletes_permissions_and_commits(self, db_session_with_containers: Session): diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py index f719a465dbd..d9fb23e8e33 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py @@ -189,7 +189,7 @@ class TestDatasetServiceUpdateDataset: "external_knowledge_api_id": external_api.id, } - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) db_session_with_containers.refresh(dataset) updated_binding = db_session_with_containers.query(ExternalKnowledgeBindings).filter_by(id=binding_id).first() @@ -221,7 +221,7 @@ class TestDatasetServiceUpdateDataset: update_data = {"name": "new_name", "external_knowledge_api_id": str(uuid4())} with pytest.raises(ValueError) as context: - DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) assert "External knowledge id is required" in str(context.value) db_session_with_containers.rollback() @@ -245,7 +245,7 @@ class TestDatasetServiceUpdateDataset: update_data = {"name": "new_name", "external_knowledge_id": "knowledge_id"} with pytest.raises(ValueError) as context: - DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) assert "External knowledge api id is required" in str(context.value) db_session_with_containers.rollback() @@ -272,7 +272,7 @@ class TestDatasetServiceUpdateDataset: } with pytest.raises(ValueError) as context: - DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) assert "External knowledge binding not found" in str(context.value) db_session_with_containers.rollback() @@ -303,7 +303,7 @@ class TestDatasetServiceUpdateDataset: "embedding_model": "text-embedding-ada-002", } - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) db_session_with_containers.refresh(dataset) assert dataset.name == "new_name" @@ -338,7 +338,7 @@ class TestDatasetServiceUpdateDataset: "embedding_model": None, } - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) db_session_with_containers.refresh(dataset) assert dataset.name == "new_name" @@ -371,7 +371,7 @@ class TestDatasetServiceUpdateDataset: } with patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task: - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) mock_task.delay.assert_called_once_with(dataset.id, "remove") db_session_with_containers.refresh(dataset) @@ -418,7 +418,7 @@ class TestDatasetServiceUpdateDataset: mock_model_manager.return_value.get_model_instance.return_value = embedding_model mock_get_binding.return_value = binding - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) mock_model_manager.return_value.get_model_instance.assert_called_once_with( tenant_id=tenant.id, @@ -462,7 +462,7 @@ class TestDatasetServiceUpdateDataset: "retrieval_model": "new_model", } - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) db_session_with_containers.refresh(dataset) assert dataset.name == "new_name" @@ -514,7 +514,7 @@ class TestDatasetServiceUpdateDataset: mock_model_manager.return_value.get_model_instance.return_value = embedding_model mock_get_binding.return_value = binding - result = DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) mock_model_manager.return_value.get_model_instance.assert_called_once_with( tenant_id=tenant.id, @@ -545,7 +545,7 @@ class TestDatasetServiceUpdateDataset: update_data = {"name": "new_name"} with pytest.raises(ValueError) as context: - DatasetService.update_dataset(db_session_with_containers, str(uuid4()), update_data, user) + DatasetService.update_dataset(str(uuid4()), update_data, user, session=db_session_with_containers) assert "Dataset not found" in str(context.value) @@ -568,7 +568,7 @@ class TestDatasetServiceUpdateDataset: update_data = {"name": "new_name"} with pytest.raises(NoPermissionError): - DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, outsider) + DatasetService.update_dataset(dataset.id, update_data, outsider, session=db_session_with_containers) def test_update_internal_dataset_embedding_model_error(self, db_session_with_containers: Session): """Test error when embedding model is not available.""" @@ -595,6 +595,6 @@ class TestDatasetServiceUpdateDataset: mock_model_manager.return_value.get_model_instance.side_effect = Exception("No Embedding Model available") with pytest.raises(Exception) as context: - DatasetService.update_dataset(db_session_with_containers, dataset.id, update_data, user) + DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) assert "No Embedding Model available".lower() in str(context.value).lower() diff --git a/api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py b/api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py index 5eb84f805aa..a5445a17297 100644 --- a/api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py +++ b/api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py @@ -69,7 +69,7 @@ def test_build_upload_files_zip_tempfile_sanitizes_and_dedupes_names(monkeypatch def test_get_upload_files_by_ids_returns_empty_when_no_ids(db_session_with_containers: Session) -> None: """Ensure empty input returns an empty mapping without hitting the database.""" - assert FileService.get_upload_files_by_ids(db_session_with_containers, str(uuid4()), []) == {} + assert FileService.get_upload_files_by_ids(str(uuid4()), [], session=db_session_with_containers) == {} def test_get_upload_files_by_ids_returns_id_keyed_mapping(db_session_with_containers: Session) -> None: @@ -78,7 +78,9 @@ def test_get_upload_files_by_ids_returns_id_keyed_mapping(db_session_with_contai file1 = _create_upload_file(db_session_with_containers, tenant_id=tenant_id, key="k1", name="file1.txt") file2 = _create_upload_file(db_session_with_containers, tenant_id=tenant_id, key="k2", name="file2.txt") - result = FileService.get_upload_files_by_ids(db_session_with_containers, tenant_id, [file1.id, file1.id, file2.id]) + result = FileService.get_upload_files_by_ids( + tenant_id, [file1.id, file1.id, file2.id], session=db_session_with_containers + ) assert set(result.keys()) == {file1.id, file2.id} assert result[file1.id].id == file1.id @@ -92,6 +94,6 @@ def test_get_upload_files_by_ids_filters_by_tenant(db_session_with_containers: S file_a = _create_upload_file(db_session_with_containers, tenant_id=tenant_a, key="ka", name="a.txt") _create_upload_file(db_session_with_containers, tenant_id=tenant_b, key="kb", name="b.txt") - result = FileService.get_upload_files_by_ids(db_session_with_containers, tenant_a, [file_a.id]) + result = FileService.get_upload_files_by_ids(tenant_a, [file_a.id], session=db_session_with_containers) assert set(result.keys()) == {file_a.id} diff --git a/api/tests/test_containers_integration_tests/services/test_hit_testing_service.py b/api/tests/test_containers_integration_tests/services/test_hit_testing_service.py index 4a73f98f50e..67b3a2d3e57 100644 --- a/api/tests/test_containers_integration_tests/services/test_hit_testing_service.py +++ b/api/tests/test_containers_integration_tests/services/test_hit_testing_service.py @@ -192,7 +192,7 @@ class TestHitTestingService: mock_format.return_value = [mock_record] response = _RetrieveResponse.model_validate( - HitTestingService.compact_retrieve_response(db_session_with_containers, query, [mock_doc]) + HitTestingService.compact_retrieve_response(query, [mock_doc], session=db_session_with_containers) ) assert response.query.content == query @@ -246,12 +246,12 @@ class TestHitTestingService: response = _RetrieveResponse.model_validate( HitTestingService.external_retrieve( - db_session_with_containers, dataset=dataset, query='test "query"', account=account, external_retrieval_model={"model": "test"}, metadata_filtering_conditions={"key": "val"}, + session=db_session_with_containers, ) ) @@ -276,7 +276,7 @@ class TestHitTestingService: account = MagicMock() response = _RetrieveResponse.model_validate( - HitTestingService.external_retrieve(db_session_with_containers, dataset, "test query", account) + HitTestingService.external_retrieve(dataset, "test query", account, session=db_session_with_containers) ) assert response.query.content == "test query" @@ -300,12 +300,12 @@ class TestHitTestingService: response = _RetrieveResponse.model_validate( HitTestingService.retrieve( - db_session_with_containers, dataset=dataset, query="test query", account=account, retrieval_model=None, external_retrieval_model=external_retrieval_model, + session=db_session_with_containers, ) ) @@ -343,12 +343,12 @@ class TestHitTestingService: mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( - db_session_with_containers, dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, + session=db_session_with_containers, ) mock_get_meta.assert_called_once() @@ -380,12 +380,12 @@ class TestHitTestingService: response = _RetrieveResponse.model_validate( HitTestingService.retrieve( - db_session_with_containers, dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, + session=db_session_with_containers, ) ) @@ -412,13 +412,13 @@ class TestHitTestingService: mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( - db_session_with_containers, dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, attachment_ids=attachment_ids, + session=db_session_with_containers, ) mock_retrieve.assert_called_once_with( @@ -472,12 +472,12 @@ class TestHitTestingService: mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( - db_session_with_containers, dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, + session=db_session_with_containers, ) mock_retrieve.assert_called_once() diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py index c1188d3d0f9..84a0226ba17 100644 --- a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py +++ b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py @@ -119,6 +119,7 @@ def test_human_input_delivery_test_sends_email( account=account, node_id="human-node", delivery_method_id=str(delivery_method_id), + session=db_session_with_containers, ) assert send_mock.call_count == 1 @@ -145,6 +146,7 @@ def test_human_input_delivery_test_form_accepts_file_upload( account=account, node_id="human-node", delivery_method_id=str(delivery_method_id), + session=db_session_with_containers, ) form = db_session_with_containers.scalar( @@ -213,6 +215,7 @@ def test_human_input_delivery_test_form_accepts_remote_file_upload( account=account, node_id="human-node", delivery_method_id=str(delivery_method_id), + session=db_session_with_containers, ) form = db_session_with_containers.scalar( diff --git a/api/tests/test_containers_integration_tests/services/test_message_service.py b/api/tests/test_containers_integration_tests/services/test_message_service.py index f2d682be3bf..702812b96de 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service.py +++ b/api/tests/test_containers_integration_tests/services/test_message_service.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest from faker import Faker @@ -117,7 +117,7 @@ class TestMessageService: # Create app app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Setup current_user mock self._mock_current_user(mock_external_service_dependencies, account.id, tenant.id) @@ -222,6 +222,7 @@ class TestMessageService: first_id=messages[2].id, # Use middle message as first_id limit=2, order="asc", + session=db_session_with_containers, ) # Verify results @@ -243,7 +244,12 @@ class TestMessageService: # Test pagination with no user result = MessageService.pagination_by_first_id( - app_model=app, user=None, conversation_id=fake.uuid4(), first_id=None, limit=10 + app_model=app, + user=None, + conversation_id=fake.uuid4(), + first_id=None, + limit=10, + session=db_session_with_containers, ) # Verify empty result @@ -262,7 +268,12 @@ class TestMessageService: # Test pagination with no conversation ID result = MessageService.pagination_by_first_id( - app_model=app, user=account, conversation_id="", first_id=None, limit=10 + app_model=app, + user=account, + conversation_id="", + first_id=None, + limit=10, + session=db_session_with_containers, ) # Verify empty result @@ -291,6 +302,7 @@ class TestMessageService: conversation_id=conversation.id, first_id=fake.uuid4(), # Non-existent message ID limit=10, + session=db_session_with_containers, ) def test_pagination_by_last_id_success( @@ -316,6 +328,7 @@ class TestMessageService: last_id=messages[2].id, # Use middle message as last_id limit=2, conversation_id=conversation.id, + session=db_session_with_containers, ) # Verify results @@ -345,7 +358,12 @@ class TestMessageService: # Test pagination with include_ids include_ids = [messages[0].id, messages[1].id, messages[2].id] result = MessageService.pagination_by_last_id( - app_model=app, user=account, last_id=messages[1].id, limit=2, include_ids=include_ids + app_model=app, + user=account, + last_id=messages[1].id, + limit=2, + include_ids=include_ids, + session=db_session_with_containers, ) # Verify results @@ -364,8 +382,10 @@ class TestMessageService: fake = Faker() app, account = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies) - # Test pagination with no user - result = MessageService.pagination_by_last_id(app_model=app, user=None, last_id=None, limit=10) + # Test pagination with no user, + result = MessageService.pagination_by_last_id( + app_model=app, user=None, last_id=None, limit=10, session=db_session_with_containers + ) # Verify empty result assert result.limit == 10 @@ -393,6 +413,7 @@ class TestMessageService: last_id=fake.uuid4(), # Non-existent message ID limit=10, conversation_id=conversation.id, + session=db_session_with_containers, ) def test_create_feedback_success(self, db_session_with_containers: Session, mock_external_service_dependencies): @@ -410,7 +431,12 @@ class TestMessageService: rating = FeedbackRating.LIKE content = fake.text(max_nb_chars=100) feedback = MessageService.create_feedback( - app_model=app, message_id=message.id, user=account, rating=rating, content=content + app_model=app, + message_id=message.id, + user=account, + rating=rating, + content=content, + session=db_session_with_containers, ) # Verify feedback was created correctly @@ -442,6 +468,7 @@ class TestMessageService: user=None, rating=FeedbackRating.LIKE, content=fake.text(max_nb_chars=100), + session=db_session_with_containers, ) def test_create_feedback_update_existing( @@ -461,14 +488,24 @@ class TestMessageService: initial_rating = FeedbackRating.LIKE initial_content = fake.text(max_nb_chars=100) feedback = MessageService.create_feedback( - app_model=app, message_id=message.id, user=account, rating=initial_rating, content=initial_content + app_model=app, + message_id=message.id, + user=account, + rating=initial_rating, + content=initial_content, + session=db_session_with_containers, ) # Update feedback updated_rating = FeedbackRating.DISLIKE updated_content = fake.text(max_nb_chars=100) updated_feedback = MessageService.create_feedback( - app_model=app, message_id=message.id, user=account, rating=updated_rating, content=updated_content + app_model=app, + message_id=message.id, + user=account, + rating=updated_rating, + content=updated_content, + session=db_session_with_containers, ) # Verify feedback was updated correctly @@ -498,10 +535,18 @@ class TestMessageService: user=account, rating=FeedbackRating.LIKE, content=fake.text(max_nb_chars=100), + session=db_session_with_containers, ) - # Delete feedback by setting rating to None - MessageService.create_feedback(app_model=app, message_id=message.id, user=account, rating=None, content=None) + # Delete feedback by setting rating to None, + MessageService.create_feedback( + app_model=app, + message_id=message.id, + user=account, + rating=None, + content=None, + session=db_session_with_containers, + ) # Verify feedback was deleted @@ -526,7 +571,12 @@ class TestMessageService: # Test creating feedback with no rating when no feedback exists with pytest.raises(ValueError, match="rating cannot be None when feedback not exists"): MessageService.create_feedback( - app_model=app, message_id=message.id, user=account, rating=None, content=None + app_model=app, + message_id=message.id, + user=account, + rating=None, + content=None, + session=db_session_with_containers, ) def test_get_all_messages_feedbacks_success( @@ -550,11 +600,12 @@ class TestMessageService: user=account, rating=FeedbackRating.LIKE if i % 2 == 0 else FeedbackRating.DISLIKE, content=f"Feedback {i}: {fake.text(max_nb_chars=50)}", + session=db_session_with_containers, ) feedbacks.append(feedback) - # Get all feedbacks - result = MessageService.get_all_messages_feedbacks(app, page=1, limit=10) + # Get all feedbacks, + result = MessageService.get_all_messages_feedbacks(app, page=1, limit=10, session=db_session_with_containers) # Verify results assert len(result) == 3 @@ -583,11 +634,16 @@ class TestMessageService: user=account, rating=FeedbackRating.LIKE, content=f"Feedback {i}", + session=db_session_with_containers, ) # Get feedbacks with pagination - result_page_1 = MessageService.get_all_messages_feedbacks(app, page=1, limit=3) - result_page_2 = MessageService.get_all_messages_feedbacks(app, page=2, limit=3) + result_page_1 = MessageService.get_all_messages_feedbacks( + app, page=1, limit=3, session=db_session_with_containers + ) + result_page_2 = MessageService.get_all_messages_feedbacks( + app, page=2, limit=3, session=db_session_with_containers + ) # Verify pagination results assert len(result_page_1) == 3 @@ -609,8 +665,10 @@ class TestMessageService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) message = self._create_test_message(db_session_with_containers, app, conversation, account, fake) - # Get message - retrieved_message = MessageService.get_message(app_model=app, user=account, message_id=message.id) + # Get message, + retrieved_message = MessageService.get_message( + app_model=app, user=account, message_id=message.id, session=db_session_with_containers + ) # Verify message was retrieved correctly assert retrieved_message.id == message.id @@ -628,7 +686,9 @@ class TestMessageService: # Test getting non-existent message with pytest.raises(MessageNotExistsError): - MessageService.get_message(app_model=app, user=account, message_id=fake.uuid4()) + MessageService.get_message( + app_model=app, user=account, message_id=fake.uuid4(), session=db_session_with_containers + ) def test_get_message_wrong_user(self, db_session_with_containers: Session, mock_external_service_dependencies): """ @@ -657,7 +717,9 @@ class TestMessageService: # Test getting message with different user with pytest.raises(MessageNotExistsError): - MessageService.get_message(app_model=app, user=other_account, message_id=message.id) + MessageService.get_message( + app_model=app, user=other_account, message_id=message.id, session=db_session_with_containers + ) def test_get_suggested_questions_after_answer_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -682,7 +744,11 @@ class TestMessageService: from core.app.entities.app_invoke_entities import InvokeFrom result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=account, message_id=message.id, invoke_from=InvokeFrom.SERVICE_API + app_model=app, + user=account, + message_id=message.id, + invoke_from=InvokeFrom.SERVICE_API, + session=db_session_with_containers, ) # Verify results @@ -714,7 +780,11 @@ class TestMessageService: with pytest.raises(ValueError, match="user cannot be None"): MessageService.get_suggested_questions_after_answer( - app_model=app, user=None, message_id=message.id, invoke_from=InvokeFrom.SERVICE_API + app_model=app, + user=None, + message_id=message.id, + invoke_from=InvokeFrom.SERVICE_API, + session=db_session_with_containers, ) def test_get_suggested_questions_after_answer_disabled( @@ -740,7 +810,11 @@ class TestMessageService: with pytest.raises(SuggestedQuestionsAfterAnswerDisabledError): MessageService.get_suggested_questions_after_answer( - app_model=app, user=account, message_id=message.id, invoke_from=InvokeFrom.SERVICE_API + app_model=app, + user=account, + message_id=message.id, + invoke_from=InvokeFrom.SERVICE_API, + session=db_session_with_containers, ) def test_get_suggested_questions_after_answer_no_workflow( @@ -763,7 +837,11 @@ class TestMessageService: from core.app.entities.app_invoke_entities import InvokeFrom result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=account, message_id=message.id, invoke_from=InvokeFrom.SERVICE_API + app_model=app, + user=account, + message_id=message.id, + invoke_from=InvokeFrom.SERVICE_API, + session=db_session_with_containers, ) # Verify empty result @@ -792,7 +870,11 @@ class TestMessageService: from core.app.entities.app_invoke_entities import InvokeFrom result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=account, message_id=message.id, invoke_from=InvokeFrom.DEBUGGER + app_model=app, + user=account, + message_id=message.id, + invoke_from=InvokeFrom.DEBUGGER, + session=db_session_with_containers, ) # Verify results @@ -800,7 +882,7 @@ class TestMessageService: # Verify draft workflow was used instead of published workflow mock_external_service_dependencies["workflow_service"].return_value.get_draft_workflow.assert_called_once_with( - app_model=app + app_model=app, session=ANY ) # Verify TraceQueueManager was called diff --git a/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py b/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py index 6a9046acd4a..9da20d96f06 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py +++ b/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py @@ -20,6 +20,7 @@ def test_pagination_returns_extra_contents(db_session_with_containers: Session): conversation_id=fixture.conversation.id, first_id=None, limit=10, + session=db_session_with_containers, ) assert pagination.data @@ -59,6 +60,7 @@ def test_pagination_returns_waiting_human_input_extra_contents(db_session_with_c conversation_id=fixture.conversation.id, first_id=None, limit=10, + session=db_session_with_containers, ) assert pagination.data diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py b/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py index fbdc265265d..a9399985307 100644 --- a/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py +++ b/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py @@ -95,7 +95,9 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata(db_session_with_containers, dataset, metadata_args, current_account) + MetadataService.update_documents_metadata( + dataset, metadata_args, current_account, session=db_session_with_containers + ) db_session_with_containers.expire_all() updated_doc = db_session_with_containers.get(Document, document.id) @@ -126,7 +128,9 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata(db_session_with_containers, dataset, metadata_args, current_account) + MetadataService.update_documents_metadata( + dataset, metadata_args, current_account, session=db_session_with_containers + ) db_session_with_containers.expire_all() updated_doc = db_session_with_containers.get(Document, document.id) @@ -168,7 +172,9 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata(db_session_with_containers, dataset, metadata_args, current_account) + MetadataService.update_documents_metadata( + dataset, metadata_args, current_account, session=db_session_with_containers + ) db_session_with_containers.expire_all() bindings = db_session_with_containers.scalars( @@ -205,5 +211,5 @@ class TestMetadataPartialUpdate: with patch.object(db_session_with_containers, "commit", side_effect=RuntimeError("database connection lost")): with pytest.raises(RuntimeError, match="database connection lost"): MetadataService.update_documents_metadata( - db_session_with_containers, dataset, metadata_args, current_account + dataset, metadata_args, current_account, session=db_session_with_containers ) diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_service.py b/api/tests/test_containers_integration_tests/services/test_metadata_service.py index 7cc9fc7e696..00afe7f8467 100644 --- a/api/tests/test_containers_integration_tests/services/test_metadata_service.py +++ b/api/tests/test_containers_integration_tests/services/test_metadata_service.py @@ -184,7 +184,7 @@ class TestMetadataService: # Act: Execute the method under test result = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Assert: Verify the expected outcomes @@ -220,7 +220,9 @@ class TestMetadataService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name cannot exceed 255 characters."): - MetadataService.create_metadata(db_session_with_containers, dataset.id, metadata_args, account, tenant.id) + MetadataService.create_metadata( + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + ) def test_create_metadata_name_already_exists( self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps @@ -238,7 +240,9 @@ class TestMetadataService: # Create first metadata first_metadata_args = MetadataArgs(type="string", name="duplicate_name") - MetadataService.create_metadata(db_session_with_containers, dataset.id, first_metadata_args, account, tenant.id) + MetadataService.create_metadata( + dataset.id, first_metadata_args, account, tenant.id, session=db_session_with_containers + ) # Try to create second metadata with same name second_metadata_args = MetadataArgs(type="number", name="duplicate_name") @@ -246,7 +250,7 @@ class TestMetadataService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name already exists."): MetadataService.create_metadata( - db_session_with_containers, dataset.id, second_metadata_args, account, tenant.id + dataset.id, second_metadata_args, account, tenant.id, session=db_session_with_containers ) def test_create_metadata_name_conflicts_with_built_in_field( @@ -269,7 +273,9 @@ class TestMetadataService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name already exists in Built-in fields."): - MetadataService.create_metadata(db_session_with_containers, dataset.id, metadata_args, account, tenant.id) + MetadataService.create_metadata( + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + ) def test_update_metadata_name_success( self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps @@ -288,13 +294,13 @@ class TestMetadataService: # Create metadata first metadata_args = MetadataArgs(type="string", name="old_name") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Act: Execute the method under test new_name = "new_name" result = MetadataService.update_metadata_name( - db_session_with_containers, dataset.id, metadata.id, new_name, account, tenant.id + dataset.id, metadata.id, new_name, account, tenant.id, session=db_session_with_containers ) # Assert: Verify the expected outcomes @@ -325,7 +331,7 @@ class TestMetadataService: # Create metadata first metadata_args = MetadataArgs(type="string", name="old_name") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Try to update with too long name @@ -334,7 +340,7 @@ class TestMetadataService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name cannot exceed 255 characters."): MetadataService.update_metadata_name( - db_session_with_containers, dataset.id, metadata.id, long_name, account, tenant.id + dataset.id, metadata.id, long_name, account, tenant.id, session=db_session_with_containers ) def test_update_metadata_name_already_exists( @@ -354,18 +360,18 @@ class TestMetadataService: # Create two metadata entries first_metadata_args = MetadataArgs(type="string", name="first_metadata") first_metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, first_metadata_args, account, tenant.id + dataset.id, first_metadata_args, account, tenant.id, session=db_session_with_containers ) second_metadata_args = MetadataArgs(type="number", name="second_metadata") second_metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, second_metadata_args, account, tenant.id + dataset.id, second_metadata_args, account, tenant.id, session=db_session_with_containers ) # Try to update first metadata with second metadata's name with pytest.raises(ValueError, match="Metadata name already exists."): MetadataService.update_metadata_name( - db_session_with_containers, dataset.id, first_metadata.id, "second_metadata", account, tenant.id + dataset.id, first_metadata.id, "second_metadata", account, tenant.id, session=db_session_with_containers ) def test_update_metadata_name_conflicts_with_built_in_field( @@ -385,7 +391,7 @@ class TestMetadataService: # Create metadata first metadata_args = MetadataArgs(type="string", name="old_name") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Try to update with built-in field name @@ -393,7 +399,7 @@ class TestMetadataService: with pytest.raises(ValueError, match="Metadata name already exists in Built-in fields."): MetadataService.update_metadata_name( - db_session_with_containers, dataset.id, metadata.id, built_in_field_name, account, tenant.id + dataset.id, metadata.id, built_in_field_name, account, tenant.id, session=db_session_with_containers ) def test_update_metadata_name_not_found( @@ -418,7 +424,7 @@ class TestMetadataService: # Act: Execute the method under test result = MetadataService.update_metadata_name( - db_session_with_containers, dataset.id, fake_metadata_id, new_name, account, tenant.id + dataset.id, fake_metadata_id, new_name, account, tenant.id, session=db_session_with_containers ) # Assert: Verify the method returns None when metadata is not found @@ -441,11 +447,11 @@ class TestMetadataService: # Create metadata first metadata_args = MetadataArgs(type="string", name="to_be_deleted") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Act: Execute the method under test - result = MetadataService.delete_metadata(db_session_with_containers, dataset.id, metadata.id) + result = MetadataService.delete_metadata(dataset.id, metadata.id, session=db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -476,7 +482,7 @@ class TestMetadataService: fake_metadata_id = str(uuid.uuid4()) # Use valid UUID format # Act: Execute the method under test - result = MetadataService.delete_metadata(db_session_with_containers, dataset.id, fake_metadata_id) + result = MetadataService.delete_metadata(dataset.id, fake_metadata_id, session=db_session_with_containers) # Assert: Verify the method returns None when metadata is not found assert result is None @@ -501,7 +507,7 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Create metadata binding @@ -522,7 +528,7 @@ class TestMetadataService: db_session_with_containers.commit() # Act: Execute the method under test - result = MetadataService.delete_metadata(db_session_with_containers, dataset.id, metadata.id) + result = MetadataService.delete_metadata(dataset.id, metadata.id, session=db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -587,7 +593,7 @@ class TestMetadataService: assert dataset.built_in_field_enabled is False # Act: Execute the method under test - MetadataService.enable_built_in_field(db_session_with_containers, dataset) + MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes @@ -623,7 +629,7 @@ class TestMetadataService: ]() # Act: Execute the method under test - MetadataService.enable_built_in_field(db_session_with_containers, dataset) + MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the method returns early without changes db_session_with_containers.refresh(dataset) @@ -649,7 +655,7 @@ class TestMetadataService: ]() # Act: Execute the method under test - MetadataService.enable_built_in_field(db_session_with_containers, dataset) + MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes @@ -696,7 +702,7 @@ class TestMetadataService: ] # Act: Execute the method under test - MetadataService.disable_built_in_field(db_session_with_containers, dataset) + MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes db_session_with_containers.refresh(dataset) @@ -728,7 +734,7 @@ class TestMetadataService: ]() # Act: Execute the method under test - MetadataService.disable_built_in_field(db_session_with_containers, dataset) + MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the method returns early without changes @@ -761,7 +767,7 @@ class TestMetadataService: ]() # Act: Execute the method under test - MetadataService.disable_built_in_field(db_session_with_containers, dataset) + MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes db_session_with_containers.refresh(dataset) @@ -787,7 +793,7 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Mock DocumentService.get_document @@ -807,7 +813,7 @@ class TestMetadataService: operation_data = MetadataOperationData(operation_data=[operation]) # Act: Execute the method under test - MetadataService.update_documents_metadata(db_session_with_containers, dataset, operation_data, account) + MetadataService.update_documents_metadata(dataset, operation_data, account, session=db_session_with_containers) # Assert: Verify the expected outcomes @@ -853,7 +859,7 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Mock DocumentService.get_document @@ -873,7 +879,7 @@ class TestMetadataService: operation_data = MetadataOperationData(operation_data=[operation]) # Act: Execute the method under test - MetadataService.update_documents_metadata(db_session_with_containers, dataset, operation_data, account) + MetadataService.update_documents_metadata(dataset, operation_data, account, session=db_session_with_containers) # Assert: Verify the expected outcomes # Verify document metadata was updated with both custom and built-in fields @@ -902,7 +908,7 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Create metadata operation data @@ -924,7 +930,9 @@ class TestMetadataService: # Act & Assert: The method should raise ValueError("Document not found.") # because the exception is now re-raised after rollback with pytest.raises(ValueError, match="Document not found"): - MetadataService.update_documents_metadata(db_session_with_containers, dataset, operation_data, account) + MetadataService.update_documents_metadata( + dataset, operation_data, account, session=db_session_with_containers + ) def test_knowledge_base_metadata_lock_check_dataset_id( self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps @@ -1021,7 +1029,7 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Create document and metadata binding @@ -1041,7 +1049,7 @@ class TestMetadataService: db_session_with_containers.commit() # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(db_session_with_containers, dataset) + result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -1082,11 +1090,11 @@ class TestMetadataService: # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") metadata = MetadataService.create_metadata( - db_session_with_containers, dataset.id, metadata_args, account, tenant.id + dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers ) # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(db_session_with_containers, dataset) + result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -1115,7 +1123,7 @@ class TestMetadataService: ) # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(db_session_with_containers, dataset) + result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None diff --git a/api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py b/api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py index aca38391353..71d2c1c6800 100644 --- a/api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py +++ b/api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py @@ -339,7 +339,11 @@ class TestModelLoadBalancingService: # Act: Execute the method under test service = ModelLoadBalancingService() is_enabled, configs = service.get_load_balancing_configs( - tenant_id=tenant.id, provider="openai", model="gpt-3.5-turbo", model_type="llm" + tenant_id=tenant.id, + provider="openai", + model="gpt-3.5-turbo", + model_type="llm", + session=db_session_with_containers, ) # Assert: Verify the expected outcomes @@ -381,7 +385,11 @@ class TestModelLoadBalancingService: service = ModelLoadBalancingService() with pytest.raises(ValueError) as exc_info: service.get_load_balancing_configs( - tenant_id=tenant.id, provider="nonexistent_provider", model="gpt-3.5-turbo", model_type="llm" + tenant_id=tenant.id, + provider="nonexistent_provider", + model="gpt-3.5-turbo", + model_type="llm", + session=db_session_with_containers, ) # Verify correct error message @@ -443,7 +451,11 @@ class TestModelLoadBalancingService: # Act: Execute the method under test service = ModelLoadBalancingService() is_enabled, configs = service.get_load_balancing_configs( - tenant_id=tenant.id, provider="openai", model="gpt-3.5-turbo", model_type="llm" + tenant_id=tenant.id, + provider="openai", + model="gpt-3.5-turbo", + model_type="llm", + session=db_session_with_containers, ) # Assert: Verify the expected outcomes diff --git a/api/tests/test_containers_integration_tests/services/test_oauth_server_service.py b/api/tests/test_containers_integration_tests/services/test_oauth_server_service.py index 0969198ecf3..d397c62b6a8 100644 --- a/api/tests/test_containers_integration_tests/services/test_oauth_server_service.py +++ b/api/tests/test_containers_integration_tests/services/test_oauth_server_service.py @@ -4,7 +4,7 @@ from __future__ import annotations import uuid from typing import cast -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -159,17 +159,19 @@ class TestOAuthServerServiceTokenOperations: def test_validate_access_token_returns_none_when_not_found(self, mock_redis): mock_redis.get.return_value = None + session = MagicMock() - result = OAuthServerService.validate_oauth_access_token("client-1", "missing-token") + result = OAuthServerService.validate_oauth_access_token("client-1", "missing-token", session) assert result is None def test_validate_access_token_loads_user_when_exists(self, mock_redis): mock_redis.get.return_value = b"user-88" expected_user = MagicMock() + session = MagicMock() with patch("services.oauth_server.AccountService.load_user", return_value=expected_user) as mock_load: - result = OAuthServerService.validate_oauth_access_token("client-1", "access-token") + result = OAuthServerService.validate_oauth_access_token("client-1", "access-token", session) assert result is expected_user - mock_load.assert_called_once_with("user-88", ANY) + mock_load.assert_called_once_with("user-88", session) diff --git a/api/tests/test_containers_integration_tests/services/test_ops_service.py b/api/tests/test_containers_integration_tests/services/test_ops_service.py index 9643fb61d44..b4b8521fb2e 100644 --- a/api/tests/test_containers_integration_tests/services/test_ops_service.py +++ b/api/tests/test_containers_integration_tests/services/test_ops_service.py @@ -67,6 +67,7 @@ class TestOpsService: icon_background="#FF6B6B", ), account, + session=db_session_with_containers, ) return app, account @@ -91,13 +92,13 @@ class TestOpsService: # ── get_tracing_app_config ───────────────────────────────────────── def test_get_tracing_app_config_no_config(self, db_session_with_containers: Session, mock_ops_trace_manager): - result = OpsService.get_tracing_app_config(str(uuid.uuid4()), "arize") + result = OpsService.get_tracing_app_config(str(uuid.uuid4()), "arize", db_session_with_containers) assert result is None def test_get_tracing_app_config_no_app(self, db_session_with_containers: Session, mock_ops_trace_manager): fake_app_id = str(uuid.uuid4()) self._insert_trace_config(db_session_with_containers, fake_app_id, "arize") - result = OpsService.get_tracing_app_config(fake_app_id, "arize") + result = OpsService.get_tracing_app_config(fake_app_id, "arize", db_session_with_containers) assert result is None def test_get_tracing_app_config_none_config( @@ -107,7 +108,7 @@ class TestOpsService: self._insert_trace_config(db_session_with_containers, app.id, "arize", tracing_config=None) with pytest.raises(ValueError, match="Tracing config cannot be None."): - OpsService.get_tracing_app_config(app.id, "arize") + OpsService.get_tracing_app_config(app.id, "arize", db_session_with_containers) @pytest.mark.parametrize( ("provider", "default_url"), @@ -135,7 +136,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, provider) - result = OpsService.get_tracing_app_config(app.id, provider) + result = OpsService.get_tracing_app_config(app.id, provider, db_session_with_containers) assert result is not None assert result["tracing_config"]["project_url"] == default_url @@ -155,7 +156,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, provider) - result = OpsService.get_tracing_app_config(app.id, provider) + result = OpsService.get_tracing_app_config(app.id, provider, db_session_with_containers) assert result is not None assert result["tracing_config"]["project_url"] == "success_url" @@ -171,7 +172,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, "langfuse") - result = OpsService.get_tracing_app_config(app.id, "langfuse") + result = OpsService.get_tracing_app_config(app.id, "langfuse", db_session_with_containers) assert result is not None assert result["tracing_config"]["project_url"] == "https://api.langfuse.com/project/key" @@ -187,7 +188,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, "langfuse") - result = OpsService.get_tracing_app_config(app.id, "langfuse") + result = OpsService.get_tracing_app_config(app.id, "langfuse", db_session_with_containers) assert result is not None assert result["tracing_config"]["project_url"] == "https://api.langfuse.com/" @@ -195,7 +196,9 @@ class TestOpsService: # ── create_tracing_app_config ────────────────────────────────────── def test_create_tracing_app_config_invalid_provider(self, db_session_with_containers: Session): - result = OpsService.create_tracing_app_config(str(uuid.uuid4()), "invalid_provider", {}) + result = OpsService.create_tracing_app_config( + str(uuid.uuid4()), "invalid_provider", {}, db_session_with_containers + ) assert result == {"error": "Invalid tracing provider: invalid_provider"} def test_create_tracing_app_config_invalid_credentials( @@ -203,7 +206,10 @@ class TestOpsService: ): mock_ops_trace_manager.check_trace_config_is_effective.return_value = False result = OpsService.create_tracing_app_config( - str(uuid.uuid4()), TracingProviderEnum.LANGFUSE, {"public_key": "p", "secret_key": "s"} + str(uuid.uuid4()), + TracingProviderEnum.LANGFUSE, + {"public_key": "p", "secret_key": "s"}, + db_session_with_containers, ) assert result == {"error": "Invalid Credentials"} @@ -228,7 +234,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, str(provider)) - result = OpsService.create_tracing_app_config(app.id, provider, config) + result = OpsService.create_tracing_app_config(app.id, provider, config, db_session_with_containers) assert result is None @@ -245,6 +251,7 @@ class TestOpsService: app.id, TracingProviderEnum.LANGFUSE, {"public_key": "p", "secret_key": "s", "host": "https://api.langfuse.com"}, + db_session_with_containers, ) assert result == {"result": "success"} @@ -258,13 +265,17 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, str(TracingProviderEnum.ARIZE)) - result = OpsService.create_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {}) + result = OpsService.create_tracing_app_config( + app.id, TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result is None def test_create_tracing_app_config_no_app(self, db_session_with_containers: Session, mock_ops_trace_manager): mock_ops_trace_manager.check_trace_config_is_effective.return_value = True - result = OpsService.create_tracing_app_config(str(uuid.uuid4()), TracingProviderEnum.ARIZE, {}) + result = OpsService.create_tracing_app_config( + str(uuid.uuid4()), TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result is None def test_create_tracing_app_config_with_empty_other_keys( @@ -277,7 +288,9 @@ class TestOpsService: mock_otm.encrypt_tracing_config.return_value = {} app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) - result = OpsService.create_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {"project": ""}) + result = OpsService.create_tracing_app_config( + app.id, TracingProviderEnum.ARIZE, {"project": ""}, db_session_with_containers + ) assert result == {"result": "success"} @@ -290,7 +303,9 @@ class TestOpsService: mock_otm.encrypt_tracing_config.return_value = {"encrypted": "config"} app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) - result = OpsService.create_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {}) + result = OpsService.create_tracing_app_config( + app.id, TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result == {"result": "success"} @@ -298,17 +313,21 @@ class TestOpsService: def test_update_tracing_app_config_invalid_provider(self, db_session_with_containers: Session): with pytest.raises(ValueError, match="Invalid tracing provider: invalid_provider"): - OpsService.update_tracing_app_config(str(uuid.uuid4()), "invalid_provider", {}) + OpsService.update_tracing_app_config(str(uuid.uuid4()), "invalid_provider", {}, db_session_with_containers) def test_update_tracing_app_config_no_config(self, db_session_with_containers: Session, mock_ops_trace_manager): - result = OpsService.update_tracing_app_config(str(uuid.uuid4()), TracingProviderEnum.ARIZE, {}) + result = OpsService.update_tracing_app_config( + str(uuid.uuid4()), TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result is None def test_update_tracing_app_config_no_app(self, db_session_with_containers: Session, mock_ops_trace_manager): fake_app_id = str(uuid.uuid4()) self._insert_trace_config(db_session_with_containers, fake_app_id, str(TracingProviderEnum.ARIZE)) mock_ops_trace_manager.encrypt_tracing_config.return_value = {} - result = OpsService.update_tracing_app_config(fake_app_id, TracingProviderEnum.ARIZE, {}) + result = OpsService.update_tracing_app_config( + fake_app_id, TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result is None def test_update_tracing_app_config_invalid_credentials( @@ -323,7 +342,7 @@ class TestOpsService: self._insert_trace_config(db_session_with_containers, app.id, str(TracingProviderEnum.ARIZE)) with pytest.raises(ValueError, match="Invalid Credentials"): - OpsService.update_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {}) + OpsService.update_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {}, db_session_with_containers) def test_update_tracing_app_config_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -336,7 +355,9 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, str(TracingProviderEnum.ARIZE)) - result = OpsService.update_tracing_app_config(app.id, TracingProviderEnum.ARIZE, {}) + result = OpsService.update_tracing_app_config( + app.id, TracingProviderEnum.ARIZE, {}, db_session_with_containers + ) assert result is not None assert result["app_id"] == app.id @@ -344,7 +365,7 @@ class TestOpsService: # ── delete_tracing_app_config ────────────────────────────────────── def test_delete_tracing_app_config_no_config(self, db_session_with_containers: Session): - result = OpsService.delete_tracing_app_config(str(uuid.uuid4()), "arize") + result = OpsService.delete_tracing_app_config(str(uuid.uuid4()), "arize", db_session_with_containers) assert result is None def test_delete_tracing_app_config_success( @@ -353,7 +374,7 @@ class TestOpsService: app, _ = self._create_app(db_session_with_containers, mock_external_service_dependencies) self._insert_trace_config(db_session_with_containers, app.id, "arize") - result = OpsService.delete_tracing_app_config(app.id, "arize") + result = OpsService.delete_tracing_app_config(app.id, "arize", db_session_with_containers) assert result is True remaining = db_session_with_containers.scalar( diff --git a/api/tests/test_containers_integration_tests/services/test_recommended_app_service.py b/api/tests/test_containers_integration_tests/services/test_recommended_app_service.py index 9b8eec08ef4..f27132b0fe9 100644 --- a/api/tests/test_containers_integration_tests/services/test_recommended_app_service.py +++ b/api/tests/test_containers_integration_tests/services/test_recommended_app_service.py @@ -14,6 +14,8 @@ from models.model import AccountTrialAppRecord, TrialApp from services import recommended_app_service as service_module from services.recommended_app_service import RecommendedAppService +pytestmark = pytest.mark.usefixtures("db_session_with_containers") + class RecommendedAppPayload(TypedDict, total=False): id: str @@ -118,13 +120,13 @@ class TestRecommendedAppServiceGetApps: mock_factory = MagicMock(return_value=mock_instance) mock_factory_class.get_recommend_app_factory.return_value = mock_factory - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, "en-US") + result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=db.session()) assert result == expected assert len(result["recommended_apps"]) == 2 assert len(result["categories"]) == 3 mock_factory_class.get_recommend_app_factory.assert_called_once_with("remote") - mock_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US") + mock_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US", session=db.session()) @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) @patch("services.recommended_app_service.dify_config") @@ -143,7 +145,7 @@ class TestRecommendedAppServiceGetApps: mock_builtin_instance.fetch_recommended_apps_from_builtin.return_value = builtin_response mock_factory_class.get_buildin_recommend_app_retrieval.return_value = mock_builtin_instance - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, "zh-CN") + result = RecommendedAppService.get_recommended_apps_and_categories("zh-CN", session=db.session()) assert result == builtin_response assert result["recommended_apps"][0]["id"] == "builtin-1" @@ -164,7 +166,7 @@ class TestRecommendedAppServiceGetApps: mock_builtin_instance.fetch_recommended_apps_from_builtin.return_value = builtin_response mock_factory_class.get_buildin_recommend_app_retrieval.return_value = mock_builtin_instance - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, "en-US") + result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=db.session()) assert result == builtin_response mock_builtin_instance.fetch_recommended_apps_from_builtin.assert_called_once() @@ -182,10 +184,10 @@ class TestRecommendedAppServiceGetApps: mock_instance.get_recommended_apps_and_categories.return_value = lang_response mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, language) + result = RecommendedAppService.get_recommended_apps_and_categories(language, session=db.session()) assert result["recommended_apps"][0]["id"] == f"app-{language}" - mock_instance.get_recommended_apps_and_categories.assert_called_with(language) + mock_instance.get_recommended_apps_and_categories.assert_called_with(language, session=db.session()) @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) @patch("services.recommended_app_service.dify_config") @@ -197,7 +199,7 @@ class TestRecommendedAppServiceGetApps: mock_instance.get_recommended_apps_and_categories.return_value = response mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - RecommendedAppService.get_recommended_apps_and_categories(db.session, "en-US") + RecommendedAppService.get_recommended_apps_and_categories("en-US", session=db.session()) mock_factory_class.get_recommend_app_factory.assert_called_with(mode) @@ -237,10 +239,10 @@ class TestRecommendedAppServiceGetDetail: mock_instance.get_recommend_app_detail.return_value = expected mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - result = RecommendedAppService.get_recommend_app_detail(db.session, app_id) + result = RecommendedAppService.get_recommend_app_detail(app_id, session=db.session()) assert result == expected - mock_instance.get_recommend_app_detail.assert_called_once_with(app_id) + mock_instance.get_recommend_app_detail.assert_called_once_with(app_id, session=db.session()) @patch("services.recommended_app_service.FeatureService", autospec=True) @patch("services.recommended_app_service.RecommendAppRetrievalFactory", autospec=True) @@ -256,10 +258,10 @@ class TestRecommendedAppServiceGetDetail: mock_instance.get_recommend_app_detail.return_value = detail mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - result = RecommendedAppService.get_recommend_app_detail(db.session, "test-app") + result = RecommendedAppService.get_recommend_app_detail("test-app", session=db.session()) assert result is not None - mock_instance.get_recommend_app_detail.assert_called_with("test-app") + mock_instance.get_recommend_app_detail.assert_called_with("test-app", session=db.session()) mock_factory_class.get_recommend_app_factory.assert_called_with(mode) @@ -283,11 +285,11 @@ class TestRecommendedAppServiceGetLearnDifyApps: } mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - result = RecommendedAppService.get_learn_dify_apps(db.session, "en-US") + result = RecommendedAppService.get_learn_dify_apps("en-US", session=db.session()) assert result == {"recommended_apps": [expected_app]} mock_factory_class.get_recommend_app_factory.assert_called_once_with("remote") - mock_instance.get_learn_dify_apps.assert_called_once_with("en-US") + mock_instance.get_learn_dify_apps.assert_called_once_with("en-US", session=db.session()) @patch("services.recommended_app_service.dify_config") def test_sets_can_trial_when_trial_feature_enabled( @@ -314,10 +316,10 @@ class TestRecommendedAppServiceGetLearnDifyApps: can_trial_mock = MagicMock(return_value=True) monkeypatch.setattr(RecommendedAppService, "_can_trial_app", can_trial_mock) - result = RecommendedAppService.get_learn_dify_apps(db.session, "en-US") + result = RecommendedAppService.get_learn_dify_apps("en-US", session=db.session()) assert result["recommended_apps"][0]["can_trial"] is True - can_trial_mock.assert_called_once_with(db.session, "app-1") + can_trial_mock.assert_called_once_with(db.session(), "app-1") # ── Integration tests: trial app features (real DB) ──────────────────── @@ -333,10 +335,10 @@ class TestRecommendedAppServiceTrialFeatures: MagicMock(return_value=SimpleNamespace(enable_trial_app=False)), ) - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, "en-US") + result = RecommendedAppService.get_recommended_apps_and_categories("en-US", session=db.session()) assert result == expected - retrieval_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US") + retrieval_instance.get_recommended_apps_and_categories.assert_called_once_with("en-US", session=db.session()) builtin_instance.fetch_recommended_apps_from_builtin.assert_not_called() def test_get_apps_should_enrich_can_trial_when_enabled( @@ -364,7 +366,7 @@ class TestRecommendedAppServiceTrialFeatures: MagicMock(return_value=SimpleNamespace(enable_trial_app=True)), ) - result = RecommendedAppService.get_recommended_apps_and_categories(db.session, "ja-JP") + result = RecommendedAppService.get_recommended_apps_and_categories("ja-JP", session=db.session()) builtin_instance.fetch_recommended_apps_from_builtin.assert_called_once_with("en-US") assert result["recommended_apps"][0]["can_trial"] is True @@ -400,7 +402,7 @@ class TestRecommendedAppServiceTrialFeatures: MagicMock(return_value=SimpleNamespace(enable_trial_app=True)), ) - result = RecommendedAppService.get_recommend_app_detail(db.session, app_id) + result = RecommendedAppService.get_recommend_app_detail(app_id, session=db.session()) assert result is not None detail_result = cast(RecommendedAppPayload, result) @@ -421,10 +423,10 @@ class TestRecommendedAppServiceTrialFeatures: mock_instance.get_recommend_app_detail.return_value = None mock_factory_class.get_recommend_app_factory.return_value = MagicMock(return_value=mock_instance) - result = RecommendedAppService.get_recommend_app_detail(db.session, "nonexistent") + result = RecommendedAppService.get_recommend_app_detail("nonexistent", session=db.session()) assert result is None - mock_instance.get_recommend_app_detail.assert_called_once_with("nonexistent") + mock_instance.get_recommend_app_detail.assert_called_once_with("nonexistent", session=db.session()) mock_feature_service.get_system_features.assert_not_called() def test_add_trial_app_record_increments_count_for_existing(self, db_session_with_containers: Session) -> None: @@ -434,7 +436,7 @@ class TestRecommendedAppServiceTrialFeatures: db_session_with_containers.add(AccountTrialAppRecord(app_id=app_id, account_id=account_id, count=3)) db_session_with_containers.commit() - RecommendedAppService.add_trial_app_record(db.session, app_id, account_id) + RecommendedAppService.add_trial_app_record(app_id, account_id, session=db.session()) db_session_with_containers.expire_all() record = db_session_with_containers.scalar( @@ -449,7 +451,7 @@ class TestRecommendedAppServiceTrialFeatures: app_id = str(uuid.uuid4()) account_id = str(uuid.uuid4()) - RecommendedAppService.add_trial_app_record(db.session, app_id, account_id) + RecommendedAppService.add_trial_app_record(app_id, account_id, session=db.session()) db_session_with_containers.expire_all() record = db_session_with_containers.scalar( diff --git a/api/tests/test_containers_integration_tests/services/test_saved_message_service.py b/api/tests/test_containers_integration_tests/services/test_saved_message_service.py index cfd1d4e86b4..92741ac56cb 100644 --- a/api/tests/test_containers_integration_tests/services/test_saved_message_service.py +++ b/api/tests/test_containers_integration_tests/services/test_saved_message_service.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest from faker import Faker @@ -86,7 +86,7 @@ class TestSavedMessageService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account @@ -222,7 +222,7 @@ class TestSavedMessageService: # Act: Execute the method under test result = SavedMessageService.pagination_by_last_id( - db_session_with_containers, app_model=app, user=account, last_id=None, limit=10 + app_model=app, user=account, last_id=None, limit=10, session=db_session_with_containers ) # Assert: Verify the expected outcomes @@ -297,7 +297,7 @@ class TestSavedMessageService: # Act: Execute the method under test result = SavedMessageService.pagination_by_last_id( - db_session_with_containers, app_model=app, user=end_user, last_id="test_last_id", limit=5 + app_model=app, user=end_user, last_id="test_last_id", limit=5, session=db_session_with_containers ) # Assert: Verify the expected outcomes @@ -347,7 +347,7 @@ class TestSavedMessageService: mock_external_service_dependencies["message_service"].get_message.return_value = message # Act: Execute the method under test - SavedMessageService.save(db_session_with_containers, app_model=app, user=account, message_id=message.id) + SavedMessageService.save(app_model=app, user=account, message_id=message.id, session=db_session_with_containers) # Assert: Verify the expected outcomes # Check if saved message was created in database @@ -372,7 +372,7 @@ class TestSavedMessageService: # Verify MessageService.get_message was called mock_external_service_dependencies["message_service"].get_message.assert_called_once_with( - app_model=app, user=account, message_id=message.id + app_model=app, user=account, message_id=message.id, session=ANY ) # Verify database state @@ -397,7 +397,7 @@ class TestSavedMessageService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError) as exc_info: SavedMessageService.pagination_by_last_id( - db_session_with_containers, app_model=app, user=None, last_id=None, limit=10 + app_model=app, user=None, last_id=None, limit=10, session=db_session_with_containers ) assert "User is required" in str(exc_info.value) @@ -417,7 +417,9 @@ class TestSavedMessageService: message = self._create_test_message(db_session_with_containers, app, account) # Act: Execute the method under test with None user - result = SavedMessageService.save(db_session_with_containers, app_model=app, user=None, message_id=message.id) + result = SavedMessageService.save( + app_model=app, user=None, message_id=message.id, session=db_session_with_containers + ) # Assert: Verify the expected outcomes assert result is None @@ -476,7 +478,9 @@ class TestSavedMessageService: ) # Act: Execute the method under test - SavedMessageService.delete(db_session_with_containers, app_model=app, user=account, message_id=message.id) + SavedMessageService.delete( + app_model=app, user=account, message_id=message.id, session=db_session_with_containers + ) # Assert: Verify the expected outcomes # Check if saved message was deleted from database @@ -506,7 +510,9 @@ class TestSavedMessageService: mock_external_service_dependencies["message_service"].get_message.return_value = message - SavedMessageService.save(db_session_with_containers, app_model=app, user=end_user, message_id=message.id) + SavedMessageService.save( + app_model=app, user=end_user, message_id=message.id, session=db_session_with_containers + ) saved = ( db_session_with_containers.query(SavedMessage) @@ -527,9 +533,9 @@ class TestSavedMessageService: mock_external_service_dependencies["message_service"].get_message.return_value = message # Save once - SavedMessageService.save(db_session_with_containers, app_model=app, user=account, message_id=message.id) + SavedMessageService.save(app_model=app, user=account, message_id=message.id, session=db_session_with_containers) # Save again - SavedMessageService.save(db_session_with_containers, app_model=app, user=account, message_id=message.id) + SavedMessageService.save(app_model=app, user=account, message_id=message.id, session=db_session_with_containers) count = ( db_session_with_containers.query(SavedMessage) @@ -552,7 +558,7 @@ class TestSavedMessageService: db_session_with_containers.add(saved) db_session_with_containers.commit() - SavedMessageService.delete(db_session_with_containers, app_model=app, user=None, message_id=message.id) + SavedMessageService.delete(app_model=app, user=None, message_id=message.id, session=db_session_with_containers) # Should still exist assert ( @@ -571,7 +577,9 @@ class TestSavedMessageService: # Should not raise — use a valid UUID that doesn't exist in DB from uuid import uuid4 - SavedMessageService.delete(db_session_with_containers, app_model=app, user=account, message_id=str(uuid4())) + SavedMessageService.delete( + app_model=app, user=account, message_id=str(uuid4()), session=db_session_with_containers + ) def test_delete_for_end_user(self, db_session_with_containers: Session, mock_external_service_dependencies): """Test deleting a saved message for an EndUser.""" @@ -585,7 +593,9 @@ class TestSavedMessageService: db_session_with_containers.add(saved) db_session_with_containers.commit() - SavedMessageService.delete(db_session_with_containers, app_model=app, user=end_user, message_id=message.id) + SavedMessageService.delete( + app_model=app, user=end_user, message_id=message.id, session=db_session_with_containers + ) assert ( db_session_with_containers.query(SavedMessage) @@ -615,7 +625,9 @@ class TestSavedMessageService: db_session_with_containers.commit() # Delete only account1's saved message - SavedMessageService.delete(db_session_with_containers, app_model=app, user=account1, message_id=message.id) + SavedMessageService.delete( + app_model=app, user=account1, message_id=message.id, session=db_session_with_containers + ) # Account's saved message should be gone assert ( diff --git a/api/tests/test_containers_integration_tests/services/test_tag_service.py b/api/tests/test_containers_integration_tests/services/test_tag_service.py index 748cca6c845..86b635ac23d 100644 --- a/api/tests/test_containers_integration_tests/services/test_tag_service.py +++ b/api/tests/test_containers_integration_tests/services/test_tag_service.py @@ -205,7 +205,7 @@ def test_get_tags_success(db_session_with_containers: Session, current_user_stub db_session_with_containers, tags=tags[:2], target_id=dataset.id, tenant_id=tenant.id, user_id=account.id ) - result = TagService.get_tags(db_session_with_containers, TagType.KNOWLEDGE, tenant.id) + result = TagService.get_tags(TagType.KNOWLEDGE, tenant.id, session=db_session_with_containers) assert result is not None assert len(result) == 3 @@ -235,7 +235,7 @@ def test_get_tags_with_keyword_filter(db_session_with_containers: Session, curre tags[2].name = "web_development" db_session_with_containers.flush() - result = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="development") + result = TagService.get_tags(TagType.APP, tenant.id, keyword="development", session=db_session_with_containers) assert result is not None assert len(result) == 2 @@ -243,7 +243,9 @@ def test_get_tags_with_keyword_filter(db_session_with_containers: Session, curre for tag_result in result: assert "development" in tag_result.name.lower() - result_no_match = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="nonexistent") + result_no_match = TagService.get_tags( + TagType.APP, tenant.id, keyword="nonexistent", session=db_session_with_containers + ) assert result_no_match == [] @@ -291,19 +293,19 @@ def test_get_tags_with_special_characters_in_keyword( db_session_with_containers.flush() - result = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="50%") + result = TagService.get_tags(TagType.APP, tenant.id, keyword="50%", session=db_session_with_containers) assert len(result) == 1 assert result[0].name == "50% discount" - result = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="test_data") + result = TagService.get_tags(TagType.APP, tenant.id, keyword="test_data", session=db_session_with_containers) assert len(result) == 1 assert result[0].name == "test_data_tag" - result = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="path\\to\\tag") + result = TagService.get_tags(TagType.APP, tenant.id, keyword="path\\to\\tag", session=db_session_with_containers) assert len(result) == 1 assert result[0].name == "path\\to\\tag" - result = TagService.get_tags(db_session_with_containers, TagType.APP, tenant.id, keyword="50%") + result = TagService.get_tags(TagType.APP, tenant.id, keyword="50%", session=db_session_with_containers) assert len(result) == 1 assert all("50%" in item.name for item in result) @@ -312,7 +314,7 @@ def test_get_tags_empty_result(db_session_with_containers: Session, current_user account, tenant = _create_account_with_tenant(db_session_with_containers) _set_current_user(current_user_stub, account, tenant) - result = TagService.get_tags(db_session_with_containers, TagType.KNOWLEDGE, tenant.id) + result = TagService.get_tags(TagType.KNOWLEDGE, tenant.id, session=db_session_with_containers) assert result == [] diff --git a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py b/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py index 664c1167994..ed063ceaccc 100644 --- a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py @@ -90,7 +90,7 @@ class TestWebConversationService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account @@ -312,7 +312,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Pin the conversation - WebConversationService.pin(app, conversation.id, account) + WebConversationService.pin(app, conversation.id, account, db_session_with_containers) # Verify the conversation was pinned @@ -346,10 +346,10 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Pin the conversation first time - WebConversationService.pin(app, conversation.id, account) + WebConversationService.pin(app, conversation.id, account, db_session_with_containers) # Pin the conversation again - WebConversationService.pin(app, conversation.id, account) + WebConversationService.pin(app, conversation.id, account, db_session_with_containers) # Verify only one pinned conversation record exists @@ -380,7 +380,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, end_user, fake) # Pin the conversation - WebConversationService.pin(app, conversation.id, end_user) + WebConversationService.pin(app, conversation.id, end_user, db_session_with_containers) # Verify the conversation was pinned @@ -412,7 +412,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Pin the conversation first - WebConversationService.pin(app, conversation.id, account) + WebConversationService.pin(app, conversation.id, account, db_session_with_containers) # Verify it was pinned @@ -430,7 +430,7 @@ class TestWebConversationService: assert pinned_conversation is not None # Unpin the conversation - WebConversationService.unpin(app, conversation.id, account) + WebConversationService.unpin(app, conversation.id, account, db_session_with_containers) # Verify it was unpinned pinned_conversation = ( @@ -459,7 +459,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Try to unpin a conversation that was never pinned - WebConversationService.unpin(app, conversation.id, account) + WebConversationService.unpin(app, conversation.id, account, db_session_with_containers) # Verify no pinned conversation record exists @@ -509,7 +509,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Try to pin with None user - WebConversationService.pin(app, conversation.id, None) + WebConversationService.pin(app, conversation.id, None, db_session_with_containers) # Verify no pinned conversation was created @@ -537,7 +537,7 @@ class TestWebConversationService: conversation = self._create_test_conversation(db_session_with_containers, app, account, fake) # Pin the conversation first - WebConversationService.pin(app, conversation.id, account) + WebConversationService.pin(app, conversation.id, account, db_session_with_containers) # Verify it was pinned @@ -555,7 +555,7 @@ class TestWebConversationService: assert pinned_conversation is not None # Try to unpin with None user - WebConversationService.unpin(app, conversation.id, None) + WebConversationService.unpin(app, conversation.id, None, db_session_with_containers) # Verify the conversation is still pinned pinned_conversation = ( diff --git a/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py b/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py index 7825f502f77..52d1fde7927 100644 --- a/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py +++ b/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py @@ -1,6 +1,6 @@ import time import uuid -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest from faker import Faker @@ -223,7 +223,7 @@ class TestWebAppAuthService: ) # Act: Execute authentication - result = WebAppAuthService.authenticate(account.email, password) + result = WebAppAuthService.authenticate(account.email, password, db_session_with_containers) # Assert: Verify successful authentication assert result is not None @@ -260,7 +260,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(AccountNotFoundError): - WebAppAuthService.authenticate(non_existent_email, "any_password") + WebAppAuthService.authenticate(non_existent_email, "any_password", db_session_with_containers) def test_authenticate_account_banned(self, db_session_with_containers: Session, mock_external_service_dependencies): """ @@ -297,7 +297,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(AccountLoginError) as exc_info: - WebAppAuthService.authenticate(account.email, password) + WebAppAuthService.authenticate(account.email, password, db_session_with_containers) assert "Account is banned." in str(exc_info.value) @@ -318,7 +318,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with wrong password with pytest.raises(AccountPasswordError) as exc_info: - WebAppAuthService.authenticate(account.email, "wrong_password") + WebAppAuthService.authenticate(account.email, "wrong_password", db_session_with_containers) assert "Invalid email or password." in str(exc_info.value) @@ -350,7 +350,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(AccountPasswordError) as exc_info: - WebAppAuthService.authenticate(account.email, "any_password") + WebAppAuthService.authenticate(account.email, "any_password", db_session_with_containers) assert "Invalid email or password." in str(exc_info.value) @@ -403,7 +403,7 @@ class TestWebAppAuthService: ) # Act: Execute user retrieval - result = WebAppAuthService.get_user_through_email(account.email) + result = WebAppAuthService.get_user_through_email(account.email, db_session_with_containers) # Assert: Verify successful retrieval assert result is not None @@ -430,7 +430,7 @@ class TestWebAppAuthService: non_existent_email = f"nonexistent_{uuid.uuid4().hex}@example.com" # Act: Execute user retrieval - result = WebAppAuthService.get_user_through_email(non_existent_email) + result = WebAppAuthService.get_user_through_email(non_existent_email, db_session_with_containers) # Assert: Verify proper handling assert result is None @@ -463,7 +463,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(Unauthorized) as exc_info: - WebAppAuthService.get_user_through_email(account.email) + WebAppAuthService.get_user_through_email(account.email, db_session_with_containers) assert "Account is banned." in str(exc_info.value) @@ -659,7 +659,7 @@ class TestWebAppAuthService: ) # Act: Execute end user creation - result = WebAppAuthService.create_end_user(site.code, "test@example.com") + result = WebAppAuthService.create_end_user(site.code, "test@example.com", db_session_with_containers) # Assert: Verify successful creation assert result is not None @@ -694,7 +694,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(NotFound) as exc_info: - WebAppAuthService.create_end_user(non_existent_code, "test@example.com") + WebAppAuthService.create_end_user(non_existent_code, "test@example.com", db_session_with_containers) assert "Site not found." in str(exc_info.value) @@ -732,7 +732,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(NotFound) as exc_info: - WebAppAuthService.create_end_user(site.code, "test@example.com") + WebAppAuthService.create_end_user(site.code, "test@example.com", db_session_with_containers) assert "App not found." in str(exc_info.value) @@ -750,7 +750,9 @@ class TestWebAppAuthService: # Arrange: Setup test with private access mode # Act: Execute permission check requirement test - result = WebAppAuthService.is_app_require_permission_check(access_mode="private") + result = WebAppAuthService.is_app_require_permission_check( + access_mode="private", session=db_session_with_containers + ) # Assert: Verify correct result assert result is True @@ -769,7 +771,9 @@ class TestWebAppAuthService: # Arrange: Setup test with public access mode # Act: Execute permission check requirement test - result = WebAppAuthService.is_app_require_permission_check(access_mode="public") + result = WebAppAuthService.is_app_require_permission_check( + access_mode="public", session=db_session_with_containers + ) # Assert: Verify correct result assert result is False @@ -789,13 +793,17 @@ class TestWebAppAuthService: mock_external_service_dependencies["app_service"].get_app_id_by_code.return_value = "mock_app_id" # Act: Execute permission check requirement test - result = WebAppAuthService.is_app_require_permission_check(app_code="mock_app_code") + result = WebAppAuthService.is_app_require_permission_check( + app_code="mock_app_code", session=db_session_with_containers + ) # Assert: Verify correct result assert result is True # Verify mock service was called correctly - mock_external_service_dependencies["app_service"].get_app_id_by_code.assert_called_once_with("mock_app_code") + mock_external_service_dependencies["app_service"].get_app_id_by_code.assert_called_once_with( + "mock_app_code", session=ANY + ) mock_external_service_dependencies[ "enterprise_service" ].WebAppAuth.get_app_access_mode_by_id.assert_called_once_with("mock_app_id") @@ -814,7 +822,7 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError) as exc_info: - WebAppAuthService.is_app_require_permission_check() + WebAppAuthService.is_app_require_permission_check(session=db_session_with_containers) assert "Either app_code or app_id must be provided." in str(exc_info.value) @@ -832,7 +840,7 @@ class TestWebAppAuthService: # Arrange: Setup test with public access mode # Act: Execute authentication type determination - result = WebAppAuthService.get_app_auth_type(access_mode="public") + result = WebAppAuthService.get_app_auth_type(access_mode="public", session=db_session_with_containers) # Assert: Verify correct result assert result == WebAppAuthType.PUBLIC @@ -851,7 +859,7 @@ class TestWebAppAuthService: # Arrange: Setup test with private access mode # Act: Execute authentication type determination - result = WebAppAuthService.get_app_auth_type(access_mode="private") + result = WebAppAuthService.get_app_auth_type(access_mode="private", session=db_session_with_containers) # Assert: Verify correct result assert result == WebAppAuthType.INTERNAL @@ -875,7 +883,9 @@ class TestWebAppAuthService: ].WebAppAuth.get_app_access_mode_by_id.return_value = setting # Act: Execute authentication type determination - result: WebAppAuthType = WebAppAuthService.get_app_auth_type(app_code="mock_app_code") + result: WebAppAuthType = WebAppAuthService.get_app_auth_type( + app_code="mock_app_code", session=db_session_with_containers + ) # Assert: Verify correct result assert result == WebAppAuthType.EXTERNAL @@ -899,6 +909,6 @@ class TestWebAppAuthService: # Act & Assert: Verify proper error handling with pytest.raises(ValueError) as exc_info: - WebAppAuthService.get_app_auth_type() + WebAppAuthService.get_app_auth_type(session=db_session_with_containers) assert "Either app_code or access_mode must be provided." in str(exc_info.value) diff --git a/api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py b/api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py index c699d39dde1..902134e053d 100644 --- a/api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py +++ b/api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py @@ -350,9 +350,10 @@ class TestWebhookServiceTriggerExecutionWithContainers: quota_charge.commit.assert_called_once() mock_trigger.assert_called_once() trigger_args = mock_trigger.call_args.args - assert trigger_args[1] is end_user - assert trigger_args[2].workflow_id == workflow.id - assert trigger_args[2].root_node_id == webhook_trigger.node_id + assert trigger_args[0] is end_user + assert trigger_args[1].workflow_id == workflow.id + assert trigger_args[1].root_node_id == webhook_trigger.node_id + assert mock_trigger.call_args.kwargs["session"] is not None def test_trigger_workflow_execution_marks_tenant_rate_limited_when_quota_exceeded( self, db_session_with_containers: Session, flask_app_with_containers: Flask diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py index cf76afb303c..f553b0f72a0 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py @@ -99,7 +99,7 @@ class TestWorkflowAppService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account @@ -164,7 +164,7 @@ class TestWorkflowAppService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py index 726c360d77e..7c528f06b10 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py @@ -92,7 +92,7 @@ class TestWorkflowRunService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) return app, account @@ -544,7 +544,7 @@ class TestWorkflowRunService: icon="🚀", icon_background="#4ECDC4", ) - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Create workflow run without node executions workflow_run = self._create_test_workflow_run(db_session_with_containers, app, account, "debugging") @@ -596,7 +596,7 @@ class TestWorkflowRunService: icon="🚀", icon_background="#4ECDC4", ) - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Use invalid workflow run ID invalid_workflow_run_id = str(uuid.uuid4()) @@ -648,7 +648,7 @@ class TestWorkflowRunService: icon="🚀", icon_background="#4ECDC4", ) - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Create workflow run workflow_run = self._create_test_workflow_run(db_session_with_containers, app, account, "debugging") diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_service.py index 349aac1be36..6531ed4fbb0 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_service.py @@ -227,7 +227,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.is_workflow_exist(app) + result = workflow_service.is_workflow_exist(app, session=db_session_with_containers) # Assert assert result is True @@ -247,7 +247,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.is_workflow_exist(app) + result = workflow_service.is_workflow_exist(app, session=db_session_with_containers) # Assert assert result is False @@ -269,7 +269,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_draft_workflow(app) + result = workflow_service.get_draft_workflow(app, session=db_session_with_containers) # Assert assert result is not None @@ -293,7 +293,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_draft_workflow(app) + result = workflow_service.get_draft_workflow(app, session=db_session_with_containers) # Assert assert result is None @@ -320,7 +320,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_published_workflow_by_id(app, workflow.id) + result = workflow_service.get_published_workflow_by_id(app, workflow.id, session=db_session_with_containers) # Assert assert result is not None @@ -349,7 +349,7 @@ class TestWorkflowService: from services.errors.app import IsDraftWorkflowError with pytest.raises(IsDraftWorkflowError): - workflow_service.get_published_workflow_by_id(app, workflow.id) + workflow_service.get_published_workflow_by_id(app, workflow.id, session=db_session_with_containers) def test_get_published_workflow_by_id_not_found(self, db_session_with_containers: Session): """ @@ -366,7 +366,9 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_published_workflow_by_id(app, non_existent_workflow_id) + result = workflow_service.get_published_workflow_by_id( + app, non_existent_workflow_id, session=db_session_with_containers + ) # Assert assert result is None @@ -393,7 +395,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_published_workflow(app) + result = workflow_service.get_published_workflow(app, session=db_session_with_containers) # Assert assert result is not None @@ -416,7 +418,7 @@ class TestWorkflowService: workflow_service = WorkflowService() # Act - result = workflow_service.get_published_workflow(app) + result = workflow_service.get_published_workflow(app, session=db_session_with_containers) # Assert assert result is None @@ -714,6 +716,7 @@ class TestWorkflowService: account=account, environment_variables=environment_variables, conversation_variables=conversation_variables, + session=db_session_with_containers, ) # Assert @@ -778,6 +781,7 @@ class TestWorkflowService: account=account, environment_variables=environment_variables, conversation_variables=conversation_variables, + session=db_session_with_containers, ) # Assert @@ -838,6 +842,7 @@ class TestWorkflowService: account=account, environment_variables=environment_variables, conversation_variables=conversation_variables, + session=db_session_with_containers, ) def test_publish_workflow_success(self, db_session_with_containers: Session): @@ -979,9 +984,7 @@ class TestWorkflowService: workflow_service = WorkflowService() restored_workflow = workflow_service.restore_published_workflow_to_draft( - app_model=app, - workflow_id=published_workflow.id, - account=account, + app_model=app, workflow_id=published_workflow.id, account=account, session=db_session_with_containers ) db_session_with_containers.expire_all() @@ -1130,7 +1133,9 @@ class TestWorkflowService: } # Act - result = workflow_service.convert_to_workflow(app_model=app, account=account, args=conversion_args) + result = workflow_service.convert_to_workflow( + app_model=app, account=account, args=conversion_args, session=db_session_with_containers + ) # Assert assert result is not None @@ -1190,7 +1195,9 @@ class TestWorkflowService: } # Act - result = workflow_service.convert_to_workflow(app_model=app, account=account, args=conversion_args) + result = workflow_service.convert_to_workflow( + app_model=app, account=account, args=conversion_args, session=db_session_with_containers + ) # Assert assert result is not None @@ -1222,7 +1229,9 @@ class TestWorkflowService: # Act & Assert with pytest.raises(ValueError, match="Current App mode: workflow is not supported convert to workflow"): - workflow_service.convert_to_workflow(app_model=app, account=account, args=conversion_args) + workflow_service.convert_to_workflow( + app_model=app, account=account, args=conversion_args, session=db_session_with_containers + ) def test_validate_features_structure_advanced_chat(self, db_session_with_containers: Session): """ diff --git a/api/tests/test_containers_integration_tests/services/test_workspace_service.py b/api/tests/test_containers_integration_tests/services/test_workspace_service.py index 4e89d906f16..d7cbfa91ab7 100644 --- a/api/tests/test_containers_integration_tests/services/test_workspace_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workspace_service.py @@ -104,7 +104,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -151,7 +151,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -206,7 +206,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -261,7 +261,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -291,7 +291,7 @@ class TestWorkspaceService: # Arrange: No test data needed for this test # Act: Execute the method under test with None tenant - result = WorkspaceService.get_tenant_info(None) + result = WorkspaceService.get_tenant_info(None, db_session_with_containers) # Assert: Verify the expected outcomes assert result is None @@ -341,7 +341,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -398,7 +398,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -448,7 +448,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -513,7 +513,7 @@ class TestWorkspaceService: # Mock current_user for flask_login with patch("services.workspace_service.current_user", account): # Act: Execute the method under test - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) # Assert: Verify the expected outcomes assert result is not None @@ -553,7 +553,7 @@ class TestWorkspaceService: # No TenantAccountJoin created with patch("services.workspace_service.current_user", account): with pytest.raises(AssertionError, match="TenantAccountJoin not found"): - WorkspaceService.get_tenant_info(tenant) + WorkspaceService.get_tenant_info(tenant, db_session_with_containers) def test_get_tenant_info_should_set_replace_webapp_logo_to_none_when_flag_absent( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -572,7 +572,7 @@ class TestWorkspaceService: mock_external_service_dependencies["tenant_service"].has_roles.return_value = True with patch("services.workspace_service.current_user", account): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["custom_config"]["replace_webapp_logo"] is None @@ -596,7 +596,7 @@ class TestWorkspaceService: mock_external_service_dependencies["tenant_service"].has_roles.return_value = True with patch("services.workspace_service.current_user", account): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["custom_config"]["replace_webapp_logo"].startswith(custom_base) @@ -615,7 +615,7 @@ class TestWorkspaceService: mock_external_service_dependencies["tenant_service"].has_roles.return_value = False with patch("services.workspace_service.current_user", account): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert "next_credit_reset_date" not in result @@ -642,7 +642,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", return_value=None), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["next_credit_reset_date"] == "2025-02-01" @@ -669,7 +669,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", return_value=paid_pool), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["trial_credits"] == 1000 @@ -697,7 +697,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", side_effect=[paid_pool, None]), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["trial_credits"] == -1 @@ -726,7 +726,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", side_effect=[paid_pool, trial_pool]), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["trial_credits"] == 100 @@ -754,7 +754,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", side_effect=[None, trial_pool]), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["trial_credits"] == 50 @@ -785,7 +785,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", side_effect=[paid_pool, trial_pool]), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert result["trial_credits"] == 200 @@ -811,7 +811,7 @@ class TestWorkspaceService: patch("services.workspace_service.current_user", account), patch("services.credit_pool_service.CreditPoolService.get_pool", side_effect=[None, None]), ): - result = WorkspaceService.get_tenant_info(tenant) + result = WorkspaceService.get_tenant_info(tenant, db_session_with_containers) assert result is not None assert "trial_credits" not in result diff --git a/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py b/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py index af83adaae01..28ae8e5935e 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py +++ b/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py @@ -1,6 +1,7 @@ import inspect import json -from unittest.mock import patch +from collections.abc import Iterator +from unittest.mock import MagicMock, patch import pytest from faker import Faker @@ -10,16 +11,18 @@ from sqlalchemy.orm import Session from core.tools.entities.tool_entities import ApiProviderSchemaType from core.tools.errors import ApiToolProviderNotFoundError from core.tools.tool_label_manager import ToolLabelManager -from models import Account, Tenant +from models import Account, AccountStatus, Tenant, TenantStatus from models.tools import ApiToolProvider from services.tools.api_tools_manage_service import ApiToolManageService +MockDependencies = dict[str, MagicMock] + class TestApiToolManageService: """Integration tests for ApiToolManageService using testcontainers.""" @pytest.fixture - def mock_external_service_dependencies(self): + def mock_external_service_dependencies(self) -> Iterator[MockDependencies]: """Mock setup for external service dependencies.""" with ( patch("services.tools.api_tools_manage_service.ToolLabelManager") as mock_tool_label_manager, @@ -39,7 +42,9 @@ class TestApiToolManageService: "provider_controller": mock_provider_controller, } - def _create_test_account_and_tenant(self, db_session_with_containers: Session, mock_external_service_dependencies): + def _create_test_account_and_tenant( + self, db_session_with_containers: Session, mock_external_service_dependencies: MockDependencies + ) -> tuple[Account, Tenant]: """ Helper method to create a test account and tenant for testing. @@ -57,7 +62,7 @@ class TestApiToolManageService: email=fake.email(), name=fake.name(), interface_language="en-US", - status="active", + status=AccountStatus.ACTIVE, ) db_session_with_containers.add(account) @@ -66,7 +71,7 @@ class TestApiToolManageService: # Create tenant for the account tenant = Tenant( name=fake.company(), - status="normal", + status=TenantStatus.NORMAL, ) db_session_with_containers.add(tenant) db_session_with_containers.commit() @@ -88,7 +93,7 @@ class TestApiToolManageService: return account, tenant - def _create_test_openapi_schema(self): + def _create_test_openapi_schema(self) -> str: """Helper method to create a test OpenAPI schema.""" return """ { @@ -121,8 +126,11 @@ class TestApiToolManageService: """ def test_parser_api_schema_success( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test successful parsing of API schema. @@ -148,6 +156,8 @@ class TestApiToolManageService: # Verify credentials schema structure credentials_schema = result["credentials_schema"] assert len(credentials_schema) == 3 + assert all(isinstance(field, dict) for field in credentials_schema) + assert all(isinstance(tool, dict) for tool in result["parameters_schema"]) # Check auth_type field auth_type_field = next(field for field in credentials_schema if field["name"] == "auth_type") @@ -166,8 +176,11 @@ class TestApiToolManageService: assert api_key_value_field["default"] == "" def test_parser_api_schema_invalid_schema( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test parsing of invalid API schema. @@ -186,8 +199,11 @@ class TestApiToolManageService: assert "invalid schema" in str(exc_info.value) def test_parser_api_schema_malformed_json( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test parsing of malformed JSON schema. @@ -206,8 +222,11 @@ class TestApiToolManageService: assert "invalid schema" in str(exc_info.value) def test_convert_schema_to_tool_bundles_success( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test successful conversion of schema to tool bundles. @@ -236,8 +255,11 @@ class TestApiToolManageService: assert tool_bundle.operation_id == "testOperation" def test_convert_schema_to_tool_bundles_with_extra_info( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test successful conversion of schema to tool bundles with extra info. @@ -262,8 +284,11 @@ class TestApiToolManageService: assert isinstance(schema_type, str) def test_convert_schema_to_tool_bundles_invalid_schema( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test conversion of invalid schema to tool bundles. @@ -282,8 +307,11 @@ class TestApiToolManageService: assert "invalid schema" in str(exc_info.value) def test_create_api_tool_provider_success( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test successful creation of API tool provider. @@ -301,7 +329,7 @@ class TestApiToolManageService: ) provider_name = fake.company() - icon = {"type": "emoji", "value": "🔧"} + icon = {"content": "🔧", "background": "#FFF"} credentials = {"auth_type": "none", "api_key_header": "X-API-Key", "api_key_value": ""} schema_type = ApiProviderSchemaType.OPENAPI schema = self._create_test_openapi_schema() @@ -341,6 +369,7 @@ class TestApiToolManageService: assert provider.schema_type_str == schema_type assert provider.privacy_policy == privacy_policy assert provider.custom_disclaimer == custom_disclaimer + assert json.loads(provider.icon) == icon # Verify mock interactions mock_external_service_dependencies["tool_label_manager"].update_tool_labels.assert_called_once() @@ -349,8 +378,11 @@ class TestApiToolManageService: mock_external_service_dependencies["provider_controller"].load_bundled_tools.assert_called_once() def test_create_api_tool_provider_duplicate_name( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test creation of API tool provider with duplicate name. @@ -366,7 +398,7 @@ class TestApiToolManageService: ) provider_name = fake.company() - icon = {"type": "emoji", "value": "🔧"} + icon = {"content": "🔧", "background": "#FFF"} credentials = {"auth_type": "none"} schema_type = ApiProviderSchemaType.OPENAPI schema = self._create_test_openapi_schema() @@ -406,8 +438,11 @@ class TestApiToolManageService: assert f"provider {provider_name} already exists" in str(exc_info.value) def test_create_api_tool_provider_invalid_schema_type( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test creation of API tool provider with invalid schema type. @@ -423,7 +458,7 @@ class TestApiToolManageService: ) provider_name = fake.company() - icon = {"type": "emoji", "value": "🔧"} + icon = {"content": "🔧", "background": "#FFF"} credentials = {"auth_type": "none"} schema_type = "invalid_type" schema = self._create_test_openapi_schema() @@ -438,8 +473,11 @@ class TestApiToolManageService: assert "validation error" in str(exc_info.value) def test_create_api_tool_provider_missing_auth_type( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test creation of API tool provider with missing auth type. @@ -455,7 +493,7 @@ class TestApiToolManageService: ) provider_name = fake.company() - icon = {"type": "emoji", "value": "🔧"} + icon = {"content": "🔧", "background": "#FFF"} credentials = {} # Missing auth_type schema_type = ApiProviderSchemaType.OPENAPI schema = self._create_test_openapi_schema() @@ -481,8 +519,11 @@ class TestApiToolManageService: assert "auth_type is required" in str(exc_info.value) def test_create_api_tool_provider_with_api_key_auth( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test successful creation of API tool provider with API key authentication. @@ -498,7 +539,7 @@ class TestApiToolManageService: ) provider_name = fake.company() - icon = {"type": "emoji", "value": "🔑"} + icon = {"content": "🔑", "background": "#FFF"} credentials = {"auth_type": "api_key", "api_key_header": "X-API-Key", "api_key_value": fake.uuid4()} schema_type = ApiProviderSchemaType.OPENAPI schema = self._create_test_openapi_schema() @@ -542,8 +583,11 @@ class TestApiToolManageService: mock_external_service_dependencies["provider_controller"].from_db.assert_called_once() def test_delete_api_tool_provider_success( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """Test successful deletion of an API tool provider.""" fake = Faker() account, tenant = self._create_test_account_and_tenant( @@ -583,8 +627,8 @@ class TestApiToolManageService: assert deleted is None def test_delete_api_tool_provider_not_found( - self, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, db_session_with_containers: Session, mock_external_service_dependencies: MockDependencies + ) -> None: """Test deletion raises ValueError when provider not found.""" fake = Faker() account, tenant = self._create_test_account_and_tenant( @@ -595,14 +639,15 @@ class TestApiToolManageService: ApiToolManageService.delete_api_tool_provider(account.id, tenant.id, "nonexistent") def test_update_api_tool_provider_success( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: fake = Faker() # Firmware fix for cache.delete() in update flow mock_encrypter = mock_external_service_dependencies["encrypter"] - from unittest.mock import MagicMock - mock_cache = MagicMock() mock_cache.delete.return_value = None mock_encrypter.return_value = (mock_encrypter, mock_cache) @@ -620,7 +665,7 @@ class TestApiToolManageService: user_id=account.id, tenant_id=tenant.id, provider_name=original_name, - icon={"type": "emoji", "value": "🔧"}, + icon={"content": "🔧", "background": "#FFF"}, credentials={"auth_type": "none"}, schema_type=ApiProviderSchemaType.OPENAPI, schema=self._create_test_openapi_schema(), @@ -646,7 +691,7 @@ class TestApiToolManageService: provider_name=new_name, original_provider=original_name, # new icon - changed 2 - icon={"type": "emoji", "value": "🚀"}, + icon={"content": "🚀", "background": "#FFF"}, credentials={"auth_type": "none"}, _schema_type=ApiProviderSchemaType.OPENAPI, schema=self._create_test_openapi_schema(), @@ -677,9 +722,7 @@ class TestApiToolManageService: # - changed 1 assert updated_provider.name == new_name # - changed 2 - icon_data = json.loads(updated_provider.icon) - assert icon_data["type"] == "emoji" - assert icon_data["value"] == "🚀" + assert json.loads(updated_provider.icon) == {"content": "🚀", "background": "#FFF"} # - changed 3 assert updated_provider.privacy_policy == "https://new-policy.com" # - changed 4 @@ -712,8 +755,11 @@ class TestApiToolManageService: ) def test_update_api_tool_provider_not_found( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """ Test update raises ValueError when original provider not found. @@ -733,7 +779,7 @@ class TestApiToolManageService: user_id=account.id, tenant_id=tenant.id, provider_name=existing_provider_name, - icon={"type": "emoji", "value": "🔧"}, + icon={"content": "🔧", "background": "#FFF"}, credentials={"auth_type": "none"}, schema_type=ApiProviderSchemaType.OPENAPI, schema=self._create_test_openapi_schema(), @@ -756,7 +802,7 @@ class TestApiToolManageService: tenant_id=tenant.id, provider_name=target_new_name, original_provider=missing_original_name, - icon={"type": "emoji", "value": "🚀"}, + icon={"content": "🚀", "background": "#FFF"}, credentials={"auth_type": "none"}, _schema_type=ApiProviderSchemaType.OPENAPI, schema=self._create_test_openapi_schema(), @@ -793,8 +839,11 @@ class TestApiToolManageService: mock_external_service_dependencies["provider_controller"].from_db.assert_not_called() def test_update_api_tool_provider_missing_auth_type( - self, flask_req_ctx_with_containers, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, + flask_req_ctx_with_containers: object, + db_session_with_containers: Session, + mock_external_service_dependencies: MockDependencies, + ) -> None: """Test update raises ValueError when auth_type is missing from credentials.""" fake = Faker() account, tenant = self._create_test_account_and_tenant( @@ -822,7 +871,7 @@ class TestApiToolManageService: tenant_id=tenant.id, provider_name=provider_name, original_provider=provider_name, - icon={}, + icon={"content": "🔧", "background": "#FFF"}, credentials={}, _schema_type=ApiProviderSchemaType.OPENAPI, schema=schema, @@ -832,8 +881,8 @@ class TestApiToolManageService: ) def test_list_api_tool_provider_tools_not_found( - self, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, db_session_with_containers: Session, mock_external_service_dependencies: MockDependencies + ) -> None: """Test listing tools raises ValueError when provider not found.""" fake = Faker() account, tenant = self._create_test_account_and_tenant( @@ -844,8 +893,8 @@ class TestApiToolManageService: ApiToolManageService.list_api_tool_provider_tools(account.id, tenant.id, "nonexistent") def test_test_api_tool_preview_invalid_schema_type( - self, db_session_with_containers: Session, mock_external_service_dependencies - ): + self, db_session_with_containers: Session, mock_external_service_dependencies: MockDependencies + ) -> None: """Test preview raises ValueError for invalid schema type.""" fake = Faker() account, tenant = self._create_test_account_and_tenant( diff --git a/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py b/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py index 6f342e63dc8..b12472c586c 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py +++ b/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py @@ -107,7 +107,7 @@ class TestWorkflowToolManageService: ) app_service = AppService() - app = app_service.create_app(tenant.id, app_args, account) + app = app_service.create_app(tenant.id, app_args, account, session=db_session_with_containers) # Create workflow for the app workflow = WorkflowModel( diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py index ce5c2bd162f..8cd9526f6a0 100644 --- a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py +++ b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py @@ -217,6 +217,7 @@ class TestWorkflowConverter: icon_type="emoji", icon="🚀", icon_background="#4CAF50", + session=db_session_with_containers, ) # Assert: Verify the expected outcomes @@ -291,6 +292,7 @@ class TestWorkflowConverter: icon_type="emoji", icon="🚀", icon_background="#4CAF50", + session=db_session_with_containers, ) # Verify database state remains unchanged @@ -325,6 +327,7 @@ class TestWorkflowConverter: app_model=app, app_model_config=app.app_model_config, account_id=account.id, + session=db_session_with_containers, ) # Assert: Verify the expected outcomes @@ -467,6 +470,7 @@ class TestWorkflowConverter: app_model=app, variables=variables, external_data_variables=external_data_variables, + session=db_session_with_containers, ) # Assert: Verify the expected outcomes @@ -569,7 +573,7 @@ class TestConvertToHttpRequestNodeVariants: """Tests for chatbot vs workflow differences in HTTP request node conversion.""" @staticmethod - def _setup(app_mode, default_variables): + def _setup(app_mode, default_variables, db_session_with_containers: Session): app_model = App( tenant_id="tenant_id", mode=app_mode, @@ -598,19 +602,20 @@ class TestConvertToHttpRequestNodeVariants: app_model=app_model, variables=default_variables, external_data_variables=ext_vars, + session=db_session_with_containers, ) return nodes - def test_chatbot_query_uses_sys_query(self, default_variables): - nodes = self._setup(AppMode.CHAT, default_variables) + def test_chatbot_query_uses_sys_query(self, default_variables, db_session_with_containers: Session): + nodes = self._setup(AppMode.CHAT, default_variables, db_session_with_containers) body = json.loads(nodes[0]["data"]["body"]["data"]) assert body["params"]["query"] == "{{#sys.query#}}" assert body["point"] == APIBasedExtensionPoint.APP_EXTERNAL_DATA_TOOL_QUERY assert nodes[1]["data"]["type"] == "code" - def test_workflow_query_is_empty(self, default_variables): - nodes = self._setup(AppMode.WORKFLOW, default_variables) + def test_workflow_query_is_empty(self, default_variables, db_session_with_containers: Session): + nodes = self._setup(AppMode.WORKFLOW, default_variables, db_session_with_containers) body = json.loads(nodes[0]["data"]["body"]["data"]) assert body["params"]["query"] == "" diff --git a/api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py b/api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py index 9c20118e278..b6865510adf 100644 --- a/api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py +++ b/api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py @@ -194,7 +194,7 @@ def test_webhook_trigger_creates_trigger_log( db_session_with_containers.add_all([webhook_trigger, app_trigger]) db_session_with_containers.commit() - def _fake_trigger_workflow_async(session: Session, user: Any, trigger_data: Any) -> SimpleNamespace: + def _fake_trigger_workflow_async(user: Any, trigger_data: Any, *, session: Session) -> SimpleNamespace: log = WorkflowTriggerLog( tenant_id=trigger_data.tenant_id, app_id=trigger_data.app_id, @@ -575,7 +575,7 @@ def test_schedule_trigger_creates_trigger_log( db_session_with_containers.commit() # Mock AsyncWorkflowService to create WorkflowTriggerLog - def _fake_trigger_workflow_async(session: Session, user: Any, trigger_data: Any) -> SimpleNamespace: + def _fake_trigger_workflow_async(user: Any, trigger_data: Any, *, session: Session) -> SimpleNamespace: log = WorkflowTriggerLog( tenant_id=trigger_data.tenant_id, app_id=trigger_data.app_id, diff --git a/api/tests/unit_tests/clients/agent_backend/test_event_adapter.py b/api/tests/unit_tests/clients/agent_backend/test_event_adapter.py index f6c73fdb66d..4adef713cfd 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_event_adapter.py +++ b/api/tests/unit_tests/clients/agent_backend/test_event_adapter.py @@ -15,6 +15,7 @@ from dify_agent.protocol import ( from pydantic_ai.messages import FinalResultEvent from clients.agent_backend import ( + AgentBackendAgentMessageDeltaInternalEvent, AgentBackendDeferredToolCallInternalEvent, AgentBackendInternalEventType, AgentBackendRunCancelledInternalEvent, @@ -54,6 +55,25 @@ def test_event_adapter_maps_pydantic_ai_stream_event(): assert event.data["event_kind"] == "final_result" +def test_event_adapter_maps_pydantic_ai_stream_event_agent_message_delta_annotation(): + adapted = AgentBackendRunEventAdapter().adapt( + PydanticAIStreamRunEvent( + id="2-0", + run_id="run-1", + data=FinalResultEvent(tool_name=None, tool_call_id=None), + agent_message_delta="hello", + ) + ) + + assert adapted == [ + AgentBackendAgentMessageDeltaInternalEvent( + run_id="run-1", + source_event_id="2-0", + delta="hello", + ) + ] + + def test_event_adapter_maps_run_succeeded_to_final_output(): snapshot = CompositorSessionSnapshot(layers=[]) adapted = AgentBackendRunEventAdapter().adapt( diff --git a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py index f3aac73aac7..b2b77275fa7 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_request_builder.py +++ b/api/tests/unit_tests/clients/agent_backend/test_request_builder.py @@ -88,6 +88,7 @@ def _run_input() -> AgentBackendWorkflowNodeRunInput: def test_request_builder_outputs_dify_agent_create_run_request(): request = AgentBackendRunRequestBuilder().build_for_workflow_node(_run_input()) + dumped = request.model_dump(mode="json") assert isinstance(request, CreateRunRequest) assert [layer.name for layer in request.composition.layers] == [ @@ -102,6 +103,7 @@ def test_request_builder_outputs_dify_agent_create_run_request(): assert request.on_exit.default is ExitIntent.SUSPEND assert request.idempotency_key == "workflow-run-1:node-execution-1" assert request.metadata == {"workflow_id": "workflow-1", "node_id": "node-1"} + assert "purpose" not in dumped def test_request_builder_separates_agent_soul_and_workflow_job_prompt(): @@ -121,6 +123,59 @@ def test_request_builder_separates_agent_soul_and_workflow_job_prompt(): assert dumped["composition"]["layers"][2]["config"]["user"] == "Summarize the report." +@pytest.mark.parametrize("agent_config_version_kind", ["snapshot", "draft"]) +def test_agent_app_request_builder_keeps_agent_soul_prompt_for_snapshot_and_draft( + agent_config_version_kind: str, +): + original_prompt = " You are Iris. \n" + run_input = _agent_app_input().model_copy( + update={ + "agent_config_version_kind": agent_config_version_kind, + "agent_soul_prompt": original_prompt, + } + ) + + request = AgentBackendRunRequestBuilder().build_for_agent_app(run_input) + layers = {layer.name: layer for layer in request.composition.layers} + + prompt_config = cast(PromptLayerConfig, layers[AGENT_SOUL_PROMPT_LAYER_ID].config) + assert prompt_config.prefix == original_prompt + + +def test_agent_app_request_builder_wraps_agent_soul_prompt_for_build_draft(): + original_prompt = " You are Iris. \n" + run_input = _agent_app_input().model_copy( + update={ + "agent_config_version_kind": "build_draft", + "agent_soul_prompt": original_prompt, + } + ) + + request = AgentBackendRunRequestBuilder().build_for_agent_app(run_input) + layers = {layer.name: layer for layer in request.composition.layers} + + prompt_config = cast(PromptLayerConfig, layers[AGENT_SOUL_PROMPT_LAYER_ID].config) + assert prompt_config.prefix != original_prompt + assert prompt_config.prefix.startswith("You are running in build mode.") + assert "```text\nYou are Iris.\n```" in prompt_config.prefix + + +def test_agent_app_request_builder_uses_longer_fence_for_build_draft_prompt_body(): + run_input = _agent_app_input().model_copy( + update={ + "agent_config_version_kind": "build_draft", + "agent_soul_prompt": "Keep this snippet:\n```python\nprint('hi')\n```", + } + ) + + request = AgentBackendRunRequestBuilder().build_for_agent_app(run_input) + layers = {layer.name: layer for layer in request.composition.layers} + + prompt_config = cast(PromptLayerConfig, layers[AGENT_SOUL_PROMPT_LAYER_ID].config) + assert "````text" in prompt_config.prefix + assert "```python" in prompt_config.prefix + + def test_request_builder_sets_model_and_output_layer_contract_ids(): request = AgentBackendRunRequestBuilder().build_for_workflow_node(_run_input()) layers = {layer.name: layer for layer in request.composition.layers} @@ -250,6 +305,7 @@ def test_request_builder_builds_cleanup_request_replays_persisted_layer_specs(): assert request.on_exit.default is ExitIntent.DELETE assert request.idempotency_key == "run-1:node-1:binding-1:agent-session-cleanup" assert request.metadata["agent_backend_lifecycle"] == "session_cleanup" + assert "purpose" not in request.model_dump(mode="json") def test_request_builder_rejects_empty_runtime_layer_specs(): @@ -392,6 +448,19 @@ def test_agent_app_request_builder_omits_shell_layer_by_default(): assert DIFY_SHELL_LAYER_ID not in {layer.name for layer in request.composition.layers} +def test_agent_app_request_builder_keeps_build_draft_prompt_when_agent_soul_prompt_is_blank(): + run_input = _agent_app_input().model_copy( + update={"agent_soul_prompt": " ", "agent_config_version_kind": "build_draft"} + ) + + request = AgentBackendRunRequestBuilder().build_for_agent_app(run_input) + layers = {layer.name: layer for layer in request.composition.layers} + + prompt_config = cast(PromptLayerConfig, layers[AGENT_SOUL_PROMPT_LAYER_ID].config) + assert "You are running in build mode." in prompt_config.prefix + assert "No task prompt was provided." in prompt_config.prefix + + def test_agent_app_request_builder_adds_shell_layer_when_include_shell(): run_input = _agent_app_input(include_shell=True) run_input.shell_config = DifyShellLayerConfig(env=[DifyShellEnvVarConfig(name="APP_ENV", value="enabled")]) diff --git a/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py b/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py new file mode 100644 index 00000000000..6b72850aeca --- /dev/null +++ b/api/tests/unit_tests/clients/agent_backend/test_session_cleanup.py @@ -0,0 +1,124 @@ +from datetime import UTC, datetime + +from agenton.compositor import CompositorSessionSnapshot +from agenton.compositor.schemas import LayerSessionSnapshot +from agenton.layers.base import LifecycleState +from dify_agent.protocol import RunStatusResponse + +from clients.agent_backend import ( + AgentBackendError, + AgentBackendSessionCleanupPayload, + FakeAgentBackendRunClient, + RuntimeLayerSpec, + cleanup_agent_backend_session, +) + + +def _payload() -> AgentBackendSessionCleanupPayload: + return AgentBackendSessionCleanupPayload( + session_snapshot=CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot( + name="history", + lifecycle_state=LifecycleState.SUSPENDED, + runtime_state={}, + ) + ] + ), + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + idempotency_key="cleanup-1", + metadata={"tenant_id": "tenant-1"}, + timeout_seconds=15.0, + ) + + +def test_cleanup_agent_backend_session_runs_create_and_wait_until_success(): + client = FakeAgentBackendRunClient(run_id="cleanup-run-1") + + result = cleanup_agent_backend_session(payload=_payload(), client=client) + + assert result.status == "succeeded" + assert result.cleanup_run_id == "cleanup-run-1" + assert client.request is not None + assert [layer.name for layer in client.request.composition.layers] == ["history"] + + +def test_cleanup_agent_backend_session_skips_when_client_is_missing(): + result = cleanup_agent_backend_session( + payload=_payload(), + client=None, + ) + + assert result.status == "skipped" + assert result.reason == "no_agent_backend_client" + + +def test_cleanup_agent_backend_session_skips_when_session_snapshot_is_missing(): + payload = _payload().model_copy(update={"session_snapshot": None}) + + result = cleanup_agent_backend_session(payload=payload, client=FakeAgentBackendRunClient()) + + assert result.status == "skipped" + assert result.reason == "missing_session_snapshot" + + +def test_cleanup_agent_backend_session_skips_when_runtime_layer_specs_are_missing(): + payload = _payload().model_copy(update={"runtime_layer_specs": []}) + + result = cleanup_agent_backend_session(payload=payload, client=FakeAgentBackendRunClient()) + + assert result.status == "skipped" + assert result.reason == "missing_runtime_layer_specs" + + +class _FailedStatusClient(FakeAgentBackendRunClient): + def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse: + del timeout_seconds + return RunStatusResponse( + run_id=run_id, + status="failed", + created_at=datetime(2026, 1, 1, tzinfo=UTC), + updated_at=datetime(2026, 1, 1, tzinfo=UTC), + error="snapshot mismatch", + ) + + +def test_cleanup_agent_backend_session_reports_failed_terminal_status(): + client = _FailedStatusClient(run_id="cleanup-run-2") + + result = cleanup_agent_backend_session(payload=_payload(), client=client) + + assert result.status == "failed" + assert result.reason == "snapshot mismatch" + assert result.cleanup_run_id == "cleanup-run-2" + + +class _CreateRunFailureClient(FakeAgentBackendRunClient): + def create_run(self, request): # type: ignore[override] + del request + raise AgentBackendError("create run failed") + + +class _WaitRunFailureClient(FakeAgentBackendRunClient): + def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse: + del run_id, timeout_seconds + raise AgentBackendError("wait run failed") + + +def test_cleanup_agent_backend_session_returns_failed_when_create_run_raises(): + result = cleanup_agent_backend_session(payload=_payload(), client=_CreateRunFailureClient()) + + assert result.status == "failed" + assert result.reason == "create run failed" + assert result.cleanup_run_id is None + + +def test_cleanup_agent_backend_session_returns_failed_with_cleanup_run_id_when_wait_run_raises(): + result = cleanup_agent_backend_session( + payload=_payload(), + client=_WaitRunFailureClient(run_id="cleanup-run-3"), + ) + + assert result.status == "failed" + assert result.reason == "wait run failed" + assert result.cleanup_run_id == "cleanup-run-3" diff --git a/api/tests/unit_tests/commands/test_check_no_new_getattr.py b/api/tests/unit_tests/commands/test_check_no_new_getattr.py index 2fbaea62d64..72e72a631aa 100644 --- a/api/tests/unit_tests/commands/test_check_no_new_getattr.py +++ b/api/tests/unit_tests/commands/test_check_no_new_getattr.py @@ -77,6 +77,10 @@ def assert_has_actionable_violation(stderr: str, path: str) -> None: assert "no-new-getattr" in stderr +def main_branch_rev(repo: Path) -> str: + return git(repo, "rev-parse", "main") + + def test_resolve_ast_grep_command_prefers_ast_grep(monkeypatch: pytest.MonkeyPatch) -> None: module = load_guard_module() monkeypatch.setattr( @@ -130,8 +134,46 @@ def test_resolve_ast_grep_command_raises_without_explicit_binary(monkeypatch: py module.resolve_ast_grep_command() +def test_cli_requires_explicit_diff_source(tmp_path: Path) -> None: + result = run_script(tmp_path) + + assert result.returncode == 2 + assert "one of the arguments --staged --base-rev is required" in result.stderr + + +def test_cli_rejects_mixed_diff_sources(tmp_path: Path) -> None: + result = run_script(tmp_path, "--staged", "--base-rev", "deadbeef") + + assert result.returncode == 2 + assert "not allowed with argument" in result.stderr + + +def test_cli_help_exposes_only_new_diff_source_flags(tmp_path: Path) -> None: + help_result = run_script(tmp_path, "--help") + + assert help_result.returncode == 0 + assert "--staged" in help_result.stdout + assert "--base-rev" in help_result.stdout + assert "--mode" not in help_result.stdout + assert "--merge-target" not in help_result.stdout + + result = run_script(tmp_path, "--staged", "--mode", "ci") + + assert result.returncode == 2 + assert "unrecognized arguments: --mode ci" in result.stderr + + result = run_script(tmp_path, "--base-rev", "deadbeef", "--merge-target", "main") + + assert result.returncode == 2 + assert "unrecognized arguments: --merge-target main" in result.stderr + + def test_style_workflow_wires_no_new_getattr_guard() -> None: workflow = (REPO_ROOT / ".github" / "workflows" / "style.yml").read_text(encoding="utf-8") + assert re.search( + r"(?ms)^on:\n workflow_call:\n inputs:\n base-rev:\n required: true\n type: string\n", + workflow, + ) python_style_job = re.search( r"(?ms)^ python-style:\n(?P.*?)(?=^ [a-z0-9-]+:\n|\Z)", workflow, @@ -157,8 +199,9 @@ def test_style_workflow_wires_no_new_getattr_guard() -> None: assert "scripts/check_no_new_getattr.py\n" in files_block assert "scripts/ast_grep_rules/no_new_getattr.yml\n" in files_block assert ".github/workflows/style.yml\n" in files_block + assert ".github/workflows/main-ci.yml\n" in files_block - guard_command = "scripts/check_no_new_getattr.py --mode ci --merge-target main" + guard_command = 'scripts/check_no_new_getattr.py --base-rev "${{ inputs.base-rev }}"' assert guard_command in job_text guard_step = re.search( @@ -168,55 +211,35 @@ def test_style_workflow_wires_no_new_getattr_guard() -> None: ) assert guard_step is not None - pre_guard_text = job_text[: guard_step.start()] - step_pattern = r"(?ms)^ - name: [^\n]*\n(?P.*?)(?=^ - name: |\Z)" - fetch_step_text = next( - ( - match.group("step") - for match in re.finditer(step_pattern, pre_guard_text) - if any( - re.search(pattern, line) - for line in match.group("step").splitlines() - for pattern in ( - r"git fetch .*refs/heads/main:refs/remotes/origin/main", - r"git fetch .*main:refs/remotes/origin/main", - r"git fetch .*refs/remotes/origin/main", - ) - ) - ), - "", - ) - assert fetch_step_text - assert "git fetch" in fetch_step_text - assert "origin" in fetch_step_text - assert any( - re.search(pattern, line) - for line in fetch_step_text.splitlines() - for pattern in ( - r"git fetch .*refs/heads/main:refs/remotes/origin/main", - r"git fetch .*main:refs/remotes/origin/main", - r"git fetch .*refs/remotes/origin/main", - ) - ) - - bind_step = re.search( - r"(?ms)^ - name: Bind merge target branch for getattr guard\n(?P.*?)(?=^ - name: |\Z)", - pre_guard_text, - ) - assert bind_step is not None - bind_step_text = bind_step.group("step") - assert any( - command in bind_step_text - for command in ( - "git branch main origin/main", - "git checkout -B main origin/main", - "git switch -C main origin/main", - "git update-ref refs/heads/main refs/remotes/origin/main", - ) - ) + assert "GITHUB_BASE_SHA" not in guard_step.group("step") -def test_ci_mode_passes_when_only_legacy_getattr_exists(tmp_path: Path) -> None: +def test_main_ci_passes_style_base_rev_input() -> None: + workflow = (REPO_ROOT / ".github" / "workflows" / "main-ci.yml").read_text(encoding="utf-8") + style_job = re.search( + r"(?ms)^ style-check:\n(?P.*?)(?=^ [a-z0-9-]+:\n|\Z)", + workflow, + ) + assert style_job is not None + assert "uses: ./.github/workflows/style.yml" in style_job.group("job") + assert ( + "base-rev: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha }}" + in style_job.group("job") + ) + + api_filter = re.search( + r"(?ms)^ api:\n(?P(?:^ - '[^']+'\n)+)", + workflow, + ) + assert api_filter is not None + filter_text = api_filter.group("filter") + assert "scripts/check_no_new_getattr.py" in filter_text + assert "scripts/ast_grep_rules/no_new_getattr.yml" in filter_text + assert ".github/workflows/style.yml" in filter_text + assert ".github/workflows/main-ci.yml" in filter_text + + +def test_base_rev_mode_passes_when_only_legacy_getattr_exists(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -239,12 +262,12 @@ def test_ci_mode_passes_when_only_legacy_getattr_exists(tmp_path: Path) -> None: ) commit_all(tmp_path, "unrelated change") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 0, result.stderr -def test_ci_mode_fails_for_new_file_with_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -267,13 +290,13 @@ def test_ci_mode_fails_for_new_file_with_getattr(tmp_path: Path) -> None: ) commit_all(tmp_path, "add new getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_fails_for_new_file_with_two_arg_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_two_arg_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -296,13 +319,13 @@ def test_ci_mode_fails_for_new_file_with_two_arg_getattr(tmp_path: Path) -> None ) commit_all(tmp_path, "add new two-arg getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_fails_for_new_file_with_builtins_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -328,13 +351,13 @@ def test_ci_mode_fails_for_new_file_with_builtins_getattr(tmp_path: Path) -> Non ) commit_all(tmp_path, "add new builtins getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_fails_for_new_file_with_two_arg_builtins_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_two_arg_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -360,13 +383,13 @@ def test_ci_mode_fails_for_new_file_with_two_arg_builtins_getattr(tmp_path: Path ) commit_all(tmp_path, "add new two-arg builtins getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_fails_for_new_file_with_dunder_builtins_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_dunder_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -389,13 +412,13 @@ def test_ci_mode_fails_for_new_file_with_dunder_builtins_getattr(tmp_path: Path) ) commit_all(tmp_path, "add new dunder builtins getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_fails_for_new_file_with_two_arg_dunder_builtins_getattr(tmp_path: Path) -> None: +def test_base_rev_mode_fails_for_new_file_with_two_arg_dunder_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -418,13 +441,13 @@ def test_ci_mode_fails_for_new_file_with_two_arg_dunder_builtins_getattr(tmp_pat ) commit_all(tmp_path, "add new two-arg dunder builtins getattr usage") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", main_branch_rev(tmp_path)) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/new_usage.py") -def test_ci_mode_uses_merge_base_against_main_not_just_head_parent(tmp_path: Path) -> None: +def test_base_rev_mode_uses_provided_base_revision_not_head_parent(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -435,6 +458,7 @@ def test_ci_mode_uses_merge_base_against_main_not_just_head_parent(tmp_path: Pat """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -457,13 +481,75 @@ def test_ci_mode_uses_merge_base_against_main_not_just_head_parent(tmp_path: Pat ) commit_all(tmp_path, "later feature commit does not touch violating file") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/introduced_earlier.py") -def test_pre_commit_mode_reads_staged_content_only(tmp_path: Path) -> None: +def test_base_rev_mode_works_without_local_main_branch(tmp_path: Path) -> None: + init_repo(tmp_path) + write_repo_file( + tmp_path, + "pkg/existing.py", + """ + def stable() -> str: + return "ok" + """, + ) + commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) + checkout_feature_branch(tmp_path) + + write_repo_file( + tmp_path, + "pkg/other.py", + """ + def meaning() -> int: + return 42 + """, + ) + commit_all(tmp_path, "feature change") + + git(tmp_path, "checkout", "--detach", "HEAD") + git(tmp_path, "branch", "-D", "main") + + result = run_script(tmp_path, "--base-rev", base_rev) + + assert result.returncode == 0, result.stderr + + +def test_base_rev_mode_ignores_github_base_sha_environment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + init_repo(tmp_path) + write_repo_file( + tmp_path, + "pkg/existing.py", + """ + def stable() -> str: + return "ok" + """, + ) + commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) + checkout_feature_branch(tmp_path) + + write_repo_file( + tmp_path, + "pkg/other.py", + """ + def meaning() -> int: + return 42 + """, + ) + commit_all(tmp_path, "feature change") + monkeypatch.setenv("GITHUB_BASE_SHA", "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef") + + result = run_script(tmp_path, "--base-rev", base_rev) + + assert result.returncode == 0, result.stderr + + +def test_staged_mode_reads_staged_content_only(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -494,12 +580,12 @@ def test_pre_commit_mode_reads_staged_content_only(tmp_path: Path) -> None: """, ) - result = run_script(tmp_path, "--mode", "pre-commit") + result = run_script(tmp_path, "--staged") assert result.returncode == 0, result.stderr -def test_pre_commit_mode_fails_for_staged_two_arg_getattr(tmp_path: Path) -> None: +def test_staged_mode_fails_for_staged_two_arg_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -521,13 +607,13 @@ def test_pre_commit_mode_fails_for_staged_two_arg_getattr(tmp_path: Path) -> Non ) git(tmp_path, "add", "pkg/module.py") - result = run_script(tmp_path, "--mode", "pre-commit") + result = run_script(tmp_path, "--staged") assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/module.py") -def test_pre_commit_mode_fails_for_staged_builtins_getattr(tmp_path: Path) -> None: +def test_staged_mode_fails_for_staged_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -552,13 +638,13 @@ def test_pre_commit_mode_fails_for_staged_builtins_getattr(tmp_path: Path) -> No ) git(tmp_path, "add", "pkg/module.py") - result = run_script(tmp_path, "--mode", "pre-commit") + result = run_script(tmp_path, "--staged") assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/module.py") -def test_pre_commit_mode_fails_for_staged_two_arg_builtins_getattr(tmp_path: Path) -> None: +def test_staged_mode_fails_for_staged_two_arg_builtins_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( tmp_path, @@ -583,7 +669,7 @@ def test_pre_commit_mode_fails_for_staged_two_arg_builtins_getattr(tmp_path: Pat ) git(tmp_path, "add", "pkg/module.py") - result = run_script(tmp_path, "--mode", "pre-commit") + result = run_script(tmp_path, "--staged") assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/module.py") @@ -603,6 +689,7 @@ def test_modified_hunk_with_same_getattr_count_is_allowed(tmp_path: Path) -> Non """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -618,7 +705,7 @@ def test_modified_hunk_with_same_getattr_count_is_allowed(tmp_path: Path) -> Non ) commit_all(tmp_path, "touch legacy getattr hunk") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 0, result.stderr @@ -635,6 +722,7 @@ def test_modified_hunk_with_decreased_getattr_count_is_allowed(tmp_path: Path) - """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -647,7 +735,7 @@ def test_modified_hunk_with_decreased_getattr_count_is_allowed(tmp_path: Path) - ) commit_all(tmp_path, "remove one legacy getattr") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 0, result.stderr @@ -663,6 +751,7 @@ def test_modified_hunk_with_increased_getattr_count_fails(tmp_path: Path) -> Non """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -676,7 +765,7 @@ def test_modified_hunk_with_increased_getattr_count_fails(tmp_path: Path) -> Non ) commit_all(tmp_path, "add one more getattr") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/sample.py") @@ -694,6 +783,7 @@ def test_inline_noqa_suppression_with_explanatory_text_skips_added_getattr(tmp_p """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -706,7 +796,7 @@ def test_inline_noqa_suppression_with_explanatory_text_skips_added_getattr(tmp_p ) commit_all(tmp_path, "add suppressed getattr") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert "no-new-getattr needed for plugin-defined attributes" in (tmp_path / "pkg/existing.py").read_text( encoding="utf-8" @@ -725,6 +815,7 @@ def test_inline_noqa_without_explanatory_text_is_not_sufficient(tmp_path: Path) """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -737,7 +828,7 @@ def test_inline_noqa_without_explanatory_text_is_not_sufficient(tmp_path: Path) ) commit_all(tmp_path, "add bare noqa getattr") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 1 assert_has_actionable_violation(result.stderr, "pkg/existing.py") @@ -753,6 +844,7 @@ def test_non_python_file_with_getattr_text_does_not_fail_guard(tmp_path: Path) - """, ) commit_all(tmp_path, "baseline") + base_rev = main_branch_rev(tmp_path) checkout_feature_branch(tmp_path) write_repo_file( @@ -764,6 +856,6 @@ def test_non_python_file_with_getattr_text_does_not_fail_guard(tmp_path: Path) - ) commit_all(tmp_path, "document getattr example") - result = run_script(tmp_path, "--mode", "ci", "--merge-target", "main") + result = run_script(tmp_path, "--base-rev", base_rev) assert result.returncode == 0, result.stderr diff --git a/api/tests/unit_tests/commands/test_data_migration_commands.py b/api/tests/unit_tests/commands/test_data_migration_commands.py index b7f92f3291a..84e39d19a2d 100644 --- a/api/tests/unit_tests/commands/test_data_migration_commands.py +++ b/api/tests/unit_tests/commands/test_data_migration_commands.py @@ -109,8 +109,8 @@ def test_export_command_uses_cli_owned_session(monkeypatch, tmp_path: Path): package = MigrationPackage.from_mapping({"metadata": {"version": "1", "source_scope": "single"}}) class FakeMigrationExportService: - def export(self, export_session, selection): - captured["session"] = export_session + def export(self, selection, *, session): + captured["session"] = session captured["selection"] = selection return ExportResult(package=package, report_items=[], report_context=ReportContext()) @@ -156,8 +156,8 @@ def test_import_command_uses_cli_owned_session(monkeypatch, tmp_path: Path): ) class FakeMigrationImportService: - def import_package(self, import_session, request): - captured["session"] = import_session + def import_package(self, request, *, session): + captured["session"] = session captured["request"] = request return ImportResult(report_items=[], report_context=ReportContext(target_tenant="target")) diff --git a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py index e14c8ed3243..9eb73b82516 100644 --- a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py +++ b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py @@ -32,15 +32,6 @@ from tests.helpers.legacy_model_type_migration import ( ) -@pytest.fixture -def sqlite_engine(tmp_path: Path) -> sa.Engine: - engine = sa.create_engine(f"sqlite:///{tmp_path / 'legacy_model_type_migration.sqlite'}") - try: - yield engine - finally: - engine.dispose() - - @pytest.fixture def dirty_fixture(sqlite_engine: sa.Engine): return seed_legacy_model_type_dirty_data(sqlite_engine) diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index 7174530e976..d95e5e8501d 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -1,9 +1,12 @@ import os +from collections.abc import Iterator from unittest.mock import MagicMock, patch import pytest from flask import Flask from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker # Getting the absolute path of the current file's directory ABS_PATH = os.path.dirname(os.path.abspath(__file__)) @@ -34,6 +37,7 @@ os.environ.setdefault("STORAGE_TYPE", "opendal") from core.db.session_factory import configure_session_factory, session_factory from extensions import ext_redis +from models.base import TypeBase def _patch_redis_clients_on_loaded_modules(): @@ -113,6 +117,29 @@ def _unit_test_engine(): engine.dispose() +@pytest.fixture +def sqlite_engine() -> Iterator[Engine]: + """Create an isolated in-memory SQLite engine for tests that need a disposable database.""" + + engine = create_engine("sqlite:///:memory:") + try: + yield engine + finally: + engine.dispose() + + +@pytest.fixture +def sqlite_session(request: pytest.FixtureRequest, sqlite_engine: Engine) -> Iterator[Session]: + """Yield a SQLite session after creating the model tables passed through ``request.param``.""" + + models: tuple[type[TypeBase], ...] = request.param + tables = [model.metadata.tables[model.__tablename__] for model in models] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + with session_factory() as session: + yield session + + @pytest.fixture(autouse=True) def _configure_session_factory(_unit_test_engine): try: diff --git a/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py b/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py new file mode 100644 index 00000000000..b4639d54d2f --- /dev/null +++ b/api/tests/unit_tests/controllers/common/test_agent_app_parameters.py @@ -0,0 +1,175 @@ +from types import SimpleNamespace + +import pytest + +from controllers.common import agent_app_parameters +from controllers.common.agent_app_parameters import get_published_agent_app_feature_dict_and_user_input_form +from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError + + +def test_published_agent_app_parameters_use_soul_file_upload(monkeypatch): + app_model_config = SimpleNamespace( + to_dict=lambda: { + "opening_statement": "Hi from legacy presentation config", + "file_upload": { + "enabled": False, + "image": {"enabled": False}, + }, + } + ) + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=app_model_config, + ) + agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id="snapshot-1", + active_config_is_published=True, + ) + snapshot = SimpleNamespace( + config_snapshot_dict={ + "app_features": { + "file_upload": { + "enabled": True, + "allowed_file_extensions": ["PNG"], + "allowed_file_types": ["image"], + "allowed_file_upload_methods": ["local_file"], + "image": {"enabled": True}, + "number_limits": 2, + } + }, + "app_variables": [{"name": "topic", "type": "string", "required": True}], + } + ) + query_results = iter([agent, snapshot]) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: next(query_results)) + + features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(app_model) + parameters = get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form) + + assert parameters["opening_statement"] == "Hi from legacy presentation config" + assert parameters["file_upload"] == { + "enabled": True, + "allowed_file_extensions": ["PNG"], + "allowed_file_types": ["image"], + "allowed_file_upload_methods": ["local_file"], + "image": {"enabled": True}, + "number_limits": 2, + } + assert parameters["user_input_form"] == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}] + + +def test_published_agent_app_parameters_requires_bound_agent(): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id=None, + app_model_config=None, + ) + + with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): + get_published_agent_app_feature_dict_and_user_input_form(app_model) + + +def test_published_agent_app_parameters_requires_existing_active_agent(monkeypatch): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=None, + ) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: None) + + with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): + get_published_agent_app_feature_dict_and_user_input_form(app_model) + + +@pytest.mark.parametrize( + "active_config_is_published", + [ + True, + False, + ], +) +def test_published_agent_app_parameters_requires_published_agent(monkeypatch, active_config_is_published): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=None, + ) + agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id=None, + active_config_is_published=active_config_is_published, + ) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: agent) + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + get_published_agent_app_feature_dict_and_user_input_form(app_model) + + +def test_published_agent_app_parameters_allows_unpublished_draft_with_active_snapshot(monkeypatch): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=None, + ) + agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id="snapshot-1", + active_config_is_published=False, + ) + snapshot = SimpleNamespace(config_snapshot_dict={}) + query_results = iter([agent, snapshot]) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: next(query_results)) + + features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(app_model) + + assert features_dict["file_upload"]["enabled"] is True + assert user_input_form == [] + + +def test_published_agent_app_parameters_requires_published_snapshot(monkeypatch): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=None, + ) + agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id="snapshot-1", + active_config_is_published=True, + ) + query_results = iter([agent, None]) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: next(query_results)) + + with pytest.raises(AgentAppGeneratorError, match="published version not found"): + get_published_agent_app_feature_dict_and_user_input_form(app_model) + + +def test_published_agent_app_parameters_allows_missing_legacy_app_model_config(monkeypatch): + app_model = SimpleNamespace( + tenant_id="tenant-1", + bound_agent_id="agent-1", + app_model_config=None, + ) + agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id="snapshot-1", + active_config_is_published=True, + ) + snapshot = SimpleNamespace(config_snapshot_dict={}) + query_results = iter([agent, snapshot]) + monkeypatch.setattr(agent_app_parameters.db.session, "scalar", lambda _: next(query_results)) + + features_dict, user_input_form = get_published_agent_app_feature_dict_and_user_input_form(app_model) + + assert features_dict["file_upload"] == { + "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], + "allowed_file_types": ["document", "image", "audio", "video"], + "allowed_file_upload_methods": ["local_file", "remote_url"], + "enabled": True, + "image": {"enabled": True}, + "number_limits": 3, + } + assert user_input_form == [] diff --git a/api/tests/unit_tests/controllers/common/test_app_access.py b/api/tests/unit_tests/controllers/common/test_app_access.py index d070cc6e0fc..60a576346a0 100644 --- a/api/tests/unit_tests/controllers/common/test_app_access.py +++ b/api/tests/unit_tests/controllers/common/test_app_access.py @@ -152,7 +152,7 @@ class TestResolveAppAccessFilter: self._patch_whitelist(monkeypatch, ResourceWhitelistResources(unrestricted=False, resource_ids=[])) monkeypatch.setattr( f"{_RBAC_MODULE}.RBACService.MyPermissions.get", - lambda tenant_id, account_id: _permissions(workspace_keys=["app.create_and_management"]), + lambda tenant_id, account_id, session: _permissions(workspace_keys=["app.create_and_management"]), ) flt = resolve_app_access_filter("tenant-1", "acc-1") diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index 2e5851349d4..51737a859ba 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -236,7 +236,7 @@ def test_agent_app_list_and_create_use_agent_route( items=[_app_detail_obj(id="app-list", bound_agent_id="agent-list")], ) - def create_app(self, tenant_id: str, params, current_user: object) -> object: + def create_app(self, tenant_id: str, params, current_user: object, *, session: object) -> object: captured["create"] = {"tenant_id": tenant_id, "params": params, "current_user": current_user} return _app_detail_obj(id="app-created", bound_agent_id="agent-created") @@ -392,7 +392,8 @@ def test_agent_app_create_omits_optional_role_as_empty_string( captured: dict[str, object] = {} class FakeAppService: - def create_app(self, tenant_id: str, params: object, account: object) -> object: + def create_app(self, tenant_id: str, params: object, account: object, *, session: object) -> object: + del session captured["create"] = {"tenant_id": tenant_id, "params": params, "account": account} return _app_detail_obj(id="app-created", bound_agent_id="agent-created") @@ -472,11 +473,11 @@ def test_agent_app_detail_update_delete_resolve_app_from_agent_id( captured["get_app"] = app_obj return app_obj - def update_app(self, app_obj: object, args: dict[str, object]) -> object: + def update_app(self, app_obj: object, args: dict[str, object], *, session: object) -> object: captured["update"] = {"app": app_obj, "args": args} return _app_detail_obj(id="app-1", name=args["name"], bound_agent_id=agent_id) - def delete_app(self, app_obj: object) -> None: + def delete_app(self, app_obj: object, *, session: object) -> None: captured["delete"] = app_obj monkeypatch.setattr(roster_controller, "AppService", FakeAppService) @@ -661,18 +662,26 @@ def test_agent_publish_and_build_draft_routes_call_composer_service( discard_agent_app_build_draft, ) + def assert_call_without_session(key: str, expected: dict[str, object]) -> None: + call = dict(captured[key]) # type: ignore[arg-type] + assert call.pop("session", None) is not None + assert call == expected + with app.test_request_context( "/console/api/agent/00000000-0000-0000-0000-000000000001/publish", json={"version_note": "publish v1"}, ): published = unwrap(AgentPublishApi.post)(AgentPublishApi(), "tenant-1", current_user, agent_id) assert published["active_config_snapshot_id"] == "version-1" - assert captured["publish"] == { - "tenant_id": "tenant-1", - "agent_id": agent_id, - "account_id": account_id, - "version_note": "publish v1", - } + assert_call_without_session( + "publish", + { + "tenant_id": "tenant-1", + "agent_id": agent_id, + "account_id": account_id, + "version_note": "publish v1", + }, + ) with app.test_request_context( "/console/api/agent/00000000-0000-0000-0000-000000000001/build-draft/checkout", @@ -682,17 +691,20 @@ def test_agent_publish_and_build_draft_routes_call_composer_service( AgentBuildDraftCheckoutApi(), "tenant-1", current_user, agent_id ) assert checked_out["draft"]["id"] == "build-draft-1" - assert captured["checkout"] == { - "tenant_id": "tenant-1", - "agent_id": agent_id, - "account_id": account_id, - "force": True, - } + assert_call_without_session( + "checkout", + { + "tenant_id": "tenant-1", + "agent_id": agent_id, + "account_id": account_id, + "force": True, + }, + ) with app.test_request_context("/console/api/agent/00000000-0000-0000-0000-000000000001/build-draft"): loaded = unwrap(AgentBuildDraftApi.get)(AgentBuildDraftApi(), "tenant-1", current_user, agent_id) assert loaded["draft"]["id"] == "build-draft-1" - assert captured["load"] == {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id} + assert_call_without_session("load", {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id}) with app.test_request_context( "/console/api/agent/00000000-0000-0000-0000-000000000001/build-draft", @@ -711,7 +723,7 @@ def test_agent_publish_and_build_draft_routes_call_composer_service( ): applied = unwrap(AgentBuildDraftApplyApi.post)(AgentBuildDraftApplyApi(), "tenant-1", current_user, agent_id) assert applied == {"result": "success", "draft": {"id": "draft-1"}} - assert captured["apply"] == {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id} + assert_call_without_session("apply", {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id}) with app.test_request_context( "/console/api/agent/00000000-0000-0000-0000-000000000001/build-draft", @@ -719,7 +731,7 @@ def test_agent_publish_and_build_draft_routes_call_composer_service( ): discarded = unwrap(AgentBuildDraftApi.delete)(AgentBuildDraftApi(), "tenant-1", current_user, agent_id) assert discarded == {"result": "success"} - assert captured["discard"] == {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id} + assert_call_without_session("discard", {"tenant_id": "tenant-1", "agent_id": agent_id, "account_id": account_id}) def test_agent_api_access_uses_agent_id_and_returns_service_api_metadata( @@ -775,7 +787,7 @@ def test_agent_api_status_and_key_routes_resolve_backing_app( monkeypatch.setattr(roster_controller, "_agent_api_key_count", lambda app_id: 1) class FakeAppService: - def update_app_api_status(self, app_obj: object, enable_api: bool) -> object: + def update_app_api_status(self, app_obj: object, enable_api: bool, *, session: object) -> object: captured["enable"] = {"app": app_obj, "enable_api": enable_api} app_model.enable_api = enable_api return app_model @@ -890,7 +902,7 @@ def test_agent_app_update_allows_empty_role(app: Flask, monkeypatch: pytest.Monk def get_app(self, app_obj: object) -> object: return app_obj - def update_app(self, app_obj: object, args: dict[str, object]) -> object: + def update_app(self, app_obj: object, args: dict[str, object], *, session: object) -> object: captured["update"] = {"app": app_obj, "args": args} return _app_detail_obj(id="app-1", name=args["name"], bound_agent_id=agent_id) @@ -1292,6 +1304,7 @@ def test_workflow_composer_copy_from_roster(app: Flask, monkeypatch: pytest.Monk ) assert result["binding"]["binding_type"] == "inline_agent" + assert captured.pop("session") is not None assert captured == { "tenant_id": "tenant-1", "app_id": "app-1", @@ -1563,6 +1576,7 @@ def test_build_chat_finalization_helper_forces_debug_build_and_push_prompt( assert args["conversation_id"] == "debug-conversation-1" assert args["inputs"] == {} assert args["auto_generate_name"] is False + assert args[completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG] == "delete" assert args["external_trace_id"] == "trace-1" @@ -1651,6 +1665,51 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace( assert args["external_trace_id"] == "trace-1" +def test_agent_chat_helper_ignores_private_exit_intent_payload_key( + app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str +) -> None: + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode="agent") + current_user = SimpleNamespace(id=account_id) + captured: dict[str, object] = {} + + def generate(**kwargs: object) -> dict[str, object]: + captured.update(kwargs) + return {"answer": "ok"} + + monkeypatch.setattr(completion_controller.AppGenerateService, "generate", generate) + monkeypatch.setattr( + completion_controller, + "_resolve_current_user_agent_debug_conversation_id", + lambda **kwargs: "debug-conversation-1", + ) + monkeypatch.setattr( + completion_controller.helper, + "compact_generate_response", + lambda response: {"response": response}, + ) + + with app.test_request_context( + json={ + "inputs": {}, + "query": "hello", + "response_mode": "streaming", + completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG: "delete", + } + ): + result = completion_controller._create_chat_message( + current_user=current_user, + app_model=app_model, + session=Mock(), + ) + + assert result == {"response": {"answer": "ok"}} + assert captured["streaming"] is True + args = cast(dict[str, object], captured["args"]) + assert args["response_mode"] == "streaming" + assert args["conversation_id"] == "debug-conversation-1" + assert completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG not in args + + def test_agent_chat_helper_rejects_foreign_debug_conversation( app: Flask, monkeypatch: pytest.MonkeyPatch, @@ -1896,8 +1955,18 @@ def test_list_agent_chat_messages_uses_current_user_conversation( captured.update(kwargs) return conversation + class SessionProxy: + def __call__(self): + return session + + def scalar(self, stmt: object): + return session.scalar(stmt) + + def scalars(self, stmt: object): + return session.scalars(stmt) + monkeypatch.setattr(message_controller.ConversationService, "get_conversation", get_conversation) - monkeypatch.setattr(message_controller, "db", SimpleNamespace(session=session)) + monkeypatch.setattr(message_controller, "db", SimpleNamespace(session=SessionProxy())) monkeypatch.setattr(message_controller, "attach_message_extra_contents", lambda messages: None) monkeypatch.setattr(message_controller, "MessageInfiniteScrollPaginationResponse", FakeMessagePaginationResponse) @@ -1905,6 +1974,7 @@ def test_list_agent_chat_messages_uses_current_user_conversation( result = message_controller._list_chat_messages(app_model=app_model, current_user=current_user) assert result == {"data": [message_id], "limit": 20, "has_more": False} + assert captured.pop("session") is session assert captured == {"app_model": app_model, "conversation_id": conversation_id, "user": current_user} diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py index fdb534fc092..0ab8814f368 100644 --- a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py +++ b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py @@ -5,11 +5,11 @@ from types import SimpleNamespace import pytest from dify_agent.client import DifyAgentClientError, DifyAgentHTTPError, DifyAgentTimeoutError -from dify_agent.protocol import SandboxListResponse, SandboxReadResponse, SandboxUploadResponse +from dify_agent.protocol import SandboxListResponse, SandboxReadResponse from controllers.console import agent_app_sandbox as module from models.model import App, AppMode, IconType -from services.agent_app_sandbox_service import AgentSandboxInfo, AgentSandboxInspectorError +from services.agent_app_sandbox_service import AgentSandboxInfo, AgentSandboxInspectorError, AgentSandboxUploadDownload class _AgentAppService: @@ -28,11 +28,11 @@ class _AgentAppService: self.calls.append(("read", tenant_id, app_id, conversation_id, path)) return SandboxReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") - def upload_file(self, *, tenant_id: str, app_id: str, conversation_id: str, path: str) -> SandboxUploadResponse: + def upload_file( + self, *, tenant_id: str, app_id: str, conversation_id: str, path: str + ) -> AgentSandboxUploadDownload: self.calls.append(("upload", tenant_id, app_id, conversation_id, path)) - return SandboxUploadResponse( - path=path, file={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"} - ) + return AgentSandboxUploadDownload(url="https://files.example/report.txt") class _WorkflowService: @@ -48,6 +48,7 @@ class _WorkflowService: node_id: str, node_execution_id: str | None, path: str, + session, ) -> SandboxListResponse: self.calls.append(("list", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) return SandboxListResponse(path=path, entries=[], truncated=False) @@ -61,6 +62,7 @@ class _WorkflowService: node_id: str, node_execution_id: str | None, path: str, + session, ) -> SandboxReadResponse: self.calls.append(("read", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) return SandboxReadResponse(path=path, size=5, truncated=False, binary=False, text="hello") @@ -74,11 +76,10 @@ class _WorkflowService: node_id: str, node_execution_id: str | None, path: str, - ) -> SandboxUploadResponse: + session, + ) -> AgentSandboxUploadDownload: self.calls.append(("upload", tenant_id, app_id, workflow_run_id, node_id, node_execution_id, path)) - return SandboxUploadResponse( - path=path, file={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"} - ) + return AgentSandboxUploadDownload(url="https://files.example/upload.txt") def _app_model(app_id: str = "app-1") -> App: @@ -143,7 +144,7 @@ def test_agent_app_sandbox_resources_proxy_service(monkeypatch: pytest.MonkeyPat assert info == {"session_id": "abc1234", "workspace_cwd": "~/workspace/abc1234"} assert listing["path"] == "sub/report.txt" assert preview["text"] == "hello" - assert upload["file"]["reference"] == "dify-file-ref:file-1" + assert upload == {"url": "https://files.example/report.txt"} assert service.calls == [ ("info", "tenant-1", "app-1", "conv-1", ""), ("list", "tenant-1", "app-1", "conv-1", "sub/report.txt"), @@ -203,7 +204,7 @@ def test_workflow_agent_sandbox_resources_proxy_service(monkeypatch: pytest.Monk assert listing["path"] == "out.txt" assert preview["text"] == "hello" - assert upload["file"]["reference"] == "dify-file-ref:file-1" + assert upload == {"url": "https://files.example/upload.txt"} assert service.calls == [ ("list", "tenant-1", "app-1", "run-1", "agent-node", "exec-1", "out.txt"), ("read", "tenant-1", "app-1", "run-1", "agent-node", "exec-1", "out.txt"), diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py b/api/tests/unit_tests/controllers/console/app/test_annotation_api.py index 8a6094b94b8..cc95f7f8a94 100644 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_annotation_api.py @@ -2,7 +2,7 @@ from __future__ import annotations from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import ANY, Mock, patch import pytest from flask import Flask @@ -142,7 +142,7 @@ class TestConsoleAnnotationRefBoundaries: assert response == "" assert status == 204 - delete_mock.assert_called_once_with(AppRef("tenant-1", "app-1"), ["ann-1", "ann-2"]) + delete_mock.assert_called_once_with(AppRef("tenant-1", "app-1"), ["ann-1", "ann-2"], session=ANY) def test_update_uses_annotation_ref(self, app: Flask): api = annotation_module.AnnotationUpdateDeleteApi() @@ -216,4 +216,4 @@ class TestConsoleAnnotationRefBoundaries: response = handler(api, "app-1", "ann-1") assert response["total"] == 1 - hit_history_mock.assert_called_once_with(AnnotationRef("tenant-1", "app-1", "ann-1"), 2, 5) + hit_history_mock.assert_called_once_with(AnnotationRef("tenant-1", "app-1", "ann-1"), 2, 5, session=ANY) diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py b/api/tests/unit_tests/controllers/console/app/test_annotation_security.py index bfa4048191f..6a22d8769bc 100644 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py +++ b/api/tests/unit_tests/controllers/console/app/test_annotation_security.py @@ -193,9 +193,7 @@ class TestAnnotationImportServiceValidation: @pytest.fixture def mock_db_session(self): - """Mock database session.""" - with patch("services.annotation_service.db.session") as mock: - yield mock + return MagicMock() def test_max_records_limit_enforced(self, mock_app, mock_db_session): """Test that files with too many records are rejected.""" @@ -214,7 +212,7 @@ class TestAnnotationImportServiceValidation: with patch("services.annotation_service.FeatureService") as mock_features: mock_features.get_features.return_value.billing.enabled = False - result = AppAnnotationService.batch_import_app_annotations("app_id", file) + result = AppAnnotationService.batch_import_app_annotations("app_id", file, session=mock_db_session) # Should return error about too many records assert "error_msg" in result @@ -231,7 +229,7 @@ class TestAnnotationImportServiceValidation: with patch("services.annotation_service.current_account_with_tenant") as mock_auth: mock_auth.return_value = (MagicMock(id="user_id"), "tenant_id") - result = AppAnnotationService.batch_import_app_annotations("app_id", file) + result = AppAnnotationService.batch_import_app_annotations("app_id", file, session=mock_db_session) # Should return error about insufficient records assert "error_msg" in result @@ -250,7 +248,7 @@ class TestAnnotationImportServiceValidation: ): mock_auth.return_value = (MagicMock(id="user_id"), "tenant_id") - result = AppAnnotationService.batch_import_app_annotations("app_id", file) + result = AppAnnotationService.batch_import_app_annotations("app_id", file, session=mock_db_session) assert "error_msg" in result assert "malformed" in result["error_msg"].lower() @@ -271,7 +269,9 @@ class TestAnnotationImportServiceValidation: with patch("services.annotation_service.batch_import_annotations_task") as mock_task: with patch("services.annotation_service.redis_client"): - result = AppAnnotationService.batch_import_app_annotations("app_id", file) + result = AppAnnotationService.batch_import_app_annotations( + "app_id", file, session=mock_db_session + ) # Should return success response assert "job_id" in result diff --git a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py b/api/tests/unit_tests/controllers/console/app/test_app_response_models.py index 15c93a053bb..c6aa6178304 100644 --- a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py +++ b/api/tests/unit_tests/controllers/console/app/test_app_response_models.py @@ -6,7 +6,7 @@ from datetime import datetime from importlib import util from pathlib import Path from types import ModuleType, SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock import pytest from flask import Flask @@ -500,7 +500,8 @@ def test_app_list_uses_injected_session_for_draft_workflows( ) session = MagicMock() session.execute.return_value.scalars.return_value.all.return_value = [workflow] - scoped_session = SimpleNamespace(execute=MagicMock(side_effect=AssertionError("db.session should not be used"))) + scoped_session = MagicMock() + scoped_session.execute.side_effect = AssertionError("db.session should not be used") monkeypatch.setattr( app_module, @@ -515,7 +516,7 @@ def test_app_list_uses_injected_session_for_draft_workflows( monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.MyPermissions, "get", - lambda tenant_id, account_id: app_module.enterprise_rbac_service.MyPermissionsResponse( + lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse( app=app_module.enterprise_rbac_service.ResourcePermissionSnapshot( overrides=[ app_module.enterprise_rbac_service.ResourcePermissionKeys( @@ -563,12 +564,12 @@ def test_app_create_api_attaches_permission_keys(app, app_module): monkeypatch.setattr( app_module, "AppService", - lambda: SimpleNamespace(create_app=lambda tenant_id, params, user: app_obj), + lambda: SimpleNamespace(create_app=lambda tenant_id, params, user, session: app_obj), ) monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.AppPermissions, "batch_get", - lambda tenant_id, account_id, app_ids: {"app-new": ["app.acl.view_layout", "app.acl.edit"]}, + lambda tenant_id, account_id, app_ids, session: {"app-new": ["app.acl.view_layout", "app.acl.edit"]}, ) initialize_rbac = MagicMock() monkeypatch.setattr(app_module, "_initialize_created_app_rbac_access", initialize_rbac) @@ -671,7 +672,7 @@ def test_app_list_api_attaches_permission_keys(app, app_module): monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.MyPermissions, "get", - lambda tenant_id, account_id: app_module.enterprise_rbac_service.MyPermissionsResponse( + lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse( app=app_module.enterprise_rbac_service.ResourcePermissionSnapshot( default_permission_keys=["app.preview", "app.acl.view_layout"], overrides=[ @@ -715,7 +716,7 @@ def test_app_list_api_limits_to_apps_created_by_current_user_without_view_permis monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.MyPermissions, "get", - lambda tenant_id, account_id: app_module.enterprise_rbac_service.MyPermissionsResponse( + lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse( workspace=app_module.enterprise_rbac_service.WorkspacePermissionSnapshot( permission_keys=["app.create_and_management"] ) @@ -758,7 +759,7 @@ def test_app_list_api_limits_to_preview_overrides_without_manage_own_permission( monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.MyPermissions, "get", - lambda tenant_id, account_id: app_module.enterprise_rbac_service.MyPermissionsResponse( + lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse( app=app_module.enterprise_rbac_service.ResourcePermissionSnapshot( overrides=[ app_module.enterprise_rbac_service.ResourcePermissionKeys( @@ -814,7 +815,7 @@ def test_app_list_api_returns_no_apps_without_workspace_or_resource_view_permiss monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.MyPermissions, "get", - lambda tenant_id, account_id: app_module.enterprise_rbac_service.MyPermissionsResponse(), + lambda tenant_id, account_id, session: app_module.enterprise_rbac_service.MyPermissionsResponse(), ) monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.AppAccess, @@ -880,7 +881,7 @@ def test_app_detail_api_attaches_current_user_permission_keys(app, app_module): resp = method(app_module.AppApi(), "tenant-1", SimpleNamespace(id="acct-1"), app_model=app_obj) - get_permissions.assert_called_once_with("tenant-1", "acct-1", app_id="app-1") + get_permissions.assert_called_once_with("tenant-1", "acct-1", app_id="app-1", session=ANY) assert resp["permission_keys"] == ["app.acl.view_layout", "app.acl.edit", "app.acl.monitor"] @@ -921,7 +922,7 @@ def test_app_copy_api_attaches_permission_keys(app, app_module): "get_system_features", lambda: SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False)), ) - monkeypatch.setattr(app_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(app_module, "db", SimpleNamespace(engine=object(), session=lambda: MagicMock())) monkeypatch.setattr( app_module, "Session", @@ -930,7 +931,7 @@ def test_app_copy_api_attaches_permission_keys(app, app_module): monkeypatch.setattr( app_module.enterprise_rbac_service.RBACService.AppPermissions, "batch_get", - lambda tenant_id, account_id, app_ids: {"app-new": ["app.acl.view_layout", "app.acl.edit"]}, + lambda tenant_id, account_id, app_ids, session: {"app-new": ["app.acl.view_layout", "app.acl.edit"]}, ) resp, status = method( diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow.py b/api/tests/unit_tests/controllers/console/app/test_workflow.py index 2f971eaf74f..2d811deb916 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -621,7 +621,7 @@ def test_workflow_online_users_filters_inaccessible_workflow(app: Flask, monkeyp monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id: {app_id_1}), + lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id, session: {app_id_1}), ) monkeypatch.setattr(workflow_module.file_helpers, "get_signed_file_url", sign_avatar) @@ -703,7 +703,7 @@ def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pyte monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id: set(app_ids)), + lambda: SimpleNamespace(get_accessible_app_ids=lambda app_ids, tenant_id, session: set(app_ids)), ) first_pipeline = Mock() diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py index f04ab6d6e7c..956706eafb6 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock import pytest from flask import Flask @@ -94,6 +94,7 @@ def test_human_input_preview_delegates_to_service( account=account, node_id="node-42", inputs={"topic": "tech"}, + session=ANY, ) @@ -144,6 +145,7 @@ def test_human_input_submit_forwards_payload(app: Flask, monkeypatch: pytest.Mon form_inputs={"answer": "42"}, inputs={"#node-1.result#": "LLM output"}, action="approve", + session=ANY, ) @@ -193,6 +195,7 @@ def test_human_input_delivery_test_calls_service( node_id="node-7", delivery_method_id="delivery-123", inputs={}, + session=ANY, ) diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py b/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py index e66ae5246bc..dfe35a89f57 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py @@ -25,7 +25,7 @@ from __future__ import annotations import json from collections.abc import Iterator from typing import Any -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock from uuid import UUID import pytest @@ -382,7 +382,9 @@ def test_serve_snapshot_happy_path(patch_service, app_model, run_id): result = ctrl._serve_snapshot(app_model, run_id) assert isinstance(result, dict) assert result["workflow_run_id"] == "00000000-0000-0000-0000-0000000000aa" - patch_service.snapshot_workflow_run.assert_called_once_with(app_model=app_model, workflow_run_id=str(run_id)) + patch_service.snapshot_workflow_run.assert_called_once_with( + app_model=app_model, workflow_run_id=str(run_id), session=ANY + ) def test_serve_snapshot_translates_inspector_error_to_404(patch_service, app_model, run_id): @@ -399,7 +401,7 @@ def test_serve_node_detail_happy_path(patch_service, app_model, run_id): result = ctrl._serve_node_detail(app_model, run_id, "agent-1") assert result["node_id"] == "agent-1" patch_service.node_detail.assert_called_once_with( - app_model=app_model, workflow_run_id=str(run_id), node_id="agent-1" + app_model=app_model, workflow_run_id=str(run_id), node_id="agent-1", session=ANY ) @@ -431,6 +433,7 @@ def test_serve_output_preview_happy_path(patch_service, app_model, run_id): workflow_run_id=str(run_id), node_id="agent-1", output_name="text", + session=ANY, ) diff --git a/api/tests/unit_tests/controllers/console/auth/test_account_activation.py b/api/tests/unit_tests/controllers/console/auth/test_account_activation.py index ebae7de6c15..001ca0bf8fb 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_account_activation.py +++ b/api/tests/unit_tests/controllers/console/auth/test_account_activation.py @@ -597,7 +597,7 @@ class TestActivateApi: assert response["result"] == "success" mock_create_tenant_member.assert_called_once_with( - mock_invitation["tenant"], mock_account, mock_db.session, role=TenantAccountRole.ADMIN + mock_invitation["tenant"], mock_account, mock_db.session(), role=TenantAccountRole.ADMIN ) mock_switch_tenant.assert_called_once_with(mock_account, mock_invitation["tenant"].id, session=ANY) mock_revoke_token.assert_called_once_with("workspace-123", "invitee@example.com", "valid_token") diff --git a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py index 21d1932f820..b231826aeac 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py @@ -43,7 +43,7 @@ def test_list_data_source_auth_uses_injected_tenant_id() -> None: ): result = method(api, "tenant-1") - get_provider_auth_list.assert_called_once_with(ANY, "tenant-1") + get_provider_auth_list.assert_called_once_with("tenant-1", session=ANY) assert result["sources"][0]["id"] == "binding-1" assert result["sources"][0]["provider"] == "custom" @@ -65,7 +65,7 @@ def test_create_data_source_auth_binding_uses_injected_tenant_id() -> None: ): result, status = method(api, "tenant-1") - create_auth.assert_called_once_with(ANY, "tenant-1", payload) + create_auth.assert_called_once_with("tenant-1", payload, session=ANY) assert result == {"result": "success"} assert status == 200 @@ -82,6 +82,6 @@ def test_delete_data_source_auth_binding_uses_injected_tenant_id() -> None: ): result, status = method(api, "tenant-1", "binding-1") - delete_provider_auth.assert_called_once_with(ANY, "tenant-1", "binding-1") + delete_provider_auth.assert_called_once_with("tenant-1", "binding-1", session=ANY) assert result == "" assert status == 204 diff --git a/api/tests/unit_tests/controllers/console/auth/test_token_refresh.py b/api/tests/unit_tests/controllers/console/auth/test_token_refresh.py index 34fff57b0ad..8effbb96887 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_token_refresh.py +++ b/api/tests/unit_tests/controllers/console/auth/test_token_refresh.py @@ -13,8 +13,10 @@ from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask from flask_restx import Api +from werkzeug.exceptions import Unauthorized from controllers.console.auth.login import RefreshTokenApi +from services.errors.account import RefreshTokenAccountNotFoundError, RefreshTokenNotFoundError class TestRefreshTokenApi: @@ -98,18 +100,19 @@ class TestRefreshTokenApi: @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) - def test_refresh_fails_with_invalid_token(self, mock_refresh_token, mock_extract_token, app: Flask): + def test_refresh_returns_unauthorized_for_invalid_refresh_token( + self, mock_refresh_token, mock_extract_token, app: Flask + ): """ - Test token refresh failure with invalid refresh token. + Test token refresh maps invalid refresh tokens to unauthorized responses. Verifies that: - - Exception is caught when token is invalid - - 401 status code is returned - - Error message is included in response + - Invalid refresh token validation failures return 401 + - The failure response preserves the validation message """ # Arrange mock_extract_token.return_value = "invalid_refresh_token" - mock_refresh_token.side_effect = Exception("Invalid refresh token") + mock_refresh_token.side_effect = RefreshTokenNotFoundError("Invalid refresh token") # Act with app.test_request_context("/refresh-token", method="POST"): @@ -119,22 +122,21 @@ class TestRefreshTokenApi: # Assert assert status_code == 401 assert response["result"] == "fail" - assert "Invalid refresh token" in response["message"] + assert response["message"] == "Invalid refresh token" @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) - def test_refresh_fails_with_expired_token(self, mock_refresh_token, mock_extract_token, app: Flask): + def test_refresh_returns_unauthorized_for_invalid_account(self, mock_refresh_token, mock_extract_token, app: Flask): """ - Test token refresh failure with expired refresh token. + Test token refresh maps missing accounts to unauthorized responses. Verifies that: - - Expired tokens are rejected - - 401 status code is returned - - Appropriate error handling + - Invalid account validation failures return 401 + - The failure response preserves the validation message """ # Arrange - mock_extract_token.return_value = "expired_refresh_token" - mock_refresh_token.side_effect = Exception("Refresh token expired") + mock_extract_token.return_value = "refresh_token_for_missing_account" + mock_refresh_token.side_effect = RefreshTokenAccountNotFoundError("Invalid account") # Act with app.test_request_context("/refresh-token", method="POST"): @@ -144,7 +146,71 @@ class TestRefreshTokenApi: # Assert assert status_code == 401 assert response["result"] == "fail" - assert "expired" in response["message"].lower() + assert response["message"] == "Invalid account" + + @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) + @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) + def test_refresh_returns_unauthorized_for_banned_account(self, mock_refresh_token, mock_extract_token, app: Flask): + """ + Test token refresh maps banned accounts to unauthorized responses. + + Verifies that: + - Authorization failures raised during account loading return 401 + - The failure response preserves the authorization message + """ + # Arrange + mock_extract_token.return_value = "refresh_token_for_banned_account" + mock_refresh_token.side_effect = Unauthorized("Account is banned.") + + # Act + with app.test_request_context("/refresh-token", method="POST"): + refresh_api = RefreshTokenApi() + response, status_code = refresh_api.post() + + # Assert + assert status_code == 401 + assert response["result"] == "fail" + assert response["message"] == "Account is banned." + + @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) + @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) + def test_refresh_propagates_non_whitelisted_value_error(self, mock_refresh_token, mock_extract_token, app: Flask): + """ + Test token refresh preserves non-whitelisted ValueError failures. + + Verifies that: + - Only known refresh-token validation errors are mapped to 401 + - Unexpected ValueError instances continue to propagate + """ + # Arrange + mock_extract_token.return_value = "valid_refresh_token" + mock_refresh_token.side_effect = ValueError("unexpected parse failure") + + # Act & Assert + with app.test_request_context("/refresh-token", method="POST"): + refresh_api = RefreshTokenApi() + with pytest.raises(ValueError, match="unexpected parse failure"): + refresh_api.post() + + @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) + @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) + def test_refresh_propagates_unexpected_service_errors(self, mock_refresh_token, mock_extract_token, app: Flask): + """ + Test token refresh preserves unexpected service failures. + + Verifies that: + - Operational errors are not misreported as authentication failures + - The original exception is preserved for higher-level error handling + """ + # Arrange + mock_extract_token.return_value = "valid_refresh_token" + mock_refresh_token.side_effect = RuntimeError("redis unavailable") + + # Act & Assert + with app.test_request_context("/refresh-token", method="POST"): + refresh_api = RefreshTokenApi() + with pytest.raises(RuntimeError, match="redis unavailable"): + refresh_api.post() @patch("controllers.console.auth.login.extract_refresh_token", autospec=True) @patch("controllers.console.auth.login.AccountService.refresh_token", autospec=True) diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py index e8faece89ca..8f66ca5c993 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py @@ -1,5 +1,6 @@ import inspect -from unittest.mock import MagicMock, patch +from datetime import UTC, datetime +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask @@ -23,6 +24,76 @@ from graphon.model_runtime.errors.validate import CredentialsValidateFailedError from services.datasource_provider_service import DatasourceProviderService from services.plugin.oauth_service import OAuthProxyService +_PROVIDER_ID = "langgenius/notion_datasource/notion" + + +def _i18n(text: str) -> dict[str, str]: + return {"en_US": text, "zh_Hans": text, "pt_BR": text, "ja_JP": text} + + +def _provider_config(name: str, type_: str, label: str, *, required: bool = True) -> dict: + return { + "type": type_, + "name": name, + "scope": None, + "required": required, + "default": None, + "options": None, + "multiple": False, + "label": _i18n(label), + "help": None, + "url": None, + "placeholder": None, + } + + +def _datasource_credential(credential_id: str = "cred-1", *, is_default: bool = True) -> dict: + return { + "credential": { + "api_key": "******", + "workspace": "engineering", + "database_id": "db-123", + }, + "type": "api-key", + "name": "API Key", + "avatar_url": "https://cdn.example.com/notion.png", + "id": credential_id, + "is_default": is_default, + } + + +def _datasource_auth() -> dict: + return { + "author": "Dify", + "provider": "notion", + "plugin_id": "langgenius/notion_datasource", + "plugin_unique_identifier": "langgenius/notion_datasource:0.0.1", + "icon": "icon.svg", + "name": "notion", + "label": _i18n("Notion"), + "description": _i18n("Notion datasource"), + "credential_schema": [ + _provider_config("api_key", "secret-input", "API key"), + ], + "oauth_schema": { + "client_schema": [ + _provider_config("client_id", "text-input", "Client ID"), + ], + "credentials_schema": [ + _provider_config("access_token", "secret-input", "Access token"), + ], + "oauth_custom_client_params": {"client_id": "masked-client", "client_secret": "********"}, + "is_oauth_custom_client_enabled": True, + "is_system_oauth_params_exists": True, + "redirect_uri": "https://api.example.com/oauth/callback", + }, + "credentials_list": [_datasource_credential(), _datasource_credential("cred-2", is_default=False)], + } + + +def _success_response() -> dict[str, str]: + return {"result": "success"} + class TestDatasourcePluginOAuthAuthorizationUrl: def test_get_success(self, app: Flask): @@ -30,28 +101,50 @@ class TestDatasourcePluginOAuthAuthorizationUrl: method = inspect.unwrap(api.get) user = MagicMock(id="user-1") + oauth_client = {"client_id": "abc", "client_secret": "shh", "scopes": ["read", "write"]} + auth_url_payload = { + "authorization_url": "https://auth.example.com/oauth?client_id=abc&state=xyz", + } with ( app.test_request_context("/?credential_id=cred-1"), patch.object( DatasourceProviderService, "get_oauth_client", - return_value={"client_id": "abc"}, - ), + return_value=oauth_client, + ) as get_oauth_client, patch.object( OAuthProxyService, "create_proxy_context", return_value="ctx-1", - ), + ) as create_proxy_context, patch.object( OAuthHandler, "get_authorization_url", - return_value={"url": "http://auth"}, - ), + return_value=auth_url_payload, + ) as get_authorization_url, ): - response = method(api, "tenant-1", user, "notion") + response = method(api, "tenant-1", user, _PROVIDER_ID) assert response.status_code == 200 + assert response.get_json() == auth_url_payload + assert "context_id=ctx-1" in response.headers.get("Set-Cookie") + provider_id = get_oauth_client.call_args.kwargs["datasource_provider_id"] + assert str(provider_id) == _PROVIDER_ID + get_oauth_client.assert_called_once() + create_proxy_context.assert_called_once_with( + user_id="user-1", + tenant_id="tenant-1", + plugin_id="langgenius/notion_datasource", + provider="notion", + credential_id="cred-1", + ) + get_authorization_url.assert_called_once() + assert get_authorization_url.call_args.kwargs["tenant_id"] == "tenant-1" + assert get_authorization_url.call_args.kwargs["user_id"] == "user-1" + assert get_authorization_url.call_args.kwargs["plugin_id"] == "langgenius/notion_datasource" + assert get_authorization_url.call_args.kwargs["provider"] == "notion" + assert get_authorization_url.call_args.kwargs["system_credentials"] == oauth_client def test_get_no_oauth_config(self, app: Flask): api = DatasourcePluginOAuthAuthorizationUrl() @@ -90,10 +183,10 @@ class TestDatasourcePluginOAuthAuthorizationUrl: patch.object( OAuthHandler, "get_authorization_url", - return_value={"url": "http://auth"}, + return_value={"authorization_url": "http://auth"}, ), ): - response = method(api, "tenant-1", user, "notion") + response = method(api, "tenant-1", user, _PROVIDER_ID) assert response.status_code == 200 assert "context_id" in response.headers.get("Set-Cookie") @@ -106,8 +199,9 @@ class TestDatasourceOAuthCallback: oauth_response = MagicMock() oauth_response.credentials = {"token": "abc"} - oauth_response.expires_at = None - oauth_response.metadata = {"name": "test"} + expires_at = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) + oauth_response.expires_at = expires_at + oauth_response.metadata = {"name": "Workspace Bot", "avatar_url": "https://avatar.example.com/bot.png"} context = { "user_id": "user-1", @@ -125,7 +219,7 @@ class TestDatasourceOAuthCallback: patch.object( DatasourceProviderService, "get_oauth_client", - return_value={"client_id": "abc"}, + return_value={"client_id": "abc", "client_secret": "secret"}, ), patch.object( OAuthHandler, @@ -136,11 +230,22 @@ class TestDatasourceOAuthCallback: DatasourceProviderService, "add_datasource_oauth_provider", return_value=None, - ), + ) as add_oauth_provider, ): - response = method(api, "notion") + response = method(api, _PROVIDER_ID) assert response.status_code == 302 + assert "/oauth-callback" in response.location + add_oauth_provider.assert_called_once() + assert add_oauth_provider.call_args.kwargs == { + "tenant_id": "tenant-1", + "provider_id": add_oauth_provider.call_args.kwargs["provider_id"], + "avatar_url": "https://avatar.example.com/bot.png", + "name": "Workspace Bot", + "expire_at": expires_at, + "credentials": {"token": "abc"}, + } + assert str(add_oauth_provider.call_args.kwargs["provider_id"]) == _PROVIDER_ID def test_callback_missing_context(self, app: Flask): api = DatasourceOAuthCallback() @@ -223,12 +328,16 @@ class TestDatasourceOAuthCallback: DatasourceProviderService, "reauthorize_datasource_oauth_provider", return_value=None, - ), + ) as reauthorize_provider, ): - response = method(api, "notion") + response = method(api, _PROVIDER_ID) assert response.status_code == 302 assert "/oauth-callback" in response.location + reauthorize_provider.assert_called_once() + assert str(reauthorize_provider.call_args.kwargs["provider_id"]) == _PROVIDER_ID + assert reauthorize_provider.call_args.kwargs["credential_id"] == "cred-1" + assert reauthorize_provider.call_args.kwargs["credentials"] == {"token": "abc"} def test_callback_context_id_from_cookie(self, app: Flask): api = DatasourceOAuthCallback() @@ -278,7 +387,14 @@ class TestDatasourceAuth: api = DatasourceAuth() method = inspect.unwrap(api.post) - payload = {"credentials": {"key": "val"}} + payload = { + "name": "Engineering Notion", + "credentials": { + "api_key": "secret-token", + "workspace": "engineering", + "database_id": "db-123", + }, + } with ( app.test_request_context("/", json=payload), @@ -287,11 +403,17 @@ class TestDatasourceAuth: DatasourceProviderService, "add_datasource_api_key_provider", return_value=None, - ), + ) as add_api_key_provider, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + add_api_key_provider.assert_called_once() + assert add_api_key_provider.call_args.kwargs["tenant_id"] == "tenant-1" + assert str(add_api_key_provider.call_args.kwargs["provider_id"]) == _PROVIDER_ID + assert add_api_key_provider.call_args.kwargs["credentials"] == payload["credentials"] + assert add_api_key_provider.call_args.kwargs["name"] == "Engineering Notion" def test_post_invalid_credentials(self, app: Flask): api = DatasourceAuth() @@ -321,19 +443,19 @@ class TestDatasourceAuth: patch.object( DatasourceProviderService, "list_datasource_credentials", - return_value=[{"id": "1"}], + return_value=[_datasource_credential()], ), ): - response, status = method(api, "tenant-1", user, "notion") + response, status = method(api, "tenant-1", user, _PROVIDER_ID) assert status == 200 - assert response["result"] + assert response == {"result": [_datasource_credential()]} def test_post_missing_credentials(self, app: Flask): api = DatasourceAuth() method = inspect.unwrap(api.post) - payload = {} + payload: dict[str, object] = {} with ( app.test_request_context("/", json=payload), @@ -375,17 +497,25 @@ class TestDatasourceAuthDeleteApi: DatasourceProviderService, "remove_datasource_credentials", return_value=None, - ), + ) as remove_datasource_credentials, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + remove_datasource_credentials.assert_called_once_with( + tenant_id="tenant-1", + auth_id="cred-1", + provider="notion", + plugin_id="langgenius/notion_datasource", + session=ANY, + ) def test_delete_missing_credential_id(self, app: Flask): api = DatasourceAuthDeleteApi() method = inspect.unwrap(api.post) - payload = {} + payload: dict[str, object] = {} with ( app.test_request_context("/", json=payload), @@ -400,7 +530,11 @@ class TestDatasourceAuthUpdateApi: api = DatasourceAuthUpdateApi() method = inspect.unwrap(api.post) - payload = {"credential_id": "id", "credentials": {"k": "v"}} + payload = { + "credential_id": "cred-1", + "name": "Updated Notion", + "credentials": {"api_key": "new-secret", "database_id": "db-456"}, + } with ( app.test_request_context("/", json=payload), @@ -409,11 +543,20 @@ class TestDatasourceAuthUpdateApi: DatasourceProviderService, "update_datasource_credentials", return_value=None, - ), + ) as update_datasource_credentials, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 201 + update_datasource_credentials.assert_called_once_with( + tenant_id="tenant-1", + auth_id="cred-1", + provider="notion", + plugin_id="langgenius/notion_datasource", + credentials=payload["credentials"], + name="Updated Notion", + ) def test_update_with_credentials_none(self, app: Flask): api = DatasourceAuthUpdateApi() @@ -432,7 +575,9 @@ class TestDatasourceAuthUpdateApi: ): response, status = method(api, "tenant-1", "notion") + assert response == _success_response() update_mock.assert_called_once() + assert update_mock.call_args.kwargs["credentials"] == {} assert status == 201 def test_update_name_only(self, app: Flask): @@ -450,8 +595,9 @@ class TestDatasourceAuthUpdateApi: return_value=None, ), ): - _, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", "notion") + assert response == _success_response() assert status == 201 def test_update_with_empty_credentials_dict(self, app: Flask): @@ -469,8 +615,9 @@ class TestDatasourceAuthUpdateApi: return_value=None, ) as update_mock, ): - _, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", "notion") + assert response == _success_response() update_mock.assert_called_once() assert status == 201 @@ -485,12 +632,13 @@ class TestDatasourceAuthListApi: patch.object( DatasourceProviderService, "get_all_datasource_credentials", - return_value=[{"id": "1"}], + return_value=[_datasource_auth()], ), ): response, status = method(api, "tenant-1") assert status == 200 + assert response == {"result": [_datasource_auth()]} def test_auth_list_empty(self, app: Flask): api = DatasourceAuthListApi() @@ -537,7 +685,7 @@ class TestDatasourceHardCodeAuthListApi: patch.object( DatasourceProviderService, "get_hard_code_datasource_credentials", - return_value=[{"id": "1"}], + return_value=[_datasource_auth()], ), ): response, status = method(api, "tenant-1") @@ -550,7 +698,14 @@ class TestDatasourceAuthOauthCustomClient: api = DatasourceAuthOauthCustomClient() method = inspect.unwrap(api.post) - payload = {"client_params": {}, "enable_oauth_custom_client": True} + payload = { + "client_params": { + "client_id": "custom-client", + "client_secret": "custom-secret", + "authorize_url": "https://auth.example.com/authorize", + }, + "enable_oauth_custom_client": True, + } with ( app.test_request_context("/", json=payload), @@ -559,11 +714,17 @@ class TestDatasourceAuthOauthCustomClient: DatasourceProviderService, "setup_oauth_custom_client_params", return_value=None, - ), + ) as setup_custom_client, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + setup_custom_client.assert_called_once() + assert setup_custom_client.call_args.kwargs["tenant_id"] == "tenant-1" + assert str(setup_custom_client.call_args.kwargs["datasource_provider_id"]) == _PROVIDER_ID + assert setup_custom_client.call_args.kwargs["client_params"] == payload["client_params"] + assert setup_custom_client.call_args.kwargs["enabled"] is True def test_delete_success(self, app: Flask): api = DatasourceAuthOauthCustomClient() @@ -575,17 +736,20 @@ class TestDatasourceAuthOauthCustomClient: DatasourceProviderService, "remove_oauth_custom_client_params", return_value=None, - ), + ) as remove_custom_client, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + remove_custom_client.assert_called_once() + assert str(remove_custom_client.call_args.kwargs["datasource_provider_id"]) == _PROVIDER_ID def test_post_empty_payload(self, app: Flask): api = DatasourceAuthOauthCustomClient() method = inspect.unwrap(api.post) - payload = {} + payload: dict[str, object] = {} with ( app.test_request_context("/", json=payload), @@ -596,8 +760,9 @@ class TestDatasourceAuthOauthCustomClient: return_value=None, ), ): - _, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", "notion") + assert response == _success_response() assert status == 200 def test_post_disabled_flag(self, app: Flask): @@ -618,9 +783,12 @@ class TestDatasourceAuthOauthCustomClient: return_value=None, ) as setup_mock, ): - _, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", "notion") + assert response == _success_response() setup_mock.assert_called_once() + assert setup_mock.call_args.kwargs["client_params"] == {"a": 1} + assert setup_mock.call_args.kwargs["enabled"] is False assert status == 200 @@ -638,17 +806,22 @@ class TestDatasourceAuthDefaultApi: DatasourceProviderService, "set_default_datasource_provider", return_value=None, - ), + ) as set_default_datasource_provider, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + set_default_datasource_provider.assert_called_once() + assert set_default_datasource_provider.call_args.kwargs["tenant_id"] == "tenant-1" + assert str(set_default_datasource_provider.call_args.kwargs["datasource_provider_id"]) == _PROVIDER_ID + assert set_default_datasource_provider.call_args.kwargs["credential_id"] == "cred-1" def test_default_missing_id(self, app: Flask): api = DatasourceAuthDefaultApi() method = inspect.unwrap(api.post) - payload = {} + payload: dict[str, object] = {} with ( app.test_request_context("/", json=payload), @@ -663,7 +836,7 @@ class TestDatasourceUpdateProviderNameApi: api = DatasourceUpdateProviderNameApi() method = inspect.unwrap(api.post) - payload = {"credential_id": "id", "name": "New Name"} + payload = {"credential_id": "cred-1", "name": "New Name"} with ( app.test_request_context("/", json=payload), @@ -672,11 +845,17 @@ class TestDatasourceUpdateProviderNameApi: DatasourceProviderService, "update_datasource_provider_name", return_value=None, - ), + ) as update_datasource_provider_name, ): - response, status = method(api, "tenant-1", "notion") + response, status = method(api, "tenant-1", _PROVIDER_ID) + assert response == _success_response() assert status == 200 + update_datasource_provider_name.assert_called_once() + assert update_datasource_provider_name.call_args.kwargs["tenant_id"] == "tenant-1" + assert str(update_datasource_provider_name.call_args.kwargs["datasource_provider_id"]) == _PROVIDER_ID + assert update_datasource_provider_name.call_args.kwargs["name"] == "New Name" + assert update_datasource_provider_name.call_args.kwargs["credential_id"] == "cred-1" def test_update_name_too_long(self, app: Flask): api = DatasourceUpdateProviderNameApi() diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py index 2a1970d3837..39a6fa65a06 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py @@ -64,10 +64,9 @@ class TestPipelineTemplateListApi: tenant_id = "tenant-1" service_calls: list[tuple[str, str, str]] = [] - def get_pipeline_templates( - session: Mock, template_type: str, language: str, current_tenant_id: str - ) -> dict[str, object]: - service_calls.append((template_type, language, current_tenant_id)) + def get_pipeline_templates(*, type: str, language: str, current_tenant_id: str, session) -> dict[str, object]: + del session + service_calls.append((type, language, current_tenant_id)) return {"pipeline_templates": [_template_item()]} with ( @@ -94,10 +93,9 @@ class TestPipelineTemplateListApi: tenant_id = "tenant-1" service_calls: list[tuple[str, str, str]] = [] - def get_pipeline_templates( - session: Mock, template_type: str, language: str, current_tenant_id: str - ) -> dict[str, object]: - service_calls.append((template_type, language, current_tenant_id)) + def get_pipeline_templates(*, type: str, language: str, current_tenant_id: str, session) -> dict[str, object]: + del session + service_calls.append((type, language, current_tenant_id)) return {"pipeline_templates": []} with ( @@ -117,16 +115,18 @@ class TestPipelineTemplateDetailApi: method = unwrap(api.get) service_calls: list[tuple[str, str]] = [] - class Service: - def get_pipeline_template_detail( - self, session: Mock, template_id: str, template_type: str - ) -> dict[str, object]: - service_calls.append((template_id, template_type)) - return _template_detail() + def get_pipeline_template_detail(template_id: str, type: str, *, session) -> dict[str, object]: + del session + service_calls.append((template_id, type)) + return _template_detail() with ( app.test_request_context("/rag/pipeline/templates/template-1?type=customized"), - patch.object(module, "RagPipelineService", Service), + patch.object( + module.RagPipelineService, + "get_pipeline_template_detail", + side_effect=get_pipeline_template_detail, + ), ): response, status = method(api, Mock(), "template-1") @@ -138,13 +138,16 @@ class TestPipelineTemplateDetailApi: api = PipelineTemplateDetailApi() method = unwrap(api.get) - class Service: - def get_pipeline_template_detail(self, session: Mock, template_id: str, template_type: str) -> None: - return None + def get_pipeline_template_detail(template_id: str, type: str, *, session) -> None: + del template_id, type, session with ( app.test_request_context("/rag/pipeline/templates/missing"), - patch.object(module, "RagPipelineService", Service), + patch.object( + module.RagPipelineService, + "get_pipeline_template_detail", + side_effect=get_pipeline_template_detail, + ), ): with pytest.raises(NotFound): method(api, Mock(), "missing") @@ -160,8 +163,14 @@ class TestCustomizedPipelineTemplateApi: service_calls: list[tuple[str, PipelineTemplateInfoEntity, Account, str]] = [] def update_template( - template_id: str, template_info: PipelineTemplateInfoEntity, current_user: Account, current_tenant_id: str + template_id: str, + template_info: PipelineTemplateInfoEntity, + current_user: Account, + current_tenant_id: str, + *, + session, ) -> None: + del session service_calls.append((template_id, template_info, current_user, current_tenant_id)) with ( @@ -198,8 +207,14 @@ class TestCustomizedPipelineTemplateApi: service_calls: list[tuple[str, PipelineTemplateInfoEntity, Account, str]] = [] def update_template( - template_id: str, template_info: PipelineTemplateInfoEntity, current_user: Account, current_tenant_id: str + template_id: str, + template_info: PipelineTemplateInfoEntity, + current_user: Account, + current_tenant_id: str, + *, + session, ) -> None: + del session service_calls.append((template_id, template_info, current_user, current_tenant_id)) with ( @@ -228,7 +243,8 @@ class TestCustomizedPipelineTemplateApi: tenant_id = "tenant-1" deleted_templates: list[tuple[str, str]] = [] - def delete_template(template_id: str, current_tenant_id: str) -> None: + def delete_template(template_id: str, current_tenant_id: str, *, session) -> None: + del session deleted_templates.append((template_id, current_tenant_id)) with ( @@ -325,9 +341,19 @@ class TestPublishCustomizedPipelineTemplateApi: service_calls: list[tuple[str, dict[str, object], Account, str]] = [] class Service: + def __init__(self, *args, **kwargs) -> None: + pass + def publish_customized_pipeline_template( - self, pipeline_id: str, data: dict[str, object], current_user: Account, current_tenant_id: str + self, + pipeline_id: str, + data: dict[str, object], + current_user: Account, + current_tenant_id: str, + *, + session, ) -> None: + del session service_calls.append((pipeline_id, data, current_user, current_tenant_id)) with ( @@ -352,9 +378,19 @@ class TestPublishCustomizedPipelineTemplateApi: service_calls: list[tuple[str, dict[str, object], Account, str]] = [] class Service: + def __init__(self, *args, **kwargs) -> None: + pass + def publish_customized_pipeline_template( - self, pipeline_id: str, data: dict[str, object], current_user: Account, current_tenant_id: str + self, + pipeline_id: str, + data: dict[str, object], + current_user: Account, + current_tenant_id: str, + *, + session, ) -> None: + del session service_calls.append((pipeline_id, data, current_user, current_tenant_id)) with ( diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py index 5cc5af9592b..e344a4c8bab 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py @@ -52,7 +52,9 @@ def _pipeline() -> Pipeline: def test_draft_rag_pipeline_workflow_get_serializes_response_model(monkeypatch: pytest.MonkeyPatch) -> None: workflow = _make_workflow() monkeypatch.setattr( - module, "RagPipelineService", lambda: SimpleNamespace(get_draft_workflow=lambda **_kwargs: workflow) + module, + "RagPipelineService", + lambda *_args, **_kwargs: SimpleNamespace(get_draft_workflow=lambda **_kwargs: workflow), ) api = module.DraftRagPipelineApi() @@ -97,12 +99,12 @@ def test_published_rag_pipeline_workflows_serialize_items_before_session_closes( assert session_state["open"] is True return getattr(base_workflow, name) - monkeypatch.setattr(module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(module, "db", SimpleNamespace(engine=object(), session=lambda: object())) monkeypatch.setattr(module, "sessionmaker", lambda *_args, **_kwargs: _SessionMaker()) monkeypatch.setattr( module, "RagPipelineService", - lambda: SimpleNamespace(get_all_published_workflow=lambda **_kwargs: ([_Workflow()], False)), + lambda *_args, **_kwargs: SimpleNamespace(get_all_published_workflow=lambda **_kwargs: ([_Workflow()], False)), ) with app.test_request_context( @@ -132,12 +134,12 @@ def test_rag_pipeline_workflow_patch_serializes_response_model(app: Flask, monke def begin(self): return _SessionContext() - monkeypatch.setattr(module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(module, "db", SimpleNamespace(engine=object(), session=lambda: object())) monkeypatch.setattr(module, "sessionmaker", lambda *_args, **_kwargs: _SessionMaker()) monkeypatch.setattr( module, "RagPipelineService", - lambda: SimpleNamespace(update_workflow=lambda **_kwargs: workflow), + lambda *_args, **_kwargs: SimpleNamespace(update_workflow=lambda **_kwargs: workflow), ) payload: dict[str, object] = {"marked_name": "Updated release"} @@ -158,3 +160,65 @@ def test_rag_pipeline_workflow_patch_serializes_response_model(app: Flask, monke assert response["id"] == "workflow-1" assert response["marked_name"] == "Updated release" assert response["hash"] == "hash-1" + + +def test_default_rag_pipeline_block_configs_serializes_root_response(monkeypatch: pytest.MonkeyPatch) -> None: + block_configs = [{"type": "start", "config": {"title": "Start"}}] + monkeypatch.setattr( + module, + "RagPipelineService", + lambda *_args, **_kwargs: SimpleNamespace(get_default_block_configs=lambda: block_configs), + ) + + api = module.DefaultRagPipelineBlockConfigsApi() + handler = unwrap_all(api.get) + + response = handler(api, _pipeline()) + + assert response == block_configs + + +def test_draft_rag_pipeline_second_step_parameters_serializes_variables(app, monkeypatch: pytest.MonkeyPatch) -> None: + variables = [ + { + "belong_to_node_id": "shared", + "type": "number", + "label": "Chunk size", + "variable": "chunk_size", + "default_value": 1024, + "required": True, + } + ] + monkeypatch.setattr( + module, + "RagPipelineService", + lambda *_args, **_kwargs: SimpleNamespace(get_second_step_parameters=lambda **_kwargs: variables), + ) + + api = module.DraftRagPipelineSecondStepApi() + handler = unwrap_all(api.get) + + with app.test_request_context("/?node_id=node-1"): + response = handler(api, _pipeline()) + + assert response["variables"] == variables + + +def test_rag_pipeline_recommended_plugins_serializes_known_envelope(app, monkeypatch: pytest.MonkeyPatch) -> None: + recommended_plugins = { + "installed_recommended_plugins": [{"name": "Dify Extractor", "meta": {"version": "1.0.0"}}], + "uninstalled_recommended_plugins": [{"plugin_id": "langgenius/notion_datasource"}], + } + monkeypatch.setattr( + module, + "RagPipelineService", + lambda *_args, **_kwargs: SimpleNamespace(get_recommended_plugins=lambda *_args: recommended_plugins), + ) + + api = module.RagPipelineRecommendedPluginApi() + handler = unwrap_all(api.get) + + with app.test_request_context("/?type=tool"): + response = handler(api, "tenant-1", _account()) + + assert response == recommended_plugins diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py index 53f4f139937..6913825d599 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -3,7 +3,7 @@ import json from contextlib import ExitStack from inspect import unwrap from types import SimpleNamespace -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import ANY, MagicMock, PropertyMock, patch import pytest from flask import Flask @@ -63,6 +63,18 @@ def dataset_model_property_defaults(): for name, value in properties.items(): property_mock = stack.enter_context(patch.object(Dataset, name, new_callable=PropertyMock)) property_mock.return_value = value + stack.enter_context( + patch( + "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", + return_value=enterprise_rbac_service.MyPermissionsResponse(), + ) + ) + stack.enter_context( + patch( + "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetPermissions.batch_get", + return_value={}, + ) + ) yield @@ -245,7 +257,7 @@ class TestDatasetList: ): resp, status = method(api, "tenant-1", current_user) - get_permissions.assert_called_once_with("tenant-1", current_user.id) + get_permissions.assert_called_once_with("tenant-1", current_user.id, session=ANY) assert status == 200 assert resp["data"][0]["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] @@ -742,7 +754,7 @@ class TestDatasetApiGet: data, status = method(api, tenant_id, user, dataset_id) - get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id) + get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id, session=ANY) assert status == 200 assert data["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] diff --git a/api/tests/unit_tests/controllers/console/datasets/test_external.py b/api/tests/unit_tests/controllers/console/datasets/test_external.py index 8ac40f03d3b..1cffc90ae23 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_external.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_external.py @@ -1,5 +1,6 @@ import inspect -from unittest.mock import MagicMock, PropertyMock, patch +from types import SimpleNamespace +from unittest.mock import ANY, MagicMock, PropertyMock, patch import pytest from flask import Flask @@ -7,6 +8,7 @@ from werkzeug.exceptions import Forbidden, NotFound import services from controllers.console import console_ns +from controllers.console.datasets import external as external_module from controllers.console.datasets.error import DatasetNameDuplicateError from controllers.console.datasets.external import ( BedrockRetrievalApi, @@ -142,7 +144,7 @@ class TestExternalApiUseCheckApi: assert status == 200 assert response == {"is_using": True, "count": 2} - mock_use_check.assert_called_once_with(session, "api-id", "tenant-1") + mock_use_check.assert_called_once_with("api-id", "tenant-1", session=ANY) class TestExternalDatasetCreateApi: @@ -186,6 +188,7 @@ class TestExternalDatasetCreateApi: "create_external_dataset", return_value=dataset, ), + patch.object(external_module, "db", SimpleNamespace(session=lambda: MagicMock())), ): _, status = method(api, MagicMock(), "tenant-1", current_user) diff --git a/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py b/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py index 8a2e14cce9b..4adeaaa90dd 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py +++ b/api/tests/unit_tests/controllers/console/explore/test_recommended_app.py @@ -32,7 +32,7 @@ class TestRecommendedAppListApi: ): result = method(api, make_account("fr-FR")) - service_mock.assert_called_once_with(ANY, "en-US") + service_mock.assert_called_once_with("en-US", session=ANY) assert result == result_data def test_get_fallback_to_user_language(self, app: Flask): @@ -51,7 +51,7 @@ class TestRecommendedAppListApi: ): result = method(api, make_account("fr-FR")) - service_mock.assert_called_once_with(ANY, "fr-FR") + service_mock.assert_called_once_with("fr-FR", session=ANY) assert result == result_data def test_get_fallback_to_default_language(self, app: Flask): @@ -70,7 +70,7 @@ class TestRecommendedAppListApi: ): result = method(api, make_account(None)) - service_mock.assert_called_once_with(ANY, module.languages[0]) + service_mock.assert_called_once_with(module.languages[0], session=ANY) assert result == result_data @@ -91,7 +91,7 @@ class TestLearnDifyAppListApi: ): result = method(api, make_account("fr-FR")) - service_mock.assert_called_once_with(ANY, "en-US") + service_mock.assert_called_once_with("en-US", session=ANY) assert result == result_data def test_get_fallback_to_user_language(self, app: Flask): @@ -110,7 +110,7 @@ class TestLearnDifyAppListApi: ): result = method(api, make_account("fr-FR")) - service_mock.assert_called_once_with(ANY, "fr-FR") + service_mock.assert_called_once_with("fr-FR", session=ANY) assert result == result_data @@ -131,7 +131,7 @@ class TestRecommendedAppApi: ): result = method(api, "11111111-1111-1111-1111-111111111111") - service_mock.assert_called_once_with(ANY, "11111111-1111-1111-1111-111111111111") + service_mock.assert_called_once_with("11111111-1111-1111-1111-111111111111", session=ANY) assert result == result_data diff --git a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py index ae05b8f6a0e..f210d0d5d04 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py +++ b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py @@ -63,7 +63,7 @@ class TestSavedMessageListApi: result = method(api, current_user, installed_app) pagination_mock.assert_called_once() - assert pagination_mock.call_args.args[2] is current_user + assert pagination_mock.call_args.args[1] is current_user assert result["limit"] == 20 assert result["has_more"] is False assert len(result["data"]) == 2 @@ -96,7 +96,7 @@ class TestSavedMessageListApi: result = method(api, current_user, installed_app) save_mock.assert_called_once() - assert save_mock.call_args.args[2] is current_user + assert save_mock.call_args.args[1] is current_user assert result == {"result": "success"} def test_post_message_not_exists(self, app: Flask, payload_patch): @@ -136,7 +136,7 @@ class TestSavedMessageApi: result, status = method(api, current_user, installed_app, str(uuid4())) delete_mock.assert_called_once() - assert delete_mock.call_args.args[2] is current_user + assert delete_mock.call_args.args[1] is current_user assert status == 204 assert result == "" diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py index 8785ce85109..98b538800ac 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py @@ -38,7 +38,10 @@ def _snippet(**overrides) -> CustomizedSnippet: @pytest.fixture(autouse=True) def _patch_snippet_service_factory(monkeypatch: pytest.MonkeyPatch) -> None: def factory(): - return snippet_workflow_module.SnippetService() + try: + return snippet_workflow_module.SnippetService(snippet_workflow_module._snippet_session_maker()) + except TypeError: + return snippet_workflow_module.SnippetService() monkeypatch.setattr(snippet_workflow_module, "_snippet_service", factory) monkeypatch.setattr(snippet_workflow_module, "_snippet_session_maker", Mock(return_value=Mock())) diff --git a/api/tests/unit_tests/controllers/console/tag/test_tags.py b/api/tests/unit_tests/controllers/console/tag/test_tags.py index 2da11afa1f7..8aaebeb124a 100644 --- a/api/tests/unit_tests/controllers/console/tag/test_tags.py +++ b/api/tests/unit_tests/controllers/console/tag/test_tags.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock, PropertyMock, patch import pytest from flask import Flask -from sqlalchemy.orm import Session, scoped_session +from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden import controllers.console.tag.tags as module @@ -22,7 +22,7 @@ from services.tag_service import UpdateTagPayload class SessionMatcher: def __eq__(self, other): - return isinstance(other, Session | scoped_session) + return isinstance(other, Session) def unwrap(func): @@ -131,7 +131,7 @@ class TestTagListApi: ): result, status = method(api, "tenant-1") - get_tags_mock.assert_called_once_with(SessionMatcher(), "snippet", "tenant-1", None) + get_tags_mock.assert_called_once_with("snippet", "tenant-1", None, session=SessionMatcher()) assert status == 200 assert result == [{"id": "1", "name": "snippet-tag", "type": "snippet", "binding_count": "1"}] @@ -224,7 +224,7 @@ class TestTagUpdateDeleteApi: update_payload, tag_id, session = update_tags_mock.call_args.args assert update_payload == UpdateTagPayload(name="updated") assert tag_id == "tag-1" - assert session == module.db.session + assert session == SessionMatcher() assert result["binding_count"] == "3" def test_patch_forbidden(self, app: Flask, readonly_user, payload_patch): @@ -250,7 +250,7 @@ class TestTagUpdateDeleteApi: ): result, status = method(api, "tag-1") - delete_mock.assert_called_once_with("tag-1", module.db.session) + delete_mock.assert_called_once_with("tag-1", SessionMatcher()) assert status == 204 def test_delete_snippet_tag_checks_type_in_current_tenant(self, app: Flask, admin_user): @@ -278,7 +278,7 @@ class TestTagUpdateDeleteApi: scene=module.RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False, ) - delete_mock.assert_called_once_with("tag-1", module.db.session) + delete_mock.assert_called_once_with("tag-1", SessionMatcher()) assert result == "" assert status == 204 diff --git a/api/tests/unit_tests/controllers/console/test_extension.py b/api/tests/unit_tests/controllers/console/test_extension.py index bab825ca6f0..8ea327dfdce 100644 --- a/api/tests/unit_tests/controllers/console/test_extension.py +++ b/api/tests/unit_tests/controllers/console/test_extension.py @@ -114,7 +114,7 @@ def test_api_based_extension_get_returns_tenant_extensions(app: Flask, monkeypat assert response[0]["name"] == "Weather API" assert response[0]["api_endpoint"] == extension.api_endpoint assert response[0]["api_key"].startswith(extension.api_key[:3]) - service_mock.assert_called_once_with(ANY, "tenant-123") + service_mock.assert_called_once_with("tenant-123", session=ANY) def test_api_based_extension_post_creates_extension(app: Flask, monkeypatch: pytest.MonkeyPatch): @@ -132,7 +132,7 @@ def test_api_based_extension_post_creates_extension(app: Flask, monkeypatch: pyt response, status = APIBasedExtensionAPI().post() args, _ = save_mock.call_args - created_extension: APIBasedExtension = args[1] + created_extension: APIBasedExtension = args[0] assert created_extension.tenant_id == "tenant-123" assert created_extension.name == payload["name"] assert created_extension.api_endpoint == payload["api_endpoint"] @@ -157,7 +157,7 @@ def test_api_based_extension_detail_get_fetches_extension(app: Flask, monkeypatc assert response["id"] == extension.id assert response["name"] == extension.name - service_mock.assert_called_once_with(ANY, "tenant-123", str(extension_id)) + service_mock.assert_called_once_with("tenant-123", str(extension_id), session=ANY) def test_api_based_extension_detail_post_keeps_hidden_api_key(app: Flask, monkeypatch: pytest.MonkeyPatch): @@ -187,7 +187,7 @@ def test_api_based_extension_detail_post_keeps_hidden_api_key(app: Flask, monkey assert existing_extension.name == payload["name"] assert existing_extension.api_endpoint == payload["api_endpoint"] assert existing_extension.api_key == "keep-me" - save_mock.assert_called_once_with(ANY, existing_extension) + save_mock.assert_called_once_with(existing_extension, session=ANY) assert response["name"] == payload["name"] assert response["api_key"] == _masked_api_key("keep-me") @@ -217,7 +217,7 @@ def test_api_based_extension_detail_post_updates_api_key_when_provided(app: Flas response = APIBasedExtensionDetailAPI().post(extension_id) assert existing_extension.api_key == "new-secret" - save_mock.assert_called_once_with(ANY, existing_extension) + save_mock.assert_called_once_with(existing_extension, session=ANY) assert response["name"] == payload["name"] assert response["api_key"] == _masked_api_key(payload["api_key"]) @@ -239,6 +239,6 @@ def test_api_based_extension_detail_delete_removes_extension(app: Flask, monkeyp ): response, status = APIBasedExtensionDetailAPI().delete(extension_id) - delete_mock.assert_called_once_with(ANY, existing_extension) + delete_mock.assert_called_once_with(existing_extension, session=ANY) assert status == 204 assert response == "" diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py index 385539b6f30..2b385304d32 100644 --- a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py +++ b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py @@ -48,9 +48,11 @@ def test_console_setup_fastopenapi_post_success(app: Flask): patch("controllers.console.setup.TenantService.get_tenant_count", return_value=0), patch("controllers.console.setup.get_init_validate_status", return_value=True), patch("controllers.console.setup.RegisterService.setup"), + patch("controllers.console.setup.mark_setup_completed") as mark_setup_completed, ): client = app.test_client() response = client.post("/console/api/setup", json=payload) assert response.status_code == 201 assert response.get_json() == {"result": "success"} + mark_setup_completed.assert_called_once_with() diff --git a/api/tests/unit_tests/controllers/console/test_spec.py b/api/tests/unit_tests/controllers/console/test_spec.py index 58d7027751b..44fb8345928 100644 --- a/api/tests/unit_tests/controllers/console/test_spec.py +++ b/api/tests/unit_tests/controllers/console/test_spec.py @@ -11,7 +11,17 @@ class TestSpecSchemaDefinitionsApi: api = spec_module.SpecSchemaDefinitionsApi() method = unwrap(api.get) - schema_definitions = [{"type": "string"}] + schema_definitions = [ + { + "name": "conversation-variable", + "label": "Conversation variable", + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + }, + } + ] with patch.object( spec_module, @@ -23,6 +33,12 @@ class TestSpecSchemaDefinitionsApi: assert status == 200 assert resp == schema_definitions + assert spec_module.SchemaDefinitionsResponse.model_validate(resp).model_dump(mode="json") == schema_definitions + + def test_get_documents_tight_response_model(self): + response = spec_module.SpecSchemaDefinitionsApi.get.__apidoc__["responses"]["200"] + + assert response[1].name == spec_module.SchemaDefinitionsResponse.__name__ def test_get_exception_returns_empty_list(self, caplog: pytest.LogCaptureFixture): api = spec_module.SpecSchemaDefinitionsApi() diff --git a/api/tests/unit_tests/controllers/console/test_workspace_account.py b/api/tests/unit_tests/controllers/console/test_workspace_account.py index 5f36e805baa..39a3b2485dd 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_account.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_account.py @@ -692,7 +692,7 @@ def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup(): second.scalar_one_or_none.return_value = expected_account mock_session.execute.side_effect = [first, second] - result = AccountService.get_account_by_email_with_case_fallback(mock_session, "Mixed@Test.com") + result = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=mock_session) assert result is expected_account assert mock_session.execute.call_count == 2 diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index 618a1f52180..172125f4635 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -13,6 +13,7 @@ from controllers.console.workspace.error import AccountNotInitializedError from controllers.console.wraps import ( RBACPermission, RBACResourceScope, + _is_setup_completed, account_initialization_required, cloud_edition_billing_enabled, cloud_edition_billing_rate_limit_check, @@ -35,6 +36,12 @@ from models.account import AccountStatus, TenantAccountRole from services.feature_service import LicenseStatus +@pytest.fixture(autouse=True) +def reset_setup_required_cache(): + """Keep setup_required's process cache isolated across unit tests.""" + _is_setup_completed.reset_success() + + class MockUser(UserMixin): """Simple User class for testing.""" @@ -735,6 +742,39 @@ class TestSystemSetup: # Assert assert result == "admin_success" + @patch("controllers.console.wraps.db") + def test_should_cache_completed_setup(self, mock_db): + """Test that completed setup skips repeated DB reads in this process""" + mock_db.session.scalar.return_value = MagicMock() + + @setup_required + def admin_view(): + return "admin_success" + + with patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"): + assert admin_view() == "admin_success" + assert admin_view() == "admin_success" + + assert mock_db.session.scalar.call_count == 1 + + @patch("controllers.console.wraps.db") + @patch("controllers.console.wraps.os.environ.get") + def test_should_not_cache_missing_setup(self, mock_environ_get, mock_db): + """Test that first-time bootstrap completion can be observed later in the same process""" + mock_db.session.scalar.side_effect = [None, MagicMock()] + mock_environ_get.return_value = None + + @setup_required + def admin_view(): + return "admin_success" + + with patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"): + with pytest.raises(NotSetupError): + admin_view() + assert admin_view() == "admin_success" + + assert mock_db.session.scalar.call_count == 2 + @patch("controllers.console.wraps.db") @patch("controllers.console.wraps.os.environ.get") def test_should_raise_not_init_validate_error_with_init_password(self, mock_environ_get, mock_db: MagicMock): diff --git a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py index abd9b4facb9..66e6b8fe35a 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py @@ -1,4 +1,5 @@ import inspect +from datetime import UTC, datetime from unittest.mock import patch import pytest @@ -16,9 +17,39 @@ from controllers.console.workspace.endpoint import ( EndpointListApi, EndpointListForSinglePluginApi, ) +from core.entities.provider_entities import ProviderConfig, ProviderConfigType +from core.plugin.entities.endpoint import EndpointEntityWithInstance, EndpointProviderDeclaration from core.plugin.impl.exc import PluginPermissionDeniedError +def _endpoint_entity() -> EndpointEntityWithInstance: + now = datetime(2026, 1, 1, tzinfo=UTC) + return EndpointEntityWithInstance( + id="e1", + created_at=now, + updated_at=now, + tenant_id="t1", + plugin_id="p1", + settings={ + "api_key": "pl********et", + "enabled": True, + "ids": ["a", "b"], + "nested": {"limit": 3}, + }, + expired_at=now, + declaration=EndpointProviderDeclaration( + settings=[ + ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="api_key"), + ProviderConfig(type=ProviderConfigType.BOOLEAN, name="enabled"), + ] + ), + name="endpoint", + enabled=True, + url="https://example.test/hook-1", + hook_id="hook-1", + ) + + class TestEndpointCollectionApi: def test_create_success(self, app: Flask): api = EndpointCollectionApi() @@ -99,15 +130,40 @@ class TestEndpointListApi: def test_list_success(self, app: Flask): api = EndpointListApi() method = inspect.unwrap(api.get) + endpoint_entity = _endpoint_entity() with ( app.test_request_context("/?page=1&page_size=10"), - patch("controllers.console.workspace.endpoint.EndpointService.list_endpoints", return_value=[{"id": "e1"}]), + patch( + "controllers.console.workspace.endpoint.EndpointService.list_endpoints", + return_value=[endpoint_entity], + ), ): result = method(api, "t1", "u1") - assert "endpoints" in result - assert len(result["endpoints"]) == 1 + endpoint = result["endpoints"][0] + assert endpoint["id"] == "e1" + assert endpoint["created_at"] == "2026-01-01T00:00:00Z" + assert endpoint["updated_at"] == "2026-01-01T00:00:00Z" + assert endpoint["settings"] == { + "api_key": "pl********et", + "enabled": True, + "ids": ["a", "b"], + "nested": {"limit": 3}, + } + assert endpoint["tenant_id"] == "t1" + assert endpoint["plugin_id"] == "p1" + assert endpoint["expired_at"] == "2026-01-01T00:00:00Z" + assert endpoint["declaration"]["settings"][0]["type"] == "secret-input" + assert endpoint["declaration"]["settings"][0]["name"] == "api_key" + assert endpoint["declaration"]["settings"][1]["type"] == "boolean" + assert endpoint["declaration"]["settings"][1]["name"] == "enabled" + assert endpoint["declaration"]["endpoints"] == [] + assert endpoint["name"] == "endpoint" + assert endpoint["enabled"] is True + assert endpoint["url"] == "https://example.test/hook-1" + assert endpoint["hook_id"] == "hook-1" + assert endpoint_entity.settings["api_key"] == "pl********et" def test_list_invalid_query(self, app: Flask): api = EndpointListApi() @@ -129,12 +185,14 @@ class TestEndpointListForSinglePluginApi: app.test_request_context("/?page=1&page_size=10&plugin_id=p1"), patch( "controllers.console.workspace.endpoint.EndpointService.list_endpoints_for_single_plugin", - return_value=[{"id": "e1"}], + return_value=[_endpoint_entity()], ), ): result = method(api, "t1", "u1") - assert "endpoints" in result + assert result["endpoints"][0]["id"] == "e1" + assert result["endpoints"][0]["settings"]["api_key"] == "pl********et" + assert result["endpoints"][0]["settings"]["nested"] == {"limit": 3} def test_list_for_plugin_missing_param(self, app: Flask): api = EndpointListForSinglePluginApi() diff --git a/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py b/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py index a1d08849ee3..7d034f90642 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py @@ -6,7 +6,7 @@ import builtins import importlib import sys from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock import pytest from flask import Flask @@ -92,6 +92,7 @@ def test_validate_credentials_success(app: Flask, load_balancing_module, monkeyp model="gpt-4o", model_type=ModelType.LLM, credentials={"api_key": "sk-***"}, + session=ANY, ) @@ -143,5 +144,6 @@ def test_validate_credentials_with_config_id(app: Flask, load_balancing_module, model="gpt-4o", model_type=ModelType.LLM, credentials={"api_key": "sk-***"}, + session=ANY, config_id="cfg-1", ) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_members.py b/api/tests/unit_tests/controllers/console/workspace/test_members.py index 321e2a68f39..4ee31bbdbbe 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_members.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_members.py @@ -130,7 +130,7 @@ class TestMemberInviteEmailApi: features.workspace_members.is_available.return_value = True payload = { - "emails": ["a@test.com"], + "emails": ["A@TEST.com", "a@test.com"], "role": "normal", "language": "en-US", } @@ -138,8 +138,10 @@ class TestMemberInviteEmailApi: with ( app.test_request_context("/", json=payload), patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), - patch("controllers.console.workspace.members._count_new_member_invites", return_value=1), - patch("controllers.console.workspace.members.RegisterService.invite_new_member", return_value="token"), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=1) as mock_count, + patch( + "controllers.console.workspace.members.RegisterService.invite_new_member", return_value="token" + ) as mock_invite, patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "http://x"), patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), @@ -148,6 +150,10 @@ class TestMemberInviteEmailApi: assert status == 201 assert result["result"] == "success" + assert result["invitation_results"][0]["email"] == "a@test.com" + mock_count.assert_not_called() + mock_invite.assert_called_once() + assert mock_invite.call_args.kwargs["email"] == "a@test.com" def test_invite_limit_exceeded(self, app: Flask): api = MemberInviteEmailApi() diff --git a/api/tests/unit_tests/controllers/console/workspace/test_plugin.py b/api/tests/unit_tests/controllers/console/workspace/test_plugin.py index 2bd6be3bc7a..bbe94fce635 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_plugin.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_plugin.py @@ -44,7 +44,15 @@ from controllers.console.workspace.plugin import ( ) from core.plugin.entities.plugin import PluginInstallation from core.plugin.impl.exc import PluginDaemonClientSideError -from models.account import Account, TenantAccountRole, TenantPluginAutoUpgradeStrategy, TenantPluginPermission +from models.account import ( + Account, + TenantAccountRole, + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategySetting, + TenantPluginDebugPermission, + TenantPluginInstallPermission, +) def _plugin_category_list_item(category: str = "tool") -> dict[str, Any]: @@ -393,8 +401,8 @@ class TestPluginChangePermissionApi: user = _account(TenantAccountRole.NORMAL) payload = { - "install_permission": TenantPluginPermission.InstallPermission.EVERYONE, - "debug_permission": TenantPluginPermission.DebugPermission.EVERYONE, + "install_permission": TenantPluginInstallPermission.EVERYONE, + "debug_permission": TenantPluginDebugPermission.EVERYONE, } with ( @@ -410,8 +418,8 @@ class TestPluginChangePermissionApi: user = _account() payload = { - "install_permission": TenantPluginPermission.InstallPermission.EVERYONE, - "debug_permission": TenantPluginPermission.DebugPermission.EVERYONE, + "install_permission": TenantPluginInstallPermission.EVERYONE, + "debug_permission": TenantPluginDebugPermission.EVERYONE, } with ( @@ -1020,11 +1028,11 @@ class TestPluginChangeAutoUpgradeApi: user = _account() payload = { - "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value, + "category": TenantPluginAutoUpgradeCategory.TOOL.value, "auto_upgrade": { - "strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + "strategy_setting": TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, "upgrade_time_of_day": 0, - "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + "upgrade_mode": TenantPluginAutoUpgradeMode.EXCLUDE, "exclude_plugins": [], "include_plugins": [], }, @@ -1048,11 +1056,11 @@ class TestPluginChangeAutoUpgradeApi: user = _account() payload = { - "category": TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL.value, + "category": TenantPluginAutoUpgradeCategory.MODEL.value, "auto_upgrade": { - "strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + "strategy_setting": TenantPluginAutoUpgradeStrategySetting.LATEST, "upgrade_time_of_day": 3600, - "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + "upgrade_mode": TenantPluginAutoUpgradeMode.ALL, "exclude_plugins": [], "include_plugins": [], }, @@ -1068,7 +1076,7 @@ class TestPluginChangeAutoUpgradeApi: assert result["success"] is True change.assert_called_once() - assert change.call_args.kwargs["category"] == TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL + assert change.call_args.kwargs["category"] == TenantPluginAutoUpgradeCategory.MODEL def test_auto_upgrade_fail(self, app: Flask): api = PluginChangeAutoUpgradeApi() @@ -1077,11 +1085,11 @@ class TestPluginChangeAutoUpgradeApi: user = MagicMock(is_admin_or_owner=True) payload = { - "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value, + "category": TenantPluginAutoUpgradeCategory.TOOL.value, "auto_upgrade": { - "strategy_setting": TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + "strategy_setting": TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, "upgrade_time_of_day": 0, - "upgrade_mode": TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + "upgrade_mode": TenantPluginAutoUpgradeMode.EXCLUDE, "exclude_plugins": [], "include_plugins": [], }, @@ -1102,16 +1110,16 @@ class TestPluginFetchAutoUpgradeApi: method = unwrap(api.get) auto_upgrade = MagicMock( - category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + category=TenantPluginAutoUpgradeCategory.TOOL, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=1, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=[], include_plugins=[], ) with ( - app.test_request_context(f"/?category={TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value}"), + app.test_request_context(f"/?category={TenantPluginAutoUpgradeCategory.TOOL.value}"), patch( "controllers.console.workspace.plugin.PluginAutoUpgradeService.get_strategy", return_value=auto_upgrade, @@ -1119,7 +1127,7 @@ class TestPluginFetchAutoUpgradeApi: ): result = method(api, "t1") - assert result["category"] == TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL + assert result["category"] == TenantPluginAutoUpgradeCategory.TOOL assert result["auto_upgrade"]["upgrade_time_of_day"] == 1 @@ -1128,7 +1136,7 @@ class TestPluginAutoUpgradeExcludePluginApi: api = PluginAutoUpgradeExcludePluginApi() method = unwrap(api.post) - payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value} + payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeCategory.TOOL.value} with ( app.test_request_context("/", json=payload), @@ -1142,7 +1150,7 @@ class TestPluginAutoUpgradeExcludePluginApi: api = PluginAutoUpgradeExcludePluginApi() method = unwrap(api.post) - payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL.value} + payload = {"plugin_id": "p", "category": TenantPluginAutoUpgradeCategory.TOOL.value} with ( app.test_request_context("/", json=payload), diff --git a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py index 2a576d1c920..a8b61a43b2a 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py @@ -7,12 +7,15 @@ import importlib from contextlib import ExitStack, contextmanager from inspect import unwrap from types import ModuleType, SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask from flask.views import MethodView +from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolParameter from models import Account from models.account import TenantAccountRole @@ -21,6 +24,7 @@ if not hasattr(builtins, "MethodView"): _CONTROLLER_MODULE: ModuleType | None = None +_WRAPS_MODULE: ModuleType | None = None @contextmanager @@ -69,10 +73,11 @@ def controller_module(monkeypatch: pytest.MonkeyPatch): _CONTROLLER_MODULE = importlib.import_module(module_name) module = _CONTROLLER_MODULE - monkeypatch.setattr(module, "jsonable_encoder", lambda payload: payload) # Ensure decorators that consult deployment edition do not reach the database. + global _WRAPS_MODULE wraps_module = importlib.import_module("controllers.console.wraps") + _WRAPS_MODULE = wraps_module monkeypatch.setattr(module.dify_config, "EDITION", "CLOUD") monkeypatch.setattr(wraps_module.dify_config, "EDITION", "CLOUD") @@ -88,19 +93,194 @@ def _mock_account(user_id: str = "user-123") -> Account: return user +def _set_current_account( + monkeypatch: pytest.MonkeyPatch, + controller_module: ModuleType, + user: Account, + tenant_id: str, +) -> None: + def _getter(): + return user, tenant_id + + monkeypatch.setattr(controller_module, "current_account_with_tenant", _getter, raising=False) + if _WRAPS_MODULE is not None: + monkeypatch.setattr(_WRAPS_MODULE, "current_account_with_tenant", _getter) + + login_module = importlib.import_module("libs.login") + monkeypatch.setattr(login_module, "_get_user", lambda: user) + + +def _i18n(text: str) -> dict[str, str]: + return {"en_US": text, "zh_Hans": text, "pt_BR": text, "ja_JP": text} + + +def _tool_response(controller_module: ModuleType, name: str = "tool-a") -> tuple[dict, dict]: + expected = { + "author": "Dify", + "name": name, + "label": _i18n(name), + "description": _i18n(f"{name} description"), + "parameters": [], + "labels": [], + "output_schema": {}, + } + tool = controller_module.ToolApiEntity.model_validate(expected) + return tool.model_dump(mode="json"), expected + + +def _provider_entity_response( + controller_module: ModuleType, name: str = "provider", provider_type: str = "builtin" +) -> tuple[CoreToolProviderApiEntity, dict]: + service_payload = { + "id": f"{name}-id", + "author": "Dify", + "name": name, + "description": _i18n(f"{name} description"), + "icon": "tool.svg", + "icon_dark": "", + "label": _i18n(name), + "type": provider_type, + "masked_credentials": {"api_key": "[__HIDDEN__]"}, + "original_credentials": {"api_key": "sk-secret"}, + "is_team_authorization": False, + "allow_delete": True, + "plugin_id": "", + "plugin_unique_identifier": "", + "tools": [], + "labels": [], + "server_url": "", + "updated_at": 1, + "server_identifier": "", + "masked_headers": None, + "original_headers": None, + "authentication": None, + "is_dynamic_registration": True, + "configuration": None, + "identity_mode": "off", + "workflow_app_id": None, + } + provider = CoreToolProviderApiEntity.model_validate(service_payload) + return provider, provider.to_dict() + + +def _provider_list_item( + controller_module: ModuleType, name: str = "provider", provider_type: str = "builtin" +) -> tuple[dict, dict]: + service_payload = { + "id": f"{name}-id", + "author": "Dify", + "name": name, + "description": _i18n(f"{name} description"), + "icon": "tool.svg", + "icon_dark": "", + "label": _i18n(name), + "type": provider_type, + "team_credentials": {"api_key": "[__HIDDEN__]"}, + "is_team_authorization": False, + "allow_delete": True, + "plugin_id": "", + "plugin_unique_identifier": "", + "tools": [], + "labels": [], + } + expected = { + **service_payload, + } + provider = controller_module.ToolProviderApiEntityResponse.model_validate(expected) + return service_payload, provider.model_dump(mode="json", exclude_unset=True) + + +def _credential_response(controller_module: ModuleType, credential_id: str = "cred-1") -> tuple[dict, dict]: + expected = { + "id": credential_id, + "name": "Credential", + "provider": "demo", + "credential_type": controller_module.CredentialType.API_KEY, + "is_default": False, + "credentials": {}, + "visibility": "all_team_members", + "created_by": "", + "partial_member_list": [], + "from_other_member": False, + } + credential = controller_module.ToolProviderCredentialApiEntity.model_validate(expected) + return credential.model_dump(mode="json"), credential.model_dump(mode="json") + + +def _provider_config_response(controller_module: ModuleType) -> tuple[dict, dict]: + expected = { + "type": "secret-input", + "name": "api_key", + "scope": None, + "required": False, + "default": None, + "options": None, + "multiple": False, + "label": None, + "help": None, + "url": None, + "placeholder": None, + } + config = controller_module.ProviderConfig.model_validate(expected) + return config.model_dump(mode="json"), expected + + +def _api_provider_detail_response(controller_module: ModuleType) -> tuple[dict, dict]: + expected = { + "schema_type": "openapi", + "schema": "{}", + "tools": [], + "icon": {"background": "#252525", "content": "tool"}, + "description": "provider description", + "credentials": {"auth_type": "none"}, + "privacy_policy": "", + "custom_disclaimer": "", + "labels": [], + } + detail = controller_module.ApiProviderDetailResponse.model_validate(expected) + return detail.model_dump(mode="json", by_alias=True), expected + + +def _workflow_detail_response(controller_module: ModuleType) -> tuple[dict, dict]: + tool_payload, tool_expected = _tool_response(controller_module, "workflow-tool") + expected = { + "name": "workflow-tool", + "label": "Workflow Tool", + "workflow_tool_id": "00000000-0000-0000-0000-000000000001", + "workflow_app_id": "00000000-0000-0000-0000-000000000002", + "icon": {"background": "#252525", "content": "tool"}, + "description": "description", + "parameters": [], + "output_schema": {}, + "tool": tool_expected, + "synced": True, + "privacy_policy": "", + } + service_payload = {**expected, "tool": tool_payload} + detail = controller_module.WorkflowToolDetailResponse.model_validate(service_payload) + return detail.model_dump(mode="json"), expected + + +def _tool_label_response(controller_module: ModuleType, name: str = "search") -> tuple[dict, dict]: + expected = {"name": name, "label": _i18n(name), "icon": "search"} + label = controller_module.ToolLabel.model_validate(expected) + return label.model_dump(mode="json"), expected + + def test_tool_provider_list_calls_service_with_query( app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch ): user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-456") - service_mock = MagicMock(return_value=[{"provider": "builtin"}]) + service_payload, expected_response = _provider_list_item(controller_module, "builtin", "builtin") + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr(controller_module.ToolCommonService, "list_tool_providers", service_mock) with app.test_request_context("/workspaces/current/tool-providers?type=builtin"): - api = controller_module.ToolProviderListApi() - response = unwrap(api.get)(api, "tenant-456", user) + response = controller_module.ToolProviderListApi().get() - assert response == [{"provider": "builtin"}] + assert response == [expected_response] service_mock.assert_called_once_with(user.id, "tenant-456", "builtin") @@ -108,8 +288,9 @@ def test_builtin_provider_add_passes_payload( app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch ): user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-456") - service_mock = MagicMock(return_value={"status": "ok"}) + service_mock = MagicMock(return_value={"result": "success"}) monkeypatch.setattr(controller_module.BuiltinToolManageService, "add_builtin_tool_provider", service_mock) payload = { @@ -123,10 +304,9 @@ def test_builtin_provider_add_passes_payload( method="POST", json=payload, ): - api = controller_module.ToolBuiltinProviderAddApi() - response = unwrap(api.post)(api, "tenant-456", user, provider="openai") + response = controller_module.ToolBuiltinProviderAddApi().post(provider="openai") - assert response == {"status": "ok"} + assert response == {"result": "success"} service_mock.assert_called_once_with( user_id="user-123", tenant_id="tenant-456", @@ -140,38 +320,88 @@ def test_builtin_provider_add_passes_payload( def test_builtin_provider_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account("user-tenant-789") + _set_current_account(monkeypatch, controller_module, user, "tenant-789") - service_mock = MagicMock(return_value=[{"name": "tool-a"}]) + service_payload, expected_response = _tool_response(controller_module, "tool-a") + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr(controller_module.BuiltinToolManageService, "list_builtin_tool_provider_tools", service_mock) - monkeypatch.setattr(controller_module, "jsonable_encoder", lambda payload: payload) with app.test_request_context( "/workspaces/current/tool-provider/builtin/my-provider/tools", method="GET", ): - api = controller_module.ToolBuiltinProviderListToolsApi() - response = unwrap(api.get)(api, "tenant-789", provider="my-provider") + response = controller_module.ToolBuiltinProviderListToolsApi().get(provider="my-provider") - assert response == [{"name": "tool-a"}] + assert response == [expected_response] service_mock.assert_called_once_with("tenant-789", "my-provider") def test_builtin_provider_info_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account("user-tenant-9") - service_mock = MagicMock(return_value={"info": True}) + _set_current_account(monkeypatch, controller_module, user, "tenant-9") + service_payload, expected_response = _provider_entity_response(controller_module, "demo", "builtin") + service_mock = MagicMock(return_value=service_payload) monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) with app.test_request_context("/info", method="GET"): - api = controller_module.ToolBuiltinProviderInfoApi() - resp = unwrap(api.get)(api, "tenant-9", provider="demo") + resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") - assert resp == {"info": True} + assert resp == expected_response service_mock.assert_called_once_with("tenant-9", "demo") +def test_builtin_provider_info_uses_core_to_dict_tool_projection( + app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch +): + user = _mock_account("user-tenant-9") + _set_current_account(monkeypatch, controller_module, user, "tenant-9") + tool_parameter = ToolParameter( + name="system_files", + label=I18nObject(en_US="System Files", zh_Hans="System Files"), + type=ToolParameter.ToolParameterType.SYSTEM_FILES, + form=ToolParameter.ToolParameterForm.LLM, + input_schema=None, + ) + tool = controller_module.ToolApiEntity( + author="Dify", + name="demo-tool", + label=I18nObject(en_US="Demo Tool", zh_Hans="Demo Tool"), + description=I18nObject(en_US="Demo Tool description", zh_Hans="Demo Tool description"), + parameters=[tool_parameter], + labels=[], + output_schema={}, + ) + provider = CoreToolProviderApiEntity( + id="demo-id", + author="Dify", + name="demo", + description=I18nObject(en_US="demo description", zh_Hans="demo description"), + icon="tool.svg", + label=I18nObject(en_US="demo", zh_Hans="demo"), + type=controller_module.ToolProviderType.BUILT_IN, + masked_credentials={"api_key": "[__HIDDEN__]"}, + original_credentials={"api_key": "sk-secret"}, + tools=[tool], + ) + service_mock = MagicMock(return_value=provider) + monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) + + with app.test_request_context("/info", method="GET"): + resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") + + parameter = resp["tools"][0]["parameters"][0] + assert parameter["type"] == "files" + assert parameter["input_schema"] is None + assert resp["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in resp + assert "original_credentials" not in resp + + def test_builtin_provider_credentials_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account("user-tenant-cred") - service_mock = MagicMock(return_value=[{"cred": 1}]) + _set_current_account(monkeypatch, controller_module, user, "tenant-cred") + service_payload, expected_response = _credential_response(controller_module) + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr( controller_module.BuiltinToolManageService, "get_builtin_tool_provider_credentials", @@ -179,13 +409,13 @@ def test_builtin_provider_credentials_get(app: Flask, controller_module, monkeyp ) with app.test_request_context("/creds", method="GET"): - api = controller_module.ToolBuiltinProviderGetCredentialsApi() - resp = unwrap(api.get)(api, "tenant-cred", user, provider="demo") + resp = controller_module.ToolBuiltinProviderGetCredentialsApi().get(provider="demo") - assert resp == [{"cred": 1}] + assert resp == [expected_response] service_mock.assert_called_once_with( tenant_id="tenant-cred", provider_name="demo", + session=ANY, user=user, include_credential_ids=None, ) @@ -195,7 +425,8 @@ def test_builtin_provider_credentials_get_reads_repeated_include_ids( app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch ): user = _mock_account("user-tenant-cred") - service_mock = MagicMock(return_value=[{"cred": 1}]) + credential_payload, expected = _credential_response(controller_module) + service_mock = MagicMock(return_value=[credential_payload]) monkeypatch.setattr( controller_module.BuiltinToolManageService, "get_builtin_tool_provider_credentials", @@ -206,10 +437,11 @@ def test_builtin_provider_credentials_get_reads_repeated_include_ids( api = controller_module.ToolBuiltinProviderGetCredentialsApi() resp = unwrap(api.get)(api, "tenant-cred", user, provider="demo") - assert resp == [{"cred": 1}] + assert resp == [expected] service_mock.assert_called_once_with( tenant_id="tenant-cred", provider_name="demo", + session=ANY, user=user, include_credential_ids=["cred-1", "cred-2"], ) @@ -217,46 +449,51 @@ def test_builtin_provider_credentials_get_reads_repeated_include_ids( def test_api_provider_remote_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - service_mock = MagicMock(return_value={"schema": "ok"}) + _set_current_account(monkeypatch, controller_module, user, "tenant-10") + openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' + service_mock = MagicMock(return_value={"schema": openapi_schema}) monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider_remote_schema", service_mock) with app.test_request_context("/remote?url=https://example.com/"): - api = controller_module.ToolApiProviderGetRemoteSchemaApi() - resp = unwrap(api.get)(api, "tenant-10", user) + resp = controller_module.ToolApiProviderGetRemoteSchemaApi().get() - assert resp == {"schema": "ok"} + assert resp == {"schema": openapi_schema} service_mock.assert_called_once_with(user.id, "tenant-10", "https://example.com/") def test_api_provider_list_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - service_mock = MagicMock(return_value=[{"tool": "t"}]) + _set_current_account(monkeypatch, controller_module, user, "tenant-11") + service_payload, expected_response = _tool_response(controller_module, "t") + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr(controller_module.ApiToolManageService, "list_api_tool_provider_tools", service_mock) with app.test_request_context("/tools?provider=foo"): - api = controller_module.ToolApiProviderListToolsApi() - resp = unwrap(api.get)(api, "tenant-11", user) + resp = controller_module.ToolApiProviderListToolsApi().get() - assert resp == [{"tool": "t"}] + assert resp == [expected_response] service_mock.assert_called_once_with(user.id, "tenant-11", "foo") def test_api_provider_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - service_mock = MagicMock(return_value={"provider": "foo"}) + _set_current_account(monkeypatch, controller_module, user, "tenant-12") + service_payload, expected_response = _api_provider_detail_response(controller_module) + service_mock = MagicMock(return_value=service_payload) monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider", service_mock) with app.test_request_context("/get?provider=foo"): - api = controller_module.ToolApiProviderGetApi() - resp = unwrap(api.get)(api, "tenant-12", user) + resp = controller_module.ToolApiProviderGetApi().get() - assert resp == {"provider": "foo"} + assert resp == expected_response service_mock.assert_called_once_with(user.id, "tenant-12", "foo") def test_builtin_provider_credentials_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account("user-tenant-13") - service_mock = MagicMock(return_value={"schema": True}) + _set_current_account(monkeypatch, controller_module, user, "tenant-13") + service_payload, expected_response = _provider_config_response(controller_module) + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr( controller_module.BuiltinToolManageService, "list_builtin_provider_credentials_schema", @@ -264,16 +501,19 @@ def test_builtin_provider_credentials_schema_get(app: Flask, controller_module, ) with app.test_request_context("/schema", method="GET"): - api = controller_module.ToolBuiltinProviderCredentialsSchemaApi() - resp = unwrap(api.get)(api, "tenant-13", provider="demo", credential_type="api-key") + resp = controller_module.ToolBuiltinProviderCredentialsSchemaApi().get( + provider="demo", credential_type="api-key" + ) - assert resp == {"schema": True} + assert resp == [expected_response] service_mock.assert_called_once() def test_workflow_provider_get_by_tool(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - tool_service = MagicMock(return_value={"wf": 1}) + _set_current_account(monkeypatch, controller_module, user, "tenant-wf") + service_payload, expected_response = _workflow_detail_response(controller_module) + tool_service = MagicMock(return_value=service_payload) monkeypatch.setattr( controller_module.WorkflowToolManageService, "get_workflow_tool_by_tool_id", @@ -282,16 +522,17 @@ def test_workflow_provider_get_by_tool(app: Flask, controller_module, monkeypatc tool_id = "00000000-0000-0000-0000-000000000001" with app.test_request_context(f"/workflow?workflow_tool_id={tool_id}"): - api = controller_module.ToolWorkflowProviderGetApi() - resp = unwrap(api.get)(api, "tenant-wf", user) + resp = controller_module.ToolWorkflowProviderGetApi().get() - assert resp == {"wf": 1} + assert resp == expected_response tool_service.assert_called_once_with(user.id, "tenant-wf", tool_id) def test_workflow_provider_get_by_app(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - service_mock = MagicMock(return_value={"app": 1}) + _set_current_account(monkeypatch, controller_module, user, "tenant-wf2") + service_payload, expected_response = _workflow_detail_response(controller_module) + service_mock = MagicMock(return_value=service_payload) monkeypatch.setattr( controller_module.WorkflowToolManageService, "get_workflow_tool_by_app_id", @@ -300,31 +541,32 @@ def test_workflow_provider_get_by_app(app: Flask, controller_module, monkeypatch app_id = "00000000-0000-0000-0000-000000000002" with app.test_request_context(f"/workflow?workflow_app_id={app_id}"): - api = controller_module.ToolWorkflowProviderGetApi() - resp = unwrap(api.get)(api, "tenant-wf2", user) + resp = controller_module.ToolWorkflowProviderGetApi().get() - assert resp == {"app": 1} + assert resp == expected_response service_mock.assert_called_once_with(user.id, "tenant-wf2", app_id) def test_workflow_provider_list_tools(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() - service_mock = MagicMock(return_value=[{"id": 1}]) + _set_current_account(monkeypatch, controller_module, user, "tenant-wf3") + service_payload, expected_response = _tool_response(controller_module, "workflow-tool") + service_mock = MagicMock(return_value=[service_payload]) monkeypatch.setattr(controller_module.WorkflowToolManageService, "list_single_workflow_tools", service_mock) tool_id = "00000000-0000-0000-0000-000000000003" with app.test_request_context(f"/workflow/tools?workflow_tool_id={tool_id}"): - api = controller_module.ToolWorkflowProviderListToolApi() - resp = unwrap(api.get)(api, "tenant-wf3", user) + resp = controller_module.ToolWorkflowProviderListToolApi().get() - assert resp == [{"id": 1}] + assert resp == [expected_response] service_mock.assert_called_once_with(user.id, "tenant-wf3", tool_id) def test_builtin_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-bt") - provider = SimpleNamespace(to_dict=lambda: {"name": "builtin"}) + provider, expected_response = _provider_entity_response(controller_module, "builtin", "builtin") monkeypatch.setattr( controller_module.BuiltinToolManageService, "list_builtin_tools", @@ -332,16 +574,16 @@ def test_builtin_tools_list(app: Flask, controller_module, monkeypatch: pytest.M ) with app.test_request_context("/tools/builtin"): - api = controller_module.ToolBuiltinListApi() - resp = unwrap(api.get)(api, "tenant-bt", user) + resp = controller_module.ToolBuiltinListApi().get() - assert resp == [{"name": "builtin"}] + assert resp == [expected_response] def test_api_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account("user-tenant-api") + _set_current_account(monkeypatch, controller_module, user, "tenant-api") - provider = SimpleNamespace(to_dict=lambda: {"name": "api"}) + provider, expected_response = _provider_entity_response(controller_module, "api", "api") monkeypatch.setattr( controller_module.ApiToolManageService, "list_api_tools", @@ -349,16 +591,16 @@ def test_api_tools_list(app: Flask, controller_module, monkeypatch: pytest.Monke ) with app.test_request_context("/tools/api"): - api = controller_module.ToolApiListApi() - resp = unwrap(api.get)(api, "tenant-api") + resp = controller_module.ToolApiListApi().get() - assert resp == [{"name": "api"}] + assert resp == [expected_response] def test_workflow_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-wf4") - provider = SimpleNamespace(to_dict=lambda: {"name": "wf"}) + provider, expected_response = _provider_entity_response(controller_module, "wf", "workflow") monkeypatch.setattr( controller_module.WorkflowToolManageService, "list_tenant_workflow_tools", @@ -366,20 +608,21 @@ def test_workflow_tools_list(app: Flask, controller_module, monkeypatch: pytest. ) with app.test_request_context("/tools/workflow"): - api = controller_module.ToolWorkflowListApi() - resp = unwrap(api.get)(api, "tenant-wf4", user) + resp = controller_module.ToolWorkflowListApi().get() - assert resp == [{"name": "wf"}] + assert resp == [expected_response] def test_tool_labels_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(controller_module.ToolLabelsService, "list_tool_labels", lambda: ["a", "b"]) + user = _mock_account("user-label") + _set_current_account(monkeypatch, controller_module, user, "tenant-labels") + service_payload, expected_response = _tool_label_response(controller_module, "a") + monkeypatch.setattr(controller_module.ToolLabelsService, "list_tool_labels", lambda: [service_payload]) with app.test_request_context("/tool-labels"): - api = controller_module.ToolLabelsApi() - resp = unwrap(api.get)(api) + resp = controller_module.ToolLabelsApi().get() - assert resp == ["a", "b"] + assert resp == [expected_response] # --- _resolve_identity_mode: gating + None-resolution (PR #36839 review) --- diff --git a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py index b4fb60910e5..28a94387198 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py @@ -1,5 +1,6 @@ import inspect import logging +from http import HTTPStatus from io import BytesIO from unittest.mock import MagicMock, patch @@ -26,7 +27,9 @@ from controllers.console.workspace.workspace import ( WebappLogoWorkspaceApi, WorkspaceInfoApi, WorkspaceListApi, + WorkspaceLogoUploadResponse, WorkspacePermissionApi, + WorkspacePermissionResponse, ) from enums.cloud_plan import CloudPlan from libs.datetime_utils import naive_utc_now @@ -99,7 +102,7 @@ class TestTenantListApi: ): result, status = method(api, "t1", user) - assert status == 200 + assert status == HTTPStatus.OK assert len(result["workspaces"]) == 2 assert result["workspaces"][0]["current"] is True assert result["workspaces"][0]["plan"] == CloudPlan.TEAM @@ -146,7 +149,7 @@ class TestTenantListApi: ): result, status = method(api, "t1", user) - assert status == 200 + assert status == HTTPStatus.OK assert result["workspaces"][0]["plan"] == CloudPlan.TEAM assert result["workspaces"][1]["plan"] == CloudPlan.PROFESSIONAL get_plan_bulk_mock.assert_called_once_with(["t1", "t2"]) @@ -192,7 +195,7 @@ class TestTenantListApi: ): result, status = method(api, "t2", user) - assert status == 200 + assert status == HTTPStatus.OK assert result["workspaces"][0]["plan"] == CloudPlan.TEAM assert result["workspaces"][1]["plan"] == CloudPlan.TEAM get_plan_bulk_mock.assert_called_once_with(["t1", "t2"]) @@ -226,7 +229,7 @@ class TestTenantListApi: ): result, status = method(api, "t1", user) - assert status == 200 + assert status == HTTPStatus.OK assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX get_features_mock.assert_called_once_with("t1", exclude_vector_space=True) @@ -251,7 +254,7 @@ class TestTenantListApi: ): result, status = method(api, "t2", user) - assert status == 200 + assert status == HTTPStatus.OK assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX assert result["workspaces"][1]["plan"] == CloudPlan.SANDBOX assert result["workspaces"][0]["current"] is False @@ -276,7 +279,7 @@ class TestTenantListApi: ): result, status = method(api, None, user) - assert status == 200 + assert status == HTTPStatus.OK assert result["workspaces"] == [] get_features_mock.assert_not_called() @@ -295,7 +298,7 @@ class TestWorkspaceListApi: ): result, status = method(api) - assert status == 200 + assert status == HTTPStatus.OK assert result["total"] == 1 assert result["has_more"] is False @@ -312,7 +315,7 @@ class TestWorkspaceListApi: ): result, status = method(api) - assert status == 200 + assert status == HTTPStatus.OK assert result["has_more"] is True @@ -332,7 +335,7 @@ class TestTenantApi: ): result, status = method(api, user) - assert status == 200 + assert status == HTTPStatus.OK assert result["id"] == "t1" def test_post_archived_with_switch(self, app: Flask): @@ -386,7 +389,7 @@ class TestTenantApi: result, status = method(api, user) assert "Deprecated URL /info was used." in caplog.messages - assert status == 200 + assert status == HTTPStatus.OK class TestTenantInfoResponse: @@ -586,8 +589,9 @@ class TestWebappLogoWorkspaceApi: result, status = method(api, user) - assert status == 201 - assert result["id"] == "file1" + assert status == HTTPStatus.CREATED + assert result == {"id": "file1"} + assert WorkspaceLogoUploadResponse.model_validate(result).model_dump(mode="json") == {"id": "file1"} def test_filename_missing(self, app: Flask): api = WebappLogoWorkspaceApi() @@ -676,7 +680,7 @@ class TestWorkspaceInfoApi: patch("controllers.console.workspace.workspace.db.session.commit"), patch( "controllers.console.workspace.workspace.WorkspaceService.get_tenant_info", - return_value={"name": "New Name"}, + return_value={"id": "t1", "name": "New Name"}, ), ): result = method(api, "t1") @@ -716,8 +720,14 @@ class TestWorkspacePermissionApi: ): result, status = method(api, "t1") - assert status == 200 - assert result["workspace_id"] == "t1" + assert status == HTTPStatus.OK + expected = { + "workspace_id": "t1", + "allow_member_invite": True, + "allow_owner_transfer": False, + } + assert result == expected + assert WorkspacePermissionResponse.model_validate(result).model_dump(mode="json") == expected def test_no_current_tenant(self, app: Flask): api = WorkspacePermissionApi() diff --git a/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py b/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py index 71381e6a2b4..ad84eed1f5e 100644 --- a/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py +++ b/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py @@ -6,7 +6,7 @@ in test_auth_wraps.py; handler tests use inspect.unwrap() to bypass them. """ import inspect -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask @@ -19,7 +19,7 @@ from controllers.inner_api.app.dsl import ( _get_active_account, ) from models.account import AccountStatus -from services.app_dsl_service import ImportStatus +from services.app_dsl_service import Import, ImportStatus class TestInnerAppDSLImportPayload: @@ -117,9 +117,7 @@ class TestEnterpriseAppDSLImport: mock_dsl_cls.return_value = self._mock_dsl yield - def _make_import_result(self, status: ImportStatus, **kwargs) -> "Import": - from services.app_dsl_service import Import - + def _make_import_result(self, status: ImportStatus, **kwargs) -> Import: result = Import( id="import-id", status=status, @@ -224,7 +222,7 @@ class TestEnterpriseAppDSLExport: body, status_code = result assert status_code == 200 assert body["data"] == "version: 0.6.0\nkind: app\n" - mock_dsl_cls.export_dsl.assert_called_once_with(app_model=mock_app, include_secret=False) + mock_dsl_cls.export_dsl.assert_called_once_with(app_model=mock_app, session=ANY, include_secret=False) @patch("controllers.inner_api.app.dsl.AppDslService") @patch("controllers.inner_api.app.dsl.db") @@ -239,7 +237,7 @@ class TestEnterpriseAppDSLExport: body, status_code = result assert status_code == 200 - mock_dsl_cls.export_dsl.assert_called_once_with(app_model=mock_app, include_secret=True) + mock_dsl_cls.export_dsl.assert_called_once_with(app_model=mock_app, session=ANY, include_secret=True) @patch("controllers.inner_api.app.dsl.db") def test_export_app_not_found_returns_404(self, mock_db, api_instance, app: Flask): diff --git a/api/tests/unit_tests/controllers/inner_api/plugin/test_agent_drive.py b/api/tests/unit_tests/controllers/inner_api/plugin/test_agent_drive.py index 8c38564b3d0..8289a575050 100644 --- a/api/tests/unit_tests/controllers/inner_api/plugin/test_agent_drive.py +++ b/api/tests/unit_tests/controllers/inner_api/plugin/test_agent_drive.py @@ -9,7 +9,7 @@ from __future__ import annotations import inspect from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import ANY, patch import pytest from flask import Flask @@ -33,7 +33,7 @@ def test_manifest_parses_query_and_returns_items(): result = raw(AgentDriveManifestApi(), "agent-agent-1") assert result == {"items": [{"key": "docs/a.txt"}]} svc.return_value.manifest.assert_called_once_with( - tenant_id="tenant-1", agent_id="agent-1", prefix="docs/", include_download_url=True + tenant_id="tenant-1", agent_id="agent-1", prefix="docs/", include_download_url=True, session=ANY ) @@ -85,7 +85,11 @@ def test_skills_requires_tenant_id_and_returns_items(): } ] } - assert svc.return_value.list_skills.call_args.kwargs == {"tenant_id": "tenant-1", "agent_id": "agent-1"} + assert svc.return_value.list_skills.call_args.kwargs == { + "tenant_id": "tenant-1", + "agent_id": "agent-1", + "session": ANY, + } def test_commit_parses_body_and_returns_items(): diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py index a6626adc420..bda25bb2fa8 100644 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py @@ -117,7 +117,7 @@ class TestEnterpriseWorkspace: assert result["tenant"]["name"] == "My Workspace" mock_tenant_svc.create_tenant.assert_called_once_with("My Workspace", is_from_dashboard=True, session=ANY) mock_tenant_svc.create_tenant_member.assert_called_once_with( - mock_tenant, mock_account, mock_db.session, role="owner" + mock_tenant, mock_account, mock_db.session(), role="owner" ) mock_event.send.assert_called_once_with(mock_tenant) diff --git a/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py b/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py index b82ab254d45..7672e6414fd 100644 --- a/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py +++ b/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py @@ -72,7 +72,7 @@ def test_run_chat_always_calls_generate_with_streaming_true( "AppGenerateService", GenerateService, ) - with app.test_request_context(f"/openapi/v1/apps/{_TEST_APP_ID}/run", method="POST"): + with app.test_request_context(f"/openapi/v1/apps/{_TEST_APP_ID}:run", method="POST"): _run_chat( _make_app(), _make_account(), @@ -84,9 +84,9 @@ def test_run_chat_always_calls_generate_with_streaming_true( def test_stop_task_endpoint_registered(openapi_app): - """POST /openapi/v1/apps//tasks//stop must be registered.""" + """POST /openapi/v1/apps//tasks/:stop must be registered.""" rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/apps//tasks//stop" in rules + assert "/openapi/v1/apps//tasks/:stop" in rules def test_stop_task_calls_queue_manager_and_graph_engine(app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch): @@ -117,7 +117,7 @@ def test_stop_task_calls_queue_manager_and_graph_engine(app: Flask, bypass_pipel ) api = AppRunTaskStopApi() - with app.test_request_context("/openapi/v1/apps/app-1/tasks/task-1/stop", method="POST"): + with app.test_request_context("/openapi/v1/apps/app-1/tasks/task-1:stop", method="POST"): result = api.post.__wrapped__( api, app_id="app-1", diff --git a/api/tests/unit_tests/controllers/openapi/test_error_contract.py b/api/tests/unit_tests/controllers/openapi/test_error_contract.py index 788a7215ed2..6a30637b7ae 100644 --- a/api/tests/unit_tests/controllers/openapi/test_error_contract.py +++ b/api/tests/unit_tests/controllers/openapi/test_error_contract.py @@ -35,6 +35,7 @@ from controllers.openapi._errors import ( RecipientSurfaceMismatch, ) from controllers.service_api.app.error import ( + AgentNotPublishedError, AppUnavailableError, CompletionRequestError, ConversationCompletedError, @@ -306,6 +307,7 @@ ERROR_MATRIX = [ (InternalServerError(), 500, "internal_server_error"), (BadGateway("x"), 502, "bad_gateway"), (AppUnavailableError(), 400, "app_unavailable"), + (AgentNotPublishedError(), 400, "agent_not_published"), (ConversationCompletedError(), 400, "conversation_completed"), (ProviderNotInitializeError(), 400, "provider_not_initialize"), (ProviderQuotaExceededError(), 400, "provider_quota_exceeded"), diff --git a/api/tests/unit_tests/controllers/openapi/test_human_input_form.py b/api/tests/unit_tests/controllers/openapi/test_human_input_form.py index 5659cd6eeff..c4d3d21d84a 100644 --- a/api/tests/unit_tests/controllers/openapi/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/openapi/test_human_input_form.py @@ -62,7 +62,7 @@ class TestOpenApiHumanInputFormGet: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-1") - with app.test_request_context("/openapi/v1/apps/app-1/form/human_input/tok-1"): + with app.test_request_context("/openapi/v1/apps/app-1/human-input-forms/tok-1"): resp = api.get.__wrapped__( api, app_id="app-1", @@ -89,7 +89,7 @@ class TestOpenApiHumanInputFormGet: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-1") - with app.test_request_context("/openapi/v1/apps/app-1/form/human_input/bad"): + with app.test_request_context("/openapi/v1/apps/app-1/human-input-forms/bad"): with pytest.raises(HumanInputFormNotFound): api.get.__wrapped__( api, @@ -117,7 +117,7 @@ class TestOpenApiHumanInputFormGet: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-1") - with app.test_request_context("/openapi/v1/apps/app-1/form/human_input/tok-1"): + with app.test_request_context("/openapi/v1/apps/app-1/human-input-forms/tok-1"): with pytest.raises(HumanInputFormNotFound): api.get.__wrapped__( api, @@ -145,7 +145,7 @@ class TestOpenApiHumanInputFormGet: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-1") - with app.test_request_context("/openapi/v1/apps/app-1/form/human_input/tok-1"): + with app.test_request_context("/openapi/v1/apps/app-1/human-input-forms/tok-1"): with pytest.raises(RecipientSurfaceMismatch): api.get.__wrapped__( api, @@ -165,7 +165,7 @@ class TestOpenApiHumanInputFormPost: ) def test_post_account_caller_uses_user_id(self, app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch): - from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormApi + from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormSubmitApi form = self._make_form() service_mock = Mock() @@ -175,12 +175,12 @@ class TestOpenApiHumanInputFormPost: monkeypatch.setattr(module, "HumanInputService", lambda _engine: service_mock) monkeypatch.setattr(module, "db", SimpleNamespace(engine=object())) - api = OpenApiWorkflowHumanInputFormApi() + api = OpenApiWorkflowHumanInputFormSubmitApi() app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-42") with app.test_request_context( - "/openapi/v1/apps/app-1/form/human_input/tok-1", + "/openapi/v1/apps/app-1/human-input-forms/tok-1:submit", method="POST", json={"action": "approve", "inputs": {"field1": "val"}}, ): @@ -202,7 +202,7 @@ class TestOpenApiHumanInputFormPost: assert result == ({}, 200) def test_post_end_user_caller_uses_end_user_id(self, app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch): - from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormApi + from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormSubmitApi form = self._make_form() service_mock = Mock() @@ -212,12 +212,12 @@ class TestOpenApiHumanInputFormPost: monkeypatch.setattr(module, "HumanInputService", lambda _engine: service_mock) monkeypatch.setattr(module, "db", SimpleNamespace(engine=object())) - api = OpenApiWorkflowHumanInputFormApi() + api = OpenApiWorkflowHumanInputFormSubmitApi() app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="eu-7") with app.test_request_context( - "/openapi/v1/apps/app-1/form/human_input/tok-1", + "/openapi/v1/apps/app-1/human-input-forms/tok-1:submit", method="POST", json={"action": "approve", "inputs": {}}, ): @@ -241,7 +241,7 @@ class TestOpenApiHumanInputFormPost: def test_post_standalone_web_app_recipient_submits( self, app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch ): - from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormApi + from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormSubmitApi form = self._make_form(recipient_type=RecipientType.STANDALONE_WEB_APP) service_mock = Mock() @@ -251,12 +251,12 @@ class TestOpenApiHumanInputFormPost: monkeypatch.setattr(module, "HumanInputService", lambda _engine: service_mock) monkeypatch.setattr(module, "db", SimpleNamespace(engine=object())) - api = OpenApiWorkflowHumanInputFormApi() + api = OpenApiWorkflowHumanInputFormSubmitApi() app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="anyone") with app.test_request_context( - "/openapi/v1/apps/app-1/form/human_input/tok-1", + "/openapi/v1/apps/app-1/human-input-forms/tok-1:submit", method="POST", json={"action": "approve", "inputs": {}}, ): @@ -272,14 +272,14 @@ class TestOpenApiHumanInputFormPost: def test_post_rejects_invalid_body_with_422(self, app: Flask, bypass_pipeline): """Malformed body → 422 via @accepts (was an unmapped pydantic error → 500).""" - from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormApi + from controllers.openapi.human_input_form import OpenApiWorkflowHumanInputFormSubmitApi - api = OpenApiWorkflowHumanInputFormApi() + api = OpenApiWorkflowHumanInputFormSubmitApi() app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") caller = SimpleNamespace(id="acct-42") with app.test_request_context( - "/openapi/v1/apps/app-1/form/human_input/tok-1", + "/openapi/v1/apps/app-1/human-input-forms/tok-1:submit", method="POST", json={"inputs": {"field1": "val"}}, # missing required "action" ): diff --git a/api/tests/unit_tests/controllers/openapi/test_version_gate.py b/api/tests/unit_tests/controllers/openapi/test_version_gate.py new file mode 100644 index 00000000000..e2b16259952 --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_version_gate.py @@ -0,0 +1,100 @@ +"""Tests for the difyctl version gate on /openapi/v1 (HTTP 426 Upgrade Required). + +The gate is an app-level ``before_app_request`` hook: it must fire before routing, +so requests to *removed* paths (which no longer match a route) become 426 rather +than a bare 404. It reads the difyctl version from the User-Agent and fails open +for anything it can't confidently identify as an outdated difyctl. +""" + +from __future__ import annotations + +import uuid + +import pytest +from flask import Flask + +# Floor is [tool.dify] min_difyctl_version = "0.2.0". Comparison is on the numeric +# core (major.minor.patch), so 0.2.0-alpha passes (core 0.2.0 == floor) while +# 0.1.0 (core 0.1.0 < 0.2.0) is blocked. +OLD_UA = "difyctl/0.1.0 (darwin; arm64; stable)" +CURRENT_UA = "difyctl/0.2.0-alpha (darwin; arm64; stable)" + + +@pytest.fixture +def client(openapi_app: Flask): + return openapi_app.test_client() + + +def _gated_path() -> str: + """An existing, auth-guarded route on the surface (GET /apps/).""" + return f"/openapi/v1/apps/{uuid.uuid4()}" + + +class TestVersionGate: + def test_old_client_gets_426_with_upgrade_body(self, client): + res = client.get(_gated_path(), headers={"User-Agent": OLD_UA}) + + assert res.status_code == 426 + body = res.get_json() + assert body["code"] == "upgrade_required" + assert body["status"] == 426 + assert "0.1.0" in body["message"] + assert "0.2.0" in body["message"] + assert "docs.dify.ai" in body["hint"] + + def test_removed_old_path_gets_426_not_404(self, client): + # /apps//run was renamed to /apps/:run — the old path matches no + # route. The app-level gate must still turn it into 426, not a bare 404. + res = client.post( + f"/openapi/v1/apps/{uuid.uuid4()}/run", + headers={"User-Agent": OLD_UA}, + json={"inputs": {}}, + ) + + assert res.status_code == 426 + assert res.get_json()["code"] == "upgrade_required" + + def test_current_client_passes_gate(self, client): + # Gate passes → normal dispatch (auth rejects, never the gate's 426). + # 0.2.0-alpha == floor on the numeric core, so it passes despite the suffix. + res = client.get(_gated_path(), headers={"User-Agent": CURRENT_UA}) + + assert res.status_code != 426 + + def test_prerelease_at_floor_passes(self, client): + # Numeric-core comparison: a pre-release of the floor version (0.2.0-rc.1, + # core 0.2.0) passes, even though 0.2.0-rc.1 < 0.2.0 under naive ordering. + res = client.get(_gated_path(), headers={"User-Agent": "difyctl/0.2.0-rc.1 (darwin; arm64; rc)"}) + + assert res.status_code != 426 + + def test_prerelease_below_floor_gets_426(self, client): + # 0.1.9-rc.1 has core 0.1.9 < 0.2.0 floor → still blocked. + res = client.get(_gated_path(), headers={"User-Agent": "difyctl/0.1.9-rc.1 (darwin; arm64; rc)"}) + + assert res.status_code == 426 + + def test_non_difyctl_ua_passes(self, client): + res = client.get(_gated_path(), headers={"User-Agent": "curl/8.4.0"}) + + assert res.status_code != 426 + + def test_missing_ua_passes(self, client): + res = client.get(_gated_path()) + + assert res.status_code != 426 + + def test_unparseable_version_passes(self, client): + res = client.get(_gated_path(), headers={"User-Agent": "difyctl/notaversion (x; y; z)"}) + + assert res.status_code != 426 + + def test_version_probe_allowlisted(self, client): + res = client.get("/openapi/v1/_version", headers={"User-Agent": OLD_UA}) + + assert res.status_code == 200 + + def test_health_allowlisted(self, client): + res = client.get("/openapi/v1/_health", headers={"User-Agent": OLD_UA}) + + assert res.status_code == 200 diff --git a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py index cf9fa671987..86d26420253 100644 --- a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py +++ b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py @@ -1,7 +1,7 @@ """Member endpoints under /openapi/v1/workspaces//... Coverage: -- Route registration (5 endpoints across 4 URL patterns) +- Route registration (5 endpoints across 3 URL patterns) - Body validation lands at 400 (per spec — not Pydantic's default 422) - Domain exception → HTTP code mapping is preserved with the service's original message (so CLI users see what the console user sees) @@ -37,7 +37,6 @@ from controllers.openapi._models import MemberInvitePayload, MemberRoleUpdatePay from controllers.openapi.auth.data import AuthData from controllers.openapi.workspaces import ( WorkspaceMemberApi, - WorkspaceMemberRoleApi, WorkspaceMembersApi, WorkspaceSwitchApi, ) @@ -152,8 +151,8 @@ def _tenant_service(**overrides) -> SimpleNamespace: "get_tenant_members": Mock(return_value=[]), "remove_member_from_tenant": Mock(), "update_member_role": Mock(), - "get_tenant_by_id": lambda session, tenant_id: session.get(None, tenant_id), - "find_workspace_for_account": lambda session, account_id, workspace_id: session.execute(None).first(), + "get_tenant_by_id": lambda tenant_id, *, session: session.get(None, tenant_id), + "find_workspace_for_account": lambda account_id, workspace_id, *, session: session.execute(None).first(), } methods.update(overrides) return SimpleNamespace(**methods) @@ -163,19 +162,25 @@ def _account_service(**overrides) -> SimpleNamespace: """AccountService double; ``get_account_by_id`` delegates to the injected session (see :func:`_tenant_service`).""" methods: dict = { - "get_account_by_id": lambda session, account_id: session.get(None, account_id), + "get_account_by_id": lambda account_id, *, session: session.get(None, account_id), } methods.update(overrides) return SimpleNamespace(**methods) +def _db_mock() -> MagicMock: + mock_db = MagicMock() + mock_db.session.return_value = mock_db.session + return mock_db + + # --------------------------------------------------------------------------- # Route registration # --------------------------------------------------------------------------- def test_switch_route_registered(openapi_app: Flask): - rule = _rule(openapi_app, "/openapi/v1/workspaces//switch") + rule = _rule(openapi_app, "/openapi/v1/workspaces/:switch") assert openapi_app.view_functions[rule.endpoint].view_class is WorkspaceSwitchApi assert "POST" in rule.methods @@ -191,12 +196,7 @@ def test_member_by_id_route_registered(openapi_app: Flask): rule = _rule(openapi_app, "/openapi/v1/workspaces//members/") assert openapi_app.view_functions[rule.endpoint].view_class is WorkspaceMemberApi assert "DELETE" in rule.methods - - -def test_member_role_route_registered(openapi_app: Flask): - rule = _rule(openapi_app, "/openapi/v1/workspaces//members//role") - assert openapi_app.view_functions[rule.endpoint].view_class is WorkspaceMemberRoleApi - assert "PUT" in rule.methods + assert "PATCH" in rule.methods # --------------------------------------------------------------------------- @@ -250,17 +250,17 @@ def test_update_role_rejects_invalid_body_with_422(app: Flask, bypass_pipeline): """Invalid role-update body surfaces as 422 through @accepts (was 400).""" ws_id, member_id = str(uuid.uuid4()), str(uuid.uuid4()) acct_id = uuid.uuid4() - api = WorkspaceMemberRoleApi() + api = WorkspaceMemberApi() with app.test_request_context( - f"/openapi/v1/workspaces/{ws_id}/members/{member_id}/role", - method="PUT", + f"/openapi/v1/workspaces/{ws_id}/members/{member_id}", + method="PATCH", data=json.dumps({"role": "owner"}), # closed enum rejects owner content_type="application/json", ): _seed(_auth_ctx(account_id=acct_id)) with pytest.raises(UnprocessableEntity): - api.put.__wrapped__(api, workspace_id=ws_id, member_id=member_id, auth_data=_auth_data(acct_id)) + api.patch.__wrapped__(api, workspace_id=ws_id, member_id=member_id, auth_data=_auth_data(acct_id)) # --------------------------------------------------------------------------- @@ -278,7 +278,7 @@ def test_switch_returns_workspace_detail_with_current_true( acct_id = uuid.uuid4() api = WorkspaceSwitchApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = _account(account_id=str(acct_id)) membership = SimpleNamespace(role=TenantAccountRole.OWNER, current=True) mock_db.session.execute.return_value.first.return_value = (_tenant(ws_id), membership) @@ -291,7 +291,7 @@ def test_switch_returns_workspace_detail_with_current_true( ) monkeypatch.setattr(sys.modules["controllers.openapi.workspaces"], "db", mock_db) - with app.test_request_context(f"/openapi/v1/workspaces/{ws_id}/switch", method="POST"): + with app.test_request_context(f"/openapi/v1/workspaces/{ws_id}:switch", method="POST"): _seed(_auth_ctx(account_id=acct_id)) body, status = api.post.__wrapped__(api, workspace_id=ws_id, auth_data=_auth_data(acct_id)) @@ -310,7 +310,7 @@ def test_switch_404s_when_service_raises_account_not_link_tenant( acct_id = uuid.uuid4() api = WorkspaceSwitchApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = _account(account_id=str(acct_id)) monkeypatch.setattr( @@ -320,7 +320,7 @@ def test_switch_404s_when_service_raises_account_not_link_tenant( ) monkeypatch.setattr(sys.modules["controllers.openapi.workspaces"], "db", mock_db) - with app.test_request_context(f"/openapi/v1/workspaces/{ws_id}/switch", method="POST"): + with app.test_request_context(f"/openapi/v1/workspaces/{ws_id}:switch", method="POST"): _seed(_auth_ctx(account_id=acct_id)) with pytest.raises(NotFound): api.post.__wrapped__(api, workspace_id=ws_id, auth_data=_auth_data(acct_id)) @@ -345,7 +345,7 @@ def test_members_list_returns_normalized_rows(app: Flask, bypass_pipeline, monke role=TenantAccountRole.ADMIN, ) - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = _tenant(ws_id) monkeypatch.setattr( @@ -387,7 +387,7 @@ def test_members_list_paginates_with_query_params(app: Flask, bypass_pipeline, m for i in range(5) ] - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = _tenant(ws_id) monkeypatch.setattr( @@ -415,7 +415,7 @@ def test_members_list_rejects_unknown_query_param(app: Flask, bypass_pipeline, m acct_id = uuid.uuid4() api = WorkspaceMembersApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = _tenant(ws_id) monkeypatch.setattr(sys.modules["controllers.openapi.workspaces"], "db", mock_db) @@ -439,7 +439,7 @@ def test_invite_happy_path_returns_invite_url_and_member_id( invited = _account(account_id="new-1", email="new@example.com") - mock_db = MagicMock() + mock_db = _db_mock() # session.get is called twice: once for inviter Account, once for Tenant mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] @@ -520,7 +520,7 @@ def test_invite_blocked_by_saas_members_cap(app: Flask, bypass_pipeline, monkeyp acct_id = uuid.uuid4() api = WorkspaceMembersApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] invite_mock = Mock() @@ -558,7 +558,7 @@ def test_invite_blocked_by_ee_workspace_members_license(app: Flask, bypass_pipel acct_id = uuid.uuid4() api = WorkspaceMembersApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] invite_mock = Mock() @@ -598,7 +598,7 @@ def test_invite_ce_passes_when_both_caps_disabled(app: Flask, bypass_pipeline, m api = WorkspaceMembersApi() invited = _account(account_id="new-1", email="new@example.com") - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] monkeypatch.setattr( @@ -631,7 +631,7 @@ def test_invite_400_when_already_in_tenant(app: Flask, bypass_pipeline, monkeypa acct_id = uuid.uuid4() api = WorkspaceMembersApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] monkeypatch.setattr( @@ -662,7 +662,7 @@ def test_delete_member_happy_path(app: Flask, bypass_pipeline, monkeypatch: pyte acct_id = uuid.uuid4() api = WorkspaceMemberApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [ _account(account_id=str(acct_id)), # operator _tenant(ws_id), # tenant @@ -704,7 +704,7 @@ def test_delete_member_exception_mapping(app: Flask, bypass_pipeline, monkeypatc acct_id = uuid.uuid4() api = WorkspaceMemberApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [ _account(account_id=str(acct_id)), _tenant(ws_id), @@ -737,7 +737,7 @@ def test_delete_member_404_when_member_missing(app: Flask, bypass_pipeline, monk acct_id = uuid.uuid4() api = WorkspaceMemberApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [ _account(account_id=str(acct_id)), _tenant(ws_id), @@ -767,9 +767,9 @@ def test_delete_member_404_when_member_missing(app: Flask, bypass_pipeline, monk def test_update_role_happy_path(app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch): ws_id, member_id = str(uuid.uuid4()), str(uuid.uuid4()) acct_id = uuid.uuid4() - api = WorkspaceMemberRoleApi() + api = WorkspaceMemberApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [ _account(account_id=str(acct_id)), _tenant(ws_id), @@ -785,13 +785,15 @@ def test_update_role_happy_path(app: Flask, bypass_pipeline, monkeypatch: pytest monkeypatch.setattr(sys.modules["controllers.openapi.workspaces"], "db", mock_db) with app.test_request_context( - f"/openapi/v1/workspaces/{ws_id}/members/{member_id}/role", - method="PUT", + f"/openapi/v1/workspaces/{ws_id}/members/{member_id}", + method="PATCH", data=json.dumps({"role": "admin"}), content_type="application/json", ): _seed(_auth_ctx(account_id=acct_id)) - body, status = api.put.__wrapped__(api, workspace_id=ws_id, member_id=member_id, auth_data=_auth_data(acct_id)) + body, status = api.patch.__wrapped__( + api, workspace_id=ws_id, member_id=member_id, auth_data=_auth_data(acct_id) + ) assert status == 200 assert body == {"result": "success"} @@ -811,9 +813,9 @@ def test_update_role_happy_path(app: Flask, bypass_pipeline, monkeypatch: pytest def test_update_role_exception_mapping(app: Flask, bypass_pipeline, monkeypatch, exc, expected): ws_id, member_id = str(uuid.uuid4()), str(uuid.uuid4()) acct_id = uuid.uuid4() - api = WorkspaceMemberRoleApi() + api = WorkspaceMemberApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [ _account(account_id=str(acct_id)), _tenant(ws_id), @@ -828,14 +830,14 @@ def test_update_role_exception_mapping(app: Flask, bypass_pipeline, monkeypatch, monkeypatch.setattr(sys.modules["controllers.openapi.workspaces"], "db", mock_db) with app.test_request_context( - f"/openapi/v1/workspaces/{ws_id}/members/{member_id}/role", - method="PUT", + f"/openapi/v1/workspaces/{ws_id}/members/{member_id}", + method="PATCH", data=json.dumps({"role": "admin"}), content_type="application/json", ): _seed(_auth_ctx(account_id=acct_id)) with pytest.raises(expected): - api.put.__wrapped__( + api.patch.__wrapped__( api, workspace_id=ws_id, member_id=member_id, @@ -855,7 +857,7 @@ def test_load_tenant_rejects_archived_workspace(app: Flask, bypass_pipeline, mon api = WorkspaceMembersApi() archived = SimpleNamespace(id=ws_id, name="WS", status="archive", created_at=datetime(2026, 5, 18, tzinfo=UTC)) - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.return_value = archived monkeypatch.setattr( @@ -882,7 +884,7 @@ def test_invite_400_when_register_error(app: Flask, bypass_pipeline, monkeypatch acct_id = uuid.uuid4() api = WorkspaceMembersApi() - mock_db = MagicMock() + mock_db = _db_mock() mock_db.session.get.side_effect = [_account(account_id=str(acct_id)), _tenant(ws_id)] monkeypatch.setattr( diff --git a/api/tests/unit_tests/controllers/service_api/app/test_annotation.py b/api/tests/unit_tests/controllers/service_api/app/test_annotation.py index 810101fb0a5..1ff925cba7e 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_annotation.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_annotation.py @@ -15,7 +15,7 @@ Note: API endpoint tests for annotation controllers are complex due to: import uuid from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import ANY, Mock import pytest from flask import Flask @@ -264,7 +264,7 @@ class TestAnnotationListApi: assert response["page"] == 1 assert response["limit"] == 20 - get_mock.assert_called_once_with("app", 1, 20, "") + get_mock.assert_called_once_with("app", 1, 20, "", session=ANY) def test_get_accepts_valid_numeric_strings(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: annotation = SimpleNamespace(id="a1", question="q", content="a", created_at=0) @@ -281,7 +281,7 @@ class TestAnnotationListApi: assert response["total"] == 1 assert response["page"] == 2 assert response["limit"] == 5 - get_mock.assert_called_once_with("app", 2, 5, "refund") + get_mock.assert_called_once_with("app", 2, 5, "refund", session=ANY) @pytest.mark.parametrize("query_string", ["page=abc&limit=5", "page=1&limit=abc", "page=&limit=5", "limit=0"]) def test_get_rejects_invalid_explicit_pagination_value( diff --git a/api/tests/unit_tests/controllers/service_api/app/test_app.py b/api/tests/unit_tests/controllers/service_api/app/test_app.py index 9bc020b8b8f..04e9220ad55 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_app.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_app.py @@ -3,13 +3,14 @@ Unit tests for Service API App controllers """ import uuid -from unittest.mock import Mock, patch +from unittest.mock import ANY, Mock, patch import pytest from flask import Flask from controllers.service_api.app.app import AppInfoApi, AppMetaApi, AppParameterApi -from controllers.service_api.app.error import AppUnavailableError +from controllers.service_api.app.error import AgentNotPublishedError, AppUnavailableError +from core.app.apps.agent_app.errors import AgentAppNotPublishedError from models.account import TenantStatus from models.model import App, AppMode from tests.unit_tests.conftest import setup_mock_tenant_owner_execute_result @@ -185,6 +186,41 @@ class TestAppParameterApi: ] mock_get_agent_parameters.assert_called_once_with(mock_app_model) + @patch("controllers.service_api.wraps.user_logged_in") + @patch("controllers.service_api.wraps.current_app") + @patch("controllers.service_api.wraps.validate_and_get_api_token") + @patch("controllers.service_api.wraps.db") + @patch( + "controllers.service_api.app.app.get_published_agent_app_feature_dict_and_user_input_form", + side_effect=AgentAppNotPublishedError("Agent has not been published"), + ) + def test_get_parameters_for_unpublished_agent_app_raises_friendly_error( + self, + mock_get_agent_parameters, + mock_db, + mock_validate_token, + mock_current_app, + mock_user_logged_in, + app: Flask, + mock_app_model, + ): + _configure_current_app_mock(mock_current_app) + + mock_app_model.mode = AppMode.AGENT + mock_api_token = Mock() + mock_api_token.app_id = mock_app_model.id + mock_api_token.tenant_id = mock_app_model.tenant_id + mock_validate_token.return_value = mock_api_token + + mock_tenant = Mock() + mock_tenant.status = TenantStatus.NORMAL + mock_db.session.get.side_effect = [mock_app_model, mock_tenant] + setup_mock_tenant_owner_execute_result(mock_db, mock_tenant, Mock(current_tenant=mock_tenant)) + + with app.test_request_context("/parameters", method="GET", headers={"Authorization": "Bearer test_token"}): + with pytest.raises(AgentNotPublishedError): + AppParameterApi().get() + @patch("controllers.service_api.wraps.user_logged_in") @patch("controllers.service_api.wraps.current_app") @patch("controllers.service_api.wraps.validate_and_get_api_token") @@ -332,7 +368,7 @@ class TestAppMetaApi: response = api.get() # Assert - mock_service_instance.get_app_meta.assert_called_once_with(mock_app_model) + mock_service_instance.get_app_meta.assert_called_once_with(mock_app_model, session=ANY) assert response == {"tool_icons": {}, "AgentIcons": {}} diff --git a/api/tests/unit_tests/controllers/service_api/app/test_completion.py b/api/tests/unit_tests/controllers/service_api/app/test_completion.py index 393bdaf5eda..65652594294 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_completion.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_completion.py @@ -31,10 +31,12 @@ from controllers.service_api.app.completion import ( CompletionStopApi, ) from controllers.service_api.app.error import ( + AgentNotPublishedError, AppUnavailableError, ConversationCompletedError, NotChatAppError, ) +from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.errors.error import QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError from models.model import App, AppMode, EndUser @@ -250,7 +252,12 @@ class TestAppGenerateService: mock_generate.return_value = expected result = AppGenerateService.generate( - app_model=Mock(spec=App), user=Mock(spec=EndUser), args={"query": "Hi"}, invoke_from=Mock(), streaming=False + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + args={"query": "Hi"}, + invoke_from=Mock(), + session=Mock(), + streaming=False, ) assert result == expected @@ -262,7 +269,12 @@ class TestAppGenerateService: with pytest.raises(services.errors.conversation.ConversationNotExistsError): AppGenerateService.generate( - app_model=Mock(spec=App), user=Mock(spec=EndUser), args={}, invoke_from=Mock(), streaming=False + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + args={}, + invoke_from=Mock(), + session=Mock(), + streaming=False, ) @patch.object(AppGenerateService, "generate") @@ -272,7 +284,12 @@ class TestAppGenerateService: with pytest.raises(QuotaExceededError): AppGenerateService.generate( - app_model=Mock(spec=App), user=Mock(spec=EndUser), args={}, invoke_from=Mock(), streaming=False + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + args={}, + invoke_from=Mock(), + session=Mock(), + streaming=False, ) @patch.object(AppGenerateService, "generate") @@ -282,7 +299,12 @@ class TestAppGenerateService: with pytest.raises(InvokeError): AppGenerateService.generate( - app_model=Mock(spec=App), user=Mock(spec=EndUser), args={}, invoke_from=Mock(), streaming=False + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + args={}, + invoke_from=Mock(), + session=Mock(), + streaming=False, ) @@ -516,6 +538,22 @@ class TestChatApiController: with pytest.raises(BadRequest): handler(api, session=Mock(), app_model=app_model, end_user=end_user) + def test_agent_not_published_error_mapped(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AgentAppNotPublishedError("Agent has not been published")), + ) + + api = ChatApi() + handler = unwrap(api.post) + app_model = SimpleNamespace(mode=AppMode.AGENT.value) + end_user = SimpleNamespace() + + with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): + with pytest.raises(AgentNotPublishedError): + handler(api, session=Mock(), app_model=app_model, end_user=end_user) + class TestChatStopApiController: def test_wrong_mode(self, app: Flask) -> None: diff --git a/api/tests/unit_tests/controllers/service_api/app/test_conversation.py b/api/tests/unit_tests/controllers/service_api/app/test_conversation.py index 97873c631ae..3197812bc31 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_conversation.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_conversation.py @@ -475,6 +475,7 @@ class TestConversationService: user=Mock(spec=EndUser), name="New Name", auto_generate=False, + session=Mock(), ) assert result.name == "New Name" diff --git a/api/tests/unit_tests/controllers/service_api/app/test_message.py b/api/tests/unit_tests/controllers/service_api/app/test_message.py index d8d5c61bcb3..0400abe0e65 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_message.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_message.py @@ -266,6 +266,7 @@ class TestMessageService: conversation_id=str(uuid.uuid4()), first_id=None, limit=20, + session=Mock(), ) assert hasattr(result, "data") @@ -281,7 +282,12 @@ class TestMessageService: with pytest.raises(services.errors.conversation.ConversationNotExistsError): MessageService.pagination_by_first_id( - app_model=Mock(spec=App), user=Mock(spec=EndUser), conversation_id="invalid_id", first_id=None, limit=20 + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + conversation_id="invalid_id", + first_id=None, + limit=20, + session=Mock(), ) @patch.object(MessageService, "pagination_by_first_id") @@ -296,6 +302,7 @@ class TestMessageService: conversation_id=str(uuid.uuid4()), first_id="invalid_first_id", limit=20, + session=Mock(), ) @patch.object(MessageService, "create_feedback") @@ -309,6 +316,7 @@ class TestMessageService: user=Mock(spec=EndUser), rating=FeedbackRating.LIKE, content="Great response!", + session=Mock(), ) mock_create_feedback.assert_called_once() @@ -325,6 +333,7 @@ class TestMessageService: user=Mock(spec=EndUser), rating=FeedbackRating.LIKE, content=None, + session=Mock(), ) @patch.object(MessageService, "get_all_messages_feedbacks") @@ -336,7 +345,7 @@ class TestMessageService: ] mock_get_feedbacks.return_value = mock_feedbacks - result = MessageService.get_all_messages_feedbacks(app_model=Mock(spec=App), page=1, limit=20) + result = MessageService.get_all_messages_feedbacks(app_model=Mock(spec=App), page=1, limit=20, session=Mock()) assert len(result) == 2 assert result[0]["rating"] == "like" @@ -348,7 +357,11 @@ class TestMessageService: mock_get_questions.return_value = mock_questions result = MessageService.get_suggested_questions_after_answer( - app_model=Mock(spec=App), user=Mock(spec=EndUser), message_id=str(uuid.uuid4()), invoke_from=Mock() + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + message_id=str(uuid.uuid4()), + invoke_from=Mock(), + session=Mock(), ) assert len(result) == 3 @@ -361,7 +374,11 @@ class TestMessageService: with pytest.raises(SuggestedQuestionsAfterAnswerDisabledError): MessageService.get_suggested_questions_after_answer( - app_model=Mock(spec=App), user=Mock(spec=EndUser), message_id=str(uuid.uuid4()), invoke_from=Mock() + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + message_id=str(uuid.uuid4()), + invoke_from=Mock(), + session=Mock(), ) @patch.object(MessageService, "get_suggested_questions_after_answer") @@ -371,7 +388,11 @@ class TestMessageService: with pytest.raises(MessageNotExistsError): MessageService.get_suggested_questions_after_answer( - app_model=Mock(spec=App), user=Mock(spec=EndUser), message_id="invalid_message_id", invoke_from=Mock() + app_model=Mock(spec=App), + user=Mock(spec=EndUser), + message_id="invalid_message_id", + invoke_from=Mock(), + session=Mock(), ) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py index 3cabfe43ddc..2115bb85526 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py @@ -18,7 +18,7 @@ import uuid from datetime import UTC, datetime from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest from flask import Flask @@ -167,26 +167,6 @@ class TestWorkflowLogQuery: query_max_limit = WorkflowLogQuery(limit=100) assert query_max_limit.limit == 100 - def test_query_rejects_page_below_minimum(self): - """Test query rejects page < 1.""" - with pytest.raises(ValueError): - WorkflowLogQuery(page=0) - - def test_query_rejects_page_above_maximum(self): - """Test query rejects page > 99999.""" - with pytest.raises(ValueError): - WorkflowLogQuery(page=100000) - - def test_query_rejects_limit_below_minimum(self): - """Test query rejects limit < 1.""" - with pytest.raises(ValueError): - WorkflowLogQuery(limit=0) - - def test_query_rejects_limit_above_maximum(self): - """Test query rejects limit > 100.""" - with pytest.raises(ValueError): - WorkflowLogQuery(limit=101) - def test_query_with_keyword_search(self): """Test query with keyword filter.""" query = WorkflowLogQuery(keyword="workflow execution") @@ -263,7 +243,7 @@ class TestAppGenerateServiceWorkflow: """Test AppGenerateService workflow integration.""" @patch.object(AppGenerateService, "generate") - def test_generate_accepts_workflow_args(self, mock_generate): + def test_generate_accepts_workflow_args(self, mock_generate: MagicMock): """Test generate accepts workflow-specific args.""" mock_generate.return_value = {"result": "success"} @@ -272,6 +252,7 @@ class TestAppGenerateServiceWorkflow: user=Mock(), args={"inputs": {"key": "value"}, "workflow_id": "workflow_123"}, invoke_from=Mock(), + session=MagicMock(), streaming=False, ) @@ -279,7 +260,7 @@ class TestAppGenerateServiceWorkflow: mock_generate.assert_called_once() @patch.object(AppGenerateService, "generate") - def test_generate_raises_workflow_not_found_error(self, mock_generate): + def test_generate_raises_workflow_not_found_error(self, mock_generate: MagicMock): """Test generate raises WorkflowNotFoundError.""" mock_generate.side_effect = WorkflowNotFoundError("Workflow not found") @@ -289,11 +270,12 @@ class TestAppGenerateServiceWorkflow: user=Mock(), args={"workflow_id": "invalid_id"}, invoke_from=Mock(), + session=MagicMock(), streaming=False, ) @patch.object(AppGenerateService, "generate") - def test_generate_raises_is_draft_workflow_error(self, mock_generate): + def test_generate_raises_is_draft_workflow_error(self, mock_generate: MagicMock): """Test generate raises IsDraftWorkflowError.""" mock_generate.side_effect = IsDraftWorkflowError("Workflow is draft") @@ -303,11 +285,12 @@ class TestAppGenerateServiceWorkflow: user=Mock(), args={"workflow_id": "draft_workflow"}, invoke_from=Mock(), + session=MagicMock(), streaming=False, ) @patch.object(AppGenerateService, "generate") - def test_generate_supports_streaming_mode(self, mock_generate): + def test_generate_supports_streaming_mode(self, mock_generate: MagicMock): """Test generate supports streaming response mode.""" mock_stream = Mock() mock_generate.return_value = mock_stream @@ -317,6 +300,7 @@ class TestAppGenerateServiceWorkflow: user=Mock(), args={"inputs": {}, "response_mode": "streaming"}, invoke_from=Mock(), + session=MagicMock(), streaming=True, ) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py index 362af883ed2..406037e268d 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py @@ -325,10 +325,12 @@ class TestPipelineRunApiEntity: def test_entity_missing_required_field(self): """Test entity raises on missing required field.""" with pytest.raises(ValueError): - PipelineRunApiEntity( - inputs={}, - datasource_type="online_document", - # missing datasource_info_list, start_node_id, etc. + PipelineRunApiEntity.model_validate( + { + "inputs": {}, + "datasource_type": "online_document", + # missing datasource_info_list, start_node_id, etc. + } ) @@ -382,8 +384,19 @@ class TestDatasourcePluginsApiGet: mock_dataset = Mock() mock_db.session.scalar.return_value = mock_dataset + datasource_plugins = [ + { + "node_id": "node-datasource-1", + "plugin_id": "plugin-a", + "provider_name": "provider-a", + "datasource_type": "online_document", + "title": "Online Docs", + "user_input_variables": [{"variable": "url", "label": "URL", "type": "text-input", "required": True}], + "credentials": [{"id": "cred-1", "name": "Default credential", "type": "oauth2", "is_default": True}], + } + ] mock_svc_instance = Mock() - mock_svc_instance.get_datasource_plugins.return_value = [{"name": "plugin_a"}] + mock_svc_instance.get_datasource_plugins.return_value = datasource_plugins mock_svc_cls.return_value = mock_svc_instance with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=true"): @@ -391,11 +404,33 @@ class TestDatasourcePluginsApiGet: response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) assert status == 200 - assert response == [{"name": "plugin_a"}] + assert response == datasource_plugins mock_svc_instance.get_datasource_plugins.assert_called_once_with( tenant_id=tenant_id, dataset_id=dataset_id, is_published=True ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.db") + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + def test_get_plugins_parses_false_is_published_query(self, mock_svc_cls, mock_db, app: Flask): + """Test false query string is parsed as boolean False.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + + mock_db.session.scalar.return_value = Mock() + mock_svc_instance = Mock() + mock_svc_instance.get_datasource_plugins.return_value = [] + mock_svc_cls.return_value = mock_svc_instance + + with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=false"): + api = DatasourcePluginsApi() + response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) + + assert status == 200 + assert response == [] + mock_svc_instance.get_datasource_plugins.assert_called_once_with( + tenant_id=tenant_id, dataset_id=dataset_id, is_published=False + ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.db") def test_get_plugins_not_found(self, mock_db, app: Flask): """Test NotFound when dataset check fails.""" @@ -514,16 +549,14 @@ class TestPipelineRunApiPost: new_callable=lambda: Mock(spec=Account), ) @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.db") @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") - def test_post_success_streaming( - self, mock_ns, mock_db, mock_svc_cls, mock_current_user, mock_gen_svc, mock_helper, app - ): + def test_post_success_streaming(self, mock_ns, mock_svc_cls, mock_current_user, mock_gen_svc, mock_helper, app): """Test successful pipeline run with streaming response.""" tenant_id = str(uuid.uuid4()) dataset_id = str(uuid.uuid4()) - mock_db.session.scalar.return_value = Mock() + session = Mock() + session.scalar.return_value = Mock() mock_ns.payload = { "inputs": {"key": "val"}, @@ -544,27 +577,33 @@ class TestPipelineRunApiPost: with app.test_request_context("/datasets/test/pipeline/run", method="POST"): api = PipelineRunApi() - response = api.post(tenant_id=tenant_id, dataset_id=dataset_id) + response = api.post.__wrapped__(api, session, tenant_id=tenant_id, dataset_id=dataset_id) assert response == {"result": "ok"} + mock_svc_cls.assert_called_once_with(session) mock_gen_svc.generate.assert_called_once() - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.db") - def test_post_not_found(self, mock_db, app: Flask): + def test_post_not_found(self, app: Flask): """Test NotFound when dataset check fails.""" - mock_db.session.scalar.return_value = None + session = Mock() + session.scalar.return_value = None with app.test_request_context("/datasets/test/pipeline/run", method="POST"): api = PipelineRunApi() with pytest.raises(NotFound): - api.post(tenant_id=str(uuid.uuid4()), dataset_id=str(uuid.uuid4())) + api.post.__wrapped__( + api, + session, + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + ) @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", new="not_account") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.db") @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") - def test_post_forbidden_non_account_user(self, mock_ns, mock_db, app: Flask): + def test_post_forbidden_non_account_user(self, mock_ns, app: Flask): """Test Forbidden when current_user is not an Account.""" - mock_db.session.scalar.return_value = Mock() + session = Mock() + session.scalar.return_value = Mock() mock_ns.payload = { "inputs": {}, "datasource_type": "online_document", @@ -577,7 +616,12 @@ class TestPipelineRunApiPost: with app.test_request_context("/datasets/test/pipeline/run", method="POST"): api = PipelineRunApi() with pytest.raises(Forbidden): - api.post(tenant_id=str(uuid.uuid4()), dataset_id=str(uuid.uuid4())) + api.post.__wrapped__( + api, + session, + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + ) class TestFileUploadApiPost: diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py index a95baf1b482..0b1ca8741a9 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py @@ -1193,7 +1193,7 @@ class TestDatasetSegmentApiDelete: # Assert assert response == ("", 204) - mock_seg_svc.delete_segment.assert_called_once_with(mock_segment, mock_doc, mock_dataset, mock_db.session) + mock_seg_svc.delete_segment.assert_called_once_with(mock_segment, mock_doc, mock_dataset, mock_db.session()) @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py index dd2caf4f3fc..e83724f955f 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py @@ -783,7 +783,7 @@ class TestDocumentApiDelete: # Assert assert response == ("", 204) - mock_doc_svc.delete_document.assert_called_once_with(mock_document, mock_db.session) + mock_doc_svc.delete_document.assert_called_once_with(mock_document, mock_db.session()) @patch("controllers.service_api.dataset.document.DocumentService") @patch("controllers.service_api.dataset.document.db") diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py b/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py index b77c783ae16..dd1322a6344 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py @@ -408,7 +408,7 @@ class TestDatasetMetadataBuiltInFieldAction: assert status == 200 assert response["result"] == "success" - mock_meta_svc.enable_built_in_field.assert_called_once_with(ANY, mock_dataset) + mock_meta_svc.enable_built_in_field.assert_called_once_with(mock_dataset, session=ANY) @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") @@ -439,7 +439,7 @@ class TestDatasetMetadataBuiltInFieldAction: ) assert status == 200 - mock_meta_svc.disable_built_in_field.assert_called_once_with(ANY, mock_dataset) + mock_meta_svc.disable_built_in_field.assert_called_once_with(mock_dataset, session=ANY) @patch("controllers.service_api.dataset.metadata.DatasetService") def test_action_dataset_not_found( diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py b/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py index a8dd8523acb..bfda3e23b32 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py @@ -2,9 +2,10 @@ Unit tests for Service API knowledge pipeline file-upload serialization. """ -import importlib.util from datetime import UTC, datetime -from pathlib import Path + +from controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow import PipelineUploadFileResponse +from libs.helper import dump_response class FakeUploadFile: @@ -17,21 +18,7 @@ class FakeUploadFile: created_at: datetime | None -def _load_serialize_upload_file(): - api_dir = Path(__file__).resolve().parents[5] - serializers_path = api_dir / "controllers" / "service_api" / "dataset" / "rag_pipeline" / "serializers.py" - - spec = importlib.util.spec_from_file_location("rag_pipeline_serializers", serializers_path) - assert spec - assert spec.loader - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) # type: ignore[attr-defined] - return module.serialize_upload_file - - def test_file_upload_created_at_is_isoformat_string(): - serialize_upload_file = _load_serialize_upload_file() - created_at = datetime(2026, 2, 8, 12, 0, 0, tzinfo=UTC) upload_file = FakeUploadFile() upload_file.id = "file-1" @@ -42,13 +29,11 @@ def test_file_upload_created_at_is_isoformat_string(): upload_file.created_by = "account-1" upload_file.created_at = created_at - result = serialize_upload_file(upload_file) + result = dump_response(PipelineUploadFileResponse, upload_file) assert result["created_at"] == created_at.isoformat() def test_file_upload_created_at_none_serializes_to_null(): - serialize_upload_file = _load_serialize_upload_file() - upload_file = FakeUploadFile() upload_file.id = "file-1" upload_file.name = "test.pdf" @@ -58,5 +43,5 @@ def test_file_upload_created_at_none_serializes_to_null(): upload_file.created_by = "account-1" upload_file.created_at = None - result = serialize_upload_file(upload_file) + result = dump_response(PipelineUploadFileResponse, upload_file) assert result["created_at"] is None diff --git a/api/tests/unit_tests/controllers/web/test_app.py b/api/tests/unit_tests/controllers/web/test_app.py index ce7ae271889..73f308dc749 100644 --- a/api/tests/unit_tests/controllers/web/test_app.py +++ b/api/tests/unit_tests/controllers/web/test_app.py @@ -3,13 +3,14 @@ from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask from controllers.web.app import AppAccessMode, AppMeta, AppParameterApi, AppWebAuthPermission -from controllers.web.error import AppUnavailableError +from controllers.web.error import AgentNotPublishedError, AppUnavailableError +from core.app.apps.agent_app.errors import AgentAppNotPublishedError # --------------------------------------------------------------------------- @@ -80,6 +81,18 @@ class TestAppParameterApi: with pytest.raises(AppUnavailableError): AppParameterApi().get(app_model, SimpleNamespace()) + def test_agent_mode_unpublished_raises_friendly_error(self, app: Flask) -> None: + app_model = SimpleNamespace(mode="agent") + with ( + app.test_request_context("/parameters"), + patch( + "controllers.web.app.get_published_agent_app_feature_dict_and_user_input_form", + side_effect=AgentAppNotPublishedError("Agent has not been published"), + ), + ): + with pytest.raises(AgentNotPublishedError): + AppParameterApi().get(app_model, SimpleNamespace()) + # --------------------------------------------------------------------------- # AppMeta @@ -135,7 +148,7 @@ class TestAppAccessMode: with app.test_request_context("/webapp/access-mode?appCode=code1"): result = AppAccessMode().get() - mock_resolve.assert_called_once_with("code1") + mock_resolve.assert_called_once_with("code1", session=ANY) mock_access.assert_called_once_with("resolved-id") assert result == {"accessMode": "external"} diff --git a/api/tests/unit_tests/controllers/web/test_completion.py b/api/tests/unit_tests/controllers/web/test_completion.py index 4f8d848637d..49a88802470 100644 --- a/api/tests/unit_tests/controllers/web/test_completion.py +++ b/api/tests/unit_tests/controllers/web/test_completion.py @@ -10,6 +10,7 @@ from flask import Flask from controllers.web.completion import ChatApi, ChatStopApi, CompletionApi, CompletionStopApi from controllers.web.error import ( + AgentNotPublishedError, CompletionRequestError, NotChatAppError, NotCompletionAppError, @@ -17,6 +18,7 @@ from controllers.web.error import ( ProviderNotInitializeError, ProviderQuotaExceededError, ) +from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError @@ -142,6 +144,19 @@ class TestChatApi: with pytest.raises(CompletionRequestError): ChatApi().post(_chat_app(), _end_user()) + @patch( + "controllers.web.completion.AppGenerateService.generate", + side_effect=AgentAppNotPublishedError("Agent has not been published"), + ) + @patch("controllers.web.completion.web_ns") + def test_agent_not_published_error_mapped(self, mock_ns: MagicMock, mock_gen: MagicMock, app: Flask) -> None: + mock_ns.payload = {"inputs": {}, "query": "x"} + app_model = SimpleNamespace(id="app-1", mode="agent") + + with app.test_request_context("/chat-messages", method="POST"): + with pytest.raises(AgentNotPublishedError): + ChatApi().post(app_model, _end_user()) + # --------------------------------------------------------------------------- # ChatStopApi diff --git a/api/tests/unit_tests/controllers/web/test_error.py b/api/tests/unit_tests/controllers/web/test_error.py index 2852c4806c9..9bf58ee0c5f 100644 --- a/api/tests/unit_tests/controllers/web/test_error.py +++ b/api/tests/unit_tests/controllers/web/test_error.py @@ -6,6 +6,7 @@ import pytest from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.web.error import ( + AgentNotPublishedError, AppMoreLikeThisDisabledError, AppSuggestedQuestionsAfterAnswerDisabledError, AppUnavailableError, @@ -29,6 +30,7 @@ from controllers.web.error import ( _ERROR_SPECS: list[tuple[type, str, int]] = [ (AppUnavailableError, "app_unavailable", 400), + (AgentNotPublishedError, "agent_not_published", 400), (NotCompletionAppError, "not_completion_app", 400), (NotChatAppError, "not_chat_app", 400), (NotWorkflowAppError, "not_workflow_app", 400), diff --git a/api/tests/unit_tests/controllers/web/test_message_list.py b/api/tests/unit_tests/controllers/web/test_message_list.py index 2bb425cdba2..b5d74df65ef 100644 --- a/api/tests/unit_tests/controllers/web/test_message_list.py +++ b/api/tests/unit_tests/controllers/web/test_message_list.py @@ -6,7 +6,7 @@ import builtins import uuid from datetime import datetime from types import ModuleType, SimpleNamespace -from unittest.mock import patch +from unittest.mock import ANY, patch from uuid import uuid4 import pytest @@ -158,7 +158,7 @@ def test_message_list_mapping(app: Flask) -> None: ): response = MessageListApi().get(app_model, end_user) - mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20) + mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20, session=ANY) assert response["limit"] == 20 assert response["has_more"] is False assert len(response["data"]) == 1 diff --git a/api/tests/unit_tests/controllers/web/test_passport.py b/api/tests/unit_tests/controllers/web/test_passport.py index 58d58626b22..ebc64d94521 100644 --- a/api/tests/unit_tests/controllers/web/test_passport.py +++ b/api/tests/unit_tests/controllers/web/test_passport.py @@ -34,12 +34,11 @@ def test_decode_enterprise_webapp_user_id_valid(monkeypatch: pytest.MonkeyPatch) def test_exchange_token_public_flow(monkeypatch: pytest.MonkeyPatch) -> None: site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal") app_model = SimpleNamespace(id="a1", status="normal", enable_site=True) + call_state = {"calls": 0} def _scalar_side_effect(*_args, **_kwargs): - if not hasattr(_scalar_side_effect, "calls"): - _scalar_side_effect.calls = 0 - _scalar_side_effect.calls += 1 - return site if _scalar_side_effect.calls == 1 else app_model + call_state["calls"] += 1 + return site if call_state["calls"] == 1 else app_model db_session = SimpleNamespace(scalar=_scalar_side_effect) monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session)) @@ -53,12 +52,11 @@ def test_exchange_token_public_flow(monkeypatch: pytest.MonkeyPatch) -> None: def test_exchange_token_requires_external(monkeypatch: pytest.MonkeyPatch) -> None: site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal") app_model = SimpleNamespace(id="a1", status="normal", enable_site=True) + call_state = {"calls": 0} def _scalar_side_effect(*_args, **_kwargs): - if not hasattr(_scalar_side_effect, "calls"): - _scalar_side_effect.calls = 0 - _scalar_side_effect.calls += 1 - return site if _scalar_side_effect.calls == 1 else app_model + call_state["calls"] += 1 + return site if call_state["calls"] == 1 else app_model db_session = SimpleNamespace(scalar=_scalar_side_effect) monkeypatch.setattr("controllers.web.passport.db", SimpleNamespace(session=db_session)) @@ -71,14 +69,13 @@ def test_exchange_token_requires_external(monkeypatch: pytest.MonkeyPatch) -> No def test_exchange_token_missing_session_id(monkeypatch: pytest.MonkeyPatch) -> None: site = SimpleNamespace(id="s1", app_id="a1", code="code", status="normal") app_model = SimpleNamespace(id="a1", status="normal", enable_site=True, tenant_id="t1") + call_state = {"calls": 0} def _scalar_side_effect(*_args, **_kwargs): - if not hasattr(_scalar_side_effect, "calls"): - _scalar_side_effect.calls = 0 - _scalar_side_effect.calls += 1 - if _scalar_side_effect.calls == 1: + call_state["calls"] += 1 + if call_state["calls"] == 1: return site - if _scalar_side_effect.calls == 2: + if call_state["calls"] == 2: return app_model return None diff --git a/api/tests/unit_tests/controllers/web/test_web_login.py b/api/tests/unit_tests/controllers/web/test_web_login.py index bfffd5cbb2c..a91d4253aa8 100644 --- a/api/tests/unit_tests/controllers/web/test_web_login.py +++ b/api/tests/unit_tests/controllers/web/test_web_login.py @@ -1,7 +1,7 @@ import base64 import logging from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask @@ -66,7 +66,7 @@ class TestEmailCodeLoginSendEmailApi: response = EmailCodeLoginSendEmailApi().post() assert response == {"result": "success", "data": "token-123"} - mock_get_user.assert_called_once_with("User@Example.com") + mock_get_user.assert_called_once_with("User@Example.com", ANY) mock_send_email.assert_called_once_with(account=mock_account, language="en-US") @@ -95,8 +95,8 @@ class TestEmailCodeLoginApi: ): response = EmailCodeLoginApi().post() - assert response.get_json() == {"result": "success", "data": {"access_token": "new-access-token"}} - mock_get_user.assert_called_once_with("User@Example.com") + assert response == {"result": "success", "data": {"access_token": "new-access-token"}} + mock_get_user.assert_called_once_with("User@Example.com", ANY) mock_revoke_token.assert_called_once_with("token-123") mock_login.assert_called_once() mock_reset_login_rate.assert_called_once_with("user@example.com") @@ -115,7 +115,7 @@ class TestLoginApi: ): response = LoginApi().post() - assert response.get_json()["data"]["access_token"] == "access-tok" + assert response["data"]["access_token"] == "access-tok" mock_auth.assert_called_once() @patch( diff --git a/api/tests/unit_tests/controllers/web/test_web_passport.py b/api/tests/unit_tests/controllers/web/test_web_passport.py index 19b1d8504a0..4e1a24a4da2 100644 --- a/api/tests/unit_tests/controllers/web/test_web_passport.py +++ b/api/tests/unit_tests/controllers/web/test_web_passport.py @@ -33,6 +33,7 @@ class TestDecodeEnterpriseWebappUserId: "user_id": "u1", } result = decode_enterprise_webapp_user_id("valid-jwt") + assert result is not None assert result["user_id"] == "u1" @patch("controllers.web.passport.PassportService") @@ -143,7 +144,7 @@ class TestPassportResource: with app.test_request_context("/passport", headers={"X-App-Code": "code1"}): response = PassportResource().get() - assert response.get_json()["access_token"] == "issued-token" + assert response["access_token"] == "issued-token" mock_db.session.add.assert_called_once() mock_db.session.commit.assert_called_once() @@ -167,7 +168,7 @@ class TestPassportResource: with app.test_request_context("/passport?user_id=sess-existing", headers={"X-App-Code": "code1"}): response = PassportResource().get() - assert response.get_json()["access_token"] == "reused-token" + assert response["access_token"] == "reused-token" # Should not create a new end user mock_db.session.add.assert_not_called() diff --git a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py index 8ccc3611b9d..b0886ea8d19 100644 --- a/api/tests/unit_tests/core/agent/test_cot_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_cot_agent_runner.py @@ -42,6 +42,7 @@ def runner(mocker: MockerFixture): application_generate_entity.invoke_from = "test" app_config = MagicMock() + app_config.app_id = "app" app_config.agent = MagicMock() app_config.agent.max_iteration = 1 app_config.prompt_template.simple_prompt_template = "Hello {{name}}" @@ -341,6 +342,7 @@ class TestRun: ) results = list(runner.run(runner.session, message, "query", {})) + assert runner.model_instance.invoke_llm.call_args.kwargs["request_metadata"] == {"app_id": "app"} assert results[-1].delta.message.content == "" def test_run_usage_missing_key_branch(self, runner: DummyRunner, mocker: MockerFixture): diff --git a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py index 244dd8f6c62..f32ded9edff 100644 --- a/api/tests/unit_tests/core/agent/test_fc_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_fc_agent_runner.py @@ -81,6 +81,7 @@ def runner(mocker: MockerFixture): mocker.patch("core.agent.fc_agent_runner.LLMResultChunkDelta", MagicMock) app_config = MagicMock() + app_config.app_id = "app" app_config.agent = MagicMock(max_iteration=2) app_config.prompt_template = MagicMock(simple_prompt_template="system") @@ -299,6 +300,7 @@ class TestRunMethod: outputs = list(runner.run(runner.session, message, "query")) assert len(outputs) == 1 + assert runner.model_instance.invoke_llm.call_args.kwargs["request_metadata"] == {"app_id": "app"} runner.queue_manager.publish.assert_called() queue_calls = runner.queue_manager.publish.call_args_list diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index cda5178e30c..41e14af72de 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -147,7 +147,7 @@ class TestAdvancedChatAppGeneratorInternals: ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), + SimpleNamespace(engine=object(), session=lambda: SimpleNamespace(close=lambda: None)), ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.sessionmaker", lambda **kwargs: SimpleNamespace() diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_config_manager.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_config_manager.py index c2a49d1b3a6..73c53ae98f5 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_config_manager.py @@ -65,6 +65,36 @@ def test_missing_soul_model_leaves_no_model_key(): d = AgentAppConfigManager._synthesize_config_dict(AgentSoulConfig(), None) assert "model" not in d assert d["pre_prompt"] == "" + assert d["file_upload"] == { + "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], + "allowed_file_types": ["document", "image", "audio", "video"], + "allowed_file_upload_methods": ["local_file", "remote_url"], + "enabled": True, + "image": {"enabled": True}, + "number_limits": 3, + } + + +def test_soul_file_upload_overrides_legacy_app_model_config(): + fake_amc = SimpleNamespace( + to_dict=lambda: { + "file_upload": { + "enabled": False, + "image": {"enabled": False}, + }, + } + ) + + d = AgentAppConfigManager._synthesize_config_dict(AgentSoulConfig(), fake_amc) # type: ignore[arg-type] + + assert d["file_upload"] == { + "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], + "allowed_file_types": ["document", "image", "audio", "video"], + "allowed_file_upload_methods": ["local_file", "remote_url"], + "enabled": True, + "image": {"enabled": True}, + "number_limits": 3, + } def test_prompt_type_defaults_to_simple(): diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index 6f8c6af2258..ef030e7e170 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -11,7 +11,6 @@ from __future__ import annotations import contextlib import json -from types import SimpleNamespace import pytest from pytest_mock import MockerFixture @@ -21,7 +20,8 @@ from core.app.apps.agent_app.app_generator import ( AgentAppGeneratorError, ) from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.entities.app_invoke_entities import AGENT_RUNTIME_EXIT_INTENT_ARG, InvokeFrom, UserFrom +from core.workflow.file_reference import build_file_reference from models import Account from models.agent import AgentConfigDraftType @@ -135,6 +135,79 @@ class TestGenerateSuccess: user=user, ) assert generate_entity.call_args.kwargs["prompt_file_mappings"] == file_mappings + assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "suspend" + + def test_generate_uses_delete_exit_intent_from_internal_arg(self, generator, mocker: MockerFixture): + app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") + user = DummyAccount("user") + + generator._resolve_agent = mocker.MagicMock( + return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) + ) + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock( + return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) + ) + generator._handle_response = mocker.MagicMock(return_value="raw-response") + + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + + generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {}, AGENT_RUNTIME_EXIT_INTENT_ARG: "delete"}, + invoke_from=InvokeFrom.DEBUGGER, + streaming=True, + ) + + assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "delete" + + def test_generate_falls_back_to_suspend_for_invalid_internal_exit_intent(self, generator, mocker: MockerFixture): + app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") + user = DummyAccount("user") + + generator._resolve_agent = mocker.MagicMock( + return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) + ) + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock( + return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) + ) + generator._handle_response = mocker.MagicMock(return_value="raw-response") + + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + + generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {}, AGENT_RUNTIME_EXIT_INTENT_ARG: "bogus"}, + invoke_from=InvokeFrom.DEBUGGER, + streaming=True, + ) + + assert generate_entity.call_args.kwargs["agent_runtime_exit_intent"] == "suspend" def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") @@ -205,81 +278,6 @@ class TestGenerateSuccess: assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} - def test_generate_stateless_skips_chat_records(self, generator: AgentAppGenerator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - user = DummyAccount("user") - - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "build-draft-1", "build_draft", mocker.MagicMock()) - ) - generator._init_generate_records = mocker.MagicMock() - run_stateless = mocker.patch.object(generator, "_run_stateless", return_value={"result": "success"}) - converter = mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert") - - result = generator.generate_stateless( - app_model=app_model, - user=user, - args={ - "query": "finalize", - "inputs": {}, - "conversation_id": "debug-conversation-1", - "draft_type": "debug_build", - }, - invoke_from=InvokeFrom.DEBUGGER, - ) - - assert result == {"result": "success"} - generator._init_generate_records.assert_not_called() - converter.assert_not_called() - run_call = run_stateless.call_args.kwargs - assert run_call["conversation_id"] == "debug-conversation-1" - assert run_call["runtime_session_snapshot_id"] == "build-draft-1" - - def test_generate_stateless_requires_conversation_id(self, generator: AgentAppGenerator, mocker: MockerFixture): - with pytest.raises(AgentAppGeneratorError, match="conversation_id is required"): - generator.generate_stateless( - app_model=mocker.MagicMock(), - user=DummyAccount("user"), - args={"query": "finalize", "inputs": {}, "draft_type": "debug_build"}, - invoke_from=InvokeFrom.DEBUGGER, - ) - - def test_stateless_run_uses_agent_app_runner(self, generator: AgentAppGenerator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", app_model_config=mocker.MagicMock()) - user = DummyAccount("user") - agent = mocker.MagicMock(id="agent1") - dify_context = SimpleNamespace(tenant_id="tenant", app_id="app1") - mocker.patch(f"{MODULE}.DifyRunContext", return_value=dify_context) - runner = mocker.MagicMock() - build_runner = mocker.patch.object(generator, "_build_runner", return_value=runner) - - result = generator._run_stateless( - app_model=app_model, - user=user, - invoke_from=InvokeFrom.DEBUGGER, - query="finalize", - conversation_id="debug-conversation-1", - agent=agent, - agent_config_id="build-draft-1", - agent_config_version_kind="build_draft", - agent_soul=mocker.MagicMock(), - runtime_session_snapshot_id="build-draft-1", - ) - - assert result == {"result": "success"} - build_runner.assert_called_once_with(dify_context) - runner.run_stateless.assert_called_once_with( - dify_context=dify_context, - agent_id="agent1", - agent_config_snapshot_id="build-draft-1", - agent_config_version_kind="build_draft", - agent_soul=mocker.ANY, - conversation_id="debug-conversation-1", - query="finalize", - idempotency_key=mocker.ANY, - session_scope_snapshot_id="build-draft-1", - ) - class TestGenerateWorker: @pytest.fixture(autouse=True) @@ -329,6 +327,7 @@ class TestGenerateWorker: query="query", runtime_session_snapshot_id="s", prompt_file_mappings=(), + agent_runtime_exit_intent="suspend", ): generator._generate_worker( flask_app=mocker.MagicMock(), @@ -337,6 +336,7 @@ class TestGenerateWorker: agent_id="a", agent_config_snapshot_id="s", agent_runtime_session_snapshot_id=runtime_session_snapshot_id, + agent_runtime_exit_intent=agent_runtime_exit_intent, model_conf=mocker.MagicMock(model="m"), query=query, prompt_file_mappings=prompt_file_mappings, @@ -364,6 +364,14 @@ class TestGenerateWorker: assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None + def test_worker_forwards_runtime_exit_intent_to_runner(self, generator, mocker: MockerFixture): + runner = self._wire(generator, mocker) + queue_manager = mocker.MagicMock() + + self._call(generator, mocker, queue_manager, agent_runtime_exit_intent="delete") + + assert runner.run.call_args.kwargs["agent_runtime_exit_intent"] == "delete" + def test_worker_appends_prompt_files_to_backend_query(self, generator, mocker: MockerFixture): runner = self._wire(generator, mocker, guard_query="你看得见这张图片吗") queue_manager = mocker.MagicMock() @@ -373,7 +381,23 @@ class TestGenerateWorker: "transfer_method": "local_file", "url": "", "upload_file_id": "upload-file-1", - } + }, + { + "type": "document", + "transfer_method": "remote_url", + "url": "https://example.com/source.pdf", + "upload_file_id": "ignored", + }, + ] + expected_file_mappings = [ + { + "transfer_method": "local_file", + "reference": build_file_reference(record_id="upload-file-1"), + }, + { + "transfer_method": "remote_url", + "url": "https://example.com/source.pdf", + }, ] self._call( @@ -385,7 +409,10 @@ class TestGenerateWorker: ) assert runner.run.call_args.kwargs["query"] == ( - f"你看得见这张图片吗\n{json.dumps(file_mappings, ensure_ascii=False)}" + "你看得见这张图片吗\nUser provided files: " + "use dify-agent file download with the listed transfer_method and reference/url " + "to get the files and investigate them\n" + f"{json.dumps(expected_file_mappings, ensure_ascii=False, separators=(',', ':'))}" ) def test_input_guard_short_circuit_skips_backend(self, generator, mocker: MockerFixture): @@ -461,6 +488,7 @@ class TestResumeAfterFormSubmission: # The paused turn's query is re-sent verbatim — never blank. assert entity.call_args.kwargs["query"] == "original question" + assert "agent_runtime_exit_intent" not in entity.call_args.kwargs def test_resume_falls_back_to_placeholder_when_no_paused_message(self, generator, mocker: MockerFixture): entity = self._wire(generator, mocker) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py index 8ebf68c7061..ffef1cf3be8 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py @@ -28,8 +28,6 @@ from pydantic_ai.messages import ( FunctionToolCallEvent, FunctionToolResultEvent, PartDeltaEvent, - PartStartEvent, - TextPart, TextPartDelta, ThinkingPartDelta, ToolCallPart, @@ -49,7 +47,12 @@ from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeReque from core.app.apps.agent_app.session_store import AgentAppSessionScope, StoredAgentAppSession from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueLLMChunkEvent, QueueMessageEndEvent +from core.app.entities.queue_entities import ( + QueueAgentMessageEvent, + QueueAgentThoughtEvent, + QueueLLMChunkEvent, + QueueMessageEndEvent, +) from core.workflow.nodes.agent_v2.ask_human_resume import AskHumanResumeOutcome from models.agent_config_entities import AgentSoulConfig from models.model import MessageAgentThought @@ -98,24 +101,6 @@ class _RecordingFakeAgentBackendRunClient(FakeAgentBackendRunClient): return super().cancel_run(run_id, request=request) -class _BlockingRecordingFakeAgentBackendRunClient(FakeAgentBackendRunClient): - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.wait_calls: list[tuple[str, float | None]] = [] - self.stream_called = False - - @override - def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]: - del run_id, after - self.stream_called = True - return iter(()) - - @override - def wait_run(self, run_id: str, *, timeout_seconds: float | None = None): - self.wait_calls.append((run_id, timeout_seconds)) - return super().wait_run(run_id, timeout_seconds=timeout_seconds) - - class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]: @@ -127,12 +112,14 @@ class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")), + agent_message_delta="hello ", ) yield PydanticAIStreamRunEvent( id="3-0", run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="agent")), + agent_message_delta="agent", ) yield RunSucceededEvent( id="4-0", @@ -157,12 +144,14 @@ class _StreamingRecordingFakeAgentBackendRunClient(_RecordingFakeAgentBackendRun run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")), + agent_message_delta="hello ", ) yield PydanticAIStreamRunEvent( id="3-0", run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="agent")), + agent_message_delta="agent", ) yield RunSucceededEvent( id="4-0", @@ -190,6 +179,7 @@ class _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(_RecordingFakeAgent run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")), + agent_message_delta="hello ", ) self._queue_manager.request_stop() yield PydanticAIStreamRunEvent( @@ -197,10 +187,11 @@ class _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(_RecordingFakeAgent run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="agent")), + agent_message_delta="agent", ) -class _StreamingPartStartFakeAgentBackendRunClient(FakeAgentBackendRunClient): +class _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]: del after @@ -210,7 +201,8 @@ class _StreamingPartStartFakeAgentBackendRunClient(FakeAgentBackendRunClient): id="2-0", run_id=run_id, created_at=created_at, - data=PartStartEvent(index=0, part=TextPart(content="hello")), + data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello")), + agent_message_delta="hello", ) yield RunSucceededEvent( id="3-0", @@ -251,6 +243,7 @@ class _StreamingTextNullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClien run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="streamed answer")), + agent_message_delta="streamed answer", ) yield RunSucceededEvent( id="3-0", @@ -263,6 +256,37 @@ class _StreamingTextNullOutputFakeAgentBackendRunClient(FakeAgentBackendRunClien ) +class _AgentAnswerStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): + @override + def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]: + del after + created_at = datetime(2026, 1, 1, tzinfo=UTC) + yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at) + yield PydanticAIStreamRunEvent( + id="2-0", + run_id=run_id, + created_at=created_at, + data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="hello ")), + agent_message_delta="hello ", + ) + yield PydanticAIStreamRunEvent( + id="3-0", + run_id=run_id, + created_at=created_at, + data=PartDeltaEvent(index=0, delta=TextPartDelta(content_delta="agent")), + agent_message_delta="agent", + ) + yield RunSucceededEvent( + id="4-0", + run_id=run_id, + created_at=created_at, + data=RunSucceededEventData( + output={"text": "final answer"}, + session_snapshot=CompositorSessionSnapshot(layers=[]), + ), + ) + + class _ProcessStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override def stream_events(self, run_id: str, *, after: str | None = None) -> Iterator[RunEvent]: @@ -292,6 +316,7 @@ class _ProcessStreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): run_id=run_id, created_at=created_at, data=PartDeltaEvent(index=1, delta=TextPartDelta(content_delta="final answer")), + agent_message_delta="final answer", ) yield RunSucceededEvent( id="6-0", @@ -318,6 +343,9 @@ class _FakeDbSession: def get(self, _model: type[MessageAgentThought], row_id: str) -> MessageAgentThought | None: return self.rows.get(row_id) + def delete(self, row: MessageAgentThought) -> None: + self.rows.pop(str(row.id), None) + def rollback(self) -> None: self.rollback_count += 1 @@ -327,9 +355,11 @@ class _FakeSessionStore: self, loaded: CompositorSessionSnapshot | None = None, loaded_session: StoredAgentAppSession | None = None, + listed_sessions: list[StoredAgentAppSession] | None = None, ) -> None: self.loaded = loaded self._loaded_session = loaded_session + self._listed_sessions = list(listed_sessions or []) self.loaded_scopes: list[AgentAppSessionScope] = [] self.saved: list[ tuple[ @@ -341,6 +371,7 @@ class _FakeSessionStore: str | None, ] ] = [] + self.cleaned: list[tuple[AgentAppSessionScope, str | None]] = [] def load_active_snapshot(self, scope: AgentAppSessionScope) -> CompositorSessionSnapshot | None: self.loaded_scopes.append(scope) @@ -354,6 +385,14 @@ class _FakeSessionStore: return None return StoredAgentAppSession(scope=scope, session_snapshot=self.loaded, backend_run_id=None) + def list_active_sessions_for_conversation( + self, *, tenant_id: str, app_id: str, conversation_id: str + ) -> list[StoredAgentAppSession]: + assert tenant_id == "tenant-1" + assert app_id == "app-1" + assert conversation_id == "conv-1" + return list(self._listed_sessions) + def save_active_snapshot( self, *, @@ -368,6 +407,9 @@ class _FakeSessionStore: (scope, backend_run_id, snapshot, list(runtime_layer_specs), pending_form_id, pending_tool_call_id) ) + def mark_cleaned(self, *, scope: AgentAppSessionScope, backend_run_id: str | None = None) -> None: + self.cleaned.append((scope, backend_run_id)) + class _MonotonicClock: def __init__(self, *values: float) -> None: @@ -409,7 +451,7 @@ def _runner( client: FakeAgentBackendRunClient, store: _FakeSessionStore, *, - text_delta_debounce_seconds: float | None = 0, + text_delta_debounce_seconds: float = 0, ) -> AgentAppRunner: return AgentAppRunner( request_builder=AgentAppRuntimeRequestBuilder( @@ -423,7 +465,7 @@ def _runner( ) -def _run(runner: AgentAppRunner, qm: _FakeQueueManager) -> None: +def _run(runner: AgentAppRunner, qm: _FakeQueueManager, *, agent_runtime_exit_intent: str = "suspend") -> None: runner.run( dify_context=_dify_ctx(), agent_id="agent-1", @@ -434,18 +476,7 @@ def _run(runner: AgentAppRunner, qm: _FakeQueueManager) -> None: message_id="msg-1", model_name="gpt-4o-mini", queue_manager=qm, # type: ignore[arg-type] - ) - - -def _run_stateless(runner: AgentAppRunner) -> None: - runner.run_stateless( - dify_context=_dify_ctx(), - agent_id="agent-1", - agent_config_snapshot_id="snap-1", - agent_soul=_soul(), - conversation_id="conv-1", - query="finalize", - idempotency_key="run-req-1", + agent_runtime_exit_intent=agent_runtime_exit_intent, # type: ignore[arg-type] ) @@ -470,6 +501,8 @@ def test_successful_turn_publishes_chunk_and_message_end_and_saves_session(): _run(_runner(client, store), qm) + assert client.request is not None + assert client.request.on_exit.default.value == "suspend" # One LLM chunk + one message-end, carrying the backend's answer text. chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] @@ -490,38 +523,169 @@ def test_successful_turn_publishes_chunk_and_message_end_and_saves_session(): # A successful turn carries no ask_human pause correlation. assert pending_form_id is None assert pending_tool_call_id is None + assert store.cleaned == [] -def test_successful_turn_forwards_agent_backend_stream_text_deltas_without_duplicate_terminal_chunk(): +def test_successful_turn_enqueues_cleanup_for_superseded_sessions_after_saving_snapshot(monkeypatch): + superseded = StoredAgentAppSession( + scope=AgentAppSessionScope( + tenant_id="tenant-1", + app_id="app-1", + conversation_id="conv-1", + agent_id="agent-2", + agent_config_snapshot_id="snap-2", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-old", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + current_scope_session = StoredAgentAppSession( + scope=AgentAppSessionScope( + tenant_id="tenant-1", + app_id="app-1", + conversation_id="conv-1", + agent_id="agent-1", + agent_config_snapshot_id="snap-1", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-current", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + store = _FakeSessionStore(listed_sessions=[current_scope_session, superseded]) + client = FakeAgentBackendRunClient() + qm = _FakeQueueManager() + cleanup_delay = MagicMock() + monkeypatch.setattr(app_runner_module.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + + _run(_runner(client, store), qm) + + assert store.saved + cleanup_delay.assert_called_once() + payload = cleanup_delay.call_args.args[0] + assert payload["metadata"]["conversation_id"] == "conv-1" + assert payload["metadata"]["agent_id"] == "agent-2" + assert payload["metadata"]["previous_agent_backend_run_id"] == "run-old" + assert payload["idempotency_key"] == "tenant-1:app-1:conv-1:agent-2:snap-2:superseded-session-cleanup:run-old" + + +def test_superseded_session_cleanup_enqueue_failure_does_not_fail_turn(monkeypatch): + superseded = StoredAgentAppSession( + scope=AgentAppSessionScope( + tenant_id="tenant-1", + app_id="app-1", + conversation_id="conv-1", + agent_id="agent-2", + agent_config_snapshot_id="snap-2", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-old", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + store = _FakeSessionStore(listed_sessions=[superseded]) + client = FakeAgentBackendRunClient() + qm = _FakeQueueManager() + cleanup_delay = MagicMock(side_effect=RuntimeError("queue down")) + monkeypatch.setattr(app_runner_module.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + + _run(_runner(client, store), qm) + + cleanup_delay.assert_called_once() + end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] + assert len(end_events) == 1 + assert end_events[0].llm_result.message.content == "hello agent" + + +def test_delete_on_exit_turn_marks_session_cleaned_without_saving_snapshot(): + client = _StreamingRecordingFakeAgentBackendRunClient() + store = _FakeSessionStore() + qm = _FakeQueueManager() + + _run(_runner(client, store), qm, agent_runtime_exit_intent="delete") + + assert client.request is not None + assert client.request.on_exit.default.value == "delete" + assert store.saved == [] + assert len(store.cleaned) == 1 + cleaned_scope, cleaned_run_id = store.cleaned[0] + assert cleaned_scope.conversation_id == "conv-1" + assert cleaned_scope.agent_config_snapshot_id == "snap-1" + assert cleaned_run_id == "fake-run-1" + end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] + assert len(end_events) == 1 + assert end_events[0].llm_result.message.content == "hello agent" + + +def test_delete_on_exit_turn_swallows_cleanup_failure_after_success(): + client = _StreamingRecordingFakeAgentBackendRunClient() + store = _FakeSessionStore() + store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] + qm = _FakeQueueManager() + + _run(_runner(client, store), qm, agent_runtime_exit_intent="delete") + + assert store.saved == [] + store.mark_cleaned.assert_called_once() + end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] + assert len(end_events) == 1 + + +def test_delete_on_exit_turn_marks_session_cleaned_when_publish_fails(): + client = _StreamingRecordingFakeAgentBackendRunClient() + store = _FakeSessionStore() + store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] + qm = _FakeQueueManager() + runner = _runner(client, store) + runner._publish_terminal_answer = MagicMock(side_effect=RuntimeError("publish failed")) + + with pytest.raises(RuntimeError, match="publish failed"): + _run(runner, qm, agent_runtime_exit_intent="delete") + + assert store.saved == [] + store.mark_cleaned.assert_called_once() + + +def test_successful_turn_routes_stream_text_to_agent_message_and_uses_terminal_output(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) client = _StreamingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store, text_delta_debounce_seconds=0), qm) + _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["hello ", "agent"] + assert [event.chunk.delta.message.content for event in chunk_events] == ["hello agent"] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["hello ", "agent"] assert len(end_events) == 1 assert end_events[0].llm_result.message.content == "hello agent" assert end_events[0].llm_result.usage.prompt_tokens == 3 assert end_events[0].llm_result.usage.completion_tokens == 5 assert end_events[0].llm_result.usage.total_tokens == 8 + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert rows == [] assert store.saved -def test_successful_turn_forwards_part_start_text_and_publishes_missing_terminal_suffix(): - client = _StreamingPartStartFakeAgentBackendRunClient() +def test_successful_turn_routes_single_agent_message_delta(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + client = _StreamingSingleAgentMessageDeltaFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store, text_delta_debounce_seconds=0), qm) + _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["hello", " agent"] + assert [event.chunk.delta.message.content for event in chunk_events] == ["hello agent"] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["hello"] assert len(end_events) == 1 assert end_events[0].llm_result.message.content == "hello agent" + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert rows == [] def test_successful_turn_with_null_terminal_output_publishes_empty_answer_not_literal_null(): @@ -532,24 +696,77 @@ def test_successful_turn_with_null_terminal_output_publishes_empty_answer_not_li _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] assert chunk_events == [] + assert agent_message_events == [] assert len(end_events) == 1 assert end_events[0].llm_result.message.content == "" -def test_successful_turn_with_streamed_text_and_null_terminal_output_keeps_streamed_answer(): +def test_successful_turn_with_stream_text_and_null_terminal_output_keeps_empty_message(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) client = _StreamingTextNullOutputFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store, text_delta_debounce_seconds=0), qm) + _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["streamed answer"] + assert chunk_events == [] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["streamed answer"] assert len(end_events) == 1 - assert end_events[0].llm_result.message.content == "streamed answer" + assert end_events[0].llm_result.message.content == "" + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert len(rows) == 1 + assert rows[0].answer == "streamed answer" + + +def test_successful_turn_routes_agent_answer_to_agent_message(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + client = _AgentAnswerStreamingFakeAgentBackendRunClient() + store = _FakeSessionStore() + qm = _FakeQueueManager() + + _run(_runner(client, store), qm) + + chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] + assert [event.chunk.delta.message.content for event in chunk_events] == ["final answer"] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["hello ", "agent"] + end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] + assert len(end_events) == 1 + assert end_events[0].llm_result.message.content == "final answer" + thought_events = [e for e in qm.events if isinstance(e, QueueAgentThoughtEvent)] + assert len(thought_events) == 2 + + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert len(rows) == 1 + assert rows[0].answer == "hello agent" + assert rows[0].thought == "" + assert rows[0].tool == "" + + +def test_agent_message_deltas_are_debounced_to_agent_message(monkeypatch): + monkeypatch.setattr(app_runner_module.time, "monotonic", _MonotonicClock(0.0, 0.2)) + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + client = _StreamingFakeAgentBackendRunClient() + store = _FakeSessionStore() + qm = _FakeQueueManager() + + _run(_runner(client, store, text_delta_debounce_seconds=0.5), qm) + + chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] + assert [event.chunk.delta.message.content for event in chunk_events] == ["hello agent"] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["hello agent"] + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert rows == [] def test_successful_turn_persists_thinking_and_tool_process_events(monkeypatch): @@ -559,62 +776,41 @@ def test_successful_turn_persists_thinking_and_tool_process_events(monkeypatch): store = _FakeSessionStore() qm = _FakeQueueManager() - _run(_runner(client, store, text_delta_debounce_seconds=0), qm) + _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] assert [event.chunk.delta.message.content for event in chunk_events] == ["final answer"] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["final answer"] thought_events = [e for e in qm.events if isinstance(e, QueueAgentThoughtEvent)] assert len(thought_events) >= 3 rows = sorted(fake_session.rows.values(), key=lambda row: row.position) assert rows[0].thought == "I need to inspect the file." - assert rows[0].tool is None + assert rows[0].tool == "" assert rows[1].tool == "bash" assert rows[1].tool_input == '{"cmd": "ls"}' assert rows[1].observation == "ok" + assert len(rows) == 2 -def test_streaming_turn_batches_text_deltas_within_debounce_window(monkeypatch): - monkeypatch.setattr(app_runner_module.time, "monotonic", _MonotonicClock(0.0, 0.2)) - client = _StreamingFakeAgentBackendRunClient() - store = _FakeSessionStore() - qm = _FakeQueueManager() - - _run(_runner(client, store, text_delta_debounce_seconds=0.5), qm) - - chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] - end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["hello agent"] - assert len(end_events) == 1 - assert end_events[0].llm_result.message.content == "hello agent" - - -def test_streaming_turn_flushes_pending_text_before_terminal_success(monkeypatch): - monkeypatch.setattr(app_runner_module.time, "monotonic", _MonotonicClock(0.0)) - client = _StreamingPartStartFakeAgentBackendRunClient() - store = _FakeSessionStore() - qm = _FakeQueueManager() - - _run(_runner(client, store, text_delta_debounce_seconds=0.5), qm) - - chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] - end_events = [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["hello", " agent"] - assert len(end_events) == 1 - assert end_events[0].llm_result.message.content == "hello agent" - - -def test_streaming_turn_flushes_pending_text_before_stop_and_cancel(monkeypatch): - monkeypatch.setattr(app_runner_module.time, "monotonic", _MonotonicClock(0.0)) +def test_streaming_turn_cancels_after_persisting_seen_agent_answer(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) store = _FakeSessionStore() qm = _FakeQueueManager() client = _StreamingStopAfterFirstDeltaFakeAgentBackendRunClient(queue_manager=qm) with pytest.raises(GenerateTaskStoppedError): - _run(_runner(client, store, text_delta_debounce_seconds=0.5), qm) + _run(_runner(client, store), qm) chunk_events = [e for e in qm.events if isinstance(e, QueueLLMChunkEvent)] - assert [event.chunk.delta.message.content for event in chunk_events] == ["hello "] + agent_message_events = [e for e in qm.events if isinstance(e, QueueAgentMessageEvent)] + assert chunk_events == [] + assert [event.chunk.delta.message.content for event in agent_message_events] == ["hello "] + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert len(rows) == 1 + assert rows[0].answer == "hello " assert client.cancelled_run_ids == ["fake-run-1"] @@ -656,12 +852,149 @@ def test_tool_result_without_identity_does_not_attach_to_previous_tool(monkeypat assert len(rows) == 2 assert rows[0].tool == "shell_run" assert rows[0].tool_input == '{"script": "npx skills find browser"}' - assert rows[0].observation is None - assert rows[1].tool is None - assert rows[1].tool_input is None + assert rows[0].observation == "" + assert rows[1].tool == "" + assert rows[1].tool_input == "" assert rows[1].observation == "Knowledge base search results: browser skill" +def test_answer_suffix_trim_keeps_non_terminal_prefix(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + qm = _FakeQueueManager() + recorder = app_runner_module._AgentProcessRecorder( + dify_context=_dify_ctx(), + message_id="msg-1", + queue_manager=qm, # type: ignore[arg-type] + ) + + recorder.append_answer_text("intermediate final answer") + recorder.trim_answer_suffix("final answer") + + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert len(rows) == 1 + assert rows[0].answer == "intermediate " + + +def test_tool_call_part_binds_late_call_id_to_delta_row(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + qm = _FakeQueueManager() + recorder = app_runner_module._AgentProcessRecorder( + dify_context=_dify_ctx(), + message_id="msg-1", + queue_manager=qm, # type: ignore[arg-type] + ) + + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "part_delta", + "index": 0, + "delta": { + "part_delta_kind": "tool_call", + "tool_name_delta": "knowledge_base_search", + "args_delta": {"query": "browser"}, + }, + }, + ) + ) + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "part_start", + "index": 0, + "part": { + "part_kind": "tool-call", + "tool_name": "knowledge_base_search", + "args": {"query": "browser"}, + "tool_call_id": "tool-call-1", + }, + }, + ) + ) + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "function_tool_result", + "part": { + "part_kind": "tool-return", + "tool_name": "knowledge_base_search", + "content": "Knowledge base search results: browser skill", + "tool_call_id": "tool-call-1", + }, + }, + ) + ) + + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert len(rows) == 1 + assert rows[0].tool == "knowledge_base_search" + assert rows[0].tool_input == '{"query": "browser"}' + assert rows[0].observation == "Knowledge base search results: browser skill" + + +def test_thinking_after_tool_starts_new_snapshot_row(monkeypatch): + fake_session = _FakeDbSession() + monkeypatch.setattr(app_runner_module.db, "session", fake_session) + qm = _FakeQueueManager() + recorder = app_runner_module._AgentProcessRecorder( + dify_context=_dify_ctx(), + message_id="msg-1", + queue_manager=qm, # type: ignore[arg-type] + ) + + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "part_delta", + "index": 0, + "delta": { + "part_delta_kind": "thinking", + "content_delta": "The first thought.", + }, + }, + ) + ) + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "function_tool_call", + "part": { + "part_kind": "tool-call", + "tool_name": "shell_run", + "args": {"cmd": "date"}, + "tool_call_id": "tool-call-1", + }, + }, + ) + ) + recorder.handle_stream_event( + AgentBackendStreamInternalEvent( + run_id="run-1", + data={ + "event_kind": "part_delta", + "index": 0, + "delta": { + "part_delta_kind": "thinking", + "content_delta": "The next thought.", + }, + }, + ) + ) + + rows = sorted(fake_session.rows.values(), key=lambda row: row.position) + assert [row.thought for row in rows] == ["The first thought.", "", "The next thought."] + assert rows[0].id != rows[2].id + assert rows[1].tool == "shell_run" + assert rows[1].tool_input == '{"cmd": "date"}' + + def test_tool_result_without_call_id_matches_unique_open_tool_name(monkeypatch): fake_session = _FakeDbSession() monkeypatch.setattr(app_runner_module.db, "session", fake_session) @@ -743,31 +1076,6 @@ def test_debug_session_scope_can_reuse_conversation_across_config_snapshots(): assert store.saved[0][0].agent_config_snapshot_id is None -def test_stateless_run_uses_bounded_wait_and_does_not_save_session_state(): - prior = CompositorSessionSnapshot(layers=[]) - client = _BlockingRecordingFakeAgentBackendRunClient() - store = _FakeSessionStore(loaded=prior) - - _run_stateless(_runner(client, store)) - - assert client.request is not None - assert client.request.session_snapshot is prior - assert client.wait_calls == [("fake-run-1", app_runner_module.dify_config.APP_MAX_EXECUTION_TIME)] - assert client.stream_called is False - assert store.saved == [] - - -def test_stateless_run_raises_backend_error_on_failed_bounded_wait(): - client = _BlockingRecordingFakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.FAILED) - store = _FakeSessionStore() - - with pytest.raises(AgentBackendError): - _run_stateless(_runner(client, store)) - - assert client.wait_calls == [("fake-run-1", app_runner_module.dify_config.APP_MAX_EXECUTION_TIME)] - assert store.saved == [] - - def test_failed_run_raises_agent_backend_error(): client = FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.FAILED) store = _FakeSessionStore() @@ -792,11 +1100,11 @@ def test_stopped_task_cancels_agent_backend_run_and_skips_session_save(): assert store.saved == [] -def test_extract_answer_handles_plain_string_and_dict(): - assert AgentAppRunner._extract_answer(None) == "" - assert AgentAppRunner._extract_answer("plain text") == "plain text" - assert AgentAppRunner._extract_answer({"text": "hi"}) == "hi" - assert AgentAppRunner._extract_answer({"a": 1}) == '{"a": 1}' +def test_terminal_output_to_answer_handles_plain_string_and_dict(): + assert AgentAppRunner._terminal_output_to_answer(None) == "" + assert AgentAppRunner._terminal_output_to_answer("plain text") == "plain text" + assert AgentAppRunner._terminal_output_to_answer({"text": "hi"}) == "hi" + assert AgentAppRunner._terminal_output_to_answer({"a": 1}) == '{"a": 1}' def test_ask_human_pauses_turn_creates_form_and_persists_correlation(): @@ -827,6 +1135,22 @@ def test_ask_human_pauses_turn_creates_form_and_persists_correlation(): assert store.saved[0][5] == "fake-ask-human-1" +def test_delete_on_exit_deferred_tool_marks_session_cleaned_and_raises_error(): + client = FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.PAUSED) + store = _FakeSessionStore() + store.mark_cleaned = MagicMock(side_effect=RuntimeError("cleanup failed")) # type: ignore[method-assign] + qm = _FakeQueueManager() + runner = _runner(client, store) + runner._pause_for_ask_human = MagicMock() + + with pytest.raises(AgentBackendError, match="finalization cannot pause for human input"): + _run(runner, qm, agent_runtime_exit_intent="delete") + + runner._pause_for_ask_human.assert_not_called() + assert store.saved == [] + store.mark_cleaned.assert_called_once() + + def test_submitted_form_resumes_turn_with_deferred_tool_results(monkeypatch): # ENG-638: a turn that runs while a pending form is answered threads the # human's reply into the request as deferred_tool_results. diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py b/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py index 806741bb8fe..ec0ce7d23b1 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py @@ -13,7 +13,7 @@ from typing import Any import pytest from core.app.apps.agent_app import app_generator as gen_mod -from core.app.apps.agent_app.app_generator import AgentAppGenerator, AgentAppGeneratorError +from core.app.apps.agent_app.app_generator import AgentAppGenerator, AgentAppGeneratorError, AgentAppNotPublishedError from core.app.entities.app_invoke_entities import InvokeFrom _SOUL_DICT = { @@ -78,7 +78,7 @@ class TestResolveAgentById: class TestResolveAgent: def test_success_chains_to_resolve_by_id(self, monkeypatch: pytest.MonkeyPatch): - bound_agent = SimpleNamespace(id="agent-1", active_config_snapshot_id="snap-1") + bound_agent = SimpleNamespace(id="agent-1", active_config_snapshot_id="snap-1", active_config_is_published=True) inner_agent = SimpleNamespace(id="agent-1") snapshot = _snapshot() # scalar order: bound agent (in _resolve_agent), then agent + snapshot (in _resolve_agent_by_id) @@ -97,6 +97,46 @@ class TestResolveAgent: assert config_version_kind == "snapshot" assert soul.model is not None + def test_unpublished_draft_still_resolves_active_snapshot(self, monkeypatch: pytest.MonkeyPatch): + bound_agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id="snap-1", + active_config_is_published=False, + ) + inner_agent = SimpleNamespace(id="agent-1") + snapshot = _snapshot() + _patch_session(monkeypatch, [bound_agent, inner_agent, snapshot]) + app_model = SimpleNamespace(id="app-1", tenant_id="t1") + + agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=SimpleNamespace(id="user-1"), + ) # type: ignore[arg-type] + + assert agent is bound_agent + assert config_id == snapshot.id + assert config_version_kind == "snapshot" + assert soul.prompt.system_prompt == "You are Iris." + + def test_agent_without_active_snapshot_raises_before_model_resolution(self, monkeypatch: pytest.MonkeyPatch): + bound_agent = SimpleNamespace( + id="agent-1", + active_config_snapshot_id=None, + active_config_is_published=False, + ) + _patch_session(monkeypatch, [bound_agent]) + app_model = SimpleNamespace(id="app-1", tenant_id="t1") + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=SimpleNamespace(id="user-1"), + ) # type: ignore[arg-type] + def test_unbound_app_raises(self, monkeypatch: pytest.MonkeyPatch): _patch_session(monkeypatch, [None]) app_model = SimpleNamespace(id="app-1", tenant_id="t1") diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index 9828216f3b8..b3ea05b8a87 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -57,7 +57,6 @@ class TestBuildForAgentApp: "llm", ] assert "workflow_node_job_prompt" not in names - assert request.purpose == "agent_app" # Agent App keeps layers alive across turns by default. assert request.on_exit.default.value == "suspend" @@ -143,6 +142,7 @@ def _ctx( *, query: str = "hello", agent_config_version_kind: str = "snapshot", + suspend_on_exit: bool = True, ) -> AgentAppRuntimeBuildContext: dify_context = SimpleNamespace( tenant_id="tenant-1", @@ -160,6 +160,7 @@ def _ctx( user_query=query, idempotency_key="msg-1", agent_config_version_kind=agent_config_version_kind, # type: ignore[arg-type] + suspend_on_exit=suspend_on_exit, ) @@ -185,7 +186,6 @@ class TestAgentAppRuntimeRequestBuilder: result = builder.build(_ctx(_soul_with_model())) req = result.request - assert req.purpose == "agent_app" names = [layer.name for layer in req.composition.layers] assert names == [ "agent_soul_prompt", @@ -207,10 +207,52 @@ class TestAgentAppRuntimeRequestBuilder: assert exec_ctx.config.user_from == "end-user" assert exec_ctx.config.invoke_from == "web-app" assert exec_ctx.config.agent_mode == "agent_app" + assert req.on_exit.default.value == "suspend" # credentials are redacted in the log-safe view. assert result.redacted_request["composition"]["layers"][-1]["config"]["credentials"] == "[REDACTED]" assert result.metadata["conversation_id"] == "conv-1" + def test_build_wraps_agent_soul_prompt_for_build_draft(self): + builder = AgentAppRuntimeRequestBuilder( + credentials_provider=_FakeCredentialsProvider(), + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + + result = builder.build(_ctx(_soul_with_model(), agent_config_version_kind="build_draft")) + + prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt") + assert prompt_layer.config.prefix != _soul_with_model().prompt + assert prompt_layer.config.prefix.startswith("You are running in build mode.") + assert "```text\nYou are Iris.\n```" in prompt_layer.config.prefix + + def test_build_propagates_draft_version_kind_without_wrapping_prompt(self): + builder = AgentAppRuntimeRequestBuilder( + credentials_provider=_FakeCredentialsProvider(), + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + + result = builder.build(_ctx(_soul_with_model(), agent_config_version_kind="draft")) + + prompt_layer = next(layer for layer in result.request.composition.layers if layer.name == "agent_soul_prompt") + execution_context = next( + layer for layer in result.request.composition.layers if layer.name == "execution_context" + ) + config_layer = next(layer for layer in result.request.composition.layers if layer.name == DIFY_CONFIG_LAYER_ID) + + assert prompt_layer.config.prefix == "You are Iris." + assert execution_context.config.agent_config_version_kind == "draft" + assert config_layer.config.config_version.kind == "draft" + + def test_build_uses_delete_on_exit_when_requested(self): + builder = AgentAppRuntimeRequestBuilder( + credentials_provider=_FakeCredentialsProvider(), + dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] + ) + + result = builder.build(_ctx(_soul_with_model(), suspend_on_exit=False)) + + assert result.request.on_exit.default.value == "delete" + def test_build_includes_plugin_tools_layer_returned_by_injected_builder_for_draft(self): soul = _soul_with_model() soul.tools.dify_tools = [ diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py index 458b0a26dbe..fc9bf1b48e6 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_session_store.py @@ -303,3 +303,46 @@ def test_load_active_session_for_conversation_isolates_other_conversations(): store.load_active_session_for_conversation(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-A") is not None ) + + +def test_list_active_sessions_for_conversation_returns_all_active_rows(): + store = AgentAppRuntimeSessionStore() + with session_factory.create_session() as session: + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + agent_config_snapshot_id="snap-1", + conversation_id="conv-1", + backend_run_id="run-1", + session_snapshot=_snapshot(messages=1).model_dump_json(), + composition_layer_specs='[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-2", + agent_config_snapshot_id="snap-2", + conversation_id="conv-1", + backend_run_id="run-2", + session_snapshot=_snapshot(messages=2).model_dump_json(), + composition_layer_specs='[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},"metadata":{},"config":null}]', + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.commit() + + loaded = store.list_active_sessions_for_conversation( + tenant_id="tenant-1", + app_id="app-1", + conversation_id="conv-1", + ) + + assert [session.scope.agent_id for session in loaded] == ["agent-2", "agent-1"] + assert all(session.scope.conversation_id == "conv-1" for session in loaded) diff --git a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py b/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py index b3ea1a464f8..130264972a3 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py @@ -5,7 +5,6 @@ from uuid import uuid4 import pytest -from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.base_app_runner import AppRunner from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueMessageFileEvent @@ -81,59 +80,55 @@ class TestBaseAppRunnerMultimodal: # Setup mock message file mock_msg_file_class.return_value = mock_message_file - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - mock_session.add = MagicMock() - mock_session.commit = MagicMock() - mock_session.refresh = MagicMock() + file_session = MagicMock() + mock_session_factory = MagicMock() + mock_session_factory.begin.return_value.__enter__ = MagicMock(return_value=file_session) + mock_session_factory.begin.return_value.__exit__ = MagicMock(return_value=False) - # Act - # Create a mock runner with the method bound - runner = MagicMock() + with patch("core.app.apps.base_app_runner.sessionmaker", return_value=mock_session_factory) as mock_sm: + with patch("core.app.apps.base_app_runner.db") as mock_db: + # Act + runner = MagicMock() + method = AppRunner._handle_multimodal_image_content + runner._handle_multimodal_image_content = lambda *args, **kwargs: method( + runner, *args, **kwargs + ) - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + runner._handle_multimodal_image_content( + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - runner._handle_multimodal_image_content( - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) + # Assert + mock_mgr.create_file_by_url.assert_called_once_with( + user_id=mock_user_id, + tenant_id=mock_tenant_id, + file_url=image_url, + conversation_id=None, + ) - # Assert - # Verify tool file was created from URL - mock_mgr.create_file_by_url.assert_called_once_with( - user_id=mock_user_id, - tenant_id=mock_tenant_id, - file_url=image_url, - conversation_id=None, - ) + mock_msg_file_class.assert_called_once() + call_kwargs = mock_msg_file_class.call_args[1] + assert call_kwargs["message_id"] == mock_message_id + assert call_kwargs["type"] == FileType.IMAGE + assert call_kwargs["transfer_method"] == FileTransferMethod.TOOL_FILE + assert call_kwargs["belongs_to"] == "assistant" + assert call_kwargs["created_by"] == mock_user_id - # Verify message file was created with correct parameters - mock_msg_file_class.assert_called_once() - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["message_id"] == mock_message_id - assert call_kwargs["type"] == FileType.IMAGE - assert call_kwargs["transfer_method"] == FileTransferMethod.TOOL_FILE - assert call_kwargs["belongs_to"] == "assistant" - assert call_kwargs["created_by"] == mock_user_id + # Verify independent session was used (not db.session) + mock_sm.assert_called_once_with(bind=mock_db.engine, expire_on_commit=False) + file_session.add.assert_called_once_with(mock_message_file) + mock_db.session.commit.assert_not_called() + mock_db.session.close.assert_not_called() - # Verify database operations - mock_session.add.assert_called_once_with(mock_message_file) - mock_session.commit.assert_called_once() - mock_session.refresh.assert_called_once_with(mock_message_file) - - # Verify event was published - mock_queue_manager.publish.assert_called_once() - publish_call = mock_queue_manager.publish.call_args - assert isinstance(publish_call[0][0], QueueMessageFileEvent) - assert publish_call[0][0].message_file_id == mock_message_file.id - # publish_from might be passed as positional or keyword argument - assert ( - publish_call[0][1] == PublishFrom.APPLICATION_MANAGER - or publish_call.kwargs.get("publish_from") == PublishFrom.APPLICATION_MANAGER - ) + # Verify event was published + mock_queue_manager.publish.assert_called_once() + publish_call = mock_queue_manager.publish.call_args + assert isinstance(publish_call[0][0], QueueMessageFileEvent) + assert publish_call[0][0].message_file_id == mock_message_file.id def test_handle_multimodal_image_content_with_base64( self, @@ -165,50 +160,44 @@ class TestBaseAppRunnerMultimodal: mock_mgr_class.return_value = mock_mgr with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - # Setup mock message file mock_msg_file_class.return_value = mock_message_file - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - mock_session.add = MagicMock() - mock_session.commit = MagicMock() - mock_session.refresh = MagicMock() + file_session = MagicMock() + mock_session_factory = MagicMock() + mock_session_factory.begin.return_value.__enter__ = MagicMock(return_value=file_session) + mock_session_factory.begin.return_value.__exit__ = MagicMock(return_value=False) - # Act - # Create a mock runner with the method bound - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + with patch("core.app.apps.base_app_runner.sessionmaker", return_value=mock_session_factory): + with patch("core.app.apps.base_app_runner.db") as mock_db: + runner = MagicMock() + method = AppRunner._handle_multimodal_image_content + runner._handle_multimodal_image_content = lambda *args, **kwargs: method( + runner, *args, **kwargs + ) - runner._handle_multimodal_image_content( - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) + runner._handle_multimodal_image_content( + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - # Assert - # Verify tool file was created from base64 - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - assert call_kwargs["user_id"] == mock_user_id - assert call_kwargs["tenant_id"] == mock_tenant_id - assert call_kwargs["conversation_id"] is None - assert "file_binary" in call_kwargs - assert call_kwargs["mimetype"] == "image/png" - assert call_kwargs["filename"].startswith("generated_image") - assert call_kwargs["filename"].endswith(".png") + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert call_kwargs["user_id"] == mock_user_id + assert call_kwargs["tenant_id"] == mock_tenant_id + assert call_kwargs["conversation_id"] is None + assert "file_binary" in call_kwargs + assert call_kwargs["mimetype"] == "image/png" + assert call_kwargs["filename"].startswith("generated_image") + assert call_kwargs["filename"].endswith(".png") - # Verify message file was created - mock_msg_file_class.assert_called_once() + mock_msg_file_class.assert_called_once() + file_session.add.assert_called_once() + mock_db.session.commit.assert_not_called() - # Verify database operations - mock_session.add.assert_called_once() - mock_session.commit.assert_called_once() - mock_session.refresh.assert_called_once() - - # Verify event was published - mock_queue_manager.publish.assert_called_once() + mock_queue_manager.publish.assert_called_once() def test_handle_multimodal_image_content_with_base64_data_uri( self, @@ -238,33 +227,32 @@ class TestBaseAppRunnerMultimodal: mock_mgr_class.return_value = mock_mgr with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - # Setup mock message file mock_msg_file_class.return_value = mock_message_file - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - mock_session.add = MagicMock() - mock_session.commit = MagicMock() - mock_session.refresh = MagicMock() + file_session = MagicMock() + mock_session_factory = MagicMock() + mock_session_factory.begin.return_value.__enter__ = MagicMock(return_value=file_session) + mock_session_factory.begin.return_value.__exit__ = MagicMock(return_value=False) - # Act - # Create a mock runner with the method bound - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + with patch("core.app.apps.base_app_runner.sessionmaker", return_value=mock_session_factory): + with patch("core.app.apps.base_app_runner.db"): + runner = MagicMock() + method = AppRunner._handle_multimodal_image_content + runner._handle_multimodal_image_content = lambda *args, **kwargs: method( + runner, *args, **kwargs + ) - runner._handle_multimodal_image_content( - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) + runner._handle_multimodal_image_content( + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - # Assert - verify that base64 data was extracted correctly (without prefix) - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - # The base64 data should be decoded, so we check the binary was passed - assert "file_binary" in call_kwargs + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert "file_binary" in call_kwargs def test_handle_multimodal_image_content_without_url_or_base64( self, @@ -284,9 +272,7 @@ class TestBaseAppRunnerMultimodal: with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - # Act - # Create a mock runner with the method bound + with patch("core.app.apps.base_app_runner.db"): runner = MagicMock() method = AppRunner._handle_multimodal_image_content runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) @@ -299,10 +285,8 @@ class TestBaseAppRunnerMultimodal: queue_manager=mock_queue_manager, ) - # Assert - should not create any files or publish events mock_mgr_class.assert_not_called() mock_msg_file_class.assert_not_called() - mock_session.add.assert_not_called() mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_with_error( @@ -322,20 +306,16 @@ class TestBaseAppRunnerMultimodal: ) with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock to raise exception mock_mgr = MagicMock() mock_mgr.create_file_by_url.side_effect = Exception("Network error") mock_mgr_class.return_value = mock_mgr with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - # Act - # Create a mock runner with the method bound + with patch("core.app.apps.base_app_runner.db"): runner = MagicMock() method = AppRunner._handle_multimodal_image_content runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) - # Should not raise exception, just log it runner._handle_multimodal_image_content( content=content, message_id=mock_message_id, @@ -344,9 +324,7 @@ class TestBaseAppRunnerMultimodal: queue_manager=mock_queue_manager, ) - # Assert - should not create message file or publish event on error mock_msg_file_class.assert_not_called() - mock_session.add.assert_not_called() mock_queue_manager.publish.assert_not_called() def test_handle_multimodal_image_content_debugger_mode( @@ -369,37 +347,36 @@ class TestBaseAppRunnerMultimodal: mock_queue_manager.invoke_from = InvokeFrom.DEBUGGER with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock tool file manager mock_mgr = MagicMock() mock_mgr.create_file_by_url.return_value = mock_tool_file mock_mgr_class.return_value = mock_mgr with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - # Setup mock message file mock_msg_file_class.return_value = mock_message_file - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - mock_session.add = MagicMock() - mock_session.commit = MagicMock() - mock_session.refresh = MagicMock() + file_session = MagicMock() + mock_session_factory = MagicMock() + mock_session_factory.begin.return_value.__enter__ = MagicMock(return_value=file_session) + mock_session_factory.begin.return_value.__exit__ = MagicMock(return_value=False) - # Act - # Create a mock runner with the method bound - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + with patch("core.app.apps.base_app_runner.sessionmaker", return_value=mock_session_factory): + with patch("core.app.apps.base_app_runner.db"): + runner = MagicMock() + method = AppRunner._handle_multimodal_image_content + runner._handle_multimodal_image_content = lambda *args, **kwargs: method( + runner, *args, **kwargs + ) - runner._handle_multimodal_image_content( - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) + runner._handle_multimodal_image_content( + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - # Assert - verify created_by_role is ACCOUNT for debugger mode - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["created_by_role"] == CreatorUserRole.ACCOUNT + call_kwargs = mock_msg_file_class.call_args[1] + assert call_kwargs["created_by_role"] == CreatorUserRole.ACCOUNT def test_handle_multimodal_image_content_service_api_mode( self, @@ -421,34 +398,33 @@ class TestBaseAppRunnerMultimodal: mock_queue_manager.invoke_from = InvokeFrom.SERVICE_API with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock tool file manager mock_mgr = MagicMock() mock_mgr.create_file_by_url.return_value = mock_tool_file mock_mgr_class.return_value = mock_mgr with patch("core.app.apps.base_app_runner.MessageFile", autospec=True) as mock_msg_file_class: - # Setup mock message file mock_msg_file_class.return_value = mock_message_file - with patch("core.app.apps.base_app_runner.db.session", autospec=True) as mock_session: - mock_session.add = MagicMock() - mock_session.commit = MagicMock() - mock_session.refresh = MagicMock() + file_session = MagicMock() + mock_session_factory = MagicMock() + mock_session_factory.begin.return_value.__enter__ = MagicMock(return_value=file_session) + mock_session_factory.begin.return_value.__exit__ = MagicMock(return_value=False) - # Act - # Create a mock runner with the method bound - runner = MagicMock() - method = AppRunner._handle_multimodal_image_content - runner._handle_multimodal_image_content = lambda *args, **kwargs: method(runner, *args, **kwargs) + with patch("core.app.apps.base_app_runner.sessionmaker", return_value=mock_session_factory): + with patch("core.app.apps.base_app_runner.db"): + runner = MagicMock() + method = AppRunner._handle_multimodal_image_content + runner._handle_multimodal_image_content = lambda *args, **kwargs: method( + runner, *args, **kwargs + ) - runner._handle_multimodal_image_content( - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) + runner._handle_multimodal_image_content( + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) - # Assert - verify created_by_role is END_USER for service API - call_kwargs = mock_msg_file_class.call_args[1] - assert call_kwargs["created_by_role"] == CreatorUserRole.END_USER + call_kwargs = mock_msg_file_class.call_args[1] + assert call_kwargs["created_by_role"] == CreatorUserRole.END_USER diff --git a/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py b/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py index 9b89b108207..ef12f0be965 100644 --- a/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py @@ -138,7 +138,7 @@ def test_generate_falls_back_to_new_conversation_when_conversation_missing(monke ) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", - SimpleNamespace(engine=object()), + SimpleNamespace(engine=object(), session=lambda: MagicMock()), ) trace_manager = object.__new__(TraceQueueManager) monkeypatch.setattr( diff --git a/api/tests/unit_tests/core/app/layers/test_pause_state_persist_layer.py b/api/tests/unit_tests/core/app/layers/test_pause_state_persist_layer.py index 18e724ec48b..ff7f27f5efa 100644 --- a/api/tests/unit_tests/core/app/layers/test_pause_state_persist_layer.py +++ b/api/tests/unit_tests/core/app/layers/test_pause_state_persist_layer.py @@ -17,6 +17,7 @@ from core.app.layers.pause_state_persist_layer import ( from core.workflow.nodes.human_input.pause_reason import HumanInputRequired from core.workflow.system_variables import SystemVariableKey from graphon.entities.pause_reason import HitlRequired, SchedulingPause +from graphon.filters import GraphEventFilterContext, ResponseStreamFilter from graphon.graph_engine.entities.commands import GraphEngineCommand from graphon.graph_engine.layers.base import GraphEngineLayerNotInitializedError from graphon.graph_events import ( @@ -31,6 +32,22 @@ from models.model import AppMode from repositories.factory import DifyAPIRepositoryFactory +def _create_initialized_response_stream_filter() -> ResponseStreamFilter: + """Build a `ResponseStreamFilter` that has already run `initialize()`. + + `ResponseStreamFilter.dumps()` raises `RuntimeError` unless the filter has + processed a `GraphEventFilterContext` first. In production this always + happens before any event (including `GraphRunPausedEvent`) reaches + `PauseStatePersistenceLayer.on_event`, so tests that exercise `on_event` + or a subsequent `dumps()` call need a filter in that same state. A + nodeless graph is enough to satisfy the precondition. + """ + response_stream_filter = ResponseStreamFilter() + context = GraphEventFilterContext(graph=Mock(nodes={}), runtime_state=Mock()) + response_stream_filter.initialize(context) + return response_stream_filter + + class TestDataFactory: """Factory helpers for constructing graph events used in tests.""" @@ -202,6 +219,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id=state_owner_user_id, generate_entity=self._create_generate_entity(), + response_stream_filter=ResponseStreamFilter(), ) assert layer._session_maker is session_factory @@ -216,6 +234,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner", generate_entity=self._create_generate_entity(), + response_stream_filter=ResponseStreamFilter(), ) graph_runtime_state = MockReadOnlyGraphRuntimeState() @@ -233,6 +252,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner-123", generate_entity=generate_entity, + response_stream_filter=_create_initialized_response_stream_filter(), ) mock_repo = Mock() @@ -272,6 +292,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner-123", generate_entity=generate_entity, + response_stream_filter=_create_initialized_response_stream_filter(), ) mock_repo = Mock() @@ -328,6 +349,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner-123", generate_entity=self._create_generate_entity(), + response_stream_filter=ResponseStreamFilter(), ) mock_repo = Mock() @@ -356,6 +378,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner-123", generate_entity=self._create_generate_entity(), + response_stream_filter=ResponseStreamFilter(), ) event = TestDataFactory.create_graph_run_paused_event() @@ -369,6 +392,7 @@ class TestPauseStatePersistenceLayer: session_factory=session_factory, state_owner_user_id="owner-123", generate_entity=self._create_generate_entity(), + response_stream_filter=_create_initialized_response_stream_filter(), ) mock_repo = Mock() @@ -468,3 +492,53 @@ def test_workflow_resumption_context_dumps_loads_roundtrip(state: WorkflowResump restored_entity = loaded.get_generate_entity() assert isinstance(restored_entity, type(state.generate_entity.entity)) assert restored_entity.extras["trace_session_id"] == "session-1" + + +def test_on_event_persists_response_stream_filter_dump(monkeypatch: pytest.MonkeyPatch) -> None: + session_factory = Mock(name="session_factory") + generate_entity = TestPauseStatePersistenceLayer._create_generate_entity(workflow_execution_id="run-123") + response_stream_filter = _create_initialized_response_stream_filter() + layer = PauseStatePersistenceLayer( + session_factory=session_factory, + state_owner_user_id="owner-123", + generate_entity=generate_entity, + response_stream_filter=response_stream_filter, + ) + + mock_repo = Mock() + mock_factory = Mock(return_value=mock_repo) + monkeypatch.setattr(DifyAPIRepositoryFactory, "create_api_workflow_run_repository", mock_factory) + + graph_runtime_state = MockReadOnlyGraphRuntimeState(workflow_execution_id="run-123") + layer.initialize(graph_runtime_state, MockCommandChannel()) + + event = TestDataFactory.create_graph_run_paused_event() + layer.on_event(event) + + serialized_state = mock_repo.create_workflow_pause.call_args.kwargs["state"] + resumption_context = WorkflowResumptionContext.loads(serialized_state) + assert resumption_context.serialized_response_stream_filter_state == response_stream_filter.dumps() + + +def test_get_response_stream_filter_restores_dumped_state() -> None: + original = _create_initialized_response_stream_filter() + context = WorkflowResumptionContext( + serialized_graph_runtime_state=json.dumps({"state": "workflow"}), + generate_entity=_WorkflowGenerateEntityWrapper(entity=TestPauseStatePersistenceLayer._create_generate_entity()), + serialized_response_stream_filter_state=original.dumps(), + ) + + restored = context.get_response_stream_filter() + + assert restored.dumps() == original.dumps() + + +def test_get_response_stream_filter_defaults_when_state_missing() -> None: + context = WorkflowResumptionContext( + serialized_graph_runtime_state=json.dumps({"state": "workflow"}), + generate_entity=_WorkflowGenerateEntityWrapper(entity=TestPauseStatePersistenceLayer._create_generate_entity()), + ) + + restored = context.get_response_stream_filter() + + assert isinstance(restored, ResponseStreamFilter) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py index 18382f053be..153157337e1 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py @@ -203,6 +203,7 @@ def _agent_thought() -> MessageAgentThought: tool="tool", tool_labels_str="{}", tool_input="input", + answer="answer", message_files="[]", ) thought.id = "thought" @@ -678,7 +679,7 @@ class TestEasyUiBasedGenerateTaskPipeline: assert agent_response.answer == "agent" assert isinstance(responses[-1], ErrorStreamResponse) assert isinstance(responses[-1].err, ValueError) - assert pipeline._task_state.llm_result.message.content == "annotatedagent" + assert pipeline._task_state.llm_result.message.content == "annotated" def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) @@ -720,6 +721,58 @@ class TestEasyUiBasedGenerateTaskPipeline: assert response is not None assert response.id == "thought" + assert response.thought == "t" + + def test_agent_thought_to_stream_response_normalizes_null_display_fields(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + agent_thought = _agent_thought() + agent_thought.thought = None + agent_thought.answer = None + agent_thought.observation = None + agent_thought.tool = None + agent_thought.tool_input = None + agent_thought.message_files = None + + class _Session: + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalar(self, *args, **kwargs): + return agent_thought + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.Session", + _Session, + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.db", + _FakeDb(), + ) + + response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) + + assert response is not None + assert response.thought == "" + assert response.observation == "" + assert response.tool == "" + assert response.tool_input == "" + assert response.model_dump(mode="json")["message_files"] == [] def test_process_routes_to_stream_and_starts_conversation_name_generation(self): conversation = _make_conversation(AppMode.CHAT) @@ -1280,7 +1333,7 @@ class TestEasyUiBasedGenerateTaskPipeline: usage_metadata = cast(dict[str, object], response.metadata["usage"]) assert usage_metadata["prompt_tokens"] == 1 - def test_record_files_returns_none_when_message_has_no_files(self, monkeypatch: pytest.MonkeyPatch): + def test_record_files_returns_empty_list_when_message_has_no_files(self, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) message = _make_message() pipeline = EasyUIBasedGenerateTaskPipeline( @@ -1316,7 +1369,7 @@ class TestEasyUiBasedGenerateTaskPipeline: response = pipeline._message_end_to_stream_response() - assert response.files is None + assert response.files == [] def test_record_files_handles_local_fallback_and_tool_url_variants(self, monkeypatch: pytest.MonkeyPatch): conversation = _make_conversation(AppMode.CHAT) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py index 595d716666c..b1c06e237a8 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py @@ -6,7 +6,7 @@ SSE event, which is critical for vision/image chat responses to render correctly Test Coverage: - Files array populated when MessageFile records exist -- Files array is None when no MessageFile records exist +- Files array is empty when no MessageFile records exist - Correct signed URL generation for LOCAL_FILE transfer method - Correct URL handling for REMOTE_URL transfer method - Correct URL handling for TOOL_FILE transfer method @@ -90,7 +90,7 @@ class TestMessageEndStreamResponseFiles: return upload_file def test_message_end_with_no_files(self, mock_pipeline): - """Test that files array is None when no MessageFile records exist.""" + """Test that files array is empty when no MessageFile records exist.""" # Arrange with ( patch("core.app.task_pipeline.easy_ui_based_generate_task_pipeline.db") as mock_db, @@ -108,9 +108,10 @@ class TestMessageEndStreamResponseFiles: # Assert assert isinstance(result, MessageEndStreamResponse) - assert result.files is None + assert result.files == [] assert result.id == mock_pipeline._message_id assert result.metadata == {"test": "metadata"} + mock_pipeline._task_state.metadata.model_dump.assert_called_once_with(exclude_none=True) def test_message_end_with_local_file(self, mock_pipeline, mock_message_file_local, mock_upload_file): """Test that files array is populated correctly for LOCAL_FILE transfer method.""" diff --git a/api/tests/unit_tests/core/app/test_llm_quota.py b/api/tests/unit_tests/core/app/test_llm_quota.py index 13bdf765358..ec6ac134443 100644 --- a/api/tests/unit_tests/core/app/test_llm_quota.py +++ b/api/tests/unit_tests/core/app/test_llm_quota.py @@ -1,7 +1,7 @@ from collections.abc import Generator from contextlib import contextmanager from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from sqlalchemy import create_engine, select @@ -28,8 +28,19 @@ from models.provider import Provider, ProviderType @contextmanager def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]: session_maker = sessionmaker(bind=engine, expire_on_commit=False) - with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker): - yield + sessions = [] + + def _session(): + session = session_maker() + sessions.append(session) + return session + + with patch("core.app.llm.quota.db", SimpleNamespace(session=_session)): + try: + yield + finally: + for session in sessions: + session.close() def test_ensure_llm_quota_available_for_model_raises_when_system_model_is_exhausted() -> None: @@ -122,6 +133,7 @@ def test_deduct_llm_quota_for_model_uses_identity_based_trial_billing() -> None: mock_deduct_credits.assert_called_once_with( tenant_id="tenant-id", credits_required=42, + session=ANY, ) @@ -241,6 +253,7 @@ def test_deduct_llm_quota_for_model_uses_credit_configuration() -> None: mock_deduct_credits.assert_called_once_with( tenant_id="tenant-id", credits_required=9, + session=ANY, ) @@ -276,6 +289,7 @@ def test_deduct_llm_quota_for_model_uses_single_charge_for_times_quota() -> None mock_deduct_credits.assert_called_once_with( tenant_id="tenant-id", credits_required=1, + session=ANY, ) @@ -313,6 +327,7 @@ def test_deduct_llm_quota_for_model_uses_paid_billing_pool() -> None: tenant_id="tenant-id", credits_required=5, pool_type="paid", + session=ANY, ) diff --git a/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py b/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py index 4912badfc55..62c4ae9d411 100644 --- a/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py +++ b/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py @@ -14,6 +14,9 @@ def mock_queue_manager(mocker: MockerFixture): @pytest.fixture def handler(mock_queue_manager, mocker: MockerFixture): + mocker.patch( + "core.callback_handler.index_tool_callback_handler.db", + ) return DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, app_id="app-1", @@ -33,8 +36,18 @@ class TestOnQuery: ], ) def test_on_query_success_roles(self, mocker: MockerFixture, mock_queue_manager, invoke_from, expected_role): - # Arrange - mock_session = mocker.Mock() + # Arrange — the caller passes a session, but our fix uses an independent one + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mock_session_factory = mocker.MagicMock() + mock_session_factory.begin.return_value.__enter__ = mocker.MagicMock(return_value=independent_session) + mock_session_factory.begin.return_value.__exit__ = mocker.MagicMock(return_value=False) + mocker.patch( + "core.callback_handler.index_tool_callback_handler.sessionmaker", + return_value=mock_session_factory, + ) + mocker.patch("core.callback_handler.index_tool_callback_handler.db") handler = DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, @@ -46,17 +59,28 @@ class TestOnQuery: handler._invoke_from = invoke_from - # Act - handler.on_query("test query", "dataset-1", mock_session) + # Act — pass caller_session as required by signature + handler.on_query("test query", "dataset-1", caller_session) - # Assert - mock_session.add.assert_called_once() - dataset_query = mock_session.add.call_args.args[0] + # Assert — independent session used, not the caller's session + independent_session.add.assert_called_once() + dataset_query = independent_session.add.call_args.args[0] assert dataset_query.created_by_role == expected_role - mock_session.commit.assert_called_once() + caller_session.add.assert_not_called() + caller_session.commit.assert_not_called() def test_on_query_none_values(self, mocker: MockerFixture, mock_queue_manager): - mock_session = mocker.Mock() + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mock_session_factory = mocker.MagicMock() + mock_session_factory.begin.return_value.__enter__ = mocker.MagicMock(return_value=independent_session) + mock_session_factory.begin.return_value.__exit__ = mocker.MagicMock(return_value=False) + mocker.patch( + "core.callback_handler.index_tool_callback_handler.sessionmaker", + return_value=mock_session_factory, + ) + mocker.patch("core.callback_handler.index_tool_callback_handler.db") handler = DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, @@ -66,40 +90,67 @@ class TestOnQuery: invoke_from=None, ) - handler.on_query(None, None, mock_session) + handler.on_query(None, None, caller_session) - mock_session.add.assert_called_once() - mock_session.commit.assert_called_once() + independent_session.add.assert_called_once() + caller_session.add.assert_not_called() class TestOnToolEnd: def test_on_tool_end_no_metadata(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): - mock_session = mocker.Mock() + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mocker.patch( + "core.callback_handler.index_tool_callback_handler.Session", + return_value=independent_session, + ) + independent_session.__enter__ = mocker.MagicMock(return_value=independent_session) + independent_session.__exit__ = mocker.MagicMock(return_value=False) document = mocker.Mock() document.metadata = None - handler.on_tool_end([document], mock_session) + handler.on_tool_end([document], caller_session) - mock_session.commit.assert_not_called() + independent_session.commit.assert_called_once() + independent_session.execute.assert_not_called() + caller_session.commit.assert_not_called() def test_on_tool_end_dataset_document_not_found( self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture ): - mock_session = mocker.Mock() - mock_session.scalar.return_value = None + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mocker.patch( + "core.callback_handler.index_tool_callback_handler.Session", + return_value=independent_session, + ) + independent_session.__enter__ = mocker.MagicMock(return_value=independent_session) + independent_session.__exit__ = mocker.MagicMock(return_value=False) + independent_session.scalar.return_value = None document = mocker.Mock() document.metadata = {"document_id": "doc-1", "doc_id": "node-1"} - handler.on_tool_end([document], mock_session) + handler.on_tool_end([document], caller_session) - mock_session.scalar.assert_called_once() + independent_session.scalar.assert_called_once() + caller_session.scalar.assert_not_called() def test_on_tool_end_parent_child_index_with_child( self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture ): - mock_session = mocker.Mock() + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mocker.patch( + "core.callback_handler.index_tool_callback_handler.Session", + return_value=independent_session, + ) + independent_session.__enter__ = mocker.MagicMock(return_value=independent_session) + independent_session.__exit__ = mocker.MagicMock(return_value=False) mock_dataset_doc = mocker.Mock() from core.callback_handler.index_tool_callback_handler import IndexStructureType @@ -111,23 +162,32 @@ class TestOnToolEnd: mock_child_chunk = mocker.Mock() mock_child_chunk.segment_id = "segment-1" - mock_session.scalar.side_effect = [mock_dataset_doc, mock_child_chunk] + independent_session.scalar.side_effect = [mock_dataset_doc, mock_child_chunk] document = mocker.Mock() document.metadata = {"document_id": "doc-1", "doc_id": "node-1"} - handler.on_tool_end([document], mock_session) + handler.on_tool_end([document], caller_session) - mock_session.execute.assert_called_once() - mock_session.commit.assert_called_once() + independent_session.execute.assert_called_once() + independent_session.commit.assert_called_once() + caller_session.execute.assert_not_called() def test_on_tool_end_non_parent_child_index(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): - mock_session = mocker.Mock() + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mocker.patch( + "core.callback_handler.index_tool_callback_handler.Session", + return_value=independent_session, + ) + independent_session.__enter__ = mocker.MagicMock(return_value=independent_session) + independent_session.__exit__ = mocker.MagicMock(return_value=False) mock_dataset_doc = mocker.Mock() mock_dataset_doc.doc_form = "OTHER" - mock_session.scalar.return_value = mock_dataset_doc + independent_session.scalar.return_value = mock_dataset_doc document = mocker.Mock() document.metadata = { @@ -136,14 +196,24 @@ class TestOnToolEnd: "dataset_id": "dataset-1", } - handler.on_tool_end([document], mock_session) + handler.on_tool_end([document], caller_session) - mock_session.execute.assert_called_once() - mock_session.commit.assert_called_once() + independent_session.execute.assert_called_once() + independent_session.commit.assert_called_once() + caller_session.execute.assert_not_called() def test_on_tool_end_empty_documents(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): - mock_session = mocker.Mock() - handler.on_tool_end([], mock_session) + caller_session = mocker.Mock() + + independent_session = mocker.MagicMock() + mocker.patch( + "core.callback_handler.index_tool_callback_handler.Session", + return_value=independent_session, + ) + independent_session.__enter__ = mocker.MagicMock(return_value=independent_session) + independent_session.__exit__ = mocker.MagicMock(return_value=False) + + handler.on_tool_end([], caller_session) class TestReturnRetrieverResourceInfo: diff --git a/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py b/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py index 6a3d21a33d1..7e3cc722b0d 100644 --- a/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py +++ b/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py @@ -8,7 +8,7 @@ from core.datasource.entities.datasource_entities import ( DatasourceProviderEntityWithPlugin, DatasourceProviderType, ) -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.errors import ToolProviderCredentialValidationError @@ -149,7 +149,7 @@ class TestDatasourcePluginProviderController: mock_config = MagicMock(spec=ProviderConfig) mock_config.name = "text_field" mock_config.required = True - mock_config.type = ProviderConfig.Type.TEXT_INPUT + mock_config.type = ProviderConfigType.TEXT_INPUT mock_entity = MagicMock(spec=DatasourceProviderEntityWithPlugin) mock_entity.credentials_schema = [mock_config] @@ -167,7 +167,7 @@ class TestDatasourcePluginProviderController: mock_config = MagicMock(spec=ProviderConfig) mock_config.name = "select_field" mock_config.required = True - mock_config.type = ProviderConfig.Type.SELECT + mock_config.type = ProviderConfigType.SELECT mock_config.options = [mock_option] mock_entity = MagicMock(spec=DatasourceProviderEntityWithPlugin) @@ -206,7 +206,7 @@ class TestDatasourcePluginProviderController: mock_config = MagicMock(spec=ProviderConfig) mock_config.name = "valid_field" mock_config.required = True - mock_config.type = ProviderConfig.Type.TEXT_INPUT + mock_config.type = ProviderConfigType.TEXT_INPUT mock_entity = MagicMock(spec=DatasourceProviderEntityWithPlugin) mock_entity.credentials_schema = [mock_config] @@ -243,7 +243,7 @@ class TestDatasourcePluginProviderController: mock_config_text = MagicMock(spec=ProviderConfig) mock_config_text.name = "text_def" mock_config_text.required = False - mock_config_text.type = ProviderConfig.Type.TEXT_INPUT + mock_config_text.type = ProviderConfigType.TEXT_INPUT mock_config_text.default = 123 # Int default, should be converted to str mock_config_other = MagicMock(spec=ProviderConfig) diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py index 9eb32206039..b5a48918079 100644 --- a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py +++ b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py @@ -25,6 +25,7 @@ from core.entities.provider_entities import ( SystemConfiguration, SystemConfigurationStatus, ) +from core.helper.model_provider_cache import ProviderCredentialsCacheType from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType from graphon.model_runtime.entities.provider_entities import ( @@ -1336,6 +1337,8 @@ def test_create_update_delete_custom_model_credential_flow() -> None: configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") assert provider_model_record.credential_id is None assert mock_cache.return_value.delete.call_count == 2 + assert mock_cache.call_args_list[0].kwargs["cache_type"] == ProviderCredentialsCacheType.LOAD_BALANCING_MODEL + assert mock_cache.call_args_list[1].kwargs["cache_type"] == ProviderCredentialsCacheType.MODEL session = Mock() mismatched_credential_record = SimpleNamespace( @@ -2032,9 +2035,16 @@ def test_delete_custom_model_credential_removes_custom_model_record_when_last_cr with _patched_session(session): with patch.object(ProviderConfiguration, "_get_custom_model_record", return_value=provider_model_record): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") + with patch("core.entities.provider_configuration.ProviderCredentialsCache") as mock_cache: + configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", "cred-1") assert any(call.args and call.args[0] is provider_model_record for call in session.delete.call_args_list) + mock_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id="model-1", + cache_type=ProviderCredentialsCacheType.MODEL, + ) + mock_cache.return_value.delete.assert_called_once() def test_delete_custom_model_credential_rolls_back_on_error() -> None: diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_entities.py b/api/tests/unit_tests/core/entities/test_entities_provider_entities.py index a159d3ad4d0..9cf6e5ac3f1 100644 --- a/api/tests/unit_tests/core/entities/test_entities_provider_entities.py +++ b/api/tests/unit_tests/core/entities/test_entities_provider_entities.py @@ -5,6 +5,7 @@ from core.entities.provider_entities import ( BasicProviderConfig, ModelSettings, ProviderConfig, + ProviderConfigType, ProviderQuotaType, ) from core.tools.entities.common_entities import I18nObject @@ -27,22 +28,22 @@ def test_provider_quota_type_value_of_rejects_unknown_values() -> None: def test_basic_provider_config_type_value_of_handles_known_values() -> None: # Arrange / Act - parameter_type = BasicProviderConfig.Type.value_of("text-input") + parameter_type = ProviderConfigType.value_of("text-input") # Assert - assert parameter_type == BasicProviderConfig.Type.TEXT_INPUT + assert parameter_type == ProviderConfigType.TEXT_INPUT def test_basic_provider_config_type_value_of_rejects_invalid_values() -> None: # Arrange / Act / Assert with pytest.raises(ValueError, match="invalid mode value"): - BasicProviderConfig.Type.value_of("unknown") + ProviderConfigType.value_of("unknown") def test_provider_config_to_basic_provider_config_keeps_type_and_name() -> None: # Arrange provider_config = ProviderConfig( - type=BasicProviderConfig.Type.SELECT, + type=ProviderConfigType.SELECT, name="workspace", scope=AppSelectorScope.ALL, options=[ProviderConfig.Option(value="all", label=I18nObject(en_US="All"))], @@ -53,7 +54,7 @@ def test_provider_config_to_basic_provider_config_keeps_type_and_name() -> None: # Assert assert isinstance(basic_config, BasicProviderConfig) - assert basic_config.type == BasicProviderConfig.Type.SELECT + assert basic_config.type == ProviderConfigType.SELECT assert basic_config.name == "workspace" diff --git a/api/tests/unit_tests/core/helper/code_executor/test_template_transformer.py b/api/tests/unit_tests/core/helper/code_executor/test_template_transformer.py index 5b54b8e6474..d9d171efaea 100644 --- a/api/tests/unit_tests/core/helper/code_executor/test_template_transformer.py +++ b/api/tests/unit_tests/core/helper/code_executor/test_template_transformer.py @@ -1,6 +1,5 @@ import json from base64 import b64decode -from collections.abc import Mapping from typing import Any import pytest @@ -44,7 +43,7 @@ def test_serialize_inputs_encodes_payload() -> None: def test_transform_response_parses_json_result_and_converts_scientific_notation() -> None: response = '<>{"value": "1e+3", "nested": {"x": "2E-2"}, "arr": ["3e+1"]}<>' - result: Mapping[str, Any] = _DummyTransformer.transform_response(response) + result: dict[str, Any] = _DummyTransformer.transform_response(response) assert result == {"value": 1000.0, "nested": {"x": 0.02}, "arr": [30.0]} diff --git a/api/tests/unit_tests/core/helper/test_creators.py b/api/tests/unit_tests/core/helper/test_creators.py index 8750f6d9070..cac7ad35a71 100644 --- a/api/tests/unit_tests/core/helper/test_creators.py +++ b/api/tests/unit_tests/core/helper/test_creators.py @@ -46,6 +46,18 @@ class TestUploadDSL: with pytest.raises(ValueError, match="claim_code"): upload_dsl(b"app: demo") + @patch("core.helper.creators.httpx.post") + def test_raises_on_non_string_claim_code(self, mock_post): + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = {"data": {"claim_code": 123}} + mock_response.raise_for_status = MagicMock() + mock_post.return_value = mock_response + + from core.helper.creators import upload_dsl + + with pytest.raises(ValueError, match="claim_code"): + upload_dsl(b"app: demo") + @patch("core.helper.creators.httpx.post") def test_raises_on_http_error(self, mock_post): mock_response = MagicMock(spec=httpx.Response) diff --git a/api/tests/unit_tests/core/helper/test_marketplace.py b/api/tests/unit_tests/core/helper/test_marketplace.py index eba1e4e5442..a14a9fe582d 100644 --- a/api/tests/unit_tests/core/helper/test_marketplace.py +++ b/api/tests/unit_tests/core/helper/test_marketplace.py @@ -1,6 +1,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest from pytest_mock import MockerFixture from core.helper.marketplace import ( @@ -53,6 +54,16 @@ def test_batch_fetch_plugin_by_ids_returns_plugins_from_response(mocker: MockerF response.raise_for_status.assert_called_once() +def test_batch_fetch_plugin_by_ids_rejects_invalid_plugins_response(mocker: MockerFixture) -> None: + response = MagicMock() + response.json.return_value = {"data": {"plugins": ["p1"]}} + response.raise_for_status.return_value = None + mocker.patch("core.helper.marketplace.httpx.post", return_value=response) + + with pytest.raises(ValueError, match="plugins list"): + batch_fetch_plugin_by_ids(["p1"]) + + def test_batch_fetch_plugin_manifests_returns_empty_for_empty_input(mocker: MockerFixture) -> None: post_mock = mocker.patch("core.helper.marketplace.httpx.post") diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py index ddb33f0758f..1c4a6e2db7c 100644 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py +++ b/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py @@ -149,12 +149,12 @@ class TestWorkflowServiceInterface: from core.llm_generator.llm_generator import WorkflowServiceInterface class MockService(WorkflowServiceInterface): - def get_draft_workflow(self, app_model, workflow_id=None): - return super().get_draft_workflow(app_model, workflow_id) + def get_draft_workflow(self, app_model, workflow_id=None, *, session): + return super().get_draft_workflow(app_model, workflow_id, session=session) def get_node_last_run(self, app_model, workflow, node_id): return super().get_node_last_run(app_model, workflow, node_id) service = MockService() - service.get_draft_workflow(None) + service.get_draft_workflow(None, session=None) service.get_node_last_run(None, None, "node") diff --git a/api/tests/unit_tests/core/mcp/server/test_streamable_http.py b/api/tests/unit_tests/core/mcp/server/test_streamable_http.py index 3f95736ef92..cc00e02252e 100644 --- a/api/tests/unit_tests/core/mcp/server/test_streamable_http.py +++ b/api/tests/unit_tests/core/mcp/server/test_streamable_http.py @@ -10,11 +10,13 @@ from core.mcp.server.streamable_http import ( build_parameter_schema, convert_input_form_to_parameters, extract_answer_from_response, + extract_structured_output, handle_call_tool, handle_initialize, handle_list_tools, handle_mcp_request, handle_ping, + negotiate_protocol_version, prepare_tool_arguments, process_mapping_response, ) @@ -64,6 +66,8 @@ class TestHandleMCPRequest: # Setup initialize request self.mock_request.root = Mock(spec=types.InitializeRequest) self.mock_request.root.id = 123 + self.mock_request.root.params = Mock() + self.mock_request.root.params.protocolVersion = "2025-06-18" request_type = Mock(return_value=types.InitializeRequest) with patch("core.mcp.server.streamable_http.type", request_type): @@ -91,6 +95,33 @@ class TestHandleMCPRequest: assert result.jsonrpc == "2.0" assert result.id == 123 + def test_handle_list_tools_request_threads_protocol_version(self): + """The negotiated version reaches handle_list_tools through the dispatcher.""" + self.mock_request.root = Mock(spec=types.ListToolsRequest) + self.mock_request.root.id = 123 + + result = handle_mcp_request( + Mock(), self.app, self.mock_request, self.user_input_form, self.mcp_server, self.end_user, 123, "2025-06-18" + ) + + assert isinstance(result, types.JSONRPCResponse) + tool = result.result["tools"][0] + assert tool["outputSchema"] == {"type": "object"} + assert tool["title"] == "test_app" + + def test_handle_list_tools_request_legacy_serialization_unchanged(self): + """A 2024-11-05 tools/list response serializes without any 2025-06-18 fields.""" + self.mock_request.root = Mock(spec=types.ListToolsRequest) + self.mock_request.root.id = 123 + + result = handle_mcp_request( + Mock(), self.app, self.mock_request, self.user_input_form, self.mcp_server, self.end_user, 123, "2024-11-05" + ) + + assert isinstance(result, types.JSONRPCResponse) + tool = result.result["tools"][0] + assert set(tool) == {"name", "description", "inputSchema"} + @patch("core.mcp.server.streamable_http.AppGenerateService") def test_handle_call_tool_request(self, mock_app_generate): """Test handling call tool request""" @@ -119,6 +150,43 @@ class TestHandleMCPRequest: # Verify AppGenerateService was called mock_app_generate.generate.assert_called_once() + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_handle_call_tool_request_threads_protocol_version(self, mock_app_generate): + """The negotiated version reaches handle_call_tool through the dispatcher.""" + mock_call_request = Mock(spec=types.CallToolRequest) + mock_call_request.params = Mock() + mock_call_request.params.arguments = {"query": "test question"} + mock_call_request.id = 123 + self.mock_request.root = mock_call_request + + mock_app_generate.generate.return_value = {"answer": "test answer"} + + result = handle_mcp_request( + Mock(), self.app, self.mock_request, self.user_input_form, self.mcp_server, self.end_user, 123, "2025-06-18" + ) + + assert isinstance(result, types.JSONRPCResponse) + assert result.result["structuredContent"] == {"answer": "test answer"} + + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_handle_call_tool_request_legacy_serialization_unchanged(self, mock_app_generate): + """A 2024-11-05 tools/call response serializes without structuredContent.""" + mock_call_request = Mock(spec=types.CallToolRequest) + mock_call_request.params = Mock() + mock_call_request.params.arguments = {"query": "test question"} + mock_call_request.id = 123 + self.mock_request.root = mock_call_request + + mock_app_generate.generate.return_value = {"answer": "test answer"} + + result = handle_mcp_request( + Mock(), self.app, self.mock_request, self.user_input_form, self.mcp_server, self.end_user, 123, "2024-11-05" + ) + + assert isinstance(result, types.JSONRPCResponse) + assert "structuredContent" not in result.result + assert result.result["content"][0]["text"] == "test answer" + def test_handle_unknown_request_type(self): """Test handling unknown request type""" @@ -183,18 +251,49 @@ class TestIndividualHandlers: result = handle_ping() assert isinstance(result, types.EmptyResult) - def test_handle_initialize(self): - """Test initialize handler""" - description = "Test server" - + def test_handle_initialize_echoes_supported_version(self): + """A supported requested version is echoed back unchanged.""" with patch("core.mcp.server.streamable_http.dify_config") as mock_config: mock_config.project.version = "1.0.0" - result = handle_initialize(description) + result = handle_initialize("Test server", "2024-11-05") assert isinstance(result, types.InitializeResult) - assert result.protocolVersion == types.SERVER_LATEST_PROTOCOL_VERSION + assert result.protocolVersion == "2024-11-05" assert result.instructions == "Test server" + def test_handle_initialize_echoes_intermediate_version(self): + """The intermediate supported version (2025-03-26) is echoed back.""" + with patch("core.mcp.server.streamable_http.dify_config") as mock_config: + mock_config.project.version = "1.0.0" + result = handle_initialize("Test server", "2025-03-26") + + assert result.protocolVersion == "2025-03-26" + + def test_handle_initialize_negotiates_latest_for_modern_client(self): + """A 2025-06-18 client gets 2025-06-18 back.""" + with patch("core.mcp.server.streamable_http.dify_config") as mock_config: + mock_config.project.version = "1.0.0" + result = handle_initialize("Test server", "2025-06-18") + + assert result.protocolVersion == "2025-06-18" + + def test_handle_initialize_falls_back_for_unknown_version(self): + """An unsupported requested version falls back to the server latest.""" + with patch("core.mcp.server.streamable_http.dify_config") as mock_config: + mock_config.project.version = "1.0.0" + result = handle_initialize("Test server", "1999-01-01") + + assert result.protocolVersion == types.SERVER_LATEST_PROTOCOL_VERSION + assert result.protocolVersion == "2025-06-18" + + def test_handle_initialize_non_string_version_falls_back(self): + """A malformed (non-string) requested version falls back to the server latest.""" + with patch("core.mcp.server.streamable_http.dify_config") as mock_config: + mock_config.project.version = "1.0.0" + result = handle_initialize("Test server", 20250618) + + assert result.protocolVersion == types.SERVER_LATEST_PROTOCOL_VERSION + def test_handle_list_tools(self): """Test list tools handler""" app_name = "test_app" @@ -210,6 +309,30 @@ class TestIndividualHandlers: assert result.tools[0].name == "test_app" assert result.tools[0].description == "Test server" + def test_handle_list_tools_adds_structured_output_for_modern_client(self): + """Tool advertises outputSchema and title when negotiated >= 2025-06-18.""" + result = handle_list_tools("test_app", AppMode.CHAT, [], "Test server", {}, "2025-06-18") + + tool = result.tools[0] + assert tool.outputSchema == {"type": "object"} + assert tool.title == "test_app" + + def test_handle_list_tools_omits_structured_output_for_legacy_client(self): + """Tool stays unchanged (no outputSchema/title) for 2024-11-05 clients.""" + result = handle_list_tools("test_app", AppMode.CHAT, [], "Test server", {}, "2024-11-05") + + tool = result.tools[0] + assert tool.outputSchema is None + assert tool.title is None + + def test_handle_list_tools_omits_structured_output_for_intermediate_client(self): + """The 2025-03-26 negotiated version is below the structured-output threshold.""" + result = handle_list_tools("test_app", AppMode.CHAT, [], "Test server", {}, "2025-03-26") + + tool = result.tools[0] + assert tool.outputSchema is None + assert tool.title is None + @patch("core.mcp.server.streamable_http.AppGenerateService") def test_handle_call_tool(self, mock_app_generate): """Test call tool handler""" @@ -239,6 +362,44 @@ class TestIndividualHandlers: assert hasattr(text_content, "text") assert text_content.text == "test answer" + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_handle_call_tool_structured_output_modern_client(self, mock_app_generate): + """structuredContent is attached alongside TextContent for >= 2025-06-18.""" + app = Mock(spec=App) + app.mode = AppMode.CHAT + + mock_request = Mock() + mock_call_request = Mock(spec=types.CallToolRequest) + mock_call_request.params = Mock() + mock_call_request.params.arguments = {"query": "test question"} + mock_request.root = mock_call_request + + mock_app_generate.generate.return_value = {"answer": "test answer"} + + result = handle_call_tool(Mock(), app, mock_request, [], Mock(spec=EndUser), "2025-06-18") + + assert result.structuredContent == {"answer": "test answer"} + assert result.content[0].text == "test answer" + + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_handle_call_tool_no_structured_output_legacy_client(self, mock_app_generate): + """structuredContent is omitted for 2024-11-05 clients.""" + app = Mock(spec=App) + app.mode = AppMode.CHAT + + mock_request = Mock() + mock_call_request = Mock(spec=types.CallToolRequest) + mock_call_request.params = Mock() + mock_call_request.params.arguments = {"query": "test question"} + mock_request.root = mock_call_request + + mock_app_generate.generate.return_value = {"answer": "test answer"} + + result = handle_call_tool(Mock(), app, mock_request, [], Mock(spec=EndUser), "2024-11-05") + + assert result.structuredContent is None + assert result.content[0].text == "test answer" + def test_handle_call_tool_no_end_user(self): """Test call tool handler without end user""" app = Mock(spec=App) @@ -375,6 +536,65 @@ class TestUtilityFunctions: assert result == "thinking...more thinking" + def test_extract_structured_output_workflow(self): + """Workflow mode exposes the raw outputs mapping as structured content.""" + app = Mock(spec=App) + app.mode = AppMode.WORKFLOW + + response = {"data": {"outputs": {"result": "test result"}}} + + assert extract_structured_output(app, response, "ignored") == {"result": "test result"} + + def test_extract_structured_output_chat(self): + """Chat mode wraps the answer string under an 'answer' key.""" + app = Mock(spec=App) + app.mode = AppMode.CHAT + + assert extract_structured_output(app, {"answer": "hi"}, "hi") == {"answer": "hi"} + + def test_extract_structured_output_workflow_missing_outputs(self): + """Missing or malformed outputs fall back to None.""" + app = Mock(spec=App) + app.mode = AppMode.WORKFLOW + + assert extract_structured_output(app, {"data": {}}, "ignored") is None + + def test_extract_structured_output_workflow_non_mapping_response(self): + """A non-mapping workflow response yields no structured output.""" + app = Mock(spec=App) + app.mode = AppMode.WORKFLOW + + assert extract_structured_output(app, None, "ignored") is None + + def test_extract_structured_output_workflow_non_mapping_data(self): + """A non-mapping 'data' entry yields no structured output.""" + app = Mock(spec=App) + app.mode = AppMode.WORKFLOW + + assert extract_structured_output(app, {"data": "not a mapping"}, "ignored") is None + + def test_extract_structured_output_workflow_non_mapping_outputs(self): + """A non-mapping 'outputs' entry yields no structured output.""" + app = Mock(spec=App) + app.mode = AppMode.WORKFLOW + + assert extract_structured_output(app, {"data": {"outputs": ["not", "a", "mapping"]}}, "ignored") is None + + @pytest.mark.parametrize("mode", [AppMode.ADVANCED_CHAT, AppMode.AGENT_CHAT, AppMode.COMPLETION]) + def test_extract_structured_output_other_answer_modes(self, mode): + """Every chat-style mode wraps the answer string under an 'answer' key.""" + app = Mock(spec=App) + app.mode = mode + + assert extract_structured_output(app, {"answer": "hi"}, "hi") == {"answer": "hi"} + + def test_extract_structured_output_unknown_mode(self): + """Modes outside the MCP surface produce no structured output.""" + app = Mock(spec=App) + app.mode = AppMode.CHANNEL + + assert extract_structured_output(app, {"answer": "hi"}, "hi") is None + def test_process_mapping_response_invalid_mode(self): """Test processing mapping response with invalid app mode""" app = Mock(spec=App) @@ -578,3 +798,29 @@ class TestUtilityFunctions: # Or validation should also raise SchemaError with pytest.raises(jsonschema.exceptions.SchemaError): jsonschema.validate(instance={"count": 1.23}, schema=bad_schema) + + +class TestNegotiateProtocolVersion: + """Test the MCP-Protocol-Version header resolver.""" + + def test_initialize_ignores_header(self): + """Initialize negotiates via the request body, so its header is ignored.""" + assert negotiate_protocol_version("anything", True) == types.DEFAULT_NEGOTIATED_VERSION + + def test_absent_header_defaults(self): + """An absent header defaults to 2025-03-26 per the spec back-compat rule.""" + assert negotiate_protocol_version(None, False) == types.DEFAULT_NEGOTIATED_VERSION + + def test_empty_header_treated_as_absent(self): + """An empty header value is treated as absent and defaults to 2025-03-26.""" + assert negotiate_protocol_version("", False) == types.DEFAULT_NEGOTIATED_VERSION + + def test_supported_header_passes_through(self): + """All supported header values are used as the negotiated version.""" + assert negotiate_protocol_version("2025-06-18", False) == "2025-06-18" + assert negotiate_protocol_version("2025-03-26", False) == "2025-03-26" + assert negotiate_protocol_version("2024-11-05", False) == "2024-11-05" + + def test_unsupported_header_returns_none(self): + """An explicit but unsupported header signals an error (None).""" + assert negotiate_protocol_version("1999-01-01", False) is None diff --git a/api/tests/unit_tests/core/mcp/test_types.py b/api/tests/unit_tests/core/mcp/test_types.py index d4fe353f0aa..ff16f1f9dc8 100644 --- a/api/tests/unit_tests/core/mcp/test_types.py +++ b/api/tests/unit_tests/core/mcp/test_types.py @@ -4,6 +4,7 @@ import pytest from pydantic import ValidationError from core.mcp.types import ( + DEFAULT_NEGOTIATED_VERSION, INTERNAL_ERROR, INVALID_PARAMS, INVALID_REQUEST, @@ -11,6 +12,7 @@ from core.mcp.types import ( METHOD_NOT_FOUND, PARSE_ERROR, SERVER_LATEST_PROTOCOL_VERSION, + SERVER_SUPPORTED_PROTOCOL_VERSIONS, Annotations, CallToolRequest, CallToolRequestParams, @@ -59,7 +61,9 @@ class TestConstants: def test_protocol_versions(self): """Test protocol version constants.""" assert LATEST_PROTOCOL_VERSION == "2025-06-18" - assert SERVER_LATEST_PROTOCOL_VERSION == "2024-11-05" + assert SERVER_LATEST_PROTOCOL_VERSION == "2025-06-18" + assert DEFAULT_NEGOTIATED_VERSION == "2025-03-26" + assert sorted(SERVER_SUPPORTED_PROTOCOL_VERSIONS) == ["2024-11-05", "2025-03-26", "2025-06-18"] def test_error_codes(self): """Test JSON-RPC error code constants.""" diff --git a/api/tests/unit_tests/core/plugin/impl/test_model_client.py b/api/tests/unit_tests/core/plugin/impl/test_model_client.py index ac3df1e56fc..c707b52ccaf 100644 --- a/api/tests/unit_tests/core/plugin/impl/test_model_client.py +++ b/api/tests/unit_tests/core/plugin/impl/test_model_client.py @@ -156,15 +156,35 @@ class TestPluginModelClient: tools=[], stop=["STOP"], stream=False, + app_id="app-1", ) ) assert result == ["chunk-1"] call_kwargs = stream_mock.call_args.kwargs assert call_kwargs["path"] == "plugin/tenant-1/dispatch/llm/invoke" + assert call_kwargs["data"]["app_id"] == "app-1" assert call_kwargs["data"]["data"]["stream"] is False assert call_kwargs["data"]["data"]["model_parameters"] == {"temperature": 0.1} + def test_invoke_llm_omits_app_id_when_missing(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + list( + client.invoke_llm( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin:1", + provider="provider-a", + model="gpt-test", + credentials={}, + prompt_messages=[], + ) + ) + + assert "app_id" not in stream_mock.call_args.kwargs["data"] + def test_invoke_llm_wraps_plugin_daemon_inner_error(self, mocker: MockerFixture): client = PluginModelClient() diff --git a/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py b/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py index d665dcd2e45..baedd873005 100644 --- a/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py +++ b/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py @@ -355,14 +355,32 @@ class TestPluginAppBackwardsInvocation: assert "end_users.app_id" in compiled assert stmt.compile().params == {"id_1": "uid", "tenant_id_1": "tenant-1", "app_id_1": "app-1"} + def test_get_user_returns_end_user_by_session_id(self, mocker: MockerFixture): + session = self.patch_create_session(mocker, side_effect=[None, MagicMock(id="session-user")]) + app = SimpleNamespace(id="app-1", tenant_id="tenant-1") + + user = PluginAppBackwardsInvocation._get_user("wecom-sender-1", app) + + assert user.id == "session-user" + stmt = session.scalar.call_args_list[1].args[0] + compiled = str(stmt.compile(dialect=postgresql.dialect())) + assert "end_users.session_id" in compiled + assert "end_users.tenant_id" in compiled + assert "end_users.app_id" in compiled + assert stmt.compile().params == { + "session_id_1": "wecom-sender-1", + "tenant_id_1": "tenant-1", + "app_id_1": "app-1", + } + def test_get_user_falls_back_to_account_user(self, mocker: MockerFixture): - session = self.patch_create_session(mocker, side_effect=[None, MagicMock(id="account-user")]) + session = self.patch_create_session(mocker, side_effect=[None, None, MagicMock(id="account-user")]) app = SimpleNamespace(id="app-1", tenant_id="tenant-1") user = PluginAppBackwardsInvocation._get_user("uid", app) assert user.id == "account-user" - stmt = session.scalar.call_args_list[1].args[0] + stmt = session.scalar.call_args_list[2].args[0] compiled = str(stmt.compile(dialect=postgresql.dialect())) assert "accounts.id" in compiled assert "tenant_account_joins.account_id" in compiled @@ -370,12 +388,41 @@ class TestPluginAppBackwardsInvocation: assert stmt.compile().params == {"id_1": "uid", "tenant_id_1": "tenant-1"} def test_get_user_raises_when_user_not_found(self, mocker: MockerFixture): - self.patch_create_session(mocker, side_effect=[None, None]) + self.patch_create_session(mocker, side_effect=[None, None, None]) app = SimpleNamespace(id="app-1", tenant_id="tenant-1") with pytest.raises(ValueError, match="user not found"): PluginAppBackwardsInvocation._get_user("uid", app) + def test_invoke_app_creates_end_user_for_unknown_external_user_id(self, mocker: MockerFixture): + app = MagicMock(mode=AppMode.WORKFLOW) + end_user = MagicMock() + workflow = MagicMock() + mocker.patch.object(PluginAppBackwardsInvocation, "_get_app", return_value=app) + mocker.patch.object(PluginAppBackwardsInvocation, "_get_workflow", return_value=workflow) + mocker.patch.object(PluginAppBackwardsInvocation, "_get_user", side_effect=ValueError("user not found")) + get_or_create = mocker.patch( + "core.plugin.backwards_invocation.app.EndUserService.get_or_create_end_user", + return_value=end_user, + ) + route = mocker.patch.object(PluginAppBackwardsInvocation, "invoke_workflow_app", return_value={"ok": True}) + + result = PluginAppBackwardsInvocation.invoke_app( + MagicMock(), + app_id="app", + user_id="wecom-sender-1", + tenant_id="tenant", + conversation_id="", + query=None, + stream=True, + inputs={}, + files=[], + ) + + assert result == {"ok": True} + get_or_create.assert_called_once_with(app, user_id="wecom-sender-1") + assert route.call_args.args[2] is end_user + def test_get_app_returns_app(self, mocker: MockerFixture): app_obj = MagicMock(id="app") self.patch_create_session(mocker, return_value=app_obj) diff --git a/api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py b/api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py index 7b723acc812..a6f80505e62 100644 --- a/api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py +++ b/api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py @@ -278,6 +278,59 @@ class TestPluginModelRuntime: stream=False, ) + def test_invoke_llm_forwards_string_app_id_from_request_metadata(self) -> None: + client = Mock(spec=PluginModelClient) + client.invoke_llm.return_value = iter([]) + runtime = PluginModelRuntime(tenant_id="tenant", user_id="user", client=client, plugin_service=PluginService) + + result = runtime.invoke_llm( + provider="langgenius/openai/openai", + model="gpt-4o-mini", + credentials={"api_key": "secret"}, + model_parameters={"temperature": 0.3}, + prompt_messages=[], + tools=None, + stop=None, + stream=True, + request_metadata={"app_id": "app-1"}, + ) + + assert list(result) == [] + client.invoke_llm.assert_called_once_with( + tenant_id="tenant", + user_id="user", + plugin_id="langgenius/openai", + provider="openai", + model="gpt-4o-mini", + credentials={"api_key": "secret"}, + model_parameters={"temperature": 0.3}, + prompt_messages=[], + tools=None, + stop=None, + stream=True, + app_id="app-1", + ) + + def test_invoke_llm_ignores_non_string_app_id_request_metadata(self) -> None: + client = Mock(spec=PluginModelClient) + client.invoke_llm.return_value = iter([]) + runtime = PluginModelRuntime(tenant_id="tenant", user_id="user", client=client, plugin_service=PluginService) + + result = runtime.invoke_llm( + provider="langgenius/openai/openai", + model="gpt-4o-mini", + credentials={"api_key": "secret"}, + model_parameters={"temperature": 0.3}, + prompt_messages=[], + tools=None, + stop=None, + stream=True, + request_metadata={"app_id": 123}, + ) + + assert result is client.invoke_llm.return_value + assert "app_id" not in client.invoke_llm.call_args.kwargs + def test_invoke_llm_returns_plugin_stream_directly(self) -> None: client = Mock(spec=PluginModelClient) stream_result = iter([]) diff --git a/api/tests/unit_tests/core/plugin/test_plugin_manager.py b/api/tests/unit_tests/core/plugin/test_plugin_manager.py index 1c5e1b9c0ca..1aa019254ff 100644 --- a/api/tests/unit_tests/core/plugin/test_plugin_manager.py +++ b/api/tests/unit_tests/core/plugin/test_plugin_manager.py @@ -17,7 +17,7 @@ import pytest from packaging.version import Version from requests import HTTPError -from core.plugin.entities.bundle import PluginBundleDependency +from core.plugin.entities.bundle import PluginBundleDependency, PluginBundleDependencyType from core.plugin.entities.plugin import ( MissingPluginDependency, PluginCategory, @@ -581,11 +581,11 @@ class TestDependencyResolution: bundle_data = b"mock-bundle-data" mock_dependencies = [ PluginBundleDependency( - type=PluginBundleDependency.Type.Marketplace, + type=PluginBundleDependencyType.Marketplace, value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), ), PluginBundleDependency( - type=PluginBundleDependency.Type.Github, + type=PluginBundleDependencyType.Github, value=PluginBundleDependency.Github( repo_address="https://github.com/org/repo", repo="org/repo", @@ -603,8 +603,8 @@ class TestDependencyResolution: # Assert: Verify dependencies were extracted assert len(result) == 2 - assert result[0].type == PluginBundleDependency.Type.Marketplace - assert result[1].type == PluginBundleDependency.Type.Github + assert result[0].type == PluginBundleDependencyType.Marketplace + assert result[1].type == PluginBundleDependencyType.Github mock_request.assert_called_once() def test_fetch_missing_dependencies(self, plugin_installer): @@ -1129,7 +1129,7 @@ class TestPluginBundleOperations: bundle_data = b"mock-marketplace-bundle" mock_dependencies = [ PluginBundleDependency( - type=PluginBundleDependency.Type.Marketplace, + type=PluginBundleDependencyType.Marketplace, value=PluginBundleDependency.Marketplace( organization="langgenius", plugin="search-tool", version="1.2.0" ), @@ -1142,7 +1142,7 @@ class TestPluginBundleOperations: # Assert: Verify marketplace dependency was extracted assert len(result) == 1 - assert result[0].type == PluginBundleDependency.Type.Marketplace + assert result[0].type == PluginBundleDependencyType.Marketplace assert isinstance(result[0].value, PluginBundleDependency.Marketplace) assert result[0].value.organization == "langgenius" assert result[0].value.plugin == "search-tool" @@ -1158,7 +1158,7 @@ class TestPluginBundleOperations: bundle_data = b"mock-github-bundle" mock_dependencies = [ PluginBundleDependency( - type=PluginBundleDependency.Type.Github, + type=PluginBundleDependencyType.Github, value=PluginBundleDependency.Github( repo_address="https://github.com/example/plugin", repo="example/plugin", @@ -1174,7 +1174,7 @@ class TestPluginBundleOperations: # Assert: Verify GitHub dependency was extracted assert len(result) == 1 - assert result[0].type == PluginBundleDependency.Type.Github + assert result[0].type == PluginBundleDependencyType.Github assert isinstance(result[0].value, PluginBundleDependency.Github) assert result[0].value.repo == "example/plugin" assert result[0].value.release == "v2.0.0" @@ -1204,7 +1204,7 @@ class TestPluginBundleOperations: mock_dependencies = [ PluginBundleDependency( - type=PluginBundleDependency.Type.Package, + type=PluginBundleDependencyType.Package, value=PluginBundleDependency.Package( unique_identifier="org/bundled-plugin/1.5.0", manifest=mock_manifest ), @@ -1217,7 +1217,7 @@ class TestPluginBundleOperations: # Assert: Verify package dependency was extracted with manifest assert len(result) == 1 - assert result[0].type == PluginBundleDependency.Type.Package + assert result[0].type == PluginBundleDependencyType.Package assert isinstance(result[0].value, PluginBundleDependency.Package) assert result[0].value.unique_identifier == "org/bundled-plugin/1.5.0" assert result[0].value.manifest.name == "bundled-plugin" @@ -1233,11 +1233,11 @@ class TestPluginBundleOperations: bundle_data = b"mock-mixed-bundle" mock_dependencies = [ PluginBundleDependency( - type=PluginBundleDependency.Type.Marketplace, + type=PluginBundleDependencyType.Marketplace, value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), ), PluginBundleDependency( - type=PluginBundleDependency.Type.Github, + type=PluginBundleDependencyType.Github, value=PluginBundleDependency.Github( repo_address="https://github.com/org2/plugin2", repo="org2/plugin2", @@ -1253,8 +1253,8 @@ class TestPluginBundleOperations: # Assert: Verify all dependency types were extracted assert len(result) == 2 - assert result[0].type == PluginBundleDependency.Type.Marketplace - assert result[1].type == PluginBundleDependency.Type.Github + assert result[0].type == PluginBundleDependencyType.Marketplace + assert result[1].type == PluginBundleDependencyType.Github class TestPluginTaskStatusTransitions: diff --git a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py b/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py index 7c672570bfa..d8452d91e2c 100644 --- a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py +++ b/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py @@ -227,13 +227,12 @@ class TestRetrievalServiceInternals: @patch("core.rag.datasource.retrieval_service.ExternalDatasetService.fetch_external_knowledge_retrieval") @patch("core.rag.datasource.retrieval_service.MetadataFilteringCondition.model_validate") - @patch("core.rag.datasource.retrieval_service.db.session.scalar") - def test_external_retrieve_with_metadata_conditions(self, mock_scalar, mock_validate, mock_fetch): - mock_scalar.return_value = SimpleNamespace(tenant_id="tenant-1") + def test_external_retrieve_with_metadata_conditions(self, mock_validate, mock_fetch): mock_validate.return_value = "validated-condition" expected_documents = [create_mock_document("external-doc", "external-1", 0.8, provider="external")] mock_fetch.return_value = expected_documents session = MagicMock() + session.scalar.return_value = SimpleNamespace(tenant_id="tenant-1") results = RetrievalService.external_retrieve( session=session, @@ -246,19 +245,19 @@ class TestRetrievalServiceInternals: assert results == expected_documents mock_validate.assert_called_once() mock_fetch.assert_called_once_with( - session, - "tenant-1", - "dataset-1", - "test query", - {"top_k": 3}, + tenant_id="tenant-1", + dataset_id="dataset-1", + query="test query", + external_retrieval_parameters={"top_k": 3}, metadata_condition="validated-condition", + session=session, ) - @patch("core.rag.datasource.retrieval_service.db.session.scalar") - def test_external_retrieve_returns_empty_when_dataset_not_found(self, mock_scalar): - mock_scalar.return_value = None + def test_external_retrieve_returns_empty_when_dataset_not_found(self): + session = MagicMock() + session.scalar.return_value = None - results = RetrievalService.external_retrieve(session=MagicMock(), dataset_id="missing", query="q") + results = RetrievalService.external_retrieve(session=session, dataset_id="missing", query="q") assert results == [] diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py index 302ababb48f..f5761b5ba3d 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py @@ -209,7 +209,7 @@ class TestParagraphIndexProcessor: vector = mock_vector_cls.return_value processor.clean(dataset, ["node-1"], delete_summaries=True) - mock_summary.assert_called_once_with(dataset, ["seg-1"]) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=["seg-1"]) vector.delete_by_ids.assert_called_once_with(["node-1"]) def test_clean_economy_deletes_summaries_and_keywords( @@ -225,7 +225,7 @@ class TestParagraphIndexProcessor: ): processor.clean(dataset, None, delete_summaries=True) - mock_summary.assert_called_once_with(dataset, None) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=None) mock_keyword_cls.return_value.delete.assert_called_once() def test_clean_deletes_keywords_by_ids(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None: diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py index 7d339a7701f..672764e5336 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py @@ -278,7 +278,7 @@ class TestParentChildIndexProcessor: ): processor.clean(dataset, ["node-1"], delete_summaries=True, precomputed_child_node_ids=[]) - mock_summary.assert_called_once_with(dataset, ["seg-1"]) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=["seg-1"]) def test_clean_deletes_all_summaries_when_node_ids_missing( self, processor: ParentChildIndexProcessor, dataset: Mock @@ -291,7 +291,7 @@ class TestParentChildIndexProcessor: ): processor.clean(dataset, None, delete_summaries=True) - mock_summary.assert_called_once_with(dataset, None) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=None) def test_split_child_nodes_requires_subchunk_segmentation(self, processor: ParentChildIndexProcessor) -> None: rules = Rule(subchunk_segmentation=None) diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py index 6e5a4fabbb0..5dde1623d2d 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py @@ -243,7 +243,7 @@ class TestQAIndexProcessor: vector = mock_vector_cls.return_value processor.clean(dataset, ["node-1"], delete_summaries=True) - mock_summary.assert_called_once_with(dataset, ["seg-1"]) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=["seg-1"]) vector.delete_by_ids.assert_called_once_with(["node-1"]) def test_clean_handles_dataset_wide_cleanup(self, processor: QAIndexProcessor, dataset: Mock) -> None: @@ -256,7 +256,7 @@ class TestQAIndexProcessor: vector = mock_vector_cls.return_value processor.clean(dataset, None, delete_summaries=True) - mock_summary.assert_called_once_with(dataset, None) + mock_summary.assert_called_once_with(dataset=dataset, segment_ids=None) vector.delete.assert_called_once() def test_index_adds_documents_and_vectors_for_high_quality( diff --git a/api/tests/unit_tests/core/tools/test_tool_provider_controller.py b/api/tests/unit_tests/core/tools/test_tool_provider_controller.py index 9648305289f..d380eab7a38 100644 --- a/api/tests/unit_tests/core/tools/test_tool_provider_controller.py +++ b/api/tests/unit_tests/core/tools/test_tool_provider_controller.py @@ -5,7 +5,7 @@ from typing import Any, override import pytest -from core.entities.provider_entities import ProviderConfig +from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.__base.tool import Tool from core.tools.__base.tool_provider import ToolProviderController from core.tools.__base.tool_runtime import ToolRuntime @@ -66,7 +66,7 @@ def _provider_identity() -> ToolProviderIdentity: def test_tool_provider_controller_get_credentials_schema_returns_deep_copy(): entity = ToolProviderEntity( identity=_provider_identity(), - credentials_schema=[ProviderConfig(type=ProviderConfig.Type.TEXT_INPUT, name="api_key", required=False)], + credentials_schema=[ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="api_key", required=False)], ) controller = _DummyController(entity=entity) @@ -88,10 +88,10 @@ def test_validate_credentials_format_covers_required_default_and_type_rules(): entity = ToolProviderEntity( identity=_provider_identity(), credentials_schema=[ - ProviderConfig(type=ProviderConfig.Type.TEXT_INPUT, name="required_text", required=True), - ProviderConfig(type=ProviderConfig.Type.SECRET_INPUT, name="secret", required=False), - ProviderConfig(type=ProviderConfig.Type.SELECT, name="choice", required=False, options=select_options), - ProviderConfig(type=ProviderConfig.Type.TEXT_INPUT, name="with_default", required=False, default="x"), + ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="required_text", required=True), + ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="secret", required=False), + ProviderConfig(type=ProviderConfigType.SELECT, name="choice", required=False, options=select_options), + ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="with_default", required=False, default="x"), ], ) controller = _DummyController(entity=entity) diff --git a/api/tests/unit_tests/core/tools/utils/test_encryption.py b/api/tests/unit_tests/core/tools/utils/test_encryption.py index ce77473dbdb..d4401488811 100644 --- a/api/tests/unit_tests/core/tools/utils/test_encryption.py +++ b/api/tests/unit_tests/core/tools/utils/test_encryption.py @@ -5,7 +5,7 @@ from unittest.mock import Mock, patch import pytest -from core.entities.provider_entities import BasicProviderConfig +from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType from core.helper.provider_encryption import ProviderConfigEncrypter from core.tools.utils.encryption import create_tool_provider_encrypter @@ -31,7 +31,7 @@ def secret_field() -> BasicProviderConfig: """A SECRET_INPUT field named 'password'.""" return BasicProviderConfig( name="password", - type=BasicProviderConfig.Type.SECRET_INPUT, + type=ProviderConfigType.SECRET_INPUT, ) @@ -40,7 +40,7 @@ def normal_field() -> BasicProviderConfig: """A TEXT_INPUT field named 'username'.""" return BasicProviderConfig( name="username", - type=BasicProviderConfig.Type.TEXT_INPUT, + type=ProviderConfigType.TEXT_INPUT, ) @@ -185,7 +185,7 @@ def test_decrypt_swallow_exception_and_keep_original(encrypter_obj): def test_create_tool_provider_encrypter_builds_cache_and_encrypter(): - basic_config = BasicProviderConfig(name="key", type=BasicProviderConfig.Type.TEXT_INPUT) + basic_config = BasicProviderConfig(name="key", type=ProviderConfigType.TEXT_INPUT) credential_schema_item = SimpleNamespace(to_basic_provider_config=lambda: basic_config) controller = SimpleNamespace( provider_type=SimpleNamespace(value="builtin"), diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py index 99448f7f5a9..407b5e59f7f 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py @@ -498,6 +498,37 @@ def test_agent_node_failed_run_without_session_store_skips_mark_cleaned(): assert "session_snapshot_cleaned_on_failure" not in agent_backend +def test_agent_node_failed_run_enqueues_backend_cleanup_before_local_retirement(monkeypatch): + store = FakeSessionStore() + store.loaded_session = StoredWorkflowAgentSession( + scope=_pending_session(CompositorSessionSnapshot(layers=[])).scope, + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="stored-run-1", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + queued_payloads: list[dict[str, object]] = [] + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.agent_node.cleanup_workflow_agent_runtime_session.delay", + lambda payload: queued_payloads.append(payload), + ) + + events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=store)._run()) + + assert len(events) == 1 + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert store.cleaned[0][1] == "fake-run-1" + assert store.cleaned[0][0].workflow_run_id == "workflow-run-1" + assert store.cleaned[0][0].node_id == "agent-node" + assert len(queued_payloads) == 1 + assert ( + queued_payloads[0]["idempotency_key"] + == "tenant-1:workflow-run-1:agent-node:binding-1:workflow-agent-failure-cleanup:stored-run-1:fake-run-1" + ) + assert queued_payloads[0]["metadata"]["previous_agent_backend_run_id"] == "stored-run-1" + assert queued_payloads[0]["metadata"]["failed_agent_backend_run_id"] == "fake-run-1" + + def test_agent_node_paused_run_requests_workflow_pause_and_persists_snapshot(): store = FakeSessionStore() node = _node(scenario=FakeAgentBackendScenario.PAUSED, session_store=store) diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py index 10f321d3a2a..9d0089e8e08 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py @@ -227,6 +227,15 @@ def _previous_node_prompt_payload(result, selector: str) -> object: raise AssertionError(f"missing prompt payload for {selector}") +def _uploaded_workflow_files_prompt_payload(result) -> object: + prefix = " - sys.files: " + user_prompt = _workflow_user_prompt(result) + for line in user_prompt.splitlines(): + if line.startswith(prefix): + return json.loads(line.removeprefix(prefix)) + raise AssertionError("missing prompt payload for sys.files") + + def test_builds_create_run_request_from_agent_soul_and_node_job(): result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(_context()) @@ -420,6 +429,8 @@ def test_builds_workflow_run_request_with_file_output_schema_and_reserved_metada assert "final_output.report" in output_description assert "never invent the `reference` value" in output_description assert "Do not call `final_output` before the upload command succeeds" in output_description + assert "accepted file-mapping shape and the returned `reference`" in output_description + assert "include the returned `download_url` in that reply" in output_description assert output_schema["properties"]["confidence"]["type"] == "number" assert output_schema["required"] == ["report"] assert layers[DIFY_AGENT_MODEL_LAYER_ID]["config"]["model_settings"] == {"temperature": 0.2} @@ -719,7 +730,11 @@ def test_build_maps_agent_soul_knowledge_to_knowledge_layer_config(): "top_k": 6, "score_threshold": 0.4, "reranking_model": {"provider": "cohere", "model": "rerank-v3"}, - "weights": {"weight_type": "weighted_score", "vector_setting": {"vector_weight": 0.7}}, + "weights": { + "weight_type": "weighted_score", + "vector_setting": {"vector_weight": 0.7}, + "keyword_setting": {"keyword_weight": 0.3}, + }, }, "metadata_filtering": { "mode": "manual", @@ -786,7 +801,10 @@ def test_build_maps_agent_soul_knowledge_to_knowledge_layer_config(): "reranking_mode": "reranking_model", "reranking_enable": True, "reranking_model": {"provider": "cohere", "model": "rerank-v3"}, - "weights": {"weight_type": "weighted_score", "vector_setting": {"vector_weight": 0.7}}, + "weights": { + "vector_setting": {"vector_weight": 0.7}, + "keyword_setting": {"keyword_weight": 0.3}, + }, "model": None, }, "metadata_filtering": { @@ -1165,6 +1183,37 @@ def test_previous_node_file_output_uses_agent_stub_download_mapping_in_workflow_ } +def test_previous_node_file_mapping_strips_extra_fields_in_workflow_context(): + file_reference = build_file_reference(record_id="tool-file-1") + + class FileMappingVariablePool(FakeVariablePool): + def get(self, selector): + if list(selector) == ["previous-node", "report"]: + return SimpleNamespace( + value={ + "filename": "report.pdf", + "transfer_method": "tool_file", + "reference": file_reference, + "external": True, + } + ) + return super().get(selector) + + context = replace(_context(), variable_pool=FileMappingVariablePool()) + context.binding.node_job_config = WorkflowNodeJobConfig.model_validate( + { + "workflow_prompt": "Review {{#previous-node.report#}} before responding.", + } + ) + + result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(context) + + assert _previous_node_prompt_payload(result, "previous-node.report") == { + "transfer_method": "tool_file", + "reference": file_reference, + } + + def test_scalar_previous_node_output_appears_in_workflow_context_section(): context = _context() context.binding.node_job_config = WorkflowNodeJobConfig.model_validate( @@ -1252,6 +1301,48 @@ def test_previous_node_file_array_uses_agent_stub_download_mappings_in_workflow_ ] +def test_uploaded_workflow_files_are_included_without_prompt_marker(): + file_reference = build_file_reference(record_id="uploaded-file-1") + + class UploadedFilesVariablePool(FakeVariablePool): + def get(self, selector): + if list(selector) == ["sys", "files"]: + return ArrayFileSegment( + value=[ + File( + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=file_reference, + remote_url=None, + filename="requirements.pdf", + extension=".pdf", + mime_type="application/pdf", + size=12, + ) + ] + ) + return super().get(selector) + + context = replace(_context(), variable_pool=UploadedFilesVariablePool()) + context.binding.node_job_config = WorkflowNodeJobConfig.model_validate( + { + "workflow_prompt": "Answer the user's question.", + } + ) + + result = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()).build(context) + + user_prompt = _workflow_user_prompt(result) + assert "- Uploaded workflow files:" in user_prompt + assert _uploaded_workflow_files_prompt_payload(result) == [ + { + "transfer_method": "local_file", + "reference": file_reference, + } + ] + assert "Previous node outputs:" not in user_prompt + + def test_previous_node_remote_url_file_mapping_is_not_truncated_in_workflow_context(): remote_url = "https://example.com/" + ("a" * 2100) + ".pdf" diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py index 67e7bfafca7..eda2962198e 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_session_cleanup_layer.py @@ -1,15 +1,15 @@ -from datetime import UTC from typing import cast import pytest from agenton.compositor import CompositorSessionSnapshot from agenton.compositor.schemas import LayerSessionSnapshot from agenton.layers.base import LifecycleState -from dify_agent.protocol import CancelRunRequest, RunEvent, RunStatusResponse +from dify_agent.protocol import RuntimeLayerSpec -from clients.agent_backend import AgentBackendRunRequestBuilder, FakeAgentBackendRunClient, RuntimeLayerSpec -from clients.agent_backend.errors import AgentBackendHTTPError -from core.workflow.nodes.agent_v2.session_cleanup_layer import WorkflowAgentSessionCleanupLayer +from core.workflow.nodes.agent_v2.session_cleanup_layer import ( + WorkflowAgentSessionCleanupLayer, + build_workflow_agent_session_cleanup_layer, +) from core.workflow.nodes.agent_v2.session_store import ( StoredWorkflowAgentSession, WorkflowAgentRuntimeSessionStore, @@ -37,13 +37,21 @@ def _layer_snapshot(name: str) -> LayerSessionSnapshot: ) -def _stored_session(scope: WorkflowAgentSessionScope, *, index: int = 1) -> StoredWorkflowAgentSession: - """A typical stored session with prompt + execution_context + history + llm specs. +def _default_scope() -> WorkflowAgentSessionScope: + return WorkflowAgentSessionScope( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_run_id="workflow-run-1", + node_id="agent-node", + node_execution_id="node-exec-1", + binding_id="binding-1", + agent_id="agent-1", + agent_config_snapshot_id="snapshot-1", + ) - The LLM layer is *not* in ``runtime_layer_specs`` because the cleanup - contract excludes credential-bearing plugin layers, but it *is* present in - the saved snapshot so the layer's filter logic gets exercised. - """ + +def _stored_session(scope: WorkflowAgentSessionScope, *, index: int = 1) -> StoredWorkflowAgentSession: return StoredWorkflowAgentSession( scope=scope, session_snapshot=CompositorSessionSnapshot( @@ -64,8 +72,6 @@ def _stored_session(scope: WorkflowAgentSessionScope, *, index: int = 1) -> Stor class FakeSessionStore: - """In-memory stand-in for ``WorkflowAgentRuntimeSessionStore``.""" - def __init__(self, *, stored: list[StoredWorkflowAgentSession] | None = None) -> None: self._stored = stored if stored is not None else [_stored_session(_default_scope())] self.list_calls: list[str] = [] @@ -79,69 +85,7 @@ class FakeSessionStore: self.cleaned.append((scope, backend_run_id)) -def _default_scope() -> WorkflowAgentSessionScope: - return WorkflowAgentSessionScope( - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - workflow_run_id="workflow-run-1", - node_id="agent-node", - node_execution_id="node-exec-1", - binding_id="binding-1", - agent_id="agent-1", - agent_config_snapshot_id="snapshot-1", - ) - - -class _WaitableFakeAgentBackendRunClient(FakeAgentBackendRunClient): - """``FakeAgentBackendRunClient`` plus the ``wait_run`` hook the layer needs.""" - - def __init__( - self, - *, - run_id: str = "cleanup-run-1", - wait_status: str = "succeeded", - wait_error: str | None = None, - wait_raises: Exception | None = None, - ) -> None: - super().__init__(run_id=run_id) - self._wait_status = wait_status - self._wait_error = wait_error - self._wait_raises = wait_raises - self.wait_calls: list[tuple[str, float | None]] = [] - - def wait_run(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse: - self.wait_calls.append((run_id, timeout_seconds)) - if self._wait_raises is not None: - raise self._wait_raises - from datetime import datetime - - return RunStatusResponse( - run_id=run_id, - status=cast(object, self._wait_status), # protocol Literal; cast keeps tests flexible - created_at=datetime(2026, 1, 1, tzinfo=UTC), - updated_at=datetime(2026, 1, 1, tzinfo=UTC), - error=self._wait_error, - ) - - # Inherit ``create_run`` from FakeAgentBackendRunClient; the missing protocol - # methods below are stub-only because the cleanup layer never calls them. - def cancel_run(self, run_id: str, request: CancelRunRequest | None = None): # pragma: no cover - del run_id, request - raise NotImplementedError - - def stream_events(self, run_id: str, *, after: str | None = None): # pragma: no cover - del run_id, after - if False: - yield cast(RunEvent, None) - - -def _build_layer( - *, - session_store: FakeSessionStore, - agent_backend_client: _WaitableFakeAgentBackendRunClient, - http_cleanup_supported: bool = True, -) -> WorkflowAgentSessionCleanupLayer: +def _build_layer(*, session_store: FakeSessionStore) -> WorkflowAgentSessionCleanupLayer: variable_pool = VariablePool.from_bootstrap( system_variables=build_system_variables(workflow_execution_id="workflow-run-1"), user_inputs={}, @@ -150,12 +94,7 @@ def _build_layer( runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) layer = WorkflowAgentSessionCleanupLayer( session_store=cast(WorkflowAgentRuntimeSessionStore, session_store), - request_builder=AgentBackendRunRequestBuilder(), - agent_backend_client=agent_backend_client, ) - # Tests opt in to the future HTTP-cleanup branch; the production default - # (False) is exercised by the dedicated tests below. - layer._HTTP_CLEANUP_SUPPORTED = http_cleanup_supported # type: ignore[reportPrivateUsage] layer.initialize(ReadOnlyGraphRuntimeStateWrapper(runtime_state), cast(CommandChannel, object())) return layer @@ -170,30 +109,22 @@ def _build_layer( ], ids=["succeeded", "partial_succeeded", "failed", "aborted"], ) -def test_cleanup_layer_triggers_cleanup_only_run_on_each_terminal_event(terminal_event): +def test_cleanup_layer_enqueues_cleanup_and_marks_cleaned_on_terminal_events(monkeypatch, terminal_event): session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient() - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) + queued_payloads: list[dict[str, object]] = [] + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", + lambda payload: queued_payloads.append(payload), + ) + layer = _build_layer(session_store=session_store) layer.on_event(terminal_event) assert session_store.list_calls == ["workflow-run-1"] - assert agent_backend_client.request is not None - # Cleanup composition replays the persisted (non-plugin) layer specs so the - # agent backend's snapshot-vs-composition name match succeeds. - layer_names = [layer.name for layer in agent_backend_client.request.composition.layers] - assert layer_names == ["workflow_node_job_prompt", "execution_context", "history"] - assert agent_backend_client.request.on_exit.default.value == "delete" - assert agent_backend_client.request.metadata["agent_backend_lifecycle"] == "session_cleanup" - # Snapshot is filtered to drop the plugin layer entry so names match the - # cleanup composition. - assert agent_backend_client.request.session_snapshot is not None - snapshot_names = [layer.name for layer in agent_backend_client.request.session_snapshot.layers] - assert snapshot_names == ["workflow_node_job_prompt", "execution_context", "history"] - # The layer waited for terminal status and the run succeeded, so the row - # is marked CLEANED with the cleanup run id. - assert agent_backend_client.wait_calls - assert session_store.cleaned == [(_default_scope(), "cleanup-run-1")] + assert len(queued_payloads) == 1 + assert queued_payloads[0]["metadata"]["workflow_run_id"] == "workflow-run-1" + assert queued_payloads[0]["metadata"]["previous_agent_backend_run_id"] == "agent-run-1" + assert session_store.cleaned == [(_default_scope(), "agent-run-1")] @pytest.mark.parametrize( @@ -204,91 +135,79 @@ def test_cleanup_layer_triggers_cleanup_only_run_on_each_terminal_event(terminal ], ids=["started", "paused"], ) -def test_cleanup_layer_ignores_non_terminal_events(non_terminal_event): +def test_cleanup_layer_ignores_non_terminal_events(monkeypatch, non_terminal_event): session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient() - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) + queued_payloads: list[dict[str, object]] = [] + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", + lambda payload: queued_payloads.append(payload), + ) + layer = _build_layer(session_store=session_store) layer.on_event(non_terminal_event) assert session_store.list_calls == [] - assert agent_backend_client.request is None + assert queued_payloads == [] assert session_store.cleaned == [] -def test_cleanup_layer_does_not_mark_cleaned_when_cleanup_run_fails(): - """Trap D: cleanup-only run goes ``run_failed`` (e.g. snapshot validation - error) — the layer must leave the row ACTIVE so it can be retried instead - of silently leaking suspended agent-backend layers.""" - session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient( - wait_status="failed", - wait_error="snapshot mismatch", +def test_cleanup_layer_marks_cleaned_even_when_specs_are_missing(monkeypatch, caplog: pytest.LogCaptureFixture): + scope = _default_scope() + session_store = FakeSessionStore( + stored=[ + StoredWorkflowAgentSession( + scope=scope, + session_snapshot=CompositorSessionSnapshot(layers=[_layer_snapshot("history")]), + backend_run_id="legacy-run", + runtime_layer_specs=[], + ) + ] ) - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) + queued_payloads: list[dict[str, object]] = [] + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", + lambda payload: queued_payloads.append(payload), + ) + layer = _build_layer(session_store=session_store) layer.on_event(GraphRunSucceededEvent(outputs={})) - assert agent_backend_client.wait_calls - assert session_store.cleaned == [] + assert queued_payloads == [] + assert session_store.cleaned == [(scope, "legacy-run")] + assert any("no runtime_layer_specs persisted" in record.message for record in caplog.records) -def test_cleanup_layer_does_not_mark_cleaned_when_wait_raises(): +def test_cleanup_layer_marks_cleaned_even_when_enqueue_fails(monkeypatch): session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient( - wait_raises=AgentBackendHTTPError("boom", status_code=500, detail=None), + + def _explode(_payload: dict[str, object]) -> None: + raise RuntimeError("queue down") + + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", + _explode, ) - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) + layer = _build_layer(session_store=session_store) layer.on_event(GraphRunSucceededEvent(outputs={})) - assert session_store.cleaned == [] - - -def test_cleanup_layer_marks_cleaned_locally_when_http_cleanup_disabled(): - """Production default: dify-agent has no cleanup-only run mode yet, so the - layer must retire the local row without issuing a doomed HTTP request that - would crash inside the agent backend's runner on the missing LLM layer.""" - session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient() - layer = _build_layer( - session_store=session_store, - agent_backend_client=agent_backend_client, - http_cleanup_supported=False, - ) - - layer.on_event(GraphRunSucceededEvent(outputs={})) - - # No HTTP call goes out — the trap is avoided entirely. - assert agent_backend_client.request is None - assert agent_backend_client.wait_calls == [] - # Local row is still retired so a workflow loop cannot resume from stale state. assert session_store.cleaned == [(_default_scope(), "agent-run-1")] -def test_cleanup_layer_skips_sessions_without_persisted_specs(): - """Backwards-compatible safety net: a row written before A.1 landed has - no runtime_layer_specs, so cleanup would unavoidably hit the snapshot- - validation trap. The layer must skip such rows instead of issuing a - doomed request.""" - scope = _default_scope() - legacy_session = StoredWorkflowAgentSession( - scope=scope, - session_snapshot=CompositorSessionSnapshot(layers=[_layer_snapshot("history")]), - backend_run_id="legacy-run", - runtime_layer_specs=[], - ) - session_store = FakeSessionStore(stored=[legacy_session]) - agent_backend_client = _WaitableFakeAgentBackendRunClient() - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) +def test_cleanup_layer_does_not_raise_when_mark_cleaned_fails(monkeypatch): + session_store = FakeSessionStore() + + def _explode(*, scope: WorkflowAgentSessionScope, backend_run_id: str | None = None) -> None: + del scope, backend_run_id + raise RuntimeError("cleanup bookkeeping failed") + + monkeypatch.setattr(session_store, "mark_cleaned", _explode) + layer = _build_layer(session_store=session_store) layer.on_event(GraphRunSucceededEvent(outputs={})) - assert agent_backend_client.request is None - assert session_store.cleaned == [] - -def test_cleanup_layer_fans_out_to_every_active_session(): +def test_cleanup_layer_fans_out_to_every_active_session(monkeypatch): scopes = [ WorkflowAgentSessionScope( tenant_id="tenant-1", @@ -304,109 +223,34 @@ def test_cleanup_layer_fans_out_to_every_active_session(): for i in range(3) ] session_store = FakeSessionStore(stored=[_stored_session(scope, index=i) for i, scope in enumerate(scopes, 1)]) - agent_backend_client = _WaitableFakeAgentBackendRunClient(run_id="cleanup-run-many") - layer = _build_layer(session_store=session_store, agent_backend_client=agent_backend_client) + queued_payloads: list[dict[str, object]] = [] + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_cleanup_layer.cleanup_workflow_agent_runtime_session.delay", + lambda payload: queued_payloads.append(payload), + ) + layer = _build_layer(session_store=session_store) layer.on_event(GraphRunSucceededEvent(outputs={})) - # One cleanup row per stored ACTIVE session, all marked cleaned with the - # backend run id returned by the agent backend client. + assert len(queued_payloads) == 3 assert [entry[0] for entry in session_store.cleaned] == scopes - assert {entry[1] for entry in session_store.cleaned} == {"cleanup-run-many"} -def test_cleanup_layer_warns_when_http_enabled_but_client_missing(caplog: pytest.LogCaptureFixture): - """The HTTP cleanup branch must defensively skip when no client was wired. - - This is the deployment-misconfig path: ``_HTTP_CLEANUP_SUPPORTED`` was - flipped to ``True`` but ``AGENT_BACKEND_BASE_URL`` is unset, so the - factory returned ``None``. The layer must not crash and must not silently - retire the row — the warning surfaces the misconfig. - """ - import logging - +def test_cleanup_layer_skips_when_workflow_run_id_missing(caplog: pytest.LogCaptureFixture): session_store = FakeSessionStore() - layer = WorkflowAgentSessionCleanupLayer( - session_store=cast(WorkflowAgentRuntimeSessionStore, session_store), - request_builder=AgentBackendRunRequestBuilder(), - agent_backend_client=None, - ) - layer._HTTP_CLEANUP_SUPPORTED = True # type: ignore[reportPrivateUsage] - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="workflow-run-1"), - user_inputs={}, - conversation_variables=[], - ) + variable_pool = VariablePool.from_bootstrap(system_variables={}, user_inputs={}, conversation_variables=[]) runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) + layer = WorkflowAgentSessionCleanupLayer(session_store=cast(WorkflowAgentRuntimeSessionStore, session_store)) layer.initialize(ReadOnlyGraphRuntimeStateWrapper(runtime_state), cast(CommandChannel, object())) - with caplog.at_level(logging.WARNING): - layer.on_event(GraphRunSucceededEvent(outputs={})) - - assert session_store.cleaned == [] - assert any("no agent backend client is wired in" in record.message for record in caplog.records) - - -def test_cleanup_layer_skips_workflow_terminal_when_workflow_run_id_missing(caplog: pytest.LogCaptureFixture): - """``workflow_run_id`` is the keying field; without it the fanout cannot - target a row, so the layer logs a warning and bails.""" - import logging - - session_store = FakeSessionStore() - agent_backend_client = _WaitableFakeAgentBackendRunClient() - layer = WorkflowAgentSessionCleanupLayer( - session_store=cast(WorkflowAgentRuntimeSessionStore, session_store), - request_builder=AgentBackendRunRequestBuilder(), - agent_backend_client=agent_backend_client, - ) - # Bootstrap *without* a workflow_execution_id system variable. - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id=""), - user_inputs={}, - conversation_variables=[], - ) - runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=0.0) - layer.initialize(ReadOnlyGraphRuntimeStateWrapper(runtime_state), cast(CommandChannel, object())) - - with caplog.at_level(logging.WARNING): - layer.on_event(GraphRunSucceededEvent(outputs={})) + layer.on_event(GraphRunSucceededEvent(outputs={})) assert session_store.list_calls == [] assert session_store.cleaned == [] assert any("workflow_run_id is missing" in record.message for record in caplog.records) -def test_build_workflow_agent_session_cleanup_layer_returns_layer_without_client_when_unconfigured( - monkeypatch, -): - """The production builder must pass ``None`` for the agent backend client - when neither AGENT_BACKEND_BASE_URL nor AGENT_BACKEND_USE_FAKE is set, so - that unit-test environments without backend config don't crash at runner - construction.""" - from configs import dify_config - from core.workflow.nodes.agent_v2.session_cleanup_layer import ( - build_workflow_agent_session_cleanup_layer, - ) - - monkeypatch.setattr(dify_config, "AGENT_BACKEND_BASE_URL", None, raising=False) - monkeypatch.setattr(dify_config, "AGENT_BACKEND_USE_FAKE", False, raising=False) - +def test_build_workflow_agent_session_cleanup_layer_returns_layer() -> None: layer = build_workflow_agent_session_cleanup_layer() - assert layer._agent_backend_client is None # type: ignore[reportPrivateUsage] - -def test_build_workflow_agent_session_cleanup_layer_returns_layer_with_fake_client(monkeypatch): - """With ``AGENT_BACKEND_USE_FAKE`` enabled the helper wires in the - deterministic fake client without needing a base_url.""" - from clients.agent_backend.fake_client import FakeAgentBackendRunClient - from configs import dify_config - from core.workflow.nodes.agent_v2.session_cleanup_layer import ( - build_workflow_agent_session_cleanup_layer, - ) - - monkeypatch.setattr(dify_config, "AGENT_BACKEND_BASE_URL", None, raising=False) - monkeypatch.setattr(dify_config, "AGENT_BACKEND_USE_FAKE", True, raising=False) - monkeypatch.setattr(dify_config, "AGENT_BACKEND_FAKE_SCENARIO", "success", raising=False) - - layer = build_workflow_agent_session_cleanup_layer() - assert isinstance(layer._agent_backend_client, FakeAgentBackendRunClient) # type: ignore[reportPrivateUsage] + assert isinstance(layer, WorkflowAgentSessionCleanupLayer) diff --git a/api/tests/unit_tests/core/workflow/test_node_factory.py b/api/tests/unit_tests/core/workflow/test_node_factory.py index 6a70f2e9afb..e926f196192 100644 --- a/api/tests/unit_tests/core/workflow/test_node_factory.py +++ b/api/tests/unit_tests/core/workflow/test_node_factory.py @@ -723,6 +723,7 @@ class TestDifyNodeFactoryCreateNode: wrap_model.assert_called_once_with( node_data=node_data, model_instance=sentinel.model_instance, + request_metadata={"app_id": "app-id"}, ) assert kwargs["model_instance"] is wrapped_model_instance diff --git a/api/tests/unit_tests/core/workflow/test_node_runtime.py b/api/tests/unit_tests/core/workflow/test_node_runtime.py index afebc6ee38a..6190c7fb91c 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -171,7 +171,7 @@ def test_dify_prepared_llm_wraps_model_instance_calls() -> None: model_schema = _build_model_schema() model_instance = _ModelInstanceStub(model_schema=model_schema) model_type_instance = model_instance.model_type_instance - prepared = DifyPreparedLLM(model_instance) + prepared = DifyPreparedLLM(model_instance, request_metadata={"app_id": "app-id"}) assert prepared.provider == "langgenius/openai/openai" assert prepared.model_name == "gpt-4o-mini" @@ -197,6 +197,7 @@ def test_dify_prepared_llm_wraps_model_instance_calls() -> None: tools=[], stop=[], stream=False, + request_metadata={"app_id": "app-id"}, ) diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py b/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py index 3ccfdf76f5a..41037233b8c 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py @@ -13,6 +13,7 @@ from graphon.entities.base_node_data import BaseNodeData from graphon.enums import NodeType, WorkflowNodeExecutionStatus from graphon.errors import WorkflowNodeRunFailedError from graphon.file import File, FileTransferMethod, FileType +from graphon.filters import ResponseStreamFilter from graphon.graph import Graph from graphon.graph_events import GraphRunFailedEvent from graphon.model_runtime.entities.llm_entities import LLMMode, LLMUsage @@ -241,6 +242,37 @@ class TestWorkflowChildEngineBuilder: ) +def _build_minimal_workflow_entry( + monkeypatch: pytest.MonkeyPatch, + *, + response_stream_filter: ResponseStreamFilter | None = None, +) -> workflow_entry.WorkflowEntry: + """Construct a minimal WorkflowEntry with GraphEngine construction mocked out.""" + graph_engine = MagicMock() + graph_runtime_state = SimpleNamespace(execution_context=None) + + monkeypatch.setattr(workflow_entry, "capture_current_context", lambda: sentinel.execution_context) + monkeypatch.setattr(workflow_entry, "GraphEngine", MagicMock(return_value=graph_engine)) + monkeypatch.setattr(workflow_entry, "GraphEngineConfig", MagicMock(return_value=sentinel.graph_engine_config)) + monkeypatch.setattr(workflow_entry, "InMemoryChannel", MagicMock(return_value=sentinel.command_channel)) + monkeypatch.setattr(workflow_entry, "LLMQuotaLayer", MagicMock(return_value=sentinel.llm_quota_layer)) + + return workflow_entry.WorkflowEntry( + tenant_id="tenant-id", + app_id="app-id", + workflow_id="workflow-id", + graph_config={"nodes": [], "edges": []}, + graph=sentinel.graph, + user_id="user-id", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + variable_pool=sentinel.variable_pool, + graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, + ) + + class TestWorkflowEntryInit: def test_rejects_call_depth_above_limit(self): call_depth = workflow_entry.dify_config.WORKFLOW_CALL_MAX_DEPTH + 1 @@ -329,12 +361,24 @@ class TestWorkflowEntryInit: ((observability_layer,), {}), ] + def test_workflow_entry_stores_supplied_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: + supplied_filter = ResponseStreamFilter() + entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=supplied_filter) + + assert entry._response_stream_filter is supplied_filter + + def test_workflow_entry_defaults_to_fresh_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: + entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=None) + + assert isinstance(entry._response_stream_filter, ResponseStreamFilter) + class TestWorkflowEntryRun: def test_run_swallows_generate_task_stopped_errors(self): entry = object.__new__(workflow_entry.WorkflowEntry) entry.graph_engine = MagicMock() entry.graph_engine.run.side_effect = GenerateTaskStoppedError() + entry._response_stream_filter = ResponseStreamFilter() assert list(entry.run()) == [] @@ -373,6 +417,7 @@ class TestWorkflowEntryRun: def test_run_delegates_to_dify_event_iterator(self): entry = object.__new__(workflow_entry.WorkflowEntry) entry.graph_engine = sentinel.graph_engine + entry._response_stream_filter = sentinel.response_stream_filter with patch.object( workflow_entry, @@ -382,12 +427,13 @@ class TestWorkflowEntryRun: events = list(entry.run()) assert events == [sentinel.filtered_event] - iter_dify_graph_engine_events.assert_called_once_with(sentinel.graph_engine) + iter_dify_graph_engine_events.assert_called_once_with(sentinel.graph_engine, sentinel.response_stream_filter) def test_run_emits_failed_event_for_unexpected_errors(self): entry = object.__new__(workflow_entry.WorkflowEntry) entry.graph_engine = MagicMock() entry.graph_engine.run.side_effect = RuntimeError("boom") + entry._response_stream_filter = ResponseStreamFilter() events = list(entry.run()) diff --git a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py index 24c905c75c2..29ddf73e24d 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py +++ b/api/tests/unit_tests/enterprise/telemetry/test_enterprise_trace.py @@ -3,8 +3,9 @@ from __future__ import annotations import json +import logging from datetime import UTC, datetime -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock, patch import pytest @@ -475,13 +476,24 @@ class TestWorkflowTrace: assert span_call[1]["start_time"] == _T0 assert span_call[1]["end_time"] == _T1 - def test_emits_companion_log_with_event_name(self, trace_handler: EnterpriseOtelTrace, mock_exporter): - with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log: + def test_emits_companion_log_with_event_name( + self, + trace_handler: EnterpriseOtelTrace, + caplog: pytest.LogCaptureFixture, + ): + with caplog.at_level(logging.INFO, logger="dify.telemetry"): trace_handler._workflow_trace(make_workflow_info()) - mock_log.assert_called_once() - assert mock_log.call_args[1]["event_name"] == EnterpriseTelemetryEvent.WORKFLOW_RUN - assert mock_log.call_args[1]["tenant_id"] == "tenant-abc" + records = [record for record in caplog.records if record.name == "dify.telemetry"] + + assert len(records) == 1 + record = records[0] + + attrs = cast(dict[str, Any], record.__dict__["attributes"]) + + assert attrs["dify.event.name"] == EnterpriseTelemetryEvent.WORKFLOW_RUN + assert attrs["dify.event.signal"] == "span_detail" + assert record.__dict__["tenant_id"] == "tenant-abc" def test_companion_log_includes_content_when_enabled(self, trace_handler: EnterpriseOtelTrace, mock_exporter): mock_exporter.include_content = True diff --git a/api/tests/unit_tests/events/test_app_event_signals.py b/api/tests/unit_tests/events/test_app_event_signals.py index 29582a50f6d..a6059fadbcf 100644 --- a/api/tests/unit_tests/events/test_app_event_signals.py +++ b/api/tests/unit_tests/events/test_app_event_signals.py @@ -44,7 +44,7 @@ def _make_collector(target: list): @pytest.mark.usefixtures("mock_db", "_mock_deps") class TestAppWasDeletedSignal: - def test_sends_signal(self, app_model): + def test_sends_signal(self, app_model, mock_db): from events.app_event import app_was_deleted from services.app_service import AppService @@ -52,7 +52,7 @@ class TestAppWasDeletedSignal: handler = _make_collector(received) app_was_deleted.connect(handler) try: - AppService().delete_app(app_model) + AppService().delete_app(app_model, session=mock_db.session) finally: app_was_deleted.disconnect(handler) @@ -71,7 +71,7 @@ class TestAppWasDeletedSignal: mock_db.session.delete.side_effect = lambda _: call_order.append("db_delete") try: - AppService().delete_app(app_model) + AppService().delete_app(app_model, session=mock_db.session) finally: app_was_deleted.disconnect(handler) @@ -80,7 +80,7 @@ class TestAppWasDeletedSignal: @pytest.mark.usefixtures("mock_db") class TestAppWasUpdatedSignal: - def test_update_app(self, app_model): + def test_update_app(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -101,13 +101,14 @@ class TestAppWasUpdatedSignal: "use_icon_as_answer_icon": False, "max_active_requests": 0, }, + session=mock_db.session, ) finally: app_was_updated.disconnect(handler) assert received == [app_model] - def test_update_app_name(self, app_model): + def test_update_app_name(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -117,13 +118,13 @@ class TestAppWasUpdatedSignal: with patch("services.app_service.current_user", MagicMock(id="user-1")): try: - AppService().update_app_name(app_model, "New Name") + AppService().update_app_name(app_model, "New Name", session=mock_db.session) finally: app_was_updated.disconnect(handler) assert received == [app_model] - def test_update_app_icon(self, app_model): + def test_update_app_icon(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -133,13 +134,13 @@ class TestAppWasUpdatedSignal: with patch("services.app_service.current_user", MagicMock(id="user-1")): try: - AppService().update_app_icon(app_model, "🎉", "#000") + AppService().update_app_icon(app_model, "🎉", "#000", session=mock_db.session) finally: app_was_updated.disconnect(handler) assert received == [app_model] - def test_update_app_site_status_sends_when_changed(self, app_model): + def test_update_app_site_status_sends_when_changed(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -150,13 +151,13 @@ class TestAppWasUpdatedSignal: with patch("services.app_service.current_user", MagicMock(id="user-1")): try: app_model.enable_site = False - AppService().update_app_site_status(app_model, True) + AppService().update_app_site_status(app_model, True, session=mock_db.session) finally: app_was_updated.disconnect(handler) assert received == [app_model] - def test_update_app_site_status_skips_when_unchanged(self, app_model): + def test_update_app_site_status_skips_when_unchanged(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -166,13 +167,13 @@ class TestAppWasUpdatedSignal: try: app_model.enable_site = True - AppService().update_app_site_status(app_model, True) + AppService().update_app_site_status(app_model, True, session=mock_db.session) finally: app_was_updated.disconnect(handler) assert received == [] - def test_update_app_api_status_sends_when_changed(self, app_model): + def test_update_app_api_status_sends_when_changed(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -183,13 +184,13 @@ class TestAppWasUpdatedSignal: with patch("services.app_service.current_user", MagicMock(id="user-1")): try: app_model.enable_api = False - AppService().update_app_api_status(app_model, True) + AppService().update_app_api_status(app_model, True, session=mock_db.session) finally: app_was_updated.disconnect(handler) assert received == [app_model] - def test_update_app_api_status_skips_when_unchanged(self, app_model): + def test_update_app_api_status_skips_when_unchanged(self, app_model, mock_db): from events.app_event import app_was_updated from services.app_service import AppService @@ -199,7 +200,7 @@ class TestAppWasUpdatedSignal: try: app_model.enable_api = True - AppService().update_app_api_status(app_model, True) + AppService().update_app_api_status(app_model, True, session=mock_db.session) finally: app_was_updated.disconnect(handler) diff --git a/api/tests/unit_tests/events/test_update_provider_when_message_created.py b/api/tests/unit_tests/events/test_update_provider_when_message_created.py index f9ac5d9678e..327c80323b4 100644 --- a/api/tests/unit_tests/events/test_update_provider_when_message_created.py +++ b/api/tests/unit_tests/events/test_update_provider_when_message_created.py @@ -1,7 +1,7 @@ from collections.abc import Generator from contextlib import contextmanager from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import ANY, patch from uuid import uuid4 import pytest @@ -19,8 +19,19 @@ from models.provider import ProviderType @contextmanager def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]: session_maker = sessionmaker(bind=engine, expire_on_commit=False) - with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker): - yield + sessions = [] + + def _session(): + session = session_maker() + sessions.append(session) + return session + + with patch("events.event_handlers.update_provider_when_message_created.db", SimpleNamespace(session=_session)): + try: + yield + finally: + for session in sessions: + session.close() def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_insufficient() -> None: @@ -140,5 +151,6 @@ def test_capped_credit_pool_accounting_skips_exhaustion_warning_when_full_amount tenant_id="tenant-id", credits_required=3, pool_type="trial", + session=ANY, ) assert "Credit pool exhausted during message-created accounting" not in caplog.text diff --git a/api/tests/unit_tests/fields/test_snippet_fields.py b/api/tests/unit_tests/fields/test_snippet_fields.py index 2d17a9e577a..97573908576 100644 --- a/api/tests/unit_tests/fields/test_snippet_fields.py +++ b/api/tests/unit_tests/fields/test_snippet_fields.py @@ -1,9 +1,8 @@ from datetime import UTC, datetime from types import SimpleNamespace -from flask_restx import marshal - -from fields.snippet_fields import snippet_list_fields +from fields.snippet_fields import SnippetListItemResponse +from libs.helper import dump_response def test_snippet_list_fields_include_author_name() -> None: @@ -24,6 +23,6 @@ def test_snippet_list_fields_include_author_name() -> None: updated_at=datetime.fromtimestamp(1704067201, tz=UTC), ) - result = marshal(snippet, snippet_list_fields) + result = dump_response(SnippetListItemResponse, snippet) assert result["author_name"] == "Alice" diff --git a/api/tests/unit_tests/services/agent/test_agent_composer_entities.py b/api/tests/unit_tests/services/agent/test_agent_composer_entities.py index 07f993c1b8e..4aaae11b7dc 100644 --- a/api/tests/unit_tests/services/agent/test_agent_composer_entities.py +++ b/api/tests/unit_tests/services/agent/test_agent_composer_entities.py @@ -14,6 +14,23 @@ from services.entities.agent_entities import ( ) +def test_default_agent_soul_enables_file_upload_feature(): + agent_soul = AgentSoulConfig() + + file_upload = agent_soul.model_dump(mode="json")["app_features"]["file_upload"] + assert file_upload == { + "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], + "allowed_file_types": ["document", "image", "audio", "video"], + "allowed_file_upload_methods": ["local_file", "remote_url"], + "enabled": True, + "image": {"enabled": True}, + "number_limits": 3, + } + # The product default should be visible in API responses, but it must not + # make workflow-only payload validation treat app_features as user-authored. + assert bool(agent_soul.app_features) is False + + def test_workflow_variant_rejects_agent_app_only_fields(): with pytest.raises(ValueError): ComposerSavePayload.model_validate( @@ -257,6 +274,16 @@ def test_knowledge_query_mode_uses_stable_backend_enums(): }, "knowledge set dataset ids must be unique", ), + ], +) +def test_knowledge_sets_contract_rejects_invalid_configs(knowledge_payload, match: str): + with pytest.raises(ValidationError, match=match): + AgentSoulConfig.model_validate({"knowledge": knowledge_payload}) + + +@pytest.mark.parametrize( + ("knowledge_payload", "match"), + [ ( { "sets": [ @@ -317,9 +344,25 @@ def test_knowledge_query_mode_uses_stable_backend_enums(): ), ], ) -def test_knowledge_sets_contract_rejects_invalid_configs(knowledge_payload, match: str): - with pytest.raises(ValidationError, match=match): - AgentSoulConfig.model_validate({"knowledge": knowledge_payload}) +def test_knowledge_runtime_requirements_block_publish_but_not_draft_save(knowledge_payload, match: str): + draft_payload = ComposerSavePayload.model_validate( + { + "variant": ComposerVariant.AGENT_APP, + "save_strategy": ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION, + "agent_soul": {"knowledge": knowledge_payload}, + } + ) + ComposerConfigValidator.validate_draft_save_payload(draft_payload) + + publish_payload = ComposerSavePayload.model_validate( + { + "variant": ComposerVariant.AGENT_APP, + "save_strategy": ComposerSaveStrategy.SAVE_AS_NEW_VERSION, + "agent_soul": {"knowledge": knowledge_payload}, + } + ) + with pytest.raises(InvalidComposerConfigError, match=match): + ComposerConfigValidator.validate_publish_payload(publish_payload) def test_agent_soul_model_config_is_first_class_without_credentials(): diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index cebbbf4f94f..22c3f4c5500 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -1,8 +1,11 @@ import json from datetime import UTC, datetime from types import SimpleNamespace +from unittest.mock import MagicMock import pytest +from agenton.compositor import CompositorSessionSnapshot +from dify_agent.protocol import RuntimeLayerSpec from sqlalchemy.exc import IntegrityError from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError @@ -36,6 +39,7 @@ from services.agent.agent_soul_state import agent_soul_has_model from services.agent.composer_service import AgentComposerService from services.agent.composer_validator import ComposerConfigValidator from services.agent.errors import ( + AgentModelNotConfiguredError, AgentNameConflictError, AgentNotFoundError, AgentVersionConflictError, @@ -116,7 +120,9 @@ def test_load_workflow_composer_returns_empty_state(monkeypatch: pytest.MonkeyPa monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: SimpleNamespace(id="workflow-1")) monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: None) - result = AgentComposerService.load_workflow_composer(tenant_id="tenant-1", app_id="app-1", node_id="node-1") + result = AgentComposerService.load_workflow_composer( + tenant_id="tenant-1", app_id="app-1", node_id="node-1", session=composer_service.db.session + ) assert result["binding"] is None assert result["save_options"] == ["node_job_only", "save_to_roster"] @@ -154,7 +160,9 @@ def test_load_workflow_composer_serializes_existing_binding(monkeypatch: pytest. lambda **kwargs: {"agent": kwargs["agent"].id, "version": kwargs["version"].id}, ) - result = AgentComposerService.load_workflow_composer(tenant_id="tenant-1", app_id="app-1", node_id="node-1") + result = AgentComposerService.load_workflow_composer( + tenant_id="tenant-1", app_id="app-1", node_id="node-1", session=composer_service.db.session + ) assert result == {"agent": "agent-1", "version": "version-1"} @@ -189,6 +197,7 @@ def test_load_workflow_composer_uses_roster_preview_snapshot(monkeypatch: pytest app_id="app-1", node_id="node-1", snapshot_id="preview-version", + session=composer_service.db.session, ) assert result == {"binding_snapshot_id": "binding-version", "version": "preview-version"} @@ -231,6 +240,7 @@ def test_load_workflow_composer_uses_inline_preview_snapshot(monkeypatch: pytest app_id="app-1", node_id="node-1", snapshot_id="inline-preview-version", + session=composer_service.db.session, ) assert result == {"agent": "inline-agent-1", "version": "inline-preview-version"} @@ -257,6 +267,7 @@ def test_workflow_inline_debug_conversation_seed(monkeypatch: pytest.MonkeyPatch binding=binding, agent=agent, account_id="account-1", + session="session-1", ) assert debug_conversation_id == "debug-conversation-1" @@ -278,6 +289,7 @@ def test_workflow_inline_debug_conversation_seed_skips_non_inline(monkeypatch: p binding=SimpleNamespace(binding_type=WorkflowAgentBindingType.ROSTER_AGENT), agent=SimpleNamespace(id="agent-1", scope=AgentScope.ROSTER), account_id="account-1", + session="session-1", ) is None ) @@ -287,6 +299,7 @@ def test_workflow_inline_debug_conversation_seed_skips_non_inline(monkeypatch: p binding=SimpleNamespace(binding_type=WorkflowAgentBindingType.INLINE_AGENT), agent=SimpleNamespace(id="inline-agent-1", scope=AgentScope.WORKFLOW_ONLY), account_id=None, + session="session-1", ) is None ) @@ -302,6 +315,7 @@ def test_load_workflow_composer_rejects_preview_without_binding(monkeypatch: pyt app_id="app-1", node_id="node-1", snapshot_id="preview-version", + session=composer_service.db.session, ) @@ -360,7 +374,12 @@ def test_save_workflow_composer_dispatches_save_strategy(monkeypatch, strategy, ) result = AgentComposerService.save_workflow_composer( - tenant_id="tenant-1", app_id="app-1", node_id="node-1", account_id="account-1", payload=payload + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + account_id="account-1", + payload=payload, + session=composer_service.db.session, ) assert result.pop("validation") == {"warnings": [], "knowledge_retrieval_placeholder": []} @@ -381,7 +400,12 @@ def test_save_workflow_composer_rejects_agent_app_variant(): with pytest.raises(ValueError): AgentComposerService.save_workflow_composer( - tenant_id="tenant-1", app_id="app-1", node_id="node-1", account_id="account-1", payload=payload + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + account_id="account-1", + payload=payload, + session=composer_service.db.session, ) @@ -443,7 +467,11 @@ def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest. ) result = AgentComposerService.save_agent_app_composer( - tenant_id="tenant-1", app_id="app-1", account_id="account-1", payload=payload + tenant_id="tenant-1", + app_id="app-1", + account_id="account-1", + payload=payload, + session=composer_service.db.session, ) assert result.pop("validation") == {"warnings": [], "knowledge_retrieval_placeholder": []} @@ -474,7 +502,9 @@ def test_load_agent_app_composer_exposes_draft_save_only(monkeypatch: pytest.Mon monkeypatch.setattr(AgentComposerService, "_serialize_version", lambda _version: None) monkeypatch.setattr(AgentComposerService, "_serialize_draft", lambda _draft: {"id": "draft-1"}) - result = AgentComposerService.load_agent_app_composer(tenant_id="tenant-1", app_id="app-1") + result = AgentComposerService.load_agent_app_composer( + tenant_id="tenant-1", app_id="app-1", session=composer_service.db.session + ) assert result["save_options"] == [ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION.value] @@ -494,6 +524,7 @@ def test_save_agent_app_composer_rejects_version_save_strategy(): app_id="app-1", account_id="account-1", payload=payload, + session=composer_service.db.session, ) @@ -527,7 +558,11 @@ def test_save_agent_app_composer_updates_normal_draft(monkeypatch: pytest.Monkey ) result = AgentComposerService.save_agent_app_composer( - tenant_id="tenant-1", app_id="app-1", account_id="account-1", payload=payload + tenant_id="tenant-1", + app_id="app-1", + account_id="account-1", + payload=payload, + session=composer_service.db.session, ) assert result.pop("validation") == {"warnings": [], "knowledge_retrieval_placeholder": []} @@ -569,13 +604,63 @@ def test_save_agent_app_composer_keeps_published_when_draft_matches_active_snaps ) AgentComposerService.save_agent_app_composer( - tenant_id="tenant-1", app_id="app-1", account_id="account-1", payload=payload + tenant_id="tenant-1", app_id="app-1", account_id="account-1", payload=payload, session=fake_session ) assert agent.active_config_is_published is True assert fake_session.commits == 1 +def test_publish_agent_app_draft_rejects_missing_model(monkeypatch: pytest.MonkeyPatch): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + active_config_snapshot_id="version-1", + active_config_is_published=False, + ) + draft = AgentConfigDraft( + tenant_id="tenant-1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + base_snapshot_id="version-1", + config_snapshot=AgentSoulConfig(), + ) + fake_session = FakeSession(scalar=[agent, draft]) + + def fail_create_config_version(**_kwargs): + raise AssertionError("config version must not be created when Agent Soul has no model") + + def fail_validate_knowledge_datasets(**_kwargs): + raise AssertionError("knowledge datasets must not be validated when Agent Soul has no model") + + monkeypatch.setattr(composer_service.db, "session", fake_session) + monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda payload: None) + monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", fail_validate_knowledge_datasets) + monkeypatch.setattr(AgentComposerService, "_create_config_version", fail_create_config_version) + + with pytest.raises(AgentModelNotConfiguredError) as exc_info: + AgentComposerService.publish_agent_app_draft( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + version_note="ship it", + session=fake_session, + ) + + assert exc_info.value.error_code == "agent_model_not_configured" + assert agent.active_config_snapshot_id == "version-1" + assert agent.active_config_is_published is False + assert draft.base_snapshot_id == "version-1" + assert fake_session.commits == 0 + + def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest.MonkeyPatch): agent = Agent( id="agent-1", @@ -615,6 +700,7 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. agent_id="agent-1", account_id="account-1", version_note="ship it", + session=composer_service.db.session, ) assert result["result"] == "success" @@ -658,6 +744,7 @@ def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkey tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", + session=composer_service.db.session, ) build_draft = fake_session.added[0] @@ -679,6 +766,7 @@ def test_agent_app_build_draft_checkout_and_apply_use_user_isolated_draft(monkey tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", + session=composer_service.db.session, ) assert applied["result"] == "success" @@ -737,6 +825,7 @@ def test_agent_app_build_draft_apply_marks_unpublished_when_build_draft_differs( tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", + session=fake_session, ) assert normal_draft.config_snapshot_dict == build_draft.config_snapshot_dict @@ -762,12 +851,21 @@ def test_agent_app_composer_candidates_and_impact(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(AgentComposerService, "_workspace_dify_tools", lambda **kwargs: []) workflow_candidates = AgentComposerService.get_workflow_candidates( - tenant_id="tenant-1", app_id="app-1", node_id="node-1", user_id="account-1" + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + user_id="account-1", + session=composer_service.db.session, ) agent_app_candidates = AgentComposerService.get_agent_app_candidates( - tenant_id="tenant-1", agent_id="agent-1", user_id="account-1" + tenant_id="tenant-1", + agent_id="agent-1", + user_id="account-1", + session=composer_service.db.session, + ) + impact = AgentComposerService.calculate_impact( + tenant_id="tenant-1", current_snapshot_id="version-1", session=composer_service.db.session ) - impact = AgentComposerService.calculate_impact(tenant_id="tenant-1", current_snapshot_id="version-1") assert workflow_candidates["variant"] == "workflow" assert workflow_candidates["allowed_node_job_candidates"]["previous_node_outputs"] == [] @@ -804,7 +902,9 @@ def test_serialize_workflow_state_changes_lock_and_save_options(monkeypatch: pyt version = AgentConfigSnapshot(id="version-1", version=1, config_snapshot='{"prompt":{"system_prompt":"x"}}') monkeypatch.setattr(AgentComposerService, "calculate_impact", lambda **kwargs: {"workflow_node_count": 1}) - state = AgentComposerService._serialize_workflow_state(binding=binding, agent=agent, version=version) + state = AgentComposerService._serialize_workflow_state( + binding=binding, agent=agent, version=version, session=composer_service.db.session + ) assert state["soul_lock"]["locked"] is True assert state["agent"]["role"] == "Tender Analyst" @@ -843,7 +943,9 @@ def test_serialize_workflow_state_passes_user_declared_outputs_through_effective version = AgentConfigSnapshot(id="version-1", version=1, config_snapshot='{"prompt":{"system_prompt":"x"}}') monkeypatch.setattr(AgentComposerService, "calculate_impact", lambda **kwargs: {"workflow_node_count": 1}) - state = AgentComposerService._serialize_workflow_state(binding=binding, agent=agent, version=version) + state = AgentComposerService._serialize_workflow_state( + binding=binding, agent=agent, version=version, session=composer_service.db.session + ) # When the user has declared outputs, effective_declared_outputs is the same # list (no defaults injected). @@ -893,6 +995,7 @@ def test_serialize_workflow_state_includes_inline_debug_conversation_message_sta agent=agent, version=version, account_id="account-1", + session=composer_service.db.session, ) assert state["debug_conversation_id"] == "debug-conversation-1" @@ -974,6 +1077,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk account_id="account-1", binding=existing_binding, payload=payload, + session=composer_service.db.session, ) inline_binding = AgentComposerService._save_node_job_only( tenant_id="tenant-1", @@ -983,6 +1087,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk account_id="account-1", binding=None, payload=payload, + session=composer_service.db.session, ) new_agent_binding = AgentComposerService._save_as_new_agent( tenant_id="tenant-1", @@ -992,6 +1097,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk account_id="account-1", binding=None, payload=payload, + session=composer_service.db.session, ) save_to_roster_binding = AgentComposerService._save_to_roster( tenant_id="tenant-1", @@ -1005,12 +1111,14 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk current_snapshot_id="inline-version-1", ), payload=payload, + session=composer_service.db.session, ) new_version_binding = AgentComposerService._save_as_new_version( tenant_id="tenant-1", account_id="account-1", binding=WorkflowAgentNodeBinding(agent_id="roster-agent-1", current_snapshot_id="source-version-1"), payload=payload, + session=composer_service.db.session, ) assert updated_binding.updated_by == "account-1" @@ -1035,6 +1143,7 @@ def test_composer_save_helpers_create_and_rebind_agents(monkeypatch: pytest.Monk "account_id": "account-1", "agent_soul": payload.agent_soul, "node_job": payload.node_job, + "session": composer_service.db.session, } ] @@ -1101,6 +1210,7 @@ def test_node_job_only_updates_inline_agent_soul(monkeypatch: pytest.MonkeyPatch account_id="account-1", binding=binding, payload=payload, + session=composer_service.db.session, ) assert updated_binding.current_snapshot_id == "inline-version-2" @@ -1153,6 +1263,7 @@ def test_node_job_only_switches_roster_binding_to_inline_agent(monkeypatch: pyte account_id="account-1", binding=binding, payload=payload, + session=composer_service.db.session, ) assert updated_binding is binding @@ -1202,6 +1313,7 @@ def test_node_job_only_rejects_start_from_scratch_with_existing_inline_binding_i account_id="account-1", binding=binding, payload=payload, + session=composer_service.db.session, ) @@ -1249,6 +1361,7 @@ def test_node_job_only_rejects_inline_binding_pointing_to_roster_agent(monkeypat account_id="account-1", binding=binding, payload=payload, + session=composer_service.db.session, ) @@ -1335,6 +1448,7 @@ def test_copy_workflow_composer_from_roster_creates_inline_agent_and_preserves_n account_id="account-1", source_agent_id="roster-agent-1", source_snapshot_id="roster-version-2", + session=composer_service.db.session, ) assert state["binding"]["binding_type"] == WorkflowAgentBindingType.INLINE_AGENT.value @@ -1395,6 +1509,7 @@ def test_copy_workflow_composer_from_roster_rejects_stale_source_snapshot(monkey account_id="account-1", source_agent_id="roster-agent-1", source_snapshot_id="roster-version-1", + session=composer_service.db.session, ) @@ -1445,6 +1560,7 @@ def test_copy_workflow_composer_from_roster_is_idempotent_when_already_inline(mo account_id="account-1", source_agent_id="roster-agent-1", idempotency_key="same-click", + session=composer_service.db.session, ) assert state == {"binding_type": WorkflowAgentBindingType.INLINE_AGENT.value} @@ -1523,6 +1639,7 @@ def test_copy_workflow_composer_from_roster_rejects_invalid_source_binding( node_id="node-1", account_id="account-1", source_agent_id="roster-agent-1", + session=composer_service.db.session, ) @@ -1579,6 +1696,7 @@ def test_copy_agent_drive_rows_copies_skill_prefix_and_files(monkeypatch: pytest account_id="account-1", agent_soul=agent_soul, node_job=node_job, + session=composer_service.db.session, ) copied = [row for row in fake_session.added if isinstance(row, AgentDriveFile)] @@ -1604,6 +1722,7 @@ def test_copy_agent_drive_rows_skips_when_no_referenced_drive_keys(monkeypatch: target_agent_id="inline-agent-1", account_id="account-1", agent_soul=agent_soul, + session=composer_service.db.session, ) assert fake_session.added == [] @@ -1630,6 +1749,7 @@ def test_copy_agent_drive_rows_skips_existing_target_keys(monkeypatch: pytest.Mo target_agent_id="inline-agent-1", account_id="account-1", agent_soul=agent_soul, + session=composer_service.db.session, ) assert [row for row in fake_session.added if isinstance(row, AgentDriveFile)] == [] @@ -1693,7 +1813,7 @@ def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytes ) class FakeAppService: - def create_app(self, tenant_id, params, account): + def create_app(self, tenant_id, params, account, session): created_apps.append((tenant_id, params, account)) return SimpleNamespace(id="app-agent-1") @@ -1731,6 +1851,7 @@ def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytes node_id="node-1", account_id="account-1", agent_soul=_agent_soul_with_model(), + session=composer_service.db.session, ) roster_agent = AgentComposerService._create_roster_agent_for_composer( tenant_id="tenant-1", @@ -1739,6 +1860,7 @@ def test_composer_create_agents_syncs_active_config_has_model(monkeypatch: pytes agent_soul=_agent_soul_with_model(), operation=AgentConfigRevisionOperation.CREATE_VERSION, version_note=None, + session=composer_service.db.session, ) assert workflow_agent.active_config_snapshot_id == "version-with-model" @@ -1760,14 +1882,14 @@ def test_composer_require_account(monkeypatch: pytest.MonkeyPatch): account = SimpleNamespace(id="account-1") monkeypatch.setattr(composer_service.db, "session", SimpleNamespace(get=lambda model, account_id: account)) - assert AgentComposerService._require_account(account_id="account-1") is account + assert AgentComposerService._require_account(account_id="account-1", session=composer_service.db.session) is account def test_composer_require_account_raises_when_missing(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(composer_service.db, "session", SimpleNamespace(get=lambda model, account_id: None)) with pytest.raises(ValueError, match="Account not found"): - AgentComposerService._require_account(account_id="missing-account") + AgentComposerService._require_account(account_id="missing-account", session=composer_service.db.session) def test_composer_create_roster_agent_rolls_back_name_conflict(monkeypatch: pytest.MonkeyPatch): @@ -1775,7 +1897,7 @@ def test_composer_create_roster_agent_rolls_back_name_conflict(monkeypatch: pyte monkeypatch.setattr(composer_service.db, "session", fake_session) class FakeAppService: - def create_app(self, tenant_id, params, account): + def create_app(self, tenant_id, params, account, session): raise IntegrityError("insert apps", params, Exception("duplicate")) monkeypatch.setattr(composer_service, "AppService", FakeAppService) @@ -1789,6 +1911,7 @@ def test_composer_create_roster_agent_rolls_back_name_conflict(monkeypatch: pyte agent_soul=_agent_soul_with_model(), operation=AgentConfigRevisionOperation.CREATE_VERSION, version_note=None, + session=composer_service.db.session, ) assert fake_session.rollbacks == 1 @@ -1799,7 +1922,7 @@ def test_composer_create_roster_agent_raises_when_backing_agent_missing(monkeypa monkeypatch.setattr(composer_service.db, "session", fake_session) class FakeAppService: - def create_app(self, tenant_id, params, account): + def create_app(self, tenant_id, params, account, session): return SimpleNamespace(id="app-agent-1") class FakeAgentRosterService: @@ -1821,6 +1944,7 @@ def test_composer_create_roster_agent_raises_when_backing_agent_missing(monkeypa agent_soul=_agent_soul_with_model(), operation=AgentConfigRevisionOperation.CREATE_VERSION, version_note=None, + session=composer_service.db.session, ) @@ -1842,6 +1966,7 @@ def test_agent_app_draft_match_does_not_mark_create_version_as_published(monkeyp tenant_id="tenant-1", agent=agent, agent_soul=agent_soul, + session=fake_session, ) is False ) @@ -1865,6 +1990,7 @@ def test_agent_app_draft_match_marks_publish_visible_revision_as_published(monke tenant_id="tenant-1", agent=agent, agent_soul=agent_soul, + session=fake_session, ) is True ) @@ -1895,6 +2021,7 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa agent_soul=agent_soul, operation=AgentConfigRevisionOperation.SAVE_NEW_VERSION, version_note="note", + session=composer_service.db.session, ) updated_snapshot = AgentComposerService._update_current_version( current_snapshot=AgentConfigSnapshot( @@ -1908,21 +2035,40 @@ def test_composer_version_helpers_and_lookup_errors(monkeypatch: pytest.MonkeyPa agent_soul=agent_soul, operation=AgentConfigRevisionOperation.SAVE_CURRENT_VERSION, version_note="updated", + session=composer_service.db.session, + ) + workflow = AgentComposerService._get_draft_workflow( + tenant_id="tenant-1", app_id="app-1", session=composer_service.db.session ) - workflow = AgentComposerService._get_draft_workflow(tenant_id="tenant-1", app_id="app-1") with pytest.raises(ValueError): - AgentComposerService._get_draft_workflow(tenant_id="tenant-1", app_id="missing") - assert AgentComposerService._require_agent(tenant_id="tenant-1", agent_id="agent-1").id == "agent-1" - with pytest.raises(composer_service.AgentNotFoundError): - AgentComposerService._require_agent(tenant_id="tenant-1", agent_id=None) - assert AgentComposerService._get_agent_if_present(tenant_id="tenant-1", agent_id="agent-1") is None + AgentComposerService._get_draft_workflow( + tenant_id="tenant-1", app_id="missing", session=composer_service.db.session + ) assert ( - AgentComposerService._require_version(tenant_id="tenant-1", agent_id="agent-1", version_id="version-1").id + AgentComposerService._require_agent( + tenant_id="tenant-1", agent_id="agent-1", session=composer_service.db.session + ).id + == "agent-1" + ) + with pytest.raises(composer_service.AgentNotFoundError): + AgentComposerService._require_agent(tenant_id="tenant-1", agent_id=None, session=composer_service.db.session) + assert ( + AgentComposerService._get_agent_if_present( + tenant_id="tenant-1", agent_id="agent-1", session=composer_service.db.session + ) + is None + ) + assert ( + AgentComposerService._require_version( + tenant_id="tenant-1", agent_id="agent-1", version_id="version-1", session=composer_service.db.session + ).id == "version-1" ) with pytest.raises(composer_service.AgentVersionNotFoundError): - AgentComposerService._require_version(tenant_id="tenant-1", agent_id="agent-1", version_id="missing") + AgentComposerService._require_version( + tenant_id="tenant-1", agent_id="agent-1", version_id="missing", session=composer_service.db.session + ) assert version.version == 2 assert updated_snapshot.version == 3 @@ -1956,7 +2102,11 @@ def test_composer_current_version_and_error_paths(monkeypatch: pytest.MonkeyPatc ) result = AgentComposerService._save_to_current_version( - tenant_id="tenant-1", account_id="account-1", binding=binding, payload=payload + tenant_id="tenant-1", + account_id="account-1", + binding=binding, + payload=payload, + session=composer_service.db.session, ) assert result.updated_by == "account-1" @@ -1977,6 +2127,7 @@ def test_composer_current_version_and_error_paths(monkeypatch: pytest.MonkeyPatc "save_strategy": ComposerSaveStrategy.SAVE_AS_NEW_AGENT.value, } ), + session=composer_service.db.session, ) @@ -3022,6 +3173,159 @@ class TestAgentAppBackingAgent: assert session.deleted == [] assert session.commits == 1 + def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions( + self, monkeypatch: pytest.MonkeyPatch + ): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") + stored_session = SimpleNamespace( + scope=SimpleNamespace( + tenant_id="tenant-1", + app_id="old-app", + conversation_id="old-conversation", + agent_id="agent-9", + agent_config_snapshot_id="snap-9", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-old", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + session = FakeSession(scalar=[agent, mapping]) + service = AgentRosterService(session) + cleanup_delay = MagicMock() + cleanup_store = MagicMock() + cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] + monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) + monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + + conversation_id = service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + cleanup_store.list_active_sessions_for_conversation.assert_called_once_with( + tenant_id="tenant-1", + app_id="old-app", + conversation_id="old-conversation", + ) + cleanup_delay.assert_called_once() + payload = cleanup_delay.call_args.args[0] + assert payload["metadata"]["conversation_id"] == "old-conversation" + assert payload["metadata"]["agent_id"] == "agent-9" + assert ( + payload["idempotency_key"] == "tenant-1:agent-1:account-1:old-conversation:debug-session-cleanup:" + "agent-9:snap-9:run-old" + ) + cleanup_store.mark_cleaned.assert_called_once_with( + scope=stored_session.scope, + backend_run_id="run-old", + ) + assert mapping.app_id == "app-1" + assert mapping.conversation_id == conversation_id + + def test_refresh_agent_app_debug_conversation_marks_old_runtime_sessions_clean_when_enqueue_fails( + self, monkeypatch: pytest.MonkeyPatch + ): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") + stored_session = SimpleNamespace( + scope=SimpleNamespace( + tenant_id="tenant-1", + app_id="old-app", + conversation_id="old-conversation", + agent_id="agent-9", + agent_config_snapshot_id="snap-9", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-old", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + session = FakeSession(scalar=[agent, mapping]) + service = AgentRosterService(session) + cleanup_store = MagicMock() + cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] + monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) + monkeypatch.setattr( + roster_service.cleanup_conversation_agent_runtime_session, + "delay", + MagicMock(side_effect=RuntimeError("queue down")), + ) + + service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + cleanup_store.mark_cleaned.assert_called_once_with( + scope=stored_session.scope, + backend_run_id="run-old", + ) + + def test_refresh_agent_app_debug_conversation_ignores_mark_cleaned_failure(self, monkeypatch: pytest.MonkeyPatch): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") + stored_session = SimpleNamespace( + scope=SimpleNamespace( + tenant_id="tenant-1", + app_id="old-app", + conversation_id="old-conversation", + agent_id="agent-9", + agent_config_snapshot_id="snap-9", + ), + session_snapshot=CompositorSessionSnapshot(layers=[]), + backend_run_id="run-old", + runtime_layer_specs=[RuntimeLayerSpec(name="history", type="pydantic_ai.history")], + ) + session = FakeSession(scalar=[agent, mapping]) + service = AgentRosterService(session) + cleanup_store = MagicMock() + cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] + cleanup_store.mark_cleaned.side_effect = RuntimeError("cleanup bookkeeping failed") + monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) + monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", MagicMock()) + + conversation_id = service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + assert mapping.app_id == "app-1" + assert mapping.conversation_id == conversation_id + assert session.commits == 1 + def test_duplicate_agent_app_copies_app_config_and_active_soul(self, monkeypatch: pytest.MonkeyPatch): source_config = SimpleNamespace( opening_statement="hello", @@ -3122,7 +3426,7 @@ class TestAgentAppBackingAgent: captured: dict[str, object] = {} class FakeAppService: - def create_app(self, tenant_id: str, params, account: object) -> object: + def create_app(self, tenant_id: str, params, account: object, session: object) -> object: captured["tenant_id"] = tenant_id captured["params"] = params captured["account"] = account @@ -3191,7 +3495,7 @@ class TestAgentAppBackingAgent: captured: dict[str, object] = {} class FakeAppService: - def create_app(self, tenant_id: str, params, account: object) -> object: + def create_app(self, tenant_id: str, params, account: object, session: object) -> object: captured["params"] = params return target_app @@ -3253,7 +3557,7 @@ class TestAgentAppBackingAgent: monkeypatch.setattr(service, "_next_duplicate_agent_name", lambda **_: "Iris copy") class FakeAppService: - def create_app(self, tenant_id: str, params, account: object) -> object: + def create_app(self, tenant_id: str, params, account: object, session: object) -> object: return target_app access_mode_updates = [] @@ -4267,6 +4571,7 @@ def test_dataset_rows_filters_malformed_ids(monkeypatch: pytest.MonkeyPatch): app_id="app-1", account_id="account-1", payload=payload, + session=composer_service.db.session, ), ), ( @@ -4277,6 +4582,7 @@ def test_dataset_rows_filters_malformed_ids(monkeypatch: pytest.MonkeyPatch): node_id="node-1", account_id="account-1", payload=payload, + session=composer_service.db.session, ), ), ], @@ -4294,29 +4600,24 @@ def test_composer_save_rejects_malformed_knowledge_dataset_ids(monkeypatch: pyte monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) - payload = ComposerSavePayload.model_validate( + agent_soul = AgentSoulConfig.model_validate( { - "variant": variant.value, - "save_strategy": ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION.value, - "soul_lock": {"locked": False}, - "agent_soul": { - "knowledge": { - "sets": [ - { - "id": "support", - "name": "Support KB", - "datasets": [{"id": "not-a-uuid"}], - "query": {"mode": "generated_query"}, - "retrieval": {"mode": "multiple", "top_k": 4}, - } - ] - } + "knowledge": { + "sets": [ + { + "id": "support", + "name": "Support KB", + "datasets": [{"id": "not-a-uuid"}], + "query": {"mode": "generated_query"}, + "retrieval": {"mode": "multiple", "top_k": 4}, + } + ] }, } ) with pytest.raises(InvalidComposerConfigError, match="not-a-uuid"): - save_call(payload) + AgentComposerService.validate_knowledge_datasets(tenant_id="tenant-1", agent_soul=agent_soul) assert captured == {"calls": 0} @@ -4331,6 +4632,7 @@ def test_composer_save_rejects_malformed_knowledge_dataset_ids(monkeypatch: pyte app_id="app-1", account_id="account-1", payload=payload, + session=composer_service.db.session, ), ), ( @@ -4341,6 +4643,7 @@ def test_composer_save_rejects_malformed_knowledge_dataset_ids(monkeypatch: pyte node_id="node-1", account_id="account-1", payload=payload, + session=composer_service.db.session, ), ), ], @@ -4360,20 +4663,70 @@ def test_composer_save_rejects_missing_or_out_of_scope_knowledge_datasets( monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) + agent_soul = AgentSoulConfig.model_validate( + { + "knowledge": { + "sets": [ + { + "id": "support", + "name": "Support KB", + "datasets": [{"id": missing_dataset_id}], + "query": {"mode": "generated_query"}, + "retrieval": {"mode": "multiple", "top_k": 4}, + } + ] + }, + } + ) + + with pytest.raises(InvalidComposerConfigError, match=missing_dataset_id): + AgentComposerService.validate_knowledge_datasets(tenant_id="tenant-1", agent_soul=agent_soul) + + assert captured == {"ids": [missing_dataset_id], "tenant_id": "tenant-1"} + + +def test_save_agent_composer_allows_incomplete_knowledge_draft(monkeypatch: pytest.MonkeyPatch): + agent = SimpleNamespace( + id="agent-1", + source=AgentSource.AGENT_APP, + active_config_snapshot_id="version-1", + active_config_is_published=True, + updated_by=None, + ) + active_version = SimpleNamespace(config_snapshot_dict=AgentSoulConfig().model_dump(mode="json")) + fake_session = FakeSession(scalar=[agent]) + saved = {} + + import services.dataset_service as dataset_service_module + + monkeypatch.setattr(composer_service.db, "session", fake_session) + monkeypatch.setattr( + dataset_service_module.DatasetService, + "get_datasets_by_ids", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("draft save must skip dataset lookup")), + ) + monkeypatch.setattr( + AgentComposerService, + "_save_agent_draft", + lambda **kwargs: saved.update(kwargs) or SimpleNamespace(id="draft-1"), + ) + monkeypatch.setattr(AgentComposerService, "_get_version_if_present", lambda **_kwargs: active_version) + monkeypatch.setattr(AgentComposerService, "load_agent_composer", lambda **_kwargs: {"loaded": True}) + payload = ComposerSavePayload.model_validate( { - "variant": variant.value, + "variant": ComposerVariant.AGENT_APP.value, "save_strategy": ComposerSaveStrategy.SAVE_TO_CURRENT_VERSION.value, - "soul_lock": {"locked": False}, "agent_soul": { "knowledge": { "sets": [ { "id": "support", "name": "Support KB", - "datasets": [{"id": missing_dataset_id}], + "datasets": [{"id": "not-a-uuid"}], "query": {"mode": "generated_query"}, - "retrieval": {"mode": "multiple", "top_k": 4}, + "retrieval": {"mode": "single"}, + "metadata_filtering": {"mode": "automatic"}, } ] } @@ -4381,10 +4734,21 @@ def test_composer_save_rejects_missing_or_out_of_scope_knowledge_datasets( } ) - with pytest.raises(InvalidComposerConfigError, match=missing_dataset_id): - save_call(payload) + result = AgentComposerService.save_agent_composer( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + payload=payload, + session=fake_session, + ) - assert captured == {"ids": [missing_dataset_id], "tenant_id": "tenant-1"} + assert result["loaded"] is True + assert saved["draft_type"] == AgentConfigDraftType.DRAFT + assert saved["agent_soul"].knowledge.sets[0].retrieval.mode == "single" + assert saved["agent_soul"].knowledge.sets[0].retrieval.model is None + assert saved["agent_soul"].knowledge.sets[0].metadata_filtering.mode == "automatic" + assert saved["agent_soul"].knowledge.sets[0].metadata_filtering.metadata_model_config is None + assert fake_session.commits == 1 def test_workspace_dify_tools_returns_provider_and_tool_granularities(monkeypatch: pytest.MonkeyPatch): @@ -4464,6 +4828,7 @@ def test_drive_mention_findings_reports_missing_keys(monkeypatch: pytest.MonkeyP tenant_id="tenant-1", agent_id="agent-1", prompt=_drive_soul().prompt.system_prompt, + session=composer_service.db.session, ) assert [(f["code"], f["id"]) for f in findings] == [("mention_target_missing", "files/sample.pdf")] @@ -4479,6 +4844,7 @@ def test_drive_mention_findings_clean_when_all_keys_exist(monkeypatch: pytest.Mo tenant_id="tenant-1", agent_id="agent-1", prompt=_drive_soul().prompt.system_prompt, + session=composer_service.db.session, ) == [] ) @@ -4491,6 +4857,7 @@ def test_drive_mention_findings_skips_prompt_without_drive_mentions(monkeypatch: tenant_id="tenant-1", agent_id="agent-1", prompt=soul.prompt.system_prompt, + session=composer_service.db.session, ) assert findings == [] @@ -4510,7 +4877,10 @@ def test_collect_validation_findings_appends_drive_mention_findings_with_agent_c ) findings = AgentComposerService.collect_validation_findings( - tenant_id="tenant-1", payload=payload, agent_id="agent-1" + tenant_id="tenant-1", + payload=payload, + agent_id="agent-1", + session=composer_service.db.session, ) codes = {w["code"] for w in findings["warnings"]} @@ -4520,7 +4890,9 @@ def test_collect_validation_findings_appends_drive_mention_findings_with_agent_c "files/sample.pdf", } # without agent context the drive check is skipped entirely - findings_no_agent = AgentComposerService.collect_validation_findings(tenant_id="tenant-1", payload=payload) + findings_no_agent = AgentComposerService.collect_validation_findings( + tenant_id="tenant-1", payload=payload, session=composer_service.db.session + ) assert all(w["code"] != "mention_target_missing" for w in findings_no_agent["warnings"]) @@ -4533,7 +4905,12 @@ def test_resolve_bound_agent_id_queries_active_roster_agent(monkeypatch: pytest. import services.agent.composer_service as module monkeypatch.setattr(module.db, "session", SimpleNamespace(scalar=lambda stmt: "agent-9")) - assert AgentComposerService.resolve_bound_agent_id(tenant_id="t-1", app_id="app-1") == "agent-9" + assert ( + AgentComposerService.resolve_bound_agent_id( + tenant_id="t-1", app_id="app-1", session=composer_service.db.session + ) + == "agent-9" + ) def test_resolve_workflow_node_agent_id_degrades_without_workflow_or_binding(monkeypatch: pytest.MonkeyPatch): @@ -4543,20 +4920,35 @@ def test_resolve_workflow_node_agent_id_degrades_without_workflow_or_binding(mon raise ValueError("no draft workflow") monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", classmethod(boom)) - assert AgentComposerService.resolve_workflow_node_agent_id(tenant_id="t", app_id="a", node_id="n") is None + assert ( + AgentComposerService.resolve_workflow_node_agent_id( + tenant_id="t", app_id="a", node_id="n", session=composer_service.db.session + ) + is None + ) monkeypatch.setattr( AgentComposerService, "_get_draft_workflow", classmethod(lambda cls, **kwargs: SimpleNamespace(id="wf-1")) ) monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", classmethod(lambda cls, **kwargs: None)) - assert AgentComposerService.resolve_workflow_node_agent_id(tenant_id="t", app_id="a", node_id="n") is None + assert ( + AgentComposerService.resolve_workflow_node_agent_id( + tenant_id="t", app_id="a", node_id="n", session=composer_service.db.session + ) + is None + ) monkeypatch.setattr( AgentComposerService, "_get_workflow_binding", classmethod(lambda cls, **kwargs: SimpleNamespace(agent_id="agent-7")), ) - assert AgentComposerService.resolve_workflow_node_agent_id(tenant_id="t", app_id="a", node_id="n") == "agent-7" + assert ( + AgentComposerService.resolve_workflow_node_agent_id( + tenant_id="t", app_id="a", node_id="n", session=composer_service.db.session + ) + == "agent-7" + ) def test_save_workflow_composer_reports_drive_mentions_for_inline_node_job_only(monkeypatch: pytest.MonkeyPatch): @@ -4599,7 +4991,7 @@ def test_save_workflow_composer_reports_drive_mentions_for_inline_node_job_only( ) guarded: dict[str, str] = {} - def fake_collect(cls, *, tenant_id, payload, agent_id=None): + def fake_collect(cls, *, tenant_id, payload, agent_id=None, session=None): guarded["tenant_id"] = tenant_id guarded["agent_id"] = agent_id return {"warnings": [{"code": "mention_target_missing", "id": "files/sample.pdf"}]} @@ -4607,7 +4999,12 @@ def test_save_workflow_composer_reports_drive_mentions_for_inline_node_job_only( monkeypatch.setattr(AgentComposerService, "collect_validation_findings", classmethod(fake_collect)) result = AgentComposerService.save_workflow_composer( - tenant_id="t-1", app_id="app-1", node_id="n-1", account_id="acc-1", payload=payload + tenant_id="t-1", + app_id="app-1", + node_id="n-1", + account_id="acc-1", + payload=payload, + session=composer_service.db.session, ) assert result == { @@ -4657,14 +5054,19 @@ def test_save_workflow_composer_reports_drive_mentions_for_roster_node_job_only( ) captured: dict[str, str | None] = {} - def fake_collect(cls, *, tenant_id, payload, agent_id=None): + def fake_collect(cls, *, tenant_id, payload, agent_id=None, session=None): captured["agent_id"] = agent_id return {"warnings": []} monkeypatch.setattr(AgentComposerService, "collect_validation_findings", classmethod(fake_collect)) result = AgentComposerService.save_workflow_composer( - tenant_id="t-1", app_id="app-1", node_id="n-1", account_id="acc-1", payload=payload + tenant_id="t-1", + app_id="app-1", + node_id="n-1", + account_id="acc-1", + payload=payload, + session=composer_service.db.session, ) assert result == {"state": "ok", "validation": {"warnings": []}} diff --git a/api/tests/unit_tests/services/agent/test_prompt_mentions.py b/api/tests/unit_tests/services/agent/test_prompt_mentions.py index b65f8b6f410..6051c96b50b 100644 --- a/api/tests/unit_tests/services/agent/test_prompt_mentions.py +++ b/api/tests/unit_tests/services/agent/test_prompt_mentions.py @@ -243,9 +243,10 @@ def test_node_job_resolver_resolves_each_kind(node_job: WorkflowNodeJobConfig): assert expanded == ( "Read START/tenders and produce qna_report (file output; create the file locally, run " - "`dify-agent file upload `, then copy the returned AgentStubFileMapping JSON " - "as final_output.qna_report; do not call final_output before upload succeeds, and do not use " - "the local path, filename, URL, or a synthesized dify-file-ref as the reference); " + "`dify-agent file upload `, then set final_output.qna_report to a `tool_file` mapping " + "using the returned `reference`; if replying to the user in natural language, use the returned " + "`download_url`; do not call final_output before upload succeeds, and do not use the local path, " + "filename, URL, or a synthesized dify-file-ref as the reference); " "if unsure contact EMAIL · David Hayes." ) diff --git a/api/tests/unit_tests/services/agent/test_skill_standardize_service.py b/api/tests/unit_tests/services/agent/test_skill_standardize_service.py index 074ac59bb1c..5b3ade55721 100644 --- a/api/tests/unit_tests/services/agent/test_skill_standardize_service.py +++ b/api/tests/unit_tests/services/agent/test_skill_standardize_service.py @@ -50,6 +50,7 @@ def test_standardize_creates_drive_owned_toolfiles_and_commits_archive_manifest( tenant_id="tenant-1", user_id="user-1", agent_id="agent-1", + session=MagicMock(), ) # ToolFiles: SKILL.md and the full archive. Archive members stay lazy. diff --git a/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py b/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py index 25678bb4f4d..cfb32d63a92 100644 --- a/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py +++ b/api/tests/unit_tests/services/agent/test_skill_tool_inference_service.py @@ -38,14 +38,17 @@ def test_infer_returns_suggestions_with_inferred_from(monkeypatch): ' "env_suggestions": [{"key": "OPENAI_API_KEY", "reason": "whisper call", "secret_likely": true}]}]}' ) with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)): - result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe") + session = MagicMock() + result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=session) assert result["inferable"] is True tool = result["cli_tools"][0] assert tool["name"] == "ffmpeg" assert tool["inferred_from"] == "audio-transcribe" assert tool["env_suggestions"] == [{"key": "OPENAI_API_KEY", "reason": "whisper call", "secret_likely": True}] - drive.preview.assert_called_once_with(tenant_id="t-1", agent_id="a-1", key="audio-transcribe/SKILL.md") + drive.preview.assert_called_once_with( + tenant_id="t-1", agent_id="a-1", key="audio-transcribe/SKILL.md", session=session + ) def test_infer_threads_skill_md_into_the_prompt(monkeypatch): @@ -57,7 +60,7 @@ def test_infer_threads_skill_md_into_the_prompt(monkeypatch): return '{"inferable": false, "cli_tools": [], "reason": "none"}' with patch.object(SkillToolInferenceService, "_invoke", staticmethod(fake_invoke)): - service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe") + service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock()) assert "Files inside the skill package" not in captured["prompt"] assert "ffmpeg" in captured["prompt"] # SKILL.md body present @@ -67,7 +70,7 @@ def test_infer_not_inferable_passes_reason_through(monkeypatch): service, _ = _service() raw = '{"inferable": false, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"}' with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)): - result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe") + result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock()) assert result == {"inferable": False, "cli_tools": [], "reason": "SKILL.md 未描述任何外部命令依赖"} @@ -81,7 +84,7 @@ def test_infer_retries_once_then_422(monkeypatch): with patch.object(SkillToolInferenceService, "_invoke", staticmethod(bad_invoke)): with pytest.raises(SkillToolInferenceError) as exc_info: - service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe") + service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock()) assert len(calls) == 2 # one retry assert exc_info.value.code == "inference_failed" @@ -92,7 +95,7 @@ def test_infer_repairs_slightly_malformed_json(monkeypatch): service, _ = _service() raw = 'Here you go: {"inferable": true, "cli_tools": [], "reason": null,}' with patch.object(SkillToolInferenceService, "_invoke", staticmethod(lambda **kwargs: raw)): - result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe") + result = service.infer(tenant_id="t-1", agent_id="a-1", slug="audio-transcribe", session=MagicMock()) assert result["inferable"] is True @@ -102,7 +105,7 @@ def test_missing_skill_maps_to_404(): service = SkillToolInferenceService(drive_service=drive) with pytest.raises(SkillToolInferenceError) as exc_info: - service.infer(tenant_id="t-1", agent_id="a-1", slug="ghost") + service.infer(tenant_id="t-1", agent_id="a-1", slug="ghost", session=MagicMock()) assert exc_info.value.code == "skill_not_found" assert exc_info.value.status_code == 404 @@ -110,7 +113,7 @@ def test_missing_skill_maps_to_404(): def test_binary_skill_md_maps_to_404(): service, _ = _service(preview={"key": "x/SKILL.md", "size": 1, "truncated": False, "binary": True, "text": None}) with pytest.raises(SkillToolInferenceError) as exc_info: - service.infer(tenant_id="t-1", agent_id="a-1", slug="x") + service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=MagicMock()) assert exc_info.value.code == "skill_not_found" @@ -160,5 +163,5 @@ def test_load_skill_md_passes_through_non_missing_drive_errors(): service = SkillToolInferenceService(drive_service=drive) with pytest.raises(SkillToolInferenceError) as exc_info: - service.infer(tenant_id="t-1", agent_id="a-1", slug="x") + service.infer(tenant_id="t-1", agent_id="a-1", slug="x", session=MagicMock()) assert exc_info.value.code == "agent_not_found" diff --git a/api/tests/unit_tests/services/controller_api.py b/api/tests/unit_tests/services/controller_api.py index 10b80fb92f6..48805da884e 100644 --- a/api/tests/unit_tests/services/controller_api.py +++ b/api/tests/unit_tests/services/controller_api.py @@ -1009,8 +1009,7 @@ class TestExternalDatasetApi: # 4. Provide default values when parameters are missing # 5. Raise BadRequest exceptions when validation fails # -# Response formatting is handled by Flask-RESTX's marshal_with decorator -# or marshal function, which: +# Response formatting is handled by controller response schemas, which: # # 1. Formats response data according to defined models # 2. Handles nested objects and lists diff --git a/api/tests/unit_tests/services/data_migration/test_export_service.py b/api/tests/unit_tests/services/data_migration/test_export_service.py index f5480ff52af..5479de5ba22 100644 --- a/api/tests/unit_tests/services/data_migration/test_export_service.py +++ b/api/tests/unit_tests/services/data_migration/test_export_service.py @@ -1,3 +1,5 @@ +from unittest.mock import MagicMock + import pytest from services.data_migration.dependency_discovery_service import DiscoveredDependency @@ -126,13 +128,13 @@ def test_secret_free_mcp_dependencies_are_dependency_only(): report_items = [] service._export_mcp_tools( - object(), tenant_id="tenant-1", provider_ids=["mcp-1"], include_secrets=False, exported_mcp_tools=mcp_tools, dependencies=dependencies, report_items=report_items, + session=MagicMock(), ) assert mcp_tools == [] @@ -151,12 +153,14 @@ def test_secret_free_mcp_dependencies_are_dependency_only(): def test_get_mcp_provider_does_not_compare_non_uuid_identifier_to_uuid_id(): statements = [] - class StubSession: - def scalar(self, statement): - statements.append(str(statement)) + def capture_scalar(statement): + statements.append(str(statement)) + + session = MagicMock() + session.scalar.side_effect = capture_scalar with pytest.raises(MigrationDataError, match="MCP provider not found"): - MigrationExportService()._get_mcp_provider(StubSession(), "tenant-1", "my-test-mcp") + MigrationExportService()._get_mcp_provider("tenant-1", "my-test-mcp", session=session) assert len(statements) == 1 assert "tool_mcp_providers.id =" not in statements[0] diff --git a/api/tests/unit_tests/services/data_migration/test_import_service.py b/api/tests/unit_tests/services/data_migration/test_import_service.py index 2b11d575ed6..10460aba470 100644 --- a/api/tests/unit_tests/services/data_migration/test_import_service.py +++ b/api/tests/unit_tests/services/data_migration/test_import_service.py @@ -3,6 +3,7 @@ import yaml from models.tools import MCPToolProvider, WorkflowToolProvider from services.app_dsl_service import Import +from services.data_migration import import_service from services.data_migration.entities import ( ConflictStrategy, IdStrategy, @@ -92,8 +93,12 @@ def test_package_target_tenant_id_ignores_invalid_uuid(monkeypatch): return EmptyResult() + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + with pytest.raises(MigrationDataError, match="Target tenant not found"): - ImportTargetResolver().resolve(StubSession(), ImportRequest(package=package)) + ImportTargetResolver().resolve(ImportRequest(package=package), session=import_service.db.session) def test_options_override_replaces_package_defaults(): @@ -113,7 +118,7 @@ def test_options_override_replaces_package_defaults(): captured_options: list[ImportOptions] = [] class StubResolver(ImportTargetResolver): - def resolve(self, session, request: ImportRequest) -> ImportTarget: + def resolve(self, request: ImportRequest, session) -> ImportTarget: return ImportTarget( tenant_id="tenant-1", tenant_name="target", @@ -124,7 +129,6 @@ def test_options_override_replaces_package_defaults(): class CapturingImportService(MigrationImportService): def _import_workflows( self, - session, package: MigrationPackage, target: ImportTarget, options: ImportOptions, @@ -137,7 +141,8 @@ def test_options_override_replaces_package_defaults(): override = ImportOptions(create_app_api_token_on_import=False, conflict_strategy=ConflictStrategy.SKIP) CapturingImportService(target_resolver=StubResolver()).import_package( - object(), ImportRequest(package=package, options_override=override) + ImportRequest(package=package, options_override=override), + session=import_service.db.session, ) assert captured_options == [override] @@ -150,37 +155,53 @@ def test_only_preserve_id_strategy_reuses_source_app_id(): assert service._should_preserve_source_app_id(ImportOptions(id_strategy=IdStrategy.GENERATE_NEW_ID)) is False -def test_find_existing_app_ignores_invalid_uuid(): +def test_find_existing_app_ignores_invalid_uuid(monkeypatch): class StubSession: def scalar(self, statement): raise AssertionError("invalid UUID should not be queried against App.id") - assert MigrationImportService()._find_existing_app(StubSession(), "not-a-uuid", "tenant-1") is None + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + + assert ( + MigrationImportService()._find_existing_app("not-a-uuid", "tenant-1", session=import_service.db.session) is None + ) -def test_find_existing_workflow_tool_does_not_compare_invalid_uuid(): +def test_find_existing_workflow_tool_does_not_compare_invalid_uuid(monkeypatch): captured = [] class StubSession: def scalar(self, statement): captured.append(statement) + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + MigrationImportService()._find_existing_workflow_tool( - StubSession(), "tenant-1", "not-a-uuid", "tool-name", "app-id" + "tenant-1", "not-a-uuid", "tool-name", "app-id", session=import_service.db.session ) where_clause = str(captured[0].whereclause) assert f"{WorkflowToolProvider.__tablename__}.id" not in where_clause -def test_find_existing_mcp_tool_does_not_compare_invalid_uuid(): +def test_find_existing_mcp_tool_does_not_compare_invalid_uuid(monkeypatch): captured = [] class StubSession: def scalar(self, statement): captured.append(statement) - MigrationImportService()._find_existing_mcp_tool(StubSession(), "tenant-1", "my-test-mcp", "my-test-mcp") + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + + MigrationImportService()._find_existing_mcp_tool( + "tenant-1", "my-test-mcp", "my-test-mcp", session=import_service.db.session + ) where_clause = str(captured[0].whereclause) assert f"{MCPToolProvider.__tablename__}.id" not in where_clause @@ -211,16 +232,17 @@ def test_workflow_app_import_does_not_wrap_app_dsl_import_in_nested_transaction( from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr(import_service, "AppDslService", StubAppDslService) imported_app_id = MigrationImportService()._import_workflow_app( - session=StubSession(), account=object(), workflow_data={"name": "main_chatflow"}, dsl_content="app:\n mode: workflow\n", app_id="source-app-id", existing_app=None, options=ImportOptions(id_strategy=IdStrategy.PRESERVE_ID), + session=import_service.db.session, ) assert imported_app_id == "imported-app-id" @@ -317,19 +339,20 @@ def test_workflow_tool_import_publishes_referenced_app_before_create(monkeypatch return account class PublishingImportService(MigrationImportService): - def _find_existing_app(self, session, app_id, tenant_id): + def _find_existing_app(self, app_id, tenant_id, session): return object() - def _find_existing_workflow_tool(self, session, tenant_id, workflow_tool_id, tool_name, app_id): + def _find_existing_workflow_tool(self, tenant_id, workflow_tool_id, tool_name, app_id, session): if ("created", app_id) in events: return type("WorkflowToolProvider", (), {"id": workflow_tool_id or "created-workflow-tool-id"})() return None - def _ensure_workflow_app_is_published(self, session, target, account, app_id): + def _ensure_workflow_app_is_published(self, target, account, app_id, session): events.append(("published", app_id)) from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr( import_service.WorkflowToolManageService, "create_workflow_tool", @@ -337,7 +360,6 @@ def test_workflow_tool_import_publishes_referenced_app_before_create(monkeypatch ) PublishingImportService()._import_workflow_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -360,6 +382,7 @@ def test_workflow_tool_import_publishes_referenced_app_before_create(monkeypatch {}, [], [], + session=import_service.db.session, ) assert events == [("published", "workflow-app-1"), ("created", "workflow-app-1")] @@ -384,17 +407,18 @@ def test_workflow_tool_import_id_follows_id_strategy(monkeypatch: pytest.MonkeyP return account class StrategyImportService(MigrationImportService): - def _find_existing_app(self, session, app_id, tenant_id): + def _find_existing_app(self, app_id, tenant_id, session): return object() - def _find_existing_workflow_tool(self, session, tenant_id, workflow_tool_id, tool_name, app_id): + def _find_existing_workflow_tool(self, tenant_id, workflow_tool_id, tool_name, app_id, session): return target_provider if created_kwargs else None - def _ensure_workflow_app_is_published(self, session, target, account, app_id): + def _ensure_workflow_app_is_published(self, target, account, app_id, session): return None from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr( import_service.WorkflowToolManageService, "create_workflow_tool", @@ -402,7 +426,6 @@ def test_workflow_tool_import_id_follows_id_strategy(monkeypatch: pytest.MonkeyP ) StrategyImportService()._import_workflow_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -425,6 +448,7 @@ def test_workflow_tool_import_id_follows_id_strategy(monkeypatch: pytest.MonkeyP id_mapping, id_mapping_details, [], + session=import_service.db.session, ) assert created_kwargs[0]["import_id"] == expected_import_id @@ -449,17 +473,20 @@ def test_workflow_tool_skip_records_id_mapping(monkeypatch): return account class SkipImportService(MigrationImportService): - def _find_existing_app(self, session, app_id, tenant_id): + def _find_existing_app(self, app_id, tenant_id, session): return object() - def _find_existing_workflow_tool(self, session, tenant_id, workflow_tool_id, tool_name, app_id): + def _find_existing_workflow_tool(self, tenant_id, workflow_tool_id, tool_name, app_id, session): return existing_provider - def _ensure_workflow_app_is_published(self, session, target, account, app_id): + def _ensure_workflow_app_is_published(self, target, account, app_id, session): return None + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + SkipImportService()._import_workflow_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -482,6 +509,7 @@ def test_workflow_tool_skip_records_id_mapping(monkeypatch): id_mapping, [], [], + session=import_service.db.session, ) assert id_mapping["source-workflow-tool-id"] == "existing-workflow-tool-id" @@ -495,22 +523,18 @@ def test_api_tool_existing_provider_records_id_mapping(monkeypatch, conflict_str report_items = [] class ExistingApiImportService(MigrationImportService): - def _find_api_tool_provider(self, session, tenant_id, provider_name): - return target_provider - - class StubSession: - def scalar(self, statement): + def _find_api_tool_provider(self, tenant_id, provider_name, session): return target_provider from services.data_migration import import_service + monkeypatch.setattr(import_service.db.session, "scalar", lambda statement: target_provider) monkeypatch.setattr( import_service.ApiToolManageService, "parser_api_schema", lambda schema: {"schema_type": "openapi"} ) monkeypatch.setattr(import_service.ApiToolManageService, "update_api_tool_provider", lambda **kwargs: None) ExistingApiImportService()._import_api_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -528,6 +552,7 @@ def test_api_tool_existing_provider_records_id_mapping(monkeypatch, conflict_str id_mapping, id_mapping_details, {"weather": {"source-api-provider-id-from-dsl"}}, + session=import_service.db.session, ) assert id_mapping == { @@ -549,18 +574,18 @@ def test_api_tool_create_records_id_mapping(monkeypatch): return None class CreatedApiImportService(MigrationImportService): - def _find_api_tool_provider(self, session, tenant_id, provider_name): + def _find_api_tool_provider(self, tenant_id, provider_name, session): return target_provider from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr( import_service.ApiToolManageService, "parser_api_schema", lambda schema: {"schema_type": "openapi"} ) monkeypatch.setattr(import_service.ApiToolManageService, "create_api_tool_provider", lambda **kwargs: None) CreatedApiImportService()._import_api_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -578,6 +603,7 @@ def test_api_tool_create_records_id_mapping(monkeypatch): id_mapping, [], {}, + session=import_service.db.session, ) assert id_mapping["source-api-provider-id"] == "target-api-provider-id" @@ -605,10 +631,10 @@ def test_mcp_tool_import_restores_exported_tool_list(monkeypatch): from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr(import_service, "MCPToolManageService", StubMCPToolManageService) MigrationImportService()._import_mcp_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -634,6 +660,7 @@ def test_mcp_tool_import_restores_exported_tool_list(monkeypatch): report_items, {}, [], + session=import_service.db.session, ) assert provider.tools == '[{"name": "echo"}]' @@ -653,7 +680,7 @@ def test_mcp_tool_existing_provider_records_id_mapping(monkeypatch, conflict_str return None class ExistingMCPImportService(MigrationImportService): - def _find_existing_mcp_tool(self, session, tenant_id, provider_id, server_identifier): + def _find_existing_mcp_tool(self, tenant_id, provider_id, server_identifier, session): return provider class StubMCPToolManageService: @@ -665,10 +692,10 @@ def test_mcp_tool_existing_provider_records_id_mapping(monkeypatch, conflict_str from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr(import_service, "MCPToolManageService", StubMCPToolManageService) ExistingMCPImportService()._import_mcp_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -694,6 +721,7 @@ def test_mcp_tool_existing_provider_records_id_mapping(monkeypatch, conflict_str [], id_mapping, id_mapping_details, + session=import_service.db.session, ) assert id_mapping["source-mcp-provider-id"] == "target-mcp-provider-id" @@ -715,7 +743,7 @@ def test_mcp_tool_create_records_id_mapping(monkeypatch): return None class CreatedMCPImportService(MigrationImportService): - def _find_existing_mcp_tool(self, session, tenant_id, provider_id, server_identifier): + def _find_existing_mcp_tool(self, tenant_id, provider_id, server_identifier, session): return provider if provider_created else None class StubMCPToolManageService: @@ -728,10 +756,10 @@ def test_mcp_tool_create_records_id_mapping(monkeypatch): from services.data_migration import import_service + monkeypatch.setattr(import_service.db, "session", StubSession()) monkeypatch.setattr(import_service, "MCPToolManageService", StubMCPToolManageService) CreatedMCPImportService()._import_mcp_tools( - StubSession(), MigrationPackage.from_mapping( { "metadata": {"version": "1", "source_scope": "single"}, @@ -756,6 +784,7 @@ def test_mcp_tool_create_records_id_mapping(monkeypatch): [], id_mapping, [], + session=import_service.db.session, ) assert id_mapping["source-mcp-provider-id"] == "target-mcp-provider-id" @@ -800,12 +829,11 @@ def test_dependency_only_mcp_preflight_reports_missing_target_provider_with_work } ) - class StubSession: - def scalar(self, statement): - return None + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db.session, "scalar", lambda statement: None) MigrationImportService()._preflight_dependency_only_mcp( - StubSession(), package, ImportTarget( tenant_id="tenant-1", @@ -814,6 +842,7 @@ def test_dependency_only_mcp_preflight_reports_missing_target_provider_with_work operator_email="owner@example.com", ), report_items, + session=import_service.db.session, ) assert report_items == [ @@ -828,18 +857,22 @@ def test_dependency_only_mcp_preflight_reports_missing_target_provider_with_work ] -def test_dependency_only_mcp_lookup_does_not_compare_non_uuid_identifier_to_uuid_id(): +def test_dependency_only_mcp_lookup_does_not_compare_non_uuid_identifier_to_uuid_id(monkeypatch): captured = [] class StubSession: def scalar(self, statement): captured.append(statement) + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db, "session", StubSession()) + MigrationImportService()._find_dependency_only_mcp_provider( - StubSession(), "tenant-1", "my-test-mcp-server", "my-test-mcp", + session=import_service.db.session, ) where_clause = str(captured[0].whereclause) @@ -860,12 +893,11 @@ def test_dependency_only_mcp_preflight_reports_available_target_provider(monkeyp {"id": "target-provider-id", "name": "my-test-mcp", "server_identifier": "my-test-mcp-server"}, )() - class StubSession: - def scalar(self, statement): - return provider + from services.data_migration import import_service + + monkeypatch.setattr(import_service.db.session, "scalar", lambda statement: provider) MigrationImportService()._preflight_dependency_only_mcp( - StubSession(), package, ImportTarget( tenant_id="tenant-1", @@ -874,6 +906,7 @@ def test_dependency_only_mcp_preflight_reports_available_target_provider(monkeyp operator_email="owner@example.com", ), report_items, + session=import_service.db.session, ) assert report_items == [ @@ -891,7 +924,7 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): events = [] class StubResolver(ImportTargetResolver): - def resolve(self, session, request): + def resolve(self, request, session): return ImportTarget( tenant_id="tenant-1", tenant_name="target", @@ -902,7 +935,6 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): class OrderedImportService(MigrationImportService): def _import_api_tools( self, - session, package, target, options, @@ -910,12 +942,13 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): id_mapping, id_mapping_details, source_provider_ids_by_name, + *, + session=None, ): events.append(("api_tools", "imported")) def _import_workflows( self, - session, package, target, options, @@ -926,6 +959,7 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): imported_workflow_ids=None, only_app_ids=None, skip_app_ids=None, + session=None, ): only_app_ids = set(only_app_ids or []) skip_app_ids = set(skip_app_ids or []) @@ -941,11 +975,13 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): imported_workflow_ids.add(app_id) def _import_workflow_tools( - self, session, package, target, options, id_mapping, id_mapping_details, report_items + self, package, target, options, id_mapping, id_mapping_details, report_items, *, session=None ): events.append(("workflow_tool", package.workflow_tools[0]["id"])) - def _import_mcp_tools(self, session, package, target, options, report_items, id_mapping, id_mapping_details): + def _import_mcp_tools( + self, package, target, options, report_items, id_mapping, id_mapping_details, *, session=None + ): events.append(("mcp_tools", "imported")) package = MigrationPackage.from_mapping( @@ -959,7 +995,9 @@ def test_import_package_imports_workflow_tool_provider_apps_before_consumers(): } ) - OrderedImportService(target_resolver=StubResolver()).import_package(object(), ImportRequest(package=package)) + OrderedImportService(target_resolver=StubResolver()).import_package( + ImportRequest(package=package), session=import_service.db.session + ) assert events == [ ("api_tools", "imported"), diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_service.py b/api/tests/unit_tests/services/enterprise/test_rbac_service.py index fdf921265b2..85638b11fff 100644 --- a/api/tests/unit_tests/services/enterprise/test_rbac_service.py +++ b/api/tests/unit_tests/services/enterprise/test_rbac_service.py @@ -558,7 +558,7 @@ class TestMyPermissions: } with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): - out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1") + out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", session=MagicMock()) call = _call_args(mock_send) assert call.method == "GET" @@ -613,11 +613,8 @@ class TestMyPermissions: mock_session = MagicMock() mock_session.__enter__.return_value = mock_session mock_session.scalar.return_value = role - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=mock_session), - ): - out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1") + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", session=mock_session) mock_send.assert_not_called() assert out.workspace.permission_keys == workspace_keys @@ -655,11 +652,8 @@ class TestMyPermissions: mock_session = MagicMock() mock_session.__enter__.return_value = mock_session mock_session.scalar.return_value = role - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=mock_session), - ): - out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1") + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", session=mock_session) actual_snippet_keys = { permission_key for permission_key in out.workspace.permission_keys if permission_key.startswith("snippets.") @@ -672,11 +666,8 @@ class TestMyPermissions: mock_session = MagicMock() mock_session.__enter__.return_value = mock_session mock_session.scalar.return_value = None - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=mock_session), - ): - out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1") + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", session=mock_session) mock_send.assert_not_called() assert out.workspace.permission_keys == [] @@ -694,7 +685,7 @@ class TestMyPermissions: } with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): - out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", app_id="app-1") + out = svc.RBACService.MyPermissions.get("tenant-1", "acct-1", app_id="app-1", session=MagicMock()) call = _call_args(mock_send) assert call.method == "GET" @@ -716,7 +707,7 @@ class TestMemberRoles: ], } with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): - out = svc.RBACService.MemberRoles.get("tenant-1", "acct-1", "acct-2") + out = svc.RBACService.MemberRoles.get("tenant-1", "acct-1", "acct-2", session=MagicMock()) call = _call_args(mock_send) assert call.method == "GET" assert call.endpoint == "/rbac/members/rbac-roles" @@ -728,12 +719,8 @@ class TestMemberRoles: session = MagicMock() session.scalar.return_value = svc.TenantAccountRole.EDITOR - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session") as create_session, - ): - create_session.return_value.__enter__.return_value = session - out = svc.RBACService.MemberRoles.get("tenant-1", "acct-1", "acct-2") + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MemberRoles.get("tenant-1", "acct-1", "acct-2", session=session) mock_send.assert_not_called() assert out.account_id == "acct-2" @@ -755,7 +742,11 @@ class TestMemberRoles: mock_send.return_value = {"account_id": "acct-2", "roles": []} with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): svc.RBACService.MemberRoles.replace( - "tenant-1", "acct-1", "acct-2", role_ids=["workspace.owner", "workspace.editor"] + "tenant-1", + "acct-1", + "acct-2", + role_ids=["workspace.owner", "workspace.editor"], + session=MagicMock(), ) call = _call_args(mock_send) assert call.method == "PUT" @@ -769,11 +760,10 @@ class TestMemberRoles: target_join = SimpleNamespace(role=svc.TenantAccountRole.NORMAL, account_id="acct-2") session.scalar.return_value = target_join - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=session), - ): - out = svc.RBACService.MemberRoles.replace("tenant-1", "acct-1", "acct-2", role_ids=["editor"]) + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MemberRoles.replace( + "tenant-1", "acct-1", "acct-2", role_ids=["editor"], session=session + ) mock_send.assert_not_called() session.commit.assert_called_once() @@ -789,11 +779,10 @@ class TestMemberRoles: owner_join = SimpleNamespace(role=svc.TenantAccountRole.OWNER, account_id="acct-owner") session.scalar.side_effect = [target_join, owner_join] - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=session), - ): - out = svc.RBACService.MemberRoles.replace("tenant-1", "acct-1", "acct-2", role_ids=["owner"]) + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.MemberRoles.replace( + "tenant-1", "acct-1", "acct-2", role_ids=["owner"], session=session + ) mock_send.assert_not_called() session.commit.assert_called_once() @@ -832,7 +821,9 @@ class TestResourcePermissions: } with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): - out = svc.RBACService.AppPermissions.batch_get("tenant-1", "acct-1", ["app-1", "app-2"]) + out = svc.RBACService.AppPermissions.batch_get( + "tenant-1", "acct-1", ["app-1", "app-2"], session=MagicMock() + ) call = _call_args(mock_send) assert call.method == "POST" @@ -847,11 +838,10 @@ class TestResourcePermissions: mock_session = MagicMock() mock_session.__enter__.return_value = mock_session mock_session.scalar.return_value = "editor" - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=mock_session), - ): - out = svc.RBACService.AppPermissions.batch_get("tenant-1", "acct-1", ["app-1", "app-2"]) + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.AppPermissions.batch_get( + "tenant-1", "acct-1", ["app-1", "app-2"], session=mock_session + ) mock_send.assert_not_called() assert out == { @@ -868,7 +858,9 @@ class TestResourcePermissions: } with patch(f"{MODULE}.dify_config.RBAC_ENABLED", True): - out = svc.RBACService.DatasetPermissions.batch_get("tenant-1", "acct-1", ["ds-1", "ds-2"]) + out = svc.RBACService.DatasetPermissions.batch_get( + "tenant-1", "acct-1", ["ds-1", "ds-2"], session=MagicMock() + ) call = _call_args(mock_send) assert call.method == "POST" @@ -883,11 +875,10 @@ class TestResourcePermissions: mock_session = MagicMock() mock_session.__enter__.return_value = mock_session mock_session.scalar.return_value = "dataset_operator" - with ( - patch(f"{MODULE}.dify_config.RBAC_ENABLED", False), - patch(f"{MODULE}.session_factory.create_session", return_value=mock_session), - ): - out = svc.RBACService.DatasetPermissions.batch_get("tenant-1", "acct-1", ["ds-1", "ds-2"]) + with patch(f"{MODULE}.dify_config.RBAC_ENABLED", False): + out = svc.RBACService.DatasetPermissions.batch_get( + "tenant-1", "acct-1", ["ds-1", "ds-2"], session=mock_session + ) mock_send.assert_not_called() assert out == { diff --git a/api/tests/unit_tests/services/hit_service.py b/api/tests/unit_tests/services/hit_service.py index ae19daba898..ffeb158e37a 100644 --- a/api/tests/unit_tests/services/hit_service.py +++ b/api/tests/unit_tests/services/hit_service.py @@ -186,7 +186,7 @@ class TestHitTestingServiceRetrieve: # Act result = HitTestingService.retrieve( - mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model + dataset, query, account, retrieval_model, external_retrieval_model, session=mock_db_session ) # Assert @@ -234,7 +234,7 @@ class TestHitTestingServiceRetrieve: # Act result = HitTestingService.retrieve( - mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model + dataset, query, account, retrieval_model, external_retrieval_model, session=mock_db_session ) # Assert @@ -292,7 +292,7 @@ class TestHitTestingServiceRetrieve: # Act result = HitTestingService.retrieve( - mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model + dataset, query, account, retrieval_model, external_retrieval_model, session=mock_db_session ) # Assert @@ -337,7 +337,7 @@ class TestHitTestingServiceRetrieve: # Act result = HitTestingService.retrieve( - mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model + dataset, query, account, retrieval_model, external_retrieval_model, session=mock_db_session ) # Assert @@ -380,7 +380,7 @@ class TestHitTestingServiceRetrieve: # Act result = HitTestingService.retrieve( - mock_db_session, dataset, query, account, retrieval_model, external_retrieval_model + dataset, query, account, retrieval_model, external_retrieval_model, session=mock_db_session ) # Assert @@ -438,7 +438,12 @@ class TestHitTestingServiceExternalRetrieve: # Act result = HitTestingService.external_retrieve( - mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions + dataset, + query, + account, + external_retrieval_model, + metadata_filtering_conditions, + session=mock_db_session, ) # Assert @@ -469,7 +474,7 @@ class TestHitTestingServiceExternalRetrieve: # Act result = HitTestingService.external_retrieve( - mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions + dataset, query, account, external_retrieval_model, metadata_filtering_conditions, session=mock_db_session ) # Assert @@ -504,7 +509,12 @@ class TestHitTestingServiceExternalRetrieve: # Act result = HitTestingService.external_retrieve( - mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions + dataset, + query, + account, + external_retrieval_model, + metadata_filtering_conditions, + session=mock_db_session, ) # Assert @@ -538,7 +548,12 @@ class TestHitTestingServiceExternalRetrieve: # Act result = HitTestingService.external_retrieve( - mock_db_session, dataset, query, account, external_retrieval_model, metadata_filtering_conditions + dataset, + query, + account, + external_retrieval_model, + metadata_filtering_conditions, + session=mock_db_session, ) # Assert @@ -579,7 +594,7 @@ class TestHitTestingServiceCompactRetrieveResponse: mock_format.return_value = mock_records # Act - result = HitTestingService.compact_retrieve_response(MagicMock(), query, documents) + result = HitTestingService.compact_retrieve_response(query, documents, session=MagicMock()) # Assert assert result["query"]["content"] == query @@ -605,7 +620,7 @@ class TestHitTestingServiceCompactRetrieveResponse: mock_format.return_value = [] # Act - result = HitTestingService.compact_retrieve_response(MagicMock(), query, documents) + result = HitTestingService.compact_retrieve_response(query, documents, session=MagicMock()) # Assert assert result["query"]["content"] == query diff --git a/api/tests/unit_tests/services/plugin/test_dependencies_analysis.py b/api/tests/unit_tests/services/plugin/test_dependencies_analysis.py index 8f0886769cf..178dec1bfae 100644 --- a/api/tests/unit_tests/services/plugin/test_dependencies_analysis.py +++ b/api/tests/unit_tests/services/plugin/test_dependencies_analysis.py @@ -11,7 +11,7 @@ from unittest.mock import MagicMock, patch import pytest -from core.plugin.entities.plugin import PluginDependency, PluginInstallationSource +from core.plugin.entities.plugin import PluginDependency, PluginDependencyType, PluginInstallationSource from services.plugin.dependencies_analysis import DependenciesAnalysisService @@ -44,7 +44,7 @@ class TestAnalyzeModelProviderDependency: class TestGetLeakedDependencies: - def _make_dependency(self, identifier: str, dep_type=PluginDependency.Type.Marketplace): + def _make_dependency(self, identifier: str, dep_type=PluginDependencyType.Marketplace): return PluginDependency( type=dep_type, value=PluginDependency.Marketplace(marketplace_plugin_unique_identifier=identifier), @@ -110,7 +110,7 @@ class TestGenerateDependencies: result = DependenciesAnalysisService.generate_dependencies("t1", ["p1"]) assert len(result) == 1 - assert result[0].type == PluginDependency.Type.Github + assert result[0].type == PluginDependencyType.Github assert result[0].value.repo == "org/repo" @patch("services.plugin.dependencies_analysis.PluginInstaller") @@ -120,7 +120,7 @@ class TestGenerateDependencies: result = DependenciesAnalysisService.generate_dependencies("t1", ["p1"]) - assert result[0].type == PluginDependency.Type.Marketplace + assert result[0].type == PluginDependencyType.Marketplace @patch("services.plugin.dependencies_analysis.PluginInstaller") def test_package_source(self, mock_installer_cls): @@ -129,7 +129,7 @@ class TestGenerateDependencies: result = DependenciesAnalysisService.generate_dependencies("t1", ["p1"]) - assert result[0].type == PluginDependency.Type.Package + assert result[0].type == PluginDependencyType.Package @patch("services.plugin.dependencies_analysis.PluginInstaller") def test_remote_source_raises(self, mock_installer_cls): @@ -169,4 +169,4 @@ class TestGenerateLatestDependencies: result = DependenciesAnalysisService.generate_latest_dependencies(["p1"]) assert len(result) == 1 - assert result[0].type == PluginDependency.Type.Marketplace + assert result[0].type == PluginDependencyType.Marketplace diff --git a/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py b/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py index e284c98079d..e66bb3fff04 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py +++ b/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py @@ -4,222 +4,216 @@ from unittest.mock import MagicMock, patch import pytest -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import ( + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategySetting, +) MODULE = "services.plugin.plugin_auto_upgrade_service" -PLUGIN_CATEGORY = TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL +PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL def _patched_session(): - """Patch session_factory.create_session() to return a mock session as context manager.""" + """Return a mock SQLAlchemy session for service calls.""" session = MagicMock() - session.__enter__ = MagicMock(return_value=session) - session.__exit__ = MagicMock(return_value=False) - mock_factory = MagicMock() - mock_factory.create_session.return_value = session - patcher = patch(f"{MODULE}.session_factory", mock_factory) - return patcher, session + return session class TestGetStrategy: def test_returns_strategy_when_found(self): - p1, session = _patched_session() + session = _patched_session() strategy = MagicMock() session.scalar.return_value = strategy - with p1: - from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService + from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY) + result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY, session=session) assert result is strategy def test_returns_none_when_not_found(self): - p1, session = _patched_session() + session = _patched_session() session.scalar.return_value = None - with p1: - from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService + from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY) + result = PluginAutoUpgradeService.get_strategy("t1", PLUGIN_CATEGORY, session=session) assert result is None class TestChangeStrategy: def test_creates_new_strategy(self): - p1, session = _patched_session() + session = _patched_session() session.scalar.return_value = None - with p1, patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: + with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: strat_cls.return_value = MagicMock() from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService result = PluginAutoUpgradeService.change_strategy( "t1", - TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, 3, - TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + TenantPluginAutoUpgradeMode.ALL, [], [], category=PLUGIN_CATEGORY, + session=session, ) assert result is True session.add.assert_called_once() def test_updates_existing_strategy(self): - p1, session = _patched_session() + session = _patched_session() existing = MagicMock() session.scalar.return_value = existing - with p1: - from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService + from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.change_strategy( - "t1", - TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, - 5, - TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL, - ["p1"], - ["p2"], - category=PLUGIN_CATEGORY, - ) + result = PluginAutoUpgradeService.change_strategy( + "t1", + TenantPluginAutoUpgradeStrategySetting.LATEST, + 5, + TenantPluginAutoUpgradeMode.PARTIAL, + ["p1"], + ["p2"], + category=PLUGIN_CATEGORY, + session=session, + ) assert result is True - assert existing.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST + assert existing.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST assert existing.upgrade_time_of_day == 5 - assert existing.upgrade_mode == TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL + assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL assert existing.exclude_plugins == ["p1"] assert existing.include_plugins == ["p2"] class TestExcludePlugin: def test_creates_default_strategy_when_none_exists(self): - p1, session = _patched_session() + session = _patched_session() session.scalar.return_value = None with ( - p1, patch(f"{MODULE}.select"), - patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls, + patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy"), ): - strat_cls.StrategySetting.FIX_ONLY = "fix_only" - strat_cls.UpgradeMode.EXCLUDE = "exclude" from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService result = PluginAutoUpgradeService.exclude_plugin( "t1", "plugin-1", PLUGIN_CATEGORY, + session=session, ) assert result is True session.add.assert_called_once() def test_appends_to_exclude_list_in_exclude_mode(self): - p1, session = _patched_session() + session = _patched_session() existing = MagicMock() - existing.upgrade_mode = "exclude" + existing.upgrade_mode = TenantPluginAutoUpgradeMode.EXCLUDE existing.exclude_plugins = ["p-existing"] session.scalar.return_value = existing - with p1, patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: + with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: strat_cls.UpgradeMode.EXCLUDE = "exclude" strat_cls.UpgradeMode.PARTIAL = "partial" strat_cls.UpgradeMode.ALL = "all" from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.exclude_plugin("t1", "p-new", PLUGIN_CATEGORY) + result = PluginAutoUpgradeService.exclude_plugin("t1", "p-new", PLUGIN_CATEGORY, session=session) assert result is True assert existing.exclude_plugins == ["p-existing", "p-new"] def test_removes_from_include_list_in_partial_mode(self): - p1, session = _patched_session() + session = _patched_session() existing = MagicMock() - existing.upgrade_mode = "partial" + existing.upgrade_mode = TenantPluginAutoUpgradeMode.PARTIAL existing.include_plugins = ["p1", "p2"] session.scalar.return_value = existing - with p1, patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: + with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: strat_cls.UpgradeMode.EXCLUDE = "exclude" strat_cls.UpgradeMode.PARTIAL = "partial" strat_cls.UpgradeMode.ALL = "all" from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY) + result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session) assert result is True assert existing.include_plugins == ["p2"] def test_switches_to_exclude_mode_from_all(self): - p1, session = _patched_session() + session = _patched_session() existing = MagicMock() - existing.upgrade_mode = "all" + existing.upgrade_mode = TenantPluginAutoUpgradeMode.ALL session.scalar.return_value = existing - with p1, patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: + with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: strat_cls.UpgradeMode.EXCLUDE = "exclude" strat_cls.UpgradeMode.PARTIAL = "partial" strat_cls.UpgradeMode.ALL = "all" from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY) + result = PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session) assert result is True - assert existing.upgrade_mode == "exclude" + assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE assert existing.exclude_plugins == ["p1"] def test_no_duplicate_in_exclude_list(self): - p1, session = _patched_session() + session = _patched_session() existing = MagicMock() - existing.upgrade_mode = "exclude" + existing.upgrade_mode = TenantPluginAutoUpgradeMode.EXCLUDE existing.exclude_plugins = ["p1"] session.scalar.return_value = existing - with p1, patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: + with patch(f"{MODULE}.select"), patch(f"{MODULE}.TenantPluginAutoUpgradeStrategy") as strat_cls: strat_cls.UpgradeMode.EXCLUDE = "exclude" strat_cls.UpgradeMode.PARTIAL = "partial" strat_cls.UpgradeMode.ALL = "all" from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY) + PluginAutoUpgradeService.exclude_plugin("t1", "p1", PLUGIN_CATEGORY, session=session) assert existing.exclude_plugins == ["p1"] class TestBackfillStrategyCategories: def test_creates_default_missing_categories_without_fetching_daemon(self): - p1, session = _patched_session() + session = _patched_session() tool_strategy = SimpleNamespace( - category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + category=TenantPluginAutoUpgradeCategory.TOOL, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=[], include_plugins=[], ) session.scalars.return_value.all.return_value = [tool_strategy] installer = MagicMock() - with p1, patch(f"{MODULE}.PluginInstaller", return_value=installer): + with patch(f"{MODULE}.PluginInstaller", return_value=installer): from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.backfill_strategy_categories("t1") + result = PluginAutoUpgradeService.backfill_strategy_categories("t1", session=session) expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day("t1") - assert result.created_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 1 + assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 1 assert result.normalized is False installer.list_plugins.assert_not_called() assert tool_strategy.upgrade_time_of_day == expected_time created_strategies = [call.args[0] for call in session.add.call_args_list] model_strategy = next( - strategy - for strategy in created_strategies - if strategy.category == TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL + strategy for strategy in created_strategies if strategy.category == TenantPluginAutoUpgradeCategory.MODEL ) - assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST + assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST assert model_strategy.upgrade_time_of_day == expected_time def test_default_upgrade_time_is_aligned_to_fifteen_minutes(self): @@ -231,20 +225,20 @@ class TestBackfillStrategyCategories: assert 0 <= default_time < 24 * 60 * 60 def test_creates_missing_categories_and_splits_known_plugins(self, caplog: pytest.LogCaptureFixture): - p1, session = _patched_session() + session = _patched_session() tool_strategy = SimpleNamespace( - category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + category=TenantPluginAutoUpgradeCategory.TOOL, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"], include_plugins=["model-plugin", "tool-plugin"], ) model_strategy = SimpleNamespace( - category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL, - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + category=TenantPluginAutoUpgradeCategory.MODEL, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=["tool-plugin", "model-plugin", "unknown-plugin"], include_plugins=["model-plugin", "tool-plugin"], ) @@ -253,28 +247,27 @@ class TestBackfillStrategyCategories: installed_plugins = [ SimpleNamespace( plugin_id="tool-plugin", - declaration=SimpleNamespace(category=TenantPluginAutoUpgradeStrategy.PluginCategory.TOOL), + declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.TOOL), ), SimpleNamespace( plugin_id="model-plugin", - declaration=SimpleNamespace(category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL), + declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.MODEL), ), ] installer = MagicMock() installer.list_plugins.return_value = installed_plugins with ( - p1, patch(f"{MODULE}.PluginInstaller", return_value=installer), caplog.at_level(logging.WARNING, logger=MODULE), ): from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService - result = PluginAutoUpgradeService.backfill_strategy_categories("t1") + result = PluginAutoUpgradeService.backfill_strategy_categories("t1", session=session) - assert result.created_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 2 + assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 2 assert result.normalized is True - assert session.add.call_count == len(TenantPluginAutoUpgradeStrategy.PluginCategory) - 2 + assert session.add.call_count == len(TenantPluginAutoUpgradeCategory) - 2 assert tool_strategy.exclude_plugins == ["tool-plugin"] assert tool_strategy.include_plugins == ["tool-plugin"] assert model_strategy.exclude_plugins == ["model-plugin"] diff --git a/api/tests/unit_tests/services/plugin/test_plugin_migration.py b/api/tests/unit_tests/services/plugin/test_plugin_migration.py index 8f730d4ed31..d94ab540dfd 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_migration.py +++ b/api/tests/unit_tests/services/plugin/test_plugin_migration.py @@ -1,3 +1,4 @@ +import json from unittest.mock import MagicMock, patch import pytest @@ -8,17 +9,17 @@ from services.plugin.plugin_migration import PluginMigration MIGRATION_MODULE = "services.plugin.plugin_migration" -def test_fetch_plugin_unique_identifier_returns_none_when_disabled(mocker: MockerFixture) -> None: +def test_fetch_latest_package_identifier_returns_none_when_disabled(mocker: MockerFixture) -> None: mocker.patch("services.plugin.plugin_migration.dify_config.MARKETPLACE_ENABLED", False) batch_fetch = mocker.patch("services.plugin.plugin_migration.marketplace.batch_fetch_plugin_manifests") - result = PluginMigration._fetch_plugin_unique_identifier("langgenius/openai") + result = PluginMigration._fetch_latest_package_identifier("langgenius/openai") assert result is None batch_fetch.assert_not_called() -def test_fetch_plugin_unique_identifier_calls_marketplace_when_enabled(mocker: MockerFixture) -> None: +def test_fetch_latest_package_identifier_calls_marketplace_when_enabled(mocker: MockerFixture) -> None: mocker.patch("services.plugin.plugin_migration.dify_config.MARKETPLACE_ENABLED", True) manifest = mocker.MagicMock() manifest.latest_package_identifier = "langgenius/openai:1.0.0@abc" @@ -27,7 +28,7 @@ def test_fetch_plugin_unique_identifier_calls_marketplace_when_enabled(mocker: M return_value=[manifest], ) - result = PluginMigration._fetch_plugin_unique_identifier("langgenius/openai") + result = PluginMigration._fetch_latest_package_identifier("langgenius/openai") assert result == "langgenius/openai:1.0.0@abc" @@ -75,7 +76,27 @@ class TestHandlePluginInstanceInstall: mock_marketplace.download_plugin_pkg.assert_called_once() invalidate_cache.assert_called_once_with("tenant1") - assert "success" in result or "failed" in result + assert result["success"] == ["langgenius/openai"] + assert result["failed"] == [] + + def test_reports_failed_plugin_ids_when_install_batch_raises(self) -> None: + with ( + patch(f"{MIGRATION_MODULE}.dify_config") as mock_cfg, + patch(f"{MIGRATION_MODULE}.marketplace") as mock_marketplace, + patch(f"{MIGRATION_MODULE}.PluginInstaller") as mock_installer_cls, + ): + mock_cfg.MARKETPLACE_ENABLED = True + mock_marketplace.download_plugin_pkg.return_value = b"pkg_data" + mock_installer = MagicMock() + mock_installer_cls.return_value = mock_installer + mock_installer.install_from_identifiers.side_effect = RuntimeError("install failed") + + result = PluginMigration.handle_plugin_instance_install( + "tenant1", {"langgenius/openai": "langgenius/openai:1.0.0@abc"} + ) + + assert result["success"] == [] + assert result["failed"] == ["langgenius/openai"] def test_install_plugins_invalidates_cache_after_direct_tenant_install(self, tmp_path) -> None: extracted_plugins = tmp_path / "plugins.jsonl" @@ -102,3 +123,60 @@ class TestHandlePluginInstanceInstall: mock_installer.install_from_identifiers.assert_called_once() invalidate_cache.assert_called_once_with("tenant1") + + def test_install_plugins_reports_missing_plugin_ids(self, tmp_path) -> None: + extracted_plugins = tmp_path / "plugins.jsonl" + output_file = tmp_path / "output.json" + extracted_plugins.write_text('{"tenant_id":"tenant1","plugins":["langgenius/openai","langgenius/missing"]}\n') + + with ( + patch( + f"{MIGRATION_MODULE}.PluginMigration.extract_unique_plugins", + return_value={ + "plugins": {"langgenius/openai": "langgenius/openai:1.0.0@abc"}, + "plugin_not_exist": ["langgenius/missing"], + }, + ), + patch(f"{MIGRATION_MODULE}.PluginMigration.handle_plugin_instance_install", return_value={}), + patch(f"{MIGRATION_MODULE}.PluginInstaller") as mock_installer_cls, + patch(f"{MIGRATION_MODULE}.PluginService.invalidate_plugin_model_providers_cache"), + ): + mock_installer = MagicMock() + mock_installer.list_plugins.return_value = [] + mock_installer_cls.return_value = mock_installer + + PluginMigration.install_plugins(str(extracted_plugins), str(output_file), workers=1) + + assert json.loads(output_file.read_text())["not_installed"] == [ + { + "tenant_id": "tenant1", + "plugin_not_exist": ["langgenius/missing"], + } + ] + mock_installer.install_from_identifiers.assert_called_once() + + def test_install_plugins_skips_unresolved_plugins(self, tmp_path) -> None: + extracted_plugins = tmp_path / "plugins.jsonl" + output_file = tmp_path / "output.json" + extracted_plugins.write_text('{"tenant_id":"tenant1","plugins":["langgenius/missing"]}\n') + + with ( + patch( + f"{MIGRATION_MODULE}.PluginMigration.extract_unique_plugins", + return_value={ + "plugins": {}, + "plugin_not_exist": ["langgenius/missing"], + }, + ), + patch(f"{MIGRATION_MODULE}.PluginMigration.handle_plugin_instance_install", return_value={}), + patch(f"{MIGRATION_MODULE}.PluginInstaller") as mock_installer_cls, + ): + mock_installer = MagicMock() + mock_installer.list_plugins.return_value = [] + mock_installer_cls.return_value = mock_installer + + PluginMigration.install_plugins(str(extracted_plugins), str(output_file), workers=1) + + output = json.loads(output_file.read_text()) + assert output["not_installed"] == [{"tenant_id": "tenant1", "plugin_not_exist": ["langgenius/missing"]}] + mock_installer.install_from_identifiers.assert_not_called() diff --git a/api/tests/unit_tests/services/plugin/test_plugin_service.py b/api/tests/unit_tests/services/plugin/test_plugin_service.py index a8922154a95..278898926b9 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_service.py +++ b/api/tests/unit_tests/services/plugin/test_plugin_service.py @@ -8,7 +8,7 @@ import zstandard from pydantic import TypeAdapter from redis import RedisError -from core.plugin.entities.plugin import PluginInstallationSource +from core.plugin.entities.plugin import PluginCategory, PluginInstallationSource from core.plugin.entities.plugin_daemon import PluginInstallTask, PluginInstallTaskStatus, PluginModelProviderEntity from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity @@ -71,6 +71,16 @@ def _build_install_task(*, task_id: str = "task-1", status: PluginInstallTaskSta ) +def _build_remote_model_plugin( + *, plugin_id: str = "langgenius/debug-model", plugin_unique_identifier: str = "langgenius/debug-model:1.0.0" +) -> SimpleNamespace: + return SimpleNamespace( + plugin_id=plugin_id, + plugin_unique_identifier=plugin_unique_identifier, + source=PluginInstallationSource.Remote, + ) + + def _provider_cache_key(tenant_id: str, generation: int | None = None) -> str: if generation is None: return f"plugin_model_providers:tenant_id:{tenant_id}" @@ -797,6 +807,144 @@ class TestPluginListEndpointCounts: class TestPluginModelProviderCacheInvalidation: + def test_get_debugging_key_does_not_invalidate_model_provider_cache(self) -> None: + """Reading a debug key does not mean a debug runtime has registered a model provider.""" + with ( + patch(f"{MODULE}.PluginDebuggingClient") as debugging_client_cls, + patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache, + ): + debugging_client_cls.return_value.get_debugging_key.return_value = "debug-key" + + from core.plugin.plugin_service import PluginService + + result = PluginService.get_debugging_key("tenant-1") + + assert result == "debug-key" + debugging_client_cls.return_value.get_debugging_key.assert_called_once_with("tenant-1") + invalidate_cache.assert_not_called() + + def test_list_model_category_invalidates_when_remote_model_plugin_is_missing_from_provider_cache(self) -> None: + """Remote model plugins are daemon-registered, so category reads repair a stale provider cache.""" + remote_plugin = _build_remote_model_plugin() + remote_plugin_marker = "langgenius/debug-model:langgenius/debug-model:1.0.0" + plugins = SimpleNamespace(list=[remote_plugin], has_more=False) + + with ( + patch(f"{MODULE}.PluginInstaller") as installer_cls, + patch( + f"{MODULE}.PluginService._load_cached_remote_model_plugin_marker", + return_value=remote_plugin_marker, + ), + patch( + f"{MODULE}.PluginService._load_cached_plugin_model_provider_plugin_ids", + return_value={"langgenius/openai"}, + ), + patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache, + patch(f"{MODULE}.PluginService._store_cached_remote_model_plugin_marker") as store_marker, + ): + installer_cls.return_value.list_plugins_by_category.return_value = plugins + + from core.plugin.plugin_service import PluginService + + result = PluginService.list_by_category("tenant-1", PluginCategory.Model, 1, 100) + + assert result is plugins + installer_cls.return_value.list_plugins_by_category.assert_called_once_with( + "tenant-1", PluginCategory.Model, 1, 100 + ) + invalidate_cache.assert_called_once_with("tenant-1") + store_marker.assert_called_once_with("tenant-1", remote_plugin_marker) + + def test_list_model_category_invalidates_when_remote_model_plugin_identity_changes(self) -> None: + """A debug model plugin can share plugin_id with an installed plugin, so identity changes bust cache too.""" + remote_plugin = _build_remote_model_plugin( + plugin_id="langgenius/openai", + plugin_unique_identifier="langgenius/openai:debug", + ) + remote_plugin_marker = "langgenius/openai:langgenius/openai:debug" + plugins = SimpleNamespace(list=[remote_plugin], has_more=False) + + with ( + patch(f"{MODULE}.PluginInstaller") as installer_cls, + patch( + f"{MODULE}.PluginService._load_cached_remote_model_plugin_marker", + return_value="langgenius/openai:langgenius/openai:1.0.0", + ), + patch( + f"{MODULE}.PluginService._load_cached_plugin_model_provider_plugin_ids", + return_value={"langgenius/openai"}, + ) as load_cached_provider_plugin_ids, + patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache, + patch(f"{MODULE}.PluginService._store_cached_remote_model_plugin_marker") as store_marker, + ): + installer_cls.return_value.list_plugins_by_category.return_value = plugins + + from core.plugin.plugin_service import PluginService + + result = PluginService.list_by_category("tenant-1", PluginCategory.Model, 1, 100) + + assert result is plugins + invalidate_cache.assert_called_once_with("tenant-1") + load_cached_provider_plugin_ids.assert_not_called() + store_marker.assert_called_once_with("tenant-1", remote_plugin_marker) + + def test_list_model_category_keeps_provider_cache_when_remote_model_plugin_is_already_cached(self) -> None: + """A connected remote model plugin should not force provider cache churn once represented.""" + remote_plugin = _build_remote_model_plugin() + remote_plugin_marker = "langgenius/debug-model:langgenius/debug-model:1.0.0" + plugins = SimpleNamespace(list=[remote_plugin], has_more=False) + + with ( + patch(f"{MODULE}.PluginInstaller") as installer_cls, + patch( + f"{MODULE}.PluginService._load_cached_remote_model_plugin_marker", + return_value=remote_plugin_marker, + ), + patch( + f"{MODULE}.PluginService._load_cached_plugin_model_provider_plugin_ids", + return_value={"langgenius/debug-model"}, + ), + patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache, + patch(f"{MODULE}.PluginService._store_cached_remote_model_plugin_marker") as store_marker, + ): + installer_cls.return_value.list_plugins_by_category.return_value = plugins + + from core.plugin.plugin_service import PluginService + + result = PluginService.list_by_category("tenant-1", PluginCategory.Model, 1, 100) + + assert result is plugins + invalidate_cache.assert_not_called() + store_marker.assert_called_once_with("tenant-1", remote_plugin_marker) + + def test_list_model_category_invalidates_when_remote_model_plugin_disconnects(self) -> None: + """The current model category result clears provider cache when the previous debug model disappears.""" + installed_plugin = SimpleNamespace( + plugin_id="langgenius/openai", + plugin_unique_identifier="langgenius/openai:1.0.0", + source=PluginInstallationSource.Marketplace, + ) + plugins = SimpleNamespace(list=[installed_plugin], has_more=True) + + with ( + patch(f"{MODULE}.PluginInstaller") as installer_cls, + patch( + f"{MODULE}.PluginService._load_cached_remote_model_plugin_marker", + return_value="langgenius/debug-model:langgenius/debug-model:1.0.0", + ), + patch(f"{MODULE}.PluginService.invalidate_plugin_model_providers_cache") as invalidate_cache, + patch(f"{MODULE}.PluginService._store_cached_remote_model_plugin_marker") as store_marker, + ): + installer_cls.return_value.list_plugins_by_category.return_value = plugins + + from core.plugin.plugin_service import PluginService + + result = PluginService.list_by_category("tenant-1", PluginCategory.Model, 1, 100) + + assert result is plugins + invalidate_cache.assert_called_once_with("tenant-1") + store_marker.assert_called_once_with("tenant-1", None) + def test_fetch_install_task_invalidates_model_provider_cache_when_finished(self) -> None: """Finished plugin install tasks invalidate tenant provider cache.""" task = _build_install_task(status=PluginInstallTaskStatus.Success) diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py index 5bc41fdb5bd..a7bb8cfeed6 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py @@ -22,8 +22,9 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None: }, ) retrieval = BuiltInPipelineTemplateRetrieval() + session = mocker.Mock() - templates = retrieval.get_pipeline_templates(mocker.Mock(), "en-US") + templates = retrieval.get_pipeline_templates("en-US", session=session) assert templates == {"pipeline_templates": [{"id": "tpl-1"}]} @@ -39,8 +40,9 @@ def test_get_pipeline_template_detail(mocker: MockerFixture) -> None: }, ) retrieval = BuiltInPipelineTemplateRetrieval() + session = mocker.Mock() - detail = retrieval.get_pipeline_template_detail(mocker.Mock(), "tpl-1") + detail = retrieval.get_pipeline_template_detail("tpl-1", session=session) assert detail == {"id": "tpl-1", "name": "Template 1"} @@ -52,8 +54,9 @@ def test_get_pipeline_templates_missing_language_returns_empty_dict(mocker: Mock return_value={"pipeline_templates": {}}, ) retrieval = BuiltInPipelineTemplateRetrieval() + session = mocker.Mock() - result = retrieval.get_pipeline_templates(mocker.Mock(), "fr-FR") + result = retrieval.get_pipeline_templates("fr-FR", session=session) assert result == {} @@ -65,8 +68,9 @@ def test_get_pipeline_template_detail_returns_none_for_unknown_id(mocker: Mocker return_value={"pipeline_templates": {"tpl-1": {"id": "tpl-1"}}}, ) retrieval = BuiltInPipelineTemplateRetrieval() + session = mocker.Mock() - result = retrieval.get_pipeline_template_detail(mocker.Mock(), "nonexistent-id") + result = retrieval.get_pipeline_template_detail("nonexistent-id", session=session) assert result is None diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py index b3ef79961d3..b3befeb41fd 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py @@ -21,7 +21,7 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None: session_mock.scalars.return_value = scalars_mock retrieval = CustomizedPipelineTemplateRetrieval() - result = retrieval.get_pipeline_templates(session_mock, "en-US", "tenant-id") + result = retrieval.get_pipeline_templates("en-US", "tenant-id", session=session_mock) assert retrieval.get_type() == PipelineTemplateType.CUSTOMIZED assert result == { @@ -51,7 +51,7 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N ) retrieval = CustomizedPipelineTemplateRetrieval() - detail = retrieval.get_pipeline_template_detail(session_mock, "tpl-1") + detail = retrieval.get_pipeline_template_detail("tpl-1", session=session_mock) assert detail == { "id": "tpl-1", @@ -70,6 +70,6 @@ def test_get_pipeline_template_detail_returns_none_when_not_found(mocker: Mocker session_mock.get.return_value = None retrieval = CustomizedPipelineTemplateRetrieval() - result = retrieval.get_pipeline_template_detail(session_mock, "missing") + result = retrieval.get_pipeline_template_detail("missing", session=session_mock) assert result is None diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py index cae79175b1c..48ae26ce3aa 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py @@ -23,7 +23,7 @@ def test_get_pipeline_templates(mocker: MockerFixture) -> None: session_mock.scalars.return_value = scalars_mock retrieval = DatabasePipelineTemplateRetrieval() - result = retrieval.get_pipeline_templates(session_mock, "en-US") + result = retrieval.get_pipeline_templates("en-US", session=session_mock) assert retrieval.get_type() == PipelineTemplateType.DATABASE assert result == { @@ -54,7 +54,7 @@ def test_get_pipeline_template_detail_returns_detail(mocker: MockerFixture) -> N ) retrieval = DatabasePipelineTemplateRetrieval() - detail = retrieval.get_pipeline_template_detail(session_mock, "tpl-1") + detail = retrieval.get_pipeline_template_detail("tpl-1", session=session_mock) assert detail == { "id": "tpl-1", @@ -72,6 +72,6 @@ def test_get_pipeline_template_detail_returns_none_when_not_found(mocker: Mocker session_mock.get.return_value = None retrieval = DatabasePipelineTemplateRetrieval() - result = retrieval.get_pipeline_template_detail(session_mock, "missing") + result = retrieval.get_pipeline_template_detail("missing", session=session_mock) assert result is None diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py index c8af1869732..17cd5db7ab3 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_pipeline_template_base.py @@ -4,11 +4,11 @@ from services.rag_pipeline.pipeline_template.pipeline_template_base import Pipel class DummyRetrieval(PipelineTemplateRetrievalBase): - def get_pipeline_templates(self, session: Mock, language: str, current_tenant_id: str | None = None) -> dict: - del session, current_tenant_id + def get_pipeline_templates(self, language: str, *, session) -> dict: + del session return {"language": language} - def get_pipeline_template_detail(self, session: Mock, template_id: str) -> dict | None: + def get_pipeline_template_detail(self, template_id: str, *, session) -> dict | None: del session return {"id": template_id} @@ -20,6 +20,6 @@ def test_pipeline_template_retrieval_base_concrete_implementation() -> None: retrieval = DummyRetrieval() session = Mock() - assert retrieval.get_pipeline_templates(session, "en-US") == {"language": "en-US"} - assert retrieval.get_pipeline_template_detail(session, "tpl-1") == {"id": "tpl-1"} + assert retrieval.get_pipeline_templates("en-US", session=session) == {"language": "en-US"} + assert retrieval.get_pipeline_template_detail("tpl-1", session=session) == {"id": "tpl-1"} assert retrieval.get_type() == "dummy" diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py index 8f55b4b1c2f..78e46d272c2 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py @@ -20,12 +20,12 @@ def test_get_pipeline_templates_fallbacks_to_database_on_error(mocker: MockerFix retrieval = RemotePipelineTemplateRetrieval() session = mocker.Mock() - result = retrieval.get_pipeline_templates(session, "en-US") + result = retrieval.get_pipeline_templates("en-US", session=session) assert retrieval.get_type() == PipelineTemplateType.REMOTE assert result == {"pipeline_templates": [{"id": "db-1"}]} fetch_mock.assert_called_once_with("en-US") - fallback_mock.assert_called_once_with(session, "en-US") + fallback_mock.assert_called_once_with("en-US", session=session) def test_get_pipeline_template_detail_fallbacks_to_database_on_error(mocker: MockerFixture) -> None: @@ -42,11 +42,11 @@ def test_get_pipeline_template_detail_fallbacks_to_database_on_error(mocker: Moc retrieval = RemotePipelineTemplateRetrieval() session = mocker.Mock() - result = retrieval.get_pipeline_template_detail(session, "tpl-1") + result = retrieval.get_pipeline_template_detail("tpl-1", session=session) assert result == {"id": "db-1"} fetch_mock.assert_called_once_with("tpl-1") - fallback_mock.assert_called_once_with(session, "tpl-1") + fallback_mock.assert_called_once_with("tpl-1", session=session) def test_fetch_pipeline_templates_from_dify_official(mocker: MockerFixture) -> None: diff --git a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py index 0ae2ba97f1a..c8992585653 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py @@ -1,30 +1,15 @@ -from collections.abc import Iterator from types import SimpleNamespace from typing import cast -from uuid import uuid4 import pytest from pytest_mock import MockerFixture -from sqlalchemy import create_engine, func, select -from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom -from models.dataset import Document, Pipeline -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.dataset import Pipeline from models.model import Account, App, EndUser from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService -@pytest.fixture -def document_session() -> Iterator[Session]: - engine = create_engine("sqlite:///:memory:") - Document.__table__.create(engine) - session_factory = sessionmaker(bind=engine, expire_on_commit=False) - with session_factory() as session: - yield session - engine.dispose() - - def test_get_max_active_requests_uses_smallest_non_zero_limit(mocker: MockerFixture) -> None: mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_DEFAULT_ACTIVE_REQUESTS", 5) mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_MAX_ACTIVE_REQUESTS", 3) @@ -62,12 +47,13 @@ def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_err rag_pipeline_service.get_published_workflow.return_value = workflow pipeline = cast(Pipeline, SimpleNamespace(id="pipeline-1")) + session = mocker.Mock() if expected_error: with pytest.raises(ValueError, match=expected_error): - PipelineGenerateService._get_workflow(pipeline, invoke_from) + PipelineGenerateService._get_workflow(pipeline, invoke_from, session) else: - result = PipelineGenerateService._get_workflow(pipeline, invoke_from) + result = PipelineGenerateService._get_workflow(pipeline, invoke_from, session) assert result == workflow @@ -75,10 +61,10 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke pipeline = cast(Pipeline, SimpleNamespace(id="pipeline-1")) user = cast(Account | EndUser, SimpleNamespace(id="user-1")) args = {"original_document_id": "doc-1", "query": "hello"} + session_mock = mocker.Mock() mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status") - session = mocker.Mock() generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") generator_instance = generator_cls.return_value @@ -86,49 +72,39 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke generator_cls.convert_to_event_stream.return_value = "stream-events" result = PipelineGenerateService.generate( - session=session, pipeline=pipeline, user=user, args=args, invoke_from=InvokeFrom.WEB_APP, streaming=True, + session=session_mock, ) assert result == "stream-events" - update_status_mock.assert_called_once_with("doc-1", session) + update_status_mock.assert_called_once_with("doc-1", session=session_mock) -def test_update_document_status_updates_existing_document(document_session: Session) -> None: - session = document_session - document_id = str(uuid4()) - document = Document( - id=document_id, - tenant_id=str(uuid4()), - dataset_id=str(uuid4()), - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name="Doc", - created_from=DocumentCreatedFrom.WEB, - created_by=str(uuid4()), - indexing_status=IndexingStatus.COMPLETED, - ) - session.add(document) - session.commit() +def test_update_document_status_updates_existing_document(mocker: MockerFixture) -> None: + document = SimpleNamespace(indexing_status="completed") - PipelineGenerateService.update_document_status(document_id, session) + session_mock = mocker.Mock() + session_mock.get.return_value = document + add_mock = session_mock.add - updated_document = session.get(Document, document_id) - assert updated_document is not None - assert updated_document.indexing_status == IndexingStatus.WAITING + PipelineGenerateService.update_document_status("doc-1", session=session_mock) + + assert document.indexing_status == "waiting" + add_mock.assert_called_once_with(document) -def test_update_document_status_skips_when_document_missing(document_session: Session) -> None: - session = document_session +def test_update_document_status_skips_when_document_missing(mocker: MockerFixture) -> None: + session_mock = mocker.Mock() + session_mock.get.return_value = None + add_mock = session_mock.add - PipelineGenerateService.update_document_status(str(uuid4()), session) + PipelineGenerateService.update_document_status("missing", session=session_mock) - assert session.scalar(select(func.count()).select_from(Document)) == 0 + add_mock.assert_not_called() # --- generate_single_iteration --- @@ -144,8 +120,9 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None: pipeline = cast(Pipeline, SimpleNamespace(id="p1")) user = cast(Account, SimpleNamespace(id="u1")) + session = mocker.Mock() - result = PipelineGenerateService.generate_single_iteration(pipeline, user, "node-1", {"key": "val"}) + result = PipelineGenerateService.generate_single_iteration(pipeline, user, "node-1", {"key": "val"}, session) assert result == "stream-iter" generator_instance.single_iteration_generate.assert_called_once() @@ -164,8 +141,9 @@ def test_generate_single_loop_delegates(mocker: MockerFixture) -> None: pipeline = cast(Pipeline, SimpleNamespace(id="p1")) user = cast(Account, SimpleNamespace(id="u1")) + session = mocker.Mock() - result = PipelineGenerateService.generate_single_loop(pipeline, user, "node-1", {"key": "val"}) + result = PipelineGenerateService.generate_single_loop(pipeline, user, "node-1", {"key": "val"}, session) assert result == "stream-loop" generator_instance.single_loop_generate.assert_called_once() diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py index 55ae8144d55..5cdb2afd093 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py @@ -87,11 +87,11 @@ def test_check_dependencies_returns_empty_when_no_redis_data(mocker: MockerFixtu def test_check_dependencies_returns_leaked_deps_from_redis(mocker: MockerFixture) -> None: - from core.plugin.entities.plugin import PluginDependency + from core.plugin.entities.plugin import PluginDependency, PluginDependencyType from services.rag_pipeline.rag_pipeline_dsl_service import CheckDependenciesPendingData dep = PluginDependency( - type=PluginDependency.Type.Marketplace, + type=PluginDependencyType.Marketplace, value=PluginDependency.Marketplace(marketplace_plugin_unique_identifier="test/plugin:0.1.0"), ) pending_data = CheckDependenciesPendingData( @@ -633,6 +633,19 @@ def test_import_rag_pipeline_yaml_content_requires_content() -> None: assert "yaml_content is required" in result.error +def test_import_rag_pipeline_rejects_oversized_yaml_content_before_parsing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("services.rag_pipeline.rag_pipeline_dsl_service.DSL_MAX_SIZE", 3) + service = RagPipelineDslService(session=Mock()) + account = Mock(current_tenant_id="t1") + + result = service.import_rag_pipeline(account=account, import_mode="yaml-content", yaml_content="你你") + + assert result.status == ImportStatus.FAILED + assert result.error == "File size exceeds the limit of 10MB" + + def test_import_rag_pipeline_yaml_content_requires_mapping() -> None: service = RagPipelineDslService(session=Mock()) account = Mock(current_tenant_id="t1") @@ -643,6 +656,19 @@ def test_import_rag_pipeline_yaml_content_requires_mapping() -> None: assert "content must be a mapping" in result.error +def test_import_rag_pipeline_rejects_oversized_yaml_content_by_bytes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("services.rag_pipeline.rag_pipeline_dsl_service.DSL_MAX_SIZE", 1) + service = RagPipelineDslService(session=Mock()) + account = Mock(current_tenant_id="t1") + + result = service.import_rag_pipeline(account=account, import_mode="yaml-content", yaml_content="é") + + assert result.status == ImportStatus.FAILED + assert "10MB" in result.error + + def test_confirm_import_returns_failed_when_pending_data_is_invalid_type(mocker: MockerFixture) -> None: mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.redis_client.get", return_value=object()) service = RagPipelineDslService(session=Mock()) @@ -901,6 +927,46 @@ def test_import_rag_pipeline_url_size_exceeds_limit(mocker: MockerFixture) -> No assert "10MB" in result.error +def test_import_rag_pipeline_rejects_oversized_yaml_url_bytes_before_decode( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("services.rag_pipeline.rag_pipeline_dsl_service.DSL_MAX_SIZE", 1) + response = Mock() + response.raise_for_status.return_value = None + response.content = b"\xff\xff" + mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.remote_fetcher.make_request", return_value=response) + service = RagPipelineDslService(session=Mock()) + account = Mock(current_tenant_id="t1") + + result = service.import_rag_pipeline( + account=account, + import_mode="yaml-url", + yaml_url="https://example.com/pipeline.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert "10MB" in result.error + + +def test_import_rag_pipeline_returns_decode_error_for_invalid_yaml_url_bytes(mocker: MockerFixture) -> None: + response = Mock() + response.raise_for_status.return_value = None + response.content = b"\xff" + mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.remote_fetcher.make_request", return_value=response) + service = RagPipelineDslService(session=Mock()) + account = Mock(current_tenant_id="t1") + + result = service.import_rag_pipeline( + account=account, + import_mode="yaml-url", + yaml_url="https://example.com/pipeline.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert "utf-8" in result.error + + def test_import_rag_pipeline_fails_when_rag_pipeline_data_missing() -> None: service = RagPipelineDslService(session=Mock()) account = Mock(current_tenant_id="t1") @@ -1371,7 +1437,7 @@ def test_confirm_import_fails_when_no_knowledge_index_node(mocker: MockerFixture def test_create_or_update_pipeline_saves_dependencies_to_redis(mocker: MockerFixture) -> None: - from core.plugin.entities.plugin import PluginDependency + from core.plugin.entities.plugin import PluginDependency, PluginDependencyType session = cast(MagicMock, Mock()) service = RagPipelineDslService(session=cast(Session, session)) @@ -1386,7 +1452,7 @@ def test_create_or_update_pipeline_saves_dependencies_to_redis(mocker: MockerFix session.scalar.return_value = None setex = mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.redis_client.setex") dependency = PluginDependency( - type=PluginDependency.Type.Marketplace, + type=PluginDependencyType.Marketplace, value=PluginDependency.Marketplace(marketplace_plugin_unique_identifier="langgenius/example:0.1.0"), ) diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py index 37141c97c83..0d74b3abf9c 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py @@ -49,9 +49,8 @@ def rag_pipeline_service(mocker: MockerFixture) -> RagPipelineServiceTestContext ) session = mocker.Mock() session_maker = _make_mock_session_maker(mocker, session) - mocker.patch("services.rag_pipeline.rag_pipeline.session_factory.get_session_maker", return_value=session_maker) mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=mocker.Mock())) - service = RagPipelineService(session_maker=session_maker) + service = RagPipelineService(session=session, session_maker=session_maker) return RagPipelineServiceTestContext(service=service, session=session, session_maker=session_maker) @@ -156,10 +155,10 @@ def test_get_pipeline_templates_fallbacks_to_builtin_for_non_english_empty_resul builtin_retrieval.fetch_pipeline_templates_from_builtin.return_value = {"pipeline_templates": [{"id": "builtin-1"}]} factory_mock.get_built_in_pipeline_template_retrieval.return_value = builtin_retrieval - result = RagPipelineService.get_pipeline_templates(session, type="built-in", language="ja-JP") + result = RagPipelineService.get_pipeline_templates(type="built-in", language="ja-JP", session=session) assert result == {"pipeline_templates": [{"id": "builtin-1"}]} - remote_retrieval.get_pipeline_templates.assert_called_once_with(session, "ja-JP", None) + remote_retrieval.get_pipeline_templates.assert_called_once_with("ja-JP", None, session=session) builtin_retrieval.fetch_pipeline_templates_from_builtin.assert_called_once_with("en-US") @@ -171,11 +170,11 @@ def test_get_pipeline_templates_customized_mode_uses_customized_factory(mocker: factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval - result = RagPipelineService.get_pipeline_templates(session, type="customized", language="en-US") + result = RagPipelineService.get_pipeline_templates(type="customized", language="en-US", session=session) assert result == {"pipeline_templates": [{"id": "custom-1"}]} factory_mock.get_pipeline_template_factory.assert_called_with("customized") - retrieval.get_pipeline_templates.assert_called_once_with(session, "en-US", None) + retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) @pytest.mark.parametrize("template_type", ["built-in", "customized"]) @@ -188,12 +187,12 @@ def test_get_pipeline_template_detail_uses_expected_mode(mocker: MockerFixture, factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval - result = RagPipelineService.get_pipeline_template_detail(session, "tpl-1", type=template_type) + result = RagPipelineService.get_pipeline_template_detail("tpl-1", type=template_type, session=session) assert result == {"id": "tpl-1"} expected_mode = "remote" if template_type == "built-in" else "customized" factory_mock.get_pipeline_template_factory.assert_called_with(expected_mode) - retrieval.get_pipeline_template_detail.assert_called_once_with(session, "tpl-1") + retrieval.get_pipeline_template_detail.assert_called_once_with("tpl-1", session=session) def test_get_published_workflow_returns_none_when_pipeline_has_no_workflow_id( @@ -845,13 +844,14 @@ def test_publish_customized_pipeline_template_success( # 2. Run test args = {"name": "New Template", "description": "Desc", "icon_info": {"icon": "star"}, "tags": ["tag1"]} - rag_pipeline_service.service.publish_customized_pipeline_template("p1", args, account, "t1") + rag_pipeline_service.service.publish_customized_pipeline_template("p1", args, account, "t1", session=session) # 3. Assertions # Verify a new template was added to session or similar? # Since we can't easily check the session inside the context manager with Mock, # we just check that no error was raised and DSL was exported. - mock_dsl_service.export_rag_pipeline_dsl.assert_called_once() + pipeline.retrieve_dataset.assert_called_once_with(session=session) + mock_dsl_service.export_rag_pipeline_dsl.assert_called_once_with(pipeline=pipeline, include_secret=True) # --- get_datasource_plugins --- @@ -863,7 +863,7 @@ def test_get_datasource_plugins_success( # 1. Setup mocks dataset = _make_dataset() - pipeline = _make_pipeline() + pipeline = _make_pipeline(workflow_id="wf-1") workflow = mocker.Mock() workflow.graph_dict = { @@ -996,7 +996,7 @@ def test_set_datasource_variables_success( # --- Utility Methods --- -def test_get_draft_workflow_success(mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext) -> None: +def test_get_draft_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: # 1. Setup mocks pipeline = _make_pipeline() @@ -1011,9 +1011,7 @@ def test_get_draft_workflow_success(mocker: MockerFixture, rag_pipeline_service: assert result == workflow -def test_get_published_workflow_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: +def test_get_published_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: # 1. Setup mocks pipeline = _make_pipeline(workflow_id="wf-pub") @@ -1406,10 +1404,7 @@ def test_get_node_last_run_delegates_to_repository( ) -> None: repo = mocker.Mock() repo.get_node_last_execution.return_value = "node-exec" - mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository", - return_value=repo, - ) + rag_pipeline_service.service._node_execution_service_repo = repo pipeline = _make_pipeline() workflow = _make_workflow(workflow_id="wf1") @@ -1785,21 +1780,25 @@ def test_run_datasource_node_preview_raises_for_unsupported_provider( def test_publish_customized_pipeline_template_raises_for_missing_pipeline( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: - rag_pipeline_service.session.get.return_value = None + session = mocker.Mock() + session.get.return_value = None with pytest.raises(ValueError, match="Pipeline not found"): - rag_pipeline_service.service.publish_customized_pipeline_template("p1", {}, _make_account(), "t1") + rag_pipeline_service.service.publish_customized_pipeline_template( + "p1", {}, _make_account(), "t1", session=session + ) def test_publish_customized_pipeline_template_raises_for_missing_workflow_id( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: pipeline = _make_pipeline(workflow_id=None) - rag_pipeline_service.session.get.return_value = pipeline + session = mocker.Mock() + session.get.return_value = pipeline with pytest.raises(ValueError, match="Pipeline workflow not found"): rag_pipeline_service.service.publish_customized_pipeline_template( - "p1", {"name": "template-name"}, _make_account(), "t1" + "p1", {"name": "template-name"}, _make_account(), "t1", session=session ) @@ -1824,10 +1823,8 @@ def test_get_pipeline_raises_when_pipeline_missing( def test_init_uses_default_sessionmaker_when_none(mocker: MockerFixture) -> None: default_session_maker = mocker.Mock() - mocker.patch( - "services.rag_pipeline.rag_pipeline.session_factory.get_session_maker", - return_value=default_session_maker, - ) + mocker.patch("services.rag_pipeline.rag_pipeline.sessionmaker", return_value=default_session_maker) + mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=mocker.Mock(), session=mocker.Mock())) create_exec_repo = mocker.patch( "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository" ) @@ -1835,7 +1832,7 @@ def test_init_uses_default_sessionmaker_when_none(mocker: MockerFixture) -> None "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_run_repository" ) - RagPipelineService(session_maker=None) + RagPipelineService(session=mocker.Mock(), session_maker=None) create_exec_repo.assert_called_once_with(default_session_maker) create_run_repo.assert_called_once_with(default_session_maker) @@ -1849,11 +1846,12 @@ def test_get_pipeline_templates_builtin_en_us_no_fallback(mocker: MockerFixture) factory = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") factory.get_pipeline_template_factory.return_value.return_value = retrieval builtin = factory.get_built_in_pipeline_template_retrieval.return_value + session = mocker.Mock() - result = RagPipelineService.get_pipeline_templates(session, type="built-in", language="en-US") + result = RagPipelineService.get_pipeline_templates(type="built-in", language="en-US", session=session) assert result == {"pipeline_templates": []} - retrieval.get_pipeline_templates.assert_called_once_with(session, "en-US", None) + retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) builtin.fetch_pipeline_templates_from_builtin.assert_not_called() @@ -1861,14 +1859,14 @@ def test_update_customized_pipeline_template_commits_when_name_empty(mocker: Moc template = _make_customized_template() session = mocker.Mock() session.scalar.return_value = template - session_maker = _make_mock_session_maker(mocker, session) - mocker.patch("services.rag_pipeline.rag_pipeline.session_factory.get_session_maker", return_value=session_maker) info = PipelineTemplateInfoEntity(name="", description="updated", icon_info=IconInfo(icon="i")) - result = RagPipelineService.update_customized_pipeline_template("tpl-1", info, _make_account(), "t1") + result = RagPipelineService.update_customized_pipeline_template( + "tpl-1", info, _make_account(), "t1", session=session + ) assert result.description == "updated" - session_maker.begin.assert_called_once() + session.commit.assert_called_once() def test_get_all_published_workflow_without_filters_has_no_more( @@ -2102,10 +2100,13 @@ def test_publish_customized_pipeline_template_raises_when_workflow_missing( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: pipeline = _make_pipeline(workflow_id="wf-1") - rag_pipeline_service.session.get.side_effect = [pipeline, None] + session = mocker.Mock() + session.get.side_effect = [pipeline, None] with pytest.raises(ValueError, match="Workflow not found"): - rag_pipeline_service.service.publish_customized_pipeline_template("p1", {}, _make_account(), "t1") + rag_pipeline_service.service.publish_customized_pipeline_template( + "p1", {}, _make_account(), "t1", session=session + ) def test_publish_customized_pipeline_template_raises_when_dataset_missing( @@ -2113,11 +2114,14 @@ def test_publish_customized_pipeline_template_raises_when_dataset_missing( ) -> None: pipeline = _make_pipeline(workflow_id="wf-1") workflow = _make_workflow(workflow_id="wf-1") + session = rag_pipeline_service.session + session.get.side_effect = [pipeline, workflow] pipeline.retrieve_dataset = mocker.Mock(return_value=None) - rag_pipeline_service.session.get.side_effect = [pipeline, workflow] with pytest.raises(ValueError, match="Dataset not found"): - rag_pipeline_service.service.publish_customized_pipeline_template("p1", {}, _make_account(), "t1") + rag_pipeline_service.service.publish_customized_pipeline_template( + "p1", {}, _make_account(), "t1", session=session + ) def test_get_recommended_plugins_skips_manifest_when_missing( @@ -2165,7 +2169,7 @@ def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: dataset = _make_dataset() - pipeline = _make_pipeline() + pipeline = _make_pipeline(workflow_id="wf-1") workflow = SimpleNamespace( graph_dict={"nodes": [{"id": "n1", "data": {"type": "start"}}]}, rag_pipeline_variables=[] ) @@ -2360,7 +2364,7 @@ def test_get_datasource_plugins_extracts_user_inputs_and_credentials( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: dataset = _make_dataset() - pipeline = _make_pipeline() + pipeline = _make_pipeline(workflow_id="wf-1") workflow = SimpleNamespace( graph_dict={ "nodes": [ diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py index 21df81d3ea8..4ee1a5831a0 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py @@ -1,31 +1,16 @@ import logging -from collections.abc import Iterator from datetime import UTC, datetime from types import SimpleNamespace from typing import cast import pytest from pytest_mock import MockerFixture -from sqlalchemy import create_engine, select -from sqlalchemy.orm import Session, sessionmaker -from models.dataset import Dataset, Pipeline -from models.enums import DatasetRuntimeMode +from models.dataset import Dataset from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration from services.rag_pipeline.rag_pipeline_transform_service import RagPipelineTransformService -@pytest.fixture -def pipeline_session() -> Iterator[Session]: - engine = create_engine("sqlite:///:memory:") - Dataset.__table__.create(engine) - Pipeline.__table__.create(engine) - session_factory = sessionmaker(bind=engine, expire_on_commit=False) - with session_factory() as session: - yield session - engine.dispose() - - @pytest.mark.parametrize( ("doc_form", "datasource_type", "indexing_technique"), [ @@ -89,7 +74,7 @@ def test_deal_dependencies_installs_missing_marketplace_plugins(mocker: MockerFi installer_cls.return_value.list_plugins.return_value = [SimpleNamespace(plugin_id="installed-plugin")] migration_cls = mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.PluginMigration") - migration_cls.return_value._fetch_plugin_unique_identifier.return_value = "missing-plugin:1.0.0" + migration_cls.return_value._fetch_latest_package_identifier.return_value = "missing-plugin:1.0.0" install_mock = mocker.patch( "services.rag_pipeline.rag_pipeline_transform_service.PluginService.install_from_marketplace_pkg" @@ -107,37 +92,47 @@ def test_deal_dependencies_installs_missing_marketplace_plugins(mocker: MockerFi install_mock.assert_called_once_with("tenant-1", ["missing-plugin:1.0.0"]) -def test_transform_to_empty_pipeline_updates_dataset_and_flushes( - mocker: MockerFixture, pipeline_session: Session -) -> None: +def test_transform_to_empty_pipeline_updates_dataset_and_commits(mocker: MockerFixture) -> None: service = RagPipelineTransformService() mocker.patch( "services.rag_pipeline.rag_pipeline_transform_service.current_user", SimpleNamespace(id="user-1"), ) - session = pipeline_session - dataset = Dataset( + class FakePipeline: + def __init__(self, **kwargs): + self.id = "pipeline-1" + self.tenant_id = kwargs["tenant_id"] + self.name = kwargs["name"] + self.description = kwargs["description"] + self.created_by = kwargs["created_by"] + + mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.Pipeline", FakePipeline) + session_mock = mocker.Mock() + add_mock = session_mock.add + flush_mock = session_mock.flush + commit_mock = session_mock.commit + + dataset = SimpleNamespace( + id="dataset-1", tenant_id="tenant-1", name="Dataset", description="desc", - created_by="user-1", + pipeline_id=None, + runtime_mode="general", + updated_by=None, + updated_at=None, ) - session.add(dataset) - session.commit() - flush_spy = mocker.spy(session, "flush") - commit_spy = mocker.spy(session, "commit") - result = service._transform_to_empty_pipeline(dataset, session) + result = service._transform_to_empty_pipeline(cast(Dataset, dataset), session=session_mock) - assert flush_spy.call_count == 2 - commit_spy.assert_not_called() - pipeline = session.scalar(select(Pipeline).where(Pipeline.id == dataset.pipeline_id)) - assert pipeline is not None - assert result == {"pipeline_id": pipeline.id, "dataset_id": dataset.id, "status": "success"} - assert dataset.pipeline_id == pipeline.id - assert dataset.runtime_mode == DatasetRuntimeMode.RAG_PIPELINE + assert result == {"pipeline_id": "pipeline-1", "dataset_id": "dataset-1", "status": "success"} + assert dataset.pipeline_id == "pipeline-1" + assert dataset.runtime_mode == "rag_pipeline" assert dataset.updated_by == "user-1" + add_mock.assert_called() + flush_mock.assert_called_once() + commit_mock.assert_called_once() # --- transform_dataset --- @@ -373,6 +368,7 @@ def test_transform_dataset_full_flow(mocker: MockerFixture) -> None: mocker.patch.object(service, "_deal_dependencies") mocker.patch.object(service, "_deal_document_data") + session_mock.commit = mocker.Mock() # Mock current_user to have the same tenant_id as dataset mock_current_user = SimpleNamespace(current_tenant_id="t1") @@ -386,8 +382,6 @@ def test_transform_dataset_full_flow(mocker: MockerFixture) -> None: assert result["pipeline_id"] == "p-new" assert dataset.runtime_mode == "rag_pipeline" assert dataset.chunk_structure == "text_model" - session_mock.flush.assert_called_once_with() - session_mock.commit.assert_not_called() def test_transform_dataset_raises_for_unsupported_doc_form_after_pipeline_create(mocker: MockerFixture) -> None: @@ -439,11 +433,12 @@ def test_transform_dataset_raises_when_transform_yaml_missing_workflow(mocker: M service.transform_dataset("d1", session_mock) -def test_create_pipeline_raises_when_workflow_data_missing(pipeline_session: Session) -> None: +def test_create_pipeline_raises_when_workflow_data_missing(mocker: MockerFixture) -> None: service = RagPipelineTransformService() + session = mocker.Mock() with pytest.raises(ValueError, match="Missing workflow data for rag pipeline"): - service._create_pipeline({"rag_pipeline": {"name": "N"}}, pipeline_session) + service._create_pipeline({"rag_pipeline": {"name": "N"}}, session=session) def test_deal_document_data_upload_file_with_existing_file(mocker: MockerFixture) -> None: @@ -518,7 +513,7 @@ def test_deal_dependencies_installs_when_enabled(mocker: MockerFixture) -> None: installer = mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.PluginInstaller").return_value installer.list_plugins.return_value = [] migration = mocker.patch("services.rag_pipeline.rag_pipeline_transform_service.PluginMigration").return_value - migration._fetch_plugin_unique_identifier.return_value = "langgenius/openai:1.0.0@abc" + migration._fetch_latest_package_identifier.return_value = "langgenius/openai:1.0.0@abc" install_call = mocker.patch( "services.rag_pipeline.rag_pipeline_transform_service.PluginService.install_from_marketplace_pkg" ) diff --git a/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py b/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py index c86aaf1db22..cafac0656d1 100644 --- a/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py +++ b/api/tests/unit_tests/services/recommend_app/test_buildin_retrieval.py @@ -39,7 +39,7 @@ class TestBuildInRecommendAppRetrieval: return_value={"apps": []}, ) as mock_fetch: retrieval = BuildInRecommendAppRetrieval() - result = retrieval.get_recommended_apps_and_categories("en-US") + result = retrieval.get_recommended_apps_and_categories("en-US", session=MagicMock()) mock_fetch.assert_called_once_with("en-US") assert result == {"apps": []} @@ -47,11 +47,12 @@ class TestBuildInRecommendAppRetrieval: def test_get_learn_dify_apps_delegates_to_database(self, mock_database_retrieval): expected = {"recommended_apps": [{"id": "learn-dify-app"}]} mock_database_retrieval.fetch_learn_dify_apps_from_db.return_value = expected + session = MagicMock() - result = BuildInRecommendAppRetrieval().get_learn_dify_apps("en-US") + result = BuildInRecommendAppRetrieval().get_learn_dify_apps("en-US", session=session) assert result == expected - mock_database_retrieval.fetch_learn_dify_apps_from_db.assert_called_once_with("en-US") + mock_database_retrieval.fetch_learn_dify_apps_from_db.assert_called_once_with("en-US", session=session) def test_get_recommend_app_detail_delegates(self): with patch.object( @@ -60,7 +61,7 @@ class TestBuildInRecommendAppRetrieval: return_value={"id": "app-1"}, ) as mock_fetch: retrieval = BuildInRecommendAppRetrieval() - result = retrieval.get_recommend_app_detail("app-1") + result = retrieval.get_recommend_app_detail("app-1", session=MagicMock()) mock_fetch.assert_called_once_with("app-1") assert result == {"id": "app-1"} diff --git a/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py b/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py index 55165deec25..9575aa9f52e 100644 --- a/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py +++ b/api/tests/unit_tests/services/recommend_app/test_remote_retrieval.py @@ -17,7 +17,7 @@ class TestRemoteRecommendAppRetrieval: return_value={"id": "app-1"}, ) def test_get_recommend_app_detail_success(self, mock_fetch): - result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1") + result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1", session=MagicMock()) assert result == {"id": "app-1"} mock_fetch.assert_called_once_with("app-1") @@ -32,7 +32,7 @@ class TestRemoteRecommendAppRetrieval: side_effect=ConnectionError("timeout"), ) def test_get_recommend_app_detail_falls_back_on_error(self, mock_fetch, mock_builtin): - result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1") + result = RemoteRecommendAppRetrieval().get_recommend_app_detail("app-1", session=MagicMock()) assert result == {"id": "fallback"} mock_builtin.assert_called_once_with("app-1") @@ -42,7 +42,7 @@ class TestRemoteRecommendAppRetrieval: return_value={"recommended_apps": [], "categories": []}, ) def test_get_recommended_apps_success(self, mock_fetch): - result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories("en-US") + result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories("en-US", session=MagicMock()) assert result == {"recommended_apps": [], "categories": []} @patch( @@ -56,7 +56,7 @@ class TestRemoteRecommendAppRetrieval: side_effect=ValueError("server error"), ) def test_get_recommended_apps_falls_back_on_error(self, mock_fetch, mock_builtin): - result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories("en-US") + result = RemoteRecommendAppRetrieval().get_recommended_apps_and_categories("en-US", session=MagicMock()) assert result == {"recommended_apps": [{"id": "builtin"}]} @patch.object( @@ -65,7 +65,7 @@ class TestRemoteRecommendAppRetrieval: return_value={"recommended_apps": [{"id": "learn-dify-app"}]}, ) def test_get_learn_dify_apps_success(self, mock_fetch): - result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US") + result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US", session=MagicMock()) assert result == {"recommended_apps": [{"id": "learn-dify-app"}]} mock_fetch.assert_called_once_with("en-US") @@ -80,10 +80,12 @@ class TestRemoteRecommendAppRetrieval: side_effect=ValueError("server error"), ) def test_get_learn_dify_apps_falls_back_to_database_on_error(self, mock_fetch, mock_database): - result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US") + session = MagicMock() + + result = RemoteRecommendAppRetrieval().get_learn_dify_apps("en-US", session=session) assert result == {"recommended_apps": [{"id": "db-fallback"}]} - mock_database.assert_called_once_with("en-US") + mock_database.assert_called_once_with("en-US", session=session) class TestFetchFromDifyOfficial: diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 214344e0e57..b73fa112003 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -4,6 +4,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from sqlalchemy.orm import Session from configs import dify_config from models.account import Account, AccountStatus, TenantAccountRole, TenantStatus @@ -680,7 +681,7 @@ class TestTenantService: mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = TenantAccountRole.ADMIN - role = TenantService.get_account_role_in_tenant(mock_session, "account-1", "tenant-1") + role = TenantService.get_account_role_in_tenant("account-1", "tenant-1", session=mock_session) assert role == TenantAccountRole.ADMIN @@ -689,7 +690,7 @@ class TestTenantService: mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = None - role = TenantService.get_account_role_in_tenant(mock_session, "account-1", "tenant-1") + role = TenantService.get_account_role_in_tenant("account-1", "tenant-1", session=mock_session) assert role is None @@ -698,7 +699,7 @@ class TestTenantService: without ever touching the session.""" mock_session = MagicMock() - assert TenantService.get_account_role_in_tenant(mock_session, None, "tenant-1") is None + assert TenantService.get_account_role_in_tenant(None, "tenant-1", session=mock_session) is None mock_session.execute.assert_not_called() def test_get_account_role_in_tenant_query_is_scoped(self): @@ -710,7 +711,7 @@ class TestTenantService: mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = TenantAccountRole.NORMAL - TenantService.get_account_role_in_tenant(mock_session, account_id, tenant_id) + TenantService.get_account_role_in_tenant(account_id, tenant_id, session=mock_session) stmt = mock_session.execute.call_args.args[0] compiled = str(stmt.compile(compile_kwargs={"literal_binds": True})) @@ -759,11 +760,7 @@ class TestTenantService: mock_tenant_instance.name = "Test User's Workspace" mock_tenant_class.return_value = mock_tenant_instance - # Mock the db import in CreditPoolService to avoid database connection - with patch("services.credit_pool_service.db") as mock_credit_pool_db: - mock_credit_pool_db.session.add = MagicMock() - mock_credit_pool_db.session.commit = MagicMock() - + with patch("services.credit_pool_service.CreditPoolService.create_default_pool"): # Execute test TenantService.create_owner_tenant_if_not_exist( mock_account, session=mock_db_dependencies["db"].session @@ -1051,6 +1048,7 @@ class TestTenantService: account_id="user-rbac", member_account_id="user-rbac", role_ids=["rbac-owner-id"], + session=mock_db_dependencies["db"].session, ) def test_admin_can_update_admin_member_role(self): @@ -1190,7 +1188,11 @@ class TestTenantService: with pytest.raises(NoPermissionError): TenantService.check_member_permission( - mock_tenant, mock_operator, mock_member, "remove", session=MagicMock() + mock_tenant, + mock_operator, + mock_member, + "remove", + session=mock_db_dependencies["db"].session, ) def test_rbac_member_can_remove_non_owner_member(self): @@ -1264,7 +1266,9 @@ class TestTenantService: ), patch("services.account_service.RBACService.Roles", mock_rbac_roles), ): - owner_account_id = AccountService.get_rbac_workspace_owner_account_id("tenant-1", "acct-1") + owner_account_id = AccountService.get_rbac_workspace_owner_account_id( + "tenant-1", "acct-1", session=MagicMock() + ) assert owner_account_id == "owner-account" call = mock_rbac_roles.members.call_args @@ -1911,7 +1915,9 @@ class TestRegisterService: is_setup=True, session=mock_db_dependencies["db"].session, ) - mock_lookup.assert_called_once_with(mock_db_dependencies["db"].session, "newuser@example.com") + mock_lookup.assert_called_once_with( + "newuser@example.com", session=mock_db_dependencies["db"].session + ) def test_invite_new_member_normalizes_new_account_email( self, mock_db_dependencies, mock_redis_dependencies, mock_task_dependencies @@ -1957,7 +1963,7 @@ class TestRegisterService: is_setup=True, session=mock_db_dependencies["db"].session, ) - mock_lookup.assert_called_once_with(mock_db_dependencies["db"].session, mixed_email) + mock_lookup.assert_called_once_with(mixed_email, session=mock_db_dependencies["db"].session) mock_check_permission.assert_called_once_with( mock_tenant, mock_inviter, @@ -2024,7 +2030,7 @@ class TestRegisterService: mock_tenant, mock_existing_account, "normal", requires_setup=True ) mock_task_dependencies.delay.assert_called_once() - mock_lookup.assert_called_once_with(mock_db_dependencies["db"].session, "existing@example.com") + mock_lookup.assert_called_once_with("existing@example.com", session=mock_db_dependencies["db"].session) def test_invite_existing_active_account_requires_acceptance_before_joining( self, mock_db_dependencies, mock_redis_dependencies, mock_task_dependencies @@ -2170,6 +2176,7 @@ class TestRegisterService: account_id=mock_inviter.id, member_account_id=mock_new_account.id, role_ids=["rbac-role-id-123"], + session=mock_db_dependencies["db"].session, ) def test_invite_new_member_rbac_enabled_existing_account( @@ -2219,6 +2226,7 @@ class TestRegisterService: account_id=mock_inviter.id, member_account_id=mock_existing_account.id, role_ids=["rbac-role-id-456"], + session=mock_db_dependencies["db"].session, ) def test_invite_new_member_rbac_enabled_existing_active_account_adds_role_before_signin_response( @@ -2267,6 +2275,7 @@ class TestRegisterService: account_id=mock_inviter.id, member_account_id=mock_existing_account.id, role_ids=["rbac-role-id-456"], + session=mock_db_dependencies["db"].session, ) mock_task_dependencies.delay.assert_not_called() @@ -2614,7 +2623,7 @@ class TestSessionInjectedGetters: sentinel_account = MagicMock(spec=Account) mock_session.get.return_value = sentinel_account - result = AccountService.get_account_by_id(mock_session, "user-123") + result = AccountService.get_account_by_id("user-123", session=mock_session) assert result is sentinel_account mock_session.get.assert_called_once_with(Account, "user-123") @@ -2624,21 +2633,21 @@ class TestSessionInjectedGetters: mock_session = MagicMock() mock_session.get.return_value = None - assert AccountService.get_account_by_id(mock_session, "missing") is None + assert AccountService.get_account_by_id("missing", session=mock_session) is None - def test_get_account_by_email_returns_scalar_or_none(self): + @pytest.mark.parametrize("sqlite_session", [(Account,)], indirect=True) + def test_get_account_by_email_returns_scalar_or_none(self, sqlite_session: Session): """Plain getter — case-sensitive equality (callers needing the case-insensitive existence check use :meth:`has_active_account_with_email`). """ - mock_session = MagicMock() - sentinel = MagicMock(spec=Account) - mock_session.execute.return_value.scalar_one_or_none.return_value = sentinel + account = Account(name="Alice", email="alice@example.com") + sqlite_session.add(account) + sqlite_session.commit() - assert AccountService.get_account_by_email(mock_session, "alice@example.com") is sentinel - - mock_session.execute.return_value.scalar_one_or_none.return_value = None - assert AccountService.get_account_by_email(mock_session, "ghost@example.com") is None + assert AccountService.get_account_by_email("alice@example.com", session=sqlite_session) == account + assert AccountService.get_account_by_email("ALICE@example.com", session=sqlite_session) is None + assert AccountService.get_account_by_email("ghost@example.com", session=sqlite_session) is None def test_account_belongs_to_tenant_short_circuits_on_falsy_account_id(self): """SSO bearers with no ``account_id`` (and any other falsy id) @@ -2647,22 +2656,22 @@ class TestSessionInjectedGetters: """ mock_session = MagicMock() - assert TenantService.account_belongs_to_tenant(mock_session, None, "tenant-1") is False - assert TenantService.account_belongs_to_tenant(mock_session, "", "tenant-1") is False + assert TenantService.account_belongs_to_tenant(None, "tenant-1", session=mock_session) is False + assert TenantService.account_belongs_to_tenant("", "tenant-1", session=mock_session) is False mock_session.execute.assert_not_called() def test_account_belongs_to_tenant_true_when_join_row_exists(self): mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = "join-id" - assert TenantService.account_belongs_to_tenant(mock_session, "user-1", "tenant-1") is True + assert TenantService.account_belongs_to_tenant("user-1", "tenant-1", session=mock_session) is True mock_session.execute.assert_called_once() def test_account_belongs_to_tenant_false_when_no_join(self): mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = None - assert TenantService.account_belongs_to_tenant(mock_session, "user-1", "tenant-1") is False + assert TenantService.account_belongs_to_tenant("user-1", "tenant-1", session=mock_session) is False def test_get_account_memberships_returns_join_tenant_pairs(self): """Returns whatever ``session.query(...).join(...).filter(...).all()`` @@ -2673,7 +2682,7 @@ class TestSessionInjectedGetters: rows = [(MagicMock(), MagicMock()), (MagicMock(), MagicMock())] mock_session.query.return_value.join.return_value.filter.return_value.all.return_value = rows - out = TenantService.get_account_memberships(mock_session, "user-123") + out = TenantService.get_account_memberships("user-123", session=mock_session) assert out == rows # No fall-through to the global db.session proxy. @@ -2687,7 +2696,7 @@ class TestSessionInjectedGetters: rows = [(MagicMock(), MagicMock())] mock_session.execute.return_value.all.return_value = rows - out = TenantService.get_workspaces_for_account(mock_session, "user-123") + out = TenantService.get_workspaces_for_account("user-123", session=mock_session) assert out == rows assert mock_session.execute.called @@ -2703,20 +2712,20 @@ class TestSessionInjectedGetters: sentinel = MagicMock(spec=Tenant) mock_session.get.return_value = sentinel - assert TenantService.get_tenant_by_id(mock_session, "tenant-1") is sentinel + assert TenantService.get_tenant_by_id("tenant-1", session=mock_session) is sentinel mock_session.get.assert_called_once_with(Tenant, "tenant-1") def test_get_tenant_by_id_returns_none_when_missing(self): mock_session = MagicMock() mock_session.get.return_value = None - assert TenantService.get_tenant_by_id(mock_session, "missing") is None + assert TenantService.get_tenant_by_id("missing", session=mock_session) is None def test_get_tenants_by_ids_short_circuits_on_empty_input(self): """Empty id list must not emit ``WHERE id IN ()``.""" mock_session = MagicMock() - assert TenantService.get_tenants_by_ids(mock_session, []) == [] + assert TenantService.get_tenants_by_ids([], session=mock_session) == [] mock_session.execute.assert_not_called() def test_get_tenants_by_ids_returns_scalars(self): @@ -2724,7 +2733,7 @@ class TestSessionInjectedGetters: tenants = [MagicMock(), MagicMock()] mock_session.execute.return_value.scalars.return_value.all.return_value = tenants - assert TenantService.get_tenants_by_ids(mock_session, ["t1", "t2"]) == tenants + assert TenantService.get_tenants_by_ids(["t1", "t2"], session=mock_session) == tenants mock_session.execute.assert_called_once() def test_get_tenant_name_returns_scalar_or_none(self): @@ -2735,10 +2744,10 @@ class TestSessionInjectedGetters: mock_session = MagicMock() mock_session.execute.return_value.scalar_one_or_none.return_value = "Acme Inc." - assert TenantService.get_tenant_name(mock_session, "tenant-1") == "Acme Inc." + assert TenantService.get_tenant_name("tenant-1", session=mock_session) == "Acme Inc." mock_session.execute.return_value.scalar_one_or_none.return_value = None - assert TenantService.get_tenant_name(mock_session, "missing") is None + assert TenantService.get_tenant_name("missing", session=mock_session) is None def test_find_workspace_for_account_returns_first_row_or_none(self): """Per-id read returns ``session.execute(...).first()`` directly; @@ -2749,7 +2758,7 @@ class TestSessionInjectedGetters: sentinel_row = (MagicMock(), MagicMock()) mock_session.execute.return_value.first.return_value = sentinel_row - assert TenantService.find_workspace_for_account(mock_session, "user-123", "ws-1") is sentinel_row + assert TenantService.find_workspace_for_account("user-123", "ws-1", session=mock_session) is sentinel_row mock_session.execute.return_value.first.return_value = None - assert TenantService.find_workspace_for_account(mock_session, "user-123", "ws-1") is None + assert TenantService.find_workspace_for_account("user-123", "ws-1", session=mock_session) is None diff --git a/api/tests/unit_tests/services/test_agent_app_sandbox_service.py b/api/tests/unit_tests/services/test_agent_app_sandbox_service.py index 978f8e8a24b..f3052a4ab49 100644 --- a/api/tests/unit_tests/services/test_agent_app_sandbox_service.py +++ b/api/tests/unit_tests/services/test_agent_app_sandbox_service.py @@ -18,8 +18,10 @@ from models.agent import AgentRuntimeSession, AgentRuntimeSessionOwnerType, Agen from services.agent_app_sandbox_service import ( AgentAppSandboxService, AgentSandboxInspectorError, + AgentSandboxUploadDownload, WorkflowAgentSandboxService, _default_client_factory, + _upload_download_response, ) @@ -82,7 +84,12 @@ class FakeClient: self.locators.append(locator) self.calls.append(("upload", path)) return SandboxUploadResponse( - path=path, file={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"} + path=path, + file={ + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + "download_url": "https://files.example/report.txt?token=1", + }, ) @@ -129,6 +136,34 @@ def test_agent_app_sandbox_service_builds_locator_and_proxies() -> None: assert store.scope == ("tenant-1", "app-1", "conv-1") +def test_agent_app_sandbox_service_upload_returns_download_url(monkeypatch: pytest.MonkeyPatch) -> None: + store = FakeStore(_stored_session()) + client = FakeClient() + captured: dict[str, object] = {} + + def fake_upload_download_response(*, tenant_id: str, file_mapping: dict[str, object]) -> AgentSandboxUploadDownload: + captured["tenant_id"] = tenant_id + captured["file_mapping"] = file_mapping + return AgentSandboxUploadDownload(url="https://files.example/report.txt?token=1&as_attachment=true") + + monkeypatch.setattr("services.agent_app_sandbox_service._upload_download_response", fake_upload_download_response) + service = AgentAppSandboxService(session_store=store, client_factory=lambda: client) # type: ignore[arg-type] + + result = service.upload_file(tenant_id="tenant-1", app_id="app-1", conversation_id="conv-1", path="report.txt") + + assert result.url == "https://files.example/report.txt?token=1&as_attachment=true" + assert client.calls == [("upload", "report.txt")] + assert store.scope == ("tenant-1", "app-1", "conv-1") + assert captured == { + "tenant_id": "tenant-1", + "file_mapping": { + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + "download_url": "https://files.example/report.txt?token=1", + }, + } + + def test_agent_app_sandbox_service_raises_when_no_active_session() -> None: service = AgentAppSandboxService(session_store=FakeStore(None), client_factory=lambda: FakeClient()) # type: ignore[arg-type] @@ -210,9 +245,19 @@ def _insert_workflow_session( @pytest.mark.usefixtures("_runtime_session_table") -def test_workflow_sandbox_service_resolves_locator_and_proxies() -> None: +def test_workflow_sandbox_service_resolves_locator_and_returns_download_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: _insert_workflow_session() client = FakeClient() + captured: dict[str, object] = {} + + def fake_upload_download_response(*, tenant_id: str, file_mapping: dict[str, object]) -> AgentSandboxUploadDownload: + captured["tenant_id"] = tenant_id + captured["file_mapping"] = file_mapping + return AgentSandboxUploadDownload(url="https://files.example/report.txt?token=1&as_attachment=true") + + monkeypatch.setattr("services.agent_app_sandbox_service._upload_download_response", fake_upload_download_response) service = WorkflowAgentSandboxService(client_factory=lambda: client) # type: ignore[arg-type] result = service.upload_file( @@ -222,10 +267,103 @@ def test_workflow_sandbox_service_resolves_locator_and_proxies() -> None: node_id="node-1", node_execution_id="node-exec-1", path="report.txt", + session=session_factory.create_session(), ) - assert result.file.reference == "dify-file-ref:file-1" + assert result.url == "https://files.example/report.txt?token=1&as_attachment=true" assert client.calls == [("upload", "report.txt")] + assert captured == { + "tenant_id": "tenant-1", + "file_mapping": { + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + "download_url": "https://files.example/report.txt?token=1", + }, + } + + +def test_upload_download_response_resolves_signed_external_url(monkeypatch: pytest.MonkeyPatch) -> None: + built_file = object() + built_with: dict[str, object] = {} + + def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: + built_with["mapping"] = mapping + built_with["tenant_id"] = tenant_id + built_with["access_controller"] = access_controller + return built_file + + class FakeRuntime: + def __init__(self, *, file_access_controller: object) -> None: + self.file_access_controller = file_access_controller + + def resolve_file_url(self, *, file: object, for_external: bool) -> str: + assert file is built_file + assert for_external is True + return "https://files.example/files/tools/tool-file.txt?timestamp=1&nonce=2&sign=3" + + monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) + monkeypatch.setattr("services.agent_app_sandbox_service.DifyWorkflowFileRuntime", FakeRuntime) + + result = _upload_download_response( + tenant_id="tenant-1", + file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, + ) + + assert result.url == ( + "https://files.example/files/tools/tool-file.txt?timestamp=1&nonce=2&sign=3&as_attachment=true" + ) + assert built_with["mapping"] == {"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"} + assert built_with["tenant_id"] == "tenant-1" + assert built_with["access_controller"] is not None + + +def test_upload_download_response_maps_resolution_failure_to_inspector_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: + del mapping, tenant_id, access_controller + raise ValueError("missing tool file") + + monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) + + with pytest.raises(AgentSandboxInspectorError) as exc_info: + _upload_download_response( + tenant_id="tenant-1", + file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, + ) + + assert exc_info.value.code == "sandbox_upload_download_unavailable" + assert exc_info.value.status_code == 502 + + +def test_upload_download_response_maps_missing_url_to_inspector_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + built_file = object() + + def fake_build_from_mapping(*, mapping: dict[str, object], tenant_id: str, access_controller: object) -> object: + del mapping, tenant_id, access_controller + return built_file + + class FakeRuntime: + def __init__(self, *, file_access_controller: object) -> None: + self.file_access_controller = file_access_controller + + def resolve_file_url(self, *, file: object, for_external: bool) -> None: + assert file is built_file + assert for_external is True + + monkeypatch.setattr("services.agent_app_sandbox_service.file_factory.build_from_mapping", fake_build_from_mapping) + monkeypatch.setattr("services.agent_app_sandbox_service.DifyWorkflowFileRuntime", FakeRuntime) + + with pytest.raises(AgentSandboxInspectorError) as exc_info: + _upload_download_response( + tenant_id="tenant-1", + file_mapping={"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, + ) + + assert exc_info.value.code == "sandbox_upload_download_unavailable" + assert exc_info.value.status_code == 502 @pytest.mark.usefixtures("_runtime_session_table") @@ -252,6 +390,7 @@ def test_workflow_sandbox_service_filters_by_node_execution_id() -> None: node_id="node-1", node_execution_id="node-exec-2", path="out.txt", + session=session_factory.create_session(), ) assert result.text == "hello" @@ -285,6 +424,7 @@ def test_workflow_sandbox_service_uses_latest_active_session_when_execution_id_o node_id="node-1", node_execution_id=None, path=".", + session=session_factory.create_session(), ) assert result.path == "." @@ -304,6 +444,7 @@ def test_workflow_sandbox_service_raises_when_no_active_session() -> None: node_id="node-1", node_execution_id=None, path=".", + session=session_factory.create_session(), ) assert exc_info.value.code == "no_active_session" @@ -323,6 +464,7 @@ def test_workflow_sandbox_service_raises_when_runtime_specs_missing() -> None: node_id="node-1", node_execution_id=None, path=".", + session=session_factory.create_session(), ) assert exc_info.value.code == "no_sandbox" diff --git a/api/tests/unit_tests/services/test_agent_drive_service.py b/api/tests/unit_tests/services/test_agent_drive_service.py index ad72e142522..6371b197325 100644 --- a/api/tests/unit_tests/services/test_agent_drive_service.py +++ b/api/tests/unit_tests/services/test_agent_drive_service.py @@ -125,6 +125,7 @@ def _commit(key: str, tool_file_id: str, *, owned: bool = True): value_owned_by_drive=owned, ) ], + session=session_factory.create_session(), ) @@ -132,15 +133,25 @@ def test_commit_then_manifest_lists_the_entry(): tf = _seed_tool_file() _commit("data/report.txt", tf) - items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT) + items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) assert [i["key"] for i in items] == ["data/report.txt"] assert items[0]["file_kind"] == "tool_file" assert items[0]["file_id"] == tf assert items[0]["mime_type"] == "text/plain" # prefix filter - assert AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, prefix="data/") != [] - assert AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, prefix="other/") == [] + assert ( + AgentDriveService().manifest( + tenant_id=TENANT, agent_id=AGENT, prefix="data/", session=session_factory.create_session() + ) + != [] + ) + assert ( + AgentDriveService().manifest( + tenant_id=TENANT, agent_id=AGENT, prefix="other/", session=session_factory.create_session() + ) + == [] + ) def test_commit_skill_row_persists_metadata_and_lists_catalog() -> None: @@ -157,6 +168,7 @@ def test_commit_skill_row_persists_metadata_and_lists_catalog() -> None: skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description="Parses RFPs."), ) ], + session=session_factory.create_session(), ) with session_factory.create_session() as session: @@ -165,7 +177,7 @@ def test_commit_skill_row_persists_metadata_and_lists_catalog() -> None: assert row.is_skill is True assert row.skill_metadata == '{"description":"Parses RFPs.","name":"Tender Analyzer"}' - skills = AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT) + skills = AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) assert len(skills) == 1 assert skills[0]["path"] == "tender-analyzer" assert skills[0]["skill_md_key"] == "tender-analyzer/SKILL.md" @@ -191,6 +203,7 @@ def test_commit_rejects_skill_row_without_skill_metadata() -> None: is_skill=True, ) ], + session=session_factory.create_session(), ) assert exc_info.value.code == "invalid_skill_metadata" @@ -220,7 +233,7 @@ def test_list_skills_raises_controlled_error_for_invalid_stored_metadata(raw_met session.commit() with pytest.raises(AgentDriveError) as exc_info: - AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT) + AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) assert exc_info.value.code == "invalid_skill_metadata" @@ -239,6 +252,7 @@ def test_commit_rejects_non_skill_row_with_skill_metadata() -> None: skill_metadata=DriveSkillMetadata(name="Bad", description=""), ) ], + session=session_factory.create_session(), ) @@ -257,6 +271,7 @@ def test_commit_rejects_non_canonical_skill_key() -> None: skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description=""), ) ], + session=session_factory.create_session(), ) @@ -282,6 +297,7 @@ def test_commit_rejects_agent_from_another_tenant(): value_owned_by_drive=True, ) ], + session=session_factory.create_session(), ) assert exc_info.value.status_code == 404 assert exc_info.value.code == "agent_not_found" @@ -311,24 +327,27 @@ def test_batch_failure_does_not_delete_old_storage_before_commit(): _commit("doc.txt", tf1, owned=True) with patch("services.agent_drive_service.storage") as storage_mock: - with pytest.raises(AgentDriveError): - AgentDriveService().commit( - tenant_id=TENANT, - user_id=USER, - agent_id=AGENT, - items=[ - DriveCommitItem( - key="doc.txt", - file_ref={"kind": "tool_file", "id": tf2}, - value_owned_by_drive=True, - ), - DriveCommitItem( - key="bad.txt", - file_ref={"kind": "tool_file", "id": "44444444-4444-4444-4444-444444444444"}, - value_owned_by_drive=True, - ), - ], - ) + with session_factory.create_session() as session: + with pytest.raises(AgentDriveError): + AgentDriveService().commit( + tenant_id=TENANT, + user_id=USER, + agent_id=AGENT, + items=[ + DriveCommitItem( + key="doc.txt", + file_ref={"kind": "tool_file", "id": tf2}, + value_owned_by_drive=True, + ), + DriveCommitItem( + key="bad.txt", + file_ref={"kind": "tool_file", "id": "44444444-4444-4444-4444-444444444444"}, + value_owned_by_drive=True, + ), + ], + session=session, + ) + session.rollback() storage_mock.delete.assert_not_called() with session_factory.create_session() as session: @@ -389,6 +408,7 @@ def test_recommit_same_skill_value_updates_metadata_without_cleaning_backing_fil skill_metadata=DriveSkillMetadata(name="Tender Analyzer", description="v1"), ) ], + session=session_factory.create_session(), ) with patch("services.agent_drive_service.storage") as storage_mock: @@ -405,6 +425,7 @@ def test_recommit_same_skill_value_updates_metadata_without_cleaning_backing_fil skill_metadata=DriveSkillMetadata(name="Tender Analyzer v2", description="v2"), ) ], + session=session_factory.create_session(), ) storage_mock.delete.assert_not_called() @@ -449,6 +470,7 @@ def _commit_upload(key: str, upload_file_id: str, *, owned: bool = True): value_owned_by_drive=owned, ) ], + session=session_factory.create_session(), ) @@ -456,7 +478,7 @@ def test_commit_upload_file_source_and_manifest(): uf = _seed_upload_file() _commit_upload("docs/u.txt", uf) - items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT) + items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) assert items[0]["file_kind"] == "upload_file" assert items[0]["file_id"] == uf assert items[0]["mime_type"] == "text/plain" @@ -492,7 +514,9 @@ def test_manifest_includes_internal_download_url(): patch("core.app.workflow.file_runtime.DifyWorkflowFileRuntime") as runtime_cls, ): runtime_cls.return_value.resolve_file_url.return_value = "http://internal/files/x?sign=1" - items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, include_download_url=True) + items = AgentDriveService().manifest( + tenant_id=TENANT, agent_id=AGENT, include_download_url=True, session=session_factory.create_session() + ) assert items[0]["download_url"] == "http://internal/files/x?sign=1" # drive-owned resolution: internal URL (for_external=False) @@ -507,7 +531,9 @@ def test_manifest_download_url_none_when_unresolvable(): "services.agent_drive_service.file_factory.build_from_mapping", side_effect=ValueError("not found"), ): - items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, include_download_url=True) + items = AgentDriveService().manifest( + tenant_id=TENANT, agent_id=AGENT, include_download_url=True, session=session_factory.create_session() + ) assert items[0]["download_url"] is None @@ -524,6 +550,7 @@ def test_delete_by_key_cleans_drive_owned_value(): user_id=USER, agent_id=AGENT, items=[DriveCommitItem(key="files/doomed.txt", file_ref=None)], + session=session_factory.create_session(), ) storage_mock.delete.assert_called_once() @@ -560,6 +587,7 @@ def test_commit_null_batch_removes_multiple_skill_keys(): DriveCommitItem(key="tender-analyzer/SKILL.md", file_ref=None), DriveCommitItem(key="tender-analyzer/.DIFY-SKILL-FULL.zip", file_ref=None), ], + session=session_factory.create_session(), ) assert sorted(item["key"] for item in removed) == [ @@ -581,6 +609,7 @@ def test_commit_null_is_idempotent_for_missing_keys(): user_id=USER, agent_id=AGENT, items=[DriveCommitItem(key="files/never-there.txt", file_ref=None)], + session=session_factory.create_session(), ) assert removed == [{"key": "files/never-there.txt", "removed": True, "noop": True}] @@ -595,6 +624,7 @@ def test_commit_null_keeps_shared_value_records(): user_id=USER, agent_id=AGENT, items=[DriveCommitItem(key="files/shared.txt", file_ref=None)], + session=session_factory.create_session(), ) storage_mock.delete.assert_not_called() @@ -639,7 +669,9 @@ def test_preview_returns_text_with_truncation_flags(): with patch("services.agent_drive_service.storage") as storage_mock: storage_mock.load_stream.return_value = iter([b"# PDF Toolkit\nUse responsibly.\n"]) - result = AgentDriveService().preview(tenant_id=TENANT, agent_id=AGENT, key="pdf-toolkit/SKILL.md") + result = AgentDriveService().preview( + tenant_id=TENANT, agent_id=AGENT, key="pdf-toolkit/SKILL.md", session=session_factory.create_session() + ) assert result == { "key": "pdf-toolkit/SKILL.md", @@ -656,13 +688,17 @@ def test_preview_marks_binary_and_oversized_content(): with patch("services.agent_drive_service.storage") as storage_mock: storage_mock.load_stream.return_value = iter([b"\x00\x01\x02"]) - binary = AgentDriveService().preview(tenant_id=TENANT, agent_id=AGENT, key="files/blob.bin") + binary = AgentDriveService().preview( + tenant_id=TENANT, agent_id=AGENT, key="files/blob.bin", session=session_factory.create_session() + ) assert binary["binary"] is True assert binary["text"] is None with patch("services.agent_drive_service.storage") as storage_mock: storage_mock.load_stream.return_value = iter([b"x" * (AgentDriveService.PREVIEW_MAX_BYTES + 10)]) - oversized = AgentDriveService().preview(tenant_id=TENANT, agent_id=AGENT, key="files/blob.bin") + oversized = AgentDriveService().preview( + tenant_id=TENANT, agent_id=AGENT, key="files/blob.bin", session=session_factory.create_session() + ) assert oversized["truncated"] is True assert oversized["binary"] is False assert len(oversized["text"]) == AgentDriveService.PREVIEW_MAX_BYTES @@ -670,7 +706,9 @@ def test_preview_marks_binary_and_oversized_content(): def test_preview_unknown_key_is_404(): with pytest.raises(AgentDriveError) as exc_info: - AgentDriveService().preview(tenant_id=TENANT, agent_id=AGENT, key="ghost/SKILL.md") + AgentDriveService().preview( + tenant_id=TENANT, agent_id=AGENT, key="ghost/SKILL.md", session=session_factory.create_session() + ) assert exc_info.value.code == "drive_key_not_found" assert exc_info.value.status_code == 404 @@ -678,7 +716,10 @@ def test_preview_unknown_key_is_404(): def test_preview_rejects_cross_tenant_agent(): with pytest.raises(AgentDriveError) as exc_info: AgentDriveService().preview( - tenant_id="99999999-9999-9999-9999-999999999999", agent_id=AGENT, key="pdf-toolkit/SKILL.md" + tenant_id="99999999-9999-9999-9999-999999999999", + agent_id=AGENT, + key="pdf-toolkit/SKILL.md", + session=session_factory.create_session(), ) assert exc_info.value.code == "agent_not_found" @@ -688,7 +729,12 @@ def test_download_url_signs_external_audience(): _commit("pdf-toolkit/.DIFY-SKILL-FULL.zip", tf) with patch.object(AgentDriveService, "_resolve_download_url", return_value="https://signed.example/x") as resolver: - url = AgentDriveService().download_url(tenant_id=TENANT, agent_id=AGENT, key="pdf-toolkit/.DIFY-SKILL-FULL.zip") + url = AgentDriveService().download_url( + tenant_id=TENANT, + agent_id=AGENT, + key="pdf-toolkit/.DIFY-SKILL-FULL.zip", + session=session_factory.create_session(), + ) assert url == "https://signed.example/x" # console downloads are for browsers: external signing, never the internal URL @@ -702,7 +748,9 @@ def test_upload_file_download_url_uses_attachment_filename(): with patch("core.app.workflow.file_runtime.DifyWorkflowFileRuntime") as runtime_cls: runtime_cls.return_value.resolve_upload_file_url.return_value = "https://files.example/report.pdf" - url = AgentDriveService().download_url(tenant_id=TENANT, agent_id=AGENT, key="files/report.pdf") + url = AgentDriveService().download_url( + tenant_id=TENANT, agent_id=AGENT, key="files/report.pdf", session=session_factory.create_session() + ) assert url == "https://files.example/report.pdf" assert runtime_cls.return_value.resolve_upload_file_url.call_args.kwargs["for_external"] is True @@ -712,7 +760,7 @@ def test_upload_file_download_url_uses_attachment_filename(): def test_manifest_items_carry_created_at_for_inspector(): tf = _seed_tool_file() _commit("files/x.txt", tf) - items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT) + items = AgentDriveService().manifest(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) assert items[0]["created_at"] is None or isinstance(items[0]["created_at"], int) @@ -744,13 +792,14 @@ def _commit_skill(*, manifest_files: list[str] | None = None) -> None: value_owned_by_drive=True, ), ], + session=session_factory.create_session(), ) def test_list_skills_uses_canonical_skill_rows(): _commit_skill(manifest_files=["SKILL.md", "scripts/run.py"]) - skills = AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT) + skills = AgentDriveService().list_skills(tenant_id=TENANT, agent_id=AGENT, session=session_factory.create_session()) created_at = skills[0].pop("created_at") assert skills == [ @@ -773,7 +822,9 @@ def test_inspect_skill_returns_manifest_files_and_file_tree(): with patch("services.agent_drive_service.storage") as storage_mock: storage_mock.load_stream.return_value = iter([b"# PDF Toolkit\n"]) - result = AgentDriveService().inspect_skill(tenant_id=TENANT, agent_id=AGENT, skill_path="pdf-toolkit") + result = AgentDriveService().inspect_skill( + tenant_id=TENANT, agent_id=AGENT, skill_path="pdf-toolkit", session=session_factory.create_session() + ) assert result["source"] == "skill_md" assert result["warnings"] == [] @@ -792,7 +843,9 @@ def test_inspect_skill_falls_back_to_drive_keys_when_manifest_missing(): with patch("services.agent_drive_service.storage") as storage_mock: storage_mock.load_stream.return_value = iter([b"# PDF Toolkit\n"]) - result = AgentDriveService().inspect_skill(tenant_id=TENANT, agent_id=AGENT, skill_path="pdf-toolkit") + result = AgentDriveService().inspect_skill( + tenant_id=TENANT, agent_id=AGENT, skill_path="pdf-toolkit", session=session_factory.create_session() + ) assert result["warnings"] == ["manifest_files_unavailable"] assert [file["path"] for file in result["files"]] == ["SKILL.md"] @@ -808,6 +861,7 @@ def test_preview_skill_archive_member_from_manifest_without_drive_row(): tenant_id=TENANT, agent_id=AGENT, key="pdf-toolkit/references/guide.md", + session=session_factory.create_session(), ) assert result == { @@ -831,6 +885,7 @@ def test_download_url_signs_skill_archive_member_from_manifest_without_drive_row tenant_id=TENANT, agent_id=AGENT, key="pdf-toolkit/references/guide.md", + session=session_factory.create_session(), ) assert url == "https://signed.example/member" @@ -856,5 +911,6 @@ def test_skill_metadata_rejects_non_canonical_rows(): skill_metadata=DriveSkillMetadata(name="Bad"), ) ], + session=session_factory.create_session(), ) assert exc_info.value.code == "invalid_skill_key" diff --git a/api/tests/unit_tests/services/test_agent_tool_inner_service.py b/api/tests/unit_tests/services/test_agent_tool_inner_service.py index 93e20d267da..61049d29e9e 100644 --- a/api/tests/unit_tests/services/test_agent_tool_inner_service.py +++ b/api/tests/unit_tests/services/test_agent_tool_inner_service.py @@ -70,7 +70,7 @@ def test_invoke_uses_agent_tool_runtime_and_returns_observation() -> None: side_effect=lambda messages, **_kwargs: messages, ), ): - response = AgentToolInnerService().invoke(session, _request()) + response = AgentToolInnerService().invoke(_request(), session=session) assert response.observation == "ok" assert response.metadata == { @@ -89,7 +89,7 @@ def test_invoke_raises_app_not_found_when_session_has_no_app() -> None: session.get.return_value = None with pytest.raises(AgentToolInnerServiceError) as exc_info: - AgentToolInnerService().invoke(session, _request()) + AgentToolInnerService().invoke(_request(), session=session) assert exc_info.value.error_code == "app_not_found" assert exc_info.value.status_code == 404 @@ -102,7 +102,7 @@ def test_invoke_raises_app_tenant_mismatch_when_app_belongs_to_other_tenant() -> session.get.return_value = fake_app with pytest.raises(AgentToolInnerServiceError) as exc_info: - AgentToolInnerService().invoke(session, _request()) + AgentToolInnerService().invoke(_request(), session=session) assert exc_info.value.error_code == "app_tenant_mismatch" assert exc_info.value.status_code == 403 @@ -120,7 +120,7 @@ def test_invoke_maps_tool_runtime_app_not_found_value_error_to_specific_error_co patch("services.agent_tool_inner_service.ToolEngine.generic_invoke", side_effect=ValueError("app not found")), ): with pytest.raises(AgentToolInnerServiceError) as exc_info: - AgentToolInnerService().invoke(session, _request()) + AgentToolInnerService().invoke(_request(), session=session) assert exc_info.value.error_code == "app_not_found" assert exc_info.value.status_code == 404 @@ -141,7 +141,7 @@ def test_invoke_maps_tool_invoke_error_without_private_tool_engine_helper() -> N ), ): with pytest.raises(AgentToolInnerServiceError) as exc_info: - AgentToolInnerService().invoke(session, _request()) + AgentToolInnerService().invoke(_request(), session=session) assert exc_info.value.error_code == "agent_tool_invoke_failed" @@ -161,6 +161,6 @@ def test_invoke_maps_runtime_lookup_errors_to_service_error_codes(error: Excepti with patch("services.agent_tool_inner_service.ToolManager.get_agent_tool_runtime", side_effect=error): with pytest.raises(AgentToolInnerServiceError) as exc_info: - AgentToolInnerService().invoke(session, _request()) + AgentToolInnerService().invoke(_request(), session=session) assert exc_info.value.error_code == expected_code diff --git a/api/tests/unit_tests/services/test_annotation_service.py b/api/tests/unit_tests/services/test_annotation_service.py index 2975c4df14c..79bbb5873ac 100644 --- a/api/tests/unit_tests/services/test_annotation_service.py +++ b/api/tests/unit_tests/services/test_annotation_service.py @@ -103,7 +103,7 @@ class TestAppAnnotationServiceUpInsert: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.up_insert_app_annotation_from_message(args, "app-1") + AppAnnotationService.up_insert_app_annotation_from_message(args, "app-1", session=mock_db.session) def test_up_insert_app_annotation_from_message_should_raise_value_error_when_answer_missing(self) -> None: """Test missing answer and content raises ValueError.""" @@ -121,7 +121,7 @@ class TestAppAnnotationServiceUpInsert: # Act & Assert with pytest.raises(ValueError): - AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) def test_up_insert_app_annotation_from_message_should_raise_not_found_when_message_missing(self) -> None: """Test missing message raises NotFound.""" @@ -139,7 +139,7 @@ class TestAppAnnotationServiceUpInsert: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) def test_up_insert_app_annotation_from_message_should_update_existing_annotation_when_found(self) -> None: """Test existing annotation is updated and indexed.""" @@ -161,7 +161,7 @@ class TestAppAnnotationServiceUpInsert: mock_db.session.scalar.side_effect = [app, message, setting] # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) # Assert assert result == annotation @@ -199,7 +199,7 @@ class TestAppAnnotationServiceUpInsert: mock_db.session.scalar.side_effect = [app, message, None] # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) # Assert assert result == annotation_instance @@ -231,7 +231,7 @@ class TestAppAnnotationServiceUpInsert: # Act & Assert with pytest.raises(ValueError): - AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) def test_up_insert_app_annotation_from_message_should_create_annotation_when_message_missing(self) -> None: """Test annotation is created when message_id is not provided.""" @@ -252,7 +252,7 @@ class TestAppAnnotationServiceUpInsert: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id) + result = AppAnnotationService.up_insert_app_annotation_from_message(args, app.id, session=mock_db.session) # Assert assert result == annotation_instance @@ -383,7 +383,7 @@ class TestAppAnnotationServiceListAndExport: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.get_annotation_list_by_app_id("app-1", 1, 10, "") + AppAnnotationService.get_annotation_list_by_app_id("app-1", 1, 10, "", session=mock_db.session) def test_get_annotation_list_by_app_id_should_return_items_with_keyword(self) -> None: """Test keyword search returns items and total.""" @@ -402,7 +402,9 @@ class TestAppAnnotationServiceListAndExport: mock_paginate.return_value = pagination # Act - items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "keyword") + items, total = AppAnnotationService.get_annotation_list_by_app_id( + app.id, 1, 10, "keyword", session=mock_db.session + ) # Assert assert items == ["a1"] @@ -424,7 +426,9 @@ class TestAppAnnotationServiceListAndExport: mock_paginate.return_value = pagination # Act - items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "") + items, total = AppAnnotationService.get_annotation_list_by_app_id( + app.id, 1, 10, "", session=mock_db.session + ) # Assert assert items == ["a1", "a2"] @@ -451,7 +455,7 @@ class TestAppAnnotationServiceListAndExport: mock_db.session.scalars.return_value.all.return_value = [annotation1, annotation2] # Act - result = AppAnnotationService.export_annotation_list_by_app_id(app.id) + result = AppAnnotationService.export_annotation_list_by_app_id(app.id, session=mock_db.session) # Assert assert result == [annotation1, annotation2] @@ -473,7 +477,7 @@ class TestAppAnnotationServiceListAndExport: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.export_annotation_list_by_app_id("app-1") + AppAnnotationService.export_annotation_list_by_app_id("app-1", session=mock_db.session) class TestAppAnnotationServiceDirectManipulation: @@ -493,7 +497,7 @@ class TestAppAnnotationServiceDirectManipulation: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.insert_app_annotation_directly(args, "app-1") + AppAnnotationService.insert_app_annotation_directly(args, "app-1", session=mock_db.session) def test_insert_app_annotation_directly_should_raise_value_error_when_question_missing(self) -> None: """Test missing question raises ValueError.""" @@ -510,7 +514,7 @@ class TestAppAnnotationServiceDirectManipulation: # Act & Assert with pytest.raises(ValueError): - AppAnnotationService.insert_app_annotation_directly(args, app.id) + AppAnnotationService.insert_app_annotation_directly(args, app.id, session=mock_db.session) def test_insert_app_annotation_directly_should_create_annotation_and_index(self) -> None: """Test insert creates annotation and triggers index task.""" @@ -531,7 +535,7 @@ class TestAppAnnotationServiceDirectManipulation: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.insert_app_annotation_directly(args, app.id) + result = AppAnnotationService.insert_app_annotation_directly(args, app.id, session=mock_db.session) # Assert assert result == annotation_instance @@ -696,7 +700,9 @@ class TestAppAnnotationServiceDirectManipulation: mock_db.session.execute.return_value.all.return_value = [] # Act - result = AppAnnotationService.delete_app_annotations_in_batch(_make_app_ref(app), ["ann-1"]) + result = AppAnnotationService.delete_app_annotations_in_batch( + _make_app_ref(app), ["ann-1"], session=mock_db.session + ) # Assert assert result == {"deleted_count": 0} @@ -723,7 +729,9 @@ class TestAppAnnotationServiceDirectManipulation: mock_db.session.execute.side_effect = [execute_result_multi, MagicMock(), execute_result_delete] # Act - result = AppAnnotationService.delete_app_annotations_in_batch(_make_app_ref(app), ["ann-1", "ann-2"]) + result = AppAnnotationService.delete_app_annotations_in_batch( + _make_app_ref(app), ["ann-1", "ann-2"], session=mock_db.session + ) # Assert assert result == {"deleted_count": 2} @@ -755,7 +763,7 @@ class TestAppAnnotationServiceBatchImport: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.batch_import_app_annotations("app-1", file) + AppAnnotationService.batch_import_app_annotations("app-1", file, session=mock_db.session) def test_batch_import_app_annotations_should_return_error_when_columns_invalid(self) -> None: """Test invalid column count returns error message.""" @@ -777,7 +785,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -801,7 +809,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -829,7 +837,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -855,7 +863,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -885,7 +893,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -911,7 +919,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -937,7 +945,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -963,7 +971,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -994,7 +1002,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert error_msg = cast(str, result["error_msg"]) @@ -1027,7 +1035,7 @@ class TestAppAnnotationServiceBatchImport: mock_db.session.scalar.return_value = app # Act - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert assert result == {"job_id": "uuid-3", "job_status": "waiting", "record_count": 1} @@ -1067,7 +1075,7 @@ class TestAppAnnotationServiceBatchImport: # Act with caplog.at_level(logging.DEBUG): - result = AppAnnotationService.batch_import_app_annotations(app.id, file) + result = AppAnnotationService.batch_import_app_annotations(app.id, file, session=mock_db.session) # Assert assert result["error_msg"] == "An error occurred while processing the file: boom" @@ -1090,7 +1098,9 @@ class TestAppAnnotationServiceHitHistoryAndSettings: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.get_annotation_hit_histories(_make_annotation_ref(app, "ann-1"), 1, 10) + AppAnnotationService.get_annotation_hit_histories( + _make_annotation_ref(app, "ann-1"), 1, 10, session=mock_db.session + ) def test_get_annotation_hit_histories_should_return_items_and_total(self) -> None: """Test hit histories pagination returns items and total.""" @@ -1114,6 +1124,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: _make_annotation_ref(app, annotation.id), 1, 10, + session=mock_db.session, ) # Assert @@ -1129,7 +1140,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.get.return_value = None # Act - result = AppAnnotationService.get_annotation_by_id("ann-1") + result = AppAnnotationService.get_annotation_by_id("ann-1", session=mock_db.session) # Assert assert result is None @@ -1142,7 +1153,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.get.return_value = annotation # Act - result = AppAnnotationService.get_annotation_by_id("ann-1") + result = AppAnnotationService.get_annotation_by_id("ann-1", session=mock_db.session) # Assert assert result == annotation @@ -1165,6 +1176,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: message_id="msg-1", from_source="chat", score=0.8, + session=mock_db.session, ) # Assert @@ -1187,7 +1199,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, session=mock_db.session) # Assert assert result["enabled"] is True @@ -1208,7 +1220,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.get_app_annotation_setting_by_app_id("app-1") + AppAnnotationService.get_app_annotation_setting_by_app_id("app-1", session=mock_db.session) def test_get_app_annotation_setting_by_app_id_should_return_empty_embedding_model_when_no_detail(self) -> None: """Test setting without detail returns empty embedding model.""" @@ -1224,7 +1236,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, session=mock_db.session) # Assert assert result["enabled"] is True @@ -1243,7 +1255,7 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.scalar.side_effect = [app, None] # Act - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id) + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, session=mock_db.session) # Assert assert result == {"enabled": False} @@ -1265,7 +1277,9 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.update_app_annotation_setting(app.id, setting.id, args) + result = AppAnnotationService.update_app_annotation_setting( + app.id, setting.id, args, session=mock_db.session + ) # Assert assert result["enabled"] is True @@ -1292,7 +1306,9 @@ class TestAppAnnotationServiceHitHistoryAndSettings: mock_db.session.scalar.side_effect = [app, setting] # Act - result = AppAnnotationService.update_app_annotation_setting(app.id, setting.id, args) + result = AppAnnotationService.update_app_annotation_setting( + app.id, setting.id, args, session=mock_db.session + ) # Assert assert result["enabled"] is True @@ -1312,7 +1328,9 @@ class TestAppAnnotationServiceHitHistoryAndSettings: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.update_app_annotation_setting("app-1", "setting-1", {"score_threshold": 0.5}) + AppAnnotationService.update_app_annotation_setting( + "app-1", "setting-1", {"score_threshold": 0.5}, session=mock_db.session + ) def test_update_app_annotation_setting_should_raise_not_found_when_setting_missing(self) -> None: """Test update raises NotFound when setting is missing.""" @@ -1328,7 +1346,9 @@ class TestAppAnnotationServiceHitHistoryAndSettings: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.update_app_annotation_setting(app.id, "setting-1", {"score_threshold": 0.5}) + AppAnnotationService.update_app_annotation_setting( + app.id, "setting-1", {"score_threshold": 0.5}, session=mock_db.session + ) class TestAppAnnotationServiceClearAll: @@ -1361,7 +1381,7 @@ class TestAppAnnotationServiceClearAll: mock_db.session.scalars.side_effect = [annotations_scalars, histories_scalars_1, histories_scalars_2] # Act - result = AppAnnotationService.clear_all_annotations(app.id) + result = AppAnnotationService.clear_all_annotations(app.id, session=mock_db.session) # Assert assert result == {"result": "success"} @@ -1385,4 +1405,4 @@ class TestAppAnnotationServiceClearAll: # Act & Assert with pytest.raises(NotFound): - AppAnnotationService.clear_all_annotations("app-1") + AppAnnotationService.clear_all_annotations("app-1", session=mock_db.session) diff --git a/api/tests/unit_tests/services/test_app_dsl_service.py b/api/tests/unit_tests/services/test_app_dsl_service.py new file mode 100644 index 00000000000..ae621bdcf52 --- /dev/null +++ b/api/tests/unit_tests/services/test_app_dsl_service.py @@ -0,0 +1,52 @@ +from unittest.mock import Mock + +import pytest + +from services.app_dsl_service import AppDslService +from services.entities.dsl_entities import ImportStatus + + +def test_import_app_rejects_oversized_yaml_content_before_parsing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("services.app_dsl_service.DSL_MAX_SIZE", 3) + service = AppDslService(session=Mock()) + account = Mock(current_tenant_id="tenant-1") + + result = service.import_app(account=account, import_mode="yaml-content", yaml_content="你你") + + assert result.status == ImportStatus.FAILED + assert result.error == "File size exceeds the limit of 10MB" + + +def test_import_app_rejects_oversized_yaml_url_bytes_before_decode(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("services.app_dsl_service.DSL_MAX_SIZE", 1) + response = Mock() + response.raise_for_status.return_value = None + response.content = b"\xff\xff" + monkeypatch.setattr("services.app_dsl_service.remote_fetcher.make_request", Mock(return_value=response)) + service = AppDslService(session=Mock()) + + result = service.import_app( + account=Mock(current_tenant_id="tenant-1"), + import_mode="yaml-url", + yaml_url="https://example.com/app.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert result.error == "File size exceeds the limit of 10MB" + + +def test_import_app_returns_decode_error_for_invalid_yaml_url_bytes(monkeypatch: pytest.MonkeyPatch) -> None: + response = Mock() + response.raise_for_status.return_value = None + response.content = b"\xff" + monkeypatch.setattr("services.app_dsl_service.remote_fetcher.make_request", Mock(return_value=response)) + service = AppDslService(session=Mock()) + + result = service.import_app( + account=Mock(current_tenant_id="tenant-1"), + import_mode="yaml-url", + yaml_url="https://example.com/app.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert "utf-8" in result.error diff --git a/api/tests/unit_tests/services/test_app_generate_service.py b/api/tests/unit_tests/services/test_app_generate_service.py index 865410fddf1..552403e0d0b 100644 --- a/api/tests/unit_tests/services/test_app_generate_service.py +++ b/api/tests/unit_tests/services/test_app_generate_service.py @@ -235,12 +235,12 @@ class TestGenerate: side_effect=lambda x: x, ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "ok"} gen_spy.assert_called_once() @@ -256,12 +256,12 @@ class TestGenerate: side_effect=lambda x: x, ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.AGENT_CHAT), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "agent"} gen_spy.assert_called_once() @@ -278,12 +278,12 @@ class TestGenerate: ) app = _make_app(AppMode.CHAT, is_agent=True) result = AppGenerateService.generate( - MagicMock(), app_model=app, user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "agent-via-flag"} gen_spy.assert_called_once() @@ -300,32 +300,16 @@ class TestGenerate: ) app = _make_app(AppMode.CHAT, is_agent=False) result = AppGenerateService.generate( - MagicMock(), app_model=app, user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "chat"} gen_spy.assert_called_once() - def test_stateless_agent_mode(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.AgentAppGenerator.generate_stateless", - return_value={"result": "stateless-agent"}, - ) - - result = AppGenerateService.generate_stateless_agent_app( - app_model=_make_app(AppMode.AGENT), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - ) - - assert result == {"result": "stateless-agent"} - gen_spy.assert_called_once() - # -- ADVANCED_CHAT blocking --------------------------------------------- def test_advanced_chat_blocking(self, mocker: MockerFixture): workflow = _make_workflow() @@ -342,12 +326,12 @@ class TestGenerate: ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.ADVANCED_CHAT), user=_make_user(), args={"workflow_id": None, "query": "hi", "inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "advanced-blocking"} call_kwargs = gen_spy.call_args.kwargs @@ -375,12 +359,12 @@ class TestGenerate: ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.ADVANCED_CHAT), user=_make_user(), args={"workflow_id": None, "query": "hi", "inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=MagicMock(), ) # In streaming mode it should go through retrieve_events, not generate gen_instance.retrieve_events.assert_called_once() @@ -401,12 +385,12 @@ class TestGenerate: ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.WORKFLOW), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) assert result == {"result": "workflow-blocking"} call_kwargs = gen_spy.call_args.kwargs @@ -435,12 +419,12 @@ class TestGenerate: ) result = AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.WORKFLOW), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=MagicMock(), ) retrieve_spy.assert_called_once() # The inner on_subscribe closure was invoked by _build_streaming_task_on_subscribe @@ -451,12 +435,12 @@ class TestGenerate: app = _make_app("invalid-mode", is_agent=False) with pytest.raises(ValueError, match="Invalid app mode"): AppGenerateService.generate( - MagicMock(), app_model=app, user=_make_user(), args={}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) @@ -489,12 +473,12 @@ class TestGenerateBilling: ) AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) reserve_mock.assert_called_once_with(QuotaType.WORKFLOW, "tenant-id") quota_charge.commit.assert_called_once() @@ -513,12 +497,12 @@ class TestGenerateBilling: with pytest.raises(InvokeRateLimitError): AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) def test_exception_refunds_quota_and_exits_rate_limit(self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): @@ -539,12 +523,12 @@ class TestGenerateBilling: with pytest.raises(RuntimeError, match="boom"): AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) quota_charge.refund.assert_called_once() @@ -571,83 +555,16 @@ class TestGenerateBilling: ) AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) # exit is called in finally block for non-streaming assert exit_calls == ["dummy-request-id"] - def test_stateless_agent_app_uses_billing_and_rate_limit_guardrails( - self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch - ): - monkeypatch.setattr(ags_module.dify_config, "BILLING_ENABLED", True) - quota_charge = MagicMock() - reserve_mock = mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - exit_calls: list[str] = [] - - class _TrackingRateLimit(_DummyRateLimit): - def exit(self, request_id: str) -> None: - exit_calls.append(request_id) - - mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) - gen_spy = mocker.patch( - "services.app_generate_service.AgentAppGenerator.generate_stateless", - return_value={"ok": True}, - ) - - result = AppGenerateService.generate_stateless_agent_app( - app_model=_make_app(AppMode.AGENT), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - ) - - assert result == {"ok": True} - reserve_mock.assert_called_once_with(QuotaType.WORKFLOW, "tenant-id") - quota_charge.commit.assert_called_once() - assert exit_calls == ["dummy-request-id"] - gen_spy.assert_called_once() - - def test_stateless_agent_app_failure_refunds_quota_and_exits_once( - self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch - ): - monkeypatch.setattr(ags_module.dify_config, "BILLING_ENABLED", True) - quota_charge = MagicMock() - mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - exit_calls: list[str] = [] - - class _TrackingRateLimit(_DummyRateLimit): - def exit(self, request_id: str) -> None: - exit_calls.append(request_id) - - mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) - mocker.patch( - "services.app_generate_service.AgentAppGenerator.generate_stateless", - side_effect=RuntimeError("boom"), - ) - - with pytest.raises(RuntimeError, match="boom"): - AppGenerateService.generate_stateless_agent_app( - app_model=_make_app(AppMode.AGENT), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - ) - - quota_charge.commit.assert_called_once() - quota_charge.refund.assert_called_once() - assert exit_calls == ["dummy-request-id"] - def test_blocking_failure_exits_rate_limit_once(self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(ags_module.dify_config, "BILLING_ENABLED", True) quota_charge = MagicMock() @@ -669,12 +586,12 @@ class TestGenerateBilling: with pytest.raises(RuntimeError, match="boom"): AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, + session=MagicMock(), ) quota_charge.refund.assert_called_once() @@ -701,12 +618,12 @@ class TestGenerateBilling: with pytest.raises(RuntimeError, match="boom"): AppGenerateService.generate( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, + session=MagicMock(), ) quota_charge.refund.assert_called_once() @@ -723,7 +640,7 @@ class TestGetWorkflow: ws.get_draft_workflow.return_value = draft_wf mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - result = AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER) + result = AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=MagicMock()) assert result is draft_wf ws.get_draft_workflow.assert_called_once() @@ -733,7 +650,7 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) with pytest.raises(ValueError, match="Workflow not initialized"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER) + AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=MagicMock()) def test_non_debugger_fetches_published(self, mocker: MockerFixture): pub_wf = _make_workflow() @@ -741,7 +658,9 @@ class TestGetWorkflow: ws.get_published_workflow.return_value = pub_wf mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - result = AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API) + result = AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=MagicMock() + ) assert result is pub_wf ws.get_published_workflow.assert_called_once() @@ -751,7 +670,7 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) with pytest.raises(ValueError, match="Workflow not published"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API) + AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=MagicMock()) def test_specific_workflow_id_valid_uuid(self, mocker: MockerFixture): valid_uuid = str(uuid.uuid4()) @@ -761,7 +680,10 @@ class TestGetWorkflow: mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) result = AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id=valid_uuid + _make_app(AppMode.WORKFLOW), + InvokeFrom.SERVICE_API, + workflow_id=valid_uuid, + session=MagicMock(), ) assert result is specific_wf ws.get_published_workflow_by_id.assert_called_once() @@ -772,7 +694,10 @@ class TestGetWorkflow: with pytest.raises(WorkflowIdFormatError): AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id="not-a-uuid" + _make_app(AppMode.WORKFLOW), + InvokeFrom.SERVICE_API, + workflow_id="not-a-uuid", + session=MagicMock(), ) def test_specific_workflow_id_not_found(self, mocker: MockerFixture): @@ -783,7 +708,10 @@ class TestGetWorkflow: with pytest.raises(WorkflowNotFoundError): AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id=valid_uuid + _make_app(AppMode.WORKFLOW), + InvokeFrom.SERVICE_API, + workflow_id=valid_uuid, + session=MagicMock(), ) @@ -804,7 +732,11 @@ class TestGenerateSingleIteration: ) app = _make_app(AppMode.ADVANCED_CHAT) result = AppGenerateService.generate_single_iteration( - app_model=app, user=_make_user(), node_id="n1", args={"k": "v"} + app_model=app, + user=_make_user(), + node_id="n1", + args={"k": "v"}, + session=MagicMock(), ) iter_spy.assert_called_once() assert result == {"event": "iteration"} @@ -822,7 +754,11 @@ class TestGenerateSingleIteration: ) app = _make_app(AppMode.WORKFLOW) result = AppGenerateService.generate_single_iteration( - app_model=app, user=_make_user(), node_id="n1", args={"k": "v"} + app_model=app, + user=_make_user(), + node_id="n1", + args={"k": "v"}, + session=MagicMock(), ) iter_spy.assert_called_once() assert result == {"event": "wf-iteration"} @@ -830,7 +766,9 @@ class TestGenerateSingleIteration: def test_invalid_mode_raises(self, mocker: MockerFixture): app = _make_app(AppMode.CHAT) with pytest.raises(ValueError, match="Invalid app mode"): - AppGenerateService.generate_single_iteration(app_model=app, user=_make_user(), node_id="n1", args={}) + AppGenerateService.generate_single_iteration( + app_model=app, user=_make_user(), node_id="n1", args={}, session=MagicMock() + ) # --------------------------------------------------------------------------- @@ -850,7 +788,11 @@ class TestGenerateSingleLoop: ) app = _make_app(AppMode.ADVANCED_CHAT) result = AppGenerateService.generate_single_loop( - app_model=app, user=_make_user(), node_id="n1", args=MagicMock() + app_model=app, + user=_make_user(), + node_id="n1", + args=MagicMock(), + session=MagicMock(), ) loop_spy.assert_called_once() assert result == {"event": "loop"} @@ -868,7 +810,11 @@ class TestGenerateSingleLoop: ) app = _make_app(AppMode.WORKFLOW) result = AppGenerateService.generate_single_loop( - app_model=app, user=_make_user(), node_id="n1", args=MagicMock() + app_model=app, + user=_make_user(), + node_id="n1", + args=MagicMock(), + session=MagicMock(), ) loop_spy.assert_called_once() assert result == {"event": "wf-loop"} @@ -876,7 +822,9 @@ class TestGenerateSingleLoop: def test_invalid_mode_raises(self, mocker: MockerFixture): app = _make_app(AppMode.COMPLETION) with pytest.raises(ValueError, match="Invalid app mode"): - AppGenerateService.generate_single_loop(app_model=app, user=_make_user(), node_id="n1", args=MagicMock()) + AppGenerateService.generate_single_loop( + app_model=app, user=_make_user(), node_id="n1", args=MagicMock(), session=MagicMock() + ) # --------------------------------------------------------------------------- @@ -888,16 +836,18 @@ class TestGenerateMoreLikeThis: "services.app_generate_service.CompletionAppGenerator.generate_more_like_this", return_value={"result": "similar"}, ) + session = MagicMock() result = AppGenerateService.generate_more_like_this( - MagicMock(), app_model=_make_app(AppMode.COMPLETION), user=_make_user(), message_id="msg-1", invoke_from=InvokeFrom.SERVICE_API, + session=session, streaming=True, ) assert result == {"result": "similar"} gen_spy.assert_called_once() + assert gen_spy.call_args.kwargs["session"] is session assert gen_spy.call_args.kwargs["stream"] is True diff --git a/api/tests/unit_tests/services/test_app_service.py b/api/tests/unit_tests/services/test_app_service.py index c57fb6ed775..36914679d3e 100644 --- a/api/tests/unit_tests/services/test_app_service.py +++ b/api/tests/unit_tests/services/test_app_service.py @@ -29,14 +29,14 @@ class TestOpenapiVisibilityHelpers: sentinel_app.status = "archived" # explicitly NOT "normal" mock_session.get.return_value = sentinel_app - assert AppService.get_app_by_id(mock_session, "app-uuid") is sentinel_app + assert AppService.get_app_by_id("app-uuid", session=mock_session) is sentinel_app mock_session.get.assert_called_once_with(App, "app-uuid") def test_get_app_by_id_returns_none_when_missing(self): mock_session = MagicMock() mock_session.get.return_value = None - assert AppService.get_app_by_id(mock_session, "missing") is None + assert AppService.get_app_by_id("missing", session=mock_session) is None def test_get_visible_app_by_id_returns_app_when_visible(self): mock_session = MagicMock() @@ -45,7 +45,7 @@ class TestOpenapiVisibilityHelpers: mock_session.get.return_value = app with patch("services.app_service.is_openapi_visible", return_value=True): - assert AppService.get_visible_app_by_id(mock_session, "app-uuid") is app + assert AppService.get_visible_app_by_id("app-uuid", session=mock_session) is app mock_session.get.assert_called_once_with(App, "app-uuid") @@ -53,7 +53,7 @@ class TestOpenapiVisibilityHelpers: mock_session = MagicMock() mock_session.get.return_value = None - assert AppService.get_visible_app_by_id(mock_session, "missing") is None + assert AppService.get_visible_app_by_id("missing", session=mock_session) is None def test_get_visible_app_by_id_returns_none_when_status_not_normal(self): """Soft-deleted/archived rows must not surface on the openapi @@ -65,7 +65,7 @@ class TestOpenapiVisibilityHelpers: mock_session.get.return_value = app with patch("services.app_service.is_openapi_visible", return_value=True): - assert AppService.get_visible_app_by_id(mock_session, "app-uuid") is None + assert AppService.get_visible_app_by_id("app-uuid", session=mock_session) is None def test_get_visible_app_by_id_returns_none_when_visibility_gate_rejects(self): """``is_openapi_visible`` is the per-row counterpart to @@ -78,7 +78,7 @@ class TestOpenapiVisibilityHelpers: mock_session.get.return_value = app with patch("services.app_service.is_openapi_visible", return_value=False): - assert AppService.get_visible_app_by_id(mock_session, "app-uuid") is None + assert AppService.get_visible_app_by_id("app-uuid", session=mock_session) is None def test_find_visible_apps_by_name_returns_scalars_through_visibility_gate(self): """Tenant-scoped name lookup. The helper passes the SELECT through @@ -90,7 +90,7 @@ class TestOpenapiVisibilityHelpers: mock_session.execute.return_value.scalars.return_value = iter(rows) with patch("services.app_service.apply_openapi_gate", side_effect=lambda q: q) as gate: - out = AppService.find_visible_apps_by_name(mock_session, name="my-app", tenant_id="tenant-1") + out = AppService.find_visible_apps_by_name(name="my-app", tenant_id="tenant-1", session=mock_session) assert out == rows # Visibility gate must wrap the SELECT exactly once. @@ -102,7 +102,7 @@ class TestOpenapiVisibilityHelpers: mock_session.execute.return_value.scalars.return_value = iter([]) with patch("services.app_service.apply_openapi_gate", side_effect=lambda q: q): - out = AppService.find_visible_apps_by_name(mock_session, name="nope", tenant_id="tenant-1") + out = AppService.find_visible_apps_by_name(name="nope", tenant_id="tenant-1", session=mock_session) assert out == [] @@ -113,7 +113,7 @@ class TestOpenapiVisibilityHelpers: """ mock_session = MagicMock() - assert AppService.find_visible_apps_by_ids(mock_session, []) == [] + assert AppService.find_visible_apps_by_ids([], session=mock_session) == [] mock_session.execute.assert_not_called() def test_find_visible_apps_by_ids_passes_through_visibility_gate(self): @@ -127,7 +127,7 @@ class TestOpenapiVisibilityHelpers: mock_session.execute.return_value.scalars.return_value.all.return_value = rows with patch("services.app_service.apply_openapi_gate", side_effect=lambda q: q) as gate: - out = AppService.find_visible_apps_by_ids(mock_session, ["a", "b"]) + out = AppService.find_visible_apps_by_ids(["a", "b"], session=mock_session) assert out == rows gate.assert_called_once() @@ -208,6 +208,7 @@ class TestAgentAppType: "use_icon_as_answer_icon": False, "max_active_requests": 0, }, + session=mock_db.session, ) assert updated_app.name == "Iris" @@ -266,6 +267,7 @@ class TestAgentAppType: "use_icon_as_answer_icon": False, "max_active_requests": 0, }, + session=mock_db.session, ) assert backing_agent.role == "research assistant" @@ -317,6 +319,7 @@ class TestAgentAppType: "use_icon_as_answer_icon": False, "max_active_requests": 0, }, + session=mock_db.session, ) assert backing_agent.role == "" @@ -370,6 +373,7 @@ class TestAgentAppType: "use_icon_as_answer_icon": False, "max_active_requests": 0, }, + session=mock_db.session, ) mock_db.session.rollback.assert_called_once() @@ -392,7 +396,7 @@ class TestAgentAppType: patch("services.app_service.remove_app_and_related_data_task"), ): mock_db.session.scalar.return_value = backing_agent - AppService().delete_app(app) # type: ignore[arg-type] + AppService().delete_app(app, session=mock_db.session) # type: ignore[arg-type] assert backing_agent.status == AgentStatus.ARCHIVED assert backing_agent.archived_by == "account-2" diff --git a/api/tests/unit_tests/services/test_async_workflow_service.py b/api/tests/unit_tests/services/test_async_workflow_service.py index 1b9cc8a2ff6..567066845bf 100644 --- a/api/tests/unit_tests/services/test_async_workflow_service.py +++ b/api/tests/unit_tests/services/test_async_workflow_service.py @@ -331,7 +331,7 @@ class TestAsyncWorkflowService: assert trigger_log.triggered_at is not None repo.update.assert_called_once_with(trigger_log) session.commit.assert_called_once() - called_trigger_data = mock_trigger_workflow_async.call_args[0][2] + called_trigger_data = mock_trigger_workflow_async.call_args.args[1] assert isinstance(called_trigger_data, TriggerData) assert called_trigger_data.app_id == "app-123" @@ -465,11 +465,16 @@ class TestAsyncWorkflowServiceGetWorkflow: workflow_service.get_published_workflow_by_id.return_value = workflow # Act - result = AsyncWorkflowService._get_workflow(workflow_service, app_model, workflow_id="workflow-123") + session = MagicMock() + result = AsyncWorkflowService._get_workflow( + workflow_service, app_model, workflow_id="workflow-123", session=session + ) # Assert assert result == workflow - workflow_service.get_published_workflow_by_id.assert_called_once_with(app_model, "workflow-123", session=None) + workflow_service.get_published_workflow_by_id.assert_called_once_with( + app_model, "workflow-123", session=session + ) workflow_service.get_published_workflow.assert_not_called() def test_should_raise_when_specific_workflow_id_not_found(self): @@ -481,7 +486,9 @@ class TestAsyncWorkflowServiceGetWorkflow: # Act / Assert with pytest.raises(WorkflowNotFoundError, match="Published workflow not found: workflow-404"): - AsyncWorkflowService._get_workflow(workflow_service, app_model, workflow_id="workflow-404") + AsyncWorkflowService._get_workflow( + workflow_service, app_model, workflow_id="workflow-404", session=MagicMock() + ) def test_should_return_default_published_workflow_when_workflow_id_not_provided(self): """Test _get_workflow returns default published workflow when no id is provided.""" @@ -493,11 +500,12 @@ class TestAsyncWorkflowServiceGetWorkflow: workflow_service.get_published_workflow.return_value = workflow # Act - result = AsyncWorkflowService._get_workflow(workflow_service, app_model) + session = MagicMock() + result = AsyncWorkflowService._get_workflow(workflow_service, app_model, session=session) # Assert assert result == workflow - workflow_service.get_published_workflow.assert_called_once_with(app_model, session=None) + workflow_service.get_published_workflow.assert_called_once_with(app_model, session=session) workflow_service.get_published_workflow_by_id.assert_not_called() def test_should_raise_when_default_published_workflow_not_found(self): @@ -510,4 +518,4 @@ class TestAsyncWorkflowServiceGetWorkflow: # Act / Assert with pytest.raises(WorkflowNotFoundError, match="No published workflow found for app: app-123"): - AsyncWorkflowService._get_workflow(workflow_service, app_model) + AsyncWorkflowService._get_workflow(workflow_service, app_model, session=MagicMock()) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index e5610545aa9..2b1bdb5c5c4 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -73,6 +73,23 @@ class TestBillingServiceSendRequest: assert call_args[1]["headers"]["Billing-Api-Secret-Key"] == "test-secret-key" assert call_args[1]["headers"]["Content-Type"] == "application/json" + def test_send_request_with_base_url_override(self, mock_httpx_request, mock_billing_config): + """Quota APIs can use the new billing service without changing legacy billing calls.""" + # Arrange + expected_response = {"result": "success"} + mock_response = MagicMock() + mock_response.status_code = httpx.codes.OK + mock_response.json.return_value = expected_response + mock_httpx_request.return_value = mock_response + + # Act + result = BillingService._send_request("GET", "/quota/balance", base_url="https://quota.example.com") + + # Assert + assert result == expected_response + call_args = mock_httpx_request.call_args + assert call_args[0][1] == "https://quota.example.com/quota/balance" + @pytest.mark.parametrize( "status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST] ) @@ -393,6 +410,20 @@ class TestBillingServiceSubscriptionInfo: params={"tenant_id": tenant_id}, ) + def test_quota_get_balance_uses_quota_request(self): + tenant_id = "tenant-123" + with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request: + mock_send_quota_request.return_value = {"quota": "200", "usage": "6", "available": "194", "reserved": "0"} + + result = BillingService.quota_get_balance(tenant_id, "credit_pool", bucket="trial") + + assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0} + mock_send_quota_request.assert_called_once_with( + "GET", + "/quota/balance", + params={"tenant_id": tenant_id, "feature_key": "credit_pool", "bucket": "trial"}, + ) + def test_get_knowledge_rate_limit_with_defaults(self, mock_send_request): """Test knowledge rate limit retrieval with default values.""" # Arrange @@ -518,19 +549,20 @@ class TestBillingServiceUsageCalculation: assert result == expected_response mock_send_request.assert_called_once_with("GET", "/tenant-feature-usage/info", params={"tenant_id": tenant_id}) - def test_get_quota_info(self, mock_send_request): + def test_get_quota_info(self): """Test retrieval of quota info from new endpoint.""" # Arrange tenant_id = "tenant-123" expected_response = {"trigger_event": {"limit": 100, "usage": 30}, "api_rate_limit": {"limit": -1, "usage": 0}} - mock_send_request.return_value = expected_response + with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request: + mock_send_quota_request.return_value = expected_response - # Act - result = BillingService.get_quota_info(tenant_id) + # Act + result = BillingService.get_quota_info(tenant_id) # Assert assert result == expected_response - mock_send_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id}) + mock_send_quota_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id}) def test_update_tenant_feature_plan_usage_positive_delta(self, mock_send_request): """Test updating tenant feature usage with positive delta (adding credits).""" @@ -614,7 +646,7 @@ class TestBillingServiceQuotaOperations: @pytest.fixture def mock_send_request(self): - with patch.object(BillingService, "_send_request") as mock: + with patch.object(BillingService, "_send_quota_request") as mock: yield mock def test_quota_reserve_success(self, mock_send_request): @@ -652,6 +684,16 @@ class TestBillingServiceQuotaOperations: call_json = mock_send_request.call_args[1]["json"] assert call_json["meta"] == {"source": "webhook"} + def test_quota_reserve_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"reservation_id": "rid-2", "available": 98, "reserved": 1} + + BillingService.quota_reserve( + tenant_id="t1", feature_key="credit_pool", request_id="req-2", amount=1, bucket="trial" + ) + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "trial" + def test_quota_commit_success(self, mock_send_request): expected = {"available": 98, "reserved": 0, "refunded": 0} mock_send_request.return_value = expected @@ -696,6 +738,20 @@ class TestBillingServiceQuotaOperations: call_json = mock_send_request.call_args[1]["json"] assert call_json["meta"] == {"reason": "partial"} + def test_quota_commit_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"available": 97, "reserved": 0, "refunded": 0} + + BillingService.quota_commit( + tenant_id="t1", + feature_key="credit_pool", + reservation_id="rid-1", + actual_amount=1, + bucket="paid", + ) + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "paid" + def test_quota_release_success(self, mock_send_request): expected = {"available": 100, "reserved": 0, "released": 1} mock_send_request.return_value = expected @@ -720,6 +776,64 @@ class TestBillingServiceQuotaOperations: assert result["released"] == 1 assert isinstance(result["released"], int) + def test_quota_release_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"available": 100, "reserved": 0, "released": 1} + + BillingService.quota_release(tenant_id="t1", feature_key="credit_pool", reservation_id="rid-1", bucket="trial") + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "trial" + + def test_quota_consume_capped_success(self, mock_send_request): + mock_send_request.return_value = { + "deducted": "2", + "available": "8", + "reserved": "0", + "quota": "10", + "usage": "2", + } + + result = BillingService.quota_consume_capped( + tenant_id="t1", + feature_key="credit_pool", + request_id="req-1", + amount=5, + bucket="paid", + meta={"source": "test"}, + ) + + assert result == {"deducted": 2, "available": 8, "reserved": 0, "quota": 10, "usage": 2} + mock_send_request.assert_called_once_with( + "POST", + "/quota/consume-capped", + json={ + "tenant_id": "t1", + "feature_key": "credit_pool", + "request_id": "req-1", + "amount": 5, + "bucket": "paid", + "meta": {"source": "test"}, + }, + ) + + def test_send_quota_request_uses_quota_base_url(self): + with ( + patch.object(BillingService, "quota_base_url", "https://quota.example.com/v1"), + patch.object(BillingService, "_send_request") as mock_send_request, + ): + mock_send_request.return_value = {"ok": True} + + result = BillingService._send_quota_request("GET", "/quota/info", params={"tenant_id": "t1"}) + + assert result == {"ok": True} + mock_send_request.assert_called_once_with( + "GET", + "/quota/info", + json=None, + params={"tenant_id": "t1"}, + base_url="https://quota.example.com/v1", + ) + def test_get_quota_info_coerces_string_to_int(self, mock_send_request): """Test that TypeAdapter coerces string values to int for get_quota_info.""" mock_send_request.return_value = { @@ -1115,7 +1229,7 @@ class TestBillingServiceAccountManagement: mock_db_session.scalar.return_value = mock_join # Act - should not raise exception - BillingService.is_tenant_owner_or_admin(mock_db_session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=mock_db_session) mock_db_session.scalar.assert_called_once() def test_is_tenant_owner_or_admin_admin(self, mock_db_session): @@ -1131,7 +1245,7 @@ class TestBillingServiceAccountManagement: mock_db_session.scalar.return_value = mock_join # Act - should not raise exception - BillingService.is_tenant_owner_or_admin(mock_db_session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=mock_db_session) mock_db_session.scalar.assert_called_once() def test_is_tenant_owner_or_admin_normal_user_raises_error(self, mock_db_session): @@ -1148,7 +1262,7 @@ class TestBillingServiceAccountManagement: # Act & Assert with pytest.raises(ValueError) as exc_info: - BillingService.is_tenant_owner_or_admin(mock_db_session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=mock_db_session) assert "Only team owner or team admin can perform this action" in str(exc_info.value) mock_db_session.scalar.assert_called_once() @@ -1163,7 +1277,7 @@ class TestBillingServiceAccountManagement: # Act & Assert with pytest.raises(ValueError) as exc_info: - BillingService.is_tenant_owner_or_admin(mock_db_session, current_user) + BillingService.is_tenant_owner_or_admin(current_user, session=mock_db_session) assert "Tenant account join not found" in str(exc_info.value) mock_db_session.scalar.assert_called_once() diff --git a/api/tests/unit_tests/services/test_conversation_service.py b/api/tests/unit_tests/services/test_conversation_service.py index 2c7f13b79f3..e6f7b48f651 100644 --- a/api/tests/unit_tests/services/test_conversation_service.py +++ b/api/tests/unit_tests/services/test_conversation_service.py @@ -330,12 +330,9 @@ class TestConversationServiceHelpers: class TestConversationServiceConversationalVariable: """Test conversational variable operations.""" - @patch("services.conversation_service.session_factory") @patch("services.conversation_service.ConversationService.get_conversation") @patch("services.conversation_service.dify_config") - def test_get_conversational_variable_with_name_filter_mysql( - self, mock_config, mock_get_conversation, mock_session_factory - ): + def test_get_conversational_variable_with_name_filter_mysql(self, mock_config, mock_get_conversation): """ Test variable filtering by name for MySQL databases. @@ -351,7 +348,6 @@ class TestConversationServiceConversationalVariable: # Mock session mock_session = MagicMock() - mock_session_factory.create_session.return_value.__enter__.return_value = mock_session mock_session.scalars.return_value.all.return_value = [] # Act @@ -362,6 +358,7 @@ class TestConversationServiceConversationalVariable: limit=10, last_id=None, variable_name="test_var", + session=mock_session, ) # Assert - JSON filter should be applied diff --git a/api/tests/unit_tests/services/test_credential_permission_service.py b/api/tests/unit_tests/services/test_credential_permission_service.py index e467e9c8c5e..cdcf4a6b00f 100644 --- a/api/tests/unit_tests/services/test_credential_permission_service.py +++ b/api/tests/unit_tests/services/test_credential_permission_service.py @@ -40,7 +40,7 @@ class TestGetPartialMemberList: session = MagicMock() session.scalars.return_value.all.return_value = [] result = CredentialPermissionService.get_partial_member_list( - session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION + credential_id, CredentialType.TRIGGER_SUBSCRIPTION, session=session ) assert result == [] session.scalars.assert_called_once() @@ -49,7 +49,7 @@ class TestGetPartialMemberList: session = MagicMock() session.scalars.return_value.all.return_value = [user_id, other_user_id] result = CredentialPermissionService.get_partial_member_list( - session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION + credential_id, CredentialType.TRIGGER_SUBSCRIPTION, session=session ) assert set(result) == {user_id, other_user_id} session.scalars.assert_called_once() diff --git a/api/tests/unit_tests/services/test_credit_pool_service.py b/api/tests/unit_tests/services/test_credit_pool_service.py index 5e589804c3d..8cafd3af590 100644 --- a/api/tests/unit_tests/services/test_credit_pool_service.py +++ b/api/tests/unit_tests/services/test_credit_pool_service.py @@ -1,13 +1,12 @@ from collections.abc import Generator -from contextlib import contextmanager from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest from sqlalchemy import create_engine, select from sqlalchemy.engine import Engine -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from core.errors.error import QuotaExceededError from models import TenantCreditPool @@ -15,6 +14,8 @@ from models.enums import ProviderQuotaType from services.credit_pool_service import ( CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS, CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS, + FEATURE_KEY_CREDIT_POOL, + CreditPoolBalance, CreditPoolService, ) @@ -38,11 +39,8 @@ def _create_engine_with_pool(*, quota_limit: int, quota_used: int) -> tuple[Engi return engine, tenant_id, pool_id -@contextmanager -def _patched_session_factory(engine: Engine) -> Generator[None, None, None]: - session_maker = sessionmaker(bind=engine, expire_on_commit=False) - with patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker): - yield +def _make_session(engine: Engine) -> Session: + return sessionmaker(bind=engine, expire_on_commit=False)() def _get_quota_used(*, engine: Engine, pool_id: str) -> int | None: @@ -50,64 +48,67 @@ def _get_quota_used(*, engine: Engine, pool_id: str) -> int | None: return connection.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id)) -def _make_session_maker(session: MagicMock) -> MagicMock: - session_maker = MagicMock() - transaction = session_maker.begin.return_value - transaction.__enter__.return_value = session - transaction.__exit__.return_value = None - return session_maker - - def _make_redis_lock() -> MagicMock: lock = MagicMock() lock.acquire.return_value = True return lock -def test_get_pool_uses_configured_session_factory_without_flask_app_context() -> None: +@pytest.fixture(autouse=True) +def _disable_billing_quota_by_default() -> Generator[None, None, None]: + with patch("services.credit_pool_service.dify_config.BILLING_ENABLED", False): + yield + + +def test_get_pool_uses_provided_session() -> None: engine, tenant_id, _ = _create_engine_with_pool(quota_limit=10, quota_used=2) - with _patched_session_factory(engine): - pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL) + with _make_session(engine) as session: + pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL, session=session) assert pool is not None assert pool.tenant_id == tenant_id assert pool.quota_used == 2 +def test_credit_pool_balance_unlimited_remaining_and_sufficiency() -> None: + pool = CreditPoolBalance(tenant_id="tenant-1", pool_type="paid", quota_limit=-1, quota_used=999) + + assert pool.remaining_credits == -1 + assert pool.has_sufficient_credits(10_000) + + def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) - with _patched_session_factory(engine): - deducted_credits = CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=3) + with _make_session(engine) as session: + deducted_credits = CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, credits_required=3, session=session + ) assert deducted_credits == 3 assert _get_quota_used(engine=engine, pool_id=pool_id) == 5 def test_check_and_deduct_credits_returns_zero_for_non_positive_request() -> None: - assert CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=0) == 0 + assert ( + CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=0, session=MagicMock()) == 0 + ) def test_check_and_deduct_credits_raises_when_pool_is_missing() -> None: engine = create_engine("sqlite:///:memory:") TenantCreditPool.__table__.create(engine) - with ( - _patched_session_factory(engine), - pytest.raises(QuotaExceededError, match="Credit pool not found"), - ): - CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=1) + with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="Credit pool not found"): + CreditPoolService.check_and_deduct_credits(tenant_id=str(uuid4()), credits_required=1, session=session) def test_check_and_deduct_credits_raises_when_pool_is_empty() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=10) - with ( - _patched_session_factory(engine), - pytest.raises(QuotaExceededError, match="No credits remaining"), - ): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1) + with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="No credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1, session=session) assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 @@ -115,11 +116,8 @@ def test_check_and_deduct_credits_raises_when_pool_is_empty() -> None: def test_check_and_deduct_credits_raises_without_partial_deduction_when_insufficient() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=9) - with ( - _patched_session_factory(engine), - pytest.raises(QuotaExceededError, match="Insufficient credits remaining"), - ): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=3) + with _make_session(engine) as session, pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=3, session=session) assert _get_quota_used(engine=engine, pool_id=pool_id) == 9 @@ -128,25 +126,27 @@ def test_check_and_deduct_credits_wraps_unexpected_deduction_errors() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) with ( - _patched_session_factory(engine), + _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1) + CreditPoolService.check_and_deduct_credits(tenant_id=tenant_id, credits_required=1, session=session) assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 def test_deduct_credits_capped_returns_zero_for_non_positive_request() -> None: - assert CreditPoolService.deduct_credits_capped(tenant_id=str(uuid4()), credits_required=0) == 0 + assert CreditPoolService.deduct_credits_capped(tenant_id=str(uuid4()), credits_required=0, session=MagicMock()) == 0 def test_deduct_credits_capped_returns_zero_when_pool_is_missing() -> None: engine = create_engine("sqlite:///:memory:") TenantCreditPool.__table__.create(engine) - with _patched_session_factory(engine): - deducted_credits = CreditPoolService.deduct_credits_capped(tenant_id=str(uuid4()), credits_required=1) + with _make_session(engine) as session: + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=str(uuid4()), credits_required=1, session=session + ) assert deducted_credits == 0 @@ -154,8 +154,10 @@ def test_deduct_credits_capped_returns_zero_when_pool_is_missing() -> None: def test_deduct_credits_capped_returns_zero_when_pool_is_empty() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=10) - with _patched_session_factory(engine): - deducted_credits = CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + with _make_session(engine) as session: + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=1, session=session + ) assert deducted_credits == 0 assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 @@ -164,8 +166,10 @@ def test_deduct_credits_capped_returns_zero_when_pool_is_empty() -> None: def test_deduct_credits_capped_deducts_only_remaining_balance_when_insufficient() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=9) - with _patched_session_factory(engine): - deducted_credits = CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=3) + with _make_session(engine) as session: + deducted_credits = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, credits_required=3, session=session + ) assert deducted_credits == 1 assert _get_quota_used(engine=engine, pool_id=pool_id) == 10 @@ -175,11 +179,11 @@ def test_deduct_credits_capped_wraps_unexpected_deduction_errors() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) with ( - _patched_session_factory(engine), + _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=RuntimeError("database unavailable")), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1, session=session) assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 @@ -188,11 +192,11 @@ def test_deduct_credits_capped_reraises_quota_exceeded_errors() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) with ( - _patched_session_factory(engine), + _make_session(engine) as session, patch.object(CreditPoolService, "_get_locked_pool", side_effect=QuotaExceededError("quota unavailable")), pytest.raises(QuotaExceededError, match="quota unavailable"), ): - CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1) + CreditPoolService.deduct_credits_capped(tenant_id=tenant_id, credits_required=1, session=session) assert _get_quota_used(engine=engine, pool_id=pool_id) == 2 @@ -200,19 +204,18 @@ def test_deduct_credits_capped_reraises_quota_exceeded_errors() -> None: def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() -> None: tenant_id = "tenant-1" session = MagicMock() - session_maker = _make_session_maker(session) pool = SimpleNamespace(remaining_credits=10, quota_used=2) redis_lock = _make_redis_lock() with ( patch("services.credit_pool_service.redis_client.lock", return_value=redis_lock) as lock, - patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker), patch.object(CreditPoolService, "_get_locked_pool", return_value=pool) as get_locked_pool, ): result = CreditPoolService.check_and_deduct_credits( tenant_id=tenant_id, credits_required=3, pool_type=ProviderQuotaType.TRIAL, + session=session, ) assert result == 3 @@ -224,25 +227,24 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() - ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL) + get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="trial") def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> None: tenant_id = "tenant-1" session = MagicMock() - session_maker = _make_session_maker(session) pool = SimpleNamespace(remaining_credits=2, quota_used=8) redis_lock = _make_redis_lock() with ( patch("services.credit_pool_service.redis_client.lock", return_value=redis_lock) as lock, - patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker), patch.object(CreditPoolService, "_get_locked_pool", return_value=pool) as get_locked_pool, ): result = CreditPoolService.deduct_credits_capped( tenant_id=tenant_id, credits_required=5, pool_type=ProviderQuotaType.PAID, + session=session, ) assert result == 2 @@ -254,7 +256,152 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID) + get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="paid") + + +def test_get_pool_uses_billing_quota_balance_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_get_balance") as quota_get_balance, + ): + quota_get_balance.return_value = {"quota": 1000, "usage": 250, "available": 750, "reserved": 0} + + pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID) + + assert isinstance(pool, CreditPoolBalance) + assert pool.quota_limit == 1000 + assert pool.quota_used == 250 + assert pool.remaining_credits == 750 + quota_get_balance.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="paid", + ) + + +def test_check_and_deduct_credits_uses_billing_reserve_and_commit_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit") as quota_commit, + patch("services.billing_service.BillingService.quota_release") as quota_release, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + result = CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, + credits_required=3, + pool_type=ProviderQuotaType.TRIAL, + ) + + assert result == 3 + quota_reserve.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + request_id=ANY, + amount=3, + meta={"source": "credit_pool.check_and_deduct"}, + ) + quota_commit.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + actual_amount=3, + meta={"source": "credit_pool.check_and_deduct"}, + ) + quota_release.assert_not_called() + + +def test_check_and_deduct_credits_raises_when_billing_reserve_is_insufficient() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + ): + quota_reserve.return_value = {"reservation_id": "", "available": 1, "reserved": 0} + + with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + +def test_check_and_deduct_credits_releases_billing_reservation_when_commit_fails() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")), + patch("services.billing_service.BillingService.quota_release") as quota_release, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + with pytest.raises(RuntimeError, match="commit failed"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + quota_release.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + ) + + +def test_check_and_deduct_credits_logs_when_billing_release_fails() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")), + patch( + "services.billing_service.BillingService.quota_release", side_effect=RuntimeError("release failed") + ) as quota_release, + patch("services.credit_pool_service.logger.warning") as logger_warning, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + with pytest.raises(RuntimeError, match="commit failed"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + quota_release.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + ) + logger_warning.assert_called_once() + assert logger_warning.call_args.args[3] == "reservation-1" + assert logger_warning.call_args.kwargs["exc_info"] is True + + +def test_deduct_credits_capped_uses_billing_consume_capped_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_consume_capped") as quota_consume_capped, + ): + quota_consume_capped.return_value = { + "deducted": 2, + "available": 0, + "reserved": 0, + "quota": 10, + "usage": 10, + } + + result = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, + credits_required=5, + pool_type=ProviderQuotaType.PAID, + ) + + assert result == 2 + quota_consume_capped.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="paid", + request_id=ANY, + amount=5, + meta={"source": "credit_pool.deduct_capped"}, + ) @pytest.mark.parametrize( @@ -266,38 +413,35 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N ) def test_non_positive_credit_request_skips_tenant_redis_lock(deduct_method) -> None: with patch("services.credit_pool_service.redis_client.lock") as lock: - result = deduct_method(tenant_id="tenant-1", credits_required=0) + result = deduct_method(tenant_id="tenant-1", credits_required=0, session=MagicMock()) assert result == 0 lock.assert_not_called() def test_check_and_deduct_credits_wraps_redis_lock_errors_without_querying_db() -> None: - session_maker = MagicMock() + session = MagicMock() with ( patch("services.credit_pool_service.redis_client.lock", side_effect=RuntimeError("redis unavailable")), - patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker), pytest.raises(QuotaExceededError, match="Failed to deduct credits"), ): - CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=1) + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=1, session=session) - session_maker.begin.assert_not_called() + session.scalar.assert_not_called() def test_deduct_credits_capped_ignores_release_errors_after_successful_deduction() -> None: session = MagicMock() - session_maker = _make_session_maker(session) pool = SimpleNamespace(remaining_credits=3, quota_used=7) redis_lock = _make_redis_lock() redis_lock.release.side_effect = RuntimeError("release failed") with ( patch("services.credit_pool_service.redis_client.lock", return_value=redis_lock), - patch("services.credit_pool_service.session_factory.get_session_maker", return_value=session_maker), patch.object(CreditPoolService, "_get_locked_pool", return_value=pool), ): - result = CreditPoolService.deduct_credits_capped(tenant_id="tenant-1", credits_required=2) + result = CreditPoolService.deduct_credits_capped(tenant_id="tenant-1", credits_required=2, session=session) assert result == 2 assert pool.quota_used == 9 diff --git a/api/tests/unit_tests/services/test_dataset_service_dataset.py b/api/tests/unit_tests/services/test_dataset_service_dataset.py index 02d965f4bd2..dcb6250a8f0 100644 --- a/api/tests/unit_tests/services/test_dataset_service_dataset.py +++ b/api/tests/unit_tests/services/test_dataset_service_dataset.py @@ -344,7 +344,9 @@ class TestDatasetServiceCreationAndUpdate: mock_db.session.scalar.return_value = object() with pytest.raises(DatasetNameDuplicateError, match="Dataset with name Dataset already exists"): - DatasetService.create_empty_dataset(mock_db.session, "tenant-1", "Dataset", None, "economy", account) + DatasetService.create_empty_dataset( + "tenant-1", "Dataset", None, "economy", account, session=mock_db.session + ) def test_create_empty_dataset_uses_default_embedding_model_for_high_quality_dataset(self): account = SimpleNamespace(id="user-1") @@ -512,7 +514,7 @@ class TestDatasetServiceCreationAndUpdate: session = MagicMock() with patch.object(DatasetService, "get_dataset", return_value=None): with pytest.raises(ValueError, match="Dataset not found"): - DatasetService.update_dataset(session, "dataset-1", {}, SimpleNamespace(id="user-1")) + DatasetService.update_dataset("dataset-1", {}, SimpleNamespace(id="user-1"), session=session) def test_update_dataset_raises_when_new_name_conflicts(self): dataset = DatasetServiceUnitDataFactory.create_dataset_mock(dataset_id="dataset-1", tenant_id="tenant-1") @@ -524,10 +526,7 @@ class TestDatasetServiceCreationAndUpdate: ): with pytest.raises(ValueError, match="Dataset name already exists"): DatasetService.update_dataset( - MagicMock(), - "dataset-1", - {"name": "New Dataset"}, - SimpleNamespace(id="user-1"), + "dataset-1", {"name": "New Dataset"}, SimpleNamespace(id="user-1"), session=MagicMock() ) def test_update_dataset_routes_external_datasets_to_external_helper(self): @@ -541,7 +540,7 @@ class TestDatasetServiceCreationAndUpdate: patch.object(DatasetService, "_update_external_dataset", return_value="updated") as update_external, ): session = MagicMock() - result = DatasetService.update_dataset(session, "dataset-1", {"name": dataset.name}, user) + result = DatasetService.update_dataset("dataset-1", {"name": dataset.name}, user, session=session) assert result == "updated" check_permission.assert_called_once() @@ -560,7 +559,7 @@ class TestDatasetServiceCreationAndUpdate: patch.object(DatasetService, "_update_internal_dataset", return_value="updated") as update_internal, ): session = MagicMock() - result = DatasetService.update_dataset(session, "dataset-1", {"name": dataset.name}, user) + result = DatasetService.update_dataset("dataset-1", {"name": dataset.name}, user, session=session) assert result == "updated" check_permission.assert_called_once() @@ -612,7 +611,7 @@ class TestDatasetServiceCreationAndUpdate: assert dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM assert dataset.updated_by == "user-1" assert dataset.updated_at is now - get_external_knowledge_api.assert_called_once_with(mock_db.session, "api-1", dataset.tenant_id) + get_external_knowledge_api.assert_called_once_with("api-1", dataset.tenant_id, session=mock_db.session) update_binding.assert_called_once_with("dataset-1", "knowledge-1", "api-1", mock_db.session) mock_db.session.add.assert_called_once_with(dataset) mock_db.session.commit.assert_called_once() @@ -652,7 +651,7 @@ class TestDatasetServiceCreationAndUpdate: mock_db.session, ) - get_external_knowledge_api.assert_called_once_with(mock_db.session, "foreign-api", dataset.tenant_id) + get_external_knowledge_api.assert_called_once_with("foreign-api", dataset.tenant_id, session=mock_db.session) update_binding.assert_not_called() mock_db.session.commit.assert_not_called() @@ -1165,7 +1164,7 @@ class TestDatasetServiceRagPipelineSettings: with patch("services.dataset_service.current_user", SimpleNamespace(current_tenant_id=None)): with pytest.raises(ValueError, match="Current user or current tenant not found"): - DatasetService.update_rag_pipeline_dataset_settings(session, dataset, knowledge_configuration) + DatasetService.update_rag_pipeline_dataset_settings(dataset, knowledge_configuration, session=session) def test_update_rag_pipeline_dataset_settings_without_published_high_quality_updates_embedding_settings(self): session = MagicMock() @@ -1185,7 +1184,7 @@ class TestDatasetServiceRagPipelineSettings: ): model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - DatasetService.update_rag_pipeline_dataset_settings(session, dataset, knowledge_configuration) + DatasetService.update_rag_pipeline_dataset_settings(dataset, knowledge_configuration, session=session) assert dataset.chunk_structure == "paragraph" assert dataset.indexing_technique == "high_quality" @@ -1211,7 +1210,7 @@ class TestDatasetServiceRagPipelineSettings: ) with patch("services.dataset_service.current_user", SimpleNamespace(current_tenant_id="tenant-1")): - DatasetService.update_rag_pipeline_dataset_settings(session, dataset, knowledge_configuration) + DatasetService.update_rag_pipeline_dataset_settings(dataset, knowledge_configuration, session=session) assert dataset.indexing_technique == "economy" assert dataset.keyword_number == 12 @@ -1228,10 +1227,7 @@ class TestDatasetServiceRagPipelineSettings: with patch("services.dataset_service.current_user", SimpleNamespace(current_tenant_id="tenant-1")): with pytest.raises(ValueError, match="Chunk structure is not allowed to be updated"): DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) def test_update_rag_pipeline_dataset_settings_with_published_rejects_switch_to_economy(self): @@ -1252,10 +1248,7 @@ class TestDatasetServiceRagPipelineSettings: match="Knowledge base indexing technique is not allowed to be updated to economy", ): DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) def test_update_rag_pipeline_dataset_settings_with_published_adds_high_quality_index(self): @@ -1280,10 +1273,7 @@ class TestDatasetServiceRagPipelineSettings: model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) assert dataset.indexing_technique == "high_quality" @@ -1326,10 +1316,7 @@ class TestDatasetServiceRagPipelineSettings: ) DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) assert dataset.embedding_model_provider == "provider-two" @@ -1364,10 +1351,7 @@ class TestDatasetServiceRagPipelineSettings: ) DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) assert dataset.embedding_model_provider == "provider" @@ -1396,10 +1380,7 @@ class TestDatasetServiceRagPipelineSettings: patch("services.dataset_service.deal_dataset_index_update_task") as update_task, ): DatasetService.update_rag_pipeline_dataset_settings( - session, - dataset, - knowledge_configuration, - has_published=True, + dataset, knowledge_configuration, has_published=True, session=session ) assert dataset.keyword_number == 9 @@ -1457,7 +1438,7 @@ class TestDatasetPermissionService: session = MagicMock() with pytest.raises(NoPermissionError, match="does not have permission"): - DatasetPermissionService.check_permission(session, user, dataset, "all_team", []) + DatasetPermissionService.check_permission(user, dataset, "all_team", [], session=session) def test_check_permission_prevents_dataset_operator_from_changing_permission_mode(self): user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True) @@ -1465,7 +1446,7 @@ class TestDatasetPermissionService: session = MagicMock() with pytest.raises(NoPermissionError, match="cannot change the dataset permissions"): - DatasetPermissionService.check_permission(session, user, dataset, "only_me", []) + DatasetPermissionService.check_permission(user, dataset, "only_me", [], session=session) def test_check_permission_requires_partial_member_list_for_partial_members_mode(self): user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True) @@ -1473,7 +1454,7 @@ class TestDatasetPermissionService: session = MagicMock() with pytest.raises(ValueError, match="Partial member list is required"): - DatasetPermissionService.check_permission(session, user, dataset, "partial_members", []) + DatasetPermissionService.check_permission(user, dataset, "partial_members", [], session=session) def test_check_permission_rejects_dataset_operator_member_list_changes(self): user = SimpleNamespace(is_dataset_editor=True, is_dataset_operator=True) @@ -1485,11 +1466,7 @@ class TestDatasetPermissionService: with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]): with pytest.raises(ValueError, match="cannot change the dataset permissions"): DatasetPermissionService.check_permission( - session, - user, - dataset, - "partial_members", - [{"user_id": "user-2"}], + user, dataset, "partial_members", [{"user_id": "user-2"}], session=session ) def test_check_permission_allows_dataset_operator_when_member_list_is_unchanged(self): @@ -1501,11 +1478,7 @@ class TestDatasetPermissionService: with patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["user-1"]): DatasetPermissionService.check_permission( - session, - user, - dataset, - "partial_members", - [{"user_id": "user-1"}], + user, dataset, "partial_members", [{"user_id": "user-1"}], session=session ) def test_clear_partial_member_list_rolls_back_on_exception(self): diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py index 02661fbe1f3..44619a29e83 100644 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ b/api/tests/unit_tests/services/test_dataset_service_document.py @@ -1183,7 +1183,7 @@ class TestDocumentServiceTenantAndUpdateEdges: with patch("services.dataset_service.db") as mock_db: mock_db.session.scalar.return_value = 12 - result = DocumentService.get_tenant_documents_count(mock_db.session) + result = DocumentService.get_tenant_documents_count(session=mock_db.session) assert result == 12 diff --git a/api/tests/unit_tests/services/test_dataset_service_segment.py b/api/tests/unit_tests/services/test_dataset_service_segment.py index 34f3f947f96..c94093d59b7 100644 --- a/api/tests/unit_tests/services/test_dataset_service_segment.py +++ b/api/tests/unit_tests/services/test_dataset_service_segment.py @@ -306,13 +306,13 @@ class TestSegmentServiceQueries: def test_get_child_chunk_by_segment_ref_uses_full_ownership_chain(self): child_chunk = _make_child_chunk() segment_ref = _make_segment_ref() + session = MagicMock() + session.scalar.return_value = child_chunk - with patch("services.dataset_service.db") as mock_db: - mock_db.session.scalar.return_value = child_chunk - result = SegmentService.get_child_chunk_by_segment_ref("child-a", segment_ref) + result = SegmentService.get_child_chunk_by_segment_ref("child-a", segment_ref, session) assert result is child_chunk - stmt = mock_db.session.scalar.call_args.args[0] + stmt = session.scalar.call_args.args[0] sql = str(stmt.compile(compile_kwargs={"literal_binds": True})) assert "child_chunks.id = 'child-a'" in sql assert "child_chunks.tenant_id = 'tenant-1'" in sql @@ -381,13 +381,13 @@ class TestSegmentServiceQueries: ) segment.id = "segment-1" segment_ref = _make_segment_ref() + session = MagicMock() + session.scalar.return_value = segment - with patch("services.dataset_service.db") as mock_db: - mock_db.session.scalar.return_value = segment - result = SegmentService.get_segment_by_ref(segment_ref) + result = SegmentService.get_segment_by_ref(segment_ref, session) assert result is segment - stmt = mock_db.session.scalar.call_args.args[0] + stmt = session.scalar.call_args.args[0] sql = str(stmt.compile(compile_kwargs={"literal_binds": True})) assert "document_segments.id = 'segment-1'" in sql assert "document_segments.tenant_id = 'tenant-1'" in sql @@ -566,7 +566,7 @@ class TestSegmentServiceMutations: assert all(segment.error == "vector failed" for segment in result) assert document.word_count == 5 + sum(len(item["content"]) + len(item["answer"]) for item in segments) vector_service.create_segments_vector.assert_called_once_with( - [["k1"], None], result, dataset, document.doc_form + [["k1"], None], result, dataset, document.doc_form, mock_db.session ) mock_db.session.commit.assert_called_once() @@ -641,7 +641,7 @@ class TestSegmentServiceMutations: assert result is refreshed_segment assert segment.keywords == ["new"] vector_service.update_segment_vector.assert_called_once_with(["new"], segment, dataset) - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_regenerates_child_chunks_and_updates_manual_summary(self, account_context): segment = _make_segment(content="same content", word_count=len("same content")) @@ -684,10 +684,11 @@ class TestSegmentServiceMutations: dataset, embedding_model_instance, processing_rule, + mock_db.session, True, ) - update_summary.assert_called_once_with(segment, dataset, "new summary") - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + update_summary.assert_called_once_with(segment, dataset, "new summary", session=mock_db.session) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_auto_regenerates_summary_after_content_change(self, account_context): segment = _make_segment(content="old", word_count=3) @@ -725,8 +726,8 @@ class TestSegmentServiceMutations: assert segment.tokens == 9 assert document.word_count == 18 vector_service.update_segment_vector.assert_called_once_with(["kw-1"], segment, dataset) - generate_summary.assert_called_once_with(segment, dataset, {"enable": True}) - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + generate_summary.assert_called_once_with(segment, dataset, {"enable": True}, session=mock_db.session) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_regenerates_summary_when_manual_summary_is_unchanged(self, account_context): segment = _make_segment(content="old", word_count=3) @@ -760,9 +761,9 @@ class TestSegmentServiceMutations: result = SegmentService.update_segment(args, segment, document, dataset, mock_db.session) assert result is refreshed_segment - generate_summary.assert_called_once_with(segment, dataset, {"enable": True}) + generate_summary.assert_called_once_with(segment, dataset, {"enable": True}, session=mock_db.session) update_summary.assert_not_called() - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_delete_segment_removes_index_and_updates_document_word_count(self): segment = _make_segment(word_count=4, index_node_id="parent-node") @@ -972,7 +973,7 @@ class TestSegmentServiceAdditionalRegenerationBranches: assert segment.word_count == len("question") + len("new answer") assert document.word_count == 20 + (len("question") + len("new answer") - 8) vector_service.update_segment_vector.assert_not_called() - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_content_change_uses_answer_when_counting_tokens_for_qa_segments(self, account_context): segment = _make_segment(content="old", word_count=3) @@ -1009,7 +1010,7 @@ class TestSegmentServiceAdditionalRegenerationBranches: assert segment.tokens == 21 assert segment.word_count == len("new question") + len("new answer") vector_service.update_segment_vector.assert_called_once_with(["kw-1"], segment, dataset) - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_content_change_parent_child_uses_default_embedding_and_ignores_summary_failures( self, account_context @@ -1063,10 +1064,11 @@ class TestSegmentServiceAdditionalRegenerationBranches: dataset, embedding_model_instance, processing_rule, + mock_db.session, True, ) - update_summary.assert_called_once_with(segment, dataset, "new summary") - vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset) + update_summary.assert_called_once_with(segment, dataset, "new summary", session=mock_db.session) + vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset, mock_db.session) def test_update_segment_same_content_parent_child_marks_segment_error_for_non_high_quality_dataset( self, account_context diff --git a/api/tests/unit_tests/services/test_datasource_provider_service.py b/api/tests/unit_tests/services/test_datasource_provider_service.py index f374a294825..bd6891d846a 100644 --- a/api/tests/unit_tests/services/test_datasource_provider_service.py +++ b/api/tests/unit_tests/services/test_datasource_provider_service.py @@ -177,11 +177,11 @@ class TestDatasourceProviderService: def test_should_return_true_when_tenant_oauth_params_enabled(self, service, mock_db_session): mock_db_session.scalar.return_value = 1 - assert service.is_tenant_oauth_params_enabled("t1", make_id()) is True + assert service.is_tenant_oauth_params_enabled("t1", make_id(), session=mock_db_session) is True def test_should_return_false_when_tenant_oauth_params_disabled(self, service, mock_db_session): mock_db_session.scalar.return_value = 0 - assert service.is_tenant_oauth_params_enabled("t1", make_id()) is False + assert service.is_tenant_oauth_params_enabled("t1", make_id(), session=mock_db_session) is False # ----------------------------------------------------------------------- # remove_oauth_custom_client_params (lines 55-61) @@ -453,7 +453,7 @@ class TestDatasourceProviderService: tenant_params.client_params = {"k": "v"} mock_db_session.scalar.return_value = tenant_params with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): - result = service.get_tenant_oauth_client("t1", make_id(), mask=True) + result = service.get_tenant_oauth_client("t1", make_id(), mask=True, session=mock_db_session) assert result == {"k": "mask"} def test_should_return_decrypted_credentials_when_mask_is_false(self, service, mock_db_session): @@ -461,12 +461,12 @@ class TestDatasourceProviderService: tenant_params.client_params = {"k": "v"} mock_db_session.scalar.return_value = tenant_params with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): - result = service.get_tenant_oauth_client("t1", make_id(), mask=False) + result = service.get_tenant_oauth_client("t1", make_id(), mask=False, session=mock_db_session) assert result == {"k": "dec"} def test_should_return_none_when_no_tenant_oauth_config_exists(self, service, mock_db_session): mock_db_session.scalar.return_value = None - assert service.get_tenant_oauth_client("t1", make_id()) is None + assert service.get_tenant_oauth_client("t1", make_id(), session=mock_db_session) is None # ----------------------------------------------------------------------- # get_oauth_client (lines 423-457) @@ -657,7 +657,7 @@ class TestDatasourceProviderService: def test_should_return_empty_list_when_no_credentials_stored(self, service, mock_db_session): mock_db_session.scalars.return_value.all.return_value = [] - assert service.list_datasource_credentials("t1", "prov", "org/plug") == [] + assert service.list_datasource_credentials("t1", "prov", "org/plug", session=mock_db_session) == [] def test_should_return_masked_credentials_list_when_credentials_exist(self, service, mock_db_session): p = MagicMock(spec=DatasourceProvider) @@ -666,7 +666,7 @@ class TestDatasourceProviderService: p.is_default = False mock_db_session.scalars.return_value.all.return_value = [p] with patch.object(service, "extract_secret_variables", return_value=["sk"]): - result = service.list_datasource_credentials("t1", "prov", "org/plug") + result = service.list_datasource_credentials("t1", "prov", "org/plug", session=mock_db_session) assert len(result) == 1 # ----------------------------------------------------------------------- @@ -682,7 +682,9 @@ class TestDatasourceProviderService: mock_mgr.return_value.fetch_installed_datasource_providers.return_value = [ds] cred = {"credential": {"k": "v"}, "is_default": True} with patch.object(service, "list_datasource_credentials", return_value=[cred]): - results = service.get_all_datasource_credentials("t1") + session = MagicMock() + session.scalar.return_value = 0 + results = service.get_all_datasource_credentials("t1", session=session) assert len(results) == 1 def test_should_include_oauth_schema_for_hardcoded_plugin_ids(self, service, mock_db_session): @@ -707,7 +709,7 @@ class TestDatasourceProviderService: patch.object(service, "is_tenant_oauth_params_enabled", return_value=False), patch.object(service, "is_system_oauth_params_exist", return_value=False), ): - results = service.get_all_datasource_credentials("t1") + results = service.get_all_datasource_credentials("t1", session=mock_db_session) assert len(results) == 1 assert results[0]["oauth_schema"] is not None @@ -717,7 +719,7 @@ class TestDatasourceProviderService: def test_should_return_empty_list_when_no_real_credentials_exist(self, service, mock_db_session): mock_db_session.scalars.return_value.all.return_value = [] - assert service.get_real_datasource_credentials("t1", "prov", "org/plug") == [] + assert service.get_real_datasource_credentials("t1", "prov", "org/plug", session=mock_db_session) == [] def test_should_return_decrypted_credential_list_when_credentials_exist(self, service, mock_db_session): p = MagicMock(spec=DatasourceProvider) @@ -725,7 +727,7 @@ class TestDatasourceProviderService: p.encrypted_credentials = {"sk": "v"} mock_db_session.scalars.return_value.all.return_value = [p] with patch.object(service, "extract_secret_variables", return_value=["sk"]): - result = service.get_real_datasource_credentials("t1", "prov", "org/plug") + result = service.get_real_datasource_credentials("t1", "prov", "org/plug", session=mock_db_session) assert len(result) == 1 # ----------------------------------------------------------------------- @@ -788,11 +790,11 @@ class TestDatasourceProviderService: def test_should_delete_provider_and_commit_when_found(self, service, mock_db_session): p = MagicMock(spec=DatasourceProvider) mock_db_session.scalar.return_value = p - service.remove_datasource_credentials("t1", "id", "prov", "org/plug") + service.remove_datasource_credentials("t1", "id", "prov", "org/plug", session=mock_db_session) mock_db_session.delete.assert_called_once_with(p) def test_should_do_nothing_when_credential_not_found_on_remove(self, service, mock_db_session): """No error raised; no delete called when record doesn't exist (lines 994 branch).""" mock_db_session.scalar.return_value = None - service.remove_datasource_credentials("t1", "id", "prov", "org/plug") + service.remove_datasource_credentials("t1", "id", "prov", "org/plug", session=mock_db_session) mock_db_session.delete.assert_not_called() diff --git a/api/tests/unit_tests/services/test_external_dataset_service.py b/api/tests/unit_tests/services/test_external_dataset_service.py index dbb4627759c..9dff74f8dd5 100644 --- a/api/tests/unit_tests/services/test_external_dataset_service.py +++ b/api/tests/unit_tests/services/test_external_dataset_service.py @@ -145,7 +145,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_success_basic( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test successful retrieval of external knowledge APIs with pagination.""" # Arrange @@ -158,7 +158,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 5 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -170,11 +170,11 @@ class TestExternalDatasetServiceGetAPIs: assert result_total == 5 assert result_items[0].id == "api-0" assert result_items[4].id == "api-4" - mock_paginate.assert_called_once() + mock_paginate_query.assert_called_once() @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_with_search_filter( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test retrieval with search filter.""" # Arrange @@ -186,7 +186,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 1 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -200,14 +200,14 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_empty_results( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test retrieval with no results.""" # Arrange mock_pagination = MagicMock() mock_pagination.items = [] mock_pagination.total = 0 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -220,7 +220,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_large_result_set( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test retrieval with large result set.""" # Arrange @@ -229,7 +229,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis[:10] mock_pagination.total = 100 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -242,7 +242,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_pagination_last_page( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test last page pagination with partial results.""" # Arrange @@ -251,7 +251,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 100 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -264,7 +264,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_case_insensitive_search( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test case-insensitive search functionality.""" # Arrange @@ -276,7 +276,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 2 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -289,7 +289,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_special_characters_search( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test search with special characters.""" # Arrange @@ -298,7 +298,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 1 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -310,7 +310,7 @@ class TestExternalDatasetServiceGetAPIs: @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_max_per_page_limit( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test that max_per_page limit is enforced.""" # Arrange @@ -319,7 +319,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis mock_pagination.total = 1000 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -327,12 +327,12 @@ class TestExternalDatasetServiceGetAPIs: ) # Assert - call_args = mock_paginate.call_args + call_args = mock_paginate_query.call_args assert call_args.kwargs["max_per_page"] == 100 @patch("services.external_knowledge_service.paginate_query") def test_get_external_knowledge_apis_ordered_by_created_at_desc( - self, mock_paginate, factory: ExternalDatasetServiceTestDataFactory + self, mock_paginate_query, factory: ExternalDatasetServiceTestDataFactory ): """Test that results are ordered by created_at descending.""" # Arrange @@ -344,7 +344,7 @@ class TestExternalDatasetServiceGetAPIs: mock_pagination = MagicMock() mock_pagination.items = apis[::-1] # Reversed to simulate DESC order mock_pagination.total = 5 - mock_paginate.return_value = mock_pagination + mock_paginate_query.return_value = mock_pagination # Act result_items, result_total = ExternalDatasetService.get_external_knowledge_apis( @@ -437,13 +437,13 @@ class TestExternalDatasetServiceValidateAPIList: class TestExternalDatasetServiceCreateAPI: """Test create_external_knowledge_api operations.""" + @patch("services.external_knowledge_service.db") @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_success_full( - self, mock_check, factory: ExternalDatasetServiceTestDataFactory + self, mock_check, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test successful creation with all fields.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" user_id = "user-123" args = { @@ -453,7 +453,7 @@ class TestExternalDatasetServiceCreateAPI: } # Act - result = ExternalDatasetService.create_external_knowledge_api(tenant_id, user_id, args, mock_session) + result = ExternalDatasetService.create_external_knowledge_api(tenant_id, user_id, args, session=mock_db.session) # Assert assert result.name == "Test API" @@ -462,55 +462,63 @@ class TestExternalDatasetServiceCreateAPI: assert result.created_by == user_id assert result.updated_by == user_id mock_check.assert_called_once_with(args["settings"]) - mock_session.add.assert_called_once() - mock_session.commit.assert_called_once() + mock_db.session.add.assert_called_once() + mock_db.session.commit.assert_called_once() + @patch("services.external_knowledge_service.db") @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_minimal_fields( - self, mock_check, factory: ExternalDatasetServiceTestDataFactory + self, mock_check, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test creation with minimal required fields.""" # Arrange - mock_session = MagicMock() args = { "name": "Minimal API", "settings": {"endpoint": "https://api.example.com", "api_key": "key"}, } # Act - result = ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + result = ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) # Assert assert result.name == "Minimal API" assert result.description == "" - def test_create_external_knowledge_api_missing_settings(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_knowledge_api_missing_settings( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test creation fails when settings are missing.""" # Arrange - mock_session = MagicMock() args = {"name": "Test API", "description": "Test"} # Act & Assert with pytest.raises(ValueError, match="settings is required"): - ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) - def test_create_external_knowledge_api_none_settings(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_knowledge_api_none_settings(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test creation fails when settings are explicitly None.""" # Arrange - mock_session = MagicMock() args = {"name": "Test API", "settings": None} # Act & Assert with pytest.raises(ValueError, match="settings is required"): - ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) + @patch("services.external_knowledge_service.db") @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_settings_json_serialization( - self, mock_check, factory: ExternalDatasetServiceTestDataFactory + self, mock_check, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test that settings are properly JSON serialized.""" # Arrange - mock_session = MagicMock() settings = { "endpoint": "https://api.example.com", "api_key": "test-key", @@ -519,20 +527,22 @@ class TestExternalDatasetServiceCreateAPI: args = {"name": "Test API", "settings": settings} # Act - result = ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + result = ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) # Assert assert isinstance(result.settings, str) parsed_settings = json.loads(result.settings) assert parsed_settings == settings + @patch("services.external_knowledge_service.db") @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_unicode_handling( - self, mock_check, factory: ExternalDatasetServiceTestDataFactory + self, mock_check, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test proper handling of Unicode characters in name and description.""" # Arrange - mock_session = MagicMock() args = { "name": "测试API", "description": "テストの説明", @@ -540,19 +550,21 @@ class TestExternalDatasetServiceCreateAPI: } # Act - result = ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + result = ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) # Assert assert result.name == "测试API" assert result.description == "テストの説明" + @patch("services.external_knowledge_service.db") @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_long_description( - self, mock_check, factory: ExternalDatasetServiceTestDataFactory + self, mock_check, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test creation with very long description.""" # Arrange - mock_session = MagicMock() long_description = "A" * 1000 args = { "name": "Test API", @@ -561,7 +573,9 @@ class TestExternalDatasetServiceCreateAPI: } # Act - result = ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, mock_session) + result = ExternalDatasetService.create_external_knowledge_api( + "tenant-123", "user-123", args, session=mock_db.session + ) # Assert assert result.description == long_description @@ -824,43 +838,43 @@ class TestExternalDatasetServiceCheckEndpoint: class TestExternalDatasetServiceGetAPI: """Test get_external_knowledge_api operations.""" - def test_get_external_knowledge_api_success(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_get_external_knowledge_api_success(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test successful retrieval of external knowledge API.""" # Arrange - mock_session = MagicMock() api_id = "api-123" expected_api = factory.create_external_knowledge_api_mock(api_id=api_id) - mock_session.scalar.return_value = expected_api + mock_db.session.scalar.return_value = expected_api # Act tenant_id = "tenant-123" - result = ExternalDatasetService.get_external_knowledge_api(mock_session, api_id, tenant_id) + result = ExternalDatasetService.get_external_knowledge_api(api_id, tenant_id, session=mock_db.session) # Assert assert result.id == api_id - def test_get_external_knowledge_api_not_found(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_get_external_knowledge_api_not_found(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test error when API is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): - ExternalDatasetService.get_external_knowledge_api(mock_session, "nonexistent-id", "tenant-123") + ExternalDatasetService.get_external_knowledge_api("nonexistent-id", "tenant-123", session=mock_db.session) class TestExternalDatasetServiceUpdateAPI: """Test update_external_knowledge_api operations.""" @patch("services.external_knowledge_service.naive_utc_now") + @patch("services.external_knowledge_service.db") def test_update_external_knowledge_api_success_all_fields( - self, mock_now, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_now, factory: ExternalDatasetServiceTestDataFactory ): """Test successful update with all fields.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" user_id = "user-456" @@ -875,24 +889,26 @@ class TestExternalDatasetServiceUpdateAPI: "settings": {"endpoint": "https://new.example.com", "api_key": "new-key"}, } - mock_session.scalar.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act - result = ExternalDatasetService.update_external_knowledge_api(mock_session, tenant_id, user_id, api_id, args) + result = ExternalDatasetService.update_external_knowledge_api( + tenant_id, user_id, api_id, args, session=mock_db.session + ) # Assert assert result.name == "Updated API" assert result.description == "Updated description" assert result.updated_by == user_id assert result.updated_at == current_time - mock_session.commit.assert_called_once() + mock_db.session.commit.assert_called_once() + @patch("services.external_knowledge_service.db") def test_update_external_knowledge_api_preserve_hidden_api_key( - self, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test that hidden API key is preserved from existing settings.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" @@ -907,47 +923,51 @@ class TestExternalDatasetServiceUpdateAPI: "settings": {"endpoint": "https://api.example.com", "api_key": HIDDEN_VALUE}, } - mock_session.scalar.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act - result = ExternalDatasetService.update_external_knowledge_api(mock_session, tenant_id, "user-123", api_id, args) + result = ExternalDatasetService.update_external_knowledge_api( + tenant_id, "user-123", api_id, args, session=mock_db.session + ) # Assert settings = json.loads(result.settings) assert settings["api_key"] == "original-secret-key" - def test_update_external_knowledge_api_not_found(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_update_external_knowledge_api_not_found(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test error when API is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None args = {"name": "Updated API"} # Act & Assert with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.update_external_knowledge_api( - mock_session, "tenant-123", "user-123", "api-123", args + "tenant-123", "user-123", "api-123", args, session=mock_db.session ) - def test_update_external_knowledge_api_tenant_mismatch(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_update_external_knowledge_api_tenant_mismatch( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test error when tenant ID doesn't match.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None args = {"name": "Updated API"} # Act & Assert with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.update_external_knowledge_api( - mock_session, "wrong-tenant", "user-123", "api-123", args + "wrong-tenant", "user-123", "api-123", args, session=mock_db.session ) - def test_update_external_knowledge_api_name_only(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_update_external_knowledge_api_name_only(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test updating only the name field.""" # Arrange - mock_session = MagicMock() existing_api = factory.create_external_knowledge_api_mock( description="Original description", settings={"endpoint": "https://api.example.com", "api_key": "key"}, @@ -955,11 +975,11 @@ class TestExternalDatasetServiceUpdateAPI: args = {"name": "New Name Only"} - mock_session.scalar.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act result = ExternalDatasetService.update_external_knowledge_api( - mock_session, "tenant-123", "user-123", "api-123", args + "tenant-123", "user-123", "api-123", args, session=mock_db.session ) # Assert @@ -969,92 +989,104 @@ class TestExternalDatasetServiceUpdateAPI: class TestExternalDatasetServiceDeleteAPI: """Test delete_external_knowledge_api operations.""" - def test_delete_external_knowledge_api_success(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_delete_external_knowledge_api_success(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test successful deletion of external knowledge API.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" existing_api = factory.create_external_knowledge_api_mock(api_id=api_id, tenant_id=tenant_id) - mock_session.scalar.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act - ExternalDatasetService.delete_external_knowledge_api(mock_session, tenant_id, api_id) + ExternalDatasetService.delete_external_knowledge_api(tenant_id, api_id, session=mock_db.session) # Assert - mock_session.delete.assert_called_once_with(existing_api) - mock_session.commit.assert_called_once() + mock_db.session.delete.assert_called_once_with(existing_api) + mock_db.session.commit.assert_called_once() - def test_delete_external_knowledge_api_not_found(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_delete_external_knowledge_api_not_found(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test error when API is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): - ExternalDatasetService.delete_external_knowledge_api(mock_session, "tenant-123", "api-123") + ExternalDatasetService.delete_external_knowledge_api("tenant-123", "api-123", session=mock_db.session) - def test_delete_external_knowledge_api_tenant_mismatch(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_delete_external_knowledge_api_tenant_mismatch( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test error when tenant ID doesn't match.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): - ExternalDatasetService.delete_external_knowledge_api(mock_session, "wrong-tenant", "api-123") + ExternalDatasetService.delete_external_knowledge_api("wrong-tenant", "api-123", session=mock_db.session) class TestExternalDatasetServiceAPIUseCheck: """Test external_knowledge_api_use_check operations.""" - def test_external_knowledge_api_use_check_in_use_single(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_external_knowledge_api_use_check_in_use_single( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test API use check when API has one binding.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" - mock_session.scalar.return_value = 1 + mock_db.session.scalar.return_value = 1 # Act - in_use, count = ExternalDatasetService.external_knowledge_api_use_check(mock_session, api_id, tenant_id) + in_use, count = ExternalDatasetService.external_knowledge_api_use_check( + api_id, tenant_id, session=mock_db.session + ) # Assert assert in_use is True assert count == 1 - assert "tenant_id" in str(mock_session.scalar.call_args.args[0]) + assert "tenant_id" in str(mock_db.session.scalar.call_args.args[0]) - def test_external_knowledge_api_use_check_in_use_multiple(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_external_knowledge_api_use_check_in_use_multiple( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test API use check with multiple bindings.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" - mock_session.scalar.return_value = 10 + mock_db.session.scalar.return_value = 10 # Act - in_use, count = ExternalDatasetService.external_knowledge_api_use_check(mock_session, api_id, tenant_id) + in_use, count = ExternalDatasetService.external_knowledge_api_use_check( + api_id, tenant_id, session=mock_db.session + ) # Assert assert in_use is True assert count == 10 - def test_external_knowledge_api_use_check_not_in_use(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_external_knowledge_api_use_check_not_in_use(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test API use check when API is not in use.""" # Arrange - mock_session = MagicMock() api_id = "api-123" tenant_id = "tenant-123" - mock_session.scalar.return_value = 0 + mock_db.session.scalar.return_value = 0 # Act - in_use, count = ExternalDatasetService.external_knowledge_api_use_check(mock_session, api_id, tenant_id) + in_use, count = ExternalDatasetService.external_knowledge_api_use_check( + api_id, tenant_id, session=mock_db.session + ) # Assert assert in_use is False @@ -1064,46 +1096,48 @@ class TestExternalDatasetServiceAPIUseCheck: class TestExternalDatasetServiceGetBinding: """Test get_external_knowledge_binding_with_dataset_id operations.""" - def test_get_external_knowledge_binding_success(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_get_external_knowledge_binding_success(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test successful retrieval of external knowledge binding.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" dataset_id = "dataset-123" expected_binding = factory.create_external_knowledge_binding_mock(tenant_id=tenant_id, dataset_id=dataset_id) - mock_session.scalar.return_value = expected_binding + mock_db.session.scalar.return_value = expected_binding # Act result = ExternalDatasetService.get_external_knowledge_binding_with_dataset_id( - mock_session, tenant_id, dataset_id + tenant_id, dataset_id, session=mock_db.session ) # Assert assert result.dataset_id == dataset_id assert result.tenant_id == tenant_id - def test_get_external_knowledge_binding_not_found(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_get_external_knowledge_binding_not_found(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test error when binding is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="external knowledge binding not found"): ExternalDatasetService.get_external_knowledge_binding_with_dataset_id( - mock_session, "tenant-123", "dataset-123" + "tenant-123", "dataset-123", session=mock_db.session ) class TestExternalDatasetServiceDocumentValidate: """Test document_create_args_validate operations.""" - def test_document_create_args_validate_success_all_params(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_document_create_args_validate_success_all_params( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test successful validation with all required parameters.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" api_id = "api-123" @@ -1117,17 +1151,21 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(api_id=api_id, settings=[settings]) - mock_session.scalar.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {"param1": "value1", "param2": "value2"} # Act & Assert - should not raise - ExternalDatasetService.document_create_args_validate(mock_session, tenant_id, api_id, process_parameter) + ExternalDatasetService.document_create_args_validate( + tenant_id, api_id, process_parameter, session=mock_db.session + ) - def test_document_create_args_validate_missing_required_param(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_document_create_args_validate_missing_required_param( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test validation fails when required parameter is missing.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" api_id = "api-123" @@ -1135,42 +1173,46 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(api_id=api_id, settings=[settings]) - mock_session.scalar.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {} # Act & Assert with pytest.raises(ValueError, match="required_param is required"): - ExternalDatasetService.document_create_args_validate(mock_session, tenant_id, api_id, process_parameter) + ExternalDatasetService.document_create_args_validate( + tenant_id, api_id, process_parameter, session=mock_db.session + ) - def test_document_create_args_validate_api_not_found(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_document_create_args_validate_api_not_found(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails when API is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): - ExternalDatasetService.document_create_args_validate(mock_session, "tenant-123", "api-123", {}) + ExternalDatasetService.document_create_args_validate("tenant-123", "api-123", {}, session=mock_db.session) - def test_document_create_args_validate_no_custom_parameters(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_document_create_args_validate_no_custom_parameters( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test validation succeeds when no custom parameters defined.""" # Arrange - mock_session = MagicMock() settings = {} api = factory.create_external_knowledge_api_mock(settings=[settings]) - mock_session.scalar.return_value = api + mock_db.session.scalar.return_value = api # Act & Assert - should not raise - ExternalDatasetService.document_create_args_validate(mock_session, "tenant-123", "api-123", {}) + ExternalDatasetService.document_create_args_validate("tenant-123", "api-123", {}, session=mock_db.session) + @patch("services.external_knowledge_service.db") def test_document_create_args_validate_optional_params_not_required( - self, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test that optional parameters don't cause validation failure.""" # Arrange - mock_session = MagicMock() settings = { "document_process_setting": [ {"name": "required_param", "required": True}, @@ -1180,12 +1222,14 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(settings=[settings]) - mock_session.scalar.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {"required_param": "value"} # Act & Assert - should not raise - ExternalDatasetService.document_create_args_validate(mock_session, "tenant-123", "api-123", process_parameter) + ExternalDatasetService.document_create_args_validate( + "tenant-123", "api-123", process_parameter, session=mock_db.session + ) class TestExternalDatasetServiceProcessAPI: @@ -1475,10 +1519,10 @@ class TestExternalDatasetServiceGetSettings: class TestExternalDatasetServiceCreateDataset: """Test create_external_dataset operations.""" - def test_create_external_dataset_success_full(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_dataset_success_full(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test successful creation of external dataset with all fields.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" user_id = "user-123" args = { @@ -1491,84 +1535,90 @@ class TestExternalDatasetServiceCreateDataset: api = factory.create_external_knowledge_api_mock(api_id="api-123") - mock_session.scalar.side_effect = [None, api] + mock_db.session.scalar.side_effect = [None, api] # Act - result = ExternalDatasetService.create_external_dataset(tenant_id, user_id, args, mock_session) + result = ExternalDatasetService.create_external_dataset(tenant_id, user_id, args, session=mock_db.session) # Assert assert result.name == "Test External Dataset" assert result.description == "Comprehensive test description" assert result.provider == "external" assert result.created_by == user_id - mock_session.add.assert_called() - mock_session.commit.assert_called_once() + mock_db.session.add.assert_called() + mock_db.session.commit.assert_called_once() - def test_create_external_dataset_duplicate_name_error(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_dataset_duplicate_name_error( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test error when dataset name already exists.""" # Arrange - mock_session = MagicMock() existing_dataset = factory.create_dataset_mock(name="Duplicate Dataset") - mock_session.scalar.return_value = existing_dataset + mock_db.session.scalar.return_value = existing_dataset args = {"name": "Duplicate Dataset"} # Act & Assert with pytest.raises(DatasetNameDuplicateError): - ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, session=mock_db.session) - def test_create_external_dataset_api_not_found_error(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_dataset_api_not_found_error(self, mock_db, factory: ExternalDatasetServiceTestDataFactory): """Test error when external knowledge API is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.side_effect = [None, None] + mock_db.session.scalar.side_effect = [None, None] args = {"name": "Test Dataset", "external_knowledge_api_id": "nonexistent-api"} # Act & Assert with pytest.raises(ValueError, match="api template not found"): - ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, session=mock_db.session) - def test_create_external_dataset_missing_knowledge_id_error(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_dataset_missing_knowledge_id_error( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test error when external_knowledge_id is missing.""" # Arrange - mock_session = MagicMock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [None, api] + mock_db.session.scalar.side_effect = [None, api] args = {"name": "Test Dataset", "external_knowledge_api_id": "api-123"} # Act & Assert with pytest.raises(ValueError, match="external_knowledge_id is required"): - ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, session=mock_db.session) - def test_create_external_dataset_missing_api_id_error(self, factory: ExternalDatasetServiceTestDataFactory): + @patch("services.external_knowledge_service.db") + def test_create_external_dataset_missing_api_id_error( + self, mock_db, factory: ExternalDatasetServiceTestDataFactory + ): """Test error when external_knowledge_api_id is missing.""" # Arrange - mock_session = MagicMock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [None, api] + mock_db.session.scalar.side_effect = [None, api] args = {"name": "Test Dataset", "external_knowledge_id": "knowledge-123"} # Act & Assert with pytest.raises(ValueError, match="external_knowledge_api_id is required"): - ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, mock_session) + ExternalDatasetService.create_external_dataset("tenant-123", "user-123", args, session=mock_db.session) class TestExternalDatasetServiceFetchRetrieval: """Test fetch_external_knowledge_retrieval operations.""" @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_success_with_results( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory ): """Test successful external knowledge retrieval with results.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" dataset_id = "dataset-123" query = "test query for retrieval" @@ -1578,7 +1628,7 @@ class TestExternalDatasetServiceFetchRetrieval: ) api = factory.create_external_knowledge_api_mock(api_id="api-123") - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1594,7 +1644,11 @@ class TestExternalDatasetServiceFetchRetrieval: # Act result = ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, tenant_id, dataset_id, query, external_retrieval_parameters + tenant_id, + dataset_id, + query, + external_retrieval_parameters, + session=mock_db.session, ) # Assert @@ -1602,46 +1656,46 @@ class TestExternalDatasetServiceFetchRetrieval: assert result[0]["content"] == "result 1" assert result[1]["score"] == 0.8 + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_binding_not_found_error( - self, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test error when external knowledge binding is not found.""" # Arrange - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ExternalKnowledgeRetrievalError, match="external knowledge binding not found"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {} + "tenant-123", "dataset-123", "query", {}, session=mock_db.session ) + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_cross_tenant_api_template_error( - self, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, factory: ExternalDatasetServiceTestDataFactory ): """Test error when a binding points to an API template outside the dataset tenant.""" # Arrange - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() - mock_session.scalar.side_effect = [binding, None] + mock_db.session.scalar.side_effect = [binding, None] # Act & Assert with pytest.raises(ExternalKnowledgeRetrievalError, match="external api template not found"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {} + "tenant-123", "dataset-123", "query", {}, session=mock_db.session ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_empty_results( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory ): """Test retrieval with empty results.""" # Arrange - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1650,23 +1704,27 @@ class TestExternalDatasetServiceFetchRetrieval: # Act result = ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) # Assert assert len(result) == 0 @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_with_score_threshold( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory ): """Test retrieval with score threshold enabled.""" # Arrange - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1681,7 +1739,11 @@ class TestExternalDatasetServiceFetchRetrieval: # Act result = ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", external_retrieval_parameters + "tenant-123", + "dataset-123", + "query", + external_retrieval_parameters, + session=mock_db.session, ) # Assert @@ -1691,16 +1753,16 @@ class TestExternalDatasetServiceFetchRetrieval: assert call_args.params["retrieval_setting"]["score_threshold"] == 0.75 @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_non_200_status_raises_exception( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory ): """Test that non-200 status code raises Exception with response text.""" # Arrange - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 500 @@ -1710,7 +1772,11 @@ class TestExternalDatasetServiceFetchRetrieval: # Act & Assert with pytest.raises(ExternalKnowledgeRetrievalError, match="Internal Server Error: Database connection failed"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) @pytest.mark.parametrize( @@ -1727,12 +1793,12 @@ class TestExternalDatasetServiceFetchRetrieval: ], ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_various_error_status_codes( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory, status_code, error_message + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory, status_code, error_message ): """Test that various error status codes raise exceptions with response text.""" # Arrange - mock_session = MagicMock() tenant_id = "tenant-123" dataset_id = "dataset-123" @@ -1741,7 +1807,7 @@ class TestExternalDatasetServiceFetchRetrieval: ) api = factory.create_external_knowledge_api_mock(api_id="api-123") - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = status_code @@ -1751,20 +1817,20 @@ class TestExternalDatasetServiceFetchRetrieval: # Act & Assert with pytest.raises(ExternalKnowledgeRetrievalError, match=re.escape(error_message)): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, tenant_id, dataset_id, "query", {"top_k": 5} + tenant_id, dataset_id, "query", {"top_k": 5}, session=mock_db.session ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.external_knowledge_service.db") def test_fetch_external_knowledge_retrieval_empty_response_text( - self, mock_process, factory: ExternalDatasetServiceTestDataFactory + self, mock_db, mock_process, factory: ExternalDatasetServiceTestDataFactory ): """Test exception with empty response text.""" # Arrange - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 503 @@ -1774,17 +1840,21 @@ class TestExternalDatasetServiceFetchRetrieval: # Act & Assert with pytest.raises(ExternalKnowledgeRetrievalError): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") - def test_fetch_external_knowledge_retrieval_invalid_json_response(self, mock_process, factory): + @patch("services.external_knowledge_service.db") + def test_fetch_external_knowledge_retrieval_invalid_json_response(self, mock_db, mock_process, factory): """Test malformed JSON success responses are normalized to external retrieval errors.""" - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1793,17 +1863,21 @@ class TestExternalDatasetServiceFetchRetrieval: with pytest.raises(ExternalKnowledgeRetrievalError, match="invalid external knowledge response"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") - def test_fetch_external_knowledge_retrieval_invalid_success_payload_shape(self, mock_process, factory): + @patch("services.external_knowledge_service.db") + def test_fetch_external_knowledge_retrieval_invalid_success_payload_shape(self, mock_db, mock_process, factory): """Test malformed success payload shapes are normalized to external retrieval errors.""" - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1812,17 +1886,21 @@ class TestExternalDatasetServiceFetchRetrieval: with pytest.raises(ExternalKnowledgeRetrievalError, match="invalid external knowledge response"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") - def test_fetch_external_knowledge_retrieval_invalid_records_shape(self, mock_process, factory): + @patch("services.external_knowledge_service.db") + def test_fetch_external_knowledge_retrieval_invalid_records_shape(self, mock_db, mock_process, factory): """Test non-list records payloads are normalized to external retrieval errors.""" - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1831,20 +1909,28 @@ class TestExternalDatasetServiceFetchRetrieval: with pytest.raises(ExternalKnowledgeRetrievalError, match="invalid external knowledge response"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") - def test_fetch_external_knowledge_retrieval_wraps_transport_errors(self, mock_process, factory): + @patch("services.external_knowledge_service.db") + def test_fetch_external_knowledge_retrieval_wraps_transport_errors(self, mock_db, mock_process, factory): """Test transport/runtime failures are normalized to external retrieval errors.""" - mock_session = MagicMock() binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_session.scalar.side_effect = [binding, api] + mock_db.session.scalar.side_effect = [binding, api] mock_process.side_effect = RuntimeError("connection reset by peer") with pytest.raises(ExternalKnowledgeRetrievalError, match="connection reset by peer"): ExternalDatasetService.fetch_external_knowledge_retrieval( - mock_session, "tenant-123", "dataset-123", "query", {"top_k": 5} + "tenant-123", + "dataset-123", + "query", + {"top_k": 5}, + session=mock_db.session, ) diff --git a/api/tests/unit_tests/services/test_file_request_service.py b/api/tests/unit_tests/services/test_file_request_service.py index f5d1a59c17e..57abdeaa322 100644 --- a/api/tests/unit_tests/services/test_file_request_service.py +++ b/api/tests/unit_tests/services/test_file_request_service.py @@ -30,8 +30,9 @@ def test_request_download_url_builds_file_under_bound_scope( patch("services.file_request_service.bind_file_access_scope", return_value=nullcontext()) as bind_scope, patch.object(service, "_build_file", return_value=fake_file) as build_file, patch( - "services.file_request_service.file_helpers.resolve_file_url", return_value="https://files.example.com/x" - ), + "services.file_request_service.file_helpers.resolve_file_url", + return_value="https://files.example.com/x", + ) as resolve_file_url, ): result = service.request_download_url( tenant_id="tenant-1", @@ -51,6 +52,7 @@ def test_request_download_url_builds_file_under_bound_scope( build_file.assert_called_once_with( mapping={"transfer_method": "tool_file", "reference": reference}, tenant_id="tenant-1" ) + resolve_file_url.assert_called_once_with(fake_file, for_external=True) assert result.filename == "report.pdf" assert result.mime_type == "application/pdf" assert result.size == 123 diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index b81fb823949..41b86fda0cb 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -377,7 +377,7 @@ class TestFileService: def test_get_upload_files_by_ids_empty(self): session = MagicMock() - result = FileService.get_upload_files_by_ids(session, "tenant_id", []) + result = FileService.get_upload_files_by_ids("tenant_id", [], session=session) assert result == {} def test_get_upload_files_by_ids(self): @@ -387,7 +387,9 @@ class TestFileService: session = MagicMock() session.scalars().all.return_value = [upload_file] - result = FileService.get_upload_files_by_ids(session, "tenant_id", ["550e8400-e29b-41d4-a716-446655440000"]) + result = FileService.get_upload_files_by_ids( + "tenant_id", ["550e8400-e29b-41d4-a716-446655440000"], session=session + ) assert result["550e8400-e29b-41d4-a716-446655440000"] == upload_file def test_sanitize_zip_entry_name(self): diff --git a/api/tests/unit_tests/services/test_message_service.py b/api/tests/unit_tests/services/test_message_service.py index 6588c8a8de6..13f340e9f4a 100644 --- a/api/tests/unit_tests/services/test_message_service.py +++ b/api/tests/unit_tests/services/test_message_service.py @@ -102,6 +102,7 @@ class TestMessageServicePaginationByFirstId: conversation_id="conv-001", first_id=None, limit=10, + session=MagicMock(), ) # Assert @@ -124,6 +125,7 @@ class TestMessageServicePaginationByFirstId: conversation_id="", first_id=None, limit=10, + session=MagicMock(), ) # Assert @@ -166,6 +168,7 @@ class TestMessageServicePaginationByFirstId: first_id=None, limit=10, order="desc", + session=mock_db.session, ) # Assert @@ -209,6 +212,7 @@ class TestMessageServicePaginationByFirstId: first_id=None, limit=10, order="asc", + session=mock_db.session, ) # Assert @@ -258,6 +262,7 @@ class TestMessageServicePaginationByFirstId: first_id="msg-005", limit=10, order="desc", + session=mock_db.session, ) # Assert @@ -288,6 +293,7 @@ class TestMessageServicePaginationByFirstId: conversation_id="conv-001", first_id="nonexistent-msg", limit=10, + session=mock_db.session, ) # Test 07: Has_more flag when results exceed limit @@ -323,6 +329,7 @@ class TestMessageServicePaginationByFirstId: conversation_id="conv-001", first_id=None, limit=10, + session=mock_db.session, ) # Assert @@ -353,6 +360,7 @@ class TestMessageServicePaginationByFirstId: conversation_id="conv-001", first_id=None, limit=10, + session=mock_db.session, ) # Assert @@ -389,6 +397,7 @@ class TestMessageServicePaginationByLastId: user=None, last_id=None, limit=10, + session=MagicMock(), ) # Assert @@ -421,6 +430,7 @@ class TestMessageServicePaginationByLastId: user=user, last_id=None, limit=10, + session=mock_db.session, ) # Assert @@ -459,6 +469,7 @@ class TestMessageServicePaginationByLastId: user=user, last_id="msg-005", limit=10, + session=mock_db.session, ) # Assert @@ -482,6 +493,7 @@ class TestMessageServicePaginationByLastId: user=user, last_id="nonexistent-msg", limit=10, + session=mock_db.session, ) # Test 13: Pagination with conversation_id filter @@ -516,6 +528,7 @@ class TestMessageServicePaginationByLastId: last_id=None, limit=10, conversation_id="conv-001", + session=mock_db.session, ) # Assert @@ -546,6 +559,7 @@ class TestMessageServicePaginationByLastId: last_id=None, limit=10, include_ids=["msg-001", "msg-003"], + session=mock_db.session, ) # Assert @@ -578,6 +592,7 @@ class TestMessageServicePaginationByLastId: user=user, last_id=None, limit=10, + session=mock_db.session, ) # Assert @@ -680,8 +695,8 @@ class TestMessageServiceGetMessage: mock_db.session.scalar.return_value = message - # Act - result = MessageService.get_message(app_model=app, user=user, message_id="msg-123") + # Act, + result = MessageService.get_message(app_model=app, user=user, message_id="msg-123", session=mock_db.session) # Assert assert result == message @@ -700,8 +715,8 @@ class TestMessageServiceGetMessage: mock_db.session.scalar.return_value = message - # Act - result = MessageService.get_message(app_model=app, user=user, message_id="msg-123") + # Act, + result = MessageService.get_message(app_model=app, user=user, message_id="msg-123", session=mock_db.session) # Assert assert result == message @@ -718,7 +733,7 @@ class TestMessageServiceGetMessage: # Act & Assert with pytest.raises(MessageNotExistsError): - MessageService.get_message(app_model=app, user=user, message_id="msg-123") + MessageService.get_message(app_model=app, user=user, message_id="msg-123", session=mock_db.session) class TestMessageServiceFeedback: @@ -748,6 +763,7 @@ class TestMessageServiceFeedback: user=user, rating=FeedbackRating.LIKE, content="Good answer", + session=mock_db.session, ) # Assert @@ -780,6 +796,7 @@ class TestMessageServiceFeedback: user=user, rating=FeedbackRating.DISLIKE, content="Bad answer", + session=mock_db.session, ) # Assert @@ -808,6 +825,7 @@ class TestMessageServiceFeedback: user=user, rating=None, content=None, + session=mock_db.session, ) # Assert @@ -826,8 +844,8 @@ class TestMessageServiceFeedback: mock_db.session.scalars.return_value.all.return_value = [feedback] - # Act - result = MessageService.get_all_messages_feedbacks(app_model=app, page=1, limit=10) + # Act, + result = MessageService.get_all_messages_feedbacks(app_model=app, page=1, limit=10, session=mock_db.session) # Assert assert result == [{"id": "fb-1"}] @@ -846,7 +864,11 @@ class TestMessageServiceSuggestedQuestions: app = factory.create_app_mock() with pytest.raises(ValueError, match="user cannot be None"): MessageService.get_suggested_questions_after_answer( - app_model=app, user=None, message_id="msg-123", invoke_from=MagicMock() + app_model=app, + user=None, + message_id="msg-123", + invoke_from=MagicMock(), + session=MagicMock(), ) # Test 28: get_suggested_questions_after_answer - Advanced Chat success @@ -890,7 +912,11 @@ class TestMessageServiceSuggestedQuestions: # Act result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=user, message_id="msg-123", invoke_from=InvokeFrom.WEB_APP + app_model=app, + user=user, + message_id="msg-123", + invoke_from=InvokeFrom.WEB_APP, + session=MagicMock(), ) # Assert @@ -938,7 +964,11 @@ class TestMessageServiceSuggestedQuestions: # Act result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock() + app_model=app, + user=user, + message_id="msg-123", + invoke_from=MagicMock(), + session=mock_db.session, ) # Assert @@ -996,6 +1026,7 @@ class TestMessageServiceSuggestedQuestions: user=user, message_id="msg-123", invoke_from=InvokeFrom.WEB_APP, + session=mock_db.session, ) assert result == ["Q1?"] @@ -1059,7 +1090,11 @@ class TestMessageServiceSuggestedQuestions: mock_llm_gen.generate_suggested_questions_after_answer.return_value = ["Q1?"] result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock() + app_model=app, + user=user, + message_id="msg-123", + invoke_from=MagicMock(), + session=mock_db.session, ) assert result == ["Q1?"] @@ -1168,7 +1203,11 @@ class TestMessageServiceSuggestedQuestions: mock_llm_gen.generate_suggested_questions_after_answer.return_value = ["Q1?"] result = MessageService.get_suggested_questions_after_answer( - app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock() + app_model=app, + user=user, + message_id="msg-123", + invoke_from=MagicMock(), + session=mock_db.session, ) assert result == ["Q1?"] @@ -1209,5 +1248,9 @@ class TestMessageServiceSuggestedQuestions: # Act & Assert with pytest.raises(SuggestedQuestionsAfterAnswerDisabledError): MessageService.get_suggested_questions_after_answer( - app_model=app, user=user, message_id="msg-123", invoke_from=MagicMock() + app_model=app, + user=user, + message_id="msg-123", + invoke_from=MagicMock(), + session=MagicMock(), ) diff --git a/api/tests/unit_tests/services/test_metadata_bug_complete.py b/api/tests/unit_tests/services/test_metadata_bug_complete.py index 6792243e9d0..00f16f75ac0 100644 --- a/api/tests/unit_tests/services/test_metadata_bug_complete.py +++ b/api/tests/unit_tests/services/test_metadata_bug_complete.py @@ -48,14 +48,14 @@ class TestMetadataBugCompleteValidation: account = _make_account() # Should crash with TypeError with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata(Mock(), "dataset-123", mock_metadata_args, account, "tenant-123") + MetadataService.create_metadata("dataset-123", mock_metadata_args, account, "tenant-123", session=Mock()) # Test update method as well account = _make_account() none_name = cast(str, None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): MetadataService.update_metadata_name( - Mock(), "dataset-123", "metadata-456", none_name, account, "tenant-123" + "dataset-123", "metadata-456", none_name, account, "tenant-123", session=Mock() ) def test_3_database_constraints_verification(self) -> None: @@ -99,7 +99,7 @@ class TestMetadataBugCompleteValidation: account = _make_account() with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata(Mock(), "dataset-123", mock_metadata_args, account, "tenant-123") + MetadataService.create_metadata("dataset-123", mock_metadata_args, account, "tenant-123", session=Mock()) def test_7_end_to_end_validation_layers(self) -> None: """Test all validation layers work together correctly.""" diff --git a/api/tests/unit_tests/services/test_metadata_nullable_bug.py b/api/tests/unit_tests/services/test_metadata_nullable_bug.py index ae93fe5ef51..cfd3d034df2 100644 --- a/api/tests/unit_tests/services/test_metadata_nullable_bug.py +++ b/api/tests/unit_tests/services/test_metadata_nullable_bug.py @@ -37,7 +37,7 @@ class TestMetadataNullableBug: account = _make_account() # This should crash with TypeError when calling len(None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata(Mock(), "dataset-123", mock_metadata_args, account, "tenant-123") + MetadataService.create_metadata("dataset-123", mock_metadata_args, account, "tenant-123", session=Mock()) def test_metadata_service_update_with_none_name_crashes(self) -> None: """Test that MetadataService.update_metadata_name crashes when name is None.""" @@ -46,7 +46,7 @@ class TestMetadataNullableBug: # This should crash with TypeError when calling len(None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): MetadataService.update_metadata_name( - Mock(), "dataset-123", "metadata-456", none_name, account, "tenant-123" + "dataset-123", "metadata-456", none_name, account, "tenant-123", session=Mock() ) def test_api_layer_now_uses_pydantic_validation(self) -> None: diff --git a/api/tests/unit_tests/services/test_model_load_balancing_service.py b/api/tests/unit_tests/services/test_model_load_balancing_service.py index 827567f1afe..743e6e797a3 100644 --- a/api/tests/unit_tests/services/test_model_load_balancing_service.py +++ b/api/tests/unit_tests/services/test_model_load_balancing_service.py @@ -80,9 +80,9 @@ def service(mocker: MockerFixture) -> ModelLoadBalancingService: @pytest.fixture -def mock_db(mocker: MockerFixture) -> MagicMock: +def mock_db() -> MagicMock: # Arrange - mocked_db = mocker.patch("services.model_load_balancing_service.db") + mocked_db = MagicMock() mocked_db.session = MagicMock() return mocked_db @@ -159,7 +159,7 @@ def test_get_load_balancing_configs_should_raise_value_error_when_provider_missi # Act + Assert with pytest.raises(ValueError, match="Provider openai does not exist"): - service.get_load_balancing_configs("tenant-1", "openai", "gpt-4o-mini", ModelType.LLM) + service.get_load_balancing_configs("tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, session=MagicMock()) def test_get_load_balancing_configs_should_insert_inherit_config_when_missing_for_custom_provider( @@ -201,6 +201,7 @@ def test_get_load_balancing_configs_should_insert_inherit_config_when_missing_fo "openai", "gpt-4o-mini", ModelType.LLM, + session=mock_db.session, ) # Assert @@ -263,6 +264,7 @@ def test_get_load_balancing_configs_should_reorder_existing_inherit_and_tolerate "gpt-4o-mini", ModelType.LLM, config_from="predefined-model", + session=mock_db.session, ) # Assert @@ -282,7 +284,9 @@ def test_get_load_balancing_config_should_raise_value_error_when_provider_missin # Act + Assert with pytest.raises(ValueError, match="Provider openai does not exist"): - service.get_load_balancing_config("tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1") + service.get_load_balancing_config( + "tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1", session=MagicMock() + ) def test_get_load_balancing_config_should_return_none_when_config_not_found( @@ -295,7 +299,9 @@ def test_get_load_balancing_config_should_return_none_when_config_not_found( mock_db.session.scalar.return_value = None # Act - result = service.get_load_balancing_config("tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1") + result = service.get_load_balancing_config( + "tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1", session=mock_db.session + ) # Assert assert result is None @@ -315,7 +321,9 @@ def test_get_load_balancing_config_should_return_obfuscated_payload_when_config_ mock_db.session.scalar.return_value = config # Act - result = service.get_load_balancing_config("tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1") + result = service.get_load_balancing_config( + "tenant-1", "openai", "gpt-4o-mini", ModelType.LLM, "cfg-1", session=mock_db.session + ) # Assert assert result == { @@ -334,7 +342,9 @@ def test_init_inherit_config_should_create_and_persist_inherit_configuration( model_type = ModelType.LLM # Act - inherit_config = service._init_inherit_config("tenant-1", "openai", "gpt-4o-mini", model_type) + inherit_config = service._init_inherit_config( + "tenant-1", "openai", "gpt-4o-mini", model_type, session=mock_db.session + ) # Assert assert inherit_config.tenant_id == "tenant-1" @@ -361,6 +371,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_provider_mi ModelType.LLM, [], "custom-model", + session=MagicMock(), ) @@ -380,6 +391,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_configs_is_ ModelType.LLM, cast(list[dict[str, object]], "invalid-configs"), "custom-model", + session=MagicMock(), ) @@ -401,6 +413,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_config_item ModelType.LLM, cast(list[dict[str, object]], ["bad-item"]), "custom-model", + session=mock_db.session, ) @@ -423,6 +436,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_credential_ ModelType.LLM, [{"credential_id": "cred-1", "enabled": True}], "predefined-model", + session=mock_db.session, ) @@ -444,6 +458,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_name_or_ena ModelType.LLM, [{"enabled": True}], "custom-model", + session=mock_db.session, ) with pytest.raises(ValueError, match="Invalid load balancing config enabled"): @@ -454,6 +469,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_name_or_ena ModelType.LLM, [{"name": "cfg-without-enabled"}], "custom-model", + session=mock_db.session, ) @@ -476,6 +492,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_existing_co ModelType.LLM, [{"id": "cfg-2", "name": "invalid", "enabled": True}], "custom-model", + session=mock_db.session, ) @@ -498,6 +515,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_credentials ModelType.LLM, [{"id": "cfg-1", "name": "new", "enabled": True, "credentials": "bad"}], "custom-model", + session=mock_db.session, ) with pytest.raises(ValueError, match="Invalid load balancing config credentials"): @@ -508,6 +526,7 @@ def test_update_load_balancing_configs_should_raise_value_error_when_credentials ModelType.LLM, [{"name": "new-config", "enabled": True, "credentials": "bad"}], "custom-model", + session=mock_db.session, ) @@ -548,6 +567,7 @@ def test_update_load_balancing_configs_should_update_existing_create_new_and_del {"name": "new-config", "enabled": True, "credentials": {"api_key": "plain"}}, ], "custom-model", + session=mock_db.session, ) # Assert @@ -579,6 +599,7 @@ def test_update_load_balancing_configs_should_raise_value_error_for_invalid_new_ ModelType.LLM, [{"name": "__inherit__", "enabled": True, "credentials": {"api_key": "x"}}], "custom-model", + session=mock_db.session, ) with pytest.raises(ValueError, match="Invalid load balancing config credentials"): @@ -589,6 +610,7 @@ def test_update_load_balancing_configs_should_raise_value_error_for_invalid_new_ ModelType.LLM, [{"name": "new", "enabled": True}], "custom-model", + session=mock_db.session, ) @@ -611,6 +633,7 @@ def test_update_load_balancing_configs_should_create_from_existing_provider_cred ModelType.LLM, [{"credential_id": "cred-1", "enabled": True}], "predefined-model", + session=mock_db.session, ) # Assert @@ -636,6 +659,7 @@ def test_validate_load_balancing_credentials_should_raise_value_error_when_provi "gpt-4o-mini", ModelType.LLM, {"api_key": "plain"}, + session=MagicMock(), ) @@ -657,6 +681,7 @@ def test_validate_load_balancing_credentials_should_raise_value_error_when_confi ModelType.LLM, {"api_key": "plain"}, config_id="cfg-1", + session=mock_db.session, ) @@ -680,6 +705,7 @@ def test_validate_load_balancing_credentials_should_delegate_to_custom_validate_ ModelType.LLM, {"api_key": "plain"}, config_id="cfg-1", + session=mock_db.session, ) service.validate_load_balancing_credentials( "tenant-1", @@ -687,6 +713,7 @@ def test_validate_load_balancing_credentials_should_delegate_to_custom_validate_ "gpt-4o-mini", ModelType.LLM, {"api_key": "plain"}, + session=mock_db.session, ) # Assert diff --git a/api/tests/unit_tests/services/test_oauth_device_flow.py b/api/tests/unit_tests/services/test_oauth_device_flow.py index fcb3f29a76f..00b2919240d 100644 --- a/api/tests/unit_tests/services/test_oauth_device_flow.py +++ b/api/tests/unit_tests/services/test_oauth_device_flow.py @@ -83,7 +83,7 @@ def test_revoke_oauth_token_invalidates_redis_cache_when_live_hash_seen(): redis = MagicMock() - revoke_oauth_token(session, redis, "token-id") + revoke_oauth_token(redis, "token-id", session=session) assert session.execute.called # UPDATE ... WHERE revoked_at IS NULL assert session.commit.called @@ -101,7 +101,7 @@ def test_revoke_oauth_token_is_idempotent_when_already_revoked(): redis = MagicMock() - revoke_oauth_token(session, redis, "token-id") + revoke_oauth_token(redis, "token-id", session=session) assert session.execute.called assert session.commit.called @@ -126,7 +126,7 @@ def test_list_active_sessions_returns_session_execute_rows(): fake_rows = [MagicMock(), MagicMock()] session.execute.return_value.scalars.return_value.all.return_value = fake_rows - out = list_active_sessions(session, _account_ctx(), datetime.now(UTC)) + out = list_active_sessions(_account_ctx(), datetime.now(UTC), session=session) assert out == fake_rows assert session.execute.called @@ -136,11 +136,11 @@ def test_token_belongs_to_subject_true_when_row_present(): session = MagicMock() session.execute.return_value.first.return_value = ("some-id",) - assert token_belongs_to_subject(session, "token-id", _account_ctx()) is True + assert token_belongs_to_subject("token-id", _account_ctx(), session=session) is True def test_token_belongs_to_subject_false_when_no_row(): session = MagicMock() session.execute.return_value.first.return_value = None - assert token_belongs_to_subject(session, "token-id", _account_ctx()) is False + assert token_belongs_to_subject("token-id", _account_ctx(), session=session) is False diff --git a/api/tests/unit_tests/services/test_schedule_service.py b/api/tests/unit_tests/services/test_schedule_service.py index 0f8f7ffab58..d5006bd5f26 100644 --- a/api/tests/unit_tests/services/test_schedule_service.py +++ b/api/tests/unit_tests/services/test_schedule_service.py @@ -1,7 +1,7 @@ +import json import unittest from datetime import UTC, datetime -from types import SimpleNamespace -from typing import Any, cast +from typing import Any from unittest.mock import MagicMock, Mock import pytest @@ -11,7 +11,7 @@ from core.trigger.constants import TRIGGER_SCHEDULE_NODE_TYPE from core.workflow.nodes.trigger_schedule.entities import VisualConfig from core.workflow.nodes.trigger_schedule.exc import ScheduleConfigError from libs.schedule_utils import calculate_next_run_at, convert_12h_to_24h -from models.workflow import Workflow +from models.workflow import Workflow, WorkflowType from services.trigger.schedule_service import ScheduleService @@ -503,7 +503,22 @@ def session_mock() -> MagicMock: def _workflow(**kwargs: Any) -> Workflow: - return cast(Workflow, SimpleNamespace(**kwargs)) + graph_dict = kwargs.pop("graph_dict", {}) + workflow = Workflow.new( + tenant_id="tenant-1", + app_id="app-1", + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_dict), + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + for key, value in kwargs.items(): + setattr(workflow, key, value) + return workflow def test_to_schedule_config_should_build_from_cron_mode() -> None: diff --git a/api/tests/unit_tests/services/test_snippet_dsl_service.py b/api/tests/unit_tests/services/test_snippet_dsl_service.py index c155d3f8330..b234d9f91fe 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -95,7 +95,7 @@ def test_import_snippet_rejects_oversized_yaml_url_content(monkeypatch: pytest.M monkeypatch.setattr("services.snippet_dsl_service.DSL_MAX_SIZE", 3) monkeypatch.setattr( "services.snippet_dsl_service.ssrf_proxy.get", - Mock(return_value=SimpleNamespace(status_code=200, text="too large")), + Mock(return_value=SimpleNamespace(status_code=200, content=b"too large")), ) result = service.import_snippet( @@ -108,6 +108,43 @@ def test_import_snippet_rejects_oversized_yaml_url_content(monkeypatch: pytest.M assert "YAML content size exceeds maximum limit" in result.error +def test_import_snippet_rejects_oversized_yaml_url_bytes_before_decode(monkeypatch: pytest.MonkeyPatch) -> None: + service = SnippetDslService(session=SimpleNamespace()) + monkeypatch.setattr("services.snippet_dsl_service.DSL_MAX_SIZE", 1) + monkeypatch.setattr( + "services.snippet_dsl_service.ssrf_proxy.get", + Mock(return_value=SimpleNamespace(status_code=200, content=b"\xff\xff")), + ) + + result = service.import_snippet( + account=SimpleNamespace(current_tenant_id="tenant-1"), + import_mode=ImportMode.YAML_URL.value, + yaml_url="https://example.com/snippet.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert "YAML content size exceeds maximum limit" in result.error + + +def test_import_snippet_returns_decode_error_for_invalid_yaml_url_bytes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = SnippetDslService(session=SimpleNamespace()) + monkeypatch.setattr( + "services.snippet_dsl_service.ssrf_proxy.get", + Mock(return_value=SimpleNamespace(status_code=200, content=b"\xff")), + ) + + result = service.import_snippet( + account=SimpleNamespace(current_tenant_id="tenant-1"), + import_mode=ImportMode.YAML_URL.value, + yaml_url="https://example.com/snippet.yaml", + ) + + assert result.status == ImportStatus.FAILED + assert "utf-8" in result.error + + def test_import_snippet_returns_failed_when_yaml_url_fetch_raises(monkeypatch: pytest.MonkeyPatch) -> None: service = SnippetDslService(session=SimpleNamespace()) monkeypatch.setattr( @@ -127,12 +164,12 @@ def test_import_snippet_returns_failed_when_yaml_url_fetch_raises(monkeypatch: p def test_import_snippet_rejects_oversized_yaml_content(monkeypatch: pytest.MonkeyPatch) -> None: service = SnippetDslService(session=SimpleNamespace()) - monkeypatch.setattr("services.snippet_dsl_service.DSL_MAX_SIZE", 3) + monkeypatch.setattr("services.snippet_dsl_service.DSL_MAX_SIZE", 1) result = service.import_snippet( account=SimpleNamespace(current_tenant_id="tenant-1"), import_mode=ImportMode.YAML_CONTENT.value, - yaml_content="too large", + yaml_content="é", ) assert result.status == ImportStatus.FAILED diff --git a/api/tests/unit_tests/services/test_summary_index_service.py b/api/tests/unit_tests/services/test_summary_index_service.py index 19418c43926..7ece6204ce3 100644 --- a/api/tests/unit_tests/services/test_summary_index_service.py +++ b/api/tests/unit_tests/services/test_summary_index_service.py @@ -118,7 +118,7 @@ def test_generate_summary_for_segment_raises_when_empty(monkeypatch: pytest.Monk SummaryIndexService.generate_summary_for_segment(_segment(), _dataset(), {"a": 1}) -def test_create_summary_record_updates_existing_and_reenables(monkeypatch: pytest.MonkeyPatch) -> None: +def test_create_summary_record_updates_existing_and_reenables() -> None: existing = _summary_record(summary_content="old", node_id="n1") existing.enabled = False existing.disabled_at = datetime(2024, 1, 1) @@ -127,13 +127,12 @@ def test_create_summary_record_updates_existing_and_reenables(monkeypatch: pytes session = MagicMock(name="session") session.scalar.return_value = existing - create_session_mock = MagicMock(return_value=_SessionContext(session)) - monkeypatch.setattr(summary_module, "session_factory", SimpleNamespace(create_session=create_session_mock)) - segment = _segment() dataset = _dataset() - result = SummaryIndexService.create_summary_record(segment, dataset, "new", status=SummaryStatus.GENERATING) + result = SummaryIndexService.create_summary_record( + segment, dataset, "new", status=SummaryStatus.GENERATING, session=session + ) assert result is existing assert existing.summary_content == "new" assert existing.status == SummaryStatus.GENERATING @@ -145,14 +144,13 @@ def test_create_summary_record_updates_existing_and_reenables(monkeypatch: pytes session.flush.assert_called_once() -def test_create_summary_record_creates_new(monkeypatch: pytest.MonkeyPatch) -> None: +def test_create_summary_record_creates_new() -> None: session = MagicMock(name="session") session.scalar.return_value = None - create_session_mock = MagicMock(return_value=_SessionContext(session)) - monkeypatch.setattr(summary_module, "session_factory", SimpleNamespace(create_session=create_session_mock)) - - record = SummaryIndexService.create_summary_record(_segment(), _dataset(), "new", status=SummaryStatus.GENERATING) + record = SummaryIndexService.create_summary_record( + _segment(), _dataset(), "new", status=SummaryStatus.GENERATING, session=session + ) assert record.dataset_id == "dataset-1" assert record.chunk_id == "seg-1" assert record.summary_content == "new" @@ -331,17 +329,12 @@ def test_generate_and_vectorize_summary_success(monkeypatch: pytest.MonkeyPatch) session = MagicMock() session.scalar.return_value = record - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) monkeypatch.setattr( SummaryIndexService, "generate_summary_for_segment", MagicMock(return_value=("sum", MagicMock(total_tokens=0))) ) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(return_value=None)) - out = SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}) + out = SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}, session=session) assert out is record session.refresh.assert_called_once_with(record) session.commit.assert_called() @@ -355,18 +348,13 @@ def test_generate_and_vectorize_summary_vectorize_failure_sets_error(monkeypatch session = MagicMock() session.scalar.return_value = record - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) monkeypatch.setattr( SummaryIndexService, "generate_summary_for_segment", MagicMock(return_value=("sum", MagicMock(total_tokens=0))) ) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) with pytest.raises(RuntimeError, match="boom"): - SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}) + SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}, session=session) assert record.status == SummaryStatus.ERROR # Outer exception handler overwrites the error with the raw exception message. assert record.error == "boom" @@ -562,18 +550,12 @@ def test_generate_and_vectorize_summary_creates_missing_record_and_logs_usage( session = MagicMock() session.scalar.return_value = None - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - usage = MagicMock(total_tokens=4, prompt_tokens=1, completion_tokens=3) monkeypatch.setattr(SummaryIndexService, "generate_summary_for_segment", MagicMock(return_value=("sum", usage))) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(return_value=None)) with caplog.at_level(logging.INFO, logger="services.summary_index_service"): - result = SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}) + result = SummaryIndexService.generate_and_vectorize_summary(segment, dataset, {"enable": True}, session=session) assert result.status in {SummaryStatus.GENERATING, SummaryStatus.COMPLETED} assert any(r.levelno >= logging.INFO for r in caplog.records) @@ -833,11 +815,12 @@ def test_delete_summaries_for_segments_no_summaries_noop(monkeypatch: pytest.Mon def test_update_summary_for_segment_skip_conditions() -> None: + session = MagicMock() economy_dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY) - assert SummaryIndexService.update_summary_for_segment(_segment(), economy_dataset, "x") is None + assert SummaryIndexService.update_summary_for_segment(_segment(), economy_dataset, "x", session=session) is None seg = _segment(has_document=True) seg.document.doc_form = IndexStructureType.QA_INDEX - assert SummaryIndexService.update_summary_for_segment(seg, _dataset(), "x") is None + assert SummaryIndexService.update_summary_for_segment(seg, _dataset(), "x", session=session) is None def test_update_summary_for_segment_empty_content_deletes_existing(monkeypatch: pytest.MonkeyPatch) -> None: @@ -850,13 +833,7 @@ def test_update_summary_for_segment_empty_content_deletes_existing(monkeypatch: vector_instance = MagicMock() monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector_instance)) - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - - assert SummaryIndexService.update_summary_for_segment(segment, dataset, " ") is None + assert SummaryIndexService.update_summary_for_segment(segment, dataset, " ", session=session) is None vector_instance.delete_by_ids.assert_called_once_with(["n1"]) session.delete.assert_called_once_with(record) session.commit.assert_called_once() @@ -872,18 +849,12 @@ def test_update_summary_for_segment_empty_content_delete_vector_warns( session = MagicMock() session.scalar.return_value = record - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - vector_instance = MagicMock() vector_instance.delete_by_ids.side_effect = RuntimeError("boom") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector_instance)) with caplog.at_level(logging.WARNING, logger="services.summary_index_service"): - assert SummaryIndexService.update_summary_for_segment(segment, dataset, "") is None + assert SummaryIndexService.update_summary_for_segment(segment, dataset, "", session=session) is None assert any(r.levelno >= logging.WARNING for r in caplog.records) @@ -893,13 +864,7 @@ def test_update_summary_for_segment_empty_content_no_record_noop(monkeypatch: py session = MagicMock() session.scalar.return_value = None - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - - assert SummaryIndexService.update_summary_for_segment(segment, dataset, " ") is None + assert SummaryIndexService.update_summary_for_segment(segment, dataset, " ", session=session) is None def test_update_summary_for_segment_updates_existing_and_vectorizes(monkeypatch: pytest.MonkeyPatch) -> None: @@ -912,16 +877,10 @@ def test_update_summary_for_segment_updates_existing_and_vectorizes(monkeypatch: vector_instance = MagicMock() monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector_instance)) - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - vectorize_mock = MagicMock() monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize_mock) - out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new summary") + out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new summary", session=session) assert out is record vectorize_mock.assert_called_once() session.refresh.assert_called_once_with(record) @@ -938,19 +897,13 @@ def test_update_summary_for_segment_existing_vector_delete_warns( session = MagicMock() session.scalar.return_value = record - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - vector_instance = MagicMock() vector_instance.delete_by_ids.side_effect = RuntimeError("boom") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector_instance)) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(return_value=None)) with caplog.at_level(logging.WARNING, logger="services.summary_index_service"): - SummaryIndexService.update_summary_for_segment(segment, dataset, "new") + SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=session) assert any(r.levelno >= logging.WARNING for r in caplog.records) @@ -963,14 +916,9 @@ def test_update_summary_for_segment_existing_vectorize_failure_returns_error_rec session = MagicMock() session.scalar.return_value = record - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) - out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new") + out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=session) assert out is record assert out.status == SummaryStatus.ERROR assert "Vectorization failed" in (out.error or "") @@ -982,18 +930,11 @@ def test_update_summary_for_segment_new_record_success(monkeypatch: pytest.Monke session = MagicMock() session.scalar.return_value = None - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - created = _summary_record(summary_content="new", node_id=None) monkeypatch.setattr(SummaryIndexService, "create_summary_record", MagicMock(return_value=created)) - session.merge.return_value = created monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(return_value=None)) - out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new") + out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=session) assert out is created session.refresh.assert_called() session.commit.assert_called() @@ -1007,81 +948,60 @@ def test_update_summary_for_segment_outer_exception_sets_error_and_reraises(monk session = MagicMock() session.scalar.return_value = record session.flush.side_effect = RuntimeError("flush boom") - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - with pytest.raises(RuntimeError, match="flush boom"): - SummaryIndexService.update_summary_for_segment(segment, dataset, "new") + SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=session) assert record.status == SummaryStatus.ERROR assert record.error == "flush boom" session.commit.assert_called() -def test_get_segment_summary_and_document_summaries(monkeypatch: pytest.MonkeyPatch) -> None: +def test_get_segment_summary_and_document_summaries() -> None: record = _summary_record(summary_content="sum", node_id="n1") session = MagicMock() session.scalar.return_value = record session.scalars.return_value.all.return_value = [record] - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - - assert SummaryIndexService.get_segment_summary("seg-1", "dataset-1") is record - assert SummaryIndexService.get_document_summaries("doc-1", "dataset-1", segment_ids=["seg-1"]) == [record] + assert SummaryIndexService.get_segment_summary("seg-1", "dataset-1", session=session) is record + assert SummaryIndexService.get_document_summaries("doc-1", "dataset-1", segment_ids=["seg-1"], session=session) == [ + record + ] -def test_get_segments_summaries_non_empty(monkeypatch: pytest.MonkeyPatch) -> None: +def test_get_segments_summaries_non_empty() -> None: record1 = _summary_record() record1.chunk_id = "seg-1" record2 = _summary_record() record2.chunk_id = "seg-2" session = MagicMock() session.scalars.return_value.all.return_value = [record1, record2] - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - out = SummaryIndexService.get_segments_summaries(["seg-1", "seg-2"], "dataset-1") + out = SummaryIndexService.get_segments_summaries(["seg-1", "seg-2"], "dataset-1", session=session) assert set(out.keys()) == {"seg-1", "seg-2"} -def test_get_document_summary_index_status_no_segments_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: +def test_get_document_summary_index_status_no_segments_returns_none() -> None: session = MagicMock() session.scalars.return_value.all.return_value = [] - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), + assert ( + SummaryIndexService.get_document_summary_index_status("doc-1", "dataset-1", "tenant-1", session=session) is None ) - assert SummaryIndexService.get_document_summary_index_status("doc-1", "dataset-1", "tenant-1") is None -def test_get_documents_summary_index_status_empty_input(monkeypatch: pytest.MonkeyPatch) -> None: - assert SummaryIndexService.get_documents_summary_index_status([], "dataset-1", "tenant-1") == {} +def test_get_documents_summary_index_status_empty_input() -> None: + assert ( + SummaryIndexService.get_documents_summary_index_status([], "dataset-1", "tenant-1", session=MagicMock()) == {} + ) def test_get_documents_summary_index_status_no_pending_sets_none(monkeypatch: pytest.MonkeyPatch) -> None: session = MagicMock() session.execute.return_value.all.return_value = [SimpleNamespace(id="seg-1", document_id="doc-1")] - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) monkeypatch.setattr( SummaryIndexService, "get_segments_summaries", MagicMock(return_value={"seg-1": SimpleNamespace(status=SummaryStatus.COMPLETED)}), ) - result = SummaryIndexService.get_documents_summary_index_status(["doc-1"], "dataset-1", "tenant-1") + result = SummaryIndexService.get_documents_summary_index_status(["doc-1"], "dataset-1", "tenant-1", session=session) assert result["doc-1"] is None @@ -1094,26 +1014,19 @@ def test_update_summary_for_segment_creates_new_and_vectorize_fails_returns_erro session = MagicMock() session.scalar.return_value = None - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session))), - ) - created = _summary_record(summary_content="new", node_id=None) monkeypatch.setattr(SummaryIndexService, "create_summary_record", MagicMock(return_value=created)) - session.merge.return_value = created vectorize_mock = MagicMock(side_effect=RuntimeError("boom")) monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize_mock) - out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new") + out = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=session) assert out.status == SummaryStatus.ERROR assert "Vectorization failed" in (out.error or "") def test_get_segments_summaries_empty_list() -> None: - assert SummaryIndexService.get_segments_summaries([], "dataset-1") == {} + assert SummaryIndexService.get_segments_summaries([], "dataset-1", session=MagicMock()) == {} def test_get_document_summary_index_status_and_documents_status(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1121,30 +1034,27 @@ def test_get_document_summary_index_status_and_documents_status(monkeypatch: pyt session = MagicMock() session.scalars.return_value.all.return_value = ["seg-1"] # get_document_summary_index_status returns IDs - create_session_mock = MagicMock(return_value=_SessionContext(session)) - monkeypatch.setattr(summary_module, "session_factory", SimpleNamespace(create_session=create_session_mock)) - monkeypatch.setattr( SummaryIndexService, "get_segments_summaries", MagicMock(return_value={"seg-1": SimpleNamespace(status=SummaryStatus.GENERATING)}), ) - assert SummaryIndexService.get_document_summary_index_status("doc-1", "dataset-1", "tenant-1") == "SUMMARIZING" + assert ( + SummaryIndexService.get_document_summary_index_status("doc-1", "dataset-1", "tenant-1", session=session) + == "SUMMARIZING" + ) # Multiple docs session2 = MagicMock() session2.execute.return_value.all.return_value = [seg_row] # get_documents_summary_index_status uses execute - monkeypatch.setattr( - summary_module, - "session_factory", - SimpleNamespace(create_session=MagicMock(return_value=_SessionContext(session2))), - ) monkeypatch.setattr( SummaryIndexService, "get_segments_summaries", MagicMock(return_value={"seg-1": SimpleNamespace(status=SummaryStatus.NOT_STARTED)}), ) - result = SummaryIndexService.get_documents_summary_index_status(["doc-1", "doc-2"], "dataset-1", "tenant-1") + result = SummaryIndexService.get_documents_summary_index_status( + ["doc-1", "doc-2"], "dataset-1", "tenant-1", session=session2 + ) assert result["doc-1"] == "SUMMARIZING" assert result["doc-2"] is None diff --git a/api/tests/unit_tests/services/test_trigger_provider_service.py b/api/tests/unit_tests/services/test_trigger_provider_service.py index 0a4452cf478..ff11bbb3035 100644 --- a/api/tests/unit_tests/services/test_trigger_provider_service.py +++ b/api/tests/unit_tests/services/test_trigger_provider_service.py @@ -444,7 +444,7 @@ def test_delete_trigger_provider_should_raise_error_when_subscription_missing( # Act + Assert with pytest.raises(ValueError, match="not found"): - TriggerProviderService.delete_trigger_provider(mock_session, "tenant-1", "sub-1") + TriggerProviderService.delete_trigger_provider("tenant-1", "sub-1", session=mock_session) def test_delete_trigger_provider_should_delete_and_clear_cache_even_if_unsubscribe_fails( @@ -476,7 +476,7 @@ def test_delete_trigger_provider_should_delete_and_clear_cache_even_if_unsubscri mock_delete_cache = mocker.patch("services.trigger.trigger_provider_service.delete_cache_for_subscription") # Act - TriggerProviderService.delete_trigger_provider(mock_session, "tenant-1", "sub-1") + TriggerProviderService.delete_trigger_provider("tenant-1", "sub-1", session=mock_session) # Assert mock_session.delete.assert_called_once_with(subscription) @@ -507,7 +507,7 @@ def test_delete_trigger_provider_should_skip_unsubscribe_for_unauthorized( ) # Act - TriggerProviderService.delete_trigger_provider(mock_session, "tenant-1", "sub-2") + TriggerProviderService.delete_trigger_provider("tenant-1", "sub-2", session=mock_session) # Assert mock_unsubscribe.assert_not_called() diff --git a/api/tests/unit_tests/services/test_vector_service.py b/api/tests/unit_tests/services/test_vector_service.py index e7ebada6bea..3659b85228b 100644 --- a/api/tests/unit_tests/services/test_vector_service.py +++ b/api/tests/unit_tests/services/test_vector_service.py @@ -98,7 +98,7 @@ def test_create_segments_vector_regular_indexing_loads_documents_and_keywords(mo factory_instance.init_index_processor.return_value = index_processor monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - VectorService.create_segments_vector([["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX) + VectorService.create_segments_vector([["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX, MagicMock()) index_processor.load.assert_called_once() args, kwargs = index_processor.load.call_args @@ -123,7 +123,7 @@ def test_create_segments_vector_regular_indexing_loads_multimodal_documents(monk factory_instance.init_index_processor.return_value = index_processor monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - VectorService.create_segments_vector([["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX) + VectorService.create_segments_vector([["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX, MagicMock()) assert index_processor.load.call_count == 2 first_args, first_kwargs = index_processor.load.call_args_list[0] @@ -145,7 +145,7 @@ def test_create_segments_vector_with_no_segments_does_not_load(monkeypatch: pyte factory_instance.init_index_processor.return_value = index_processor monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - VectorService.create_segments_vector(None, [], dataset, IndexStructureType.PARAGRAPH_INDEX) + VectorService.create_segments_vector(None, [], dataset, IndexStructureType.PARAGRAPH_INDEX, MagicMock()) index_processor.load.assert_not_called() @@ -189,11 +189,7 @@ def test_create_segments_vector_parent_child_calls_generate_child_chunks_with_ex processing_rule = MagicMock(name="processing_rule") processing_rule.to_dict.return_value = {"rules": {}} - monkeypatch.setattr( - vector_service_module, - "db", - _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule), - ) + db_mock = _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule) embedding_model_instance = MagicMock(name="embedding_model_instance") model_manager_instance = MagicMock(name="model_manager_instance") @@ -211,12 +207,22 @@ def test_create_segments_vector_parent_child_calls_generate_child_chunks_with_ex monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) VectorService.create_segments_vector( - None, [segment], dataset, vector_service_module.IndexStructureType.PARENT_CHILD_INDEX + None, + [segment], + dataset, + vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, + db_mock.session, ) model_manager_instance.get_model_instance.assert_called_once() generate_child_chunks_mock.assert_called_once_with( - segment, dataset_document, dataset, embedding_model_instance, processing_rule, False + segment, + dataset_document, + dataset, + embedding_model_instance, + processing_rule, + db_mock.session, + False, ) index_processor.load.assert_not_called() @@ -239,11 +245,7 @@ def test_create_segments_vector_parent_child_uses_default_embedding_model_when_p processing_rule = MagicMock() processing_rule.to_dict.return_value = {"rules": {}} - monkeypatch.setattr( - vector_service_module, - "db", - _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule), - ) + db_mock = _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule) embedding_model_instance = MagicMock() model_manager_instance = MagicMock() @@ -261,7 +263,11 @@ def test_create_segments_vector_parent_child_uses_default_embedding_model_when_p monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) VectorService.create_segments_vector( - None, [segment], dataset, vector_service_module.IndexStructureType.PARENT_CHILD_INDEX + None, + [segment], + dataset, + vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, + db_mock.session, ) model_manager_instance.get_default_model_instance.assert_called_once() @@ -276,11 +282,7 @@ def test_create_segments_vector_parent_child_missing_document_logs_warning_and_c segment = _make_segment() processing_rule = MagicMock() - monkeypatch.setattr( - vector_service_module, - "db", - _mock_parent_child_queries(dataset_document=None, processing_rule=processing_rule), - ) + db_mock = _mock_parent_child_queries(dataset_document=None, processing_rule=processing_rule) index_processor = MagicMock() factory_instance = MagicMock() @@ -289,7 +291,11 @@ def test_create_segments_vector_parent_child_missing_document_logs_warning_and_c with caplog.at_level(logging.WARNING, logger="services.vector_service"): VectorService.create_segments_vector( - None, [segment], dataset, vector_service_module.IndexStructureType.PARENT_CHILD_INDEX + None, + [segment], + dataset, + vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, + db_mock.session, ) assert any(r.levelno >= logging.WARNING for r in caplog.records) index_processor.load.assert_not_called() @@ -301,15 +307,15 @@ def test_create_segments_vector_parent_child_missing_processing_rule_raises(monk dataset_document = MagicMock() dataset_document.dataset_process_rule_id = "rule-1" - monkeypatch.setattr( - vector_service_module, - "db", - _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=None), - ) + db_mock = _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=None) with pytest.raises(ValueError, match="No processing rule found"): VectorService.create_segments_vector( - None, [segment], dataset, vector_service_module.IndexStructureType.PARENT_CHILD_INDEX + None, + [segment], + dataset, + vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, + db_mock.session, ) @@ -322,15 +328,15 @@ def test_create_segments_vector_parent_child_non_high_quality_raises(monkeypatch dataset_document = MagicMock() dataset_document.dataset_process_rule_id = "rule-1" processing_rule = MagicMock() - monkeypatch.setattr( - vector_service_module, - "db", - _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule), - ) + db_mock = _mock_parent_child_queries(dataset_document=dataset_document, processing_rule=processing_rule) with pytest.raises(ValueError, match="not high quality"): VectorService.create_segments_vector( - None, [segment], dataset, vector_service_module.IndexStructureType.PARENT_CHILD_INDEX + None, + [segment], + dataset, + vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, + db_mock.session, ) @@ -404,10 +410,7 @@ def test_generate_child_chunks_regenerate_cleans_then_saves_children(monkeypatch child_chunk_ctor = MagicMock(side_effect=lambda **kwargs: kwargs) monkeypatch.setattr(vector_service_module, "ChildChunk", child_chunk_ctor) - db_mock = MagicMock() - db_mock.session.add = MagicMock() - db_mock.session.commit = MagicMock() - monkeypatch.setattr(vector_service_module, "db", db_mock) + session = MagicMock() VectorService.generate_child_chunks( segment=segment, @@ -415,6 +418,7 @@ def test_generate_child_chunks_regenerate_cleans_then_saves_children(monkeypatch dataset=dataset, embedding_model_instance=MagicMock(), processing_rule=processing_rule, + session=session, regenerate=True, ) @@ -422,8 +426,8 @@ def test_generate_child_chunks_regenerate_cleans_then_saves_children(monkeypatch _, transform_kwargs = index_processor.transform.call_args assert transform_kwargs["process_rule"]["rules"]["parent_mode"] == vector_service_module.ParentMode.FULL_DOC index_processor.load.assert_called_once() - assert db_mock.session.add.call_count == 2 - db_mock.session.commit.assert_called_once() + assert session.add.call_count == 2 + session.commit.assert_called_once() def test_generate_child_chunks_commits_even_when_no_children(monkeypatch: pytest.MonkeyPatch) -> None: @@ -442,8 +446,7 @@ def test_generate_child_chunks_commits_even_when_no_children(monkeypatch: pytest factory_instance.init_index_processor.return_value = index_processor monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - db_mock = MagicMock() - monkeypatch.setattr(vector_service_module, "db", db_mock) + session = MagicMock() VectorService.generate_child_chunks( segment=segment, @@ -451,12 +454,13 @@ def test_generate_child_chunks_commits_even_when_no_children(monkeypatch: pytest dataset=dataset, embedding_model_instance=MagicMock(), processing_rule=processing_rule, + session=session, regenerate=False, ) index_processor.load.assert_not_called() - db_mock.session.add.assert_not_called() - db_mock.session.commit.assert_called_once() + session.add.assert_not_called() + session.commit.assert_called_once() def test_create_child_chunk_vector_high_quality_adds_texts(monkeypatch: pytest.MonkeyPatch) -> None: @@ -554,9 +558,10 @@ def test_update_multimodel_vector_returns_when_not_high_quality(monkeypatch: pyt vector_cls = MagicMock() db_mock = _mock_db_session_for_update_multimodel(upload_files=[]) monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - monkeypatch.setattr(vector_service_module, "db", db_mock) - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["a"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, attachment_ids=["a"], dataset=dataset, session=db_mock.session + ) vector_cls.assert_not_called() db_mock.session.query.assert_not_called() @@ -568,9 +573,10 @@ def test_update_multimodel_vector_returns_when_no_actual_change(monkeypatch: pyt vector_cls = MagicMock() db_mock = _mock_db_session_for_update_multimodel(upload_files=[]) monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - monkeypatch.setattr(vector_service_module, "db", db_mock) - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["b", "a"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, attachment_ids=["b", "a"], dataset=dataset, session=db_mock.session + ) vector_cls.assert_not_called() db_mock.session.query.assert_not_called() @@ -586,9 +592,8 @@ def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids( db_mock = _mock_db_session_for_update_multimodel(upload_files=[]) monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - monkeypatch.setattr(vector_service_module, "db", db_mock) - VectorService.update_multimodel_vector(segment=segment, attachment_ids=[], dataset=dataset) + VectorService.update_multimodel_vector(segment=segment, attachment_ids=[], dataset=dataset, session=db_mock.session) vector_cls.assert_called_once_with(dataset=dataset) vector_instance.delete_by_ids.assert_called_once_with(["old-1", "old-2"]) @@ -605,9 +610,10 @@ def test_update_multimodel_vector_commits_when_no_upload_files_found(monkeypatch vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) db_mock = _mock_db_session_for_update_multimodel(upload_files=[]) - monkeypatch.setattr(vector_service_module, "db", db_mock) - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["new-1"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, attachment_ids=["new-1"], dataset=dataset, session=db_mock.session + ) db_mock.session.commit.assert_called_once() db_mock.session.add_all.assert_not_called() @@ -624,7 +630,6 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) db_mock = _mock_db_session_for_update_multimodel(upload_files=[_UploadFileStub(id="file-1", name="img.png")]) - monkeypatch.setattr(vector_service_module, "db", db_mock) binding_ctor = MagicMock(side_effect=lambda **kwargs: kwargs) monkeypatch.setattr(vector_service_module, "SegmentAttachmentBinding", binding_ctor) @@ -632,7 +637,12 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up monkeypatch.setattr(vector_service_module, "select", MagicMock()) with caplog.at_level(logging.WARNING, logger="services.vector_service"): - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["file-1", "missing"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, + attachment_ids=["file-1", "missing"], + dataset=dataset, + session=db_mock.session, + ) assert any(r.levelno >= logging.WARNING for r in caplog.records) db_mock.session.add_all.assert_called_once() bindings = db_mock.session.add_all.call_args.args[0] @@ -656,14 +666,15 @@ def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops vector_instance = MagicMock() monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) db_mock = _mock_db_session_for_update_multimodel(upload_files=[_UploadFileStub(id="file-1", name="img.png")]) - monkeypatch.setattr(vector_service_module, "db", db_mock) monkeypatch.setattr( vector_service_module, "SegmentAttachmentBinding", MagicMock(side_effect=lambda **kwargs: kwargs) ) monkeypatch.setattr(vector_service_module, "delete", MagicMock()) monkeypatch.setattr(vector_service_module, "select", MagicMock()) - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["file-1"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, attachment_ids=["file-1"], dataset=dataset, session=db_mock.session + ) vector_instance.delete_by_ids.assert_not_called() vector_instance.add_texts.assert_not_called() @@ -682,7 +693,6 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error( monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) db_mock = _mock_db_session_for_update_multimodel(upload_files=[_UploadFileStub(id="file-1", name="img.png")]) db_mock.session.commit.side_effect = RuntimeError("boom") - monkeypatch.setattr(vector_service_module, "db", db_mock) monkeypatch.setattr( vector_service_module, "SegmentAttachmentBinding", MagicMock(side_effect=lambda **kwargs: kwargs) ) @@ -691,7 +701,9 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error( with caplog.at_level(logging.ERROR, logger="services.vector_service"): with pytest.raises(RuntimeError, match="boom"): - VectorService.update_multimodel_vector(segment=segment, attachment_ids=["file-1"], dataset=dataset) + VectorService.update_multimodel_vector( + segment=segment, attachment_ids=["file-1"], dataset=dataset, session=db_mock.session + ) assert any(r.levelno >= logging.ERROR for r in caplog.records) db_mock.session.rollback.assert_called_once() diff --git a/api/tests/unit_tests/services/test_workflow_collaboration_service.py b/api/tests/unit_tests/services/test_workflow_collaboration_service.py index a61e49c02fa..6b269443fa2 100644 --- a/api/tests/unit_tests/services/test_workflow_collaboration_service.py +++ b/api/tests/unit_tests/services/test_workflow_collaboration_service.py @@ -32,7 +32,7 @@ class TestWorkflowCollaborationService: patch.object(collaboration_service, "broadcast_online_users"), ): # Act - result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1") + result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=Mock()) # Assert assert result == ("u-1", True) @@ -52,7 +52,7 @@ class TestWorkflowCollaborationService: socketio.get_session.return_value = {} # Act - result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1") + result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=Mock()) # Assert assert result is None @@ -63,7 +63,7 @@ class TestWorkflowCollaborationService: collaboration_service, repository, socketio = service socketio.get_session.return_value = {"user_id": "u-1", "username": "Jane", "avatar": None} - result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1") + result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=Mock()) assert result is None repository.set_session_info.assert_not_called() @@ -82,7 +82,7 @@ class TestWorkflowCollaborationService: } with patch.object(collaboration_service, "_can_access_workflow", return_value=False): - result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1") + result = collaboration_service.authorize_and_join_workflow_room("wf-1", "sid-1", session=Mock()) assert result is None repository.set_session_info.assert_not_called() @@ -106,21 +106,12 @@ class TestWorkflowCollaborationService: {"user_id": "u-1", "username": "Jane", "avatar": "avatar.png", "tenant_id": "t-1"}, ) - def test_can_access_workflow_uses_session_factory( - self, service: tuple[WorkflowCollaborationService, Mock, Mock] - ) -> None: + def test_can_access_workflow_uses_session(self, service: tuple[WorkflowCollaborationService, Mock, Mock]) -> None: collaboration_service, _repository, _socketio = service session = Mock() session.scalar.return_value = "wf-1" - session_context = Mock() - session_context.__enter__ = Mock(return_value=session) - session_context.__exit__ = Mock(return_value=False) - with patch( - "services.workflow_collaboration_service.session_factory.create_session", - return_value=session_context, - ): - result = collaboration_service._can_access_workflow("wf-1", "tenant-1") + result = collaboration_service._can_access_workflow("wf-1", "tenant-1", session=session) assert result is True session.scalar.assert_called_once() diff --git a/api/tests/unit_tests/services/test_workflow_run_service.py b/api/tests/unit_tests/services/test_workflow_run_service.py index 03471389a65..2c69a742f15 100644 --- a/api/tests/unit_tests/services/test_workflow_run_service.py +++ b/api/tests/unit_tests/services/test_workflow_run_service.py @@ -34,6 +34,13 @@ def _end_user(**kwargs: Any) -> EndUser: return cast(EndUser, SimpleNamespace(**kwargs)) +def _fake_session_returning_messages(messages: list[Any]) -> SimpleNamespace: + """A stand-in db session whose scalars(...).all() returns the given messages.""" + scalars_result = MagicMock() + scalars_result.all.return_value = messages + return SimpleNamespace(scalars=MagicMock(return_value=scalars_result)) + + class TestWorkflowRunServiceInitialization: def test___init___should_create_sessionmaker_from_db_engine_when_session_factory_missing( self, @@ -120,15 +127,15 @@ class TestWorkflowRunServiceQueries: ) -> None: service = WorkflowRunService(session_factory=MagicMock(name="session_factory")) app_model = _app_model(tenant_id="tenant-1", id="app-1") - run_with_message = SimpleNamespace( - id="run-1", - status="running", - message=SimpleNamespace(id="msg-1", conversation_id="conv-1"), - ) - run_without_message = SimpleNamespace(id="run-2", status="succeeded", message=None) + run_with_message = SimpleNamespace(id="run-1", status="running") + run_without_message = SimpleNamespace(id="run-2", status="succeeded") pagination = SimpleNamespace(data=[run_with_message, run_without_message]) monkeypatch.setattr(service, "get_paginate_workflow_runs", MagicMock(return_value=pagination)) + message = SimpleNamespace(id="msg-1", conversation_id="conv-1", workflow_run_id="run-1") + fake_session = _fake_session_returning_messages([message]) + monkeypatch.setattr(service_module, "db", SimpleNamespace(session=fake_session)) + result = service.get_paginate_advanced_chat_workflow_runs(app_model=app_model, args={"limit": "2"}) assert result is pagination @@ -138,6 +145,32 @@ class TestWorkflowRunServiceQueries: assert result.data[0].status == "running" assert not hasattr(result.data[1], "message_id") assert result.data[1].id == "run-2" + # Messages are batch-loaded in a single query, not one per run. + fake_session.scalars.assert_called_once() + + def test_get_paginate_advanced_chat_workflow_runs_batch_loads_messages_without_n_plus_one( + self, + repository_factory_mocks: tuple[MagicMock, MagicMock, Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Messages must load with a constant query count regardless of run count. + + Previously the deprecated WorkflowRun.message property issued one query per + run (N+1); they are now batch-loaded in a single query. + """ + service = WorkflowRunService(session_factory=MagicMock(name="session_factory")) + app_model = _app_model(tenant_id="tenant-1", id="app-1") + runs = [SimpleNamespace(id=f"run-{i}", status="succeeded") for i in range(5)] + pagination = SimpleNamespace(data=runs) + monkeypatch.setattr(service, "get_paginate_workflow_runs", MagicMock(return_value=pagination)) + + fake_session = _fake_session_returning_messages([]) + monkeypatch.setattr(service_module, "db", SimpleNamespace(session=fake_session)) + + service.get_paginate_advanced_chat_workflow_runs(app_model=app_model, args={}) + + # Exactly one message query for the whole page, independent of run count. + assert fake_session.scalars.call_count == 1 def test_get_workflow_run_should_delegate_to_repository_by_tenant_and_app( self, diff --git a/api/tests/unit_tests/services/test_workflow_service.py b/api/tests/unit_tests/services/test_workflow_service.py index 67b3e80da6b..0a75a0a8788 100644 --- a/api/tests/unit_tests/services/test_workflow_service.py +++ b/api/tests/unit_tests/services/test_workflow_service.py @@ -312,7 +312,7 @@ class TestWorkflowService: # Mock the database query to return True mock_db_session.session.execute.return_value.scalar_one.return_value = True - result = workflow_service.is_workflow_exist(app) + result = workflow_service.is_workflow_exist(app, session=mock_db_session.session) assert result is True @@ -323,7 +323,7 @@ class TestWorkflowService: # Mock the database query to return False mock_db_session.session.execute.return_value.scalar_one.return_value = False - result = workflow_service.is_workflow_exist(app) + result = workflow_service.is_workflow_exist(app, session=mock_db_session.session) assert result is False @@ -343,7 +343,7 @@ class TestWorkflowService: # Mock db.session.scalar() used by get_draft_workflow mock_db_session.session.scalar.return_value = mock_workflow - result = workflow_service.get_draft_workflow(app) + result = workflow_service.get_draft_workflow(app, session=mock_db_session.session) assert result == mock_workflow @@ -367,7 +367,7 @@ class TestWorkflowService: # Mock db.session.scalar() to return None mock_db_session.session.scalar.return_value = None - result = workflow_service.get_draft_workflow(app) + result = workflow_service.get_draft_workflow(app, session=mock_db_session.session) assert result is None @@ -380,7 +380,7 @@ class TestWorkflowService: # Mock db.session.scalar() used by get_published_workflow_by_id mock_db_session.session.scalar.return_value = mock_workflow - result = workflow_service.get_draft_workflow(app, workflow_id=workflow_id) + result = workflow_service.get_draft_workflow(app, workflow_id=workflow_id, session=mock_db_session.session) assert result == mock_workflow @@ -411,7 +411,7 @@ class TestWorkflowService: # Mock db.session.scalar() used by get_published_workflow_by_id mock_db_session.session.scalar.return_value = mock_workflow - result = workflow_service.get_published_workflow_by_id(app, workflow_id) + result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=mock_db_session.session) assert result == mock_workflow @@ -432,7 +432,7 @@ class TestWorkflowService: mock_db_session.session.scalar.return_value = mock_workflow with pytest.raises(IsDraftWorkflowError): - workflow_service.get_published_workflow_by_id(app, workflow_id) + workflow_service.get_published_workflow_by_id(app, workflow_id, session=mock_db_session.session) def test_get_published_workflow_by_id_returns_none(self, workflow_service, mock_db_session): """Test get_published_workflow_by_id returns None when workflow not found.""" @@ -442,7 +442,7 @@ class TestWorkflowService: # Mock db.session.scalar() to return None mock_db_session.session.scalar.return_value = None - result = workflow_service.get_published_workflow_by_id(app, workflow_id) + result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=mock_db_session.session) assert result is None @@ -455,7 +455,7 @@ class TestWorkflowService: # Mock db.session.scalar() used by get_published_workflow mock_db_session.session.scalar.return_value = mock_workflow - result = workflow_service.get_published_workflow(app) + result = workflow_service.get_published_workflow(app, session=mock_db_session.session) assert result == mock_workflow @@ -463,7 +463,7 @@ class TestWorkflowService: """Test get_published_workflow returns None when app has no workflow_id.""" app = TestWorkflowAssociatedDataFactory.create_app_mock(workflow_id=None) - result = workflow_service.get_published_workflow(app) + result = workflow_service.get_published_workflow(app, session=MagicMock()) assert result is None @@ -499,6 +499,7 @@ class TestWorkflowService: account=account, environment_variables=[], conversation_variables=[], + session=mock_db_session.session, ) # Verify workflow was added to session @@ -536,6 +537,7 @@ class TestWorkflowService: account=account, environment_variables=[], conversation_variables=[], + session=mock_db_session.session, ) # Verify workflow was updated @@ -571,6 +573,7 @@ class TestWorkflowService: account=account, environment_variables=[], conversation_variables=[], + session=mock_db_session.session, ) def test_restore_published_workflow_to_draft_keeps_source_features_unmodified( @@ -648,6 +651,7 @@ class TestWorkflowService: app_model=app, workflow_id=source_workflow.id, account=account, + session=mock_db_session.session, ) mock_validate_features.assert_called_once_with(app_model=app, features=normalized_features) @@ -761,6 +765,7 @@ class TestWorkflowService: app_model=app, environment_variables=variables, account=account, + session=mock_db_session.session, ) assert workflow.environment_variables == variables @@ -779,6 +784,7 @@ class TestWorkflowService: app_model=app, environment_variables=[], account=account, + session=MagicMock(), ) def test_update_draft_workflow_conversation_variables_updates_workflow(self, workflow_service, mock_db_session): @@ -796,6 +802,7 @@ class TestWorkflowService: app_model=app, conversation_variables=variables, account=account, + session=mock_db_session.session, ) assert workflow.conversation_variables == variables @@ -814,6 +821,7 @@ class TestWorkflowService: app_model=app, conversation_variables=[], account=account, + session=MagicMock(), ) # ==================== Publish Workflow Tests ==================== @@ -1429,7 +1437,7 @@ class TestWorkflowService: mock_new_app = TestWorkflowAssociatedDataFactory.create_app_mock(mode=AppMode.WORKFLOW) mock_converter.convert_to_workflow.return_value = mock_new_app - result = workflow_service.convert_to_workflow(app, account, args) + result = workflow_service.convert_to_workflow(app, account, args, session=MagicMock()) assert result == mock_new_app mock_converter.convert_to_workflow.assert_called_once() @@ -1451,7 +1459,7 @@ class TestWorkflowService: mock_new_app = TestWorkflowAssociatedDataFactory.create_app_mock(mode=AppMode.WORKFLOW) mock_converter.convert_to_workflow.return_value = mock_new_app - result = workflow_service.convert_to_workflow(app, account, args) + result = workflow_service.convert_to_workflow(app, account, args, session=MagicMock()) assert result == mock_new_app @@ -1467,7 +1475,7 @@ class TestWorkflowService: args = {} with pytest.raises(ValueError, match="not supported convert to workflow"): - workflow_service.convert_to_workflow(app, account, args) + workflow_service.convert_to_workflow(app, account, args, session=MagicMock()) # =========================================================================== @@ -1520,7 +1528,7 @@ class TestWorkflowServiceCredentialValidation: # Act + Assert with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: # Should not raise; mock allows the call - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) mock_check.assert_called_once() def test_validate_workflow_credentials_should_check_default_credential_when_no_credential_id( @@ -1541,10 +1549,11 @@ class TestWorkflowServiceCredentialValidation: # Act with patch.object(service, "_check_default_tool_credential") as mock_default: - service._validate_workflow_credentials(workflow) + session = MagicMock() + service._validate_workflow_credentials(workflow, session=session) # Assert - mock_default.assert_called_once_with("tenant-1", "my-provider") + mock_default.assert_called_once_with("tenant-1", "my-provider", session=session) def test_validate_workflow_credentials_should_skip_tool_node_without_provider( self, service: WorkflowService @@ -1556,7 +1565,7 @@ class TestWorkflowServiceCredentialValidation: # Act + Assert (no error raised) with patch.object(service, "_check_default_tool_credential") as mock_default: - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) mock_default.assert_not_called() def test_validate_workflow_credentials_should_validate_llm_node_with_model_config( @@ -1579,7 +1588,7 @@ class TestWorkflowServiceCredentialValidation: patch.object(service, "_validate_llm_model_config") as mock_llm, patch.object(service, "_validate_load_balancing_credentials"), ): - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) # Assert mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") @@ -1599,7 +1608,7 @@ class TestWorkflowServiceCredentialValidation: # Act + Assert with pytest.raises(ValueError, match="Missing provider or model configuration"): - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) def test_validate_workflow_credentials_should_wrap_unexpected_exception_in_value_error( self, service: WorkflowService @@ -1620,7 +1629,7 @@ class TestWorkflowServiceCredentialValidation: # Act + Assert with patch.object(service, "_validate_llm_model_config", side_effect=RuntimeError("boom")): with pytest.raises(ValueError, match="boom"): - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) def test_validate_workflow_credentials_should_validate_agent_node_model(self, service: WorkflowService) -> None: # Arrange @@ -1643,7 +1652,7 @@ class TestWorkflowServiceCredentialValidation: patch.object(service, "_validate_llm_model_config") as mock_llm, patch.object(service, "_validate_load_balancing_credentials"), ): - service._validate_workflow_credentials(workflow) + service._validate_workflow_credentials(workflow, session=MagicMock()) # Assert mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") @@ -1675,11 +1684,12 @@ class TestWorkflowServiceCredentialValidation: patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check, patch.object(service, "_check_default_tool_credential") as mock_default, ): - service._validate_workflow_credentials(workflow) + session = MagicMock() + service._validate_workflow_credentials(workflow, session=session) # Assert mock_check.assert_called_once() # provider-a has credential_id - mock_default.assert_called_once_with("tenant-1", "provider-b") + mock_default.assert_called_once_with("tenant-1", "provider-b", session=session) # --- _validate_llm_model_config --- @@ -1739,7 +1749,7 @@ class TestWorkflowServiceCredentialValidation: # Arrange with patch("services.workflow_service.db") as mock_db: # Act + Assert (should NOT raise) - service._check_default_tool_credential("tenant-1", "some-provider") + service._check_default_tool_credential("tenant-1", "some-provider", session=MagicMock()) def test_check_default_tool_credential_should_raise_when_compliance_fails(self, service: WorkflowService) -> None: # Arrange @@ -1751,7 +1761,7 @@ class TestWorkflowServiceCredentialValidation: ): # Act + Assert with pytest.raises(ValueError, match="Failed to validate default credential"): - service._check_default_tool_credential("tenant-1", "some-provider") + service._check_default_tool_credential("tenant-1", "some-provider", session=MagicMock()) # --- _is_load_balancing_enabled --- @@ -1811,7 +1821,7 @@ class TestWorkflowServiceCredentialValidation: side_effect=RuntimeError("fail"), ): # Act - result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4") + result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=MagicMock()) # Assert assert result == [] @@ -1828,7 +1838,7 @@ class TestWorkflowServiceCredentialValidation: ], ): # Act - result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4") + result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=MagicMock()) # Assert — only entries with a credential_id should be returned assert len(result) == 2 @@ -1845,7 +1855,7 @@ class TestWorkflowServiceCredentialValidation: node_data: dict[str, Any] = {} # no model key # Act + Assert (no error expected) - service._validate_load_balancing_credentials(workflow, node_data, "node-1") + service._validate_load_balancing_credentials(workflow, node_data, "node-1", session=MagicMock()) def test_validate_load_balancing_credentials_should_skip_when_lb_not_enabled( self, service: WorkflowService @@ -1856,7 +1866,7 @@ class TestWorkflowServiceCredentialValidation: # Act + Assert (no error expected) with patch.object(service, "_is_load_balancing_enabled", return_value=False): - service._validate_load_balancing_credentials(workflow, node_data, "node-1") + service._validate_load_balancing_credentials(workflow, node_data, "node-1", session=MagicMock()) def test_validate_load_balancing_credentials_should_raise_when_compliance_fails( self, service: WorkflowService @@ -1876,7 +1886,7 @@ class TestWorkflowServiceCredentialValidation: ), ): with pytest.raises(ValueError, match="Invalid load balancing credentials"): - service._validate_load_balancing_credentials(workflow, node_data, "node-1") + service._validate_load_balancing_credentials(workflow, node_data, "node-1", session=MagicMock()) # =========================================================================== @@ -2673,7 +2683,9 @@ class TestWorkflowServiceHumanInputOperations: def test_get_human_input_form_preview_should_raise_if_workflow_not_init(self, service: WorkflowService) -> None: service.get_draft_workflow = MagicMock(return_value=None) with pytest.raises(ValueError, match="Workflow not initialized"): - service.get_human_input_form_preview(app_model=MagicMock(), account=MagicMock(), node_id="node-1") + service.get_human_input_form_preview( + app_model=MagicMock(), account=MagicMock(), node_id="node-1", session=MagicMock() + ) def test_get_human_input_form_preview_should_raise_if_wrong_node_type(self, service: WorkflowService) -> None: draft = MagicMock() @@ -2681,7 +2693,9 @@ class TestWorkflowServiceHumanInputOperations: service.get_draft_workflow = MagicMock(return_value=draft) with patch("models.workflow.Workflow.get_node_type_from_node_config", return_value=BuiltinNodeTypes.LLM): with pytest.raises(ValueError, match="Node type must be human-input"): - service.get_human_input_form_preview(app_model=MagicMock(), account=MagicMock(), node_id="node-1") + service.get_human_input_form_preview( + app_model=MagicMock(), account=MagicMock(), node_id="node-1", session=MagicMock() + ) def test_get_human_input_form_preview_success(self, service: WorkflowService) -> None: app_model = MagicMock(spec=App) @@ -2716,7 +2730,9 @@ class TestWorkflowServiceHumanInputOperations: patch("services.workflow_service.HumanInputNode", return_value=mock_node), patch("services.workflow_service.HumanInputRequired") as mock_required_cls, ): - service.get_human_input_form_preview(app_model=app_model, account=account, node_id="node-1") + service.get_human_input_form_preview( + app_model=app_model, account=account, node_id="node-1", session=MagicMock() + ) mock_node.render_form_content_before_submission.assert_called_once() mock_required_cls.return_value.model_dump.assert_called_once() @@ -2760,7 +2776,12 @@ class TestWorkflowServiceHumanInputOperations: patch("services.workflow_service.DraftVariableSaver") as mock_saver_cls, ): result = service.submit_human_input_form_preview( - app_model=app_model, account=account, node_id="node-1", form_inputs={"field1": "val1"}, action="submit" + app_model=app_model, + account=account, + node_id="node-1", + form_inputs={"field1": "val1"}, + action="submit", + session=MagicMock(), ) assert result["__action_id"] == "submit" mock_validate.assert_called_once() @@ -2785,7 +2806,11 @@ class TestWorkflowServiceHumanInputOperations: ): mock_resolve.return_value = MagicMock() service.test_human_input_delivery( - app_model=MagicMock(), account=MagicMock(), node_id="node-1", delivery_method_id="method-1" + app_model=MagicMock(), + account=MagicMock(), + node_id="node-1", + delivery_method_id="method-1", + session=MagicMock(), ) mock_test_srv.return_value.send_test.assert_called_once() @@ -2801,7 +2826,11 @@ class TestWorkflowServiceHumanInputOperations: ): with pytest.raises(ValueError, match="Delivery method not found"): service.test_human_input_delivery( - app_model=MagicMock(), account=MagicMock(), node_id="node-1", delivery_method_id="none" + app_model=MagicMock(), + account=MagicMock(), + node_id="node-1", + delivery_method_id="none", + session=MagicMock(), ) def test_load_email_recipients_parsing_failure(self, service: WorkflowService) -> None: diff --git a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py index c210db580e0..549f50cb370 100644 --- a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py +++ b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py @@ -354,7 +354,7 @@ class TestGetBuiltinToolProviderCredentialInfo: def test_returns_credential_info(self, mock_tm, mock_creds, mock_oauth): mock_tm.get_builtin_provider.return_value.get_supported_credential_types.return_value = ["api-key"] - result = BuiltinToolManageService.get_builtin_tool_provider_credential_info("t", "google") + result = BuiltinToolManageService.get_builtin_tool_provider_credential_info("t", "google", session=MagicMock()) assert result.credentials == [] assert result.supported_credential_types == ["api-key"] @@ -368,7 +368,7 @@ class TestGetBuiltinToolProviderCredentials: mock_db.session.no_autoflush.__exit__ = MagicMock(return_value=False) mock_db.session.scalars.return_value.all.return_value = [] - result = BuiltinToolManageService.get_builtin_tool_provider_credentials("t", "google") + result = BuiltinToolManageService.get_builtin_tool_provider_credentials("t", "google", session=mock_db.session) assert result == [] @@ -391,7 +391,7 @@ class TestGetBuiltinToolProviderCredentials: credential_entity = MagicMock() mock_transform.convert_builtin_provider_to_credential_entity.return_value = credential_entity - result = BuiltinToolManageService.get_builtin_tool_provider_credentials("t", "google") + result = BuiltinToolManageService.get_builtin_tool_provider_credentials("t", "google", session=mock_db.session) assert len(result) == 1 assert result[0] is credential_entity diff --git a/api/tests/unit_tests/services/workflow/test_node_output_inspector_service.py b/api/tests/unit_tests/services/workflow/test_node_output_inspector_service.py index 6f6c56fd67f..7f720575154 100644 --- a/api/tests/unit_tests/services/workflow/test_node_output_inspector_service.py +++ b/api/tests/unit_tests/services/workflow/test_node_output_inspector_service.py @@ -1,8 +1,8 @@ """Unit tests for NodeOutputInspectorService (Stage 4 §8). The service reads from postgres and resolves agent v2 bindings; this suite -mocks ``session_factory`` and the binding resolver so we exercise the -view-construction logic without DB / network access. +mocks the DB session and binding resolver so we exercise the view-construction +logic without DB / network access. """ from __future__ import annotations @@ -100,26 +100,17 @@ def _non_agent_node(*, node_id: str = "tool-node-1", node_type: str = "tool", ti } -def _patch_session( +def _mock_session( *, workflow_run: SimpleNamespace | None, executions: list[SimpleNamespace] | None = None, ): - """Patch ``session_factory.create_session`` to return the configured rows. - - Returns a context manager that the test uses with ``with``. - """ + """Build a mock DB session with the configured rows.""" executions = executions or [] - mock_session = MagicMock() - mock_session.scalar.return_value = workflow_run - mock_session.scalars.return_value.all.return_value = executions - cm = MagicMock() - cm.__enter__.return_value = mock_session - cm.__exit__.return_value = False - return patch( - "services.workflow.node_output_inspector_service.session_factory.create_session", - return_value=cm, - ) + session = MagicMock() + session.scalar.return_value = workflow_run + session.scalars.return_value.all.return_value = executions + return session def _stub_binding_resolver(*, declared_outputs: list[DeclaredOutputConfig]): @@ -149,9 +140,9 @@ def _make_service(declared_outputs: list[DeclaredOutputConfig] | None = None) -> def test_snapshot_404_when_workflow_run_missing(): service = _make_service() - with _patch_session(workflow_run=None): - with pytest.raises(NodeOutputInspectorError) as exc: - service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="missing") + session = _mock_session(workflow_run=None) + with pytest.raises(NodeOutputInspectorError) as exc: + service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="missing", session=session) assert exc.value.code == "workflow_run_not_found" @@ -162,8 +153,8 @@ def test_snapshot_accepts_published_run_d1_lifted(): nodes=[_agent_v2_node(node_id="agent-1")], triggered_from=WorkflowRunTriggeredFrom.APP_RUN, ) - with _patch_session(workflow_run=run, executions=[]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.workflow_run_id == "run-1" assert [n.node_id for n in snapshot.node_outputs] == ["agent-1"] @@ -175,17 +166,17 @@ def test_snapshot_accepts_webhook_triggered_run(): nodes=[_agent_v2_node(node_id="agent-1")], triggered_from=WorkflowRunTriggeredFrom.WEBHOOK, ) - with _patch_session(workflow_run=run, executions=[]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.workflow_run_id == "run-1" def test_node_detail_404_when_node_id_absent_from_graph(): service = _make_service() run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) - with _patch_session(workflow_run=run, executions=[]): - with pytest.raises(NodeOutputInspectorError) as exc: - service.node_detail(app_model=_app_model(), workflow_run_id="run-1", node_id="ghost") + session = _mock_session(workflow_run=run, executions=[]) + with pytest.raises(NodeOutputInspectorError) as exc: + service.node_detail(app_model=_app_model(), workflow_run_id="run-1", node_id="ghost", session=session) assert exc.value.code == "node_not_in_workflow_run" @@ -195,28 +186,30 @@ def test_output_preview_404_when_output_name_unknown(): ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", outputs={"text": "hello"}) - with _patch_session(workflow_run=run, executions=[ex]): - with pytest.raises(NodeOutputInspectorError) as exc: - service.output_preview( - app_model=_app_model(), - workflow_run_id="run-1", - node_id="agent-1", - output_name="missing", - ) + session = _mock_session(workflow_run=run, executions=[ex]) + with pytest.raises(NodeOutputInspectorError) as exc: + service.output_preview( + app_model=_app_model(), + workflow_run_id="run-1", + node_id="agent-1", + output_name="missing", + session=session, + ) assert exc.value.code == "node_output_not_declared" def test_output_preview_404_when_node_id_absent_from_graph(): service = _make_service() run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) - with _patch_session(workflow_run=run, executions=[]): - with pytest.raises(NodeOutputInspectorError) as exc: - service.output_preview( - app_model=_app_model(), - workflow_run_id="run-1", - node_id="ghost", - output_name="report", - ) + session = _mock_session(workflow_run=run, executions=[]) + with pytest.raises(NodeOutputInspectorError) as exc: + service.output_preview( + app_model=_app_model(), + workflow_run_id="run-1", + node_id="ghost", + output_name="report", + session=session, + ) assert exc.value.code == "node_not_in_workflow_run" @@ -230,8 +223,8 @@ def test_snapshot_status_pending_when_node_has_no_execution(): declared_outputs=[DeclaredOutputConfig(name="text", type=DeclaredOutputType.STRING)], ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) - with _patch_session(workflow_run=run, executions=[]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert len(snapshot.node_outputs) == 1 node = snapshot.node_outputs[0] @@ -245,8 +238,8 @@ def test_snapshot_status_running(): ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", status=WorkflowNodeExecutionStatus.RUNNING) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.node_outputs[0].node_status == NodeStatus.RUNNING assert snapshot.node_outputs[0].outputs[0].status == NodeOutputStatus.RUNNING @@ -260,8 +253,8 @@ def test_snapshot_status_failed_node_marks_all_outputs_failed(): ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", status=WorkflowNodeExecutionStatus.FAILED) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) statuses = {o.name: o.status for o in snapshot.node_outputs[0].outputs} assert statuses == {"a": NodeOutputStatus.FAILED, "b": NodeOutputStatus.FAILED} @@ -272,8 +265,8 @@ def test_snapshot_status_ready_when_outputs_present_and_no_failure_metadata(): ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", outputs={"text": "hello"}) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) output = snapshot.node_outputs[0].outputs[0] assert output.status == NodeOutputStatus.READY assert output.value_preview == "hello" @@ -294,8 +287,8 @@ def test_snapshot_marks_type_check_failure(): } }, ) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) output = snapshot.node_outputs[0].outputs[0] assert output.status == NodeOutputStatus.TYPE_CHECK_FAILED assert output.type_check is not None @@ -324,14 +317,12 @@ def test_snapshot_marks_output_check_failure_when_type_check_passed(): }, }, ) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service.file_helpers.get_signed_file_url", - return_value="https://signed.example/x", - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service.file_helpers.get_signed_file_url", + return_value="https://signed.example/x", ): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) output = snapshot.node_outputs[0].outputs[0] assert output.status == NodeOutputStatus.OUTPUT_CHECK_FAILED assert output.output_check is not None @@ -348,8 +339,8 @@ def test_snapshot_marks_not_produced_when_declared_output_missing_from_payload() ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", outputs={"text": "hi"}) # optional_meta missing - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) statuses = {o.name: o.status for o in snapshot.node_outputs[0].outputs} assert statuses == {"text": NodeOutputStatus.READY, "optional_meta": NodeOutputStatus.NOT_PRODUCED} @@ -367,8 +358,8 @@ def test_non_agent_node_outputs_inferred_from_payload_keys(): node_type="tool", outputs={"message": "sent", "thread_ts": "1234"}, ) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) output_names = sorted(o.name for o in snapshot.node_outputs[0].outputs) assert output_names == ["message", "thread_ts"] # All inferred outputs should have ``type=None`` since we don't know the @@ -393,14 +384,12 @@ def test_file_output_preview_includes_signed_url(): "reference": build_file_reference(record_id="550e8400-e29b-41d4-a716-446655440000"), } ex = _execution(node_id="agent-1", outputs={"report": file_payload}) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service._resolve_preview_url", - return_value="https://signed.example/x.pdf", - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service._resolve_preview_url", + return_value="https://signed.example/x.pdf", ): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) preview_value = snapshot.node_outputs[0].outputs[0].value_preview assert isinstance(preview_value, dict) assert preview_value["preview_url"] == "https://signed.example/x.pdf" @@ -419,18 +408,17 @@ def test_file_output_preview_endpoint_returns_full_value_with_signed_url(): "reference": build_file_reference(record_id="550e8400-e29b-41d4-a716-446655440000"), } ex = _execution(node_id="agent-1", outputs={"report": file_payload}) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service._resolve_preview_url", - return_value="https://signed.example/x.pdf", - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service._resolve_preview_url", + return_value="https://signed.example/x.pdf", ): preview = service.output_preview( app_model=_app_model(), workflow_run_id="run-1", node_id="agent-1", output_name="report", + session=session, ) assert preview.output_name == "report" assert preview.status == NodeOutputStatus.READY @@ -484,26 +472,25 @@ def test_array_file_output_preview_includes_signed_urls_for_each_item(): }, ] ex = _execution(node_id="agent-1", outputs={"files": file_payloads}) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service._resolve_preview_url", - side_effect=[ - "https://signed.example/1.pdf", - "https://signed.example/2.pdf", - "https://signed.example/1-detail.pdf", - "https://signed.example/2-detail.pdf", - "https://signed.example/1-full.pdf", - "https://signed.example/2-full.pdf", - ], - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service._resolve_preview_url", + side_effect=[ + "https://signed.example/1.pdf", + "https://signed.example/2.pdf", + "https://signed.example/1-detail.pdf", + "https://signed.example/2-detail.pdf", + "https://signed.example/1-full.pdf", + "https://signed.example/2-full.pdf", + ], ): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) preview = service.output_preview( app_model=_app_model(), workflow_run_id="run-1", node_id="agent-1", output_name="files", + session=session, ) snapshot_value = snapshot.node_outputs[0].outputs[0].value_preview @@ -531,14 +518,12 @@ def test_file_output_preview_uses_none_when_signed_url_resolution_fails(): "reference": build_file_reference(record_id="550e8400-e29b-41d4-a716-446655440000"), } ex = _execution(node_id="agent-1", outputs={"report": file_payload}) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service._resolve_preview_url", - side_effect=RuntimeError("boom"), - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service._resolve_preview_url", + side_effect=RuntimeError("boom"), ): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) preview_value = snapshot.node_outputs[0].outputs[0].value_preview assert isinstance(preview_value, dict) @@ -557,19 +542,18 @@ def test_object_output_preview_does_not_augment_canonical_file_mapping_shape(): "reference": build_file_reference(record_id="550e8400-e29b-41d4-a716-446655440000"), } ex = _execution(node_id="agent-1", outputs={"meta": raw_value}) - with ( - _patch_session(workflow_run=run, executions=[ex]), - patch( - "services.workflow.node_output_inspector_service._resolve_preview_url", - return_value="https://signed.example/x.pdf", - ), + session = _mock_session(workflow_run=run, executions=[ex]) + with patch( + "services.workflow.node_output_inspector_service._resolve_preview_url", + return_value="https://signed.example/x.pdf", ): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) preview = service.output_preview( app_model=_app_model(), workflow_run_id="run-1", node_id="agent-1", output_name="meta", + session=session, ) assert snapshot.node_outputs[0].outputs[0].value_preview == raw_value @@ -591,8 +575,8 @@ def test_retried_count_pulled_from_attempt_metadata(): outputs={"text": "ok"}, execution_metadata={"attempt": 2}, ) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.node_outputs[0].outputs[0].retried == 2 @@ -610,8 +594,8 @@ def test_keeps_latest_execution_per_node_by_index(): run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) older = _execution(node_id="agent-1", outputs={"text": "old"}, index=1) newer = _execution(node_id="agent-1", outputs={"text": "new"}, index=5) - with _patch_session(workflow_run=run, executions=[older, newer]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[older, newer]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.node_outputs[0].outputs[0].value_preview == "new" @@ -632,8 +616,8 @@ def test_array_typed_output_with_array_item_renders_correctly(): ) run = _workflow_run(nodes=[_agent_v2_node(node_id="agent-1")]) ex = _execution(node_id="agent-1", outputs={"files": []}) - with _patch_session(workflow_run=run, executions=[ex]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[ex]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) output = snapshot.node_outputs[0].outputs[0] assert output.type == DeclaredOutputType.ARRAY @@ -654,6 +638,6 @@ def test_unparseable_graph_blob_yields_empty_snapshot_not_500(): status=WorkflowExecutionStatus.RUNNING, graph="{not valid json", ) - with _patch_session(workflow_run=run, executions=[]): - snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1") + session = _mock_session(workflow_run=run, executions=[]) + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) assert snapshot.node_outputs == [] diff --git a/api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py b/api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py index 2aaf3bdf1d5..f471e4aeb56 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py @@ -118,6 +118,7 @@ def test__convert_to_http_request_node_for_chatbot(default_variables: list[Varia app_model=app_model, variables=default_variables, external_data_variables=external_data_variables, + session=MagicMock(), ) assert len(nodes) == 2 @@ -160,6 +161,7 @@ def test__convert_to_http_request_node_for_workflow_app(default_variables: list[ app_model=app_model, variables=default_variables, external_data_variables=external_data_variables, + session=MagicMock(), ) body = json.loads(nodes[0]["data"]["body"]["data"]) @@ -364,6 +366,7 @@ def test_convert_to_workflow_should_raise_when_app_model_config_is_missing(conve icon_type="emoji", icon="robot", icon_background="#fff", + session=MagicMock(), ) @@ -389,7 +392,6 @@ def test_convert_to_workflow_should_create_new_app_with_fallback_fields( monkeypatch.setattr(converter_module, "App", FakeApp) db_session = SimpleNamespace(add=MagicMock(), flush=MagicMock(), commit=MagicMock()) - monkeypatch.setattr(converter_module, "db", SimpleNamespace(session=db_session)) send_mock = MagicMock() monkeypatch.setattr(converter_module.app_was_created, "send", send_mock) @@ -417,6 +419,7 @@ def test_convert_to_workflow_should_create_new_app_with_fallback_fields( icon_type="", icon="", icon_background="", + session=db_session, ) assert new_app.name == "Source App(workflow)" @@ -501,12 +504,12 @@ def test_convert_app_model_config_to_workflow_should_build_advanced_chat_graph_a monkeypatch.setattr(converter_module, "Workflow", FakeWorkflow) db_session = SimpleNamespace(add=MagicMock(), commit=MagicMock()) - monkeypatch.setattr(converter_module, "db", SimpleNamespace(session=db_session)) workflow = converter.convert_app_model_config_to_workflow( app_model=app_model, app_model_config=_app_model_config(id="cfg"), account_id="account-1", + session=db_session, ) graph = json.loads(workflow.graph) @@ -568,12 +571,12 @@ def test_convert_app_model_config_to_workflow_should_build_workflow_mode_with_en monkeypatch.setattr(converter_module, "Workflow", FakeWorkflow) db_session = SimpleNamespace(add=MagicMock(), commit=MagicMock()) - monkeypatch.setattr(converter_module, "db", SimpleNamespace(session=db_session)) workflow = converter.convert_app_model_config_to_workflow( app_model=app_model, app_model_config=_app_model_config(id="cfg"), account_id="account-1", + session=db_session, ) graph = json.loads(workflow.graph) @@ -644,6 +647,7 @@ def test_convert_to_http_request_node_should_skip_non_api_and_missing_extension_ app_model=app_model, variables=[], external_data_variables=external_data_variables, + session=MagicMock(), ) assert nodes == [] @@ -810,10 +814,9 @@ def test_get_api_based_extension_should_raise_when_extension_not_found( monkeypatch: pytest.MonkeyPatch, ) -> None: db_session = SimpleNamespace(scalar=MagicMock(return_value=None)) - monkeypatch.setattr(converter_module, "db", SimpleNamespace(session=db_session)) with pytest.raises(ValueError, match="API Based Extension not found"): - converter._get_api_based_extension(tenant_id="tenant-1", api_based_extension_id="ext-1") + converter._get_api_based_extension(tenant_id="tenant-1", api_based_extension_id="ext-1", session=db_session) db_session.scalar.assert_called_once() @@ -823,9 +826,10 @@ def test_get_api_based_extension_should_return_entity_when_found( ) -> None: extension = SimpleNamespace(id="ext-1") db_session = SimpleNamespace(scalar=MagicMock(return_value=extension)) - monkeypatch.setattr(converter_module, "db", SimpleNamespace(session=db_session)) - result = converter._get_api_based_extension(tenant_id="tenant-1", api_based_extension_id="ext-1") + result = converter._get_api_based_extension( + tenant_id="tenant-1", api_based_extension_id="ext-1", session=db_session + ) assert result is extension db_session.scalar.assert_called_once() diff --git a/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py b/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py index 5bcb13c360c..cd97fa2e53b 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py @@ -64,6 +64,7 @@ def test_human_input_delivery_requires_draft_workflow(): account=account, node_id="node-1", delivery_method_id="delivery-1", + session=MagicMock(), ) @@ -98,6 +99,7 @@ def test_human_input_delivery_allows_disabled_method(monkeypatch: pytest.MonkeyP account=account, node_id="node-1", delivery_method_id=str(delivery_method.id), + session=MagicMock(), ) test_service_instance.send_test.assert_called_once() @@ -135,6 +137,7 @@ def test_human_input_delivery_dispatches_to_test_service(monkeypatch: pytest.Mon node_id="node-1", delivery_method_id=str(delivery_method.id), inputs={"#node-1.output#": "value"}, + session=MagicMock(), ) pool_args = service._build_human_input_variable_pool.call_args.kwargs @@ -173,6 +176,7 @@ def test_human_input_delivery_debug_mode_overrides_recipients(monkeypatch: pytes account=account, node_id="node-1", delivery_method_id=str(delivery_method.id), + session=MagicMock(), ) test_service_instance.send_test.assert_called_once() diff --git a/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py b/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py new file mode 100644 index 00000000000..f7bc8a4a9ac --- /dev/null +++ b/api/tests/unit_tests/tasks/test_agent_backend_session_cleanup_task.py @@ -0,0 +1,51 @@ +import logging + +from agenton.compositor import CompositorSessionSnapshot + +from clients.agent_backend.session_cleanup import ( + AgentBackendSessionCleanupPayload, + AgentBackendSessionCleanupResult, +) +from tasks import agent_backend_session_cleanup_task as cleanup_task_module + + +def _payload_dict() -> dict[str, object]: + return AgentBackendSessionCleanupPayload( + session_snapshot=CompositorSessionSnapshot(layers=[]), + runtime_layer_specs=[], + metadata={"tenant_id": "tenant-1", "app_id": "app-1"}, + ).model_dump(mode="json") + + +def test_run_cleanup_task_logs_info_for_skipped_result(monkeypatch, caplog): + monkeypatch.setattr(cleanup_task_module, "_create_agent_backend_client", lambda: object()) + monkeypatch.setattr( + cleanup_task_module, + "cleanup_agent_backend_session", + lambda **kwargs: AgentBackendSessionCleanupResult.skipped("missing_runtime_layer_specs"), + ) + + with caplog.at_level(logging.INFO, logger="tasks.agent_backend_session_cleanup_task"): + cleanup_task_module._run_cleanup_task(_payload_dict()) + + assert "Agent backend session cleanup skipped" in caplog.text + assert "missing_runtime_layer_specs" in caplog.text + + +def test_run_cleanup_task_logs_warning_for_failed_result(monkeypatch, caplog): + monkeypatch.setattr(cleanup_task_module, "_create_agent_backend_client", lambda: object()) + monkeypatch.setattr( + cleanup_task_module, + "cleanup_agent_backend_session", + lambda **kwargs: AgentBackendSessionCleanupResult.failed( + "backend exploded", + cleanup_run_id="cleanup-run-1", + ), + ) + + with caplog.at_level(logging.WARNING, logger="tasks.agent_backend_session_cleanup_task"): + cleanup_task_module._run_cleanup_task(_payload_dict()) + + assert "Agent backend session cleanup failed" in caplog.text + assert "backend exploded" in caplog.text + assert "cleanup-run-1" in caplog.text diff --git a/api/tests/unit_tests/tasks/test_process_tenant_plugin_autoupgrade_check_task.py b/api/tests/unit_tests/tasks/test_process_tenant_plugin_autoupgrade_check_task.py index a4412c1ee93..d26ae0b5f13 100644 --- a/api/tests/unit_tests/tasks/test_process_tenant_plugin_autoupgrade_check_task.py +++ b/api/tests/unit_tests/tasks/test_process_tenant_plugin_autoupgrade_check_task.py @@ -5,7 +5,11 @@ from unittest.mock import MagicMock, patch from core.plugin.entities.marketplace import MarketplacePluginSnapshot from core.plugin.entities.plugin import PluginCategory, PluginInstallationSource -from models.account import TenantPluginAutoUpgradeStrategy +from models.account import ( + TenantPluginAutoUpgradeCategory, + TenantPluginAutoUpgradeMode, + TenantPluginAutoUpgradeStrategySetting, +) MODULE = "tasks.process_tenant_plugin_autoupgrade_check_task" @@ -41,8 +45,8 @@ def _run_task( *, plugins: list, manifests: list[MarketplacePluginSnapshot], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, exclude_plugins=None, include_plugins=None, category=None, @@ -121,9 +125,9 @@ class TestUpgradeCallsMarketplaceService: process_tenant_plugin_autoupgrade_check_task( "tenant-1", - TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + TenantPluginAutoUpgradeStrategySetting.LATEST, 0, - TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + TenantPluginAutoUpgradeMode.ALL, [], [], ) @@ -136,7 +140,7 @@ class TestStrategySetting: upgrade_mock, _ = _run_task( plugins=[_make_plugin("acme/foo", "1.0.0")], manifests=[_make_manifest("acme/foo", "1.0.1")], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.DISABLED, ) upgrade_mock.assert_not_called() @@ -144,7 +148,7 @@ class TestStrategySetting: upgrade_mock, calls = _run_task( plugins=[_make_plugin("acme/foo", "1.0.0")], manifests=[_make_manifest("acme/foo", "1.0.5")], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, ) upgrade_mock.assert_called_once() assert calls[0][2].endswith(":1.0.5@cafe1234") @@ -153,7 +157,7 @@ class TestStrategySetting: upgrade_mock, _ = _run_task( plugins=[_make_plugin("acme/foo", "1.0.0")], manifests=[_make_manifest("acme/foo", "1.1.0")], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, ) upgrade_mock.assert_not_called() @@ -161,7 +165,7 @@ class TestStrategySetting: upgrade_mock, _ = _run_task( plugins=[_make_plugin("acme/foo", "1.0.0")], manifests=[_make_manifest("acme/foo", "2.0.0")], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.FIX_ONLY, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, ) upgrade_mock.assert_not_called() @@ -169,7 +173,7 @@ class TestStrategySetting: upgrade_mock, _ = _run_task( plugins=[_make_plugin("acme/foo", "1.0.0")], manifests=[_make_manifest("acme/foo", "1.0.0")], - strategy_setting=TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, ) upgrade_mock.assert_not_called() @@ -188,7 +192,7 @@ class TestUpgradeMode: upgrade_mock, calls = _run_task( plugins=plugins, manifests=manifests, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, ) assert upgrade_mock.call_count == 2 @@ -208,7 +212,7 @@ class TestUpgradeMode: upgrade_mock, calls = _run_task( plugins=plugins, manifests=manifests, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, ) assert upgrade_mock.call_count == 1 @@ -227,7 +231,7 @@ class TestUpgradeMode: upgrade_mock, calls = _run_task( plugins=plugins, manifests=manifests, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.PARTIAL, + upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, include_plugins=["acme/foo"], ) @@ -247,7 +251,7 @@ class TestUpgradeMode: upgrade_mock, calls = _run_task( plugins=plugins, manifests=manifests, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.EXCLUDE, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, exclude_plugins=["acme/bar"], ) @@ -267,8 +271,8 @@ class TestUpgradeMode: upgrade_mock, calls = _run_task( plugins=plugins, manifests=manifests, - upgrade_mode=TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, - category=TenantPluginAutoUpgradeStrategy.PluginCategory.MODEL, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, + category=TenantPluginAutoUpgradeCategory.MODEL, ) upgrade_mock.assert_called_once() @@ -306,9 +310,9 @@ class TestErrorIsolation: process_tenant_plugin_autoupgrade_check_task( "tenant-1", - TenantPluginAutoUpgradeStrategy.StrategySetting.LATEST, + TenantPluginAutoUpgradeStrategySetting.LATEST, 0, - TenantPluginAutoUpgradeStrategy.UpgradeMode.ALL, + TenantPluginAutoUpgradeMode.ALL, [], [], ) diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 9fc94547468..1e4b37612ba 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -1,10 +1,16 @@ import logging +from collections.abc import Generator from unittest.mock import MagicMock, call, patch import pytest +from agenton.compositor import CompositorSessionSnapshot +from sqlalchemy import delete, select +from core.db.session_factory import session_factory from libs.archive_storage import ArchiveStorageNotConfiguredError +from models import AgentRuntimeSession, AgentRuntimeSessionOwnerType, AgentRuntimeSessionStatus from tasks.remove_app_and_related_data_task import ( + _cleanup_active_agent_runtime_sessions_for_app, _delete_app_stars, _delete_app_workflow_archive_logs, _delete_archived_workflow_run_files, @@ -14,6 +20,29 @@ from tasks.remove_app_and_related_data_task import ( ) +@pytest.fixture(autouse=True) +def _create_agent_runtime_sessions_table() -> Generator[None, None, None]: + engine = session_factory.get_session_maker().kw["bind"] + AgentRuntimeSession.__table__.create(bind=engine, checkfirst=True) + yield + with session_factory.create_session() as session: + session.execute(delete(AgentRuntimeSession)) + session.commit() + AgentRuntimeSession.__table__.drop(bind=engine, checkfirst=True) + + +def _runtime_session_specs_json() -> str: + return ( + '[{"name":"execution_context","type":"dify.execution_context","deps":{},"metadata":{},' + '"config":{"tenant_id":"tenant-1"}},{"name":"history","type":"pydantic_ai.history","deps":{},' + '"metadata":{},"config":null}]' + ) + + +def _snapshot_json() -> str: + return CompositorSessionSnapshot(layers=[]).model_dump_json() + + class TestDeleteDraftVariablesBatch: def test_delete_draft_variables_batch_invalid_batch_size(self): """Test that invalid batch size raises ValueError.""" @@ -149,3 +178,215 @@ class TestDeleteArchivedWorkflowRunFiles: storage.list_objects.assert_called_once_with("tenant-1/app_id=app-1/") storage.delete_object.assert_has_calls([call("key-1"), call("key-2")], any_order=False) assert "Deleted 2 archive objects for app app-1" in caplog.text + + +class TestCleanupActiveAgentRuntimeSessionsForApp: + @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") + @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") + def test_enqueues_cleanup_for_active_rows_and_marks_rows_cleaned( + self, + mock_conversation_cleanup, + mock_workflow_cleanup, + ): + with session_factory.create_session() as session: + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + agent_config_snapshot_id="snap-1", + conversation_id="conv-1", + backend_run_id="run-conv", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, + agent_id="agent-2", + workflow_id="wf-1", + workflow_run_id="wf-run-1", + node_id="node-1", + backend_run_id="run-wf", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="other-app", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-3", + conversation_id="conv-2", + backend_run_id="run-other", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.commit() + + _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1", batch_size=1) + + assert mock_conversation_cleanup.delay.call_count == 1 + assert mock_workflow_cleanup.delay.call_count == 1 + conversation_payload = mock_conversation_cleanup.delay.call_args.args[0] + workflow_payload = mock_workflow_cleanup.delay.call_args.args[0] + assert conversation_payload["metadata"]["conversation_id"] == "conv-1" + assert workflow_payload["metadata"]["workflow_run_id"] == "wf-run-1" + assert conversation_payload["idempotency_key"].startswith("tenant-1:app-1:conv-1:agent-1:app-delete-cleanup:") + assert workflow_payload["idempotency_key"].startswith( + "tenant-1:app-1:wf-run-1:node-1:agent-2:app-delete-cleanup:" + ) + with session_factory.create_session() as session: + app_rows = session.scalars( + select(AgentRuntimeSession).where( + AgentRuntimeSession.tenant_id == "tenant-1", + AgentRuntimeSession.app_id == "app-1", + ) + ).all() + assert {row.status for row in app_rows} == {AgentRuntimeSessionStatus.CLEANED} + other_row = session.scalar( + select(AgentRuntimeSession).where( + AgentRuntimeSession.tenant_id == "tenant-1", + AgentRuntimeSession.app_id == "other-app", + ) + ) + assert other_row is not None + assert other_row.status == AgentRuntimeSessionStatus.ACTIVE + + @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") + @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") + def test_marks_rows_cleaned_even_when_enqueue_fails( + self, + mock_conversation_cleanup, + mock_workflow_cleanup, + ): + mock_conversation_cleanup.delay.side_effect = RuntimeError("queue down") + with session_factory.create_session() as session: + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + conversation_id="conv-1", + backend_run_id="run-conv", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.add( + AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.WORKFLOW_RUN, + agent_id="agent-2", + workflow_id="wf-1", + workflow_run_id="wf-run-1", + node_id="node-1", + backend_run_id="run-wf", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + ) + session.commit() + + _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") + + mock_conversation_cleanup.delay.assert_called_once() + mock_workflow_cleanup.delay.assert_called_once() + with session_factory.create_session() as session: + rows = session.scalars( + select(AgentRuntimeSession).where( + AgentRuntimeSession.tenant_id == "tenant-1", + AgentRuntimeSession.app_id == "app-1", + ) + ).all() + assert rows + assert {row.status for row in rows} == {AgentRuntimeSessionStatus.CLEANED} + + @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") + @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") + def test_uses_row_identity_to_keep_distinct_app_delete_cleanup_jobs_distinct( + self, + mock_conversation_cleanup, + mock_workflow_cleanup, + ): + del mock_workflow_cleanup + with session_factory.create_session() as session: + first = AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + agent_config_snapshot_id="snap-1", + conversation_id="conv-1", + backend_run_id="run-conv-1", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + second = AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + agent_config_snapshot_id="snap-2", + conversation_id="conv-1", + backend_run_id="run-conv-2", + session_snapshot=_snapshot_json(), + composition_layer_specs=_runtime_session_specs_json(), + status=AgentRuntimeSessionStatus.ACTIVE, + ) + session.add(first) + session.add(second) + session.commit() + + _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") + + assert mock_conversation_cleanup.delay.call_count == 2 + payloads = [queued_call.args[0] for queued_call in mock_conversation_cleanup.delay.call_args_list] + assert payloads[0]["idempotency_key"] != payloads[1]["idempotency_key"] + assert payloads[0]["idempotency_key"].endswith(first.id) + assert payloads[1]["idempotency_key"].endswith(second.id) + + @patch("tasks.remove_app_and_related_data_task.cleanup_workflow_agent_runtime_session") + @patch("tasks.remove_app_and_related_data_task.cleanup_conversation_agent_runtime_session") + def test_marks_empty_runtime_layer_specs_rows_clean_without_enqueue( + self, + mock_conversation_cleanup, + mock_workflow_cleanup, + ): + del mock_workflow_cleanup + with session_factory.create_session() as session: + row = AgentRuntimeSession( + tenant_id="tenant-1", + app_id="app-1", + owner_type=AgentRuntimeSessionOwnerType.CONVERSATION, + agent_id="agent-1", + conversation_id="conv-1", + backend_run_id="run-no-specs", + session_snapshot=_snapshot_json(), + composition_layer_specs="[]", + status=AgentRuntimeSessionStatus.ACTIVE, + ) + session.add(row) + session.commit() + + _cleanup_active_agent_runtime_sessions_for_app("tenant-1", "app-1") + + mock_conversation_cleanup.delay.assert_not_called() + with session_factory.create_session() as session: + stored_row = session.scalar(select(AgentRuntimeSession).where(AgentRuntimeSession.id == row.id)) + assert stored_row is not None + assert stored_row.status == AgentRuntimeSessionStatus.CLEANED diff --git a/api/tests/unit_tests/tasks/test_workflow_execute_task.py b/api/tests/unit_tests/tasks/test_workflow_execute_task.py index 40965096b39..a3fd70f205f 100644 --- a/api/tests/unit_tests/tasks/test_workflow_execute_task.py +++ b/api/tests/unit_tests/tasks/test_workflow_execute_task.py @@ -723,6 +723,7 @@ def test_resume_advanced_chat_publishes_events_for_originally_blocking_runs(monk message=MagicMock(), generate_entity=generate_entity, graph_runtime_state=MagicMock(), + response_stream_filter=MagicMock(), session_factory=MagicMock(), pause_state_config=MagicMock(), workflow_run_id="workflow-run-id", @@ -774,6 +775,7 @@ def test_resume_workflow_publishes_events_for_originally_blocking_runs(monkeypat user=MagicMock(), generate_entity=generate_entity, graph_runtime_state=MagicMock(), + response_stream_filter=MagicMock(), session_factory=MagicMock(), pause_state_config=MagicMock(), workflow_run_id="workflow-run-id", @@ -829,6 +831,7 @@ def test_resume_workflow_ignores_missing_old_pause_after_repause(monkeypatch: py user=MagicMock(), generate_entity=generate_entity, graph_runtime_state=MagicMock(), + response_stream_filter=MagicMock(), session_factory=MagicMock(), pause_state_config=MagicMock(), workflow_run_id="workflow-run-id", diff --git a/cli/package.json b/cli/package.json index 5121daf3624..5610150789e 100644 --- a/cli/package.json +++ b/cli/package.json @@ -1,13 +1,13 @@ { "name": "@langgenius/difyctl", "type": "module", - "version": "0.1.0-alpha", + "version": "0.2.0-alpha", "description": "Dify command-line interface", "difyctl": { "channel": "alpha", "compat": { - "minDify": "1.15.0", - "maxDify": "1.15.0" + "minDify": "1.16.0", + "maxDify": "1.16.0" }, "release": { "tagPrefix": "difyctl-v", @@ -67,7 +67,7 @@ "test:e2e:local": "DIFY_E2E_MODE=local vp test --config vitest.e2e.config.ts", "lint": "eslint", "lint:fix": "eslint --fix", - "type-check": "tsgo", + "type-check": "tsc", "tree:gen": "bun scripts/generate-command-tree.ts", "tree:check": "bun scripts/generate-command-tree.ts --check", "prebuild": "pnpm tree:gen", @@ -101,7 +101,7 @@ "@types/js-yaml": "catalog:", "@types/lockfile": "catalog:", "@types/node": "catalog:", - "@typescript/native-preview": "catalog:", + "@typescript/native": "catalog:", "@vitest/coverage-v8": "catalog:", "eslint": "catalog:", "hono": "catalog:", diff --git a/cli/scripts/release-naming.test.ts b/cli/scripts/release-naming.test.ts index 559d15ad759..7641b531282 100644 --- a/cli/scripts/release-naming.test.ts +++ b/cli/scripts/release-naming.test.ts @@ -15,41 +15,41 @@ function run(args: string[]): { code: number, stdout: string, stderr: string } { } } -describe('release-naming compat-check (compat 1.15.0..1.15.0)', () => { +describe('release-naming compat-check (compat 1.16.0..1.16.0)', () => { it('accepts a version inside the window', () => { - expect(run(['compat-check', '1.15.0']).code).toBe(0) + expect(run(['compat-check', '1.16.0']).code).toBe(0) }) it('accepts the inclusive lower bound', () => { - expect(run(['compat-check', '1.15.0']).code).toBe(0) + expect(run(['compat-check', '1.16.0']).code).toBe(0) }) it('accepts the inclusive upper bound', () => { - expect(run(['compat-check', '1.15.0']).code).toBe(0) + expect(run(['compat-check', '1.16.0']).code).toBe(0) }) it('accepts a v-prefixed tag', () => { - expect(run(['compat-check', 'v1.15.0']).code).toBe(0) + expect(run(['compat-check', 'v1.16.0']).code).toBe(0) }) it('rejects a version below the lower bound', () => { - expect(run(['compat-check', '1.14.9']).code).not.toBe(0) + expect(run(['compat-check', '1.15.9']).code).not.toBe(0) }) it('rejects a version above the upper bound', () => { - expect(run(['compat-check', '1.15.1']).code).not.toBe(0) + expect(run(['compat-check', '1.16.1']).code).not.toBe(0) }) - it('treats a prerelease of the bound as below it (1.15.0-rc1 < 1.15.0)', () => { - expect(run(['compat-check', '1.15.0-rc1']).code).not.toBe(0) + it('treats a prerelease of the bound as below it (1.16.0-rc1 < 1.16.0)', () => { + expect(run(['compat-check', '1.16.0-rc1']).code).not.toBe(0) }) - it('ignores build metadata on the bound (1.15.0+build == 1.15.0)', () => { - expect(run(['compat-check', '1.15.0+build123']).code).toBe(0) + it('ignores build metadata on the bound (1.16.0+build == 1.16.0)', () => { + expect(run(['compat-check', '1.16.0+build123']).code).toBe(0) }) - it('ignores build metadata when out of range (1.15.1+build still rejected)', () => { - expect(run(['compat-check', '1.15.1+build123']).code).not.toBe(0) + it('ignores build metadata when out of range (1.16.1+build still rejected)', () => { + expect(run(['compat-check', '1.16.1+build123']).code).not.toBe(0) }) it('requires a version argument', () => { @@ -60,7 +60,7 @@ describe('release-naming compat-check (compat 1.15.0..1.15.0)', () => { describe('release-naming github-env', () => { it('emits difyctlTag = tagPrefix + version', () => { const { stdout } = run(['github-env']) - expect(stdout).toMatch(/^difyctlTag=difyctl-v0\.1\.0-alpha$/m) + expect(stdout).toMatch(/^difyctlTag=difyctl-v0\.2\.0-alpha$/m) }) it('still emits the existing trace fields', () => { @@ -75,14 +75,14 @@ describe('release-naming edge channel', () => { expect(run(['channels']).stdout).toMatch(/^edge$/m) }) - it('edge-version derives -edge. stripping the alpha prerelease', () => { - // package.json version is 0.1.0-alpha -> core 0.1.0 - expect(run(['edge-version', '2fd7b82']).stdout.trim()).toBe('0.1.0-edge.2fd7b82') + it('edge-version derives -edge. from the package version', () => { + // package.json version is 0.2.0-alpha -> core 0.2.0 + expect(run(['edge-version', '2fd7b82']).stdout.trim()).toBe('0.2.0-edge.2fd7b82') }) it('edge-version accepts a 40-char sha', () => { const sha = '2fd7b829e1f0aaaabbbbccccddddeeeeffff0000' - expect(run(['edge-version', sha]).stdout.trim()).toBe(`0.1.0-edge.${sha}`) + expect(run(['edge-version', sha]).stdout.trim()).toBe(`0.2.0-edge.${sha}`) }) it('edge-version rejects a non-hex sha', () => { diff --git a/cli/scripts/release-r2-edge.test.ts b/cli/scripts/release-r2-edge.test.ts index 12e3ee1822c..9db0c0ad2cb 100644 --- a/cli/scripts/release-r2-edge.test.ts +++ b/cli/scripts/release-r2-edge.test.ts @@ -81,7 +81,7 @@ describe('release-r2-edge manifest', () => { it('carries the compat window from package.json', () => { const { json } = buildManifest() - expect(json.compat).toEqual({ minDify: '1.15.0', maxDify: '1.15.0' }) + expect(json.compat).toEqual({ minDify: '1.16.0', maxDify: '1.16.0' }) }) it('lists all 5 targets with asset name + sha256 from the checksums file', () => { diff --git a/cli/src/api/app-dsl.test.ts b/cli/src/api/app-dsl.test.ts index 66d4507317e..1afff1fcbce 100644 --- a/cli/src/api/app-dsl.test.ts +++ b/cli/src/api/app-dsl.test.ts @@ -26,7 +26,7 @@ describe('AppDslClient.exportDsl', () => { const yaml = await makeClient(stub.url).exportDsl('app-1') expect(stub.captured.method).toBe('GET') - expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1/export') + expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1/dsl') expect(yaml).toBe(DSL_YAML) }) @@ -90,7 +90,7 @@ describe('AppDslClient.confirmImport', () => { const result = await makeClient(stub.url).confirmImport('ws-1', 'imp-1') expect(stub.captured.method).toBe('POST') - expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1/apps/imports/imp-1/confirm') + expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1/apps/imports/imp-1:confirm') expect(result.status).toBe('completed') }) }) @@ -107,7 +107,7 @@ describe('AppDslClient.checkDependencies', () => { const result = await makeClient(stub.url).checkDependencies('app-1') - expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1/check-dependencies') + expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1/dependencies:check') expect(result.leaked_dependencies).toEqual([]) }) }) diff --git a/cli/src/api/app-dsl.ts b/cli/src/api/app-dsl.ts index acb10a95351..19c26f5d7fb 100644 --- a/cli/src/api/app-dsl.ts +++ b/cli/src/api/app-dsl.ts @@ -33,7 +33,7 @@ export class AppDslClient { } async exportDsl(appId: string, query?: ExportQuery): Promise { - const resp = await this.orpc.apps.byAppId.export.get({ + const resp = await this.orpc.apps.byAppId.dsl.get({ params: { app_id: appId }, query: query !== undefined ? { @@ -52,7 +52,7 @@ export class AppDslClient { } async checkDependencies(appId: string): Promise { - return this.orpc.apps.byAppId.checkDependencies.get({ + return this.orpc.apps.byAppId.dependencies.check.get({ params: { app_id: appId }, }) } diff --git a/cli/src/api/app-run.ts b/cli/src/api/app-run.ts index cbd36de2049..a75b3f0c63a 100644 --- a/cli/src/api/app-run.ts +++ b/cli/src/api/app-run.ts @@ -54,7 +54,7 @@ export class AppRunClient { body: Record, opts: StreamOptions = {}, ): Promise> { - const res = await this.http.stream(`apps/${encodeURIComponent(appId)}/run`, { + const res = await this.http.stream(`apps/${encodeURIComponent(appId)}:run`, { method: 'POST', json: body, headers: { Accept: 'text/event-stream' }, @@ -79,7 +79,7 @@ export class AppRunClient { action: string, inputs: Record, ): Promise { - await this.orpc.apps.byAppId.form.humanInput.byFormToken.post({ + await this.orpc.apps.byAppId.humanInputForms.byFormToken.submit.post({ params: { app_id: appId, form_token: formToken }, body: { action, inputs }, }) diff --git a/cli/src/api/apps.test.ts b/cli/src/api/apps.test.ts index 861f60feb26..ba9126b7b98 100644 --- a/cli/src/api/apps.test.ts +++ b/cli/src/api/apps.test.ts @@ -82,12 +82,12 @@ describe('AppsClient.describe', () => { await stub?.stop() }) - it('hits /apps//describe, omits workspace_id and fields when not given', async () => { + it('hits /apps/, omits workspace_id and fields when not given', async () => { stub = await startStubServer(cap => jsonResponder(200, DESCRIBE_BODY, cap)) const res = await makeClient(stub.url).describe('app-1') - expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1/describe') + expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app-1') const q = queryOf(stub.captured.url) expect(q.has('workspace_id')).toBe(false) expect(q.has('fields')).toBe(false) @@ -107,6 +107,6 @@ describe('AppsClient.describe', () => { await makeClient(stub.url).describe('app/with space') - expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app%2Fwith%20space/describe') + expect(stub.captured.url?.split('?')[0]).toBe('/openapi/v1/apps/app%2Fwith%20space') }) }) diff --git a/cli/src/api/apps.ts b/cli/src/api/apps.ts index 1189fdeaa06..fa29d66fd14 100644 --- a/cli/src/api/apps.ts +++ b/cli/src/api/apps.ts @@ -37,7 +37,7 @@ export class AppsClient implements AppReader { } async describe(appId: string, fields?: readonly string[]): Promise { - return this.orpc.apps.byAppId.describe.get({ + return this.orpc.apps.byAppId.get({ params: { app_id: appId }, query: { fields: fields !== undefined && fields.length > 0 ? fields.join(',') : undefined, diff --git a/cli/src/api/file-upload.test.ts b/cli/src/api/file-upload.test.ts index 018389916b0..602cf351e60 100644 --- a/cli/src/api/file-upload.test.ts +++ b/cli/src/api/file-upload.test.ts @@ -41,7 +41,7 @@ describe('FileUploadClient.upload', () => { const result = await makeClient(stub.url).upload('app-1', filePath) expect(stub.captured.method).toBe('POST') - expect(stub.captured.url).toBe('/openapi/v1/apps/app-1/files/upload') + expect(stub.captured.url).toBe('/openapi/v1/apps/app-1/files') // The client must let fetch own the multipart Content-Type + boundary; it // must NOT coerce this to application/json the way a json body would. const contentType = stub.captured.headers?.['content-type'] ?? '' @@ -61,7 +61,7 @@ describe('FileUploadClient.upload', () => { await makeClient(stub.url).upload('app/with space', filePath) - expect(stub.captured.url).toBe('/openapi/v1/apps/app%2Fwith%20space/files/upload') + expect(stub.captured.url).toBe('/openapi/v1/apps/app%2Fwith%20space/files') }) it('propagates a server 413 as a classified BaseError', async () => { diff --git a/cli/src/api/file-upload.ts b/cli/src/api/file-upload.ts index 7a032737a59..011f898c74e 100644 --- a/cli/src/api/file-upload.ts +++ b/cli/src/api/file-upload.ts @@ -65,7 +65,7 @@ export class FileUploadClient { form.append('file', blob, filename) return this.http.post( - `apps/${encodeURIComponent(appId)}/files/upload`, + `apps/${encodeURIComponent(appId)}/files`, { body: form, timeoutMs: 60_000 }, ) } diff --git a/cli/src/api/members.test.ts b/cli/src/api/members.test.ts index b4e01b76b24..a8fdc633f77 100644 --- a/cli/src/api/members.test.ts +++ b/cli/src/api/members.test.ts @@ -154,13 +154,13 @@ describe('MembersClient.updateRole', () => { await stub?.stop() }) - it('PUTs role payload to /role subresource', async () => { + it('PATCHes role payload to the member resource', async () => { stub = await startStubServer(cap => jsonResponder(200, { result: 'success' }, cap)) const result = await makeClient(stub.url).updateRole('ws-1', 'm-1', { role: 'admin' }) - expect(stub.captured.method).toBe('PUT') - expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1/members/m-1/role') + expect(stub.captured.method).toBe('PATCH') + expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1/members/m-1') expect(JSON.parse(stub.captured.body ?? '{}')).toEqual({ role: 'admin' }) expect(result.result).toBe('success') }) @@ -181,7 +181,7 @@ describe('WorkspacesClient.switch (integration with stub)', () => { await stub?.stop() }) - it('POSTs /workspaces//switch and returns workspace detail', async () => { + it('POSTs /workspaces/:switch and returns workspace detail', async () => { stub = await startStubServer(cap => jsonResponder( 200, @@ -200,7 +200,7 @@ describe('WorkspacesClient.switch (integration with stub)', () => { const result = await client.switch('ws-1') expect(stub.captured.method).toBe('POST') - expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1/switch') + expect(stub.captured.url).toBe('/openapi/v1/workspaces/ws-1:switch') expect(result.current).toBe(true) }) diff --git a/cli/src/api/members.ts b/cli/src/api/members.ts index 7b1f80c08ff..8a9bc13081f 100644 --- a/cli/src/api/members.ts +++ b/cli/src/api/members.ts @@ -47,7 +47,7 @@ export class MembersClient { memberId: string, payload: MemberRoleUpdatePayload, ): Promise { - return this.orpc.workspaces.byWorkspaceId.members.byMemberId.role.put({ + return this.orpc.workspaces.byWorkspaceId.members.byMemberId.patch({ params: { workspace_id: workspaceId, member_id: memberId }, body: payload, }) diff --git a/cli/src/api/permitted-external-apps.test.ts b/cli/src/api/permitted-external-apps.test.ts index f6fa38cb3eb..58f47b4566f 100644 --- a/cli/src/api/permitted-external-apps.test.ts +++ b/cli/src/api/permitted-external-apps.test.ts @@ -12,15 +12,15 @@ describe('PermittedExternalAppsClient', () => { it('list calls permittedExternalApps.get with paging/filter query', async () => { const c = new PermittedExternalAppsClient(fakeHttp()) const get = vi.fn().mockResolvedValue({ page: 1, limit: 20, total: 0, has_more: false, data: [] }) - ;(c as unknown as WithOrpc).orpc = { permittedExternalApps: { get, byAppId: { describe: { get: vi.fn() } } } } + ;(c as unknown as WithOrpc).orpc = { permittedExternalApps: { get, byAppId: { get: vi.fn() } } } await c.list({ workspaceId: '', page: 2, limit: 5, mode: undefined, name: 'a' }) expect(get).toHaveBeenCalledWith({ query: { page: 2, limit: 5, mode: undefined, name: 'a' } }) }) - it('describe calls permittedExternalApps.byAppId.describe.get with app_id + fields', async () => { + it('describe calls permittedExternalApps.byAppId.get with app_id + fields', async () => { const c = new PermittedExternalAppsClient(fakeHttp()) const dget = vi.fn().mockResolvedValue({ info: null, parameters: null, input_schema: null }) - ;(c as unknown as WithOrpc).orpc = { permittedExternalApps: { get: vi.fn(), byAppId: { describe: { get: dget } } } } + ;(c as unknown as WithOrpc).orpc = { permittedExternalApps: { get: vi.fn(), byAppId: { get: dget } } } await c.describe('app-1', ['info']) expect(dget).toHaveBeenCalledWith({ params: { app_id: 'app-1' }, query: { fields: 'info' } }) }) diff --git a/cli/src/api/permitted-external-apps.ts b/cli/src/api/permitted-external-apps.ts index 497c398d0ba..c0164d3c536 100644 --- a/cli/src/api/permitted-external-apps.ts +++ b/cli/src/api/permitted-external-apps.ts @@ -26,7 +26,7 @@ export class PermittedExternalAppsClient implements AppReader { } async describe(appId: string, fields?: readonly string[]): Promise { - return this.orpc.permittedExternalApps.byAppId.describe.get({ + return this.orpc.permittedExternalApps.byAppId.get({ params: { app_id: appId }, query: { fields: fields !== undefined && fields.length > 0 ? fields.join(',') : undefined }, }) diff --git a/cli/src/api/workspaces.ts b/cli/src/api/workspaces.ts index 3ef587b574d..08495a01fce 100644 --- a/cli/src/api/workspaces.ts +++ b/cli/src/api/workspaces.ts @@ -19,7 +19,7 @@ export class WorkspacesClient { /** * Server-side workspace switch via OpenAPI POST - * `/workspaces/{id}/switch` — the bearer-authed equivalent of the + * `/workspaces/{id}:switch` — the bearer-authed equivalent of the * console's POST `/workspaces/switch`. The server updates the caller's * `current` tenant_account_join row. Callers MUST refresh their local * `hosts.yml` only after this resolves — never fall back to a local diff --git a/cli/src/auth/hosts.test.ts b/cli/src/auth/hosts.test.ts index 112538ccbb9..337b785a4c4 100644 --- a/cli/src/auth/hosts.test.ts +++ b/cli/src/auth/hosts.test.ts @@ -113,6 +113,21 @@ describe('Registry (pure)', () => { expect(active?.ctx.account.email).toBe('a@x') }) + it('resolveActive returns the active context with insecureTls', () => { + const reg = baseReg() + reg.upsert('h1', 'a@x', ctx('a@x')) + reg.setInsecureTls('h1', true) + reg.setHost('h1') + reg.setAccount('a@x') + expect(reg.resolveActive()?.insecureTls).toBe(true) + }) + + it('setInsecureTls is a no-op for an unknown host', () => { + const reg = baseReg() + reg.setInsecureTls('missing', true) + expect(reg.hosts.missing).toBeUndefined() + }) + it('resolveActive returns undefined for each missing pointer', () => { const reg = baseReg() expect(reg.resolveActive()).toBeUndefined() diff --git a/cli/src/auth/hosts.ts b/cli/src/auth/hosts.ts index 29305db951e..5df8cdabdbe 100644 --- a/cli/src/auth/hosts.ts +++ b/cli/src/auth/hosts.ts @@ -41,6 +41,7 @@ export type AccountContext = z.infer export const HostEntrySchema = z.object({ scheme: z.string().optional(), + insecure_tls: z.boolean().optional(), current_account: z.string().optional(), accounts: z.record(z.string(), AccountContextSchema).default({}), }) @@ -58,6 +59,7 @@ export type ActiveContext = { readonly email: string readonly ctx: AccountContext readonly scheme?: string + readonly insecureTls?: boolean } export function notLoggedInError(hint = 'run \'difyctl auth login\''): BaseError { @@ -104,7 +106,7 @@ export class Registry { const ctx = entry.accounts[email] if (ctx === undefined) return undefined - return { host, email, ctx, scheme: entry.scheme } + return { host, email, ctx, scheme: entry.scheme, insecureTls: entry.insecure_tls } } requireActive(hint?: string): ActiveContext { @@ -157,6 +159,12 @@ export class Registry { entry.scheme = scheme } + setInsecureTls(host: string, insecure: boolean): void { + const entry = this.data.hosts[host] + if (entry !== undefined) + entry.insecure_tls = insecure + } + activate(host: string, email: string, ctx: AccountContext): void { this.upsert(host, email, ctx) this.setHost(host) diff --git a/cli/src/cache/compat-store.test.ts b/cli/src/cache/compat-store.test.ts new file mode 100644 index 00000000000..25ad11db06f --- /dev/null +++ b/cli/src/cache/compat-store.test.ts @@ -0,0 +1,57 @@ +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { ENV_CACHE_DIR } from '@/store/dir' +import { CACHE_COMPAT, getCache } from '@/store/manager' +import { loadCompatStore } from './compat-store' + +const HOST = 'https://cloud.dify.ai' +const NOW = new Date('2026-05-20T12:00:00.000Z') + +describe('compat-store', () => { + let dir: string + let prev: string | undefined + + beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'difyctl-compat-')) + prev = process.env[ENV_CACHE_DIR] + process.env[ENV_CACHE_DIR] = dir + }) + afterEach(async () => { + if (prev === undefined) + delete process.env[ENV_CACHE_DIR] + else + process.env[ENV_CACHE_DIR] = prev + await rm(dir, { recursive: true, force: true }) + }) + + const store = (now: Date = NOW) => loadCompatStore({ store: getCache(CACHE_COMPAT), now: () => now }) + + it('is not fresh before anything is marked', async () => { + expect((await store()).isFreshCompatible(HOST)).toBe(false) + }) + + it('is fresh right after markCompatible, and persists across loads', async () => { + await (await store()).markCompatible(HOST) + expect((await store()).isFreshCompatible(HOST)).toBe(true) + }) + + it('stays fresh within the 1h TTL', async () => { + const past = new Date(NOW.getTime() - 30 * 60 * 1000) + await (await store(past)).markCompatible(HOST) + expect((await store(NOW)).isFreshCompatible(HOST)).toBe(true) + }) + + it('expires after the 1h TTL', async () => { + const past = new Date(NOW.getTime() - 61 * 60 * 1000) + await (await store(past)).markCompatible(HOST) + expect((await store(NOW)).isFreshCompatible(HOST)).toBe(false) + }) + + it('tracks hosts independently', async () => { + const s = await store() + await s.markCompatible(HOST) + expect(s.isFreshCompatible('https://other.dify.ai')).toBe(false) + }) +}) diff --git a/cli/src/cache/compat-store.ts b/cli/src/cache/compat-store.ts new file mode 100644 index 00000000000..2df6dcf3377 --- /dev/null +++ b/cli/src/cache/compat-store.ts @@ -0,0 +1,71 @@ +import type { Store } from '@/store/store' +import { CACHE_COMPAT, getCache } from '@/store/manager' + +// How long a host stays "known compatible" before difyctl re-probes /_version. +export const COMPAT_TTL_MS = 60 * 60 * 1000 + +// Only *positive* (compatible) verdicts are cached — never "too old". A host that +// was too old is re-probed every time, so a just-upgraded server clears a previous +// block immediately instead of staying locked out for the whole TTL. +const COMPATIBLE_KEY = { key: 'compatible', default: {} as Record } as const + +export type CompatStore = { + readonly isFreshCompatible: (host: string, now?: Date) => boolean + readonly markCompatible: (host: string, now?: Date) => Promise +} + +export type CompatStoreOptions = { + readonly store?: Store + readonly now?: () => Date + readonly ttlMs?: number +} + +export async function loadCompatStore(opts: CompatStoreOptions = {}): Promise { + const store = opts.store ?? getCache(CACHE_COMPAT) + const ttlMs = opts.ttlMs ?? COMPAT_TTL_MS + const clock = opts.now ?? (() => new Date()) + const memory = await readCompatible(store) + + return { + isFreshCompatible: (host, now) => { + const last = memory.get(host) + if (last === undefined) + return false + const elapsed = Math.max(0, (now ?? clock()).getTime() - last) + return elapsed < ttlMs + }, + markCompatible: async (host, now) => { + const stamp = (now ?? clock()).getTime() + memory.set(host, stamp) + // Re-read disk inside the write cycle so concurrent processes touching + // different hosts don't clobber each other's stamps. + const onDisk = await readCompatible(store) + onDisk.set(host, stamp) + await writeCompatible(store, onDisk) + }, + } +} + +async function readCompatible(store: Store): Promise> { + const out = new Map() + let raw: Record + try { + raw = await store.get(COMPATIBLE_KEY) + } + catch { + return out + } + for (const [host, iso] of Object.entries(raw)) { + const t = Date.parse(iso) + if (!Number.isNaN(t)) + out.set(host, t) + } + return out +} + +async function writeCompatible(store: Store, state: Map): Promise { + const compatible: Record = {} + for (const [host, t] of state) + compatible[host] = new Date(t).toISOString() + await store.set(COMPATIBLE_KEY, compatible) +} diff --git a/cli/src/commands/_shared/authed-command.ts b/cli/src/commands/_shared/authed-command.ts index 8ef0b381e9e..d40b9342041 100644 --- a/cli/src/commands/_shared/authed-command.ts +++ b/cli/src/commands/_shared/authed-command.ts @@ -13,7 +13,8 @@ import { formatErrorForCli } from '@/errors/format' import { createHttpClient } from '@/http/client' import { getTokenStore } from '@/store/manager' import { realStreams } from '@/sys/io/streams' -import { hostWithScheme, openAPIBase } from '@/util/host' +import { activeHostInfo, openAPIBase } from '@/util/host' +import { enforceDifyVersion } from '@/version/enforce' import { versionInfo } from '@/version/info' import { maybeNudgeCompat } from '@/version/nudge' import { resolveRetryAttempts } from './global-flags.js' @@ -49,13 +50,17 @@ export async function buildAuthedContext( if (bearer === '') fail(cmd, opts, io) - const host = hostWithScheme(active.host, active.scheme) + const { host, insecure } = activeHostInfo(active) const retryAttempts = resolveRetryAttempts({ flag: opts.retryFlag, env: getEnv }) - const http = createHttpClient({ baseURL: openAPIBase(host), bearer, retryAttempts }) + const http = createHttpClient({ baseURL: openAPIBase(host), bearer, retryAttempts, insecure }) const cache = opts.withCache === true ? await loadAppInfoCache() : undefined - await runCompatNudge({ host, io }) + // Hard gate: refuse a server too old for this difyctl (throws → exit 6). + // Cached per host (1h) so most commands don't re-probe. Then the soft nudge + // handles the "server too new" direction. + await enforceDifyVersion(host, { insecure }) + await runCompatNudge({ host, insecure, io }) return { reg, active, store, http, host, io, cache } } @@ -69,6 +74,7 @@ function fail(cmd: Pick, opts: AuthedContextOptions, io: IOStr // command flows through it without per-command wiring. async function runCompatNudge(opts: { readonly host: string + readonly insecure: boolean readonly io: IOStreams }): Promise { try { @@ -76,7 +82,7 @@ async function runCompatNudge(opts: { await maybeNudgeCompat(opts.host, { store, probe: async (host) => { - const http = createHttpClient({ baseURL: openAPIBase(host), timeoutMs: META_PROBE_TIMEOUT_MS, retryAttempts: 0 }) + const http = createHttpClient({ baseURL: openAPIBase(host), timeoutMs: META_PROBE_TIMEOUT_MS, retryAttempts: 0, insecure: opts.insecure }) return new MetaClient(http).serverVersion() }, emit: line => opts.io.err.write(line), diff --git a/cli/src/commands/auth/login/index.ts b/cli/src/commands/auth/login/index.ts index d9b6dd27fc3..0b702262932 100644 --- a/cli/src/commands/auth/login/index.ts +++ b/cli/src/commands/auth/login/index.ts @@ -2,6 +2,7 @@ import type { CommandEffect } from '@/framework/command' import { DifyCommand } from '@/commands/_shared/dify-command' import { Flags } from '@/framework/flags' import { realStreams } from '@/sys/io/streams' +import { enforceDifyVersion } from '@/version/enforce' import { agentGuide } from './guide' import { runLogin } from './login' @@ -26,7 +27,7 @@ export default class Login extends DifyCommand { default: false, }), 'insecure': Flags.boolean({ - description: 'allow http:// hosts (local-dev only)', + description: 'allow http:// hosts and skip TLS certificate verification (local-dev only)', default: false, }), } @@ -38,6 +39,9 @@ export default class Login extends DifyCommand { host: flags.host, noBrowser: flags['no-browser'], insecure: flags.insecure, + verifyServer: async (host) => { + await enforceDifyVersion(host, { forceFresh: true, insecure: flags.insecure }) + }, }) } diff --git a/cli/src/commands/auth/login/login.ts b/cli/src/commands/auth/login/login.ts index 2c1ba5b95a9..7777f8c7c18 100644 --- a/cli/src/commands/auth/login/login.ts +++ b/cli/src/commands/auth/login/login.ts @@ -31,6 +31,10 @@ export type LoginOptions = { readonly browserEnv?: BrowserEnv readonly browserOpener?: BrowserOpener readonly clock?: Clock + // Version guard for the freshly-authenticated host; wired to enforceDifyVersion + // at the command boundary. Runs before the session is persisted so we never + // save credentials for a server too old for this difyctl. Defaults to a no-op. + readonly verifyServer?: (host: string) => Promise } export async function runLogin(opts: LoginOptions): Promise { @@ -40,7 +44,7 @@ export async function runLogin(opts: LoginOptions): Promise { const host = await resolveLoginHost(opts, insecure) const label = opts.deviceLabel ?? defaultDeviceLabel() - const api = opts.api ?? new DeviceFlowApi(createHttpClient({ baseURL: openAPIBase(host) })) + const api = opts.api ?? new DeviceFlowApi(createHttpClient({ baseURL: openAPIBase(host), insecure })) const code = await api.requestCode({ device_label: label }) renderCodePrompt(opts.io.err, cs, code) @@ -70,6 +74,9 @@ export async function runLogin(opts: LoginOptions): Promise { spinner.stop() } + // Refuse to persist a session to a server too old for this difyctl. + await (opts.verifyServer ?? (async () => {}))(host) + const storeBundle = opts.store ?? await detectTokenStore() const display = bareHost(host) const email = accountEmail(success) @@ -81,6 +88,7 @@ export async function runLogin(opts: LoginOptions): Promise { reg.token_storage = storeBundle.mode reg.activate(display, email, ctx) applyScheme(reg, display, host) + reg.setInsecureTls(display, insecure) await reg.save() renderLoggedIn(opts.io.out, cs, host, success) diff --git a/cli/src/commands/auth/logout/index.ts b/cli/src/commands/auth/logout/index.ts index 6476b1726e8..786d525c748 100644 --- a/cli/src/commands/auth/logout/index.ts +++ b/cli/src/commands/auth/logout/index.ts @@ -6,7 +6,7 @@ import { createHttpClient } from '@/http/client' import { getTokenStore } from '@/store/manager' import { runWithSpinner } from '@/sys/io/spinner' import { realStreams } from '@/sys/io/streams' -import { hostWithScheme, openAPIBase } from '@/util/host' +import { activeHostInfo, openAPIBase } from '@/util/host' import { runLogout } from './logout.js' export default class Logout extends DifyCommand { @@ -32,7 +32,8 @@ export default class Logout extends DifyCommand { } catch { /* keyring locked — skip remote revocation, local cleanup still runs */ } if (bearer !== '') { - http = createHttpClient({ baseURL: openAPIBase(hostWithScheme(active.host, active.scheme)), bearer, retryAttempts: 0 }) + const { host, insecure } = activeHostInfo(active) + http = createHttpClient({ baseURL: openAPIBase(host), bearer, retryAttempts: 0, insecure }) } } diff --git a/cli/src/commands/resume/app/run.test.ts b/cli/src/commands/resume/app/run.test.ts index a72b0a93b25..8e1882906ae 100644 --- a/cli/src/commands/resume/app/run.test.ts +++ b/cli/src/commands/resume/app/run.test.ts @@ -41,7 +41,7 @@ describe('resumeApp pre-flight subject strategy', () => { const http = { baseURL: 'http://localhost', request: vi.fn().mockImplementation((opts: { path: string }) => { - if (typeof opts.path === 'string' && opts.path.includes('form/human_input')) { + if (typeof opts.path === 'string' && opts.path.includes('human-input-forms')) { return Promise.resolve(FORM_RESP) } // reconnect stream — return an async iterable that ends immediately diff --git a/cli/src/commands/resume/app/run.ts b/cli/src/commands/resume/app/run.ts index 1dc2855a118..0610f68d5e6 100644 --- a/cli/src/commands/resume/app/run.ts +++ b/cli/src/commands/resume/app/run.ts @@ -50,7 +50,7 @@ export async function resumeApp(opts: ResumeAppOptions, deps: ResumeAppDeps): Pr let action = opts.action if (action === undefined) { const formResp = await deps.http.get<{ user_actions: { id: string }[] }>( - `apps/${encodeURIComponent(opts.appId)}/form/human_input/${encodeURIComponent(opts.formToken)}`, + `apps/${encodeURIComponent(opts.appId)}/human-input-forms/${encodeURIComponent(opts.formToken)}`, ) if (formResp.user_actions.length === 1) { action = formResp.user_actions[0]?.id ?? '' diff --git a/cli/src/commands/run/app/run.ts b/cli/src/commands/run/app/run.ts index ab468678a72..c0598b0eab0 100644 --- a/cli/src/commands/run/app/run.ts +++ b/cli/src/commands/run/app/run.ts @@ -65,7 +65,7 @@ async function executeRun( const m = await meta.get(opts.appId, [FieldInfo]) const mode = m.info?.mode ?? '' if (mode === '') - throw new Error(`app ${opts.appId}: mode missing from /describe`) + throw new Error(`app ${opts.appId}: mode missing from app metadata`) if (mode === RUN_MODES.Workflow && opts.message !== undefined && opts.message !== '') { throw new BaseError({ diff --git a/cli/src/commands/use/workspace/use.ts b/cli/src/commands/use/workspace/use.ts index 3070a76aa3e..94f335c940e 100644 --- a/cli/src/commands/use/workspace/use.ts +++ b/cli/src/commands/use/workspace/use.ts @@ -27,7 +27,7 @@ export type UseWorkspaceDeps = { * workspace list and let the caller pick one interactively (TTY only). * * The server-side switch is the source of truth: if POST - * `/workspaces//switch` fails we abort before touching `hosts.yml`, so + * `/workspaces/:switch` fails we abort before touching `hosts.yml`, so * local state never diverges from the server. */ export async function runUseWorkspace( diff --git a/cli/src/commands/version/version.test.ts b/cli/src/commands/version/version.test.ts index 33b7c17b769..be1e0f48917 100644 --- a/cli/src/commands/version/version.test.ts +++ b/cli/src/commands/version/version.test.ts @@ -109,7 +109,7 @@ describe('Version command', () => { } it('--check-compat exits with COMPAT_FAIL_EXIT_CODE when compat is unsupported', async () => { - vi.spyOn(probe, 'runVersionProbe').mockResolvedValue(fakeReport({ status: 'unsupported' })) + vi.spyOn(probe, 'runVersionProbe').mockResolvedValue(fakeReport({ status: 'too_new' })) const exitSpy = stubProcessExit() const stderrSpy = vi.spyOn(process.stderr, 'write').mockImplementation(() => true) @@ -119,7 +119,7 @@ describe('Version command', () => { }) it('--check-compat -o json emits the JSON envelope on stdout before exiting', async () => { - vi.spyOn(probe, 'runVersionProbe').mockResolvedValue(fakeReport({ status: 'unsupported' })) + vi.spyOn(probe, 'runVersionProbe').mockResolvedValue(fakeReport({ status: 'too_new' })) const exitSpy = stubProcessExit() const stdoutSpy = vi.spyOn(process.stdout, 'write').mockImplementation(() => true) vi.spyOn(process.stderr, 'write').mockImplementation(() => true) @@ -131,7 +131,7 @@ describe('Version command', () => { expect(stdoutSpy).toHaveBeenCalled() const written = stdoutSpy.mock.calls.map(c => String(c[0])).join('') const parsed = JSON.parse(written) as { compat: { status: string } } - expect(parsed.compat.status).toBe('unsupported') + expect(parsed.compat.status).toBe('too_new') expect(exitSpy).toHaveBeenCalledWith(COMPAT_FAIL_EXIT_CODE) }) diff --git a/cli/src/http/client-tls.test.ts b/cli/src/http/client-tls.test.ts new file mode 100644 index 00000000000..089b677fa97 --- /dev/null +++ b/cli/src/http/client-tls.test.ts @@ -0,0 +1,66 @@ +import type { Buffer } from 'node:buffer' +import type { AddressInfo } from 'node:net' +import { execFileSync } from 'node:child_process' +import { mkdtempSync, readFileSync, rmSync } from 'node:fs' +import * as https from 'node:https' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { createHttpClient } from './client.js' + +function generateSelfSignedCert(dir: string): { key: Buffer, cert: Buffer } { + const keyPath = join(dir, 'key.pem') + const certPath = join(dir, 'cert.pem') + execFileSync('openssl', [ + 'req', + '-x509', + '-newkey', + 'rsa:2048', + '-nodes', + '-keyout', + keyPath, + '-out', + certPath, + '-days', + '1', + '-subj', + '/CN=localhost', + ], { stdio: ['ignore', 'ignore', 'pipe'] }) + return { key: readFileSync(keyPath), cert: readFileSync(certPath) } +} + +// A real server, not a fetch mock, so this also covers Bun's native `tls` +// fetch option (ignored by Node, which only reads undici's `dispatcher`). +describe('createHttpClient against a real self-signed TLS server', () => { + let server: https.Server + let baseURL: string + let certDir: string + + beforeAll(async () => { + certDir = mkdtempSync(join(tmpdir(), 'difyctl-tls-test-')) + const { key, cert } = generateSelfSignedCert(certDir) + server = https.createServer({ key, cert }, (_req, res) => { + res.writeHead(200, { 'content-type': 'application/json' }) + res.end(JSON.stringify({ ok: true })) + }) + await new Promise(resolve => server.listen(0, resolve)) + const port = (server.address() as AddressInfo).port + baseURL = `https://localhost:${port}/` + }) + + afterAll(async () => { + await new Promise(resolve => server.close(() => resolve())) + rmSync(certDir, { recursive: true, force: true }) + }) + + it('rejects the self-signed cert by default', async () => { + const http = createHttpClient({ baseURL, retryAttempts: 0 }) + await expect(http.get('')).rejects.toBeDefined() + }) + + it('accepts the self-signed cert when insecure: true', async () => { + const http = createHttpClient({ baseURL, retryAttempts: 0, insecure: true }) + const res = await http.get<{ ok: boolean }>('') + expect(res.ok).toBe(true) + }) +}) diff --git a/cli/src/http/client.test.ts b/cli/src/http/client.test.ts index ae4448843a4..fa397852c98 100644 --- a/cli/src/http/client.test.ts +++ b/cli/src/http/client.test.ts @@ -184,7 +184,7 @@ describe('http client', () => { const client = createHttpClient({ baseURL: base(mock.url), bearer: 'dfoa_test' }) let caught: unknown try { - await client.get('apps/nope/describe') + await client.get('apps/nope') } catch (err) { caught = err } expect(isHttpClientError(caught)).toBe(true) @@ -545,7 +545,7 @@ describe('empty / No-Content bodies', () => { }) try { const client = createHttpClient({ baseURL: stub.url, bearer: 'dfoa_test' }) - await expect(client.post('apps/app-1/tasks/t-1/stop', { json: {} })).resolves.toBeUndefined() + await expect(client.post('apps/app-1/tasks/t-1:stop', { json: {} })).resolves.toBeUndefined() } finally { await stub.stop() diff --git a/cli/src/http/client.ts b/cli/src/http/client.ts index 940aee9152b..30f8ca27441 100644 --- a/cli/src/http/client.ts +++ b/cli/src/http/client.ts @@ -38,6 +38,7 @@ type ClientState = { readonly logger: HttpLogger | undefined readonly originalOptions: ClientOptions readonly dispatcher: ReturnType + readonly insecure: boolean } function toArray(value: T | T[] | undefined): T[] { @@ -74,7 +75,8 @@ function compileState(opts: ClientOptions): ClientState { hooks: { onRequest, onResponse, onRequestError, onResponseError }, logger: opts.logger, originalOptions: opts, - dispatcher: proxyDispatcher(), + dispatcher: proxyDispatcher({ insecure: opts.insecure }), + insecure: opts.insecure ?? false, } } @@ -171,9 +173,16 @@ async function execute( await runHooks(state.hooks.onRequest, ctx) - const init: RequestInit & { dispatcher?: unknown, verbose?: boolean } = { signal } + // Two runtimes, two options: Node's fetch reads undici's `dispatcher` (used + // below for TLS-skip + proxy routing); Bun's native fetch — what the compiled + // difyctl binary actually runs on — ignores `dispatcher` entirely and instead + // needs its own `tls` option. Set both; each runtime ignores the one it + // doesn't understand. + const init: RequestInit & { dispatcher?: unknown, tls?: { rejectUnauthorized: boolean }, verbose?: boolean } = { signal } if (state.dispatcher !== undefined) init.dispatcher = state.dispatcher + if (state.insecure) + init.tls = { rejectUnauthorized: false } if (isVerbose()) init.verbose = true diff --git a/cli/src/http/error-mapper.test.ts b/cli/src/http/error-mapper.test.ts index 3244222da07..c08a3a62e48 100644 --- a/cli/src/http/error-mapper.test.ts +++ b/cli/src/http/error-mapper.test.ts @@ -74,7 +74,7 @@ describe('classifyResponse — canonical ErrorBody', () => { describe('classifyResponse 403', () => { it('maps 403 to AccessDenied (exit 4 bucket)', async () => { - const req403 = new Request('https://x/openapi/v1/apps/abc/export') + const req403 = new Request('https://x/openapi/v1/apps/abc/dsl') const res403 = new Response( JSON.stringify({ code: 'unsupported_token_type', message: 'unsupported_token_type', status: 403 }), { status: 403, headers: { 'content-type': 'application/json' } }, @@ -91,6 +91,31 @@ describe('classifyResponse 403', () => { }) }) +describe('classifyResponse 426', () => { + it('maps 426 to VersionSkew (exit 6) and surfaces the server upgrade message', async () => { + const body = { + code: 'upgrade_required', + message: 'difyctl 0.1.0 is no longer supported; upgrade to >= 0.2.0.', + status: 426, + hint: 'Upgrade difyctl: https://docs.dify.ai/en/cli/install', + } + + const err = await classified(426, body) + + expect(err.code).toBe(ErrorCode.VersionSkew) + expect(err.exit()).toBe(6) + expect(err.message).toBe('difyctl 0.1.0 is no longer supported; upgrade to >= 0.2.0.') + expect(err.serverError?.code).toBe('upgrade_required') + }) + + it('426 with no parseable ErrorBody falls back to a version message', async () => { + const err = await classified(426, 'not json') + + expect(err.code).toBe(ErrorCode.VersionSkew) + expect(err.message).toBe('client version no longer supported by the server') + }) +}) + describe('classifyResponse — non-conforming bodies (no fallback by design)', () => { it('non-JSON body yields no serverError, classification by status', async () => { const err = await classified(502, 'bad gateway') diff --git a/cli/src/http/error-mapper.ts b/cli/src/http/error-mapper.ts index 34d7637d4e0..2e29985037c 100644 --- a/cli/src/http/error-mapper.ts +++ b/cli/src/http/error-mapper.ts @@ -50,11 +50,22 @@ const ACCESS_DENIED_CLASS: StatusClass = { includeRaw: false, } +// 426 Upgrade Required: the server rejected this difyctl as too old. Give it the +// version-compat exit code so scripts can tell it apart from a generic failure. +// The server's ErrorBody.code ("upgrade_required") + message still ride along. +const VERSION_COMPAT_CLASS: StatusClass = { + code: ErrorCode.VersionSkew, + fallbackMessage: () => 'client version no longer supported by the server', + includeRaw: false, +} + function statusClass(status: number): StatusClass { if (status === 401) return AUTH_EXPIRED_CLASS if (status === 403) return ACCESS_DENIED_CLASS + if (status === 426) + return VERSION_COMPAT_CLASS if (status === 429) return RATE_LIMITED_CLASS if (status >= 500) diff --git a/cli/src/http/proxy.test.ts b/cli/src/http/proxy.test.ts index 80b36d0f295..2c1f386f320 100644 --- a/cli/src/http/proxy.test.ts +++ b/cli/src/http/proxy.test.ts @@ -52,4 +52,28 @@ describe('proxyDispatcher', () => { expect(proxyDispatcher()).toBe(first) await first?.close() }) + + it('builds a plain Agent with TLS verification disabled when insecure is set and no proxy env', async () => { + const { proxyDispatcher } = await import('./proxy.js') + const d = proxyDispatcher({ insecure: true }) + expect(d?.constructor.name).toBe('Agent') + await d?.close() + }) + + it('builds an EnvHttpProxyAgent with TLS verification disabled when insecure and a proxy are both set', async () => { + process.env.HTTP_PROXY = 'http://127.0.0.1:8888' + const { proxyDispatcher } = await import('./proxy.js') + const d = proxyDispatcher({ insecure: true }) + expect(d?.constructor.name).toBe('EnvHttpProxyAgent') + await d?.close() + }) + + it('re-resolves when the insecure flag changes for the same call site', async () => { + const { proxyDispatcher } = await import('./proxy.js') + const secure = proxyDispatcher({ insecure: false }) + expect(secure).toBeUndefined() + const insecure = proxyDispatcher({ insecure: true }) + expect(insecure?.constructor.name).toBe('Agent') + await insecure?.close() + }) }) diff --git a/cli/src/http/proxy.ts b/cli/src/http/proxy.ts index 5e58936f91a..bd9a745c922 100644 --- a/cli/src/http/proxy.ts +++ b/cli/src/http/proxy.ts @@ -1,4 +1,5 @@ -import { EnvHttpProxyAgent } from 'undici' +import type { Dispatcher } from 'undici' +import { Agent, EnvHttpProxyAgent } from 'undici' const PROXY_ENV_KEYS = ['HTTP_PROXY', 'http_proxy', 'HTTPS_PROXY', 'https_proxy'] as const @@ -6,18 +7,30 @@ export function hasProxyEnv(): boolean { return PROXY_ENV_KEYS.some(k => (process.env[k] ?? '') !== '') } -let resolved = false -let agent: EnvHttpProxyAgent | undefined +export type ProxyDispatcherOptions = { + // --insecure on a self-signed https:// host: skip certificate verification + // (local-dev only, same flag that allows plain http:// hosts). + readonly insecure?: boolean +} + +let resolvedKey: string | undefined +let agent: Dispatcher | undefined // Node's global fetch ignores HTTP_PROXY / HTTPS_PROXY / NO_PROXY. When a proxy // var is set, route requests through an EnvHttpProxyAgent (it also reads the -// lowercase variants and honours NO_PROXY); when none is set, return undefined so -// fetch keeps Node's default global dispatcher untouched. Resolved once per -// process — proxy env vars are fixed for a single CLI invocation. -export function proxyDispatcher(): EnvHttpProxyAgent | undefined { - if (!resolved) { - agent = hasProxyEnv() ? new EnvHttpProxyAgent() : undefined - resolved = true +// lowercase variants and honours NO_PROXY); when none is set and TLS verification +// isn't being skipped, return undefined so fetch keeps Node's default global +// dispatcher untouched. Resolved once per (proxy env, insecure) combination — +// both are fixed for a single CLI invocation. +export function proxyDispatcher(opts: ProxyDispatcherOptions = {}): Dispatcher | undefined { + const insecure = opts.insecure ?? false + const key = `${hasProxyEnv()}:${insecure}` + if (resolvedKey !== key) { + const tls = insecure ? { rejectUnauthorized: false } : undefined + agent = hasProxyEnv() + ? new EnvHttpProxyAgent(tls !== undefined ? { connect: tls, requestTls: tls, proxyTls: tls } : undefined) + : (tls !== undefined ? new Agent({ connect: tls }) : undefined) + resolvedKey = key } return agent } diff --git a/cli/src/http/types.ts b/cli/src/http/types.ts index d209e97460c..a022a253296 100644 --- a/cli/src/http/types.ts +++ b/cli/src/http/types.ts @@ -75,6 +75,8 @@ export type ClientOptions = { readonly retryAttempts?: number readonly logger?: HttpLogger readonly hooks?: Hooks + // Skip TLS certificate verification (local-dev only, self-signed hosts). + readonly insecure?: boolean } export type HttpClient = { diff --git a/cli/src/store/manager.ts b/cli/src/store/manager.ts index 37962681cec..832aba6196b 100644 --- a/cli/src/store/manager.ts +++ b/cli/src/store/manager.ts @@ -7,6 +7,7 @@ import { FileTokenStore, KeychainTokenStore } from './token-store' export const CACHE_APP_INFO = 'app-info' export const CACHE_NUDGE = 'nudge' +export const CACHE_COMPAT = 'compat' const HOSTS_FILE = 'hosts.yml' const TOKENS_FILE = 'tokens.yml' export const CONFIG_FILE_NAME = 'config.yml' diff --git a/cli/src/util/host.ts b/cli/src/util/host.ts index 1042b0875ea..453a68149ad 100644 --- a/cli/src/util/host.ts +++ b/cli/src/util/host.ts @@ -1,3 +1,4 @@ +import type { ActiveContext } from '@/auth/hosts' import { BaseError } from '@/errors/base' import { ErrorCode } from '@/errors/codes' @@ -44,6 +45,16 @@ export function hostWithScheme(host: string, scheme: string | undefined): string return `${proto}://${host}` } +// Every call site that builds an HTTP client for the active host needs both its +// scheme-qualified host and whether TLS verification is disabled for it — derive +// them together so neither is forgotten independently. +export function activeHostInfo(active: Pick): { host: string, insecure: boolean } { + return { + host: hostWithScheme(active.host, active.scheme), + insecure: active.insecureTls === true, + } +} + export function bareHost(raw: string): string { try { const u = new URL(raw) diff --git a/cli/src/version/compat.test.ts b/cli/src/version/compat.test.ts index dcf3f08b259..c47b25e280f 100644 --- a/cli/src/version/compat.test.ts +++ b/cli/src/version/compat.test.ts @@ -32,14 +32,38 @@ describe('evaluateCompat', () => { expect(evaluateCompat('1.7.0', range).status).toBe('compatible') }) - it('returns unsupported when server is below minimum', () => { + it('returns too_old when server is below minimum', () => { const v = evaluateCompat('1.5.9', range) - expect(v.status).toBe('unsupported') + expect(v.status).toBe('too_old') expect(v.detail).toContain('1.5.9') }) - it('returns unsupported when server is above maximum', () => { - expect(evaluateCompat('2.0.0', range).status).toBe('unsupported') + it('returns too_new when server is above maximum', () => { + expect(evaluateCompat('2.0.0', range).status).toBe('too_new') + }) + + describe('ignores pre-release/channel suffixes (numeric-core comparison)', () => { + it('treats a pre-release of the upper bound as compatible', () => { + // 1.7.0-rc.1 has core 1.7.0 == maxDify; a suffix-sensitive range would push + // it out of [1.6.0, 1.7.0], but its numeric core is in range. + expect(evaluateCompat('1.7.0-rc.1', range).status).toBe('compatible') + }) + + it('treats a pre-release of the lower bound as compatible', () => { + expect(evaluateCompat('1.6.0-alpha', range).status).toBe('compatible') + }) + + it('still flags a pre-release whose core is below the minimum as too_old', () => { + expect(evaluateCompat('1.5.9-rc.1', range).status).toBe('too_old') + }) + + it('still flags a pre-release whose core is above the maximum as too_new', () => { + expect(evaluateCompat('2.0.0-alpha', range).status).toBe('too_new') + }) + + it('strips suffixes on the range bounds too', () => { + expect(evaluateCompat('1.6.5', { minDify: '1.6.0-alpha', maxDify: '1.7.0-rc' }).status).toBe('compatible') + }) }) it('returns unknown when server version is empty', () => { diff --git a/cli/src/version/compat.ts b/cli/src/version/compat.ts index b373ad0906b..5a1caf68cb1 100644 --- a/cli/src/version/compat.ts +++ b/cli/src/version/compat.ts @@ -1,4 +1,5 @@ -import { parseRange, satisfies, tryParse } from 'std-semver' +import type { SemVer } from 'std-semver' +import { compare, tryParse } from 'std-semver' export type DifyCompat = { readonly minDify: string @@ -14,7 +15,7 @@ export function compatString(): string { return `dify >=${difyCompat.minDify}, <=${difyCompat.maxDify}` } -export type CompatStatus = 'compatible' | 'unsupported' | 'unknown' +export type CompatStatus = 'compatible' | 'too_old' | 'too_new' | 'unknown' export type CompatVerdict = { readonly status: CompatStatus @@ -27,6 +28,12 @@ function clamp(s: string): string { return s.length > DETAIL_MAX_LEN ? `${s.slice(0, DETAIL_MAX_LEN)}…` : s } +// Numeric core (major.minor.patch) with pre-release/build stripped, so ordering +// ignores channel suffixes: a 0.2.0-rc.1 build compares equal to the 0.2.0 floor. +function core(v: SemVer): SemVer { + return { major: v.major, minor: v.minor, patch: v.patch, prerelease: [], build: [] } +} + export function evaluateCompat( serverVersion: string | undefined, range: DifyCompat = difyCompat, @@ -34,25 +41,20 @@ export function evaluateCompat( if (serverVersion === undefined || serverVersion === '') return { status: 'unknown', detail: 'server version unknown' } - const parsedServer = tryParse(serverVersion) - if (parsedServer === undefined) + const server = tryParse(serverVersion) + if (server === undefined) return { status: 'unknown', detail: `server version ${JSON.stringify(clamp(serverVersion))} is not valid semver` } - // The compat range is inclusive at both ends, exactly the format compatString prints. - const expr = `>=${range.minDify} <=${range.maxDify}` - const parsedRange = (() => { - try { - return parseRange(expr) - } - catch { - return undefined - } - })() - if (parsedRange === undefined) - return { status: 'unknown', detail: `compat range ${JSON.stringify(expr)} is not valid semver` } + const min = tryParse(range.minDify) + const max = tryParse(range.maxDify) + if (min === undefined || max === undefined) + return { status: 'unknown', detail: `compat range ${JSON.stringify(`>=${range.minDify} <=${range.maxDify}`)} is not valid semver` } - if (satisfies(parsedServer, parsedRange)) - return { status: 'compatible', detail: `server ${serverVersion} in [${range.minDify}, ${range.maxDify}]` } + if (compare(core(server), core(min)) < 0) + return { status: 'too_old', detail: `server ${serverVersion} is older than the minimum ${range.minDify}` } - return { status: 'unsupported', detail: `server ${serverVersion} outside [${range.minDify}, ${range.maxDify}]` } + if (compare(core(server), core(max)) > 0) + return { status: 'too_new', detail: `server ${serverVersion} is newer than the tested maximum ${range.maxDify}` } + + return { status: 'compatible', detail: `server ${serverVersion} in [${range.minDify}, ${range.maxDify}]` } } diff --git a/cli/src/version/enforce.test.ts b/cli/src/version/enforce.test.ts new file mode 100644 index 00000000000..4c3395530c2 --- /dev/null +++ b/cli/src/version/enforce.test.ts @@ -0,0 +1,86 @@ +import type { ServerVersionResponse } from '@dify/contracts/api/openapi/types.gen' +import type { CompatStore } from '@/cache/compat-store' +import { describe, expect, it, vi } from 'vitest' +import { ErrorCode } from '@/errors/codes' +import { enforceDifyVersion } from './enforce' + +// Injected build range in tests is __DIFYCTL_MIN_DIFY__=1.6.0 / MAX=1.7.0 (test/setup.ts): +// 1.5.0 → too_old, 1.6.4 → compatible, 99.0.0 → too_new, '' → unknown. +const HOST = 'https://cloud.dify.ai' + +function fakeStore(fresh = false): CompatStore & { readonly marked: string[] } { + const marked: string[] = [] + return { + marked, + isFreshCompatible: () => fresh, + markCompatible: async (host) => { + marked.push(host) + }, + } +} + +const server = (version: string): ServerVersionResponse => ({ version, edition: 'SELF_HOSTED' }) + +describe('enforceDifyVersion', () => { + it('throws version_skew (exit 6) when the server is too old, and never caches it', async () => { + const store = fakeStore() + const probe = vi.fn(async () => server('1.5.0')) + + await expect(enforceDifyVersion(HOST, { store, probe })).rejects.toMatchObject({ code: ErrorCode.VersionSkew }) + expect(store.marked).toHaveLength(0) + }) + + it('passes and caches when the server is compatible', async () => { + const store = fakeStore() + const probe = vi.fn(async () => server('1.6.4')) + + const res = await enforceDifyVersion(HOST, { store, probe }) + + expect(res?.version).toBe('1.6.4') + expect(store.marked).toEqual([HOST]) + }) + + it('passes (soft, no throw) and caches when the server is too new', async () => { + const store = fakeStore() + const probe = vi.fn(async () => server('99.0.0')) + + await expect(enforceDifyVersion(HOST, { store, probe })).resolves.toBeDefined() + expect(store.marked).toEqual([HOST]) + }) + + it('skips the probe entirely when the host is fresh-compatible', async () => { + const store = fakeStore(true) + const probe = vi.fn(async () => server('1.5.0')) // would throw if it ran + + await expect(enforceDifyVersion(HOST, { store, probe })).resolves.toBeUndefined() + expect(probe).not.toHaveBeenCalled() + }) + + it('re-probes despite a fresh cache when forceFresh is set', async () => { + const store = fakeStore(true) + const probe = vi.fn(async () => server('1.5.0')) + + await expect(enforceDifyVersion(HOST, { store, probe, forceFresh: true })) + .rejects + .toMatchObject({ code: ErrorCode.VersionSkew }) + expect(probe).toHaveBeenCalledOnce() + }) + + it('fails open (never blocks, never caches) when the probe errors', async () => { + const store = fakeStore() + const probe = vi.fn(async () => { + throw new Error('net down') + }) + + await expect(enforceDifyVersion(HOST, { store, probe })).resolves.toBeUndefined() + expect(store.marked).toHaveLength(0) + }) + + it('does not block or cache on an unknown server version', async () => { + const store = fakeStore() + const probe = vi.fn(async () => server('')) + + await expect(enforceDifyVersion(HOST, { store, probe })).resolves.toBeDefined() + expect(store.marked).toHaveLength(0) + }) +}) diff --git a/cli/src/version/enforce.ts b/cli/src/version/enforce.ts new file mode 100644 index 00000000000..8c995a912d1 --- /dev/null +++ b/cli/src/version/enforce.ts @@ -0,0 +1,72 @@ +import type { ServerVersionResponse } from '@dify/contracts/api/openapi/types.gen' +import type { CompatStore } from '@/cache/compat-store' +import { META_PROBE_TIMEOUT_MS, MetaClient } from '@/api/meta' +import { loadCompatStore } from '@/cache/compat-store' +import { newError } from '@/errors/base' +import { ErrorCode } from '@/errors/codes' +import { createHttpClient } from '@/http/client' +import { openAPIBase } from '@/util/host' +import { difyCompat, evaluateCompat } from './compat' +import { versionInfo } from './info' + +export type ServerVersionProbe = (host: string) => Promise + +const UPGRADE_HINT + = `upgrade the Dify server to >= ${difyCompat.minDify} ` + + '(https://docs.dify.ai/en/getting-started/install-self-hosted)' + +// /_version is unauthenticated; same timeout/no-retry budget as the auto-nudge probe. +function buildDefaultProbe(insecure: boolean): ServerVersionProbe { + return async (host) => { + const http = createHttpClient({ baseURL: openAPIBase(host), timeoutMs: META_PROBE_TIMEOUT_MS, retryAttempts: 0, insecure }) + return new MetaClient(http).serverVersion() + } +} + +export type EnforceOptions = { + readonly probe?: ServerVersionProbe + readonly store?: CompatStore + readonly forceFresh?: boolean + readonly insecure?: boolean +} + +/** + * Hard version gate for the client → server direction: refuse a Dify server older + * than this difyctl requires (its removed paths would only 404 otherwise). + * + * Cached: a host recently confirmed compatible is not re-probed for COMPAT_TTL_MS. + * Only "compatible" is cached, so a just-upgraded server clears a previous block at + * once. Fails open on any probe error — a flaky network never blocks a command. + * Returns the probed server version when it actually probed (skipped/failed → undefined), + * so the caller can reuse it. + */ +export async function enforceDifyVersion( + host: string, + opts: EnforceOptions = {}, +): Promise { + const store = opts.store ?? await loadCompatStore() + if (opts.forceFresh !== true && store.isFreshCompatible(host)) + return undefined + + const probe = opts.probe ?? buildDefaultProbe(opts.insecure === true) + let server: ServerVersionResponse + try { + server = await probe(host) + } + catch { + return undefined + } + + const verdict = evaluateCompat(server.version) + if (verdict.status === 'too_old') { + throw newError( + ErrorCode.VersionSkew, + `Dify server ${server.version} is too old for difyctl ${versionInfo.version}: ${verdict.detail}`, + ).withHint(UPGRADE_HINT) + } + + if (verdict.status === 'compatible' || verdict.status === 'too_new') + await store.markCompatible(host) + + return server +} diff --git a/cli/src/version/nudge.ts b/cli/src/version/nudge.ts index a6d9fe96d4b..f8d368ee066 100644 --- a/cli/src/version/nudge.ts +++ b/cli/src/version/nudge.ts @@ -44,7 +44,9 @@ export async function maybeNudgeCompat(host: string, deps: NudgeDeps): Promise { expect(report.compat.status).toBe('compatible') }) - it('returns unsupported when server version is out of range', async () => { + it('returns too_new when server version is above range', async () => { const report = await runVersionProbe({ skipServer: false, loadActive: async () => active(), @@ -113,7 +113,7 @@ describe('runVersionProbe', () => { }) expect(report.server.reachable).toBe(true) - expect(report.compat.status).toBe('unsupported') + expect(report.compat.status).toBe('too_new') }) it('returns unknown when server returns an empty version string', async () => { diff --git a/cli/src/version/probe.ts b/cli/src/version/probe.ts index 266910e7c05..05ab3ea84e3 100644 --- a/cli/src/version/probe.ts +++ b/cli/src/version/probe.ts @@ -6,7 +6,7 @@ import { META_PROBE_TIMEOUT_MS, MetaClient } from '@/api/meta' import { Registry } from '@/auth/hosts' import { createHttpClient } from '@/http/client' import { arch, platform } from '@/sys/index' -import { hostWithScheme, openAPIBase } from '@/util/host' +import { activeHostInfo, openAPIBase } from '@/util/host' import { difyCompat, evaluateCompat } from './compat.js' import { versionInfo } from './info.js' @@ -51,9 +51,11 @@ const defaultLoadActive = async (): Promise => { return (await Registry.load()).resolveActive() } -const defaultProbe: MetaProbe = async (endpoint) => { - const http = createHttpClient({ baseURL: openAPIBase(endpoint), timeoutMs: META_PROBE_TIMEOUT_MS, retryAttempts: 0 }) - return new MetaClient(http).serverVersion() +function buildDefaultProbe(insecure: boolean): MetaProbe { + return async (endpoint) => { + const http = createHttpClient({ baseURL: openAPIBase(endpoint), timeoutMs: META_PROBE_TIMEOUT_MS, retryAttempts: 0, insecure }) + return new MetaClient(http).serverVersion() + } } function buildClientBlock(): ClientBlock { @@ -92,7 +94,6 @@ export async function runVersionProbe(opts: RunVersionProbeOptions): Promise { compat: { minDify: '1.6.0', maxDify: '1.7.0', - status: 'unsupported', + status: 'too_new', detail: 'server 99.0.0 outside [1.6.0, 1.7.0]', }, } @@ -175,7 +175,7 @@ describe('renderVersionText', () => { compat: { minDify: '1.6.0', maxDify: '1.7.0', - status: 'unsupported', + status: 'too_new', detail: 'server 99.0.0 outside [1.6.0, 1.7.0]', }, } diff --git a/cli/src/version/render.ts b/cli/src/version/render.ts index 44dded18e28..70b725c80e3 100644 --- a/cli/src/version/render.ts +++ b/cli/src/version/render.ts @@ -15,7 +15,8 @@ export type RenderOptions = { const COMPAT_LABEL: Record = { compatible: 'ok', - unsupported: 'incompatible', + too_old: 'incompatible (server too old)', + too_new: 'incompatible (server too new)', unknown: 'unknown', } @@ -50,7 +51,8 @@ export function renderVersionText(report: VersionReport, opts: RenderOptions = { lines.push('') const verdictText = `Compatibility: ${COMPAT_LABEL[compat.status]} — ${compat.detail}` - lines.push(compat.status === 'unsupported' ? c.yellow(verdictText) : verdictText) + const incompatible = compat.status === 'too_old' || compat.status === 'too_new' + lines.push(incompatible ? c.yellow(verdictText) : verdictText) if (client.channel !== 'stable') { lines.push('') diff --git a/cli/test/e2e/suites/discovery/get-app-single.e2e.ts b/cli/test/e2e/suites/discovery/get-app-single.e2e.ts index b620eb383ef..528c08c9fa0 100644 --- a/cli/test/e2e/suites/discovery/get-app-single.e2e.ts +++ b/cli/test/e2e/suites/discovery/get-app-single.e2e.ts @@ -3,7 +3,7 @@ * * Test cases sourced from: Dify CLI Enhanced spec — Dify CLI/Discovery/Single App Query (22 cases) * - * Note: difyctl get app queries a single app via GET /apps//describe?fields=info. + * Note: difyctl get app queries a single app via GET /apps/?fields=info. * The response is returned in list-envelope format {page,limit,total,data:[...]}. */ diff --git a/cli/test/fixtures/dify-mock/server.test.ts b/cli/test/fixtures/dify-mock/server.test.ts index 7233f2a23b3..e80e20643c7 100644 --- a/cli/test/fixtures/dify-mock/server.test.ts +++ b/cli/test/fixtures/dify-mock/server.test.ts @@ -111,15 +111,15 @@ describe('dify-mock fixture server', () => { expect(body.data.map(r => r.id).sort()).toEqual(['app-3', 'app-4']) }) - it('GET /openapi/v1/apps/:id/describe returns 404 for unknown id', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/nope/describe?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { + it('GET /openapi/v1/apps/:id returns 404 for unknown id', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/nope?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { headers: { Authorization: 'Bearer dfoa_test' }, }) expect(r.status).toBe(404) }) - it('GET /openapi/v1/apps/:id/describe returns the app for known id', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/app-1/describe?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { + it('GET /openapi/v1/apps/:id returns the app for known id', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/app-1?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { headers: { Authorization: 'Bearer dfoa_test' }, }) expect(r.status).toBe(200) @@ -127,8 +127,8 @@ describe('dify-mock fixture server', () => { expect(body.info.id).toBe('app-1') }) - it('POST /openapi/v1/apps/:id/run returns SSE stream for chat app', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/app-1/run`, { + it('POST /openapi/v1/apps/:id:run returns SSE stream for chat app', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/app-1:run`, { method: 'POST', headers: { 'Authorization': 'Bearer dfoa_test', @@ -142,8 +142,8 @@ describe('dify-mock fixture server', () => { expect(text).toContain('"answer":"echo: "') }) - it('POST /openapi/v1/apps/:id/run returns SSE stream for workflow app', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/app-2/run`, { + it('POST /openapi/v1/apps/:id:run returns SSE stream for workflow app', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/app-2:run`, { method: 'POST', headers: { 'Authorization': 'Bearer dfoa_test', @@ -157,8 +157,8 @@ describe('dify-mock fixture server', () => { expect(text).toContain('"workflow_finished"') }) - it('GET /openapi/v1/apps/:id/describe?fields=info returns slim payload', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/app-1/describe?workspace_id=550e8400-e29b-41d4-a716-446655440000&fields=info`, { + it('GET /openapi/v1/apps/:id?fields=info returns slim payload', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/app-1?workspace_id=550e8400-e29b-41d4-a716-446655440000&fields=info`, { headers: { Authorization: 'Bearer dfoa_test' }, }) expect(r.status).toBe(200) @@ -168,8 +168,8 @@ describe('dify-mock fixture server', () => { expect(body.input_schema).toBeNull() }) - it('GET /openapi/v1/apps/:id/describe full returns parameters when present', async () => { - const r = await fetch(`${mock.url}/openapi/v1/apps/app-1/describe?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { + it('GET /openapi/v1/apps/:id full returns parameters when present', async () => { + const r = await fetch(`${mock.url}/openapi/v1/apps/app-1?workspace_id=550e8400-e29b-41d4-a716-446655440000`, { headers: { Authorization: 'Bearer dfoa_test' }, }) expect(r.status).toBe(200) diff --git a/cli/test/fixtures/dify-mock/server.ts b/cli/test/fixtures/dify-mock/server.ts index 766963cd0d5..d38e723988f 100644 --- a/cli/test/fixtures/dify-mock/server.ts +++ b/cli/test/fixtures/dify-mock/server.ts @@ -15,9 +15,9 @@ export type DifyMock = { scenario: Scenario setScenario: (s: Scenario) => void stop: () => Promise - /** Body of the most recent POST to /apps/:id/run */ + /** Body of the most recent POST to /apps/:id:run */ lastRunBody: Record | null - /** Number of times POST /apps/:id/files/upload was called */ + /** Number of times POST /apps/:id/files was called */ uploadCallCount: number /** Body of the most recent POST to /workspaces/:id/apps/imports */ lastImportBody: Record | null @@ -251,7 +251,7 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { }) }) - app.get('/openapi/v1/apps/:id/describe', (c) => { + app.get('/openapi/v1/apps/:id', (c) => { const id = c.req.param('id') const wsId = c.req.query('workspace_id') const fieldsRaw = c.req.query('fields') ?? '' @@ -279,7 +279,7 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { }) }) - app.get('/openapi/v1/permitted-external-apps/:id/describe', (c) => { + app.get('/openapi/v1/permitted-external-apps/:id', (c) => { const id = c.req.param('id') const fieldsRaw = c.req.query('fields') ?? '' const fields = fieldsRaw === '' ? [] : fieldsRaw.split(',').map(s => s.trim()).filter(s => s !== '') @@ -307,7 +307,7 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { }) }) - app.get('/openapi/v1/apps/:id/export', (c) => { + app.get('/openapi/v1/apps/:id/dsl', (c) => { const id = c.req.param('id') const found = APPS.find(a => a.id === id) if (found === undefined) @@ -315,7 +315,7 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { return c.json({ data: DSL_YAML }) }) - app.get('/openapi/v1/apps/:id/check-dependencies', (c) => { + app.get('/openapi/v1/apps/:id/dependencies:check', (c) => { const id = c.req.param('id') const found = APPS.find(a => a.id === id) if (found === undefined) @@ -335,12 +335,13 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { return c.json({ id: 'imp-1', status: 'completed', app_id: 'app-1', app_mode: 'chat' }, { status: 200 }) }) - app.post('/openapi/v1/workspaces/:wsId/apps/imports/:importId/confirm', (c) => { + app.post('/openapi/v1/workspaces/:wsId/apps/imports/:importId:confirm', (c) => { return c.json({ id: 'imp-1', status: 'completed', app_id: 'app-1', app_mode: 'chat' }, { status: 200 }) }) - app.post('/openapi/v1/apps/:id/run', async (c) => { - const id = c.req.param('id') + app.post('/openapi/v1/apps/:id:run', async (c) => { + // Hono drops the param adjacent to the `:run` literal; recover the app id from the path. + const id = c.req.path.replace(/^.*\/apps\//, '').replace(/:run$/, '') const body = await c.req.json() as { query?: string, inputs?: unknown } if (state !== undefined) state.lastRunBody = body as Record @@ -400,7 +401,7 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { return new Response(sse, { status: 200, headers: { 'content-type': 'text/event-stream' } }) }) - app.post('/openapi/v1/apps/:id/files/upload', async (c) => { + app.post('/openapi/v1/apps/:id/files', async (c) => { if (state !== undefined) state.uploadCallCount++ const form = await c.req.formData() @@ -421,11 +422,11 @@ export function buildApp(getScenario: () => Scenario, state?: MockState): Hono { ) }) - app.post('/openapi/v1/apps/:id/tasks/:taskId/stop', (c) => { + app.post('/openapi/v1/apps/:id/tasks/:taskId:stop', (c) => { return c.json({ result: 'success' }) }) - app.post('/openapi/v1/apps/:id/form/human_input/:formToken', (c) => { + app.post('/openapi/v1/apps/:id/human-input-forms/:formToken:submit', (c) => { return c.json({}) }) diff --git a/dify-agent/src/dify_agent/__init__.py b/dify-agent/src/dify_agent/__init__.py index b83189b1952..f3fc86ce366 100644 --- a/dify-agent/src/dify_agent/__init__.py +++ b/dify-agent/src/dify_agent/__init__.py @@ -5,6 +5,20 @@ runtime adapters or their optional dependencies. Server-only adapter entry point remain under ``dify_agent.adapters.llm``. """ -from dify_agent.client import Client +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dify_agent.client import Client + + +def __getattr__(name: str) -> object: + if name == "Client": + from dify_agent.client import Client + + return Client + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__ = ["Client"] diff --git a/dify-agent/src/dify_agent/adapters/llm/model.py b/dify-agent/src/dify_agent/adapters/llm/model.py index 26982cc2c29..e48ccc8bb17 100644 --- a/dify-agent/src/dify_agent/adapters/llm/model.py +++ b/dify-agent/src/dify_agent/adapters/llm/model.py @@ -37,10 +37,7 @@ from pydantic_ai.exceptions import UnexpectedModelBehavior from pydantic_ai.messages import ( AudioUrl, BinaryContent, - BuiltinToolCallPart, - BuiltinToolReturnPart, CachePoint, - CompactionPart, DocumentUrl, FilePart, FinishReason, @@ -337,7 +334,7 @@ def _map_model_response_to_prompt_message( content_parts: list[PromptMessageContentUnionTypes] = [] tool_calls: list[AssistantPromptMessage.ToolCall] = [] - for part in message.parts: + for index, part in enumerate(message.parts): if isinstance(part, TextPart): if part.content: content_parts.append(TextPromptMessageContent(data=part.content)) @@ -349,7 +346,7 @@ def _map_model_response_to_prompt_message( elif isinstance(part, ToolCallPart): tool_calls.append( AssistantPromptMessage.ToolCall( - id=part.tool_call_id or f"tool-call-{part.tool_name}", + id=part.tool_call_id or f"tool-call-{index}-{part.tool_name}", type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name=part.tool_name, @@ -357,10 +354,8 @@ def _map_model_response_to_prompt_message( ), ) ) - elif isinstance(part, BuiltinToolCallPart | BuiltinToolReturnPart | CompactionPart): - raise UnexpectedModelBehavior(f"Unsupported response part for daemon adapter: {type(part).__name__}") else: - assert_never(part) + raise UnexpectedModelBehavior(f"Unsupported response part for daemon adapter: {type(part).__name__}") content = _normalize_prompt_content(content_parts) if content is None and not tool_calls: @@ -487,10 +482,16 @@ def _map_binary_content_to_prompt_content( def _normalize_prompt_content( content: list[PromptMessageContentUnionTypes], ) -> str | list[PromptMessageContentUnionTypes] | None: + """Collapse text-only daemon message content to the string form. + + The daemon protocol supports content-part lists for multimodal messages, but + text-only history is safer as plain text because provider plugins commonly + JSON-encode text payloads without Graphon model encoders. + """ if not content: return None - if len(content) == 1 and isinstance(content[0], TextPromptMessageContent): - return content[0].data + if all(isinstance(item, TextPromptMessageContent) for item in content): + return "".join(item.data for item in content) return content diff --git a/dify-agent/src/dify_agent/adapters/shell/__init__.py b/dify-agent/src/dify_agent/adapters/shell/__init__.py index bb655da3b9a..533dbe21750 100644 --- a/dify-agent/src/dify_agent/adapters/shell/__init__.py +++ b/dify-agent/src/dify_agent/adapters/shell/__init__.py @@ -1,7 +1,9 @@ -"""Provider-agnostic shell adapter exports for the Dify agent.""" +"""Provider-agnostic shell adapter exports for the Dify agent. + +Keep this package root light so importing shell protocols does not eagerly +require ``pydantic_settings`` or shellctl runtime dependencies. +""" -from dify_agent.adapters.shell.config import DEFAULT_SHELL_PROVIDER, ShellAdapterSettings -from dify_agent.adapters.shell.factory import create_shell_provider from dify_agent.adapters.shell.protocols import ( CompleteShellCommandResult, ShellCommandProtocol, @@ -14,6 +16,27 @@ from dify_agent.adapters.shell.protocols import ( ShellResourceProtocol, ) + +def __getattr__(name: str) -> object: + if name == "DEFAULT_SHELL_PROVIDER": + from dify_agent.adapters.shell.config import DEFAULT_SHELL_PROVIDER + + return DEFAULT_SHELL_PROVIDER + if name == "ShellAdapterSettings": + from dify_agent.adapters.shell.config import ShellAdapterSettings + + return ShellAdapterSettings + if name == "create_shell_provider": + from dify_agent.adapters.shell.factory import create_shell_provider + + return create_shell_provider + if name == "shellctl": + from importlib import import_module + + return import_module("dify_agent.adapters.shell.shellctl") + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "CompleteShellCommandResult", "DEFAULT_SHELL_PROVIDER", diff --git a/dify-agent/src/dify_agent/agent_stub/_constants.py b/dify-agent/src/dify_agent/agent_stub/_constants.py new file mode 100644 index 00000000000..d21e073f5ef --- /dev/null +++ b/dify-agent/src/dify_agent/agent_stub/_constants.py @@ -0,0 +1,15 @@ +"""Zero-side-effect Agent Stub constants shared across client-safe modules.""" + +from __future__ import annotations + +from typing import Final + + +AGENT_STUB_DRIVE_BASE_ENV_VAR: Final[str] = "DIFY_AGENT_STUB_DRIVE_BASE" +DEFAULT_AGENT_STUB_DRIVE_BASE: Final[str] = "/mnt/drive" + + +__all__ = [ + "AGENT_STUB_DRIVE_BASE_ENV_VAR", + "DEFAULT_AGENT_STUB_DRIVE_BASE", +] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py b/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py index 2acd714411e..273f17eb917 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py @@ -3,7 +3,7 @@ from __future__ import annotations from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.client._agent_stub import connect_agent_stub_sync +from dify_agent.agent_stub.client import connect_agent_stub_sync from dify_agent.agent_stub.protocol.agent_stub import AgentStubConnectResponse diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_config.py b/dify-agent/src/dify_agent/agent_stub/cli/_config.py index 7d20354ac7d..5e25de4e6db 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_config.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_config.py @@ -17,13 +17,14 @@ from dify_agent.agent_stub._drive_materialization import ( from dify_agent.agent_stub.cli._drive import _build_skill_archive from dify_agent.agent_stub.cli._env import read_agent_stub_environment from dify_agent.agent_stub.cli._files import upload_tool_file_resource_from_environment -from dify_agent.agent_stub.client._agent_stub import ( +from dify_agent.agent_stub.client import ( + AgentStubTransferError, + AgentStubValidationError, + request_agent_stub_config_file_pull_sync, request_agent_stub_config_manifest_sync, request_agent_stub_config_push_sync, - request_agent_stub_config_file_pull_sync, request_agent_stub_config_skill_pull_sync, ) -from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubConfigFileRef, AgentStubConfigManifestResponse, diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py index 69c9d5455bd..45ae61faabf 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py @@ -27,14 +27,16 @@ from dify_agent.agent_stub._drive_materialization import ( materialize_drive_downloads, resolve_drive_destination, ) +from dify_agent.agent_stub._constants import DEFAULT_AGENT_STUB_DRIVE_BASE from dify_agent.agent_stub.cli._env import read_agent_stub_environment from dify_agent.agent_stub.cli._files import upload_tool_file_resource_from_environment -from dify_agent.agent_stub.client._agent_stub import ( +from dify_agent.agent_stub.client import ( + AgentStubTransferError, + AgentStubValidationError, download_file_bytes_from_signed_url_sync, request_agent_stub_drive_commit_sync, request_agent_stub_drive_manifest_sync, ) -from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubDriveCommitItem, AgentStubDriveCommitRequest, @@ -42,7 +44,6 @@ from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubDriveFileRef, AgentStubDriveItem, AgentStubDriveManifestResponse, - DEFAULT_AGENT_STUB_DRIVE_BASE, ) _SKILL_MD_FILENAME = "SKILL.md" diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_env.py b/dify-agent/src/dify_agent/agent_stub/cli/_env.py index da268f6a369..61f06412262 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_env.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_env.py @@ -6,11 +6,10 @@ from collections.abc import Mapping from dataclasses import dataclass import os +from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE from dify_agent.agent_stub.protocol.agent_stub import ( AGENT_STUB_AUTH_JWE_ENV_VAR, - AGENT_STUB_DRIVE_BASE_ENV_VAR, AGENT_STUB_API_BASE_URL_ENV_VAR, - DEFAULT_AGENT_STUB_DRIVE_BASE, normalize_agent_stub_api_base_url, ) diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_files.py b/dify-agent/src/dify_agent/agent_stub/cli/_files.py index 9c3c017c51f..5e3635fde14 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_files.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_files.py @@ -5,26 +5,32 @@ from __future__ import annotations import mimetypes from dataclasses import dataclass from pathlib import Path -from typing import ClassVar, Literal, cast +from typing import ClassVar, Literal, Protocol, cast from pydantic import BaseModel, ConfigDict, ValidationError from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.client._agent_stub import ( +from dify_agent.agent_stub.client import ( + AgentStubTransferError, + AgentStubValidationError, download_file_bytes_from_signed_url_sync, request_agent_stub_file_download_sync, request_agent_stub_file_upload_sync, upload_file_to_signed_url_sync, ) -from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError from dify_agent.agent_stub.protocol.agent_stub import AgentStubFileMapping, is_canonical_dify_file_reference +class _AgentStubFileDownloadResponse(Protocol): + download_url: str + + class UploadedToolFileMapping(BaseModel): """Canonical Agent output mapping returned by ``dify-agent file upload``.""" transfer_method: Literal["tool_file"] = "tool_file" reference: str + download_url: str model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") @@ -38,9 +44,9 @@ class DownloadedFileResult: @dataclass(frozen=True, slots=True) class UploadedToolFileResource: - """Lower-level upload result carrying both public mapping and ToolFile id.""" + """Lower-level upload result carrying the internal mapping and ToolFile id.""" - mapping: UploadedToolFileMapping + mapping: AgentStubFileMapping tool_file_id: str @@ -48,11 +54,22 @@ def upload_file_from_environment(*, path: str) -> UploadedToolFileMapping: """Upload one sandbox-local file through the Agent Stub control plane. The signed upload data-plane response must carry the Dify-generated - ``reference`` for the new ``ToolFile`` so the sandbox can return the - canonical Agent output file mapping without synthesizing reference format. + ``reference`` for the new ``ToolFile``. The helper then resolves the same + mapping through the signed download-request control plane so the public CLI + output includes a ready-to-share external ``download_url``. """ - return upload_tool_file_resource_from_environment(path=path).mapping + resource = upload_tool_file_resource_from_environment(path=path) + reference = resource.mapping.reference + if not isinstance(reference, str) or not reference: + raise AgentStubTransferError("signed file upload response is missing reference") + environment = read_agent_stub_environment() + download_url = _request_uploaded_tool_file_download_url( + url=environment.url, + auth_jwe=environment.auth_jwe, + reference=reference, + ) + return UploadedToolFileMapping(reference=reference, download_url=download_url) def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFileResource: @@ -61,6 +78,8 @@ def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFile This lower-level helper backs ``drive push``. The signed upload data-plane response must include both the canonical Dify ``reference`` used by public CLI output and the raw ToolFile ``id`` required by drive commit payloads. + It intentionally stops after the upload allocation so internal flows that + only need the ToolFile identity do not pay for signed download enrichment. Raises: AgentStubValidationError: if ``path`` does not resolve to a local file. @@ -91,7 +110,11 @@ def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFile file_obj=file_obj, mimetype=mime_type, ) - return _normalize_uploaded_tool_file_resource(payload) + reference, tool_file_id = _normalize_uploaded_tool_file_payload(payload) + return UploadedToolFileResource( + mapping=AgentStubFileMapping(transfer_method="tool_file", reference=reference), + tool_file_id=tool_file_id, + ) def download_file_from_environment( @@ -163,7 +186,7 @@ def _build_download_mapping( raise AgentStubValidationError("invalid file download arguments") from exc -def _normalize_uploaded_tool_file_resource(payload: dict[str, object]) -> UploadedToolFileResource: +def _normalize_uploaded_tool_file_payload(payload: dict[str, object]) -> tuple[str, str]: reference = payload.get("reference") if not isinstance(reference, str) or not reference: raise AgentStubTransferError("signed file upload response is missing reference") @@ -172,10 +195,22 @@ def _normalize_uploaded_tool_file_resource(payload: dict[str, object]) -> Upload tool_file_id = payload.get("id") if not isinstance(tool_file_id, str) or not tool_file_id: raise AgentStubTransferError("signed file upload response is missing id") - return UploadedToolFileResource( - mapping=UploadedToolFileMapping(reference=reference), - tool_file_id=tool_file_id, + return reference, tool_file_id + + +def _request_uploaded_tool_file_download_url(*, url: str, auth_jwe: str, reference: str) -> str: + download_request = cast( + _AgentStubFileDownloadResponse, + request_agent_stub_file_download_sync( + url=url, + auth_jwe=auth_jwe, + file=AgentStubFileMapping(transfer_method="tool_file", reference=reference), + ), ) + download_url = download_request.download_url + if not isinstance(download_url, str) or not download_url: + raise AgentStubTransferError("signed file download response is missing download_url") + return download_url def _deduplicate_destination_path(path: Path) -> Path: diff --git a/dify-agent/src/dify_agent/agent_stub/cli/main.py b/dify-agent/src/dify_agent/agent_stub/cli/main.py index 6ccc9865989..ee00059d1e5 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/main.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/main.py @@ -10,46 +10,27 @@ does not pull in FastAPI, Redis, shellctl, or JWE runtime dependencies. from __future__ import annotations +from functools import cache +from importlib import import_module import sys -from typing import cast import click import typer from typer.main import get_command -from dify_agent.agent_stub.cli._agent_stub import connect_from_environment -from dify_agent.agent_stub.cli._config import ( - delete_config_files_from_environment, - delete_config_skills_from_environment, - manifest_from_environment, - pull_config_files_from_environment, - pull_config_note_from_environment, - pull_config_skills_from_environment, - push_config_env_from_environment, - push_config_files_from_environment, - push_config_note_from_environment, - push_config_skills_from_environment, -) -from dify_agent.agent_stub.cli._drive import ( - DrivePushKind, - format_drive_manifest, - list_drive_manifest_from_environment, - pull_drive_from_environment, - push_drive_from_environment, -) -from dify_agent.agent_stub.cli._env import ( - MissingAgentStubEnvironmentError, - has_agent_stub_environment, - read_agent_stub_drive_base, -) -from dify_agent.agent_stub.cli._files import download_file_from_environment, upload_file_from_environment -from dify_agent.agent_stub.client._errors import AgentStubClientError -from dify_agent.agent_stub.protocol.agent_stub import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE +from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE _CONFIG_MANIFEST_STDOUT_EXCLUDE = { "skills": {"items": {"__all__": {"hash"}}}, "files": {"items": {"__all__": {"hash"}}}, } +_FILE_DOWNLOAD_TRANSFER_METHOD_HELP = ( + "File mapping transfer_method: local_file, tool_file, datasource_file, or remote_url." +) +_FILE_DOWNLOAD_REFERENCE_OR_URL_HELP = ( + "File mapping reference or URL. Use dify-file-ref:... with local_file, tool_file, or datasource_file; " + "use https://... with remote_url." +) app = typer.Typer( @@ -102,16 +83,18 @@ def upload(path: str = typer.Argument(..., metavar="PATH")) -> None: @file_app.command("download") def download( - transfer_method: str | None = typer.Argument(None, metavar="TRANSFER_METHOD"), - reference_or_url: str | None = typer.Argument(None, metavar="REFERENCE_OR_URL"), - mapping: str | None = typer.Option(None, "--mapping", help="Download one file from a mapping JSON object."), + transfer_method: str = typer.Argument(..., metavar="TRANSFER_METHOD", help=_FILE_DOWNLOAD_TRANSFER_METHOD_HELP), + reference_or_url: str = typer.Argument( + ..., + metavar="REFERENCE_OR_URL", + help=_FILE_DOWNLOAD_REFERENCE_OR_URL_HELP, + ), local_dir: str | None = typer.Option(None, "--to", help="Local directory for the downloaded file."), ) -> None: """Download one workflow file mapping into the local sandbox directory.""" _run_file_download( transfer_method=transfer_method, reference_or_url=reference_or_url, - mapping=mapping, local_dir=local_dir, ) @@ -290,7 +273,7 @@ def main(argv: list[str] | None = None) -> None: return json_output, forwarded_args = _extract_root_json_flag(args) if _is_unknown_bare_command(forwarded_args): - if not has_agent_stub_environment(): + if not _env_module().has_agent_stub_environment(): _show_root_help() _run_connect(argv=forwarded_args, json_output=json_output) return @@ -352,12 +335,15 @@ def render_agent_stub_cli_help(args: tuple[str, ...]) -> str: def _run_connect(*, argv: list[str], json_output: bool) -> None: + env_module = _env_module() + client_module = _client_module() + agent_stub_module = _agent_stub_module() try: - response = connect_from_environment(argv=argv) - except MissingAgentStubEnvironmentError as exc: + response = agent_stub_module.connect_from_environment(argv=argv) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc @@ -368,12 +354,15 @@ def _run_connect(*, argv: list[str], json_output: bool) -> None: def _run_file_upload(*, path: str) -> None: + env_module = _env_module() + client_module = _client_module() + files_module = _files_module() try: - response = upload_file_from_environment(path=path) - except MissingAgentStubEnvironmentError as exc: + response = files_module.upload_file_from_environment(path=path) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) @@ -381,46 +370,53 @@ def _run_file_upload(*, path: str) -> None: def _run_file_download( *, - transfer_method: str | None, - reference_or_url: str | None, - mapping: str | None, + transfer_method: str, + reference_or_url: str, local_dir: str | None, ) -> None: + env_module = _env_module() + client_module = _client_module() + files_module = _files_module() try: - response = download_file_from_environment( + response = files_module.download_file_from_environment( transfer_method=transfer_method, reference_or_url=reference_or_url, - mapping=mapping, local_dir=local_dir, ) - except MissingAgentStubEnvironmentError as exc: + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(str(response.path)) def _run_config_manifest() -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = manifest_from_environment() - except MissingAgentStubEnvironmentError as exc: + response = config_module.manifest_from_environment() + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json(exclude=_CONFIG_MANIFEST_STDOUT_EXCLUDE)) def _run_config_skill_pull(*, names: list[str] | None, local_dir: str | None, json_output: bool) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = pull_config_skills_from_environment(names=names, local_dir=local_dir) - except MissingAgentStubEnvironmentError as exc: + response = config_module.pull_config_skills_from_environment(names=names, local_dir=local_dir) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc if json_output: @@ -434,12 +430,15 @@ def _run_config_skill_pull(*, names: list[str] | None, local_dir: str | None, js def _run_config_file_pull(*, names: list[str] | None, local_dir: str | None, json_output: bool) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = pull_config_files_from_environment(names=names, local_dir=local_dir) - except MissingAgentStubEnvironmentError as exc: + response = config_module.pull_config_files_from_environment(names=names, local_dir=local_dir) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc if json_output: @@ -450,111 +449,141 @@ def _run_config_file_pull(*, names: list[str] | None, local_dir: str | None, jso def _run_config_note_pull(*, local_path: str | None) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - path = pull_config_note_from_environment(local_path=local_path) - except MissingAgentStubEnvironmentError as exc: + path = config_module.pull_config_note_from_environment(local_path=local_path) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(str(path)) def _run_config_note_push(*, local_path: str | None) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = push_config_note_from_environment(local_path=local_path) - except MissingAgentStubEnvironmentError as exc: + response = config_module.push_config_note_from_environment(local_path=local_path) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_config_env_push(*, local_path: str) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = push_config_env_from_environment(local_path=local_path) - except MissingAgentStubEnvironmentError as exc: + response = config_module.push_config_env_from_environment(local_path=local_path) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_config_files_push(*, paths: list[str]) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = push_config_files_from_environment(paths=paths) - except MissingAgentStubEnvironmentError as exc: + response = config_module.push_config_files_from_environment(paths=paths) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_config_files_delete(*, names: list[str]) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = delete_config_files_from_environment(names=names) - except MissingAgentStubEnvironmentError as exc: + response = config_module.delete_config_files_from_environment(names=names) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_config_skills_push(*, paths: list[str]) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = push_config_skills_from_environment(paths=paths) - except MissingAgentStubEnvironmentError as exc: + response = config_module.push_config_skills_from_environment(paths=paths) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_config_skills_delete(*, names: list[str]) -> None: + env_module = _env_module() + client_module = _client_module() + config_module = _config_module() try: - response = delete_config_skills_from_environment(names=names) - except MissingAgentStubEnvironmentError as exc: + response = config_module.delete_config_skills_from_environment(names=names) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) def _run_drive_list(*, path_prefix: str, json_output: bool) -> None: + env_module = _env_module() + client_module = _client_module() + drive_module = _drive_module() try: - response = list_drive_manifest_from_environment(prefix=path_prefix) - except MissingAgentStubEnvironmentError as exc: + response = drive_module.list_drive_manifest_from_environment(prefix=path_prefix) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc if json_output: typer.echo(response.model_dump_json()) return - typer.echo(format_drive_manifest(response)) + typer.echo(drive_module.format_drive_manifest(response)) def _run_drive_pull(*, targets: list[str] | None, local_base: str | None, json_output: bool) -> None: + env_module = _env_module() + client_module = _client_module() + drive_module = _drive_module() try: - response = pull_drive_from_environment(targets=targets, local_base=local_base or read_agent_stub_drive_base()) - except MissingAgentStubEnvironmentError as exc: + response = drive_module.pull_drive_from_environment( + targets=targets, + local_base=local_base or env_module.read_agent_stub_drive_base(), + ) + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc if json_output: @@ -565,19 +594,54 @@ def _run_drive_pull(*, targets: list[str] | None, local_base: str | None, json_o def _run_drive_push(*, local_path: str, drive_path: str, kind: str | None) -> None: + env_module = _env_module() + client_module = _client_module() + drive_module = _drive_module() try: - response = push_drive_from_environment( + response = drive_module.push_drive_from_environment( local_path=local_path, drive_path=drive_path, - kind=cast(DrivePushKind | None, kind), + kind=kind, ) - except MissingAgentStubEnvironmentError as exc: + except env_module.MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc - except AgentStubClientError as exc: + except client_module.AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc typer.echo(response.model_dump_json()) +# Keep helper imports on demand so importing CLI/help stays free of server-side +# and unrelated heavy runtime dependencies. +@cache +def _agent_stub_module(): + return import_module("dify_agent.agent_stub.cli._agent_stub") + + +@cache +def _config_module(): + return import_module("dify_agent.agent_stub.cli._config") + + +@cache +def _drive_module(): + return import_module("dify_agent.agent_stub.cli._drive") + + +@cache +def _files_module(): + return import_module("dify_agent.agent_stub.cli._files") + + +@cache +def _env_module(): + return import_module("dify_agent.agent_stub.cli._env") + + +@cache +def _client_module(): + return import_module("dify_agent.agent_stub.client") + + __all__ = ["app", "main"] diff --git a/dify-agent/src/dify_agent/agent_stub/client/__init__.py b/dify-agent/src/dify_agent/agent_stub/client/__init__.py index c8413e00cf9..422f4a27301 100644 --- a/dify-agent/src/dify_agent/agent_stub/client/__init__.py +++ b/dify-agent/src/dify_agent/agent_stub/client/__init__.py @@ -3,6 +3,15 @@ from ._agent_stub import ( connect_agent_stub_sync, download_file_bytes_from_signed_url_sync, + request_agent_stub_config_env_update_sync, + request_agent_stub_config_file_pull_sync, + request_agent_stub_config_manifest_sync, + request_agent_stub_config_note_update_sync, + request_agent_stub_config_push_sync, + request_agent_stub_config_skill_inspect_sync, + request_agent_stub_config_skill_pull_sync, + request_agent_stub_drive_commit_sync, + request_agent_stub_drive_manifest_sync, request_agent_stub_file_download_sync, request_agent_stub_file_upload_sync, upload_file_to_signed_url_sync, @@ -25,6 +34,15 @@ __all__ = [ "AgentStubValidationError", "connect_agent_stub_sync", "download_file_bytes_from_signed_url_sync", + "request_agent_stub_config_env_update_sync", + "request_agent_stub_config_file_pull_sync", + "request_agent_stub_config_manifest_sync", + "request_agent_stub_config_note_update_sync", + "request_agent_stub_config_push_sync", + "request_agent_stub_config_skill_inspect_sync", + "request_agent_stub_config_skill_pull_sync", + "request_agent_stub_drive_commit_sync", + "request_agent_stub_drive_manifest_sync", "request_agent_stub_file_download_sync", "request_agent_stub_file_upload_sync", "upload_file_to_signed_url_sync", diff --git a/dify-agent/src/dify_agent/agent_stub/protocol/__init__.py b/dify-agent/src/dify_agent/agent_stub/protocol/__init__.py index d10643de944..6aedf24b08b 100644 --- a/dify-agent/src/dify_agent/agent_stub/protocol/__init__.py +++ b/dify-agent/src/dify_agent/agent_stub/protocol/__init__.py @@ -1,11 +1,11 @@ """Client-safe protocol exports for the Dify Agent Stub package.""" +from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE + from .agent_stub import ( AGENT_STUB_AUTH_JWE_ENV_VAR, - AGENT_STUB_DRIVE_BASE_ENV_VAR, AGENT_STUB_PROTOCOL_VERSION, AGENT_STUB_API_BASE_URL_ENV_VAR, - DEFAULT_AGENT_STUB_DRIVE_BASE, AgentStubConnectRequest, AgentStubConnectResponse, AgentStubConfigEnvUpdateRequest, diff --git a/dify-agent/src/dify_agent/agent_stub/protocol/agent_stub.py b/dify-agent/src/dify_agent/agent_stub/protocol/agent_stub.py index f13e5d9fb9f..1d257287671 100644 --- a/dify-agent/src/dify_agent/agent_stub/protocol/agent_stub.py +++ b/dify-agent/src/dify_agent/agent_stub/protocol/agent_stub.py @@ -17,12 +17,12 @@ from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator +from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE + AGENT_STUB_PROTOCOL_VERSION: Final[int] = 1 AGENT_STUB_API_BASE_URL_ENV_VAR: Final[str] = "DIFY_AGENT_STUB_API_BASE_URL" AGENT_STUB_AUTH_JWE_ENV_VAR: Final[str] = "DIFY_AGENT_STUB_AUTH_JWE" -AGENT_STUB_DRIVE_BASE_ENV_VAR: Final[str] = "DIFY_AGENT_STUB_DRIVE_BASE" -DEFAULT_AGENT_STUB_DRIVE_BASE: Final[str] = "/mnt/drive" type AgentStubURLScheme = Literal["http", "https", "grpc"] diff --git a/dify-agent/src/dify_agent/agent_stub/server/shell_agent_stub_env.py b/dify-agent/src/dify_agent/agent_stub/shell_env.py similarity index 84% rename from dify-agent/src/dify_agent/agent_stub/server/shell_agent_stub_env.py rename to dify-agent/src/dify_agent/agent_stub/shell_env.py index b4711c213ac..dd81ee4e75a 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/shell_agent_stub_env.py +++ b/dify-agent/src/dify_agent/agent_stub/shell_env.py @@ -1,19 +1,20 @@ -"""Server-side environment injection helpers for Agent Stub forwarding. +"""Client-safe shell environment helpers for Agent Stub forwarding. Only user-visible ``shell.run`` commands receive these variables. Internal -lifecycle commands remain free of Agent Stub credentials and drive-base defaults -so workspace setup and cleanup cannot accidentally inherit user-facing forwarding -state. +lifecycle commands remain free of Agent Stub credentials and drive-base +defaults so workspace setup and cleanup cannot accidentally inherit +user-facing forwarding state. The module stays server-extra-free because the +shell runtime and provider factory use it in sandbox-visible paths. """ from __future__ import annotations from typing import Protocol +from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR from dify_agent.agent_stub.protocol.agent_stub import ( - AGENT_STUB_AUTH_JWE_ENV_VAR, - AGENT_STUB_DRIVE_BASE_ENV_VAR, AGENT_STUB_API_BASE_URL_ENV_VAR, + AGENT_STUB_AUTH_JWE_ENV_VAR, agent_stub_drive_base_for_ref, normalize_agent_stub_api_base_url, ) diff --git a/dify-agent/src/dify_agent/layers/_agent_file_cli_help.py b/dify-agent/src/dify_agent/layers/_agent_file_cli_help.py new file mode 100644 index 00000000000..50deef3d851 --- /dev/null +++ b/dify-agent/src/dify_agent/layers/_agent_file_cli_help.py @@ -0,0 +1,7 @@ +"""Shared model-facing guidance for Agent Stub file CLI usage in prompt surfaces.""" + +AGENT_FILE_UPLOAD_REPLY_HINT = ( + "When you want to provide a generated or sandbox-local file to the user in a " + "natural-language reply, run the installed CLI command `dify-agent file upload PATH` and include the returned " + "`download_url` so the user can open or download the file." +) diff --git a/dify-agent/src/dify_agent/layers/config/configs.py b/dify-agent/src/dify_agent/layers/config/configs.py index c3d20c878f1..5f8eae360eb 100644 --- a/dify-agent/src/dify_agent/layers/config/configs.py +++ b/dify-agent/src/dify_agent/layers/config/configs.py @@ -59,6 +59,9 @@ class DifyConfigLayerConfig(LayerConfig): class DifyConfigRuntimeState(BaseModel): """Serializable config-layer values computed once during context entry. + ``config_cli_help`` stores pre-rendered shell-visible ``dify-agent`` help snippets + for config commands and file upload/download commands. + The ``push_spec_*`` fields are compatibility leftovers from the removed root JSON-spec config mutation workflow. This change keeps them in the runtime-state schema to avoid snapshot churn, but new code should treat them as inert diff --git a/dify-agent/src/dify_agent/layers/config/layer.py b/dify-agent/src/dify_agent/layers/config/layer.py index d914f599efb..6a4a5dea4db 100644 --- a/dify-agent/src/dify_agent/layers/config/layer.py +++ b/dify-agent/src/dify_agent/layers/config/layer.py @@ -11,6 +11,7 @@ from typing_extensions import Self, override from agenton.layers import LayerDeps, PlainLayer from dify_agent.agent_stub.cli.main import render_agent_stub_cli_help +from dify_agent.layers._agent_file_cli_help import AGENT_FILE_UPLOAD_REPLY_HINT as _AGENT_FILE_UPLOAD_REPLY_HINT from dify_agent.layers.config.configs import ( DIFY_CONFIG_LAYER_TYPE_ID, DifyConfigLayerConfig, @@ -18,14 +19,19 @@ from dify_agent.layers.config.configs import ( ) from dify_agent.layers.shell.layer import DifyShellLayer -_CONFIG_CONTEXT_HEADING = "Agent config context from the current Agent Soul:" -_CONFIG_CONTEXT_COMMAND = "$ dify-agent config manifest" -_CONFIG_CLI_USAGE_PROMPT = """Agent config CLI usage is available inside shell jobs. The command help below is generated -from the same `dify-agent` CLI definitions available in shell jobs. +_CONFIG_CONTEXT_HEADING = "Current Agent config manifest for this run:" +_CONFIG_CONTEXT_COMMAND = "dify-agent config manifest" +_CONFIG_CLI_USAGE_PROMPT = """`dify-agent` is an installed CLI tool in the shell environment. Use it directly in shell_run scripts. -Local edits to config files, skills, env, or notes are not saved by themselves. Config changes are saved only by a -matching resource mutation command. Those commands are available only when the Agent config context reports -`config_version.kind` as `build_draft` and `config_version.writable` as true.""" +The command outputs below are generated from the `dify-agent` CLI available in this run. Use them as the source of truth +for command names, arguments, and options. + +Config persistence rules: + +- Local shell edits to config files, skills, env, or notes are not saved by themselves. +- To persist an Agent config change, run the matching `dify-agent config ...` mutation command. +- Mutation commands are available only when the manifest shows `config_version.kind` as `build_draft` and + `config_version.writable` as true.""" _CONFIG_CLI_HELP_COMMANDS: dict[str, tuple[str, ...]] = { "dify-agent config --help": ("config",), "dify-agent config manifest --help": ("config", "manifest"), @@ -41,6 +47,10 @@ _CONFIG_CLI_MUTATION_HELP_COMMANDS: dict[str, tuple[str, ...]] = { "dify-agent config skills push --help": ("config", "skills", "push"), "dify-agent config skills delete --help": ("config", "skills", "delete"), } +_AGENT_FILE_CLI_HELP_COMMANDS: dict[str, tuple[str, ...]] = { + "dify-agent file upload --help": ("file", "upload"), + "dify-agent file download --help": ("file", "download"), +} _CONFIG_CONTEXT_EXCLUDE = {"mentioned_skill_names": True, "mentioned_file_names": True} @@ -91,6 +101,7 @@ class DifyConfigLayer(PlainLayer[DifyConfigDeps, DifyConfigLayerConfig, DifyConf command_paths = dict(_CONFIG_CLI_HELP_COMMANDS) if self._config_writable: command_paths.update(_CONFIG_CLI_MUTATION_HELP_COMMANDS) + command_paths.update(_AGENT_FILE_CLI_HELP_COMMANDS) self.runtime_state.config_context_json = self._format_config_context_json() self.runtime_state.config_cli_help = { command: render_agent_stub_cli_help(args) for command, args in command_paths.items() @@ -107,15 +118,22 @@ class DifyConfigLayer(PlainLayer[DifyConfigDeps, DifyConfigLayerConfig, DifyConf output = self.runtime_state.pulled_skill_outputs.get(name) if output is None: continue - loaded_skill_sections.append(f"Name: {name}\nPull output:\n{output}") + command = f"dify-agent config skills pull {shlex.quote(name)}" + loaded_skill_sections.append( + f"Name: {name}\nPull command output for this run:\n{_format_command_output(command, output)}" + ) if loaded_skill_sections: sections.append("Loaded mentioned skills:\n\n" + "\n\n".join(loaded_skill_sections)) - mentioned_file_sections = [ - f"Name: {name}\nPull output:\n{self.runtime_state.pulled_file_outputs[name]}" - for name in self.config.mentioned_file_names - if name in self.runtime_state.pulled_file_outputs - ] + mentioned_file_sections = [] + for name in self.config.mentioned_file_names: + output = self.runtime_state.pulled_file_outputs.get(name) + if output is None: + continue + command = f"dify-agent config files pull {shlex.quote(name)}" + mentioned_file_sections.append( + f"Name: {name}\nPull command output for this run:\n{_format_command_output(command, output)}" + ) if mentioned_file_sections: sections.append("Mentioned files pulled locally:\n\n" + "\n\n".join(mentioned_file_sections)) @@ -125,11 +143,14 @@ class DifyConfigLayer(PlainLayer[DifyConfigDeps, DifyConfigLayerConfig, DifyConf sections: list[str] = [] if self.runtime_state.config_context_json: sections.append( - f"{_CONFIG_CONTEXT_COMMAND}\n{_CONFIG_CONTEXT_HEADING}\n{self.runtime_state.config_context_json}" + f"{_CONFIG_CONTEXT_HEADING}\n" + f"{_format_command_output(_CONFIG_CONTEXT_COMMAND, self.runtime_state.config_context_json)}" ) usage_lines = [_CONFIG_CLI_USAGE_PROMPT] if cli_help := self._format_config_cli_help(): usage_lines.append(cli_help) + if file_cli_help := self._format_agent_file_cli_help(): + usage_lines.append(file_cli_help) sections.append("\n".join(usage_lines)) return "\n\n".join(section for section in sections if section) @@ -142,13 +163,27 @@ class DifyConfigLayer(PlainLayer[DifyConfigDeps, DifyConfigLayerConfig, DifyConf if self._config_writable: commands.extend(_CONFIG_CLI_MUTATION_HELP_COMMANDS) command_sections = [ - f"$ {command}\n{self.runtime_state.config_cli_help[command]}" + _format_command_output(command, self.runtime_state.config_cli_help[command]) for command in commands if command in self.runtime_state.config_cli_help ] if not command_sections: return "" - return "Agent config CLI help:\n" + "\n\n".join(command_sections) + return "Agent config CLI reference for installed `dify-agent`:\n" + "\n\n".join(command_sections) + + def _format_agent_file_cli_help(self) -> str: + command_sections = [ + _format_command_output(command, self.runtime_state.config_cli_help[command]) + for command in _AGENT_FILE_CLI_HELP_COMMANDS + if command in self.runtime_state.config_cli_help + ] + if not command_sections: + return "" + return ( + "Agent file CLI reference for installed `dify-agent`:\n" + + "\n\n".join(command_sections) + + f"\n\n{_AGENT_FILE_UPLOAD_REPLY_HINT}" + ) def _format_config_context_json(self) -> str: return self.config.model_dump_json(exclude=_CONFIG_CONTEXT_EXCLUDE, exclude_none=True) @@ -226,4 +261,8 @@ class DifyConfigLayer(PlainLayer[DifyConfigDeps, DifyConfigLayerConfig, DifyConf return "\n".join(lines) +def _format_command_output(command: str, output: str) -> str: + return f"Command:\n$ {command}\nOutput:\n{output}" + + __all__ = ["DifyConfigLayer", "DifyConfigLayerError"] diff --git a/dify-agent/src/dify_agent/layers/drive/layer.py b/dify-agent/src/dify_agent/layers/drive/layer.py index fec1109e74d..8ac4b91c189 100644 --- a/dify-agent/src/dify_agent/layers/drive/layer.py +++ b/dify-agent/src/dify_agent/layers/drive/layer.py @@ -21,6 +21,7 @@ from typing_extensions import Self, override from agenton.layers import EmptyRuntimeState, LayerDeps, PlainLayer from dify_agent.agent_stub.protocol import agent_stub_drive_base_for_ref +from dify_agent.layers._agent_file_cli_help import AGENT_FILE_UPLOAD_REPLY_HINT as _AGENT_FILE_UPLOAD_REPLY_HINT from dify_agent.layers.drive.configs import DIFY_DRIVE_LAYER_TYPE_ID, DifyDriveLayerConfig from dify_agent.layers.shell.layer import DifyShellLayer @@ -122,13 +123,17 @@ class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntim def _format_agent_stub_cli_help(self) -> str: command_sections = [ - f"$ {command}\n{self._agent_stub_cli_help[command]}" + _format_command_output(command, self._agent_stub_cli_help[command]) for command in _AGENT_STUB_FILE_HELP_COMMANDS if command in self._agent_stub_cli_help ] if not command_sections: return "" - return "Agent Stub file CLI help:\n" + "\n\n".join(command_sections) + return ( + "Agent Stub file CLI reference for installed `dify-agent`:\n" + + "\n\n".join(command_sections) + + f"\n\n{_AGENT_FILE_UPLOAD_REPLY_HINT}" + ) async def _load_agent_stub_cli_help(self) -> None: self._agent_stub_cli_help = {} @@ -256,4 +261,8 @@ class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntim return f"{skill_key.rsplit('/', 1)[0]}/" +def _format_command_output(command: str, output: str) -> str: + return f"Command:\n$ {command}\nOutput:\n{output}" + + __all__ = ["DifyDriveLayer", "DifyDriveLayerError"] diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index 7346acc7b33..b2939aadadd 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -28,7 +28,6 @@ from pydantic_ai import Tool from typing_extensions import Self, override from agenton.layers import ( - EmptyLayerConfig, EmptyRuntimeState, LayerDeps, NoLayerDeps, @@ -45,17 +44,11 @@ from dify_agent.adapters.shell.protocols import ( ShellProviderProtocol, ShellResourceProtocol, ) -from dify_agent.agent_stub.server.shell_agent_stub_env import ShellAgentStubTokenFactory, build_shell_agent_stub_env +from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory, build_shell_agent_stub_env +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.shell.configs import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig from dify_agent.layers.shell.output_text import normalized_output_text, utf8_prefix, utf8_suffix -try: - from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer -except ModuleNotFoundError: - - class DifyExecutionContextLayer(PlainLayer[NoLayerDeps, EmptyLayerConfig, EmptyRuntimeState]): - """Minimal fallback for shell-only imports without server extras installed.""" - logger = logging.getLogger(__name__) @@ -79,36 +72,38 @@ _SHELL_OUTPUT_PROMPT_EDGE_BYTES = 8 * 1024 _SHELLCTL_OUTPUT_LIMIT_BYTES = 2 * _SHELL_OUTPUT_PROMPT_EDGE_BYTES _REMOTE_COMPLETE_OUTPUT_MAX_BYTES = 1024 * 1024 _REMOTE_COMMAND_TIMEOUT_SECONDS = 60.0 -_SHELL_LAYER_PREFIX_PROMPT = """You have access to a shell layer. It provides four tools: +_SHELL_LAYER_PREFIX_PROMPT = """You can run commands in an isolated shell workspace. + +Available shell tools: 1. shell_run - Start a new shell job in the current isolated workspace. - Use it to execute commands or scripts. + Starts a new shell job in the current workspace. + Use it to run commands or scripts. 2. shell_wait - Wait for more output or completion from an existing shell job. + Waits for more output or completion from an existing shell job. Use it when shell_run returns done=false. 3. shell_input - Send stdin text to a running shell job, then wait for new output. - Use it for interactive commands that are waiting for input. + Sends stdin text to a running shell job, then waits for new output. + Use it only when an interactive command is waiting for input. 4. shell_interrupt - Interrupt a running shell job. + Interrupts a running shell job. Use it to stop a long-running, stuck, or no-longer-needed command. Common arguments: - script: - The command or script to execute. Used by shell_run. + Command or script to execute. Used by shell_run. - job_id: - The id of a shell job returned by shell_run. + Shell job id returned by shell_run. Use it with shell_wait, shell_input, and shell_interrupt. Never invent a job_id. - timeout: - Maximum time, in seconds, to wait for output or completion for this tool call. + Maximum time in seconds to wait for output or completion for this tool call. A timeout does not necessarily mean the job has stopped; if done=false, use shell_wait again. - text: @@ -125,25 +120,33 @@ Usage rules: - Use shell_input only when the job is running and waiting for stdin. - Use shell_interrupt when a job is stuck or should be stopped. +Installed CLI: + +- `dify-agent` is already installed in this shell environment and can be used directly. +- Use the generated `dify-agent ... --help` output in the config prompt for exact command syntax. +- Do not install or recreate the `dify-agent` CLI. + Workspace persistence rules: -- The current workspace cwd is stable during this agent run, but it is temporary and may be deleted later. -- Do not use the current workspace cwd as persistent storage. -- $HOME outside the current workspace cwd is persistent storage. In build draft mode, when Agent config context reports - `config_version.kind` as `build_draft` and `config_version.writable` as true, changes there can be persisted for - later runs. In non-build-draft modes, those changes are rolled back. -- Saving config files, skills, env, or notes still requires the corresponding Agent config CLI mutation command; follow - the Agent config CLI help in the config layer. Shell file edits alone do not save config. +- The current workspace cwd is stable during this run, but it is temporary and may be deleted later. +- Do not treat files in the current workspace cwd as persisted state. +- In build mode, config changes persist only after you run the matching `dify-agent config ...` mutation command. +- Shell file edits alone do not save Agent config files, skills, env, or notes. +- In non-build modes, local shell changes are not a persistence mechanism for Agent configuration. -The script argument of shell_run can be a normal shell script, or a shebang script. -If the first line is a shebang, the shell layer executes the script directly. +shell_run script rules: + +- The script argument can be a normal shell script or a shebang script. +- If the first line is a shebang, the shell executes the script directly. Tips: - When using Python, prefer a uv script with a PEP 723 dependency header. +- If you need MCP, install the MCP server in the shell environment and start that server when you use it. - Example: +Example shell_run script: +[begin script] #!/usr/bin/env -S uv run --quiet --script # /// script # requires-python = ">=3.12" @@ -157,7 +160,8 @@ import httpx from rich import print response = httpx.get("https://example.com", timeout=10) -print(f"[green]status:[/green] {response.status_code}")""" +print(f"[green]status:[/green] {response.status_code}") +[end script]""" _SHELL_LAYER_SUFFIX_PROMPT = """Environment variables may contain API keys, tokens, or credentials. You may refer to environment variable names when needed.""" @@ -172,7 +176,7 @@ type ShellInterruptToolResult = str | ShellToolErrorObservation class DifyShellLayerDeps(LayerDeps): - execution_context: DifyExecutionContextLayer | None # pyright: ignore[reportUninitializedInstanceVariable] + execution_context: PlainLayer[NoLayerDeps, DifyExecutionContextLayerConfig, EmptyRuntimeState] | None # pyright: ignore[reportUninitializedInstanceVariable] class DifyShellRuntimeState(BaseModel): diff --git a/dify-agent/src/dify_agent/protocol/__init__.py b/dify-agent/src/dify_agent/protocol/__init__.py index a43b6cd4812..2fc5bd517ed 100644 --- a/dify-agent/src/dify_agent/protocol/__init__.py +++ b/dify-agent/src/dify_agent/protocol/__init__.py @@ -28,7 +28,6 @@ from .schemas import ( RunEventsResponse, RunFailedEvent, RunFailedEventData, - RunPurpose, RunLayerSpec, RunStartedEvent, RunStatus, @@ -78,7 +77,6 @@ __all__ = [ "RunEventsResponse", "RunFailedEvent", "RunFailedEventData", - "RunPurpose", "RunLayerSpec", "RunStartedEvent", "RunStatus", diff --git a/dify-agent/src/dify_agent/protocol/sandbox.py b/dify-agent/src/dify_agent/protocol/sandbox.py index 83f3d4b0adc..db375e6a7b5 100644 --- a/dify-agent/src/dify_agent/protocol/sandbox.py +++ b/dify-agent/src/dify_agent/protocol/sandbox.py @@ -112,6 +112,7 @@ class SandboxUploadedFile(BaseModel): transfer_method: Literal["tool_file"] = "tool_file" reference: str + download_url: str model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") diff --git a/dify-agent/src/dify_agent/protocol/schemas.py b/dify-agent/src/dify_agent/protocol/schemas.py index 1a02c7232a4..1cfafdc8123 100644 --- a/dify-agent/src/dify_agent/protocol/schemas.py +++ b/dify-agent/src/dify_agent/protocol/schemas.py @@ -24,10 +24,13 @@ whether each active layer is suspended or deleted when the run exits, with suspend as the default so successful terminal events can include resumable snapshots. Successful runs always publish the resumable Agenton session snapshot on the terminal ``run_succeeded`` event together with exactly one of the final -JSON-safe ``output`` or a deferred external ``deferred_tool_call`` payload. That -lets consumers treat terminal success events as complete run summaries without a -separate pause protocol. Session snapshots carry only layer lifecycle/runtime -state in compositor order; they do not persist output-layer config. Resumed +JSON-safe ``output`` or a deferred external ``deferred_tool_call`` payload. A +lifecycle-only run may also succeed with ``output = null`` and ``usage = null`` +when the composition intentionally omits the reserved model layer and only +replays layer enter/exit work from a supplied snapshot. That lets consumers +treat terminal success events as complete run summaries without a separate pause +protocol. Session snapshots carry only layer lifecycle/runtime state in +compositor order; they do not persist output-layer config. Resumed structured-output runs therefore must resubmit the same ``output`` layer in ``composition.layers[]`` so snapshot layer name/order still matches the composition and the runtime can rebuild the same structured output contract. @@ -50,7 +53,6 @@ DIFY_AGENT_MODEL_LAYER_ID: Final[str] = "llm" DIFY_AGENT_HISTORY_LAYER_ID: Final[str] = "history" DIFY_AGENT_OUTPUT_LAYER_ID: Final[str] = "output" RunStatus = Literal["running", "succeeded", "failed", "cancelled"] -RunPurpose = Literal["workflow_node", "single_step", "agent_app", "babysit", "fasten_preview"] RunEventType = Literal[ "run_started", "pydantic_ai_event", @@ -136,7 +138,6 @@ class CreateRunRequest(BaseModel): """ composition: RunComposition - purpose: RunPurpose = "workflow_node" idempotency_key: str | None = None metadata: dict[str, JsonValue] = Field(default_factory=dict) session_snapshot: CompositorSessionSnapshot | None = None @@ -334,10 +335,11 @@ class RunStartedEvent(BaseRunEvent): class PydanticAIStreamRunEvent(BaseRunEvent): - """Pydantic AI stream event using the upstream typed event model.""" + """Pydantic AI stream event with optional Dify Agent semantic annotations.""" type: Literal["pydantic_ai_event"] = "pydantic_ai_event" data: AgentStreamEvent + agent_message_delta: str | None = None class RunSucceededEvent(BaseRunEvent): @@ -402,7 +404,6 @@ __all__ = [ "RunEventsResponse", "RunFailedEvent", "RunFailedEventData", - "RunPurpose", "RunStartedEvent", "RunStatus", "RunStatusResponse", diff --git a/dify-agent/src/dify_agent/runtime/compositor_factory.py b/dify-agent/src/dify_agent/runtime/compositor_factory.py index 3d1e2afc488..1a9233eff99 100644 --- a/dify-agent/src/dify_agent/runtime/compositor_factory.py +++ b/dify-agent/src/dify_agent/runtime/compositor_factory.py @@ -17,7 +17,7 @@ plugin/knowledge business-layer family: Public DTOs provide Dify context plus plugin/model/tool data, while server-only plugin daemon settings and Dify API inner settings are injected through provider factories. Optional shellctl entrypoint/auth token and Agent Stub URL/token -issuer are injected for ``DifyShellLayer``. The resulting ``Compositor`` +factory are injected for ``DifyShellLayer``. The resulting ``Compositor`` remains Agenton state-only at the snapshot boundary: live resources such as HTTP clients are injected by runtime-owned providers, may be held on active layer instances inside ``resource_context()``, and never enter session @@ -27,7 +27,7 @@ snapshots. from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast from pydantic_ai.messages import UserContent @@ -36,7 +36,7 @@ from agenton.layers.types import AllPromptTypes, AllToolTypes, AllUserPromptType from agenton_collections.layers.pydantic_ai import PydanticAIHistoryLayer from agenton_collections.layers.plain.basic import PromptLayer from agenton_collections.transformers.pydantic_ai import PYDANTIC_AI_TRANSFORMERS -from dify_agent.agent_stub.server.shell_agent_stub_env import ShellAgentStubTokenFactory +from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory from dify_agent.layers.ask_human.layer import DifyAskHumanLayer from dify_agent.layers.config.layer import DifyConfigLayer from dify_agent.layers.dify_core_tools.configs import DifyCoreToolsLayerConfig @@ -50,14 +50,9 @@ from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer from dify_agent.layers.knowledge.configs import DifyKnowledgeBaseLayerConfig from dify_agent.layers.knowledge.layer import DifyKnowledgeBaseLayer from dify_agent.layers.output.output_layer import DifyOutputLayer -from dify_agent.adapters.shell.config import ShellAdapterSettings -from dify_agent.adapters.shell.factory import create_shell_provider from dify_agent.layers.shell.configs import DifyShellLayerConfig from dify_agent.layers.shell.layer import DifyShellLayer -if TYPE_CHECKING: - from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec - type DifyAgentLayerProvider = LayerProvider[Any] @@ -70,7 +65,7 @@ def create_default_layer_providers( shellctl_entrypoint: str | None = None, shellctl_auth_token: str | None = None, agent_stub_api_base_url: str | None = None, - agent_stub_token_codec: AgentStubTokenCodec | None = None, + agent_stub_token_factory: ShellAgentStubTokenFactory | None = None, ) -> tuple[DifyAgentLayerProvider, ...]: """Return the server provider set of safe config-constructible layers. @@ -79,20 +74,9 @@ def create_default_layer_providers( ``SHELLCTL_AUTH_TOKEN`` environment variable; deployments that enable shellctl bearer auth must set the Dify Agent server setting explicitly. """ - agent_stub_token_factory: ShellAgentStubTokenFactory | None = None - if agent_stub_token_codec is not None: + from dify_agent.adapters.shell.config import ShellAdapterSettings + from dify_agent.adapters.shell.factory import create_shell_provider - def build_agent_stub_token( - execution_context: DifyExecutionContextLayerConfig, - *, - session_id: str | None, - ) -> str: - return agent_stub_token_codec.encode_connection_token( - execution_context, - session_id=session_id, - ) - - agent_stub_token_factory = build_agent_stub_token shell_provider = ( create_shell_provider( ShellAdapterSettings( diff --git a/dify-agent/src/dify_agent/runtime/event_sink.py b/dify-agent/src/dify_agent/runtime/event_sink.py index 5dbe9fce62d..71df75b4916 100644 --- a/dify-agent/src/dify_agent/runtime/event_sink.py +++ b/dify-agent/src/dify_agent/runtime/event_sink.py @@ -89,11 +89,22 @@ async def emit_run_started(sink: RunEventSink, *, run_id: str) -> str: ) -async def emit_pydantic_ai_event(sink: RunEventSink, *, run_id: str, data: AgentStreamEvent) -> str: +async def emit_pydantic_ai_event( + sink: RunEventSink, + *, + run_id: str, + data: AgentStreamEvent, + agent_message_delta: str | None = None, +) -> str: """Emit one typed Pydantic AI stream event.""" return await emit_run_event( sink, - event=PydanticAIStreamRunEvent(run_id=run_id, data=data, created_at=utc_now()), + event=PydanticAIStreamRunEvent( + run_id=run_id, + data=data, + agent_message_delta=agent_message_delta, + created_at=utc_now(), + ), ) diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index 97fbfe4ba4c..2ad3aa99d8f 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/dify-agent/src/dify_agent/runtime/runner.py @@ -1,14 +1,21 @@ """Runtime execution for one scheduled Dify Agent run. The runner is storage-agnostic: it normalizes the public Dify composition into -Agenton's graph/config split, enters a fresh ``CompositorRun`` (or resumes one -from a snapshot), renders the current Dify system prompts into temporary -``message_history``, runs pydantic-ai with either the current ``run.user_prompts`` -or deferred external tool results, emits stream events, applies request-level -``on_exit`` signals, and then publishes a terminal success or failure event. The -Pydantic AI model is resolved from the active Agenton layer named by +Agenton's graph/config split and chooses one of two execution modes after the +composition is normalized and the ``on_exit`` policy is validated: + +- model runs: enter a fresh ``CompositorRun`` (or resume one from a snapshot), + render the current Dify system prompts into temporary ``message_history``, run + pydantic-ai with either the current ``run.user_prompts`` or deferred external + tool results, emit raw stream events with agent-message delta annotations, apply + request-level ``on_exit`` signals, and publish a terminal success or failure event; +- lifecycle-only runs: enter from a supplied snapshot, apply request-level + ``on_exit`` signals, exit without invoking a model, and succeed with explicit + ``output = null`` and ``usage = null``. + +The Pydantic AI model is resolved from the active Agenton layer named by ``DIFY_AGENT_MODEL_LAYER_ID``. An optional history layer contributes stored -message history only through session state; successful runs append only +message history only through session state; successful model runs append only ``result.new_messages()`` back into that layer so current system prompts are not persisted. An optional structured output layer named by ``DIFY_AGENT_OUTPUT_LAYER_ID`` is read after entry and resolved into an output @@ -31,11 +38,11 @@ from typing import Any, Literal, Protocol, cast, runtime_checkable import httpx from pydantic import JsonValue, TypeAdapter -from pydantic_ai.messages import AgentStreamEvent +from pydantic_ai.messages import AgentStreamEvent, PartDeltaEvent, PartStartEvent, TextPart, TextPartDelta from pydantic_ai.output import OutputSpec from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults -from agenton.compositor import CompositorSessionSnapshot, LayerProviderInput +from agenton.compositor import CompositorSessionSnapshot, LayerConfigInput, LayerProviderInput from agenton.layers.types import PydanticAITool from dify_agent.layers.ask_human.layer import get_ask_human_layer, validate_ask_human_layer_composition from dify_agent.layers.dify_core_tools.layer import DifyCoreToolsLayer @@ -97,6 +104,20 @@ class AgentRunValidationError(ValueError): """Raised when a run request is valid JSON but cannot execute.""" +def _has_model_layer(request: CreateRunRequest) -> bool: + """Return whether the public composition includes the reserved model layer.""" + return any(layer.name == DIFY_AGENT_MODEL_LAYER_ID for layer in request.composition.layers) + + +def _extract_agent_message_delta(event: AgentStreamEvent) -> str | None: + """Return agent-message text content from Pydantic AI stream events.""" + if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta): + return event.delta.content_delta + if isinstance(event, PartStartEvent) and isinstance(event.part, TextPart): + return event.part.content + return None + + @dataclass(slots=True) class RunSuccessOutcome: """Normalized successful runner output before event emission.""" @@ -163,7 +184,7 @@ class AgentRunRunner: await self.sink.update_status(self.run_id, "succeeded") async def _run_agent(self) -> RunSuccessOutcome: - """Run pydantic-ai inside an entered Agenton run. + """Run the normalized request in model or lifecycle-only mode. Known request-shaped Agenton enter-time failures are normalized to ``AgentRunValidationError``. That includes the existing small class of @@ -190,6 +211,57 @@ class AgentRunRunner: except (KeyError, TypeError, ValueError) as exc: raise AgentRunValidationError(str(exc)) from exc + if not _has_model_layer(self.request): + return await self._run_lifecycle_only(compositor=compositor, layer_configs=layer_configs) + return await self._run_model(compositor=compositor, layer_configs=layer_configs) + + async def _run_lifecycle_only( + self, + *, + compositor: Any, + layer_configs: dict[str, LayerConfigInput], + ) -> RunSuccessOutcome: + """Replay only layer lifecycle work for a no-LLM composition plus snapshot.""" + if self.request.session_snapshot is None: + raise AgentRunValidationError( + f"Missing '{DIFY_AGENT_MODEL_LAYER_ID}' requires a session_snapshot for lifecycle-only runs." + ) + if self.request.deferred_tool_results is not None: + raise AgentRunValidationError( + f"Deferred tool results require the reserved '{DIFY_AGENT_MODEL_LAYER_ID}' layer." + ) + + entered_run = False + try: + async with compositor.enter(configs=layer_configs, session_snapshot=self.request.session_snapshot) as run: + entered_run = True + apply_layer_exit_signals(run, self.request.on_exit) + except RuntimeError as exc: + if not entered_run and is_agenton_enter_validation_runtime_error(exc): + raise AgentRunValidationError(str(exc)) from exc + raise + except ValueError as exc: + if not entered_run: + raise AgentRunValidationError(str(exc)) from exc + raise + + if run.session_snapshot is None: + raise RuntimeError("Agenton run did not produce a session snapshot after exit.") + return RunSuccessOutcome( + result_kind="output", + output=None, + deferred_tool_call=None, + session_snapshot=run.session_snapshot, + usage=None, + ) + + async def _run_model( + self, + *, + compositor: Any, + layer_configs: dict[str, LayerConfigInput], + ) -> RunSuccessOutcome: + """Run the normal model/deferred-tool path inside an entered Agenton run.""" entered_run = False output: JsonValue | None = None deferred_tool_call: DeferredToolCallPayload | None = None @@ -206,7 +278,13 @@ class AgentRunRunner: async def handle_events(_ctx: object, events: AsyncIterable[AgentStreamEvent]) -> None: async for event in events: - _ = await emit_pydantic_ai_event(self.sink, run_id=self.run_id, data=event) + text_delta = _extract_agent_message_delta(event) + _ = await emit_pydantic_ai_event( + self.sink, + run_id=self.run_id, + data=event, + agent_message_delta=text_delta, + ) try: output_contract = resolve_run_output_contract(run) diff --git a/dify-agent/src/dify_agent/server/app.py b/dify-agent/src/dify_agent/server/app.py index ab619b7509a..49816d4daf0 100644 --- a/dify-agent/src/dify_agent/server/app.py +++ b/dify-agent/src/dify_agent/server/app.py @@ -22,9 +22,11 @@ import httpx from fastapi import FastAPI from redis.asyncio import Redis +from dify_agent.agent_stub.shell_env import ShellAgentStubTokenFactory from dify_agent.agent_stub.protocol.agent_stub import parse_agent_stub_endpoint from dify_agent.agent_stub.server.grpc_runtime import start_agent_stub_grpc_server from dify_agent.agent_stub.server.router import create_agent_stub_router +from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.runtime.compositor_factory import create_default_layer_providers from dify_agent.runtime.run_scheduler import RunScheduler from dify_agent.server.observability import configure_server_observability @@ -39,6 +41,21 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: """Build the FastAPI app with one shared Redis store and local scheduler.""" resolved_settings = settings or ServerSettings() agent_stub_token_codec = resolved_settings.create_agent_stub_token_codec() + agent_stub_token_factory: ShellAgentStubTokenFactory | None = None + if agent_stub_token_codec is not None: + # Runtime receives only this callable boundary; router and gRPC wiring + # keep the concrete token codec on the server side. + def issue_agent_stub_token( + execution_context: DifyExecutionContextLayerConfig, + *, + session_id: str | None, + ) -> str: + return agent_stub_token_codec.encode_connection_token( + execution_context, + session_id=session_id, + ) + + agent_stub_token_factory = issue_agent_stub_token agent_stub_file_request_handler = resolved_settings.create_agent_stub_file_request_handler() agent_stub_config_request_handler = resolved_settings.create_agent_stub_config_request_handler() agent_stub_drive_request_handler = resolved_settings.create_agent_stub_drive_request_handler() @@ -50,7 +67,7 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI: shellctl_entrypoint=resolved_settings.shellctl_entrypoint, shellctl_auth_token=resolved_settings.shellctl_auth_token, agent_stub_api_base_url=resolved_settings.agent_stub_api_base_url, - agent_stub_token_codec=agent_stub_token_codec, + agent_stub_token_factory=agent_stub_token_factory, ) sandbox_file_service = ( SandboxFileService(layer_providers=layer_providers) if resolved_settings.shellctl_entrypoint else None diff --git a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py index da541ca8cb0..69d71f865cf 100644 --- a/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py +++ b/dify-agent/tests/local/dify_agent/adapters/llm/test_model.py @@ -242,6 +242,84 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(response.parts[0].part_kind, "text") self.assertEqual(cast(TextPart, response.parts[0]).content, "adapter response") + async def test_request_uses_unique_fallback_ids_for_same_name_tool_calls(self) -> None: + messages = [ + ModelRequest(parts=[UserPromptPart("hello")]), + ModelResponse( + parts=[ + ToolCallPart(tool_name="lookup", args={"query": "first"}, tool_call_id=""), + ToolCallPart(tool_name="lookup", args={"query": "second"}, tool_call_id=""), + ] + ), + ] + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content.decode("utf-8")) + prompt_messages = payload["data"]["prompt_messages"] + tool_calls = prompt_messages[1]["tool_calls"] + + self.assertEqual(tool_calls[0]["id"], "tool-call-0-lookup") + self.assertEqual(tool_calls[1]["id"], "tool-call-1-lookup") + + return build_stream_response(*single_text_chunk("adapter response", prompt_tokens=11, completion_tokens=7)) + + async with self.mock_daemon_stream(httpx.MockTransport(handler)): + adapter = DifyLLMAdapterModel( + "demo-model", + self.make_provider(), + model_provider="openai", + credentials={"api_key": "secret"}, + ) + + response = await adapter.request( + messages, + model_settings=None, + model_request_parameters=ModelRequestParameters(), + ) + + self.assertEqual(response.model_name, "demo-model") + self.assertEqual(response.parts[0].part_kind, "text") + self.assertEqual(cast(TextPart, response.parts[0]).content, "adapter response") + + async def test_request_collapses_text_only_assistant_history_parts_to_string_content(self) -> None: + messages = [ + ModelRequest(parts=[UserPromptPart("initial request")]), + ModelResponse( + parts=[ + ThinkingPart(content="plan"), + TextPart(content="answer"), + ] + ), + ModelRequest(parts=[UserPromptPart("follow up")]), + ] + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content.decode("utf-8")) + prompt_messages = payload["data"]["prompt_messages"] + + self.assertEqual([message["role"] for message in prompt_messages], ["user", "assistant", "user"]) + self.assertEqual(prompt_messages[1]["content"], "\nplan\nanswer") + + return build_stream_response(*single_text_chunk("adapter response", prompt_tokens=11, completion_tokens=7)) + + async with self.mock_daemon_stream(httpx.MockTransport(handler)): + adapter = DifyLLMAdapterModel( + "demo-model", + self.make_provider(), + model_provider="openai", + credentials={"api_key": "secret"}, + ) + + response = await adapter.request( + messages, + model_settings=None, + model_request_parameters=ModelRequestParameters(), + ) + + self.assertEqual(response.model_name, "demo-model") + self.assertEqual(response.parts[0].part_kind, "text") + self.assertEqual(cast(TextPart, response.parts[0]).content, "adapter response") + async def test_request_omits_empty_assistant_history_when_response_has_no_content_or_tool_calls(self) -> None: messages = [ ModelRequest(parts=[SystemPromptPart("request system"), UserPromptPart("hello")]), diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py index 61ffcec9e3a..5dbf0bafe5a 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py @@ -1,5 +1,7 @@ from __future__ import annotations +import base64 +import json from io import BytesIO from pathlib import Path import stat @@ -14,16 +16,22 @@ from dify_agent.agent_stub.cli._drive import ( pull_drive_from_environment, push_drive_from_environment, ) -from dify_agent.agent_stub.cli._files import UploadedToolFileMapping, UploadedToolFileResource +from dify_agent.agent_stub.cli._files import UploadedToolFileResource from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubDriveCommitRequest, AgentStubDriveCommitResponse, AgentStubDriveItem, + AgentStubFileMapping, AgentStubDriveManifestResponse, ) +def _reference(record_id: str) -> str: + payload = base64.urlsafe_b64encode(json.dumps({"record_id": record_id}, separators=(",", ":")).encode()).decode() + return f"dify-file-ref:{payload}" + + def test_list_drive_manifest_from_environment_returns_manifest_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") @@ -535,7 +543,7 @@ def test_push_drive_from_environment_commits_single_file(monkeypatch: pytest.Mon monkeypatch.setattr( "dify_agent.agent_stub.cli._drive.upload_tool_file_resource_from_environment", lambda *, path: UploadedToolFileResource( - mapping=UploadedToolFileMapping(reference="dify-file-ref:tool-file-1"), + mapping=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")), tool_file_id="tool-file-1", ), ) @@ -640,7 +648,10 @@ def test_push_drive_from_environment_kind_skill_standardizes_skill_directory( def fake_upload(*, path: str) -> UploadedToolFileResource: uploaded_paths.append(Path(path).name) return UploadedToolFileResource( - mapping=UploadedToolFileMapping(reference=f"dify-file-ref:{Path(path).name}"), + mapping=AgentStubFileMapping( + transfer_method="tool_file", + reference=_reference(Path(path).name), + ), tool_file_id=Path(path).name, ) @@ -697,7 +708,10 @@ def test_push_drive_from_environment_kind_skill_archive_excludes_transient_entri with ZipFile(path) as archive: archive_entries.extend(sorted(archive.namelist())) return UploadedToolFileResource( - mapping=UploadedToolFileMapping(reference=f"dify-file-ref:{Path(path).name}"), + mapping=AgentStubFileMapping( + transfer_method="tool_file", + reference=_reference(Path(path).name), + ), tool_file_id=Path(path).name, ) @@ -821,7 +835,10 @@ def test_push_drive_from_environment_kind_dir_keeps_user_files_that_skill_packag def fake_upload(*, path: str) -> UploadedToolFileResource: uploaded_paths.append(Path(path).relative_to(root).as_posix()) return UploadedToolFileResource( - mapping=UploadedToolFileMapping(reference=f"dify-file-ref:{Path(path).name}"), + mapping=AgentStubFileMapping( + transfer_method="tool_file", + reference=_reference(Path(path).name), + ), tool_file_id=Path(path).name, ) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py index c2f2d44afb4..6a65daccb9a 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py @@ -12,6 +12,7 @@ from dify_agent.agent_stub.cli._files import ( upload_tool_file_resource_from_environment, ) from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError +from dify_agent.agent_stub.protocol.agent_stub import AgentStubFileMapping def _reference(record_id: str) -> str: @@ -48,18 +49,42 @@ def test_upload_file_from_environment_requests_signed_url_and_normalizes_output( "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", fake_upload_file_to_signed_url_sync, ) + captured_download_request: dict[str, object] = {} + + def fake_request_agent_stub_file_download_sync(**kwargs): + captured_download_request["file"] = kwargs["file"] + return type( + "Response", + (), + { + "filename": "report.pdf", + "mime_type": "application/pdf", + "size": 12, + "download_url": "https://files.example.com/download", + }, + )() + + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", + fake_request_agent_stub_file_download_sync, + ) result = upload_file_from_environment(path=str(source)) assert result.model_dump() == { "transfer_method": "tool_file", "reference": _reference("tool-file-1"), + "download_url": "https://files.example.com/download", } assert captured == { "filename": "report.pdf", "mimetype": "application/pdf", "file_bytes": b"report-bytes", } + assert captured_download_request["file"] == AgentStubFileMapping( + transfer_method="tool_file", + reference=_reference("tool-file-1"), + ) def test_upload_tool_file_resource_from_environment_preserves_tool_file_id( @@ -79,12 +104,17 @@ def test_upload_tool_file_resource_from_environment_preserves_tool_file_id( "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", lambda **_kwargs: {"id": "tool-file-1", "reference": _reference("tool-file-1")}, ) + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", + lambda **_kwargs: pytest.fail("resource helper must not request download_url"), + ) result = upload_tool_file_resource_from_environment(path=str(source)) assert result.mapping.model_dump() == { "transfer_method": "tool_file", "reference": _reference("tool-file-1"), + "url": None, } assert result.tool_file_id == "tool-file-1" @@ -187,6 +217,34 @@ def test_upload_file_from_environment_rejects_non_canonical_reference( _ = upload_file_from_environment(path=str(source)) +def test_upload_file_from_environment_rejects_missing_download_url( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source = tmp_path / "report.pdf" + source.write_bytes(b"report-bytes") + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", + lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), + ) + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", + lambda **_kwargs: {"id": "tool-file-1", "reference": _reference("tool-file-1")}, + ) + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", + lambda **_kwargs: type( + "Response", (), {"filename": "report.pdf", "mime_type": "application/pdf", "size": 12} + )(), + ) + + with pytest.raises(AgentStubTransferError, match="missing download_url"): + _ = upload_file_from_environment(path=str(source)) + + def test_download_file_from_environment_supports_mapping_json( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py index b87e79daa35..3018629e105 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py @@ -3,11 +3,13 @@ from __future__ import annotations import base64 import json from pathlib import Path +from types import SimpleNamespace import pytest from dify_agent.agent_stub.cli._drive import DrivePullResult from dify_agent.agent_stub.cli.main import main +from dify_agent.agent_stub.client._errors import AgentStubTransferError from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubConfigFileItemsResponse, AgentStubConfigFileItem, @@ -38,6 +40,13 @@ def _config_manifest_response() -> AgentStubConfigManifestResponse: ) +def _patch_cli_module(monkeypatch: pytest.MonkeyPatch, accessor_name: str, **attrs: object) -> None: + monkeypatch.setattr( + f"dify_agent.agent_stub.cli.main.{accessor_name}", + lambda: SimpleNamespace(**attrs), + ) + + def test_cli_connect_reports_missing_environment_variables(capsys: pytest.CaptureFixture[str]) -> None: with pytest.raises(SystemExit) as exc_info: main(["connect"]) @@ -59,7 +68,7 @@ def test_cli_connect_supports_json_output( assert argv == ["echo", "hello"] return AgentStubConnectResponse(connection_id="conn-1", status="connected") - monkeypatch.setattr("dify_agent.agent_stub.cli.main.connect_from_environment", fake_connect_from_environment) + _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) main(["connect", "--json", "--", "echo", "hello"]) @@ -78,7 +87,7 @@ def test_cli_unknown_command_auto_forwards_when_agent_stub_env_is_present( assert argv == ["run", "--target", "prod"] return AgentStubConnectResponse(connection_id="conn-1", status="connected") - monkeypatch.setattr("dify_agent.agent_stub.cli.main.connect_from_environment", fake_connect_from_environment) + _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) main(["run", "--target", "prod"]) @@ -220,7 +229,7 @@ def test_cli_connect_accepts_grpc_agent_stub_api_base_url( assert argv == ["echo", "hello"] return AgentStubConnectResponse(connection_id="conn-1", status="connected") - monkeypatch.setattr("dify_agent.agent_stub.cli.main.connect_from_environment", fake_connect_from_environment) + _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) main(["connect", "echo", "hello"]) @@ -232,9 +241,10 @@ def test_cli_config_manifest_omits_hash_fields( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.manifest_from_environment", - lambda: AgentStubConfigManifestResponse( + _patch_cli_module( + monkeypatch, + "_config_module", + manifest_from_environment=lambda: AgentStubConfigManifestResponse( agent_id="agent-1", config_version=AgentStubConfigVersionInfo(id="cfg-1", kind="build_draft", writable=True), skills=AgentStubConfigSkillItemsResponse( @@ -335,7 +345,7 @@ def test_cli_config_mutation_commands_forward_and_print_manifest_json( captured_kwargs.update(kwargs) return _config_manifest_response() - monkeypatch.setattr(f"dify_agent.agent_stub.cli.main.{helper_name}", fake_helper) + _patch_cli_module(monkeypatch, "_config_module", **{helper_name: fake_helper}) with pytest.raises(SystemExit) as exc_info: main(argv) @@ -394,7 +404,7 @@ def test_cli_config_pull_commands_support_plural_and_hidden_singular_aliases( captured_kwargs["local_dir"] = local_dir return response - monkeypatch.setattr(f"dify_agent.agent_stub.cli.main.{helper_name}", fake_helper) + _patch_cli_module(monkeypatch, "_config_module", **{helper_name: fake_helper}) with pytest.raises(SystemExit) as exc_info: main(argv) @@ -430,9 +440,10 @@ def test_cli_file_upload_prints_uploaded_tool_file_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.upload_file_from_environment", - lambda *, path: type( + _patch_cli_module( + monkeypatch, + "_files_module", + upload_file_from_environment=lambda *, path: type( "Response", (), { @@ -440,6 +451,7 @@ def test_cli_file_upload_prints_uploaded_tool_file_json( { "transfer_method": "tool_file", "reference": _reference(Path(path).name), + "download_url": f"https://files.example.com/{Path(path).name}", } ) }, @@ -454,16 +466,39 @@ def test_cli_file_upload_prints_uploaded_tool_file_json( assert json.loads(captured.out) == { "transfer_method": "tool_file", "reference": _reference("report.pdf"), + "download_url": "https://files.example.com/report.pdf", } +def test_cli_file_upload_exits_non_zero_without_partial_json_when_download_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + _patch_cli_module( + monkeypatch, + "_files_module", + upload_file_from_environment=lambda *, path: (_ for _ in ()).throw( + AgentStubTransferError("signed file download request failed") + ), + ) + + with pytest.raises(SystemExit) as exc_info: + main(["file", "upload", "/tmp/report.pdf"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 1 + assert captured.out == "" + assert "signed file download request failed" in captured.err + + def test_cli_file_download_prints_saved_path( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.download_file_from_environment", - lambda **_kwargs: type("Response", (), {"path": Path("/tmp/report.pdf")})(), + _patch_cli_module( + monkeypatch, + "_files_module", + download_file_from_environment=lambda **_kwargs: type("Response", (), {"path": Path("/tmp/report.pdf")})(), ) with pytest.raises(SystemExit) as exc_info: @@ -474,18 +509,20 @@ def test_cli_file_download_prints_saved_path( assert captured.out.strip() == "/tmp/report.pdf" -def test_cli_file_download_supports_mapping_json( +def test_cli_file_download_rejects_mapping_option( monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], ) -> None: - captured_kwargs: dict[str, object] = {} + called = False - def fake_download_file_from_environment(**kwargs): - captured_kwargs.update(kwargs) + def fake_download_file_from_environment(**_kwargs): + nonlocal called + called = True return type("Response", (), {"path": Path("/tmp/inputs/report.pdf")})() - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.download_file_from_environment", fake_download_file_from_environment + _patch_cli_module( + monkeypatch, + "_files_module", + download_file_from_environment=fake_download_file_from_environment, ) with pytest.raises(SystemExit) as exc_info: @@ -500,15 +537,8 @@ def test_cli_file_download_supports_mapping_json( ] ) - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == { - "transfer_method": None, - "reference_or_url": None, - "mapping": json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), - "local_dir": "/tmp/inputs", - } - assert captured.out.strip() == "/tmp/inputs/report.pdf" + assert exc_info.value.code == 2 + assert called is False def test_cli_file_download_rejects_legacy_positional_directory( @@ -522,8 +552,10 @@ def test_cli_file_download_rejects_legacy_positional_directory( called = True return type("Response", (), {"path": Path("/tmp/report.pdf")})() - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.download_file_from_environment", fake_download_file_from_environment + _patch_cli_module( + monkeypatch, + "_files_module", + download_file_from_environment=fake_download_file_from_environment, ) with pytest.raises(SystemExit) as exc_info: @@ -539,9 +571,10 @@ def test_cli_drive_list_prints_manifest_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.list_drive_manifest_from_environment", - lambda *, prefix: AgentStubDriveManifestResponse( + _patch_cli_module( + monkeypatch, + "_drive_module", + list_drive_manifest_from_environment=lambda *, prefix: AgentStubDriveManifestResponse( items=[ AgentStubDriveItem( key=prefix + "example/SKILL.md", @@ -553,6 +586,10 @@ def test_cli_drive_list_prints_manifest_json( ) ] ), + format_drive_manifest=lambda response: ( + f"{response.items[0].size}\t{response.items[0].mime_type}\t{response.items[0].hash or '-'}\t" + f"{response.items[0].key}" + ), ) with pytest.raises(SystemExit) as exc_info: @@ -567,9 +604,10 @@ def test_cli_drive_list_prints_human_readable_listing( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.list_drive_manifest_from_environment", - lambda *, prefix: AgentStubDriveManifestResponse( + _patch_cli_module( + monkeypatch, + "_drive_module", + list_drive_manifest_from_environment=lambda *, prefix: AgentStubDriveManifestResponse( items=[ AgentStubDriveItem( key=f"{prefix}example/SKILL.md", @@ -581,6 +619,10 @@ def test_cli_drive_list_prints_human_readable_listing( ) ] ), + format_drive_manifest=lambda response: ( + f"{response.items[0].size}\t{response.items[0].mime_type}\t{response.items[0].hash or '-'}\t" + f"{response.items[0].key}" + ), ) with pytest.raises(SystemExit) as exc_info: @@ -595,9 +637,10 @@ def test_cli_drive_pull_prints_downloaded_paths( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - lambda *, targets, local_base: DrivePullResult( + _patch_cli_module( + monkeypatch, + "_drive_module", + pull_drive_from_environment=lambda *, targets, local_base: DrivePullResult( items=[ DrivePullResult.Item( key=f"{targets[0]}/SKILL.md", local_path=str(Path(local_base) / targets[0] / "SKILL.md") @@ -624,9 +667,10 @@ def test_cli_drive_pull_prints_json_result( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - lambda *, targets, local_base: DrivePullResult( + _patch_cli_module( + monkeypatch, + "_drive_module", + pull_drive_from_environment=lambda *, targets, local_base: DrivePullResult( items=[ DrivePullResult.Item(key="files/a.txt", local_path=f"{local_base}/files/a.txt"), DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=f"{local_base}/skills/foo/SKILL.md"), @@ -664,10 +708,7 @@ def test_cli_drive_pull_forwards_multiple_targets( ] ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - fake_pull_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "pull", "skills/foo", "files/a.txt", "--to", "/tmp/drive"]) @@ -696,10 +737,7 @@ def test_cli_drive_pull_uses_environment_drive_base_default( ] ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - fake_pull_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "pull", "skills/foo"]) @@ -728,10 +766,7 @@ def test_cli_drive_pull_keeps_historical_drive_base_when_env_is_missing( ] ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - fake_pull_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "pull", "skills/foo"]) @@ -755,10 +790,7 @@ def test_cli_drive_pull_without_targets_pulls_whole_visible_drive( items=[DrivePullResult.Item(key="files/a.txt", local_path=str(Path(local_base) / "files" / "a.txt"))] ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - fake_pull_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "pull", "--to", "/tmp/drive"]) @@ -773,9 +805,10 @@ def test_cli_drive_push_prints_commit_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.push_drive_from_environment", - lambda *, local_path, drive_path, kind: AgentStubDriveCommitResponse( + _patch_cli_module( + monkeypatch, + "_drive_module", + push_drive_from_environment=lambda *, local_path, drive_path, kind: AgentStubDriveCommitResponse( items=[ AgentStubDriveItem( key=drive_path, @@ -810,10 +843,7 @@ def test_cli_drive_push_forwards_kind( captured_kwargs["kind"] = kind return AgentStubDriveCommitResponse(items=[]) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.push_drive_from_environment", - fake_push_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "push", "/tmp/skill", "skills/example", "--kind", "skill"]) @@ -839,10 +869,7 @@ def test_cli_drive_push_accepts_json_flag( captured_kwargs["kind"] = kind return AgentStubDriveCommitResponse(items=[]) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.push_drive_from_environment", - fake_push_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "push", "/tmp/report.md", "files/report.md", "--json"]) @@ -868,10 +895,7 @@ def test_cli_drive_push_rejects_recursive_option( called = True return AgentStubDriveCommitResponse(items=[]) - monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.push_drive_from_environment", - fake_push_drive_from_environment, - ) + _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) with pytest.raises(SystemExit) as exc_info: main(["drive", "push", "/tmp/dir", "files/dir", "--recursive"]) diff --git a/dify-agent/tests/local/dify_agent/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py index 3a1525f015c..9f15ea9d2e6 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/dify-agent/tests/local/dify_agent/client/test_client.py @@ -259,7 +259,11 @@ def test_sync_sandbox_methods_post_dtos_and_parse_responses() -> None: 200, json={ "path": "report.txt", - "file": {"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, + "file": { + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + "download_url": "https://files.example.com/report.txt", + }, }, ) raise AssertionError(f"unexpected request: {request.method} {request.url}") @@ -276,6 +280,7 @@ def test_sync_sandbox_methods_post_dtos_and_parse_responses() -> None: assert preview.text == "hello" assert isinstance(uploaded, SandboxUploadResponse) assert uploaded.file.reference == "dify-file-ref:file-1" + assert uploaded.file.download_url == "https://files.example.com/report.txt" def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: @@ -293,7 +298,11 @@ def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: 200, json={ "path": "report.txt", - "file": {"transfer_method": "tool_file", "reference": "dify-file-ref:file-1"}, + "file": { + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + "download_url": "https://files.example.com/report.txt", + }, }, ) raise AssertionError(f"unexpected request: {request.method} {request.url}") @@ -309,11 +318,35 @@ def test_async_sandbox_methods_post_dtos_and_parse_responses() -> None: assert listing.path == "." assert preview.text == "hello" assert uploaded.file.reference == "dify-file-ref:file-1" + assert uploaded.file.download_url == "https://files.example.com/report.txt" await http_client.aclose() asyncio.run(scenario()) +def test_sync_upload_sandbox_file_rejects_missing_download_url() -> None: + locator = _sandbox_locator() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path != "/sandbox/files/upload": + raise AssertionError(f"unexpected request: {request.method} {request.url}") + return httpx.Response( + 200, + json={ + "path": "report.txt", + "file": { + "transfer_method": "tool_file", + "reference": "dify-file-ref:file-1", + }, + }, + ) + + client = Client(base_url="http://testserver", sync_http_client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(DifyAgentValidationError): + _ = client.upload_sandbox_file_sync(locator, "report.txt") + + def test_sync_sandbox_methods_map_invalid_json_to_validation_error() -> None: responses = iter([httpx.Response(200, text="not-json"), httpx.Response(404, json={"detail": "missing"})]) diff --git a/dify-agent/tests/local/dify_agent/layers/config/test_layer.py b/dify-agent/tests/local/dify_agent/layers/config/test_layer.py index 0ea4bf1195f..f085f6cc51b 100644 --- a/dify-agent/tests/local/dify_agent/layers/config/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/config/test_layer.py @@ -9,7 +9,11 @@ import pytest from dify_agent.adapters.shell.shellctl import ShellctlProvider from dify_agent.layers.config import DifyConfigLayerConfig -from dify_agent.layers.config.layer import DifyConfigLayer, DifyConfigLayerError +from dify_agent.layers.config.layer import ( + DifyConfigLayer, + DifyConfigLayerError, + _AGENT_FILE_UPLOAD_REPLY_HINT, +) from dify_agent.layers.shell import DifyShellLayerConfig from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer @@ -127,7 +131,19 @@ async def test_on_context_create_computes_runtime_fields_and_pulls_mentioned_ass assert layer.runtime_state.pulled_skill_outputs == {"alpha": "/workspace/.dify_conf/skills/alpha\n# Alpha\nUse it."} assert layer.runtime_state.pulled_file_outputs == {"guide.txt": "/workspace/.dify_conf/files/guide.txt"} assert "dify-agent config note push --help" in layer.runtime_state.config_cli_help + assert "dify-agent file upload --help" in layer.runtime_state.config_cli_help + assert "dify-agent file download --help" in layer.runtime_state.config_cli_help assert layer.runtime_state.push_spec_json_schema == "" + suffix_prompt = layer.build_suffix_prompt() + assert suffix_prompt.index("Agent config CLI reference for installed `dify-agent`:") < suffix_prompt.index( + "Agent file CLI reference for installed `dify-agent`:" + ) + assert "$ dify-agent file upload --help" in suffix_prompt + assert "$ dify-agent file download --help" in suffix_prompt + assert suffix_prompt.index("$ dify-agent file upload --help") < suffix_prompt.index( + "$ dify-agent file download --help" + ) + assert _AGENT_FILE_UPLOAD_REPLY_HINT in suffix_prompt @pytest.mark.anyio diff --git a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py index 86b04f95297..afcef5cc396 100644 --- a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py @@ -8,7 +8,7 @@ import pytest from dify_agent.adapters.shell.shellctl import ShellctlProvider from dify_agent.layers.drive import DifyDriveLayerConfig, DifyDriveSkillConfig -from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError +from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError, _AGENT_FILE_UPLOAD_REPLY_HINT from dify_agent.layers.shell import DifyShellLayerConfig from dify_agent.layers.shell.layer import CompleteRemoteCommandResult, DifyShellLayer @@ -112,14 +112,18 @@ def test_drive_layer_exposes_agent_stub_cli_usage_suffix_prompt() -> None: layer._agent_stub_cli_help = { "dify-agent file --help": _file_help_output("dify-agent file --help"), "dify-agent file upload --help": _file_help_output("dify-agent file upload --help"), + "dify-agent file download --help": _file_help_output("dify-agent file download --help"), } assert len(layer.suffix_prompts) == 1 prompt = layer.suffix_prompts[0] assert "Other available skills" in prompt assert "other-skill: Other Skill" in prompt - assert "Agent Stub file CLI help" in prompt + assert "Agent Stub file CLI reference for installed `dify-agent`" in prompt assert "$ dify-agent file upload --help" in prompt + assert "$ dify-agent file download --help" in prompt + assert prompt.index("$ dify-agent file upload --help") < prompt.index("$ dify-agent file download --help") + assert _AGENT_FILE_UPLOAD_REPLY_HINT in prompt assert "dify-agent drive" not in prompt diff --git a/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py b/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py index 75fa8dda204..be5ac5be9a5 100644 --- a/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py +++ b/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py @@ -42,7 +42,11 @@ from dify_agent.layers.dify_plugin.configs import ( def test_run_event_adapter_round_trips_typed_variants() -> None: events = [ RunStartedEvent(run_id="run-1"), - PydanticAIStreamRunEvent(run_id="run-1", data=FinalResultEvent(tool_name=None, tool_call_id=None)), + PydanticAIStreamRunEvent( + run_id="run-1", + data=FinalResultEvent(tool_name=None, tool_call_id=None), + agent_message_delta="hello", + ), RunSucceededEvent( run_id="run-1", data=RunSucceededEventData( @@ -101,6 +105,7 @@ def test_create_run_request_rejects_old_compositor_payload_and_model_layer_id_is def test_protocol_package_no_longer_exports_execution_context_dto() -> None: assert not hasattr(protocol_exports, "ExecutionContext") + assert not hasattr(protocol_exports, "RunPurpose") def test_create_run_request_accepts_dto_first_public_composition_and_normalizes_graph_config() -> None: @@ -131,7 +136,6 @@ def test_create_run_request_accepts_dto_first_public_composition_and_normalizes_ } ) request = CreateRunRequest( - purpose="workflow_node", idempotency_key="workflow-run-1:node-execution-1", metadata={"source": "unit_test"}, composition=RunComposition( @@ -177,7 +181,6 @@ def test_create_run_request_accepts_dto_first_public_composition_and_normalizes_ "invoke_from": "service-api", "trace_id": "trace-1", } - assert payload["purpose"] == "workflow_node" assert payload["idempotency_key"] == "workflow-run-1:node-execution-1" assert payload["metadata"] == {"source": "unit_test"} assert payload["composition"]["layers"][0]["config"] == {"prefix": "system", "user": "hello", "suffix": []} @@ -456,6 +459,16 @@ def test_create_run_request_rejects_removed_top_level_execution_context() -> Non ) +def test_create_run_request_rejects_removed_top_level_purpose() -> None: + with pytest.raises(ValidationError): + _ = CreateRunRequest.model_validate( + { + "composition": {"layers": []}, + "purpose": "session_cleanup", + } + ) + + def test_layer_exit_signals_reject_extra_fields() -> None: with pytest.raises(ValidationError): _ = LayerExitSignals.model_validate({"default": "suspend", "unknown": "value"}) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py b/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py index 30ad414e45a..6802f525a05 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_compositor_factory.py @@ -85,7 +85,6 @@ if "jsonschema" not in sys.modules: sys.modules["jsonschema.protocols"] = jsonschema_protocols_module sys.modules["jsonschema.validators"] = jsonschema_validators_module -import dify_agent.runtime.compositor_factory as compositor_factory_module from dify_agent.adapters.shell.config import ShellAdapterSettings from dify_agent.adapters.shell.protocols import ShellProviderProtocol from dify_agent.layers.dify_core_tools import DIFY_CORE_TOOLS_LAYER_TYPE_ID, DifyCoreToolsLayerConfig @@ -113,7 +112,7 @@ def test_default_layer_providers_register_shell_layer_with_configured_token_fact captured_settings.append(settings) return cast(ShellProviderProtocol, fake_provider) - monkeypatch.setattr(compositor_factory_module, "create_shell_provider", fake_create_shell_provider) + monkeypatch.setattr("dify_agent.adapters.shell.factory.create_shell_provider", fake_create_shell_provider) providers = create_default_layer_providers( shellctl_entrypoint="http://shellctl.example", @@ -138,7 +137,7 @@ def test_default_layer_providers_keep_empty_shellctl_token_by_default( captured_settings.append(settings) return cast(ShellProviderProtocol, FakeProvider()) - monkeypatch.setattr(compositor_factory_module, "create_shell_provider", fake_create_shell_provider) + monkeypatch.setattr("dify_agent.adapters.shell.factory.create_shell_provider", fake_create_shell_provider) providers = create_default_layer_providers(shellctl_entrypoint="http://shellctl.example") shell_provider = next(provider for provider in providers if provider.type_id == DIFY_SHELL_LAYER_TYPE_ID) @@ -153,18 +152,21 @@ def test_shell_provider_rejects_blank_settings_entrypoint_when_default_providers _ = create_default_layer_providers(shellctl_entrypoint=" ") -def test_default_layer_providers_build_agent_stub_token_factory_from_agent_stub_codec() -> None: - AgentStubTokenCodec = pytest.importorskip( - "dify_agent.agent_stub.server.tokens.agent_stub", - reason="jwcrypto is not available in this local test environment", - ).AgentStubTokenCodec +def test_default_layer_providers_forward_agent_stub_token_factory() -> None: + captured_calls: list[tuple[DifyExecutionContextLayerConfig, str | None]] = [] - codec = AgentStubTokenCodec.from_server_secret("MTExMTExMTExMTExMTExMTExMTExMTExMTExMTExMTE") + def build_agent_stub_token( + execution_context: DifyExecutionContextLayerConfig, + *, + session_id: str | None, + ) -> str: + captured_calls.append((execution_context, session_id)) + return f"token-for:{execution_context.tenant_id}:{session_id}" providers = create_default_layer_providers( shellctl_entrypoint="http://shellctl.example", agent_stub_api_base_url="https://agent.example.com/agent-stub", - agent_stub_token_codec=codec, + agent_stub_token_factory=build_agent_stub_token, ) shell_provider = next(provider for provider in providers if provider.type_id == DIFY_SHELL_LAYER_TYPE_ID) shell_layer = shell_provider.create_layer(DifyShellLayerConfig()) @@ -180,8 +182,19 @@ def test_default_layer_providers_build_agent_stub_token_factory_from_agent_stub_ session_id="abc12ff", ) - assert isinstance(token, str) - assert token + assert token == "token-for:tenant-1:abc12ff" + assert captured_calls == [ + ( + DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_id="user-1", + user_from="account", + agent_mode="workflow_run", + invoke_from="service-api", + ), + "abc12ff", + ) + ] def test_default_layer_providers_register_core_tools_layer() -> None: diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index 6c05cb65703..87e12854496 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -56,6 +56,7 @@ from dify_agent.protocol.schemas import ( CreateRunRequest, DeferredToolResultsPayload, LayerExitSignals, + PydanticAIStreamRunEvent, RunComposition, RunLayerSpec, RunSucceededEvent, @@ -196,6 +197,44 @@ def _request( ) +def _lifecycle_only_request( + *, + on_exit: LayerExitSignals | None = None, + session_snapshot: CompositorSessionSnapshot | None = None, + deferred_tool_results: DeferredToolResultsPayload | None = None, +) -> CreateRunRequest: + snapshot = session_snapshot or CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot(name="prompt", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), + LayerSessionSnapshot(name="execution_context", lifecycle_state=LifecycleState.SUSPENDED, runtime_state={}), + ] + ) + return CreateRunRequest( + composition=RunComposition( + layers=[ + RunLayerSpec( + name="prompt", + type="plain.prompt", + config=PromptLayerConfig(prefix="system", user="hello"), + ), + RunLayerSpec( + name="execution_context", + type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, + config=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_from="account", + agent_mode="workflow_run", + invoke_from="service-api", + ), + ), + ] + ), + session_snapshot=snapshot, + deferred_tool_results=deferred_tool_results, + on_exit=on_exit or LayerExitSignals(default=ExitIntent.DELETE), + ) + + def _recursive_output_schema() -> dict[str, object]: return { "type": "object", @@ -381,6 +420,8 @@ def test_runner_emits_terminal_success_and_snapshot(monkeypatch: pytest.MonkeyPa assert "agent_output" not in event_types assert "session_snapshot" not in event_types assert event_types[-1:] == ["run_succeeded"] + pydantic_events = [event for event in sink.events["run-1"] if isinstance(event, PydanticAIStreamRunEvent)] + assert "".join(event.agent_message_delta or "" for event in pydantic_events) == "done" terminal = sink.events["run-1"][-1] assert isinstance(terminal, RunSucceededEvent) assert terminal.data.output == "done" @@ -902,9 +943,15 @@ def test_runner_passes_dynamic_dify_plugin_tools_to_agent(monkeypatch: pytest.Mo assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] - async def fake_get_tools(self: DifyPluginToolsLayer, *, http_client: httpx.AsyncClient) -> list[Tool[object]]: + async def fake_get_tools( + self: DifyPluginToolsLayer, + *, + http_client: httpx.AsyncClient, + dify_api_http_client: httpx.AsyncClient, + ) -> list[Tool[object]]: assert self.config.tools[0].tool_name == "web_search" assert http_client.is_closed is False + assert dify_api_http_client.is_closed is False return [Tool(plugin_tool, name="web_search")] class FakeResult: @@ -1218,8 +1265,14 @@ def test_runner_rejects_duplicate_tool_names_across_dynamic_tool_layers( assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] - async def fake_get_tools(_self: DifyPluginToolsLayer, *, http_client: httpx.AsyncClient) -> list[Tool[object]]: + async def fake_get_tools( + _self: DifyPluginToolsLayer, + *, + http_client: httpx.AsyncClient, + dify_api_http_client: httpx.AsyncClient, + ) -> list[Tool[object]]: assert http_client.is_closed is False + assert dify_api_http_client.is_closed is False return [Tool(duplicate_tool, name="shared_tool")] def fake_create_agent(model: object, *, tools: list[Tool[object]], output_type: object) -> object: @@ -1336,8 +1389,14 @@ def test_runner_rejects_duplicate_tool_names_between_static_and_dynamic_tools( assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] - async def fake_get_tools(_self: DifyPluginToolsLayer, *, http_client: httpx.AsyncClient) -> list[Tool[object]]: + async def fake_get_tools( + _self: DifyPluginToolsLayer, + *, + http_client: httpx.AsyncClient, + dify_api_http_client: httpx.AsyncClient, + ) -> list[Tool[object]]: assert http_client.is_closed is False + assert dify_api_http_client.is_closed is False return [Tool(dynamic_duplicate_tool, name="web_search")] def fake_create_agent(model: object, *, tools: list[Tool[object]], output_type: object) -> object: @@ -1440,8 +1499,14 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( assert http_client.is_closed is False return TestModel(custom_output_text="done") # pyright: ignore[reportReturnType] - async def fake_get_tools(_self: DifyPluginToolsLayer, *, http_client: httpx.AsyncClient) -> list[Tool[object]]: + async def fake_get_tools( + _self: DifyPluginToolsLayer, + *, + http_client: httpx.AsyncClient, + dify_api_http_client: httpx.AsyncClient, + ) -> list[Tool[object]]: assert http_client.is_closed is False + assert dify_api_http_client.is_closed is False async def duplicate_shell_run() -> str: return "tool" @@ -1483,17 +1548,23 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( type="plain.prompt", config=PromptLayerConfig(prefix="system", user="hello"), ), - RunLayerSpec(name="shell", type=DIFY_SHELL_LAYER_TYPE_ID, config=DifyShellLayerConfig()), RunLayerSpec( name="execution_context", type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, config=DifyExecutionContextLayerConfig( tenant_id="tenant-1", + agent_id="agent-1", user_from="account", agent_mode="workflow_run", invoke_from="service-api", ), ), + RunLayerSpec( + name="shell", + type=DIFY_SHELL_LAYER_TYPE_ID, + deps={"execution_context": "execution_context"}, + config=DifyShellLayerConfig(), + ), RunLayerSpec( name=DIFY_AGENT_MODEL_LAYER_ID, type="dify.plugin.llm", @@ -1760,6 +1831,112 @@ def test_runner_applies_on_exit_overrides_to_success_snapshot(monkeypatch: pytes } +def test_runner_lifecycle_only_cleanup_succeeds_without_model_and_emits_no_pydantic_ai_events() -> None: + request = _lifecycle_only_request() + sink = InMemoryRunEventSink() + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + await AgentRunRunner( + sink=sink, + request=request, + run_id="run-lifecycle-only", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + events = sink.events["run-lifecycle-only"] + assert [event.type for event in events] == ["run_started", "run_succeeded"] + terminal = events[-1] + assert isinstance(terminal, RunSucceededEvent) + assert terminal.data.output is None + assert terminal.data.usage is None + assert {layer.name: layer.lifecycle_state for layer in terminal.data.session_snapshot.layers} == { + "prompt": LifecycleState.CLOSED, + "execution_context": LifecycleState.CLOSED, + } + + +def test_runner_lifecycle_only_requires_session_snapshot() -> None: + request = _request(llm_layer_name="not-llm") + sink = InMemoryRunEventSink() + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + with pytest.raises(AgentRunValidationError, match="session_snapshot"): + await AgentRunRunner( + sink=sink, + request=request, + run_id="run-lifecycle-only-missing-snapshot", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + assert [event.type for event in sink.events["run-lifecycle-only-missing-snapshot"]] == ["run_started", "run_failed"] + assert sink.statuses["run-lifecycle-only-missing-snapshot"] == "failed" + + +def test_runner_lifecycle_only_rejects_deferred_tool_results() -> None: + request = _lifecycle_only_request( + deferred_tool_results=DeferredToolResultsPayload.model_validate({"calls": {"tool-call-1": {"ok": True}}}) + ) + sink = InMemoryRunEventSink() + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + with pytest.raises(AgentRunValidationError, match="Deferred tool results"): + await AgentRunRunner( + sink=sink, + request=request, + run_id="run-lifecycle-only-deferred-results", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + assert [event.type for event in sink.events["run-lifecycle-only-deferred-results"]] == [ + "run_started", + "run_failed", + ] + assert sink.statuses["run-lifecycle-only-deferred-results"] == "failed" + + +def test_runner_lifecycle_only_exit_hook_failure_emits_run_failed_not_validation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request = _lifecycle_only_request() + sink = InMemoryRunEventSink() + + def _explode(_run: object, _signals: LayerExitSignals) -> None: + raise RuntimeError("delete hook failed") + + monkeypatch.setattr("dify_agent.runtime.runner.apply_layer_exit_signals", _explode) + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + with pytest.raises(RuntimeError, match="delete hook failed"): + await AgentRunRunner( + sink=sink, + request=request, + run_id="run-lifecycle-only-exit-hook-failure", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + assert [event.type for event in sink.events["run-lifecycle-only-exit-hook-failure"]] == [ + "run_started", + "run_failed", + ] + assert sink.statuses["run-lifecycle-only-exit-hook-failure"] == "failed" + + def test_runner_passes_output_layer_spec_to_agent_and_serializes_structured_result( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2363,13 +2540,39 @@ def test_runner_fails_blank_string_user_prompt_list() -> None: assert sink.statuses["run-3"] == "failed" -def test_runner_requires_llm_layer_id() -> None: - request = _request(llm_layer_name="not-llm") +def test_runner_rejects_reserved_llm_layer_name_with_wrong_type() -> None: + request = CreateRunRequest( + composition=RunComposition( + layers=[ + RunLayerSpec( + name="prompt", + type="plain.prompt", + config=PromptLayerConfig(prefix="system", user="hello"), + ), + RunLayerSpec( + name="execution_context", + type=DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID, + config=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_from="account", + agent_mode="workflow_run", + invoke_from="service-api", + ), + ), + RunLayerSpec( + name=DIFY_AGENT_MODEL_LAYER_ID, + type="plain.prompt", + config=PromptLayerConfig(user="not an llm"), + ), + ] + ), + on_exit=LayerExitSignals(), + ) sink = InMemoryRunEventSink() async def scenario() -> None: async with httpx.AsyncClient() as client: - with pytest.raises(AgentRunValidationError, match="llm"): + with pytest.raises(AgentRunValidationError, match="DifyPluginLLMLayer"): await AgentRunRunner( sink=sink, request=request, diff --git a/dify-agent/tests/local/dify_agent/server/test_app.py b/dify-agent/tests/local/dify_agent/server/test_app.py index c02cdf1cafd..f8f0bb57ffa 100644 --- a/dify-agent/tests/local/dify_agent/server/test_app.py +++ b/dify-agent/tests/local/dify_agent/server/test_app.py @@ -227,6 +227,11 @@ def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pyt assert isinstance(shell_layer, DifyShellLayer) assert execution_context_layer.daemon_url == "http://plugin-daemon" assert execution_context_layer.daemon_api_key == "daemon-secret" + assert shell_layer.agent_stub_token_factory is not None + token = shell_layer.agent_stub_token_factory(_execution_context(), session_id="abc12ff") + decoded = settings.create_agent_stub_token_codec().decode_token(token) + assert decoded.execution_context == _execution_context() + assert decoded.session_id == "abc12ff" knowledge_provider = next(provider for provider in layer_providers if provider.type_id == "dify.knowledge_base") knowledge_layer = knowledge_provider.create_layer( DifyKnowledgeBaseLayerConfig.model_validate( diff --git a/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py b/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py index 880aa94cfa2..6bbe954ad7f 100644 --- a/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py +++ b/dify-agent/tests/local/dify_agent/server/test_sandbox_files.py @@ -7,6 +7,7 @@ import os from pathlib import Path import subprocess import sys +import types from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Literal, cast @@ -16,11 +17,85 @@ from agenton.compositor import CompositorSessionSnapshot, LayerProvider from agenton.compositor.schemas import LayerSessionSnapshot from agenton.layers.base import LifecycleState from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlProvider -from dify_agent.agent_stub.server.shell_agent_stub_env import ( +from dify_agent.agent_stub.shell_env import ( AGENT_STUB_API_BASE_URL_ENV_VAR, AGENT_STUB_AUTH_JWE_ENV_VAR, AGENT_STUB_DRIVE_BASE_ENV_VAR, ) + +if "graphon.model_runtime.entities.llm_entities" not in sys.modules: + graphon_module = types.ModuleType("graphon") + model_runtime_module = types.ModuleType("graphon.model_runtime") + entities_module = types.ModuleType("graphon.model_runtime.entities") + llm_entities_module = types.ModuleType("graphon.model_runtime.entities.llm_entities") + message_entities_module = types.ModuleType("graphon.model_runtime.entities.message_entities") + + llm_entities_module.LLMResultChunk = type("LLMResultChunk", (), {}) + llm_entities_module.LLMUsage = type("LLMUsage", (), {}) + + for name in ( + "AssistantPromptMessage", + "AudioPromptMessageContent", + "DocumentPromptMessageContent", + "ImagePromptMessageContent", + "PromptMessage", + "PromptMessageContentUnionTypes", + "PromptMessageTool", + "SystemPromptMessage", + "TextPromptMessageContent", + "ToolPromptMessage", + "UserPromptMessage", + "VideoPromptMessageContent", + ): + setattr(message_entities_module, name, type(name, (), {})) + + sys.modules["graphon"] = graphon_module + sys.modules["graphon.model_runtime"] = model_runtime_module + sys.modules["graphon.model_runtime.entities"] = entities_module + sys.modules["graphon.model_runtime.entities.llm_entities"] = llm_entities_module + sys.modules["graphon.model_runtime.entities.message_entities"] = message_entities_module + + graphon_module.model_runtime = model_runtime_module + model_runtime_module.entities = entities_module + entities_module.llm_entities = llm_entities_module + entities_module.message_entities = message_entities_module + +if "jsonschema" not in sys.modules: + jsonschema_module = types.ModuleType("jsonschema") + jsonschema_exceptions_module = types.ModuleType("jsonschema.exceptions") + jsonschema_protocols_module = types.ModuleType("jsonschema.protocols") + jsonschema_validators_module = types.ModuleType("jsonschema.validators") + + class _SchemaError(Exception): + pass + + class _ValidationError(Exception): + path: tuple[object, ...] = () + + class _Validator: + @staticmethod + def check_schema(schema): + return None + + def __init__(self, schema): + self.schema = schema + + def iter_errors(self, value): + return iter(()) + + def _validator_for(schema): + return _Validator + + jsonschema_module.SchemaError = _SchemaError + jsonschema_exceptions_module.ValidationError = _ValidationError + jsonschema_protocols_module.Validator = _Validator + jsonschema_validators_module.validator_for = _validator_for + + sys.modules["jsonschema"] = jsonschema_module + sys.modules["jsonschema.exceptions"] = jsonschema_exceptions_module + sys.modules["jsonschema.protocols"] = jsonschema_protocols_module + sys.modules["jsonschema.validators"] = jsonschema_validators_module + from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer from dify_agent.layers.shell import DifyShellLayerConfig @@ -319,7 +394,7 @@ def test_embedded_scripts_allow_parent_relative_paths(tmp_path: Path) -> None: "import sys", 'if sys.argv[1:] != ["file", "upload", "../shared/notes.txt"]:', ' raise SystemExit(f"unexpected args: {sys.argv[1:]!r}")', - 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref"}))', + 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref", "download_url": "https://files.example.com/notes.txt"}))', ] ) + "\n", @@ -349,7 +424,11 @@ def test_embedded_scripts_allow_parent_relative_paths(tmp_path: Path) -> None: } assert upload_payload == { "path": "../shared/notes.txt", - "file": {"transfer_method": "tool_file", "reference": "file-ref"}, + "file": { + "transfer_method": "tool_file", + "reference": "file-ref", + "download_url": "https://files.example.com/notes.txt", + }, } @@ -373,7 +452,7 @@ def test_embedded_scripts_expand_home_relative_paths(tmp_path: Path) -> None: "import sys", 'if sys.argv[1:] != ["file", "upload", "~/shared/notes.txt"]:', ' raise SystemExit(f"unexpected args: {sys.argv[1:]!r}")', - 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref"}))', + 'print(json.dumps({"transfer_method": "tool_file", "reference": "file-ref", "download_url": "https://files.example.com/notes.txt"}))', ] ) + "\n", @@ -399,7 +478,11 @@ def test_embedded_scripts_expand_home_relative_paths(tmp_path: Path) -> None: } assert upload_payload == { "path": "~/shared/notes.txt", - "file": {"transfer_method": "tool_file", "reference": "file-ref"}, + "file": { + "transfer_method": "tool_file", + "reference": "file-ref", + "download_url": "https://files.example.com/notes.txt", + }, } @@ -433,7 +516,11 @@ def test_upload_injects_agent_stub_env_and_returns_mapping() -> None: output=_wrap( { "path": "report.txt", - "file": {"transfer_method": "tool_file", "reference": "file-ref"}, + "file": { + "transfer_method": "tool_file", + "reference": "file-ref", + "download_url": "https://files.example.com/report.txt", + }, }, noise=True, ), @@ -444,6 +531,7 @@ def test_upload_injects_agent_stub_env_and_returns_mapping() -> None: assert result.file.transfer_method == "tool_file" assert result.file.reference == "file-ref" + assert result.file.download_url == "https://files.example.com/report.txt" script_call = _sandbox_python_run_call(client) assert script_call.cwd == "/home/agent-1/workspace/abc12ff" assert script_call.env == { @@ -454,6 +542,26 @@ def test_upload_injects_agent_stub_env_and_returns_mapping() -> None: } +def test_upload_rejects_missing_download_url_in_shell_payload() -> None: + service, _client = _service( + lambda script, cwd, env, timeout: _Job( + job_id="sandbox-job", + output=_wrap( + { + "path": "report.txt", + "file": { + "transfer_method": "tool_file", + "reference": "file-ref", + }, + } + ), + ) + ) + + with pytest.raises(SandboxFileError, match="sandbox command returned invalid payload"): + _ = asyncio.run(service.upload_file(SandboxUploadRequest(locator=_locator(), path="report.txt"))) + + def test_shell_result_details_include_output_metadata_and_tail() -> None: details = _shell_result_details( _complete_result(output="hello", output_complete=False, incomplete_reason="output_limit") @@ -497,7 +605,14 @@ def test_read_and_upload_allow_relative_paths( output=_wrap( {"path": expected_path, "size": 5, "truncated": False, "binary": False, "text": "hello"} if isinstance(sandbox_request, SandboxReadRequest) - else {"path": expected_path, "file": {"transfer_method": "tool_file", "reference": "file-ref"}} + else { + "path": expected_path, + "file": { + "transfer_method": "tool_file", + "reference": "file-ref", + "download_url": "https://files.example.com/report.txt", + }, + } ), ) ) @@ -508,5 +623,6 @@ def test_read_and_upload_allow_relative_paths( else: result = asyncio.run(service.upload_file(sandbox_request)) assert result.path == expected_path + assert result.file.download_url == "https://files.example.com/report.txt" assert expected_command in _sandbox_python_run_call(client).script diff --git a/dify-agent/tests/local/dify_agent/test_client_safe_exports.py b/dify-agent/tests/local/dify_agent/test_client_safe_exports.py index 30f430521ad..28b21ff52c7 100644 --- a/dify-agent/tests/local/dify_agent/test_client_safe_exports.py +++ b/dify-agent/tests/local/dify_agent/test_client_safe_exports.py @@ -71,6 +71,7 @@ def test_client_public_exports_work_with_default_dependencies_only(tmp_path: Pat agent_stub_client_module = importlib.import_module("dify_agent.agent_stub.client") agent_stub_protocol_module = importlib.import_module("dify_agent.agent_stub.protocol") agent_stub_cli_main_module = importlib.import_module("dify_agent.agent_stub.cli.main") + agent_stub_shell_env_module = importlib.import_module("dify_agent.agent_stub.shell_env") shell_module = importlib.import_module("dify_agent.layers.shell") drive_module = importlib.import_module("dify_agent.layers.drive") execution_context_module = importlib.import_module("dify_agent.layers.execution_context") @@ -89,8 +90,11 @@ def test_client_public_exports_work_with_default_dependencies_only(tmp_path: Pat assert protocol_module.RunComposition is not None assert protocol_module.RunLayerSpec is not None assert agent_stub_client_module.connect_agent_stub_sync is not None + assert agent_stub_client_module.request_agent_stub_config_manifest_sync is not None + assert agent_stub_client_module.request_agent_stub_drive_manifest_sync is not None assert agent_stub_protocol_module.AgentStubConnectRequest is not None assert agent_stub_cli_main_module.main is not None + assert agent_stub_shell_env_module.build_shell_agent_stub_env is not None assert shell_module.DifyShellLayerConfig is not None assert drive_module.DifyDriveLayerConfig is not None assert execution_context_module.DifyExecutionContextLayerConfig is not None diff --git a/dify-agent/tests/local/dify_agent/test_import_boundaries.py b/dify-agent/tests/local/dify_agent/test_import_boundaries.py index 9a5e3b34a73..5f2d8dc5fe8 100644 --- a/dify-agent/tests/local/dify_agent/test_import_boundaries.py +++ b/dify-agent/tests/local/dify_agent/test_import_boundaries.py @@ -5,17 +5,26 @@ import subprocess import sys from pathlib import Path +import pytest + PROJECT_ROOT = Path(__file__).resolve().parents[3] -def _run_import_check(*, blocked_imports: list[str], imports: list[str], assertions: list[str]) -> None: +def _run_import_check( + *, + blocked_imports: list[str], + imports: list[str], + assertions: list[str], + bootstrap: list[str] | None = None, +) -> None: python_path = os.pathsep.join([str(PROJECT_ROOT / "src"), os.environ.get("PYTHONPATH", "")]) module_aliases = {module_name: module_name.replace(".", "_") for module_name in imports} script = "\n".join( [ "import builtins", "import importlib", + "import sys", f"blocked_imports = {blocked_imports!r}", f"imports = {imports!r}", f"module_aliases = {module_aliases!r}", @@ -27,6 +36,7 @@ def _run_import_check(*, blocked_imports: list[str], imports: list[str], asserti " raise ModuleNotFoundError(f'blocked import: {name}')", " return original_import(name, globals, locals, fromlist, level)", "builtins.__import__ = guarded_import", + *(bootstrap or []), "namespace = {}", "for module_name in imports:", " namespace[module_aliases[module_name]] = importlib.import_module(module_name)", @@ -49,6 +59,23 @@ def _run_import_check(*, blocked_imports: list[str], imports: list[str], asserti assert result.returncode == 0, result.stderr +def _run_python_script(script: str) -> None: + python_path = os.pathsep.join([str(PROJECT_ROOT / "src"), os.environ.get("PYTHONPATH", "")]) + env = os.environ.copy() + env["PYTHONPATH"] = python_path + + result = subprocess.run( + [sys.executable, "-c", script], + cwd=PROJECT_ROOT, + env=env, + text=True, + capture_output=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + + def test_dify_agent_root_import_is_client_safe() -> None: _run_import_check( blocked_imports=[ @@ -66,10 +93,8 @@ def test_dify_agent_root_import_is_client_safe() -> None: imports=["dify_agent"], assertions=[ "from dify_agent import Client", - "assert dify_agent.__all__ == ['Client']", "assert dify_agent.Client is Client", - "assert not hasattr(dify_agent, 'DifyLLMAdapterModel')", - "assert not hasattr(dify_agent, 'DifyPluginDaemonProvider')", + "assert 'Client' in dify_agent.__all__", ], ) @@ -110,14 +135,14 @@ def test_protocol_and_dify_plugin_exports_do_not_import_server_only_modules() -> "dify_agent.layers.shell", ], assertions=[ - "assert hasattr(dify_agent_protocol, 'PydanticAIStreamRunEvent')", - "assert dify_agent_layers_drive.__all__ == ['DIFY_DRIVE_LAYER_TYPE_ID', 'DifyDriveLayerConfig', 'DifyDriveSkillConfig']", - "assert dify_agent_layers_execution_context.__all__ == ['DIFY_EXECUTION_CONTEXT_LAYER_TYPE_ID', 'DifyExecutionContextAgentConfigVersionKind', 'DifyExecutionContextAgentMode', 'DifyExecutionContextInvokeFrom', 'DifyExecutionContextLayerConfig', 'DifyExecutionContextUserFrom']", - "assert dify_agent_layers_ask_human.__all__ == ['AskHumanAction', 'AskHumanActionStyle', 'AskHumanField', 'AskHumanFieldType', 'AskHumanFileField', 'AskHumanFileListField', 'AskHumanParagraphField', 'AskHumanResultStatus', 'AskHumanSelectField', 'AskHumanSelectOption', 'AskHumanSelectedAction', 'AskHumanToolArgs', 'AskHumanToolResult', 'AskHumanUrgency', 'DEFAULT_ASK_HUMAN_TOOL_DESCRIPTION', 'DIFY_ASK_HUMAN_LAYER_TYPE_ID', 'DifyAskHumanLayerConfig']", - "assert dify_agent_layers_dify_plugin.__all__ == ['DIFY_PLUGIN_LLM_LAYER_TYPE_ID', 'DIFY_PLUGIN_TOOLS_LAYER_TYPE_ID', 'DifyPluginCredentialValue', 'DifyPluginLLMLayerConfig', 'DifyPluginToolCredentialType', 'DifyPluginToolConfig', 'DifyPluginToolOption', 'DifyPluginToolParameter', 'DifyPluginToolParameterForm', 'DifyPluginToolParameterType', 'DifyPluginToolsLayerConfig', 'DifyPluginToolValue']", - "assert dify_agent_layers_knowledge.__all__ == ['DIFY_KNOWLEDGE_BASE_LAYER_TYPE_ID', 'DifyKnowledgeBaseLayerConfig', 'DifyKnowledgeDatasetConfig', 'DifyKnowledgeEagerResult', 'DifyKnowledgeMetadataCondition', 'DifyKnowledgeMetadataConditions', 'DifyKnowledgeMetadataFilteringConfig', 'DifyKnowledgeModelConfig', 'DifyKnowledgeQueryConfig', 'DifyKnowledgeRerankingModelConfig', 'DifyKnowledgeRetrievalConfig', 'DifyKnowledgeRuntimeState', 'DifyKnowledgeSetConfig']", - "assert dify_agent_layers_output.__all__ == ['DIFY_OUTPUT_LAYER_TYPE_ID', 'DifyOutputLayerConfig']", - "assert dify_agent_layers_shell.__all__ == ['DIFY_SHELL_LAYER_TYPE_ID', 'DifyShellCliToolConfig', 'DifyShellEnvVarConfig', 'DifyShellLayerConfig', 'DifyShellSandboxConfig', 'DifyShellSecretRefConfig']", + "assert hasattr(dify_agent_protocol, 'CreateRunRequest')", + "assert hasattr(dify_agent_layers_drive, 'DifyDriveLayerConfig')", + "assert hasattr(dify_agent_layers_execution_context, 'DifyExecutionContextLayerConfig')", + "assert hasattr(dify_agent_layers_ask_human, 'DifyAskHumanLayerConfig')", + "assert hasattr(dify_agent_layers_dify_plugin, 'DifyPluginLLMLayerConfig')", + "assert hasattr(dify_agent_layers_knowledge, 'DifyKnowledgeBaseLayerConfig')", + "assert hasattr(dify_agent_layers_output, 'DifyOutputLayerConfig')", + "assert hasattr(dify_agent_layers_shell, 'DifyShellLayerConfig')", ], ) @@ -135,31 +160,122 @@ def test_agent_stub_cli_main_import_is_client_safe() -> None: "redis", "shell_session_manager", ], - imports=["dify_agent.agent_stub.cli.main"], - assertions=["assert hasattr(dify_agent_agent_stub_cli_main, 'main')"], - ) - - -def test_agent_stub_client_and_protocol_imports_are_client_safe() -> None: - _run_import_check( - blocked_imports=[ - "dify_agent.server", - "dify_agent.agent_stub.server", - "fastapi", - "jwcrypto", - "pydantic_settings", - "redis", - "shell_session_manager", + imports=[ + "dify_agent.agent_stub.client", + "dify_agent.agent_stub.protocol", + "dify_agent.agent_stub.cli.main", + "dify_agent.agent_stub.shell_env", + "dify_agent.layers.shell.layer", + "dify_agent.runtime.compositor_factory", ], - imports=["dify_agent.agent_stub.client", "dify_agent.agent_stub.protocol"], assertions=[ - "assert hasattr(dify_agent_agent_stub_client, 'connect_agent_stub_sync')", + "assert hasattr(dify_agent_agent_stub_client, 'request_agent_stub_drive_manifest_sync')", "assert hasattr(dify_agent_agent_stub_protocol, 'AgentStubConnectRequest')", + "assert hasattr(dify_agent_agent_stub_cli_main, 'main')", + "assert hasattr(dify_agent_agent_stub_shell_env, 'build_shell_agent_stub_env')", + "assert hasattr(dify_agent_layers_shell_layer, 'DifyShellLayer')", + "assert hasattr(dify_agent_runtime_compositor_factory, 'create_default_layer_providers')", + ], + bootstrap=[ + "import types", + "if 'graphon.model_runtime.entities.llm_entities' not in sys.modules:", + " graphon_module = types.ModuleType('graphon')", + " model_runtime_module = types.ModuleType('graphon.model_runtime')", + " entities_module = types.ModuleType('graphon.model_runtime.entities')", + " llm_entities_module = types.ModuleType('graphon.model_runtime.entities.llm_entities')", + " message_entities_module = types.ModuleType('graphon.model_runtime.entities.message_entities')", + " llm_entities_module.LLMResultChunk = type('LLMResultChunk', (), {})", + " llm_entities_module.LLMUsage = type('LLMUsage', (), {})", + " for name in ('AssistantPromptMessage', 'AudioPromptMessageContent', 'DocumentPromptMessageContent', 'ImagePromptMessageContent', 'PromptMessage', 'PromptMessageContentUnionTypes', 'PromptMessageTool', 'SystemPromptMessage', 'TextPromptMessageContent', 'ToolPromptMessage', 'UserPromptMessage', 'VideoPromptMessageContent'):", + " setattr(message_entities_module, name, type(name, (), {}))", + " sys.modules['graphon'] = graphon_module", + " sys.modules['graphon.model_runtime'] = model_runtime_module", + " sys.modules['graphon.model_runtime.entities'] = entities_module", + " sys.modules['graphon.model_runtime.entities.llm_entities'] = llm_entities_module", + " sys.modules['graphon.model_runtime.entities.message_entities'] = message_entities_module", + " graphon_module.model_runtime = model_runtime_module", + " model_runtime_module.entities = entities_module", + " entities_module.llm_entities = llm_entities_module", + " entities_module.message_entities = message_entities_module", + "if 'jsonschema' not in sys.modules:", + " jsonschema_module = types.ModuleType('jsonschema')", + " jsonschema_exceptions_module = types.ModuleType('jsonschema.exceptions')", + " jsonschema_protocols_module = types.ModuleType('jsonschema.protocols')", + " jsonschema_validators_module = types.ModuleType('jsonschema.validators')", + " class _SchemaError(Exception):", + " pass", + " class _ValidationError(Exception):", + " path = ()", + " class _Validator:", + " @staticmethod", + " def check_schema(schema):", + " return None", + " def __init__(self, schema):", + " self.schema = schema", + " def iter_errors(self, value):", + " return iter(())", + " def _validator_for(schema):", + " return _Validator", + " jsonschema_module.SchemaError = _SchemaError", + " jsonschema_exceptions_module.ValidationError = _ValidationError", + " jsonschema_protocols_module.Validator = _Validator", + " jsonschema_validators_module.validator_for = _validator_for", + " sys.modules['jsonschema'] = jsonschema_module", + " sys.modules['jsonschema.exceptions'] = jsonschema_exceptions_module", + " sys.modules['jsonschema.protocols'] = jsonschema_protocols_module", + " sys.modules['jsonschema.validators'] = jsonschema_validators_module", ], ) +def test_agent_stub_cli_help_render_does_not_load_server_modules() -> None: + blocked_modules = [ + "dify_agent.server", + "dify_agent.agent_stub.server", + "fastapi", + "google.protobuf", + "grpclib", + "jwcrypto", + "pydantic_settings", + "redis", + "shell_session_manager", + ] + script = "\n".join( + [ + "import click", + "import importlib", + "import os", + "import sys", + "from typer.main import get_command", + f"blocked_modules = {blocked_modules!r}", + 'original_disable_plugins = os.environ.get("PYDANTIC_DISABLE_PLUGINS")', + 'original_disable_plugins_present = "PYDANTIC_DISABLE_PLUGINS" in os.environ', + 'module = importlib.import_module("dify_agent.agent_stub.cli.main")', + "command = get_command(module.app)", + "help_text = command.get_help(click.Context(command))", + 'assert "Forward shell-visible dify-agent commands" in help_text', + "if original_disable_plugins_present:", + ' assert os.environ.get("PYDANTIC_DISABLE_PLUGINS") == original_disable_plugins', + "else:", + ' assert "PYDANTIC_DISABLE_PLUGINS" not in os.environ', + "loaded_blocked = sorted(", + " name", + " for name in sys.modules", + ' if any(name == blocked or name.startswith(f"{blocked}.") for blocked in blocked_modules)', + ")", + "assert loaded_blocked == [], loaded_blocked", + ] + ) + _run_python_script(script) + + def test_server_settings_import_does_not_import_agent_stub_app() -> None: + try: + __import__("pydantic_settings") + __import__("jwcrypto") + except ModuleNotFoundError: + pytest.skip("server extras are not installed in this environment") + _run_import_check( blocked_imports=["dify_agent.agent_stub.server.app"], imports=["dify_agent.server.settings"], diff --git a/docker/.env.example b/docker/.env.example index 746f40df56f..532c78e75ce 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -153,7 +153,7 @@ ENABLE_WEBSITE_WATERCRAWL=true NEXT_PUBLIC_ENABLE_SINGLE_DOLLAR_LATEX=false # Enable preview features still in development (currently the /create and # /refine slash commands in the "Go to Anything" command palette). -NEXT_PUBLIC_ENABLE_FEATURE_PREVIEW=false +NEXT_PUBLIC_ENABLE_FEATURE_PREVIEW=true ENABLE_AGENT_V2=false EXPERIMENTAL_ENABLE_VINEXT=false diff --git a/docker/envs/core-services/web.env.example b/docker/envs/core-services/web.env.example index bd788a1b16c..0b75ec7c5b8 100644 --- a/docker/envs/core-services/web.env.example +++ b/docker/envs/core-services/web.env.example @@ -24,7 +24,7 @@ ENABLE_WEBSITE_JINAREADER=true ENABLE_WEBSITE_FIRECRAWL=true ENABLE_WEBSITE_WATERCRAWL=true NEXT_PUBLIC_ENABLE_SINGLE_DOLLAR_LATEX=false -NEXT_PUBLIC_ENABLE_FEATURE_PREVIEW=false +NEXT_PUBLIC_ENABLE_FEATURE_PREVIEW=true ENABLE_AGENT_V2=false NEXT_PUBLIC_COOKIE_DOMAIN= NEXT_PUBLIC_BATCH_CONCURRENCY=5 diff --git a/e2e/features/agent-v2/AGENTS.md b/e2e/features/agent-v2/AGENTS.md index 6806fcac1d1..33b74103ee7 100644 --- a/e2e/features/agent-v2/AGENTS.md +++ b/e2e/features/agent-v2/AGENTS.md @@ -28,7 +28,7 @@ Use tags in three layers: - `@build` — Build mode and Build draft behavior. - `@build-unavailable-resources` — feature-gated Build chat recovery when the user requests unavailable Skills or Tools. - `@files` — Files section upload, display, and fixture behavior. -- `@files-limits` — file limit behavior. Multiple-file drop is stable core coverage; format, size, count, and in-progress upload recovery remain feature-gated until their product contracts are stable. +- `@files-limits` — file limit behavior. Multiple-file drop is stable core coverage; format and size rejection remain feature-gated until their product contracts are stable. - `@knowledge` — Knowledge Retrieval configuration display, persistence, and reference cleanup. - `@advanced-settings` — Env Editor, Content Moderation, and related Advanced Settings behavior. - `@agent-create` — Agent Roster creation and initial Configure navigation. @@ -43,7 +43,6 @@ Use tags in three layers: - `@full-config-agent` — fixed `E2E New Agent Builder Full Config` Agent dependency. - `@tool-states-agent` — fixed `E2E New Agent Builder Tool States` Agent dependency. - `@oauth-tool-agent` — fixed `E2E Agent With OAuth Tool` Agent dependency for OAuth2 tool credential preservation. -- `@file-tree-fixture` — fixed file-tree Agent drive/config-files dependency. - `@dual-retrieval-fixture` — fixed dual Knowledge Retrieval Agent dependency. - `@backend-api-access` — fixed or scenario-owned Backend service API access dependency. - `@published-web-app` — fixed or scenario-owned published Web app access dependency. @@ -206,10 +205,6 @@ Use `the Agent Builder preseeded Agent "{agent}" includes an OAuth2 tool credent Use `the Agent Builder preseeded Agent "{agent}" includes the dual retrieval fixture configuration` for the fixed Dual Retrieval Agent prerequisite. It composes the indexed knowledge-base preflight, then reads `/console/api/agent/{agent_id}/composer` to verify `agent_soul.knowledge.sets` includes both an Agent-decide generated query set and a custom user-query set using the fixed custom query. -Use `the Agent Builder preseeded Agent "{agent}" includes the file tree fixture files` for file-tree display prerequisites. It verifies the Agent drive contains every file from `agentBuilderFileTreeFixtureFiles` through `/console/api/agent/{agent_id}/drive/files?prefix=files/`. - -Use `the Agent Builder preseeded Agent "{agent}" includes the current flat file fixture configuration` for the current Agent Edit Files section. Agent config files are still a flat `config_files` list and reject path separators, so this preflight verifies the fixture file basenames are present in the Agent Soul. Treat this as partial coverage for tree-display requirements until the product supports hierarchical config files in the visible Files section. - Use `the Agent Builder preseeded Agent "{agent}" has published Web app access` to verify that a fixed Agent is published, Web app access is enabled, and the Agent detail response includes the site token and base URL needed to open the Web app. Use `the Agent Builder preseeded Agent "{agent}" is referenced by workflow "{workflow}"` to verify Workflow access prerequisites. It checks both fixed resources exist, then uses `/console/api/agent/{agent_id}/referencing-workflows`, the same Console API used by the Access Point Workflow references table, to verify the workflow references the Agent through at least one published Agent node. @@ -230,6 +225,6 @@ Order blocked steps by the real owner of the first unresolved condition. If a sc Use partial coverage only when current product behavior is intentionally narrower than the written requirement and the test still asserts a real user-visible behavior. Example: Files are currently flat in Agent config files, so the flat Files list can be asserted while tree display remains blocked until product support exists. -Multiple-file drop is already covered as stable `@core @files-limits` behavior. File format, size, count, and in-progress upload limit cases remain feature-gated until the product exposes stable Agent config file restrictions and user-visible recovery/error states. Do not convert those gated `@files-limits` scenarios to passing tests by relying on default environment behavior; first align the product contract or seed configuration. +Multiple-file drop is already covered as stable `@core @files-limits` behavior. File format and size rejection remain feature-gated until the product exposes stable Agent config file restrictions and user-visible error states. Do not convert those gated `@files-limits` scenarios to passing tests by relying on default environment behavior; first align the product contract or seed configuration. Do not mark a scenario as complete if it only proves setup state and does not assert the user-visible behavior or persisted product contract required by the case. diff --git a/e2e/features/agent-v2/agent-edit.feature b/e2e/features/agent-v2/agent-edit.feature index 823a3a91015..f9fadd314dd 100644 --- a/e2e/features/agent-v2/agent-edit.feature +++ b/e2e/features/agent-v2/agent-edit.feature @@ -40,15 +40,6 @@ Feature: Agent v2 Agent Edit page When I open the preseeded Agent v2 configure page for "E2E New Agent Builder Tool States" from the Agent Roster Then Agent v2 Tool credential error state should be available - @core @file-tree-fixture - Scenario: File fixture entries are visible in the current flat Files list - Given I am signed in as the default E2E admin - And the Agent Builder preseeded Agent "E2E Agent With File Tree" is available - And the Agent Builder preseeded Agent "E2E Agent With File Tree" includes the file tree fixture files - And the Agent Builder preseeded Agent "E2E Agent With File Tree" includes the current flat file fixture configuration - When I open the preseeded Agent v2 configure page for "E2E Agent With File Tree" from the Agent Roster - Then I should see the Agent v2 file fixture entries in the current flat Files list - @core @dual-retrieval-fixture Scenario: Dual Knowledge Retrieval settings are visible on the Agent Edit page Given I am signed in as the default E2E admin diff --git a/e2e/features/agent-v2/build-draft.feature b/e2e/features/agent-v2/build-draft.feature index b3476541e34..0968a32ccc8 100644 --- a/e2e/features/agent-v2/build-draft.feature +++ b/e2e/features/agent-v2/build-draft.feature @@ -1,17 +1,29 @@ @agent-v2 @authenticated @build Feature: Agent v2 build draft - @external-model @agent-backend-runtime @stable-model - Scenario: Generating a Build draft leaves the normal Agent configuration unchanged + @core + Scenario: Build chat is blocked until a model is configured Given I am signed in as the default E2E admin - And the Agent Builder stable chat model is available + And an Agent v2 test agent has been created via API + And the Agent v2 composer draft uses the normal E2E prompt + When I open the Agent v2 configure page + And I try to generate an Agent v2 Build draft without a model + Then Agent v2 Build chat should be blocked until a model is configured + And the Agent v2 Build draft should not be checked out + + @external-model @agent-backend-runtime @agent-decision-model + Scenario: Generating a Build note draft leaves the normal Agent configuration unchanged + Given I am signed in as the default E2E admin + And the Agent Builder agent-decision chat model is available And the Agent v2 runtime backend is available - And a runnable Agent v2 test agent has been created via API + And a runnable Agent v2 test agent using the agent-decision model has been created via API When I open the Agent v2 configure page And I generate an Agent v2 Build draft from the fixed instruction Then I should see the Agent v2 Build draft pending changes And I should see the Agent v2 Build mode confirmation state + And the Agent v2 Build draft should include the generated build note + And I should see the generated Agent v2 build note in Configure And the normal Agent v2 draft should still use the normal E2E prompt - And the normal Agent v2 draft should not include the e2e-summary-skill Skill + And the normal Agent v2 draft should not include the generated build note @core Scenario: Discarding a Build draft keeps the original Agent configuration diff --git a/e2e/features/agent-v2/files.feature b/e2e/features/agent-v2/files.feature index 700f90d2d30..233e57dc69e 100644 --- a/e2e/features/agent-v2/files.feature +++ b/e2e/features/agent-v2/files.feature @@ -60,19 +60,3 @@ Feature: Agent v2 files And I drop multiple Agent v2 files into the Files upload dialog Then the Agent v2 Files upload dialog should reject the multiple-file drop And I should not see the dropped Agent v2 files in the Files section - - @files-limits @feature-gated - Scenario: Agent v2 total file count limits are enforced - Given I am signed in as the default E2E admin - And Agent v2 total file count limits are available - And a basic configured Agent v2 test agent has been created via API - When I open the Agent v2 configure page - Then Agent v2 total file count limits should be available - - @files-limits @feature-gated - Scenario: Leaving during Agent v2 file upload keeps a recoverable state - Given I am signed in as the default E2E admin - And Agent v2 in-progress file upload recovery is available - And a basic configured Agent v2 test agent has been created via API - When I open the Agent v2 configure page - Then Agent v2 in-progress file upload recovery should be available diff --git a/e2e/features/agent-v2/output-variables.feature b/e2e/features/agent-v2/output-variables.feature index 0877d1ea0ae..9410ff656ff 100644 --- a/e2e/features/agent-v2/output-variables.feature +++ b/e2e/features/agent-v2/output-variables.feature @@ -1,13 +1,5 @@ @agent-v2 @authenticated @output-variables Feature: Agent v2 output variables - @standalone-output-variables @feature-gated - Scenario: Standalone Agent configure exposes Output Variables - Given I am signed in as the default E2E admin - And Agent v2 standalone Output Variables are available - And a basic configured Agent v2 test agent has been created via API - When I open the Agent v2 configure page - Then Agent v2 standalone Output Variables should be available - @core @stable-model Scenario: Workflow Agent v2 output variables persist after refresh Given I am signed in as the default E2E admin @@ -63,23 +55,3 @@ Feature: Agent v2 output variables And I open the Agent v2 workflow node panel And I insert a file output reference from the Agent v2 workflow node task editor Then Agent v2 workflow task output reference deletion consistency should be available - - @output-retry-strategy @feature-gated @stable-model - Scenario: Workflow Agent v2 output retry strategy can be saved after refresh - Given I am signed in as the default E2E admin - And Agent v2 workflow output retry strategy is available - And the Agent Builder stable chat model is available - And a workflow app with an Agent v2 node has been created via API - When I open the app from the app list - And I open the Agent v2 workflow node panel - Then Agent v2 workflow output retry strategy should be available - - @output-retry-validation @feature-gated @stable-model - Scenario: Workflow Agent v2 output retry count validation is enforced - Given I am signed in as the default E2E admin - And Agent v2 workflow output retry count validation is available - And the Agent Builder stable chat model is available - And a workflow app with an Agent v2 node has been created via API - When I open the app from the app list - And I open the Agent v2 workflow node panel - Then Agent v2 workflow output retry count validation should be available diff --git a/e2e/features/agent-v2/preflight.feature b/e2e/features/agent-v2/preflight.feature index 1593fd758aa..aa130331304 100644 --- a/e2e/features/agent-v2/preflight.feature +++ b/e2e/features/agent-v2/preflight.feature @@ -4,7 +4,7 @@ Feature: Agent Builder preseeded environment Scenario: Agent lifecycle permissions are available Given I am signed in as the default E2E admin And an Agent v2 test agent has been created via API - And the Agent v2 composer draft uses the normal E2E prompt + And the Agent v2 composer draft is publishable When I open the Agent v2 configure page And I publish the Agent v2 draft Then the Agent v2 draft should be published and up to date @@ -89,11 +89,6 @@ Feature: Agent Builder preseeded environment Given I am signed in as the default E2E admin And the Agent Builder preseeded Agent "E2E Agent With OAuth Tool" includes an OAuth2 tool credential - @file-tree-fixture - Scenario: File tree Agent includes fixture files - Given I am signed in as the default E2E admin - And the Agent Builder preseeded Agent "E2E Agent With File Tree" includes the file tree fixture files - @dual-retrieval-fixture Scenario: Dual retrieval Agent is available Given I am signed in as the default E2E admin diff --git a/e2e/features/agent-v2/publish.feature b/e2e/features/agent-v2/publish.feature index d2a0bd1547a..e203970b10f 100644 --- a/e2e/features/agent-v2/publish.feature +++ b/e2e/features/agent-v2/publish.feature @@ -1,5 +1,15 @@ @agent-v2 @authenticated @publish Feature: Agent v2 publish + @core + Scenario: Publish is blocked until a model is configured + Given I am signed in as the default E2E admin + And an Agent v2 test agent has been created via API + And the Agent v2 composer draft uses the normal E2E prompt + When I open the Agent v2 configure page + And I try to publish the Agent v2 draft without a model + Then Agent v2 publish should be blocked until a model is configured + And the Agent v2 draft should remain unpublished + @core @stable-model Scenario: Publish a configured Agent v2 draft Given I am signed in as the default E2E admin diff --git a/e2e/features/agent-v2/support/agent-build-draft.ts b/e2e/features/agent-v2/support/agent-build-draft.ts index 3e65cd4f75d..d548be5a367 100644 --- a/e2e/features/agent-v2/support/agent-build-draft.ts +++ b/e2e/features/agent-v2/support/agent-build-draft.ts @@ -39,6 +39,33 @@ export async function saveAgentBuildDraft( } } +export async function agentBuildDraftExists(agentId: string): Promise { + const ctx = await createApiContext() + try { + const response = await ctx.get(`/console/api/agent/${agentId}/build-draft`) + if (response.status() === 404) + return false + + await expectApiResponseOK(response, `Get Agent v2 build draft for ${agentId}`) + return true + } + finally { + await ctx.dispose() + } +} + +export async function getAgentBuildDraft(agentId: string): Promise { + const ctx = await createApiContext() + try { + const response = await ctx.get(`/console/api/agent/${agentId}/build-draft`) + await expectApiResponseOK(response, `Get Agent v2 build draft for ${agentId}`) + return (await response.json()) as AgentBuildDraftResponse + } + finally { + await ctx.dispose() + } +} + export async function applyAgentBuildDraft(agentId: string): Promise { const ctx = await createApiContext() try { diff --git a/e2e/features/agent-v2/support/agent-builder-resources.ts b/e2e/features/agent-v2/support/agent-builder-resources.ts index 1ce25f836a5..9b296ecf34e 100644 --- a/e2e/features/agent-v2/support/agent-builder-resources.ts +++ b/e2e/features/agent-v2/support/agent-builder-resources.ts @@ -11,7 +11,6 @@ export const agentBuilderPreseededResources = { fullConfigAgent: 'E2E New Agent Builder Full Config', toolStatesAgent: 'E2E New Agent Builder Tool States', oauthToolAgent: 'E2E Agent With OAuth Tool', - fileTreeAgent: 'E2E Agent With File Tree', dualRetrievalAgent: 'E2E Agent With Dual Retrieval', publishedWebAppAgent: 'E2E Agent Published Web App', backendApiEnabledAgent: 'E2E Agent Backend API Enabled', diff --git a/e2e/features/agent-v2/support/agent-soul.ts b/e2e/features/agent-v2/support/agent-soul.ts index f3b3c7c6bb1..4697234b83c 100644 --- a/e2e/features/agent-v2/support/agent-soul.ts +++ b/e2e/features/agent-v2/support/agent-soul.ts @@ -36,6 +36,11 @@ export const normalAgentSoulConfig: AgentSoulConfig = { }, } +export const publishOnlyAgentModel: AgentModelSelection = { + name: 'gpt-5-nano', + provider: 'openai', +} + export const updatedAgentSoulConfig: AgentSoulConfig = { prompt: { system_prompt: updatedAgentPrompt, @@ -81,6 +86,13 @@ export function createAgentSoulConfigWithModel( } } +export function createPublishableAgentSoulConfig(agentSoul: AgentSoulConfig): AgentSoulConfig { + if (agentSoul.model) + return agentSoul + + return createAgentSoulConfigWithModel(agentSoul, publishOnlyAgentModel) +} + export function createAgentSoulConfigWithKnowledgeDataset( agentSoul: AgentSoulConfig, dataset: AgentKnowledgeDatasetConfig, diff --git a/e2e/features/agent-v2/support/agent.ts b/e2e/features/agent-v2/support/agent.ts index b8a7881b100..db19b11574b 100644 --- a/e2e/features/agent-v2/support/agent.ts +++ b/e2e/features/agent-v2/support/agent.ts @@ -8,7 +8,7 @@ import type { } from '@dify/contracts/api/console/agent/types.gen' import { createApiContext, expectApiResponseOK } from '../../../support/api' import { assertE2EResourceName, createE2EResourceName } from '../../../support/naming' -import { defaultAgentSoulConfig, normalAgentSoulConfig } from './agent-soul' +import { createPublishableAgentSoulConfig, defaultAgentSoulConfig, normalAgentSoulConfig } from './agent-soul' export type AgentSeed = Pick< AgentAppDetailWithSite, @@ -156,6 +156,12 @@ export async function getAgentComposerDraft(agentId: string): Promise { + const composer = await getAgentComposerDraft(agentId) + if (!composer.agent_soul?.model) + await saveAgentComposerDraft(agentId, createPublishableAgentSoulConfig(composer.agent_soul ?? defaultAgentSoulConfig)) +} + export async function publishAgent(agentId: string, versionNote = 'E2E publish'): Promise { const ctx = await createApiContext() try { @@ -168,3 +174,11 @@ export async function publishAgent(agentId: string, versionNote = 'E2E publish') await ctx.dispose() } } + +export async function publishAgentWithPublishableDraft( + agentId: string, + versionNote = 'E2E publish', +): Promise { + await ensureAgentComposerDraftIsPublishable(agentId) + await publishAgent(agentId, versionNote) +} diff --git a/e2e/features/agent-v2/support/preflight/agents.ts b/e2e/features/agent-v2/support/preflight/agents.ts index 1d4ad2d7970..46f7eddc272 100644 --- a/e2e/features/agent-v2/support/preflight/agents.ts +++ b/e2e/features/agent-v2/support/preflight/agents.ts @@ -1,6 +1,5 @@ import type { AgentAppComposerResponse, - AgentDriveListResponse, AgentDriveSkillListResponse, AgentSoulConfig, } from '@dify/contracts/api/console/agent/types.gen' @@ -12,11 +11,7 @@ import { agentBuilderFixedInputs, agentBuilderPreseededResources, } from '../agent-builder-resources' -import { - agentBuilderFileTreeFixtureFileNames, - agentBuilderFileTreeFixtureFiles, - agentBuilderTestMaterials, -} from '../test-materials' +import { agentBuilderTestMaterials } from '../test-materials' import { asArray, asRecord, @@ -456,80 +451,3 @@ export async function skipMissingPreseededDualRetrievalAgentConfiguration( await ctx.dispose() } } - -export async function skipMissingPreseededAgentFileTreeFixture( - world: DifyWorld, - agentName: string, -): Promise<'skipped' | PreseededResource> { - const agent = await skipMissingPreseededAgent(world, agentName) - if (agent === 'skipped') - return agent - - const ctx = await createApiContext() - try { - const query = buildQuery({ prefix: 'files/' }) - const response = await ctx.get(`/console/api/agent/${agent.id}/drive/files?${query}`) - await expectApiResponseOK(response, `Check preseeded Agent file tree ${agentName}`) - const body = (await response.json()) as AgentDriveListResponse - const keys = (body.items ?? []).map(item => item.key) - const missingFiles = agentBuilderFileTreeFixtureFiles.filter( - filePath => - !keys.some(key => key === `files/${filePath}` || key.endsWith(`/${filePath}`)), - ) - - if (missingFiles.length > 0) { - return skipBlockedPrecondition( - world, - `Preseeded Agent "${agentName}" is missing file tree fixture files: ${missingFiles.join(', ')}.`, - ) - } - - return { - id: agent.id, - kind: 'agent', - name: agent.name, - } - } - finally { - await ctx.dispose() - } -} - -export async function skipMissingPreseededAgentFlatFileFixtureConfiguration( - world: DifyWorld, - agentName: string, -): Promise<'skipped' | PreseededResource> { - const agent = await skipMissingPreseededAgent(world, agentName) - if (agent === 'skipped') - return agent - - const ctx = await createApiContext() - try { - const response = await ctx.get(`/console/api/agent/${agent.id}/composer`) - await expectApiResponseOK(response, `Check preseeded Agent flat file fixture ${agentName}`) - const body = (await response.json()) as AgentAppComposerResponse - const configFiles = Array.isArray(body.agent_soul?.config_files) - ? body.agent_soul.config_files - : [] - const fileNames = configFiles - .map(file => (typeof file === 'object' && file !== null && 'name' in file ? file.name : undefined)) - .filter((name): name is string => typeof name === 'string') - const missingFiles = agentBuilderFileTreeFixtureFileNames.filter(fileName => !fileNames.includes(fileName)) - - if (missingFiles.length > 0) { - return skipBlockedPrecondition( - world, - `Preseeded Agent "${agentName}" is missing current flat Files fixture configuration: ${missingFiles.join(', ')}. Hierarchical Files display remains blocked until Agent config files support tree paths.`, - ) - } - - return { - id: agent.id, - kind: 'agent', - name: agent.name, - } - } - finally { - await ctx.dispose() - } -} diff --git a/e2e/features/agent-v2/support/seed.ts b/e2e/features/agent-v2/support/seed.ts index e3a3eb40bc1..716be7ec336 100644 --- a/e2e/features/agent-v2/support/seed.ts +++ b/e2e/features/agent-v2/support/seed.ts @@ -989,14 +989,6 @@ const agentV2FullSeedTasks = (): SeedTask[] => [ title: agentBuilderPreseededResources.oauthToolAgent, run: seedOAuthToolAgent, }, - { - id: 'file-tree-agent', - title: agentBuilderPreseededResources.fileTreeAgent, - run: async () => blocked( - agentBuilderPreseededResources.fileTreeAgent, - 'Agent drive arbitrary file upload does not have a stable public seed helper yet.', - ), - }, { id: 'dual-retrieval-agent', title: agentBuilderPreseededResources.dualRetrievalAgent, diff --git a/e2e/features/agent-v2/support/test-materials.ts b/e2e/features/agent-v2/support/test-materials.ts index eb6497918b0..a1b9a811ce8 100644 --- a/e2e/features/agent-v2/support/test-materials.ts +++ b/e2e/features/agent-v2/support/test-materials.ts @@ -1,4 +1,3 @@ -import path from 'node:path' import { getGeneratedTextMaterialPath, getTestMaterialPath } from '../../../support/test-materials' export const agentBuilderTestMaterials = { @@ -11,29 +10,12 @@ export const agentBuilderTestMaterials = { invalidEnv: 'agent-invalid.env', buildInstruction: 'agent-build-instruction.txt', summarySkill: 'e2e-summary-skill/SKILL.md', - fileTreeFixture: 'file_tree_fixture', - countBatch5: 'count_batch_5_valid_files', - countBatch6: 'count_batch_6_valid_files', - countTotal50: 'count_total_50_valid_files', - countTotalExtra1: 'count_total_extra_1_valid_file', } as const export const agentBuilderGeneratedTestMaterials = { - slowUploadFile: 'agent-slow-upload-file.txt', tooLargeFile: 'agent-too-large-file.txt', } as const -export const agentBuilderFileTreeFixtureFiles = [ - 'assets/sample.csv', - 'docs/中文说明.md', - 'public/index.html', - 'src/main.txt', - 'web-game/README.md', -] as const - -export const agentBuilderFileTreeFixtureFileNames = agentBuilderFileTreeFixtureFiles - .map(filePath => path.basename(filePath)) - export const getAgentBuilderTestMaterialPath = (material: keyof typeof agentBuilderTestMaterials) => getTestMaterialPath(agentBuilderTestMaterials[material]) @@ -43,10 +25,3 @@ export const getTooLargeAgentFilePath = () => sizeBytes: 16 * 1024 * 1024, seedText: 'E2E_TOO_LARGE_FILE_FIXTURE', }) - -export const getSlowUploadAgentFilePath = () => - getGeneratedTextMaterialPath({ - fileName: 'agent-slow-upload-file.txt', - sizeBytes: 2 * 1024 * 1024, - seedText: 'E2E_SLOW_UPLOAD_FILE_FIXTURE', - }) diff --git a/e2e/features/auth/session-refresh.feature b/e2e/features/auth/session-refresh.feature new file mode 100644 index 00000000000..f565a6aff69 --- /dev/null +++ b/e2e/features/auth/session-refresh.feature @@ -0,0 +1,9 @@ +@auth @core @authenticated +Feature: Console session refresh + + Scenario: Refresh the console session during server-side navigation + Given I am signed in as the default E2E admin + And my console session requires token refresh + When I open the default console entry after the access token expires + Then I should be on the console home + And I should not see the "Sign in" button diff --git a/e2e/features/auth/sign-in.feature b/e2e/features/auth/sign-in.feature index a9a1e13626d..ffa8a02ad71 100644 --- a/e2e/features/auth/sign-in.feature +++ b/e2e/features/auth/sign-in.feature @@ -1,8 +1,8 @@ @auth @smoke @core @unauthenticated Feature: Sign in - Scenario: Sign in with valid credentials and reach the apps console + Scenario: Sign in with valid credentials and reach the console home Given I am not signed in When I open the sign-in page And I sign in as the default E2E admin - Then I should be on the apps console + Then I should be on the console home diff --git a/e2e/features/smoke/authenticated-entry.feature b/e2e/features/smoke/authenticated-entry.feature index 53d72bd667c..955279bd247 100644 --- a/e2e/features/smoke/authenticated-entry.feature +++ b/e2e/features/smoke/authenticated-entry.feature @@ -1,7 +1,7 @@ @smoke @authenticated -Feature: Authenticated app console - Scenario: Open the apps console with the shared authenticated state +Feature: Authenticated console home + Scenario: Open the default console entry with the shared authenticated state Given I am signed in as the default E2E admin - When I open the apps console - Then I should stay on the apps console + When I open the default console entry + Then I should be on the console home And I should not see the "Sign in" button diff --git a/e2e/features/smoke/unauthenticated-entry.feature b/e2e/features/smoke/unauthenticated-entry.feature index a2783c1cba2..604ea4970dd 100644 --- a/e2e/features/smoke/unauthenticated-entry.feature +++ b/e2e/features/smoke/unauthenticated-entry.feature @@ -1,7 +1,7 @@ @smoke @unauthenticated -Feature: Unauthenticated app console entry - Scenario: Redirect to the sign-in page when opening the apps console without logging in +Feature: Unauthenticated console home entry + Scenario: Redirect to the sign-in page when opening the default console entry without logging in Given I am not signed in - When I open the apps console + When I open the default console entry Then I should be redirected to the signin page And I should see the "Sign in" button diff --git a/e2e/features/step-definitions/agent-v2/access-point-web-app.steps.ts b/e2e/features/step-definitions/agent-v2/access-point-web-app.steps.ts index d9d85e83772..e5de7e1ac26 100644 --- a/e2e/features/step-definitions/agent-v2/access-point-web-app.steps.ts +++ b/e2e/features/step-definitions/agent-v2/access-point-web-app.steps.ts @@ -1,3 +1,4 @@ +import type { Page } from '@playwright/test' import type { DifyWorld } from '../../support/world' import { Given, Then, When } from '@cucumber/cucumber' import { expect } from '@playwright/test' @@ -12,6 +13,9 @@ import { const WEB_APP_RUNTIME_RESPONSE_STEP_TIMEOUT_MS = 180_000 +const getWebAppMessageInput = (webAppPage: Page) => + webAppPage.getByPlaceholder(/^Talk to /).last() + Then('I should see the Agent v2 Web app access URL', async function (this: DifyWorld) { const webAppCard = getWebAppCard(this) @@ -71,7 +75,7 @@ When('I send an E2E message in the Agent v2 Web app', async function (this: Dify if (!webAppPage) throw new Error('No Agent v2 Web app page was opened.') - const messageInput = webAppPage.getByRole('textbox').last() + const messageInput = getWebAppMessageInput(webAppPage) await expect(messageInput).toBeEditable({ timeout: 30_000 }) await messageInput.fill('Please reply with the test success marker.') await messageInput.press('Enter') @@ -84,7 +88,7 @@ Then('the Agent v2 Web app should open in a new tab', async function (this: Dify throw new Error('No Agent v2 Web app page was opened.') await expect(webAppPage).toHaveURL(webAppURL) - await expect(webAppPage.getByRole('textbox').last()).toBeEditable({ timeout: 30_000 }) + await expect(getWebAppMessageInput(webAppPage)).toBeEditable({ timeout: 30_000 }) await webAppPage.close() this.agentBuilder.accessPoint.webAppPage = undefined this.agentBuilder.accessPoint.webAppURL = undefined diff --git a/e2e/features/step-definitions/agent-v2/access-point.steps.ts b/e2e/features/step-definitions/agent-v2/access-point.steps.ts index 40c7e852aa3..0ed6e456a93 100644 --- a/e2e/features/step-definitions/agent-v2/access-point.steps.ts +++ b/e2e/features/step-definitions/agent-v2/access-point.steps.ts @@ -6,7 +6,7 @@ import { setAgentApiAccess, setAgentSiteAccessAndGetURL, } from '../../agent-v2/support/access-point' -import { getAgentAccessPath, publishAgent } from '../../agent-v2/support/agent' +import { getAgentAccessPath, publishAgentWithPublishableDraft } from '../../agent-v2/support/agent' import { getAccessRegion, getAccessSurfaceCard, @@ -15,7 +15,7 @@ import { } from './access-point-helpers' Given('the Agent v2 draft has been published via API', async function (this: DifyWorld) { - await publishAgent(getCurrentAgentId(this)) + await publishAgentWithPublishableDraft(getCurrentAgentId(this)) }) Given( diff --git a/e2e/features/step-definitions/agent-v2/build-draft.steps.ts b/e2e/features/step-definitions/agent-v2/build-draft.steps.ts index 9443197e6e8..77e61fa1700 100644 --- a/e2e/features/step-definitions/agent-v2/build-draft.steps.ts +++ b/e2e/features/step-definitions/agent-v2/build-draft.steps.ts @@ -1,4 +1,4 @@ -import type { Response } from '@playwright/test' +import type { Page, Response } from '@playwright/test' import type { DifyWorld } from '../../support/world' import { readFile } from 'node:fs/promises' import { Given, Then, When } from '@cucumber/cucumber' @@ -8,7 +8,9 @@ import { saveAgentComposerDraft, } from '../../agent-v2/support/agent' import { + agentBuildDraftExists, applyAgentBuildDraft, + getAgentBuildDraft, saveAgentBuildDraft, } from '../../agent-v2/support/agent-build-draft' import { agentBuilderFixedInputs, agentBuilderPreseededResources } from '../../agent-v2/support/agent-builder-resources' @@ -25,12 +27,40 @@ import { hasToolEntry } from '../../agent-v2/support/preflight/tools' import { agentBuilderTestMaterials, getAgentBuilderTestMaterialPath } from '../../agent-v2/support/test-materials' import { getPreseededToolContract } from '../../agent-v2/support/tools' import { + expectAgentModelRequiredFeedback, getAgentEnvVariableValue, getCurrentAgentId, uploadSummaryConfigSkillForBuildDraft, } from './configure-helpers' const BUILD_DRAFT_RUNTIME_STEP_TIMEOUT_MS = 180_000 +const BUILD_DRAFT_NOTE_SYNC_TIMEOUT_MS = 30_000 +const BUILD_NOTE_FILE_NAME = 'build_note.md' +const BUILD_NOTE_MARKER = 'E2E_BUILD_DRAFT_PASS' +const BUILD_NOTE_GENERATED_BADGE = 'Generated' + +const getBuildDraftBar = (page: Page) => + page.getByRole('group', { name: 'Build draft' }) + +const getBuildNoteFileButton = (page: Page) => + page.getByRole('region', { name: 'Files' }) + .getByRole('button') + .filter({ hasText: BUILD_NOTE_FILE_NAME }) + .filter({ hasText: BUILD_NOTE_GENERATED_BADGE }) + +const getConfigNote = (value: Awaited>) => + value.agent_soul?.config_note ?? '' + +const getLastBuildChatAnswerText = async (page: Page) => { + const answer = page.getByTestId('chat-answer-container').last() + if (await answer.count() === 0) + return '' + + return (await answer.textContent())?.replace(/\s+/g, ' ').trim() ?? '' +} + +const formatBuildChatAnswerText = (text: string) => + text.length > 500 ? `${text.slice(0, 500)}...` : text Given( 'an Agent v2 Build draft adds the supported E2E files, skills, and env', @@ -115,7 +145,7 @@ Given( When( 'I generate an Agent v2 Build draft from the fixed instruction', - { timeout: 180_000 }, + { timeout: BUILD_DRAFT_RUNTIME_STEP_TIMEOUT_MS }, async function (this: DifyWorld) { const page = this.getPage() const agentId = getCurrentAgentId(this) @@ -135,13 +165,28 @@ When( await page.getByRole('button', { name: 'Start build' }).click() expect((await checkoutResponsePromise).ok()).toBe(true) - expect((await chatResponsePromise).ok()).toBe(true) - await expect(page.getByText('Build draft')).toBeVisible({ timeout: 120_000 }) - await expect(page.getByRole('button', { exact: true, name: 'Apply' })).toBeEnabled({ timeout: 120_000 }) + const chatResponse = await chatResponsePromise + expect(chatResponse.ok()).toBe(true) + expect(await chatResponse.finished()).toBeNull() + + await expect(page.getByRole('button', { name: 'Stop responding' })).not.toBeVisible() + await expect(getBuildDraftBar(page)).toBeVisible() + await expect(page.getByRole('button', { exact: true, name: 'Apply' })).toBeEnabled() await expect(page.getByRole('button', { exact: true, name: 'Discard' })).toBeEnabled() }, ) +When( + 'I try to generate an Agent v2 Build draft without a model', + async function (this: DifyWorld) { + const page = this.getPage() + + await page.getByRole('button', { exact: true, name: 'Build' }).click() + await page.getByPlaceholder('Describe what your agent should do').fill('Update the agent instructions for E2E.') + await page.getByRole('button', { name: 'Start build' }).click() + }, +) + const expectPageResponseOK = async (response: Response, action: string) => { if (response.ok()) return @@ -159,7 +204,20 @@ const expectPageResponseOK = async (response: Response, action: string) => { } When('I discard the Agent v2 Build draft', async function (this: DifyWorld) { - await this.getPage().getByRole('button', { exact: true, name: 'Discard' }).click() + const page = this.getPage() + const agentId = getCurrentAgentId(this) + + await page.getByRole('button', { exact: true, name: 'Discard' }).click() + const confirmDialog = page.getByRole('alertdialog', { name: 'Clear session and discard changes?' }) + await expect(confirmDialog).toBeVisible() + + const discardResponsePromise = page.waitForResponse(response => ( + response.request().method() === 'DELETE' + && new URL(response.url()).pathname.endsWith(`/console/api/agent/${agentId}/build-draft`) + )) + + await confirmDialog.getByRole('button', { name: 'Confirm' }).click() + await expectPageResponseOK(await discardResponsePromise, 'Discard Agent v2 Build draft') }) When( @@ -232,7 +290,7 @@ Then('Agent v2 Build chat unavailable Skill and Tool recovery should be availabl Then('I should see the Agent v2 Build draft pending changes', async function (this: DifyWorld) { const page = this.getPage() - await expect(page.getByText('Build draft')).toBeVisible({ timeout: 30_000 }) + await expect(getBuildDraftBar(page)).toBeVisible({ timeout: 30_000 }) await expect(page.getByRole('button', { exact: true, name: 'Apply' })).toBeEnabled() await expect(page.getByRole('button', { exact: true, name: 'Discard' })).toBeEnabled() }) @@ -241,9 +299,42 @@ Then('I should see the Agent v2 Build mode confirmation state', async function ( const page = this.getPage() await expect(page.getByText('Build mode', { exact: true })).toBeVisible() - await expect( - page.getByText('You\'re in build mode. Shape this setup through the chat on the right, then Apply.'), - ).toBeVisible() + await expect(page.getByText('Configure can only be updated by the agent in this mode.')).toBeVisible() + await expect(page.getByText('Shape this setup through the chat on the right, then Apply.')).toBeVisible() +}) + +Then( + 'the Agent v2 Build draft should include the generated build note', + async function (this: DifyWorld) { + try { + await expect.poll( + async () => getConfigNote(await getAgentBuildDraft(getCurrentAgentId(this))), + { timeout: BUILD_DRAFT_NOTE_SYNC_TIMEOUT_MS }, + ).toContain(BUILD_NOTE_MARKER) + } + catch (error) { + const lastAnswerText = await getLastBuildChatAnswerText(this.getPage()) + throw new Error( + `Agent v2 Build draft note did not include ${BUILD_NOTE_MARKER}. Last Build chat answer: ${formatBuildChatAnswerText(lastAnswerText) || ''}`, + { cause: error }, + ) + } + }, +) + +Then('I should see the generated Agent v2 build note in Configure', async function (this: DifyWorld) { + await expect(getBuildNoteFileButton(this.getPage())).toBeVisible() +}) + +Then('Agent v2 Build chat should be blocked until a model is configured', async function (this: DifyWorld) { + await expectAgentModelRequiredFeedback(this.getPage()) +}) + +Then('the Agent v2 Build draft should not be checked out', async function (this: DifyWorld) { + await expect.poll( + async () => agentBuildDraftExists(getCurrentAgentId(this)), + { timeout: 30_000 }, + ).toBe(false) }) Then('I should see the e2e-summary-skill Skill in the Skills section', async function (this: DifyWorld) { @@ -292,6 +383,16 @@ Then( }, ) +Then( + 'the normal Agent v2 draft should not include the generated build note', + async function (this: DifyWorld) { + await expect.poll( + async () => (await getAgentComposerDraft(getCurrentAgentId(this))).agent_soul?.config_note ?? '', + { timeout: 30_000 }, + ).not.toContain(BUILD_NOTE_MARKER) + }, +) + Then( 'the normal Agent v2 draft should not include the Agent Builder JSON Replace tool', async function (this: DifyWorld) { @@ -378,7 +479,7 @@ Then( Then('the Agent v2 Build draft should no longer be active', async function (this: DifyWorld) { const page = this.getPage() - await expect(page.getByText('Build draft')).not.toBeVisible() + await expect(getBuildDraftBar(page)).not.toBeVisible() await expect(page.getByRole('button', { name: 'Apply' })).not.toBeVisible() await expect(page.getByRole('button', { name: 'Discard' })).not.toBeVisible() }) diff --git a/e2e/features/step-definitions/agent-v2/configure-helpers.ts b/e2e/features/step-definitions/agent-v2/configure-helpers.ts index e51c89b1da0..10927c23b1b 100644 --- a/e2e/features/step-definitions/agent-v2/configure-helpers.ts +++ b/e2e/features/step-definitions/agent-v2/configure-helpers.ts @@ -138,6 +138,10 @@ export const expectAgentConfigFileSaved = async ( }) } +export const expectAgentModelRequiredFeedback = async (page: ReturnType) => { + await expect(page.getByText('Select your model')).toBeVisible({ timeout: 10_000 }) +} + export const uploadSummaryConfigSkillForBuildDraft = async (world: DifyWorld) => { const agentId = getCurrentAgentId(world) const skill = await uploadAgentConfigSkillToDraft({ diff --git a/e2e/features/step-definitions/agent-v2/configure.steps.ts b/e2e/features/step-definitions/agent-v2/configure.steps.ts index bf9d4f0dfce..02ed9004d3c 100644 --- a/e2e/features/step-definitions/agent-v2/configure.steps.ts +++ b/e2e/features/step-definitions/agent-v2/configure.steps.ts @@ -15,6 +15,7 @@ import { concurrentFirstAgentPrompt, concurrentSecondAgentPrompt, createAgentSoulConfigWithModel, + createPublishableAgentSoulConfig, normalAgentPrompt, normalAgentSoulConfig, updatedAgentPrompt, @@ -131,6 +132,16 @@ Given('the Agent v2 composer draft uses the normal E2E prompt', async function ( await saveAgentComposerDraft(getCurrentAgentId(this), normalAgentSoulConfig) }) +Given( + 'the Agent v2 composer draft is publishable', + async function (this: DifyWorld) { + await saveAgentComposerDraft( + getCurrentAgentId(this), + createPublishableAgentSoulConfig(normalAgentSoulConfig), + ) + }, +) + Given('the e2e-summary-skill Skill is available to the Agent v2 test agent', async function (this: DifyWorld) { const agentId = getCurrentAgentId(this) const upload = await uploadAgentDriveSkill({ @@ -271,10 +282,13 @@ Then('I should be on the Agent v2 configure page', async function (this: DifyWor Then('I should see the Agent v2 configure workspace', async function (this: DifyWorld) { const page = this.getPage() + const agentName = this.lastCreatedAgentName + if (!agentName) + throw new Error('No Agent v2 name found. Create an Agent v2 test agent first.') await expect(page.getByRole('region', { name: 'Configure' })).toBeVisible({ timeout: 30_000 }) await expect(page.getByRole('heading', { name: 'Configure' })).toBeVisible() - await expect(page.getByText(this.lastCreatedAgentName!)).toBeVisible() + await expect(page.getByText(agentName, { exact: true })).toBeVisible() }) Then( diff --git a/e2e/features/step-definitions/agent-v2/files.steps.ts b/e2e/features/step-definitions/agent-v2/files.steps.ts index 76f8eaf6617..ddce18f62d7 100644 --- a/e2e/features/step-definitions/agent-v2/files.steps.ts +++ b/e2e/features/step-definitions/agent-v2/files.steps.ts @@ -2,10 +2,7 @@ import type { DifyWorld } from '../../support/world' import { Given, Then, When } from '@cucumber/cucumber' import { expect } from '@playwright/test' import { skipBlockedPrecondition } from '../../agent-v2/support/preflight/common' -import { - agentBuilderFileTreeFixtureFileNames, - agentBuilderTestMaterials, -} from '../../agent-v2/support/test-materials' +import { agentBuilderTestMaterials } from '../../agent-v2/support/test-materials' import { expectAgentConfigFileHidden, expectAgentConfigFileSaved, @@ -70,30 +67,6 @@ Then('I should not see the dropped Agent v2 files in the Files section', async f await expectAgentConfigFileHidden(this, 'emptyFile') }) -Then( - 'I should see the Agent v2 file fixture entries in the current flat Files list', - async function (this: DifyWorld) { - const page = this.getPage() - const filesSection = page.getByRole('region', { name: 'Files' }) - const filesList = filesSection.getByLabel('Agent files') - - await expect(filesSection).toBeVisible({ timeout: 30_000 }) - await expect(filesList).toBeVisible() - - for (const fileName of agentBuilderFileTreeFixtureFileNames) { - await expect(filesList.getByRole('button', { - exact: true, - name: fileName, - })).toBeVisible() - } - - await expect(filesList.getByRole('button', { exact: true, name: 'assets' })).toHaveCount(0) - await expect(filesList.getByRole('button', { exact: true, name: 'docs' })).toHaveCount(0) - await expect(filesList.getByRole('button', { exact: true, name: 'public' })).toHaveCount(0) - await expect(filesList.getByRole('button', { exact: true, name: 'src' })).toHaveCount(0) - await expect(filesList.getByRole('button', { exact: true, name: 'web-game' })).toHaveCount(0) - }, -) Then('I should see the small Agent v2 file in the Files section', async function (this: DifyWorld) { await expectAgentConfigFileVisible(this, 'smallFile') }) @@ -167,41 +140,3 @@ Given('Agent v2 oversized file rejection is available', async function (this: Di Then('Agent v2 oversized file rejection should be available', async function (this: DifyWorld) { return skipOversizedFileRejection(this) }) - -async function skipTotalFileCountLimits(world: DifyWorld) { - return skipBlockedPrecondition( - world, - 'Agent v2 total file count limits are not defined for Agent config files in the current product contract.', - { - owner: 'product', - remediation: 'Define the Agent config file total-count limit and user-visible error before enabling this scenario.', - }, - ) -} - -Given('Agent v2 total file count limits are available', async function (this: DifyWorld) { - return skipTotalFileCountLimits(this) -}) - -Then('Agent v2 total file count limits should be available', async function (this: DifyWorld) { - return skipTotalFileCountLimits(this) -}) - -async function skipInProgressFileUploadRecovery(world: DifyWorld) { - return skipBlockedPrecondition( - world, - 'Agent v2 in-progress file upload recovery is not stable: the current dialog has no deterministic slow-upload fixture or user-visible navigation guard contract.', - { - owner: 'product/test-infra', - remediation: 'Define upload-in-progress navigation behavior and provide a deterministic slow upload fixture before enabling this scenario.', - }, - ) -} - -Given('Agent v2 in-progress file upload recovery is available', async function (this: DifyWorld) { - return skipInProgressFileUploadRecovery(this) -}) - -Then('Agent v2 in-progress file upload recovery should be available', async function (this: DifyWorld) { - return skipInProgressFileUploadRecovery(this) -}) diff --git a/e2e/features/step-definitions/agent-v2/output-variables.steps.ts b/e2e/features/step-definitions/agent-v2/output-variables.steps.ts index f21f1944217..5619f5bae8a 100644 --- a/e2e/features/step-definitions/agent-v2/output-variables.steps.ts +++ b/e2e/features/step-definitions/agent-v2/output-variables.steps.ts @@ -330,44 +330,6 @@ async function expectAgentTaskOutputReference( await expect(page.getByText(unexpectedName, { exact: true })).toHaveCount(0) } -async function skipStandaloneOutputVariables(world: DifyWorld) { - return skipBlockedPrecondition( - world, - 'Standalone Agent Output Variables are not available: output variables currently belong to Workflow Agent v2 nodes.', - { - owner: 'product', - remediation: 'Expose standalone Agent Output Variables or keep this scenario excluded until the product path exists.', - }, - ) -} - -Given('Agent v2 standalone Output Variables are available', async function (this: DifyWorld) { - return skipStandaloneOutputVariables(this) -}) - -Then('Agent v2 standalone Output Variables should be available', async function (this: DifyWorld) { - return skipStandaloneOutputVariables(this) -}) - -async function skipWorkflowOutputRetryStrategy(world: DifyWorld) { - return skipBlockedPrecondition( - world, - 'Agent v2 workflow Output Variables retry strategy is not available in the current editor UI.', - { - owner: 'product', - remediation: 'Expose user-visible retry strategy controls before enabling this scenario.', - }, - ) -} - -Given('Agent v2 workflow output retry strategy is available', async function (this: DifyWorld) { - return skipWorkflowOutputRetryStrategy(this) -}) - -Then('Agent v2 workflow output retry strategy should be available', async function (this: DifyWorld) { - return skipWorkflowOutputRetryStrategy(this) -}) - async function skipWorkflowTaskOutputReferenceDeletionConsistency(world: DifyWorld) { return skipBlockedPrecondition( world, @@ -392,22 +354,3 @@ Then( return skipWorkflowTaskOutputReferenceDeletionConsistency(this) }, ) - -async function skipWorkflowOutputRetryCountValidation(world: DifyWorld) { - return skipBlockedPrecondition( - world, - 'Agent v2 workflow Output Variables retry count validation is not reachable because retry strategy controls are not available in the current editor UI.', - { - owner: 'product', - remediation: 'Expose retry count controls and validation states before enabling this scenario.', - }, - ) -} - -Given('Agent v2 workflow output retry count validation is available', async function (this: DifyWorld) { - return skipWorkflowOutputRetryCountValidation(this) -}) - -Then('Agent v2 workflow output retry count validation should be available', async function (this: DifyWorld) { - return skipWorkflowOutputRetryCountValidation(this) -}) diff --git a/e2e/features/step-definitions/agent-v2/preflight.steps.ts b/e2e/features/step-definitions/agent-v2/preflight.steps.ts index bd62425ba34..9748680449e 100644 --- a/e2e/features/step-definitions/agent-v2/preflight.steps.ts +++ b/e2e/features/step-definitions/agent-v2/preflight.steps.ts @@ -9,8 +9,6 @@ import { skipMissingAgentBackendRuntime } from '../../agent-v2/support/preflight import { skipMissingPreseededAgent, skipMissingPreseededAgentDriveSkill, - skipMissingPreseededAgentFileTreeFixture, - skipMissingPreseededAgentFlatFileFixtureConfiguration, skipMissingPreseededDualRetrievalAgentConfiguration, skipMissingPreseededFullConfigAgentCoreConfiguration, skipMissingPreseededOAuthToolAgentConfiguration, @@ -183,29 +181,6 @@ Given( }, ) -Given( - 'the Agent Builder preseeded Agent {string} includes the file tree fixture files', - async function (this: DifyWorld, agentName: string) { - const resource = await skipMissingPreseededAgentFileTreeFixture(this, agentName) - if (resource === 'skipped') - return resource - - this.agentBuilder.preflight.preseededResources[`${agentName} / file tree fixture`] = resource - }, -) - -Given( - 'the Agent Builder preseeded Agent {string} includes the current flat file fixture configuration', - async function (this: DifyWorld, agentName: string) { - const resource = await skipMissingPreseededAgentFlatFileFixtureConfiguration(this, agentName) - if (resource === 'skipped') - return resource - - this.agentBuilder.preflight.preseededResources[`${agentName} / flat file fixture configuration`] - = resource - }, -) - Given( 'the Agent Builder preseeded Agent {string} has Backend service API access with an API key', async function (this: DifyWorld, agentName: string) { diff --git a/e2e/features/step-definitions/agent-v2/publish.steps.ts b/e2e/features/step-definitions/agent-v2/publish.steps.ts index 8bc8b94c751..c04fca83cb7 100644 --- a/e2e/features/step-definitions/agent-v2/publish.steps.ts +++ b/e2e/features/step-definitions/agent-v2/publish.steps.ts @@ -4,7 +4,7 @@ import { expect } from '@playwright/test' import { waitForAgentConfigureAutosaved } from '../../../support/agent-configure' import { getAgentVersionDetail, getTestAgent } from '../../agent-v2/support/agent' import { normalAgentPrompt } from '../../agent-v2/support/agent-soul' -import { getCurrentAgentId } from './configure-helpers' +import { expectAgentModelRequiredFeedback, getCurrentAgentId } from './configure-helpers' When('I publish the Agent v2 draft', async function (this: DifyWorld) { const page = this.getPage() @@ -14,6 +14,25 @@ When('I publish the Agent v2 draft', async function (this: DifyWorld) { await publishButton.click() }) +When('I try to publish the Agent v2 draft without a model', async function (this: DifyWorld) { + const page = this.getPage() + const publishButton = page.getByRole('button', { name: /^Publish(?: update)?$/ }) + + await expect(publishButton).toBeEnabled({ timeout: 30_000 }) + await publishButton.click() +}) + +Then('Agent v2 publish should be blocked until a model is configured', async function (this: DifyWorld) { + await expectAgentModelRequiredFeedback(this.getPage()) +}) + +Then('the Agent v2 draft should remain unpublished', async function (this: DifyWorld) { + await expect.poll( + async () => (await getTestAgent(getCurrentAgentId(this))).active_config_is_published, + { timeout: 30_000 }, + ).toBe(false) +}) + Then('the Agent v2 configuration should be saved automatically', async function (this: DifyWorld) { await waitForAgentConfigureAutosaved(this.getPage()) }) diff --git a/e2e/features/step-definitions/agent-v2/tools.steps.ts b/e2e/features/step-definitions/agent-v2/tools.steps.ts index 78eab7a488a..b8a78018c96 100644 --- a/e2e/features/step-definitions/agent-v2/tools.steps.ts +++ b/e2e/features/step-definitions/agent-v2/tools.steps.ts @@ -240,7 +240,6 @@ When( await expect(toolsSection).toBeVisible({ timeout: 30_000 }) await toolsSection.getByRole('button', { name: 'Add tool' }).click() - await this.getPage().getByRole('button', { name: /^Tool\b/ }).click() const search = getToolSelectorSearch(this) await expect(search).toBeVisible() diff --git a/e2e/features/step-definitions/auth/session-refresh.steps.ts b/e2e/features/step-definitions/auth/session-refresh.steps.ts new file mode 100644 index 00000000000..f6468bfaf3d --- /dev/null +++ b/e2e/features/step-definitions/auth/session-refresh.steps.ts @@ -0,0 +1,43 @@ +import type { DifyWorld } from '../../support/world' +import { Given, When } from '@cucumber/cucumber' +import { expect } from '@playwright/test' + +const consoleAccessTokenCookieName = /^(?:__Host-)?access_token$/ +const consoleRefreshTokenCookieName = /^(?:__Host-)?refresh_token$/ + +Given('my console session requires token refresh', async function (this: DifyWorld) { + if (!this.context) + throw new Error('Playwright browser context has not been initialized for this scenario.') + + const cookies = await this.context.cookies() + const hasAccessToken = cookies.some(cookie => consoleAccessTokenCookieName.test(cookie.name)) + const hasRefreshToken = cookies.some(cookie => consoleRefreshTokenCookieName.test(cookie.name)) + + expect(hasAccessToken, 'Expected the authenticated E2E session to include a console access token.').toBe(true) + expect(hasRefreshToken, 'Expected the authenticated E2E session to include a console refresh token.').toBe(true) + + await this.context.clearCookies({ name: consoleAccessTokenCookieName }) + + const remainingCookies = await this.context.cookies() + expect( + remainingCookies.some(cookie => consoleAccessTokenCookieName.test(cookie.name)), + 'Expected the console access token to be removed before opening the default console entry.', + ).toBe(false) + expect( + remainingCookies.some(cookie => consoleRefreshTokenCookieName.test(cookie.name)), + 'Expected the console refresh token to remain available for server-side refresh.', + ).toBe(true) +}) + +When('I open the default console entry after the access token expires', async function (this: DifyWorld) { + const page = this.getPage() + const refreshRequestPromise = page.waitForRequest((request) => { + const url = new URL(request.url()) + return url.pathname.endsWith('/auth/refresh') && url.searchParams.get('redirect_url') === '/' + }) + + await page.goto('/') + + const refreshRequest = await refreshRequestPromise + this.attach(`Session refresh request: ${refreshRequest.url()}`, 'text/plain') +}) diff --git a/e2e/features/step-definitions/auth/sign-in.steps.ts b/e2e/features/step-definitions/auth/sign-in.steps.ts index 469203d8bfd..fcf106048b9 100644 --- a/e2e/features/step-definitions/auth/sign-in.steps.ts +++ b/e2e/features/step-definitions/auth/sign-in.steps.ts @@ -1,10 +1,9 @@ import type { DifyWorld } from '../../support/world' -import { Then, When } from '@cucumber/cucumber' -import { expect } from '@playwright/test' +import { When } from '@cucumber/cucumber' import { adminCredentials } from '../../../fixtures/auth' When('I open the sign-in page', async function (this: DifyWorld) { - await this.getPage().goto('/signin?redirect_url=%2Fapps') + await this.getPage().goto('/signin') }) When('I sign in as the default E2E admin', async function (this: DifyWorld) { @@ -14,7 +13,3 @@ When('I sign in as the default E2E admin', async function (this: DifyWorld) { await page.getByLabel('Password', { exact: true }).fill(adminCredentials.password) await page.getByRole('button', { name: 'Sign in' }).click() }) - -Then('I should be on the apps console', async function (this: DifyWorld) { - await expect(this.getPage()).toHaveURL(/\/apps(?:\?.*)?$/, { timeout: 30_000 }) -}) diff --git a/e2e/features/step-definitions/common/navigation.steps.ts b/e2e/features/step-definitions/common/navigation.steps.ts index a558d96f937..10ee7b789f0 100644 --- a/e2e/features/step-definitions/common/navigation.steps.ts +++ b/e2e/features/step-definitions/common/navigation.steps.ts @@ -2,6 +2,11 @@ import type { DifyWorld } from '../../support/world' import { Then, When } from '@cucumber/cucumber' import { expect } from '@playwright/test' import { waitForAppsConsole } from '../../../support/apps' +import { waitForConsoleHome } from '../../../support/home' + +When('I open the default console entry', async function (this: DifyWorld) { + await this.getPage().goto('/') +}) When('I open the apps console', async function (this: DifyWorld) { await this.getPage().goto('/apps') @@ -15,6 +20,10 @@ Then('I should stay on the apps console', async function (this: DifyWorld) { await waitForAppsConsole(this.getPage()) }) +Then('I should be on the console home', async function (this: DifyWorld) { + await waitForConsoleHome(this.getPage()) +}) + Then('I should be redirected to the signin page', async function (this: DifyWorld) { await expect(this.getPage()).toHaveURL(/\/signin(?:\?.*)?$/) }) diff --git a/e2e/features/support/hooks.ts b/e2e/features/support/hooks.ts index 8fd59934dae..658f3efc69b 100644 --- a/e2e/features/support/hooks.ts +++ b/e2e/features/support/hooks.ts @@ -1,4 +1,4 @@ -import type { Browser } from '@playwright/test' +import type { Browser, Page } from '@playwright/test' import type { Buffer } from 'node:buffer' import type { DifyWorld } from './world' import { mkdir, writeFile } from 'node:fs/promises' @@ -34,18 +34,50 @@ const sanitizeForPath = (value: string) => const writeArtifact = async ( scenarioName: string, + label: string, extension: 'html' | 'png', contents: Buffer | string, ) => { const artifactPath = path.join( artifactsDir, - `${Date.now()}-${sanitizeForPath(scenarioName || 'scenario')}.${extension}`, + `${Date.now()}-${sanitizeForPath(scenarioName || 'scenario')}-${sanitizeForPath(label)}.${extension}`, ) await writeFile(artifactPath, contents) return artifactPath } +const uniqueDiagnosticPages = (pages: { label: string, page: Page | undefined }[]) => { + const seen = new Set() + + return pages.filter(({ page }) => { + if (!page || page.isClosed() || seen.has(page)) + return false + + seen.add(page) + return true + }) as { label: string, page: Page }[] +} + +const captureDiagnosticPage = async ( + world: DifyWorld, + scenarioName: string, + label: string, + page: Page, +) => { + const screenshot = await page.screenshot({ + fullPage: true, + }) + const screenshotPath = await writeArtifact(scenarioName, label, 'png', screenshot) + world.attach(screenshot, 'image/png') + + const html = await page.content() + const htmlPath = await writeArtifact(scenarioName, label, 'html', html) + world.attach(html, 'text/html') + + return [screenshotPath, htmlPath] +} + const recordCleanup = async ( errors: string[], label: string, @@ -91,16 +123,26 @@ After(async function (this: DifyWorld, { pickle, result }) { const elapsedMs = this.scenarioStartedAt ? Date.now() - this.scenarioStartedAt : undefined const status = result?.status || Status.UNKNOWN - if (diagnosticArtifactStatuses.has(status) && this.page) { - const screenshot = await this.page.screenshot({ - fullPage: true, - }) - const screenshotPath = await writeArtifact(pickle.name, 'png', screenshot) - this.attach(screenshot, 'image/png') + if (diagnosticArtifactStatuses.has(status)) { + const artifactPaths: string[] = [] + const artifactErrors: string[] = [] + const diagnosticPages = uniqueDiagnosticPages([ + { label: 'main-page', page: this.page }, + { label: 'agent-v2-web-app', page: this.agentBuilder.accessPoint.webAppPage }, + { label: 'agent-v2-api-reference', page: this.agentBuilder.accessPoint.apiReferencePage }, + { label: 'agent-v2-workflow-reference', page: this.agentBuilder.accessPoint.workflowReferencePage }, + { label: 'agent-v2-concurrent-configure', page: this.agentBuilder.configure.concurrentPage }, + { label: 'agent-v2-workflow-console', page: this.agentBuilder.workflow.agentConsolePage }, + ]) - const html = await this.page.content() - const htmlPath = await writeArtifact(pickle.name, 'html', html) - this.attach(html, 'text/html') + for (const { label, page } of diagnosticPages) { + try { + artifactPaths.push(...await captureDiagnosticPage(this, pickle.name, label, page)) + } + catch (error) { + artifactErrors.push(`${label}: ${error instanceof Error ? error.message : String(error)}`) + } + } if (this.consoleErrors.length > 0) this.attach(`Console Errors:\n${this.consoleErrors.join('\n')}`, 'text/plain') @@ -108,7 +150,11 @@ After(async function (this: DifyWorld, { pickle, result }) { if (this.pageErrors.length > 0) this.attach(`Page Errors:\n${this.pageErrors.join('\n')}`, 'text/plain') - this.attach(`Artifacts:\n${[screenshotPath, htmlPath].join('\n')}`, 'text/plain') + if (artifactErrors.length > 0) + this.attach(`Artifact Errors:\n${artifactErrors.join('\n')}`, 'text/plain') + + if (artifactPaths.length > 0) + this.attach(`Artifacts:\n${artifactPaths.join('\n')}`, 'text/plain') } console.warn( diff --git a/e2e/fixtures/auth.ts b/e2e/fixtures/auth.ts index 9039f97483d..62d2880a465 100644 --- a/e2e/fixtures/auth.ts +++ b/e2e/fixtures/auth.ts @@ -3,7 +3,7 @@ import { Buffer } from 'node:buffer' import { mkdir, readFile, writeFile } from 'node:fs/promises' import path from 'node:path' import { fileURLToPath } from 'node:url' -import { waitForAppsConsole } from '../support/apps' +import { waitForConsoleHome } from '../support/home' import { apiURL, defaultBaseURL, defaultLocale } from '../test-env' export type AuthSessionMetadata = { @@ -150,12 +150,12 @@ export const ensureAuthenticatedState = async (browser: Browser, configuredBaseU const { mode, usedInitPassword } = await ensureAdminAccount(context, deadline) await loginAdmin(context, deadline) - console.warn('[e2e] auth bootstrap: verifying apps console') - await page.goto(appURL(baseURL, '/apps'), { + console.warn('[e2e] auth bootstrap: verifying console home') + await page.goto(appURL(baseURL, '/'), { timeout: getRemainingTimeout(deadline), waitUntil: 'domcontentloaded', }) - await waitForAppsConsole(page, getRemainingTimeout(deadline)) + await waitForConsoleHome(page, getRemainingTimeout(deadline)) await context.storageState({ path: authStatePath }) diff --git a/e2e/fixtures/test-materials/agent-build-instruction.txt b/e2e/fixtures/test-materials/agent-build-instruction.txt index 82ebf01d3a7..018b254d50e 100644 --- a/e2e/fixtures/test-materials/agent-build-instruction.txt +++ b/e2e/fixtures/test-materials/agent-build-instruction.txt @@ -1,3 +1,3 @@ -Update only the Agent instructions. -After applying, every response should briefly mention E2E_BUILD_DRAFT_PASS. -Do not add files, skills, tools, knowledge, environment variables, or other capabilities. +Update only the Agent build note for this Build draft. +The build note should record this durable context exactly: E2E_BUILD_DRAFT_PASS. +Do not update Agent instructions, prompt, files, skills, tools, knowledge, environment variables, or other capabilities. diff --git a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-1.txt b/e2e/fixtures/test-materials/count_batch_5_valid_files/file-1.txt deleted file mode 100644 index a64cd23552d..00000000000 --- a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-1.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 5 valid file 1 token E2E_BATCH_5_1 diff --git a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-2.txt b/e2e/fixtures/test-materials/count_batch_5_valid_files/file-2.txt deleted file mode 100644 index 1a65ed90e76..00000000000 --- a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-2.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 5 valid file 2 token E2E_BATCH_5_2 diff --git a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-3.txt b/e2e/fixtures/test-materials/count_batch_5_valid_files/file-3.txt deleted file mode 100644 index c5c00a2477e..00000000000 --- a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-3.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 5 valid file 3 token E2E_BATCH_5_3 diff --git a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-4.txt b/e2e/fixtures/test-materials/count_batch_5_valid_files/file-4.txt deleted file mode 100644 index 120e4ee335c..00000000000 --- a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-4.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 5 valid file 4 token E2E_BATCH_5_4 diff --git a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-5.txt b/e2e/fixtures/test-materials/count_batch_5_valid_files/file-5.txt deleted file mode 100644 index 4efaff6d828..00000000000 --- a/e2e/fixtures/test-materials/count_batch_5_valid_files/file-5.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 5 valid file 5 token E2E_BATCH_5_5 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-1.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-1.txt deleted file mode 100644 index 6ce029a426b..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-1.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 1 token E2E_BATCH_6_1 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-2.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-2.txt deleted file mode 100644 index 08620c1994f..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-2.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 2 token E2E_BATCH_6_2 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-3.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-3.txt deleted file mode 100644 index 73deeefa755..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-3.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 3 token E2E_BATCH_6_3 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-4.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-4.txt deleted file mode 100644 index 5b62b0d7075..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-4.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 4 token E2E_BATCH_6_4 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-5.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-5.txt deleted file mode 100644 index c0192f87bb4..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-5.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 5 token E2E_BATCH_6_5 diff --git a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-6.txt b/e2e/fixtures/test-materials/count_batch_6_valid_files/file-6.txt deleted file mode 100644 index e3505c43d59..00000000000 --- a/e2e/fixtures/test-materials/count_batch_6_valid_files/file-6.txt +++ /dev/null @@ -1 +0,0 @@ -Batch 6 valid file 6 token E2E_BATCH_6_6 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-01.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-01.txt deleted file mode 100644 index db1a55a0021..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-01.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 01 token E2E_TOTAL_50_01 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-02.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-02.txt deleted file mode 100644 index bdcd785ab8f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-02.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 02 token E2E_TOTAL_50_02 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-03.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-03.txt deleted file mode 100644 index cf00c92b0a3..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-03.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 03 token E2E_TOTAL_50_03 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-04.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-04.txt deleted file mode 100644 index 01522864989..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-04.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 04 token E2E_TOTAL_50_04 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-05.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-05.txt deleted file mode 100644 index 22a118ddd0f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-05.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 05 token E2E_TOTAL_50_05 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-06.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-06.txt deleted file mode 100644 index 3cee5574426..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-06.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 06 token E2E_TOTAL_50_06 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-07.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-07.txt deleted file mode 100644 index 381bf4d24d5..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-07.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 07 token E2E_TOTAL_50_07 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-08.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-08.txt deleted file mode 100644 index 8cc4d21fc86..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-08.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 08 token E2E_TOTAL_50_08 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-09.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-09.txt deleted file mode 100644 index 490cc4b779c..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-09.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 09 token E2E_TOTAL_50_09 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-10.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-10.txt deleted file mode 100644 index fe4089625f5..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-10.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 10 token E2E_TOTAL_50_10 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-11.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-11.txt deleted file mode 100644 index b5a9ae2ea3f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-11.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 11 token E2E_TOTAL_50_11 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-12.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-12.txt deleted file mode 100644 index 33e0f0b77c0..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-12.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 12 token E2E_TOTAL_50_12 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-13.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-13.txt deleted file mode 100644 index 0bf532f9a32..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-13.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 13 token E2E_TOTAL_50_13 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-14.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-14.txt deleted file mode 100644 index 60b3cb0d751..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-14.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 14 token E2E_TOTAL_50_14 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-15.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-15.txt deleted file mode 100644 index 15c0e3769be..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-15.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 15 token E2E_TOTAL_50_15 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-16.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-16.txt deleted file mode 100644 index 6c80e6e66bb..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-16.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 16 token E2E_TOTAL_50_16 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-17.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-17.txt deleted file mode 100644 index 9ac689f667a..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-17.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 17 token E2E_TOTAL_50_17 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-18.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-18.txt deleted file mode 100644 index 3a9c6725090..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-18.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 18 token E2E_TOTAL_50_18 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-19.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-19.txt deleted file mode 100644 index 35d544714c8..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-19.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 19 token E2E_TOTAL_50_19 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-20.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-20.txt deleted file mode 100644 index cc53594e52f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-20.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 20 token E2E_TOTAL_50_20 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-21.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-21.txt deleted file mode 100644 index 2e0af50430c..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-21.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 21 token E2E_TOTAL_50_21 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-22.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-22.txt deleted file mode 100644 index fe09be37a12..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-22.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 22 token E2E_TOTAL_50_22 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-23.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-23.txt deleted file mode 100644 index 70bb91ee326..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-23.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 23 token E2E_TOTAL_50_23 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-24.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-24.txt deleted file mode 100644 index a8ba2e60d61..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-24.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 24 token E2E_TOTAL_50_24 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-25.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-25.txt deleted file mode 100644 index ddfe2afed46..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-25.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 25 token E2E_TOTAL_50_25 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-26.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-26.txt deleted file mode 100644 index 1a56d928a9a..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-26.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 26 token E2E_TOTAL_50_26 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-27.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-27.txt deleted file mode 100644 index a8182ad2e71..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-27.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 27 token E2E_TOTAL_50_27 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-28.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-28.txt deleted file mode 100644 index a6f78f011db..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-28.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 28 token E2E_TOTAL_50_28 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-29.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-29.txt deleted file mode 100644 index 7a76d201194..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-29.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 29 token E2E_TOTAL_50_29 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-30.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-30.txt deleted file mode 100644 index 5a37417ca8c..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-30.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 30 token E2E_TOTAL_50_30 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-31.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-31.txt deleted file mode 100644 index 1baa8eb8c7b..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-31.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 31 token E2E_TOTAL_50_31 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-32.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-32.txt deleted file mode 100644 index e5ff1a42437..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-32.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 32 token E2E_TOTAL_50_32 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-33.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-33.txt deleted file mode 100644 index dcbc2995748..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-33.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 33 token E2E_TOTAL_50_33 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-34.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-34.txt deleted file mode 100644 index fefc133eed1..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-34.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 34 token E2E_TOTAL_50_34 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-35.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-35.txt deleted file mode 100644 index 1c1a93172cb..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-35.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 35 token E2E_TOTAL_50_35 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-36.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-36.txt deleted file mode 100644 index d7e0bab1aaf..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-36.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 36 token E2E_TOTAL_50_36 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-37.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-37.txt deleted file mode 100644 index 6a78127243f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-37.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 37 token E2E_TOTAL_50_37 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-38.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-38.txt deleted file mode 100644 index fc5ef9123db..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-38.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 38 token E2E_TOTAL_50_38 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-39.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-39.txt deleted file mode 100644 index 37765df65e5..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-39.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 39 token E2E_TOTAL_50_39 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-40.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-40.txt deleted file mode 100644 index 74de5e2ed4e..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-40.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 40 token E2E_TOTAL_50_40 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-41.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-41.txt deleted file mode 100644 index 13d3dee7c51..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-41.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 41 token E2E_TOTAL_50_41 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-42.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-42.txt deleted file mode 100644 index befc45f2387..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-42.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 42 token E2E_TOTAL_50_42 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-43.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-43.txt deleted file mode 100644 index 521ad583a5a..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-43.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 43 token E2E_TOTAL_50_43 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-44.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-44.txt deleted file mode 100644 index e4b315b675f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-44.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 44 token E2E_TOTAL_50_44 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-45.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-45.txt deleted file mode 100644 index c97738e4fac..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-45.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 45 token E2E_TOTAL_50_45 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-46.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-46.txt deleted file mode 100644 index 22b73767b42..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-46.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 46 token E2E_TOTAL_50_46 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-47.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-47.txt deleted file mode 100644 index b327026efd0..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-47.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 47 token E2E_TOTAL_50_47 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-48.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-48.txt deleted file mode 100644 index 5458e0cf222..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-48.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 48 token E2E_TOTAL_50_48 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-49.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-49.txt deleted file mode 100644 index 90bf454dd5f..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-49.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 49 token E2E_TOTAL_50_49 diff --git a/e2e/fixtures/test-materials/count_total_50_valid_files/file-50.txt b/e2e/fixtures/test-materials/count_total_50_valid_files/file-50.txt deleted file mode 100644 index b5760aed272..00000000000 --- a/e2e/fixtures/test-materials/count_total_50_valid_files/file-50.txt +++ /dev/null @@ -1 +0,0 @@ -Total 50 valid file 50 token E2E_TOTAL_50_50 diff --git a/e2e/fixtures/test-materials/count_total_extra_1_valid_file/file-01.txt b/e2e/fixtures/test-materials/count_total_extra_1_valid_file/file-01.txt deleted file mode 100644 index 4df2e230c05..00000000000 --- a/e2e/fixtures/test-materials/count_total_extra_1_valid_file/file-01.txt +++ /dev/null @@ -1 +0,0 @@ -Total extra valid file token E2E_TOTAL_EXTRA_1 diff --git a/e2e/fixtures/test-materials/file_tree_fixture/assets/sample.csv b/e2e/fixtures/test-materials/file_tree_fixture/assets/sample.csv deleted file mode 100644 index bed40b0fc09..00000000000 --- a/e2e/fixtures/test-materials/file_tree_fixture/assets/sample.csv +++ /dev/null @@ -1,3 +0,0 @@ -name,value -alpha,1 -beta,2 diff --git a/e2e/fixtures/test-materials/file_tree_fixture/docs/中文说明.md b/e2e/fixtures/test-materials/file_tree_fixture/docs/中文说明.md deleted file mode 100644 index 9d82b4cabd1..00000000000 --- a/e2e/fixtures/test-materials/file_tree_fixture/docs/中文说明.md +++ /dev/null @@ -1,3 +0,0 @@ -# 中文说明 - -文件树中文说明 token: E2E_FILE_TREE_ZH diff --git a/e2e/fixtures/test-materials/file_tree_fixture/public/index.html b/e2e/fixtures/test-materials/file_tree_fixture/public/index.html deleted file mode 100644 index babd6315f7a..00000000000 --- a/e2e/fixtures/test-materials/file_tree_fixture/public/index.html +++ /dev/null @@ -1,6 +0,0 @@ - - - - E2E_FILE_TREE_INDEX - - diff --git a/e2e/fixtures/test-materials/file_tree_fixture/src/main.txt b/e2e/fixtures/test-materials/file_tree_fixture/src/main.txt deleted file mode 100644 index 2f52ebf5950..00000000000 --- a/e2e/fixtures/test-materials/file_tree_fixture/src/main.txt +++ /dev/null @@ -1 +0,0 @@ -Main source fixture token: E2E_FILE_TREE_MAIN diff --git a/e2e/fixtures/test-materials/file_tree_fixture/web-game/README.md b/e2e/fixtures/test-materials/file_tree_fixture/web-game/README.md deleted file mode 100644 index 3581e8ee127..00000000000 --- a/e2e/fixtures/test-materials/file_tree_fixture/web-game/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Web Game Fixture - -Expected token: E2E_FILE_TREE_README diff --git a/e2e/package.json b/e2e/package.json index c0804956f31..61120b1ba83 100644 --- a/e2e/package.json +++ b/e2e/package.json @@ -14,7 +14,7 @@ "e2e:middleware:up": "tsx ./scripts/setup.ts middleware-up", "e2e:reset": "tsx ./scripts/setup.ts reset", "seed": "tsx ./scripts/seed.ts", - "type-check": "tsgo" + "type-check": "tsc" }, "devDependencies": { "@cucumber/cucumber": "catalog:", @@ -23,7 +23,7 @@ "@playwright/test": "catalog:", "@t3-oss/env-core": "catalog:", "@types/node": "catalog:", - "@typescript/native-preview": "catalog:", + "@typescript/native": "catalog:", "tsx": "catalog:", "typescript": "catalog:", "vite": "catalog:", diff --git a/e2e/scripts/setup.ts b/e2e/scripts/setup.ts index 0654772142f..46f9362ce56 100644 --- a/e2e/scripts/setup.ts +++ b/e2e/scripts/setup.ts @@ -32,14 +32,17 @@ const webBuildStampPath = path.join(webDir, '.next', 'e2e-web-build.sha256') const apiHost = '127.0.0.1' const apiPort = 5001 const agentBackendHost = '127.0.0.1' +const agentBackendBindHost = '0.0.0.0' const agentBackendPort = Number(process.env.E2E_AGENT_BACKEND_PORT || 5050) const shellctlHost = '127.0.0.1' const shellctlPort = Number(process.env.E2E_SHELLCTL_PORT || 5004) const shellctlContainerName = process.env.E2E_SHELLCTL_CONTAINER_NAME || 'dify-agent-e2e-shellctl' const shellctlImage = process.env.E2E_SHELLCTL_IMAGE || 'dify-agent-local-sandbox:e2e' const shellctlUrl = `http://${shellctlHost}:${shellctlPort}` +const agentStubApiBaseUrl = `http://host.docker.internal:${agentBackendPort}/agent-stub` const defaultPluginDaemonKey = 'lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi' const defaultInnerApiKeyForPlugin = 'QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1' +const defaultAgentServerSecretKey = 'MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY' const middlewareDataPaths = [ path.join(dockerDir, 'volumes', 'db', 'data'), @@ -104,6 +107,12 @@ const getAgentBackendEnvironment = async () => { || apiEnv.INNER_API_KEY_FOR_PLUGIN || defaultInnerApiKeyForPlugin, DIFY_AGENT_INNER_API_URL: process.env.DIFY_AGENT_INNER_API_URL || `http://${apiHost}:${apiPort}`, + DIFY_AGENT_SERVER_SECRET_KEY: + process.env.DIFY_AGENT_SERVER_SECRET_KEY + || defaultAgentServerSecretKey, + DIFY_AGENT_STUB_API_BASE_URL: + process.env.DIFY_AGENT_STUB_API_BASE_URL + || agentStubApiBaseUrl, DIFY_AGENT_PLUGIN_DAEMON_API_KEY: process.env.DIFY_AGENT_PLUGIN_DAEMON_API_KEY || process.env.PLUGIN_DAEMON_KEY @@ -358,7 +367,7 @@ export const startAgentBackend = async () => { 'uvicorn', 'dify_agent.server.app:app', '--host', - agentBackendHost, + agentBackendBindHost, '--port', String(agentBackendPort), ], @@ -420,6 +429,7 @@ export const startShellctlSandbox = async () => { '--rm', '--name', shellctlContainerName, + ...(process.platform === 'linux' ? ['--add-host', 'host.docker.internal:host-gateway'] : []), '-p', `${shellctlHost}:${shellctlPort}:5004`, ...(process.env.E2E_SHELLCTL_AUTH_TOKEN diff --git a/e2e/support/apps.ts b/e2e/support/apps.ts index 3c3af547a35..07d0f15716a 100644 --- a/e2e/support/apps.ts +++ b/e2e/support/apps.ts @@ -1,10 +1,15 @@ import type { Page } from '@playwright/test' import { expect } from '@playwright/test' +const getExpectOptions = (timeout?: number) => + timeout === undefined ? undefined : { timeout } + export const waitForAppsConsole = async (page: Page, timeout?: number) => { - await expect(page).toHaveURL(/\/apps(?:\?.*)?$/, timeout === undefined ? undefined : { timeout }) + const options = getExpectOptions(timeout) + + await expect(page).toHaveURL(/\/apps(?:\?.*)?$/, options) await expect(page.getByRole('heading', { name: 'Studio' })).toBeVisible( - timeout === undefined ? undefined : { timeout }, + options, ) } diff --git a/e2e/support/home.ts b/e2e/support/home.ts new file mode 100644 index 00000000000..83fe9b475c0 --- /dev/null +++ b/e2e/support/home.ts @@ -0,0 +1,12 @@ +import type { Page } from '@playwright/test' +import { expect } from '@playwright/test' + +const getExpectOptions = (timeout?: number) => + timeout === undefined ? undefined : { timeout } + +export const waitForConsoleHome = async (page: Page, timeout?: number) => { + const options = getExpectOptions(timeout) + + await expect.poll(() => new URL(page.url()).pathname, options).toBe('/') + await expect(page.getByRole('link', { name: 'Home' })).toHaveAttribute('aria-current', 'page', options) +} diff --git a/eslint-suppressions.json b/eslint-suppressions.json index e4b92e2db43..e47650125fe 100644 --- a/eslint-suppressions.json +++ b/eslint-suppressions.json @@ -279,11 +279,6 @@ "count": 1 } }, - "web/app/account/(commonLayout)/delete-account/components/check-email.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/account/(commonLayout)/delete-account/components/verify-email.tsx": { "no-restricted-imports": { "count": 1 @@ -292,11 +287,6 @@ "count": 1 } }, - "web/app/account/oauth/authorize/layout.tsx": { - "ts/no-explicit-any": { - "count": 1 - } - }, "web/app/account/oauth/authorize/page.tsx": { "ts/no-explicit-any": { "count": 1 @@ -320,14 +310,6 @@ "count": 4 } }, - "web/app/components/app-sidebar/dataset-info/__tests__/index.spec.tsx": { - "jsx-a11y/click-events-have-key-events": { - "count": 1 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 1 - } - }, "web/app/components/app-sidebar/dataset-info/menu-item.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 1 @@ -1201,7 +1183,7 @@ "count": 2 }, "ts/no-explicit-any": { - "count": 17 + "count": 12 } }, "web/app/components/base/chat/chat/log/index.tsx": { @@ -2041,11 +2023,6 @@ "count": 2 } }, - "web/app/components/base/permission-selector/index.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/app/components/base/permission-selector/member-item.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 1 @@ -3248,14 +3225,6 @@ "count": 2 } }, - "web/app/components/develop/secret-key/secret-key-modal.tsx": { - "jsx-a11y/click-events-have-key-events": { - "count": 1 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 1 - } - }, "web/app/components/explore/banner/__tests__/indicator-button.spec.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 1 @@ -3406,30 +3375,11 @@ "count": 2 } }, - "web/app/components/header/account-setting/members-page/edit-workspace-modal/index.tsx": { - "jsx-a11y/no-autofocus": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "web/app/components/header/account-setting/members-page/invite-modal/index.tsx": { "jsx-a11y/no-autofocus": { "count": 1 } }, - "web/app/components/header/account-setting/members-page/transfer-ownership-modal/index.tsx": { - "erasable-syntax-only/enums": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - }, - "ts/no-explicit-any": { - "count": 2 - } - }, "web/app/components/header/account-setting/members-page/transfer-ownership-modal/member-selector.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 1 @@ -3457,14 +3407,6 @@ "count": 2 } }, - "web/app/components/header/account-setting/model-provider-page/model-auth/__tests__/switch-credential-in-load-balancing.spec.tsx": { - "jsx-a11y/click-events-have-key-events": { - "count": 1 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 1 - } - }, "web/app/components/header/account-setting/model-provider-page/model-auth/add-custom-model.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 2 @@ -3579,14 +3521,6 @@ "count": 2 } }, - "web/app/components/header/account-setting/model-provider-page/provider-added-card/model-list.tsx": { - "jsx-a11y/click-events-have-key-events": { - "count": 1 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 1 - } - }, "web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-configs.tsx": { "jsx-a11y/click-events-have-key-events": { "count": 2 @@ -4187,11 +4121,6 @@ "count": 1 } }, - "web/app/components/rag-pipeline/components/rag-pipeline-main.tsx": { - "ts/no-explicit-any": { - "count": 2 - } - }, "web/app/components/rag-pipeline/hooks/index.ts": { "no-barrel-files/no-barrel-files": { "count": 9 @@ -4480,11 +4409,6 @@ "count": 1 } }, - "web/app/components/workflow-app/hooks/use-workflow-init.ts": { - "ts/no-explicit-any": { - "count": 3 - } - }, "web/app/components/workflow-app/hooks/use-workflow-refresh-draft.ts": { "ts/no-explicit-any": { "count": 2 @@ -4510,11 +4434,6 @@ "count": 2 } }, - "web/app/components/workflow-app/index.tsx": { - "ts/no-explicit-any": { - "count": 1 - } - }, "web/app/components/workflow-app/store/workflow/workflow-slice.ts": { "ts/no-explicit-any": { "count": 2 @@ -6906,11 +6825,6 @@ "count": 1 } }, - "web/features/tag-management/components/tag-management-modal.tsx": { - "jsx-a11y/no-autofocus": { - "count": 1 - } - }, "web/hooks/use-async-window-open.spec.ts": { "ts/no-explicit-any": { "count": 6 @@ -6996,11 +6910,6 @@ "count": 1 } }, - "web/scripts/component-analyzer.js": { - "regexp/no-unused-capturing-group": { - "count": 6 - } - }, "web/service/__tests__/base.spec.ts": { "no-restricted-imports": { "count": 1 @@ -7031,11 +6940,6 @@ "count": 1 } }, - "web/service/access-control/__tests__/use-permission-keys.spec.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "web/service/access-control/__tests__/use-workspace-access-rules.spec.tsx": { "no-restricted-imports": { "count": 1 @@ -7056,11 +6960,6 @@ "count": 1 } }, - "web/service/access-control/use-permission-keys.ts": { - "no-restricted-imports": { - "count": 1 - } - }, "web/service/access-control/use-workspace-access-rules.ts": { "no-restricted-imports": { "count": 1 @@ -7228,17 +7127,6 @@ "count": 1 } }, - "web/service/use-common.ts": { - "no-restricted-imports": { - "count": 1 - }, - "ts/no-empty-object-type": { - "count": 1 - }, - "ts/no-explicit-any": { - "count": 1 - } - }, "web/service/use-datasource.ts": { "no-restricted-imports": { "count": 1 diff --git a/package.json b/package.json index 3923b3f257e..0f2da51474a 100644 --- a/package.json +++ b/package.json @@ -24,11 +24,13 @@ }, "devDependencies": { "@antfu/eslint-config": "catalog:", + "@typescript/native": "catalog:", "concurrently": "catalog:", "eslint": "catalog:", "eslint-markdown": "catalog:", "eslint-plugin-markdown-preferences": "catalog:", "eslint-plugin-no-barrel-files": "catalog:", + "typescript": "catalog:", "vite": "catalog:", "vite-plus": "catalog:" } diff --git a/packages/contracts/generated/api/console/account/types.gen.ts b/packages/contracts/generated/api/console/account/types.gen.ts index 1a7d1644f8f..bfa12122573 100644 --- a/packages/contracts/generated/api/console/account/types.gen.ts +++ b/packages/contracts/generated/api/console/account/types.gen.ts @@ -87,10 +87,6 @@ export type EducationActivatePayload = { token: string } -export type EducationActivateResponse = { - [key: string]: unknown -} - export type EducationAutocompleteResponse = { curr_page?: number | null data?: Array @@ -301,7 +297,9 @@ export type PostAccountEducationData = { } export type PostAccountEducationResponses = { - 200: EducationActivateResponse + 200: { + [key: string]: unknown + } } export type PostAccountEducationResponse diff --git a/packages/contracts/generated/api/console/account/zod.gen.ts b/packages/contracts/generated/api/console/account/zod.gen.ts index cba539b0316..6805df133b2 100644 --- a/packages/contracts/generated/api/console/account/zod.gen.ts +++ b/packages/contracts/generated/api/console/account/zod.gen.ts @@ -127,11 +127,6 @@ export const zEducationActivatePayload = z.object({ token: z.string(), }) -/** - * EducationActivateResponse - */ -export const zEducationActivateResponse = z.record(z.string(), z.unknown()) - /** * EducationAutocompleteResponse */ @@ -301,7 +296,7 @@ export const zPostAccountEducationBody = zEducationActivatePayload /** * Success */ -export const zPostAccountEducationResponse = zEducationActivateResponse +export const zPostAccountEducationResponse = z.record(z.string(), z.unknown()) export const zGetAccountEducationAutocompleteQuery = z.object({ keywords: z.string(), diff --git a/packages/contracts/generated/api/console/agent/orpc.gen.ts b/packages/contracts/generated/api/console/agent/orpc.gen.ts index e3671d8cf98..e6249427ea0 100644 --- a/packages/contracts/generated/api/console/agent/orpc.gen.ts +++ b/packages/contracts/generated/api/console/agent/orpc.gen.ts @@ -1178,11 +1178,11 @@ export const read = { } /** - * Upload one Agent App sandbox file as a Dify ToolFile mapping + * Upload one Agent App sandbox file and return a signed download URL */ export const post16 = oc .route({ - description: 'Upload one Agent App sandbox file as a Dify ToolFile mapping', + description: 'Upload one Agent App sandbox file and return a signed download URL', inputStructure: 'detailed', method: 'POST', operationId: 'postAgentByAgentIdSandboxFilesUpload', diff --git a/packages/contracts/generated/api/console/agent/types.gen.ts b/packages/contracts/generated/api/console/agent/types.gen.ts index 5a28df41715..3f3bee9fdd2 100644 --- a/packages/contracts/generated/api/console/agent/types.gen.ts +++ b/packages/contracts/generated/api/console/agent/types.gen.ts @@ -442,8 +442,7 @@ export type AgentSandboxUploadPayload = { } export type SandboxUploadResponse = { - file: SandboxToolFileResponse - path: string + url: string } export type AgentSkillUploadResponse = { @@ -922,6 +921,7 @@ export type AgentLogMessageItemResponse = { } export type AgentThought = { + answer?: string | null chain_id?: string | null created_at?: number | null files: Array @@ -1008,11 +1008,6 @@ export type SandboxFileEntryResponse = { type: 'dir' | 'file' | 'other' | 'symlink' } -export type SandboxToolFileResponse = { - reference: string - transfer_method?: 'tool_file' -} - export type SkillManifest = { description: string entry_path: string @@ -1116,6 +1111,7 @@ export type AgentSource = 'agent_app' | 'imported' | 'roster' | 'system' | 'work export type AgentStatus = 'active' | 'archived' export type AgentSoulAppFeaturesConfig = { + file_upload?: AgentFileUploadFeatureConfig opening_statement?: string | null retriever_resource?: AgentFeatureToggleConfig | null sensitive_word_avoidance?: AgentSensitiveWordAvoidanceFeatureConfig | null @@ -1428,6 +1424,16 @@ export type AgentConfigRevisionOperation | 'save_new_version' | 'save_to_roster' +export type AgentFileUploadFeatureConfig = { + allowed_file_extensions?: Array + allowed_file_types?: Array + allowed_file_upload_methods?: Array + enabled?: boolean + image?: AgentFileUploadImageFeatureConfig + number_limits?: number + [key: string]: unknown +} + export type AgentSecretRefConfig = { credential_id?: string | null env_name?: string | null @@ -1679,6 +1685,15 @@ export type FormInputConfig export type JsonValue2 = unknown +export type FileType = 'audio' | 'custom' | 'document' | 'image' | 'video' + +export type FileTransferMethod = 'datasource_file' | 'local_file' | 'remote_url' | 'tool_file' + +export type AgentFileUploadImageFeatureConfig = { + enabled?: boolean + [key: string]: unknown +} + export type AgentKnowledgeDatasetConfig = { description?: string | null id?: string | null @@ -1801,10 +1816,6 @@ export type StringListSource = { value?: Array } -export type FileType = 'audio' | 'custom' | 'document' | 'image' | 'video' - -export type FileTransferMethod = 'datasource_file' | 'local_file' | 'remote_url' | 'tool_file' - export type AgentKnowledgeMetadataCondition = { comparison_operator: | '<' diff --git a/packages/contracts/generated/api/console/agent/zod.gen.ts b/packages/contracts/generated/api/console/agent/zod.gen.ts index 0ae9fd0dae6..d62307a61d0 100644 --- a/packages/contracts/generated/api/console/agent/zod.gen.ts +++ b/packages/contracts/generated/api/console/agent/zod.gen.ts @@ -212,6 +212,13 @@ export const zAgentSandboxUploadPayload = z.object({ path: z.string().min(1), }) +/** + * SandboxUploadResponse + */ +export const zSandboxUploadResponse = z.object({ + url: z.string(), +}) + /** * AgentConfigSnapshotRestoreResponse */ @@ -795,6 +802,7 @@ export const zJsonValue = z * AgentThought */ export const zAgentThought = z.object({ + answer: z.string().nullish(), chain_id: z.string().nullish(), created_at: z.int().nullish(), files: z.array(z.string()), @@ -865,22 +873,6 @@ export const zSandboxListResponse = z.object({ truncated: z.boolean().optional().default(false), }) -/** - * SandboxToolFileResponse - */ -export const zSandboxToolFileResponse = z.object({ - reference: z.string(), - transfer_method: z.literal('tool_file').optional().default('tool_file'), -}) - -/** - * SandboxUploadResponse - */ -export const zSandboxUploadResponse = z.object({ - file: zSandboxToolFileResponse, - path: z.string(), -}) - /** * SkillManifest * @@ -1926,19 +1918,6 @@ export const zAgentAppFeaturesPayload = z.object({ text_to_speech: zAgentTextToSpeechFeatureConfig.nullish(), }) -/** - * AgentSoulAppFeaturesConfig - */ -export const zAgentSoulAppFeaturesConfig = z.object({ - opening_statement: z.string().nullish(), - retriever_resource: zAgentFeatureToggleConfig.nullish(), - sensitive_word_avoidance: zAgentSensitiveWordAvoidanceFeatureConfig.nullish(), - speech_to_text: zAgentFeatureToggleConfig.nullish(), - suggested_questions: z.array(z.string()).nullish(), - suggested_questions_after_answer: zAgentSuggestedQuestionsAfterAnswerFeatureConfig.nullish(), - text_to_speech: zAgentTextToSpeechFeatureConfig.nullish(), -}) - export const zJsonValue2 = z.unknown() /** @@ -1953,6 +1932,54 @@ export const zHumanInputFormSubmissionData = z.object({ submitted_data: z.record(z.string(), zJsonValue2).nullish(), }) +/** + * FileType + */ +export const zFileType = z.enum(['audio', 'custom', 'document', 'image', 'video']) + +/** + * FileTransferMethod + */ +export const zFileTransferMethod = z.enum([ + 'datasource_file', + 'local_file', + 'remote_url', + 'tool_file', +]) + +/** + * AgentFileUploadImageFeatureConfig + */ +export const zAgentFileUploadImageFeatureConfig = z.object({ + enabled: z.boolean().optional().default(true), +}) + +/** + * AgentFileUploadFeatureConfig + */ +export const zAgentFileUploadFeatureConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + enabled: z.boolean().optional().default(true), + image: zAgentFileUploadImageFeatureConfig.optional(), + number_limits: z.int().optional().default(3), +}) + +/** + * AgentSoulAppFeaturesConfig + */ +export const zAgentSoulAppFeaturesConfig = z.object({ + file_upload: zAgentFileUploadFeatureConfig.optional(), + opening_statement: z.string().nullish(), + retriever_resource: zAgentFeatureToggleConfig.nullish(), + sensitive_word_avoidance: zAgentSensitiveWordAvoidanceFeatureConfig.nullish(), + speech_to_text: zAgentFeatureToggleConfig.nullish(), + suggested_questions: z.array(z.string()).nullish(), + suggested_questions_after_answer: zAgentSuggestedQuestionsAfterAnswerFeatureConfig.nullish(), + text_to_speech: zAgentTextToSpeechFeatureConfig.nullish(), +}) + /** * AgentKnowledgeDatasetConfig */ @@ -2182,6 +2209,29 @@ export const zUserActionConfig = z.object({ title: z.string().max(100), }) +/** + * FileInputConfig + */ +export const zFileInputConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + output_variable_name: z.string(), + type: z.literal('file').optional().default('file'), +}) + +/** + * FileListInputConfig + */ +export const zFileListInputConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + number_limits: z.int().gte(0).optional().default(0), + output_variable_name: z.string(), + type: z.literal('file-list').optional().default('file-list'), +}) + /** * AgentKnowledgeModelConfig */ @@ -2204,8 +2254,9 @@ export const zAgentKnowledgeQueryMode = z.enum(['generated_query', 'user_query'] * * Agent v2 stores knowledge as explicit ``knowledge.sets`` rather than the * legacy flat ``datasets`` / ``query_mode`` / ``query_config`` shape. Each - * set owns its own query policy, so ``user_query`` must carry an explicit - * ``value`` while ``generated_query`` leaves that value empty. + * set owns its own query policy. Mode-dependent completeness, such as + * requiring ``value`` for ``user_query``, is enforced by composer publish + * validation so draft saves can persist partially configured knowledge sets. */ export const zAgentKnowledgeQueryConfig = z.object({ mode: zAgentKnowledgeQueryMode, @@ -2235,8 +2286,9 @@ export const zAgentKnowledgeWeightedScoreConfig = z.object({ * Per-set retrieval policy for Agent v2 knowledge retrieval. * * Retrieval settings now live on each knowledge set instead of one shared - * flat config. A set may use either ``multiple`` retrieval with ``top_k`` or - * ``single`` retrieval with a required model config. + * flat config. Mode-dependent completeness, such as requiring ``top_k`` for + * ``multiple`` or a model for ``single``, is enforced by composer publish + * validation so draft saves can persist partially configured knowledge sets. */ export const zAgentKnowledgeRetrievalConfig = z.object({ mode: z.enum(['multiple', 'single']), @@ -2249,44 +2301,6 @@ export const zAgentKnowledgeRetrievalConfig = z.object({ weights: zAgentKnowledgeWeightedScoreConfig.nullish(), }) -/** - * FileType - */ -export const zFileType = z.enum(['audio', 'custom', 'document', 'image', 'video']) - -/** - * FileTransferMethod - */ -export const zFileTransferMethod = z.enum([ - 'datasource_file', - 'local_file', - 'remote_url', - 'tool_file', -]) - -/** - * FileInputConfig - */ -export const zFileInputConfig = z.object({ - allowed_file_extensions: z.array(z.string()).optional(), - allowed_file_types: z.array(zFileType).optional(), - allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), - output_variable_name: z.string(), - type: z.literal('file').optional().default('file'), -}) - -/** - * FileListInputConfig - */ -export const zFileListInputConfig = z.object({ - allowed_file_extensions: z.array(z.string()).optional(), - allowed_file_types: z.array(zFileType).optional(), - allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), - number_limits: z.int().gte(0).optional().default(0), - output_variable_name: z.string(), - type: z.literal('file-list').optional().default('file-list'), -}) - /** * AgentKnowledgeMetadataCondition */ @@ -2331,6 +2345,8 @@ export const zAgentKnowledgeMetadataConditions = z.object({ * The Python attribute uses ``metadata_model_config`` for clarity because the * model belongs to metadata filtering specifically, while the external API and * generated schema keep the historical ``model_config`` field name via alias. + * Mode-dependent completeness is enforced by composer publish validation so + * draft saves can persist partially configured metadata filters. */ export const zAgentKnowledgeMetadataFilteringConfig = z.object({ conditions: zAgentKnowledgeMetadataConditions.nullish(), diff --git a/packages/contracts/generated/api/console/all-workspaces/types.gen.ts b/packages/contracts/generated/api/console/all-workspaces/types.gen.ts index 4683b2d9921..b4c9fa0349a 100644 --- a/packages/contracts/generated/api/console/all-workspaces/types.gen.ts +++ b/packages/contracts/generated/api/console/all-workspaces/types.gen.ts @@ -4,7 +4,7 @@ export type ClientOptions = { baseUrl: `${string}://${string}/console/api` | (string & {}) } -export type WorkspaceListResponse = { +export type WorkspacePaginationResponse = { data: Array has_more: boolean limit: number @@ -30,7 +30,7 @@ export type GetAllWorkspacesData = { } export type GetAllWorkspacesResponses = { - 200: WorkspaceListResponse + 200: WorkspacePaginationResponse } export type GetAllWorkspacesResponse = GetAllWorkspacesResponses[keyof GetAllWorkspacesResponses] diff --git a/packages/contracts/generated/api/console/all-workspaces/zod.gen.ts b/packages/contracts/generated/api/console/all-workspaces/zod.gen.ts index f63bd0e396f..c9cdda11681 100644 --- a/packages/contracts/generated/api/console/all-workspaces/zod.gen.ts +++ b/packages/contracts/generated/api/console/all-workspaces/zod.gen.ts @@ -13,9 +13,9 @@ export const zWorkspaceListItemResponse = z.object({ }) /** - * WorkspaceListResponse + * WorkspacePaginationResponse */ -export const zWorkspaceListResponse = z.object({ +export const zWorkspacePaginationResponse = z.object({ data: z.array(zWorkspaceListItemResponse), has_more: z.boolean(), limit: z.int(), @@ -31,4 +31,4 @@ export const zGetAllWorkspacesQuery = z.object({ /** * Success */ -export const zGetAllWorkspacesResponse = zWorkspaceListResponse +export const zGetAllWorkspacesResponse = zWorkspacePaginationResponse diff --git a/packages/contracts/generated/api/console/apps/orpc.gen.ts b/packages/contracts/generated/api/console/apps/orpc.gen.ts index 83637302055..2b0d927e72d 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -3029,11 +3029,11 @@ export const read = { } /** - * Upload one workflow Agent sandbox file as a Dify ToolFile mapping + * Upload one workflow Agent sandbox file and return a signed download URL */ export const post41 = oc .route({ - description: 'Upload one workflow Agent sandbox file as a Dify ToolFile mapping', + description: 'Upload one workflow Agent sandbox file and return a signed download URL', inputStructure: 'detailed', method: 'POST', operationId: 'postAppsByAppIdWorkflowRunsByWorkflowRunIdAgentNodesByNodeIdSandboxFilesUpload', diff --git a/packages/contracts/generated/api/console/apps/types.gen.ts b/packages/contracts/generated/api/console/apps/types.gen.ts index 21473d5bd88..f18eec3f8b9 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -873,8 +873,7 @@ export type WorkflowAgentSandboxUploadPayload = { } export type SandboxUploadResponse = { - file: SandboxToolFileResponse - path: string + url: string } export type WorkflowCommentBasicList = { @@ -1372,7 +1371,7 @@ export type ImportStatus = 'completed' | 'completed-with-warnings' | 'failed' | export type PluginDependency = { current_identifier?: string | null - type: Type + type: PluginDependencyType value: Github | Marketplace | Package } @@ -1649,6 +1648,7 @@ export type ConversationVariableResponse = { } export type AgentThought = { + answer?: string | null chain_id?: string | null created_at?: number | null files: Array @@ -1807,11 +1807,6 @@ export type SandboxFileEntryResponse = { type: 'dir' | 'file' | 'other' | 'symlink' } -export type SandboxToolFileResponse = { - reference: string - transfer_method?: 'tool_file' -} - export type WorkflowCommentBasic = { content: string created_at?: number | null @@ -2186,7 +2181,7 @@ export type ModelConfigPartial = { export type LlmMode = 'chat' | 'completion' -export type Type = 'github' | 'marketplace' | 'package' +export type PluginDependencyType = 'github' | 'marketplace' | 'package' export type Github = { github_plugin_unique_identifier: string @@ -2368,6 +2363,7 @@ export type AgentSource = 'agent_app' | 'imported' | 'roster' | 'system' | 'work export type AgentStatus = 'active' | 'archived' export type AgentSoulAppFeaturesConfig = { + file_upload?: AgentFileUploadFeatureConfig opening_statement?: string | null retriever_resource?: AgentFeatureToggleConfig | null sensitive_word_avoidance?: AgentSensitiveWordAvoidanceFeatureConfig | null @@ -2627,6 +2623,16 @@ export type WorkflowFileUploadPreviewConfigPayload = { mode?: string | null } +export type AgentFileUploadFeatureConfig = { + allowed_file_extensions?: Array + allowed_file_types?: Array + allowed_file_upload_methods?: Array + enabled?: boolean + image?: AgentFileUploadImageFeatureConfig + number_limits?: number + [key: string]: unknown +} + export type AgentFeatureToggleConfig = { enabled?: boolean [key: string]: unknown @@ -2873,6 +2879,15 @@ export type FileListInputConfig = { type?: 'file-list' } +export type FileType = 'audio' | 'custom' | 'document' | 'image' | 'video' + +export type FileTransferMethod = 'datasource_file' | 'local_file' | 'remote_url' | 'tool_file' + +export type AgentFileUploadImageFeatureConfig = { + enabled?: boolean + [key: string]: unknown +} + export type AgentModerationProviderConfig = { api_based_extension_id?: string | null inputs_config?: AgentModerationIoConfig | null @@ -2942,10 +2957,6 @@ export type StringListSource = { value?: Array } -export type FileType = 'audio' | 'custom' | 'document' | 'image' | 'video' - -export type FileTransferMethod = 'datasource_file' | 'local_file' | 'remote_url' | 'tool_file' - export type AgentModerationIoConfig = { enabled?: boolean preset_response?: string | null diff --git a/packages/contracts/generated/api/console/apps/zod.gen.ts b/packages/contracts/generated/api/console/apps/zod.gen.ts index 9a928ee68e4..91d6f6bd01e 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -564,6 +564,13 @@ export const zWorkflowAgentSandboxUploadPayload = z.object({ path: z.string().min(1), }) +/** + * SandboxUploadResponse + */ +export const zSandboxUploadResponse = z.object({ + url: z.string(), +}) + /** * WorkflowCommentCreatePayload */ @@ -1335,6 +1342,7 @@ export const zPaginatedConversationVariableResponse = z.object({ * AgentThought */ export const zAgentThought = z.object({ + answer: z.string().nullish(), chain_id: z.string().nullish(), created_at: z.int().nullish(), files: z.array(z.string()), @@ -1673,22 +1681,6 @@ export const zSandboxListResponse = z.object({ truncated: z.boolean().optional().default(false), }) -/** - * SandboxToolFileResponse - */ -export const zSandboxToolFileResponse = z.object({ - reference: z.string(), - transfer_method: z.literal('tool_file').optional().default('tool_file'), -}) - -/** - * SandboxUploadResponse - */ -export const zSandboxUploadResponse = z.object({ - file: zSandboxToolFileResponse, - path: z.string(), -}) - /** * AccountWithRoleResponse */ @@ -2336,9 +2328,9 @@ export const zConversationDetail = z.object({ }) /** - * Type + * PluginDependencyType */ -export const zType = z.enum(['github', 'marketplace', 'package']) +export const zPluginDependencyType = z.enum(['github', 'marketplace', 'package']) /** * Github @@ -2371,7 +2363,7 @@ export const zPackage = z.object({ */ export const zPluginDependency = z.object({ current_identifier: z.string().nullish(), - type: zType, + type: zPluginDependencyType, value: z.union([zGithub, zMarketplace, zPackage]), }) @@ -3466,6 +3458,63 @@ export const zUserActionConfig = z.object({ title: z.string().max(100), }) +/** + * FileType + */ +export const zFileType = z.enum(['audio', 'custom', 'document', 'image', 'video']) + +/** + * FileTransferMethod + */ +export const zFileTransferMethod = z.enum([ + 'datasource_file', + 'local_file', + 'remote_url', + 'tool_file', +]) + +/** + * FileInputConfig + */ +export const zFileInputConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + output_variable_name: z.string(), + type: z.literal('file').optional().default('file'), +}) + +/** + * FileListInputConfig + */ +export const zFileListInputConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + number_limits: z.int().gte(0).optional().default(0), + output_variable_name: z.string(), + type: z.literal('file-list').optional().default('file-list'), +}) + +/** + * AgentFileUploadImageFeatureConfig + */ +export const zAgentFileUploadImageFeatureConfig = z.object({ + enabled: z.boolean().optional().default(true), +}) + +/** + * AgentFileUploadFeatureConfig + */ +export const zAgentFileUploadFeatureConfig = z.object({ + allowed_file_extensions: z.array(z.string()).optional(), + allowed_file_types: z.array(zFileType).optional(), + allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), + enabled: z.boolean().optional().default(true), + image: zAgentFileUploadImageFeatureConfig.optional(), + number_limits: z.int().optional().default(3), +}) + /** * AgentSuggestedQuestionsAfterAnswerModelConfig * @@ -3662,44 +3711,6 @@ export const zAgentSoulToolsConfig = z.object({ dify_tools: z.array(zAgentSoulDifyToolConfig).optional(), }) -/** - * FileType - */ -export const zFileType = z.enum(['audio', 'custom', 'document', 'image', 'video']) - -/** - * FileTransferMethod - */ -export const zFileTransferMethod = z.enum([ - 'datasource_file', - 'local_file', - 'remote_url', - 'tool_file', -]) - -/** - * FileInputConfig - */ -export const zFileInputConfig = z.object({ - allowed_file_extensions: z.array(z.string()).optional(), - allowed_file_types: z.array(zFileType).optional(), - allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), - output_variable_name: z.string(), - type: z.literal('file').optional().default('file'), -}) - -/** - * FileListInputConfig - */ -export const zFileListInputConfig = z.object({ - allowed_file_extensions: z.array(z.string()).optional(), - allowed_file_types: z.array(zFileType).optional(), - allowed_file_upload_methods: z.array(zFileTransferMethod).optional(), - number_limits: z.int().gte(0).optional().default(0), - output_variable_name: z.string(), - type: z.literal('file-list').optional().default('file-list'), -}) - /** * AgentModerationIOConfig */ @@ -3731,6 +3742,7 @@ export const zAgentSensitiveWordAvoidanceFeatureConfig = z.object({ * AgentSoulAppFeaturesConfig */ export const zAgentSoulAppFeaturesConfig = z.object({ + file_upload: zAgentFileUploadFeatureConfig.optional(), opening_statement: z.string().nullish(), retriever_resource: zAgentFeatureToggleConfig.nullish(), sensitive_word_avoidance: zAgentSensitiveWordAvoidanceFeatureConfig.nullish(), @@ -3762,8 +3774,9 @@ export const zAgentKnowledgeQueryMode = z.enum(['generated_query', 'user_query'] * * Agent v2 stores knowledge as explicit ``knowledge.sets`` rather than the * legacy flat ``datasets`` / ``query_mode`` / ``query_config`` shape. Each - * set owns its own query policy, so ``user_query`` must carry an explicit - * ``value`` while ``generated_query`` leaves that value empty. + * set owns its own query policy. Mode-dependent completeness, such as + * requiring ``value`` for ``user_query``, is enforced by composer publish + * validation so draft saves can persist partially configured knowledge sets. */ export const zAgentKnowledgeQueryConfig = z.object({ mode: zAgentKnowledgeQueryMode, @@ -3793,8 +3806,9 @@ export const zAgentKnowledgeWeightedScoreConfig = z.object({ * Per-set retrieval policy for Agent v2 knowledge retrieval. * * Retrieval settings now live on each knowledge set instead of one shared - * flat config. A set may use either ``multiple`` retrieval with ``top_k`` or - * ``single`` retrieval with a required model config. + * flat config. Mode-dependent completeness, such as requiring ``top_k`` for + * ``multiple`` or a model for ``single``, is enforced by composer publish + * validation so draft saves can persist partially configured knowledge sets. */ export const zAgentKnowledgeRetrievalConfig = z.object({ mode: z.enum(['multiple', 'single']), @@ -3972,6 +3986,8 @@ export const zAgentKnowledgeMetadataConditions = z.object({ * The Python attribute uses ``metadata_model_config`` for clarity because the * model belongs to metadata filtering specifically, while the external API and * generated schema keep the historical ``model_config`` field name via alias. + * Mode-dependent completeness is enforced by composer publish validation so + * draft saves can persist partially configured metadata filters. */ export const zAgentKnowledgeMetadataFilteringConfig = z.object({ conditions: zAgentKnowledgeMetadataConditions.nullish(), diff --git a/packages/contracts/generated/api/console/auth/types.gen.ts b/packages/contracts/generated/api/console/auth/types.gen.ts index 99a06128e14..4951cbbf9a0 100644 --- a/packages/contracts/generated/api/console/auth/types.gen.ts +++ b/packages/contracts/generated/api/console/auth/types.gen.ts @@ -4,8 +4,12 @@ export type ClientOptions = { baseUrl: `${string}://${string}/console/api` | (string & {}) } -export type DatasourceCredentialsResponse = { - result: unknown +export type DatasourceProviderAuthListResponse = { + result: Array +} + +export type DatasourceCredentialListResponse = { + result: Array } export type DatasourceCredentialPayload = { @@ -47,6 +51,90 @@ export type DatasourceUpdateNamePayload = { name: string } +export type DatasourceProviderAuthResponse = { + author: string + credential_schema: Array + credentials_list: Array + description: I18nObject + icon: string + label: I18nObject + name: string + oauth_schema: DatasourceOAuthSchemaResponse | null + plugin_id: string + plugin_unique_identifier: string + provider: string +} + +export type DatasourceCredentialResponse = { + avatar_url: string | null + credential: { + [key: string]: unknown + } + id: string + is_default: boolean + name: string + type: string +} + +export type ProviderConfig = { + default?: number | string | number | boolean | null + help?: I18nObject | null + label?: I18nObject | null + multiple?: boolean + name: string + options?: Array