diff --git a/.agents/skills/frontend-code-review/references/dify-ui.md b/.agents/skills/frontend-code-review/references/dify-ui.md index 93484a6a0fc..7da481ef6e6 100644 --- a/.agents/skills/frontend-code-review/references/dify-ui.md +++ b/.agents/skills/frontend-code-review/references/dify-ui.md @@ -55,7 +55,8 @@ For non-string `Select` and `RadioGroup` values, prefer explicit domain generics Flag: - Form-like UI using unrelated `Input` and `Button` pieces without a submit boundary. -- Text-like fields not composed through `Field`, `FieldLabel`, and `FieldControl` when using Dify UI form semantics. +- Text-like fields not composed through `Field`, `FieldLabel`, and `Input` or `InputGroupInput` when using Dify UI form semantics. +- Prefixes, suffixes, or actions manually layered over `Input` instead of using the canonical `InputGroup` composition. - Select fields using `FieldLabel` instead of `SelectLabel`. - Slider fields using a generic label instead of `SliderLabel`. - Checkbox/radio groups missing `Fieldset` and `FieldsetLegend`. diff --git a/.github/workflows/accessibility-e2e.yml b/.github/workflows/accessibility-e2e.yml index 03a41f587ea..1840e0676b6 100644 --- a/.github/workflows/accessibility-e2e.yml +++ b/.github/workflows/accessibility-e2e.yml @@ -52,7 +52,7 @@ jobs: uses: ./.github/actions/setup-web - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: '3.12' diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index 54009a040e3..1b1afe96529 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -35,7 +35,7 @@ jobs: persist-credentials: false - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: ${{ matrix.python-version }} @@ -94,7 +94,7 @@ jobs: persist-credentials: false - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: ${{ matrix.python-version }} @@ -145,7 +145,7 @@ jobs: persist-credentials: false - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: '3.12' diff --git a/.github/workflows/autofix.yml b/.github/workflows/autofix.yml index 005326a6380..435d3cad3d7 100644 --- a/.github/workflows/autofix.yml +++ b/.github/workflows/autofix.yml @@ -90,7 +90,7 @@ jobs: python-version: '3.11' - if: github.event_name != 'merge_group' - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 - name: Generate Docker Compose if: github.event_name != 'merge_group' && steps.docker-compose-changes.outputs.any_changed == 'true' diff --git a/.github/workflows/build-push.yml b/.github/workflows/build-push.yml index 61d654bd375..dd570b9f4ef 100644 --- a/.github/workflows/build-push.yml +++ b/.github/workflows/build-push.yml @@ -229,3 +229,32 @@ jobs: IMAGE_VERSION: ${{ steps.meta.outputs.version }} run: | docker buildx imagetools inspect "$IMAGE_NAME:$IMAGE_VERSION" + + sync-e2b-template: + needs: create-manifest + runs-on: ubuntu-24.04 + if: github.repository == 'langgenius/dify' && startsWith(github.ref, 'refs/tags/') + strategy: + matrix: + include: + - project: dev + api_key_secret: E2B_API_KEY_DEV + - project: prod + api_key_secret: E2B_API_KEY_PROD + steps: + - name: Checkout + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 + + - name: Login to Docker Hub + uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0 + with: + username: ${{ env.DOCKERHUB_USER }} + password: ${{ env.DOCKERHUB_TOKEN }} + + - name: Sync E2B Template (${{ matrix.project }}) + env: + E2B_API_KEY: ${{ secrets[matrix.api_key_secret] }} + run: ./dify-agent-runtime/docker/sync-e2b-template.sh diff --git a/.github/workflows/db-migration-test.yml b/.github/workflows/db-migration-test.yml index c68db72d226..6d2b62f92f1 100644 --- a/.github/workflows/db-migration-test.yml +++ b/.github/workflows/db-migration-test.yml @@ -19,7 +19,7 @@ jobs: persist-credentials: false - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: '3.12' @@ -69,7 +69,7 @@ jobs: persist-credentials: false - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: '3.12' diff --git a/.github/workflows/pyrefly-diff.yml b/.github/workflows/pyrefly-diff.yml index 27b04f030e1..74ccd858cec 100644 --- a/.github/workflows/pyrefly-diff.yml +++ b/.github/workflows/pyrefly-diff.yml @@ -22,7 +22,7 @@ jobs: fetch-depth: 0 - name: Setup Python & UV - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true diff --git a/.github/workflows/pyrefly-type-coverage-comment.yml b/.github/workflows/pyrefly-type-coverage-comment.yml index eacf485c7a1..7039e06ed8f 100644 --- a/.github/workflows/pyrefly-type-coverage-comment.yml +++ b/.github/workflows/pyrefly-type-coverage-comment.yml @@ -24,7 +24,7 @@ jobs: uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Python & UV - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true diff --git a/.github/workflows/pyrefly-type-coverage.yml b/.github/workflows/pyrefly-type-coverage.yml index 19a2e18d48b..66944e24520 100644 --- a/.github/workflows/pyrefly-type-coverage.yml +++ b/.github/workflows/pyrefly-type-coverage.yml @@ -22,7 +22,7 @@ jobs: fetch-depth: 0 - name: Setup Python & UV - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index 843bc9a0e13..e2dc5504205 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -47,7 +47,7 @@ jobs: - name: Setup UV and Python if: steps.changed-files.outputs.any_changed == 'true' - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: false python-version: '3.12' diff --git a/.github/workflows/translate-i18n-claude.yml b/.github/workflows/translate-i18n-claude.yml index fe589fe04a5..283ec34ff98 100644 --- a/.github/workflows/translate-i18n-claude.yml +++ b/.github/workflows/translate-i18n-claude.yml @@ -162,7 +162,7 @@ jobs: - name: Run Claude Code for Translation Sync if: steps.context.outputs.CHANGED_FILES != '' - uses: anthropics/claude-code-action@e63208cb983318a44e3f945e959ef894b707dcfa # v1.0.192 + uses: anthropics/claude-code-action@9d7150bc8a3dae8149739a88019d192b579ad90c # v1.0.193 with: anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} github_token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/vdb-tests-full.yml b/.github/workflows/vdb-tests-full.yml index e5680416058..3240d97a0f5 100644 --- a/.github/workflows/vdb-tests-full.yml +++ b/.github/workflows/vdb-tests-full.yml @@ -36,7 +36,7 @@ jobs: remove_tool_cache: true - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: ${{ matrix.python-version }} diff --git a/.github/workflows/vdb-tests.yml b/.github/workflows/vdb-tests.yml index 0fa646ad8e5..fbbfab90713 100644 --- a/.github/workflows/vdb-tests.yml +++ b/.github/workflows/vdb-tests.yml @@ -33,7 +33,7 @@ jobs: remove_tool_cache: true - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: ${{ matrix.python-version }} diff --git a/.github/workflows/web-e2e.yml b/.github/workflows/web-e2e.yml index 536fe160d74..096553af0ce 100644 --- a/.github/workflows/web-e2e.yml +++ b/.github/workflows/web-e2e.yml @@ -35,7 +35,7 @@ jobs: uses: ./.github/actions/setup-web - name: Setup UV and Python - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: enable-cache: true python-version: '3.12' diff --git a/api/clients/agent_backend/client.py b/api/clients/agent_backend/client.py index 6738b33bcf4..297bb523c0e 100644 --- a/api/clients/agent_backend/client.py +++ b/api/clients/agent_backend/client.py @@ -23,6 +23,7 @@ from dify_agent.protocol import ( CancelRunResponse, CreateRunRequest, CreateRunResponse, + RunCancelledEvent, RunEvent, RunStatusResponse, ) @@ -45,6 +46,15 @@ class AgentBackendRunClient(Protocol): def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse: """Request explicit cancellation for one Agent backend run.""" + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Request cancellation and wait for runner cleanup to finish.""" + def stream_events( self, run_id: str, @@ -67,6 +77,15 @@ class _DifyAgentSyncClient(Protocol): def cancel_run_sync(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse: """Cancel one run synchronously.""" + def cancel_run_and_wait_sync( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Cancel one run and wait for its terminal event synchronously.""" + def stream_events_sync( self, run_id: str, @@ -110,6 +129,19 @@ class DifyAgentBackendRunClient: except Exception as exc: raise _normalize_dify_agent_error(exc) from exc + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Cancel one run, then wait for the cleanup-complete terminal event.""" + try: + return self.client.cancel_run_and_wait_sync(run_id, request=request, after=after) + except Exception as exc: + raise _normalize_dify_agent_error(exc) from exc + def stream_events( self, run_id: str, diff --git a/api/clients/agent_backend/event_adapter.py b/api/clients/agent_backend/event_adapter.py index e72d2fbfb51..14f52742d2c 100644 --- a/api/clients/agent_backend/event_adapter.py +++ b/api/clients/agent_backend/event_adapter.py @@ -99,6 +99,7 @@ class AgentBackendRunFailedInternalEvent(AgentBackendInternalEventBase): error: str error_type: RunFailureType | None = None reason: str | None = None + session_snapshot: CompositorSessionSnapshot | None = None class AgentBackendRunCancelledInternalEvent(AgentBackendInternalEventBase): @@ -107,6 +108,7 @@ class AgentBackendRunCancelledInternalEvent(AgentBackendInternalEventBase): type: Literal[AgentBackendInternalEventType.RUN_CANCELLED] = AgentBackendInternalEventType.RUN_CANCELLED reason: str | None = None message: str | None = None + session_snapshot: CompositorSessionSnapshot | None = None type AgentBackendInternalEvent = Annotated[ @@ -184,6 +186,7 @@ class AgentBackendRunEventAdapter: error=event.data.error, error_type=event.data.error_type, reason=event.data.reason, + session_snapshot=event.data.session_snapshot, ) ] case RunCancelledEvent(): @@ -193,6 +196,7 @@ class AgentBackendRunEventAdapter: source_event_id=event.id, reason=event.data.reason, message=event.data.message, + session_snapshot=event.data.session_snapshot, ) ] raise TypeError(f"unsupported agent backend run event: {type(event).__name__}") diff --git a/api/clients/agent_backend/fake_client.py b/api/clients/agent_backend/fake_client.py index 2d0881c03ff..c016f2ce6bf 100644 --- a/api/clients/agent_backend/fake_client.py +++ b/api/clients/agent_backend/fake_client.py @@ -18,6 +18,8 @@ from dify_agent.protocol import ( CreateRunRequest, CreateRunResponse, DeferredToolCallPayload, + RunCancelledEvent, + RunCancelledEventData, RunEvent, RunFailedEvent, RunFailedEventData, @@ -69,6 +71,28 @@ class FakeAgentBackendRunClient: del request return CancelRunResponse(run_id=run_id, status="cancelled") + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Return a deterministic cleanup-complete cancellation event.""" + del after + request = request or CancelRunRequest() + _ = self.cancel_run(run_id, request) + return RunCancelledEvent( + id="cancel-0", + run_id=run_id, + created_at=_FIXED_TIME, + data=RunCancelledEventData( + reason=request.reason, + message=request.message, + session_snapshot=CompositorSessionSnapshot(layers=[]), + ), + ) + def stream_events( self, run_id: str, @@ -133,7 +157,11 @@ class FakeAgentBackendRunClient: id="2-0", run_id=run_id, created_at=_FIXED_TIME, - data=RunFailedEventData(error="fake failure", reason="unit_test"), + data=RunFailedEventData( + error="fake failure", + reason="unit_test", + session_snapshot=CompositorSessionSnapshot(layers=[]), + ), ), ) case FakeAgentBackendScenario.PAUSED: diff --git a/api/clients/agent_backend/request_builder.py b/api/clients/agent_backend/request_builder.py index 984aeef2174..57cbd3be926 100644 --- a/api/clients/agent_backend/request_builder.py +++ b/api/clients/agent_backend/request_builder.py @@ -150,6 +150,7 @@ class AgentBackendModelConfig(BaseModel): model_provider: str model: str model_settings: dict[str, JsonValue] = Field(default_factory=dict) + context_window_tokens: int | None = Field(default=None, gt=0) model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") @@ -413,6 +414,7 @@ class AgentBackendRunRequestBuilder: model_provider=run_input.model.model_provider, model=run_input.model.model, model_settings=_agent_model_settings(run_input.model.model_settings), + context_window_tokens=run_input.model.context_window_tokens, ), ) ) @@ -605,6 +607,7 @@ class AgentBackendRunRequestBuilder: model_provider=run_input.model.model_provider, model=run_input.model.model, model_settings=_agent_model_settings(run_input.model.model_settings), + context_window_tokens=run_input.model.context_window_tokens, ), ), ] diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 0772b1d08d4..a4cefb4ba40 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1324,6 +1324,17 @@ class MultiModalTransferConfig(BaseSettings): ) +class NewAgentBetaConfig(BaseSettings): + NEW_AGENT_BETA_ACTIVITY_START_AT: datetime | None = Field( + description="New Agent Beta Publish window start in RFC3339 UTC (inclusive)", + default=None, + ) + NEW_AGENT_BETA_ACTIVITY_END_AT: datetime | None = Field( + description="New Agent Beta Publish window end in RFC3339 UTC (exclusive)", + default=None, + ) + + class OpsTraceConfig(BaseSettings): OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES: PositiveInt = Field( description="Maximum retry attempts for transient ops trace provider dispatch failures.", @@ -1641,6 +1652,7 @@ class FeatureConfig( ModelLoadBalanceConfig, ModerationConfig, MultiModalTransferConfig, + NewAgentBetaConfig, OpsTraceConfig, PositionConfig, RagEtlConfig, diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index b3377bb7019..c335fb852f4 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -184,13 +184,15 @@ class EmailRegisterResetApi(Resource): if account: raise EmailAlreadyInUseError() + ip_address = extract_remote_ip(request) account = self._create_new_account( email=normalized_email, password=req_data.password_confirm, timezone=req_data.timezone, language=req_data.language, + ip_address=ip_address, ) - 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=ip_address) AccountService.reset_login_error_rate_limit(normalized_email) return {"result": "success", "data": token_pair.model_dump()} @@ -201,6 +203,7 @@ class EmailRegisterResetApi(Resource): password: str, timezone: str | None = None, language: str | None = None, + ip_address: str | None = None, ) -> Account: try: return AccountService.create_account_and_tenant( @@ -209,6 +212,7 @@ class EmailRegisterResetApi(Resource): password=password, interface_language=get_valid_language(language), timezone=timezone, + ip_address=ip_address, session=db.session(), ) except SeatsLimitExceededError: diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 78d7038cc76..15558a0c146 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -346,6 +346,7 @@ class EmailCodeLoginApi(Resource): else: TenantService.create_owner_tenant(account, session=db.session()) + ip_address = extract_remote_ip(request) if account is None: try: account = AccountService.create_account_and_tenant( @@ -353,6 +354,7 @@ class EmailCodeLoginApi(Resource): name=user_email, interface_language=get_valid_language(language), timezone=req_data.timezone, + ip_address=ip_address, session=db.session(), ) except WorkSpaceNotAllowedCreateError: @@ -364,7 +366,7 @@ 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=ip_address) AccountService.reset_login_error_rate_limit(user_email) # Create response with cookies instead of returning tokens in body diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 0e1b28c58f7..41b94431f9a 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -233,7 +233,13 @@ class OAuthCallback(Resource): return _redirect_with_console_session(account, target_url) try: - account, oauth_new_user = _generate_account(provider, user_info, timezone=timezone, language=language) + account, oauth_new_user = _generate_account( + provider, + user_info, + timezone=timezone, + language=language, + ip_address=extract_remote_ip(request), + ) except AccountNotFoundError: return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message=Account not found.") except (WorkSpaceNotFoundError, WorkSpaceNotAllowedCreateError): @@ -285,6 +291,7 @@ def _generate_account( user_info: OAuthUserInfo, timezone: str | None = None, language: str | None = None, + ip_address: str | None = None, ) -> tuple[Account, bool]: # Get account by openid or email. account = _get_account_by_openid_or_email(provider, user_info) @@ -322,6 +329,7 @@ def _generate_account( provider=provider, language=interface_language, timezone=timezone, + ip_address=ip_address, session=db.session(), ) diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py index df87c4dd2ee..a28b44d0346 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py @@ -30,6 +30,7 @@ from graphon.model_runtime.errors.validate import CredentialsValidateFailedError from libs.helper import dump_response from libs.login import login_required from models import Account +from models.enums import PermissionEnum from models.provider_ids import DatasourceProviderID from services.datasource_provider_service import DatasourceProviderService from services.plugin.oauth_service import OAuthProxyService @@ -74,6 +75,13 @@ class DatasourceUpdateNamePayload(BaseModel): class DatasourceOAuthAuthorizationQuery(BaseModel): credential_id: str | None = Field(default=None, description="Credential ID to reauthorize") + visibility: str | None = Field( + default=None, + description=( + "Visibility for the credential to be created. Accepts 'only_me' or 'all_team_members'; " + "any other value falls back to 'only_me'. Ignored on reauthorization (credential_id set)." + ), + ) class DatasourceOAuthCallbackQuery(BaseModel): @@ -175,12 +183,27 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): if not oauth_config: raise ValueError(f"No OAuth Client Config for {provider_id}") + # Visibility is chosen by the user in the frontend before the redirect, + # then read back in the callback below when the credential is created. + # Only ONLY_ME / ALL_TEAM are accepted; anything else falls back to + # ONLY_ME (OAuth tokens are personal by nature). + # For reauthorization (credential_id set), visibility is ignored — we + # keep whatever the credential was created with. + raw_visibility = request.args.get("visibility") + try: + requested_visibility = PermissionEnum(raw_visibility) if raw_visibility else PermissionEnum.ONLY_ME + except ValueError: + requested_visibility = PermissionEnum.ONLY_ME + if requested_visibility not in (PermissionEnum.ONLY_ME, PermissionEnum.ALL_TEAM): + requested_visibility = PermissionEnum.ONLY_ME + context_id = OAuthProxyService.create_proxy_context( user_id=current_user.id, tenant_id=tenant_id, plugin_id=plugin_id, provider=provider_name, credential_id=credential_id, + extra_data={"visibility": requested_visibility.value}, ) oauth_handler = OAuthHandler() redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/datasource/callback" @@ -253,6 +276,17 @@ class DatasourceOAuthCallback(Resource): credential_id=credential_id, ) else: + # Visibility was chosen by the user before the redirect and stashed + # in the proxy context. Fall back to ONLY_ME for older cookies (or + # anything that somehow wrote an unexpected value) — OAuth tokens + # are personal by default. + stored_visibility = context.get("visibility") + try: + visibility = PermissionEnum(stored_visibility) if stored_visibility else PermissionEnum.ONLY_ME + except ValueError: + visibility = PermissionEnum.ONLY_ME + if visibility not in (PermissionEnum.ONLY_ME, PermissionEnum.ALL_TEAM): + visibility = PermissionEnum.ONLY_ME datasource_provider_service.add_datasource_oauth_provider( tenant_id=tenant_id, provider_id=datasource_provider_id, @@ -260,6 +294,8 @@ class DatasourceOAuthCallback(Resource): name=oauth_response.metadata.get("name") or None, expire_at=oauth_response.expires_at, credentials=dict(oauth_response.credentials), + user_id=user_id, + visibility=visibility, ) return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback") @@ -381,11 +417,12 @@ class DatasourceAuthListApi(Resource): @setup_required @login_required @account_initialization_required + @with_current_user @with_current_tenant_id - def get(self, current_tenant_id: str): + def get(self, current_tenant_id: str, user: Account): datasource_provider_service = DatasourceProviderService() datasources = datasource_provider_service.get_all_datasource_credentials( - tenant_id=current_tenant_id, session=db.session() + tenant_id=current_tenant_id, session=db.session(), user=user ) return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 @@ -400,11 +437,12 @@ class DatasourceHardCodeAuthListApi(Resource): @setup_required @login_required @account_initialization_required + @with_current_user @with_current_tenant_id - def get(self, current_tenant_id: str): + def get(self, current_tenant_id: str, user: Account): datasource_provider_service = DatasourceProviderService() datasources = datasource_provider_service.get_hard_code_datasource_credentials( - tenant_id=current_tenant_id, session=db.session() + tenant_id=current_tenant_id, session=db.session(), user=user ) return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index badfb21eed6..a02e603373d 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -67,6 +67,7 @@ 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.enums import PermissionEnum from models.provider_ids import ToolProviderID # from models.provider_ids import ToolProviderID @@ -97,6 +98,13 @@ class ToolProviderListQuery(BaseModel): type: Literal["builtin", "model", "api", "workflow", "mcp"] | None = None +class ToolOAuthAuthorizationQuery(BaseModel): + visibility: Literal["only_me", "all_team_members"] | None = Field( + default=None, + description="Visibility for the OAuth credential. Defaults to 'only_me'.", + ) + + class BuiltinToolCredentialDeletePayload(BaseModel): credential_id: str @@ -436,6 +444,7 @@ class WorkflowToolDetailResponse(ResponseModel): register_schema_models( console_ns, ToolProviderListQuery, + ToolOAuthAuthorizationQuery, UrlQuery, ProviderQuery, BuiltinCredentialListQuery, @@ -1107,6 +1116,7 @@ class ToolLabelsApi(Resource): @console_ns.route("/oauth/plugin//tool/authorization-url") class ToolPluginOAuthApi(Resource): + @console_ns.doc(params=query_params_from_model(ToolOAuthAuthorizationQuery)) @console_ns.response( 200, "Tool OAuth authorization URL generated successfully", @@ -1128,9 +1138,25 @@ class ToolPluginOAuthApi(Resource): if oauth_client_params is None: raise Forbidden("no oauth available client config found for this tool provider") + # Visibility is chosen by the user in the frontend before the redirect, + # then read back in the callback below when the credential is created. + # Only ONLY_ME / ALL_TEAM are accepted; anything else falls back to + # ONLY_ME (OAuth tokens are personal by nature). + raw_visibility = request.args.get("visibility") + try: + requested_visibility = PermissionEnum(raw_visibility) if raw_visibility else PermissionEnum.ONLY_ME + except ValueError: + requested_visibility = PermissionEnum.ONLY_ME + if requested_visibility not in (PermissionEnum.ONLY_ME, PermissionEnum.ALL_TEAM): + requested_visibility = PermissionEnum.ONLY_ME + oauth_handler = OAuthHandler() context_id = OAuthProxyService.create_proxy_context( - user_id=user.id, tenant_id=tenant_id, plugin_id=plugin_id, provider=provider_name + user_id=user.id, + tenant_id=tenant_id, + plugin_id=plugin_id, + provider=provider_name, + extra_data={"visibility": requested_visibility.value}, ) redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider}/tool/callback" authorization_url_response = oauth_handler.get_authorization_url( @@ -1194,7 +1220,17 @@ class ToolOAuthCallback(Resource): if not credentials: raise Exception("the plugin credentials failed") - # add credentials to database — OAuth tokens default to only_me since they're personal + # Visibility was chosen by the user before the redirect and stashed in + # the proxy context. Fall back to ONLY_ME for older cookies (or for + # anything that somehow wrote an unexpected value) — OAuth tokens are + # personal by default. + stored_visibility = context.get("visibility") + try: + visibility = PermissionEnum(stored_visibility) if stored_visibility else PermissionEnum.ONLY_ME + except ValueError: + visibility = PermissionEnum.ONLY_ME + if visibility not in (PermissionEnum.ONLY_ME, PermissionEnum.ALL_TEAM): + visibility = PermissionEnum.ONLY_ME BuiltinToolManageService.add_builtin_tool_provider( user_id=user_id, tenant_id=tenant_id, @@ -1202,7 +1238,7 @@ class ToolOAuthCallback(Resource): credentials=dict(credentials), expires_at=expires_at, api_type=CredentialType.OAUTH2, - visibility="only_me", + visibility=visibility.value, ) # response-contract:ignore redirect response return redirect(f"{dify_config.CONSOLE_WEB_URL}/oauth-callback") diff --git a/api/controllers/files/upload.py b/api/controllers/files/upload.py index 82a7ae4fb65..aac1c3b20c8 100644 --- a/api/controllers/files/upload.py +++ b/api/controllers/files/upload.py @@ -32,6 +32,7 @@ class PluginUploadQuery(BaseModel): user_id: str | None = Field(default=None, description="User identifier") user_from: Literal["account", "end-user"] | None = Field(default=None, description="User identity type") conversation_id: str | None = Field(default=None, description="Conversation identifier") + max_size: int | None = Field(default=None, ge=0, description="Signed maximum file size in bytes") register_schema_models(files_ns, PluginUploadQuery) @@ -113,14 +114,22 @@ class PluginUploadFileApi(Resource): timestamp=timestamp, nonce=nonce, sign=sign, + max_size=args.max_size, ): raise Forbidden("Invalid request.") try: + if args.max_size is None: + file_binary = file.stream.read() + else: + file_binary = file.stream.read(args.max_size + 1) + if len(file_binary) > args.max_size: + raise FileTooLargeError("File size exceeds the signed upload limit.") + tool_file = ToolFileManager().create_file_by_raw( user_id=owner_id, tenant_id=tenant_id, - file_binary=file.stream.read(), + file_binary=file_binary, mimetype=mimetype, filename=filename, conversation_id=args.conversation_id, diff --git a/api/controllers/inner_api/agent/files.py b/api/controllers/inner_api/agent/files.py index 8838ba2d4fe..d4e224622f7 100644 --- a/api/controllers/inner_api/agent/files.py +++ b/api/controllers/inner_api/agent/files.py @@ -5,7 +5,7 @@ from __future__ import annotations from typing import Literal from flask_restx import Resource -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, Field, ValidationError from sqlalchemy.orm import Session from configs import dify_config @@ -39,6 +39,7 @@ class AgentFileUploadRequestPayload(RequestRequestUploadFile): tenant_id: str user_id: str user_from: Literal["account", "end-user"] | None = None + max_size: int = Field(ge=0, description="Maximum upload size in bytes") model_config = ConfigDict(extra="forbid") @@ -127,6 +128,7 @@ class AgentFileUploadRequestApi(Resource): user_id=owner_id, conversation_id=payload.conversation_id, user_from=payload.user_from, + max_size=payload.max_size, ) except ValueError as exc: raise AgentFileRequestHttpError( diff --git a/api/controllers/trigger/trigger.py b/api/controllers/trigger/trigger.py index c10b94050c5..203bfdf895d 100644 --- a/api/controllers/trigger/trigger.py +++ b/api/controllers/trigger/trigger.py @@ -39,5 +39,5 @@ def trigger_endpoint(endpoint_id: str): except ValueError as e: return jsonify({"error": "Endpoint processing failed", "message": str(e)}), 400 except Exception: - logger.exception("Webhook processing failed for {endpoint_id}") + logger.exception("Webhook processing failed for %s", endpoint_id) return jsonify({"error": "Internal server error"}), 500 diff --git a/api/controllers/web/forgot_password.py b/api/controllers/web/forgot_password.py index a9374555ed4..0be96287dec 100644 --- a/api/controllers/web/forgot_password.py +++ b/api/controllers/web/forgot_password.py @@ -15,7 +15,12 @@ from controllers.console.auth.error import ( PasswordMismatchError, ) from controllers.console.error import EmailSendIpLimitError -from controllers.console.wraps import email_password_login_enabled, only_edition_enterprise, setup_required +from controllers.console.wraps import ( + email_password_login_enabled, + model_validate, + only_edition_enterprise, + setup_required, +) from controllers.web import web_ns from extensions.ext_database import db from libs.helper import extract_remote_ip @@ -54,9 +59,8 @@ class ForgotPasswordSendEmailApi(Resource): } ) @web_ns.response(200, "Password reset email sent successfully", web_ns.models[SimpleResultDataResponse.__name__]) - def post(self): - payload = ForgotPasswordSendPayload.model_validate(web_ns.payload or {}) - + @model_validate(ForgotPasswordSendPayload) + def post(self, payload: ForgotPasswordSendPayload): request_email = payload.email normalized_email = request_email.lower() @@ -90,9 +94,8 @@ class ForgotPasswordCheckApi(Resource): responses={200: "Token is valid", 400: "Bad request - invalid token format", 401: "Invalid or expired token"} ) @web_ns.response(200, "Token is valid", web_ns.models[VerificationTokenResponse.__name__]) - def post(self): - payload = ForgotPasswordCheckPayload.model_validate(web_ns.payload or {}) - + @model_validate(ForgotPasswordCheckPayload) + def post(self, payload: ForgotPasswordCheckPayload): user_email = payload.email.lower() is_forgot_password_error_rate_limit = AccountService.is_forgot_password_error_rate_limit(user_email) @@ -144,9 +147,8 @@ class ForgotPasswordResetApi(Resource): } ) @web_ns.response(200, "Password reset successfully", web_ns.models[SimpleResultResponse.__name__]) - def post(self): - payload = ForgotPasswordResetPayload.model_validate(web_ns.payload or {}) - + @model_validate(ForgotPasswordResetPayload) + def post(self, payload: ForgotPasswordResetPayload): # Validate passwords match if payload.new_password != payload.password_confirm: raise PasswordMismatchError() diff --git a/api/controllers/web/human_input_form.py b/api/controllers/web/human_input_form.py index a897508e798..5668b97abaa 100644 --- a/api/controllers/web/human_input_form.py +++ b/api/controllers/web/human_input_form.py @@ -262,7 +262,7 @@ class HumanInputFormApi(Resource): raise NotFoundError("Form not found") if (recipient_type := form.recipient_type) is None: - logger.warning("Recipient type is None for form, form_id=%", form.id) + logger.warning("Recipient type is None for form, form_id=%s", form.id) raise AssertionError("Recipient type is None") try: diff --git a/api/controllers/web/login.py b/api/controllers/web/login.py index b841056743d..6d99d290b0b 100644 --- a/api/controllers/web/login.py +++ b/api/controllers/web/login.py @@ -25,6 +25,7 @@ from controllers.console.error import AccountBannedError from controllers.console.wraps import ( decrypt_code_field, decrypt_password_field, + model_validate, only_edition_enterprise, setup_required, ) @@ -100,9 +101,9 @@ class LoginApi(Resource): ) @web_ns.response(200, "Authentication successful", web_ns.models[AccessTokenResultResponse.__name__]) @decrypt_password_field - def post(self): + @model_validate(LoginPayload) + def post(self, payload: LoginPayload): """Authenticate user and login.""" - payload = LoginPayload.model_validate(web_ns.payload or {}) normalized_email = payload.email.lower() try: @@ -139,8 +140,8 @@ class LoginStatusApi(Resource): } ) @web_ns.response(200, "Login status", web_ns.models[LoginStatusResponse.__name__]) - def get(self): - query = LoginStatusQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(LoginStatusQuery) + def get(self, query: LoginStatusQuery): app_code = query.app_code user_id = query.user_id token = extract_webapp_access_token(request) @@ -206,9 +207,8 @@ class EmailCodeLoginSendEmailApi(Resource): } ) @web_ns.response(200, "Email code sent successfully", web_ns.models[SimpleResultDataResponse.__name__]) - def post(self): - payload = EmailCodeLoginSendPayload.model_validate(web_ns.payload or {}) - + @model_validate(EmailCodeLoginSendPayload) + def post(self, payload: EmailCodeLoginSendPayload): if payload.language == "zh-Hans": language = "zh-Hans" else: @@ -242,9 +242,8 @@ class EmailCodeLoginApi(Resource): web_ns.models[AccessTokenResultResponse.__name__], ) @decrypt_code_field - def post(self): - payload = EmailCodeLoginVerifyPayload.model_validate(web_ns.payload or {}) - + @model_validate(EmailCodeLoginVerifyPayload) + def post(self, payload: EmailCodeLoginVerifyPayload): user_email = payload.email.lower() token_data = WebAppAuthService.get_email_code_login_data(payload.token) diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py index da1008ddfe8..191c4223f70 100644 --- a/api/core/app/apps/agent_app/app_runner.py +++ b/api/core/app/apps/agent_app/app_runner.py @@ -24,6 +24,7 @@ from clients.agent_backend import ( AgentBackendDeferredToolCallInternalEvent, AgentBackendError, AgentBackendInternalEventType, + AgentBackendRunCancelledInternalEvent, AgentBackendRunClient, AgentBackendRunEventAdapter, AgentBackendRunFailedError, @@ -681,6 +682,8 @@ class AgentAppRunner: queue_manager=queue_manager, model_name=model_name, query=query, + session_scope=scope, + binding_id=runtime.binding_id, ) if isinstance(terminal, AgentBackendDeferredToolCallInternalEvent): @@ -700,6 +703,15 @@ class AgentAppRunner: ) return + if isinstance(terminal, AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent): + # None means no post-exit snapshot was produced; leave the previously stored session snapshot untouched. + if terminal.session_snapshot is not None: + self._save_session( + scope=scope, + binding_id=runtime.binding_id, + snapshot=terminal.session_snapshot, + ) + if not isinstance(terminal, AgentBackendRunSucceededInternalEvent): if isinstance(terminal, AgentBackendRunFailedInternalEvent): reason = terminal.reason @@ -895,6 +907,8 @@ class AgentAppRunner: queue_manager: AppQueueManager, model_name: str, query: str | None, + session_scope: AgentAppSessionScope, + binding_id: str, ): """Consume backend events while preserving raw recorder granularity.""" terminal = None @@ -904,6 +918,7 @@ class AgentAppRunner: queue_manager=queue_manager, ) text_delta_debouncer = _TextDeltaDebouncer(debounce_seconds=self._text_delta_debounce_seconds) + last_event_id: str | None = None def persist_answer_text(content_delta: str) -> None: try: @@ -934,14 +949,26 @@ class AgentAppRunner: should_stop=queue_manager.is_stopped, ) for public_event in public_events: + if public_event.id is not None: + last_event_id = public_event.id if queue_manager.is_stopped(): flush_pending_agent_message_text() - self._cancel_run(run_id) + self._cancel_run( + run_id, + after=last_event_id, + session_scope=session_scope, + binding_id=binding_id, + ) raise GenerateTaskStoppedError() for internal_event in self._event_adapter.adapt(public_event): if queue_manager.is_stopped(): flush_pending_agent_message_text() - self._cancel_run(run_id) + self._cancel_run( + run_id, + after=last_event_id, + session_scope=session_scope, + binding_id=binding_id, + ) raise GenerateTaskStoppedError() if internal_event.type in ( AgentBackendInternalEventType.RUN_STARTED, @@ -978,21 +1005,52 @@ class AgentAppRunner: raise except Exception as error: flush_pending_agent_message_text() - self._cancel_run(run_id) + self._cancel_run( + run_id, + after=last_event_id, + session_scope=session_scope, + binding_id=binding_id, + ) if queue_manager.is_stopped(): raise GenerateTaskStoppedError() from error raise flush_pending_agent_message_text() if queue_manager.is_stopped(): - self._cancel_run(run_id) + self._cancel_run( + run_id, + after=last_event_id, + session_scope=session_scope, + binding_id=binding_id, + ) raise GenerateTaskStoppedError() return terminal, process_recorder - def _cancel_run(self, run_id: str) -> None: + def _cancel_run( + self, + run_id: str, + *, + after: str | None, + session_scope: AgentAppSessionScope, + binding_id: str, + ) -> None: try: - self._agent_backend_client.cancel_run(run_id) + public_event = self._agent_backend_client.cancel_run_and_wait(run_id, after=after) + for internal_event in self._event_adapter.adapt(public_event): + if ( + isinstance(internal_event, AgentBackendRunCancelledInternalEvent) + and internal_event.session_snapshot is not None + ): + self._save_session( + scope=session_scope, + binding_id=binding_id, + snapshot=internal_event.session_snapshot, + ) except Exception: - logger.warning("Failed to cancel stopped Agent App backend run: run_id=%s", run_id, exc_info=True) + logger.warning( + "Failed to finish cancelling stopped Agent App backend run: run_id=%s", + run_id, + exc_info=True, + ) def _publish_answer( self, *, queue_manager: AppQueueManager, model_name: str, answer: str, query: str | None 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 c07859d44f6..5e9e6953448 100644 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ b/api/core/app/apps/agent_app/runtime_request_builder.py @@ -29,6 +29,7 @@ from clients.agent_backend import ( ) from configs import dify_config from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom +from core.app.llm.model_access import resolve_model_context_window from core.plugin.provider_identity import normalize_plugin_daemon_provider_identity from core.workflow.nodes.agent_v2.dify_tools_builder import ( WorkflowAgentDifyToolLayersBuilder, @@ -129,6 +130,11 @@ class AgentAppRuntimeRequestBuilder: append_runtime_warnings(metadata, config_warnings) soul_prompt_resolver = build_config_aware_soul_mention_resolver(agent_soul) knowledge_config = build_knowledge_layer_config(agent_soul) + context_window_tokens = resolve_model_context_window( + run_context=context.dify_context, + provider_name=agent_soul.model.model_provider, + model_name=agent_soul.model.model, + ) model_plugin_id, model_provider = normalize_plugin_daemon_provider_identity( ModelProviderID(agent_soul.model.model_provider), agent_soul.model.plugin_id, @@ -141,6 +147,7 @@ class AgentAppRuntimeRequestBuilder: model_provider=model_provider, model=agent_soul.model.model, model_settings=agent_soul.model.model_settings.model_dump(mode="json", exclude_none=True), + context_window_tokens=context_window_tokens, ), execution_context=DifyExecutionContextLayerConfig( tenant_id=context.dify_context.tenant_id, diff --git a/api/core/app/apps/execution_coordinator.py b/api/core/app/apps/execution_coordinator.py index ce8e8986307..1372d55155e 100644 --- a/api/core/app/apps/execution_coordinator.py +++ b/api/core/app/apps/execution_coordinator.py @@ -8,6 +8,7 @@ from enum import Enum, auto from configs import dify_config from extensions.ext_redis import redis_client +from graphon.graph_engine.command_channels import RedisChannel from graphon.graph_engine.manager import GraphEngineManager logger = logging.getLogger(__name__) @@ -20,6 +21,11 @@ class AppExecutionState(Enum): TERMINAL = auto() +def app_task_command_channel_key(task_id: str) -> str: + """Redis key of the GraphEngine command channel for one app task.""" + return f"workflow:{task_id}:commands" + + def set_app_task_stop_flag(task_id: str) -> None: if not task_id: return @@ -27,6 +33,40 @@ def set_app_task_stop_flag(task_id: str) -> None: redis_client.setex(f"generate_task_stopped:{task_id}", 600, 1) +def clear_app_task_cancellation_signals(task_id: str) -> None: + """Discard cancellation signals left over from earlier attempts of one task. + + Both cancellation channels are keyed by task ID and outlive the attempt that + armed them: the stop flag lives for 600 seconds and a queued ``AbortCommand`` + for an hour, and neither is consumed while no engine is running. A resumed + workflow deliberately reuses the paused run's task ID, so without this reset + it inherits those signals and aborts itself as soon as it starts. Call this + only when starting a new attempt that is meant to run, never mid-execution. + """ + if not task_id: + return + + try: + redis_client.delete(f"generate_task_stopped:{task_id}") + except Exception: + logger.exception("Failed to clear stop flag for app task %s", task_id) + + channel_key = app_task_command_channel_key(task_id) + try: + # fetch_commands() drains the queue and its pending marker together; the + # explicit delete covers a queue whose marker was already consumed. + discarded = RedisChannel(redis_client, channel_key).fetch_commands() + redis_client.delete(channel_key) + if discarded: + logger.info( + "Discarded %s stale GraphEngine command(s) for app task %s", + len(discarded), + task_id, + ) + except Exception: + logger.exception("Failed to clear pending GraphEngine commands for app task %s", task_id) + + class AppExecutionCoordinator: """Own cancellation policy for one app execution attempt. diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index f6cae29ce35..c3e1be5a95b 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/api/core/app/apps/pipeline/pipeline_generator.py @@ -654,8 +654,6 @@ class PipelineGenerator(BaseAppGenerator): except Exception as e: logger.exception("Unknown Error when generating") queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) - finally: - db.session.close() def _handle_response( self, diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index d7427408792..1ea7b26da29 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from typing import cast from core.app.apps.base_app_queue_manager import AppQueueManager +from core.app.apps.execution_coordinator import app_task_command_channel_key from core.app.apps.workflow.app_config_manager import WorkflowAppConfig from core.app.apps.workflow.command_channels import ( CelerySignalCommandChannel, @@ -153,7 +154,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): # RUN WORKFLOW # Create Redis command channel for this workflow execution task_id = self.application_generate_entity.task_id - channel_key = f"workflow:{task_id}:commands" + channel_key = app_task_command_channel_key(task_id) celery_signal_channel = CelerySignalCommandChannel( shutdown_state_getter=celery_warm_shutdown_started, abort_reason=WORKFLOW_WARM_SHUTDOWN_ABORT_REASON, diff --git a/api/core/app/llm/model_access.py b/api/core/app/llm/model_access.py index d2b8e3539fa..52fb9add70b 100644 --- a/api/core/app/llm/model_access.py +++ b/api/core/app/llm/model_access.py @@ -9,7 +9,7 @@ from core.errors.error import ProviderTokenNotInitError from core.model_manager import ModelInstance, ModelManager from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager from core.provider_manager import ProviderManager -from graphon.model_runtime.entities.model_entities import ModelType +from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ModelType from graphon.nodes.llm.entities import ModelConfig from graphon.nodes.llm.exc import LLMModeRequiredError, ModelNotExistError from graphon.nodes.llm.protocols import CredentialsProvider @@ -128,6 +128,27 @@ def build_dify_model_access(run_context: DifyRunContext) -> tuple[CredentialsPro ) +def resolve_model_context_window( + *, + run_context: DifyRunContext, + provider_name: str, + model_name: str, +) -> int | None: + """Return the selected model's credential-bound context-window capability. + + The ``ModelInstance`` and its schema are resolved with the current + tenant/user ``DifyRunContext``. A positive, non-boolean plugin-declared + ``CONTEXT_SIZE`` is returned; a missing or invalid value returns ``None``. + Model lookup and schema errors propagate. This function does not infer a + window from the model name or fall back to a model registry or cache. + """ + model_instance = DifyModelFactory(run_context=run_context).init_model_instance(provider_name, model_name) + context_window = model_instance.get_model_schema().model_properties.get(ModelPropertyKey.CONTEXT_SIZE) + if isinstance(context_window, bool) or not isinstance(context_window, int) or context_window <= 0: + return None + return context_window + + def _normalize_completion_params(completion_params: dict[str, Any]) -> tuple[dict[str, Any], list[str]]: """ Split node-level completion params into provider parameters and stop sequences. diff --git a/api/core/tools/signature.py b/api/core/tools/signature.py index 725160aaf8a..fc5a7642ac5 100644 --- a/api/core/tools/signature.py +++ b/api/core/tools/signature.py @@ -1,6 +1,7 @@ import base64 import hashlib import hmac +import json import os import time import urllib.parse @@ -95,6 +96,7 @@ def get_signed_file_uri_for_plugin( user_id: str, conversation_id: str | None = None, user_from: Literal["account", "end-user"] | None = None, + max_size: int | None = None, ) -> str: """Build a signed plugin-upload URI without selecting a network origin.""" @@ -109,6 +111,7 @@ def get_signed_file_uri_for_plugin( timestamp=timestamp, nonce=nonce, user_from=user_from, + max_size=max_size, ) sign = hmac.new(_secret_key(), data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() @@ -123,6 +126,8 @@ def get_signed_file_uri_for_plugin( query_params["conversation_id"] = conversation_id if user_from is not None: query_params["user_from"] = user_from + if max_size is not None: + query_params["max_size"] = str(max_size) query = urllib.parse.urlencode(query_params) return f"/files/upload/for-plugin?{query}" @@ -138,6 +143,7 @@ def verify_plugin_file_signature( timestamp: str, nonce: str, sign: str, + max_size: int | None = None, ) -> bool: """Verify the signature used by the plugin-facing file upload endpoint.""" @@ -150,6 +156,7 @@ def verify_plugin_file_signature( timestamp=timestamp, nonce=nonce, user_from=user_from, + max_size=max_size, ) recalculated_sign = hmac.new(_secret_key(), data_to_sign.encode(), hashlib.sha256).digest() recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode() @@ -171,14 +178,34 @@ def _plugin_upload_signature_payload( timestamp: str, nonce: str, user_from: Literal["account", "end-user"] | None, + max_size: int | None, ) -> str: - """Build the compatible upload signature payload with optional identity ownership. + """Build the compatible upload signature payload with optional protected claims. - Omitting ``user_from`` preserves the legacy payload. When present, the - identity kind is appended and HMAC-protected so account/end-user ownership - cannot be altered. + Omitting ``max_size`` preserves the legacy payload. Size-limited tickets use + a versioned JSON payload so unconstrained string fields cannot absorb or + impersonate optional trailing claims. """ + if max_size is not None: + return json.dumps( + { + "conversation_id": conversation_id or "", + "filename": filename, + "max_size": max_size, + "mimetype": mimetype, + "nonce": nonce, + "tenant_id": tenant_id, + "timestamp": timestamp, + "user_from": user_from, + "user_id": user_id, + "version": 2, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + payload = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{conversation_id or ''}|{timestamp}|{nonce}" if user_from is not None: payload = f"{payload}|{user_from}" diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index d0dcde5b0c1..bedbb5765c4 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -47,6 +47,7 @@ from core.workflow.nodes.agent.plugin_strategy_adapter import ( from core.workflow.nodes.agent.runtime_support import AgentRuntimeSupport from core.workflow.nodes.agent_v2 import DifyAgentNode from core.workflow.nodes.agent_v2.binding_resolver import WorkflowAgentBindingResolver +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.agent_v2.output_adapter import WorkflowAgentOutputAdapter from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeRequestBuilder from core.workflow.nodes.human_input.callback import DifyHITLCallback @@ -134,12 +135,27 @@ def get_node_type_classes_mapping() -> Mapping[NodeType, Mapping[str, type[Node] return Node.get_node_type_classes_mapping() -def resolve_workflow_node_class(*, node_type: NodeType, node_version: str) -> type[Node]: +def resolve_workflow_node_class( + *, + node_type: NodeType, + node_version: str, + node_data: Mapping[str, Any] | BaseNodeData | None = None, +) -> type[Node]: """Resolve the production node class for the requested type/version.""" node_mapping = get_node_type_classes_mapping().get(node_type) if not node_mapping: raise ValueError(f"No class mapping found for node type: {node_type}") + # Historical Agent nodes used version=2 for their tool-parameter format. + # Only the explicit kind marker identifies the newer Dify Agent node. + if ( + node_data is not None + and node_type == BuiltinNodeTypes.AGENT + and node_version == "2" + and not is_dify_agent_node_data(node_data) + ): + node_version = "1" + latest_node_class = node_mapping.get(LATEST_VERSION) matched_node_class = node_mapping.get(node_version) node_class = matched_node_class or latest_node_class @@ -400,7 +416,11 @@ class DifyNodeFactory(NodeFactory): typed_node_config = NodeConfigDictAdapter.validate_python(adapted_node_config) node_id = typed_node_config["id"] node_data = typed_node_config["data"] - node_class = self._resolve_node_class(node_type=node_data.type, node_version=str(node_data.version)) + node_class = self._resolve_node_class( + node_type=node_data.type, + node_version=str(node_data.version), + node_data=node_data, + ) # Graph configs are initially validated against permissive shared node data. # Re-validate using the resolved node class so workflow-local node schemas # stay explicit and constructors receive the concrete typed payload. @@ -493,10 +513,19 @@ class DifyNodeFactory(NodeFactory): return node_data @staticmethod - def _resolve_node_class(*, node_type: NodeType, node_version: str) -> type[Node]: + def _resolve_node_class( + *, + node_type: NodeType, + node_version: str, + node_data: Mapping[str, Any] | BaseNodeData | None = None, + ) -> type[Node]: if node_type == BuiltinNodeTypes.LLM: return DifyLLMNode - return resolve_workflow_node_class(node_type=node_type, node_version=node_version) + return resolve_workflow_node_class( + node_type=node_type, + node_version=node_version, + node_data=node_data, + ) def _resolve_llm_model_reference(self, node_data: LLMNodeData) -> LLMNodeData: """Resolve an optional shared model selector from the workflow variable pool.""" diff --git a/api/core/workflow/nodes/agent_v2/agent_node.py b/api/core/workflow/nodes/agent_v2/agent_node.py index 9911ef57ad7..73288b4bcc6 100644 --- a/api/core/workflow/nodes/agent_v2/agent_node.py +++ b/api/core/workflow/nodes/agent_v2/agent_node.py @@ -340,6 +340,17 @@ class DifyAgentNode(Node[DifyAgentNodeData]): process_data=process_data, metadata=metadata, ) + # None means no post-exit snapshot was produced; leave the previously stored session snapshot untouched. + if ( + isinstance(terminal_event, AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent) + and terminal_event.session_snapshot is not None + ): + self._save_session_snapshot( + session_scope=session_scope, + binding_id=stored_session.binding_id, + snapshot=terminal_event.session_snapshot, + metadata=metadata, + ) if exhausted is not None: # Streaming error / unexpected end — surface immediately without # retrying because the failure is transport-level. @@ -516,16 +527,20 @@ class DifyAgentNode(Node[DifyAgentNodeData]): - ``terminal_event``: the first non-stream/non-started internal event, or ``None`` if the stream ended without one. - ``transport_failure``: a populated ``StreamCompletedEvent`` when the - stream itself errored (backend/HTTP/protocol fault). Mutually - exclusive with ``terminal_event``. + stream itself errored (backend/HTTP/protocol fault). A cancellation + terminal may accompany it so the caller can persist the final session + snapshot while preserving the original transport failure. """ stream_event_count = 0 + last_event_id: str | None = None try: for public_event in self._agent_backend_client.stream_events( run_id, should_stop=self._is_graph_aborted, ): stream_event_count += 1 + if public_event.id is not None: + last_event_id = public_event.id for internal_event in self._event_adapter.adapt(public_event): if internal_event.type == AgentBackendInternalEventType.RUN_STARTED: continue @@ -552,8 +567,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]): | AgentBackendDeferredToolCallInternalEvent, ): return internal_event, None - self._cancel_backend_run(run_id, reason="unexpected_event") - return None, self._failure_event( + cancellation = self._cancel_backend_run( + run_id, + reason="unexpected_event", + after=last_event_id, + ) + return cancellation, self._failure_event( inputs=inputs, process_data=process_data, metadata=metadata, @@ -561,8 +580,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]): error_type="agent_backend_stream_error", ) except AgentBackendError as error: - self._cancel_backend_run(run_id, reason=self._stream_stop_reason()) - return None, self._failure_event( + cancellation = self._cancel_backend_run( + run_id, + reason=self._stream_stop_reason(), + after=last_event_id, + ) + return cancellation, self._failure_event( inputs=inputs, process_data=process_data, metadata=metadata, @@ -570,8 +593,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]): error_type=self._agent_backend_error_type(error), ) except Exception as error: - self._cancel_backend_run(run_id, reason=self._stream_stop_reason()) - return None, self._failure_event( + cancellation = self._cancel_backend_run( + run_id, + reason=self._stream_stop_reason(), + after=last_event_id, + ) + return cancellation, self._failure_event( inputs=inputs, process_data=process_data, metadata=metadata, @@ -579,8 +606,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]): error_type="agent_backend_stream_error", ) - self._cancel_backend_run(run_id, reason="stream_ended_without_terminal_event") - return None, None + cancellation = self._cancel_backend_run( + run_id, + reason=self._stream_stop_reason() if self._is_graph_aborted() else "stream_ended_without_terminal_event", + after=last_event_id, + ) + return cancellation, None def _is_graph_aborted(self) -> bool: """Let Agent SSE consumption observe GraphEngine's cooperative abort state.""" @@ -592,14 +623,25 @@ class DifyAgentNode(Node[DifyAgentNodeData]): def _stream_stop_reason(self) -> str: return "workflow_graph_aborted" if self._is_graph_aborted() else "event_stream_failed" - def _cancel_backend_run(self, run_id: str, *, reason: str) -> None: + def _cancel_backend_run( + self, + run_id: str, + *, + reason: str, + after: str | None, + ) -> AgentBackendRunCancelledInternalEvent | None: try: - self._agent_backend_client.cancel_run( + public_event = self._agent_backend_client.cancel_run_and_wait( run_id, CancelRunRequest(reason=reason, message="Workflow Agent event consumption stopped"), + after=after, ) + for internal_event in self._event_adapter.adapt(public_event): + if isinstance(internal_event, AgentBackendRunCancelledInternalEvent): + return internal_event except Exception: - logger.warning("Failed to cancel Workflow Agent backend run: run_id=%s", run_id, exc_info=True) + logger.warning("Failed to finish cancelling Workflow Agent backend run: run_id=%s", run_id, exc_info=True) + return None @staticmethod def _record_type_check_metadata(metadata: dict[str, Any], outcome: OutputTypeCheckOutcome) -> None: diff --git a/api/core/workflow/nodes/agent_v2/discriminator.py b/api/core/workflow/nodes/agent_v2/discriminator.py new file mode 100644 index 00000000000..8df480ca5e7 --- /dev/null +++ b/api/core/workflow/nodes/agent_v2/discriminator.py @@ -0,0 +1,28 @@ +from collections.abc import Mapping +from typing import Any + +from graphon.entities.base_node_data import BaseNodeData +from graphon.enums import BuiltinNodeTypes + +AGENT_NODE_KIND = "dify_agent" +AGENT_NODE_VERSION = "2" + + +def is_dify_agent_node_data(node_data: Mapping[str, Any] | BaseNodeData) -> bool: + """Return whether node data explicitly identifies the new Dify Agent node.""" + + if isinstance(node_data, Mapping): + node_type = node_data.get("type") + node_version = node_data.get("version") + agent_node_kind = node_data.get("agent_node_kind") + else: + serialized_node_data = node_data.model_dump(mode="python") + node_type = serialized_node_data.get("type") + node_version = serialized_node_data.get("version") + agent_node_kind = serialized_node_data.get("agent_node_kind") + + return ( + node_type == BuiltinNodeTypes.AGENT + and str(node_version) == AGENT_NODE_VERSION + and agent_node_kind == AGENT_NODE_KIND + ) diff --git a/api/core/workflow/nodes/agent_v2/entities.py b/api/core/workflow/nodes/agent_v2/entities.py index eb36b9cf1ed..a9fba739985 100644 --- a/api/core/workflow/nodes/agent_v2/entities.py +++ b/api/core/workflow/nodes/agent_v2/entities.py @@ -8,7 +8,7 @@ from graphon.enums import BuiltinNodeTypes, NodeType class DifyAgentNodeData(BaseNodeData): type: NodeType = BuiltinNodeTypes.AGENT - agent_node_kind: Literal["dify_agent"] = "dify_agent" + agent_node_kind: Literal["dify_agent"] @model_validator(mode="after") def validate_version(self) -> "DifyAgentNodeData": 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 e42686b42b9..3b617de5f15 100644 --- a/api/core/workflow/nodes/agent_v2/runtime_request_builder.py +++ b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py @@ -47,6 +47,7 @@ from clients.agent_backend import ( ) from configs import dify_config from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom +from core.app.llm.model_access import resolve_model_context_window from core.plugin.provider_identity import normalize_plugin_daemon_provider_identity from core.workflow.system_variables import SystemVariableKey, get_system_text, get_system_value from graphon.file import File, FileTransferMethod @@ -211,6 +212,11 @@ class WorkflowAgentRuntimeRequestBuilder: soul_prompt_resolver = build_config_aware_soul_mention_resolver(agent_soul) soul_prompt = expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip() knowledge_config = build_knowledge_layer_config(agent_soul) + context_window_tokens = resolve_model_context_window( + run_context=context.dify_context, + provider_name=agent_soul.model.model_provider, + model_name=agent_soul.model.model, + ) model_plugin_id, model_provider = normalize_plugin_daemon_provider_identity( ModelProviderID(agent_soul.model.model_provider), agent_soul.model.plugin_id, @@ -223,6 +229,7 @@ class WorkflowAgentRuntimeRequestBuilder: model_provider=model_provider, model=agent_soul.model.model, model_settings=agent_soul.model.model_settings.model_dump(mode="json", exclude_none=True), + context_window_tokens=context_window_tokens, ), # The execution-context layer is now the only public protocol # carrier for Dify tenant/user/run identifiers. ``user_id`` and diff --git a/api/core/workflow/nodes/agent_v2/validators.py b/api/core/workflow/nodes/agent_v2/validators.py index 7e3969d8a16..df0c413ff7e 100644 --- a/api/core/workflow/nodes/agent_v2/validators.py +++ b/api/core/workflow/nodes/agent_v2/validators.py @@ -20,6 +20,7 @@ from models.model import UploadFile from models.workflow import Workflow from services.agent.knowledge_datasets import list_missing_tenant_knowledge_dataset_ids +from .discriminator import is_dify_agent_node_data from .entities import DifyAgentNodeData @@ -243,7 +244,7 @@ class WorkflowAgentNodeValidator: node_data = node.get("data") if not isinstance(node_id, str) or not isinstance(node_data, Mapping): continue - if node_data.get("type") == BuiltinNodeTypes.AGENT and str(node_data.get("version")) == "2": + if is_dify_agent_node_data(node_data): yield node_id, node_data @staticmethod diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py index 866bc73fcf6..372bbd4e7f8 100644 --- a/api/core/workflow/workflow_entry.py +++ b/api/core/workflow/workflow_entry.py @@ -215,7 +215,11 @@ class WorkflowEntry: if node_type in {BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION}: raise ValueError("Loop and Iteration nodes must use their engine-backed debug endpoints") node_version = str(node_config_data.version) - node_cls = resolve_workflow_node_class(node_type=node_type, node_version=node_version) + node_cls = resolve_workflow_node_class( + node_type=node_type, + node_version=node_version, + node_data=node_config_data, + ) # init graph context and runtime state run_context = build_dify_run_context( diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 87ad8bfa4c5..1c42a08939e 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -173,6 +173,7 @@ def init_app(app: DifyApp) -> Celery: "tasks.regenerate_summary_index_task", # summary index regeneration "tasks.initialize_created_app_rbac_access_task", # app access initialization "tasks.install_default_plugins_task", # tenant default plugin installation + "tasks.new_agent_beta_task", # New Agent Beta eligibility checks "tasks.refresh_billing_vector_space_task", # billing vector-space cache refresh "tasks.app_generate.resume_agent_app_task", # ENG-635: Agent v2 chat ask_human resume "tasks.workflow_run_archive_download_tasks", # workflow-run archive download preparation diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 29e349e6065..d218c64b250 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -9701,6 +9701,7 @@ Initiate OAuth login process | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | | credential_id | query | Credential ID to reauthorize | No | string | +| visibility | query | Visibility for the credential to be created. Accepts 'only_me' or 'all_team_members'; any other value falls back to 'only_me'. Ignored on reauthorization (credential_id set). | No | string | | provider_id | path | | Yes | string | #### Responses @@ -9714,6 +9715,7 @@ Initiate OAuth login process | Name | Located in | Description | Required | Schema | | ---- | ---------- | ----------- | -------- | ------ | +| visibility | query | Visibility for the OAuth credential. Defaults to 'only_me'. | No | string,
**Available values:** "all_team_members", "only_me" | | provider | path | | Yes | string | #### Responses @@ -19184,6 +19186,7 @@ Model class for provider custom model configuration. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | credential_id | string | Credential ID to reauthorize | No | +| visibility | string | Visibility for the credential to be created. Accepts 'only_me' or 'all_team_members'; any other value falls back to 'only_me'. Ignored on reauthorization (credential_id set). | No | #### DatasourceOAuthCallbackQuery @@ -26779,6 +26782,12 @@ Tool label | ---- | ---- | ----------- | -------- | | ToolLabelListResponse | array | | | +#### ToolOAuthAuthorizationQuery + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| visibility | string | Visibility for the OAuth credential. Defaults to 'only_me'. | No | + #### ToolOAuthCustomClientPayload | Name | Type | Description | Required | diff --git a/api/services/account_service.py b/api/services/account_service.py index 5f098135370..0387e000769 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -448,6 +448,7 @@ class AccountService: interface_theme: str = "light", is_setup: bool | None = False, timezone: str | None = None, + ip_address: str | None = None, *, session: Session, ) -> Account: @@ -500,6 +501,7 @@ class AccountService: interface_language=interface_language, interface_theme=interface_theme, timezone=resolved_timezone, + last_login_ip=ip_address, ) session.add(account) @@ -513,6 +515,7 @@ class AccountService: interface_language: str, password: str | None = None, timezone: str | None = None, + ip_address: str | None = None, *, session: Session, ) -> Account: @@ -523,6 +526,7 @@ class AccountService: interface_language=interface_language, password=password, timezone=timezone, + ip_address=ip_address, session=session, ) @@ -2009,10 +2013,10 @@ class RegisterService: interface_language=get_valid_language(language), password=password, is_setup=True, + ip_address=ip_address, session=session, ) - account.last_login_ip = ip_address account.initialized_at = naive_utc_now() TenantService.create_owner_tenant_if_not_exist(account=account, is_setup=True, session=session) @@ -2048,6 +2052,7 @@ class RegisterService: is_setup: bool | None = False, create_workspace_required: bool | None = True, timezone: str | None = None, + ip_address: str | None = None, *, session: Session, ) -> Account: @@ -2062,6 +2067,7 @@ class RegisterService: password=password, is_setup=is_setup, timezone=timezone, + ip_address=ip_address, session=session, ) account.status = status or AccountStatus.ACTIVE diff --git a/api/services/agent/composer_candidates.py b/api/services/agent/composer_candidates.py index a650b16e9bc..efa0e422324 100644 --- a/api/services/agent/composer_candidates.py +++ b/api/services/agent/composer_candidates.py @@ -21,6 +21,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping from typing import Any +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from models.agent_config_entities import ( AgentSoulConfig, DeclaredOutputConfig, @@ -97,7 +98,7 @@ def previous_node_output_candidates( continue declared: list[DeclaredOutputConfig] | None = None - if kind == "agent" and str(data.get("version", "")) == "2": + if is_dify_agent_node_data(data): declared = declared_outputs_loader(nid) if declared is not None: for output in declared: diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index 13924aae4b1..d29d8ed6b41 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -67,6 +67,7 @@ from services.entities.agent_entities import ( WorkflowNodeJobConfig, ) from tasks.collect_agent_resources_task import enqueue_agent_resource_collection +from tasks.new_agent_beta_task import register_new_agent_beta_publish_after_commit # WorkflowAgentNodeBinding.workflow_version tag for the draft workflow row. # Mirrors Workflow.version when it is "draft" (see models/workflow.py). @@ -683,6 +684,12 @@ class AgentComposerService: app.enable_api = True app.updated_by = account_id session.flush() + register_new_agent_beta_publish_after_commit( + session=session, + tenant_id=tenant_id, + agent_id=agent.id, + snapshot_id=version.id, + ) return { "result": "success", "active_config_snapshot_id": version.id, diff --git a/api/services/agent/dsl_service.py b/api/services/agent/dsl_service.py index e1f9afd0142..1b7a09cf493 100644 --- a/api/services/agent/dsl_service.py +++ b/api/services/agent/dsl_service.py @@ -18,8 +18,8 @@ from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import func, select from sqlalchemy.orm import Session +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidator -from graphon.enums import BuiltinNodeTypes from models import Account from models.agent import ( APP_BACKED_AGENT_SOURCES, @@ -636,7 +636,7 @@ class AgentDslService: def is_agent_v2_graph(graph: Mapping[str, Any]) -> bool: return any( - node.get("data", {}).get("type") == BuiltinNodeTypes.AGENT and node.get("data", {}).get("version") == "2" + isinstance(node.get("data"), Mapping) and is_dify_agent_node_data(node["data"]) for node in graph.get("nodes", []) if isinstance(node, Mapping) ) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 0519f66f2c6..e4c38f8874c 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -208,6 +208,10 @@ class BillingService: # Cache TTL: 10 minutes _PLAN_CACHE_TTL = 600 + @classmethod + def ensure_new_agent_beta_revision(cls, revision_id: str) -> None: + cls._send_request("POST", f"/new-agent-beta/revisions/{revision_id}/ensure") + @classmethod def get_info(cls, tenant_id: str, exclude_vector_space: bool = False) -> BillingInfo: params = {"tenant_id": tenant_id} diff --git a/api/services/datasource_provider_service.py b/api/services/datasource_provider_service.py index 61fc74ec50d..9721e8aa0e2 100644 --- a/api/services/datasource_provider_service.py +++ b/api/services/datasource_provider_service.py @@ -22,6 +22,7 @@ from core.tools.utils.encryption import ProviderConfigCache, ProviderConfigEncry from extensions.ext_database import db from extensions.ext_redis import redis_client from graphon.model_runtime.entities.provider_entities import FormType +from models.enums import PermissionEnum from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider from models.provider_ids import DatasourceProviderID @@ -649,10 +650,23 @@ class DatasourceProviderService: avatar_url: str | None, expire_at: int, credentials: dict[str, Any], + user_id: str | None = None, + visibility: PermissionEnum | None = None, ) -> None: """ - add datasource oauth provider + add datasource oauth provider. + + ``user_id`` is the creator whose visibility choice this credential is + scoped to. When ``visibility`` is only_me the row is visible only to + this creator; ``all_team_members`` shares it workspace-wide. + Callers that omit ``user_id``/``visibility`` fall back to the previous + team-wide default (matches how the DB column defaults). """ + # partial_members isn't supported for plugin credentials (matches the + # tool + api-key paths); collapse it to ALL_TEAM so we never persist + # an unreachable value here. + if visibility == PermissionEnum.PARTIAL_TEAM: + visibility = PermissionEnum.ALL_TEAM credential_type = CredentialType.OAUTH2 with sessionmaker(bind=db.engine).begin() as session: lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{credential_type.value}" @@ -699,16 +713,21 @@ class DatasourceProviderService: if key in provider_credential_secret_variables: credentials[key] = encrypter.encrypt_token(tenant_id, value) - datasource_provider = DatasourceProvider( - tenant_id=tenant_id, - name=db_provider_name, - provider=provider_id.provider_name, - plugin_id=provider_id.plugin_id, - auth_type=credential_type.value, - encrypted_credentials=credentials, - avatar_url=avatar_url or "default", - expires_at=expire_at, - ) + datasource_provider_kwargs: dict[str, Any] = { + "tenant_id": tenant_id, + "name": db_provider_name, + "provider": provider_id.provider_name, + "plugin_id": provider_id.plugin_id, + "auth_type": credential_type.value, + "encrypted_credentials": credentials, + "avatar_url": avatar_url or "default", + "expires_at": expire_at, + } + if user_id is not None: + datasource_provider_kwargs["user_id"] = user_id + if visibility is not None: + datasource_provider_kwargs["visibility"] = visibility + datasource_provider = DatasourceProvider(**datasource_provider_kwargs) session.add(datasource_provider) def add_datasource_api_key_provider( @@ -886,11 +905,17 @@ class DatasourceProviderService: return copy_credentials_list - def get_all_datasource_credentials(self, tenant_id: str, *, session: Session) -> list[dict]: + def get_all_datasource_credentials( + self, tenant_id: str, *, session: Session, user: "Account | None" = None + ) -> list[dict]: """ get datasource credentials. - :return: + ``user`` is threaded through to ``list_datasource_credentials`` so the + embedded ``credentials_list`` per datasource is filtered by + per-credential visibility. Callers that omit it (background / + maintenance jobs) get the pre-visibility behavior of returning every + credential in the workspace. """ # get all plugin providers manager = PluginDatasourceManager() @@ -902,6 +927,7 @@ class DatasourceProviderService: tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id, + user=user, session=session, ) redirect_uri = ( @@ -945,11 +971,14 @@ class DatasourceProviderService: ) return datasource_credentials - def get_hard_code_datasource_credentials(self, tenant_id: str, *, session: Session) -> list[dict]: + def get_hard_code_datasource_credentials( + self, tenant_id: str, *, session: Session, user: "Account | None" = None + ) -> list[dict]: """ get hard code datasource credentials. - :return: + ``user`` is threaded through to ``list_datasource_credentials`` so + credentials in the returned envelope are visibility-filtered. """ # get all plugin providers manager = PluginDatasourceManager() @@ -967,6 +996,7 @@ class DatasourceProviderService: tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id, + user=user, session=session, ) redirect_uri = "{}/console/api/oauth/plugin/{}/datasource/callback".format( diff --git a/api/services/workflow/node_output_inspector_service.py b/api/services/workflow/node_output_inspector_service.py index 5d6a8f1c675..dee806c9464 100644 --- a/api/services/workflow/node_output_inspector_service.py +++ b/api/services/workflow/node_output_inspector_service.py @@ -59,15 +59,12 @@ from core.workflow.nodes.agent_v2.binding_resolver import ( WorkflowAgentBindingError, WorkflowAgentBindingResolver, ) +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data from core.workflow.nodes.agent_v2.runtime_request_builder import ( WorkflowAgentRuntimeRequestBuilder, ) from factories.file_factory.builders import build_from_mapping -from graphon.enums import ( - BuiltinNodeTypes, - WorkflowExecutionStatus, - WorkflowNodeExecutionStatus, -) +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.file import helpers as file_helpers from models import App from models.agent_config_entities import DeclaredOutputConfig, DeclaredOutputType @@ -182,18 +179,12 @@ class _ResolvedDeclaration: def _is_agent_v2_node(node: Mapping[str, Any]) -> bool: - """A graph node is Agent v2 iff its ``data.type`` is the AGENT builtin - AND its ``data.version`` is ``"2"``. + """Return whether a graph node explicitly identifies the new Dify Agent node.""" - ``BuiltinNodeTypes.AGENT`` is a ``ClassVar[NodeType]`` (plain string), not - a StrEnum, so we compare against it directly without ``.value``. - """ data = node.get("data") or {} if not isinstance(data, Mapping): return False - if data.get("type") != BuiltinNodeTypes.AGENT: - return False - return str(data.get("version", "")) == "2" + return is_dify_agent_node_data(data) def _graph_nodes(workflow_run: WorkflowRun) -> list[Mapping[str, Any]]: diff --git a/api/tasks/app_generate/workflow_execute_task.py b/api/tasks/app_generate/workflow_execute_task.py index 8a88ff4dfa6..3cd30f3b5b7 100644 --- a/api/tasks/app_generate/workflow_execute_task.py +++ b/api/tasks/app_generate/workflow_execute_task.py @@ -12,6 +12,7 @@ from sqlalchemy import Engine, select from sqlalchemy.orm import Session, sessionmaker from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator +from core.app.apps.execution_coordinator import clear_app_task_cancellation_signals from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.apps.workflow.app_generator import WorkflowAppGenerator from core.app.entities.app_invoke_entities import ( @@ -558,6 +559,12 @@ def _resume_app_execution(payload: dict[str, Any]) -> None: ) return + # The resumed attempt reuses the paused run's task ID, so cancellation + # signals armed against that ID before or during the pause would abort it + # immediately and report it as stopped by the user. This attempt is starting + # deliberately, so drop them before any engine can observe them. + clear_app_task_cancellation_signals(generate_entity.task_id) + workflow_run_repo.resume_workflow_pause(workflow_run_id, pause_entity) pause_config = PauseStateLayerConfig( diff --git a/api/tasks/new_agent_beta_task.py b/api/tasks/new_agent_beta_task.py new file mode 100644 index 00000000000..e8febca9e45 --- /dev/null +++ b/api/tasks/new_agent_beta_task.py @@ -0,0 +1,114 @@ +import logging +from datetime import UTC, datetime + +from celery import shared_task +from sqlalchemy import event, select +from sqlalchemy.orm import Session + +from configs import dify_config +from enums import DeploymentEdition +from models.agent import AgentConfigRevision, AgentConfigRevisionOperation +from services.billing_service import BillingService + +logger = logging.getLogger(__name__) + +_MAX_RETRIES = 8 +_RETRY_DELAY_SECONDS = 30 +_MAX_RETRY_DELAY_SECONDS = 900 +NEW_AGENT_BETA_QUEUE = "new_agent_beta" + + +def _is_publish_in_activity_window(published_at: datetime) -> bool: + start = dify_config.NEW_AGENT_BETA_ACTIVITY_START_AT + end = dify_config.NEW_AGENT_BETA_ACTIVITY_END_AT + if start is None or end is None or start.tzinfo is None or end.tzinfo is None or start >= end: + logger.error("New Agent Beta Publish window must be a valid RFC3339 interval") + return False + if published_at.tzinfo is None: + published_at = published_at.replace(tzinfo=UTC) + else: + published_at = published_at.astimezone(UTC) + return start.astimezone(UTC) <= published_at < end.astimezone(UTC) + + +def register_new_agent_beta_publish_after_commit( + *, session: Session, tenant_id: str, agent_id: str, snapshot_id: str +) -> None: + """Best-effort registration that never changes the Publish result.""" + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: + return + + try: + revision = session.scalar( + select(AgentConfigRevision) + .where( + AgentConfigRevision.tenant_id == tenant_id, + AgentConfigRevision.agent_id == agent_id, + AgentConfigRevision.current_snapshot_id == snapshot_id, + AgentConfigRevision.operation == AgentConfigRevisionOperation.PUBLISH_DRAFT, + ) + .limit(1) + ) + if revision is None: + logger.error( + "New Agent Beta publish revision was not found, tenant_id=%s, agent_id=%s, snapshot_id=%s", + tenant_id, + agent_id, + snapshot_id, + ) + return + if not _is_publish_in_activity_window(revision.created_at): + return + cancelled = False + + def cancel_on_rollback(_session: Session) -> None: + nonlocal cancelled + cancelled = True + + def dispatch_after_commit(_session: Session) -> None: + if not cancelled: + schedule_new_agent_beta_ensure(revision.id) + + event.listen(session, "after_rollback", cancel_on_rollback, once=True) + event.listen(session, "after_commit", dispatch_after_commit, once=True) + except Exception: + logger.exception( + "Failed to register New Agent Beta publish event, tenant_id=%s, agent_id=%s, snapshot_id=%s", + tenant_id, + agent_id, + snapshot_id, + ) + + +@shared_task( + queue=NEW_AGENT_BETA_QUEUE, + bind=True, + max_retries=_MAX_RETRIES, + default_retry_delay=_RETRY_DELAY_SECONDS, + acks_late=True, + reject_on_worker_lost=True, +) +def ensure_new_agent_beta_participation_task(self, revision_id: str) -> None: + try: + BillingService.ensure_new_agent_beta_revision(revision_id) + except Exception as exc: + if self.request.retries >= _MAX_RETRIES: + logger.exception("New Agent Beta eligibility retry budget exhausted, revision_id=%s", revision_id) + raise + + logger.warning( + "New Agent Beta eligibility request failed, scheduling retry %d/%d, revision_id=%s", + self.request.retries + 1, + _MAX_RETRIES, + revision_id, + exc_info=True, + ) + countdown = min(_RETRY_DELAY_SECONDS * (2**self.request.retries), _MAX_RETRY_DELAY_SECONDS) + raise self.retry(exc=exc, countdown=countdown) + + +def schedule_new_agent_beta_ensure(revision_id: str) -> None: + try: + ensure_new_agent_beta_participation_task.delay(revision_id) + except Exception: + logger.exception("Failed to dispatch New Agent Beta eligibility task, revision_id=%s", revision_id) 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 c2253a20c15..889ee56dfc4 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 @@ -133,7 +133,12 @@ def seeded_run( "nodes": [ { "id": "agent-node-1", - "data": {"type": "agent", "version": "2", "title": "My Agent"}, + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + "title": "My Agent", + }, }, { "id": "tool-node-1", @@ -328,7 +333,14 @@ def test_snapshot_404s_for_published_run_per_decision_d1(flask_req_ctx, fake_app def test_snapshot_surfaces_type_check_failure_from_metadata(flask_req_ctx, fake_app_model): """Per-output ``TYPE_CHECK_FAILED`` derived from the metadata blob the Stage 4 §5 stack records on the execution row.""" - graph = {"nodes": [{"id": "agent-1", "data": {"type": "agent", "version": "2"}}]} + graph = { + "nodes": [ + { + "id": "agent-1", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + } + ] + } workflow_run = _make_workflow_run(app_id=fake_app_model.id, tenant_id=fake_app_model.tenant_id, graph=graph) execution = _make_execution( app_id=fake_app_model.id, @@ -375,7 +387,14 @@ def test_snapshot_surfaces_type_check_failure_from_metadata(flask_req_ctx, fake_ def test_snapshot_surfaces_output_check_failure_from_metadata(flask_req_ctx, fake_app_model): """When ``output_type_check.passed`` but ``output_check.passed=False``, the output is flagged ``OUTPUT_CHECK_FAILED``.""" - graph = {"nodes": [{"id": "agent-1", "data": {"type": "agent", "version": "2"}}]} + graph = { + "nodes": [ + { + "id": "agent-1", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + } + ] + } workflow_run = _make_workflow_run(app_id=fake_app_model.id, tenant_id=fake_app_model.tenant_id, graph=graph) execution = _make_execution( app_id=fake_app_model.id, @@ -470,7 +489,14 @@ def test_keeps_latest_execution_per_node_by_index(flask_req_ctx, fake_app_model) """Multiple executions for the same node_id → service keeps the highest ``index`` (matches the agent_v2 retry pattern that re-emits node executions).""" - graph = {"nodes": [{"id": "agent-1", "data": {"type": "agent", "version": "2"}}]} + graph = { + "nodes": [ + { + "id": "agent-1", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + } + ] + } workflow_run = _make_workflow_run(app_id=fake_app_model.id, tenant_id=fake_app_model.tenant_id, graph=graph) older = _make_execution( app_id=fake_app_model.id, 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 8ea3a7b8c81..e61b22e611a 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 @@ -1010,6 +1010,7 @@ class TestAppDslService: "data": { "type": BuiltinNodeTypes.AGENT, "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": WorkflowAgentBindingType.ROSTER_AGENT.value, AGENT_PACKAGE_REF_KEY: "agent_1", @@ -1021,6 +1022,7 @@ class TestAppDslService: "data": { "type": BuiltinNodeTypes.AGENT, "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": WorkflowAgentBindingType.INLINE_AGENT.value, AGENT_PACKAGE_REF_KEY: "agent_1", diff --git a/api/tests/unit_tests/clients/agent_backend/test_client.py b/api/tests/unit_tests/clients/agent_backend/test_client.py index 5fb54117306..34b66b83211 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_client.py +++ b/api/tests/unit_tests/clients/agent_backend/test_client.py @@ -9,6 +9,8 @@ from dify_agent.protocol import ( CancelRunResponse, CreateRunRequest, CreateRunResponse, + RunCancelledEvent, + RunCancelledEventData, RunEvent, RunStartedEvent, RunStatusResponse, @@ -51,6 +53,7 @@ def _request() -> CreateRunRequest: class _SuccessfulClient: stream_options: tuple[int | None, object, Callable[[], bool] | None] | None = None + cancel_after: str | None = None def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse: assert isinstance(request, CreateRunRequest) @@ -60,6 +63,21 @@ class _SuccessfulClient: del request return CancelRunResponse(run_id=run_id, status="cancelled") + def cancel_run_and_wait_sync( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + self.cancel_after = after + request = request or CancelRunRequest() + return RunCancelledEvent( + id="2-0", + run_id=run_id, + data=RunCancelledEventData(reason=request.reason, message=request.message), + ) + def stream_events_sync( self, run_id: str, @@ -94,13 +112,20 @@ def test_dify_agent_backend_run_client_delegates_sync_methods() -> None: created = client.create_run(_request()) cancelled = client.cancel_run(created.run_id) + cancelled_event = client.cancel_run_and_wait( + created.run_id, + CancelRunRequest(reason="stopped"), + after="1-0", + ) events = list(client.stream_events(created.run_id, should_stop=should_stop)) status = client.wait_run(created.run_id) assert created.run_id == "run-1" assert cancelled.status == "cancelled" + assert cancelled_event.data.reason == "stopped" assert events[0].type == "run_started" assert status.status == "succeeded" + assert wrapped.cancel_after == "1-0" assert wrapped.stream_options == (2, _STREAM_TIMEOUT_UNSET, should_stop) diff --git a/api/tests/unit_tests/configs/_isolated_settings.py b/api/tests/unit_tests/configs/_isolated_settings.py new file mode 100644 index 00000000000..d9ed65db4c2 --- /dev/null +++ b/api/tests/unit_tests/configs/_isolated_settings.py @@ -0,0 +1,19 @@ +"""Test-only settings source that ignores process and dotenv configuration.""" + +from pydantic_settings import BaseSettings, PydanticBaseSettingsSource + + +class InitSettingsOnly: + """Mixin for settings tests whose inputs should be entirely explicit.""" + + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + del settings_cls, env_settings, dotenv_settings, file_secret_settings + return (init_settings,) diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index af671c77600..f606b731e4a 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/api/tests/unit_tests/configs/test_dify_config.py @@ -1,80 +1,63 @@ -import os +from typing import override import pytest from flask import Flask from packaging.version import Version from pydantic import SecretStr +from pydantic_settings import BaseSettings, PydanticBaseSettingsSource from yarl import URL from configs.app_config import DifyConfig from enums import DeploymentEdition -def _clear_environment(monkeypatch: pytest.MonkeyPatch) -> None: - for name in tuple(os.environ): - monkeypatch.delenv(name) +class _IsolatedDifyConfig(DifyConfig): + """Load explicit test values and packaging metadata without consulting process state.""" + + @classmethod + @override + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + production_sources = super().settings_customise_sources( + settings_cls, + init_settings, + env_settings, + dotenv_settings, + file_secret_settings, + ) + return init_settings, production_sources[-1] -def _set_basic_config_env(monkeypatch: pytest.MonkeyPatch) -> None: - _clear_environment(monkeypatch) - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") +def _make_config(**values: object) -> DifyConfig: + return _IsolatedDifyConfig(**values) -def test_dify_config_keeps_secret_key_empty_when_missing( - monkeypatch: pytest.MonkeyPatch, - tmp_path, -) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.delenv("SECRET_KEY", raising=False) - monkeypatch.setenv("OPENDAL_FS_ROOT", str(tmp_path)) - - config = DifyConfig(_env_file=None) +def test_dify_config_keeps_secret_key_empty_when_missing(tmp_path) -> None: + config = _make_config(OPENDAL_FS_ROOT=str(tmp_path)) assert config.SECRET_KEY == "" assert not hasattr(config, "OPENDAL_FS_ROOT") assert not (tmp_path / ".dify_secret_key").exists() -def test_dify_config_preserves_explicit_secret_key( - monkeypatch: pytest.MonkeyPatch, - tmp_path, -) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv("SECRET_KEY", "explicit") - monkeypatch.setenv("OPENDAL_FS_ROOT", str(tmp_path)) - - config = DifyConfig(_env_file=None) +def test_dify_config_preserves_explicit_secret_key(tmp_path) -> None: + config = _make_config(SECRET_KEY="explicit", OPENDAL_FS_ROOT=str(tmp_path)) assert config.SECRET_KEY == "explicit" assert not (tmp_path / ".dify_secret_key").exists() -def test_dify_config(monkeypatch: pytest.MonkeyPatch): - # clear system environment variables - _clear_environment(monkeypatch) - - # Set environment variables using monkeypatch - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("HTTP_REQUEST_MAX_WRITE_TIMEOUT", "30") # Custom value for testing - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("HTTP_REQUEST_MAX_READ_TIMEOUT", "300") # Custom value for testing - - # load dotenv file with pydantic-settings - # Disable `.env` loading to ensure test stability across environments - config = DifyConfig(_env_file=None) +def test_dify_config(): + config = _make_config( + HTTP_REQUEST_MAX_WRITE_TIMEOUT="30", + HTTP_REQUEST_MAX_READ_TIMEOUT="300", + ) # constant values assert config.COMMIT_SHA == "" @@ -106,93 +89,66 @@ def test_dify_config(monkeypatch: pytest.MonkeyPatch): pytest.param("pässwörd-🔐", "pässwörd-🔐", id="unicode"), ], ) -def test_init_password_defaults_to_empty_and_preserves_environment_value( - monkeypatch: pytest.MonkeyPatch, +def test_init_password_defaults_to_empty_and_preserves_explicit_value( environment_value: str | None, expected: str, ) -> None: - _set_basic_config_env(monkeypatch) - if environment_value is None: - monkeypatch.delenv("INIT_PASSWORD", raising=False) - else: - monkeypatch.setenv("INIT_PASSWORD", environment_value) - - config = DifyConfig(_env_file=None) + values = {} if environment_value is None else {"INIT_PASSWORD": environment_value} + config = _make_config(**values) assert expected == config.INIT_PASSWORD @pytest.mark.parametrize("edition", list(DeploymentEdition)) -def test_deployment_edition_is_loaded_from_environment( - monkeypatch: pytest.MonkeyPatch, - edition: DeploymentEdition, -) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv("DEPLOYMENT_EDITION", edition.value) - - config = DifyConfig(_env_file=None) +def test_deployment_edition_accepts_every_supported_value(edition: DeploymentEdition) -> None: + config = _make_config(DEPLOYMENT_EDITION=edition.value) assert config.DEPLOYMENT_EDITION is edition -def test_new_user_default_plugin_ids_are_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv( - "NEW_USER_DEFAULT_PLUGIN_IDS", - "langgenius/openai, langgenius/gemini", +def test_new_user_default_plugin_ids_are_parsed() -> None: + config = _make_config( + NEW_USER_DEFAULT_PLUGIN_IDS="langgenius/openai, langgenius/gemini", ) - config = DifyConfig(_env_file=None) - assert config.NEW_USER_DEFAULT_PLUGIN_ID_LIST == [ "langgenius/openai", "langgenius/gemini", ] -def test_turnstile_config_is_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv("TURNSTILE_SECRET_KEY", " test-secret ") - monkeypatch.setenv("TURNSTILE_ALLOWED_HOSTNAMES", "dify.dev, Login.Example.COM. ") - - config = DifyConfig(_env_file=None) +def test_turnstile_config_is_parsed() -> None: + config = _make_config( + TURNSTILE_SECRET_KEY=" test-secret ", + TURNSTILE_ALLOWED_HOSTNAMES="dify.dev, Login.Example.COM. ", + ) assert isinstance(config.TURNSTILE_SECRET_KEY, SecretStr) assert config.TURNSTILE_SECRET_KEY.get_secret_value() == "test-secret" assert frozenset({"dify.dev", "login.example.com"}) == config.TURNSTILE_ALLOWED_HOSTNAME_SET -def test_plugin_remote_install_port_rejects_host_port_spec(monkeypatch: pytest.MonkeyPatch) -> None: +def test_plugin_remote_install_port_rejects_host_port_spec() -> None: """A 'host:port' compose publish spec must produce an actionable error, not an opaque int_parsing traceback.""" - _set_basic_config_env(monkeypatch) - monkeypatch.setenv("PLUGIN_REMOTE_INSTALL_PORT", "127.0.0.1:5003") - with pytest.raises(ValueError, match="must be a bare port number"): - DifyConfig(_env_file=None) + _make_config(PLUGIN_REMOTE_INSTALL_PORT="127.0.0.1:5003") -def test_plugin_remote_install_port_accepts_bare_port(monkeypatch: pytest.MonkeyPatch) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv("PLUGIN_REMOTE_INSTALL_PORT", "5003") - - config = DifyConfig(_env_file=None) +def test_plugin_remote_install_port_accepts_bare_port() -> None: + config = _make_config(PLUGIN_REMOTE_INSTALL_PORT="5003") assert config.PLUGIN_REMOTE_INSTALL_PORT == 5003 -def test_new_user_default_models_are_parsed_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv( - "NEW_USER_DEFAULT_MODELS", - ( +def test_new_user_default_models_are_parsed() -> None: + config = _make_config( + NEW_USER_DEFAULT_MODELS=( "llm:langgenius/openai/openai:gpt-4o-mini, " "text-embedding:langgenius/openai/openai:text-embedding-3-small, " "rerank:langgenius/ollama/ollama:reranker:latest" ), ) - config = DifyConfig(_env_file=None) - assert config.NEW_USER_DEFAULT_MODEL_LIST == [ ("llm", "langgenius/openai/openai", "gpt-4o-mini"), ("text-embedding", "langgenius/openai/openai", "text-embedding-3-small"), @@ -200,34 +156,20 @@ def test_new_user_default_models_are_parsed_from_env(monkeypatch: pytest.MonkeyP ] -def test_new_user_default_models_reject_duplicate_model_types(monkeypatch: pytest.MonkeyPatch) -> None: - _set_basic_config_env(monkeypatch) - monkeypatch.setenv( - "NEW_USER_DEFAULT_MODELS", - "llm:langgenius/openai/openai:gpt-4o-mini,llm:langgenius/anthropic/anthropic:claude-sonnet-4", +def test_new_user_default_models_reject_duplicate_model_types() -> None: + config = _make_config( + NEW_USER_DEFAULT_MODELS=( + "llm:langgenius/openai/openai:gpt-4o-mini,llm:langgenius/anthropic/anthropic:claude-sonnet-4" + ), ) - config = DifyConfig(_env_file=None) - with pytest.raises(ValueError, match="duplicate model type: llm"): _ = config.NEW_USER_DEFAULT_MODEL_LIST -def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch): +def test_http_timeout_defaults(): """Test that HTTP timeout defaults are correctly set""" - # clear system environment variables - _clear_environment(monkeypatch) - - # Set minimal required env vars - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - - # Disable `.env` loading to ensure test stability across environments - config = DifyConfig(_env_file=None) + config = _make_config() # Verify default timeout values assert config.HTTP_REQUEST_MAX_CONNECT_TIMEOUT == 10 @@ -235,56 +177,43 @@ def test_http_timeout_defaults(monkeypatch: pytest.MonkeyPatch): assert config.HTTP_REQUEST_MAX_WRITE_TIMEOUT == 600 -def test_internal_files_url_falls_back_to_server_console_api_url(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001") - - config = DifyConfig(_env_file=None) +def test_internal_files_url_falls_back_to_server_console_api_url(): + config = _make_config(SERVER_CONSOLE_API_URL="http://api:5001") assert config.INTERNAL_FILES_URL == "http://api:5001" -def test_internal_files_url_prefers_explicit_value(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - monkeypatch.setenv("INTERNAL_FILES_URL", "http://files-internal:5001") - monkeypatch.setenv("SERVER_CONSOLE_API_URL", "http://api:5001") - - config = DifyConfig(_env_file=None) +def test_internal_files_url_prefers_explicit_value(): + config = _make_config( + INTERNAL_FILES_URL="http://files-internal:5001", + SERVER_CONSOLE_API_URL="http://api:5001", + ) assert config.INTERNAL_FILES_URL == "http://files-internal:5001" -def test_empty_files_url_overrides_console_api_url_for_relative_browser_uris(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - monkeypatch.setenv("FILES_URL", "") - monkeypatch.setenv("CONSOLE_API_URL", "http://api:5001") - - config = DifyConfig(_env_file=None) +def test_empty_files_url_overrides_console_api_url_for_relative_browser_uris(): + config = _make_config(FILES_URL="", CONSOLE_API_URL="http://api:5001") assert config.FILES_URL == "" -# NOTE: If there is a `.env` file in your Workspace, this test might not succeed as expected. -# This is due to `pymilvus` loading all the variables from the `.env` file into `os.environ`. -def test_flask_configs(monkeypatch: pytest.MonkeyPatch): +def test_flask_configs(): flask_app = Flask("app") - # clear system environment variables - _clear_environment(monkeypatch) - - # Set environment variables using monkeypatch - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("WEB_API_CORS_ALLOW_ORIGINS", "http://127.0.0.1:3000,*") - monkeypatch.setenv("CODE_EXECUTION_ENDPOINT", "http://127.0.0.1:8194/") - - # Disable `.env` loading to ensure test stability across environments - flask_app.config.from_mapping(DifyConfig(_env_file=None).model_dump()) + flask_app.config.from_mapping( + _make_config( + CONSOLE_API_URL="https://example.com", + CONSOLE_WEB_URL="https://example.com", + DB_TYPE="postgresql", + DB_USERNAME="postgres", + DB_PASSWORD="postgres", + DB_HOST="localhost", + DB_PORT="5432", + DB_DATABASE="dify", + WEB_API_CORS_ALLOW_ORIGINS="http://127.0.0.1:3000,*", + CODE_EXECUTION_ENDPOINT="http://127.0.0.1:8194/", + ).model_dump() + ) config = flask_app.config # configs read from pydantic-settings @@ -321,147 +250,67 @@ def test_flask_configs(monkeypatch: pytest.MonkeyPatch): assert str(URL(str(config["CODE_EXECUTION_ENDPOINT"])) / "v1") == "http://127.0.0.1:8194/v1" -def test_inner_api_config_exist(monkeypatch: pytest.MonkeyPatch): - # Set environment variables using monkeypatch - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("INNER_API_KEY", "test-inner-api-key") - - config = DifyConfig() +def test_inner_api_config_exist(): + config = _make_config(INNER_API_KEY="test-inner-api-key") assert config.INNER_API is False assert isinstance(config.INNER_API_KEY, str) assert len(config.INNER_API_KEY) > 0 -def test_db_extras_options_merging(monkeypatch: pytest.MonkeyPatch): +def test_db_extras_options_merging(): """Test that DB_EXTRAS options are merged with the default timezone startup option.""" - # Set environment variables - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("DB_EXTRAS", "options=-c search_path=myschema") - - # Create config - config = DifyConfig() + config = _make_config(DB_EXTRAS="options=-c search_path=myschema") options = config.SQLALCHEMY_ENGINE_OPTIONS["connect_args"]["options"] assert "search_path=myschema" in options assert "timezone=UTC" in options -def test_db_session_timezone_override_can_disable_app_level_timezone_injection(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("DB_EXTRAS", "options=-c search_path=myschema") - monkeypatch.setenv("DB_SESSION_TIMEZONE_OVERRIDE", "") - - config = DifyConfig() +def test_db_session_timezone_override_can_disable_app_level_timezone_injection(): + config = _make_config( + DB_EXTRAS="options=-c search_path=myschema", + DB_SESSION_TIMEZONE_OVERRIDE="", + ) assert config.SQLALCHEMY_ENGINE_OPTIONS["connect_args"] == { "options": "-c search_path=myschema", } -def test_pubsub_redis_url_default(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("REDIS_HOST", "redis.example.com") - monkeypatch.setenv("REDIS_PORT", "6380") - monkeypatch.setenv("REDIS_USERNAME", "user") - monkeypatch.setenv("REDIS_PASSWORD", "pass@word") - monkeypatch.setenv("REDIS_DB", "2") - monkeypatch.setenv("REDIS_USE_SSL", "true") - - config = DifyConfig() +def test_pubsub_redis_url_default(): + config = _make_config( + REDIS_HOST="redis.example.com", + REDIS_PORT="6380", + REDIS_USERNAME="user", + REDIS_PASSWORD="pass@word", + REDIS_DB="2", + REDIS_USE_SSL="true", + ) assert config.normalized_pubsub_redis_url == "rediss://user:pass%40word@redis.example.com:6380/2" assert config.PUBSUB_REDIS_CHANNEL_TYPE == "pubsub" -def test_pubsub_redis_url_override(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("PUBSUB_REDIS_URL", "redis://pubsub-host:6381/5") - - config = DifyConfig() +def test_pubsub_redis_url_override(): + config = _make_config(PUBSUB_REDIS_URL="redis://pubsub-host:6381/5") assert config.normalized_pubsub_redis_url == "redis://pubsub-host:6381/5" -def test_pubsub_redis_url_required_when_default_unavailable(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("REDIS_HOST", "") - +def test_pubsub_redis_url_required_when_default_unavailable(): + config = _make_config(REDIS_HOST="") with pytest.raises(ValueError, match="PUBSUB_REDIS_URL must be set"): - _ = DifyConfig().normalized_pubsub_redis_url + _ = config.normalized_pubsub_redis_url -def test_dify_config_exposes_redis_key_prefix_default(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - - config = DifyConfig(_env_file=None) +def test_dify_config_exposes_redis_key_prefix_default(): + config = _make_config() assert config.REDIS_KEY_PREFIX == "" -def test_dify_config_reads_redis_key_prefix_from_env(monkeypatch: pytest.MonkeyPatch): - _clear_environment(monkeypatch) - - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - monkeypatch.setenv("REDIS_KEY_PREFIX", "enterprise-a") - - config = DifyConfig(_env_file=None) +def test_dify_config_accepts_redis_key_prefix(): + config = _make_config(REDIS_KEY_PREFIX="enterprise-a") assert config.REDIS_KEY_PREFIX == "enterprise-a" @@ -489,7 +338,6 @@ def test_dify_config_reads_redis_key_prefix_from_env(monkeypatch: pytest.MonkeyP ], ) def test_celery_broker_url_with_special_chars_password( - monkeypatch: pytest.MonkeyPatch, broker_url, expected_host, expected_port, @@ -500,24 +348,7 @@ def test_celery_broker_url_with_special_chars_password( """Test that CELERY_BROKER_URL with various formats are handled correctly.""" from kombu.utils.url import parse_url - # clear system environment variables - _clear_environment(monkeypatch) - - # Set up basic required environment variables (following existing pattern) - monkeypatch.setenv("CONSOLE_API_URL", "https://example.com") - monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com") - monkeypatch.setenv("DB_TYPE", "postgresql") - monkeypatch.setenv("DB_USERNAME", "postgres") - monkeypatch.setenv("DB_PASSWORD", "postgres") - monkeypatch.setenv("DB_HOST", "localhost") - monkeypatch.setenv("DB_PORT", "5432") - monkeypatch.setenv("DB_DATABASE", "dify") - - # Set the CELERY_BROKER_URL to test - monkeypatch.setenv("CELERY_BROKER_URL", broker_url) - - # Create config and verify the URL is stored correctly - config = DifyConfig() + config = _make_config(CELERY_BROKER_URL=broker_url) assert broker_url == config.CELERY_BROKER_URL # Test actual parsing behavior using kombu's parse_url (same as production) diff --git a/api/tests/unit_tests/configs/test_file_upload_config.py b/api/tests/unit_tests/configs/test_file_upload_config.py index 666ff3f666e..64c76cb2fb5 100644 --- a/api/tests/unit_tests/configs/test_file_upload_config.py +++ b/api/tests/unit_tests/configs/test_file_upload_config.py @@ -1,23 +1,23 @@ -import pytest - from configs.feature import FileUploadConfig +from tests.unit_tests.configs._isolated_settings import InitSettingsOnly -def test_paid_plan_file_size_limit_uses_its_default(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("UPLOAD_FILE_SIZE_LIMIT", "23") - monkeypatch.delenv("KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN", raising=False) +class _IsolatedFileUploadConfig(InitSettingsOnly, FileUploadConfig): + pass - config = FileUploadConfig() + +def test_paid_plan_file_size_limit_uses_its_default() -> None: + config = _IsolatedFileUploadConfig(UPLOAD_FILE_SIZE_LIMIT="23") assert config.UPLOAD_FILE_SIZE_LIMIT == 23 assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 15 -def test_paid_plan_file_size_limit_can_be_configured_separately(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("UPLOAD_FILE_SIZE_LIMIT", "23") - monkeypatch.setenv("KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN", "50") - - config = FileUploadConfig() +def test_paid_plan_file_size_limit_can_be_configured_separately() -> None: + config = _IsolatedFileUploadConfig( + UPLOAD_FILE_SIZE_LIMIT="23", + KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN="50", + ) assert config.UPLOAD_FILE_SIZE_LIMIT == 23 assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 50 diff --git a/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py b/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py index 5e7ff57dde3..7c2ae5c1f84 100644 --- a/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py +++ b/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py @@ -1,19 +1,20 @@ -import pytest - from configs.middleware.vdb.tidb_on_qdrant_config import TidbOnQdrantConfig +from tests.unit_tests.configs._isolated_settings import InitSettingsOnly -def test_estimated_storage_limits_default(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB", raising=False) +class _IsolatedTidbOnQdrantConfig(InitSettingsOnly, TidbOnQdrantConfig): + pass - config = TidbOnQdrantConfig() + +def test_estimated_storage_limits_default() -> None: + config = _IsolatedTidbOnQdrantConfig() assert config.TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB == "sandbox:60,professional:6400,team:25600" -def test_estimated_storage_limits_custom(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB", "sandbox:61,professional:6500,team:26000") - - config = TidbOnQdrantConfig() +def test_estimated_storage_limits_custom() -> None: + config = _IsolatedTidbOnQdrantConfig( + TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB="sandbox:61,professional:6500,team:26000" + ) assert config.TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB == "sandbox:61,professional:6500,team:26000" diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index 3870d9b1d94..ce90f8f380a 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -1,6 +1,6 @@ import os import shutil -from collections.abc import Iterator +from collections.abc import Callable, Iterator from pathlib import Path from unittest.mock import MagicMock, patch @@ -112,6 +112,26 @@ def reset_secret_key() -> Iterator[None]: dify_config.SECRET_KEY = original +@pytest.fixture +def config_overrides(monkeypatch: pytest.MonkeyPatch) -> Callable[..., None]: + """Temporarily override fields on the shared typed application config. + + Application modules import the same config instance, so mutating known + field names keeps tests scoped without replacing that instance with an + unconstrained mock. ``monkeypatch`` restores every value after the test. + """ + from configs import dify_config + + def apply(**values: object) -> None: + unknown_fields = values.keys() - type(dify_config).model_fields.keys() + if unknown_fields: + raise ValueError(f"Unknown DifyConfig fields: {sorted(unknown_fields)}") + for name, value in values.items(): + monkeypatch.setattr(dify_config, name, value) + + return apply + + @pytest.fixture def _sqlite_engine(_sqlite_database_template: Path, tmp_path: Path) -> Iterator[Engine]: """Create an engine over a pristine per-test copy of the SQLite schema.""" diff --git a/api/tests/unit_tests/controllers/console/app/test_conversation_api.py b/api/tests/unit_tests/controllers/console/app/test_conversation_api.py index 01cc0c1edab..1f563f6c7b4 100644 --- a/api/tests/unit_tests/controllers/console/app/test_conversation_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_conversation_api.py @@ -1,7 +1,7 @@ from __future__ import annotations +from decimal import Decimal from inspect import unwrap -from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -11,13 +11,41 @@ from werkzeug.exceptions import BadRequest, NotFound from controllers.console.app import conversation as conversation_module from core.app.entities.app_invoke_entities import InvokeFrom -from models.enums import ConversationFromSource -from models.model import AppMode, Conversation +from graphon.enums import WorkflowExecutionStatus +from models import Account, App, EndUser, Message, MessageAnnotation, MessageFeedback +from models.enums import ( + ConversationFromSource, + CreatorUserRole, + EndUserType, + FeedbackFromSource, + FeedbackRating, + WorkflowRunTriggeredFrom, +) +from models.model import AppMode, Conversation, IconType +from models.workflow import WorkflowRun, WorkflowType from services.errors.conversation import ConversationNotExistsError -def _make_account(): - return SimpleNamespace(timezone="UTC", id="u1") +def _make_account() -> Account: + account = Account(name="Account", email="account@example.com", timezone="UTC") + account.id = "u1" + return account + + +def _app(*, mode: AppMode = AppMode.CHAT) -> App: + return App( + id="app-1", + tenant_id="tenant-1", + name="Conversation app", + description="", + mode=mode, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) def _conversation(*, conversation_id: str = "c1", app_id: str = "app-1") -> Conversation: @@ -63,7 +91,7 @@ def test_completion_conversation_list_returns_paginated_result( conversation_module.CompletionConversationQuery(), unbound_session, account, - app_model=SimpleNamespace(id="app-1"), + app_model=_app(mode=AppMode.COMPLETION), ) assert response == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []} @@ -88,7 +116,7 @@ def test_completion_conversation_list_invalid_time_range( conversation_module.CompletionConversationQuery(), unbound_session, account, - app_model=SimpleNamespace(id="app-1"), + app_model=_app(mode=AppMode.COMPLETION), ) @@ -112,7 +140,7 @@ def test_chat_conversation_list_advanced_chat_calls_paginate( conversation_module.ChatConversationQuery(), unbound_session, account, - app_model=SimpleNamespace(id="app-1", mode=AppMode.ADVANCED_CHAT), + app_model=_app(mode=AppMode.ADVANCED_CHAT), ) assert response == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []} @@ -122,7 +150,7 @@ def test_get_conversation_updates_read_at(sqlite_session: Session) -> None: sqlite_session.add(conversation) sqlite_session.flush() session = sqlite_session - result = conversation_module._get_conversation(session, _make_account(), SimpleNamespace(id="app-1"), "c1") + result = conversation_module._get_conversation(session, _make_account(), _app(), "c1") assert result is conversation assert conversation.read_at is not None assert conversation.read_account_id == "u1" @@ -131,58 +159,120 @@ def test_get_conversation_updates_read_at(sqlite_session: Session) -> None: def test_get_conversation_missing_raises_not_found(sqlite_session: Session) -> None: session = sqlite_session with pytest.raises(NotFound): - conversation_module._get_conversation(session, _make_account(), SimpleNamespace(id="app-1"), "missing") + conversation_module._get_conversation(session, _make_account(), _app(), "missing") -def test_conversation_response_source_uses_caller_session(unbound_session: Session) -> None: - session = unbound_session - account = object() - annotation = MagicMock() - annotation.account_with_session.return_value = account - message = MagicMock() - conversation = MagicMock() - conversation.inputs_with_session.return_value = {"topic": "support"} - conversation.model_config_with_session.return_value = {"model_id": "model-1"} - conversation.summary_or_query_with_session.return_value = "summary" - conversation.annotated_with_session.return_value = True - conversation.annotation_with_session.return_value = annotation - conversation.message_count_with_session.return_value = 3 - conversation.user_feedback_stats_with_session.return_value = {"like": 2, "dislike": 1} - conversation.admin_feedback_stats_with_session.return_value = {"like": 1, "dislike": 0} - conversation.status_count_with_session.return_value = {"success": 1, "failed": 0} - conversation.first_message_with_session.return_value = message - conversation.from_end_user_session_id_with_session.return_value = "end-user-session" - conversation.from_account_name_with_session.return_value = "Account" +def test_conversation_response_source_uses_caller_session(sqlite_session: Session) -> None: + account = _make_account() + end_user = EndUser( + id="end-user-1", + tenant_id="tenant-1", + app_id="app-1", + type=EndUserType.SERVICE_API, + external_user_id="external-user-1", + name="End user", + session_id="end-user-session", + ) + conversation = _conversation() + conversation.mode = AppMode.ADVANCED_CHAT + conversation.override_model_configs = "{}" + conversation.summary = "summary" + conversation.inputs = {"topic": "support"} + conversation.from_end_user_id = end_user.id + conversation.from_account_id = account.id + message = Message( + id="message-1", + app_id=conversation.app_id, + model_provider=None, + model_id=None, + override_model_configs=None, + conversation_id=conversation.id, + inputs={}, + query="first question", + message={}, + message_tokens=1, + message_unit_price=Decimal(0), + message_price_unit=Decimal("0.001"), + answer="answer", + answer_tokens=1, + answer_unit_price=Decimal(0), + answer_price_unit=Decimal("0.001"), + parent_message_id=None, + provider_response_latency=0, + total_price=Decimal(0), + currency="USD", + error=None, + message_metadata=None, + invoke_from=InvokeFrom.EXPLORE, + from_source=ConversationFromSource.CONSOLE, + from_end_user_id=end_user.id, + from_account_id=account.id, + workflow_run_id="run-1", + app_mode=AppMode.ADVANCED_CHAT, + ) + annotation = MessageAnnotation( + app_id=conversation.app_id, + question="question", + content="annotation", + account_id=account.id, + conversation_id=conversation.id, + message_id=message.id, + ) + feedbacks = [ + MessageFeedback( + app_id=conversation.app_id, + conversation_id=conversation.id, + message_id=message.id, + rating=rating, + from_source=source, + from_account_id=account.id, + ) + for source, rating in ( + (FeedbackFromSource.USER, FeedbackRating.LIKE), + (FeedbackFromSource.USER, FeedbackRating.DISLIKE), + (FeedbackFromSource.ADMIN, FeedbackRating.LIKE), + ) + ] + workflow_run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id=conversation.app_id, + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="1", + graph="{}", + inputs="{}", + status=WorkflowExecutionStatus.SUCCEEDED, + outputs="{}", + error=None, + elapsed_time=1, + total_tokens=1, + total_steps=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account.id, + finished_at=None, + exceptions_count=0, + ) + sqlite_session.add_all([account, end_user, conversation, message, annotation, workflow_run, *feedbacks]) + sqlite_session.commit() - source = conversation_module.ConversationResponseSource(conversation, session=session) + source = conversation_module.ConversationResponseSource(conversation, session=sqlite_session) assert source.inputs == {"topic": "support"} - assert source.model_config == {"model_id": "model-1"} + assert source.model_config == {"model_id": None, "provider": None} assert source.summary_or_query == "summary" assert source.annotated is True annotation_source = source.annotation assert annotation_source is not None assert annotation_source.account is account - assert source.message_count == 3 - assert source.user_feedback_stats == {"like": 2, "dislike": 1} + assert source.message_count == 1 + assert source.user_feedback_stats == {"like": 1, "dislike": 1} assert source.admin_feedback_stats == {"like": 1, "dislike": 0} - assert source.status_count == {"success": 1, "failed": 0} + assert source.status_count == {"success": 1, "failed": 0, "partial_success": 0, "paused": 0} assert source.first_message is not None assert source.from_end_user_session_id == "end-user-session" assert source.from_account_name == "Account" - conversation.model_config_with_session.assert_called_once_with(session=session) - conversation.inputs_with_session.assert_called_once_with(session=session) - conversation.summary_or_query_with_session.assert_called_once_with(session=session) - conversation.annotated_with_session.assert_called_once_with(session=session) - conversation.annotation_with_session.assert_called_once_with(session=session) - conversation.message_count_with_session.assert_called_once_with(session=session) - conversation.user_feedback_stats_with_session.assert_called_once_with(session=session) - conversation.admin_feedback_stats_with_session.assert_called_once_with(session=session) - conversation.status_count_with_session.assert_called_once_with(session=session) - conversation.first_message_with_session.assert_called_once_with(session=session) - conversation.from_end_user_session_id_with_session.assert_called_once_with(session=session) - conversation.from_account_name_with_session.assert_called_once_with(session=session) - annotation.account_with_session.assert_called_once_with(session=session) def test_completion_conversation_delete_maps_not_found( @@ -197,4 +287,4 @@ def test_completion_conversation_delete_maps_not_found( ) session = unbound_session with pytest.raises(NotFound): - method(api, session, _make_account(), app_model=SimpleNamespace(id="app-1"), conversation_id="c1") + method(api, session, _make_account(), app_model=_app(), conversation_id="c1") diff --git a/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py b/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py index f44358e6897..07249da37c8 100644 --- a/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py @@ -2,7 +2,7 @@ from __future__ import annotations from datetime import UTC, datetime from inspect import unwrap -from types import SimpleNamespace +from unittest.mock import PropertyMock, patch import pytest from flask import Flask @@ -13,12 +13,28 @@ from sqlalchemy.orm import Session from controllers.console.app import conversation_variables as conversation_variables_module from factories import variable_factory from graphon.variables.types import SegmentType -from models import ConversationVariable +from models import App, AppMode, ConversationVariable +from models.model import IconType + + +def _app() -> App: + return App( + id="app-1", + tenant_id="tenant-1", + name="Conversation variables app", + description="", + mode=AppMode.ADVANCED_CHAT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) def test_get_conversation_variables_returns_paginated_response( app: Flask, - monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session, ) -> None: @@ -44,17 +60,19 @@ def test_get_conversation_variables_returns_paginated_response( sqlite_session.expire(row) expected_created_at = int(row.created_at.timestamp()) expected_updated_at = int(row.updated_at.timestamp()) - monkeypatch.setattr(conversation_variables_module, "db", SimpleNamespace(engine=sqlite_engine)) - - with app.test_request_context( - "/console/api/apps/app-1/conversation-variables", - method="GET", - query_string={"conversation_id": "conv-1"}, + with ( + app.test_request_context( + "/console/api/apps/app-1/conversation-variables", + method="GET", + query_string={"conversation_id": "conv-1"}, + ), + patch.object(type(conversation_variables_module.db), "engine", new_callable=PropertyMock) as engine, ): + engine.return_value = sqlite_engine response = method( api, conversation_variables_module.ConversationVariablesQuery(conversation_id="conv-1"), - app_model=SimpleNamespace(id="app-1"), + app_model=_app(), ) assert response["page"] == 1 @@ -69,7 +87,6 @@ def test_get_conversation_variables_returns_paginated_response( @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) def test_get_conversation_variables_normalizes_value_type_and_value( app: Flask, - monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session, ) -> None: @@ -87,17 +104,19 @@ def test_get_conversation_variables_normalizes_value_type_and_value( ) sqlite_session.add(ConversationVariable.from_variable(app_id="app-1", conversation_id="conv-1", variable=variable)) sqlite_session.commit() - monkeypatch.setattr(conversation_variables_module, "db", SimpleNamespace(engine=sqlite_engine)) - - with app.test_request_context( - "/console/api/apps/app-1/conversation-variables", - method="GET", - query_string={"conversation_id": "conv-1"}, + with ( + app.test_request_context( + "/console/api/apps/app-1/conversation-variables", + method="GET", + query_string={"conversation_id": "conv-1"}, + ), + patch.object(type(conversation_variables_module.db), "engine", new_callable=PropertyMock) as engine, ): + engine.return_value = sqlite_engine response = method( api, conversation_variables_module.ConversationVariablesQuery(conversation_id="conv-1"), - app_model=SimpleNamespace(id="app-1"), + app_model=_app(), ) assert response["data"][0]["value_type"] == "number" diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py index 0978a79b27c..692b19dc091 100644 --- a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py +++ b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py @@ -1,6 +1,5 @@ import datetime from inspect import unwrap -from types import SimpleNamespace from unittest.mock import patch import pytest @@ -18,16 +17,32 @@ from controllers.console.app.mcp_server import ( MCPServerUpdatePayload, ) from controllers.console.wraps import RBACPermission, RBACResourceScope +from models import Account +from models.account import AccountStatus from models.enums import AppMCPServerStatus -from models.model import AppMCPServer +from models.model import App, AppMCPServer, AppMode, IconType -class _ValidatedResponse: - def __init__(self, payload: dict[str, str]) -> None: - self._payload = payload - - def model_dump(self, mode: str = "json") -> dict[str, str]: - return self._payload +def _app( + *, + app_id: str = "app-1", + tenant_id: str = "tenant-1", + name: str = "Demo App", + description: str = "App description", +) -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name=name, + description=description, + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) def _server( @@ -136,7 +151,7 @@ class TestAppMCPServerController: method = unwrap(api.get) with patch("controllers.console.app.mcp_server.db.session", sqlite_session): - response = method(api, app_model=SimpleNamespace(id="app-1")) + response = method(api, app_model=_app()) assert response == {} @@ -157,7 +172,7 @@ class TestAppMCPServerController: api, req_data, "tenant-1", - app_model=SimpleNamespace(id="app-1", name="Demo App", description="App description"), + app_model=_app(), ) server = sqlite_session.scalar(select(AppMCPServer)) @@ -185,9 +200,7 @@ class TestAppMCPServerController: response = method( api, req_data, - app_model=SimpleNamespace( - id="app-1", tenant_id="tenant-1", name="Demo App", description="App description" - ), + app_model=_app(), ) sqlite_session.expire_all() @@ -233,9 +246,7 @@ class TestAppMCPServerController: method( api, req_data, - app_model=SimpleNamespace( - id="app-1", tenant_id="tenant-1", name="Demo App", description="App description" - ), + app_model=_app(), ) sqlite_session.expire_all() @@ -245,32 +256,36 @@ class TestAppMCPServerController: class TestAppMCPServerRefreshController: - def test_post_refreshes_server_bound_to_app_and_tenant(self): + def test_post_refreshes_server_bound_to_app_and_tenant(self, sqlite_session: Session) -> None: api = AppMCPServerRefreshController() method = unwrap(api.post) - server = SimpleNamespace(server_code="old-code") + server = _server(server_code="old-code") + server.id = "server-1" + tenant_decoy = _server(tenant_id="tenant-2", server_code="tenant-decoy-code") + tenant_decoy.id = "server-2" + app_decoy = _server(app_id="app-2", server_code="app-decoy-code") + app_decoy.id = "server-3" + sqlite_session.add_all([server, tenant_decoy, app_decoy]) + sqlite_session.commit() with ( - patch("controllers.console.app.mcp_server.db.session.scalar", return_value=server) as scalar, - patch("controllers.console.app.mcp_server.db.session.commit") as commit, + patch("controllers.console.app.mcp_server.db.session", sqlite_session), patch("controllers.console.app.mcp_server.AppMCPServer.generate_server_code", return_value="new-code"), - patch( - "controllers.console.app.mcp_server.AppMCPServerResponse.model_validate", - return_value=_ValidatedResponse({"id": "server-1", "server_code": "new-code"}), - ), ): - response = method(api, "tenant-1", app_model=SimpleNamespace(id="app-1")) + response = method(api, "tenant-1", app_model=_app()) - stmt = scalar.call_args.args[0] - compiled = stmt.compile() - statement = str(compiled) - assert "app_mcp_servers.tenant_id" in statement - assert "app_mcp_servers.app_id" in statement - assert "tenant-1" in compiled.params.values() - assert "app-1" in compiled.params.values() - assert server.server_code == "new-code" - commit.assert_called_once() - assert response == {"id": "server-1", "server_code": "new-code"} + sqlite_session.expire_all() + refreshed_server = sqlite_session.get(AppMCPServer, "server-1") + persisted_tenant_decoy = sqlite_session.get(AppMCPServer, "server-2") + persisted_app_decoy = sqlite_session.get(AppMCPServer, "server-3") + assert refreshed_server is not None + assert persisted_tenant_decoy is not None + assert persisted_app_decoy is not None + assert refreshed_server.server_code == "new-code" + assert persisted_tenant_decoy.server_code == "tenant-decoy-code" + assert persisted_app_decoy.server_code == "app-decoy-code" + assert response["id"] == "server-1" + assert response["server_code"] == "new-code" def test_route_is_app_scoped_post(self): route_map = { @@ -291,7 +306,8 @@ class TestAppMCPServerRefreshController: class PermissionCheckedError(Exception): pass - current_user = SimpleNamespace(id="account-1") + current_user = Account(name="Current user", email="user@example.com", status=AccountStatus.ACTIVE) + current_user.id = "account-1" with ( patch("controllers.common.wraps.dify_config.RBAC_ENABLED", True), patch( diff --git a/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py b/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py index a1f2147a895..08b3799de34 100644 --- a/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py @@ -1,7 +1,6 @@ from __future__ import annotations from contextlib import nullcontext -from types import SimpleNamespace from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -16,6 +15,7 @@ from controllers.console.app import ops_trace as ops_trace_module from controllers.console.app import wraps as app_wraps from enums import DeploymentEdition from libs import login as login_lib +from models import Tenant from models.account import Account, AccountStatus, TenantAccountRole from models.model import App, AppMode, IconType @@ -25,19 +25,33 @@ def _make_account(role: TenantAccountRole) -> Account: account.id = "account-123" # type: ignore[assignment] account.status = AccountStatus.ACTIVE account.role = role - account._current_tenant = SimpleNamespace(id="tenant-123") # type: ignore[assignment] + tenant = Tenant(name="Test tenant") + tenant.id = "tenant-123" + account._current_tenant = tenant account._get_current_object = lambda: account # type: ignore[attr-defined] return account -def _make_app() -> SimpleNamespace: - return SimpleNamespace(id="app-123", tenant_id="tenant-123", status="normal", mode="chat") +def _make_app() -> App: + return App( + id="app-123", + tenant_id="tenant-123", + name="Trace app", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) def _patch_console_guards( monkeypatch: pytest.MonkeyPatch, account: Account, - app_model: SimpleNamespace, + app_model: App, *, rbac_enabled: bool = False, ) -> None: 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 82d6add152b..3677b0cbb80 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -650,6 +650,7 @@ def test_draft_workflow_get_projects_agent_node_job_to_graph(monkeypatch: pytest "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", }, } ], @@ -663,6 +664,7 @@ def test_draft_workflow_get_projects_agent_node_job_to_graph(monkeypatch: pytest "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": "Summarize it.", "agent_declared_outputs": [{"name": "summary", "type": "string"}], }, diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py index d568744292c..56049936f27 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py @@ -3,7 +3,6 @@ from __future__ import annotations from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime -from types import SimpleNamespace from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -16,7 +15,9 @@ from controllers.console.app import workflow_comment as workflow_comment_module from controllers.console.app import wraps as app_wraps from enums import DeploymentEdition from libs import login as login_lib +from models import App, Tenant, WorkflowComment, WorkflowCommentMention, WorkflowCommentReply from models.account import Account, AccountStatus, TenantAccountRole +from models.model import AppMode, IconType JAN_1_2024_NOON = datetime(2024, 1, 1, 12, 0, 0) JAN_1_2024_NOON_TS = int(JAN_1_2024_NOON.timestamp()) @@ -33,16 +34,30 @@ def _make_account(role: TenantAccountRole) -> Account: account.status = AccountStatus.ACTIVE account.role = role account.id = "account-123" # type: ignore[assignment] - account._current_tenant = SimpleNamespace(id="tenant-123") # type: ignore[attr-defined] + tenant = Tenant(name="Test tenant") + tenant.id = "tenant-123" + account._current_tenant = tenant account._get_current_object = lambda: account # type: ignore[attr-defined] return account -def _make_app() -> SimpleNamespace: - return SimpleNamespace(id="app-123", tenant_id="tenant-123", status="normal", mode="workflow") +def _make_app() -> App: + return App( + id="app-123", + tenant_id="tenant-123", + name="Workflow comments app", + description="", + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) -def _patch_console_guards(monkeypatch: pytest.MonkeyPatch, account: Account, app_model: SimpleNamespace) -> None: +def _patch_console_guards(monkeypatch: pytest.MonkeyPatch, account: Account, app_model: App) -> None: monkeypatch.setattr(login_lib.dify_config, "LOGIN_DISABLED", True) monkeypatch.setattr(login_lib, "current_user", account) monkeypatch.setattr(login_lib, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) @@ -246,29 +261,29 @@ def test_list_comments_serializes_response_model(app: Flask, monkeypatch: pytest app_model = _make_app() _patch_console_guards(monkeypatch, account, app_model) - comment_author = SimpleNamespace( - id="account-123", + comment_author = Account( name="tester", email="tester@example.com", avatar="https://example.com/avatar.png", + status=AccountStatus.ACTIVE, ) - comment = SimpleNamespace( - id="comment-1", + comment_author.id = "account-123" + comment = WorkflowComment( + tenant_id="tenant-123", + app_id="app-123", position_x=1.5, position_y=2.5, content="hello", created_by="account-123", - created_by_account=comment_author, - created_at=1_700_000_000, - updated_at=1_700_000_001, resolved=False, resolved_at=None, resolved_by=None, - resolved_by_account=None, - reply_count=0, - mention_count=0, - participants=[comment_author], ) + comment.id = "comment-1" + comment.created_at = datetime.fromtimestamp(1_700_000_000) + comment.updated_at = datetime.fromtimestamp(1_700_000_001) + comment.cache_created_by_account(comment_author) + comment.cache_resolved_by_account(None) get_comments_mock = MagicMock(return_value=[comment]) monkeypatch.setattr(workflow_comment_module.WorkflowCommentService, "get_comments", get_comments_mock) @@ -317,48 +332,49 @@ def test_get_comment_serializes_detail_response_model(app: Flask, monkeypatch: p app_model = _make_app() _patch_console_guards(monkeypatch, account, app_model) - comment_author = SimpleNamespace( - id="account-123", + comment_author = Account( name="tester", email="tester@example.com", avatar="https://example.com/avatar.png", + status=AccountStatus.ACTIVE, ) - mentioned_user = SimpleNamespace( - id="account-456", + comment_author.id = "account-123" + mentioned_user = Account( name="mentioned", email="mentioned@example.com", avatar=None, + status=AccountStatus.ACTIVE, ) - comment = SimpleNamespace( - id="comment-1", + mentioned_user.id = "account-456" + comment = WorkflowComment( + tenant_id="tenant-123", + app_id="app-123", position_x=1.5, position_y=2.5, content="hello", created_by="account-123", - created_by_account=comment_author, - created_at=JAN_1_2024_NOON, - updated_at=JAN_1_2024_1201, resolved=True, resolved_at=JAN_1_2024_1202, resolved_by="account-123", - resolved_by_account=comment_author, - replies=[ - SimpleNamespace( - id="reply-1", - content="reply", - created_by="account-456", - created_by_account=mentioned_user, - created_at=JAN_1_2024_1203, - ) - ], - mentions=[ - SimpleNamespace( - mentioned_user_id="account-456", - mentioned_user_account=mentioned_user, - reply_id="reply-1", - ) - ], ) + comment.id = "comment-1" + comment.created_at = JAN_1_2024_NOON + comment.updated_at = JAN_1_2024_1201 + comment.cache_created_by_account(comment_author) + comment.cache_resolved_by_account(comment_author) + reply = WorkflowCommentReply(comment_id=comment.id, content="reply", created_by="account-456") + reply.id = "reply-1" + reply.created_at = JAN_1_2024_1203 + reply.updated_at = JAN_1_2024_1203 + reply.cache_created_by_account(mentioned_user) + mention = WorkflowCommentMention( + comment_id=comment.id, + mentioned_user_id="account-456", + reply_id=reply.id, + ) + mention.cache_mentioned_user_account(mentioned_user) + comment.replies.append(reply) + comment.mentions.append(mention) get_comment_mock = MagicMock(return_value=comment) monkeypatch.setattr(workflow_comment_module.WorkflowCommentService, "get_comment", get_comment_mock) diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py index 942698db3bb..73ed8471caf 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py @@ -3,13 +3,30 @@ from __future__ import annotations from inspect import unwrap -from types import SimpleNamespace import pytest from flask import Flask from controllers.console.app import workflow as workflow_module from controllers.console.app.workflow import ConvertToWorkflowApi +from models import Account, App, AppMode +from models.model import IconType + + +def _app(app_id: str) -> App: + return App( + id=app_id, + tenant_id="tenant-1", + name=f"App {app_id}", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) class TestConvertToWorkflowApi: @@ -21,11 +38,12 @@ class TestConvertToWorkflowApi: self, api: ConvertToWorkflowApi, app: Flask, monkeypatch: pytest.MonkeyPatch ) -> None: method = unwrap(api.post) + new_app = _app("new-app-1") monkeypatch.setattr( workflow_module, "WorkflowService", - lambda: SimpleNamespace(convert_to_workflow=lambda **_kwargs: SimpleNamespace(id="new-app-1")), + lambda: type("WorkflowServiceStub", (), {"convert_to_workflow": lambda self, **_kwargs: new_app})(), ) monkeypatch.setattr( workflow_module, @@ -38,11 +56,13 @@ class TestConvertToWorkflowApi: method="POST", json={}, ): + current_user = Account(name="Current user", email="user@example.com") + current_user.id = "u1" response = method( api, current_tenant_id="tenant-1", - current_user=SimpleNamespace(id="u1"), - app_model=SimpleNamespace(id="app-1"), + current_user=current_user, + app_model=_app("app-1"), ) assert response["new_app_id"] == "new-app-1" 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 507a38e0155..a9f38ac504e 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 @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass -from types import SimpleNamespace from unittest.mock import ANY, MagicMock import pytest @@ -13,8 +12,9 @@ from controllers.console.app import workflow as workflow_module from controllers.console.app import wraps as app_wraps from enums import DeploymentEdition from libs import login as login_lib +from models import App, Tenant from models.account import Account, AccountStatus, TenantAccountRole -from models.model import AppMode +from models.model import AppMode, IconType def _make_account() -> Account: @@ -22,16 +22,30 @@ def _make_account() -> Account: account.status = AccountStatus.ACTIVE account.role = TenantAccountRole.OWNER account.id = "account-123" # type: ignore[assignment] - account._current_tenant = SimpleNamespace(id="tenant-123") # type: ignore[attr-defined] + tenant = Tenant(name="Test tenant") + tenant.id = "tenant-123" + account._current_tenant = tenant account._get_current_object = lambda: account # type: ignore[attr-defined] return account -def _make_app(mode: AppMode) -> SimpleNamespace: - return SimpleNamespace(id="app-123", tenant_id="tenant-123", mode=mode.value) +def _make_app(mode: AppMode) -> App: + return App( + id="app-123", + tenant_id="tenant-123", + name="Human input app", + description="", + mode=mode, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) -def _patch_console_guards(monkeypatch: pytest.MonkeyPatch, account: Account, app_model: SimpleNamespace) -> None: +def _patch_console_guards(monkeypatch: pytest.MonkeyPatch, account: Account, app_model: App) -> None: # Skip setup and auth guardrails monkeypatch.setattr("configs.dify_config.DEPLOYMENT_EDITION", DeploymentEdition.CLOUD) monkeypatch.setattr(login_lib.dify_config, "LOGIN_DISABLED", True) 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 8cca6ab9f41..ab60743a4e8 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 @@ -24,7 +24,6 @@ from __future__ import annotations import json from collections.abc import Iterator -from typing import Any from unittest.mock import ANY, MagicMock from uuid import UUID @@ -32,6 +31,8 @@ import pytest from controllers.console.app import workflow_node_output_inspector as ctrl from graphon.enums import WorkflowExecutionStatus +from models import App, AppMode +from models.model import IconType from services.workflow.inspector_events import InspectorMessage from services.workflow.node_output_inspector_service import ( NodeOutputInspectorError, @@ -47,13 +48,25 @@ from services.workflow.node_output_inspector_service import ( @pytest.fixture -def app_model() -> Any: - """A minimal ``App`` stub the controller passes through to the service. +def app_model() -> App: + """A real transient ``App`` the controller passes through to the service. The SSE generator never reads its attributes — just forwards it — so a - sentinel object is enough. + transient mapped instance is sufficient. """ - return MagicMock(name="App", tenant_id="tenant-1", id="app-1") + return App( + id="app-1", + tenant_id="tenant-1", + name="Inspector app", + description="", + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) @pytest.fixture diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py index 4986331e5e7..c1176d4155b 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py @@ -1,16 +1,26 @@ from __future__ import annotations +import json from datetime import UTC, datetime from inspect import unwrap -from types import SimpleNamespace from typing import Any import pytest from flask import Flask from flask_restx import marshal +from sqlalchemy.orm import Session from controllers.console.app import workflow_run as workflow_run_module -from models import Account +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from models import Account, App, AppMode +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.model import IconType +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowType, +) def _serialize_200_response(handler, payload: Any) -> Any: @@ -24,75 +34,112 @@ def _serialize_200_response(handler, payload: Any) -> Any: return payload -def _account() -> SimpleNamespace: - return SimpleNamespace(id="account-1", name="Alice", email="alice@example.com") - - -def _current_account() -> Account: +def _account(session: Session) -> Account: account = Account(name="Alice", email="alice@example.com") account.id = "account-1" + session.add(account) + session.commit() return account -def _workflow_run_summary(**overrides) -> SimpleNamespace: +def _app() -> App: + return App( + id="app-1", + tenant_id="tenant-1", + name="Workflow run app", + description="", + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, + ) + + +def _workflow_run_summary(session: Session, **overrides: object) -> WorkflowRun: created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - payload = { - "id": "run-1", - "version": "v1", - "status": "succeeded", - "elapsed_time": 1.5, - "total_tokens": 10, - "total_steps": 2, - "created_by_account": _account(), - "created_at": created_at, - "finished_at": created_at, - "exceptions_count": 0, - "retry_index": 0, - } - payload.update(overrides) - return SimpleNamespace(**payload) + workflow_run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="v1", + graph='{"nodes": []}', + inputs='{"query": "hello"}', + status=WorkflowExecutionStatus.SUCCEEDED, + outputs='{"answer": "world"}', + error=None, + elapsed_time=1.5, + total_tokens=10, + total_steps=2, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=created_at, + finished_at=created_at, + exceptions_count=0, + ) + for name, value in overrides.items(): + setattr(workflow_run, name, value) + workflow_run.retry_index = 0 + session.add(workflow_run) + session.commit() + return workflow_run -def _workflow_run_node_execution(**overrides) -> SimpleNamespace: +def _workflow_run_node_execution(session: Session) -> WorkflowNodeExecutionModel: created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - payload = { - "id": "node-exec-1", - "index": 1, - "predecessor_node_id": None, - "node_id": "node-1", - "node_type": "start", - "title": "Start", - "inputs_dict": {"query": "hello"}, - "process_data_dict": {"step": "prepared"}, - "outputs_dict": {"answer": "world"}, - "status": "succeeded", - "error": None, - "elapsed_time": 1.0, - "execution_metadata_dict": {"total_tokens": 3}, - "extras": {}, - "created_at": created_at, - "created_by_role": "account", - "created_by_account": _account(), - "created_by_end_user": None, - "finished_at": created_at, - "inputs_truncated": False, - "outputs_truncated": False, - "process_data_truncated": False, - } - payload.update(overrides) - return SimpleNamespace(**payload) + execution = WorkflowNodeExecutionModel( + id="node-exec-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id="run-1", + index=1, + predecessor_node_id=None, + node_execution_id="node-execution-1", + node_id="node-1", + node_type="start", + title="Start", + agent_workspace_binding_id=None, + inputs=json.dumps({"query": "hello"}), + process_data=json.dumps({"step": "prepared"}), + outputs=json.dumps({"answer": "world"}), + status=WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=1.0, + execution_metadata=json.dumps({"total_tokens": 3}), + created_at=created_at, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + finished_at=created_at, + ) + execution.offload_data = [] + session.add(execution) + session.commit() + return execution -def test_workflow_run_list_returns_frontend_history_contract(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_workflow_run_list_returns_frontend_history_contract( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary(sqlite_session) + class WorkflowRunService: def get_paginate_workflow_runs(self, **_kwargs): return { "limit": 10, "has_more": False, - "data": [_workflow_run_summary()], + "data": [workflow_run], } monkeypatch.setattr(workflow_run_module, "WorkflowRunService", WorkflowRunService) + monkeypatch.setattr(workflow_run_module.db, "session", sqlite_session) api = workflow_run_module.WorkflowRunListApi() handler = unwrap(api.get) @@ -101,7 +148,7 @@ def test_workflow_run_list_returns_frontend_history_contract(app: Flask, monkeyp payload = handler( api, workflow_run_module.WorkflowRunListQuery(limit=10), - app_model=SimpleNamespace(id="app-1", tenant_id="tenant-1"), + app_model=_app(), ) response = _serialize_200_response(api.get, payload) @@ -123,21 +170,26 @@ def test_workflow_run_list_returns_frontend_history_contract(app: Flask, monkeyp } -def test_advanced_chat_workflow_run_list_keeps_message_fields(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: +def test_advanced_chat_workflow_run_list_keeps_message_fields( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary( + sqlite_session, + conversation_id="conversation-1", + message_id="message-1", + ) + class WorkflowRunService: def get_paginate_advanced_chat_workflow_runs(self, **_kwargs): return { "limit": 1, "has_more": True, - "data": [ - _workflow_run_summary( - conversation_id="conversation-1", - message_id="message-1", - ) - ], + "data": [workflow_run], } monkeypatch.setattr(workflow_run_module, "WorkflowRunService", WorkflowRunService) + monkeypatch.setattr(workflow_run_module.db, "session", sqlite_session) api = workflow_run_module.AdvancedChatAppWorkflowRunListApi() handler = unwrap(api.get) @@ -146,7 +198,7 @@ def test_advanced_chat_workflow_run_list_keeps_message_fields(app: Flask, monkey payload = handler( api, workflow_run_module.WorkflowRunListQuery(limit=1), - app_model=SimpleNamespace(id="app-1", tenant_id="tenant-1"), + app_model=_app(), ) response = _serialize_200_response(api.get, payload) @@ -155,38 +207,24 @@ def test_advanced_chat_workflow_run_list_keeps_message_fields(app: Flask, monkey assert response["data"][0]["message_id"] == "message-1" -def test_workflow_run_detail_returns_frontend_detail_contract(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - workflow_run = SimpleNamespace( - id="run-1", - version="v1", - graph_dict={"nodes": []}, - inputs_dict={"query": "hello"}, - status="succeeded", - outputs_dict={"answer": "world"}, - error=None, - elapsed_time=1.5, - total_tokens=10, - total_steps=2, - created_by_role="account", - created_by_account=_account(), - created_by_end_user=None, - created_at=created_at, - finished_at=created_at, - exceptions_count=0, - ) +def test_workflow_run_detail_returns_frontend_detail_contract( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary(sqlite_session) class WorkflowRunService: def get_workflow_run(self, **_kwargs): return workflow_run monkeypatch.setattr(workflow_run_module, "WorkflowRunService", WorkflowRunService) + monkeypatch.setattr(workflow_run_module.db, "session", sqlite_session) api = workflow_run_module.WorkflowRunDetailApi() handler = unwrap(api.get) with app.test_request_context("/apps/app-1/workflow-runs/run-1", method="GET"): - payload = handler(api, app_model=SimpleNamespace(id="app-1", tenant_id="tenant-1"), run_id="run-1") + payload = handler(api, app_model=_app(), run_id="run-1") response = _serialize_200_response(api.get, payload) @@ -211,21 +249,23 @@ def test_workflow_run_detail_returns_frontend_detail_contract(app: Flask, monkey def test_workflow_run_node_executions_return_frontend_trace_contract( - app: Flask, monkeypatch: pytest.MonkeyPatch + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: + account = _account(sqlite_session) + execution = _workflow_run_node_execution(sqlite_session) + class WorkflowRunService: def get_workflow_run_node_executions(self, **_kwargs): - return [_workflow_run_node_execution()] + return [execution] monkeypatch.setattr(workflow_run_module, "WorkflowRunService", WorkflowRunService) + monkeypatch.setattr(workflow_run_module.db, "session", sqlite_session) api = workflow_run_module.WorkflowRunNodeExecutionListApi() handler = unwrap(api.get) with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions", method="GET"): - payload = handler( - api, _current_account(), app_model=SimpleNamespace(id="app-1", tenant_id="tenant-1"), run_id="run-1" - ) + payload = handler(api, account, app_model=_app(), run_id="run-1") response = _serialize_200_response(api.get, payload) diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_register.py b/api/tests/unit_tests/controllers/console/auth/test_email_register.py index 6c945dc5f57..9a8b4db80d0 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_email_register.py +++ b/api/tests/unit_tests/controllers/console/auth/test_email_register.py @@ -154,6 +154,7 @@ class TestEmailRegisterResetApi: password="ValidPass123!", timezone=None, language=None, + ip_address="127.0.0.1", ) mock_reset_login_rate.assert_called_once_with("invitee@example.com") mock_revoke_token.assert_called_once_with("token-123") @@ -211,6 +212,7 @@ class TestEmailRegisterResetApi: password="ValidPass123!", timezone="Asia/Shanghai", language=None, + ip_address="127.0.0.1", ) mock_reset_login_rate.assert_called_once_with("invitee@example.com") mock_revoke_token.assert_called_once_with("token-123") @@ -268,6 +270,7 @@ class TestEmailRegisterResetApi: password="ValidPass123!", timezone=None, language="zh-Hans", + ip_address="127.0.0.1", ) mock_reset_login_rate.assert_called_once_with("invitee@example.com") mock_revoke_token.assert_called_once_with("token-123") diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py b/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py index 14af718590a..385ea0834fe 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py +++ b/api/tests/unit_tests/controllers/console/auth/test_email_register_language.py @@ -25,6 +25,7 @@ def test_create_new_account_uses_requested_language(mock_create_account): password="ValidPass123!", interface_language="zh-Hans", timezone="Asia/Shanghai", + ip_address=None, session=ANY, ) diff --git a/api/tests/unit_tests/controllers/console/auth/test_email_verification.py b/api/tests/unit_tests/controllers/console/auth/test_email_verification.py index 4c8173ce80d..0279936f635 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_email_verification.py +++ b/api/tests/unit_tests/controllers/console/auth/test_email_verification.py @@ -462,16 +462,19 @@ class TestEmailCodeLoginApi: mock_login.return_value = mock_token_pair # Act - with app.test_request_context( - "/email-code-login/validity", - method="POST", - json={ - "email": "newuser@example.com", - "code": encode_code("123456"), - "token": "valid_token", - "language": "en-US", - "timezone": "Asia/Shanghai", - }, + with ( + patch("controllers.console.auth.login.extract_remote_ip", return_value="203.0.113.10"), + app.test_request_context( + "/email-code-login/validity", + method="POST", + json={ + "email": "newuser@example.com", + "code": encode_code("123456"), + "token": "valid_token", + "language": "en-US", + "timezone": "Asia/Shanghai", + }, + ), ): api = EmailCodeLoginApi() response = api.post() @@ -483,6 +486,7 @@ class TestEmailCodeLoginApi: name="newuser@example.com", interface_language="en-US", timezone="Asia/Shanghai", + ip_address="203.0.113.10", session=ANY, ) diff --git a/api/tests/unit_tests/controllers/console/auth/test_oauth.py b/api/tests/unit_tests/controllers/console/auth/test_oauth.py index b8ec042f341..114612fe5e8 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_oauth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_oauth.py @@ -211,11 +211,21 @@ class TestOAuthCallback: mock_generate_account.return_value = (oauth_setup["account"], True) mock_account_service.login.return_value = oauth_setup["token_pair"] - with app.test_request_context("/auth/oauth/github/callback?code=test_code"): + with ( + patch("controllers.console.auth.oauth.extract_remote_ip", return_value="203.0.113.10"), + app.test_request_context("/auth/oauth/github/callback?code=test_code"), + ): resource.get("github") oauth_setup["provider"].get_access_token.assert_called_once_with("test_code") oauth_setup["provider"].get_user_info.assert_called_once_with("access_token") + mock_generate_account.assert_called_once_with( + "github", + oauth_setup["provider"].get_user_info.return_value, + timezone=None, + language=None, + ip_address="203.0.113.10", + ) mock_redirect.assert_called_once_with("http://localhost:3000?oauth_new_user=true") @pytest.mark.parametrize( @@ -529,6 +539,7 @@ class TestAccountGeneration: provider="github", language="en-US", timezone=None, + ip_address=None, session=ANY, ) else: @@ -563,6 +574,7 @@ class TestAccountGeneration: provider="github", language="en-US", timezone=None, + ip_address=None, session=ANY, ) @@ -595,6 +607,7 @@ class TestAccountGeneration: provider="github", language="zh-Hans", timezone="Asia/Shanghai", + ip_address=None, session=ANY, ) @@ -627,6 +640,7 @@ class TestAccountGeneration: provider="github", language="zh-Hans", timezone=None, + ip_address=None, session=ANY, ) diff --git a/api/tests/unit_tests/controllers/console/auth/test_oauth_timezone.py b/api/tests/unit_tests/controllers/console/auth/test_oauth_timezone.py index eaf339b62a6..bb21d6f8ce9 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_oauth_timezone.py +++ b/api/tests/unit_tests/controllers/console/auth/test_oauth_timezone.py @@ -56,7 +56,9 @@ def test_generate_account_registers_with_browser_timezone( user_info = OAuthUserInfo(id="github-123", name="Test User", email="User@Example.com") with app.test_request_context(headers={"Accept-Language": "zh-Hans,zh;q=0.9"}): - result, oauth_new_user = _generate_account("github", user_info, timezone="Asia/Shanghai") + result, oauth_new_user = _generate_account( + "github", user_info, timezone="Asia/Shanghai", ip_address="203.0.113.10" + ) assert result is account assert oauth_new_user is True @@ -68,6 +70,7 @@ def test_generate_account_registers_with_browser_timezone( provider="github", language="zh-Hans", timezone="Asia/Shanghai", + ip_address="203.0.113.10", session=ANY, ) mock_link_account.assert_called_once_with("github", "github-123", account, session=ANY) @@ -100,6 +103,7 @@ def test_generate_account_prefers_state_language_over_accept_language( provider="github", language="zh-Hans", timezone=None, + ip_address=None, session=ANY, ) mock_link_account.assert_called_once_with("github", "github-123", account, session=ANY) 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 4bdc862b651..63a61586ef9 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 @@ -151,6 +151,7 @@ class TestDatasourcePluginOAuthAuthorizationUrl: plugin_id="langgenius/notion_datasource", provider="notion", credential_id="cred-1", + extra_data={"visibility": "only_me"}, ) get_authorization_url.assert_called_once() assert get_authorization_url.call_args.kwargs["tenant_id"] == "tenant-1" @@ -250,6 +251,10 @@ class TestDatasourceOAuthCallback: assert response.status_code == 302 assert "/oauth-callback" in response.location add_oauth_provider.assert_called_once() + # Legacy context without a visibility key falls back to ONLY_ME, and + # the callback now also propagates the creator's user_id. + from models.enums import PermissionEnum + assert add_oauth_provider.call_args.kwargs == { "tenant_id": "tenant-1", "provider_id": add_oauth_provider.call_args.kwargs["provider_id"], @@ -257,6 +262,8 @@ class TestDatasourceOAuthCallback: "name": "Workspace Bot", "expire_at": expires_at, "credentials": {"token": "abc"}, + "user_id": "user-1", + "visibility": PermissionEnum.ONLY_ME, } assert str(add_oauth_provider.call_args.kwargs["provider_id"]) == _PROVIDER_ID @@ -645,6 +652,7 @@ class TestDatasourceAuthListApi: def test_list_success(self, app: Flask): api = DatasourceAuthListApi() method = inspect.unwrap(api.get) + user = MagicMock(id="user-1") with ( app.test_request_context("/"), @@ -652,16 +660,20 @@ class TestDatasourceAuthListApi: DatasourceProviderService, "get_all_datasource_credentials", return_value=[_datasource_auth()], - ), + ) as get_all, ): - response, status = method(api, "tenant-1") + response, status = method(api, "tenant-1", user) assert status == 200 assert response == {"result": [_datasource_auth()]} + # user is threaded through so list_datasource_credentials applies the + # visibility filter for the current viewer. + assert get_all.call_args.kwargs["user"] is user def test_auth_list_empty(self, app: Flask): api = DatasourceAuthListApi() method = inspect.unwrap(api.get) + user = MagicMock(id="user-1") with ( app.test_request_context("/"), @@ -671,7 +683,7 @@ class TestDatasourceAuthListApi: return_value=[], ), ): - response, status = method(api, "tenant-1") + response, status = method(api, "tenant-1", user) assert status == 200 assert response["result"] == [] @@ -679,6 +691,7 @@ class TestDatasourceAuthListApi: def test_hardcode_list_empty(self, app: Flask): api = DatasourceHardCodeAuthListApi() method = inspect.unwrap(api.get) + user = MagicMock(id="user-1") with ( app.test_request_context("/"), @@ -688,7 +701,7 @@ class TestDatasourceAuthListApi: return_value=[], ), ): - response, status = method(api, "tenant-1") + response, status = method(api, "tenant-1", user) assert status == 200 assert response["result"] == [] @@ -698,6 +711,7 @@ class TestDatasourceHardCodeAuthListApi: def test_list_success(self, app: Flask): api = DatasourceHardCodeAuthListApi() method = inspect.unwrap(api.get) + user = MagicMock(id="user-1") with ( app.test_request_context("/"), @@ -705,11 +719,12 @@ class TestDatasourceHardCodeAuthListApi: DatasourceProviderService, "get_hard_code_datasource_credentials", return_value=[_datasource_auth()], - ), + ) as get_hardcode, ): - response, status = method(api, "tenant-1") + response, status = method(api, "tenant-1", user) assert status == 200 + assert get_hardcode.call_args.kwargs["user"] is user class TestDatasourceAuthOauthCustomClient: diff --git a/api/tests/unit_tests/controllers/console/explore/test_wraps.py b/api/tests/unit_tests/controllers/console/explore/test_wraps.py index f2eb8523bbf..a1da7916f02 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/explore/test_wraps.py @@ -1,8 +1,9 @@ -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from unittest.mock import patch from uuid import uuid4 import pytest -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, scoped_session from werkzeug.exceptions import Forbidden, NotFound import controllers.console.explore.wraps as wraps_module @@ -20,14 +21,20 @@ from controllers.console.explore.wraps import ( trial_feature_enable, user_allowed_to_access_app, ) -from models import AccountTrialAppRecord, App, AppMode, InstalledApp, TrialApp +from models import Account, AccountTrialAppRecord, App, AppMode, InstalledApp, TrialApp def _bind_database(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - session_proxy = MagicMock(wraps=sqlite_session) - session_proxy.return_value = sqlite_session - monkeypatch.setattr(wraps_module.db, "session", session_proxy) - monkeypatch.setattr(model_module.db, "session", session_proxy) + session_registry = scoped_session(lambda: sqlite_session) + monkeypatch.setattr(wraps_module.db, "session", session_registry) + monkeypatch.setattr(model_module.db, "session", session_registry) + + +def _account(*, account_id: str | None = None) -> Account: + account = Account(name="Explore user", email="user@example.com") + if account_id is not None: + account.id = account_id + return account def _app() -> App: @@ -67,7 +74,7 @@ def test_installed_app_required_not_found( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(), tenant_id), + return_value=(_account(), tenant_id), ): with pytest.raises(NotFound): view(str(uuid4())) @@ -91,7 +98,7 @@ def test_installed_app_required_app_deleted( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(), tenant_id), + return_value=(_account(), tenant_id), ): with pytest.raises(NotFound): view(installed_app_id) @@ -116,7 +123,7 @@ def test_installed_app_required_success( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(), app.tenant_id), + return_value=(_account(), app.tenant_id), ): result = view(installed_app.id) @@ -126,19 +133,18 @@ def test_installed_app_required_success( def test_user_allowed_to_access_app_denied(): - installed_app = MagicMock(app_id="app-1") + installed_app = _installed_app(app_id="app-1", tenant_id="tenant-1") @user_allowed_to_access_app def view(installed_app): return "ok" - feature = MagicMock() - feature.webapp_auth.enabled = True + feature = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=True)) with ( patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(id="user-1"), None), + return_value=(_account(account_id="user-1"), None), ), patch( "controllers.console.explore.wraps.FeatureService.get_system_features", @@ -154,19 +160,18 @@ def test_user_allowed_to_access_app_denied(): def test_user_allowed_to_access_app_success(): - installed_app = MagicMock(app_id="app-1") + installed_app = _installed_app(app_id="app-1", tenant_id="tenant-1") @user_allowed_to_access_app def view(installed_app): return "ok" - feature = MagicMock() - feature.webapp_auth.enabled = True + feature = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=True)) with ( patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(id="user-1"), None), + return_value=(_account(account_id="user-1"), None), ), patch( "controllers.console.explore.wraps.FeatureService.get_system_features", @@ -193,7 +198,7 @@ def test_trial_app_required_not_allowed( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(id=str(uuid4())), None), + return_value=(_account(account_id=str(uuid4())), None), ): with pytest.raises(TrialAppNotAllowed): view(str(uuid4())) @@ -218,7 +223,7 @@ def test_trial_app_required_limit_exceeded( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(id=account_id), None), + return_value=(_account(account_id=account_id), None), ): with pytest.raises(TrialAppLimitExceeded): view(app.id) @@ -243,7 +248,7 @@ def test_trial_app_required_success( with patch( "controllers.console.explore.wraps.current_account_with_tenant", - return_value=(MagicMock(id=account_id), None), + return_value=(_account(account_id=account_id), None), ): result = view(app.id) 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 8184adaad97..3e8df467ad1 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 @@ -764,3 +764,187 @@ def test_resolve_identity_mode_off_is_passthrough_when_not_enterprise( monkeypatch.setattr(controller_module.dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY) assert controller_module._resolve_identity_mode(None, current=identity_mode.OFF) == identity_mode.OFF + + +# --- OAuth start + callback: visibility passthrough (PR #39840 review) --- +# +# The OAuth flow stashes the user-chosen visibility in the proxy context on the +# start endpoint, then reads it back on the callback to persist the credential. +# These tests pin down: (a) the value round-trips for the two supported levels, +# and (b) unknown/missing values collapse to ONLY_ME on both ends. + + +_OAUTH_PROVIDER_PATH = "langgenius/github/github" + + +def _invoke_oauth_start( + controller_module: ModuleType, + app: Flask, + user: Account, + tenant_id: str, + *, + query_string: str = "", + provider: str = _OAUTH_PROVIDER_PATH, +): + """Call ToolPluginOAuthApi.get bypassing decorators (matches the pattern + used by test_builtin_provider_credentials_get_reads_repeated_include_ids).""" + api = controller_module.ToolPluginOAuthApi() + path = f"/oauth/plugin/{provider}/tool/authorization-url" + if query_string: + path = f"{path}?{query_string}" + with app.test_request_context(path, method="GET"): + return unwrap(api.get)(api, tenant_id, user, provider=provider) + + +@contextmanager +def _mock_oauth_start_deps(controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch): + """Mock everything ToolPluginOAuthApi.get calls out to and capture the + create_proxy_context kwargs so tests can assert on extra_data.""" + monkeypatch.setattr( + controller_module.BuiltinToolManageService, + "get_oauth_client", + MagicMock(return_value={"client_id": "abc"}), + ) + create_proxy_context = MagicMock(return_value="ctx-id") + monkeypatch.setattr(controller_module.OAuthProxyService, "create_proxy_context", create_proxy_context) + auth_response = MagicMock(authorization_url="https://oauth.example.com/authorize", state="s") + oauth_handler_instance = MagicMock() + oauth_handler_instance.get_authorization_url.return_value = auth_response + monkeypatch.setattr(controller_module, "OAuthHandler", MagicMock(return_value=oauth_handler_instance)) + # dump_response is called with the pydantic response class; short-circuit it. + monkeypatch.setattr( + controller_module, "dump_response", lambda _cls, obj: {"authorization_url": obj.authorization_url} + ) + yield create_proxy_context + + +@pytest.mark.parametrize( + ("query_visibility", "expected_stored"), + [ + ("only_me", "only_me"), + ("all_team_members", "all_team_members"), + ], +) +def test_tool_plugin_oauth_url_stashes_valid_visibility( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + query_visibility: str, + expected_stored: str, +): + user = _mock_account("user-oauth") + _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") + + with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: + _invoke_oauth_start( + controller_module, + app, + user, + "tenant-oauth", + query_string=f"visibility={query_visibility}", + ) + + _, kwargs = create_proxy_context.call_args + assert kwargs["extra_data"] == {"visibility": expected_stored} + + +@pytest.mark.parametrize("query_string", ["", "visibility=partial_members", "visibility=nonsense"]) +def test_tool_plugin_oauth_url_unknown_visibility_falls_back_to_only_me( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + query_string: str, +): + """Missing, partial_members (rejected for plugin creds), and garbage inputs + all collapse to only_me — OAuth tokens are personal by default.""" + user = _mock_account("user-oauth-fallback") + _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") + + with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: + _invoke_oauth_start( + controller_module, + app, + user, + "tenant-oauth", + query_string=query_string, + ) + + _, kwargs = create_proxy_context.call_args + assert kwargs["extra_data"] == {"visibility": "only_me"} + + +def _invoke_oauth_callback( + controller_module: ModuleType, + app: Flask, + *, + stored_visibility: str | None, + provider: str = _OAUTH_PROVIDER_PATH, +): + api = controller_module.ToolOAuthCallback() + ctx = {"user_id": "user-oauth", "tenant_id": "tenant-oauth"} + if stored_visibility is not None: + ctx["visibility"] = stored_visibility + with app.test_request_context( + f"/oauth/plugin/{provider}/tool/callback", + method="GET", + headers={"Cookie": "context_id=ctx-id"}, + ): + with patch.object(controller_module.OAuthProxyService, "use_proxy_context", return_value=ctx): + return unwrap(api.get)(api, provider=provider) + + +@contextmanager +def _mock_oauth_callback_deps(controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch): + """Mock everything ToolOAuthCallback.get calls out to and capture the + add_builtin_tool_provider kwargs so tests can assert visibility.""" + monkeypatch.setattr( + controller_module.BuiltinToolManageService, + "get_oauth_client", + MagicMock(return_value={"client_id": "abc"}), + ) + add_provider = MagicMock() + monkeypatch.setattr(controller_module.BuiltinToolManageService, "add_builtin_tool_provider", add_provider) + credentials_response = MagicMock(credentials={"access_token": "tok"}, expires_at=-1) + oauth_handler_instance = MagicMock() + oauth_handler_instance.get_credentials.return_value = credentials_response + monkeypatch.setattr(controller_module, "OAuthHandler", MagicMock(return_value=oauth_handler_instance)) + monkeypatch.setattr(controller_module, "redirect", lambda url: {"redirect": url}) + yield add_provider + + +@pytest.mark.parametrize( + ("stored", "expected"), + [ + ("only_me", "only_me"), + ("all_team_members", "all_team_members"), + ], +) +def test_tool_oauth_callback_persists_stored_visibility( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + stored: str, + expected: str, +): + with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: + _invoke_oauth_callback(controller_module, app, stored_visibility=stored) + + _, kwargs = add_provider.call_args + assert kwargs["visibility"] == expected + assert kwargs["api_type"] == controller_module.CredentialType.OAUTH2 + + +@pytest.mark.parametrize("stored", [None, "partial_members", "garbage"]) +def test_tool_oauth_callback_unknown_stored_visibility_falls_back_to_only_me( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + stored: str | None, +): + """Older proxy contexts (pre-feature) and anything unexpected must default + to only_me — matches the safe-default we apply on the start endpoint.""" + with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: + _invoke_oauth_callback(controller_module, app, stored_visibility=stored) + + _, kwargs = add_provider.call_args + assert kwargs["visibility"] == "only_me" diff --git a/api/tests/unit_tests/controllers/files/test_upload.py b/api/tests/unit_tests/controllers/files/test_upload.py index fdbad62d593..dec841e35fd 100644 --- a/api/tests/unit_tests/controllers/files/test_upload.py +++ b/api/tests/unit_tests/controllers/files/test_upload.py @@ -2,7 +2,7 @@ import io import types from contextlib import contextmanager from inspect import unwrap -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session @@ -206,6 +206,7 @@ class TestPluginUploadFileApi: timestamp="123", nonce="abc", sign="sig", + max_size=None, ) tool_file_manager.create_file_by_raw.assert_called_once_with( user_id="account-1", @@ -336,6 +337,65 @@ class TestPluginUploadFileApi: with pytest.raises(module.FileTooLargeError): post_fn(api) + @patch.object(module, "get_user", return_value=_end_user()) + @patch.object(module, "verify_plugin_file_signature", return_value=True) + @patch.object(module, "ToolFileManager") + def test_signed_max_size_bounds_file_read( + self, + mock_tool_file_manager, + mock_verify, + mock_get_user, + ): + dummy_file = DummyFile(content=b"data") + dummy_file.stream = MagicMock() + dummy_file.stream.read.return_value = b"data" + module.request = fake_request( + { + "timestamp": "123", + "nonce": "abc", + "sign": "sig", + "tenant_id": "tenant-1", + "user_id": "user-1", + "max_size": "4", + }, + file=dummy_file, + ) + mock_tool_file_manager.return_value.create_file_by_raw.return_value = _tool_file() + mock_tool_file_manager.sign_file.return_value = "signed-url" + + unwrap(module.PluginUploadFileApi().post)(module.PluginUploadFileApi()) + + dummy_file.stream.read.assert_called_once_with(5) + assert mock_verify.call_args.kwargs["max_size"] == 4 + assert mock_tool_file_manager.return_value.create_file_by_raw.call_args.kwargs["file_binary"] == b"data" + + @patch.object(module, "get_user", return_value=_end_user()) + @patch.object(module, "verify_plugin_file_signature", return_value=True) + @patch.object(module, "ToolFileManager") + def test_signed_max_size_rejects_oversized_file_before_creation( + self, + mock_tool_file_manager, + mock_verify, + mock_get_user, + ): + dummy_file = DummyFile(content=b"oversized") + module.request = fake_request( + { + "timestamp": "123", + "nonce": "abc", + "sign": "sig", + "tenant_id": "tenant-1", + "user_id": "user-1", + "max_size": "4", + }, + file=dummy_file, + ) + + with pytest.raises(module.FileTooLargeError): + unwrap(module.PluginUploadFileApi().post)(module.PluginUploadFileApi()) + + mock_tool_file_manager.assert_not_called() + @patch.object(module, "get_user", return_value=_end_user()) @patch.object(module, "verify_plugin_file_signature", return_value=True) @patch.object(module, "ToolFileManager") diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_files.py b/api/tests/unit_tests/controllers/inner_api/test_agent_files.py index d6a90456129..8a7d970a889 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_files.py +++ b/api/tests/unit_tests/controllers/inner_api/test_agent_files.py @@ -11,6 +11,7 @@ from sqlalchemy.orm import Session from controllers.inner_api.agent.files import ( AgentFileDownloadRequestApi, AgentFileUploadRequestApi, + AgentFileUploadRequestPayload, ) from core.workflow.file_reference import build_file_reference from services.file_request_service import DownloadFileRequestResult @@ -29,6 +30,7 @@ def test_upload_request_returns_origin_free_uri(app: Flask, unbound_session: Ses "filename": "report.pdf", "mimetype": "application/pdf", "conversation_id": "conversation-1", + "max_size": 64 * 1024 * 1024, } tenant = SimpleNamespace(id="tenant-1") user = SimpleNamespace(id="canonical-end-user-1") @@ -51,9 +53,26 @@ def test_upload_request_returns_origin_free_uri(app: Flask, unbound_session: Ses user_id="canonical-end-user-1", conversation_id="conversation-1", user_from=None, + max_size=64 * 1024 * 1024, ) +def test_upload_request_payload_requires_non_negative_max_size() -> None: + payload = { + "tenant_id": "tenant-1", + "user_id": "user-1", + "filename": "report.pdf", + "mimetype": "application/pdf", + } + + with pytest.raises(ValueError): + AgentFileUploadRequestPayload.model_validate(payload) + with pytest.raises(ValueError): + AgentFileUploadRequestPayload.model_validate({**payload, "max_size": -1}) + + assert AgentFileUploadRequestPayload.model_validate({**payload, "max_size": 0}).max_size == 0 + + def test_download_request_returns_origin_free_uri_for_sandbox(app: Flask, unbound_session: Session) -> None: reference = build_file_reference(record_id="tool-file-1") payload = { diff --git a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py index 92121ddda78..e0a59fe2f7b 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py +++ b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py @@ -7,7 +7,8 @@ from uuid import NAMESPACE_URL, uuid5 import pytest from flask import Flask -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy import Engine, event +from sqlalchemy.orm import Session from werkzeug.exceptions import HTTPException from configs import dify_config @@ -265,7 +266,7 @@ class TestEnterpriseInnerApiUserAuth: # Assert assert result == "no_user" - def test_should_pass_through_when_hmac_signature_invalid(self, app: Flask): + def test_should_pass_through_when_hmac_signature_invalid(self, app: Flask, sqlite_engine: Engine): """Invalid HMAC auth passes through without opening a database session.""" # Arrange @@ -273,17 +274,20 @@ class TestEnterpriseInnerApiUserAuth: def protected_view(**kwargs): return kwargs.get("user", "no_user") - # Act - use wrong signature - with app.test_request_context( - headers={"Authorization": "Bearer user123:wrong_signature", "X-Inner-Api-Key": "valid_key"} - ): - with patch.object(dify_config, "INNER_API", True): - with patch("controllers.inner_api.wraps.session_factory.create_session") as mock_create_session: - result = protected_view() + def fail_on_query(*_args, **_kwargs): + pytest.fail("invalid HMAC must not access the database") + + event.listen(sqlite_engine, "before_cursor_execute", fail_on_query) + try: + with app.test_request_context( + headers={"Authorization": "Bearer user123:wrong_signature", "X-Inner-Api-Key": "valid_key"} + ): + with patch.object(dify_config, "INNER_API", True): + result = protected_view() + finally: + event.remove(sqlite_engine, "before_cursor_execute", fail_on_query) - # Assert assert result == "no_user" - mock_create_session.assert_not_called() @pytest.mark.parametrize("sqlite_session", [(EndUser,)], indirect=True) def test_should_inject_user_when_hmac_signature_valid(self, app: Flask, sqlite_session: Session): @@ -313,21 +317,13 @@ class TestEnterpriseInnerApiUserAuth: ) sqlite_session.add(end_user) sqlite_session.commit() - database_session_factory = sessionmaker( - bind=sqlite_session.get_bind(), - expire_on_commit=False, - ) # Act with app.test_request_context( headers={"Authorization": f"Bearer {user_id}:{valid_signature}", "X-Inner-Api-Key": inner_api_key} ): with patch.object(dify_config, "INNER_API", True): - with patch( - "controllers.inner_api.wraps.session_factory.create_session", - database_session_factory, - ): - result = protected_view() + result = protected_view() # Assert assert isinstance(result, EndUser) diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_prepare.py b/api/tests/unit_tests/controllers/openapi/auth/test_prepare.py index 0fc691152f2..96d823b178e 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_prepare.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_prepare.py @@ -1,11 +1,12 @@ import uuid -from contextlib import nullcontext -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from unittest.mock import PropertyMock, patch import pytest from flask import Flask +from sqlalchemy import Engine, select from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound, Unauthorized +from werkzeug.exceptions import Forbidden, InternalServerError, NotFound, Unauthorized from controllers.openapi.auth.data import AuthData, ExternalIdentity from controllers.openapi.auth.prepare import ( @@ -18,346 +19,347 @@ from controllers.openapi.auth.prepare import ( resolve_external_user, ) from libs.oauth_bearer import TokenType -from models.account import TenantAccountRole +from models import Account, App, EndUser, Tenant, TenantAccountJoin +from models.account import AccountStatus, TenantAccountRole, TenantStatus +from models.enums import AppStatus +from models.model import AppMode, IconType +from services import end_user_service +from services.enterprise.enterprise_service import WebAppAccessMode + +APP_ID = "00000000-0000-0000-0000-000000000001" +TENANT_ID = "00000000-0000-0000-0000-000000000002" +ACCOUNT_ID = "00000000-0000-0000-0000-000000000003" -def _make_auth_data(**kwargs) -> AuthData: - mock_fields = {k: kwargs.pop(k) for k in ("app", "tenant", "caller") if k in kwargs} - data = AuthData( +def _make_auth_data(**kwargs: object) -> AuthData: + return AuthData( token_type=kwargs.pop("token_type", TokenType.OAUTH_ACCOUNT), token_hash=kwargs.pop("token_hash", "testhash"), scopes=kwargs.pop("scopes", frozenset()), **kwargs, ) - for k, v in mock_fields.items(): - setattr(data, k, v) - return data -_VALID_APP_UUID = "00000000-0000-0000-0000-000000000001" - - -def test_load_app_writes_app_to_data(): - app = MagicMock() - app.status = "normal" - app.enable_api = True - data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID}) - with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app): - load_app(data) - assert data.app is app - - -def test_load_app_raises_not_found_for_non_uuid_app_id(): - data = _make_auth_data(path_params={"app_id": "not-a-uuid"}) - with pytest.raises(NotFound): - load_app(data) - - -def test_load_app_raises_not_found_when_missing(): - data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID}) - with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=None): - with pytest.raises(NotFound): - load_app(data) - - -def test_load_app_raises_not_found_when_not_normal(): - app = MagicMock() - app.status = "archived" - data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID}) - with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app): - with pytest.raises(NotFound): - load_app(data) - - -def test_load_app_stashes_app_even_when_api_disabled(): - app = MagicMock() - app.status = "normal" - app.enable_api = False - data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID}) - with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app): - load_app(data) - assert data.app is app - - -def test_load_app_skips_when_already_set(): - existing_app = MagicMock() - data = _make_auth_data(app=existing_app, path_params={"app_id": "abc"}) - load_app(data) - assert data.app is existing_app - - -def test_load_tenant_writes_tenant(): - app = MagicMock() - app.tenant_id = uuid.uuid4() - tenant = MagicMock() - tenant.status = "normal" - data = _make_auth_data(app=app) - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant): - load_tenant(data) - assert data.tenant is tenant - - -def test_load_tenant_skips_when_already_set(): - existing_tenant = MagicMock() - data = _make_auth_data(app=MagicMock(), tenant=existing_tenant) - load_tenant(data) - assert data.tenant is existing_tenant - - -def test_load_tenant_raises_forbidden_when_archived(): - from models.account import TenantStatus - - app = MagicMock() - app.tenant_id = uuid.uuid4() - tenant = MagicMock() - tenant.status = TenantStatus.ARCHIVE - data = _make_auth_data(app=app) - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant): - with pytest.raises(Forbidden): - load_tenant(data) - - -def test_load_tenant_raises_forbidden_when_missing(): - app = MagicMock() - app.tenant_id = uuid.uuid4() - data = _make_auth_data(app=app) - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=None): - with pytest.raises(Forbidden): - load_tenant(data) - - -def test_load_tenant_raises_500_when_app_not_loaded(): - from werkzeug.exceptions import InternalServerError - - data = _make_auth_data() - with pytest.raises(InternalServerError): - load_tenant(data) - - -def test_load_account_writes_caller(): - account = MagicMock() - account_id = uuid.uuid4() - data = _make_auth_data(account_id=account_id) - with patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=account): - load_account(data) - assert data.caller is account - assert data.caller_kind == "account" - - -def test_load_account_skips_when_already_set(): - existing_caller = MagicMock() - data = _make_auth_data(account_id=uuid.uuid4(), caller=existing_caller) - load_account(data) - assert data.caller is existing_caller - - -def test_load_account_sets_current_tenant_when_tenant_present(sqlite_session: Session): - account = MagicMock() - tenant = MagicMock() - session = sqlite_session - data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant) - with ( - patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=account), - patch("controllers.openapi.auth.prepare.session_factory.create_session", return_value=nullcontext(session)), - ): - load_account(data) - account.set_current_tenant_with_session.assert_called_once_with(tenant, session=session) - - -def test_load_account_raises_unauthorized_when_not_found(): - data = _make_auth_data(account_id=uuid.uuid4()) - with patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=None): - with pytest.raises(Unauthorized): - load_account(data) - - -def test_resolve_external_user_writes_caller(): - tenant = MagicMock() - app = MagicMock() - end_user = MagicMock() - ext = ExternalIdentity(email="user@sso.com") - data = _make_auth_data(tenant=tenant, app=app, external_identity=ext) - with patch("controllers.openapi.auth.prepare.EndUserService.get_or_create_end_user_by_type", return_value=end_user): - resolve_external_user(data) - assert data.caller is end_user - assert data.caller_kind == "end_user" - - -def test_resolve_external_user_raises_unauthorized_when_context_missing(): - data = _make_auth_data(tenant=None, app=MagicMock(), external_identity=ExternalIdentity(email="u@s.com")) - with pytest.raises(Unauthorized): - resolve_external_user(data) - - -def test_load_app_access_mode_writes_mode(): - from services.enterprise.enterprise_service import WebAppAccessMode - - app = MagicMock() - app.id = "app-1" - settings = MagicMock() - settings.access_mode = "public" - data = _make_auth_data(app=app) - with patch( - "controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id", - return_value=settings, - ): - load_app_access_mode(data) - assert data.app_access_mode == WebAppAccessMode.PUBLIC - - -def test_load_app_access_mode_writes_none_when_value_error(): - app = MagicMock() - app.id = "app-1" - data = _make_auth_data(app=app) - with patch( - "controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id", - side_effect=ValueError("No data found."), - ): - load_app_access_mode(data) - assert data.app_access_mode is None - - -def test_load_app_access_mode_no_op_when_app_missing(): - data = _make_auth_data() - load_app_access_mode(data) - assert data.app_access_mode is None - - -@pytest.fixture -def flask_app(): - return Flask(__name__) - - -def test_load_tenant_from_request_from_path_params(flask_app): - tenant = MagicMock() - tenant.status = "normal" - wid = str(uuid.uuid4()) - data = _make_auth_data(path_params={"workspace_id": wid}) - with flask_app.test_request_context("/test"): - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant): - load_tenant_from_request(data) - assert data.tenant is tenant - - -def test_load_tenant_from_request_from_query_param(flask_app): - tenant = MagicMock() - tenant.status = "normal" - wid = str(uuid.uuid4()) - data = _make_auth_data(path_params={}) - with flask_app.test_request_context(f"/test?workspace_id={wid}"): - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant): - load_tenant_from_request(data) - assert data.tenant is tenant - - -def test_load_tenant_from_request_skips_when_already_set(flask_app): - existing_tenant = MagicMock() - data = _make_auth_data(tenant=existing_tenant, path_params={}) - with flask_app.test_request_context("/test"): - load_tenant_from_request(data) - assert data.tenant is existing_tenant - - -def test_load_tenant_from_request_raises_not_found_when_no_id(flask_app): - data = _make_auth_data(path_params={}) - with flask_app.test_request_context("/test"): - with pytest.raises(NotFound): - load_tenant_from_request(data) - - -def test_load_tenant_from_request_raises_not_found_when_missing(flask_app): - wid = str(uuid.uuid4()) - data = _make_auth_data(path_params={"workspace_id": wid}) - with flask_app.test_request_context("/test"): - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=None): - with pytest.raises(NotFound): - load_tenant_from_request(data) - - -def test_load_tenant_from_request_raises_not_found_when_archived(flask_app): - from models.account import TenantStatus - - tenant = MagicMock() - tenant.status = TenantStatus.ARCHIVE - wid = str(uuid.uuid4()) - data = _make_auth_data(path_params={"workspace_id": wid}) - with flask_app.test_request_context("/test"): - with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant): - with pytest.raises(NotFound): - load_tenant_from_request(data) - - -def test_load_tenant_from_request_raises_not_found_when_invalid_uuid(flask_app): - data = _make_auth_data(path_params={"workspace_id": "not-a-uuid"}) - with flask_app.test_request_context("/test"): - with pytest.raises(NotFound): - load_tenant_from_request(data) - - -# --- load_workspace_role --- - - -def test_load_workspace_role_stashes_role(): - tenant = MagicMock() - tenant.id = uuid.uuid4() - caller = MagicMock() - caller.status = "active" - data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller) - with patch( - "controllers.openapi.auth.prepare.TenantService.get_account_role_in_tenant", - return_value=TenantAccountRole.ADMIN, - ): - load_workspace_role(data) - assert data.tenant_role == TenantAccountRole.ADMIN - - -def test_load_workspace_role_none_when_not_member(): - tenant = MagicMock() - tenant.id = uuid.uuid4() - caller = MagicMock() - caller.status = "active" - data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller) - with patch( - "controllers.openapi.auth.prepare.TenantService.get_account_role_in_tenant", - return_value=None, - ): - load_workspace_role(data) - assert data.tenant_role is None - - -def test_load_workspace_role_none_when_account_inactive(): - tenant = MagicMock() - tenant.id = uuid.uuid4() - caller = MagicMock() - caller.status = "banned" - data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller) - load_workspace_role(data) - assert data.tenant_role is None - - -def test_load_workspace_role_skips_when_already_set(): - tenant = MagicMock() - tenant.id = uuid.uuid4() - caller = MagicMock() - caller.status = "active" - data = _make_auth_data( - account_id=uuid.uuid4(), - tenant=tenant, - caller=caller, - tenant_role=TenantAccountRole.OWNER, +def _app( + *, + app_id: str = APP_ID, + tenant_id: str = TENANT_ID, + enable_api: bool = True, +) -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="OpenAPI app", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + status=AppStatus.NORMAL, + enable_site=True, + enable_api=enable_api, + max_active_requests=None, ) - load_workspace_role(data) - assert data.tenant_role == TenantAccountRole.OWNER -def test_load_workspace_role_skips_when_tenant_missing(): - data = _make_auth_data(account_id=uuid.uuid4()) - load_workspace_role(data) - assert data.tenant_role is None +def _tenant(*, tenant_id: str = TENANT_ID, status: TenantStatus = TenantStatus.NORMAL) -> Tenant: + tenant = Tenant(name="OpenAPI tenant", status=status) + tenant.id = tenant_id + return tenant -def test_load_workspace_role_skips_when_account_id_missing(): - tenant = MagicMock() - data = _make_auth_data(tenant=tenant, account_id=None) - load_workspace_role(data) - assert data.tenant_role is None +def _account(*, status: AccountStatus = AccountStatus.ACTIVE) -> Account: + account = Account(name="OpenAPI account", email="account@example.com", status=status) + account.id = ACCOUNT_ID + return account + + +def _persist(session: Session, *models: object) -> None: + session.add_all(models) + session.commit() + + +class TestLoadApp: + def test_writes_persisted_app_to_data(self, sqlite_session: Session) -> None: + _persist(sqlite_session, _app()) + data = _make_auth_data(path_params={"app_id": APP_ID}) + + load_app(data) + + assert data.app is not None + assert data.app.id == APP_ID + + def test_rejects_non_uuid_and_missing_app(self) -> None: + with pytest.raises(NotFound, match="app not found"): + load_app(_make_auth_data(path_params={"app_id": "not-a-uuid"})) + with pytest.raises(NotFound, match="app not found"): + load_app(_make_auth_data(path_params={"app_id": APP_ID})) + + def test_rejects_non_normal_app(self) -> None: + app = _app() + app.status = "archived" # type: ignore[assignment] + + with ( + patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app), + pytest.raises(NotFound, match="app not found"), + ): + load_app(_make_auth_data(path_params={"app_id": APP_ID})) + + def test_stashes_app_even_when_api_disabled(self, sqlite_session: Session) -> None: + _persist(sqlite_session, _app(enable_api=False)) + data = _make_auth_data(path_params={"app_id": APP_ID}) + + load_app(data) + + assert data.app is not None + assert data.app.enable_api is False + + def test_skips_when_already_set(self) -> None: + existing_app = _app() + data = _make_auth_data(app=existing_app, path_params={"app_id": "invalid"}) + + load_app(data) + + assert data.app is existing_app + + +class TestLoadTenant: + def test_writes_persisted_tenant(self, sqlite_session: Session) -> None: + app = _app() + _persist(sqlite_session, app, _tenant()) + data = _make_auth_data(app=app) + + load_tenant(data) + + assert data.tenant is not None + assert data.tenant.id == TENANT_ID + + def test_skips_when_already_set(self) -> None: + tenant = _tenant() + data = _make_auth_data(app=_app(), tenant=tenant) + + load_tenant(data) + + assert data.tenant is tenant + + @pytest.mark.parametrize("persist_archived", [True, False]) + def test_rejects_archived_or_missing_tenant(self, sqlite_session: Session, persist_archived: bool) -> None: + app = _app() + models: list[object] = [app] + if persist_archived: + models.append(_tenant(status=TenantStatus.ARCHIVE)) + _persist(sqlite_session, *models) + + with pytest.raises(Forbidden, match="workspace unavailable"): + load_tenant(_make_auth_data(app=app)) + + def test_rejects_missing_app_context(self) -> None: + with pytest.raises(InternalServerError, match="app not loaded"): + load_tenant(_make_auth_data()) + + +class TestLoadAccount: + def test_writes_persisted_caller(self, sqlite_session: Session) -> None: + _persist(sqlite_session, _account()) + data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID)) + + load_account(data) + + assert data.caller is not None + assert data.caller.id == ACCOUNT_ID + assert data.caller_kind == "account" + + def test_sets_current_tenant_from_real_membership(self, sqlite_session: Session) -> None: + account = _account() + tenant = _tenant() + membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.ADMIN, + ) + _persist(sqlite_session, account, tenant, membership) + data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), tenant=tenant) + + load_account(data) + + assert isinstance(data.caller, Account) + assert data.caller.current_tenant_id == TENANT_ID + assert data.caller.role == TenantAccountRole.ADMIN + + def test_skips_when_caller_already_set(self) -> None: + account = _account() + data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), caller=account) + + load_account(data) + + assert data.caller is account + + def test_rejects_missing_account(self) -> None: + with pytest.raises(Unauthorized, match="account not found"): + load_account(_make_auth_data(account_id=uuid.UUID(ACCOUNT_ID))) + + +class TestResolveExternalUser: + def test_persists_and_writes_end_user( + self, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + app = _app() + tenant = _tenant() + _persist(sqlite_session, app, tenant) + data = _make_auth_data( + tenant=tenant, + app=app, + external_identity=ExternalIdentity(email="user@sso.com"), + ) + + with patch.object(type(end_user_service.db), "engine", new_callable=PropertyMock) as engine: + engine.return_value = sqlite_engine + resolve_external_user(data) + + assert isinstance(data.caller, EndUser) + assert data.caller_kind == "end_user" + with Session(sqlite_engine) as observer: + persisted = observer.scalar(select(EndUser).where(EndUser.session_id == "user@sso.com")) + assert persisted is not None + assert persisted.tenant_id == TENANT_ID + assert persisted.app_id == APP_ID + + def test_rejects_missing_context(self) -> None: + data = _make_auth_data(app=_app(), external_identity=ExternalIdentity(email="u@s.com")) + + with pytest.raises(Unauthorized, match="missing context"): + resolve_external_user(data) + + +class TestLoadAppAccessMode: + def test_writes_mode(self) -> None: + data = _make_auth_data(app=_app()) + settings = SimpleNamespace(access_mode="public") + + with patch( + "controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id", + return_value=settings, + ): + load_app_access_mode(data) + + assert data.app_access_mode == WebAppAccessMode.PUBLIC + + def test_writes_none_when_provider_raises(self) -> None: + data = _make_auth_data(app=_app()) + with patch( + "controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id", + side_effect=ValueError("No data found."), + ): + load_app_access_mode(data) + assert data.app_access_mode is None + + def test_noop_without_app(self) -> None: + data = _make_auth_data() + load_app_access_mode(data) + assert data.app_access_mode is None + + +class TestLoadTenantFromRequest: + def test_loads_from_path_or_query( + self, + app: Flask, + sqlite_session: Session, + ) -> None: + _persist(sqlite_session, _tenant()) + for path_params, path in (({"workspace_id": TENANT_ID}, "/test"), ({}, f"/test?workspace_id={TENANT_ID}")): + data = _make_auth_data(path_params=path_params) + with app.test_request_context(path): + load_tenant_from_request(data) + assert data.tenant is not None + assert data.tenant.id == TENANT_ID + + def test_skips_when_already_set(self, app: Flask) -> None: + tenant = _tenant() + data = _make_auth_data(tenant=tenant) + with app.test_request_context("/test"): + load_tenant_from_request(data) + assert data.tenant is tenant + + def test_rejects_missing_or_invalid_id(self, app: Flask) -> None: + for path_params in ({}, {"workspace_id": "not-a-uuid"}): + with app.test_request_context("/test"), pytest.raises(NotFound, match="workspace not found"): + load_tenant_from_request(_make_auth_data(path_params=path_params)) + + @pytest.mark.parametrize("tenant_status", [None, TenantStatus.ARCHIVE]) + def test_rejects_missing_or_archived_tenant( + self, + app: Flask, + sqlite_session: Session, + tenant_status: TenantStatus | None, + ) -> None: + if tenant_status is not None: + _persist(sqlite_session, _tenant(status=tenant_status)) + data = _make_auth_data(path_params={"workspace_id": TENANT_ID}) + + with app.test_request_context("/test"), pytest.raises(NotFound, match="workspace not found"): + load_tenant_from_request(data) + + +class TestLoadWorkspaceRole: + def test_loads_real_membership_role(self, sqlite_session: Session) -> None: + account = _account() + tenant = _tenant() + membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.ADMIN, + ) + _persist(sqlite_session, account, tenant, membership) + data = _make_auth_data( + account_id=uuid.UUID(ACCOUNT_ID), + tenant=tenant, + caller=account, + ) + + load_workspace_role(data) + + assert data.tenant_role == TenantAccountRole.ADMIN + + def test_none_when_not_member(self, sqlite_session: Session) -> None: + account = _account() + tenant = _tenant() + _persist(sqlite_session, account, tenant) + data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), tenant=tenant, caller=account) + + load_workspace_role(data) + + assert data.tenant_role is None + + def test_none_when_account_inactive(self) -> None: + data = _make_auth_data( + account_id=uuid.UUID(ACCOUNT_ID), + tenant=_tenant(), + caller=_account(status=AccountStatus.BANNED), + ) + load_workspace_role(data) + assert data.tenant_role is None + + def test_skips_when_already_set(self) -> None: + data = _make_auth_data( + account_id=uuid.UUID(ACCOUNT_ID), + tenant=_tenant(), + caller=_account(), + tenant_role=TenantAccountRole.OWNER, + ) + load_workspace_role(data) + assert data.tenant_role == TenantAccountRole.OWNER + + @pytest.mark.parametrize( + "data", + [ + _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID)), + _make_auth_data(tenant=_tenant(), account_id=None), + ], + ) + def test_skips_without_tenant_or_account(self, data: AuthData) -> None: + load_workspace_role(data) + assert data.tenant_role is None diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py index 8d306a85dcf..fc17d166994 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py @@ -1,25 +1,21 @@ """Unit tests for Service API dataset controller behavior. Service boundaries stay mocked, while ORM collaborators are concrete model instances -persisted in one in-memory SQLite session. The controller's ``db.session`` and the -session passed to unwrapped ``@with_session`` endpoints both use that same session, -so model properties and service call contracts exercise real SQLAlchemy behavior. +persisted through the shared SQLite session fixture. """ import uuid from datetime import UTC, datetime from inspect import unwrap -from typing import cast from unittest.mock import MagicMock, patch import pytest from flask import Flask -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound import services from controllers.service_api.dataset.error import DatasetInUseError, DatasetNameDuplicateError, InvalidActionError -from extensions.ext_database import db from models.account import Account, Tenant, TenantAccountRole from models.dataset import AppDatasetJoin, Dataset, DatasetMetadata, Document from models.enums import PermissionEnum @@ -42,15 +38,9 @@ DATASET_MODEL_TABLES = ( pytestmark = pytest.mark.parametrize("sqlite_session", [DATASET_MODEL_TABLES], indirect=True) -@pytest.fixture(autouse=True) -def controller_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Route controller and model database access through the test's SQLite session.""" - - # Flask-SQLAlchemy exposes a callable registry that also proxies Session methods. - # Seed that registry with this fixture's Session so both access styles share one transaction. - existing_session_factory = cast(sessionmaker[Session], lambda: sqlite_session) - session_registry = scoped_session(existing_session_factory) - monkeypatch.setattr(db, "session", session_registry) +@pytest.fixture +def controller_session(sqlite_session: Session) -> Session: + """Expose the shared SQLite session under the controller-focused fixture name.""" return sqlite_session diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py index 420b1fa1bf7..1f20b74180f 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py @@ -7,15 +7,13 @@ cover the concrete objects and session passed across the controller boundary. import uuid from inspect import unwrap -from typing import cast from unittest.mock import MagicMock, patch import pytest from flask import Flask -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden -from extensions.ext_database import db from models.account import Account, Tenant, TenantAccountRole from models.enums import TagType from models.model import Tag @@ -24,15 +22,9 @@ TAG_MODEL_TABLES = (Account, Tenant, Tag) pytestmark = pytest.mark.parametrize("sqlite_session", [TAG_MODEL_TABLES], indirect=True) -@pytest.fixture(autouse=True) -def controller_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Route controller database access through the test's SQLite session.""" - - # Flask-SQLAlchemy exposes a callable registry that also proxies Session methods. - # Seed that registry with this fixture's Session so both access styles share one transaction. - existing_session_factory = cast(sessionmaker[Session], lambda: sqlite_session) - session_registry = scoped_session(existing_session_factory) - monkeypatch.setattr(db, "session", session_registry) +@pytest.fixture +def controller_session(sqlite_session: Session) -> Session: + """Expose the shared SQLite session under the controller-focused fixture name.""" return sqlite_session 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 19d37038421..a2e8d760feb 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 @@ -18,16 +18,19 @@ Focus on: import inspect import json import uuid +from collections.abc import Iterator from dataclasses import dataclass from datetime import UTC, datetime -from unittest.mock import Mock, patch +from unittest.mock import Mock, PropertyMock, patch import pytest from flask import Flask +from sqlalchemy import Engine from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound from controllers.common.errors import FileTooLargeError as FileTooLargeHTTPError +from controllers.service_api.dataset import document as document_module from controllers.service_api.dataset.document import ( DeprecatedDocumentAddByTextApi, DeprecatedDocumentUpdateByFileApi, @@ -45,8 +48,19 @@ from controllers.service_api.dataset.document import ( ) from controllers.service_api.dataset.error import ArchivedDocumentImmutableError from core.rag.index_processor.constant.index_type import IndexStructureType +from extensions.storage.storage_type import StorageType +from models.account import Account from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, DocumentDocType, IndexingStatus, SegmentStatus +from models.enums import ( + ApiTokenType, + CreatorUserRole, + DataSourceType, + DocumentCreatedFrom, + DocumentDocType, + IndexingStatus, + SegmentStatus, +) +from models.model import ApiToken, UploadFile from services.dataset_ref_service import DatasetRef from services.dataset_service import DocumentService from services.entities.knowledge_entities.knowledge_entities import ProcessRule, RetrievalModel @@ -57,6 +71,30 @@ def _document_data_source_info() -> dict[str, str]: return {"type": "website_crawl", "url": "https://example.com/docs", "title": "Docs"} +def _account() -> Account: + account = Account(name="Document API User", email=f"document-api-{uuid.uuid4()}@example.com") + account.id = "user-1" + return account + + +def _upload_file() -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="documents/file.txt", + name="file.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + created_at=datetime.now(UTC), + used=True, + ) + upload_file.id = str(uuid.uuid4()) + return upload_file + + def _unwrap_non_wrapped_controller(view): while view.__closure__: inner_functions = [cell.cell_contents for cell in view.__closure__ if inspect.isfunction(cell.cell_contents)] @@ -603,8 +641,8 @@ class TestDocumentServiceFileOperations: @patch("services.dataset_service.DocumentService._get_upload_file_for_upload_file_document") def test_get_document_download_url(self, mock_get_file, mock_signed_url, sqlite_session: Session): """Test generation of download URL.""" - mock_doc = Mock() - mock_file = Mock() + mock_doc = make_serializable_document() + mock_file = _upload_file() mock_file.id = "file_id" mock_get_file.return_value = mock_file mock_signed_url.return_value = "https://example.com/download" @@ -622,11 +660,9 @@ class TestDocumentServiceSaveValidation: @patch("services.dataset_service.DatasetService.check_doc_form") @patch("services.dataset_service.FeatureService.get_features") - @patch("services.dataset_service.current_user") - def test_save_document_validates_doc_form(self, mock_user, mock_features, mock_check_form, sqlite_session: Session): + def test_save_document_validates_doc_form(self, mock_features, mock_check_form, sqlite_session: Session): """Test that doc_form is validated during save.""" - mock_user.current_tenant_id = "tenant_id" - dataset = Mock() + dataset = make_dataset(tenant_id="tenant_id") config = Mock() features = Mock() features.billing.enabled = False @@ -641,7 +677,7 @@ class TestDocumentServiceSaveValidation: # Skip actual logic by mocking dependent calls or raising error to stop early with pytest.raises(TestStopError): # We just want to check check_doc_form is called early - DocumentService.save_document_with_dataset_id(dataset, config, Mock(), session=session) + DocumentService.save_document_with_dataset_id(dataset, config, _account(), session=session) # This will fail if we raise exception before check_doc_form, # but check_doc_form is the first thing called. @@ -972,7 +1008,24 @@ class TestDocumentApiGet: ) -class TestDocumentApiDelete: +class SQLiteControllerTest: + session: Session + + @pytest.fixture(autouse=True) + def _use_sqlite_session(self, sqlite_session: Session, sqlite_engine: Engine) -> Iterator[None]: + self.session = sqlite_session + with ( + patch.object(type(document_module.db), "engine", new_callable=PropertyMock, return_value=sqlite_engine), + patch.object(document_module, "current_user", _account()), + ): + yield + + def _persist_dataset(self, dataset: Dataset) -> None: + self.session.add(dataset) + self.session.commit() + + +class TestDocumentApiDelete(SQLiteControllerTest): """Test suite for DocumentApi.delete() endpoint. ``delete`` is wrapped by ``@cloud_edition_billing_rate_limit_check`` which @@ -982,13 +1035,12 @@ class TestDocumentApiDelete: """ @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_delete_document_success(self, mock_db, mock_doc_svc, app: Flask, mock_tenant, mock_document): + def test_delete_document_success(self, mock_doc_svc, app: Flask, mock_tenant, mock_document): """Test successful document deletion.""" # Arrange dataset_id = str(uuid.uuid4()) mock_dataset = make_dataset(id=dataset_id, tenant_id=mock_tenant) - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) mock_doc_svc.get_document.return_value = mock_document mock_doc_svc.check_archived.return_value = False @@ -1003,7 +1055,7 @@ class TestDocumentApiDelete: delete = inspect.unwrap(type(api).delete) response = delete( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=dataset_id, document_id=mock_document.id, @@ -1011,17 +1063,16 @@ 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, self.session) @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_delete_document_not_found(self, mock_db, mock_doc_svc, app: Flask, mock_tenant): + def test_delete_document_not_found(self, mock_doc_svc, app: Flask, mock_tenant): """Test 404 when document not found.""" # Arrange dataset_id = str(uuid.uuid4()) document_id = str(uuid.uuid4()) mock_dataset = make_dataset(id=dataset_id, tenant_id=mock_tenant) - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) mock_doc_svc.get_document.return_value = None @@ -1035,20 +1086,19 @@ class TestDocumentApiDelete: with pytest.raises(NotFound): delete( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=dataset_id, document_id=document_id, ) @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_delete_document_archived_forbidden(self, mock_db, mock_doc_svc, app: Flask, mock_tenant, mock_document): + def test_delete_document_archived_forbidden(self, mock_doc_svc, app: Flask, mock_tenant, mock_document): """Test ArchivedDocumentImmutableError when deleting archived document.""" # Arrange dataset_id = str(uuid.uuid4()) mock_dataset = make_dataset(id=dataset_id, tenant_id=mock_tenant) - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) mock_doc_svc.get_document.return_value = mock_document mock_doc_svc.check_archived.return_value = True @@ -1063,20 +1113,18 @@ class TestDocumentApiDelete: with pytest.raises(ArchivedDocumentImmutableError): delete( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=dataset_id, document_id=mock_document.id, ) @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_delete_document_dataset_not_found(self, mock_db, mock_doc_svc, app: Flask, mock_tenant): + def test_delete_document_dataset_not_found(self, mock_doc_svc, app: Flask, mock_tenant): """Test ValueError when dataset not found.""" # Arrange dataset_id = str(uuid.uuid4()) document_id = str(uuid.uuid4()) - mock_db.session.scalar.return_value = None # Act & Assert with app.test_request_context( @@ -1088,23 +1136,22 @@ class TestDocumentApiDelete: with pytest.raises(ValueError, match="Dataset does not exist."): delete( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=dataset_id, document_id=document_id, ) -class TestDocumentListApi: +class TestDocumentListApi(SQLiteControllerTest): """Test suite for DocumentListApi endpoint.""" @patch("controllers.service_api.dataset.document.paginate_query") @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_list_documents_success(self, mock_db, mock_doc_svc, mock_paginate, app: Flask, mock_tenant, mock_dataset): + def test_list_documents_success(self, mock_doc_svc, mock_paginate, app: Flask, mock_tenant, mock_dataset): """Test successful document list retrieval.""" # Arrange - mock_db.session.scalar.side_effect = [mock_dataset, 0, 0] + self._persist_dataset(mock_dataset) documents = [ make_serializable_document( @@ -1128,7 +1175,7 @@ class TestDocumentListApi: ): api = DocumentListApi() response = inspect.unwrap(type(api).get)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) # Assert @@ -1142,11 +1189,9 @@ class TestDocumentListApi: assert "data_source_info_dict" not in response["data"][0] assert "doc_metadata_details" not in response["data"][0] - @patch("controllers.service_api.dataset.document.db") - def test_list_documents_dataset_not_found(self, mock_db, app: Flask, mock_tenant, mock_dataset): + def test_list_documents_dataset_not_found(self, app: Flask, mock_tenant, mock_dataset): """Test 404 when dataset not found.""" # Arrange - mock_db.session.scalar.return_value = None # Act & Assert with app.test_request_context( @@ -1155,15 +1200,14 @@ class TestDocumentListApi: ): api = DocumentListApi() with pytest.raises(NotFound): - inspect.unwrap(type(api).get)(api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id) + inspect.unwrap(type(api).get)(api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id) -class TestDocumentIndexingStatusApi: +class TestDocumentIndexingStatusApi(SQLiteControllerTest): """Test suite for DocumentIndexingStatusApi endpoint.""" @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_get_indexing_status_success(self, mock_db, mock_doc_svc, app: Flask, mock_tenant, mock_dataset): + def test_get_indexing_status_success(self, mock_doc_svc, app: Flask, mock_tenant, mock_dataset): """Test successful indexing status retrieval.""" # Arrange batch_id = "batch_123" @@ -1176,8 +1220,25 @@ class TestDocumentIndexingStatusApi: mock_doc_svc.get_batch_documents.return_value = [document] - # scalar() called 3 times: dataset lookup, completed_segments count, total_segments count - mock_db.session.scalar.side_effect = [mock_dataset, 5, 5] + self._persist_dataset(mock_dataset) + self.session.add_all( + [ + DocumentSegment( + tenant_id=mock_tenant, + dataset_id=mock_dataset.id, + document_id=document.id, + position=position, + content=f"Segment {position}", + word_count=2, + tokens=2, + created_by="user-1", + status=SegmentStatus.COMPLETED, + completed_at=datetime(2021, 1, 1, tzinfo=UTC), + ) + for position in range(1, 6) + ] + ) + self.session.commit() # Act with app.test_request_context( @@ -1187,7 +1248,7 @@ class TestDocumentIndexingStatusApi: api = DocumentIndexingStatusApi() response = inspect.unwrap(type(api).get)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, batch=batch_id, @@ -1216,12 +1277,10 @@ class TestDocumentIndexingStatusApi: ] } - @patch("controllers.service_api.dataset.document.db") - def test_get_indexing_status_dataset_not_found(self, mock_db, app: Flask, mock_tenant, mock_dataset): + def test_get_indexing_status_dataset_not_found(self, app: Flask, mock_tenant, mock_dataset): """Test 404 when dataset not found.""" # Arrange batch_id = "batch_123" - mock_db.session.scalar.return_value = None # Act & Assert with app.test_request_context( @@ -1232,21 +1291,18 @@ class TestDocumentIndexingStatusApi: with pytest.raises(NotFound): inspect.unwrap(type(api).get)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, batch=batch_id, ) @patch("controllers.service_api.dataset.document.DocumentService") - @patch("controllers.service_api.dataset.document.db") - def test_get_indexing_status_documents_not_found( - self, mock_db, mock_doc_svc, app: Flask, mock_tenant, mock_dataset - ): + def test_get_indexing_status_documents_not_found(self, mock_doc_svc, app: Flask, mock_tenant, mock_dataset): """Test 404 when no documents found for batch.""" # Arrange batch_id = "batch_empty" - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) mock_doc_svc.get_batch_documents.return_value = [] # Act & Assert @@ -1258,14 +1314,14 @@ class TestDocumentIndexingStatusApi: with pytest.raises(NotFound): inspect.unwrap(type(api).get)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, batch=batch_id, ) -class TestDocumentAddByTextApi: +class TestDocumentAddByTextApi(SQLiteControllerTest): """Test suite for DocumentAddByTextApi.post() endpoint. ``post`` is wrapped by ``@cloud_edition_billing_resource_check`` and @@ -1286,9 +1342,8 @@ class TestDocumentAddByTextApi: ``FeatureService.get_knowledge_rate_limit``. Both call ``validate_and_get_api_token`` first. """ - mock_api_token = Mock() - mock_api_token.tenant_id = tenant_id - mock_validate_token.return_value = mock_api_token + api_token = ApiToken(tenant_id=tenant_id, type=ApiTokenType.DATASET, token="dataset-token") + mock_validate_token.return_value = api_token mock_features = Mock() mock_features.billing.enabled = False @@ -1306,16 +1361,12 @@ class TestDocumentAddByTextApi: @patch("controllers.service_api.dataset.document.DocumentService") @patch("controllers.service_api.dataset.document.KnowledgeConfig") @patch("controllers.service_api.dataset.document.FileService") - @patch("controllers.service_api.dataset.document.current_user") - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_document_by_text_success( self, mock_validate_token, mock_feature_svc, - mock_db, - mock_current_user, mock_file_svc_cls, mock_knowledge_config, mock_doc_svc, @@ -1327,12 +1378,9 @@ class TestDocumentAddByTextApi: # Arrange — neutralise billing decorators self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) - mock_db.session.scalar.side_effect = [mock_dataset, None, 0] + self._persist_dataset(mock_dataset) mock_dataset.indexing_technique = "economy" - mock_current_user.id = str(uuid.uuid4()) - - mock_upload_file = Mock() - mock_upload_file.id = str(uuid.uuid4()) + mock_upload_file = _upload_file() mock_file_svc = Mock() mock_file_svc.upload_text.return_value = mock_upload_file mock_file_svc_cls.return_value = mock_file_svc @@ -1357,7 +1405,7 @@ class TestDocumentAddByTextApi: ): api = DocumentAddByTextApi() response, status = _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) # Assert @@ -1366,20 +1414,17 @@ class TestDocumentAddByTextApi: 200, ) assert "data_source_info_dict" not in response["document"] - assert mock_doc_svc.save_document_with_dataset_id.call_args.kwargs["session"] is mock_db.session + assert mock_doc_svc.save_document_with_dataset_id.call_args.kwargs["session"] is self.session @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") - @patch("controllers.service_api.dataset.document.db") def test_create_document_dataset_not_found( - self, mock_db, mock_validate_token, mock_feature_svc, app: Flask, mock_tenant, mock_dataset + self, mock_validate_token, mock_feature_svc, app: Flask, mock_tenant, mock_dataset ): """Test ValueError when dataset not found.""" # Arrange — neutralise billing decorators self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) - mock_db.session.scalar.return_value = None - # Act & Assert with app.test_request_context( f"/datasets/{mock_dataset.id}/document/create-by-text", @@ -1390,14 +1435,13 @@ class TestDocumentAddByTextApi: api = DocumentAddByTextApi() with pytest.raises(ValueError, match="Dataset does not exist."): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") - @patch("controllers.service_api.dataset.document.db") def test_create_document_missing_indexing_technique( - self, mock_db, mock_validate_token, mock_feature_svc, app: Flask, mock_tenant, mock_dataset + self, mock_validate_token, mock_feature_svc, app: Flask, mock_tenant, mock_dataset ): """Test error when both dataset and payload lack indexing_technique. @@ -1409,7 +1453,7 @@ class TestDocumentAddByTextApi: self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) mock_dataset.indexing_technique = None - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) # Act & Assert with app.test_request_context( @@ -1421,7 +1465,7 @@ class TestDocumentAddByTextApi: api = DocumentAddByTextApi() with pytest.raises(ValueError, match="indexing_technique is required."): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) @@ -1474,9 +1518,8 @@ class TestDocumentRouteDeprecation: def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): """Configure mocks to neutralise billing/auth decorators.""" - mock_api_token = Mock() - mock_api_token.tenant_id = tenant_id - mock_validate_token.return_value = mock_api_token + api_token = ApiToken(tenant_id=tenant_id, type=ApiTokenType.DATASET, token="dataset-token") + mock_validate_token.return_value = api_token mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features @@ -1489,7 +1532,7 @@ def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit -class TestDocumentUpdateByTextApiPost: +class TestDocumentUpdateByTextApiPost(SQLiteControllerTest): """Test suite for DocumentUpdateByTextApi.post() endpoint. ``post`` is wrapped by ``@cloud_edition_billing_resource_check`` and @@ -1498,16 +1541,12 @@ class TestDocumentUpdateByTextApiPost: @patch("controllers.service_api.dataset.document.DocumentService") @patch("controllers.service_api.dataset.document.FileService") - @patch("controllers.service_api.dataset.document.current_user") - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_by_text_success( self, mock_validate_token, mock_feature_svc, - mock_db, - mock_current_user, mock_file_svc_cls, mock_doc_svc, app: Flask, @@ -1518,11 +1557,9 @@ class TestDocumentUpdateByTextApiPost: """Test successful document update by text.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) mock_dataset.indexing_technique = "economy" - mock_db.session.scalar.side_effect = [mock_dataset, None, 0] + self._persist_dataset(mock_dataset) - mock_current_user.id = "user-1" - mock_upload = Mock() - mock_upload.id = str(uuid.uuid4()) + mock_upload = _upload_file() mock_file_svc_cls.return_value.upload_text.return_value = mock_upload mock_document = make_serializable_document(id="doc-update-text", name="Updated Doc") @@ -1539,7 +1576,7 @@ class TestDocumentUpdateByTextApiPost: api = DocumentUpdateByTextApi() response, status = _unwrap_non_wrapped_controller(type(api).post)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, document_id=doc_id, @@ -1550,21 +1587,18 @@ class TestDocumentUpdateByTextApiPost: 200, ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_by_text_dataset_not_found( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, ): """Test ValueError when dataset not found.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) - mock_db.session.scalar.return_value = None doc_id = str(uuid.uuid4()) with app.test_request_context( @@ -1577,14 +1611,14 @@ class TestDocumentUpdateByTextApiPost: with pytest.raises(ValueError, match="Dataset does not exist"): _unwrap_non_wrapped_controller(type(api).post)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, document_id=doc_id, ) -class TestDocumentAddByFileApiPost: +class TestDocumentAddByFileApiPost(SQLiteControllerTest): """Test suite for DocumentAddByFileApi.post() endpoint. ``post`` is wrapped by two ``@cloud_edition_billing_resource_check`` @@ -1593,16 +1627,12 @@ class TestDocumentAddByFileApiPost: @patch("controllers.service_api.dataset.document.DocumentService") @patch("controllers.service_api.dataset.document.FileService") - @patch("controllers.service_api.dataset.document.current_user") - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_add_by_file_success_serializes_document_and_batch_shape( self, mock_validate_token, mock_feature_svc, - mock_db, - mock_current_user, mock_file_svc_cls, mock_doc_svc, app: Flask, @@ -1614,11 +1644,9 @@ class TestDocumentAddByFileApiPost: mock_dataset.provider = "vendor" mock_dataset.indexing_technique = "economy" mock_dataset.chunk_structure = None - mock_db.session.scalar.side_effect = [mock_dataset, 0] + self._persist_dataset(mock_dataset) - mock_current_user.id = "user-1" - mock_upload = Mock() - mock_upload.id = str(uuid.uuid4()) + mock_upload = _upload_file() mock_file_svc_cls.return_value.upload_file.return_value = mock_upload mock_document = make_serializable_document(id="doc-create-file", name="File Document") @@ -1640,7 +1668,7 @@ class TestDocumentAddByFileApiPost: ): api = DocumentAddByFileApi() response, status = _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) assert (response, status) == ( @@ -1653,12 +1681,8 @@ class TestDocumentAddByFileApiPost: return_value=15, ) @patch("controllers.service_api.dataset.document.FileService") - @patch("controllers.service_api.dataset.document.current_user") - @patch("controllers.service_api.dataset.document.db") def test_add_by_file_too_large_returns_http_413( self, - mock_db, - mock_current_user, mock_file_svc_cls, mock_get_limit, app: Flask, @@ -1668,8 +1692,7 @@ class TestDocumentAddByFileApiPost: mock_dataset.provider = "vendor" mock_dataset.indexing_technique = "economy" mock_dataset.chunk_structure = None - mock_db.session.scalar.return_value = mock_dataset - mock_current_user.__bool__ = Mock(return_value=True) + self._persist_dataset(mock_dataset) mock_file_svc_cls.return_value.upload_file.side_effect = FileTooLargeServiceError() from io import BytesIO @@ -1687,28 +1710,25 @@ class TestDocumentAddByFileApiPost: api = DocumentAddByFileApi() with pytest.raises(FileTooLargeHTTPError) as exc_info: _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) assert exc_info.value.code == 413 assert exc_info.value.error_code == "file_too_large" mock_get_limit.assert_called_once_with(mock_tenant) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_add_by_file_dataset_not_found( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, ): """Test ValueError when dataset not found.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) - mock_db.session.scalar.return_value = None from io import BytesIO @@ -1723,17 +1743,15 @@ class TestDocumentAddByFileApiPost: api = DocumentAddByFileApi() with pytest.raises(ValueError, match="Dataset does not exist"): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_add_by_file_external_dataset( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, @@ -1741,7 +1759,7 @@ class TestDocumentAddByFileApiPost: """Test ValueError when dataset is external.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) mock_dataset.provider = "external" - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) from io import BytesIO @@ -1756,17 +1774,15 @@ class TestDocumentAddByFileApiPost: api = DocumentAddByFileApi() with pytest.raises(ValueError, match="External datasets"): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_add_by_file_no_file_uploaded( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, @@ -1778,7 +1794,7 @@ class TestDocumentAddByFileApiPost: mock_dataset.provider = "vendor" mock_dataset.indexing_technique = "economy" mock_dataset.chunk_structure = None - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) with app.test_request_context( f"/datasets/{mock_dataset.id}/document/create_by_file", @@ -1790,17 +1806,15 @@ class TestDocumentAddByFileApiPost: api = DocumentAddByFileApi() with pytest.raises(NoFileUploadedError): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_add_by_file_missing_indexing_technique( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, @@ -1810,7 +1824,7 @@ class TestDocumentAddByFileApiPost: mock_dataset.provider = "vendor" mock_dataset.indexing_technique = None mock_dataset.chunk_structure = None - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) from io import BytesIO @@ -1825,11 +1839,11 @@ class TestDocumentAddByFileApiPost: api = DocumentAddByFileApi() with pytest.raises(ValueError, match="indexing_technique is required"): _unwrap_non_wrapped_controller(type(api).post)( - api, mock_db.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id + api, self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id ) -class TestDocumentUpdateByFileApiPatch: +class TestDocumentUpdateByFileApiPatch(SQLiteControllerTest): """Test suite for the canonical document file update endpoint. ``patch`` is wrapped by ``@cloud_edition_billing_resource_check`` and @@ -1885,21 +1899,18 @@ class TestDocumentUpdateByFileApiPatch: document_id=doc_id, ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_by_file_dataset_not_found( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, ): """Test ValueError when dataset not found.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) - mock_db.session.scalar.return_value = None from io import BytesIO @@ -1916,20 +1927,18 @@ class TestDocumentUpdateByFileApiPatch: with pytest.raises(ValueError, match="Dataset does not exist"): _unwrap_non_wrapped_controller(type(api).patch)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, document_id=doc_id, ) - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_by_file_external_dataset( self, mock_validate_token, mock_feature_svc, - mock_db, app: Flask, mock_tenant, mock_dataset, @@ -1937,7 +1946,7 @@ class TestDocumentUpdateByFileApiPatch: """Test ValueError when dataset is external.""" _setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant) mock_dataset.provider = "external" - mock_db.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset) from io import BytesIO @@ -1954,7 +1963,7 @@ class TestDocumentUpdateByFileApiPatch: with pytest.raises(ValueError, match="External datasets"): _unwrap_non_wrapped_controller(type(api).patch)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, document_id=doc_id, @@ -1962,16 +1971,12 @@ class TestDocumentUpdateByFileApiPatch: @patch("controllers.service_api.dataset.document.DocumentService") @patch("controllers.service_api.dataset.document.FileService") - @patch("controllers.service_api.dataset.document.current_user") - @patch("controllers.service_api.dataset.document.db") @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_by_file_success( self, mock_validate_token, mock_feature_svc, - mock_db, - mock_current_user, mock_file_svc_cls, mock_doc_svc, app: Flask, @@ -1983,11 +1988,9 @@ class TestDocumentUpdateByFileApiPatch: mock_dataset.indexing_technique = "economy" mock_dataset.provider = "vendor" mock_dataset.chunk_structure = None - mock_db.session.scalar.side_effect = [mock_dataset, None, 0] + self._persist_dataset(mock_dataset) - mock_current_user.id = "user-1" - mock_upload = Mock() - mock_upload.id = str(uuid.uuid4()) + mock_upload = _upload_file() mock_file_svc_cls.return_value.upload_file.return_value = mock_upload mock_document = make_serializable_document(id="doc-update-file", name="File Document", batch="batch-1") @@ -2008,7 +2011,7 @@ class TestDocumentUpdateByFileApiPatch: api = DocumentApi() response, status = _unwrap_non_wrapped_controller(type(api).patch)( api, - mock_db.session, + self.session, tenant_id=mock_tenant, dataset_id=mock_dataset.id, document_id=doc_id, 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 3df48fdbfe5..8d6c3536550 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 @@ -17,13 +17,14 @@ Decorator strategy: import uuid from inspect import unwrap -from unittest.mock import ANY, MagicMock, Mock, patch +from unittest.mock import ANY, patch import pytest from flask import Flask from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound +from controllers.service_api.dataset import metadata as metadata_module from controllers.service_api.dataset.metadata import ( DatasetMetadataBuiltInFieldActionServiceApi, DatasetMetadataBuiltInFieldServiceApi, @@ -31,21 +32,43 @@ from controllers.service_api.dataset.metadata import ( DatasetMetadataServiceApi, DocumentMetadataEditServiceApi, ) +from models.account import Account, Tenant +from models.dataset import Dataset +from models.enums import PermissionEnum from services.errors.metadata import MetadataResourceNotFoundError @pytest.fixture -def mock_tenant(): - tenant = Mock() +def mock_tenant() -> Tenant: + tenant = Tenant(name="Metadata API Tenant") tenant.id = str(uuid.uuid4()) return tenant @pytest.fixture -def mock_dataset(): - dataset = Mock() - dataset.id = str(uuid.uuid4()) - return dataset +def account() -> Account: + account = Account(name="Metadata API User", email=f"metadata-api-{uuid.uuid4()}@example.com") + account.id = str(uuid.uuid4()) + return account + + +@pytest.fixture +def mock_dataset(mock_tenant: Tenant, account: Account) -> Dataset: + return Dataset( + id=str(uuid.uuid4()), + tenant_id=mock_tenant.id, + name="Metadata Dataset", + description="", + provider="vendor", + permission=PermissionEnum.ONLY_ME, + indexing_technique="economy", + created_by=account.id, + ) + + +@pytest.fixture(autouse=True) +def _use_current_user(monkeypatch: pytest.MonkeyPatch, account: Account) -> None: + monkeypatch.setattr(metadata_module, "current_user", account) # --------------------------------------------------------------------------- @@ -79,10 +102,8 @@ class TestDatasetMetadataCreatePost(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_create_metadata_success( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, @@ -213,15 +234,14 @@ class TestDatasetMetadataServiceApiPatch(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_update_metadata_name_success( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, mock_tenant, mock_dataset, + account: Account, ): """Test successful metadata name update.""" metadata_id = str(uuid.uuid4()) @@ -250,7 +270,7 @@ class TestDatasetMetadataServiceApiPatch(_UsesSQLiteSession): str(mock_dataset.id), mock_tenant.id, session=session ) mock_meta_svc.update_metadata_name.assert_called_once_with( - mock_dataset, metadata_id, "New Name", mock_current_user, session=session + mock_dataset, metadata_id, "New Name", account, session=session ) @patch("controllers.service_api.dataset.metadata.DatasetService") @@ -294,10 +314,8 @@ class TestDatasetMetadataServiceApiDelete(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_delete_metadata_success( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, @@ -410,10 +428,8 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_enable_built_in_field( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, @@ -444,10 +460,8 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_disable_built_in_field( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, @@ -519,15 +533,14 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_update_documents_metadata_success( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, mock_tenant, mock_dataset, + account: Account, ): """Test successful documents metadata update.""" mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset @@ -553,7 +566,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): mock_meta_svc.update_documents_metadata.assert_called_once_with( mock_dataset, ANY, - mock_current_user, + account, session=session, ) @@ -585,10 +598,8 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): @patch("controllers.service_api.dataset.metadata.MetadataService") @patch("controllers.service_api.dataset.metadata.DatasetService") - @patch("controllers.service_api.dataset.metadata.current_user") def test_update_documents_metadata_translates_missing_resource( self, - mock_current_user, mock_dataset_svc, mock_meta_svc, app: Flask, @@ -607,7 +618,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): with pytest.raises(NotFound) as exc_info: self._call_post( api, - MagicMock(), + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, ) diff --git a/api/tests/unit_tests/controllers/service_api/test_wraps.py b/api/tests/unit_tests/controllers/service_api/test_wraps.py index a6d502c6613..9b058fc889c 100644 --- a/api/tests/unit_tests/controllers/service_api/test_wraps.py +++ b/api/tests/unit_tests/controllers/service_api/test_wraps.py @@ -4,12 +4,12 @@ Unit tests for Service API wraps (authentication decorators) import uuid from types import SimpleNamespace -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import Mock, patch import pytest from flask import Flask from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, scoped_session from werkzeug.exceptions import Forbidden, NotFound, ServiceUnavailable, Unauthorized from controllers.service_api.wraps import ( @@ -36,11 +36,9 @@ def _configure_current_app_mock(mock_current_app): mock_current_app._get_current_object = Mock(return_value=Mock()) -def _session_proxy(session: Session) -> MagicMock: - """Emulate Flask-SQLAlchemy's callable scoped-session proxy around a test session.""" - proxy = MagicMock(wraps=session) - proxy.return_value = session - return proxy +def _session_proxy(session: Session) -> scoped_session[Session]: + """Expose the real SQLite session through Flask-SQLAlchemy's callable shape.""" + return scoped_session(lambda: session) def _api_token(*, tenant_id: str, app_id: str | None = None, token_type: ApiTokenType) -> ApiToken: diff --git a/api/tests/unit_tests/controllers/trigger/test_trigger.py b/api/tests/unit_tests/controllers/trigger/test_trigger.py index 1d6db9e232b..c7079d332c1 100644 --- a/api/tests/unit_tests/controllers/trigger/test_trigger.py +++ b/api/tests/unit_tests/controllers/trigger/test_trigger.py @@ -1,3 +1,4 @@ +import logging from unittest.mock import patch import pytest @@ -71,3 +72,11 @@ class TestTriggerEndpoint: assert status == 500 assert response["error"] == "Internal server error" + + @patch.object(module.TriggerService, "process_endpoint", side_effect=Exception("boom")) + def test_unexpected_exception_logs_endpoint_id(self, mock_trigger, caplog): + with caplog.at_level(logging.ERROR, logger=module.logger.name): + module.trigger_endpoint(VALID_UUID) + + assert VALID_UUID in caplog.text + assert "{endpoint_id}" not in caplog.text 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 b351f49c37f..5de8b8bd9e7 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 @@ -18,6 +18,8 @@ from dify_agent.protocol import ( CancelRunRequest, CancelRunResponse, PydanticAIStreamRunEvent, + RunCancelledEvent, + RunCancelledEventData, RunEvent, RunFailedEvent, RunFailedEventData, @@ -39,6 +41,7 @@ from sqlalchemy import event, select from sqlalchemy.orm import Session from clients.agent_backend import ( + AgentBackendError, AgentBackendRunEventAdapter, AgentBackendRunFailedError, AgentBackendRunFailedInternalEvent, @@ -67,9 +70,13 @@ from models.model import MessageAgentThought @pytest.fixture(autouse=True) -def bind_agent_db(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - """Bind the runner's ORM writes to the shared SQLite session.""" +def bind_agent_dependencies(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + """Bind local runner dependencies without reaching external services.""" monkeypatch.setattr(app_runner_module.db, "session", sqlite_session) + monkeypatch.setattr( + "core.app.apps.agent_app.runtime_request_builder.resolve_model_context_window", + lambda **_kwargs: None, + ) def _thought_rows(session: Session) -> list[MessageAgentThought]: @@ -108,12 +115,38 @@ class _RecordingFakeAgentBackendRunClient(FakeAgentBackendRunClient): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.cancelled_run_ids: list[str] = [] + self.cancel_after: list[str | None] = [] @override def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse: self.cancelled_run_ids.append(run_id) return super().cancel_run(run_id, request=request) + @override + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + self.cancel_after.append(after) + return super().cancel_run_and_wait(run_id, request=request, after=after) + + +class _CancelAndWaitFailingClient(_RecordingFakeAgentBackendRunClient): + @override + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + del request + self.cancel_after.append(after) + raise RuntimeError(f"failed to finish cancelling {run_id}") + class _RunLimitBindingLostFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override @@ -139,6 +172,38 @@ class _RunLimitBindingLostFakeAgentBackendRunClient(FakeAgentBackendRunClient): ) +class _TerminalWithoutSnapshotFakeAgentBackendRunClient(FakeAgentBackendRunClient): + def __init__(self, *, terminal_type: str) -> None: + super().__init__() + self.terminal_type = terminal_type + + @override + def stream_events( + self, + run_id: str, + *, + after: str | None = None, + should_stop: Callable[[], bool] | None = None, + ) -> Iterator[RunEvent]: + del after, should_stop + created_at = datetime(2026, 1, 1, tzinfo=UTC) + yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at) + if self.terminal_type == "failed": + yield RunFailedEvent( + id="2-0", + run_id=run_id, + created_at=created_at, + data=RunFailedEventData(error="failed without snapshot"), + ) + else: + yield RunCancelledEvent( + id="2-0", + run_id=run_id, + created_at=created_at, + data=RunCancelledEventData(reason="cancelled without snapshot"), + ) + + class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient): @override def stream_events( @@ -439,6 +504,26 @@ class _FakeSessionStore: self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id)) +class _ExplodingSessionStore(_FakeSessionStore): + def __init__(self, loaded: CompositorSessionSnapshot | None = None) -> None: + super().__init__(loaded=loaded) + self.save_attempts: list[CompositorSessionSnapshot | None] = [] + + @override + def save_active_snapshot( + self, + *, + scope: AgentAppSessionScope, + binding_id: str, + snapshot: CompositorSessionSnapshot | None, + pending_form_id: str | None = None, + pending_tool_call_id: str | None = None, + ) -> None: + del scope, binding_id, pending_form_id, pending_tool_call_id + self.save_attempts.append(snapshot) + raise RuntimeError("session save failed") + + class _MonotonicClock: def __init__(self, *values: float) -> None: self._values = list(values) @@ -768,6 +853,7 @@ def test_streaming_turn_cancels_after_persisting_seen_agent_answer( assert len(rows) == 1 assert rows[0].answer == "hello " assert client.cancelled_run_ids == ["fake-run-1"] + assert client.cancel_after == ["3-0"] def test_tool_result_without_identity_does_not_attach_to_previous_tool( @@ -1158,8 +1244,40 @@ def test_failed_run_raises_agent_backend_error() -> None: with pytest.raises(AgentBackendRunFailedError, match="fake failure .*agent_run_id=fake-run-1"): _run(_runner(client, store), qm) - # No message-end on failure; no snapshot saved. + # No message-end on failure; post-exit session state is still saved. assert not [e for e in qm.events if isinstance(e, QueueMessageEndEvent)] + assert store.saved[0][2] == CompositorSessionSnapshot(layers=[]) + + +@pytest.mark.parametrize("outcome", ["failed", "stopped"]) +def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str) -> None: + store = _ExplodingSessionStore() + queue_manager: _FakeQueueManager = _FakeQueueManager() if outcome == "failed" else _StoppedQueueManager() + client = FakeAgentBackendRunClient( + scenario=FakeAgentBackendScenario.FAILED if outcome == "failed" else FakeAgentBackendScenario.SUCCESS + ) + expected_error = AgentBackendRunFailedError if outcome == "failed" else GenerateTaskStoppedError + + with pytest.raises(expected_error, match="fake failure" if outcome == "failed" else None): + _run(_runner(client, store), queue_manager) + + assert store.save_attempts == [CompositorSessionSnapshot(layers=[])] + + +@pytest.mark.parametrize( + ("terminal_type", "expected_error"), + [("failed", AgentBackendRunFailedError), ("cancelled", AgentBackendError)], +) +def test_terminal_without_snapshot_preserves_prior_app_session_without_write( + terminal_type: str, + expected_error: type[Exception], +) -> None: + store = _FakeSessionStore() + client = _TerminalWithoutSnapshotFakeAgentBackendRunClient(terminal_type=terminal_type) + + with pytest.raises(expected_error): + _run(_runner(client, store), _FakeQueueManager()) + assert store.saved == [] @@ -1229,7 +1347,7 @@ def test_agent_backend_failure_to_exception_prefers_run_failure_type_over_known_ } -def test_stopped_task_cancels_agent_backend_run_and_skips_session_save() -> None: +def test_stopped_task_waits_for_cancelled_snapshot_and_saves_session() -> None: client = _RecordingFakeAgentBackendRunClient() store = _FakeSessionStore() qm = _StoppedQueueManager() @@ -1238,6 +1356,18 @@ def test_stopped_task_cancels_agent_backend_run_and_skips_session_save() -> None _run(_runner(client, store), qm) assert client.cancelled_run_ids == ["fake-run-1"] + assert len(store.saved) == 1 + assert store.saved[0][2] == CompositorSessionSnapshot(layers=[]) + + +def test_cancel_and_wait_failure_preserves_stopped_app_outcome() -> None: + client = _CancelAndWaitFailingClient() + store = _FakeSessionStore() + + with pytest.raises(GenerateTaskStoppedError): + _run(_runner(client, store), _StoppedQueueManager()) + + assert client.cancel_after == [None] assert store.saved == [] 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 518929c6ad6..142dabb9dc6 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 @@ -28,6 +28,21 @@ from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from models.agent_config_entities import AgentSoulConfig +@pytest.fixture(autouse=True) +def model_context_window_calls(monkeypatch: pytest.MonkeyPatch) -> list[tuple[object, str, str]]: + calls: list[tuple[object, str, str]] = [] + + def resolve(*, run_context: object, provider_name: str, model_name: str) -> int: + calls.append((run_context, provider_name, model_name)) + return 32_768 + + monkeypatch.setattr( + "core.app.apps.agent_app.runtime_request_builder.resolve_model_context_window", + resolve, + ) + return calls + + def _exec_ctx() -> DifyExecutionContextLayerConfig: return DifyExecutionContextLayerConfig( tenant_id="tenant-1", @@ -176,11 +191,12 @@ def _soul_with_model() -> AgentSoulConfig: class TestAgentAppRuntimeRequestBuilder: - def test_build_maps_soul_to_run_request(self): + def test_build_maps_soul_to_run_request(self, model_context_window_calls: list[tuple[object, str, str]]): builder = AgentAppRuntimeRequestBuilder( dify_tools_builder=_NoToolsBuilder(), # type: ignore[arg-type] ) - result = builder.build(_ctx(_soul_with_model())) + context = _ctx(_soul_with_model()) + result = builder.build(context) req = result.request names = [layer.name for layer in req.composition.layers] @@ -198,6 +214,8 @@ class TestAgentAppRuntimeRequestBuilder: llm = next(layer for layer in req.composition.layers if layer.name == "llm") assert llm.config.plugin_id == "langgenius/openai" assert llm.config.model_provider == "openai" + assert llm.config.context_window_tokens == 32_768 + assert model_context_window_calls == [(context.dify_context, "langgenius/openai/openai", "gpt-4o-mini")] # execution context carries conversation + agent_app invoke source. exec_ctx = next(layer for layer in req.composition.layers if layer.name == "execution_context") assert exec_ctx.config.conversation_id == "conv-1" diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py index 6c1ee20ffbc..a24aec4cb92 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_config_manager.py @@ -1,16 +1,31 @@ +import json from types import SimpleNamespace -from unittest.mock import MagicMock from pytest_mock import MockerFixture import core.app.apps.pipeline.pipeline_config_manager as module from core.app.apps.pipeline.pipeline_config_manager import PipelineConfigManager +from models.dataset import Pipeline from models.model import AppMode +from models.workflow import Workflow, WorkflowType def test_get_pipeline_config(mocker: MockerFixture): - pipeline = MagicMock(tenant_id="tenant", id="pipe1") - workflow = MagicMock(id="wf1") + pipeline = Pipeline(tenant_id="tenant", name="Pipeline", description="") + pipeline.id = "pipe1" + workflow = Workflow.new( + tenant_id="tenant", + app_id="pipe1", + type=WorkflowType.RAG_PIPELINE.value, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by="user", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "wf1" mocker.patch.object( module.WorkflowVariablesConfigManager, diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py index f1edcd1a1b1..6234297fc3e 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py @@ -1,4 +1,5 @@ import contextlib +import json from types import SimpleNamespace from unittest.mock import MagicMock, PropertyMock @@ -6,13 +7,13 @@ import pytest from pytest_mock import MockerFixture from sqlalchemy import select from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker +from sqlalchemy.orm import Session import core.app.apps.pipeline.pipeline_generator as module from core.app.apps.exc import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom from core.datasource.entities.datasource_entities import DatasourceProviderType -from models.dataset import Document, DocumentPipelineExecutionLog +from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline from models.enums import DataSourceType, EndUserType from models.model import EndUser from models.workflow import Workflow, WorkflowType @@ -49,28 +50,51 @@ def generator(mocker: MockerFixture, sqlite_engine: Engine): def _build_pipeline_dataset(): - return SimpleNamespace( + return Dataset( id=DATASET_ID, + tenant_id=TENANT_ID, name="dataset", description="desc", + created_by=USER_ID, + pipeline_id=PIPELINE_ID, chunk_structure="text_model", built_in_field_enabled=True, - tenant_id=TENANT_ID, ) def _build_pipeline(): - pipeline = MagicMock(tenant_id=TENANT_ID, id=PIPELINE_ID) - pipeline.retrieve_dataset.return_value = _build_pipeline_dataset() + pipeline = Pipeline(tenant_id=TENANT_ID, name="Pipeline", description="desc") + pipeline.id = PIPELINE_ID + pipeline.workflow_id = WORKFLOW_ID return pipeline -def _build_workflow(): - return MagicMock(id=WORKFLOW_ID, graph_dict={"nodes": [], "edges": []}, tenant_id=TENANT_ID) +def _build_workflow(*, graph: dict | None = None): + workflow = Workflow.new( + tenant_id=TENANT_ID, + app_id=PIPELINE_ID, + type=WorkflowType.RAG_PIPELINE.value, + version=Workflow.VERSION_DRAFT, + graph=json.dumps(graph if graph is not None else {"nodes": [], "edges": []}), + features="{}", + created_by=USER_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = WORKFLOW_ID + return workflow def _build_user(): - return SimpleNamespace(id=USER_ID, name="User", session_id="session") + return EndUser( + id=USER_ID, + tenant_id=TENANT_ID, + app_id=PIPELINE_ID, + type=EndUserType.BROWSER, + name="User", + session_id="session", + ) def _build_args(): @@ -86,10 +110,19 @@ def _patch_sqlite_engine(mocker: MockerFixture, sqlite_engine: Engine) -> None: mocker.patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=sqlite_engine) -def _patch_db_session(mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session]) -> scoped_session[Session]: - session_proxy = scoped_session(sqlite_session_factory) - mocker.patch.object(module.db, "session", session_proxy) - return session_proxy +def _persist_pipeline_scope( + session: Session, + *, + pipeline: Pipeline | None = None, + dataset: Dataset | None = None, + workflow: Workflow | None = None, +) -> tuple[Pipeline, Dataset, Workflow]: + pipeline = pipeline or _build_pipeline() + dataset = dataset or _build_pipeline_dataset() + workflow = workflow or _build_workflow() + session.add_all([pipeline, dataset, workflow]) + session.commit() + return pipeline, dataset, workflow def _persist_worker_records(session: Session) -> None: @@ -122,7 +155,6 @@ def _dummy_preserve(*args, **kwargs): def test_generate_dataset_missing(generator, sqlite_session: Session): pipeline = _build_pipeline() - pipeline.retrieve_dataset.return_value = None with pytest.raises(ValueError): generator.generate( @@ -137,8 +169,7 @@ def test_generate_dataset_missing(generator, sqlite_session: Session): def test_generate_debugger_calls_generate(generator, mocker: MockerFixture, sqlite_session: Session): - pipeline = _build_pipeline() - workflow = _build_workflow() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) mocker.patch.object( generator, @@ -181,8 +212,7 @@ def test_generate_debugger_calls_generate(generator, mocker: MockerFixture, sqli def test_generate_published_pipeline_creates_documents_and_delay( generator, mocker: MockerFixture, sqlite_session: Session ): - pipeline = _build_pipeline() - workflow = _build_workflow() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) datasource_info_list = [{"name": "file1"}, {"name": "file2"}] @@ -245,8 +275,7 @@ def test_generate_published_pipeline_creates_documents_and_delay( def test_generate_published_pipeline_rejects_when_document_creation_limits_exceeded( generator, mocker: MockerFixture, sqlite_session: Session ): - pipeline = _build_pipeline() - workflow = _build_workflow() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) datasource_info_list = [{"name": "file1"}, {"name": "file2"}] mocker.patch.object( @@ -283,8 +312,7 @@ def test_generate_published_pipeline_rejects_when_document_creation_limits_excee def test_generate_is_retry_calls_generate(generator, mocker: MockerFixture, sqlite_session: Session): - pipeline = _build_pipeline() - workflow = _build_workflow() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) mocker.patch.object( generator, @@ -336,14 +364,12 @@ def test_generate_worker_handles_errors( mocker: MockerFixture, sqlite_session: Session, sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], ): flask_app = MagicMock() flask_app.app_context.return_value = contextlib.nullcontext() mocker.patch.object(module, "preserve_flask_contexts", _dummy_preserve) _persist_worker_records(sqlite_session) _patch_sqlite_engine(mocker, sqlite_engine) - _patch_db_session(mocker, sqlite_session_factory) application_generate_entity = FakeRagPipelineGenerateEntity( app_config=SimpleNamespace(tenant_id=TENANT_ID, app_id=PIPELINE_ID, workflow_id=WORKFLOW_ID), @@ -374,14 +400,12 @@ def test_generate_worker_sets_system_user_id_for_external_call( mocker: MockerFixture, sqlite_session: Session, sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], ): flask_app = MagicMock() flask_app.app_context.return_value = contextlib.nullcontext() mocker.patch.object(module, "preserve_flask_contexts", _dummy_preserve) _persist_worker_records(sqlite_session) _patch_sqlite_engine(mocker, sqlite_engine) - _patch_db_session(mocker, sqlite_session_factory) application_generate_entity = FakeRagPipelineGenerateEntity( app_config=SimpleNamespace(tenant_id=TENANT_ID, app_id=PIPELINE_ID, workflow_id=WORKFLOW_ID), @@ -503,7 +527,6 @@ def test_single_iteration_generate_validates_inputs(generator, sqlite_session: S def test_single_iteration_generate_dataset_required(generator, sqlite_session: Session): pipeline = _build_pipeline() - pipeline.retrieve_dataset.return_value = None with pytest.raises(ValueError): generator.single_iteration_generate( @@ -520,9 +543,8 @@ def test_single_iteration_generate_success( generator, mocker: MockerFixture, sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], ): - pipeline = _build_pipeline() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) mocker.patch.object( module.PipelineConfigManager, @@ -539,8 +561,6 @@ def test_single_iteration_generate_success( "create_workflow_node_execution_repository", return_value=MagicMock(), ) - _patch_db_session(mocker, sqlite_session_factory) - mocker.patch.object(module, "WorkflowDraftVariableService", return_value=MagicMock()) mocker.patch.object(module, "DraftVarLoader", return_value=MagicMock()) @@ -548,7 +568,7 @@ def test_single_iteration_generate_success( result = generator.single_iteration_generate( pipeline, - _build_workflow(), + workflow, "node", _build_user(), {"inputs": {"a": 1}}, @@ -563,9 +583,8 @@ def test_single_loop_generate_success( generator, mocker: MockerFixture, sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], ): - pipeline = _build_pipeline() + pipeline, _, workflow = _persist_pipeline_scope(sqlite_session) mocker.patch.object( module.PipelineConfigManager, @@ -582,8 +601,6 @@ def test_single_loop_generate_success( "create_workflow_node_execution_repository", return_value=MagicMock(), ) - _patch_db_session(mocker, sqlite_session_factory) - mocker.patch.object(module, "WorkflowDraftVariableService", return_value=MagicMock()) mocker.patch.object(module, "DraftVarLoader", return_value=MagicMock()) @@ -591,7 +608,7 @@ def test_single_loop_generate_success( result = generator.single_loop_generate( pipeline, - _build_workflow(), + workflow, "node", _build_user(), {"inputs": {"a": 1}}, @@ -622,12 +639,7 @@ def test_handle_response_value_error_triggers_generate_task_stopped(generator, m ) -def test_build_document_sets_metadata_for_builtin_fields(generator, mocker: MockerFixture): - class DummyDocument(SimpleNamespace): - pass - - mocker.patch.object(module, "Document", side_effect=lambda **kwargs: DummyDocument(**kwargs)) - +def test_build_document_sets_metadata_for_builtin_fields(generator): document = generator._build_document( tenant_id="tenant", dataset_id="ds", @@ -693,7 +705,7 @@ def test_format_datasource_info_list_non_online_drive(generator): def test_format_datasource_info_list_missing_node_data(generator): - workflow = MagicMock(graph_dict={"nodes": []}) + workflow = _build_workflow() with pytest.raises(ValueError): generator._format_datasource_info_list( @@ -707,8 +719,8 @@ def test_format_datasource_info_list_missing_node_data(generator): def test_format_datasource_info_list_online_drive_folder(generator, mocker: MockerFixture): - workflow = MagicMock( - graph_dict={ + workflow = _build_workflow( + graph={ "nodes": [ { "id": "start", @@ -720,7 +732,7 @@ def test_format_datasource_info_list_online_drive_folder(generator, mocker: Mock }, } ] - } + }, ) runtime = MagicMock() diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index f81620d6606..3e847f65d73 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -1,10 +1,3 @@ -""" -Unit tests for PipelineRunner behavior. -Asserts correct event handling, error propagation, and user invocation logic. -Primary collaborators: PipelineRunner, InvokeFrom, GraphRunFailedEvent, UserFrom, and mocked dependencies. -Cross-references: core.app.apps.pipeline.pipeline_runner, core.app.entities.app_invoke_entities. -""" - """Unit tests for PipelineRunner behavior. This module validates core control-flow outcomes for @@ -18,16 +11,100 @@ Primary collaborators include ``PipelineRunner``, ``UserFrom``, and patched DB/runtime dependencies used by the runner. """ +import json from types import SimpleNamespace from unittest.mock import MagicMock import pytest from pytest_mock import MockerFixture +from sqlalchemy import event +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session import core.app.apps.pipeline.pipeline_runner as module from core.app.apps.pipeline.pipeline_runner import PipelineRunner from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from graphon.graph_events import GraphRunFailedEvent +from models.dataset import Dataset, Document, Pipeline +from models.enums import DataSourceType, DocumentCreatedFrom, EndUserType +from models.model import EndUser +from models.workflow import Workflow, WorkflowType + + +def _pipeline(*, tenant_id: str = "tenant", pipeline_id: str = "pipe") -> Pipeline: + pipeline = Pipeline(tenant_id=tenant_id, name="Pipeline", description="") + pipeline.id = pipeline_id + pipeline.workflow_id = "wf" + return pipeline + + +def _dataset(*, tenant_id: str = "tenant", dataset_id: str = "ds", pipeline_id: str = "pipe") -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + description="", + created_by="user", + pipeline_id=pipeline_id, + ) + + +def _workflow(*, tenant_id: str = "tenant", pipeline_id: str = "pipe", graph: dict | None = None) -> Workflow: + return Workflow.new( + tenant_id=tenant_id, + app_id=pipeline_id, + type=WorkflowType.RAG_PIPELINE.value, + version="v1", + graph=json.dumps(graph if graph is not None else {"nodes": [], "edges": []}), + features="{}", + created_by="user", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + +def _end_user() -> EndUser: + return EndUser( + id="user", + tenant_id="tenant", + app_id="pipe", + type=EndUserType.BROWSER, + name="User", + session_id="sess", + ) + + +def _document(*, document_id: str = "doc", dataset_id: str = "ds", tenant_id: str = "tenant") -> Document: + return Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="user", + ) + + +def _persist_scope( + session: Session, + *, + pipeline: Pipeline | None = None, + dataset: Dataset | None = None, + workflow: Workflow | None = None, + end_user: EndUser | None = None, + documents: tuple[Document, ...] = (), +) -> tuple[Pipeline, Dataset, Workflow]: + pipeline = pipeline or _pipeline() + dataset = dataset or _dataset(tenant_id=pipeline.tenant_id, pipeline_id=pipeline.id) + workflow = workflow or _workflow(tenant_id=pipeline.tenant_id, pipeline_id=pipeline.id) + workflow.id = "wf" + session.add_all([pipeline, dataset, workflow, *(documents or ()), *([end_user] if end_user else [])]) + session.commit() + return pipeline, dataset, workflow def _build_app_generate_entity() -> SimpleNamespace: @@ -53,38 +130,12 @@ def _build_app_generate_entity() -> SimpleNamespace: ) -def _patch_create_session(mocker: MockerFixture, session: MagicMock, *, events: list[str] | None = None): - """Patch create_session() to yield ``session`` inside its ``with`` body and ``begin()`` block. - - The runner now obtains short-lived sessions via ``create_session()`` instead of the - Flask scoped ``db.session``, so tests patch the module-level ``create_session`` and - hand back a context manager that yields the mock session. - """ - session_context = MagicMock() - - def enter_session(): - if events is not None: - events.append("session_enter") - return session - - def exit_session(*args): - if events is not None: - events.append("session_exit") - return False - - session_context.__enter__.side_effect = enter_session - session_context.__exit__.side_effect = exit_session - session.begin.return_value.__enter__.return_value = session - session.begin.return_value.__exit__.return_value = False - return mocker.patch.object(module, "create_session", return_value=session_context) - - @pytest.fixture def runner(): app_generate_entity = _build_app_generate_entity() queue_manager = MagicMock() variable_loader = MagicMock() - workflow = MagicMock() + workflow = _workflow() workflow_execution_repository = MagicMock() workflow_node_execution_repository = MagicMock() @@ -103,105 +154,91 @@ def test_get_app_id(runner): assert runner._get_app_id() == "pipe" -def test_get_workflow_returns_workflow(runner): - pipeline = MagicMock(tenant_id="tenant", id="pipe") - workflow = MagicMock(id="wf") +def test_get_workflow_returns_workflow(runner, sqlite_session: Session): + pipeline, _, workflow = _persist_scope(sqlite_session) - session = MagicMock() - session.scalar.return_value = workflow - - result = runner.get_workflow(session=session, pipeline=pipeline, workflow_id="wf") + result = runner.get_workflow(session=sqlite_session, pipeline=pipeline, workflow_id="wf") assert result == workflow def test_init_rag_pipeline_graph_invalid_config(mocker, runner): - workflow = MagicMock(id="wf", tenant_id="tenant", graph_dict={}) + workflow = _workflow(graph={}) with pytest.raises(ValueError): runner._init_rag_pipeline_graph(workflow=workflow, graph_runtime_state=MagicMock()) - workflow.graph_dict = {"nodes": "bad", "edges": []} + workflow.graph = json.dumps({"nodes": "bad", "edges": []}) with pytest.raises(ValueError): runner._init_rag_pipeline_graph(workflow=workflow, graph_runtime_state=MagicMock()) - workflow.graph_dict = {"nodes": [], "edges": "bad"} + workflow.graph = json.dumps({"nodes": [], "edges": "bad"}) with pytest.raises(ValueError): runner._init_rag_pipeline_graph(workflow=workflow, graph_runtime_state=MagicMock()) def test_init_rag_pipeline_graph_not_found(mocker, runner): - workflow = MagicMock(id="wf", tenant_id="tenant", graph_dict={"nodes": [], "edges": []}) + workflow = _workflow() mocker.patch.object(module.Graph, "init", return_value=None) with pytest.raises(ValueError): runner._init_rag_pipeline_graph(workflow=workflow, graph_runtime_state=MagicMock()) -def test_update_document_status_on_failure(mocker, runner): - document = MagicMock() - document_ref = MagicMock() - - session = MagicMock() - _patch_create_session(mocker, session) - get_document_by_ref = mocker.patch.object( - module.DatasetRefService, - "get_document_by_ref", - return_value=document, - ) +def test_update_document_status_on_failure(runner, sqlite_session: Session): + document = _document() + _, dataset, _ = _persist_scope(sqlite_session, documents=(document,)) + dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) + document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, document.id) event = GraphRunFailedEvent(error="boom") runner._update_document_status(event, document_ref) - get_document_by_ref.assert_called_once_with(document_ref, session=session) - assert document.indexing_status == "error" - assert document.error == "boom" - session.add.assert_called_once_with(document) - session.begin.assert_called_once() - session.begin.return_value.__enter__.assert_called_once() - session.begin.return_value.__exit__.assert_called_once() + sqlite_session.expire_all() + updated = sqlite_session.get(Document, document.id) + assert updated is not None + assert updated.indexing_status == "error" + assert updated.error == "boom" -def test_update_document_status_skips_when_document_not_found(mocker, runner): - document_ref = MagicMock() - session = MagicMock() - _patch_create_session(mocker, session) - get_document_by_ref = mocker.patch.object( - module.DatasetRefService, - "get_document_by_ref", - return_value=None, - ) +def test_update_document_status_skips_when_document_not_found(runner, sqlite_session: Session): + _, dataset, _ = _persist_scope(sqlite_session) + dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) + document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, "missing") runner._update_document_status(GraphRunFailedEvent(error="boom"), document_ref) - get_document_by_ref.assert_called_once_with(document_ref, session=session) - session.add.assert_not_called() + assert sqlite_session.get(Document, "missing") is None -def test_update_document_status_skips_without_document_ref(mocker, runner): - create_session = mocker.patch.object(module, "create_session") +def test_update_document_status_skips_without_document_ref(runner, sqlite_engine: Engine): + checkouts = 0 - runner._update_document_status(GraphRunFailedEvent(error="boom"), None) + def record_checkout(*_args) -> None: + nonlocal checkouts + checkouts += 1 - create_session.assert_not_called() + event.listen(sqlite_engine, "checkout", record_checkout) + try: + runner._update_document_status(GraphRunFailedEvent(error="boom"), None) + finally: + event.remove(sqlite_engine, "checkout", record_checkout) + + assert checkouts == 0 -def test_run_pipeline_not_found(mocker: MockerFixture): +def test_run_pipeline_not_found(): app_generate_entity = _build_app_generate_entity() app_generate_entity.invoke_from = InvokeFrom.WEB_APP app_generate_entity.single_iteration_run = None app_generate_entity.single_loop_run = None - session = MagicMock() - session.get.side_effect = [None, None] - _patch_create_session(mocker, session) - runner = PipelineRunner( application_generate_entity=app_generate_entity, queue_manager=MagicMock(), variable_loader=MagicMock(), - workflow=MagicMock(), + workflow=_workflow(), system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), @@ -211,177 +248,116 @@ def test_run_pipeline_not_found(mocker: MockerFixture): runner.run() -def test_run_pipeline_from_other_tenant_is_not_found(mocker: MockerFixture, runner: PipelineRunner): - pipeline = MagicMock(id="pipe", tenant_id="other-tenant") - session = MagicMock() - session.get.side_effect = [None, pipeline] - _patch_create_session(mocker, session) +def test_run_pipeline_from_other_tenant_is_not_found(runner: PipelineRunner, sqlite_session: Session): + pipeline = _pipeline(tenant_id="other-tenant") + sqlite_session.add(pipeline) + sqlite_session.commit() with pytest.raises(ValueError, match="Pipeline not found"): runner.run() - pipeline.retrieve_dataset.assert_not_called() - @pytest.mark.parametrize( "dataset", [ pytest.param(None, id="missing"), - pytest.param(SimpleNamespace(id="ds", tenant_id="other-tenant"), id="other-tenant"), - pytest.param(SimpleNamespace(id="other-dataset", tenant_id="tenant"), id="other-dataset"), + pytest.param(_dataset(tenant_id="other-tenant"), id="other-tenant"), + pytest.param(_dataset(dataset_id="other-dataset"), id="other-dataset"), ], ) def test_run_rejects_unowned_pipeline_dataset( - mocker: MockerFixture, runner: PipelineRunner, - dataset: SimpleNamespace | None, + dataset: Dataset | None, + sqlite_session: Session, ): - pipeline = MagicMock(id="pipe", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = dataset - session = MagicMock() - session.get.side_effect = [None, pipeline] - _patch_create_session(mocker, session) + pipeline = _pipeline() + sqlite_session.add(pipeline) + if dataset is not None: + sqlite_session.add(dataset) + sqlite_session.commit() runner.get_workflow = MagicMock() with pytest.raises(ValueError, match="Pipeline dataset not found"): runner.run() - pipeline.retrieve_dataset.assert_called_once_with(session) runner.get_workflow.assert_not_called() def test_run_rejects_document_outside_pipeline_dataset_after_async_boundary( - mocker: MockerFixture, runner: PipelineRunner, + sqlite_session: Session, ): runner.application_generate_entity.document_id = "foreign-doc" runner.application_generate_entity.original_document_id = "foreign-doc" - pipeline = MagicMock(id="pipe", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") - session = MagicMock() - session.get.side_effect = [None, pipeline] - _patch_create_session(mocker, session) + _persist_scope(sqlite_session) runner.get_workflow = MagicMock() - get_document_by_ref = mocker.patch.object( - module.DatasetRefService, - "get_document_by_ref", - return_value=None, - ) with pytest.raises(ValueError, match="Pipeline document not found"): runner.run() - document_ref = get_document_by_ref.call_args.args[0] - assert document_ref.dataset.tenant_id == "tenant" - assert document_ref.dataset.dataset_id == "ds" - assert document_ref.document_id == "foreign-doc" - get_document_by_ref.assert_called_once_with(document_ref, session=session) runner.get_workflow.assert_not_called() def test_run_rejects_original_document_outside_pipeline_dataset_after_async_boundary( - mocker: MockerFixture, runner: PipelineRunner, + sqlite_session: Session, ): runner.application_generate_entity.document_id = "doc" runner.application_generate_entity.original_document_id = "foreign-doc" - pipeline = MagicMock(id="pipe", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") - session = MagicMock() - session.get.side_effect = [None, pipeline] - _patch_create_session(mocker, session) + _persist_scope(sqlite_session, documents=(_document(),)) runner.get_workflow = MagicMock() - get_document_by_ref = mocker.patch.object( - module.DatasetRefService, - "get_document_by_ref", - side_effect=[SimpleNamespace(id="doc"), None], - ) with pytest.raises(ValueError, match="Pipeline original document not found"): runner.run() - document_refs = [call.args[0] for call in get_document_by_ref.call_args_list] - assert [document_ref.document_id for document_ref in document_refs] == ["doc", "foreign-doc"] - for document_ref in document_refs: - assert document_ref.dataset.tenant_id == "tenant" - assert document_ref.dataset.dataset_id == "ds" runner.get_workflow.assert_not_called() -def test_run_workflow_not_initialized(mocker: MockerFixture): +def test_run_workflow_not_initialized(sqlite_session: Session): app_generate_entity = _build_app_generate_entity() - pipeline = MagicMock(id="pipe", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") - - session = MagicMock() - session.get.side_effect = [None, pipeline] - _patch_create_session(mocker, session) + pipeline = _pipeline() + dataset = _dataset() + document = _document() + sqlite_session.add_all([pipeline, dataset, document]) + sqlite_session.commit() runner = PipelineRunner( application_generate_entity=app_generate_entity, queue_manager=MagicMock(), variable_loader=MagicMock(), - workflow=MagicMock(), + workflow=_workflow(), system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), ) - runner.get_workflow = MagicMock(return_value=None) - with pytest.raises(ValueError): runner.run() -def test_run_single_iteration_path(mocker: MockerFixture): +def test_run_single_iteration_path(mocker: MockerFixture, sqlite_session: Session): app_generate_entity = _build_app_generate_entity() app_generate_entity.single_iteration_run = MagicMock() - pipeline = MagicMock(id="pipe", tenant_id="tenant") - dataset = SimpleNamespace(id="ds", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = dataset - session = MagicMock() - session.get.return_value = pipeline - _patch_create_session(mocker, session) + _, dataset, _ = _persist_scope(sqlite_session, documents=(_document(),)) + dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) + document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, "doc") runner = PipelineRunner( application_generate_entity=app_generate_entity, queue_manager=MagicMock(), variable_loader=MagicMock(), - workflow=MagicMock(), + workflow=_workflow(), system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), ) runner._resolve_user_from = MagicMock(return_value=UserFrom.ACCOUNT) - runner.get_workflow = MagicMock( - return_value=MagicMock( - id="wf", - tenant_id="tenant", - app_id="pipe", - graph_dict={}, - type="rag-pipeline", - version="v1", - ) - ) runner._prepare_single_node_execution = MagicMock(return_value=("graph", "pool", "state")) runner._update_document_status = MagicMock() runner._handle_event = MagicMock() - dataset_ref = MagicMock() - document_ref = MagicMock() - create_dataset_ref = mocker.patch.object( - module.DatasetRefService, - "create_dataset_ref", - return_value=dataset_ref, - ) - create_document_ref_from_id = mocker.patch.object( - module.DatasetRefService, - "create_document_ref_from_id", - return_value=document_ref, - ) - event = MagicMock() workflow_entry = MagicMock() workflow_entry.graph_engine = MagicMock() @@ -392,34 +368,30 @@ def test_run_single_iteration_path(mocker: MockerFixture): runner.run() - create_dataset_ref.assert_called_once_with(dataset) - create_document_ref_from_id.assert_called_once_with(dataset_ref, "doc") runner._prepare_single_node_execution.assert_called_once() runner._update_document_status.assert_called_once_with(event, document_ref) runner._handle_event.assert_called() -def test_run_normal_path_builds_graph(mocker: MockerFixture): +def test_run_normal_path_builds_graph(mocker: MockerFixture, sqlite_session: Session, sqlite_engine: Engine): app_generate_entity = _build_app_generate_entity() - pipeline = MagicMock(id="pipe", tenant_id="tenant") - pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") - end_user = MagicMock(session_id="sess") events = [] - - session = MagicMock() - session.get.side_effect = [end_user, pipeline] - _patch_create_session(mocker, session, events=events) - - workflow = MagicMock( - id="wf", - tenant_id="tenant", - app_id="pipe", - graph_dict={"nodes": [], "edges": []}, - environment_variables=[], - rag_pipeline_variables=[{"variable": "input1", "belong_to_node_id": "start"}], - type="rag-pipeline", - version="v1", + workflow = _workflow() + workflow.rag_pipeline_variables = [ + { + "variable": "input1", + "belong_to_node_id": "start", + "type": "text-input", + "label": "Input", + } + ] + workflow.id = "wf" + _persist_scope( + sqlite_session, + workflow=workflow, + end_user=_end_user(), + documents=(_document(),), ) runner = PipelineRunner( @@ -433,18 +405,10 @@ def test_run_normal_path_builds_graph(mocker: MockerFixture): ) runner._resolve_user_from = MagicMock(return_value=UserFrom.ACCOUNT) - runner.get_workflow = MagicMock(return_value=workflow) runner._init_rag_pipeline_graph = MagicMock(return_value="graph") runner._update_document_status = MagicMock() runner._handle_event = MagicMock() - mocker.patch.object( - module.RAGPipelineVariable, - "model_validate", - return_value=SimpleNamespace(belong_to_node_id="start", variable="input1"), - ) - mocker.patch.object(module, "RAGPipelineVariableInput", side_effect=lambda **kwargs: SimpleNamespace(**kwargs)) - class FakeVariablePool: def add(self, selector, value): return None @@ -457,7 +421,15 @@ def test_run_normal_path_builds_graph(mocker: MockerFixture): mocker.patch.object(module, "WorkflowEntry", return_value=workflow_entry) mocker.patch.object(module, "WorkflowPersistenceLayer", return_value=MagicMock()) - runner.run() + def record_checkin(*_args) -> None: + events.append("session_checkin") - assert events == ["session_enter", "session_exit", "workflow_run"] + event.listen(sqlite_engine, "checkin", record_checkin) + try: + runner.run() + finally: + event.remove(sqlite_engine, "checkin", record_checkin) + + assert events[-1] == "workflow_run" + assert "session_checkin" in events[:-1] runner._init_rag_pipeline_graph.assert_called_once() diff --git a/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py index 4e25532650d..b926e694202 100644 --- a/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py +++ b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py @@ -2,7 +2,12 @@ from unittest.mock import Mock, patch import pytest -from core.app.apps.execution_coordinator import AppExecutionCoordinator, AppExecutionState +from core.app.apps.execution_coordinator import ( + AppExecutionCoordinator, + AppExecutionState, + app_task_command_channel_key, + clear_app_task_cancellation_signals, +) def test_listener_close_does_not_abort_running_attempt() -> None: @@ -74,6 +79,54 @@ def test_pausing_started_attempt_cancels_watchdog() -> None: assert coordinator.state is AppExecutionState.PAUSED +def test_clearing_cancellation_signals_drops_stop_flag_and_queued_commands() -> None: + channel = Mock() + channel.fetch_commands.return_value = [Mock()] + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.RedisChannel", return_value=channel) as redis_channel, + ): + clear_app_task_cancellation_signals("task") + + redis_channel.assert_called_once_with(redis_client, "workflow:task:commands") + channel.fetch_commands.assert_called_once_with() + assert redis_client.delete.call_args_list == [ + (("generate_task_stopped:task",), {}), + (("workflow:task:commands",), {}), + ] + + +def test_clearing_cancellation_signals_ignores_empty_task_id() -> None: + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.RedisChannel") as redis_channel, + ): + clear_app_task_cancellation_signals("") + + redis_client.delete.assert_not_called() + redis_channel.assert_not_called() + + +def test_clearing_cancellation_signals_survives_command_channel_failure( + caplog: pytest.LogCaptureFixture, +) -> None: + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.RedisChannel") as redis_channel, + ): + redis_channel.return_value.fetch_commands.side_effect = RuntimeError("redis read failed") + + clear_app_task_cancellation_signals("task") + + # The stop flag is cleared first, so a command-channel failure cannot leave it armed. + redis_client.delete.assert_called_once_with("generate_task_stopped:task") + assert "Failed to clear pending GraphEngine commands for app task task" in caplog.text + + +def test_command_channel_key_matches_the_channel_the_stop_command_targets() -> None: + assert app_task_command_channel_key("task") == "workflow:task:commands" + + def test_stop_flag_failure_does_not_block_graph_stop(caplog: pytest.LogCaptureFixture) -> None: on_timeout = Mock() with ( diff --git a/api/tests/unit_tests/core/app/apps/test_pause_resume.py b/api/tests/unit_tests/core/app/apps/test_pause_resume.py index 6bfda02eb4c..2b1690c2f5e 100644 --- a/api/tests/unit_tests/core/app/apps/test_pause_resume.py +++ b/api/tests/unit_tests/core/app/apps/test_pause_resume.py @@ -1,5 +1,6 @@ import sys import time +from collections.abc import Mapping from types import ModuleType, SimpleNamespace from typing import Any @@ -106,10 +107,19 @@ class _StubToolNode(Node[_StubToolNodeData]): def _patch_tool_node(mocker: MockerFixture): original_resolve_node_class = node_factory_module.resolve_workflow_node_class - def _patched_resolve_node_class(*, node_type: NodeType, node_version: str) -> type[Node]: + def _patched_resolve_node_class( + *, + node_type: NodeType, + node_version: str, + node_data: Mapping[str, Any] | BaseNodeData | None = None, + ) -> type[Node]: if node_type == BuiltinNodeTypes.TOOL: return _StubToolNode - return original_resolve_node_class(node_type=node_type, node_version=node_version) + return original_resolve_node_class( + node_type=node_type, + node_version=node_version, + node_data=node_data, + ) mocker.patch.object(node_factory_module, "resolve_workflow_node_class", side_effect=_patched_resolve_node_class) diff --git a/api/tests/unit_tests/core/app/llm/test_model_access.py b/api/tests/unit_tests/core/app/llm/test_model_access.py new file mode 100644 index 00000000000..8c22d56223e --- /dev/null +++ b/api/tests/unit_tests/core/app/llm/test_model_access.py @@ -0,0 +1,60 @@ +from types import SimpleNamespace +from typing import cast + +import pytest + +from core.app.entities.app_invoke_entities import DifyRunContext +from core.app.llm import model_access +from graphon.model_runtime.entities.model_entities import ModelPropertyKey + + +def _stub_model_factory(monkeypatch: pytest.MonkeyPatch, context_window: object) -> dict[str, object]: + calls: dict[str, object] = {} + + class FakeModelFactory: + def __init__(self, *, run_context: DifyRunContext) -> None: + calls["run_context"] = run_context + + def init_model_instance(self, provider_name: str, model_name: str) -> object: + calls["provider_name"] = provider_name + calls["model_name"] = model_name + schema = SimpleNamespace(model_properties={ModelPropertyKey.CONTEXT_SIZE: context_window}) + return SimpleNamespace(get_model_schema=lambda: schema) + + monkeypatch.setattr(model_access, "DifyModelFactory", FakeModelFactory) + return calls + + +def test_resolve_model_context_window_reads_selected_model_schema(monkeypatch: pytest.MonkeyPatch) -> None: + calls = _stub_model_factory(monkeypatch, 128_000) + run_context = cast(DifyRunContext, object()) + + context_window = model_access.resolve_model_context_window( + run_context=run_context, + provider_name="langgenius/openai/openai", + model_name="gpt-4o", + ) + + assert context_window == 128_000 + assert calls == { + "run_context": run_context, + "provider_name": "langgenius/openai/openai", + "model_name": "gpt-4o", + } + + +@pytest.mark.parametrize("context_window", [None, 0, -1, True, False, "128000", 128_000.0]) +def test_resolve_model_context_window_ignores_invalid_schema_values( + monkeypatch: pytest.MonkeyPatch, + context_window: object, +) -> None: + _ = _stub_model_factory(monkeypatch, context_window) + + assert ( + model_access.resolve_model_context_window( + run_context=cast(DifyRunContext, object()), + provider_name="openai", + model_name="gpt-4o", + ) + is None + ) diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py index 51d817d8f53..436d9b1280b 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch from sqlalchemy import event from sqlalchemy.orm import Session, sessionmaker +from core.rag.index_processor import index_processor as index_processor_module from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.index_processor import IndexProcessor from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem @@ -237,12 +238,10 @@ class TestIndexProcessor: "core.rag.index_processor.index_processor.current_app", SimpleNamespace(_get_current_object=lambda: flask_app), ), - patch( - "core.rag.index_processor.index_processor.session_factory", - SimpleNamespace(create_session=sqlite_session_factory), - ), - patch( - "core.rag.index_processor.index_processor.ParagraphIndexProcessor.generate_summary", + patch.object(index_processor_module.session_factory, "create_session", sqlite_session_factory), + patch.object( + index_processor_module.ParagraphIndexProcessor, + "generate_summary", side_effect=generate_summary, ) as generate_summary, ): diff --git a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py b/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py index ecb2985dc4b..460a8491125 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py +++ b/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py @@ -49,11 +49,13 @@ for the full indexing pipeline are handled separately in the integration test su import json import uuid -from types import SimpleNamespace +from datetime import UTC, datetime from typing import Any from unittest.mock import MagicMock, Mock, patch import pytest +from sqlalchemy import event +from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm.exc import ObjectDeletedError from core.entities.knowledge_entities import PreviewDetail @@ -66,12 +68,13 @@ from core.indexing_runner import ( from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.models.document import ChildDocument, Document from enums import DeploymentEdition +from extensions.storage.storage_type import StorageType from graphon.model_runtime.entities.model_entities import ModelType from libs.datetime_utils import naive_utc_now -from models.dataset import Dataset, DatasetProcessRule, DocumentSegment +from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment from models.dataset import Document as DatasetDocument -from models.enums import SegmentStatus -from models.model import Account +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, SegmentStatus +from models.model import Account, UploadFile from services.vector_space_admission_service import VectorSpaceAdmissionError # ============================================================================ @@ -156,6 +159,33 @@ def create_mock_dataset_document( ) +def persist_indexing_scope(session: Session) -> tuple[Dataset, DatasetDocument]: + tenant_id = str(uuid.uuid4()) + created_by = str(uuid.uuid4()) + dataset = Dataset( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + name="Dataset", + description="", + created_by=created_by, + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + ) + document = DatasetDocument( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by=created_by, + ) + session.add_all([dataset, document]) + session.commit() + return dataset, document + + def create_sample_documents( count: int = 3, include_children: bool = False, @@ -263,14 +293,14 @@ class TestIndexingRunnerExtract: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, sqlite_session: Session): """Mock all external dependencies for extract tests.""" with ( patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, patch("core.indexing_runner.storage") as mock_storage, ): yield { - "session": MagicMock(), + "session": sqlite_session, "factory": mock_factory, "storage": mock_storage, } @@ -329,16 +359,28 @@ class TestIndexingRunnerExtract: ), ] mock_processor.extract.return_value = extracted_docs + file_id = json.loads(sample_dataset_document.data_source_info)["upload_file_id"] + upload_file = UploadFile( + tenant_id=sample_dataset_document.tenant_id, + storage_type=StorageType.LOCAL, + key="uploads/test.pdf", + name="test.pdf", + size=10, + extension="pdf", + mime_type="application/pdf", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-id", + created_at=datetime.now(UTC), + used=True, + ) + upload_file.id = file_id + mock_dependencies["session"].add(upload_file) + mock_dependencies["session"].commit() - # Mock the entire _extract method to avoid ExtractSetting validation - # This is necessary because ExtractSetting uses Pydantic validation - with patch.object(runner, "_update_document_index_status"): - with patch("core.indexing_runner.select"): - with patch("core.indexing_runner.ExtractSetting"): - # Act: Call the extract method - result = runner._extract( - mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] - ) + with patch.object(runner, "_update_document_index_status"), patch("core.indexing_runner.ExtractSetting"): + result = runner._extract( + mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] + ) # Assert: Verify the extraction results assert len(result) == 2, "Should extract 2 documents from the PDF" @@ -460,13 +502,13 @@ class TestIndexingRunnerTransform: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, unbound_session: Session): """Mock all external dependencies for transform tests.""" with ( patch("core.indexing_runner.ModelManager.for_tenant") as mock_model_manager, ): yield { - "session": MagicMock(), + "session": unbound_session, "model_manager": mock_model_manager, } @@ -625,7 +667,7 @@ class TestIndexingRunnerLoad: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, unbound_session: Session): """Mock all external dependencies for load tests.""" with ( patch("core.indexing_runner.ModelManager.for_tenant") as mock_model_manager, @@ -635,7 +677,7 @@ class TestIndexingRunnerLoad: ): mock_app._get_current_object = Mock(return_value=Mock()) yield { - "session": MagicMock(), + "session": unbound_session, "model_manager": mock_model_manager, "app": mock_app, "thread": mock_thread, @@ -830,7 +872,7 @@ class TestIndexingRunnerRun: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, sqlite_session: Session): """Mock all external dependencies for run tests.""" with ( patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, @@ -839,7 +881,7 @@ class TestIndexingRunnerRun: patch("core.indexing_runner.threading.Thread") as mock_thread, ): yield { - "session": MagicMock(), + "session": sqlite_session, "factory": mock_factory, "model_manager": mock_model_manager, "storage": mock_storage, @@ -847,22 +889,46 @@ class TestIndexingRunnerRun: } @pytest.fixture - def sample_dataset_documents(self): + def sample_dataset_documents(self, sqlite_session: Session): """Create sample dataset documents for testing.""" docs = [] for i in range(2): - docs.append( - DatasetDocument( - id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - data_source_type="upload_file", - data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}), - dataset_process_rule_id=str(uuid.uuid4()), - ) + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + account = Account(name=f"Account {i}", email=f"account-{i}@example.com") + account.id = str(uuid.uuid4()) + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name=f"Dataset {i}", + description="", + created_by=account.id, + indexing_technique=IndexTechniqueType.HIGH_QUALITY, ) + process_rule = DatasetProcessRule( + dataset_id=dataset_id, + mode="automatic", + rules="{}", + created_by=account.id, + ) + document = DatasetDocument( + id=str(uuid.uuid4()), + dataset_id=dataset_id, + tenant_id=tenant_id, + position=1, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + doc_language="English", + data_source_type=DataSourceType.UPLOAD_FILE, + data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}), + dataset_process_rule_id=process_rule.id, + batch="batch", + name=f"Document {i}", + created_from=DocumentCreatedFrom.API, + created_by=account.id, + ) + sqlite_session.add_all([account, dataset, process_rule, document]) + docs.append(document) + sqlite_session.commit() return docs def test_run_in_indexing_status_loads_child_chunks_with_caller_session( @@ -871,9 +937,15 @@ class TestIndexingRunnerRun: runner = IndexingRunner() dataset_document = sample_dataset_documents[0] dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX - dataset = Dataset() + session = mock_dependencies["session"] + dataset = session.get(Dataset, dataset_document.dataset_id) + assert dataset is not None + process_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) + assert process_rule is not None + process_rule.mode = "hierarchical" + process_rule.rules = json.dumps({"parent_mode": "paragraph"}) segment = DocumentSegment( - tenant_id="tenant-id", + tenant_id=dataset_document.tenant_id, dataset_id=dataset_document.dataset_id, document_id=dataset_document.id, position=1, @@ -885,29 +957,39 @@ class TestIndexingRunnerRun: index_node_id="parent-node", index_node_hash="parent-hash", ) - child_chunks = [SimpleNamespace(content="child", index_node_id="child-node", index_node_hash="child-hash")] - session = mock_dependencies["session"] - session.get.side_effect = lambda model, _: dataset_document if model is DatasetDocument else dataset - session.scalars.return_value.all.return_value = [segment] + child_chunk = ChildChunk( + tenant_id=dataset_document.tenant_id, + dataset_id=dataset_document.dataset_id, + document_id=dataset_document.id, + segment_id=segment.id, + position=1, + content="child", + word_count=1, + created_by=dataset_document.created_by, + index_node_id="child-node", + index_node_hash="child-hash", + ) + session.add_all([segment, child_chunk]) + session.commit() with ( - patch.object(DocumentSegment, "get_child_chunks", return_value=child_chunks) as get_child_chunks, patch.object(runner, "_load") as load, ): runner.run_in_indexing_status(dataset_document, session) - get_child_chunks.assert_called_once_with(session=session) assert load.call_args.kwargs["documents"][0].children[0].page_content == "child" assert load.call_args.kwargs["total_tokens"] == 12 def test_run_in_indexing_status_uses_tokens_from_all_segments(self, mock_dependencies, sample_dataset_documents): runner = IndexingRunner() dataset_document = sample_dataset_documents[0] - dataset = Dataset() + session = mock_dependencies["session"] + dataset = session.get(Dataset, dataset_document.dataset_id) + assert dataset is not None completed_segment = DocumentSegment( - tenant_id="tenant-id", - dataset_id="dataset-id", - document_id="document-id", + tenant_id=dataset_document.tenant_id, + dataset_id=dataset_document.dataset_id, + document_id=dataset_document.id, position=1, content="", word_count=0, @@ -928,9 +1010,8 @@ class TestIndexingRunnerRun: index_node_id="pending-node", index_node_hash="pending-hash", ) - session = mock_dependencies["session"] - session.get.side_effect = lambda model, _: dataset_document if model is DatasetDocument else dataset - session.scalars.return_value.all.return_value = [completed_segment, incomplete_segment] + session.add_all([completed_segment, incomplete_segment]) + session.commit() with patch.object(runner, "_load") as load: runner.run_in_indexing_status(dataset_document, session) @@ -944,23 +1025,13 @@ class TestIndexingRunnerRun: # Arrange runner = IndexingRunner() doc = sample_dataset_documents[0] - - # Mock database queries - mock_dataset = Dataset( - id=doc.dataset_id, - tenant_id=doc.tenant_id, - indexing_technique=IndexTechniqueType.ECONOMY, - ) - - mock_current_user = Account(name="Test Account", email="test@example.com") - - get_dispatch = {"Document": doc, "Dataset": mock_dataset, "Account": mock_current_user} - mock_dependencies["session"].get.side_effect = lambda model, id: get_dispatch.get(model.__name__) - - mock_process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = mock_process_rule + session = mock_dependencies["session"] + mock_dataset = session.get(Dataset, doc.dataset_id) + mock_current_user = session.get(Account, doc.created_by) + assert mock_dataset is not None + assert mock_current_user is not None + mock_dataset.indexing_technique = IndexTechniqueType.ECONOMY + session.commit() # Mock processor mock_processor = MagicMock() @@ -979,30 +1050,13 @@ class TestIndexingRunnerRun: mock_thread_instance = MagicMock() mock_dependencies["thread"].return_value = mock_thread_instance - # Mock all internal methods that interact with database - with ( - patch.object(runner, "_extract", return_value=[Document(page_content="Test", metadata={})]), - patch.object( - runner, - "_transform", - return_value=[Document(page_content="Chunk", metadata={"doc_id": "c1", "doc_hash": "h1"})], - ), - patch.object(runner, "_load_segments"), - patch.object(runner, "_load"), - ): - # Act - mock_dependencies["session"].commit.reset_mock() - runner.run([doc], mock_dependencies["session"]) + commits = 0 - # Assert - verify the methods were called - # Since we're mocking the internal methods, we just verify no exceptions were raised - set_tenant_id.assert_called_once_with( - mock_current_user, - mock_dataset.tenant_id, - session=mock_dependencies["session"], - ) + def record_commit(_session) -> None: + nonlocal commits + commits += 1 - mock_dependencies["session"].commit.reset_mock() + event.listen(session, "after_commit", record_commit) with ( patch.object(runner, "_extract", return_value=[Document(page_content="Test", metadata={})]) as mock_extract, patch.object( @@ -1014,14 +1068,16 @@ class TestIndexingRunnerRun: patch.object(runner, "_load") as mock_load, ): # Act - runner.run([doc], mock_dependencies["session"]) + runner.run([doc], session) + event.remove(session, "after_commit", record_commit) # Assert - verify the methods were called mock_extract.assert_called_once() mock_transform.assert_called_once() mock_load_segments.assert_called_once() mock_load.assert_called_once() - assert mock_dependencies["session"].commit.call_count == 2 + assert commits == 2 + set_tenant_id.assert_called_once_with(mock_current_user, mock_dataset.tenant_id, session=session) mock_processor = MagicMock() mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor @@ -1038,26 +1094,15 @@ class TestIndexingRunnerRun: ): runner = IndexingRunner() dataset_document = sample_dataset_documents[0] - dataset = Dataset( - id=dataset_document.dataset_id, - tenant_id=dataset_document.tenant_id, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - current_user = Account(name="Test Account", email="test@example.com") + session = mock_dependencies["session"] + dataset = session.get(Dataset, dataset_document.dataset_id) + current_user = session.get(Account, dataset_document.created_by) + assert dataset is not None + assert current_user is not None transformed_documents = [ Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}), Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}), ] - model_dispatch = { - DatasetDocument: dataset_document, - Dataset: dataset, - Account: current_user, - } - mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model) - process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = process_rule with ( patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]), @@ -1093,22 +1138,11 @@ class TestIndexingRunnerRun: runner = IndexingRunner(enforce_vector_space_admission=True) dataset_document = sample_dataset_documents[0] dataset_document.need_summary = False - dataset = Dataset( - id=dataset_document.dataset_id, - tenant_id=dataset_document.tenant_id, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - current_user = Account(name="Test Account", email="test@example.com") - model_dispatch = { - DatasetDocument: dataset_document, - Dataset: dataset, - Account: current_user, - } - mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model) - process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = process_rule + session = mock_dependencies["session"] + dataset = session.get(Dataset, dataset_document.dataset_id) + current_user = session.get(Account, dataset_document.created_by) + assert dataset is not None + assert current_user is not None transformed_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1", "doc_hash": "h1"})] admission_error = VectorSpaceAdmissionError("estimated storage exceeds capacity") admission_service = Mock() @@ -1151,28 +1185,15 @@ class TestIndexingRunnerRun: ): runner = IndexingRunner() dataset_document = sample_dataset_documents[0] - dataset_document.created_by = "user-1" - dataset = Dataset( - id=dataset_document.dataset_id, - tenant_id=dataset_document.tenant_id, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - current_user = Account(name="Test Account", email="test@example.com") + session = mock_dependencies["session"] + dataset = session.get(Dataset, dataset_document.dataset_id) + current_user = session.get(Account, dataset_document.created_by) + assert dataset is not None + assert current_user is not None transformed_documents = [ Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}), Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}), ] - model_dispatch = { - DatasetDocument: dataset_document, - Dataset: dataset, - Account: current_user, - } - mock_dependencies["session"].get.side_effect = lambda model, _: model_dispatch.get(model) - mock_dependencies["session"].scalars.return_value.all.return_value = [] - process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = process_rule with ( patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]), @@ -1207,19 +1228,6 @@ class TestIndexingRunnerRun: runner = IndexingRunner() doc = sample_dataset_documents[0] - # Mock database - mock_dataset = Dataset( - tenant_id=doc.tenant_id, - ) - - get_dispatch = {"Document": doc, "Dataset": mock_dataset} - mock_dependencies["session"].get.side_effect = lambda model, id: get_dispatch.get(model.__name__) - - mock_process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = mock_process_rule - mock_processor = MagicMock() mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor mock_processor.extract.side_effect = ProviderTokenNotInitError("Token not initialized") @@ -1228,9 +1236,8 @@ class TestIndexingRunnerRun: with patch.object(runner, "_extract", side_effect=ProviderTokenNotInitError("Token not initialized")): runner.run([doc], mock_dependencies["session"]) - # Assert - # Verify document status was updated to error - assert mock_dependencies["session"].flush.called + mock_dependencies["session"].refresh(doc) + assert doc.indexing_status == "error" def test_run_handles_object_deleted_error(self, mock_dependencies, sample_dataset_documents): """Test run handles ObjectDeletedError gracefully.""" @@ -1238,19 +1245,6 @@ class TestIndexingRunnerRun: runner = IndexingRunner() doc = sample_dataset_documents[0] - # Mock database - mock_dataset = Dataset( - tenant_id=doc.tenant_id, - ) - - get_dispatch = {"Document": doc, "Dataset": mock_dataset} - mock_dependencies["session"].get.side_effect = lambda model, id: get_dispatch.get(model.__name__) - - mock_process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = mock_process_rule - mock_processor = MagicMock() mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor @@ -1268,28 +1262,11 @@ class TestIndexingRunnerRun: # Arrange runner = IndexingRunner() docs = sample_dataset_documents - - # Mock database - mock_dataset = Dataset( - indexing_technique=IndexTechniqueType.ECONOMY, - ) - mock_current_user = Account(name="Test Account", email="test@example.com") - - doc_map = {doc.id: doc for doc in docs} - model_dispatch = {"Dataset": mock_dataset, "Account": mock_current_user} - - def get_side_effect(model_class, id): - name = model_class.__name__ - if name == "Document": - return doc_map.get(id) - return model_dispatch.get(name) - - mock_dependencies["session"].get.side_effect = get_side_effect - - mock_process_rule = DatasetProcessRule( - dataset_id="dataset-id", mode="automatic", rules="{}", created_by="account-id" - ) - mock_dependencies["session"].scalar.return_value = mock_process_rule + for doc in docs: + dataset = mock_dependencies["session"].get(Dataset, doc.dataset_id) + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.ECONOMY + mock_dependencies["session"].commit() mock_processor = MagicMock() mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor @@ -1329,13 +1306,13 @@ class TestIndexingRunnerRetryLogic: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, sqlite_session: Session): """Mock all external dependencies.""" with ( patch("core.indexing_runner.redis_client") as mock_redis, ): yield { - "session": MagicMock(), + "session": sqlite_session, "redis": mock_redis, } @@ -1360,41 +1337,62 @@ class TestIndexingRunnerRetryLogic: def test_update_document_index_status_success(self, mock_dependencies): """Test successful document status update.""" - # Arrange document_id = str(uuid.uuid4()) - mock_document = DatasetDocument(id=document_id) - - mock_dependencies["session"].scalar.return_value = 0 - mock_dependencies["session"].get.return_value = mock_document + session = mock_dependencies["session"] + document = DatasetDocument( + id=document_id, + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by=str(uuid.uuid4()), + is_paused=False, + ) + session.add(document) + session.commit() # Act IndexingRunner._update_document_index_status( document_id, "completed", {"tokens": 100, "completed_at": naive_utc_now()}, - session=mock_dependencies["session"], + session=session, ) - # Assert - mock_dependencies["session"].flush.assert_called() + session.refresh(document) + assert document.indexing_status == "completed" + assert document.tokens == 100 def test_update_document_index_status_paused(self, mock_dependencies): """Test document status update when document is paused.""" - # Arrange document_id = str(uuid.uuid4()) - mock_dependencies["session"].scalar.return_value = 1 + session = mock_dependencies["session"] + document = DatasetDocument( + id=document_id, + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch", + name="Paused document", + created_from=DocumentCreatedFrom.API, + created_by=str(uuid.uuid4()), + is_paused=True, + ) + session.add(document) + session.commit() # Act & Assert with pytest.raises(DocumentIsPausedError): - IndexingRunner._update_document_index_status(document_id, "completed", session=mock_dependencies["session"]) + IndexingRunner._update_document_index_status(document_id, "completed", session=session) def test_update_document_index_status_deleted(self, mock_dependencies): """Test document status update when document is deleted.""" # Arrange document_id = str(uuid.uuid4()) - mock_dependencies["session"].scalar.return_value = 0 - mock_dependencies["session"].get.return_value = None - # Act & Assert with pytest.raises(DocumentIsDeletedPausedError): IndexingRunner._update_document_index_status(document_id, "completed", session=mock_dependencies["session"]) @@ -1610,13 +1608,13 @@ class TestIndexingRunnerLoadSegments: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, unbound_session: Session): """Mock all external dependencies.""" with ( patch("core.indexing_runner.DatasetDocumentStore") as mock_docstore, ): yield { - "session": MagicMock(), + "session": unbound_session, "docstore": mock_docstore, } @@ -1774,13 +1772,13 @@ class TestIndexingRunnerEstimate: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, sqlite_session: Session): """Mock all external dependencies.""" with ( patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, ): yield { - "session": MagicMock(), + "session": sqlite_session, "factory": mock_factory, } @@ -1813,7 +1811,8 @@ class TestIndexingRunnerEstimate: mock_processor = MagicMock() mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor phase_events: list[str] = [] - mock_dependencies["session"].commit.side_effect = lambda: phase_events.append("commit") + session = mock_dependencies["session"] + event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) preview_doc = Document( page_content="![image](http://files.local/files/image-1/file-preview)", @@ -1825,8 +1824,23 @@ class TestIndexingRunnerEstimate: phase_events.append("summary") or [PreviewDetail(content=preview_doc.page_content)] ) - image_file = SimpleNamespace(key="image_files/tenant-1/source-file-1/image.png") - mock_dependencies["session"].scalar.return_value = image_file + image_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key="image_files/tenant-1/source-file-1/image.png", + name="image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-id", + created_at=datetime.now(UTC), + used=True, + ) + image_file.id = "image-1" + session.add(image_file) + session.commit() + phase_events.clear() with ( patch("core.indexing_runner.get_image_upload_file_ids", return_value=["image-1"]), @@ -1844,12 +1858,12 @@ class TestIndexingRunnerEstimate: "summary_index_setting": {"enable": True}, }, doc_form=IndexStructureType.PARAGRAPH_INDEX, - session=mock_dependencies["session"], + session=session, ) assert result.total_segments == 1 mock_storage.delete.assert_called_once_with(image_file.key) - mock_dependencies["session"].delete.assert_called_once_with(image_file) + assert session.get(UploadFile, image_file.id) is None assert phase_events == ["commit", "summary"] @@ -1863,13 +1877,13 @@ class TestIndexingRunnerProcessChunk: """ @pytest.fixture - def mock_dependencies(self): + def mock_dependencies(self, sqlite_session: Session): """Mock all external dependencies.""" with ( patch("core.indexing_runner.redis_client") as mock_redis, ): yield { - "session": MagicMock(), + "session": sqlite_session, "redis": mock_redis, } @@ -1881,7 +1895,12 @@ class TestIndexingRunnerProcessChunk: app.app_context.return_value.__exit__ = MagicMock() return app - def test_process_chunk_loads_index_and_completes_segments(self, mock_dependencies, mock_flask_app): + def test_process_chunk_loads_index_and_completes_segments( + self, + mock_dependencies, + mock_flask_app, + sqlite_session_factory: sessionmaker[Session], + ): """Test process chunk loads the index and completes segments without counting tokens.""" # Arrange from core.indexing_runner import IndexingRunner @@ -1893,33 +1912,36 @@ class TestIndexingRunnerProcessChunk: Document(page_content="Chunk 2", metadata={"doc_id": "c2"}), ] - mock_dataset = Dataset( - id=str(uuid.uuid4()), - ) - - mock_dataset_document = DatasetDocument(id=str(uuid.uuid4())) + session = mock_dependencies["session"] + mock_dataset, mock_dataset_document = persist_indexing_scope(session) + segments = [ + DocumentSegment( + tenant_id=mock_dataset.tenant_id, + dataset_id=mock_dataset.id, + document_id=mock_dataset_document.id, + position=position, + content=f"Chunk {position}", + word_count=1, + tokens=1, + created_by=mock_dataset.created_by, + index_node_id=doc.metadata["doc_id"], + status=SegmentStatus.INDEXING, + ) + for position, doc in enumerate(chunk_documents, start=1) + ] + session.add_all(segments) + session.commit() mock_dependencies["redis"].get.return_value = None - # Mock database update for segment status - mock_dependencies["session"].execute.return_value = None - mock_dependencies["session"].get.side_effect = lambda model, _id: { - Dataset: mock_dataset, - DatasetDocument: mock_dataset_document, - }.get(model) - # Create a proper context manager mock mock_context = MagicMock() mock_context.__enter__ = MagicMock(return_value=None) mock_context.__exit__ = MagicMock(return_value=None) mock_flask_app.app_context.return_value = mock_context - session_context = MagicMock() - session_context.__enter__.return_value = mock_dependencies["session"] - session_context.__exit__.return_value = None - with ( - patch("core.indexing_runner.session_factory.create_session", return_value=session_context), + patch("core.indexing_runner.session_factory.create_session", sqlite_session_factory), patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, ): mock_factory.return_value.init_index_processor.return_value = mock_processor @@ -1936,10 +1958,15 @@ class TestIndexingRunnerProcessChunk: # Assert assert result is None mock_processor.load.assert_called_once() - mock_dependencies["session"].execute.assert_called_once() - mock_dependencies["session"].commit.assert_called_once() + session.expire_all() + assert all(session.get(DocumentSegment, segment.id).status == SegmentStatus.COMPLETED for segment in segments) - def test_process_chunk_detects_pause(self, mock_dependencies, mock_flask_app): + def test_process_chunk_detects_pause( + self, + mock_dependencies, + mock_flask_app, + sqlite_session_factory: sessionmaker[Session], + ): """Test process chunk detects document pause.""" # Arrange from core.indexing_runner import IndexingRunner @@ -1947,27 +1974,17 @@ class TestIndexingRunnerProcessChunk: runner = IndexingRunner() chunk_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1"})] - mock_dataset = Dataset() - mock_dataset_document = DatasetDocument(id=str(uuid.uuid4())) + mock_dataset, mock_dataset_document = persist_indexing_scope(mock_dependencies["session"]) # Mock Redis to return paused status mock_dependencies["redis"].get.return_value = "1" - mock_dependencies["session"].get.side_effect = lambda model, _id: { - Dataset: mock_dataset, - DatasetDocument: mock_dataset_document, - }.get(model) - # Create a proper context manager mock mock_context = MagicMock() mock_context.__enter__ = MagicMock(return_value=None) mock_context.__exit__ = MagicMock(return_value=None) mock_flask_app.app_context.return_value = mock_context - session_context = MagicMock() - session_context.__enter__.return_value = mock_dependencies["session"] - session_context.__exit__.return_value = None - - with patch("core.indexing_runner.session_factory.create_session", return_value=session_context): + with patch("core.indexing_runner.session_factory.create_session", sqlite_session_factory): # Act & Assert - the method creates its own app_context and session with pytest.raises(DocumentIsPausedError): runner._process_chunk( diff --git a/api/tests/unit_tests/core/tools/test_signature.py b/api/tests/unit_tests/core/tools/test_signature.py index 142b82902bf..4985fc16368 100644 --- a/api/tests/unit_tests/core/tools/test_signature.py +++ b/api/tests/unit_tests/core/tools/test_signature.py @@ -2,6 +2,7 @@ from __future__ import annotations +from typing import Literal from urllib.parse import parse_qs, urlparse import pytest @@ -172,6 +173,7 @@ def test_get_signed_file_uri_for_plugin_and_verify_roundtrip(monkeypatch: pytest assert query["tenant_id"] == ["tenant-id"] assert query["user_id"] == ["user-id"] assert query["conversation_id"] == ["conversation-id"] + assert "max_size" not in query assert ( verify_plugin_file_signature( filename="report.pdf", @@ -187,6 +189,50 @@ def test_get_signed_file_uri_for_plugin_and_verify_roundtrip(monkeypatch: pytest ) +@pytest.mark.parametrize( + ("user_from", "forged_nonce_suffix"), + [ + (None, "|1024"), + ("account", "|account|1024"), + ], +) +def test_plugin_upload_signature_binds_max_size_without_legacy_payload_ambiguity( + monkeypatch: pytest.MonkeyPatch, + user_from: Literal["account", "end-user"] | None, + forged_nonce_suffix: str, +) -> None: + monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) + monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x0a" * 16) + monkeypatch.setattr("core.tools.signature.dify_config.SECRET_KEY", "unit-secret") + monkeypatch.setattr("core.tools.signature.dify_config.FILES_ACCESS_TIMEOUT", 60) + + uri = get_signed_file_uri_for_plugin( + filename="report.pdf", + mimetype="application/pdf", + tenant_id="tenant-id", + user_id="user-id", + user_from=user_from, + max_size=1024, + ) + query = parse_qs(urlparse(uri).query) + signed = { + "filename": "report.pdf", + "mimetype": "application/pdf", + "tenant_id": "tenant-id", + "user_id": "user-id", + "timestamp": query["timestamp"][0], + "nonce": query["nonce"][0], + "sign": query["sign"][0], + } + + assert query["max_size"] == ["1024"] + assert verify_plugin_file_signature(**signed, user_from=user_from, max_size=1024) is True + assert verify_plugin_file_signature(**signed, user_from=user_from, max_size=2048) is False + assert verify_plugin_file_signature(**signed, user_from=user_from) is False + forged = {**signed, "nonce": f"{signed['nonce']}{forged_nonce_suffix}"} + assert verify_plugin_file_signature(**forged) is False + + def test_plugin_upload_signature_binds_account_user_from(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("core.tools.signature.time.time", lambda: 1700000000) monkeypatch.setattr("core.tools.signature.os.urandom", lambda _: b"\x09" * 16) 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 915cdaa8ec1..a8c254ae604 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 @@ -12,7 +12,11 @@ from dify_agent.protocol import ( CancelRunRequest, CancelRunResponse, PydanticAIStreamRunEvent, + RunCancelledEvent, + RunCancelledEventData, RunEvent, + RunFailedEvent, + RunFailedEventData, RunStartedEvent, RunSucceededEvent, RunSucceededEventData, @@ -21,6 +25,7 @@ from pydantic_ai.messages import PartDeltaEvent, TextPartDelta from clients.agent_backend import ( AgentBackendInternalEventType, + AgentBackendRunCancelledInternalEvent, AgentBackendRunEventAdapter, AgentBackendStreamError, AgentBackendStreamInternalEvent, @@ -62,6 +67,14 @@ from models.agent_config_entities import ( from services.agent.workspace_service import AgentWorkspaceNotFoundError +@pytest.fixture(autouse=True) +def _stub_model_context_window(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.runtime_request_builder.resolve_model_context_window", + lambda **_kwargs: None, + ) + + def _restored_file(*, transfer_method: FileTransferMethod, reference: str) -> File: return File( type=FileType.DOCUMENT, @@ -197,6 +210,25 @@ class FakeSessionStore: self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id)) +class ExplodingSessionStore(FakeSessionStore): + def __init__(self, snapshot: CompositorSessionSnapshot | None = None) -> None: + super().__init__(snapshot=snapshot) + self.save_attempts: list[CompositorSessionSnapshot | None] = [] + + def save_active_snapshot( + self, + *, + scope: WorkflowAgentSessionScope, + binding_id: str, + snapshot: CompositorSessionSnapshot | None, + pending_form_id: str | None = None, + pending_tool_call_id: str | None = None, + ) -> None: + del scope, binding_id, pending_form_id, pending_tool_call_id + self.save_attempts.append(snapshot) + raise RuntimeError("simulated DB failure") + + class FileOutputBackendClient(FakeAgentBackendRunClient): output_payload: dict[str, object] @@ -251,6 +283,7 @@ class FailingStreamBackendClient(FakeAgentBackendRunClient): def __init__(self) -> None: super().__init__() self.cancel_requests: list[CancelRunRequest | None] = [] + self.cancel_after: list[str | None] = [] def stream_events( self, @@ -267,6 +300,50 @@ class FailingStreamBackendClient(FakeAgentBackendRunClient): self.cancel_requests.append(request) return CancelRunResponse(run_id=run_id, status="cancelled") + def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + self.cancel_after.append(after) + return super().cancel_run_and_wait(run_id, request=request, after=after) + + +class FailingAfterStartedStreamBackendClient(FailingStreamBackendClient): + def stream_events( + self, + run_id: str, + *, + after: str | None = None, + should_stop: Callable[[], bool] | None = None, + ) -> Iterator[RunEvent]: + del after, should_stop + yield RunStartedEvent(id="cursor-1", run_id=run_id) + raise AgentBackendStreamError("stream failed after started") + + +class TerminalWithoutSnapshotBackendClient(FakeAgentBackendRunClient): + def __init__(self, *, terminal_type: str) -> None: + super().__init__() + self.terminal_type = terminal_type + + def _events(self, run_id: str): + if self.terminal_type == "failed": + terminal: RunEvent = RunFailedEvent( + id="2-0", + run_id=run_id, + data=RunFailedEventData(error="failed without snapshot"), + ) + else: + terminal = RunCancelledEvent( + id="2-0", + run_id=run_id, + data=RunCancelledEventData(reason="cancelled without snapshot"), + ) + return (RunStartedEvent(id="1-0", run_id=run_id), terminal) + class EmptyStreamBackendClient(FailingStreamBackendClient): def stream_events( @@ -343,7 +420,9 @@ def _node( node = DifyAgentNode( node_id="agent-node", - data=DifyAgentNodeData.model_validate({"type": BuiltinNodeTypes.AGENT, "version": "2"}), + data=DifyAgentNodeData.model_validate( + {"type": BuiltinNodeTypes.AGENT, "version": "2", "agent_node_kind": "dify_agent"} + ), graph_init_params=graph_init_params, graph_runtime_state=cast( GraphRuntimeState, @@ -591,13 +670,15 @@ def test_agent_node_run_normalizes_declared_array_file_output_with_canonical_map def test_agent_node_run_maps_failed_agent_backend_run_to_node_result(): - events = list(_node(scenario=FakeAgentBackendScenario.FAILED)._run()) + store = FakeSessionStore() + 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 result.error == "fake failure" assert result.error_type == "unit_test" + assert store.saved[0][2] == CompositorSessionSnapshot(layers=[]) def test_agent_node_saves_success_snapshot_and_reuses_existing_snapshot(): @@ -626,12 +707,7 @@ def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_ ``session_snapshot_persist_error`` in the agent_backend metadata so the incident is observable from the workflow_node_executions record.""" - class _ExplodingSessionStore(FakeSessionStore): - def save_active_snapshot(self, **kwargs): # type: ignore[override] - del kwargs - raise RuntimeError("simulated DB failure") - - store = _ExplodingSessionStore() + store = ExplodingSessionStore() events = list(_node(session_store=store)._run()) assert len(events) == 1 @@ -642,6 +718,46 @@ def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_ assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_workspace_store_error" +@pytest.mark.parametrize("failure_kind", ["backend", "transport"]) +def test_agent_node_snapshot_save_failure_preserves_original_failure(failure_kind: str) -> None: + store = ExplodingSessionStore() + client = ( + FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.FAILED) + if failure_kind == "backend" + else FailingStreamBackendClient() + ) + + events = list(_node(agent_backend_client=client, session_store=store)._run()) + + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + if failure_kind == "backend": + assert (result.error, result.error_type) == ("fake failure", "unit_test") + else: + assert result.error == "stream reconnect attempts exhausted" + assert result.error_type == "agent_backend_stream_error" + agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"] + assert agent_backend["session_snapshot_persisted"] is False + assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_workspace_store_error" + assert store.save_attempts == [CompositorSessionSnapshot(layers=[])] + + +@pytest.mark.parametrize("terminal_type", ["failed", "cancelled"]) +def test_agent_node_terminal_without_snapshot_preserves_prior_session_without_write(terminal_type: str) -> None: + store = FakeSessionStore() + + events = list( + _node( + agent_backend_client=TerminalWithoutSnapshotBackendClient(terminal_type=terminal_type), + session_store=store, + )._run() + ) + + result = cast(StreamCompletedEvent, events[0]).node_run_result + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert store.saved == [] + + def test_agent_node_paused_run_requests_workflow_pause_and_persists_snapshot(): store = FakeSessionStore() node = _node(scenario=FakeAgentBackendScenario.PAUSED, session_store=store) @@ -809,7 +925,7 @@ def test_agent_node_cancels_backend_run_when_stream_fails(): metadata={"agent_backend": {}}, ) - assert terminal is None + assert isinstance(terminal, AgentBackendRunCancelledInternalEvent) assert failure is not None assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} assert len(client.cancel_requests) == 1 @@ -817,6 +933,23 @@ def test_agent_node_cancels_backend_run_when_stream_fails(): assert client.cancel_requests[0].reason == "event_stream_failed" +def test_agent_node_forwards_last_stream_cursor_when_cancelling_after_failure() -> None: + client = FailingAfterStartedStreamBackendClient() + node = _node(agent_backend_client=client) + + terminal, failure = node._consume_event_stream( + "run-1", + inputs={}, + process_data={"workflow_agent_binding_id": "binding-1"}, + metadata={"agent_backend": {}}, + ) + + assert isinstance(terminal, AgentBackendRunCancelledInternalEvent) + assert failure is not None + assert failure.node_run_result.error == "stream failed after started" + assert client.cancel_after == ["cursor-1"] + + def test_agent_node_cancels_backend_run_when_stream_ends_without_terminal_event(): client = EmptyStreamBackendClient() node = _node(agent_backend_client=client) @@ -828,7 +961,7 @@ def test_agent_node_cancels_backend_run_when_stream_ends_without_terminal_event( metadata={"agent_backend": {}}, ) - assert terminal is None + assert isinstance(terminal, AgentBackendRunCancelledInternalEvent) assert failure is None assert client.cancel_requests[0] is not None assert client.cancel_requests[0].reason == "stream_ended_without_terminal_event" @@ -845,7 +978,7 @@ def test_agent_node_cancels_backend_run_when_stream_raises_unexpected_error(): metadata={"agent_backend": {}}, ) - assert terminal is None + assert isinstance(terminal, AgentBackendRunCancelledInternalEvent) assert failure is not None assert failure.node_run_result.error == "unexpected stream failure" assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} @@ -867,16 +1000,18 @@ def test_agent_node_uses_graph_abort_reason_when_cancel_request_fails(caplog): assert terminal is None assert failure is not None + assert failure.node_run_result.error == "stream reconnect attempts exhausted" + assert failure.node_run_result.error_type == "agent_backend_stream_error" assert client.cancel_requests[0] is not None assert client.cancel_requests[0].reason == "workflow_graph_aborted" - assert "Failed to cancel Workflow Agent backend run" in caplog.text + assert "Failed to finish cancelling Workflow Agent backend run" in caplog.text def test_agent_node_cancels_backend_run_for_unexpected_internal_event(): client = FakeAgentBackendRunClient() node = _node(agent_backend_client=client) - node._agent_backend_client.cancel_run = MagicMock( # type: ignore[method-assign] - return_value=CancelRunResponse(run_id="run-1", status="cancelled") + node._agent_backend_client.cancel_run_and_wait = MagicMock( # type: ignore[method-assign] + return_value=RunCancelledEvent(run_id="run-1") ) node._event_adapter.adapt = MagicMock( # type: ignore[method-assign] return_value=[SimpleNamespace(type=AgentBackendInternalEventType.RUN_FAILED)] @@ -895,7 +1030,7 @@ def test_agent_node_cancels_backend_run_for_unexpected_internal_event(): "Unexpected internal event type " ) assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"} - node._agent_backend_client.cancel_run.assert_called_once() + node._agent_backend_client.cancel_run_and_wait.assert_called_once() def test_agent_node_records_stream_usage_metadata(): diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py new file mode 100644 index 00000000000..1476764d37f --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py @@ -0,0 +1,31 @@ +import pytest +from pydantic import ValidationError + +from core.workflow.nodes.agent_v2.discriminator import is_dify_agent_node_data +from core.workflow.nodes.agent_v2.entities import DifyAgentNodeData +from graphon.entities.base_node_data import BaseNodeData + + +@pytest.mark.parametrize( + ("node_data", "expected"), + [ + ({"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, True), + ({"type": "agent", "version": 2, "agent_node_kind": "dify_agent"}, True), + ({"type": "agent", "version": "2"}, False), + ({"type": "agent", "version": "1", "agent_node_kind": "dify_agent"}, False), + ({"type": "llm", "version": "2", "agent_node_kind": "dify_agent"}, False), + ], +) +def test_is_dify_agent_node_data_mapping(node_data: dict[str, object], expected: bool) -> None: + assert is_dify_agent_node_data(node_data) is expected + + +def test_is_dify_agent_node_data_supports_base_node_data() -> None: + node_data = BaseNodeData.model_validate({"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}) + + assert is_dify_agent_node_data(node_data) is True + + +def test_dify_agent_node_data_requires_explicit_kind_marker() -> None: + with pytest.raises(ValidationError, match="agent_node_kind"): + DifyAgentNodeData.model_validate({"type": "agent", "version": "2"}) 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 ed30234537a..79301604f53 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 @@ -40,6 +40,21 @@ from models.agent_config_entities import ( ) +@pytest.fixture(autouse=True) +def model_context_window_calls(monkeypatch: pytest.MonkeyPatch) -> list[tuple[object, str, str]]: + calls: list[tuple[object, str, str]] = [] + + def resolve(*, run_context: object, provider_name: str, model_name: str) -> int: + calls.append((run_context, provider_name, model_name)) + return 32_768 + + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.runtime_request_builder.resolve_model_context_window", + resolve, + ) + return calls + + def test_agent_soul_round_trip_preserves_existing_app_feature_fields(): config = AgentSoulConfig.model_validate( { @@ -147,7 +162,7 @@ def _context() -> WorkflowAgentRuntimeBuildContext: prompt={"system_prompt": "You are careful."}, model=AgentSoulModelConfig( plugin_id="langgenius/openai", - model_provider="openai", + model_provider="langgenius/openai/openai", model="gpt-test", model_settings={"temperature": 0}, ), @@ -221,8 +236,11 @@ def _uploaded_workflow_files_prompt_payload(result) -> object: raise AssertionError("missing prompt payload for sys.files") -def test_builds_create_run_request_from_agent_soul_and_node_job(): - result = WorkflowAgentRuntimeRequestBuilder().build(_context()) +def test_builds_create_run_request_from_agent_soul_and_node_job( + model_context_window_calls: list[tuple[object, str, str]], +): + context = _context() + result = WorkflowAgentRuntimeRequestBuilder().build(context) dumped = result.request.model_dump(mode="json") layers = {layer["name"]: layer for layer in dumped["composition"]["layers"]} @@ -239,6 +257,9 @@ def test_builds_create_run_request_from_agent_soul_and_node_job(): assert "Previous node outputs:" not in dumped["composition"]["layers"][2]["config"]["user"] assert dumped["composition"]["layers"][-1]["config"]["json_schema"]["properties"]["summary"]["type"] == "string" assert DIFY_AGENT_HISTORY_LAYER_ID in layers + assert layers[DIFY_AGENT_MODEL_LAYER_ID]["config"]["model_provider"] == "openai" + assert layers[DIFY_AGENT_MODEL_LAYER_ID]["config"]["context_window_tokens"] == 32_768 + assert model_context_window_calls == [(context.dify_context, "langgenius/openai/openai", "gpt-test")] redacted_layers = {layer["name"]: layer for layer in result.redacted_request["composition"]["layers"]} assert "credentials" not in redacted_layers[DIFY_AGENT_MODEL_LAYER_ID]["config"] diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py index acf27ae2b03..ce4107d41d1 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py @@ -87,7 +87,10 @@ def _graph(edges: list[dict]) -> dict: "nodes": [ {"id": "start", "data": {"type": "start"}}, {"id": "previous-node", "data": {"type": "llm"}}, - {"id": "agent-node", "data": {"type": "agent", "version": "2"}}, + { + "id": "agent-node", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + }, {"id": "later-node", "data": {"type": "llm"}}, ], "edges": edges, @@ -118,6 +121,18 @@ def _tool_graph(tool_data: dict) -> dict: } +def test_historical_agent_version_two_is_not_validated_as_dify_agent() -> None: + graph = { + "nodes": [{"id": "legacy-agent", "data": {"type": "agent", "version": "2"}}], + "edges": [], + } + session = Mock() + + WorkflowAgentNodeValidator.validate_published_workflow(session=session, workflow=_workflow(graph)) + + session.scalar.assert_not_called() + + def test_publish_validation_accepts_upstream_previous_output_ref(): node_job = WorkflowNodeJobConfig.model_validate( {"previous_node_output_refs": [{"node_id": "previous-node", "output": "text"}]} diff --git a/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py b/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py index c7ff7e5a340..c512eb04637 100644 --- a/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py +++ b/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py @@ -23,7 +23,12 @@ def test_moved_core_nodes_resolve_after_importing_production_entrypoints(): from core.app.apps import workflow_app_runner from core.workflow import workflow_entry from core.workflow.nodes.knowledge_index import KNOWLEDGE_INDEX_NODE_TYPE - from core.workflow.node_factory import DifyNodeFactory, NODE_TYPE_CLASSES_MAPPING + from core.workflow.node_factory import ( + DifyNodeFactory, + NODE_TYPE_CLASSES_MAPPING, + resolve_workflow_node_class, + ) + from core.workflow.nodes.agent import AgentNode from core.workflow.nodes.agent_v2 import DifyAgentNode from graphon.enums import BuiltinNodeTypes from services import workflow_service @@ -46,6 +51,16 @@ def test_moved_core_nodes_resolve_after_importing_production_entrypoints(): node_type=BuiltinNodeTypes.AGENT, node_version="2", ) is DifyAgentNode + assert resolve_workflow_node_class( + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_data={"type": "agent", "version": "2"}, + ) is AgentNode + assert resolve_workflow_node_class( + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_data={"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + ) is DifyAgentNode """ ) completed = subprocess.run( diff --git a/api/tests/unit_tests/libs/test_workspace_permission.py b/api/tests/unit_tests/libs/test_workspace_permission.py index 2d3523e1fad..9afbcffef9e 100644 --- a/api/tests/unit_tests/libs/test_workspace_permission.py +++ b/api/tests/unit_tests/libs/test_workspace_permission.py @@ -14,11 +14,10 @@ from libs.workspace_permission import ( class TestWorkspacePermissionHelper: """Test workspace permission helper functions.""" - @patch("libs.workspace_permission.dify_config") @patch("libs.workspace_permission.EnterpriseService") - def test_community_edition_allows_invite(self, mock_enterprise_service, mock_config): + def test_community_edition_allows_invite(self, mock_enterprise_service, config_overrides): """Community edition should always allow invitations without calling any service.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) # Should not raise check_workspace_member_invite_permission("test-workspace-id") @@ -26,11 +25,10 @@ class TestWorkspacePermissionHelper: # EnterpriseService should NOT be called in community edition mock_enterprise_service.WorkspacePermissionService.get_permission.assert_not_called() - @patch("libs.workspace_permission.dify_config") @patch("libs.workspace_permission.FeatureService") - def test_community_edition_allows_transfer(self, mock_feature_service, mock_config): + def test_community_edition_allows_transfer(self, mock_feature_service, config_overrides): """Community edition should check billing plan but not call enterprise service.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) mock_features = Mock() mock_features.is_allow_transfer_workspace = True mock_feature_service.get_features.return_value = mock_features @@ -41,10 +39,9 @@ class TestWorkspacePermissionHelper: mock_feature_service.get_features.assert_called_once_with("test-workspace-id", exclude_vector_space=True) @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") - def test_enterprise_blocks_invite_when_disabled(self, mock_config, mock_enterprise_service): + def test_enterprise_blocks_invite_when_disabled(self, mock_enterprise_service, config_overrides): """Enterprise edition should block invitations when workspace policy is False.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) mock_permission = Mock() mock_permission.allow_member_invite = False @@ -56,10 +53,9 @@ class TestWorkspacePermissionHelper: mock_enterprise_service.WorkspacePermissionService.get_permission.assert_called_once_with("test-workspace-id") @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") - def test_enterprise_allows_invite_when_enabled(self, mock_config, mock_enterprise_service): + def test_enterprise_allows_invite_when_enabled(self, mock_enterprise_service, config_overrides): """Enterprise edition should allow invitations when workspace policy is True.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) mock_permission = Mock() mock_permission.allow_member_invite = True @@ -71,11 +67,10 @@ class TestWorkspacePermissionHelper: mock_enterprise_service.WorkspacePermissionService.get_permission.assert_called_once_with("test-workspace-id") @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") @patch("libs.workspace_permission.FeatureService") - def test_billing_plan_blocks_transfer(self, mock_feature_service, mock_config, mock_enterprise_service): + def test_billing_plan_blocks_transfer(self, mock_feature_service, mock_enterprise_service, config_overrides): """SANDBOX billing plan should block owner transfer before checking enterprise policy.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) mock_features = Mock() mock_features.is_allow_transfer_workspace = False # SANDBOX plan mock_feature_service.get_features.return_value = mock_features @@ -87,11 +82,12 @@ class TestWorkspacePermissionHelper: mock_enterprise_service.WorkspacePermissionService.get_permission.assert_not_called() @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") @patch("libs.workspace_permission.FeatureService") - def test_enterprise_blocks_transfer_when_disabled(self, mock_feature_service, mock_config, mock_enterprise_service): + def test_enterprise_blocks_transfer_when_disabled( + self, mock_feature_service, mock_enterprise_service, config_overrides + ): """Enterprise edition should block transfer when workspace policy is False.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) mock_features = Mock() mock_features.is_allow_transfer_workspace = True # Billing plan allows mock_feature_service.get_features.return_value = mock_features @@ -106,13 +102,12 @@ class TestWorkspacePermissionHelper: mock_enterprise_service.WorkspacePermissionService.get_permission.assert_called_once_with("test-workspace-id") @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") @patch("libs.workspace_permission.FeatureService") def test_enterprise_allows_transfer_when_both_enabled( - self, mock_feature_service, mock_config, mock_enterprise_service + self, mock_feature_service, mock_enterprise_service, config_overrides ): """Enterprise edition should allow transfer when both billing and workspace policy allow.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) mock_features = Mock() mock_features.is_allow_transfer_workspace = True # Billing plan allows mock_feature_service.get_features.return_value = mock_features @@ -127,12 +122,11 @@ class TestWorkspacePermissionHelper: mock_enterprise_service.WorkspacePermissionService.get_permission.assert_called_once_with("test-workspace-id") @patch("libs.workspace_permission.EnterpriseService") - @patch("libs.workspace_permission.dify_config") def test_enterprise_service_error_fails_open( - self, mock_config, mock_enterprise_service, caplog: pytest.LogCaptureFixture + self, mock_enterprise_service, config_overrides, caplog: pytest.LogCaptureFixture ): """On enterprise service error, should fail-open (allow) and log error.""" - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.ENTERPRISE + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) # Simulate enterprise service error mock_enterprise_service.WorkspacePermissionService.get_permission.side_effect = Exception("Service unavailable") diff --git a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py index 4ebf30c25ad..fa188ecb8f8 100644 --- a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py +++ b/api/tests/unit_tests/services/agent/test_agent_dsl_service.py @@ -60,7 +60,7 @@ def _snapshot(*, snapshot_id: str = "snapshot-1", soul: AgentSoulConfig | None = def _agent_node(node_id: str, binding: object | None = None) -> dict: - data = {"type": BuiltinNodeTypes.AGENT, "version": "2"} + data = {"type": BuiltinNodeTypes.AGENT, "version": "2", "agent_node_kind": "dify_agent"} if binding is not None: data["agent_binding"] = binding return {"id": node_id, "data": data} @@ -734,4 +734,17 @@ def test_require_helpers_and_graph_detection() -> None: assert AgentDslService._agent_icon_type(AgentIconType.EMOJI.value) == AgentIconType.EMOJI assert AgentDslService._agent_icon_type(None) is None assert is_agent_v2_graph({"nodes": [_agent_node("agent")]}) is True + assert is_agent_v2_graph({"nodes": [{"id": "legacy-agent", "data": {"type": "agent", "version": "2"}}]}) is False assert is_agent_v2_graph({"nodes": ["invalid", {"data": {"type": "start"}}]}) is False + + +def test_export_workflow_packages_ignores_historical_agent_version_two() -> None: + session = Mock() + service = AgentDslService(session) + graph = {"nodes": [{"id": "legacy-agent", "data": {"type": "agent", "version": "2"}}]} + + portable_graph, packages = service.export_workflow_packages(workflow=Mock(), graph=graph) + + assert portable_graph == graph + assert packages == {} + session.scalars.assert_not_called() 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 479324c3167..5dbce1c2549 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -930,6 +930,7 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. session.commit() created: dict[str, object] = {} calls: list[str] = [] + register_publish_event = MagicMock() monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda payload: None) monkeypatch.setattr(composer_service, "agent_has_workflow_callable_active_snapshot", lambda **_kwargs: False) @@ -945,6 +946,7 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. lambda **kwargs: calls.append("create_version") or created.update(kwargs) or version, ) monkeypatch.setattr(AgentComposerService, "_serialize_version", lambda _version: {"id": _version.id}) + monkeypatch.setattr(composer_service, "register_new_agent_beta_publish_after_commit", register_publish_event) result = AgentComposerService.publish_agent_app_draft( session=session, @@ -967,6 +969,12 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest. assert app.enable_site is True assert app.enable_api is True assert app.updated_by == "account-1" + register_publish_event.assert_called_once_with( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_id="version-2", + ) def test_repeated_publish_reuses_normal_draft_home_without_creating_resources( @@ -5671,7 +5679,12 @@ class TestWorkflowAgentDraftBindingSync: workflow = _workflow() workflow.graph = json.dumps( { - "nodes": [{"id": "agent-node", "data": {"type": "agent", "version": "2"}}], + "nodes": [ + { + "id": "agent-node", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + } + ], "edges": [], } ) @@ -5739,6 +5752,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": agent_task, "agent_binding": { "binding_type": "roster_agent", @@ -5882,6 +5896,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "roster_agent", "agent_id": "agent-1", @@ -5963,6 +5978,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", }, @@ -6017,6 +6033,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", }, @@ -6066,6 +6083,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": "Summarize the upstream result.", "agent_declared_outputs": [ { @@ -6165,6 +6183,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": "Use the current node context.", "agent_binding": { "binding_type": "inline_agent", @@ -6232,6 +6251,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", }, @@ -6280,6 +6300,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", "agent_id": "inline-agent-1", @@ -6338,6 +6359,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "unknown", "agent_id": "agent-1", @@ -6371,6 +6393,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", "agent_id": "inline-agent-1", @@ -6405,6 +6428,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", "agent_id": "inline-agent-1", @@ -6454,6 +6478,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": "Use the latest tender context.", "agent_binding": { "binding_type": "roster_agent", @@ -6521,6 +6546,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_task": "Keep the prompt.", "agent_declared_outputs": [], "agent_binding": { @@ -6601,6 +6627,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", "agent_id": "inline-kept", @@ -6613,6 +6640,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "roster_agent", "agent_id": "roster-new", @@ -6624,6 +6652,7 @@ class TestWorkflowAgentDraftBindingSync: "data": { "type": "agent", "version": "2", + "agent_node_kind": "dify_agent", "agent_binding": { "binding_type": "inline_agent", "agent_id": "inline-new", diff --git a/api/tests/unit_tests/services/agent/test_composer_candidates.py b/api/tests/unit_tests/services/agent/test_composer_candidates.py index d1edf2cfa60..990302e9f91 100644 --- a/api/tests/unit_tests/services/agent/test_composer_candidates.py +++ b/api/tests/unit_tests/services/agent/test_composer_candidates.py @@ -2,6 +2,8 @@ from __future__ import annotations +from unittest.mock import Mock + from fields.agent_fields import AgentComposerCandidatesResponse from models.agent_config_entities import AgentSoulConfig, DeclaredOutputConfig, DeclaredOutputType from models.dataset import Dataset @@ -22,8 +24,24 @@ _GRAPH = { }, }, {"id": "llm-1", "data": {"type": "llm", "title": "LLM"}}, - {"id": "agent-up", "data": {"type": "agent", "version": "2", "title": "Upstream Agent"}}, - {"id": "agent-target", "data": {"type": "agent", "version": "2", "title": "Target Agent"}}, + { + "id": "agent-up", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + "title": "Upstream Agent", + }, + }, + { + "id": "agent-target", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + "title": "Target Agent", + }, + }, {"id": "end", "data": {"type": "end", "title": "END"}}, ], "edges": [ @@ -97,6 +115,31 @@ def test_results_differ_per_node_id(): assert {e["node_id"] for e in entries_llm} == {"start-1"} +def test_historical_agent_version_two_uses_inferred_outputs() -> None: + graph = { + "nodes": [ + {"id": "legacy", "data": {"type": "agent", "version": "2", "title": "Legacy"}}, + { + "id": "target", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}, + }, + ], + "edges": [{"source": "legacy", "target": "target"}], + } + declared_loader = Mock(return_value=[DeclaredOutputConfig(name="declared", type=DeclaredOutputType.STRING)]) + + entries, _ = previous_node_output_candidates( + graph=graph, + node_id="target", + declared_outputs_loader=declared_loader, + draft_variables_loader=lambda node_id: [("legacy_output", "string")] if node_id == "legacy" else [], + system_variables_loader=lambda: [], + ) + + assert [entry["output"] for entry in entries] == ["legacy_output"] + declared_loader.assert_not_called() + + def test_previous_outputs_capped_and_flagged(): graph = { "nodes": [{"id": "start-1", "data": {"type": "start", "title": "S", "variables": []}}, {"id": "t"}], diff --git a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py index cb577eed287..70cd220877a 100644 --- a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py +++ b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py @@ -140,7 +140,10 @@ def test_publish_binding_replacement_returns_only_previous_inline_agent( sqlite_session: Session, ) -> None: draft_workflow = _workflow() - draft_workflow.graph = '{"nodes":[{"id":"agent-node","data":{"type":"agent","version":"2"}}],"edges":[]}' + draft_workflow.graph = ( + '{"nodes":[{"id":"agent-node","data":{"type":"agent","version":"2",' + '"agent_node_kind":"dify_agent"}}],"edges":[]}' + ) published_workflow = _workflow(workflow_id="published-new", version="published-new") app = App( id="app-1", diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index bd1378209f8..ce747c5c563 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -224,6 +224,7 @@ class TestAccountService: interface_language="en-US", password="password123", interface_theme="light", + ip_address="203.0.113.10", session=service_session, ) account_id = result.id @@ -235,6 +236,7 @@ class TestAccountService: assert result.password is not None assert result.password_salt is not None assert result.timezone == "America/New_York" + assert result.last_login_ip == "203.0.113.10" with sqlite_session_factory() as assertion_session: persisted_account = assertion_session.get(Account, account_id) @@ -246,6 +248,7 @@ class TestAccountService: assert persisted_account.password is not None assert persisted_account.password_salt is not None assert persisted_account.timezone == "America/New_York" + assert persisted_account.last_login_ip == "203.0.113.10" def test_create_account_uses_explicit_timezone( self, @@ -338,6 +341,7 @@ class TestAccountService: assert result.password is None assert result.password_salt is None assert result.timezone is not None + assert result.last_login_ip is None with sqlite_session_factory() as assertion_session: persisted_account = assertion_session.get(Account, account_id) @@ -349,6 +353,33 @@ class TestAccountService: assert persisted_account.password is None assert persisted_account.password_salt is None assert persisted_account.timezone is not None + assert persisted_account.last_login_ip is None + + def test_update_login_info_overwrites_initial_registration_ip( + self, + sqlite_session_factory: sessionmaker[Session], + mock_external_service_dependencies: _MockDependencies, + ) -> None: + mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True + mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False + + with sqlite_session_factory() as service_session: + account = AccountService.create_account( + email="test@example.com", + name="Test User", + interface_language="en-US", + ip_address="203.0.113.10", + session=service_session, + ) + account_id = account.id + + AccountService.update_login_info(account, service_session, ip_address="203.0.113.11") + + with sqlite_session_factory() as assertion_session: + persisted_account = assertion_session.get(Account, account_id) + assert persisted_account is not None + assert persisted_account.last_login_ip == "203.0.113.11" + assert persisted_account.last_login_at is not None # ==================== Password Management Tests ==================== @@ -1448,6 +1479,7 @@ class TestRegisterService: interface_language="en-US", password="password123", is_setup=True, + ip_address="192.168.1.1", session=service_session, ) mock_create_tenant.assert_called_once_with( @@ -1555,10 +1587,20 @@ class TestRegisterService: name="Test User", interface_language="en-US", password=None, + ip_address="203.0.113.10", session=sqlite_session, ) assert result == mock_account + mock_create_account.assert_called_once_with( + email="test@example.com", + name="Test User", + interface_language="en-US", + password=None, + timezone=None, + ip_address="203.0.113.10", + session=sqlite_session, + ) mock_create_workspace.assert_called_once_with(account=mock_account, session=sqlite_session) mock_join_default_workspace.assert_called_once_with(mock_account.id) @@ -1658,6 +1700,7 @@ class TestRegisterService: name="Test User", password="password123", language="en-US", + ip_address="203.0.113.10", session=sqlite_session, ) @@ -1672,6 +1715,7 @@ class TestRegisterService: password="password123", is_setup=False, timezone=None, + ip_address="203.0.113.10", session=sqlite_session, ) mock_create_owner_tenant.assert_called_once_with(mock_account, session=sqlite_session) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index b6d3a5c150e..5b1ce38ad4a 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -220,6 +220,30 @@ class TestBillingServiceSendRequest: call_args = mock_httpx_request.call_args assert call_args[0][0] == method + def test_new_agent_beta_ensure_uses_secret_authenticated_v1_base(self, mock_httpx_request, mock_billing_config): + mock_response = MagicMock() + mock_response.status_code = httpx.codes.OK + mock_response.json.return_value = {"status": "issued"} + mock_httpx_request.return_value = mock_response + + BillingService.ensure_new_agent_beta_revision("revision-1") + + call_args = mock_httpx_request.call_args + assert call_args.args == ( + "POST", + "https://billing-api.example.com/new-agent-beta/revisions/revision-1/ensure", + ) + assert call_args.kwargs["headers"]["Billing-Api-Secret-Key"] == "test-secret-key" + + def test_new_agent_beta_ensure_requires_json_response(self, mock_httpx_request, mock_billing_config): + mock_response = MagicMock() + mock_response.status_code = httpx.codes.OK + mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + mock_httpx_request.return_value = mock_response + + with pytest.raises(json.JSONDecodeError): + BillingService.ensure_new_agent_beta_revision("revision-1") + @pytest.mark.parametrize( "status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND] ) 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 8cd907ed705..6a52ebe5002 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 @@ -138,7 +138,7 @@ def _execution( def _agent_v2_node(*, node_id: str = "agent-node-1", title: str = "My Agent") -> dict[str, Any]: return { "id": node_id, - "data": {"type": "agent", "version": "2", "title": title}, + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": title}, } @@ -243,6 +243,19 @@ def test_snapshot_accepts_published_run_d1_lifted(session_for: SessionFor) -> No assert [n.node_id for n in snapshot.node_outputs] == ["agent-1"] +def test_historical_agent_version_two_uses_inferred_outputs(session_for: SessionFor) -> None: + resolver = MagicMock() + service = NodeOutputInspectorService(binding_resolver=resolver) + run = _workflow_run(nodes=[{"id": "legacy-agent", "data": {"type": "agent", "version": "2", "title": "Legacy"}}]) + execution = _execution(node_id="legacy-agent", outputs={"text": "legacy output"}) + session = session_for(workflow_run=run, executions=[execution]) + + snapshot = service.snapshot_workflow_run(app_model=_app_model(), workflow_run_id="run-1", session=session) + + assert [(output.name, output.type) for output in snapshot.node_outputs[0].outputs] == [("text", None)] + resolver.resolve.assert_not_called() + + def test_snapshot_accepts_webhook_triggered_run(session_for: SessionFor) -> None: """Webhook / schedule / plugin triggers are also published-side.""" service = _make_service() diff --git a/api/tests/unit_tests/tasks/test_new_agent_beta_task.py b/api/tests/unit_tests/tasks/test_new_agent_beta_task.py new file mode 100644 index 00000000000..d3833f7b269 --- /dev/null +++ b/api/tests/unit_tests/tasks/test_new_agent_beta_task.py @@ -0,0 +1,191 @@ +from datetime import UTC, datetime +from typing import Protocol, cast +from unittest.mock import MagicMock + +import pytest +from sqlalchemy.orm import Session + +from enums import DeploymentEdition +from models.agent import AgentConfigRevision, AgentConfigRevisionOperation +from services.billing_service import BillingService +from tasks import new_agent_beta_task as task_module +from tasks.new_agent_beta_task import ( + NEW_AGENT_BETA_QUEUE, + ensure_new_agent_beta_participation_task, + register_new_agent_beta_publish_after_commit, + schedule_new_agent_beta_ensure, +) + + +class _TaskWithQueue(Protocol): + queue: str + + +def _configure_cloud_publish(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(task_module.dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD) + monkeypatch.setattr(task_module.dify_config, "NEW_AGENT_BETA_ACTIVITY_START_AT", datetime(2026, 8, 12, tzinfo=UTC)) + monkeypatch.setattr(task_module.dify_config, "NEW_AGENT_BETA_ACTIVITY_END_AT", datetime(2026, 8, 13, tzinfo=UTC)) + + +@pytest.mark.parametrize("sqlite_session", [(AgentConfigRevision,)], indirect=True) +def test_publish_event_is_dispatched_only_after_commit( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + published_at = datetime(2026, 8, 12, 1, 0) + _configure_cloud_publish(monkeypatch) + revision = AgentConfigRevision( + id="revision-1", + tenant_id="tenant-1", + agent_id="agent-1", + current_snapshot_id="snapshot-1", + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + created_at=published_at, + ) + sqlite_session.add(revision) + sqlite_session.flush() + dispatch = MagicMock() + monkeypatch.setattr(task_module, "schedule_new_agent_beta_ensure", dispatch) + + register_new_agent_beta_publish_after_commit( + session=sqlite_session, + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_id="snapshot-1", + ) + + dispatch.assert_not_called() + sqlite_session.commit() + dispatch.assert_called_once_with("revision-1") + + +@pytest.mark.parametrize("sqlite_session", [(AgentConfigRevision,)], indirect=True) +def test_rolled_back_publish_is_never_dispatched(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + published_at = datetime(2026, 8, 12, 1, 0) + _configure_cloud_publish(monkeypatch) + sqlite_session.add( + AgentConfigRevision( + id="revision-1", + tenant_id="tenant-1", + agent_id="agent-1", + current_snapshot_id="snapshot-1", + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + created_at=published_at, + ) + ) + sqlite_session.flush() + dispatch = MagicMock() + monkeypatch.setattr(task_module, "schedule_new_agent_beta_ensure", dispatch) + + register_new_agent_beta_publish_after_commit( + session=sqlite_session, + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_id="snapshot-1", + ) + sqlite_session.rollback() + sqlite_session.commit() + + dispatch.assert_not_called() + + +def test_non_cloud_publish_skips_revision_lookup(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(task_module.dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY) + session = MagicMock() + + register_new_agent_beta_publish_after_commit( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_id="snapshot-1", + ) + + session.scalar.assert_not_called() + + +@pytest.mark.parametrize( + ("start", "end", "published_at", "expected"), + [ + (datetime(2026, 8, 12, tzinfo=UTC), datetime(2026, 8, 13, tzinfo=UTC), datetime(2026, 8, 12), True), + (datetime(2026, 8, 12, tzinfo=UTC), datetime(2026, 8, 13, tzinfo=UTC), datetime(2026, 8, 13), False), + (None, datetime(2026, 8, 13, tzinfo=UTC), datetime(2026, 8, 12), False), + (datetime(2026, 8, 13, tzinfo=UTC), datetime(2026, 8, 12, tzinfo=UTC), datetime(2026, 8, 12), False), + ], +) +def test_publish_activity_window_is_inclusive_start_exclusive_end( + monkeypatch: pytest.MonkeyPatch, + start: datetime | None, + end: datetime | None, + published_at: datetime, + expected: bool, +) -> None: + monkeypatch.setattr(task_module.dify_config, "NEW_AGENT_BETA_ACTIVITY_START_AT", start) + monkeypatch.setattr(task_module.dify_config, "NEW_AGENT_BETA_ACTIVITY_END_AT", end) + + assert task_module._is_publish_in_activity_window(published_at) is expected + + +def test_broker_failure_does_not_propagate(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ensure_new_agent_beta_participation_task, + "delay", + MagicMock(side_effect=RuntimeError("broker unavailable")), + ) + + schedule_new_agent_beta_ensure("revision-1") + + +def test_task_calls_billing_with_revision_id(monkeypatch: pytest.MonkeyPatch) -> None: + ensure = MagicMock() + monkeypatch.setattr(BillingService, "ensure_new_agent_beta_revision", ensure) + + ensure_new_agent_beta_participation_task.run("revision-1") + + ensure.assert_called_once_with("revision-1") + + +def test_task_is_redelivered_when_worker_is_lost() -> None: + task = cast(_TaskWithQueue, ensure_new_agent_beta_participation_task) + + assert task.queue == NEW_AGENT_BETA_QUEUE + assert ensure_new_agent_beta_participation_task.acks_late is True + assert ensure_new_agent_beta_participation_task.reject_on_worker_lost is True + assert ensure_new_agent_beta_participation_task.max_retries == 8 + + +def test_task_retries_billing_failure(monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("billing unavailable") + monkeypatch.setattr(BillingService, "ensure_new_agent_beta_revision", MagicMock(side_effect=error)) + retry = MagicMock(side_effect=RuntimeError("retry scheduled")) + monkeypatch.setattr(ensure_new_agent_beta_participation_task, "retry", retry) + + with pytest.raises(RuntimeError, match="retry scheduled"): + ensure_new_agent_beta_participation_task.run("revision-1") + + retry.assert_called_once_with(exc=error, countdown=30) + + +def test_task_caps_exponential_retry_delay(monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("billing unavailable") + monkeypatch.setattr(BillingService, "ensure_new_agent_beta_revision", MagicMock(side_effect=error)) + monkeypatch.setattr(ensure_new_agent_beta_participation_task.request, "retries", 7) + retry = MagicMock(side_effect=RuntimeError("retry scheduled")) + monkeypatch.setattr(ensure_new_agent_beta_participation_task, "retry", retry) + + with pytest.raises(RuntimeError, match="retry scheduled"): + ensure_new_agent_beta_participation_task.run("revision-1") + + retry.assert_called_once_with(exc=error, countdown=900) + + +def test_billing_contract_uses_internal_ensure_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + send_request = MagicMock(return_value={"status": "issued"}) + monkeypatch.setattr(BillingService, "_send_request", send_request) + + BillingService.ensure_new_agent_beta_revision("revision-1") + + send_request.assert_called_once_with( + "POST", + "/new-agent-beta/revisions/revision-1/ensure", + ) 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 7c7f8f34b08..c319d307c56 100644 --- a/api/tests/unit_tests/tasks/test_workflow_execute_task.py +++ b/api/tests/unit_tests/tasks/test_workflow_execute_task.py @@ -852,6 +852,108 @@ def test_resume_app_execution_returns_early_when_advanced_chat_missing_conversat resume_advanced_chat.assert_not_called() +def test_resume_app_execution_clears_stale_cancellation_signals_before_resuming( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], +): + """A resumed run reuses the paused task ID, so it must not inherit its cancellation signals. + + Regression test for #40878: a stop flag or queued AbortCommand left over from + an earlier attempt of the same task aborted the resumed run, which then + finished as "Stopped by user". + """ + workflow_run_id = "run-id" + _persist_resumption_models(sqlite_session_factory, workflow_run_id=workflow_run_id) + + monkeypatch.setattr("tasks.app_generate.workflow_execute_task.db", SimpleNamespace(engine=sqlite_engine)) + + pause_entity = MagicMock() + pause_entity.get_state.return_value = b"state" + + workflow_run_repo = MagicMock() + workflow_run_repo.get_workflow_pause.return_value = pause_entity + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.DifyAPIRepositoryFactory.create_api_workflow_run_repository", + lambda *_args, **_kwargs: workflow_run_repo, + ) + + generate_entity = _build_workflow_generate_entity(stream=False) + resumption_context = MagicMock() + resumption_context.serialized_graph_runtime_state = "{}" + resumption_context.get_generate_entity.return_value = generate_entity + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.WorkflowResumptionContext.loads", + lambda *_args, **_kwargs: resumption_context, + ) + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.GraphRuntimeState.from_snapshot", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task._resolve_user_for_run", lambda *_args, **_kwargs: MagicMock() + ) + + calls: list[str] = [] + clear_signals = MagicMock(side_effect=lambda task_id: calls.append(f"clear:{task_id}")) + resume_workflow = MagicMock(side_effect=lambda **_kwargs: calls.append("resume")) + monkeypatch.setattr("tasks.app_generate.workflow_execute_task.clear_app_task_cancellation_signals", clear_signals) + monkeypatch.setattr("tasks.app_generate.workflow_execute_task._resume_workflow", resume_workflow) + + _resume_app_execution({"workflow_run_id": workflow_run_id}) + + clear_signals.assert_called_once_with(generate_entity.task_id) + # Clearing after the engine started would let it observe the stale abort first. + assert calls == [f"clear:{generate_entity.task_id}", "resume"] + + +def test_resume_app_execution_keeps_cancellation_signals_when_resume_is_abandoned( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], +): + """No attempt is starting, so nothing may clear the signals guarding this task.""" + workflow_run_id = "run-id" + _persist_resumption_models(sqlite_session_factory, workflow_run_id=workflow_run_id) + + monkeypatch.setattr("tasks.app_generate.workflow_execute_task.db", SimpleNamespace(engine=sqlite_engine)) + + pause_entity = MagicMock() + pause_entity.get_state.return_value = b"state" + + workflow_run_repo = MagicMock() + workflow_run_repo.get_workflow_pause.return_value = pause_entity + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.DifyAPIRepositoryFactory.create_api_workflow_run_repository", + lambda *_args, **_kwargs: workflow_run_repo, + ) + + # Missing conversation id makes the advanced-chat resume bail out before running. + generate_entity = _build_advanced_chat_generate_entity(conversation_id=None) + resumption_context = MagicMock() + resumption_context.serialized_graph_runtime_state = "{}" + resumption_context.get_generate_entity.return_value = generate_entity + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.WorkflowResumptionContext.loads", + lambda *_args, **_kwargs: resumption_context, + ) + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task.GraphRuntimeState.from_snapshot", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "tasks.app_generate.workflow_execute_task._resolve_user_for_run", lambda *_args, **_kwargs: MagicMock() + ) + + clear_signals = MagicMock() + monkeypatch.setattr("tasks.app_generate.workflow_execute_task.clear_app_task_cancellation_signals", clear_signals) + monkeypatch.setattr("tasks.app_generate.workflow_execute_task._resume_advanced_chat", MagicMock()) + + _resume_app_execution({"workflow_run_id": workflow_run_id}) + + clear_signals.assert_not_called() + + def test_resume_advanced_chat_publishes_events_for_originally_blocking_runs( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, diff --git a/api/tests/unit_tests/test_config_overrides.py b/api/tests/unit_tests/test_config_overrides.py new file mode 100644 index 00000000000..39588272246 --- /dev/null +++ b/api/tests/unit_tests/test_config_overrides.py @@ -0,0 +1,19 @@ +"""Contract tests for the shared unit-test config override fixture.""" + +from collections.abc import Callable + +import pytest + +from configs import dify_config +from enums import DeploymentEdition + + +def test_config_overrides_updates_shared_config(config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + assert dify_config.DEPLOYMENT_EDITION is DeploymentEdition.CLOUD + + +def test_config_overrides_rejects_unknown_fields(config_overrides: Callable[..., None]) -> None: + with pytest.raises(ValueError, match=r"Unknown DifyConfig fields: \['NOT_A_CONFIG_FIELD'\]"): + config_overrides(NOT_A_CONFIG_FIELD=True) diff --git a/api/uv.lock b/api/uv.lock index 2046ba30321..124a43d31f0 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1302,6 +1302,7 @@ dependencies = [ { name = "httpx" }, { name = "httpx2" }, { name = "pydantic" }, + { name = "pydantic-ai-harness" }, { name = "pydantic-ai-slim" }, { name = "typing-extensions" }, ] @@ -1317,8 +1318,9 @@ requires-dist = [ { name = "jwcrypto", marker = "extra == 'server'", specifier = ">=1.5.6,<2" }, { name = "logfire", extras = ["fastapi", "httpx", "redis"], marker = "extra == 'server'", specifier = ">=4.37.0,<5.0.0" }, { name = "pydantic", specifier = ">=2.12.5,<2.13" }, - { name = "pydantic-ai-slim", specifier = ">=1.106.0,<2.0.0" }, - { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" }, + { name = "pydantic-ai-harness", specifier = ">=0.20.0,<0.21.0" }, + { name = "pydantic-ai-slim", specifier = ">=2.30.0,<3.0.0" }, + { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=2.30.0,<3.0.0" }, { name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" }, { name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" }, { name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" }, @@ -2674,15 +2676,15 @@ wheels = [ [[package]] name = "genai-prices" -version = "0.0.67" +version = "0.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx2" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/9e/f96ad08d62f7bd33a5b24e65d4eb220569714b9a2a8813ada2e1fa47b4dd/genai_prices-0.0.67.tar.gz", hash = "sha256:54e07eb6541fda377187a471c5dba21a81b439c57f8dc44d89db3103c29ca343", size = 80015, upload-time = "2026-06-24T20:16:23.661Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/14/a188df294f013ec9cd97fc6b145f5427f89067bfb2c260fc3fb5c8d1fb34/genai_prices-0.1.3.tar.gz", hash = "sha256:62c30cddd6c2d2199d878d1a70521c3e37347cd9394446d107dc774a78ed3780", size = 92638, upload-time = "2026-08-15T00:10:31.771Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/05/d1ca6b960a3305f86d1c5f4274f2ddf8c94611ec7edc436a01cd38a01742/genai_prices-0.0.67-py3-none-any.whl", hash = "sha256:08977f1e83b4132abcfc60dabf21ff13c2d25958afb9199e59c4407bf5c9ed3f", size = 82495, upload-time = "2026-06-24T20:16:22.4Z" }, + { url = "https://files.pythonhosted.org/packages/4e/cd/d94b47c26d6367e0b949edfe2da5a47fb74037e799a0edd5b825e049f2b9/genai_prices-0.1.3-py3-none-any.whl", hash = "sha256:a2603841429c843da91c987d9ef598c73bd940caf44e844ab046d551791c04bb", size = 96892, upload-time = "2026-08-15T00:10:30.595Z" }, ] [[package]] @@ -5204,10 +5206,25 @@ wheels = [ ] [[package]] -name = "pydantic-ai-slim" -version = "1.107.0" +name = "pydantic-ai-harness" +version = "0.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "genai-prices" }, + { name = "httpx" }, + { name = "pydantic-ai-slim" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3f/6c/a066644a3a3bff22bfdddd745fd245e2b4e3148fc895be0020f03bd7470d/pydantic_ai_harness-0.20.0.tar.gz", hash = "sha256:18ec7d6f90873a8038d094280e50af5e877320b975f334268b700a266d35f522", size = 1846014, upload-time = "2026-08-14T03:36:31.045Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/85/32cab39557e338abbd4ecb844028a31110c03096e6678bf5a1266ed201d1/pydantic_ai_harness-0.20.0-py3-none-any.whl", hash = "sha256:e1164ae4d653bd2ae257e3816ee1776b27dbe8f7eaeb29b36d86a49d6fe98168", size = 623606, upload-time = "2026-08-14T03:36:29.094Z" }, +] + +[[package]] +name = "pydantic-ai-slim" +version = "2.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, { name = "genai-prices" }, { name = "griffelib" }, { name = "httpx" }, @@ -5216,9 +5233,9 @@ dependencies = [ { name = "pydantic-graph" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4c/26/ced63dfaabbc77f3beb86d59689cdea748e7ccffb6b419dbaf4780f211e8/pydantic_ai_slim-1.107.0.tar.gz", hash = "sha256:4616f689a92fcfecfecf2a7af27aca22f139a873cf6d7a8929eaeee9c0eedbb4", size = 779902, upload-time = "2026-06-10T14:53:10.574Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/4e/2165d9b90edcd5dfc8e9465b3bdc0aa66a67760843ddb0ac99ee396898f0/pydantic_ai_slim-2.31.0.tar.gz", hash = "sha256:a9310d2464154b028096f1d680f17837f16e5c6cd209b4542e4f60ca5d344789", size = 1214087, upload-time = "2026-08-15T03:17:28.353Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/57/71044e17f931b08cc3930bc0fe5a1e1fd37fa474ae826be004729ef1cb4a/pydantic_ai_slim-1.107.0-py3-none-any.whl", hash = "sha256:1af49bbae06a6c598f72c54d4734ba377100cac493c9a05fa8e089bebeae0da6", size = 964046, upload-time = "2026-06-10T14:53:03.333Z" }, + { url = "https://files.pythonhosted.org/packages/db/09/233e529fadbece38580c3a390783f55fa196afad875ce535ee9a57a5ad71/pydantic_ai_slim-2.31.0-py3-none-any.whl", hash = "sha256:cb809ad949ca68be6bb9a0e0b994fc73a95f4cd405e8609a71034f2e2080e2a1", size = 1432053, upload-time = "2026-08-15T03:17:21.208Z" }, ] [[package]] @@ -5265,17 +5282,18 @@ wheels = [ [[package]] name = "pydantic-graph" -version = "1.107.0" +version = "2.31.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "anyio" }, { name = "httpx" }, { name = "logfire-api" }, { name = "pydantic" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dd/c3/6e8c2d13b8701041f1b3eac5deb41f25d4dbfa479a190d5c6becc23f2a49/pydantic_graph-1.107.0.tar.gz", hash = "sha256:278dd89b3e33f3a2963ac949f27a53aef705c5d883a8ce5d06d23e6e3cfbd972", size = 62564, upload-time = "2026-06-10T14:53:13.366Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8a/2c/3817ae318ecc729a258fc85aaad84fc110b9467a292b98bb10e65ae71183/pydantic_graph-2.31.0.tar.gz", hash = "sha256:a19919408dfaa5a1b8713618bcce7a5135d83664321862b0d9be1f4216c432ef", size = 45180, upload-time = "2026-08-15T03:17:30.398Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fc/72/621556e3f5068400d43a0375d38e5963de30256eaa5a702aba12e82ed0ff/pydantic_graph-1.107.0-py3-none-any.whl", hash = "sha256:71add94fe7e14c703977a895117c475aae6c0b02a774a036c4d00d9a63c78b00", size = 80106, upload-time = "2026-06-10T14:53:06.543Z" }, + { url = "https://files.pythonhosted.org/packages/54/ef/a3217caed3189cfcc6857316e06cb900990e4d5cb59e42ec517cfdda6a7f/pydantic_graph-2.31.0-py3-none-any.whl", hash = "sha256:062555c89b1d5699ddaaa6f09ae62fc1d0f5cc189d0867e2fafca68c24797ba5", size = 52662, upload-time = "2026-08-15T03:17:24.42Z" }, ] [[package]] diff --git a/dify-agent-runtime/Makefile b/dify-agent-runtime/Makefile index 76e853a03f9..533248f5e1d 100644 --- a/dify-agent-runtime/Makefile +++ b/dify-agent-runtime/Makefile @@ -1,4 +1,4 @@ -.PHONY: build clean test lint gen-cli-help integration integration-up integration-test integration-down +.PHONY: build clean test lint gen-cli-help integration integration-up integration-test integration-down sync-e2b-template-dev sync-e2b-template-prod sync-e2b-template-dry-run BIN_DIR := bin AGENT_CLI_HELP_JSON := ../dify-agent/src/dify_agent/layers/_agent_cli_help.json @@ -128,3 +128,18 @@ integration: @$(MAKE) integration-up @trap '$(MAKE) integration-down' EXIT; \ $(MAKE) integration-test + +# --- E2B Template sync --- +# +# Both E2B projects publish the same template name (dify-agent-local-sandbox); +# the project itself is the isolation boundary. Export E2B_API_KEY_DEV / +# E2B_API_KEY_PROD in your shell before running the non-dry-run targets. + +sync-e2b-template-dry-run: + ./docker/sync-e2b-template.sh --dry-run + +sync-e2b-template-dev: + E2B_API_KEY=$(E2B_API_KEY_DEV) ./docker/sync-e2b-template.sh + +sync-e2b-template-prod: + E2B_API_KEY=$(E2B_API_KEY_PROD) ./docker/sync-e2b-template.sh diff --git a/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go index a8a0c328f8d..f176b8337c7 100644 --- a/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go +++ b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go @@ -59,6 +59,7 @@ func TestDumpCLIHelpCoversPromptCommands(t *testing.T) { "config skills delete", "file upload", "file download", + "file public-url", } for _, key := range wantKeys { if _, ok := table[key]; !ok { diff --git a/dify-agent-runtime/cmd/dify-agent-cli/main.go b/dify-agent-runtime/cmd/dify-agent-cli/main.go index 2111d6cd4f1..08bf33f3589 100644 --- a/dify-agent-runtime/cmd/dify-agent-cli/main.go +++ b/dify-agent-runtime/cmd/dify-agent-cli/main.go @@ -118,6 +118,14 @@ func newFileCommand() *cobra.Command { false, "Skip creating a public download link after upload.", ) + publicURL := &cobra.Command{ + Use: "public-url REFERENCE", + Short: "Create a browser-visible download URL for an existing ToolFile reference.", + Args: cobra.ExactArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunFilePublicURL(env, args[0]) + }), + } var downloadTo string download := &cobra.Command{ @@ -130,7 +138,7 @@ func newFileCommand() *cobra.Command { } download.Flags().StringVar(&downloadTo, "to", "", "Local directory for the downloaded file.") - cmd.AddCommand(upload, download) + cmd.AddCommand(upload, download, publicURL) return cmd } diff --git a/dify-agent-runtime/cmd/dify-agent-cli/main_test.go b/dify-agent-runtime/cmd/dify-agent-cli/main_test.go index 8e024122c6a..11326c9a8c5 100644 --- a/dify-agent-runtime/cmd/dify-agent-cli/main_test.go +++ b/dify-agent-runtime/cmd/dify-agent-cli/main_test.go @@ -65,6 +65,11 @@ func TestCommandHelp(t *testing.T) { args: []string{"file", "download", "--help"}, want: []string{"dify-agent file download", "Download one workflow file", "--to"}, }, + { + name: "file public-url", + args: []string{"file", "public-url", "--help"}, + want: []string{"dify-agent file public-url", "Create a browser-visible download URL"}, + }, { name: "drive", args: []string{"drive", "--help"}, diff --git a/dify-agent-runtime/internal/agentcli/client.go b/dify-agent-runtime/internal/agentcli/client.go index 07c80515954..8e84991ab72 100644 --- a/dify-agent-runtime/internal/agentcli/client.go +++ b/dify-agent-runtime/internal/agentcli/client.go @@ -7,6 +7,7 @@ type StubClient interface { // HTTP control-plane Connect(ctx context.Context, argv []string, metadataJSON string) (*ConnectResponse, error) CreateFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) + CreateToolFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) CreateFileDownloadURL(ctx context.Context, transferMethod string, reference, url *string, forFrontend bool) (*FileDownloadResponse, error) // Drive operations (HTTP-only control-plane) diff --git a/dify-agent-runtime/internal/agentcli/client_http.go b/dify-agent-runtime/internal/agentcli/client_http.go index fa01c818711..14abb3f368b 100644 --- a/dify-agent-runtime/internal/agentcli/client_http.go +++ b/dify-agent-runtime/internal/agentcli/client_http.go @@ -36,8 +36,8 @@ func (c *httpStubClient) Connect(_ context.Context, argv []string, metadataJSON if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "connect"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub connect failed: %w", err) } var resp ConnectResponse @@ -72,6 +72,31 @@ func (c *httpStubClient) CreateFileUploadURL(_ context.Context, filename, mimety return resp.UploadURL, nil } +func (c *httpStubClient) CreateToolFileUploadURL(_ context.Context, filename, mimetype string) (string, error) { + payload := map[string]string{ + "filename": filename, + "mimetype": mimetype, + } + body, statusCode, err := c.http.postJSON("/files/upload-request?expose_expiration=true", payload) + if err != nil { + return "", err + } + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return "", fmt.Errorf("agent stub file upload request failed: %w", err) + } + + var resp struct { + UploadURL string `json:"upload_url"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return "", fmt.Errorf("parse upload response: %w", err) + } + if resp.UploadURL == "" { + return "", fmt.Errorf("signed file upload response is missing upload_url") + } + return resp.UploadURL, nil +} + func (c *httpStubClient) CreateFileDownloadURL(_ context.Context, transferMethod string, reference, url *string, forFrontend bool) (*FileDownloadResponse, error) { fileMapping := map[string]any{ "transfer_method": transferMethod, @@ -91,8 +116,8 @@ func (c *httpStubClient) CreateFileDownloadURL(_ context.Context, transferMethod if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "file download request"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub file download request failed: %w", err) } var resp FileDownloadResponse @@ -146,8 +171,8 @@ func (c *httpStubClient) GetConfigManifest(_ context.Context) ([]byte, error) { if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "config manifest"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub config manifest failed: %w", err) } return body, nil } @@ -167,8 +192,8 @@ func (c *httpStubClient) CreateConfigDownloadURL( if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "config download request"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub config download request failed: %w", err) } var resp FileDownloadResponse @@ -186,8 +211,8 @@ func (c *httpStubClient) PushConfig(_ context.Context, payload any) ([]byte, err if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "config push"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub config push failed: %w", err) } return body, nil } @@ -198,8 +223,8 @@ func (c *httpStubClient) PatchConfigEnv(_ context.Context, envText string) ([]by if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "config env update"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub config env update failed: %w", err) } return body, nil } @@ -210,8 +235,8 @@ func (c *httpStubClient) PutConfigNote(_ context.Context, note string) ([]byte, if err != nil { return nil, err } - if err := checkHTTPError(body, statusCode, "config note update"); err != nil { - return nil, err + if err := checkAgentStubHTTPError(body, statusCode); err != nil { + return nil, fmt.Errorf("agent stub config note update failed: %w", err) } return body, nil } diff --git a/dify-agent-runtime/internal/agentcli/config.go b/dify-agent-runtime/internal/agentcli/config.go index d627625c542..69c038a8e2f 100644 --- a/dify-agent-runtime/internal/agentcli/config.go +++ b/dify-agent-runtime/internal/agentcli/config.go @@ -76,7 +76,7 @@ func RunConfigSkillsPull(env *Environment, names []string, localDir string, json for _, name := range names { download, err := client.CreateConfigDownloadURL(ctx, "skill", name) if err != nil { - return err + return fmt.Errorf("request config skill %q download URL: %w", name, err) } archiveBytes, err := client.DownloadFromURL(download.DownloadURL) if err != nil { @@ -174,7 +174,7 @@ func RunConfigFilesPull(env *Environment, names []string, localDir string, jsonO for _, name := range names { download, err := client.CreateConfigDownloadURL(ctx, "file", name) if err != nil { - return err + return fmt.Errorf("request config file %q download URL: %w", name, err) } payload, err := client.DownloadFromURL(download.DownloadURL) if err != nil { @@ -243,14 +243,14 @@ func RunConfigSkillsPush(env *Environment, paths []string) error { defer func() { _ = os.Remove(archivePath) }() name := filepath.Base(absPath) - commitItem, err := uploadAndPrepareCommitItem(client, archivePath, name) + fileRef, err := uploadAndPrepareConfigItem(client, archivePath) if err != nil { - return err + return fmt.Errorf("upload config skill %q: %w", name, err) } skills = append(skills, skillPushItem{ Name: name, - FileRef: &DriveFileRef{Kind: commitItem.FileRef.Kind, ID: commitItem.FileRef.ID}, + FileRef: fileRef, }) } @@ -261,7 +261,7 @@ func RunConfigSkillsPush(env *Environment, paths []string) error { body, err := client.PushConfig(context.Background(), payload) if err != nil { - return err + return fmt.Errorf("push config skills: %w", err) } fmt.Println(string(body)) return nil @@ -296,14 +296,14 @@ func RunConfigFilesPush(env *Environment, paths []string) error { } name := filepath.Base(absPath) - commitItem, err := uploadAndPrepareCommitItem(client, absPath, name) + fileRef, err := uploadAndPrepareConfigItem(client, absPath) if err != nil { - return err + return fmt.Errorf("upload config file %q: %w", name, err) } files = append(files, filePushItem{ Name: name, - FileRef: &DriveFileRef{Kind: commitItem.FileRef.Kind, ID: commitItem.FileRef.ID}, + FileRef: fileRef, }) } @@ -314,12 +314,35 @@ func RunConfigFilesPush(env *Environment, paths []string) error { body, err := client.PushConfig(context.Background(), payload) if err != nil { - return err + return fmt.Errorf("push config files: %w", err) } fmt.Println(string(body)) return nil } +func uploadAndPrepareConfigItem(client StubClient, filePath string) (*DriveFileRef, error) { + filename := filepath.Base(filePath) + mimetype := guessMIMEType(filename) + uploadURL, err := client.CreateToolFileUploadURL(context.Background(), filename, mimetype) + if err != nil { + return nil, fmt.Errorf("request upload URL: %w", err) + } + uploadBody, err := client.UploadFileToURL(uploadURL, filePath, filename, mimetype) + if err != nil { + return nil, fmt.Errorf("upload data: %w", err) + } + + var uploadResult map[string]any + if err := json.Unmarshal(uploadBody, &uploadResult); err != nil { + return nil, fmt.Errorf("parse upload result: %w", err) + } + toolFileID, _ := uploadResult["id"].(string) + if toolFileID == "" { + return nil, fmt.Errorf("upload response is missing id") + } + return &DriveFileRef{Kind: "tool_file", ID: toolFileID}, nil +} + // RunConfigSkillsDelete executes the `config skills delete` command. func RunConfigSkillsDelete(env *Environment, names []string) error { if len(names) == 0 { diff --git a/dify-agent-runtime/internal/agentcli/config_test.go b/dify-agent-runtime/internal/agentcli/config_test.go index 8f67d5c0863..833932ff545 100644 --- a/dify-agent-runtime/internal/agentcli/config_test.go +++ b/dify-agent-runtime/internal/agentcli/config_test.go @@ -4,6 +4,7 @@ import ( "archive/zip" "bytes" "encoding/json" + "io" "net/http" "net/http/httptest" "os" @@ -111,15 +112,282 @@ func TestConfigPullRequestsURLThenDownloadsFromDataPlane(t *testing.T) { } } -func TestConfigPullReportsControlAndDataPlaneFailures(t *testing.T) { +func TestConfigPullMultiItemFailuresIdentifyItemAndStage(t *testing.T) { + skillArchive := zipFixture(t, map[string]string{"SKILL.md": "# Alpha\n"}) tests := []struct { - name string - controlStatus int - dataStatus int - want string + name string + kind string + names []string + payload []byte + failureStage string + wantStage string }{ - {name: "control", controlStatus: http.StatusNotFound, dataStatus: http.StatusOK, want: "config download request failed"}, - {name: "data plane", controlStatus: http.StatusOK, dataStatus: http.StatusBadGateway, want: "download config file"}, + { + name: "file control-plane URL acquisition", + kind: "file", + names: []string{"first.txt", "second.txt"}, + payload: []byte("first file"), + failureStage: "control", + wantStage: `request config file "second.txt" download URL`, + }, + { + name: "file signed-URL download", + kind: "file", + names: []string{"first.txt", "second.txt"}, + payload: []byte("first file"), + failureStage: "data", + wantStage: `download config file "second.txt"`, + }, + { + name: "skill control-plane URL acquisition", + kind: "skill", + names: []string{"alpha", "beta"}, + payload: skillArchive, + failureStage: "control", + wantStage: `request config skill "beta" download URL`, + }, + { + name: "skill signed-URL download", + kind: "skill", + names: []string{"alpha", "beta"}, + payload: skillArchive, + failureStage: "data", + wantStage: `download config skill "beta"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + controlCalls := 0 + dataPlaneCalls := 0 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/agent-stub/files/download-request": + controlCalls++ + var request struct { + Config struct { + Kind string `json:"kind"` + Name string `json:"name"` + } `json:"config"` + } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Errorf("decode download request: %v", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + if controlCalls > len(test.names) { + t.Errorf("unexpected control-plane call %d", controlCalls) + http.Error(w, "unexpected request", http.StatusInternalServerError) + return + } + wantName := test.names[controlCalls-1] + if request.Config.Kind != test.kind || request.Config.Name != wantName { + t.Errorf("config request = (%q, %q), want (%q, %q)", request.Config.Kind, request.Config.Name, test.kind, wantName) + } + if test.failureStage == "control" && controlCalls == 2 { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "filename": wantName, + "size": len(test.payload), + "download_url": server.URL + "/files/config-asset", + }) + case "/files/config-asset": + dataPlaneCalls++ + if test.failureStage == "data" && dataPlaneCalls == 2 { + http.Error(w, "data plane unavailable", http.StatusBadGateway) + return + } + _, _ = w.Write(test.payload) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + targetDir := t.TempDir() + env := &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"} + var err error + if test.kind == "file" { + err = RunConfigFilesPull(env, test.names, targetDir, true) + } else { + err = RunConfigSkillsPull(env, test.names, targetDir, true) + } + if err == nil { + t.Fatal("config pull succeeded, want second-item failure") + } + if !strings.Contains(err.Error(), test.wantStage) { + t.Errorf("error = %q, want stage %q", err, test.wantStage) + } + if controlCalls != 2 { + t.Errorf("control-plane calls = %d, want 2", controlCalls) + } + wantDataPlaneCalls := 2 + if test.failureStage == "control" { + wantDataPlaneCalls = 1 + for _, want := range []string{ + "expired after 5 minutes", + "will not refresh automatically", + "start a new shell tool call", + "retry the command", + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error = %q, want expiration guidance %q", err, want) + } + } + } + if dataPlaneCalls != wantDataPlaneCalls { + t.Errorf("data-plane calls = %d, want %d", dataPlaneCalls, wantDataPlaneCalls) + } + + if test.kind == "file" { + data, readErr := os.ReadFile(filepath.Join(targetDir, test.names[0])) + if readErr != nil || !bytes.Equal(data, test.payload) { + t.Errorf("first file was not completed: data=%q err=%v", data, readErr) + } + } else { + data, readErr := os.ReadFile(filepath.Join(targetDir, test.names[0], "SKILL.md")) + if readErr != nil || string(data) != "# Alpha\n" { + t.Errorf("first skill was not completed: data=%q err=%v", data, readErr) + } + } + }) + } +} + +func TestConfigPushMultiItemUploadFailuresIdentifyItemAndStage(t *testing.T) { + tests := []struct { + name string + failureStage string + makeSources func(*testing.T) []string + run func(*Environment, []string) error + wantItem string + }{ + { + name: "file control-plane URL acquisition", + failureStage: "request upload URL", + makeSources: makeConfigFileSources, + run: RunConfigFilesPush, + wantItem: `config file "second.txt"`, + }, + { + name: "file signed URL data transfer", + failureStage: "upload data", + makeSources: makeConfigFileSources, + run: RunConfigFilesPush, + wantItem: `config file "second.txt"`, + }, + { + name: "skill control-plane URL acquisition", + failureStage: "request upload URL", + makeSources: makeConfigSkillSources, + run: RunConfigSkillsPush, + wantItem: `config skill "beta"`, + }, + { + name: "skill signed URL data transfer", + failureStage: "upload data", + makeSources: makeConfigSkillSources, + run: RunConfigSkillsPush, + wantItem: `config skill "beta"`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + uploadRequestCalls := 0 + dataPlaneCalls := 0 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/agent-stub/files/upload-request": + uploadRequestCalls++ + if test.failureStage == "request upload URL" && uploadRequestCalls == 2 { + http.Error(w, "control plane unavailable", http.StatusServiceUnavailable) + return + } + _ = json.NewEncoder(w).Encode(map[string]string{"upload_url": server.URL + "/upload"}) + case "/upload": + _, _ = io.Copy(io.Discard, r.Body) + dataPlaneCalls++ + if test.failureStage == "upload data" && dataPlaneCalls == 2 { + http.Error(w, "data plane unavailable", http.StatusBadGateway) + return + } + _, _ = w.Write([]byte(`{"id":"tool-file-1"}`)) + case "/agent-stub/config/push": + t.Error("final config push must not run after an item upload failure") + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + err := test.run( + &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}, + test.makeSources(t), + ) + if err == nil { + t.Fatal("config push succeeded, want second-item upload failure") + } + if !strings.Contains(err.Error(), test.wantItem) { + t.Errorf("error = %q, want current item %q", err, test.wantItem) + } + if !strings.Contains(err.Error(), test.failureStage) { + t.Errorf("error = %q, want stage %q", err, test.failureStage) + } + if uploadRequestCalls != 2 { + t.Errorf("upload request calls = %d, want 2", uploadRequestCalls) + } + wantDataPlaneCalls := 2 + if test.failureStage == "request upload URL" { + wantDataPlaneCalls = 1 + } + if dataPlaneCalls != wantDataPlaneCalls { + t.Errorf("data-plane calls = %d, want %d", dataPlaneCalls, wantDataPlaneCalls) + } + }) + } +} + +func TestConfigPushFinalFailureIdentifiesOperationAndExplainsExpiry(t *testing.T) { + tests := []struct { + name string + makeSource func(*testing.T) string + run func(*Environment, string) error + want string + }{ + { + name: "file", + makeSource: func(t *testing.T) string { + path := filepath.Join(t.TempDir(), "guide.txt") + if err := os.WriteFile(path, []byte("guide"), 0o600); err != nil { + t.Fatalf("write config file: %v", err) + } + return path + }, + run: func(env *Environment, path string) error { return RunConfigFilesPush(env, []string{path}) }, + want: "push config files", + }, + { + name: "skill", + makeSource: func(t *testing.T) string { + dir := filepath.Join(t.TempDir(), "alpha") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("create config skill: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte("# Alpha\n"), 0o600); err != nil { + t.Fatalf("write config skill: %v", err) + } + return dir + }, + run: func(env *Environment, path string) error { return RunConfigSkillsPush(env, []string{path}) }, + want: "push config skills", + }, } for _, test := range tests { @@ -127,36 +395,61 @@ func TestConfigPullReportsControlAndDataPlaneFailures(t *testing.T) { var server *httptest.Server server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { - case "/agent-stub/files/download-request": - if test.controlStatus != http.StatusOK { - w.WriteHeader(test.controlStatus) - _, _ = w.Write([]byte(`{"detail":"missing"}`)) - return - } - _ = json.NewEncoder(w).Encode(map[string]any{ - "filename": "guide.txt", "size": 5, "download_url": server.URL + "/files/config-asset", - }) - case "/files/config-asset": - w.WriteHeader(test.dataStatus) + case "/agent-stub/files/upload-request": + _ = json.NewEncoder(w).Encode(map[string]string{"upload_url": server.URL + "/upload"}) + case "/upload": + _, _ = w.Write([]byte(`{"id":"tool-file-1"}`)) + case "/agent-stub/config/push": + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`)) default: http.NotFound(w, r) } })) defer server.Close() - err := RunConfigFilesPull( + err := test.run( &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}, - []string{"guide.txt"}, - t.TempDir(), - true, + test.makeSource(t), ) if err == nil || !strings.Contains(err.Error(), test.want) { - t.Fatalf("error = %v, want substring %q", err, test.want) + t.Fatalf("error = %v, want operation %q", err, test.want) + } + if !strings.Contains(err.Error(), "start a new shell tool call") { + t.Fatalf("error = %v, want expiry recovery guidance", err) } }) } } +func makeConfigFileSources(t *testing.T) []string { + t.Helper() + dir := t.TempDir() + paths := []string{filepath.Join(dir, "first.txt"), filepath.Join(dir, "second.txt")} + for _, path := range paths { + if err := os.WriteFile(path, []byte(filepath.Base(path)), 0o600); err != nil { + t.Fatalf("write config file: %v", err) + } + } + return paths +} + +func makeConfigSkillSources(t *testing.T) []string { + t.Helper() + dir := t.TempDir() + paths := []string{filepath.Join(dir, "alpha"), filepath.Join(dir, "beta")} + for _, path := range paths { + if err := os.MkdirAll(path, 0o755); err != nil { + t.Fatalf("create config skill: %v", err) + } + if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("# "+filepath.Base(path)+"\n"), 0o600); err != nil { + t.Fatalf("write config skill: %v", err) + } + } + return paths +} + func zipFixture(t *testing.T, files map[string]string) []byte { t.Helper() var buffer bytes.Buffer diff --git a/dify-agent-runtime/internal/agentcli/connect_test.go b/dify-agent-runtime/internal/agentcli/connect_test.go new file mode 100644 index 00000000000..9d83ea5d4fa --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/connect_test.go @@ -0,0 +1,53 @@ +package agentcli + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestRunConnectExplainsExpiredAuthorization(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/agent-stub/connections" { + http.NotFound(w, r) + return + } + var payload struct { + Argv []string `json:"argv"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode connect request: %v", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + if got, want := strings.Join(payload.Argv, " "), "echo hello"; got != want { + t.Errorf("argv = %q, want %q", got, want) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"Agent Stub authorization expired after 5 minutes; start a new shell tool call and retry the command."}}`)) + })) + defer server.Close() + + err := RunConnect( + &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}, + []string{"echo", "hello"}, + false, + ) + if err == nil { + t.Fatal("RunConnect succeeded, want expired authorization failure") + } + for _, want := range []string{ + "agent stub connect failed", + "expired after 5 minutes", + "will not refresh automatically", + "start a new shell tool call", + "retry the command", + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error = %q, want substring %q", err, want) + } + } +} diff --git a/dify-agent-runtime/internal/agentcli/file.go b/dify-agent-runtime/internal/agentcli/file.go index b1f5179bb60..cbc1c9c3096 100644 --- a/dify-agent-runtime/internal/agentcli/file.go +++ b/dify-agent-runtime/internal/agentcli/file.go @@ -39,8 +39,12 @@ func RunFileUpload(env *Environment, path string, noDownloadLink bool) error { } type fileUploadClient interface { - CreateFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) + filePublicURLClient + CreateToolFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error) +} + +type filePublicURLClient interface { CreateFileDownloadURL( ctx context.Context, transferMethod string, @@ -64,15 +68,15 @@ func runFileUpload(client fileUploadClient, path string, noDownloadLink bool, ou ctx := context.Background() // Step 1: Request a signed upload URL - uploadURL, err := client.CreateFileUploadURL(ctx, filename, mimetype) + uploadURL, err := client.CreateToolFileUploadURL(ctx, filename, mimetype) if err != nil { - return err + return fmt.Errorf("request file upload URL: %w", err) } // Step 2: Upload the file to the signed URL (data-plane) uploadBody, err := client.UploadFileToURL(uploadURL, absPath, filename, mimetype) if err != nil { - return err + return fmt.Errorf("upload file data: %w", err) } var uploadResult map[string]any @@ -84,7 +88,7 @@ func runFileUpload(client fileUploadClient, path string, noDownloadLink bool, ou if reference == "" { return fmt.Errorf("signed file upload response is missing reference") } - if noDownloadLink && !isCanonicalDifyFileReference(reference) { + if !isCanonicalDifyFileReference(reference) { return fmt.Errorf("signed file upload response has invalid reference") } @@ -98,13 +102,63 @@ func runFileUpload(client fileUploadClient, path string, noDownloadLink bool, ou ref := reference dlResp, err := client.CreateFileDownloadURL(ctx, "tool_file", &ref, nil, true) if err != nil { - return err + writeFileUploadResponse(output, result) + return fmt.Errorf( + "request public download URL: %w; retry without uploading again: dify-agent file public-url %s", + err, + shellQuoteArgument(reference), + ) + } + if dlResp.DownloadURL == "" { + writeFileUploadResponse(output, result) + return fmt.Errorf( + "public file download response is missing download_url; retry without uploading again: dify-agent file public-url %s", + shellQuoteArgument(reference), + ) } result.PublicDownloadURL = dlResp.DownloadURL } + writeFileUploadResponse(output, result) + return nil +} + +// RunFilePublicURL requests a browser-visible URL for an existing ToolFile reference. +func RunFilePublicURL(env *Environment, reference string) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + return runFilePublicURL(client, reference, os.Stdout) +} + +func runFilePublicURL(client filePublicURLClient, reference string, output io.Writer) error { + if reference == "" { + return fmt.Errorf("file reference must not be empty") + } + download, err := client.CreateFileDownloadURL(context.Background(), "tool_file", &reference, nil, true) + if err != nil { + return fmt.Errorf("request public download URL: %w", err) + } + if download.DownloadURL == "" { + return fmt.Errorf("public file download response is missing download_url") + } + writeFileUploadResponse(output, FileUploadResponse{ + TransferMethod: "tool_file", + Reference: reference, + PublicDownloadURL: download.DownloadURL, + }) + return nil +} + +func writeFileUploadResponse(output io.Writer, result FileUploadResponse) { out, _ := json.Marshal(result) _, _ = fmt.Fprintln(output, string(out)) - return nil +} + +func shellQuoteArgument(value string) string { + return "'" + strings.ReplaceAll(value, "'", `'"'"'`) + "'" } func isCanonicalDifyFileReference(reference string) bool { @@ -144,7 +198,7 @@ func RunFileDownload(env *Environment, transferMethod string, referenceOrURL str dlResp, err := client.CreateFileDownloadURL(ctx, transferMethod, reference, url, false) if err != nil { - return err + return fmt.Errorf("request file download URL: %w", err) } if dlResp.DownloadURL == "" { return fmt.Errorf("signed file download response is missing download_url") @@ -156,7 +210,7 @@ func RunFileDownload(env *Environment, transferMethod string, referenceOrURL str // Download the file (data-plane) data, err := client.DownloadFromURL(dlResp.DownloadURL) if err != nil { - return err + return fmt.Errorf("download file data: %w", err) } // Determine target directory diff --git a/dify-agent-runtime/internal/agentcli/file_test.go b/dify-agent-runtime/internal/agentcli/file_test.go index a610457e6c2..bde80828a9f 100644 --- a/dify-agent-runtime/internal/agentcli/file_test.go +++ b/dify-agent-runtime/internal/agentcli/file_test.go @@ -15,6 +15,7 @@ import ( type fakeFileUploadClient struct { forFrontend bool downloadRequestCall int + downloadMethod string uploadResponse []byte calls []string filename string @@ -22,9 +23,11 @@ type fakeFileUploadClient struct { uploadURL string uploadedBytes []byte downloadReference string + downloadErr error + downloadURL *string } -func (f *fakeFileUploadClient) CreateFileUploadURL(_ context.Context, filename, mimetype string) (string, error) { +func (f *fakeFileUploadClient) CreateToolFileUploadURL(_ context.Context, filename, mimetype string) (string, error) { f.calls = append(f.calls, "upload-request") f.filename = filename f.mimetype = mimetype @@ -45,21 +48,29 @@ func (f *fakeFileUploadClient) UploadFileToURL(uploadURL, filePath, filename, mi func (f *fakeFileUploadClient) CreateFileDownloadURL( _ context.Context, - _ string, + transferMethod string, reference, _ *string, forFrontend bool, ) (*FileDownloadResponse, error) { f.calls = append(f.calls, "download-request") f.forFrontend = forFrontend + f.downloadMethod = transferMethod f.downloadRequestCall++ if reference != nil { f.downloadReference = *reference } + if f.downloadErr != nil { + return nil, f.downloadErr + } + downloadURL := "/files/tools/report.pdf?sign=2" + if f.downloadURL != nil { + downloadURL = *f.downloadURL + } return &FileDownloadResponse{ Filename: "report.pdf", MimeType: "application/pdf", Size: 123, - DownloadURL: "/files/tools/report.pdf?sign=2", + DownloadURL: downloadURL, }, nil } @@ -125,43 +136,146 @@ func TestRunFileUploadWithoutDownloadLinkReturnsOnlyCanonicalMapping(t *testing. } } -func TestRunFileUploadDefaultAcceptsLegacyNonemptyReference(t *testing.T) { +func TestRunFileUploadPreservesReferenceWhenPublicURLRequestFails(t *testing.T) { filePath := filepath.Join(t.TempDir(), "report.pdf") if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil { t.Fatalf("write fixture: %v", err) } - client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":"raw-id"}`)} + client := &fakeFileUploadClient{downloadErr: &agentStubHTTPError{ + statusCode: http.StatusUnauthorized, + code: agentStubAuthorizationExpiredCode, + message: "expired", + }} var output bytes.Buffer err := runFileUpload(client, filePath, false, &output) - if err != nil { - t.Fatalf("run file upload: %v", err) + if err == nil { + t.Fatal("run file upload succeeded, want public URL failure") } - if client.downloadRequestCall != 1 || client.downloadReference != "raw-id" { - t.Fatalf("download request = (%d, %q), want legacy reference", client.downloadRequestCall, client.downloadReference) + const reference = "dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ==" + if got, want := strings.TrimSpace(output.String()), `{"transfer_method":"tool_file","reference":"`+reference+`"}`; got != want { + t.Fatalf("partial output = %s, want %s", got, want) } - if !strings.Contains(output.String(), `"reference":"raw-id"`) { - t.Fatalf("output = %q, want legacy reference", output.String()) + for _, want := range []string{ + "request public download URL", + "expired after 5 minutes", + "will not refresh automatically", + "start a new shell tool call", + "retry the command", + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error = %q, want substring %q", err, want) + } + } + recoveryCommand := "dify-agent file public-url '" + reference + "'" + if !strings.HasSuffix(err.Error(), "retry without uploading again: "+recoveryCommand) { + t.Fatalf("error = %q, want exact recovery command %q", err, recoveryCommand) + } + if got, want := strings.Join(client.calls, ","), "upload-request,multipart-upload,download-request"; got != want { + t.Fatalf("call order = %s, want %s", got, want) } } -func TestRunFileUploadWithoutDownloadLinkRejectsNonCanonicalReference(t *testing.T) { +func TestRunFileUploadPreservesReferenceWhenPublicURLResponseIsIncomplete(t *testing.T) { filePath := filepath.Join(t.TempDir(), "report.pdf") if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil { t.Fatalf("write fixture: %v", err) } - client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":"raw-id"}`)} + emptyURL := "" + client := &fakeFileUploadClient{downloadURL: &emptyURL} var output bytes.Buffer - err := runFileUpload(client, filePath, true, &output) - if err == nil || !strings.Contains(err.Error(), "invalid reference") { - t.Fatalf("error = %v, want invalid reference", err) + err := runFileUpload(client, filePath, false, &output) + if err == nil || !strings.Contains(err.Error(), "missing download_url") { + t.Fatalf("error = %v, want incomplete public URL response", err) } - if client.downloadRequestCall != 0 { - t.Fatalf("download request calls = %d, want 0", client.downloadRequestCall) + if !strings.Contains(output.String(), `"reference":"dify-file-ref:`) { + t.Fatalf("partial output = %q, want uploaded reference", output.String()) } - if output.Len() != 0 { - t.Fatalf("output = %q, want empty", output.String()) + if !strings.Contains(err.Error(), "dify-agent file public-url") { + t.Fatalf("error = %q, want recovery command", err) + } +} + +func TestRunFilePublicURLUsesExistingReferenceWithoutUploading(t *testing.T) { + client := &fakeFileUploadClient{} + var output bytes.Buffer + const reference = "dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ==" + + if err := runFilePublicURL(client, reference, &output); err != nil { + t.Fatalf("run file public URL: %v", err) + } + + if got, want := strings.Join(client.calls, ","), "download-request"; got != want { + t.Fatalf("calls = %s, want %s", got, want) + } + if !client.forFrontend || client.downloadMethod != "tool_file" || client.downloadReference != reference { + t.Fatalf( + "download request = (method=%q, forFrontend=%t, reference=%q)", + client.downloadMethod, + client.forFrontend, + client.downloadReference, + ) + } + want := `{"transfer_method":"tool_file","reference":"` + reference + `","public_download_url":"/files/tools/report.pdf?sign=2"}` + if got := strings.TrimSpace(output.String()); got != want { + t.Fatalf("output = %s, want %s", got, want) + } +} + +func TestRunFilePublicURLExplainsExpiredAuthorization(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/agent-stub/files/download-request" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`)) + })) + defer server.Close() + + err := RunFilePublicURL( + &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}, + "dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ==", + ) + if err == nil { + t.Fatal("RunFilePublicURL succeeded, want expired authorization failure") + } + for _, want := range []string{ + "request public download URL", + "expired after 5 minutes", + "will not refresh automatically", + "start a new shell tool call", + "retry the command", + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error = %q, want substring %q", err, want) + } + } +} + +func TestRunFileUploadRejectsNonCanonicalReferenceInBothModes(t *testing.T) { + filePath := filepath.Join(t.TempDir(), "report.pdf") + if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil { + t.Fatalf("write fixture: %v", err) + } + + for _, noDownloadLink := range []bool{false, true} { + client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":"raw-id"}`)} + var output bytes.Buffer + err := runFileUpload(client, filePath, noDownloadLink, &output) + if err == nil || !strings.Contains(err.Error(), "invalid reference") { + t.Fatalf("noDownloadLink=%t error = %v, want invalid reference", noDownloadLink, err) + } + if client.downloadRequestCall != 0 || output.Len() != 0 { + t.Fatalf( + "noDownloadLink=%t download calls = %d, output = %q; want no download or output", + noDownloadLink, + client.downloadRequestCall, + output.String(), + ) + } } } diff --git a/dify-agent-runtime/internal/agentcli/httpclient.go b/dify-agent-runtime/internal/agentcli/httpclient.go index 0c2c156c16b..93cd0d073e2 100644 --- a/dify-agent-runtime/internal/agentcli/httpclient.go +++ b/dify-agent-runtime/internal/agentcli/httpclient.go @@ -25,6 +25,24 @@ type HTTPClient struct { var errUploadRequestAborted = errors.New("upload request aborted") +const agentStubAuthorizationExpiredCode = "agent_stub_authorization_expired" + +type agentStubHTTPError struct { + statusCode int + code string + message string +} + +func (e *agentStubHTTPError) Error() string { + if e.code == agentStubAuthorizationExpiredCode { + return fmt.Sprintf( + "HTTP %d: Agent Stub authorization expired after 5 minutes; the authorization in this process will not refresh automatically; start a new shell tool call and retry the command", + e.statusCode, + ) + } + return fmt.Sprintf("HTTP %d: %s", e.statusCode, e.message) +} + // NewHTTPClient creates a new HTTP client for the Agent Stub API. func NewHTTPClient(env *Environment) *HTTPClient { return &HTTPClient{ @@ -297,3 +315,37 @@ func checkHTTPError(body []byte, statusCode int, operation string) error { } return fmt.Errorf("agent stub %s failed (HTTP %d): %s", operation, statusCode, string(body)) } + +// checkAgentStubHTTPError decodes structured errors for Agent-visible +// connect, file, and config commands. Drive retains its legacy error contract. +func checkAgentStubHTTPError(body []byte, statusCode int) error { + if statusCode < 400 { + return nil + } + + message := string(body) + code := "" + var response struct { + Detail json.RawMessage `json:"detail"` + } + if json.Unmarshal(body, &response) == nil && len(response.Detail) > 0 { + var detailMessage string + if json.Unmarshal(response.Detail, &detailMessage) == nil { + message = detailMessage + } else { + var detail struct { + Code string `json:"code"` + Message string `json:"message"` + } + if json.Unmarshal(response.Detail, &detail) == nil { + code = detail.Code + if detail.Message != "" { + message = detail.Message + } else { + message = string(response.Detail) + } + } + } + } + return &agentStubHTTPError{statusCode: statusCode, code: code, message: message} +} diff --git a/dify-agent-runtime/internal/agentcli/httpclient_test.go b/dify-agent-runtime/internal/agentcli/httpclient_test.go index 67ac9dbd11e..33d558f5378 100644 --- a/dify-agent-runtime/internal/agentcli/httpclient_test.go +++ b/dify-agent-runtime/internal/agentcli/httpclient_test.go @@ -2,6 +2,7 @@ package agentcli import ( "bytes" + "context" "errors" "fmt" "io" @@ -418,6 +419,88 @@ func TestUploadFileRejectsOversizedResponse(t *testing.T) { } } +func TestCheckAgentStubHTTPErrorParsesSupportedResponseShapes(t *testing.T) { + tests := []struct { + name string + body string + wantMessage string + }{ + { + name: "string detail", + body: `{"detail":"invalid request"}`, + wantMessage: "invalid request", + }, + { + name: "non JSON body", + body: "gateway unavailable", + wantMessage: "gateway unavailable", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := checkAgentStubHTTPError([]byte(test.body), http.StatusUnauthorized) + var httpErr *agentStubHTTPError + if !errors.As(err, &httpErr) { + t.Fatalf("error = %v, want *agentStubHTTPError", err) + } + if httpErr.statusCode != http.StatusUnauthorized || httpErr.code != "" { + t.Fatalf("parsed error = %#v", httpErr) + } + if !strings.Contains(err.Error(), test.wantMessage) { + t.Fatalf("error = %q, want substring %q", err, test.wantMessage) + } + }) + } +} + +func TestAgentStubAuthorizationExpiryErrorExplainsHowToRefresh(t *testing.T) { + err := checkAgentStubHTTPError( + []byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`), + http.StatusUnauthorized, + ) + var httpErr *agentStubHTTPError + if !errors.As(err, &httpErr) { + t.Fatalf("error = %v, want *agentStubHTTPError", err) + } + if httpErr.statusCode != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", httpErr.statusCode, http.StatusUnauthorized) + } + if httpErr.code != agentStubAuthorizationExpiredCode || httpErr.message != "expired" { + t.Fatalf("parsed error = %#v", httpErr) + } + + for _, want := range []string{ + "expired after 5 minutes", + "will not refresh automatically", + "start a new shell tool call", + "retry the command", + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error = %q, want substring %q", err, want) + } + } +} + +func TestToolFileUploadURLAloneOptsIntoStructuredExpiration(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("expose_expiration") != "true" { + t.Errorf("expose_expiration = %q, want true", r.URL.Query().Get("expose_expiration")) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`)) + })) + defer server.Close() + + client := newHTTPStubClient(&Environment{URL: server.URL, AuthJWE: "token"}) + _, toolFileErr := client.CreateToolFileUploadURL(context.Background(), "report.pdf", "application/pdf") + + if toolFileErr == nil || !strings.Contains(toolFileErr.Error(), "start a new shell tool call") { + t.Fatalf("ToolFile error = %v, want structured expiry guidance", toolFileErr) + } +} + func TestUploadFileReturnsNonSuccessStatus(t *testing.T) { filePath := filepath.Join(t.TempDir(), "payload.txt") if err := os.WriteFile(filePath, []byte("payload"), 0o600); err != nil { diff --git a/dify-agent/.example.env b/dify-agent/.example.env index bc7a12262df..44c20ad0432 100644 --- a/dify-agent/.example.env +++ b/dify-agent/.example.env @@ -46,7 +46,6 @@ DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox # Binding resources pause; temporary Home initialization resources are killed. # This is not a retention TTL for paused resources or immutable snapshots. DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 -DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= DIFY_AGENT_E2B_SHELLCTL_PORT=5004 # JSON array of regex patterns to redact from shell output shown to the agent. DIFY_AGENT_SHELL_REDACT_PATTERNS= @@ -59,6 +58,8 @@ DIFY_AGENT_STUB_API_BASE_URL=http://localhost:5050/agent-stub # Dify API base URL reachable from the Sandbox for the signed /files/* data plane, # including Config file and skill pulls. DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://localhost:5001 +# Maximum Agent Stub upload size in MiB; forwarded to Dify API as a signed byte limit. +DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT=50 # Server-wide root secret used to derive Agent Stub JWE keys. # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. diff --git a/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md b/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md index bbc9bd2fc31..e4a3593d986 100644 --- a/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md +++ b/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md @@ -76,8 +76,9 @@ current run. Callers control whether each layer is suspended or deleted through `CreateRunRequest.on_exit`. Exit signals control the **layer lifecycle state**, not the execution state of an -`agent run`. The default policy is `suspend`, so a successful `agent run` returns -a reusable `session_snapshot`. +`agent run`. The default policy is `suspend`, so any run that enters and exits its +compositor context can return a reusable `session_snapshot`, including failed or +cancelled runs. Failures or cancellations before entry have no new snapshot. ### Default: suspend layers diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md index 348d3ec97d6..68a69772984 100644 --- a/dify-agent/docs/dify-agent/guide/index.md +++ b/dify-agent/docs/dify-agent/guide/index.md @@ -54,11 +54,11 @@ also reads `.env` and `dify-agent/.env` when present. | `DIFY_AGENT_E2B_API_KEY` | empty | E2B API key; required for E2B. | | `DIFY_AGENT_E2B_TEMPLATE` | `difys-default-team/dify-agent-local-sandbox` | Prepared E2B template containing shellctl and the deployment-default Home environment. | | `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` | `3600` | Maximum continuous active time for the RuntimeLease spanning one complete Agent run. Its default intentionally matches `DIFY_AGENT_RUN_TIMEOUT_SECONDS`, but the settings are independently configurable. Binding resources pause on timeout; this setting does not own the run terminal state and is not a retention TTL. | -| `DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token expected by shellctl inside the E2B template. | | `DIFY_AGENT_E2B_SHELLCTL_PORT` | `5004` | shellctl port exposed by the E2B template. | | `DIFY_AGENT_SHELL_REDACT_PATTERNS` | empty | JSON array of additional regex patterns redacted from Shell output. | | `DIFY_AGENT_STUB_API_BASE_URL` | empty | HTTP(S) Agent Stub API base URL reachable from the Sandbox. It may be the service root or `/agent-stub`. Enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. | | `DIFY_AGENT_SANDBOX_FILES_BASE_URL` | empty | Dify API base URL reachable from the Sandbox for signed `/files/*` upload/download bytes, including Config file and skill pulls. Required when Agent Stub file operations are enabled. May include an ingress path prefix, but not a query or fragment. | +| `DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT` | `50` | Agent service-owned maximum Agent Stub upload size in MiB. The file-request handler factory converts it to bytes and sends it to Dify API as the required `max_size` used to sign a size-limited upload URL. | | `DIFY_AGENT_SERVER_SECRET_KEY` | empty | Security-sensitive server-wide root secret used to derive the JWE encryption key for Agent Stub bearer tokens; required when `DIFY_AGENT_STUB_API_BASE_URL` is set. The supplied default config uses a development value; set a unique unpadded base64url 32-byte secret in production. | | `DIFY_AGENT_OUTBOUND_HTTP_CONNECT_TIMEOUT` | `10` | Shared outbound HTTP connect timeout in seconds. | | `DIFY_AGENT_OUTBOUND_HTTP_READ_TIMEOUT` | `600` | Shared outbound HTTP read timeout in seconds. | @@ -90,6 +90,7 @@ DIFY_AGENT_LOCAL_SANDBOX_WORKSPACE_ROOT=/tmp/dify-agent/workspaces DIFY_AGENT_LOCAL_SANDBOX_HOME_SNAPSHOT_ROOT=/tmp/dify-agent/home-snapshots DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub DIFY_AGENT_SANDBOX_FILES_BASE_URL=https://dify.example.com +DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT=50 # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. # Generate one with: python -c 'import secrets; print(secrets.token_urlsafe(32))' @@ -111,10 +112,11 @@ Removing Agent Stub gRPC is a breaking transport migration: replace every For a remote Sandbox, expose only `/agent-stub/*` from Agent Backend and the existing `/files/*` Dify API data plane. The `/files/*` ingress must preserve -the complete signed query string, allow the configured upload body size, and -use response streaming and timeouts suitable for large downloads. Do not expose -Agent Backend `/runs`, Workspace, or Binding management routes through the -Sandbox ingress. +the complete signed query string and set its request-body limit above the +configured file-size limit to leave room for multipart framing and headers; the +two limits need not be numerically equal. Use response streaming and timeouts +suitable for large downloads. Do not expose Agent Backend `/runs`, Workspace, +or Binding management routes through the Sandbox ingress. Browser presentation URLs are independent. Configure Dify API `FILES_URL` to a browser-reachable public origin, or leave it empty so responses use same-origin @@ -280,26 +282,29 @@ run as failed with `error_type: "agent_run_limit_exceeded"`. During FastAPI shutdown the scheduler rejects new runs, waits up to `DIFY_AGENT_SHUTDOWN_GRACE_SECONDS` for active tasks, then cancels remaining tasks -and attempts to finalize them as failed. Success, failure, cancellation, and this -shutdown path all use one atomic Redis transition: only the first transition from -`running` appends a terminal event and updates the run record. A later terminal -attempt leaves both the record and event stream unchanged. A hard process crash -can still leave active runs stuck as `running`; there is no in-service recovery -or worker handoff. +and attempts to finalize them as failed. Success and failure use an atomic Redis +transition. Cancellation first atomically records a private intent; after the +owner exits the runner, a second atomic transition appends `run_cancelled`, +updates the run record, and deletes the intent. The first accepted success, +failure, or cancellation intent wins. A hard process crash can still leave +active runs, including runs with accepted cancellation intent, stuck as +`running`; there is no in-service recovery or worker handoff. Horizontal scaling is possible by running multiple API processes against the same Redis prefix, but each process executes only the runs it accepted. Redis provides shared status/event visibility, not load balancing or queued-job recovery. The cancel endpoint can atomically accept a running run on any process. The process -that owns the runner observes the shared `run_cancelled` event, then cancels and -cleans up its local task. The HTTP response confirms that logical cancellation is -durable; local runner cleanup may still be in progress. Retrying a cancellation -after the run is already `cancelled` is idempotent. +that owns the runner observes the private cancellation-intent stream, cancels +and cleans up its local task, and only then emits `run_cancelled`. The HTTP +response confirms that cancellation intent is durable; `GET /runs/{run_id}` may +still report `running` until cleanup finishes. Retrying an accepted or completed +cancellation is idempotent. Atomic terminal finalization currently assumes the configured Redis URL targets -one Redis deployment that can execute both run keys in a Lua script. The existing -record and event key names are unchanged and do not contain a shared Redis -Cluster hash tag, so Redis Cluster is not supported for this transition. During +one Redis deployment that can execute all run-coordination keys in a Lua script. +The record and event key names are unchanged, and cancellation adds a private +cancel-intent key. These keys do not contain a shared Redis Cluster hash tag, so +Redis Cluster is not supported for this transition. During a rolling upgrade, older processes can still use the former split event/status writes; treat the single-terminal invariant as active only after those processes have exited. Deploy atomic terminal finalization everywhere first, then ensure @@ -337,8 +342,10 @@ effective prompts are rejected during create-run validation before the run is persisted or scheduled. There is no Pydantic AI history layer. To resume Agenton layer state, pass the -`session_snapshot` from a previous `run_succeeded.data` payload together with a -composition that has the same layer names and order. +`session_snapshot` from a previous terminal event together with a composition +that has the same layer names and order. Success always contains a snapshot. +Failure and cancellation contain one only when compositor entry succeeded and +layer exit completed; otherwise callers should retain their previous snapshot. ## Observing runs @@ -350,8 +357,11 @@ progress: Failed records can also expose a stable machine-readable `error_type` alongside the diagnostic `error` text. - `POST /runs/{run_id}/cancel` atomically accepts cancellation on any API process - and emits `run_cancelled`; it returns `409` only when a success/failure terminal - already won. Runner cleanup continues asynchronously on the owner process. + and returns immediately. `CancelRunResponse.status == "cancelled"` acknowledges + a durably accepted cancellation intent, not completed runner cleanup. Callers + that require cleanup-complete state or its session snapshot must await the + public `run_cancelled` event or use `cancel_run_and_wait`. The endpoint returns + `409` only when a success/failure terminal already won. - `GET /runs/{run_id}/events` polls the Redis Stream event log with `after` and `next_cursor` cursors. - `GET /runs/{run_id}/events/sse` replays and streams events over SSE. The SSE @@ -367,12 +377,15 @@ end with `run_cancelled`. Each run can append at most one of these terminal events. Event envelopes retain `id`, `run_id`, `type`, `data`, and `created_at`; `data` is typed per event type, including Pydantic AI's `AgentStreamEvent` payload for `pydantic_ai_event` and a -terminal `run_succeeded.data` object containing a `CompositorSessionSnapshot` for -resumption. A successful run has exactly one active result branch: JSON-safe +terminal event may contain a `CompositorSessionSnapshot` for resumption. +`run_succeeded` always contains it; `run_failed` and `run_cancelled` contain it +only when compositor entry succeeded, layer exit completed, and a post-exit +snapshot was actually produced. A successful run has exactly one active result branch: JSON-safe `output` for final answers, or `deferred_tool_call` when a layer such as `dify.ask_human` ends the current agent run with an external deferred tool call. Failed event payloads contain the diagnostic `error`, optional source-specific -`reason`, and optional stable `error_type`. Pydantic AI request/step budget +`reason`, optional stable `error_type`, and optional `session_snapshot`. +Cancelled payloads likewise may contain `session_snapshot`. Pydantic AI request/step budget exhaustion enforced by Dify Agent is reported as `error_type: "agent_run_limit_exceeded"`; consumers should branch on that value rather than parsing the error text. The Dify Agent-owned wall-clock run deadline diff --git a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md index e962a23f945..9bda3401d58 100644 --- a/dify-agent/docs/dify-agent/user-manual/history-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/history-layer/index.md @@ -34,6 +34,24 @@ history_layer = RunLayerSpec( Include this layer in the same composition as your prompt, plugin, and LLM layers. +## Compaction and persistence + +When the LLM layer supplies `context_window_tokens`, Dify Agent sets the Harness +target to `min(floor(window * 0.8), window - max_tokens)` for a positive +`model_settings.max_tokens`; otherwise it uses `floor(window * 0.8)`. A target +that is not positive rejects the run before model invocation. + +Harness estimates and, when needed, rewrites history immediately before model +requests. It clears older tool results first, retaining the latest three +tool-call/result pairs and their inputs. If the history is still over target, the +same current model incrementally summarizes older messages while retaining the +latest twenty messages and the first user message. + +With a history layer, a successful run replaces its stored messages with the +rewritten complete history in the returned session snapshot. Without this layer, +compaction affects only the current run. Failed runs do not write a resumable +success snapshot, so their history rewrites do not persist across runs. + ## Resume a conversation Successful runs return a terminal event with both final output and a resumable @@ -65,12 +83,13 @@ terminal snapshot resumable. Keep that default for normal memory flows. Dify Agent handles memory conservatively: -1. Current system prompts are rendered into temporary `message_history` before - stored history. -2. Stored history is then sent to the model. -3. Current user prompts are sent after the stored history. -4. Only newly produced pydantic-ai messages are appended after a successful run. -5. Current system prompts are not persisted into the history layer. +1. Current system prompts are passed as run-level pydantic-ai instructions. +2. Stored history is sent to the model before the current user prompt. +3. When the LLM layer includes `context_window_tokens`, Harness may rewrite + over-target history immediately before a model request as described above. +4. After a successful run, the complete possibly compacted history is written + back to the layer. +5. Run-level system instructions are removed before history is persisted. 6. Failed runs emit `run_failed` and do not return a success snapshot to resume. ## Persist snapshots outside the client process diff --git a/dify-agent/docs/dify-agent/user-manual/plugin-llm-layer/index.md b/dify-agent/docs/dify-agent/user-manual/plugin-llm-layer/index.md index 92c0923af8a..4fe7f4c0849 100644 --- a/dify-agent/docs/dify-agent/user-manual/plugin-llm-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/plugin-llm-layer/index.md @@ -15,6 +15,7 @@ because that layer supplies the caller identity required by the API gateway. | `model_provider` | `str` | Provider name inside `plugin_id`. Use the value of `DIFY_AGENT_PROVIDER` from `dify-agent/.env`. | | `model` | `str` | Model name. Use the value of `DIFY_AGENT_MODEL_NAME` from `dify-agent/.env`. | | `model_settings` | `ModelSettings \| None` | Optional pydantic-ai model settings. | +| `context_window_tokens` | `int \| None` | Positive effective context-window capability metadata. Enables window-based compaction when present; omission disables it. | The plugin LLM layer type id is `dify.plugin.llm`. @@ -48,6 +49,32 @@ dependency field named `execution_context` to the composition layer named Set `MODEL_PROVIDER` and `MODEL_NAME` to the same values as `DIFY_AGENT_PROVIDER` and `DIFY_AGENT_MODEL_NAME` in `dify-agent/.env`. +## Context compaction + +Dify product request builders resolve `context_window_tokens` from the selected +model plugin schema using the current tenant and user credentials. A client that +constructs `DifyPluginLLMLayerConfig` directly is responsible for supplying an +accurate positive value. The field is model capability metadata: Dify Agent does +not forward it as a Provider parameter or merge it into `model_settings`. + +For a known window, Dify Agent computes the Harness compaction target as: + +```text +min(floor(context_window_tokens * 0.8), context_window_tokens - max_tokens) +``` + +The second term applies only when `model_settings.max_tokens` is positive. A +non-positive target rejects the run before model invocation. Immediately before +model requests, Harness estimates the message history and rewrites it when it is +over target: it first clears old tool results while retaining the three most +recent tool-call/result pairs and their inputs; if still over target, the current +model incrementally summarizes older history while retaining the latest twenty +messages and the first user message. + +Compaction affects later runs only when the composition has a +[history layer](../history-layer/index.md) and a successful run writes the +rewritten history into its session snapshot. + ## Complete minimal model composition Most runs include a prompt, execution-context layer, and LLM layer: @@ -106,3 +133,5 @@ composition = RunComposition( calls. The shared execution-context layer carries the Dify caller context. - Model credentials are never accepted from the Agent request. Dify API resolves the tenant's current provider configuration and owns quota accounting. +- Omitting `context_window_tokens` disables window-based compaction. It does not + limit or otherwise change the Provider's own context-window enforcement. diff --git a/dify-agent/docs/dify-agent/user-manual/prompt-layer/index.md b/dify-agent/docs/dify-agent/user-manual/prompt-layer/index.md index 377ac368e74..f31b2e41fca 100644 --- a/dify-agent/docs/dify-agent/user-manual/prompt-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/prompt-layer/index.md @@ -68,5 +68,4 @@ prompt_layer = RunLayerSpec( - Prompt layer names are not reserved by the runtime, but `prompt` is the recommended conventional name. - When a [history layer](../history-layer/index.md) is present, current system - prompts are sent as a temporary prefix before stored history and are not saved - into memory. + prompts are passed as run-level instructions and are not saved into memory. diff --git a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md index 210e59f43dd..9520b6e29cd 100644 --- a/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md +++ b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md @@ -77,8 +77,11 @@ DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN=replace-with-shellctl-token # DIFY_AGENT_LOCAL_SANDBOX_HOME_SNAPSHOT_ROOT=/tmp/dify-agent/home-snapshots ``` -The auth token may be empty when shellctl authentication is disabled. E2B uses -`DIFY_AGENT_E2B_API_KEY`, the prepared template, and its shellctl settings. +The auth token may be empty when shellctl authentication is disabled. Dify-created +E2B Sandboxes disable public traffic at creation and access shellctl through the +E2B port proxy with its `traffic_access_token`. Acquiring a RuntimeLease fails if +E2B does not provide a non-empty token. This policy applies only to newly created +Sandboxes and does not retrofit existing ones. To let shell jobs call the Agent Stub with `dify-agent ...`, configure a Sandbox-reachable Agent Stub URL and a unique production secret. Remote @@ -90,6 +93,7 @@ topology. ```env DIFY_AGENT_STUB_API_BASE_URL=https://agent.example.com/agent-stub DIFY_AGENT_SANDBOX_FILES_BASE_URL=https://dify.example.com +DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT=50 DIFY_AGENT_SERVER_SECRET_KEY=replace-with-unpadded-base64url-for-32-random-bytes ``` @@ -97,6 +101,11 @@ HTTP URLs may be either the service root or the explicit `/agent-stub` root. The server normalizes a service root and rejects unrelated paths. The separate Sandbox file base must point to the Dify API ingress serving `/files/*`; it is used for CLI upload/download bytes, including Config file and skill pulls. +`DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT` belongs to the Agent service, is in +MiB, and defaults to `50`. The Agent service converts it to the signed upload +URL's byte limit. Any ingress or proxy in front of `/files/*` must allow that +file limit plus multipart framing and header overhead; its request-body limit +does not need to be numerically identical. After `dify-agent file upload ` succeeds, the CLI prints JSON such as: @@ -121,6 +130,22 @@ the same streaming ToolFile upload but skips the download-request step and prints only `transfer_method` plus the canonical `reference`. The regular `file upload` command keeps the link-producing behavior shown above. +If the upload succeeds but creating the public URL fails, the command still +prints the canonical mapping and exits with an error containing an exact +`dify-agent file public-url ` retry command. Run that command from a +new shell tool call to create the public URL without uploading the file again. +On success, `file public-url` prints the complete JSON mapping containing +`transfer_method`, `reference`, and `public_download_url`. + +The Agent Stub authorization injected into a shell job is valid for five +minutes. It does not refresh inside an already-running process. If a command +reports that the authorization expired, start a new shell tool call and retry +the command (or its reported `file public-url` recovery command). + +The injected JWE is masked as `***` in model-facing `shell_run`, `shell_wait`, +and `shell_input` observations. Raw shellctl output and files referenced by +`output_path` remain unchanged. + ## Request graph A shell-enabled run contains Execution Context, Runtime, and Shell layers: diff --git a/dify-agent/examples/agenton/agenton_examples/pydantic_ai_bridge.py b/dify-agent/examples/agenton/agenton_examples/pydantic_ai_bridge.py index 0b6ff588d11..490e54100b9 100644 --- a/dify-agent/examples/agenton/agenton_examples/pydantic_ai_bridge.py +++ b/dify-agent/examples/agenton/agenton_examples/pydantic_ai_bridge.py @@ -10,12 +10,10 @@ the repository root with: from __future__ import annotations import asyncio -import json import os from dataclasses import dataclass from pydantic_ai import Agent, RunContext -from pydantic_ai.messages import BuiltinToolCallPart, ModelMessage, ToolCallPart from pydantic_ai.models.openai import OpenAIChatModel # pyright: ignore[reportDeprecated] from pydantic_ai.models.test import TestModel @@ -113,20 +111,8 @@ async def main() -> None: bridge_layer = run.get_layer("pydantic_ai_bridge", PydanticAIBridgeLayer) result = await agent.run(run.user_prompts, deps=bridge_layer.run_deps) - for line in _format_messages(result.all_messages()): - print(line) - - -def _format_messages(messages: list[ModelMessage]) -> list[str]: - lines: list[str] = [] - for message in messages: - for part in message.parts: - if isinstance(part, ToolCallPart | BuiltinToolCallPart): - args = json.dumps(part.args, ensure_ascii=False) - lines.append(f"{type(part).__name__}: {part.tool_name}({args})") - else: - lines.append(f"{type(part).__name__}: {part.content}") - return lines + for message in result.all_messages(): + print(message) if __name__ == "__main__": diff --git a/dify-agent/examples/dify_agent/dify_agent_examples/run_pydantic_ai_agent.py b/dify-agent/examples/dify_agent/dify_agent_examples/run_pydantic_ai_agent.py index 6c50ca6e2d0..471f98f8563 100644 --- a/dify-agent/examples/dify_agent/dify_agent_examples/run_pydantic_ai_agent.py +++ b/dify-agent/examples/dify_agent/dify_agent_examples/run_pydantic_ai_agent.py @@ -78,7 +78,7 @@ async def main() -> None: async with agent.run_stream("Explain the theory of relativity") as run: async for piece in run.stream_output(): print(piece, end="", flush=True) - print(run.usage()) + print(run.usage) if __name__ == "__main__": diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index 73bfcfd9e0f..89e1fe776c2 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -8,7 +8,8 @@ dependencies = [ "httpx==0.28.1", "httpx2>=2.5.0,<3.0.0", "pydantic>=2.12.5,<2.13", - "pydantic-ai-slim>=1.106.0,<2.0.0", + "pydantic-ai-harness>=0.20.0,<0.21.0", + "pydantic-ai-slim>=2.30.0,<3.0.0", "typing-extensions>=4.12.2,<5.0.0", ] @@ -24,7 +25,7 @@ server = [ "jsonschema>=4.23.0,<5.0.0", "jwcrypto>=1.5.6,<2", "logfire[fastapi,httpx,redis]>=4.37.0,<5.0.0", - "pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0", + "pydantic-ai-slim[anthropic,google,openai]>=2.30.0,<3.0.0", "pydantic-settings>=2.12.0,<3.0.0", "redis>=7.4.0,<8.0.0", "uvicorn[standard]==0.46.0", diff --git a/dify-agent/src/dify_agent/adapters/llm/model.py b/dify-agent/src/dify_agent/adapters/llm/model.py index d9a17af2379..1c77e6c8ded 100644 --- a/dify-agent/src/dify_agent/adapters/llm/model.py +++ b/dify-agent/src/dify_agent/adapters/llm/model.py @@ -51,11 +51,13 @@ from pydantic_ai.messages import ( ModelResponseStreamEvent, MultiModalContent, RetryPromptPart, + SpeechPart, SystemPromptPart, TextContent, TextPart, ThinkingPart, ToolCallPart, + ToolAvailabilityDeltaPart, ToolReturnPart, UploadedFile, UserContent, @@ -333,6 +335,8 @@ def _map_model_request_to_prompt_messages(message: ModelRequest) -> list[PromptM name=part.tool_name, ) ) + elif isinstance(part, SpeechPart | ToolAvailabilityDeltaPart): + raise UnexpectedModelBehavior(f"Unsupported request part for daemon adapter: {type(part).__name__}") else: assert_never(part) diff --git a/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py b/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py index 636a88247d1..ff7c409b8c2 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py +++ b/dify-agent/src/dify_agent/agent_stub/server/agent_stub_files.py @@ -103,7 +103,8 @@ class DifyApiAgentStubFileRequestHandler: The upload path calls ``/inner/api/agent/files/upload-request`` and injects the authenticated execution context's ``tenant_id``, ``user_id``, ``user_from``, and optional - ``conversation_id`` along with the requested filename and mimetype. The download path calls + ``conversation_id`` along with the requested filename, mimetype, and configured upload-size + limit. The download path calls ``/inner/api/agent/files/download-request`` and injects ``tenant_id``, ``user_id``, ``user_from``, and ``invoke_from`` plus the validated public file mapping. @@ -119,6 +120,7 @@ class DifyApiAgentStubFileRequestHandler: inner_api_url: str inner_api_key: str sandbox_files_base_url: str + max_upload_size_bytes: int timeout: httpx.Timeout | float = 30.0 async def create_upload_request( @@ -147,6 +149,7 @@ class DifyApiAgentStubFileRequestHandler: "filename": request.filename, "mimetype": request.mimetype, "conversation_id": execution_context.conversation_id, + "max_size": self.max_upload_size_bytes, } data = await self._post_inner_api("/inner/api/agent/files/upload-request", payload) upload_uri = data.get("upload_uri") diff --git a/dify-agent/src/dify_agent/agent_stub/server/control_plane.py b/dify-agent/src/dify_agent/agent_stub/server/control_plane.py index 05eba042288..badf83dc41c 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/control_plane.py +++ b/dify-agent/src/dify_agent/agent_stub/server/control_plane.py @@ -26,7 +26,18 @@ from dify_agent.agent_stub.protocol.agent_stub import ( from dify_agent.agent_stub.server.agent_stub_config import AgentStubConfigRequestError, AgentStubConfigRequestHandler from dify_agent.agent_stub.server.agent_stub_drive import AgentStubDriveRequestError, AgentStubDriveRequestHandler from dify_agent.agent_stub.server.agent_stub_files import AgentStubFileRequestError, AgentStubFileRequestHandler -from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubPrincipal, AgentStubTokenCodec, AgentStubTokenError +from dify_agent.agent_stub.server.tokens.agent_stub import ( + AgentStubPrincipal, + AgentStubTokenCodec, + AgentStubTokenError, + AgentStubTokenExpiredError, +) + + +_AGENT_STUB_AUTHORIZATION_EXPIRED_DETAIL = { + "code": "agent_stub_authorization_expired", + "message": "Agent Stub authorization expired after 5 minutes; start a new shell tool call and retry the command.", +} class AgentStubControlPlaneError(RuntimeError): @@ -65,7 +76,7 @@ class AgentStubControlPlaneService: async def connect(self, *, authorization: str | None) -> AgentStubConnectResponse: """Authenticate and handle one connect request.""" - _ = self._authenticate(authorization) + _ = self._authenticate(authorization, expose_expiration=True) return AgentStubConnectResponse(connection_id=self.connection_id_factory(), status="connected") async def create_file_upload_request( @@ -73,9 +84,10 @@ class AgentStubControlPlaneService: *, request: AgentStubFileUploadRequest, authorization: str | None, + expose_expiration: bool = False, ) -> AgentStubFileUploadResponse: """Authenticate and delegate one already-validated file-upload request.""" - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=expose_expiration) handler = self._require_file_request_handler() try: return await handler.create_upload_request(principal=principal, request=request) @@ -89,7 +101,7 @@ class AgentStubControlPlaneService: authorization: str | None, ) -> AgentStubFileDownloadResponse: """Authenticate and delegate one already-validated file-download request.""" - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) if request.config is not None: handler = self._require_config_request_handler() try: @@ -127,7 +139,7 @@ class AgentStubControlPlaneService: *, authorization: str | None, ) -> AgentStubConfigManifestResponse: - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) handler = self._require_config_request_handler() try: return await handler.manifest(principal=principal) @@ -140,7 +152,7 @@ class AgentStubControlPlaneService: name: str, authorization: str | None, ) -> dict[str, object]: - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) handler = self._require_config_request_handler() try: return await handler.inspect_skill(principal=principal, name=name) @@ -153,7 +165,7 @@ class AgentStubControlPlaneService: request: AgentStubConfigPushRequest, authorization: str | None, ) -> AgentStubConfigPushResponse: - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) handler = self._require_config_request_handler() try: return await handler.push(principal=principal, request=request) @@ -166,7 +178,7 @@ class AgentStubControlPlaneService: env_text: str, authorization: str | None, ) -> dict[str, object]: - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) handler = self._require_config_request_handler() try: return await handler.update_env(principal=principal, env_text=env_text) @@ -179,7 +191,7 @@ class AgentStubControlPlaneService: note: str, authorization: str | None, ) -> dict[str, object]: - principal = self._authenticate(authorization) + principal = self._authenticate(authorization, expose_expiration=True) handler = self._require_config_request_handler() try: return await handler.update_note(principal=principal, note=note) @@ -200,12 +212,19 @@ class AgentStubControlPlaneService: except AgentStubDriveRequestError as exc: raise AgentStubControlPlaneError(exc.status_code, exc.detail) from exc - def _authenticate(self, authorization: str | None) -> AgentStubPrincipal: + def _authenticate(self, authorization: str | None, *, expose_expiration: bool = False) -> AgentStubPrincipal: token_codec = self.token_codec if token_codec is None: raise AgentStubConfigurationError(503, "Agent Stub is not configured") try: return token_codec.decode_authorization_header(authorization) + except AgentStubTokenExpiredError as exc: + detail = ( + _AGENT_STUB_AUTHORIZATION_EXPIRED_DETAIL + if expose_expiration + else "invalid or missing Agent Stub authorization" + ) + raise AgentStubAuthenticationError(401, detail) from exc except AgentStubTokenError as exc: raise AgentStubAuthenticationError(401, "invalid or missing Agent Stub authorization") from exc diff --git a/dify-agent/src/dify_agent/agent_stub/server/routes/agent_stub.py b/dify-agent/src/dify_agent/agent_stub/server/routes/agent_stub.py index 22e8fa9f94f..5bcaa978b6a 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/routes/agent_stub.py +++ b/dify-agent/src/dify_agent/agent_stub/server/routes/agent_stub.py @@ -62,9 +62,14 @@ def create_agent_stub_http_router( async def create_file_upload_request( request: AgentStubFileUploadRequest, authorization: str | None = Header(default=None, alias="Authorization"), + expose_expiration: bool = False, ) -> AgentStubFileUploadResponse: try: - return await service.create_file_upload_request(request=request, authorization=authorization) + return await service.create_file_upload_request( + request=request, + authorization=authorization, + expose_expiration=expose_expiration, + ) except AgentStubControlPlaneError as exc: raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc diff --git a/dify-agent/src/dify_agent/agent_stub/server/tokens/__init__.py b/dify-agent/src/dify_agent/agent_stub/server/tokens/__init__.py index 278b3aaccb9..ca261562e76 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/tokens/__init__.py +++ b/dify-agent/src/dify_agent/agent_stub/server/tokens/__init__.py @@ -9,6 +9,7 @@ from dify_agent.agent_stub.server.tokens.agent_stub import ( AgentStubTokenClaims, AgentStubTokenCodec, AgentStubTokenError, + AgentStubTokenExpiredError, decode_server_secret_key, derive_agent_stub_jwe_key, ) @@ -22,6 +23,7 @@ __all__ = [ "AgentStubTokenClaims", "AgentStubTokenCodec", "AgentStubTokenError", + "AgentStubTokenExpiredError", "decode_server_secret_key", "derive_agent_stub_jwe_key", ] diff --git a/dify-agent/src/dify_agent/agent_stub/server/tokens/agent_stub.py b/dify-agent/src/dify_agent/agent_stub/server/tokens/agent_stub.py index 7dc87e8d426..466fbe2581e 100644 --- a/dify-agent/src/dify_agent/agent_stub/server/tokens/agent_stub.py +++ b/dify-agent/src/dify_agent/agent_stub/server/tokens/agent_stub.py @@ -30,7 +30,7 @@ from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig AGENT_STUB_TOKEN_ISSUER = "dify-agent-server" AGENT_STUB_TOKEN_AUDIENCE = "dify-agent-agent-stub" AGENT_STUB_TOKEN_SCOPE_CONNECT = "agent_stub:connect" -AGENT_STUB_TOKEN_TTL_SECONDS = 24 * 60 * 60 +AGENT_STUB_TOKEN_TTL_SECONDS = 5 * 60 _AGENT_STUB_JWE_PURPOSE = b"dify-agent:agent-stub:jwe:v1" _REQUIRED_SERVER_SECRET_BYTES = 32 _BASE64URL_TEXT_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$") @@ -41,6 +41,10 @@ class AgentStubTokenError(RuntimeError): """Raised when an Agent Stub bearer token is missing or invalid.""" +class AgentStubTokenExpiredError(AgentStubTokenError): + """Raised when an otherwise valid Agent Stub bearer token has expired.""" + + class AgentStubShellClaims(BaseModel): """Optional shell-session claims embedded in Agent Stub tokens.""" @@ -100,7 +104,7 @@ class AgentStubTokenCodec: session_id: str | None = None, now: int | None = None, ) -> AgentStubTokenClaims: - """Build the fixed-24h claim set for one Agent Stub connection token.""" + """Build the fixed-five-minute claim set for one Agent Stub connection token.""" issued_at = _timestamp(now) shell_claims = AgentStubShellClaims(session_id=session_id) if session_id is not None else None return AgentStubTokenClaims( @@ -122,7 +126,7 @@ class AgentStubTokenCodec: session_id: str | None = None, now: int | None = None, ) -> str: - """Encode one fixed-24h Agent Stub compact JWE bearer token.""" + """Encode one fixed-five-minute Agent Stub compact JWE bearer token.""" return self.encode_claims(self.build_connection_claims(execution_context, session_id=session_id, now=now)) def encode_claims(self, claims: AgentStubTokenClaims) -> str: @@ -198,12 +202,16 @@ def _validate_claims(claims: AgentStubTokenClaims, *, now: int) -> None: raise AgentStubTokenError(f"Agent Stub bearer token issuer must be {AGENT_STUB_TOKEN_ISSUER!r}") if claims.aud != AGENT_STUB_TOKEN_AUDIENCE: raise AgentStubTokenError(f"Agent Stub bearer token audience must be {AGENT_STUB_TOKEN_AUDIENCE!r}") + if claims.exp <= claims.iat: + raise AgentStubTokenError("Agent Stub bearer token expiration must be after its issue time") + if claims.exp - claims.iat > AGENT_STUB_TOKEN_TTL_SECONDS: + raise AgentStubTokenError("Agent Stub bearer token lifetime exceeds the maximum allowed lifetime") + if AGENT_STUB_TOKEN_SCOPE_CONNECT not in claims.scope: + raise AgentStubTokenError(f"Agent Stub bearer token scope must include {AGENT_STUB_TOKEN_SCOPE_CONNECT!r}") if now < claims.nbf: raise AgentStubTokenError("Agent Stub bearer token is not valid yet") if now >= claims.exp: - raise AgentStubTokenError("Agent Stub bearer token is expired") - if AGENT_STUB_TOKEN_SCOPE_CONNECT not in claims.scope: - raise AgentStubTokenError(f"Agent Stub bearer token scope must include {AGENT_STUB_TOKEN_SCOPE_CONNECT!r}") + raise AgentStubTokenExpiredError("Agent Stub bearer token is expired") def _hkdf_sha256(input_key_material: bytes, *, info: bytes, length: int) -> bytes: @@ -249,6 +257,7 @@ __all__ = [ "AgentStubTokenClaims", "AgentStubTokenCodec", "AgentStubTokenError", + "AgentStubTokenExpiredError", "decode_server_secret_key", "derive_agent_stub_jwe_key", ] diff --git a/dify-agent/src/dify_agent/client/_client.py b/dify-agent/src/dify_agent/client/_client.py index 0cd14a68711..80af37d12fe 100644 --- a/dify-agent/src/dify_agent/client/_client.py +++ b/dify-agent/src/dify_agent/client/_client.py @@ -43,6 +43,7 @@ from dify_agent.protocol import ( DestroyExecutionBindingRequest, HomeSnapshotResponse, RUN_EVENT_ADAPTER, + RunCancelledEvent, RunEvent, RunEventsResponse, RunStatusResponse, @@ -384,8 +385,8 @@ class Client: async def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse: """Request explicit cancellation for ``run_id``. - Acceptance atomically persists the cancelled state. The process executing - the run observes that state and performs runner cleanup asynchronously. + Acceptance atomically persists cancellation intent. The process executing + the run publishes ``run_cancelled`` after runner cleanup completes. """ request_model = request or CancelRunRequest() try: @@ -417,6 +418,44 @@ class Client: raise DifyAgentClientError(f"cancel_run_sync request failed: {exc}") from exc return _parse_model_response(response, CancelRunResponse) + async def cancel_run_and_wait( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Request cancellation and wait for its public terminal event.""" + _ = await self.cancel_run(run_id, request) + resume_after = after + if after is not None and (await self.get_run(run_id)).status == "cancelled": + resume_after = None + async for event in self.stream_events(run_id, after=resume_after): + if isinstance(event, RunCancelledEvent): + return event + if event.type in _TERMINAL_EVENT_TYPES: + raise DifyAgentClientError(f"run {run_id!r} finished with {event.type!r} before cancellation") + raise DifyAgentStreamError(f"run {run_id!r} stream ended before run_cancelled") + + def cancel_run_and_wait_sync( + self, + run_id: str, + request: CancelRunRequest | None = None, + *, + after: str | None = None, + ) -> RunCancelledEvent: + """Synchronous variant of ``cancel_run_and_wait``.""" + _ = self.cancel_run_sync(run_id, request) + resume_after = after + if after is not None and self.get_run_sync(run_id).status == "cancelled": + resume_after = None + for event in self.stream_events_sync(run_id, after=resume_after): + if isinstance(event, RunCancelledEvent): + return event + if event.type in _TERMINAL_EVENT_TYPES: + raise DifyAgentClientError(f"run {run_id!r} finished with {event.type!r} before cancellation") + raise DifyAgentStreamError(f"run {run_id!r} stream ended before run_cancelled") + async def get_run(self, run_id: str) -> RunStatusResponse: """Return the current status for ``run_id`` or raise a mapped client error.""" try: diff --git a/dify-agent/src/dify_agent/layers/_agent_cli_help.json b/dify-agent/src/dify_agent/layers/_agent_cli_help.json index 66578c3d9e0..c8942b2ae11 100644 --- a/dify-agent/src/dify_agent/layers/_agent_cli_help.json +++ b/dify-agent/src/dify_agent/layers/_agent_cli_help.json @@ -19,7 +19,8 @@ "drive list": "List drive files visible to the current sandbox execution.\n\nUsage:\n dify-agent drive list [REMOTE_PREFIX] [flags]\n\nFlags:\n -h, --help help for list\n --json Emit the drive manifest as JSON.", "drive pull": "Pull one or more drive keys/prefixes into one local directory tree.\n\nUsage:\n dify-agent drive pull [REMOTE]... [flags]\n\nFlags:\n -h, --help help for pull\n --json Emit the pull result as JSON.\n --to string Local base directory for pulled drive files.", "drive push": "Upload one local file or directory into the agent drive.\n\nUsage:\n dify-agent drive push LOCAL_PATH REMOTE_PATH [flags]\n\nFlags:\n -h, --help help for push\n --json Accepted for consistency; drive push output is already emitted as JSON.\n --kind string Directory upload kind: skill or dir.", - "file": "Upload or download workflow files through the Agent Stub.\n\nUsage:\n dify-agent file [command]\n\nAvailable Commands:\n download Download one workflow file mapping into the local sandbox directory.\n upload Upload one sandbox-local file as a ToolFile output reference.\n\nFlags:\n -h, --help help for file\n\nUse \"dify-agent file [command] --help\" for more information about a command.", + "file": "Upload or download workflow files through the Agent Stub.\n\nUsage:\n dify-agent file [command]\n\nAvailable Commands:\n download Download one workflow file mapping into the local sandbox directory.\n public-url Create a browser-visible download URL for an existing ToolFile reference.\n upload Upload one sandbox-local file as a ToolFile output reference.\n\nFlags:\n -h, --help help for file\n\nUse \"dify-agent file [command] --help\" for more information about a command.", "file download": "Download one workflow file mapping into the local sandbox directory.\n\nUsage:\n dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [flags]\n\nFlags:\n -h, --help help for download\n --to string Local directory for the downloaded file.", + "file public-url": "Create a browser-visible download URL for an existing ToolFile reference.\n\nUsage:\n dify-agent file public-url REFERENCE [flags]\n\nFlags:\n -h, --help help for public-url", "file upload": "Upload one sandbox-local file as a ToolFile output reference.\n\nUsage:\n dify-agent file upload PATH [flags]\n\nFlags:\n -h, --help help for upload\n --no-download-link Skip creating a public download link after upload." } diff --git a/dify-agent/src/dify_agent/layers/config/layer.py b/dify-agent/src/dify_agent/layers/config/layer.py index 7d68202d2a0..730cdc6b15c 100644 --- a/dify-agent/src/dify_agent/layers/config/layer.py +++ b/dify-agent/src/dify_agent/layers/config/layer.py @@ -49,6 +49,7 @@ _CONFIG_CLI_MUTATION_HELP_COMMANDS: dict[str, tuple[str, ...]] = { } _AGENT_FILE_CLI_HELP_COMMANDS: dict[str, tuple[str, ...]] = { "dify-agent file upload --help": ("file", "upload"), + "dify-agent file public-url --help": ("file", "public-url"), "dify-agent file download --help": ("file", "download"), } _CONFIG_CONTEXT_EXCLUDE = {"mentioned_skill_names": True, "mentioned_file_names": True} diff --git a/dify-agent/src/dify_agent/layers/dify_core_tools/layer.py b/dify-agent/src/dify_agent/layers/dify_core_tools/layer.py index 86e60ba86d6..fcd4e7eede4 100644 --- a/dify-agent/src/dify_agent/layers/dify_core_tools/layer.py +++ b/dify-agent/src/dify_agent/layers/dify_core_tools/layer.py @@ -11,9 +11,10 @@ from __future__ import annotations from copy import deepcopy from dataclasses import dataclass -from typing import ClassVar +from typing import ClassVar, cast import httpx +from pydantic import JsonValue from pydantic_ai import RunContext, Tool from pydantic_ai.tools import ToolDefinition from typing_extensions import Self, override @@ -109,7 +110,7 @@ class DifyCoreToolsLayer(PlainLayer[DifyCoreToolsDeps, DifyCoreToolsLayerConfig] response = await client.invoke( execution_context=execution_context, tool_config=tool_config, - tool_parameters=tool_arguments, + tool_parameters=cast(dict[str, JsonValue], tool_arguments), ) return response.observation except DifyCoreToolsClientConfigurationError: diff --git a/dify-agent/src/dify_agent/layers/dify_plugin/configs.py b/dify-agent/src/dify_agent/layers/dify_plugin/configs.py index 2ff75e14010..3d2ac615e5a 100644 --- a/dify-agent/src/dify_agent/layers/dify_plugin/configs.py +++ b/dify-agent/src/dify_agent/layers/dify_plugin/configs.py @@ -109,6 +109,7 @@ class DifyPluginLLMLayerConfig(LayerConfig): model_provider: str model: str model_settings: ModelSettings | None = None + context_window_tokens: int | None = Field(default=None, gt=0) model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid", arbitrary_types_allowed=True) diff --git a/dify-agent/src/dify_agent/layers/shell/layer.py b/dify-agent/src/dify_agent/layers/shell/layer.py index 87f033c68d3..e4f4987e19c 100644 --- a/dify-agent/src/dify_agent/layers/shell/layer.py +++ b/dify-agent/src/dify_agent/layers/shell/layer.py @@ -219,6 +219,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC shell_redact_patterns: list[str] = field(default_factory=list) agent_stub_api_base_url: str | None = None agent_stub_token_factory: ShellAgentStubTokenFactory | None = None + _job_agent_stub_tokens: dict[str, str] = field(default_factory=dict, init=False, repr=False) @classmethod @override @@ -300,10 +301,12 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC async def _tool_run(self, script: str, timeout: float = DEFAULT_TIMEOUT_SECONDS) -> ShellRunToolResult: try: + env = self._build_shell_command_env(include_agent_stub_env=True) + agent_stub_token = env.get(AGENT_STUB_AUTH_JWE_ENV_VAR) result = await self._require_resource().commands.run( _wrap_user_script(script, self.config), cwd=self._require_workspace_cwd(), - env=self._build_shell_command_env(include_agent_stub_env=True), + env=env, timeout=timeout, ) observation = await render_prompt_observation_from_result( @@ -313,6 +316,10 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC ) self._remember_job_id(result.job_id) self._remember_job_offset(result.job_id, observation.offset) + if agent_stub_token is not None and not result.done: + self._job_agent_stub_tokens[result.job_id] = agent_stub_token + else: + self._job_agent_stub_tokens.pop(result.job_id, None) return _tagged_shell_observation( _metadata_dict( job_id=result.job_id, @@ -321,7 +328,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - self._redact_output(observation.text), + self._redact_output(observation.text, sensitive_values=(agent_stub_token,)), ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc) @@ -339,6 +346,12 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC ) self._remember_job_id(result.job_id) self._remember_job_offset(result.job_id, observation.offset) + redacted_output = self._redact_output( + observation.text, + sensitive_values=(self._job_agent_stub_tokens.get(job_id),), + ) + if result.done: + self._job_agent_stub_tokens.pop(job_id, None) return _tagged_shell_observation( _metadata_dict( job_id=result.job_id, @@ -347,7 +360,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - self._redact_output(observation.text), + redacted_output, ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) @@ -365,6 +378,12 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC ) self._remember_job_id(result.job_id) self._remember_job_offset(result.job_id, observation.offset) + redacted_output = self._redact_output( + observation.text, + sensitive_values=(self._job_agent_stub_tokens.get(job_id),), + ) + if result.done: + self._job_agent_stub_tokens.pop(job_id, None) return _tagged_shell_observation( _metadata_dict( job_id=result.job_id, @@ -373,7 +392,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC exit_code=result.exit_code, output_path=observation.output_path, ), - self._redact_output(observation.text), + redacted_output, ) except (RuntimeError, ValueError) as exc: return _tool_error_from_exception(exc, job_id=job_id) @@ -390,6 +409,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC result = await self._require_resource().commands.interrupt(job_id, grace_seconds=grace_seconds) self._remember_job_id(result.job_id) self._remember_job_offset(result.job_id, result.offset) + self._job_agent_stub_tokens.pop(result.job_id, None) output_path: str | None = None try: # Once the interrupt itself succeeds, resolving the output path is @@ -509,6 +529,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC def _clear_tracked_jobs(self) -> None: self.runtime_state.job_offsets = {} self.runtime_state.job_ids = [] + self._job_agent_stub_tokens.clear() def _build_shell_command_env( self, @@ -536,7 +557,7 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC env.update(agent_stub_env) return env - def _redact_output(self, text: str) -> str: + def _redact_output(self, text: str, *, sensitive_values: Sequence[str | None] = ()) -> str: """Redact sensitive content from shell output before the model sees it. Two layers of redaction are applied: @@ -550,11 +571,11 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC """ if not text: return text - # Built-in: always redact the JWE token value. - env = self._build_shell_command_env(include_agent_stub_env=True) - jwe_value = env.get(AGENT_STUB_AUTH_JWE_ENV_VAR) - if jwe_value and len(jwe_value) > 8: - text = text.replace(jwe_value, "***") + # Built-in: always redact actual sensitive values supplied by the + # command owner. Redaction must never mint replacement credentials. + for value in sensitive_values: + if value and len(value) > 8: + text = text.replace(value, "***") # Server-level + per-agent regex patterns. for pattern in (*self.shell_redact_patterns, *self.config.redact_patterns): text = re.sub(pattern, "***", text) diff --git a/dify-agent/src/dify_agent/protocol/schemas.py b/dify-agent/src/dify_agent/protocol/schemas.py index ca23f8bfd59..a691650bfba 100644 --- a/dify-agent/src/dify_agent/protocol/schemas.py +++ b/dify-agent/src/dify_agent/protocol/schemas.py @@ -21,9 +21,10 @@ by ``DIFY_AGENT_MODEL_LAYER_ID``, the optional history layer named by ``DIFY_AGENT_HISTORY_LAYER_ID``, and the optional structured output layer named by ``DIFY_AGENT_OUTPUT_LAYER_ID``. Request-level ``on_exit`` signals decide 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 either the final JSON-safe +suspend as the default so terminal events can include resumable snapshots. +Successful runs always publish the resumable Agenton session snapshot on the +terminal ``run_succeeded`` event; failed and cancelled runs publish it when the +compositor context was entered and exited. Success includes either the final JSON-safe ``output`` or a deferred external ``deferred_tool_call`` payload. Session snapshots carry only layer lifecycle/runtime state in compositor order; they do not persist output-layer config. Resumed @@ -329,6 +330,7 @@ class RunFailedEventData(BaseModel): error: str error_type: RunFailureType | None = None reason: str | None = None + session_snapshot: CompositorSessionSnapshot | None = None model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") @@ -338,6 +340,7 @@ class RunCancelledEventData(BaseModel): reason: str | None = None message: str | None = None + session_snapshot: CompositorSessionSnapshot | None = None model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") diff --git a/dify-agent/src/dify_agent/runtime/agent_factory.py b/dify-agent/src/dify_agent/runtime/agent_factory.py index ad6bf7a5510..e82f246d866 100644 --- a/dify-agent/src/dify_agent/runtime/agent_factory.py +++ b/dify-agent/src/dify_agent/runtime/agent_factory.py @@ -4,9 +4,9 @@ The run request carries model/provider selection in the layer graph. This helper keeps Agent construction details out of ``AgentRunRunner`` while accepting an already resolved Pydantic AI model from the configured model layer. Tool values arriving here are already transformed by Agenton's -``PYDANTIC_AI_TRANSFORMERS`` preset, while Dify system prompts are rendered into -temporary ``message_history`` before the call reaches this helper. The caller -also passes the already resolved ``output_type`` so legacy text output and the +``PYDANTIC_AI_TRANSFORMERS`` preset. The runner passes Dify system prompts as +run-level instructions and the context compaction capability directly to +``Agent.run``. The caller also passes the already resolved ``output_type`` so legacy text output and the optional JSON Schema output layer share the same ``Agent`` construction path. """ diff --git a/dify-agent/src/dify_agent/runtime/cancellation.py b/dify-agent/src/dify_agent/runtime/cancellation.py new file mode 100644 index 00000000000..60ad2750734 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime/cancellation.py @@ -0,0 +1,19 @@ +"""Private cancellation coordination types shared by schedulers and run stores.""" + +from datetime import datetime +from typing import ClassVar + +from pydantic import BaseModel, ConfigDict + + +class RunCancellationIntent(BaseModel): + """The first accepted request to cancel one running run.""" + + reason: str | None = None + message: str | None = None + requested_at: datetime + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +__all__ = ["RunCancellationIntent"] diff --git a/dify-agent/src/dify_agent/runtime/compaction.py b/dify-agent/src/dify_agent/runtime/compaction.py new file mode 100644 index 00000000000..b1e40e5f436 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime/compaction.py @@ -0,0 +1,45 @@ +"""Build the standard Dify Agent context-compaction capability. + +``TieredCompaction`` owns the target-budget check and invokes child compactors +without evaluating their individual triggers. Each child constructor requires +at least one configured trigger (``max_messages``, ``max_tokens``, or +``max_fraction``) and validates the selected trigger. Dify uses the otherwise +unused ``max_tokens=1`` values below solely to satisfy that validation; they are +not one-token Dify policy thresholds. +""" + +from pydantic_ai.settings import ModelSettings +from pydantic_ai_harness.compaction import ClearToolResults, SummarizingCompaction, TieredCompaction + + +def build_compaction_capability( + *, + context_window_tokens: int | None, + model_settings: ModelSettings | None, +) -> TieredCompaction[None] | None: + """Build compaction for the effective model window, or disable it when unknown.""" + if context_window_tokens is None: + return None + + input_budget = context_window_tokens * 4 // 5 + max_tokens = model_settings.get("max_tokens") if model_settings is not None else None + if max_tokens is not None and max_tokens > 0: + input_budget = min(input_budget, context_window_tokens - max_tokens) + if input_budget <= 0: + raise ValueError("Model max_tokens must leave a positive input context budget.") + + return TieredCompaction( + tiers=[ + ClearToolResults(max_tokens=1, keep_pairs=3, clear_tool_inputs=False), + SummarizingCompaction( + max_tokens=1, + keep_messages=20, + preserve_first_user_message=True, + incremental=True, + ), + ], + target_tokens=input_budget, + ) + + +__all__ = ["build_compaction_capability"] diff --git a/dify-agent/src/dify_agent/runtime/event_sink.py b/dify-agent/src/dify_agent/runtime/event_sink.py index 121d58cecde..5e5995045f0 100644 --- a/dify-agent/src/dify_agent/runtime/event_sink.py +++ b/dify-agent/src/dify_agent/runtime/event_sink.py @@ -1,9 +1,10 @@ """Event sink contracts used by the runner and storage adapters. -Non-terminal events remain append-only. Terminal events use ``finalize_run`` so -the event and matching run status are committed as one compare-and-set -transition. Tests can use ``InMemoryRunEventSink`` without Redis; production -storage implements the same contract with Redis streams in +Non-terminal events remain append-only. Successful and failed terminal events +use ``finalize_run`` so the event and matching run status are committed as one +compare-and-set transition. Cancellation has a dedicated intent-aware finalizer. +Tests can use ``InMemoryRunEventSink`` without Redis; production storage +implements the same contract with Redis streams in ``dify_agent.storage.redis_run_store``. """ @@ -21,8 +22,6 @@ from dify_agent.protocol.schemas import ( EmptyRunEventData, PydanticAIStreamRunEvent, RunEvent, - RunCancelledEvent, - RunCancelledEventData, RunFailedEvent, RunFailedEventData, RunFailureType, @@ -35,7 +34,7 @@ from dify_agent.protocol.schemas import ( _UNSET = object() -TerminalRunEvent: TypeAlias = RunSucceededEvent | RunFailedEvent | RunCancelledEvent +TerminalRunEvent: TypeAlias = RunSucceededEvent | RunFailedEvent NonTerminalRunEvent: TypeAlias = RunStartedEvent | PydanticAIStreamRunEvent @@ -105,8 +104,6 @@ def terminal_event_status_fields( return "succeeded", None, None case RunFailedEvent(): return "failed", event.data.error, event.data.error_type - case RunCancelledEvent(): - return "cancelled", event.data.message or event.data.reason, None async def emit_run_event( @@ -188,29 +185,18 @@ async def emit_run_failed( error: str, error_type: RunFailureType | None = None, reason: str | None = None, + session_snapshot: CompositorSessionSnapshot | None = None, ) -> RunFinalizationResult: """Finalize a run with a failed terminal event.""" return await sink.finalize_run( RunFailedEvent( run_id=run_id, - data=RunFailedEventData(error=error, error_type=error_type, reason=reason), - created_at=utc_now(), - ), - ) - - -async def emit_run_cancelled( - sink: RunEventSink, - *, - run_id: str, - reason: str | None = None, - message: str | None = None, -) -> RunFinalizationResult: - """Finalize a run with a cancelled terminal event.""" - return await sink.finalize_run( - RunCancelledEvent( - run_id=run_id, - data=RunCancelledEventData(reason=reason, message=message), + data=RunFailedEventData( + error=error, + error_type=error_type, + reason=reason, + session_snapshot=session_snapshot, + ), created_at=utc_now(), ), ) @@ -223,7 +209,6 @@ __all__ = [ "RunFinalizationResult", "TerminalRunEvent", "emit_pydantic_ai_event", - "emit_run_cancelled", "emit_run_event", "emit_run_failed", "emit_run_started", diff --git a/dify-agent/src/dify_agent/runtime/history.py b/dify-agent/src/dify_agent/runtime/history.py index 1026b42bb9f..8b2d4d11e73 100644 --- a/dify-agent/src/dify_agent/runtime/history.py +++ b/dify-agent/src/dify_agent/runtime/history.py @@ -1,24 +1,19 @@ """Helpers for optional Dify Agent history-layer integration. Dify Agent keeps pydantic-ai conversation history as an optional Agenton layer -named ``history``. The runner always injects the current Dify system prompt via -temporary ``message_history`` instead of ``Agent.system_prompt(...)`` so the -model sees ``current system prompt -> stored history -> current user prompt`` -even when persisted history is present. Only zero-argument system prompt -callables are supported here because the prompts are rendered outside -pydantic-ai's normal run context; this matches Dify's current plain-prompt -compositions and fails fast for unsupported context-dependent prompt shapes. +named ``history``. Current system instructions belong to each run and are never +stored; successful runs replace the layer with Pydantic AI's complete, possibly +compacted history. """ from __future__ import annotations -import inspect -from collections.abc import Awaitable, Callable, Sequence -from typing import Protocol, cast +from collections.abc import Sequence +from dataclasses import replace +from typing import Protocol -from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart +from pydantic_ai.messages import ModelMessage, ModelRequest -from agenton.layers.types import PydanticAIPrompt from agenton_collections.layers.pydantic_ai import PYDANTIC_AI_HISTORY_LAYER_TYPE_ID, PydanticAIHistoryLayer from dify_agent.protocol import DIFY_AGENT_HISTORY_LAYER_ID from dify_agent.protocol.schemas import RunComposition @@ -68,66 +63,22 @@ def get_history_layer(run: SupportsHistoryLayerLookup) -> PydanticAIHistoryLayer return None -async def build_run_message_history( - *, - system_prompts: Sequence[PydanticAIPrompt[object]], - stored_history: Sequence[ModelMessage], -) -> list[ModelMessage] | None: - """Build temporary pydantic-ai history for one Dify Agent loop. - - Current system prompts are rendered first into one transient - ``ModelRequest`` prefix, followed by any already stored history messages. - When both inputs are empty, the helper returns ``None`` so callers can omit - the ``message_history`` argument entirely and preserve pydantic-ai's empty - history behavior. - """ - rendered_system_parts: list[SystemPromptPart] = [] - for prompt in system_prompts: - prompt_text = await _render_system_prompt(prompt) - if prompt_text is None: - continue - rendered_system_parts.append(SystemPromptPart(content=prompt_text)) - - message_history: list[ModelMessage] = [] - if rendered_system_parts: - message_history.append(ModelRequest(parts=rendered_system_parts)) - message_history.extend(stored_history) - return message_history or None - - -def append_successful_run_history( +def replace_successful_run_history( history_layer: PydanticAIHistoryLayer | None, - new_messages: Sequence[ModelMessage], + messages: Sequence[ModelMessage], ) -> None: - """Append only newly produced pydantic-ai messages after successful runs.""" - if history_layer is None or not new_messages: + """Persist a successful run's complete history without transient instructions.""" + if history_layer is None: return - history_layer.append_messages(new_messages) - - -async def _render_system_prompt(prompt: PydanticAIPrompt[object]) -> str | None: - signature = inspect.signature(prompt) - if signature.parameters: - raise ValueError( - "Dify Agent runtime currently supports only zero-argument system prompts when rendering temporary " - "message history." - ) - - prompt_without_context = cast(Callable[[], str | None | Awaitable[str | None]], prompt) - prompt_value = prompt_without_context() - if inspect.isawaitable(prompt_value): - prompt_value = await prompt_value - if prompt_value is None: - return None - if not isinstance(prompt_value, str): - raise TypeError(f"System prompt callables must return str | None, got '{type(prompt_value).__name__}'.") - return prompt_value + persistent_messages = [ + replace(message, instructions=None) if isinstance(message, ModelRequest) else message for message in messages + ] + history_layer.replace_messages(persistent_messages) __all__ = [ "SupportsHistoryLayerLookup", - "append_successful_run_history", - "build_run_message_history", "get_history_layer", + "replace_successful_run_history", "validate_history_layer_composition", ] diff --git a/dify-agent/src/dify_agent/runtime/run_scheduler.py b/dify-agent/src/dify_agent/runtime/run_scheduler.py index b22d48e5626..b197d7dcb3e 100644 --- a/dify-agent/src/dify_agent/runtime/run_scheduler.py +++ b/dify-agent/src/dify_agent/runtime/run_scheduler.py @@ -20,10 +20,11 @@ from typing import Protocol import httpx -from agenton.compositor import LayerProviderInput -from dify_agent.protocol.schemas import CancelRunRequest, CancelRunResponse, CreateRunRequest +from agenton.compositor import CompositorSessionSnapshot, LayerProviderInput +from dify_agent.protocol.schemas import CancelRunRequest, CancelRunResponse, CreateRunRequest, RunStatus +from dify_agent.runtime.cancellation import RunCancellationIntent from dify_agent.runtime.compositor_factory import create_default_layer_providers -from dify_agent.runtime.event_sink import RunEventSink, emit_run_cancelled, emit_run_failed +from dify_agent.runtime.event_sink import RunEventSink, RunFinalizationResult, emit_run_failed from dify_agent.runtime.runner import DEFAULT_AGENT_RUN_TIMEOUT_SECONDS, AgentRunRunner from dify_agent.server.schemas import RunRecord @@ -45,14 +46,37 @@ class RunStore(RunEventSink, Protocol): """Persist a new run record and return it with status ``running``.""" ... - async def wait_for_cancellation(self, run_id: str) -> bool: - """Wait for a terminal state and report whether cancellation won.""" + async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus: + """Persist the first cancellation intent and return the current status.""" + ... + + async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None: + """Return the accepted cancellation intent, if one exists.""" + ... + + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: + """Wait for a cancellation intent or a different terminal state.""" + ... + + async def finalize_cancellation( + self, + run_id: str, + intent: RunCancellationIntent, + *, + session_snapshot: CompositorSessionSnapshot | None = None, + ) -> RunFinalizationResult: + """Publish cancellation after the owner runner has exited.""" ... class RunnableRun(Protocol): """Executable unit for one scheduled run.""" + @property + def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None: + """Return the post-exit snapshot for the current invocation, if available.""" + ... + async def run(self) -> None: """Run until terminal status/events have been written or cancellation occurs.""" ... @@ -125,14 +149,9 @@ class RunScheduler: async def cancel_run(self, run_id: str, request: CancelRunRequest) -> CancelRunResponse: """Persist an idempotent cancellation without relying on local task ownership.""" - finalization = await emit_run_cancelled( - self.store, - run_id=run_id, - reason=request.reason, - message=request.message, - ) - if finalization.status != "cancelled": - raise RunCancellationConflictError(f"run already finished with status {finalization.status!r}") + status = await self.store.request_cancellation(run_id, request) + if status in {"succeeded", "failed"}: + raise RunCancellationConflictError(f"run already finished with status {status!r}") return CancelRunResponse(run_id=run_id, status="cancelled") async def shutdown(self) -> None: @@ -141,18 +160,14 @@ class RunScheduler: self.stopping = True if not self.active_tasks: return - tasks_by_run_id = dict(self.active_tasks) - done, pending = await asyncio.wait(tasks_by_run_id.values(), timeout=self.shutdown_grace_seconds) - del done + tasks = tuple(self.active_tasks.values()) + _done, pending = await asyncio.wait(tasks, timeout=self.shutdown_grace_seconds) if not pending: return - pending_run_ids = [run_id for run_id, task in tasks_by_run_id.items() if task in pending] for task in pending: _ = task.cancel() _ = await asyncio.gather(*pending, return_exceptions=True) - for run_id in pending_run_ids: - await self._mark_cancelled_run_failed(run_id) async def _run_record(self, record: RunRecord, request: CreateRunRequest) -> None: """Supervise one local runner and its durable cancellation observer.""" @@ -163,41 +178,92 @@ class RunScheduler: self.store.wait_for_cancellation(record.run_id), name=f"dify-agent-cancellation-observer-{record.run_id}", ) + + async def cancel_runner_and_wait() -> None: + if not cancel_requested.is_set() and not runner_task.done(): + cancel_requested.set() + _ = runner_task.cancel() + _ = await asyncio.shield(asyncio.gather(runner_task, return_exceptions=True)) + try: _ = await asyncio.wait((runner_task, observer_task), return_when=asyncio.FIRST_COMPLETED) if observer_task.done(): try: - cancellation_won = observer_task.result() + intent = observer_task.result() except Exception as exc: - cancel_requested.set() - await self._cancel_and_wait(runner_task, reinject=True) - _ = await emit_run_failed( + await cancel_runner_and_wait() + finalization = await emit_run_failed( self.store, run_id=record.run_id, error=f"run cancellation observer failed: {exc}", reason="cancellation_observer", + session_snapshot=runner.terminal_session_snapshot, ) + if not finalization.applied and finalization.status == "running": + intent = await self.store.get_cancellation_intent(record.run_id) + if intent is not None: + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) raise - if cancellation_won: - cancel_requested.set() - await self._cancel_and_wait(runner_task, reinject=True) + if intent is not None: + await cancel_runner_and_wait() + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) else: await runner_task else: - await runner_task + runner_error: Exception | None = None + try: + await runner_task + except Exception as exc: + runner_error = exc + + intent = await self.store.get_cancellation_intent(record.run_id) + if intent is not None: + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) + if runner_error is not None: + raise runner_error except asyncio.CancelledError: - cancel_requested.set() await self._cancel_and_wait(observer_task) - await self._cancel_and_wait(runner_task, reinject=True) + await cancel_runner_and_wait() + intent = await self.store.get_cancellation_intent(record.run_id) + if intent is not None: + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) + else: + finalization = await self._mark_cancelled_run_failed( + record.run_id, + session_snapshot=runner.terminal_session_snapshot, + ) + if finalization is not None and not finalization.applied and finalization.status == "running": + intent = await self.store.get_cancellation_intent(record.run_id) + if intent is not None: + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) raise except Exception: logger.exception("scheduled run failed", extra={"run_id": record.run_id}) finally: await self._cancel_and_wait(observer_task) if not runner_task.done(): - cancel_requested.set() - await self._cancel_and_wait(runner_task, reinject=True) + await cancel_runner_and_wait() def _create_runner( self, @@ -234,25 +300,31 @@ class RunScheduler: _ = self.active_tasks.pop(run_id, None) @staticmethod - async def _cancel_and_wait(task: asyncio.Task[object], *, reinject: bool = False) -> None: - """Cancel and reap a child task, with bounded reinjection for runners.""" + async def _cancel_and_wait(task: asyncio.Task[object]) -> None: + """Cancel a child task once and await its complete exit.""" if not task.done(): _ = task.cancel() - if reinject: - for _attempt in range(2): - await asyncio.sleep(0) - if task.done(): - break - _ = task.cancel() _ = await asyncio.gather(task, return_exceptions=True) - async def _mark_cancelled_run_failed(self, run_id: str) -> None: + async def _mark_cancelled_run_failed( + self, + run_id: str, + *, + session_snapshot: CompositorSessionSnapshot | None = None, + ) -> RunFinalizationResult | None: """Best-effort failure event/status for shutdown-cancelled runs.""" message = "run cancelled during server shutdown" try: - _ = await emit_run_failed(self.store, run_id=run_id, error=message, reason="shutdown") + return await emit_run_failed( + self.store, + run_id=run_id, + error=message, + reason="shutdown", + session_snapshot=session_snapshot, + ) except Exception: logger.exception("failed to mark cancelled run failed", extra={"run_id": run_id}) + return None __all__ = ["RunCancellationConflictError", "RunScheduler", "SchedulerStoppingError"] diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index 6eafa456b20..c6c8d829312 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/dify-agent/src/dify_agent/runtime/runner.py @@ -5,15 +5,16 @@ Agenton's graph/config split and executes one model run after 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 + pass the current Dify system prompts as run-level instructions, 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; 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 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 +message history only through session state; successful model runs replace that +state with ``result.all_messages()`` after transient instructions are cleared so +compaction rewrites persist without saving current system prompts. An optional +structured output layer named by ``DIFY_AGENT_OUTPUT_LAYER_ID`` is read after entry and resolved into an output contract whose type both exposes the output schema to the model and performs runtime JSON Schema validation through custom Pydantic hooks. When the ask-human @@ -61,6 +62,7 @@ from dify_agent.protocol.schemas import ( from dify_agent.runtime.agent_factory import create_agent, normalize_user_input from dify_agent.runtime.agenton_validation import is_agenton_enter_validation_runtime_error from dify_agent.runtime.compositor_factory import build_pydantic_ai_compositor, create_default_layer_providers +from dify_agent.runtime.compaction import build_compaction_capability from dify_agent.runtime_backend import BindingLostError from dify_agent.runtime.event_sink import ( RunEventSink, @@ -70,9 +72,8 @@ from dify_agent.runtime.event_sink import ( emit_run_succeeded, ) from dify_agent.runtime.history import ( - append_successful_run_history, - build_run_message_history, get_history_layer, + replace_successful_run_history, validate_history_layer_composition, ) from dify_agent.runtime.layer_exit_signals import apply_layer_exit_signals, validate_layer_exit_signals @@ -183,6 +184,7 @@ class AgentRunRunner: dify_api_http_client: httpx.AsyncClient is_cancelled: Callable[[], bool] run_timeout_seconds: float + _terminal_session_snapshot: CompositorSessionSnapshot | None def __init__( self, @@ -204,9 +206,16 @@ class AgentRunRunner: self.layer_providers = layer_providers if layer_providers is not None else create_default_layer_providers() self.is_cancelled = is_cancelled or (lambda: False) self.run_timeout_seconds = run_timeout_seconds + self._terminal_session_snapshot = None + + @property + def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None: + """Return the snapshot captured after the current compositor context exited.""" + return self._terminal_session_snapshot async def run(self) -> None: """Execute the run and emit the documented event sequence.""" + self._terminal_session_snapshot = None if self.is_cancelled(): return _ = await emit_run_started(self.sink, run_id=self.run_id) @@ -223,6 +232,7 @@ class AgentRunRunner: error=message, error_type=error_type, reason=reason, + session_snapshot=self._terminal_session_snapshot, ) if finalization.applied: raise @@ -286,6 +296,7 @@ class AgentRunRunner: deferred_tool_call: DeferredToolCallPayload | None = None result_kind: Literal["output", "deferred_tool_call"] | None = None usage: AgentRunUsage | None = None + run = None try: async with compositor.enter(configs=layer_configs, session_snapshot=self.request.session_snapshot) as run: entered_run = True @@ -310,12 +321,13 @@ class AgentRunRunner: try: output_contract = resolve_run_output_contract(run) history_layer = get_history_layer(run) - message_history = await build_run_message_history( - system_prompts=run.prompts, - stored_history=history_layer.message_history if history_layer is not None else (), - ) + message_history = history_layer.message_history if history_layer is not None else None ask_human_layer = get_ask_human_layer(run) llm_layer = run.get_layer(DIFY_AGENT_MODEL_LAYER_ID, DifyPluginLLMLayer) + compaction = build_compaction_capability( + context_window_tokens=llm_layer.config.context_window_tokens, + model_settings=llm_layer.config.model_settings, + ) model = llm_layer.get_model( http_client=self.dify_api_http_client, agent_run_id=self.run_id, @@ -346,6 +358,8 @@ class AgentRunRunner: message_history=message_history, deferred_tool_results=deferred_tool_results, event_stream_handler=handle_events, + instructions=run.prompts or None, + capabilities=[compaction] if compaction is not None else None, usage_limits=UsageLimits(request_limit=_MAX_AGENT_STEPS_PER_RUN), ) except TimeoutError as exc: @@ -356,7 +370,7 @@ class AgentRunRunner: ) from exc complete_usage = model.accumulated_usage if isinstance(model, _HasAccumulatedUsage) else None usage = _serialize_agent_usage(complete_usage if complete_usage is not None else _result_usage(result)) - append_successful_run_history(history_layer, result.new_messages()) + replace_successful_run_history(history_layer, result.all_messages()) if isinstance(result.output, DeferredToolRequests): if ask_human_layer is None: raise AgentRunValidationError( @@ -379,8 +393,11 @@ class AgentRunRunner: if not entered_run: raise AgentRunValidationError(str(exc)) from exc raise + finally: + if entered_run and run is not None: + self._terminal_session_snapshot = run.session_snapshot - if run.session_snapshot is None: + if run is None or run.session_snapshot is None: raise RuntimeError("Agenton run did not produce a session snapshot after exit.") if result_kind is None: raise RuntimeError("Agent run did not resolve either a final output or a deferred tool call.") diff --git a/dify-agent/src/dify_agent/runtime_backend/e2b.py b/dify-agent/src/dify_agent/runtime_backend/e2b.py index 930ea317dcf..376a43a54b7 100644 --- a/dify-agent/src/dify_agent/runtime_backend/e2b.py +++ b/dify-agent/src/dify_agent/runtime_backend/e2b.py @@ -126,6 +126,7 @@ class E2BSDKControlPlane: template, timeout=timeout, metadata=metadata, + network={"allow_public_traffic": False}, lifecycle={"on_timeout": on_timeout, "auto_resume": False}, **self._options(), ), @@ -209,7 +210,6 @@ class E2BExecutionBindingBackend: control_plane: E2BControlPlane template: str active_timeout_seconds: int - shellctl_auth_token: str = "" shellctl_port: int = 5004 layout: RuntimeLayout = field( default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/home/dify/workspace") @@ -308,14 +308,17 @@ class E2BExecutionBindingBackend: async def _lease(self, sandbox: _E2BSandbox) -> "E2BRuntimeLease": entrypoint = f"https://{sandbox.get_host(self.shellctl_port)}" traffic_token = sandbox.traffic_access_token - headers = {"X-Access-Token": traffic_token} if isinstance(traffic_token, str) and traffic_token else {} + if not isinstance(traffic_token, str) or not traffic_token: + raise BindingAcquireError("E2B sandbox did not provide a traffic access token") http_client = httpx.AsyncClient( base_url=entrypoint, - headers=headers, + headers={"X-Access-Token": traffic_token}, follow_redirects=True, timeout=httpx.Timeout(60.0), ) + # Explicit token="" prevents process-level SHELLCTL_AUTH_TOKEN fallback; + # E2B port access is authenticated only by X-Access-Token above. def client_factory() -> ShellctlClientProtocol: from shellctl.client import ShellctlClient @@ -323,7 +326,7 @@ class E2BExecutionBindingBackend: ShellctlClientProtocol, cast( object, - ShellctlClient(entrypoint, token=self.shellctl_auth_token, client=http_client), + ShellctlClient(entrypoint, token="", client=http_client), ), ) @@ -331,7 +334,7 @@ class E2BExecutionBindingBackend: handle=sandbox.sandbox_id, layout=self.layout, entrypoint=entrypoint, - token=self.shellctl_auth_token, + token="", client_factory=client_factory, owned_transport=http_client, ) diff --git a/dify-agent/src/dify_agent/runtime_backend/profile.py b/dify-agent/src/dify_agent/runtime_backend/profile.py index 28bcf6bd01e..9bbbf3ee3f7 100644 --- a/dify-agent/src/dify_agent/runtime_backend/profile.py +++ b/dify-agent/src/dify_agent/runtime_backend/profile.py @@ -60,7 +60,6 @@ class RuntimeBackendSettings(BaseSettings): ge=1, le=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, ) - e2b_shellctl_auth_token: str = "" e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535) model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict( @@ -146,7 +145,6 @@ def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeB control_plane=control_plane, template=settings.e2b_template, active_timeout_seconds=settings.e2b_active_timeout_seconds, - shellctl_auth_token=settings.e2b_shellctl_auth_token, shellctl_port=settings.e2b_shellctl_port, ), ) diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index ed7425b84d4..197592c8815 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -71,13 +71,18 @@ class ServerSettings(BaseSettings): ge=1, le=E2B_MAX_ACTIVE_TIMEOUT_SECONDS, ) - e2b_shellctl_auth_token: str = "" e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535) agent_stub_api_base_url: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_API_BASE_URL") sandbox_files_base_url: str | None = Field( default=None, validation_alias="DIFY_AGENT_SANDBOX_FILES_BASE_URL", ) + stub_upload_file_size_limit: int = Field( + default=50, + ge=0, + description="Maximum Agent Stub upload size in MiB", + validation_alias="DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT", + ) server_secret_key: str | None = None api_token: str | None = None shell_redact_patterns: str = "" @@ -205,7 +210,6 @@ class ServerSettings(BaseSettings): e2b_api_key=self.e2b_api_key, e2b_template=self.e2b_template, e2b_active_timeout_seconds=self.e2b_active_timeout_seconds, - e2b_shellctl_auth_token=self.e2b_shellctl_auth_token, e2b_shellctl_port=self.e2b_shellctl_port, ) ) @@ -224,6 +228,7 @@ class ServerSettings(BaseSettings): inner_api_url=self.inner_api_url, inner_api_key=self.inner_api_key, sandbox_files_base_url=self.sandbox_files_base_url, + max_upload_size_bytes=self.stub_upload_file_size_limit * 1024 * 1024, timeout=self.create_outbound_http_timeout(), ) diff --git a/dify-agent/src/dify_agent/storage/redis_keys.py b/dify-agent/src/dify_agent/storage/redis_keys.py index c93e69ba825..bc8d10ef739 100644 --- a/dify-agent/src/dify_agent/storage/redis_keys.py +++ b/dify-agent/src/dify_agent/storage/redis_keys.py @@ -11,4 +11,9 @@ def run_events_key(prefix: str, run_id: str) -> str: return f"{prefix}:runs:{run_id}:events" -__all__ = ["run_events_key", "run_record_key"] +def run_cancel_intent_key(prefix: str, run_id: str) -> str: + """Return the private Redis stream key holding one cancellation intent.""" + return f"{prefix}:runs:{run_id}:cancel-intent" + + +__all__ = ["run_cancel_intent_key", "run_events_key", "run_record_key"] diff --git a/dify-agent/src/dify_agent/storage/redis_run_store.py b/dify-agent/src/dify_agent/storage/redis_run_store.py index f488749c0ff..70a6adc85d6 100644 --- a/dify-agent/src/dify_agent/storage/redis_run_store.py +++ b/dify-agent/src/dify_agent/storage/redis_run_store.py @@ -1,8 +1,8 @@ -"""Redis-backed run records and per-run event streams. +"""Redis-backed run records, event streams, and private cancellation intents. The store writes status-only run records as JSON strings and events as Redis streams. HTTP event cursors are Redis stream ids; ``0-0`` means replay from the -beginning for polling and SSE. Records and streams share one retention window +beginning for polling and SSE. Records, event streams, and intents share one retention window that is refreshed when status or event data is written. Execution is scheduled in-process by ``dify_agent.runtime.run_scheduler``; Redis is not a job queue, and create-run payloads are never persisted because layer config may include @@ -14,7 +14,18 @@ from typing import cast from redis.asyncio import Redis -from dify_agent.protocol.schemas import RUN_EVENT_ADAPTER, RunEvent, RunEventsResponse, RunStatus +from agenton.compositor import CompositorSessionSnapshot +from dify_agent.protocol.schemas import ( + RUN_EVENT_ADAPTER, + CancelRunRequest, + RunCancelledEvent, + RunCancelledEventData, + RunEvent, + RunEventsResponse, + RunStatus, + utc_now, +) +from dify_agent.runtime.cancellation import RunCancellationIntent from dify_agent.runtime.event_sink import ( NonTerminalRunEvent, RunEventSink, @@ -24,7 +35,7 @@ from dify_agent.runtime.event_sink import ( ) from dify_agent.server.schemas import RunRecord, new_run_id from dify_agent.server.settings import DEFAULT_RUN_RETENTION_SECONDS -from dify_agent.storage.redis_keys import run_events_key, run_record_key +from dify_agent.storage.redis_keys import run_cancel_intent_key, run_events_key, run_record_key _TERMINAL_RUN_EVENT_TYPES = {"run_succeeded", "run_failed", "run_cancelled"} @@ -44,6 +55,10 @@ if record.status ~= "running" then return {0, tostring(record.status), ""} end +if redis.call("EXISTS", KEYS[3]) == 1 then + return {-2, "running", ""} +end + record.status = ARGV[1] record.updated_at = ARGV[2] if ARGV[3] == "1" then @@ -66,6 +81,67 @@ return {1, ARGV[1], event_id} """ +_REQUEST_CANCELLATION_SCRIPT = """ +local record_json = redis.call("GET", KEYS[1]) +if not record_json then + return {-1, ""} +end + +local record = cjson.decode(record_json) +if record.status == "succeeded" or record.status == "failed" then + return {0, tostring(record.status)} +end +if record.status == "cancelled" then + return {1, "cancelled"} +end +if redis.call("EXISTS", KEYS[2]) == 1 then + return {1, "running"} +end + +local ttl = tonumber(ARGV[2]) +redis.call("XADD", KEYS[2], "*", "payload", ARGV[1]) +redis.call("EXPIRE", KEYS[2], ttl) +redis.call("EXPIRE", KEYS[1], ttl) +redis.call("EXPIRE", KEYS[3], ttl) +return {1, "running"} +""" + + +_FINALIZE_CANCELLATION_SCRIPT = """ +local record_json = redis.call("GET", KEYS[1]) +if not record_json then + return {-1, "", ""} +end + +local record = cjson.decode(record_json) +if record.status == "cancelled" then + return {0, "cancelled", ""} +end +if record.status ~= "running" then + return {0, tostring(record.status), ""} +end +if redis.call("EXISTS", KEYS[2]) == 0 then + return {-2, "running", ""} +end + +record.status = "cancelled" +record.updated_at = ARGV[1] +if ARGV[2] == "1" then + record.error = ARGV[3] +else + record.error = cjson.null +end +record.error_type = cjson.null + +local ttl = tonumber(ARGV[5]) +local event_id = redis.call("XADD", KEYS[3], "*", "payload", ARGV[4]) +redis.call("DEL", KEYS[2]) +redis.call("EXPIRE", KEYS[3], ttl) +redis.call("SET", KEYS[1], cjson.encode(record), "EX", ttl) +return {1, "cancelled", event_id} +""" + + class RedisRunStore(RunEventSink): """Async Redis implementation for run records and event logs. @@ -130,16 +206,17 @@ class RedisRunStore(RunEventSink): return event_id.decode() if isinstance(event_id, bytes) else str(event_id) async def finalize_run(self, event: TerminalRunEvent) -> RunFinalizationResult: - """Atomically append the first terminal event and update its run record.""" + """Atomically append the first success/failure event and update its run record.""" status, error, error_type = terminal_event_status_fields(event) payload = RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode() evaluation = cast( Awaitable[object], self.redis.eval( _FINALIZE_RUN_SCRIPT, - 2, + 3, run_record_key(self.prefix, event.run_id), run_events_key(self.prefix, event.run_id), + run_cancel_intent_key(self.prefix, event.run_id), status, event.created_at.isoformat(), "1" if error is not None else "0", @@ -164,8 +241,39 @@ class RedisRunStore(RunEventSink): event_id=event_id, ) - async def wait_for_cancellation(self, run_id: str) -> bool: - """Wait until cancellation or another terminal state wins for one run. + async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus: + """Atomically persist the first cancellation intent for a running run.""" + intent = RunCancellationIntent( + reason=request.reason, + message=request.message, + requested_at=utc_now(), + ) + evaluation = cast( + Awaitable[object], + self.redis.eval( + _REQUEST_CANCELLATION_SCRIPT, + 3, + run_record_key(self.prefix, run_id), + run_cancel_intent_key(self.prefix, run_id), + run_events_key(self.prefix, run_id), + intent.model_dump_json(), + str(self.run_retention_seconds), + ), + ) + result = cast(list[object], await evaluation) + if int(cast(int | bytes | str, result[0])) == -1: + raise RunNotFoundError(run_id) + return cast(RunStatus, _decode_redis_text(result[1])) + + async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None: + """Return the accepted private cancellation intent, if one exists.""" + entries = await self.redis.xrange(run_cancel_intent_key(self.prefix, run_id), count=1) + if not entries: + return None + return self._decode_cancellation_intent(entries[0][1]) + + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: + """Wait until cancellation intent or another terminal state wins. The stream cursor is captured before reading the record so a terminal transition cannot fall between the initial status check and blocking @@ -176,19 +284,76 @@ class RedisRunStore(RunEventSink): cursor = _decode_redis_text(latest_events[0][0]) if latest_events else "0-0" record = await self.get_run(run_id) if record.status != "running": - return record.status == "cancelled" + return None + + intent = await self.get_cancellation_intent(run_id) + if intent is not None: + return intent while True: - response = await self.redis.xread({events_key: cursor}, block=0, count=100) - for _stream_name, entries in response: + response = await self.redis.xread( + { + run_cancel_intent_key(self.prefix, run_id): "0-0", + events_key: cursor, + }, + block=0, + count=100, + ) + for stream_name, entries in response: + if _decode_redis_text(stream_name) == run_cancel_intent_key(self.prefix, run_id): + return self._decode_cancellation_intent(entries[0][1]) for raw_id, fields in entries: event = self._decode_event(run_id, raw_id, fields) if event.id is not None: cursor = event.id if event.type == "run_cancelled": - return True + return None if event.type in {"run_succeeded", "run_failed"}: - return False + return None + + async def finalize_cancellation( + self, + run_id: str, + intent: RunCancellationIntent, + *, + session_snapshot: CompositorSessionSnapshot | None = None, + ) -> RunFinalizationResult: + """Atomically publish cancellation after the owner runner has exited.""" + event = RunCancelledEvent( + run_id=run_id, + data=RunCancelledEventData( + reason=intent.reason, + message=intent.message, + session_snapshot=session_snapshot, + ), + created_at=utc_now(), + ) + payload = RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode() + error = event.data.message or event.data.reason + evaluation = cast( + Awaitable[object], + self.redis.eval( + _FINALIZE_CANCELLATION_SCRIPT, + 3, + run_record_key(self.prefix, run_id), + run_cancel_intent_key(self.prefix, run_id), + run_events_key(self.prefix, run_id), + event.created_at.isoformat(), + "1" if error is not None else "0", + error or "", + payload, + str(self.run_retention_seconds), + ), + ) + result = cast(list[object], await evaluation) + applied = int(cast(int | bytes | str, result[0])) + if applied == -1: + raise RunNotFoundError(run_id) + return RunFinalizationResult( + applied=applied == 1, + status=cast(RunStatus, _decode_redis_text(result[1])), + event_id=_decode_redis_text(result[2]) or None, + ) async def get_events(self, run_id: str, *, after: str = "0-0", limit: int = 100) -> RunEventsResponse: """Read a bounded page of events after ``after`` cursor.""" @@ -235,6 +400,13 @@ class RedisRunStore(RunEventSink): event = RUN_EVENT_ADAPTER.validate_json(cast(str, payload)) return event.model_copy(update={"id": event_id, "run_id": run_id}) + @staticmethod + def _decode_cancellation_intent(fields: dict[object, object]) -> RunCancellationIntent: + payload = fields.get(b"payload") or fields.get("payload") + if isinstance(payload, bytes): + payload = payload.decode() + return RunCancellationIntent.model_validate_json(cast(str, payload)) + def _decode_redis_text(value: object) -> str: return value.decode() if isinstance(value, bytes) else str(value) diff --git a/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py b/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py index 396e7736f68..af3a77c687d 100644 --- a/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py +++ b/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py @@ -16,24 +16,36 @@ from agenton.compositor import CompositorSessionSnapshot from dify_agent.protocol.schemas import ( CancelRunRequest, CreateRunRequest, - RunCancelledEvent, - RunCancelledEventData, RunComposition, RunFailedEvent, RunFailedEventData, RunFailureType, + RunStartedEvent, RunSucceededEvent, RunSucceededEventData, + utc_now, ) -from dify_agent.runtime.event_sink import TerminalRunEvent, terminal_event_status_fields +from dify_agent.runtime.cancellation import RunCancellationIntent from dify_agent.runtime.run_scheduler import RunScheduler -from dify_agent.storage.redis_keys import run_events_key, run_record_key +from dify_agent.storage.redis_keys import run_cancel_intent_key, run_events_key, run_record_key from dify_agent.storage.redis_run_store import RedisRunStore pytestmark = pytest.mark.integration +def _success_or_failure_event(kind: str, run_id: str) -> RunSucceededEvent | RunFailedEvent: + if kind == "succeeded": + return RunSucceededEvent( + run_id=run_id, + data=RunSucceededEventData( + output="done", + session_snapshot=CompositorSessionSnapshot(layers=[]), + ), + ) + return RunFailedEvent(run_id=run_id, data=RunFailedEventData(error="model failed")) + + @pytest.fixture def redis_url() -> Iterator[str]: """Start an isolated Redis when the binary is available locally.""" @@ -82,7 +94,7 @@ def redis_url() -> Iterator[str]: process.wait(timeout=5) -def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str) -> None: +def test_success_and_cancel_intent_commit_exactly_one_matching_terminal(redis_url: str) -> None: async def scenario() -> None: first_client = Redis.from_url(redis_url) second_client = Redis.from_url(redis_url) @@ -92,43 +104,42 @@ def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str) second_store = RedisRunStore(second_client, prefix=prefix, run_retention_seconds=retention_seconds) try: record = await first_store.create_run() - terminal_events: tuple[TerminalRunEvent, TerminalRunEvent] = ( - RunSucceededEvent( - run_id=record.run_id, - data=RunSucceededEventData( - output="done", - session_snapshot=CompositorSessionSnapshot(layers=[]), - ), - ), - RunCancelledEvent( - run_id=record.run_id, - data=RunCancelledEventData( - reason="concurrent_cancel", - message="cancel accepted", - ), + success_event = RunSucceededEvent( + run_id=record.run_id, + data=RunSucceededEventData( + output="done", + session_snapshot=CompositorSessionSnapshot(layers=[]), ), ) - results = await asyncio.gather( - first_store.finalize_run(terminal_events[0]), - second_store.finalize_run(terminal_events[1]), + success_result, cancellation_status = await asyncio.gather( + first_store.finalize_run(success_event), + second_store.request_cancellation( + record.run_id, + CancelRunRequest(reason="concurrent_cancel", message="cancel accepted"), + ), ) - assert sum(result.applied for result in results) == 1 - winner_index = next(index for index, result in enumerate(results) if result.applied) - winner_event = terminal_events[winner_index] - winner_result = results[winner_index] - expected_status, expected_error, expected_error_type = terminal_event_status_fields(winner_event) + if success_result.applied: + assert cancellation_status == "succeeded" + winner_result = success_result + expected_status = "succeeded" + expected_event_type = "run_succeeded" + else: + assert success_result.status == "running" + assert cancellation_status == "running" + intent = await first_store.get_cancellation_intent(record.run_id) + assert intent is not None + winner_result = await first_store.finalize_cancellation(record.run_id, intent) + assert winner_result.applied is True + expected_status = "cancelled" + expected_event_type = "run_cancelled" persisted = await first_store.get_run(record.run_id) page = await second_store.get_events(record.run_id) assert persisted.status == expected_status - assert persisted.error == expected_error - assert persisted.error_type == expected_error_type - assert persisted.updated_at == winner_event.created_at assert len(page.events) == 1 - assert page.events[0].type == winner_event.type - assert page.events[0].created_at == winner_event.created_at + assert page.events[0].type == expected_event_type assert page.events[0].id == winner_result.event_id record_ttl = await first_client.ttl(run_record_key(prefix, record.run_id)) @@ -142,6 +153,116 @@ def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str) asyncio.run(scenario()) +@pytest.mark.parametrize("terminal_status", ["succeeded", "failed"]) +def test_terminal_first_rejects_late_cancellation(redis_url: str, terminal_status: str) -> None: + async def scenario() -> None: + client = Redis.from_url(redis_url) + prefix = f"terminal-first-{terminal_status}-{uuid4().hex}" + store = RedisRunStore(client, prefix=prefix, run_retention_seconds=60) + try: + record = await store.create_run() + result = await store.finalize_run(_success_or_failure_event(terminal_status, record.run_id)) + + cancellation_status = await store.request_cancellation( + record.run_id, + CancelRunRequest(reason="late_cancel"), + ) + + assert result.applied is True + assert cancellation_status == terminal_status + assert await store.get_cancellation_intent(record.run_id) is None + events = await store.get_events(record.run_id) + assert [event.type for event in events.events] == [f"run_{terminal_status}"] + finally: + await client.aclose() + + asyncio.run(scenario()) + + +def test_cancellation_intent_lifecycle_and_terminal_exclusion(redis_url: str) -> None: + async def scenario() -> None: + client = Redis.from_url(redis_url) + prefix = f"cancel-intent-lifecycle-{uuid4().hex}" + retention_seconds = 60 + store = RedisRunStore(client, prefix=prefix, run_retention_seconds=retention_seconds) + try: + record = await store.create_run() + _ = await store.append_event(RunStartedEvent(run_id=record.run_id)) + record_key = run_record_key(prefix, record.run_id) + events_key = run_events_key(prefix, record.run_id) + intent_key = run_cancel_intent_key(prefix, record.run_id) + _ = await client.expire(record_key, 1) + _ = await client.expire(events_key, 1) + + first_status = await store.request_cancellation( + record.run_id, + CancelRunRequest(reason="first", message="first message"), + ) + duplicate_status = await store.request_cancellation( + record.run_id, + CancelRunRequest(reason="second", message="second message"), + ) + intent = await store.get_cancellation_intent(record.run_id) + + assert first_status == duplicate_status == "running" + assert intent is not None + assert (intent.reason, intent.message) == ("first", "first message") + for key in (record_key, events_key, intent_key): + assert 0 < await client.ttl(key) <= retention_seconds + + success = await store.finalize_run(_success_or_failure_event("succeeded", record.run_id)) + failure = await store.finalize_run(_success_or_failure_event("failed", record.run_id)) + assert (success.applied, success.status) == (False, "running") + assert (failure.applied, failure.status) == (False, "running") + assert [event.type for event in (await store.get_events(record.run_id)).events] == ["run_started"] + + first_finalization = await store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=CompositorSessionSnapshot(layers=[]), + ) + repeated_finalization = await store.finalize_cancellation(record.run_id, intent) + post_terminal_status = await store.request_cancellation( + record.run_id, + CancelRunRequest(reason="after_finalization"), + ) + events = await store.get_events(record.run_id) + + assert first_finalization.applied is True + assert repeated_finalization.applied is False + assert repeated_finalization.status == "cancelled" + assert post_terminal_status == "cancelled" + assert [event.type for event in events.events].count("run_cancelled") == 1 + assert await client.exists(intent_key) == 0 + for key in (record_key, events_key): + assert 0 < await client.ttl(key) <= retention_seconds + finally: + await client.aclose() + + asyncio.run(scenario()) + + +def test_cancellation_finalization_without_intent_is_unapplied(redis_url: str) -> None: + async def scenario() -> None: + client = Redis.from_url(redis_url) + store = RedisRunStore(client, prefix=f"cancel-without-intent-{uuid4().hex}", run_retention_seconds=60) + try: + record = await store.create_run() + result = await store.finalize_cancellation( + record.run_id, + RunCancellationIntent(reason="not-accepted", requested_at=utc_now()), + ) + + assert result.applied is False + assert result.status == "running" + assert (await store.get_events(record.run_id)).events == [] + assert (await store.get_run(record.run_id)).status == "running" + finally: + await client.aclose() + + asyncio.run(scenario()) + + def test_classified_failure_persists_matching_record_and_event_error_type(redis_url: str) -> None: async def scenario() -> None: client = Redis.from_url(redis_url) @@ -178,6 +299,10 @@ def test_non_owner_scheduler_cancellation_stops_owner_runner(redis_url: str) -> self.started = started self.stopped = stopped + @property + def terminal_session_snapshot(self) -> None: + return None + async def run(self) -> None: self.started.set() try: 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 4fe17b399b0..084cdad4670 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 @@ -1,3 +1,4 @@ +import asyncio import json import unittest from contextlib import asynccontextmanager @@ -6,16 +7,19 @@ from typing import cast from unittest.mock import patch import httpx +import pytest from graphon.model_runtime.entities.message_entities import TextPromptMessageContent -from pydantic_ai.exceptions import ModelHTTPError, UserError +from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior, UserError from pydantic_ai.messages import ( InstructionPart, ModelRequest, ModelResponse, RetryPromptPart, + SpeechPart, SystemPromptPart, TextPart, ThinkingPart, + ToolAvailabilityDeltaPart, ToolCallPart, ToolReturnPart, UserPromptPart, @@ -617,7 +621,7 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase): content="", tool_calls=[ AssistantPromptMessage.ToolCall( - id=None, + id=None, # pyright: ignore[reportArgumentType] type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name="shell_run", @@ -636,7 +640,7 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase): content="", tool_calls=[ AssistantPromptMessage.ToolCall( - id=None, + id=None, # pyright: ignore[reportArgumentType] type="function", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name="shell_run", @@ -762,3 +766,42 @@ class DifyLLMAdapterModelTests(unittest.IsolatedAsyncioTestCase): ) self.assertEqual(str(context.exception), "missing endpoint config") + + +@pytest.mark.parametrize( + "part", + [ + pytest.param(SpeechPart(speaker="user", transcript="hello"), id="speech"), + pytest.param(ToolAvailabilityDeltaPart(tools_added=["lookup"]), id="tool-availability-delta"), + ], +) +def test_request_rejects_unsupported_pydantic_ai_request_parts( + part: SpeechPart | ToolAvailabilityDeltaPart, +) -> None: + async def scenario() -> None: + async with httpx.AsyncClient(trust_env=False) as http_client: + provider = DifyApiLLMProvider( + plugin_id="langgenius/openai", + inner_api_url="http://dify-api", + inner_api_key="inner-secret", + execution_context=DifyExecutionContextLayerConfig( + tenant_id="tenant-1", + user_id="user-123", + user_from="account", + app_id="app-1", + agent_mode="single_step", + invoke_from="debugger", + ), + agent_run_id="run-1", + http_client=http_client, + ) + adapter = DifyLLMAdapterModel("demo-model", provider, model_provider="openai") + + with pytest.raises(UnexpectedModelBehavior, match=type(part).__name__): + _ = await adapter.request( + [ModelRequest(parts=[part])], + model_settings=None, + model_request_parameters=ModelRequestParameters(), + ) + + asyncio.run(scenario()) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/protocol/test_agent_stub_protocol.py b/dify-agent/tests/local/dify_agent/agent_stub/protocol/test_agent_stub_protocol.py index 8a99158b20f..d49a66e44ef 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/protocol/test_agent_stub_protocol.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/protocol/test_agent_stub_protocol.py @@ -15,6 +15,7 @@ from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubConfigDownloadSource, AgentStubFileDownloadRequest, AgentStubFileMapping, + AgentStubFileUploadRequest, agent_stub_connections_url, agent_stub_drive_base_for_ref, agent_stub_drive_commit_url, @@ -54,6 +55,13 @@ def test_agent_stub_file_request_urls_handle_trailing_slash() -> None: ) +def test_agent_stub_file_upload_request_rejects_client_max_size() -> None: + with pytest.raises(ValidationError, match="extra_forbidden"): + AgentStubFileUploadRequest.model_validate( + {"filename": "report.pdf", "mimetype": "application/pdf", "max_size": 1024} + ) + + def test_agent_stub_drive_request_urls_handle_trailing_slash() -> None: assert agent_stub_drive_manifest_url("https://agent.example.com/agent-stub/") == ( "https://agent.example.com/agent-stub/drive/manifest" diff --git a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py index 2f204ac2b71..13e8f34089f 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_files.py @@ -47,6 +47,7 @@ def _file_handler(*, sandbox_files_base_url: str = "https://sandbox-files.exampl inner_api_url="https://api.internal.example.com", inner_api_key="inner-secret", sandbox_files_base_url=sandbox_files_base_url, + max_upload_size_bytes=50 * 1024 * 1024, ) @@ -66,6 +67,7 @@ def test_upload_request_uses_agent_inner_endpoint_and_binds_sandbox_base(monkeyp "filename": "report.pdf", "mimetype": "application/pdf", "conversation_id": "conversation-1", + "max_size": 50 * 1024 * 1024, } return httpx.Response(200, json={"upload_uri": "/files/upload/for-plugin?signed=yes"}) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_routes.py b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_routes.py index 155835f5d25..ef6a4bf6849 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_routes.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/server/test_agent_stub_routes.py @@ -18,7 +18,7 @@ from dify_agent.agent_stub.protocol.agent_stub import ( from dify_agent.agent_stub.server.agent_stub_drive import AgentStubDriveRequestError, AgentStubDriveRequestHandler from dify_agent.agent_stub.server.agent_stub_files import AgentStubFileRequestError, AgentStubFileRequestHandler from dify_agent.agent_stub.server.routes.agent_stub import create_agent_stub_http_router -from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec +from dify_agent.agent_stub.server.tokens.agent_stub import AGENT_STUB_TOKEN_TTL_SECONDS, AgentStubTokenCodec from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig @@ -91,6 +91,81 @@ def test_agent_stub_connections_route_returns_401_for_invalid_bearer_token() -> assert response.json()["detail"] == "invalid or missing Agent Stub authorization" +def test_agent_stub_connections_route_returns_structured_401_for_expired_bearer_token() -> None: + codec = _token_codec() + token = codec.encode_connection_token( + _execution_context(), + now=int(time.time()) - AGENT_STUB_TOKEN_TTL_SECONDS - 1, + ) + app = FastAPI() + app.include_router(create_agent_stub_http_router(codec)) + client = TestClient(app) + + response = client.post( + "/agent-stub/connections", + headers={"Authorization": f"Bearer {token}"}, + json={"protocol_version": 1, "argv": []}, + ) + + assert response.status_code == 401 + assert response.json()["detail"] == { + "code": "agent_stub_authorization_expired", + "message": "Agent Stub authorization expired after 5 minutes; start a new shell tool call and retry the command.", + } + + +def test_agent_stub_file_upload_route_exposes_expiration_when_requested() -> None: + codec = _token_codec() + token = codec.encode_connection_token( + _execution_context(), + now=int(time.time()) - AGENT_STUB_TOKEN_TTL_SECONDS - 1, + ) + app = FastAPI() + app.include_router(create_agent_stub_http_router(codec)) + client = TestClient(app) + headers = {"Authorization": f"Bearer {token}"} + payload = {"filename": "report.pdf", "mimetype": "application/pdf"} + + opted_in_response = client.post( + "/agent-stub/files/upload-request", + headers=headers, + params={"expose_expiration": "true"}, + json=payload, + ) + + assert opted_in_response.status_code == 401 + assert opted_in_response.json()["detail"]["code"] == "agent_stub_authorization_expired" + + +def test_agent_stub_file_download_and_config_routes_return_structured_401_for_expired_token() -> None: + codec = _token_codec() + token = codec.encode_connection_token( + _execution_context(), + now=int(time.time()) - AGENT_STUB_TOKEN_TTL_SECONDS - 1, + ) + app = FastAPI() + app.include_router(create_agent_stub_http_router(codec)) + client = TestClient(app) + headers = {"Authorization": f"Bearer {token}"} + expected_detail = { + "code": "agent_stub_authorization_expired", + "message": "Agent Stub authorization expired after 5 minutes; start a new shell tool call and retry the command.", + } + + responses = [ + client.post( + "/agent-stub/files/download-request", + headers=headers, + json={"file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")}}, + ), + client.get("/agent-stub/config/manifest", headers=headers), + ] + + for response in responses: + assert response.status_code == 401 + assert response.json()["detail"] == expected_detail + + def test_agent_stub_connections_route_returns_503_when_server_has_no_token_codec() -> None: app = FastAPI() app.include_router(create_agent_stub_http_router(None)) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/server/tokens/test_agent_stub.py b/dify-agent/tests/local/dify_agent/agent_stub/server/tokens/test_agent_stub.py index 75823ef501f..ba1f5d6d452 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/server/tokens/test_agent_stub.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/server/tokens/test_agent_stub.py @@ -12,6 +12,7 @@ from dify_agent.agent_stub.server.tokens.agent_stub import ( AGENT_STUB_TOKEN_TTL_SECONDS, AgentStubTokenCodec, AgentStubTokenError, + AgentStubTokenExpiredError, decode_server_secret_key, ) from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig @@ -54,13 +55,26 @@ def test_agent_stub_token_codec_rejects_expired_tokens() -> None: codec = _codec() token = codec.encode_connection_token(_execution_context(), now=1_780_395_720) - with pytest.raises(AgentStubTokenError, match="expired"): + with pytest.raises(AgentStubTokenExpiredError, match="expired"): _ = codec.decode_authorization_header( f"Bearer {token}", - now=1_780_395_720 + AGENT_STUB_TOKEN_TTL_SECONDS + 1, + now=1_780_395_720 + AGENT_STUB_TOKEN_TTL_SECONDS, ) +@pytest.mark.parametrize("expiration_offset", [0, AGENT_STUB_TOKEN_TTL_SECONDS + 1]) +def test_agent_stub_token_codec_rejects_invalid_declared_lifetime(expiration_offset: int) -> None: + codec = _codec() + issued_at = 1_780_395_720 + claims = codec.build_connection_claims(_execution_context(), now=issued_at) + token = codec.encode_claims(claims.model_copy(update={"exp": issued_at + expiration_offset})) + + with pytest.raises(AgentStubTokenError) as exc_info: + _ = codec.decode_authorization_header(f"Bearer {token}", now=issued_at) + + assert not isinstance(exc_info.value, AgentStubTokenExpiredError) + + def test_agent_stub_token_codec_rejects_tokens_before_nbf() -> None: codec = _codec() claims = codec.build_connection_claims(_execution_context(), now=1_780_395_720) @@ -86,6 +100,13 @@ def test_agent_stub_token_codec_rejects_wrong_audience_and_scope() -> None: with pytest.raises(AgentStubTokenError, match=AGENT_STUB_TOKEN_SCOPE_CONNECT): _ = codec.decode_authorization_header(f"Bearer {wrong_scope_token}", now=1_780_395_720) + with pytest.raises(AgentStubTokenError, match=AGENT_STUB_TOKEN_SCOPE_CONNECT) as exc_info: + _ = codec.decode_authorization_header( + f"Bearer {wrong_scope_token}", + now=1_780_395_720 + AGENT_STUB_TOKEN_TTL_SECONDS, + ) + assert not isinstance(exc_info.value, AgentStubTokenExpiredError) + def test_agent_stub_token_codec_rejects_wrong_key_and_malformed_authorization_header() -> None: codec = _codec() @@ -107,6 +128,7 @@ def test_agent_stub_token_codec_builds_fixed_server_claims() -> None: assert claims.iss == AGENT_STUB_TOKEN_ISSUER assert claims.aud == AGENT_STUB_TOKEN_AUDIENCE assert claims.scope == [AGENT_STUB_TOKEN_SCOPE_CONNECT] + assert AGENT_STUB_TOKEN_TTL_SECONDS == 300 assert claims.exp - claims.iat == AGENT_STUB_TOKEN_TTL_SECONDS assert claims.shell is not None assert claims.shell.session_id == "abc12ff" 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 e69bec36518..54c905625ac 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/dify-agent/tests/local/dify_agent/client/test_client.py @@ -34,6 +34,7 @@ from dify_agent.protocol import ( DestroyExecutionBindingRequest, RUN_EVENT_ADAPTER, RunCancelledEvent, + RunCancelledEventData, RunEvent, RunEventsResponse, RunFailedEvent, @@ -155,7 +156,7 @@ def test_sse_decoder_accepts_function_tool_result_part_alias(monkeypatch: pytest assert event is not None assert event.type == "pydantic_ai_event" assert event.data.event_kind == "function_tool_result" - assert event.data.result.tool_name == "shell_run" + assert event.data.part.tool_name == "shell_run" def test_function_tool_result_payload_normalization_supports_old_part_schema( @@ -245,6 +246,109 @@ def test_async_methods_and_wait_run_parse_protocol_dtos() -> None: asyncio.run(scenario()) +def test_cancel_run_and_wait_sync_resumes_after_cursor_and_returns_cancelled_snapshot() -> None: + snapshot = CompositorSessionSnapshot(layers=[]) + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"}) + if request.url.path == "/runs/run-1": + return httpx.Response( + 200, + json={ + "run_id": "run-1", + "status": "running", + "created_at": "2026-08-17T00:00:00Z", + "updated_at": "2026-08-17T00:00:00Z", + }, + ) + assert request.url.params["after"] == "3-0" + event = RunCancelledEvent( + id="4-0", + run_id="run-1", + data=RunCancelledEventData(reason="stopped", session_snapshot=snapshot), + ) + return httpx.Response(200, content=_event_frame(event)) + + client = Client( + base_url="http://testserver", + sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + event = client.cancel_run_and_wait_sync( + "run-1", + CancelRunRequest(reason="stopped"), + after="3-0", + ) + + assert event.data.session_snapshot == snapshot + + +def test_cancel_run_and_wait_sync_replays_when_cursor_already_points_to_cancelled_event() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"}) + if request.url.path == "/runs/run-1": + return httpx.Response( + 200, + json={ + "run_id": "run-1", + "status": "cancelled", + "created_at": "2026-08-17T00:00:00Z", + "updated_at": "2026-08-17T00:00:01Z", + }, + ) + assert request.url.params["after"] == "0-0" + return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="4-0", run_id="run-1"))) + + client = Client( + base_url="http://testserver", + sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + event = client.cancel_run_and_wait_sync("run-1", after="4-0") + + assert event.id == "4-0" + + +def test_cancel_run_and_wait_async_returns_cancelled_terminal() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"}) + if request.url.path == "/runs/run-1": + return httpx.Response(200, json=_run_status_json("running")) + assert request.url.params["after"] == "1-0" + return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="2-0", run_id="run-1"))) + + async def scenario() -> None: + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + client = Client(base_url="http://testserver", async_http_client=http_client) + event = await client.cancel_run_and_wait("run-1", after="1-0") + assert event.type == "run_cancelled" + await http_client.aclose() + + asyncio.run(scenario()) + + +def test_cancel_run_and_wait_async_replays_when_cursor_already_points_to_cancelled_event() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"}) + if request.url.path == "/runs/run-1": + return httpx.Response(200, json=_run_status_json("cancelled")) + assert request.url.params["after"] == "0-0" + return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="4-0", run_id="run-1"))) + + async def scenario() -> None: + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + client = Client(base_url="http://testserver", async_http_client=http_client) + event = await client.cancel_run_and_wait("run-1", after="4-0") + assert event.id == "4-0" + await http_client.aclose() + + asyncio.run(scenario()) + + def test_sync_binding_file_methods_post_dtos_and_parse_responses() -> None: def handler(request: httpx.Request) -> httpx.Response: if request.url.path == "/execution-bindings/files/list": 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 990a19eb477..ce9435a4b46 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 @@ -133,6 +133,7 @@ async def test_on_context_create_computes_runtime_fields_and_pulls_mentioned_ass 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 public-url --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() @@ -140,8 +141,12 @@ async def test_on_context_create_computes_runtime_fields_and_pulls_mentioned_ass "Agent file CLI reference for installed `dify-agent`:" ) assert "$ dify-agent file upload --help" in suffix_prompt + assert "$ dify-agent file public-url --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 public-url --help" + ) + assert suffix_prompt.index("$ dify-agent file public-url --help") < suffix_prompt.index( "$ dify-agent file download --help" ) assert _AGENT_FILE_UPLOAD_REPLY_HINT in suffix_prompt diff --git a/dify-agent/tests/local/dify_agent/layers/dify_core_tools/test_layer.py b/dify-agent/tests/local/dify_agent/layers/dify_core_tools/test_layer.py index 4f46cd0a394..f5813c0f7f8 100644 --- a/dify-agent/tests/local/dify_agent/layers/dify_core_tools/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/dify_core_tools/test_layer.py @@ -18,8 +18,8 @@ def _install_graphon_stubs() -> None: 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", (), {}) + setattr(llm_entities_module, "LLMResultChunk", type("LLMResultChunk", (), {})) + setattr(llm_entities_module, "LLMUsage", type("LLMUsage", (), {})) for name in ( "AssistantPromptMessage", @@ -43,10 +43,10 @@ def _install_graphon_stubs() -> None: 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 + setattr(graphon_module, "model_runtime", model_runtime_module) + setattr(model_runtime_module, "entities", entities_module) + setattr(entities_module, "llm_entities", llm_entities_module) + setattr(entities_module, "message_entities", message_entities_module) _install_graphon_stubs() diff --git a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_configs.py b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_configs.py index 898249702a4..2f9f6880564 100644 --- a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_configs.py +++ b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_configs.py @@ -44,6 +44,7 @@ def test_dify_plugin_llm_config_discards_legacy_credentials() -> None: "model": "gpt-4o-mini", "credentials": {"api_key": "secret", "nested": {"legacy": True}}, "model_settings": {"temperature": 0.2, "max_tokens": 64}, + "context_window_tokens": 128_000, } ) @@ -52,6 +53,7 @@ def test_dify_plugin_llm_config_discards_legacy_credentials() -> None: assert not hasattr(config, "credentials") assert "credentials" not in config.model_dump(mode="json") assert config.model_settings == {"temperature": 0.2, "max_tokens": 64} + assert config.context_window_tokens == 128_000 def test_dify_plugin_llm_config_rejects_old_provider_field() -> None: @@ -65,6 +67,16 @@ def test_dify_plugin_llm_config_rejects_old_provider_field() -> None: ) +def test_dify_plugin_llm_config_rejects_non_positive_context_window() -> None: + with pytest.raises(ValidationError): + _ = DifyPluginLLMLayerConfig( + plugin_id="langgenius/openai", + model_provider="openai", + model="gpt-4o-mini", + context_window_tokens=0, + ) + + def test_dify_plugin_tools_layer_config_accepts_prepared_parameters_and_schema() -> None: runtime_value: DifyPluginToolValue = {"locale": "en-US", "max_results": 5} credential_type: DifyPluginToolCredentialType = "api-key" diff --git a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py index abe020cc792..a47efd75437 100644 --- a/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py +++ b/dify-agent/tests/local/dify_agent/layers/dify_plugin/test_layers.py @@ -848,9 +848,9 @@ def test_plugin_tool_file_context_uploads_sandbox_path_and_resolves_signed_url() async def scenario() -> None: shell = FakeShell() context = _PluginToolFileContext( - file_client=FakeFileClient(), # type: ignore[arg-type] + file_client=FakeFileClient(), # pyright: ignore[reportArgumentType] execution_context=_execution_context_config(), - shell=shell, # type: ignore[arg-type] + shell=shell, # pyright: ignore[reportArgumentType] ) result = await context.to_plugin_file_parameter("outputs/report.pdf") diff --git a/dify-agent/tests/local/dify_agent/layers/knowledge/test_layer.py b/dify-agent/tests/local/dify_agent/layers/knowledge/test_layer.py index ed6c798b409..17134bf3844 100644 --- a/dify-agent/tests/local/dify_agent/layers/knowledge/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/knowledge/test_layer.py @@ -99,8 +99,10 @@ def test_knowledge_layer_exposes_one_set_scoped_tool_definition() -> None: tool_def = await tool.prepare_tool_def(None) # pyright: ignore[reportArgumentType] assert isinstance(tool, Tool) assert tool.name == "knowledge_base_search" + assert tool.description is not None assert "Pick one configured set_name" in tool.description assert tool_def is not None + assert tool_def.description is not None assert "Pick one configured set_name" in tool_def.description assert tool_def.parameters_json_schema == { "type": "object", @@ -140,7 +142,8 @@ def test_knowledge_layer_rejects_blank_query_locally() -> None: knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": " "}, None + {"set_name": "Support KB", "query": " "}, + None, # pyright: ignore[reportArgumentType] ) assert result == BLANK_QUERY_OBSERVATION @@ -313,7 +316,8 @@ def test_knowledge_layer_formats_results_and_truncates_observation() -> None: knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert result.startswith("Knowledge base search results:\n1. Title: Guide") assert "Dataset: Docs" in result @@ -345,7 +349,8 @@ def test_knowledge_layer_returns_no_results_observation() -> None: knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert result == NO_RESULTS_OBSERVATION @@ -374,7 +379,8 @@ def test_knowledge_layer_converts_retryable_failures_into_observation() -> None: knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert result == TEMPORARY_UNAVAILABLE_OBSERVATION @@ -409,7 +415,8 @@ def test_knowledge_layer_converts_retryable_transport_failures_into_observation( knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert result == TEMPORARY_UNAVAILABLE_OBSERVATION @@ -439,7 +446,8 @@ def test_knowledge_layer_raises_non_retryable_client_errors() -> None: tool = (await knowledge_layer.get_tools(http_client=http_client))[0] with pytest.raises(DifyKnowledgeBaseClientError) as exc_info: await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert exc_info.value.status_code == 403 @@ -467,7 +475,8 @@ def test_knowledge_layer_raises_for_malformed_success_responses() -> None: tool = (await knowledge_layer.get_tools(http_client=http_client))[0] with pytest.raises(DifyKnowledgeBaseClientError) as exc_info: await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert exc_info.value.error_code == "invalid_response" assert exc_info.value.retryable is False @@ -537,7 +546,8 @@ def test_knowledge_layer_sends_execution_context_and_static_config_to_inner_api( knowledge_layer = run.get_layer("knowledge", DifyKnowledgeBaseLayer) tool = (await knowledge_layer.get_tools(http_client=http_client))[0] result = await tool.function_schema.call( # pyright: ignore[reportArgumentType] - {"set_name": "Support KB", "query": "reset"}, None + {"set_name": "Support KB", "query": "reset"}, + None, # pyright: ignore[reportArgumentType] ) assert result == NO_RESULTS_OBSERVATION diff --git a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py index d8b07395774..b3362d2c033 100644 --- a/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/shell/test_layer.py @@ -354,6 +354,7 @@ def test_shell_layer_suspend_cleans_tracked_jobs_without_owning_sandbox() -> Non commands = FakeCommands() layer, _provider = _layer(commands=commands) layer.runtime_state = _runtime_state() + layer._job_agent_stub_tokens.update({"job-1": "token-1", "job-2": "token-2"}) async def scenario() -> None: async with layer.resource_context(): @@ -361,6 +362,7 @@ def test_shell_layer_suspend_cleans_tracked_jobs_without_owning_sandbox() -> Non asyncio.run(scenario()) assert commands.delete_calls == [] + assert layer._job_agent_stub_tokens == {} def test_shell_layer_resume_requires_active_sandbox_lease_only() -> None: @@ -383,6 +385,7 @@ def test_shell_layer_delete_cleans_tracked_jobs_without_deleting_workspace() -> layer, _provider = _layer(commands=commands) _bind_execution_context(layer) layer.runtime_state = _runtime_state(job_ids=["user-job"], job_offsets={"user-job": 9}) + layer._job_agent_stub_tokens["user-job"] = "long-lived-test-token" async def scenario() -> None: async with layer.resource_context(): @@ -393,6 +396,7 @@ def test_shell_layer_delete_cleans_tracked_jobs_without_deleting_workspace() -> assert [call.job_id for call in commands.delete_calls] == ["user-job"] assert layer.runtime_state.job_ids == [] assert layer.runtime_state.job_offsets == {} + assert layer._job_agent_stub_tokens == {} def test_shell_layer_tools_map_inputs_and_maintain_offsets_with_tail_end() -> None: @@ -582,6 +586,7 @@ def test_shell_interrupt_succeeds_when_tail_lookup_fails() -> None: layer, _provider = _layer(commands=commands) tools = {tool.name: tool for tool in layer.tools} layer.runtime_state = _runtime_state(job_ids=["user-job"], job_offsets={"user-job": 22}) + layer._job_agent_stub_tokens["user-job"] = "actual-long-jwe-token" async def scenario() -> None: async with layer.resource_context(): @@ -596,6 +601,7 @@ def test_shell_interrupt_succeeds_when_tail_lookup_fails() -> None: assert output == "Job was interrupted." asyncio.run(scenario()) + assert layer._job_agent_stub_tokens == {} def test_shell_run_returns_provider_timeout_error_observation_without_unexpected_logging( @@ -1193,11 +1199,20 @@ def _layer_with_redaction( return layer, provider -def test_redact_output_replaces_jwe_token_value() -> None: - """The JWE token value should always be redacted from shell output.""" - token = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0.super-secret-token-12345" +def test_shell_run_issues_one_jwe_and_redaction_does_not_issue_another() -> None: + """The output filter must use the exact JWE injected into the shell job.""" + issued_tokens: list[str] = [] + + def token_factory(execution_context: DifyExecutionContextLayerConfig, session_id: str | None) -> str: + del execution_context, session_id + token = f"actual-long-jwe-token-{len(issued_tokens) + 1}" + issued_tokens.append(token) + return token def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult: + del script, cwd, timeout + assert env is not None + token = env["DIFY_AGENT_STUB_AUTH_JWE"] return _command_result( "job-1", status="exited", @@ -1207,11 +1222,9 @@ def test_redact_output_replaces_jwe_token_value() -> None: offset=100, ) - commands = FakeCommands( - run_handler=run_handler, - tail_handler=lambda _: _command_result("job-1", done=True, exit_code=0, status="exited", offset=100), - ) - layer, _provider = _layer_with_redaction(commands=commands, token_value=token) + commands = FakeCommands(run_handler=run_handler) + layer, _provider = _layer_with_redaction(commands=commands) + layer.agent_stub_token_factory = token_factory _bind_execution_context(layer) layer.runtime_state = _runtime_state() tools = {tool.name: tool for tool in layer.tools} @@ -1220,8 +1233,80 @@ def test_redact_output_replaces_jwe_token_value() -> None: async with layer.resource_context(): result = await tools["shell_run"].function_schema.call({"script": "env"}, None) # pyright: ignore[reportArgumentType] _, output = _parse_tagged_observation(result) - assert token not in output - assert "***" in output + assert issued_tokens[0] not in output + assert "DIFY_AGENT_STUB_AUTH_JWE=***" in output + + asyncio.run(scenario()) + assert issued_tokens == ["actual-long-jwe-token-1"] + assert layer._job_agent_stub_tokens == {} + + +def test_two_running_jobs_redact_and_clean_up_their_own_jwes_independently() -> None: + tokens = ["actual-long-jwe-token-for-job-1", "actual-long-jwe-token-for-job-2"] + issued_tokens: list[str] = [] + + def token_factory(execution_context: DifyExecutionContextLayerConfig, session_id: str | None) -> str: + del execution_context, session_id + token = tokens[len(issued_tokens)] + issued_tokens.append(token) + return token + + def run_handler( + script: str, + cwd: str | None, + env: Mapping[str, str] | None, + timeout: float, + ) -> ShellCommandResult: + del script, cwd, timeout + assert env is not None + token = env["DIFY_AGENT_STUB_AUTH_JWE"] + job_id = "job-1" if token == tokens[0] else "job-2" + return _command_result(job_id, status="running", done=False, output="started\n", offset=8) + + commands = FakeCommands( + run_handler=run_handler, + wait_handler=lambda job_id, offset, timeout: _command_result( + job_id, + status="exited", + done=True, + exit_code=0, + output=f"wait={tokens[0]}\n", + offset=20, + ), + input_handler=lambda job_id, text, offset, timeout: _command_result( + job_id, + status="exited", + done=True, + exit_code=0, + output=f"input={tokens[1]}\n", + offset=40, + ), + ) + layer, _provider = _layer_with_redaction(commands=commands) + layer.agent_stub_token_factory = token_factory + _bind_execution_context(layer) + layer.runtime_state = _runtime_state() + tools = {tool.name: tool for tool in layer.tools} + + async def scenario() -> None: + async with layer.resource_context(): + await tools["shell_run"].function_schema.call({"script": "start first"}, None) # pyright: ignore[reportArgumentType] + await tools["shell_run"].function_schema.call({"script": "start second"}, None) # pyright: ignore[reportArgumentType] + assert issued_tokens == tokens + assert layer._job_agent_stub_tokens == {"job-1": tokens[0], "job-2": tokens[1]} + + wait_result = await tools["shell_wait"].function_schema.call({"job_id": "job-1"}, None) # pyright: ignore[reportArgumentType] + _, wait_output = _parse_tagged_observation(wait_result) + assert wait_output == "wait=***\n" + assert layer._job_agent_stub_tokens == {"job-2": tokens[1]} + + input_result = await tools["shell_input"].function_schema.call( + {"job_id": "job-2", "text": "finish\n"}, + None, # pyright: ignore[reportArgumentType] + ) + _, input_output = _parse_tagged_observation(input_result) + assert input_output == "input=***\n" + assert layer._job_agent_stub_tokens == {} asyncio.run(scenario()) 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 8eac3a6023e..fbbf9077d73 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 @@ -77,7 +77,12 @@ def test_run_event_adapter_round_trips_typed_variants() -> None: reason="shutdown", ), ), - RunCancelledEvent(run_id="run-1", data=RunCancelledEventData(reason="user_cancelled")), + RunCancelledEvent( + run_id="run-1", + data=RunCancelledEventData( + reason="user_cancelled", + ), + ), ] for event in events: @@ -113,6 +118,27 @@ def test_run_failed_event_error_type_is_optional_and_round_trips() -> None: assert protocol_exports.RunFailureType is RunFailureType +@pytest.mark.parametrize("event_type", ["run_failed", "run_cancelled"]) +def test_non_success_terminal_event_round_trips_optional_snapshot(event_type: str) -> None: + snapshot = CompositorSessionSnapshot(layers=[]) + event: RunFailedEvent | RunCancelledEvent + if event_type == "run_failed": + event = RunFailedEvent( + run_id="run-1", + data=RunFailedEventData(error="boom", session_snapshot=snapshot), + ) + else: + event = RunCancelledEvent( + run_id="run-1", + data=RunCancelledEventData(reason="stopped", session_snapshot=snapshot), + ) + + decoded = RUN_EVENT_ADAPTER.validate_json(RUN_EVENT_ADAPTER.dump_json(event)) + + assert isinstance(decoded, RunFailedEvent | RunCancelledEvent) + assert decoded.data.session_snapshot == snapshot + + def test_pydantic_ai_event_data_uses_agent_stream_event_model() -> None: event = RUN_EVENT_ADAPTER.validate_python( { diff --git a/dify-agent/tests/local/dify_agent/runtime/test_compaction.py b/dify-agent/tests/local/dify_agent/runtime/test_compaction.py new file mode 100644 index 00000000000..4bd7ba54758 --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime/test_compaction.py @@ -0,0 +1,132 @@ +import pytest +from pydantic_ai import Agent +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + SystemPromptPart, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models.test import TestModel +from pydantic_ai_harness.compaction import ClearToolResults, SummarizingCompaction, TieredCompaction + +from dify_agent.runtime.compaction import build_compaction_capability + + +def test_build_compaction_capability_uses_effective_input_budget_and_standard_tiers() -> None: + capability = build_compaction_capability( + context_window_tokens=10_000, + model_settings={"max_tokens": 3_000}, + ) + + assert isinstance(capability, TieredCompaction) + assert capability.target_tokens == 7_000 + assert len(capability.tiers) == 2 + assert isinstance(capability.tiers[0], ClearToolResults) + assert capability.tiers[0].keep_pairs == 3 + assert capability.tiers[0].clear_tool_inputs is False + assert isinstance(capability.tiers[1], SummarizingCompaction) + assert capability.tiers[1].model is None + assert capability.tiers[1].keep_messages == 20 + assert capability.tiers[1].preserve_first_user_message is True + assert capability.tiers[1].incremental is True + + +def test_build_compaction_capability_uses_default_budget_and_handles_unknown_window() -> None: + capability = build_compaction_capability(context_window_tokens=10_001, model_settings=None) + + assert isinstance(capability, TieredCompaction) + assert capability.target_tokens == 8_000 + assert build_compaction_capability(context_window_tokens=None, model_settings=None) is None + + +@pytest.mark.parametrize( + "max_tokens", + [ + pytest.param(1_000, id="default-budget-wins"), + pytest.param(0, id="zero-is-ignored"), + pytest.param(-1, id="negative-is-ignored"), + ], +) +def test_build_compaction_capability_uses_default_budget_when_output_reservation_is_smaller( + max_tokens: int, +) -> None: + capability = build_compaction_capability( + context_window_tokens=10_000, + model_settings={"max_tokens": max_tokens}, + ) + + assert isinstance(capability, TieredCompaction) + assert capability.target_tokens == 8_000 + + +def test_build_compaction_capability_rejects_output_budget_that_consumes_window() -> None: + with pytest.raises(ValueError, match="Model max_tokens must leave a positive input context budget"): + _ = build_compaction_capability( + context_window_tokens=1_000, + model_settings={"max_tokens": 1_000}, + ) + + +def test_compaction_clears_only_tool_results_older_than_the_last_three_pairs() -> None: + history: list[ModelRequest | ModelResponse] = [] + for index in range(4): + tool_call_id = f"call-{index}" + history.extend( + [ + ModelResponse(parts=[ToolCallPart("lookup", {"query": index}, tool_call_id)]), + ModelRequest(parts=[ToolReturnPart("lookup", "x" * 4_000, tool_call_id)]), + ] + ) + + capability = build_compaction_capability(context_window_tokens=4_100, model_settings=None) + assert capability is not None + agent = Agent[None, str](TestModel(call_tools=[]), deps_type=type(None)) + result = agent.run_sync("next", message_history=history, capabilities=[capability]) + + tool_returns = [ + part + for message in result.all_messages() + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, ToolReturnPart) + ] + assert [part.content for part in tool_returns] == ["[tool result cleared]", *("x" * 4_000 for _ in range(3))] + + +def test_compaction_summary_is_present_in_full_run_history() -> None: + history: list[ModelRequest | ModelResponse] = [] + for index in range(30): + history.extend( + [ + ModelRequest(parts=[UserPromptPart(f"user-{index}-" + "u" * 120)]), + ModelResponse(parts=[TextPart(f"assistant-{index}-" + "a" * 120)], model_name="test"), + ] + ) + + capability = build_compaction_capability(context_window_tokens=1_000, model_settings=None) + assert capability is not None + agent = Agent[None, str]( + TestModel(call_tools=[], custom_output_text="summary body"), + deps_type=type(None), + ) + result = agent.run_sync( + "next", + message_history=history, + capabilities=[capability], + ) + + messages = result.all_messages() + assert len(messages) < len(history) + assert isinstance(messages[0], ModelRequest) + assert len(messages[0].parts) == 1 + assert isinstance(messages[0].parts[0], SystemPromptPart) + assert messages[0].parts[0].content == "Summary of previous conversation:\n\nsummary body" + assert any( + isinstance(part, UserPromptPart) and str(part.content).startswith("user-0-") + for message in messages + if isinstance(message, ModelRequest) + for part in message.parts + ) 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 11b40b20bc7..afff9519fc9 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 @@ -9,8 +9,8 @@ if "graphon.model_runtime.entities.llm_entities" not in sys.modules: 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", (), {}) + setattr(llm_entities_module, "LLMResultChunk", type("LLMResultChunk", (), {})) + setattr(llm_entities_module, "LLMUsage", type("LLMUsage", (), {})) for name in ( "AssistantPromptMessage", @@ -34,10 +34,10 @@ if "graphon.model_runtime.entities.llm_entities" not in sys.modules: 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 + setattr(graphon_module, "model_runtime", model_runtime_module) + setattr(model_runtime_module, "entities", entities_module) + setattr(entities_module, "llm_entities", llm_entities_module) + setattr(entities_module, "message_entities", message_entities_module) if "jsonschema" not in sys.modules: jsonschema_module = types.ModuleType("jsonschema") @@ -65,10 +65,10 @@ if "jsonschema" not in sys.modules: 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 + setattr(jsonschema_module, "SchemaError", _SchemaError) + setattr(jsonschema_exceptions_module, "ValidationError", _ValidationError) + setattr(jsonschema_protocols_module, "Validator", _Validator) + setattr(jsonschema_validators_module, "validator_for", _validator_for) sys.modules["jsonschema"] = jsonschema_module sys.modules["jsonschema.exceptions"] = jsonschema_exceptions_module @@ -95,8 +95,8 @@ class FakeProvider: def _runtime_backend_profile() -> RuntimeBackendProfile: return RuntimeBackendProfile( - home_snapshots=cast(HomeSnapshotBackend, FakeProvider()), - execution_bindings=cast(ExecutionBindingBackend, FakeProvider()), + home_snapshots=cast(HomeSnapshotBackend, cast(object, FakeProvider())), + execution_bindings=cast(ExecutionBindingBackend, cast(object, FakeProvider())), ) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py index f7fc7a1e94e..44643d6a312 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py @@ -15,11 +15,16 @@ from dify_agent.protocol import DIFY_AGENT_MODEL_LAYER_ID, DIFY_AGENT_OUTPUT_LAY from dify_agent.protocol.schemas import ( CancelRunRequest, CreateRunRequest, + RunCancelledEvent, + RunCancelledEventData, RunComposition, RunEvent, + RunFailedEvent, RunLayerSpec, RunStatus, + utc_now, ) +from dify_agent.runtime.cancellation import RunCancellationIntent from dify_agent.runtime.event_sink import ( NonTerminalRunEvent, RunFinalizationResult, @@ -95,6 +100,7 @@ class FakeStore: errors: dict[str, str | None] error_types: dict[str, RunFailureType | None] terminal_changes: dict[str, asyncio.Event] + cancellation_intents: dict[str, RunCancellationIntent] def __init__(self) -> None: self.records = {} @@ -103,6 +109,7 @@ class FakeStore: self.errors = {} self.error_types = {} self.terminal_changes = {} + self.cancellation_intents = {} async def create_run(self) -> RunRecord: run_id = f"run-{len(self.records) + 1}" @@ -130,6 +137,8 @@ class FakeStore: current_status = self.statuses[event.run_id] if current_status != "running": return RunFinalizationResult(applied=False, status=current_status) + if event.run_id in self.cancellation_intents: + return RunFinalizationResult(applied=False, status="running") status, error, error_type = terminal_event_status_fields(event) event_id = str(len(self.events[event.run_id]) + 1) @@ -140,10 +149,55 @@ class FakeStore: self.terminal_changes[event.run_id].set() return RunFinalizationResult(applied=True, status=status, event_id=event_id) - async def wait_for_cancellation(self, run_id: str) -> bool: - while self.statuses[run_id] == "running": + async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus: + status = self.statuses[run_id] + if status != "running": + return status + if run_id not in self.cancellation_intents: + self.cancellation_intents[run_id] = RunCancellationIntent( + reason=request.reason, + message=request.message, + requested_at=utc_now(), + ) + self.terminal_changes[run_id].set() + return "running" + + async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None: + return self.cancellation_intents.get(run_id) + + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: + while self.statuses[run_id] == "running" and run_id not in self.cancellation_intents: await self.terminal_changes[run_id].wait() - return self.statuses[run_id] == "cancelled" + return self.cancellation_intents.get(run_id) + + async def finalize_cancellation( + self, + run_id: str, + intent: RunCancellationIntent, + *, + session_snapshot: CompositorSessionSnapshot | None = None, + ) -> RunFinalizationResult: + current_status = self.statuses[run_id] + if current_status != "running": + return RunFinalizationResult(applied=False, status=current_status) + if run_id not in self.cancellation_intents: + return RunFinalizationResult(applied=False, status="running") + event = RunCancelledEvent( + run_id=run_id, + data=RunCancelledEventData( + reason=intent.reason, + message=intent.message, + session_snapshot=session_snapshot, + ), + ) + event_id = str(len(self.events[run_id]) + 1) + self.events[run_id].append(event.model_copy(update={"id": event_id})) + self.statuses[run_id] = "cancelled" + self.errors[run_id] = intent.message or intent.reason + self.error_types[run_id] = None + del self.cancellation_intents[run_id] + self.terminal_changes[run_id].set() + return RunFinalizationResult(applied=True, status="cancelled", event_id=event_id) class SlowCreateStore(FakeStore): @@ -174,7 +228,7 @@ class TrackingStore(FakeStore): if not pause_observer: self.release_observer.set() - async def wait_for_cancellation(self, run_id: str) -> bool: + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: self.observer_started.set() try: await self.release_observer.wait() @@ -192,7 +246,7 @@ class FailingObserverStore(FakeStore): self.fail_observer = fail_observer self.observer_finished = asyncio.Event() - async def wait_for_cancellation(self, run_id: str) -> bool: + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: del run_id try: await self.fail_observer.wait() @@ -201,10 +255,27 @@ class FailingObserverStore(FakeStore): self.observer_finished.set() +class CancellationDuringShutdownFailureStore(FakeStore): + async def finalize_run(self, event: TerminalRunEvent) -> RunFinalizationResult: + if isinstance(event, RunFailedEvent) and event.data.reason == "shutdown": + _ = await self.request_cancellation( + event.run_id, + CancelRunRequest(reason="concurrent_shutdown_cancel"), + ) + return await super().finalize_run(event) + + +class SnapshotlessRunner: + @property + def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None: + return None + + class ControlledRunner: started: asyncio.Event release: asyncio.Event finished: asyncio.Event | None + _terminal_session_snapshot: CompositorSessionSnapshot def __init__( self, @@ -216,6 +287,11 @@ class ControlledRunner: self.started = started self.release = release self.finished = finished + self._terminal_session_snapshot = CompositorSessionSnapshot(layers=[]) + + @property + def terminal_session_snapshot(self) -> CompositorSessionSnapshot: + return self._terminal_session_snapshot async def run(self) -> None: _ = self.started.set() @@ -226,24 +302,16 @@ class ControlledRunner: self.finished.set() -class SwallowOneCancellationRunner: - started: asyncio.Event - first_cancellation: asyncio.Event - - def __init__(self, *, started: asyncio.Event, first_cancellation: asyncio.Event) -> None: +class PreEnterBlockingRunner(SnapshotlessRunner): + def __init__(self, *, started: asyncio.Event) -> None: self.started = started - self.first_cancellation = first_cancellation async def run(self) -> None: - _ = self.started.set() - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - _ = self.first_cancellation.set() - await asyncio.Event().wait() + self.started.set() + await asyncio.Event().wait() -class SuccessThenWaitRunner: +class SuccessThenWaitRunner(SnapshotlessRunner): def __init__( self, *, @@ -269,7 +337,7 @@ class SuccessThenWaitRunner: await self.release.wait() -class IgnoreCancellationThenSucceedRunner: +class IgnoreCancellationThenSucceedRunner(SnapshotlessRunner): def __init__( self, *, @@ -300,12 +368,12 @@ class IgnoreCancellationThenSucceedRunner: session_snapshot=CompositorSessionSnapshot(layers=[]), ) assert result.applied is False - assert result.status == "cancelled" + assert result.status == "running" finally: self.finished.set() -class ReleaseThenSucceedRunner: +class ReleaseThenSucceedRunner(SnapshotlessRunner): def __init__( self, *, @@ -336,7 +404,7 @@ class ReleaseThenSucceedRunner: self.finished.set() -class CompetingFailureRunner: +class CompetingFailureRunner(SnapshotlessRunner): def __init__( self, *, @@ -362,7 +430,7 @@ class CompetingFailureRunner: self.failure_attempted.set() -class FinalizeSuccessOnCancellationRunner: +class FinalizeSuccessOnCancellationRunner(SnapshotlessRunner): def __init__(self, *, store: FakeStore, run_id: str, started: asyncio.Event) -> None: self.store = store self.run_id = run_id @@ -454,6 +522,37 @@ def test_shutdown_marks_unfinished_runs_failed_and_appends_event() -> None: asyncio.run(scenario()) +def test_shutdown_failure_finalization_yields_to_concurrent_cancellation_intent() -> None: + async def scenario() -> None: + store = CancellationDuringShutdownFailureStore() + started = asyncio.Event() + async with httpx.AsyncClient() as client: + scheduler = RunScheduler( + store=store, + plugin_daemon_http_client=client, + dify_api_http_client=client, + shutdown_grace_seconds=0, + runner_factory=lambda _record, _request: ControlledRunner( + started=started, + release=asyncio.Event(), + ), + ) + record = await scheduler.create_run(_request()) + await asyncio.wait_for(started.wait(), timeout=1) + + await scheduler.shutdown() + + assert store.statuses[record.run_id] == "cancelled" + assert record.run_id not in store.cancellation_intents + assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"] + terminal = store.events[record.run_id][0] + assert isinstance(terminal, RunCancelledEvent) + assert terminal.data.reason == "concurrent_shutdown_cancel" + assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[]) + + asyncio.run(scenario()) + + def test_cancellation_observer_failure_stops_runner_and_finalizes_failed() -> None: async def scenario() -> None: fail_observer = asyncio.Event() @@ -489,6 +588,46 @@ def test_cancellation_observer_failure_stops_runner_and_finalizes_failed() -> No asyncio.run(scenario()) +def test_cancellation_observer_failure_finalizes_concurrent_intent_after_runner_exit() -> None: + async def scenario() -> None: + fail_observer = asyncio.Event() + store = FailingObserverStore(fail_observer=fail_observer) + runner_started = asyncio.Event() + runner_finished = asyncio.Event() + async with httpx.AsyncClient() as client: + scheduler = RunScheduler( + store=store, + plugin_daemon_http_client=client, + dify_api_http_client=client, + runner_factory=lambda _record, _request: ControlledRunner( + started=runner_started, + release=asyncio.Event(), + finished=runner_finished, + ), + ) + record = await scheduler.create_run(_request()) + supervisor_task = scheduler.active_tasks[record.run_id] + await asyncio.wait_for(runner_started.wait(), timeout=1) + + response = await scheduler.cancel_run( + record.run_id, + CancelRunRequest(reason="workflow_aborted", message="outer workflow stopped"), + ) + fail_observer.set() + await asyncio.wait_for(supervisor_task, timeout=1) + + assert response.status == "cancelled" + assert runner_finished.is_set() + assert store.statuses[record.run_id] == "cancelled" + assert record.run_id not in store.cancellation_intents + assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"] + terminal = store.events[record.run_id][0] + assert isinstance(terminal, RunCancelledEvent) + assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[]) + + asyncio.run(scenario()) + + def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal() -> None: async def scenario() -> None: store = TrackingStore() @@ -522,10 +661,13 @@ def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal() assert response.status == "cancelled" assert remote_scheduler.active_tasks == {} + await asyncio.wait_for(owner_task, timeout=1) assert store.statuses[record.run_id] == "cancelled" assert store.errors[record.run_id] == "outer workflow stopped" assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"] - await asyncio.wait_for(owner_task, timeout=1) + terminal = store.events[record.run_id][0] + assert isinstance(terminal, RunCancelledEvent) + assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[]) assert runner_finished.is_set() assert store.observer_finished.is_set() await asyncio.sleep(0) @@ -538,37 +680,38 @@ def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal() asyncio.run(scenario()) -def test_owner_observer_reinjects_cancellation_consumed_by_runner() -> None: +def test_pre_enter_cancellation_does_not_copy_input_session_snapshot() -> None: async def scenario() -> None: store = FakeStore() started = asyncio.Event() - first_cancellation = asyncio.Event() + request = _request() + request.session_snapshot = CompositorSessionSnapshot( + layers=[ + LayerSessionSnapshot( + name="prior", + lifecycle_state=LifecycleState.SUSPENDED, + runtime_state={"value": "prior"}, + ) + ] + ) async with httpx.AsyncClient() as client: scheduler = RunScheduler( store=store, plugin_daemon_http_client=client, dify_api_http_client=client, - runner_factory=lambda _record, _request: SwallowOneCancellationRunner( - started=started, - first_cancellation=first_cancellation, - ), + runner_factory=lambda _record, _request: PreEnterBlockingRunner(started=started), ) - record = await scheduler.create_run(_request()) - supervisor_task = scheduler.active_tasks[record.run_id] + record = await scheduler.create_run(request) + supervisor = scheduler.active_tasks[record.run_id] await asyncio.wait_for(started.wait(), timeout=1) - response = await asyncio.wait_for( - scheduler.cancel_run(record.run_id, CancelRunRequest(reason="workflow_aborted")), - timeout=1, - ) + _ = await scheduler.cancel_run(record.run_id, CancelRunRequest(reason="pre_enter_cancel")) + await asyncio.wait_for(supervisor, timeout=1) - assert response.status == "cancelled" - assert store.statuses[record.run_id] == "cancelled" - assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"] - await asyncio.wait_for(first_cancellation.wait(), timeout=1) - await asyncio.wait_for(supervisor_task, timeout=1) - await asyncio.sleep(0) - assert scheduler.active_tasks == {} + terminal = store.events[record.run_id][0] + assert isinstance(terminal, RunCancelledEvent) + assert request.session_snapshot is not None + assert terminal.data.session_snapshot is None asyncio.run(scenario()) @@ -634,8 +777,9 @@ def test_cancelled_terminal_survives_shutdown_while_runner_cleanup_is_pending() response = await scheduler.cancel_run(record.run_id, CancelRunRequest(reason="workflow_aborted")) assert response.status == "cancelled" - assert store.statuses[record.run_id] == "cancelled" - assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"] + assert store.statuses[record.run_id] == "running" + assert store.events[record.run_id] == [] + assert record.run_id in store.cancellation_intents await asyncio.wait_for(store.observer_finished.wait(), timeout=1) assert supervisor_task.done() is False shutdown_task = asyncio.create_task(scheduler.shutdown()) 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 c33e02b701c..c3199922922 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -24,6 +24,7 @@ from pydantic_ai.models.test import TestModel from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults from pydantic_ai.usage import UsageLimits from pydantic_ai.settings import ModelSettings +from pydantic_ai_harness.compaction import TieredCompaction from agenton.compositor import CompositorSessionSnapshot, LayerProvider, LayerSessionSnapshot from agenton.layers import ExitIntent, LifecycleState @@ -66,7 +67,7 @@ from dify_agent.protocol.schemas import ( RunLayerSpec, RunSucceededEvent, ) -from dify_agent.runtime.event_sink import InMemoryRunEventSink, emit_run_cancelled +from dify_agent.runtime.event_sink import InMemoryRunEventSink from dify_agent.runtime.compositor_factory import create_default_layer_providers from dify_agent.runtime.runner import ( AgentRunRunner, @@ -218,12 +219,12 @@ def test_run_failed_error_payload_classifies_usage_limit() -> None: message, error_type, reason = _run_failed_error_payload(exc) - assert message == "The next request would exceed the request_limit of 500" + assert message.startswith("The next request would exceed the request_limit of 500") assert error_type is RunFailureType.AGENT_RUN_LIMIT_EXCEEDED assert reason is None -def test_cancelled_runner_does_not_overwrite_cancelled_status_with_late_failure( +def test_cancelled_runner_does_not_emit_late_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: async def scenario() -> None: @@ -242,20 +243,13 @@ def test_cancelled_runner_does_not_overwrite_cancelled_status_with_late_failure( async def fail_after_cancel() -> RunSuccessOutcome: nonlocal cancelled cancelled = True - _ = await emit_run_cancelled( - sink, - run_id="run-cancelled", - reason="workflow_aborted", - message="workflow stopped", - ) raise RuntimeError("late model failure") monkeypatch.setattr(runner, "_run_agent", fail_after_cancel) await runner.run() - assert sink.statuses["run-cancelled"] == "cancelled" - assert sink.errors["run-cancelled"] == "workflow stopped" - assert [event.type for event in sink.events["run-cancelled"]] == ["run_started", "run_cancelled"] + assert "run-cancelled" not in sink.statuses + assert [event.type for event in sink.events["run-cancelled"]] == ["run_started"] asyncio.run(scenario()) @@ -270,6 +264,8 @@ def _request( execution_context_layer_name: str = "execution_context", on_exit: LayerExitSignals | None = None, output_config: Mapping[str, object] | DifyOutputLayerConfig | None = None, + model_settings: ModelSettings | None = None, + context_window_tokens: int | None = None, ) -> CreateRunRequest: layers = [ RunLayerSpec( @@ -311,6 +307,8 @@ def _request( plugin_id="langgenius/openai", model_provider="openai", model="demo-model", + model_settings=model_settings, + context_window_tokens=context_window_tokens, ), ), ] @@ -427,11 +425,13 @@ class SequenceOutputTestModel(TestModel): class RecordingTestModel(TestModel): seen_requests: list[list[ModelMessage]] + seen_instructions: list[list[str]] failure: Exception | None def __init__(self, *, custom_output_text: str = "done", failure: Exception | None = None) -> None: super().__init__(call_tools=[], custom_output_text=custom_output_text) self.seen_requests = [] + self.seen_instructions = [] self.failure = failure def _request( @@ -441,6 +441,7 @@ class RecordingTestModel(TestModel): model_request_parameters: ModelRequestParameters, ) -> ModelResponse: self.seen_requests.append(list(messages)) + self.seen_instructions.append([part.content for part in model_request_parameters.instruction_parts or []]) if self.failure is not None: raise self.failure return super()._request(messages, model_settings, model_request_parameters) @@ -489,14 +490,14 @@ def _flatten_message_parts(messages: list[ModelMessage]) -> list[object]: class FakeAgentRunResult: output: object - _new_messages: list[ModelMessage] + _all_messages: list[ModelMessage] - def __init__(self, output: object, new_messages: list[ModelMessage]) -> None: + def __init__(self, output: object, all_messages: list[ModelMessage]) -> None: self.output = output - self._new_messages = new_messages + self._all_messages = all_messages - def new_messages(self) -> list[ModelMessage]: - return list(self._new_messages) + def all_messages(self) -> list[ModelMessage]: + return list(self._all_messages) def test_runner_emits_terminal_success_and_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: @@ -671,6 +672,95 @@ def test_runner_passes_explicit_step_limit_to_agent(monkeypatch: pytest.MonkeyPa assert sink.statuses["run-explicit-step-limit"] == "succeeded" +def test_runner_passes_context_compaction(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): + assert http_client.is_closed is False + return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] + + class FakeAgent: + async def run(self, *_args: object, **kwargs: object) -> FakeAgentRunResult: + capabilities = cast(list[object], kwargs["capabilities"]) + assert len(capabilities) == 1 + capability = capabilities[0] + assert isinstance(capability, TieredCompaction) + assert capability.target_tokens == 7_000 + return FakeAgentRunResult("done", []) + + monkeypatch.setattr(DifyPluginLLMLayer, "get_model", fake_get_model) + monkeypatch.setattr("dify_agent.runtime.runner.create_agent", lambda *_args, **_kwargs: FakeAgent()) + sink = InMemoryRunEventSink() + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + await AgentRunRunner( + sink=sink, + request=_request( + model_settings={"max_tokens": 3_000}, + context_window_tokens=10_000, + ), + run_id="run-compaction", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + assert sink.statuses["run-compaction"] == "succeeded" + + +def test_runner_rejects_compaction_budget_before_model_resolution_or_invocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_resolution_called = False + agent_creation_called = False + model_invocation_called = False + + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): + nonlocal model_resolution_called + model_resolution_called = True + return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] + + class FakeAgent: + async def run(self, *_args: object, **_kwargs: object) -> FakeAgentRunResult: + nonlocal model_invocation_called + model_invocation_called = True + return FakeAgentRunResult("unused", []) + + def fake_create_agent(*_args: object, **_kwargs: object) -> FakeAgent: + nonlocal agent_creation_called + agent_creation_called = True + return FakeAgent() + + monkeypatch.setattr(DifyPluginLLMLayer, "get_model", fake_get_model) + monkeypatch.setattr("dify_agent.runtime.runner.create_agent", fake_create_agent) + sink = InMemoryRunEventSink() + + async def scenario() -> None: + async with httpx.AsyncClient() as client: + with pytest.raises( + AgentRunValidationError, + match="Model max_tokens must leave a positive input context budget", + ): + await AgentRunRunner( + sink=sink, + request=_request( + model_settings={"max_tokens": 1_000}, + context_window_tokens=1_000, + ), + run_id="run-invalid-compaction-budget", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ).run() + + asyncio.run(scenario()) + + assert model_resolution_called is False + assert agent_creation_called is False + assert model_invocation_called is False + assert [event.type for event in sink.events["run-invalid-compaction-budget"]] == ["run_started", "run_failed"] + assert sink.statuses["run-invalid-compaction-budget"] == "failed" + + def test_runner_timeout_excludes_tool_preparation_and_runtime_cleanup(monkeypatch: pytest.MonkeyPatch) -> None: shell_client = FakeRunnerShellctlClient() tools_prepared = False @@ -827,6 +917,44 @@ def test_runner_does_not_classify_nested_timeout_as_agent_limit(monkeypatch: pyt assert sink.statuses["run-provider-timeout"] == "failed" +def test_runner_captures_post_exit_snapshot_when_task_is_cancelled(monkeypatch: pytest.MonkeyPatch) -> None: + started = asyncio.Event() + + def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): + return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType] + + class FakeAgent: + async def run(self, *_args: object, **_kwargs: object) -> None: + started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(DifyPluginLLMLayer, "get_model", fake_get_model) + monkeypatch.setattr("dify_agent.runtime.runner.create_agent", lambda *_args, **_kwargs: FakeAgent()) + sink = InMemoryRunEventSink() + + async def scenario() -> AgentRunRunner: + async with httpx.AsyncClient() as client: + runner = AgentRunRunner( + sink=sink, + request=_request(), + run_id="run-cancel-snapshot", + plugin_daemon_http_client=client, + dify_api_http_client=client, + ) + task = asyncio.create_task(runner.run()) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + return runner + + runner = asyncio.run(scenario()) + + assert runner.terminal_session_snapshot is not None + assert all(layer.lifecycle_state is LifecycleState.SUSPENDED for layer in runner.terminal_session_snapshot.layers) + assert [event.type for event in sink.events["run-cancel-snapshot"]] == ["run_started"] + + def test_runner_emits_deferred_tool_call_and_persists_pending_history(monkeypatch: pytest.MonkeyPatch) -> None: captured_output_types: list[object] = [] captured_user_prompts: list[object] = [] @@ -937,9 +1065,11 @@ def test_runner_resumes_with_deferred_tool_results_and_no_user_prompt(monkeypatc assert deferred_tool_results is not None submitted_result = cast(dict[str, object], deferred_tool_results.calls["tool-call-1"]) assert submitted_result["status"] == "submitted" + message_history = cast(list[ModelMessage], kwargs["message_history"]) return FakeAgentRunResult( "done after human", [ + *message_history, ModelRequest( parts=[ ToolReturnPart( @@ -1038,9 +1168,11 @@ def test_runner_can_emit_second_deferred_tool_call_after_resume(monkeypatch: pyt ], ) + message_history = cast(list[ModelMessage], kwargs["message_history"]) return FakeAgentRunResult( DeferredToolRequests(calls=[second_pending_tool_call]), [ + *message_history, ModelRequest( parts=[ ToolReturnPart( @@ -1311,7 +1443,7 @@ def test_runner_passes_dynamic_dify_plugin_tools_to_agent(monkeypatch: pytest.Mo class FakeResult: output: str = "done" - def new_messages(self) -> list[ModelMessage]: + def all_messages(self) -> list[ModelMessage]: return [] class FakeAgent: @@ -1413,7 +1545,7 @@ def test_runner_passes_dynamic_dify_knowledge_tools_to_agent(monkeypatch: pytest class FakeResult: output: str = "done" - def new_messages(self) -> list[ModelMessage]: + def all_messages(self) -> list[ModelMessage]: return [] class FakeAgent: @@ -1519,7 +1651,7 @@ def test_runner_passes_dynamic_dify_core_tools_to_agent(monkeypatch: pytest.Monk class FakeResult: output: str = "done" - def new_messages(self) -> list[ModelMessage]: + def all_messages(self) -> list[ModelMessage]: return [] class FakeAgent: @@ -1965,7 +2097,9 @@ def test_runner_rejects_duplicate_tool_names_between_shell_and_other_layers( assert sink.statuses["run-shell-duplicate-tools"] == "failed" -def test_runner_passes_temporary_system_prompt_prefix_without_history_layer(monkeypatch: pytest.MonkeyPatch) -> None: +def test_runner_passes_system_prompt_as_run_instructions_without_history_layer( + monkeypatch: pytest.MonkeyPatch, +) -> None: model = RecordingTestModel(custom_output_text="done") def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str): @@ -1987,11 +2121,11 @@ def test_runner_passes_temporary_system_prompt_prefix_without_history_layer(monk asyncio.run(scenario()) + assert model.seen_instructions[0] == ["system"] request_parts = _flatten_message_parts(model.seen_requests[0]) - assert isinstance(request_parts[0], SystemPromptPart) - assert request_parts[0].content == "system" - assert isinstance(request_parts[1], UserPromptPart) - assert request_parts[1].content == "current user" + assert len(request_parts) == 1 + assert isinstance(request_parts[0], UserPromptPart) + assert request_parts[0].content == "current user" terminal = sink.events["run-no-history"][-1] assert isinstance(terminal, RunSucceededEvent) assert [layer.name for layer in terminal.data.session_snapshot.layers] == [ @@ -2001,7 +2135,7 @@ def test_runner_passes_temporary_system_prompt_prefix_without_history_layer(monk ] -def test_runner_prepends_current_system_prompt_to_stored_history_and_appends_only_new_messages( +def test_runner_passes_stored_history_with_current_instructions_and_replaces_full_history( monkeypatch: pytest.MonkeyPatch, ) -> None: model = RecordingTestModel(custom_output_text="done") @@ -2031,15 +2165,14 @@ def test_runner_prepends_current_system_prompt_to_stored_history_and_appends_onl asyncio.run(scenario()) + assert model.seen_instructions[0] == ["system"] request_parts = _flatten_message_parts(model.seen_requests[0]) - assert isinstance(request_parts[0], SystemPromptPart) - assert request_parts[0].content == "system" - assert isinstance(request_parts[1], UserPromptPart) - assert request_parts[1].content == "old user" - assert isinstance(request_parts[2], TextPart) - assert request_parts[2].content == "old assistant" - assert isinstance(request_parts[3], UserPromptPart) - assert request_parts[3].content == "current user" + assert isinstance(request_parts[0], UserPromptPart) + assert request_parts[0].content == "old user" + assert isinstance(request_parts[1], TextPart) + assert request_parts[1].content == "old assistant" + assert isinstance(request_parts[2], UserPromptPart) + assert request_parts[2].content == "current user" terminal = sink.events["run-history"][-1] assert isinstance(terminal, RunSucceededEvent) @@ -2054,9 +2187,10 @@ def test_runner_prepends_current_system_prompt_to_stored_history_and_appends_onl assert isinstance(saved_history[3].parts[0], TextPart) assert saved_history[3].parts[0].content == "done" assert all(not any(isinstance(part, SystemPromptPart) for part in message.parts) for message in saved_history) + assert all(not isinstance(message, ModelRequest) or message.instructions is None for message in saved_history) -def test_runner_with_empty_history_layer_still_sends_system_prompt_and_saves_only_new_messages( +def test_runner_with_empty_history_layer_uses_instructions_and_saves_full_history( monkeypatch: pytest.MonkeyPatch, ) -> None: model = RecordingTestModel(custom_output_text="done") @@ -2082,11 +2216,11 @@ def test_runner_with_empty_history_layer_still_sends_system_prompt_and_saves_onl asyncio.run(scenario()) + assert model.seen_instructions[0] == ["system"] request_parts = _flatten_message_parts(model.seen_requests[0]) - assert isinstance(request_parts[0], SystemPromptPart) - assert request_parts[0].content == "system" - assert isinstance(request_parts[1], UserPromptPart) - assert request_parts[1].content == "current user" + assert len(request_parts) == 1 + assert isinstance(request_parts[0], UserPromptPart) + assert request_parts[0].content == "current user" terminal = sink.events["run-empty-history"][-1] assert isinstance(terminal, RunSucceededEvent) @@ -2100,9 +2234,10 @@ def test_runner_with_empty_history_layer_still_sends_system_prompt_and_saves_onl assert isinstance(saved_history[1].parts[0], TextPart) assert saved_history[1].parts[0].content == "done" assert all(not any(isinstance(part, SystemPromptPart) for part in message.parts) for message in saved_history) + assert all(not isinstance(message, ModelRequest) or message.instructions is None for message in saved_history) -def test_runner_failure_with_history_layer_emits_failed_terminal_event_without_success_snapshot( +def test_runner_failure_with_history_layer_emits_post_exit_snapshot_without_new_history( monkeypatch: pytest.MonkeyPatch, ) -> None: model = RecordingTestModel(failure=RuntimeError("boom")) @@ -2135,6 +2270,10 @@ def test_runner_failure_with_history_layer_emits_failed_terminal_event_without_s assert [event.type for event in sink.events["run-history-failure"]] == ["run_started", "run_failed"] assert sink.statuses["run-history-failure"] == "failed" + terminal = sink.events["run-history-failure"][-1] + assert isinstance(terminal, RunFailedEvent) + assert terminal.data.session_snapshot is not None + assert _history_messages_from_snapshot(terminal.data.session_snapshot) == stored_history assert request.session_snapshot is not None assert _history_messages_from_snapshot(request.session_snapshot) == stored_history @@ -2911,6 +3050,10 @@ def test_runner_rejects_closed_session_snapshot_as_validation_error() -> None: assert [event.type for event in sink.events["run-closed-snapshot"]] == ["run_started", "run_failed"] assert sink.statuses["run-closed-snapshot"] == "failed" + terminal = sink.events["run-closed-snapshot"][-1] + assert isinstance(terminal, RunFailedEvent) + assert request.session_snapshot is not None + assert terminal.data.session_snapshot is None def test_runner_treats_missing_runtime_dependency_as_validation_error() -> None: diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runtime_history.py b/dify-agent/tests/local/dify_agent/runtime/test_runtime_history.py index 05c2b9622d6..3072a0c6a1f 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runtime_history.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runtime_history.py @@ -12,9 +12,8 @@ from dify_agent.protocol import DIFY_AGENT_HISTORY_LAYER_ID from dify_agent.protocol.schemas import RunComposition, RunLayerSpec from dify_agent.runtime.compositor_factory import create_default_layer_providers from dify_agent.runtime.history import ( - append_successful_run_history, - build_run_message_history, get_history_layer, + replace_successful_run_history, validate_history_layer_composition, ) @@ -89,63 +88,27 @@ def test_get_history_layer_returns_optional_active_history_layer() -> None: asyncio.run(scenario()) -def test_build_run_message_history_renders_current_system_prompts_before_stored_history() -> None: - stored_history = [ - ModelRequest(parts=[UserPromptPart(content="old user")]), - ModelResponse(parts=[TextPart(content="old assistant")]), +def test_replace_successful_run_history_persists_full_history_without_instructions() -> None: + history_layer = PydanticAIHistoryLayer() + history_layer.replace_messages([ModelRequest(parts=[UserPromptPart(content="stale")])]) + messages = [ + ModelRequest( + parts=[SystemPromptPart(content="Summary of previous conversation:\n\nsummary")], + instructions="current instructions", + ), + ModelRequest(parts=[UserPromptPart(content="new user")]), + ModelResponse(parts=[TextPart(content="new assistant")]), ] - async def scenario() -> None: - message_history = await build_run_message_history( - system_prompts=[lambda: "current system", lambda: "current suffix"], - stored_history=stored_history, - ) + replace_successful_run_history(history_layer, messages) - assert message_history is not None - assert isinstance(message_history[0], ModelRequest) - assert [part.content for part in message_history[0].parts] == ["current system", "current suffix"] - assert message_history[1:] == stored_history - - asyncio.run(scenario()) - - -def test_build_run_message_history_returns_none_without_system_prompt_or_history() -> None: - async def scenario() -> None: - assert await build_run_message_history(system_prompts=[], stored_history=[]) is None - - asyncio.run(scenario()) - - -def test_build_run_message_history_renders_system_prompt_without_history_layer() -> None: - async def scenario() -> None: - message_history = await build_run_message_history(system_prompts=[lambda: "current system"], stored_history=[]) - - assert message_history is not None - assert len(message_history) == 1 - assert isinstance(message_history[0], ModelRequest) - assert isinstance(message_history[0].parts[0], SystemPromptPart) - assert message_history[0].parts[0].content == "current system" - - asyncio.run(scenario()) - - -def test_build_run_message_history_rejects_context_dependent_prompt_functions() -> None: - def unsupported_prompt(_ctx: object) -> str: - return "current system" - - async def scenario() -> None: - with pytest.raises(ValueError, match="zero-argument system prompts"): - await build_run_message_history(system_prompts=[unsupported_prompt], stored_history=[]) - - asyncio.run(scenario()) - - -def test_append_successful_run_history_preserves_existing_message_order() -> None: - history_layer = PydanticAIHistoryLayer() - stored_history = [ModelRequest(parts=[UserPromptPart(content="old user")])] - new_messages = [ModelResponse(parts=[TextPart(content="new assistant")])] - - history_layer.replace_messages(stored_history) - append_successful_run_history(history_layer, new_messages) - - assert history_layer.message_history == [*stored_history, *new_messages] + persisted = history_layer.message_history + assert len(persisted) == 3 + persisted_request = persisted[0] + assert isinstance(persisted_request, ModelRequest) + assert persisted_request.instructions is None + assert persisted_request.parts == messages[0].parts + assert persisted[1:] == messages[1:] + source_request = messages[0] + assert isinstance(source_request, ModelRequest) + assert source_request.instructions == "current instructions" diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py index 247491b0b9c..60b09224344 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_e2b.py @@ -23,6 +23,7 @@ from dify_agent.runtime_backend.e2b import ( E2BExecutionBindingBackend, E2BHomeSnapshotBackend, E2BRuntimeLease, + E2BSDKControlPlane, ) from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease @@ -140,6 +141,36 @@ def _connected_backend(*, pause_error: Exception | None = None) -> tuple[E2BExec ) +@pytest.mark.anyio +async def test_e2b_sdk_create_disables_public_traffic(monkeypatch: pytest.MonkeyPatch) -> None: + from e2b import AsyncSandbox + + sandbox = _Sandbox(sandbox_id="sandbox-1") + create_options: dict[str, object] = {} + + async def create( + _cls: type[AsyncSandbox], + template: str, + **options: object, + ) -> _Sandbox: + assert template == "prepared-template" + create_options.update(options) + return sandbox + + monkeypatch.setattr(AsyncSandbox, "create", classmethod(create)) + control_plane = E2BSDKControlPlane(api_key="e2b-secret") + + created = await control_plane.create( + "prepared-template", + timeout=120, + metadata={"dify.resource": "runtime-sandbox"}, + on_timeout="pause", + ) + + assert created is sandbox + assert create_options["network"] == {"allow_public_traffic": False} + + @pytest.mark.anyio async def test_e2b_binding_uses_default_template_or_exact_snapshot_and_couples_refs() -> None: control = _ControlPlane() @@ -279,7 +310,7 @@ async def test_e2b_checkpoint_uses_exact_source_runtime() -> None: control = _ControlPlane() source_sandbox = _Sandbox(sandbox_id="source") source = E2BRuntimeLease( - sandbox=source_sandbox, + sandbox=source_sandbox, # pyright: ignore[reportArgumentType] data_plane=cast(ShellctlRuntimeLease, object()), ) backend = E2BHomeSnapshotBackend( @@ -299,11 +330,14 @@ async def test_e2b_checkpoint_uses_exact_source_runtime() -> None: async def test_e2b_acquire_retries_transient_shellctl_failures_until_ready( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "ambient-shellctl-token") attempts = 0 def handler(request: httpx.Request) -> httpx.Response: nonlocal attempts attempts += 1 + assert request.headers["X-Access-Token"] == "traffic-token" + assert "Authorization" not in request.headers if attempts == 1: raise httpx.ReadTimeout("shellctl starting", request=request) if attempts == 2: @@ -328,6 +362,26 @@ async def test_e2b_acquire_retries_transient_shellctl_failures_until_ready( assert clients[0].is_closed +@pytest.mark.anyio +@pytest.mark.parametrize("traffic_access_token", [None, ""]) +async def test_e2b_acquire_fails_closed_without_traffic_token( + monkeypatch: pytest.MonkeyPatch, + traffic_access_token: str | None, +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + raise AssertionError("shellctl must not be called without an E2B traffic access token") + + clients = _mock_http(monkeypatch, handler) + backend, sandbox = _connected_backend() + sandbox.traffic_access_token = traffic_access_token + + with pytest.raises(BindingAcquireError, match="traffic access token"): + _ = await backend.acquire(sandbox.sandbox_id) + + assert clients == [] + assert sandbox.pauses == [True] + + @pytest.mark.anyio async def test_e2b_acquire_closes_transport_and_pauses_after_readiness_retries_exhausted( monkeypatch: pytest.MonkeyPatch, diff --git a/dify-agent/tests/local/dify_agent/server/test_settings.py b/dify-agent/tests/local/dify_agent/server/test_settings.py index b4190c29082..29908851d85 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -97,12 +97,31 @@ def test_server_settings_defaults_shellctl_auth_token_to_none( def test_server_settings_reads_agent_stub_settings_from_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub/") monkeypatch.setenv("DIFY_AGENT_SANDBOX_FILES_BASE_URL", "https://dify.example.com/prefix/") + monkeypatch.setenv("DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT", "72") monkeypatch.setenv("DIFY_AGENT_SERVER_SECRET_KEY", _base64url_secret(secrets.token_bytes(32))) settings = ServerSettings() assert settings.agent_stub_api_base_url == "https://agent.example.com/agent-stub" assert settings.sandbox_files_base_url == "https://dify.example.com/prefix" + assert settings.stub_upload_file_size_limit == 72 + + +def test_server_settings_defaults_stub_upload_file_size_limit_to_50_mib( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.delenv("DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT", raising=False) + monkeypatch.chdir(tmp_path) + + assert ServerSettings().stub_upload_file_size_limit == 50 + + +def test_server_settings_accepts_zero_and_rejects_negative_stub_upload_file_size_limit() -> None: + assert ServerSettings(stub_upload_file_size_limit=0).stub_upload_file_size_limit == 0 + + with pytest.raises(ValidationError): + _ = ServerSettings(stub_upload_file_size_limit=-1) def test_server_settings_normalizes_agent_stub_service_root_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -240,6 +259,7 @@ def test_server_settings_create_agent_stub_file_request_handler_returns_handler_ inner_api_url="https://api.example.com", inner_api_key="inner-secret", sandbox_files_base_url="https://sandbox-files.example.com/dify", + stub_upload_file_size_limit=72, ) handler = settings.create_agent_stub_file_request_handler() @@ -248,6 +268,7 @@ def test_server_settings_create_agent_stub_file_request_handler_returns_handler_ assert handler.inner_api_url == "https://api.example.com" assert handler.inner_api_key == "inner-secret" assert handler.sandbox_files_base_url == "https://sandbox-files.example.com/dify" + assert handler.max_upload_size_bytes == 72 * 1024 * 1024 def test_server_settings_create_agent_stub_drive_request_handler_returns_none_without_full_settings() -> None: diff --git a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py index e33e1ec8188..13930624dd1 100644 --- a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py +++ b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py @@ -9,16 +9,19 @@ from pydantic import JsonValue from agenton.compositor import CompositorSessionSnapshot, LayerSessionSnapshot from agenton.layers import LifecycleState from dify_agent.protocol.schemas import ( + RUN_EVENT_ADAPTER, + CancelRunRequest, RunCancelledEvent, RunCancelledEventData, RunFailedEvent, RunFailedEventData, RunFailureType, RunStartedEvent, - RunStatus, RunSucceededEvent, RunSucceededEventData, + utc_now, ) +from dify_agent.runtime.cancellation import RunCancellationIntent from dify_agent.runtime.event_sink import RunFinalizationResult from dify_agent.storage.redis_run_store import DEFAULT_RUN_RETENTION_SECONDS, RedisRunStore, RunNotFoundError @@ -27,12 +30,14 @@ class FakeRedis: commands: list[tuple[object, ...]] values: dict[str, object] streams: dict[str, list[tuple[str, dict[str, object]]]] + eval_result: list[object] | None def __init__(self) -> None: self.commands = [] self.values = {} self.streams = {} self.stream_changed = asyncio.Event() + self.eval_result = None async def set(self, key: str, value: object, *, ex: int | None = None) -> None: self.commands.append(("set", key, value, ex)) @@ -109,29 +114,9 @@ class FakeRedis: async def eval(self, script: str, numkeys: int, *keys_and_args: object) -> list[object]: self.commands.append(("eval", script, numkeys, *keys_and_args)) - assert numkeys == 2 - record_key = str(keys_and_args[0]) - events_key = str(keys_and_args[1]) - status = str(keys_and_args[2]) - updated_at = str(keys_and_args[3]) - has_error = str(keys_and_args[4]) == "1" - error = str(keys_and_args[5]) if has_error else None - has_error_type = str(keys_and_args[6]) == "1" - error_type = str(keys_and_args[7]) if has_error_type else None - payload = str(keys_and_args[8]) - record_json = self.values.get(record_key) - if record_json is None: - return [-1, "", ""] - if isinstance(record_json, bytes): - record_json = record_json.decode() - record = json.loads(cast(str, record_json)) - if record["status"] != "running": - return [0, record["status"], ""] - - record.update({"status": status, "updated_at": updated_at, "error": error, "error_type": error_type}) - event_id = self._append_stream_entry(events_key, {"payload": payload}) - self.values[record_key] = json.dumps(record, separators=(",", ":")) - return [1, status, event_id] + if self.eval_result is None: + raise AssertionError("test must configure FakeRedis.eval_result") + return list(self.eval_result) @staticmethod def _is_after_min(event_id: str, min_id: str) -> bool: @@ -224,75 +209,76 @@ def test_get_run_accepts_legacy_record_without_error_type() -> None: assert loaded.error_type is None -def test_finalize_run_atomically_writes_terminal_event_and_status() -> None: +def test_request_cancellation_maps_eval_result_and_arguments() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType] - record = asyncio.run(store.create_run()) - redis.commands.clear() - event = RunCancelledEvent( - run_id=record.run_id, - data=RunCancelledEventData(reason="workflow_aborted", message="workflow stopped"), + redis.eval_result = [1, "running"] + + status = asyncio.run( + store.request_cancellation( + "run-1", + CancelRunRequest(reason="workflow_aborted", message="workflow stopped"), + ) ) - result = asyncio.run(store.finalize_run(event)) - updated = asyncio.run(store.get_run(record.run_id)) + assert status == "running" + eval_command = redis.commands[-1] + assert eval_command[0] == "eval" + assert eval_command[2] == 3 + assert eval_command[3:6] == ( + "test:runs:run-1:record", + "test:runs:run-1:cancel-intent", + "test:runs:run-1:events", + ) + intent_payload = json.loads(cast(str, eval_command[6])) + assert intent_payload["reason"] == "workflow_aborted" + assert intent_payload["message"] == "workflow stopped" + assert eval_command[7] == "60" - assert result.applied is True - assert result.status == "cancelled" - assert result.event_id == "1-0" - assert updated.status == "cancelled" - assert updated.error == "workflow stopped" - assert updated.error_type is None - assert updated.updated_at == event.created_at - stream_entry_id, stream_fields = redis.streams[f"test:runs:{record.run_id}:events"][0] - assert stream_entry_id == result.event_id - payload = json.loads(cast(str, stream_fields["payload"])) + +def test_finalize_cancellation_maps_eval_result_and_arguments() -> None: + redis = FakeRedis() + store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType] + redis.eval_result = [1, "cancelled", "7-0"] + intent = RunCancellationIntent( + reason="workflow_aborted", + message="workflow stopped", + requested_at=utc_now(), + ) + + result = asyncio.run( + store.finalize_cancellation( + "run-1", + intent, + session_snapshot=CompositorSessionSnapshot(layers=[]), + ) + ) + + assert result == RunFinalizationResult(applied=True, status="cancelled", event_id="7-0") + eval_command = redis.commands[-1] + assert eval_command[2] == 3 + assert eval_command[3:6] == ( + "test:runs:run-1:record", + "test:runs:run-1:cancel-intent", + "test:runs:run-1:events", + ) + payload = json.loads(cast(str, eval_command[9])) assert "id" not in payload assert payload["type"] == "run_cancelled" - assert payload["data"] == {"reason": "workflow_aborted", "message": "workflow stopped"} - assert payload["created_at"] == event.created_at.isoformat().replace("+00:00", "Z") - eval_command = redis.commands[0] - assert eval_command[0] == "eval" - assert eval_command[2] == 2 - assert eval_command[-1] == "60" + assert payload["data"] == { + "reason": "workflow_aborted", + "message": "workflow stopped", + "session_snapshot": {"schema_version": 1, "layers": []}, + } + assert eval_command[10] == "60" -def test_finalize_run_rejects_a_second_terminal_without_appending_event() -> None: - redis = FakeRedis() - store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType] - record = asyncio.run(store.create_run()) - snapshot = CompositorSessionSnapshot(layers=[]) - - first = asyncio.run( - store.finalize_run( - RunSucceededEvent( - run_id=record.run_id, - data=RunSucceededEventData(output="done", session_snapshot=snapshot), - ) - ) - ) - second = asyncio.run( - store.finalize_run( - RunCancelledEvent( - run_id=record.run_id, - data=RunCancelledEventData(reason="late_cancel"), - ) - ) - ) - - assert first.applied is True - assert second.applied is False - assert second.status == "succeeded" - assert second.event_id is None - assert len(redis.streams[f"test:runs:{record.run_id}:events"]) == 1 - - -def test_finalize_failed_run_derives_error_and_timestamp_from_event() -> None: +def test_finalize_failed_run_maps_eval_result_and_arguments() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - record = asyncio.run(store.create_run()) + redis.eval_result = [1, "failed", "8-0"] event = RunFailedEvent( - run_id=record.run_id, + run_id="run-1", data=RunFailedEventData( error="model failed", error_type=RunFailureType.AGENT_RUN_LIMIT_EXCEEDED, @@ -301,117 +287,46 @@ def test_finalize_failed_run_derives_error_and_timestamp_from_event() -> None: ) result = asyncio.run(store.finalize_run(event)) - updated = asyncio.run(store.get_run(record.run_id)) - assert result.applied is True - assert result.status == "failed" - assert updated.status == "failed" - assert updated.error == "model failed" - assert updated.error_type is RunFailureType.AGENT_RUN_LIMIT_EXCEEDED - assert updated.updated_at == event.created_at - stream_entry = redis.streams[f"test:runs:{record.run_id}:events"][0] - payload = json.loads(cast(str, stream_entry[1]["payload"])) + assert result == RunFinalizationResult(applied=True, status="failed", event_id="8-0") + eval_command = redis.commands[-1] + assert eval_command[3:6] == ( + "test:runs:run-1:record", + "test:runs:run-1:events", + "test:runs:run-1:cancel-intent", + ) + assert eval_command[6] == "failed" + assert eval_command[8:12] == ("1", "model failed", "1", "agent_run_limit_exceeded") + payload = json.loads(cast(str, eval_command[12])) assert payload["data"]["error_type"] == "agent_run_limit_exceeded" -def test_two_store_instances_choose_exactly_one_terminal_winner() -> None: - redis = FakeRedis() - first_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - second_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - - async def scenario() -> tuple[list[RunFinalizationResult], RunStatus, list[str]]: - record = await first_store.create_run() - snapshot = CompositorSessionSnapshot(layers=[]) - results = await asyncio.gather( - first_store.finalize_run( - RunSucceededEvent( - run_id=record.run_id, - data=RunSucceededEventData(output="done", session_snapshot=snapshot), - ) - ), - second_store.finalize_run( - RunCancelledEvent( - run_id=record.run_id, - data=RunCancelledEventData(reason="concurrent_cancel"), - ) - ), - ) - persisted = await first_store.get_run(record.run_id) - page = await second_store.get_events(record.run_id) - return list(results), persisted.status, [event.type for event in page.events] - - results, status, event_types = asyncio.run(scenario()) - - assert sum(result.applied for result in results) == 1 - assert len(event_types) == 1 - assert (status, event_types[0]) in { - ("succeeded", "run_succeeded"), - ("cancelled", "run_cancelled"), - } - - -def test_failure_and_cancellation_compete_for_one_terminal() -> None: - redis = FakeRedis() - failure_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - cancellation_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - - async def scenario() -> tuple[list[RunFinalizationResult], RunStatus, list[str]]: - record = await failure_store.create_run() - results = await asyncio.gather( - failure_store.finalize_run( - RunFailedEvent( - run_id=record.run_id, - data=RunFailedEventData(error="model failed", reason="model_error"), - ) - ), - cancellation_store.finalize_run( - RunCancelledEvent( - run_id=record.run_id, - data=RunCancelledEventData(reason="concurrent_cancel"), - ) - ), - ) - persisted = await failure_store.get_run(record.run_id) - page = await cancellation_store.get_events(record.run_id) - return list(results), persisted.status, [event.type for event in page.events] - - results, status, event_types = asyncio.run(scenario()) - - assert sum(result.applied for result in results) == 1 - assert len(event_types) == 1 - assert (status, event_types[0]) in { - ("failed", "run_failed"), - ("cancelled", "run_cancelled"), - } - - -def test_finalize_run_raises_when_record_is_missing() -> None: +def test_request_cancellation_raises_when_record_is_missing() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] + redis.eval_result = [-1, ""] with pytest.raises(RunNotFoundError): - asyncio.run( - store.finalize_run(RunCancelledEvent(run_id="missing", data=RunCancelledEventData(reason="cancelled"))) - ) + asyncio.run(store.request_cancellation("missing", CancelRunRequest(reason="cancelled"))) def test_wait_for_cancellation_observes_terminal_record_before_starting() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - async def scenario() -> bool: + async def scenario() -> object: record = await store.create_run() - _ = await store.finalize_run( - RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled")) - ) + redis.values[f"test:runs:{record.run_id}:record"] = record.model_copy( + update={"status": "cancelled"} + ).model_dump_json() redis.commands.clear() return await store.wait_for_cancellation(record.run_id) - assert asyncio.run(scenario()) is True + assert asyncio.run(scenario()) is None assert [command[0] for command in redis.commands] == ["xrevrange", "get"] -def test_wait_for_cancellation_covers_terminal_transition_during_initialization() -> None: +def test_wait_for_cancellation_covers_intent_transition_during_initialization() -> None: class PausingRecordReadRedis(FakeRedis): record_read_started: asyncio.Event release_record_read: asyncio.Event @@ -432,62 +347,75 @@ def test_wait_for_cancellation_covers_terminal_transition_during_initialization( redis = PausingRecordReadRedis() observer_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - cancelling_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - async def scenario() -> bool: + async def scenario() -> object: record = await observer_store.create_run() observer = asyncio.create_task(observer_store.wait_for_cancellation(record.run_id)) await asyncio.wait_for(redis.record_read_started.wait(), timeout=1) - _ = await cancelling_store.finalize_run( - RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled")) + _ = redis._append_stream_entry( + f"test:runs:{record.run_id}:cancel-intent", + { + "payload": RunCancellationIntent( + reason="cancelled", + requested_at=utc_now(), + ).model_dump_json() + }, ) redis.release_record_read.set() return await asyncio.wait_for(observer, timeout=1) - assert asyncio.run(scenario()) is True + assert asyncio.run(scenario()) is not None def test_wait_for_cancellation_advances_past_non_terminal_events() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - async def scenario() -> bool: + async def scenario() -> object: record = await store.create_run() observer = asyncio.create_task(store.wait_for_cancellation(record.run_id)) await asyncio.sleep(0) _ = await store.append_event(RunStartedEvent(run_id=record.run_id)) await asyncio.sleep(0) - _ = await store.finalize_run( - RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled")) + _ = redis._append_stream_entry( + f"test:runs:{record.run_id}:cancel-intent", + { + "payload": RunCancellationIntent( + reason="cancelled", + requested_at=utc_now(), + ).model_dump_json() + }, ) return await asyncio.wait_for(observer, timeout=1) - assert asyncio.run(scenario()) is True + assert asyncio.run(scenario()) is not None cursors = [command[1] for command in redis.commands if command[0] == "xread"] assert any("0-0" in streams.values() for streams in cursors if isinstance(streams, dict)) assert any("1-0" in streams.values() for streams in cursors if isinstance(streams, dict)) -def test_wait_for_cancellation_returns_false_when_success_wins() -> None: +def test_wait_for_cancellation_returns_none_when_success_wins() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - async def scenario() -> bool: + async def scenario() -> object: record = await store.create_run() observer = asyncio.create_task(store.wait_for_cancellation(record.run_id)) await asyncio.sleep(0) - _ = await store.finalize_run( - RunSucceededEvent( - run_id=record.run_id, - data=RunSucceededEventData( - output="done", - session_snapshot=CompositorSessionSnapshot(layers=[]), - ), - ) + event = RunSucceededEvent( + run_id=record.run_id, + data=RunSucceededEventData( + output="done", + session_snapshot=CompositorSessionSnapshot(layers=[]), + ), + ) + _ = redis._append_stream_entry( + f"test:runs:{record.run_id}:events", + {"payload": RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()}, ) return await asyncio.wait_for(observer, timeout=1) - assert asyncio.run(scenario()) is False + assert asyncio.run(scenario()) is None def test_append_event_serializes_typed_event_without_id_and_expires_run_keys() -> None: @@ -530,19 +458,20 @@ def test_get_events_round_trips_run_succeeded_output_and_session_snapshot() -> N async def scenario() -> tuple[str, RunSucceededEvent]: record = await store.create_run() - result = await store.finalize_run( - RunSucceededEvent( - id="local-only", - run_id=record.run_id, - data=RunSucceededEventData(output=output, session_snapshot=session_snapshot), - ) + event = RunSucceededEvent( + id="local-only", + run_id=record.run_id, + data=RunSucceededEventData(output=output, session_snapshot=session_snapshot), + ) + event_id = redis._append_stream_entry( + f"test:runs:{record.run_id}:events", + {"payload": RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()}, ) - assert result.event_id is not None page = await store.get_events(record.run_id, after="0-0", limit=10) decoded = page.events[0] assert isinstance(decoded, RunSucceededEvent) - assert page.next_cursor == result.event_id - return result.event_id, decoded + assert page.next_cursor == event_id + return event_id, decoded event_id, decoded = asyncio.run(scenario()) @@ -559,7 +488,11 @@ def test_iter_events_ends_after_replaying_terminal_event(terminal_type: str) -> async def scenario() -> list[str]: record = await store.create_run() _ = await store.append_event(RunStartedEvent(run_id=record.run_id)) - _ = await store.finalize_run(_terminal_event(terminal_type, record.run_id)) + terminal = _terminal_event(terminal_type, record.run_id) + _ = redis._append_stream_entry( + f"test:runs:{record.run_id}:events", + {"payload": RUN_EVENT_ADAPTER.dump_json(terminal, exclude={"id"}).decode()}, + ) redis.commands.clear() async def collect_events() -> list[str]: @@ -586,7 +519,11 @@ def test_iter_events_ends_after_live_terminal_event(terminal_type: str) -> None: assert not next_event.done() assert "xread" in [command[0] for command in redis.commands] - _ = await store.finalize_run(_terminal_event(terminal_type, record.run_id)) + terminal = _terminal_event(terminal_type, record.run_id) + _ = redis._append_stream_entry( + f"test:runs:{record.run_id}:events", + {"payload": RUN_EVENT_ADAPTER.dump_json(terminal, exclude={"id"}).decode()}, + ) event = await asyncio.wait_for(next_event, timeout=1) with pytest.raises(StopAsyncIteration): _ = await anext(events) diff --git a/dify-agent/tests/local/examples/test_agenton_examples.py b/dify-agent/tests/local/examples/test_agenton_examples.py index 545f5e6ef57..2b0b63a573e 100644 --- a/dify-agent/tests/local/examples/test_agenton_examples.py +++ b/dify-agent/tests/local/examples/test_agenton_examples.py @@ -44,11 +44,11 @@ def test_agenton_pydantic_ai_example_smoke() -> None: result = _run_example("agenton_examples.pydantic_ai_bridge") assert result.returncode == 0, result.stderr - assert "SystemPromptPart: Prefer concrete details." in result.stdout - assert "UserPromptPart: [\"Use the tools for 'layer composition'.\"]" in result.stdout - assert "ToolCallPart: count_words(" in result.stdout - assert "ToolCallPart: write_tagline(" in result.stdout - assert "TextPart:" in result.stdout + assert "SystemPromptPart(content='Prefer concrete details.'," in result.stdout + assert "UserPromptPart(content=[\"Use the tools for 'layer composition'.\"]," in result.stdout + assert "ToolCallPart(tool_name='count_words'" in result.stdout + assert "ToolCallPart(tool_name='write_tagline'" in result.stdout + assert "TextPart(content=" in result.stdout def test_agenton_session_snapshot_example_smoke() -> None: diff --git a/dify-agent/tests/local/test_packaging.py b/dify-agent/tests/local/test_packaging.py index 13cdc058b49..d382da2d405 100644 --- a/dify-agent/tests/local/test_packaging.py +++ b/dify-agent/tests/local/test_packaging.py @@ -10,7 +10,8 @@ CLIENT_SHARED_DTO_DEPENDENCIES = { "httpx==0.28.1", "httpx2>=2.5.0,<3.0.0", "pydantic>=2.12.5,<2.13", - "pydantic-ai-slim>=1.102.0,<2.0.0", + "pydantic-ai-harness>=0.20.0,<0.21.0", + "pydantic-ai-slim>=2.30.0,<3.0.0", "typing-extensions>=4.12.2,<5.0.0", } @@ -21,7 +22,7 @@ SERVER_RUNTIME_DEPENDENCIES = { "jsonschema>=4.23.0,<5.0.0", "jwcrypto>=1.5.6,<2", "logfire[fastapi,httpx,redis]>=4.37.0,<5.0.0", - "pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0", + "pydantic-ai-slim[anthropic,google,openai]>=2.30.0,<3.0.0", "pydantic-settings>=2.12.0,<3.0.0", "redis>=7.4.0,<8.0.0", "uvicorn[standard]==0.46.0", diff --git a/dify-agent/uv.lock b/dify-agent/uv.lock index 6dcca54cc6f..3b2039f7144 100644 --- a/dify-agent/uv.lock +++ b/dify-agent/uv.lock @@ -593,6 +593,7 @@ dependencies = [ { name = "httpx" }, { name = "httpx2" }, { name = "pydantic" }, + { name = "pydantic-ai-harness" }, { name = "pydantic-ai-slim" }, { name = "typing-extensions" }, ] @@ -638,8 +639,9 @@ requires-dist = [ { name = "jwcrypto", marker = "extra == 'server'", specifier = ">=1.5.6,<2" }, { name = "logfire", extras = ["fastapi", "httpx", "redis"], marker = "extra == 'server'", specifier = ">=4.37.0,<5.0.0" }, { name = "pydantic", specifier = ">=2.12.5,<2.13" }, - { name = "pydantic-ai-slim", specifier = ">=1.106.0,<2.0.0" }, - { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" }, + { name = "pydantic-ai-harness", specifier = ">=0.20.0,<0.21.0" }, + { name = "pydantic-ai-slim", specifier = ">=2.30.0,<3.0.0" }, + { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=2.30.0,<3.0.0" }, { name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" }, { name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" }, { name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" }, @@ -784,15 +786,15 @@ wheels = [ [[package]] name = "genai-prices" -version = "0.0.57" +version = "0.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "httpx" }, + { name = "httpx2" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/be/30/11f3d683cf3b1d9612475ad8bfffe3423ce9f50fc617733109033e73a038/genai_prices-0.0.57.tar.gz", hash = "sha256:6e101e9c53975557ceffa237b0995787d81fe75aac12410f2898504188bcad89", size = 66555, upload-time = "2026-04-21T13:42:52.554Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/14/a188df294f013ec9cd97fc6b145f5427f89067bfb2c260fc3fb5c8d1fb34/genai_prices-0.1.3.tar.gz", hash = "sha256:62c30cddd6c2d2199d878d1a70521c3e37347cd9394446d107dc774a78ed3780", size = 92638, upload-time = "2026-08-15T00:10:31.771Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a9/fe/d0095040c120d97cb63d055224ecd4e913dc5655315c203c8e83bf13aa86/genai_prices-0.0.57-py3-none-any.whl", hash = "sha256:14e50fb69cdc5a06ddb2a6df5a7fe06741b9e44304ce3f1728f56abdf1856cca", size = 69654, upload-time = "2026-04-21T13:42:51.236Z" }, + { url = "https://files.pythonhosted.org/packages/4e/cd/d94b47c26d6367e0b949edfe2da5a47fb74037e799a0edd5b825e049f2b9/genai_prices-0.1.3-py3-none-any.whl", hash = "sha256:a2603841429c843da91c987d9ef598c73bd940caf44e844ab046d551791c04bb", size = 96892, upload-time = "2026-08-15T00:10:30.595Z" }, ] [[package]] @@ -1819,7 +1821,7 @@ wheels = [ [[package]] name = "openai" -version = "2.32.0" +version = "2.54.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -1831,9 +1833,9 @@ dependencies = [ { name = "tqdm" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ed/59/bdcc6b759b8c42dd73afaf5bf8f902c04b37987a5514dbc1c64dba390fef/openai-2.32.0.tar.gz", hash = "sha256:c54b27a9e4cb8d51f0dd94972ffd1a04437efeb259a9e60d8922b8bd26fe55e0", size = 693286, upload-time = "2026-04-15T22:28:19.434Z" } +sdist = { url = "https://files.pythonhosted.org/packages/50/9a/8c75e8c8a5b407a0586faeb2afac91674ff955c191ecc1d6d3b6669f6788/openai-2.54.0.tar.gz", hash = "sha256:e3e6f8bc1ba30ddf381ace1a14340eed381cb984a1a59bd0f34b5be3b5d49cfa", size = 1100285, upload-time = "2026-08-11T18:46:59.035Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/c1/d6e64ccd0536bf616556f0cad2b6d94a8125f508d25cfd814b1d2db4e2f1/openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f", size = 1162570, upload-time = "2026-04-15T22:28:17.714Z" }, + { url = "https://files.pythonhosted.org/packages/64/a8/bb76c7356de8ad57f59d5ff993d434df0607f07f08bcc9c9a5c275e399c0/openai-2.54.0-py3-none-any.whl", hash = "sha256:89089789197ccdb87f173a03145ed1598d00795220c93e96cf712b1cbf5e5f2b", size = 1660351, upload-time = "2026-08-11T18:46:56.684Z" }, ] [[package]] @@ -2381,10 +2383,25 @@ wheels = [ ] [[package]] -name = "pydantic-ai-slim" -version = "1.106.0" +name = "pydantic-ai-harness" +version = "0.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "genai-prices" }, + { name = "httpx" }, + { name = "pydantic-ai-slim" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3f/6c/a066644a3a3bff22bfdddd745fd245e2b4e3148fc895be0020f03bd7470d/pydantic_ai_harness-0.20.0.tar.gz", hash = "sha256:18ec7d6f90873a8038d094280e50af5e877320b975f334268b700a266d35f522", size = 1846014, upload-time = "2026-08-14T03:36:31.045Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/85/32cab39557e338abbd4ecb844028a31110c03096e6678bf5a1266ed201d1/pydantic_ai_harness-0.20.0-py3-none-any.whl", hash = "sha256:e1164ae4d653bd2ae257e3816ee1776b27dbe8f7eaeb29b36d86a49d6fe98168", size = 623606, upload-time = "2026-08-14T03:36:29.094Z" }, +] + +[[package]] +name = "pydantic-ai-slim" +version = "2.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, { name = "genai-prices" }, { name = "griffelib" }, { name = "httpx" }, @@ -2393,9 +2410,9 @@ dependencies = [ { name = "pydantic-graph" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2b/45/2afc9100a7c370d8ac37bdfccfb54f46fc99da3bdce63f07c32c37807ebc/pydantic_ai_slim-1.106.0.tar.gz", hash = "sha256:e265598c8ee0e903ebb02d0494bb232be4cc8aa463ba1a55aa743cf34135dacf", size = 773504, upload-time = "2026-06-05T01:29:09.129Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/4e/2165d9b90edcd5dfc8e9465b3bdc0aa66a67760843ddb0ac99ee396898f0/pydantic_ai_slim-2.31.0.tar.gz", hash = "sha256:a9310d2464154b028096f1d680f17837f16e5c6cd209b4542e4f60ca5d344789", size = 1214087, upload-time = "2026-08-15T03:17:28.353Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/72/d9/a2785c576e3519a72a5bbc0e12027c542b265ef6eea1aa72b9c440ac2531/pydantic_ai_slim-1.106.0-py3-none-any.whl", hash = "sha256:0dd7a99ea3fa89b490098406c2240ba7d75c327eea094c3fd057dd7aa9f3d163", size = 957617, upload-time = "2026-06-05T01:28:59.979Z" }, + { url = "https://files.pythonhosted.org/packages/db/09/233e529fadbece38580c3a390783f55fa196afad875ce535ee9a57a5ad71/pydantic_ai_slim-2.31.0-py3-none-any.whl", hash = "sha256:cb809ad949ca68be6bb9a0e0b994fc73a95f4cd405e8609a71034f2e2080e2a1", size = 1432053, upload-time = "2026-08-15T03:17:21.208Z" }, ] [package.optional-dependencies] @@ -2496,17 +2513,18 @@ wheels = [ [[package]] name = "pydantic-graph" -version = "1.106.0" +version = "2.31.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "anyio" }, { name = "httpx" }, { name = "logfire-api" }, { name = "pydantic" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/42/9b/dd6826cf21eedd96a7482302be51ba6087095acbe828362135de2a505092/pydantic_graph-1.106.0.tar.gz", hash = "sha256:55afa33df4f699ed5c1185f81b6a06e2161958f1aa0c20742b2dae5745e84cce", size = 62567, upload-time = "2026-06-05T01:29:11.833Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8a/2c/3817ae318ecc729a258fc85aaad84fc110b9467a292b98bb10e65ae71183/pydantic_graph-2.31.0.tar.gz", hash = "sha256:a19919408dfaa5a1b8713618bcce7a5135d83664321862b0d9be1f4216c432ef", size = 45180, upload-time = "2026-08-15T03:17:30.398Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9b/e9/0058f0b98f5992e715a0a50128f6c3cc7946cc242d471f6e850efdf03f0c/pydantic_graph-1.106.0-py3-none-any.whl", hash = "sha256:e6bb61aef0fdb49185a81142d311f94fc3315329345471d12cab85ab5845221f", size = 80099, upload-time = "2026-06-05T01:29:04.219Z" }, + { url = "https://files.pythonhosted.org/packages/54/ef/a3217caed3189cfcc6857316e06cb900990e4d5cb59e42ec517cfdda6a7f/pydantic_graph-2.31.0-py3-none-any.whl", hash = "sha256:062555c89b1d5699ddaaa6f09ae62fc1d0f5cc189d0867e2fafca68c24797ba5", size = 52662, upload-time = "2026-08-15T03:17:24.42Z" }, ] [[package]] diff --git a/docker/.env.example b/docker/.env.example index 64513b3aea0..b614314b91a 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -283,7 +283,6 @@ DIFY_AGENT_E2B_API_KEY= DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox # RuntimeLease active limit; its default matches the independently configurable Agent run deadline. DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 -DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= DIFY_AGENT_E2B_SHELLCTL_PORT=5004 # Sandbox-reachable Dify API base for dify-agent CLI /files/* transfers. # Remote Sandboxes should use the public Dify ingress; local Compose uses api via agent_ssrf_proxy. diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index c81b53186fa..4e080afc849 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -744,7 +744,6 @@ services: DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} # RuntimeLease active limit; its default matches the independently configurable Agent run deadline. DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600} - DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-} DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004} DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub} DIFY_AGENT_SANDBOX_FILES_BASE_URL: ${DIFY_AGENT_SANDBOX_FILES_BASE_URL:-http://api:5001} diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 035012e7c8a..9b71a655108 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -754,7 +754,6 @@ services: DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} # RuntimeLease active limit; its default matches the independently configurable Agent run deadline. DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600} - DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-} DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004} DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub} DIFY_AGENT_SANDBOX_FILES_BASE_URL: ${DIFY_AGENT_SANDBOX_FILES_BASE_URL:-http://api:5001} diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index 3c63b32af85..44d5f4ebb0f 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -32,10 +32,11 @@ DIFY_AGENT_E2B_API_KEY= DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox # RuntimeLease active limit; its default matches the independently configurable Agent run deadline. DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600 -DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN= DIFY_AGENT_E2B_SHELLCTL_PORT=5004 # Sandbox-reachable Dify API base for signed /files/* transfers. DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://api:5001 +# Maximum Agent Stub upload size in MiB; forwarded to Dify API as a signed byte limit. +DIFY_AGENT_STUB_UPLOAD_FILE_SIZE_LIMIT=50 DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub # This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens. # Replace this development default in production. diff --git a/e2e/features/step-definitions/agent-v2/configure.steps.ts b/e2e/features/step-definitions/agent-v2/configure.steps.ts index b4ada8ef204..174fcc8820a 100644 --- a/e2e/features/step-definitions/agent-v2/configure.steps.ts +++ b/e2e/features/step-definitions/agent-v2/configure.steps.ts @@ -42,10 +42,10 @@ async function fillAgentPromptEditor(page: Page, prompt: string) { } async function selectAgentModel(page: Page, modelName: string) { - await page.getByRole('combobox').first().click() + await page.getByRole('button', { name: 'Configure model' }).click() await page.getByPlaceholder('Search model').fill(modelName) const escapedModelName = modelName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') - await page.getByRole('option', { name: new RegExp(`${escapedModelName}(?:\\s|$)`) }).click() + await page.getByRole('button', { name: new RegExp(`${escapedModelName}(?:\\s|$)`) }).click() } async function expectAgentComposerPrompt(world: DifyWorld, agentId: string, prompt: string) { diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 9804b81f55d..945ab9b2865 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -1642,11 +1642,6 @@ "count": 2 } }, - "web/app/components/base/prompt-editor/plugins/hitl-input-block/input-field.tsx": { - "jsx-a11y/no-autofocus": { - "count": 1 - } - }, "web/app/components/base/prompt-editor/plugins/hitl-input-block/pre-populate.tsx": { "jsx-a11y/no-autofocus": { "count": 1 @@ -2774,9 +2769,6 @@ } }, "web/app/components/plugins/install-plugin/install-from-github/index.tsx": { - "jsx-a11y/no-autofocus": { - "count": 1 - }, "typescript/no-explicit-any": { "count": 2 } @@ -3290,11 +3282,6 @@ "count": 6 } }, - "web/app/components/snippets/create-snippet-dialog.tsx": { - "jsx-a11y/no-autofocus": { - "count": 1 - } - }, "web/app/components/snippets/hooks/use-snippet-run.ts": { "no-restricted-imports": { "count": 2 @@ -3309,12 +3296,6 @@ "eslint-react/set-state-in-effect": { "count": 4 }, - "jsx-a11y/click-events-have-key-events": { - "count": 1 - }, - "jsx-a11y/no-static-element-interactions": { - "count": 1 - }, "typescript/no-explicit-any": { "count": 2 } diff --git a/packages/contracts/generated/api/console/oauth/orpc.gen.ts b/packages/contracts/generated/api/console/oauth/orpc.gen.ts index 31feb8abd4d..e65bb04c043 100644 --- a/packages/contracts/generated/api/console/oauth/orpc.gen.ts +++ b/packages/contracts/generated/api/console/oauth/orpc.gen.ts @@ -14,6 +14,7 @@ import { zGetOauthPluginByProviderIdDatasourceGetAuthorizationUrlQuery, zGetOauthPluginByProviderIdDatasourceGetAuthorizationUrlResponse, zGetOauthPluginByProviderToolAuthorizationUrlPath, + zGetOauthPluginByProviderToolAuthorizationUrlQuery, zGetOauthPluginByProviderToolAuthorizationUrlResponse, zPostOauthProviderAccountBody, zPostOauthProviderAccountResponse, @@ -137,7 +138,12 @@ export const get5 = oc path: '/oauth/plugin/{provider}/tool/authorization-url', tags: ['console'], }) - .input(z.object({ params: zGetOauthPluginByProviderToolAuthorizationUrlPath })) + .input( + z.object({ + params: zGetOauthPluginByProviderToolAuthorizationUrlPath, + query: zGetOauthPluginByProviderToolAuthorizationUrlQuery.optional(), + }), + ) .output(zGetOauthPluginByProviderToolAuthorizationUrlResponse) export const authorizationUrl = { diff --git a/packages/contracts/generated/api/console/oauth/types.gen.ts b/packages/contracts/generated/api/console/oauth/types.gen.ts index 2108c56280d..2120c2a3283 100644 --- a/packages/contracts/generated/api/console/oauth/types.gen.ts +++ b/packages/contracts/generated/api/console/oauth/types.gen.ts @@ -138,6 +138,7 @@ export type GetOauthPluginByProviderIdDatasourceGetAuthorizationUrlData = { } query?: { credential_id?: string + visibility?: string } url: '/oauth/plugin/{provider_id}/datasource/get-authorization-url' } @@ -154,7 +155,9 @@ export type GetOauthPluginByProviderToolAuthorizationUrlData = { path: { provider: string } - query?: never + query?: { + visibility?: 'all_team_members' | 'only_me' + } url: '/oauth/plugin/{provider}/tool/authorization-url' } diff --git a/packages/contracts/generated/api/console/oauth/zod.gen.ts b/packages/contracts/generated/api/console/oauth/zod.gen.ts index 960d1e429c3..d96544bcdc0 100644 --- a/packages/contracts/generated/api/console/oauth/zod.gen.ts +++ b/packages/contracts/generated/api/console/oauth/zod.gen.ts @@ -134,6 +134,7 @@ export const zGetOauthPluginByProviderIdDatasourceGetAuthorizationUrlPath = z.ob export const zGetOauthPluginByProviderIdDatasourceGetAuthorizationUrlQuery = z.object({ credential_id: z.string().optional(), + visibility: z.string().optional(), }) /** @@ -146,6 +147,10 @@ export const zGetOauthPluginByProviderToolAuthorizationUrlPath = z.object({ provider: z.string(), }) +export const zGetOauthPluginByProviderToolAuthorizationUrlQuery = z.object({ + visibility: z.enum(['all_team_members', 'only_me']).optional(), +}) + /** * Tool OAuth authorization URL generated successfully */ diff --git a/packages/dify-ui/AGENTS.md b/packages/dify-ui/AGENTS.md index 0cb66d50e4c..9b877e2a8dd 100644 --- a/packages/dify-ui/AGENTS.md +++ b/packages/dify-ui/AGENTS.md @@ -16,6 +16,7 @@ This package owns shared design tokens, CSS-first Tailwind styles, the `cn()` ut Use the README sections as the detailed owners: - [Button and icon-button contracts] +- [Form and input composition] - [Imports and public boundaries] - [Typed value contracts] - [Search and picker selection] @@ -25,6 +26,7 @@ Use the README sections as the detailed owners: [Button and icon-button contracts]: README.md#button-loading-and-disabled-contract [Development and test boundaries]: README.md#development +[Form and input composition]: README.md#form-contract [Imports and public boundaries]: README.md#imports [Overlay and portal contracts]: README.md#overlay--portal-contract [Search and picker selection]: README.md#search-and-picker-selection diff --git a/packages/dify-ui/README.md b/packages/dify-ui/README.md index f85108ad7a4..75c35061387 100644 --- a/packages/dify-ui/README.md +++ b/packages/dify-ui/README.md @@ -30,9 +30,11 @@ import { Button } from '@langgenius/dify-ui/button' import { cn } from '@langgenius/dify-ui/cn' import { Dialog, DialogContent, DialogTrigger } from '@langgenius/dify-ui/dialog' import { Drawer, DrawerPopup, DrawerTrigger } from '@langgenius/dify-ui/drawer' -import { Field, FieldControl, FieldLabel } from '@langgenius/dify-ui/field' +import { Field, FieldLabel } from '@langgenius/dify-ui/field' import { Form } from '@langgenius/dify-ui/form' import { IconButton } from '@langgenius/dify-ui/icon-button' +import { InputGroup, InputGroupAddon, InputGroupInput } from '@langgenius/dify-ui/input-group' +import { Input } from '@langgenius/dify-ui/input' import { Kbd, KbdGroup } from '@langgenius/dify-ui/kbd' import { Popover, PopoverContent, PopoverTrigger } from '@langgenius/dify-ui/popover' import { SegmentedControl, SegmentedControlItem } from '@langgenius/dify-ui/segmented-control' @@ -63,18 +65,18 @@ Keep implementation-only render helpers, context values, styling helpers, and up ## Primitives -| Category | Subpath | Notes | -| ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------- | -| Actions | `./button`, `./icon-button`, `./toggle` | Visible-label actions, icon-only commands, and persistent toggles. | -| Controls | `./segmented-control` | SegmentedControl for mode, filter, and view selection. | -| Display | `./collapsible`, `./kbd` | Collapsible disclosure primitive; keyboard input and shortcut keycap primitives. | -| Feedback | `./meter`, `./progress`, `./status-dot`, `./toast` | Inline and asynchronous status primitives; Toast owns the `z-60` layer. | -| Form | `./form`, `./field`, `./fieldset`, `./input`, `./textarea`, `./checkbox`, `./checkbox-group`, `./radio`, `./number-field`, `./select`, `./slider`, `./switch` | Native form boundary, field semantics, and controls. | -| Layout | `./scroll-area` | Custom-styled scrollbar over the host viewport. | -| Media | `./avatar` | Avatar root, image, and fallback primitives. | -| Navigation | `./file-tree`, `./pagination`, `./tabs` | FileTree for preview-oriented file disclosure lists; Pagination for page navigation; Tabs for panels. | -| Overlay / menu | `./alert-dialog`, `./context-menu`, `./dialog`, `./drawer`, `./dropdown-menu`, `./popover`, `./preview-card`, `./tooltip` | Portalled. See [Overlay & portal contract] below. | -| Search / pickers | `./autocomplete`, `./combobox`, `./select` | Search input, searchable picker, and closed picker. | +| Category | Subpath | Notes | +| ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------- | +| Actions | `./button`, `./icon-button`, `./toggle` | Visible-label actions, icon-only commands, and persistent toggles. | +| Controls | `./segmented-control` | SegmentedControl for mode, filter, and view selection. | +| Display | `./collapsible`, `./kbd` | Collapsible disclosure primitive; keyboard input and shortcut keycap primitives. | +| Feedback | `./meter`, `./progress`, `./status-dot`, `./toast` | Inline and asynchronous status primitives; Toast owns the `z-60` layer. | +| Form | `./form`, `./field`, `./fieldset`, `./input`, `./input-group`, `./textarea`, `./checkbox`, `./checkbox-group`, `./radio`, `./number-field`, `./select`, `./slider`, `./switch` | Native form boundary, field semantics, and controls. | +| Layout | `./scroll-area` | Custom-styled scrollbar over the host viewport. | +| Media | `./avatar` | Avatar root, image, and fallback primitives. | +| Navigation | `./file-tree`, `./pagination`, `./tabs` | FileTree for preview-oriented file disclosure lists; Pagination for page navigation; Tabs for panels. | +| Overlay / menu | `./alert-dialog`, `./context-menu`, `./dialog`, `./drawer`, `./dropdown-menu`, `./popover`, `./preview-card`, `./tooltip` | Portalled. See [Overlay & portal contract] below. | +| Search / pickers | `./autocomplete`, `./combobox`, `./select` | Search input, searchable picker, and closed picker. | Utilities: @@ -122,7 +124,9 @@ Dify UI's form primitives are a Base UI composition layer for native form semant Use `Form` for the submit boundary. It renders a native `
`, preserves Enter-to-submit and submit-button behavior, and adds Base UI's `onFormSubmit`, `errors`, `actionsRef`, and `validationMode` APIs for structured values and consolidated field validation. Prefer it over a bare `` when the form is composed with Dify UI fields. -Use `Field` for each standalone named field. A field must have a stable `name`, a label relationship, and either a `FieldControl` or another control that participates in the same Base UI field context. Prefer a visible label for normal form rows; when the surrounding UI already supplies the visible text, use the matching label primitive visually hidden or put `aria-label` on the actual interactive control. `FieldDescription` and `FieldError` provide the message relationships that screen readers need, while the Dify wrapper adds the default Form Input Set styling from the design system. +Use `Field` when a text control needs Base UI field semantics such as shared name, label, validation, description, or error state. A standalone `Input` may instead use a native `