-
+
- {virtualizer.getVirtualItems().map((virtualItem) => {
- const option = filteredItems[virtualItem.index]
+
+
+ {virtualizer.getVirtualItems().map((virtualItem) => {
+ const option = filteredItems[virtualItem.index]
- if (!option) return null
+ if (!option) return null
- return (
-
-
+ return (
+
- {renderVirtualizedOptionItem(option, virtualItem.index)}
-
- )
- })}
-
+ {renderVirtualizedOptionItem(option, virtualItem.index, filteredItems.length)}
+
+ )
+ })}
+
+
+
+
- {t(($) => $.noData, { ns: 'common' })}
-
-
+ {t(($) => $.noData, { ns: 'common' })}
+
+ )}
+
-
+
)
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/__tests__/index.spec.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/__tests__/index.spec.tsx
index 1e4a4368810..a5eb555b769 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/__tests__/index.spec.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/__tests__/index.spec.tsx
@@ -7,6 +7,17 @@ import { createConsoleQueryClient } from '@/test/console/query-data'
import { ConfigurationMethodEnum, ModelStatusEnum, ModelTypeEnum } from '../../declarations'
import { ModelSelector, SplitModelSelector } from '../index'
+const makeModelItem = (overrides: Partial
-
+
+ onMarketplaceCollapsedChange(!open)}
+ render={}
+ >
-
+
- {!marketplaceCollapsed && (
-
+
+
+
)
}
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/model-selector-trigger.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/model-selector-trigger.tsx
index c7e55f5a7ff..dafc380ca20 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/model-selector-trigger.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/model-selector-trigger.tsx
@@ -1,6 +1,6 @@
import type { ModelSelectorModel, ModelSelectorProvider, ModelSelectorValue } from './types'
import { cn } from '@langgenius/dify-ui/cn'
-import { ComboboxTrigger } from '@langgenius/dify-ui/combobox'
+import { PopoverTrigger } from '@langgenius/dify-ui/popover'
import { Tooltip, TooltipContent, TooltipTrigger } from '@langgenius/dify-ui/tooltip'
import { useQuery } from '@tanstack/react-query'
import { useTranslation } from 'react-i18next'
@@ -12,7 +12,7 @@ import {
} from '../derive-model-status'
import ModelIcon from '../model-icon'
import ModelName from '../model-name'
-import { useCredentialPanelState } from '../provider-added-card/use-credential-panel-state'
+import { useCredentialPanelState as useCredentialPanelInfo } from '../provider-added-card/use-credential-panel-state'
type ModelSelectorTriggerProps = {
currentProvider?: ModelSelectorProvider
@@ -52,7 +52,7 @@ function ModelSelectorTrigger({
enabled: !!providerId,
select: ({ data }) => data.find((provider) => provider.provider === providerId),
})
- const credentialPanel = useCredentialPanelState(resolvedProvider)
+ const credentialPanel = useCredentialPanelInfo(resolvedProvider)
const status = deriveModelStatus(
isSelected ? currentModel?.model : defaultModel?.model,
@@ -93,26 +93,33 @@ function ModelSelectorTrigger({
$['detailPanel.configureModel'], { ns: 'plugin' })}
- data-deprecated={isDeprecated ? '' : undefined}
- data-model-status={status}
- data-shape={shape}
- data-size={size}
- data-surface={surface}
+ span]:opacity-50',
- className,
- )}
+ render={
+
}
/>
{triggerTooltipLabel && (
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-empty-state.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-empty-state.tsx
index 2b4d28a2336..2568587a548 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-empty-state.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-empty-state.tsx
@@ -11,7 +11,7 @@ function ModelSelectorEmptyState({ onConfigure }: ModelSelectorEmptyStateProps)
return (
)
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-item.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-item.tsx
index c51acca8702..adf4449a3f0 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-item.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-item.tsx
@@ -6,13 +6,14 @@ import type {
ModelSelectorProvider,
ModelSelectorValue,
} from './types'
+import { Button } from '@langgenius/dify-ui/button'
import { cn } from '@langgenius/dify-ui/cn'
-import { ComboboxGroup, ComboboxItem, ComboboxItemIndicator } from '@langgenius/dify-ui/combobox'
-import { Popover, PopoverContent, PopoverTrigger } from '@langgenius/dify-ui/popover'
+import { Collapsible, CollapsiblePanel, CollapsibleTrigger } from '@langgenius/dify-ui/collapsible'
+import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from '@langgenius/dify-ui/popover'
import { PreviewCardTrigger } from '@langgenius/dify-ui/preview-card'
import { StatusDot } from '@langgenius/dify-ui/status-dot'
import { Tooltip, TooltipContent, TooltipTrigger } from '@langgenius/dify-ui/tooltip'
-import { useCallback, useMemo, useState } from 'react'
+import { useCallback, useId, useState } from 'react'
import { useTranslation } from 'react-i18next'
import { useModalContext } from '@/context/modal-context'
import { useProviderContext } from '@/context/provider-context'
@@ -45,8 +46,10 @@ type PopupItemProps = {
modelSuggestionPredicate?: ModelSelectorModelPredicate
previewCardHandle: PreviewCardHandle
onPreviewCardClose: () => void
+ onSelect: (provider: string, model: ModelSelectorModel) => void
onHide: () => void
}
+
function PopupItem({
defaultModel,
model,
@@ -54,12 +57,15 @@ function PopupItem({
modelSuggestionPredicate,
previewCardHandle,
onPreviewCardClose,
+ onSelect,
onHide,
}: PopupItemProps) {
- const [collapsed, setCollapsed] = useState(false)
+ const [modelsOpen, setModelsOpen] = useState(true)
const [dropdownOpen, setDropdownOpen] = useState(false)
+ const providerHeadingId = useId()
const { t } = useTranslation()
const language = useLanguage()
+ const providerLabel = renderI18nObject(model.label, language)
const suggestionTip = t(($) => $['modelProvider.selector.suggestionTip'], { ns: 'common' })
const { setShowModelModal } = useModalContext()
const { modelProviders } = useProviderContext()
@@ -69,10 +75,16 @@ function PopupItem({
const { providerDetail, loadProviderDetail } = useLazyModelProviderDetail(model.provider)
const { canUseCredential, canCreateCredential, canManageCredential } = useCredentialPermissions()
const canOpenCredentialDropdown = canUseCredential || canCreateCredential || canManageCredential
- const handleOpenModelModal = async () => {
- if (!canCreateCredential) return
+ const state = useCredentialPanelInfo(currentProvider)
+ const { isChangingPriority, handleChangePriority } = useChangeProviderPriority(currentProvider)
+ const isUsingCredits = state.priority === 'credits'
+ const hasCredits = !state.isCreditsExhausted
+ const isApiKeyActive = state.variant === 'api-active' || state.variant === 'api-fallback'
+ const { credentialName } = state
+
+ const handleOpenModelModal = async () => {
+ if (!canCreateCredential || !currentProvider) return
- if (!currentProvider) return
const detail = await loadProviderDetail()
if (!detail) return
setShowModelModal({
@@ -84,34 +96,16 @@ function PopupItem({
updateModelProviders()
const modelType = model.models[0]!.model_type
-
if (modelType) updateModelList(modelType as ModelType)
},
})
}
- const state = useCredentialPanelInfo(currentProvider)
- const { isChangingPriority, handleChangePriority } = useChangeProviderPriority(currentProvider)
- const groupItems = useMemo(
- () =>
- model.models
- .filter((modelItem) => modelItem.status !== ModelStatusEnum.noConfigure)
- .map((modelItem) => ({
- provider: model.provider,
- model: modelItem.model,
- })),
- [model.models, model.provider],
- )
-
- const isUsingCredits = state.priority === 'credits'
- const hasCredits = !state.isCreditsExhausted
- const isApiKeyActive = state.variant === 'api-active' || state.variant === 'api-fallback'
- const { credentialName } = state
-
const handleCloseDropdown = useCallback(() => {
setDropdownOpen(false)
onHide()
}, [onHide])
+
const handleDropdownOpenChange = async (nextOpen: boolean) => {
if (!nextOpen) {
setDropdownOpen(false)
@@ -124,40 +118,51 @@ function PopupItem({
if (!currentProvider) return null
return (
-
-
-
{marketplaceProviders.map((key) => {
const Icon = providerIconMap[key]
const isInstalling = installingProvider === key
return (
-
-
+
)
})}
+
+
{onOpenMarketplace ? (
-
+
) : (
)}
- )}
-
-
+
@@ -23,7 +23,7 @@ function ModelSelectorEmptyState({ onConfigure }: ModelSelectorEmptyStateProps)
-
)
- return (
-
+
+ )
+ })}
+
+
+
)
}
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-layout.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-layout.tsx
index b58d0c2cd6b..76601878272 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/popup-layout.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/popup-layout.tsx
@@ -1,6 +1,5 @@
import type { ReactNode } from 'react'
-import { cn } from '@langgenius/dify-ui/cn'
-import { ComboboxInput, ComboboxInputGroup } from '@langgenius/dify-ui/combobox'
+import { Button } from '@langgenius/dify-ui/button'
import {
ScrollArea,
ScrollAreaContent,
@@ -9,6 +8,7 @@ import {
ScrollAreaViewport,
} from '@langgenius/dify-ui/scroll-area'
import { useTranslation } from 'react-i18next'
+import { SearchInput } from '@/app/components/base/search-input'
type ModelSelectorSearchHeaderProps = {
inputValue: string
@@ -22,37 +22,13 @@ export function ModelSelectorSearchHeader({
const { t } = useTranslation()
return (
- $['form.searchModel'], { ns: 'datasetSettings' }) || ''}
+ className="mx-2 mt-2 mb-1 shrink-0"
+ placeholder={t(($) => $['form.searchModel'], { ns: 'datasetSettings' }) || ''}
+ value={inputValue}
+ onValueChange={onInputValueChange}
+ />
)
}
@@ -101,9 +77,10 @@ export function ShowIncompatibleModelsButton({
const { t } = useTranslation()
return (
-
+
)
}
@@ -124,16 +101,17 @@ export function ModelProviderSettingsFooter({ onOpenSettings }: ModelProviderSet
return (
-
+
-
-
- $['form.searchModel'], { ns: 'datasetSettings' }) || ''}
- className="block h-4.5 grow px-1 py-0 text-[13px] text-text-primary"
- placeholder={t(($) => $['form.searchModel'], { ns: 'datasetSettings' }) || ''}
- />
- {inputValue && (
-
- )}
-
-
+
-
+
)
}
diff --git a/web/app/components/header/account-setting/model-provider-page/model-selector/popup.tsx b/web/app/components/header/account-setting/model-provider-page/model-selector/popup.tsx
index a3ec1ed0bf5..c95d2edd395 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-selector/popup.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-selector/popup.tsx
@@ -1,11 +1,11 @@
import type { ModelSelectorPreviewPayload } from './popup-item'
import type {
+ ModelSelectorModel,
ModelSelectorModelPredicate,
ModelSelectorProvider,
ModelSelectorValue,
} from './types'
import type { ModelProviderQuotaGetPaid } from '@/types/model-provider'
-import { ComboboxList } from '@langgenius/dify-ui/combobox'
import {
createPreviewCardHandle,
PreviewCard,
@@ -61,6 +61,7 @@ export type PopupProps = {
modelSuggestionPredicate?: ModelSelectorModelPredicate
onConfigureEmptyState?: () => void
onInputValueChange: (value: string) => void
+ onSelect: (provider: string, model: ModelSelectorModel) => void
onOpenMarketplace?: () => void
onHide: () => void
}
@@ -74,6 +75,7 @@ function Popup({
modelSuggestionPredicate,
onConfigureEmptyState,
onInputValueChange,
+ onSelect,
onOpenMarketplace,
onHide,
}: PopupProps) {
@@ -239,6 +241,13 @@ function Popup({
)
}, [enableMarketplace, modelProviderPlugins])
+ const searchStatus =
+ !filteredModelList.length && installedModelList.length > 0
+ ? t(($) => $['modelProvider.selector.noModelFoundForSearch'], {
+ ns: 'common',
+ query: inputValue,
+ })
+ : null
const handleClosePreviewCard = useCallback(() => {
previewCardHandle.close()
}, [previewCardHandle])
@@ -251,34 +260,37 @@ function Popup({
{showCreditsExhaustedAlert && (
- {filteredModelList.map((model) => (
-
-
+ {filteredModelList.map((model) => (
+
{!filteredModelList.length && !installedModelList.length && (
boolean
-
-export const isSameModelSelectorValue = (
- itemValue: ModelSelectorValue,
- value: ModelSelectorValue,
-) => itemValue.provider === value.provider && itemValue.model === value.model
-
-export const getModelSelectorValueLabel = (value: ModelSelectorValue) => value.model
diff --git a/web/app/components/plugins/plugin-detail-panel/app-selector/__tests__/app-picker.spec.tsx b/web/app/components/plugins/plugin-detail-panel/app-selector/__tests__/app-picker.spec.tsx
new file mode 100644
index 00000000000..360436f187b
--- /dev/null
+++ b/web/app/components/plugins/plugin-detail-panel/app-selector/__tests__/app-picker.spec.tsx
@@ -0,0 +1,126 @@
+import type { AppPartial } from '@dify/contracts/api/console/apps/types.gen'
+import { render, screen } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { useState } from 'react'
+import { AppModeEnum } from '@/types/app'
+import { AppPicker } from '../app-picker'
+
+vi.mock('@/app/components/base/app-icon', () => ({
+ default: () => ,
+}))
+
+const app = {
+ id: 'app-1',
+ name: 'Workflow App',
+ mode: AppModeEnum.WORKFLOW,
+ icon_type: 'emoji',
+ icon: 'W',
+ icon_background: '#FFFFFF',
+ icon_url: null,
+} satisfies AppPartial
+
+function renderAppPicker(
+ overrides: Partial, 'isShow' | 'onShowChange'>> = {},
+) {
+ const props = {
+ disabled: false,
+ trigger: Choose app,
+ onSelect: vi.fn(),
+ apps: [app],
+ isLoading: false,
+ hasMore: false,
+ onLoadMore: vi.fn(),
+ searchText: '',
+ onSearchChange: vi.fn(),
+ ...overrides,
+ }
+
+ function AppPickerHarness() {
+ const [open, setOpen] = useState(false)
+
+ return Choose app}
+ isShow={open}
+ onShowChange={setOpen}
+ onSelect={onSelect}
+ apps={[app]}
+ isLoading={false}
+ hasMore={false}
+ onLoadMore={vi.fn()}
+ searchText={searchText}
+ onSearchChange={setSearchText}
+ />
+ )
+ }
+
+ render((null)
const handleValueChange = useCallback(
(app: AppPartial | null) => {
@@ -116,6 +127,11 @@ export function AppPicker({
[onSelect, onShowChange],
)
+ const handleClearSearch = () => {
+ onSearchChange('')
+ inputRef.current?.focus()
+ }
+
return (
items={apps}
@@ -136,56 +152,76 @@ export function AppPicker({
>
{trigger}
-
-
- {t(($) => $['modelProvider.selector.noModelFoundForSearch'], {
- ns: 'common',
- query: inputValue,
- })}
-
- )}
+
+ {searchStatus}
+
{scopeFeatures.length > 0 &&
-
({
@@ -33,14 +34,20 @@ vi.mock('@/service/client', () => ({
const manageInConsoleLabel = /manageInAgentConsole/
const startFromScratchLabel = /startFromScratch/
-const renderSelector = async ({ onStartFromScratch }: { onStartFromScratch?: () => void } = {}) => {
+const renderSelector = async ({
+ onOpenChange = vi.fn(),
+ onStartFromScratch,
+}: {
+ onOpenChange?: (open: boolean) => void
+ onStartFromScratch?: () => void
+} = {}) => {
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } })
render(
-
- {label}
-
-
-
- $['appSelector.placeholder'], { ns: 'app' })}
- placeholder={t(($) => $['appSelector.placeholder'], { ns: 'app' })}
- className="block h-4.5 grow px-1 py-0 text-[13px] text-text-primary"
- />
- {searchText && (
-
+
+
+
)
}
diff --git a/web/app/components/plugins/plugin-detail-panel/app-selector/index.tsx b/web/app/components/plugins/plugin-detail-panel/app-selector/index.tsx
index f2236035520..b732fcd6a06 100644
--- a/web/app/components/plugins/plugin-detail-panel/app-selector/index.tsx
+++ b/web/app/components/plugins/plugin-detail-panel/app-selector/index.tsx
@@ -130,7 +130,10 @@ export function AppSelector({
diff --git a/web/app/components/snippet-list/components/snippet-card.tsx b/web/app/components/snippet-list/components/snippet-card.tsx
index 438f03d61e9..6872bf0f555 100644
--- a/web/app/components/snippet-list/components/snippet-card.tsx
+++ b/web/app/components/snippet-list/components/snippet-card.tsx
@@ -177,7 +177,6 @@ const SnippetCard = ({
-
+
filter={null}
+ inline
inputValue={searchText}
- items={options}
- itemToStringLabel={getOptionLabel}
- itemToStringValue={getAgentSelectorOptionValue}
+ items={agents}
+ itemToStringLabel={(agent) => agent.name}
+ itemToStringValue={(agent) => agent.id}
open={open}
value={null}
onInputValueChange={handleInputValueChange}
@@ -130,84 +114,84 @@ export function AgentSelectorContent({
/>
-
- {isLoading && (
- $.loading, { ns: 'common' })} />
+ {statusText}
+ {isLoading ? (
+
+ {!agentsQuery.isError &&
+ agents.map((agent) =>
+ )}
+ {hasActions && (
+
)
}
-function AgentSelectorLoadingSkeleton({ label }: { label: string }) {
+function AgentSelectorLoadingSkeleton() {
return (
-
+
+ ) : (
+
+ {onStartFromScratch && (
+
)}
- {!isLoading && agentsQuery.isError && (
-
- {t(($) => $['roster.loadingError'], { ns: 'agentV2' })}
-
- )}
- {!isLoading && !agentsQuery.isError && (
- <>
- {agents.length === 0 && (
-
- {debouncedSearchText
- ? t(($) => $['roster.emptySearch'], { ns: 'agentV2' })
- : t(($) => $['roster.empty'], { ns: 'agentV2' })}
-
+ {canManageAgents && (
+ (
-
- {actionOptions.length > 0 && (
-
- {actionOptions.map((option) => (
-
- )}
-
+ )}
- (null)
const workspacePermissionKeys = useAtomValue(workspacePermissionKeysAtom)
const canManageTags = hasPermission(workspacePermissionKeys, getTagManagePermissionKey(type))
const filteredItems = useComboboxFilteredItems()
const realItemCount = filteredItems.filter((tag) => !isCreateTagOption(tag)).length
const placeholder = t(($) => $['tag.selectorPlaceholder'], { ns: 'common' }) || ''
+ const handleClearInput = () => {
+ onInputValueChange('')
+ inputRef.current?.focus()
+ }
+
return (
& {
targetId: string
type: TagType
@@ -77,13 +68,6 @@ export const TagSelector = ({
onClick,
onOpenTagManagement = () => {},
onTagsChange,
- placement = 'bottom-start',
- sideOffset = 4,
- alignOffset = 0,
- portalProps,
- positionerProps,
- popupProps,
- popupClassName,
...rootProps
}: TagSelectorProps) => {
const { t } = useTranslation()
@@ -274,27 +258,23 @@ export const TagSelector = ({
className="pointer-events-none absolute top-0 right-0 h-full w-20 bg-tag-selector-mask-bg group-hover/tag-area:hidden group-focus-visible/tag-area:hidden group-data-popup-open/tag-area:hidden"
/>
-
- handleOpenChange(false)}
- />
-
+
+
+
+ handleOpenChange(false)}
+ />
+
+
+
)
}
From 219596c4dfc2f07d831515766462910e914a6834 Mon Sep 17 00:00:00 2001
From: YungLe <143590624+ForeignKeyCN@users.noreply.github.com>
Date: Mon, 17 Aug 2026 02:41:26 +0000
Subject: [PATCH 02/32] feat(plugin-auth): visibility picker for OAuth
credentials (#39840)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>
Co-authored-by: twwu
---
.../datasets/rag_pipeline/datasource_auth.py | 46 ++-
.../console/workspace/tool_providers.py | 42 ++-
api/openapi/markdown/console-openapi.md | 9 +
api/services/datasource_provider_service.py | 60 +++-
.../rag_pipeline/test_datasource_auth.py | 27 +-
.../console/workspace/test_tool_providers.py | 184 ++++++++++++
.../generated/api/console/oauth/orpc.gen.ts | 8 +-
.../generated/api/console/oauth/types.gen.ts | 5 +-
.../generated/api/console/oauth/zod.gen.ts | 5 +
.../__tests__/index.spec.tsx | 19 ++
.../base/credential-selector/index.tsx | 7 +-
.../__tests__/card.spec.tsx | 27 +-
.../__tests__/configure.spec.tsx | 55 +++-
.../data-source-page-new/configure.tsx | 43 +--
.../__tests__/add-oauth-button.spec.tsx | 279 +++++++++++++++++-
.../__tests__/api-key-modal.spec.tsx | 21 +-
.../__tests__/authorize-components.spec.tsx | 49 +--
.../__tests__/oauth-client-settings.spec.tsx | 31 +-
.../authorize/add-oauth-button.tsx | 199 ++++++++-----
.../plugin-auth/authorize/api-key-modal.tsx | 2 +-
.../authorize/oauth-client-settings.tsx | 18 +-
.../authorize/oauth-visibility-dialog.tsx | 77 +++++
.../authorize/permission-selector.tsx | 90 +++---
.../hooks/__tests__/use-credential.spec.ts | 6 +
.../plugin-auth/hooks/use-credential.ts | 4 +-
.../tools/__tests__/index.spec.tsx | 19 +-
web/i18n/ar-TN/plugin.json | 2 +
web/i18n/de-DE/plugin.json | 2 +
web/i18n/en-US/plugin.json | 2 +
web/i18n/es-ES/plugin.json | 2 +
web/i18n/fa-IR/plugin.json | 2 +
web/i18n/fr-FR/plugin.json | 2 +
web/i18n/hi-IN/plugin.json | 2 +
web/i18n/id-ID/plugin.json | 2 +
web/i18n/it-IT/plugin.json | 2 +
web/i18n/ja-JP/plugin.json | 2 +
web/i18n/ko-KR/plugin.json | 2 +
web/i18n/lo-LA/plugin.json | 2 +
web/i18n/nl-NL/plugin.json | 2 +
web/i18n/pl-PL/plugin.json | 2 +
web/i18n/pt-BR/plugin.json | 2 +
web/i18n/ro-RO/plugin.json | 2 +
web/i18n/ru-RU/plugin.json | 2 +
web/i18n/sl-SI/plugin.json | 2 +
web/i18n/th-TH/plugin.json | 2 +
web/i18n/tr-TR/plugin.json | 2 +
web/i18n/uk-UA/plugin.json | 2 +
web/i18n/vi-VN/plugin.json | 2 +
web/i18n/zh-Hans/plugin.json | 2 +
web/i18n/zh-Hant/plugin.json | 2 +
web/models/permission.ts | 8 +
web/service/use-plugins-auth.ts | 6 +-
52 files changed, 1148 insertions(+), 246 deletions(-)
create mode 100644 web/app/components/plugins/plugin-auth/authorize/oauth-visibility-dialog.tsx
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/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md
index 9fff14ab3b1..b403b6713da 100644
--- a/api/openapi/markdown/console-openapi.md
+++ b/api/openapi/markdown/console-openapi.md
@@ -7723,6 +7723,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
@@ -7736,6 +7737,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 @@ -17205,6 +17207,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 @@ -22580,6 +22583,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/datasource_provider_service.py b/api/services/datasource_provider_service.py index 5de194dd262..ed394bd0a29 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 @@ -642,10 +643,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}" @@ -692,16 +706,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( @@ -879,11 +898,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() @@ -895,6 +920,7 @@ class DatasourceProviderService: tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id, + user=user, session=session, ) redirect_uri = ( @@ -938,11 +964,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() @@ -960,6 +989,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/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/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/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/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx index b639f7e2f6b..2720b42237d 100644 --- a/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx +++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx @@ -12,12 +12,14 @@ const credentials = [ name: 'First credential', avatar_url: 'https://example.com/first.png', type: 'oauth2', + is_default: false, }, { id: 'credential-2', name: 'Second credential', avatar_url: 'https://example.com/second.png', type: 'oauth2', + is_default: false, }, ] as DataSourceCredential[] @@ -51,6 +53,23 @@ describe('CredentialSelector', () => { await waitFor(() => expect(onCredentialChange).toHaveBeenCalledWith('credential-1')) }) + it('selects the visible default credential before falling back to the first one', async () => { + const onCredentialChange = vi.fn() + const credentialsWithDefault = credentials.map((credential, index) => ({ + ...credential, + is_default: index === 1, + })) + render( +
-
- >
- {(item) =>
+
+
+
+ >
+ {(item) =>
+
+
- >
- {(item) =>
+
+
+
+ >
+ {(item) =>
+
+
),
diff --git a/packages/dify-ui/src/autocomplete/index.tsx b/packages/dify-ui/src/autocomplete/index.tsx
index 7d5bbe90aab..d878f500f11 100644
--- a/packages/dify-ui/src/autocomplete/index.tsx
+++ b/packages/dify-ui/src/autocomplete/index.tsx
@@ -250,56 +250,46 @@ function AutocompleteIcon({ className, children, ...props }: AutocompleteIconPro
)
}
-type AutocompleteContentProps = {
- children: React.ReactNode
- placement?: Placement
- sideOffset?: number
- alignOffset?: number
+const AutocompletePortal = BaseAutocomplete.Portal
+type AutocompletePortalProps = BaseAutocomplete.Portal.Props
+
+type AutocompletePositionerProps = Omit<
+ BaseAutocomplete.Positioner.Props,
+ 'className' | 'side' | 'align'
+> & {
className?: string
- popupClassName?: string
- portalProps?: Omit
- positionerProps?: Omit<
- BaseAutocomplete.Positioner.Props,
- 'children' | 'className' | 'side' | 'align' | 'sideOffset' | 'alignOffset'
- >
- popupProps?: Omit
+ placement?: Placement
}
-function AutocompleteContent({
- children,
+function AutocompletePositioner({
+ className,
placement = 'bottom-start',
sideOffset = 4,
- alignOffset = 0,
- className,
- popupClassName,
- portalProps,
- positionerProps,
- popupProps,
-}: AutocompleteContentProps) {
+ ...props
+}: AutocompletePositionerProps) {
const { side, align } = parsePlacement(placement)
return (
-
-
-
- {children}
-
-
-
+ & {
+ className?: string
+}
+
+function AutocompletePopup({ className, ...props }: AutocompletePopupProps) {
+ return (
+
-
- >
- {(institution) => (
-
- {institution}
-
- )}
-
- {footerState !== null ? (
-
- ) : null}
-
-
- {shouldShowEmpty ? t(($) => $['form.schoolName.noResults'], { ns: 'education' }) : null}
-
-
-
- {shouldOpenPopup && isLoading
- ? t(($) => $.loading, { ns: 'common' })
- : footerState === 'error'
- ? t(($) => $['dynamicSelect.error'], { ns: 'common' })
+
+
+ {shouldShowEmpty
+ ? t(($) => $['form.schoolName.noResults'], { ns: 'education' })
: null}
-
-
-
+
+
+
+ {shouldOpenPopup && isLoading
+ ? t(($) => $.loading, { ns: 'common' })
+ : footerState === 'error'
+ ? t(($) => $['dynamicSelect.error'], { ns: 'common' })
+ : null}
+
+
+
+
+
)
From 30c304c6a732a327a36bf41054102224d7742798 Mon Sep 17 00:00:00 2001
From: Amogh Bhat Nagri <146640569+amogh-nagri-11@users.noreply.github.com>
Date: Tue, 18 Aug 2026 02:10:59 +0000
Subject: [PATCH 28/32] fix(api): repair two malformed logger format strings
(#40901)
---
api/controllers/trigger/trigger.py | 2 +-
api/controllers/web/human_input_form.py | 2 +-
api/tests/unit_tests/controllers/trigger/test_trigger.py | 9 +++++++++
3 files changed, 11 insertions(+), 2 deletions(-)
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/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/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
From 6d89557355839a8546f79d2951577cbae636226c Mon Sep 17 00:00:00 2001
From: yyh <92089059+lyzno1@users.noreply.github.com>
Date: Tue, 18 Aug 2026 02:14:39 +0000
Subject: [PATCH 29/32] feat(dify-ui): add input group composition (#40890)
---
.../references/dify-ui.md | 3 +-
oxlint-suppressions.json | 19 --
packages/dify-ui/AGENTS.md | 2 +
packages/dify-ui/README.md | 33 +--
packages/dify-ui/package.json | 4 +
.../src/autocomplete/__tests__/index.spec.tsx | 15 ++
packages/dify-ui/src/autocomplete/index.tsx | 4 +-
.../src/combobox/__tests__/index.spec.tsx | 37 +++
packages/dify-ui/src/combobox/index.tsx | 4 +-
packages/dify-ui/src/dialog/index.stories.tsx | 8 +-
.../src/field/__tests__/index.spec.tsx | 11 +-
packages/dify-ui/src/field/index.stories.tsx | 36 +--
packages/dify-ui/src/field/index.tsx | 15 +-
packages/dify-ui/src/form-control-shared.ts | 32 +--
.../dify-ui/src/form/__tests__/index.spec.tsx | 9 +-
packages/dify-ui/src/form/index.stories.tsx | 7 +-
packages/dify-ui/src/input-group/README.md | 11 +
.../src/input-group/__tests__/index.spec.tsx | 149 ++++++++++++
.../dify-ui/src/input-group/index.stories.tsx | 212 ++++++++++++++++++
packages/dify-ui/src/input-group/index.tsx | 110 +++++++++
.../src/input/__tests__/index.spec.tsx | 12 +
packages/dify-ui/src/input/index.stories.tsx | 38 +---
packages/dify-ui/src/input/index.tsx | 32 ++-
.../src/number-field/__tests__/index.spec.tsx | 15 ++
packages/dify-ui/src/number-field/index.tsx | 4 +-
web/AGENTS.md | 2 +
.../components/mail-and-password-auth.tsx | 23 +-
.../(commonLayout)/account-page/index.tsx | 139 +++++++-----
.../delete-account/components/check-email.tsx | 6 +-
.../add-access-subject-popover.tsx | 20 +-
.../app-sidebar/app-info/app-info-modals.tsx | 43 ++--
.../app/app-publisher/version-info-modal.tsx | 5 +-
.../app/create-app-dialog/app-list/index.tsx | 16 +-
.../components/app/create-app-modal/index.tsx | 6 +-
.../app/create-from-dsl-modal/index.tsx | 5 +-
.../app/overview/settings/index.tsx | 9 +-
.../components/app/switch-app-modal/index.tsx | 11 +-
.../apps/__tests__/creators-filter.spec.tsx | 4 +-
.../apps/app-card/action-bar/index.tsx | 28 ++-
web/app/components/apps/creators-filter.tsx | 21 +-
.../plugins/hitl-input-block/input-field.tsx | 10 +-
.../search-input/__tests__/index.spec.tsx | 31 ++-
.../components/base/search-input/index.tsx | 57 ++---
.../__tests__/index.spec.tsx | 4 +-
.../settings/permission-selector/index.tsx | 36 +--
.../explore/create-app-modal/index.tsx | 28 ++-
.../workspace-role-checkbox-list.spec.tsx | 22 ++
.../api-based-extension-page/modal.tsx | 15 +-
.../data-source-page-new/item.tsx | 1 +
.../invite-modal/email-recipients-field.tsx | 11 +-
.../transfer-ownership-modal/index.tsx | 12 +-
.../workspace-role-checkbox-list.tsx | 31 +--
.../install-from-github/index.tsx | 8 +-
.../marketplace/search-box/tags-filter.tsx | 23 +-
.../components/reasoning-config-form.tsx | 1 +
.../snippet-list/__tests__/index.spec.tsx | 10 +-
web/app/components/snippet-list/index.tsx | 32 +--
.../__tests__/create-snippet-dialog.spec.tsx | 11 +-
.../snippets/create-snippet-dialog.tsx | 24 +-
.../snippets/import-snippet-dsl-dialog.tsx | 11 +-
.../config-credentials.tsx | 32 ++-
.../edit-custom-collection-modal/index.tsx | 40 +++-
web/app/components/tools/mcp/modal.tsx | 29 ++-
.../mcp/sections/authentication-section.tsx | 17 +-
.../mcp/sections/configurations-section.tsx | 13 +-
.../__tests__/snippet-tags-filter.spec.tsx | 4 +-
.../snippets/snippet-tags-filter.tsx | 25 ++-
.../agent-output-variables/edit-card.tsx | 13 +-
.../human-input/components/user-action.tsx | 8 +-
web/app/signin/check-code/page.tsx | 3 +
.../signin/components/mail-and-code-auth.tsx | 5 +-
.../components/mail-and-password-auth.tsx | 18 +-
.../orchestrate/tools/cli-tool/dialog.tsx | 7 +-
.../roster/components/agent-form-fields.tsx | 7 +-
.../components/duplicate-agent-dialog.tsx | 7 +-
.../new-rag/create-knowledge-page.tsx | 11 +-
76 files changed, 1224 insertions(+), 553 deletions(-)
create mode 100644 packages/dify-ui/src/input-group/README.md
create mode 100644 packages/dify-ui/src/input-group/__tests__/index.spec.tsx
create mode 100644 packages/dify-ui/src/input-group/index.stories.tsx
create mode 100644 packages/dify-ui/src/input-group/index.tsx
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/oxlint-suppressions.json b/oxlint-suppressions.json
index 9acd454609f..eae0a944da4 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
@@ -2787,9 +2782,6 @@
}
},
"web/app/components/plugins/install-plugin/install-from-github/index.tsx": {
- "jsx-a11y/no-autofocus": {
- "count": 1
- },
"typescript/no-explicit-any": {
"count": 2
}
@@ -3303,11 +3295,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
@@ -3322,12 +3309,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/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 8597cf73c9e..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 `
- {['skeleton-1', 'skeleton-2', 'skeleton-3', 'skeleton-4'].map((key, index) => (
-
-
-
- {isStartFromScratch
- ? t(($) => $['roster.nodeSelector.startFromScratch'])
- : t(($) => $['roster.nodeSelector.manageInAgentConsole'])}
-
-
-
- )
-}
-
export function AgentBlockItem({
block,
onSelect,
@@ -303,9 +255,10 @@ export function AgentBlockItem({
{
expect(workflowOption).toHaveAttribute('data-selected')
expect(workflowOption.querySelector('.i-ri-check-line')).toBeInTheDocument()
})
+
+ it('should show one named popup state and keep retry outside the listbox', async () => {
+ const user = userEvent.setup()
+ const onRetry = vi.fn()
+
+ render(
+
-
- {t(($) => $['agentDetail.logs.filters.source.loading'])}
-
- )}
- {isError && (
-
- {t(($) => $['agentDetail.logs.filters.source.loadFailed'])}
-
-
- )}
- {!isLoading && !isError && (
- <>
- className="max-h-69 p-2 pt-1">
- {(group) => (
-
-
- {getSourceGroupLabel(group, t)}
-
- >
- {(source) => (
- (
-
+
+
+ $['agentDetail.logs.filters.source.label'])}
+ className="w-80 p-0"
+ >
+ className="max-h-69 p-2 pt-1">
+ {(group) => (
+
+
+ {getSourceGroupLabel(group, t)}
+
+ >
+ {(source) => (
+ (
+
+ )}
+
+ )}
- {t(($) => $['agentDetail.logs.filters.source.empty'])}
+ {!isLoading && !isError ? t(($) => $['agentDetail.logs.filters.source.empty']) : null}
-
- )}
-
+
+
+
)
}
-function SourcePickerStatus({ children, className }: { children: ReactNode; className?: string }) {
- return (
- (
{!!value.length && (
-
+
)}
-
- setOpen(false)}
- />
-
+
+
+
+ setOpen(false)}
+ />
+
+
+
)
diff --git a/web/features/tag-management/components/tag-search-content.tsx b/web/features/tag-management/components/tag-search-content.tsx
index 2604609230d..2f52c967bea 100644
--- a/web/features/tag-management/components/tag-search-content.tsx
+++ b/web/features/tag-management/components/tag-search-content.tsx
@@ -1,5 +1,6 @@
import type { TagType } from '@dify/contracts/api/console/tags/types.gen'
import type { TagComboboxItem } from './tag-combobox-item'
+import { Button } from '@langgenius/dify-ui/button'
import {
ComboboxEmpty,
ComboboxInput,
@@ -11,8 +12,9 @@ import {
ComboboxSeparator,
useComboboxFilteredItems,
} from '@langgenius/dify-ui/combobox'
+import { IconButton } from '@langgenius/dify-ui/icon-button'
import { useAtomValue } from 'jotai'
-import { Fragment } from 'react'
+import { Fragment, useRef } from 'react'
import { useTranslation } from 'react-i18next'
import { workspacePermissionKeysAtom } from '@/context/permission-state'
import { hasPermission } from '@/utils/permission'
@@ -37,12 +39,18 @@ export const TagSearchContent = ({
canBindOrUnbindTags = false,
}: TagSearchContentProps) => {
const { t } = useTranslation()
+ const inputRef = useRef
-
-
)
}
-function getAgentSelectorOptionValue(option: AgentSelectorOption) {
- if (isAgentSelectorActionOption(option)) return option
-
- return option.id
-}
-
-function isAgentSelectorActionOption(
- option: AgentSelectorOption,
-): option is AgentSelectorActionOption {
- return typeof option === 'string'
-}
-
function toAgentRosterNodeData(agent: AgentInviteOptionResponse): AgentRosterNodeData {
return {
description: agent.description,
@@ -250,38 +234,6 @@ function AgentSelectorItem({ agent }: { agent: AgentInviteOptionResponse }) {
)
}
-function AgentSelectorActionItem({ option }: { option: AgentSelectorActionOption }) {
- const { t } = useTranslation('agentV2')
- const isStartFromScratch = option === 'start-from-scratch'
-
- return (
-
-
-
-
+
+
-
+
+
+ {['skeleton-1', 'skeleton-2', 'skeleton-3', 'skeleton-4'].map((key, index) => (
+
+ ))}
+
+
-
+
+
+
- ))}
-
-
-
- $['agentDetail.logs.filters.source.searchLabel'])}
- placeholder={t(($) => $['agentDetail.logs.filters.source.searchPlaceholder'])}
- className="block h-4.5 grow px-1 py-0 system-sm-regular text-components-input-text-filled"
- />
-
-
- {isLoading && (
-
-
-
-
- )}
- />
- )}
-
-
+
+
+ $['agentDetail.logs.filters.source.searchLabel'])}
+ placeholder={t(($) => $['agentDetail.logs.filters.source.searchPlaceholder'])}
+ className="block h-4.5 grow px-1 py-0 system-sm-regular text-components-input-text-filled"
+ />
+
+
+
+ >
+
+ {isLoading
+ ? t(($) => $['agentDetail.logs.filters.source.loading'])
+ : isError
+ ? t(($) => $['agentDetail.logs.filters.source.loadFailed'])
+ : null}
+
+ {isError && (
+
+ )}
+
+ {!isLoading && !isError && (
+
+
+
+
+ )}
+ />
+ )}
+
+
- {children}
-
- )
-}
-
function SourceCheckbox({ checked }: { checked: boolean }) {
return (
{
expect(onOpenTagManagement).toHaveBeenCalledTimes(1)
})
+ it('should expose popup actions after the combobox in keyboard order', async () => {
+ const user = userEvent.setup()
+ render(
@@ -52,21 +60,22 @@ export const TagSearchContent = ({
className="ml-2 i-ri-search-line size-4 shrink-0 text-text-tertiary"
/>
@@ -119,9 +128,10 @@ export const TagSearchContent = ({
<>
-
+
)}
diff --git a/web/features/tag-management/components/tag-selector.tsx b/web/features/tag-management/components/tag-selector.tsx
index 9373debe15e..f69a2264e3c 100644
--- a/web/features/tag-management/components/tag-selector.tsx
+++ b/web/features/tag-management/components/tag-selector.tsx
@@ -1,12 +1,14 @@
import type { TagResponse as Tag, TagType } from '@dify/contracts/api/console/tags/types.gen'
-import type {
- ComboboxContentProps,
- ComboboxProps,
- ComboboxTriggerProps,
-} from '@langgenius/dify-ui/combobox'
+import type { ComboboxProps, ComboboxTriggerProps } from '@langgenius/dify-ui/combobox'
import type { TagComboboxItem } from './tag-combobox-item'
import { cn } from '@langgenius/dify-ui/cn'
-import { Combobox, ComboboxContent, ComboboxTrigger } from '@langgenius/dify-ui/combobox'
+import {
+ Combobox,
+ ComboboxPopup,
+ ComboboxPortal,
+ ComboboxPositioner,
+ ComboboxTrigger,
+} from '@langgenius/dify-ui/combobox'
import { toast } from '@langgenius/dify-ui/toast'
import { useMutation, useQuery } from '@tanstack/react-query'
import { useAtomValue } from 'jotai'
@@ -47,18 +49,7 @@ type TagSelectorRootProps = Omit<
| 'onOpenChangeComplete'
| 'children'
>
-type TagSelectorContentProps = Pick<
- ComboboxContentProps,
- | 'placement'
- | 'sideOffset'
- | 'alignOffset'
- | 'portalProps'
- | 'positionerProps'
- | 'popupProps'
- | 'popupClassName'
->
export type TagSelectorProps = TagSelectorRootProps &
- TagSelectorContentProps &
Pick**Available values:** "all_team_members", "only_me" | | provider | path | | Yes | string | #### Responses @@ -17205,6 +17207,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 @@ -22580,6 +22583,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/datasource_provider_service.py b/api/services/datasource_provider_service.py index 5de194dd262..ed394bd0a29 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 @@ -642,10 +643,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}" @@ -692,16 +706,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( @@ -879,11 +898,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() @@ -895,6 +920,7 @@ class DatasourceProviderService: tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id, + user=user, session=session, ) redirect_uri = ( @@ -938,11 +964,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() @@ -960,6 +989,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/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/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/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/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx index b639f7e2f6b..2720b42237d 100644 --- a/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx +++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/base/credential-selector/__tests__/index.spec.tsx @@ -12,12 +12,14 @@ const credentials = [ name: 'First credential', avatar_url: 'https://example.com/first.png', type: 'oauth2', + is_default: false, }, { id: 'credential-2', name: 'Second credential', avatar_url: 'https://example.com/second.png', type: 'oauth2', + is_default: false, }, ] as DataSourceCredential[] @@ -51,6 +53,23 @@ describe('CredentialSelector', () => { await waitFor(() => expect(onCredentialChange).toHaveBeenCalledWith('credential-1')) }) + it('selects the visible default credential before falling back to the first one', async () => { + const onCredentialChange = vi.fn() + const credentialsWithDefault = credentials.map((credential, index) => ({ + ...credential, + is_default: index === 1, + })) + render( +
- {!!canOAuth && (
- renderConfigurePopover(trigger)}
+ />
)
}
diff --git a/web/app/components/plugins/plugin-auth/authorize/__tests__/add-oauth-button.spec.tsx b/web/app/components/plugins/plugin-auth/authorize/__tests__/add-oauth-button.spec.tsx
index 75729d8ef75..71f86ff6b93 100644
--- a/web/app/components/plugins/plugin-auth/authorize/__tests__/add-oauth-button.spec.tsx
+++ b/web/app/components/plugins/plugin-auth/authorize/__tests__/add-oauth-button.spec.tsx
@@ -1,8 +1,10 @@
import type { OAuthClientSettingsProps } from '../oauth-client-settings'
import type { FormSchema } from '@/app/components/base/form/types'
-import { fireEvent, render, screen, waitFor } from '@testing-library/react'
+import { act, fireEvent, screen, waitFor, within } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
import * as React from 'react'
import { beforeEach, describe, expect, it, vi } from 'vite-plus/test'
+import { renderWithAccountProfile } from '@/test/console/account-profile'
import { AuthCategory } from '../../types'
const mockGetPluginOAuthUrl = vi
@@ -22,9 +24,21 @@ vi.mock('@/hooks/use-oauth', () => ({
}))
vi.mock('../../hooks/use-credential', () => ({
- useGetPluginOAuthUrlHook: () => ({
- mutateAsync: mockGetPluginOAuthUrl,
- }),
+ useGetPluginOAuthUrlHook: () => {
+ const [isPending, setIsPending] = React.useState(false)
+
+ return {
+ isPending,
+ mutateAsync: async (params?: { visibility?: string }) => {
+ setIsPending(true)
+ try {
+ return await mockGetPluginOAuthUrl(params)
+ } finally {
+ setIsPending(false)
+ }
+ },
+ }
+ },
useGetPluginOAuthClientSchemaHook: () => ({
data: {
schema: [],
@@ -40,7 +54,7 @@ vi.mock('../../hooks/use-credential', () => ({
vi.mock('../oauth-client-settings', () => ({
default: (props: OAuthClientSettingsProps) => {
mockOAuthClientSettingsProps.push(props)
- const { open = true, onClose, onOpenChange, schemas } = props
+ const { open = true, onClose, onOpenChange, onRequestAuthorization, schemas } = props
if (!open) return null
@@ -48,12 +62,19 @@ vi.mock('../oauth-client-settings', () => ({
onOpenChange?.(false)
onClose?.()
}
+ const handleSaveAndAuthorize = async () => {
+ handleClose()
+ await onRequestAuthorization?.()
+ }
return (
- >
- {(item) => No tag suggestion. The custom text remains valid.
-
+
+
+
+ >
+ {(item) =>
+ No tag suggestion. The custom text remains valid.
+
+
+
+
),
@@ -933,11 +966,15 @@ export const DisabledAndReadOnly: Story = {
@@ -92,7 +79,25 @@ const Configure = ({ item, pluginPayload, onUpdate, disabled }: ConfigureProps)
-
+ )
+ }
+
+ if (!canOAuth) return renderConfigurePopover()
+
+ return (
+
+
{schemas.map((schema) => (
- {permissionLabel}
-
+
+
+
value={permission}
onValueChange={onChange}
aria-label={permissionLabel}
- className="flex-col items-stretch gap-0 p-1"
>
- value={PermissionLevel.onlyMe} />}
- className={optionClassName}
+
+ aria-label={onlyMeLabel}
+ value={PermissionLevel.onlyMe}
+ closeOnClick
+ className="touch-manipulation"
>
- No suggestion. Free-form text is still valid.
-
+
+
+
+ No suggestion. Free-form text is still valid.
+
+
+
)
@@ -907,12 +934,18 @@ export const Empty: Story = {
@@ -95,13 +116,25 @@ describe('AddOAuthButton', () => {
})
it('should render OAuth button when configured (system params exist)', () => {
- render( {
})
})
+ it('should choose visibility after saving custom OAuth settings for authorization', async () => {
+ const user = userEvent.setup()
+ renderWithAccountProfile(
+ {
},
] as FormSchema[]
- render(
+ renderWithAccountProfile(
{
})
})
- it('selects credential visibility through the native trigger', async () => {
+ it('selects credential visibility through the dropdown menu', async () => {
const user = userEvent.setup()
render((
+ PermissionLevel.onlyMe,
+ )
+ const { mutateAsync: getPluginOAuthUrl, isPending: isGettingOAuthUrl } =
+ useGetPluginOAuthUrlHook(pluginPayload)
+ const { data, isLoading } = useGetPluginOAuthClientSchemaHook(pluginPayload, !oAuthData)
const mergedOAuthData = useMemo(() => {
if (oAuthData) return oAuthData
@@ -78,12 +92,50 @@ const AddOAuthButton = ({
setIsOAuthSettingsOpen(true)
}, [])
const handleOAuth = useCallback(async () => {
- const { authorization_url } = await getPluginOAuthUrl()
+ try {
+ const { authorization_url } = await getPluginOAuthUrl(
+ isVisibilityPickerSupported ? { visibility: pendingVisibility } : undefined,
+ )
+
+ if (!authorization_url) return false
- if (authorization_url) {
openOAuthPopup(authorization_url, () => onUpdate?.())
+ return true
+ } catch {
+ // The request layer surfaces the error. Keep the current UI state so the
+ // user can retry without losing their visibility selection.
+ return false
}
- }, [getPluginOAuthUrl, onUpdate])
+ }, [getPluginOAuthUrl, onUpdate, pendingVisibility, isVisibilityPickerSupported])
+ // Providers without a usable OAuth client first open settings. Once those
+ // settings are saved, authorization continues through the same visibility
+ // dialog used by configured providers.
+ const openVisibilityModal = useCallback(() => {
+ if (!isVisibilityPickerSupported) {
+ if (isConfigured) void handleOAuth()
+ else openOAuthSettings()
+ return
+ }
+ if (isConfigured) {
+ setPendingVisibility(PermissionLevel.onlyMe)
+ setIsVisibilityModalOpen(true)
+ } else {
+ openOAuthSettings()
+ }
+ }, [isConfigured, isVisibilityPickerSupported, openOAuthSettings, handleOAuth])
+ const handleAuthorizationRequest = useCallback(async () => {
+ if (!isVisibilityPickerSupported) {
+ await handleOAuth()
+ return
+ }
+
+ setPendingVisibility(PermissionLevel.onlyMe)
+ setIsVisibilityModalOpen(true)
+ }, [handleOAuth, isVisibilityPickerSupported])
+ const handleVisibilityConfirm = useCallback(async () => {
+ const didOpenOAuthPopup = await handleOAuth()
+ if (didOpenOAuthPopup) setIsVisibilityModalOpen(false)
+ }, [handleOAuth])
const renderCustomLabel = useCallback(
(item: FormSchema) => {
@@ -181,71 +233,72 @@ const AddOAuthButton = ({
}
}, [isConfigured, is_oauth_custom_client_enabled, is_system_oauth_params_exists])
+ const trigger = isConfigured ? (
+ setIsOAuthSettingsOpen(false)}
disabled={disabled || isLoading}
schemas={memorizedSchemas}
- onAuth={handleOAuth}
+ onRequestAuthorization={handleAuthorizationRequest}
editValues={{
...client_params,
__oauth_client__: __auth_client__,
@@ -263,6 +316,14 @@ const AddOAuthButton = ({
onUpdate={onUpdate}
/>
)}
+
disabled?: boolean
schemas: FormSchema[]
- onAuth?: () => Promise
+ onRequestAuthorization?: () => Promise | void
hasOriginalClientParams?: boolean
onUpdate?: () => void
}
@@ -34,7 +34,7 @@ const OAuthClientSettings = ({
editValues,
disabled,
schemas,
- onAuth,
+ onRequestAuthorization,
hasOriginalClientParams,
onUpdate,
}: OAuthClientSettingsProps) => {
@@ -64,7 +64,7 @@ const OAuthClientSettings = ({
const invalidPluginOAuthClientSchema = useInvalidPluginOAuthClientSchemaHook(pluginPayload)
const formRef = useRef(null)
const handleConfirm = useCallback(async () => {
- if (doingActionRef.current) return
+ if (doingActionRef.current) return false
try {
const { isCheckValidated, values } = formRef.current?.getFormValues({
@@ -85,6 +85,7 @@ const OAuthClientSettings = ({
onClose?.()
onUpdate?.()
invalidPluginOAuthClientSchema()
+ return true
} finally {
handleSetDoingAction(false)
}
@@ -99,9 +100,14 @@ const OAuthClientSettings = ({
])
const handleConfirmAndAuthorize = useCallback(async () => {
- await handleConfirm()
- if (onAuth) await onAuth()
- }, [handleConfirm, onAuth])
+ try {
+ const isSaved = await handleConfirm()
+ if (isSaved) await onRequestAuthorization?.()
+ } catch {
+ // The request layer reports the save error. Keep settings open and stop
+ // before opening the permission dialog.
+ }
+ }, [handleConfirm, onRequestAuthorization])
const { mutateAsync: deletePluginOAuthCustomClient } =
useDeletePluginOAuthCustomClientHook(pluginPayload)
const handleRemove = useCallback(async () => {
diff --git a/web/app/components/plugins/plugin-auth/authorize/oauth-visibility-dialog.tsx b/web/app/components/plugins/plugin-auth/authorize/oauth-visibility-dialog.tsx
new file mode 100644
index 00000000000..ce892e8609d
--- /dev/null
+++ b/web/app/components/plugins/plugin-auth/authorize/oauth-visibility-dialog.tsx
@@ -0,0 +1,77 @@
+import type { CredentialPermission } from '@/models/permission'
+import { Button } from '@langgenius/dify-ui/button'
+import {
+ Dialog,
+ DialogCloseButton,
+ DialogContent,
+ DialogDescription,
+ DialogTitle,
+} from '@langgenius/dify-ui/dialog'
+import { memo } from 'react'
+import { useTranslation } from 'react-i18next'
+import PermissionSelector from './permission-selector'
+
+type OAuthVisibilityDialogProps = {
+ open: boolean
+ onOpenChange: (open: boolean) => void
+ permission: CredentialPermission
+ onPermissionChange: (permission: CredentialPermission) => void
+ onConfirm: () => Promise | void
+ loading?: boolean
+}
+
+const OAuthVisibilityDialog = ({
+ open,
+ onOpenChange,
+ permission,
+ onPermissionChange,
+ onConfirm,
+ loading = false,
+}: OAuthVisibilityDialogProps) => {
+ const { t } = useTranslation()
+ const handleOpenChange = (nextOpen: boolean) => {
+ if (loading && !nextOpen) return
+ onOpenChange(nextOpen)
+ }
+
+ return (
+
+ )
+}
+
+export default memo(OAuthVisibilityDialog)
diff --git a/web/app/components/plugins/plugin-auth/authorize/permission-selector.tsx b/web/app/components/plugins/plugin-auth/authorize/permission-selector.tsx
index f38c993bb06..c0f4acc8205 100644
--- a/web/app/components/plugins/plugin-auth/authorize/permission-selector.tsx
+++ b/web/app/components/plugins/plugin-auth/authorize/permission-selector.tsx
@@ -1,31 +1,25 @@
+import type { CredentialPermission } from '@/models/permission'
import { Avatar } from '@langgenius/dify-ui/avatar'
import { cn } from '@langgenius/dify-ui/cn'
import {
- Popover,
- PopoverClose,
- PopoverContent,
- PopoverTitle,
- PopoverTrigger,
-} from '@langgenius/dify-ui/popover'
-import { RadioGroup, RadioItem } from '@langgenius/dify-ui/radio'
+ DropdownMenu,
+ DropdownMenuContent,
+ DropdownMenuRadioGroup,
+ DropdownMenuRadioItem,
+ DropdownMenuRadioItemIndicator,
+ DropdownMenuTrigger,
+} from '@langgenius/dify-ui/dropdown-menu'
import { useSuspenseQuery } from '@tanstack/react-query'
import { useTranslation } from 'react-i18next'
import { userProfileQueryOptions } from '@/features/account-profile/client'
import { PermissionLevel } from '@/models/permission'
-export type CredentialPermission =
- | typeof PermissionLevel.onlyMe
- | typeof PermissionLevel.allTeamMembers
-
type PermissionSelectorProps = {
disabled?: boolean
permission: CredentialPermission
onChange: (permission: CredentialPermission) => void
}
-const optionClassName =
- 'flex w-full touch-manipulation cursor-pointer items-center gap-x-1 rounded-lg border-none bg-transparent px-2 py-1 text-left outline-hidden hover:bg-state-base-hover focus-visible:ring-2 focus-visible:ring-state-accent-solid'
-
const PermissionSelector = ({ disabled, permission, onChange }: PermissionSelectorProps) => {
const { t } = useTranslation()
const { data: userProfile } = useSuspenseQuery({
@@ -35,10 +29,16 @@ const PermissionSelector = ({ disabled, permission, onChange }: PermissionSelect
const isOnlyMe = permission === PermissionLevel.onlyMe
const isAllTeamMembers = permission === PermissionLevel.allTeamMembers
const permissionLabel = t(($) => $['auth.whoCanUse'], { ns: 'plugin' })
+ const onlyMeLabel = t(($) => $['form.permissionsOnlyMe'], { ns: 'datasetSettings' })
+ const allTeamMembersLabel = t(($) => $['form.permissionsAllMember'], {
+ ns: 'datasetSettings',
+ })
+ const selectedPermissionLabel = isOnlyMe ? onlyMeLabel : allTeamMembersLabel
return (
-
-
+
{trigger}
}
+ />,
+ )
+
+ expect(within(screen.getByTestId('external-owner')).getByText('Use OAuth')).toBeInTheDocument()
+ })
+
it('should open OAuth settings modal when settings icon clicked', () => {
- render(
+
+
+ $['auth.oauthClientSettings'], { ns: 'plugin' })}
+ size="lg"
+ className={cn(
+ 'shrink-0 rounded-l-none hover:bg-components-button-primary-bg-hover',
+ buttonRightClassName,
+ )}
+ disabled={disabled}
+ onClick={openOAuthSettings}
+ >
+
+
+
+ ) : (
+
+ )
+
return (
<>
- {renderTrigger?.({
- disabled,
- isConfigured,
- onClick: isConfigured ? handleOAuth : openOAuthSettings,
- })}
- {!renderTrigger && isConfigured && (
-
-
-
- $['auth.oauthClientSettings'], { ns: 'plugin' })}
- size="lg"
- className={cn(
- 'shrink-0 rounded-l-none hover:bg-components-button-primary-bg-hover',
- buttonRightClassName,
- )}
- disabled={disabled}
- onClick={() => {
- openOAuthSettings()
- }}
- >
-
-
-
- )}
- {!renderTrigger && !isConfigured && (
-
- )}
+ {renderTrigger
+ ? renderTrigger({
+ disabled,
+ isConfigured,
+ onClick: openVisibilityModal,
+ trigger,
+ })
+ : trigger}
{isOAuthSettingsMounted && (
- {t(($) => $['form.permissionsOnlyMe'], { ns: 'datasetSettings' })}
+ {onlyMeLabel}
)}
@@ -61,7 +61,7 @@ const PermissionSelector = ({ disabled, permission, onChange }: PermissionSelect
- {t(($) => $['form.permissionsAllMember'], { ns: 'datasetSettings' })}
+ {allTeamMembersLabel}
)}
@@ -69,18 +69,22 @@ const PermissionSelector = ({ disabled, permission, onChange }: PermissionSelect
aria-hidden="true"
className="i-ri-arrow-down-s-line size-4 shrink-0 text-text-quaternary group-data-disabled/permission-trigger:text-components-input-text-placeholder! group-data-popup-open/permission-trigger:text-text-secondary"
/>
-
-
- {t(($) => $['form.permissionsOnlyMe'], { ns: 'datasetSettings' })}
+ value={PermissionLevel.allTeamMembers} />}
- className={optionClassName}
+
+ aria-label={allTeamMembersLabel}
+ value={PermissionLevel.allTeamMembers}
+ closeOnClick
+ className="touch-manipulation"
>
-
-
-
- >
- {(item) => No suggestion. Submit the typed text instead.
-
+
+
+
+
+
+ >
+ {(item) => No suggestion. Submit the typed text instead.
+
+
+
),
@@ -867,11 +890,15 @@ const VirtualizedLongSuggestionsDemo = () => {
+ {onlyMeLabel}
- {isOnlyMe && (
-
- )}
-
-
- {t(($) => $['form.permissionsAllMember'], { ns: 'datasetSettings' })}
+
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
.github/workflows/accessibility-e2e.yml | 2 +-
.github/workflows/api-tests.yml | 6 +++---
.github/workflows/autofix.yml | 2 +-
.github/workflows/db-migration-test.yml | 4 ++--
.github/workflows/pyrefly-diff.yml | 2 +-
.github/workflows/pyrefly-type-coverage-comment.yml | 2 +-
.github/workflows/pyrefly-type-coverage.yml | 2 +-
.github/workflows/style.yml | 2 +-
.github/workflows/translate-i18n-claude.yml | 2 +-
.github/workflows/vdb-tests-full.yml | 2 +-
.github/workflows/vdb-tests.yml | 2 +-
.github/workflows/web-e2e.yml | 2 +-
12 files changed, 15 insertions(+), 15 deletions(-)
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/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'
From 45be4620fead02189d69faa03b0ca5025f1e9f58 Mon Sep 17 00:00:00 2001
From: Weike Zhang <66246918+weike-zhang@users.noreply.github.com>
Date: Mon, 17 Aug 2026 05:58:40 +0000
Subject: [PATCH 04/32] refactor(controllers/web): dep-inject login and
forgot-password payloads (#36659) (#40856)
---
api/controllers/web/forgot_password.py | 22 ++++++++++++----------
api/controllers/web/login.py | 19 +++++++++----------
2 files changed, 21 insertions(+), 20 deletions(-)
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/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)
From f6a472d948cab25f5b9301facb2b1ca1de56d61b Mon Sep 17 00:00:00 2001
From: Bond Zhu <37842169+MRZHUH@users.noreply.github.com>
Date: Mon, 17 Aug 2026 06:06:46 +0000
Subject: [PATCH 05/32] feat(ci): sync E2B template on release tag push
(#40858)
---
.github/workflows/build-push.yml | 29 +++++++++++++++++++++++++++++
dify-agent-runtime/Makefile | 17 ++++++++++++++++-
2 files changed, 45 insertions(+), 1 deletion(-)
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/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
From b02072baee2aed3009825f6cb6187777589aa76a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=97=A0=E8=83=BD=E7=9A=84=E8=BF=AA=E6=8B=9C=E5=B7=A7?=
=?UTF-8?q?=E5=85=8B=E5=8A=9B?=
Date: Mon, 17 Aug 2026 06:10:06 +0000
Subject: [PATCH 06/32] fix(web): let multimodal guidance height fit content
(#40846)
---
.../datasets/common/multimodal-retrieval-guidance/tip-card.tsx | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/web/app/components/datasets/common/multimodal-retrieval-guidance/tip-card.tsx b/web/app/components/datasets/common/multimodal-retrieval-guidance/tip-card.tsx
index 69fe4890fea..28bb7f17e02 100644
--- a/web/app/components/datasets/common/multimodal-retrieval-guidance/tip-card.tsx
+++ b/web/app/components/datasets/common/multimodal-retrieval-guidance/tip-card.tsx
@@ -25,7 +25,7 @@ const renderDescription = (description: string): ReactNode => {
export const TipCard = ({ title, description, dismissLabel, onDismiss }: TipCardProps) => {
return (
-
- No suggestion. Use the text as entered.
-
+
+
+
+ No suggestion. Use the text as entered.
+
+
+
),
@@ -796,15 +815,19 @@ export const LimitResults: Story = {
+ {allTeamMembersLabel}
- {isAllTeamMembers && (
-
- )}
-
-
-
-
+
+
Date: Mon, 17 Aug 2026 06:37:23 +0000
Subject: [PATCH 07/32] fix(web): update agent model request limit copy
(#40863)
---
web/i18n/en-US/agent-v-2.json | 2 +-
web/i18n/ja-JP/agent-v-2.json | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/web/i18n/en-US/agent-v-2.json b/web/i18n/en-US/agent-v-2.json
index 2079cc500ee..5b69c74764a 100644
--- a/web/i18n/en-US/agent-v-2.json
+++ b/web/i18n/en-US/agent-v-2.json
@@ -156,7 +156,7 @@
"agentDetail.configure.preview.empty.description": "Run the agent as a finished chat, exactly how people will experience it once published.",
"agentDetail.configure.preview.empty.title": "Preview {{name}}",
"agentDetail.configure.preview.endUserAuth": "End-user authentication",
- "agentDetail.configure.preview.errors.agentModelRequestLimitExceeded": "This agent run stopped after reaching the maximum number of model requests. Try again, or break the task into smaller steps.",
+ "agentDetail.configure.preview.errors.agentModelRequestLimitExceeded": "This agent run was stopped at the model request limit. Try again, or break the task into smaller steps.",
"agentDetail.configure.preview.errors.agentRunLimitExceeded": "This agent run was stopped at the time limit. Try again, or break the task into smaller steps.",
"agentDetail.configure.preview.inputPlaceholder": "Message {{name}}",
"agentDetail.configure.preview.restart": "Start fresh",
diff --git a/web/i18n/ja-JP/agent-v-2.json b/web/i18n/ja-JP/agent-v-2.json
index 67638994d62..60d1042d83f 100644
--- a/web/i18n/ja-JP/agent-v-2.json
+++ b/web/i18n/ja-JP/agent-v-2.json
@@ -198,7 +198,7 @@
"agentDetail.configure.rightPanel.build": "ビルド",
"agentDetail.configure.rightPanel.buildTipBody": "Build はチャットを通じてエージェントを設定します。やりたいことを説明すると、左側の設定が入力されます。",
"agentDetail.configure.rightPanel.buildTipTitle": "チャットでエージェントを構築",
- "agentDetail.configure.rightPanel.learnMore": "詳しく見る",
+ "agentDetail.configure.rightPanel.learnMore": "詳細を見る",
"agentDetail.configure.rightPanel.modeLabel": "エージェント設定モード",
"agentDetail.configure.rightPanel.preview": "プレビュー",
"agentDetail.configure.rightPanel.previewDisabledTip": "Preview は Dify Cloud と Enterprise でのみ利用できます。",
From 4ce681d42fc56e0dcf7fbc033c0d3aaba5624591 Mon Sep 17 00:00:00 2001
From: yyh <92089059+lyzno1@users.noreply.github.com>
Date: Mon, 17 Aug 2026 06:48:00 +0000
Subject: [PATCH 08/32] test(e2e): align agent model selector locators (#40862)
---
e2e/features/step-definitions/agent-v2/configure.steps.ts | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
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) {
From 64587eae4328c7c2e6577abbc0cc755440b2f1fe Mon Sep 17 00:00:00 2001
From: zyssyz123 <916125788@qq.com>
Date: Mon, 17 Aug 2026 07:31:38 +0000
Subject: [PATCH 09/32] fix(workflow): preserve historical Agent node
compatibility (#40859)
---
api/core/workflow/node_factory.py | 37 +++++++++++++--
.../workflow/nodes/agent_v2/discriminator.py | 28 +++++++++++
api/core/workflow/nodes/agent_v2/entities.py | 2 +-
.../workflow/nodes/agent_v2/validators.py | 3 +-
api/core/workflow/workflow_entry.py | 6 ++-
api/services/agent/composer_candidates.py | 3 +-
api/services/agent/dsl_service.py | 4 +-
.../workflow/node_output_inspector_service.py | 17 ++-----
.../test_node_output_inspector_service.py | 34 ++++++++++++--
.../services/test_app_dsl_service.py | 2 +
.../controllers/console/app/test_workflow.py | 2 +
.../core/app/apps/test_pause_resume.py | 14 +++++-
.../nodes/agent_v2/test_agent_node.py | 4 +-
.../nodes/agent_v2/test_discriminator.py | 31 ++++++++++++
.../nodes/agent_v2/test_validators.py | 17 ++++++-
.../workflow/test_node_mapping_bootstrap.py | 17 ++++++-
.../services/agent/test_agent_dsl_service.py | 15 +++++-
.../services/agent/test_agent_services.py | 23 ++++++++-
.../agent/test_composer_candidates.py | 47 ++++++++++++++++++-
.../agent/test_workflow_publish_service.py | 5 +-
.../test_node_output_inspector_service.py | 15 +++++-
21 files changed, 288 insertions(+), 38 deletions(-)
create mode 100644 api/core/workflow/nodes/agent_v2/discriminator.py
create mode 100644 api/tests/unit_tests/core/workflow/nodes/agent_v2/test_discriminator.py
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/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/validators.py b/api/core/workflow/nodes/agent_v2/validators.py
index 6e56e28f829..03cdf70b78a 100644
--- a/api/core/workflow/nodes/agent_v2/validators.py
+++ b/api/core/workflow/nodes/agent_v2/validators.py
@@ -28,6 +28,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
@@ -258,7 +259,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/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/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/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/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/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/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/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py
index 915cdaa8ec1..ffd7bbb8f0d 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
@@ -343,7 +343,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,
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_validators.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py
index c95f30b81c7..c59b455c96a 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/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..459361f744d 100644
--- a/api/tests/unit_tests/services/agent/test_agent_services.py
+++ b/api/tests/unit_tests/services/agent/test_agent_services.py
@@ -5671,7 +5671,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 +5744,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 +5888,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 +5970,7 @@ class TestWorkflowAgentDraftBindingSync:
"data": {
"type": "agent",
"version": "2",
+ "agent_node_kind": "dify_agent",
"agent_binding": {
"binding_type": "inline_agent",
},
@@ -6017,6 +6025,7 @@ class TestWorkflowAgentDraftBindingSync:
"data": {
"type": "agent",
"version": "2",
+ "agent_node_kind": "dify_agent",
"agent_binding": {
"binding_type": "inline_agent",
},
@@ -6066,6 +6075,7 @@ class TestWorkflowAgentDraftBindingSync:
"data": {
"type": "agent",
"version": "2",
+ "agent_node_kind": "dify_agent",
"agent_task": "Summarize the upstream result.",
"agent_declared_outputs": [
{
@@ -6165,6 +6175,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 +6243,7 @@ class TestWorkflowAgentDraftBindingSync:
"data": {
"type": "agent",
"version": "2",
+ "agent_node_kind": "dify_agent",
"agent_binding": {
"binding_type": "inline_agent",
},
@@ -6280,6 +6292,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 +6351,7 @@ class TestWorkflowAgentDraftBindingSync:
"data": {
"type": "agent",
"version": "2",
+ "agent_node_kind": "dify_agent",
"agent_binding": {
"binding_type": "unknown",
"agent_id": "agent-1",
@@ -6371,6 +6385,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 +6420,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 +6470,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 +6538,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 +6619,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 +6632,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 +6644,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/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()
From 4249a0d05291c7b72cb820cad58b5afd762ea734 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?=
Date: Mon, 17 Aug 2026 08:35:53 +0000
Subject: [PATCH 10/32] fix(dify-agent): require E2B traffic authentication
(#40871)
---
dify-agent/.example.env | 1 -
dify-agent/docs/dify-agent/guide/index.md | 1 -
.../user-manual/shell-layer/index.md | 7 ++-
.../src/dify_agent/runtime_backend/e2b.py | 13 +++--
.../src/dify_agent/runtime_backend/profile.py | 2 -
dify-agent/src/dify_agent/server/settings.py | 2 -
.../dify_agent/runtime_backend/test_e2b.py | 54 +++++++++++++++++++
docker/.env.example | 1 -
docker/docker-compose-template.yaml | 1 -
docker/docker-compose.yaml | 1 -
.../envs/core-services/dify-agent.env.example | 1 -
11 files changed, 67 insertions(+), 17 deletions(-)
diff --git a/dify-agent/.example.env b/dify-agent/.example.env
index bc7a12262df..80efaf1692c 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=
diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md
index 348d3ec97d6..fd436822d77 100644
--- a/dify-agent/docs/dify-agent/guide/index.md
+++ b/dify-agent/docs/dify-agent/guide/index.md
@@ -54,7 +54,6 @@ 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. |
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..4066cc5c866 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
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..d46cc29a31f 100644
--- a/dify-agent/src/dify_agent/server/settings.py
+++ b/dify-agent/src/dify_agent/server/settings.py
@@ -71,7 +71,6 @@ 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(
@@ -205,7 +204,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,
)
)
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..0b4e638f044 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()
@@ -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/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 1a256c7ec15..a826cb4e5e7 100644
--- a/docker/docker-compose-template.yaml
+++ b/docker/docker-compose-template.yaml
@@ -677,7 +677,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 4d95add637f..d4f2154307e 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -683,7 +683,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..a3e61fbed6f 100644
--- a/docker/envs/core-services/dify-agent.env.example
+++ b/docker/envs/core-services/dify-agent.env.example
@@ -32,7 +32,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 signed /files/* transfers.
DIFY_AGENT_SANDBOX_FILES_BASE_URL=http://api:5001
From 5ef8becd259e799ead0e9bfcb7338fa01da08a0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E7=8E=AE=20=28Jade=20Lin=29?=
Date: Mon, 17 Aug 2026 09:29:57 +0000
Subject: [PATCH 11/32] feat: create accounts with client IP (#40875)
---
.../console/auth/email_register.py | 6 ++-
api/controllers/console/auth/login.py | 4 +-
api/controllers/console/auth/oauth.py | 10 ++++-
api/services/account_service.py | 8 +++-
.../console/auth/test_email_register.py | 3 ++
.../auth/test_email_register_language.py | 1 +
.../console/auth/test_email_verification.py | 24 +++++-----
.../controllers/console/auth/test_oauth.py | 16 ++++++-
.../console/auth/test_oauth_timezone.py | 6 ++-
.../services/test_account_service.py | 44 +++++++++++++++++++
10 files changed, 106 insertions(+), 16 deletions(-)
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/services/account_service.py b/api/services/account_service.py
index 3f460275e20..66ca44342da 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,
)
@@ -1961,10 +1965,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)
@@ -2000,6 +2004,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:
@@ -2014,6 +2019,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/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/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py
index 1467f154f28..21f183f6dda 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)
From 96809473fc14d6d7b68a31ef72b11f4db7761511 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?=
Date: Mon, 17 Aug 2026 09:36:21 +0000
Subject: [PATCH 12/32] feat(dify-agent): add context-aware history compaction
(#40872)
---
api/clients/agent_backend/request_builder.py | 3 +
.../apps/agent_app/runtime_request_builder.py | 7 +
api/core/app/llm/model_access.py | 23 ++-
.../nodes/agent_v2/runtime_request_builder.py | 7 +
.../app/apps/agent_app/test_app_runner.py | 8 +-
.../agent_app/test_runtime_request_builder.py | 22 ++-
.../core/app/llm/test_model_access.py | 60 +++++++
.../nodes/agent_v2/test_agent_node.py | 8 +
.../agent_v2/test_runtime_request_builder.py | 27 ++-
api/uv.lock | 42 +++--
.../user-manual/history-layer/index.md | 31 +++-
.../user-manual/plugin-llm-layer/index.md | 29 ++++
.../user-manual/prompt-layer/index.md | 3 +-
.../agenton_examples/pydantic_ai_bridge.py | 18 +-
.../run_pydantic_ai_agent.py | 2 +-
dify-agent/pyproject.toml | 5 +-
.../src/dify_agent/adapters/llm/model.py | 4 +
.../layers/dify_core_tools/client.py | 8 +-
.../layers/dify_core_tools/layer.py | 5 +-
.../dify_agent/layers/dify_plugin/configs.py | 1 +
.../src/dify_agent/runtime/agent_factory.py | 6 +-
.../src/dify_agent/runtime/compaction.py | 45 +++++
dify-agent/src/dify_agent/runtime/history.py | 81 ++-------
dify-agent/src/dify_agent/runtime/runner.py | 26 +--
.../dify_agent/adapters/llm/test_model.py | 49 +++++-
.../local/dify_agent/client/test_client.py | 2 +-
.../layers/dify_core_tools/test_layer.py | 12 +-
.../layers/dify_plugin/test_configs.py | 12 ++
.../layers/dify_plugin/test_layers.py | 4 +-
.../dify_agent/layers/knowledge/test_layer.py | 26 ++-
.../dify_agent/runtime/test_compaction.py | 132 +++++++++++++++
.../runtime/test_compositor_factory.py | 24 +--
.../local/dify_agent/runtime/test_runner.py | 160 +++++++++++++++---
.../runtime/test_runtime_history.py | 81 +++------
.../dify_agent/runtime_backend/test_e2b.py | 2 +-
.../local/examples/test_agenton_examples.py | 10 +-
dify-agent/tests/local/test_packaging.py | 5 +-
dify-agent/uv.lock | 50 ++++--
38 files changed, 765 insertions(+), 275 deletions(-)
create mode 100644 api/tests/unit_tests/core/app/llm/test_model_access.py
create mode 100644 dify-agent/src/dify_agent/runtime/compaction.py
create mode 100644 dify-agent/tests/local/dify_agent/runtime/test_compaction.py
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/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py
index 6fc85435c5d..6c58ce57823 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/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/workflow/nodes/agent_v2/runtime_request_builder.py b/api/core/workflow/nodes/agent_v2/runtime_request_builder.py
index 392c04a9125..a9390d005e0 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/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..495f8073f92 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
@@ -67,9 +67,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]:
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 7c06cc400d5..65e3d8f84d3 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",
@@ -175,11 +190,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]
@@ -197,6 +213,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/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/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py
index ffd7bbb8f0d..71f03fb92a7 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
@@ -62,6 +62,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,
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 2d2e25ca919..292d1adf31b 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},
),
@@ -220,8 +235,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"]}
@@ -238,6 +256,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/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/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/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/layers/dify_core_tools/client.py b/dify-agent/src/dify_agent/layers/dify_core_tools/client.py
index 387185724ef..0707cbac175 100644
--- a/dify-agent/src/dify_agent/layers/dify_core_tools/client.py
+++ b/dify-agent/src/dify_agent/layers/dify_core_tools/client.py
@@ -4,7 +4,7 @@ from __future__ import annotations
import json
from dataclasses import dataclass, field
-from typing import ClassVar
+from typing import ClassVar, cast
import httpx
from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError
@@ -111,9 +111,9 @@ class DifyCoreToolsClient:
request_payload = _DifyCoreToolsInvokeRequest(
caller=_DifyCoreToolsCaller(
tenant_id=execution_context.tenant_id,
- user_id=execution_context.user_id,
- user_from=execution_context.user_from,
- app_id=execution_context.app_id,
+ user_id=cast(str, execution_context.user_id),
+ user_from=cast(str, execution_context.user_from),
+ app_id=cast(str, execution_context.app_id),
invoke_from=execution_context.invoke_from,
conversation_id=execution_context.conversation_id,
workflow_id=execution_context.workflow_id,
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/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/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/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/runner.py b/dify-agent/src/dify_agent/runtime/runner.py
index 6eafa456b20..ef4228ac982 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
@@ -310,12 +311,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 +348,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 +360,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(
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/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py
index e69bec36518..7589eaa53e2 100644
--- a/dify-agent/tests/local/dify_agent/client/test_client.py
+++ b/dify-agent/tests/local/dify_agent/client/test_client.py
@@ -155,7 +155,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(
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/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_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py
index c33e02b701c..371a94e129c 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
@@ -218,7 +219,7 @@ 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
@@ -270,6 +271,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 +314,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 +432,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 +448,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 +497,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 +679,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
@@ -937,9 +1034,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 +1137,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 +1412,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 +1514,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 +1620,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 +2066,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 +2090,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 +2104,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 +2134,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 +2156,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 +2185,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,6 +2203,7 @@ 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(
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 0b4e638f044..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
@@ -310,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(
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]]
From a679ee56444cb15b7b0eccbe6ffbbb4b918dedd6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?=
Date: Mon, 17 Aug 2026 10:11:22 +0000
Subject: [PATCH 13/32] fix(dify-agent): preserve snapshots for failed and
cancelled runs (#40876)
---
api/clients/agent_backend/client.py | 32 ++
api/clients/agent_backend/event_adapter.py | 4 +
api/clients/agent_backend/fake_client.py | 30 +-
api/core/app/apps/agent_app/app_runner.py | 72 +++-
.../workflow/nodes/agent_v2/agent_node.py | 68 +++-
.../clients/agent_backend/test_client.py | 25 ++
.../app/apps/agent_app/test_app_runner.py | 130 ++++++-
.../nodes/agent_v2/test_agent_node.py | 153 +++++++-
.../concepts/run-lifecycle/index.md | 5 +-
dify-agent/docs/dify-agent/guide/index.md | 51 ++-
dify-agent/src/dify_agent/client/_client.py | 43 ++-
dify-agent/src/dify_agent/protocol/schemas.py | 9 +-
.../src/dify_agent/runtime/cancellation.py | 19 +
.../src/dify_agent/runtime/event_sink.py | 41 +--
.../src/dify_agent/runtime/run_scheduler.py | 154 +++++---
dify-agent/src/dify_agent/runtime/runner.py | 15 +-
.../src/dify_agent/storage/redis_keys.py | 7 +-
.../src/dify_agent/storage/redis_run_store.py | 198 +++++++++-
.../storage/test_terminal_finalization.py | 189 ++++++++--
.../local/dify_agent/client/test_client.py | 104 ++++++
.../protocol/test_protocol_schemas.py | 28 +-
.../dify_agent/runtime/test_run_scheduler.py | 234 +++++++++---
.../local/dify_agent/runtime/test_runner.py | 63 +++-
.../storage/test_redis_run_store.py | 341 +++++++-----------
24 files changed, 1575 insertions(+), 440 deletions(-)
create mode 100644 dify-agent/src/dify_agent/runtime/cancellation.py
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/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/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/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/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py
index 495f8073f92..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,
@@ -112,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
@@ -143,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(
@@ -443,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)
@@ -772,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(
@@ -1162,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 == []
@@ -1233,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()
@@ -1242,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/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py
index 71f03fb92a7..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,
@@ -205,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]
@@ -259,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,
@@ -275,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(
@@ -601,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():
@@ -636,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
@@ -652,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)
@@ -819,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
@@ -827,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)
@@ -838,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"
@@ -855,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"}
@@ -877,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)]
@@ -905,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/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 fd436822d77..f24fd846e76 100644
--- a/dify-agent/docs/dify-agent/guide/index.md
+++ b/dify-agent/docs/dify-agent/guide/index.md
@@ -279,26 +279,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
@@ -336,8 +339,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
@@ -349,8 +354,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
@@ -366,12 +374,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/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/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/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/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/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 ef4228ac982..c6c8d829312 100644
--- a/dify-agent/src/dify_agent/runtime/runner.py
+++ b/dify-agent/src/dify_agent/runtime/runner.py
@@ -184,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,
@@ -205,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)
@@ -224,6 +232,7 @@ class AgentRunRunner:
error=message,
error_type=error_type,
reason=reason,
+ session_snapshot=self._terminal_session_snapshot,
)
if finalization.applied:
raise
@@ -287,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
@@ -383,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/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/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py
index 7589eaa53e2..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,
@@ -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/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_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 371a94e129c..c3199922922 100644
--- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py
+++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py
@@ -67,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,
@@ -224,7 +224,7 @@ def test_run_failed_error_payload_classifies_usage_limit() -> None:
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:
@@ -243,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())
@@ -924,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] = []
@@ -2206,7 +2237,7 @@ def test_runner_with_empty_history_layer_uses_instructions_and_saves_full_histor
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"))
@@ -2239,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
@@ -3015,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/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)
From e3ab88f6dea631f872d1af619be4ab136b7d504c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?=
Date: Mon, 17 Aug 2026 10:33:50 +0000
Subject: [PATCH 14/32] feat: dispatch New Agent Beta eligibility checks
(#40725)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
api/configs/feature/__init__.py | 12 ++
api/extensions/ext_celery.py | 1 +
api/services/agent/composer_service.py | 7 +
api/services/billing_service.py | 4 +
api/tasks/new_agent_beta_task.py | 114 +++++++++++
.../services/agent/test_agent_services.py | 8 +
.../services/test_billing_service.py | 24 +++
.../tasks/test_new_agent_beta_task.py | 191 ++++++++++++++++++
8 files changed, 361 insertions(+)
create mode 100644 api/tasks/new_agent_beta_task.py
create mode 100644 api/tests/unit_tests/tasks/test_new_agent_beta_task.py
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/extensions/ext_celery.py b/api/extensions/ext_celery.py
index 184dd5bddfa..724fb4774cc 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/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/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/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/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py
index 459361f744d..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(
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/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",
+ )
From bec3890689835364f48206c79fb4d14e434982b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?=
Date: Mon, 17 Aug 2026 11:48:12 +0000
Subject: [PATCH 15/32] fix(dify-agent): enforce short-lived Agent Stub
authorization (#40884)
---
.../cmd/dify-agent-cli/dumphelp_test.go | 1 +
dify-agent-runtime/cmd/dify-agent-cli/main.go | 10 +-
.../cmd/dify-agent-cli/main_test.go | 5 +
.../internal/agentcli/client.go | 1 +
.../internal/agentcli/client_http.go | 53 ++-
.../internal/agentcli/config.go | 43 ++-
.../internal/agentcli/config_test.go | 339 ++++++++++++++++--
.../internal/agentcli/connect_test.go | 53 +++
dify-agent-runtime/internal/agentcli/file.go | 72 +++-
.../internal/agentcli/file_test.go | 154 ++++++--
.../internal/agentcli/httpclient.go | 52 +++
.../internal/agentcli/httpclient_test.go | 83 +++++
.../user-manual/shell-layer/index.md | 16 +
.../agent_stub/server/control_plane.py | 39 +-
.../agent_stub/server/routes/agent_stub.py | 7 +-
.../agent_stub/server/tokens/__init__.py | 2 +
.../agent_stub/server/tokens/agent_stub.py | 21 +-
.../dify_agent/layers/_agent_cli_help.json | 3 +-
.../src/dify_agent/layers/config/layer.py | 1 +
.../src/dify_agent/layers/shell/layer.py | 41 ++-
.../server/test_agent_stub_routes.py | 77 +++-
.../server/tokens/test_agent_stub.py | 26 +-
.../dify_agent/layers/config/test_layer.py | 5 +
.../dify_agent/layers/shell/test_layer.py | 105 +++++-
24 files changed, 1091 insertions(+), 118 deletions(-)
create mode 100644 dify-agent-runtime/internal/agentcli/connect_test.go
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/docs/dify-agent/user-manual/shell-layer/index.md b/dify-agent/docs/dify-agent/user-manual/shell-layer/index.md
index 4066cc5c866..97155355f62 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
@@ -124,6 +124,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/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/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/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/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/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/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())
From 2aaac4dadb6d9cf111a8f205bd10357b8d2607f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?=
Date: Mon, 17 Aug 2026 11:59:42 +0000
Subject: [PATCH 16/32] feat(dify-agent): enforce signed upload size limits
(#40883)
---
api/controllers/files/upload.py | 11 +++-
api/controllers/inner_api/agent/files.py | 4 +-
api/core/tools/signature.py | 35 +++++++++--
.../controllers/files/test_upload.py | 62 ++++++++++++++++++-
.../controllers/inner_api/test_agent_files.py | 19 ++++++
.../unit_tests/core/tools/test_signature.py | 46 ++++++++++++++
dify-agent/.example.env | 2 +
dify-agent/docs/dify-agent/guide/index.md | 11 ++--
.../user-manual/shell-layer/index.md | 6 ++
.../agent_stub/server/agent_stub_files.py | 5 +-
dify-agent/src/dify_agent/server/settings.py | 7 +++
.../protocol/test_agent_stub_protocol.py | 8 +++
.../server/test_agent_stub_files.py | 2 +
.../local/dify_agent/server/test_settings.py | 21 +++++++
.../envs/core-services/dify-agent.env.example | 2 +
15 files changed, 229 insertions(+), 12 deletions(-)
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/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/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/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/dify-agent/.example.env b/dify-agent/.example.env
index 80efaf1692c..44c20ad0432 100644
--- a/dify-agent/.example.env
+++ b/dify-agent/.example.env
@@ -58,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/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md
index f24fd846e76..68a69772984 100644
--- a/dify-agent/docs/dify-agent/guide/index.md
+++ b/dify-agent/docs/dify-agent/guide/index.md
@@ -58,6 +58,7 @@ also reads `.env` and `dify-agent/.env` when present.
| `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. |
@@ -89,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))'
@@ -110,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
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 97155355f62..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
@@ -93,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
```
@@ -100,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:
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/server/settings.py b/dify-agent/src/dify_agent/server/settings.py
index d46cc29a31f..197592c8815 100644
--- a/dify-agent/src/dify_agent/server/settings.py
+++ b/dify-agent/src/dify_agent/server/settings.py
@@ -77,6 +77,12 @@ class ServerSettings(BaseSettings):
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 = ""
@@ -222,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/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/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/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example
index a3e61fbed6f..44d5f4ebb0f 100644
--- a/docker/envs/core-services/dify-agent.env.example
+++ b/docker/envs/core-services/dify-agent.env.example
@@ -35,6 +35,8 @@ DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600
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.
From a575ae5ef295a6ffee64e2ba45447ef02504e7ea Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:10:44 +0000
Subject: [PATCH 17/32] test: migrate controller auth wrapper sessions and ORM
models to SQLite (#40529)
---
.../controllers/console/explore/test_wraps.py | 47 +-
.../controllers/inner_api/test_auth_wraps.py | 36 +-
.../controllers/openapi/auth/test_prepare.py | 666 +++++++++---------
.../controllers/service_api/test_wraps.py | 12 +-
4 files changed, 381 insertions(+), 380 deletions(-)
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/inner_api/test_auth_wraps.py b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py
index 324d66c0b64..0774da3f4d6 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
@@ -264,7 +265,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
@@ -272,17 +273,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):
@@ -312,21 +316,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/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:
From 0f358f5bb3ab496a808ca7f619c49a8a55585248 Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:10:57 +0000
Subject: [PATCH 18/32] test: migrate console conversation sessions and ORM
models to SQLite (#40528)
---
.../console/app/test_conversation_api.py | 186 +++++++++++++-----
1 file changed, 138 insertions(+), 48 deletions(-)
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")
From 4e02075f8dfa0f774abf9043bd3437ac19d8c4f8 Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:11:12 +0000
Subject: [PATCH 19/32] test: migrate console app MCP sessions and ORM models
to SQLite (#40526)
---
.../console/app/test_mcp_server_response.py | 88 +++++++++++--------
1 file changed, 52 insertions(+), 36 deletions(-)
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(
From ccf4d25321ba7f492f503687d33c30058ff34a26 Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:11:34 +0000
Subject: [PATCH 20/32] test: migrate console workflow controller sessions and
ORM models to SQLite (#40527)
---
.../app/test_conversation_variables_api.py | 55 +++--
.../console/app/test_ops_trace_api.py | 24 +-
.../console/app/test_workflow_comment_api.py | 100 ++++----
.../console/app/test_workflow_convert_api.py | 28 ++-
.../test_workflow_human_input_debug_api.py | 26 ++-
.../test_workflow_node_output_inspector.py | 23 +-
.../console/app/test_workflow_run_api.py | 214 +++++++++++-------
7 files changed, 303 insertions(+), 167 deletions(-)
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_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_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)
From d92cf5d19894b51f77d2257793f46278ac70c84b Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:11:53 +0000
Subject: [PATCH 21/32] test: migrate Service API dataset document sessions and
ORM models to SQLite (#40519)
---
.../service_api/dataset/test_document.py | 291 +++++++++---------
1 file changed, 147 insertions(+), 144 deletions(-)
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,
From c3d090fe4e4b6061f10507e10498b9064791c865 Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:12:06 +0000
Subject: [PATCH 22/32] test: migrate Service API dataset metadata sessions and
ORM models to SQLite (#40518)
---
.../service_api/dataset/test_dataset_apis.py | 20 ++-----
.../dataset/test_dataset_tag_apis.py | 16 ++---
.../service_api/dataset/test_metadata.py | 59 +++++++++++--------
3 files changed, 44 insertions(+), 51 deletions(-)
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_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,
)
From ce1aed8df58e354e54f66e5de4b6520bdf1a0808 Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:12:46 +0000
Subject: [PATCH 23/32] test: migrate RAG pipeline sessions and ORM models to
SQLite (#40511)
---
.../app/apps/pipeline/pipeline_generator.py | 2 -
.../pipeline/test_pipeline_config_manager.py | 21 +-
.../apps/pipeline/test_pipeline_generator.py | 106 ++---
.../app/apps/pipeline/test_pipeline_runner.py | 380 ++++++++----------
4 files changed, 253 insertions(+), 256 deletions(-)
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/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()
From 3f0a9c5060c09d5c7977db36b623928966e5fd1c Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 12:13:02 +0000
Subject: [PATCH 24/32] test: migrate RAG indexing sessions and ORM models to
SQLite (#40513)
---
.../core/rag/indexing/test_index_processor.py | 11 +-
.../core/rag/indexing/test_indexing_runner.py | 517 +++++++++---------
2 files changed, 272 insertions(+), 256 deletions(-)
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="",
@@ -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(
From 02f0e8cffaaa3e5033b29404c5746597acc14b9e Mon Sep 17 00:00:00 2001
From: Asuka Minato
Date: Mon, 17 Aug 2026 13:09:17 +0000
Subject: [PATCH 25/32] test: isolate config inputs from process environment
(#40847)
---
.../unit_tests/configs/_isolated_settings.py | 19 +
.../unit_tests/configs/test_dify_config.py | 395 +++++-------------
.../configs/test_file_upload_config.py | 22 +-
.../configs/test_tidb_on_qdrant_config.py | 19 +-
api/tests/unit_tests/conftest.py | 22 +-
.../libs/test_workspace_permission.py | 42 +-
api/tests/unit_tests/test_config_overrides.py | 19 +
7 files changed, 211 insertions(+), 327 deletions(-)
create mode 100644 api/tests/unit_tests/configs/_isolated_settings.py
create mode 100644 api/tests/unit_tests/test_config_overrides.py
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/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/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)
From 394b43376ab143c32f8aed8dd76dd82bedd4a370 Mon Sep 17 00:00:00 2001
From: Chengxiao
Date: Tue, 18 Aug 2026 01:57:21 +0000
Subject: [PATCH 26/32] fix: suggest deepseek-v4-flash for Agent mode (#40852)
---
.../agent-detail/configure/__tests__/model-compatibility.spec.ts | 1 +
.../agent-v2/agent-detail/configure/model-compatibility.ts | 1 +
2 files changed, 2 insertions(+)
diff --git a/web/features/agent-v2/agent-detail/configure/__tests__/model-compatibility.spec.ts b/web/features/agent-v2/agent-detail/configure/__tests__/model-compatibility.spec.ts
index 25372c6224d..e58ca550e8c 100644
--- a/web/features/agent-v2/agent-detail/configure/__tests__/model-compatibility.spec.ts
+++ b/web/features/agent-v2/agent-detail/configure/__tests__/model-compatibility.spec.ts
@@ -189,6 +189,7 @@ describe('isAgentSuggestedModel', () => {
)
expect(isAgentSuggestedModel(provider, createModelItem('grok-4.3'))).toBe(true)
expect(isAgentSuggestedModel(provider, createModelItem('deepseek-v4-pro'))).toBe(true)
+ expect(isAgentSuggestedModel(provider, createModelItem('deepseek-v4-flash'))).toBe(true)
expect(isAgentSuggestedModel(provider, createModelItem('kimi-k2.6'))).toBe(true)
expect(isAgentSuggestedModel(provider, createModelItem('minimax-m3'))).toBe(true)
expect(isAgentSuggestedModel(provider, createModelItem('qwen3.7-max'))).toBe(true)
diff --git a/web/features/agent-v2/agent-detail/configure/model-compatibility.ts b/web/features/agent-v2/agent-detail/configure/model-compatibility.ts
index 4d52d7ecdbc..ba67ab80772 100644
--- a/web/features/agent-v2/agent-detail/configure/model-compatibility.ts
+++ b/web/features/agent-v2/agent-detail/configure/model-compatibility.ts
@@ -79,6 +79,7 @@ const agentSuggestedModelPatterns: RegExp[] = [
// deepseek
/^deepseek[ .-]v4[ .-]pro$/i,
+ /^deepseek[ .-]v4[ .-]flash$/i,
// qwen
/^qwen[ .-]?3\.7[ .-]max$/i,
From 76a91eb43316d9cac1d0d4daee431dc033de42ec Mon Sep 17 00:00:00 2001
From: yyh <92089059+lyzno1@users.noreply.github.com>
Date: Tue, 18 Aug 2026 02:06:03 +0000
Subject: [PATCH 27/32] refactor(dify-ui): expose autocomplete popup anatomy
(#40895)
---
packages/dify-ui/README.md | 2 +-
.../src/autocomplete/__tests__/index.spec.tsx | 138 ++++++------
.../src/autocomplete/index.stories.tsx | 197 +++++++++++-------
packages/dify-ui/src/autocomplete/index.tsx | 82 ++++----
web/app/education/apply/institution-field.tsx | 109 +++++-----
5 files changed, 292 insertions(+), 236 deletions(-)
diff --git a/packages/dify-ui/README.md b/packages/dify-ui/README.md
index f85108ad7a4..8597cf73c9e 100644
--- a/packages/dify-ui/README.md
+++ b/packages/dify-ui/README.md
@@ -229,7 +229,7 @@ Convert Figma output such as `rounded-[var(--radius/sm, 6px)]` to the mapped Tai
## Overlay & portal contract
-Overlay primitives render their floating surfaces inside a [Base UI Portal] attached to `document.body`. This is the Base UI default — see the upstream [Portals][Base UI Portal] docs for the underlying behavior. Convenience content components such as `DialogContent`, `PopoverContent`, and `SelectContent` own their portal internally; primitives with explicit portal anatomy such as `Drawer` expose the matching `DrawerPortal` part so consumers can compose the full Base UI structure.
+Overlay primitives render their floating surfaces inside a [Base UI Portal] attached to `document.body`. This is the Base UI default — see the upstream [Portals][Base UI Portal] docs for the underlying behavior. Convenience content components such as `DialogContent`, `PopoverContent`, and `SelectContent` own their portal internally. Primitives with explicit anatomy expose their constituent parts so consumers can compose the Base UI structure and put behavior on its owner.
### Root isolation requirement
diff --git a/packages/dify-ui/src/autocomplete/__tests__/index.spec.tsx b/packages/dify-ui/src/autocomplete/__tests__/index.spec.tsx
index fb783cbb346..a026df78d6b 100644
--- a/packages/dify-ui/src/autocomplete/__tests__/index.spec.tsx
+++ b/packages/dify-ui/src/autocomplete/__tests__/index.spec.tsx
@@ -3,7 +3,6 @@ import { render } from 'vitest-browser-react'
import {
Autocomplete,
AutocompleteClear,
- AutocompleteContent,
AutocompleteEmpty,
AutocompleteGroup,
AutocompleteGroupLabel,
@@ -13,6 +12,9 @@ import {
AutocompleteItemIndicator,
AutocompleteItemText,
AutocompleteList,
+ AutocompletePopup,
+ AutocompletePortal,
+ AutocompletePositioner,
AutocompleteSeparator,
AutocompleteStatus,
AutocompleteTrigger,
@@ -41,28 +43,23 @@ const renderAutocomplete = ({
- 2 suggestions
-
-
- Workflow
-
-
- Dataset
-
-
- No suggestions
-
+
+
+
+ 2 suggestions
+
+
+ Workflow
+
+
+ Dataset
+
+
+ No suggestions
+
+
+
)}
,
@@ -182,30 +179,34 @@ describe('Autocomplete wrappers', () => {
.toHaveAttribute('data-align', 'start')
})
- it('should apply custom placement side and passthrough popup props', async () => {
+ it('should apply custom placement and popup props to their owning parts', async () => {
const onPopupClick = vi.fn()
const screen = await renderWithSafeViewport(
-
-
-
- Workflow
-
-
-
+
+
+
+
+
+ Workflow
+
+
+
+
+
,
)
@@ -223,18 +224,27 @@ describe('Autocomplete wrappers', () => {
-
-
-
- Resources
-
- Workflow
-
-
-
-
+
+
+
+
+
+
+ Resources
+
+
+ Workflow
+
+
+
+
+
+
,
)
@@ -253,15 +263,19 @@ describe('Autocomplete wrappers', () => {
-
- >
- {(item) => (
-
- {item}
-
- )}
-
-
+
+
+
+ >
+ {(item) => (
+
+ {item}
+
+ )}
+
+
+
+
,
)
diff --git a/packages/dify-ui/src/autocomplete/index.stories.tsx b/packages/dify-ui/src/autocomplete/index.stories.tsx
index 6afe19dd732..1b8e429b177 100644
--- a/packages/dify-ui/src/autocomplete/index.stories.tsx
+++ b/packages/dify-ui/src/autocomplete/index.stories.tsx
@@ -6,7 +6,6 @@ import {
Autocomplete,
AutocompleteClear,
AutocompleteCollection,
- AutocompleteContent,
AutocompleteEmpty,
AutocompleteGroup,
AutocompleteGroupLabel,
@@ -15,6 +14,9 @@ import {
AutocompleteItem,
AutocompleteItemText,
AutocompleteList,
+ AutocompletePopup,
+ AutocompletePortal,
+ AutocompletePositioner,
AutocompleteSeparator,
AutocompleteStatus,
AutocompleteTrigger,
@@ -369,12 +371,16 @@ const BasicTagAutocomplete = ({ size = 'medium' }: { size?: 'small' | 'medium' |
- >
- {(item) => No tag suggestion. Keep the typed value.
-
+
+
+
+ >
+ {(item) => No tag suggestion. Keep the typed value.
+
+
+
)
@@ -511,15 +517,16 @@ const AsyncSearchDemo = () => {
- {status}
- >
- {(item) =>
+
+
+
+ {status}
+ >
+ {(item) =>
+
+
)
@@ -642,29 +649,33 @@ const FuzzyMatchingDemo = () => {
- >
- {(item) => (
-
- {item.icon && (
-
+
+
+
+ >
+ {(item) => (
+
+ {item.icon && (
+
+ )}
+
)}
-
- )}
-
- No workflow suggestion. Keep typing freely.
-
+
+ No workflow suggestion. Keep typing freely.
+
+
+
)
@@ -729,12 +740,16 @@ export const InlineAutocomplete: Story = {
- >
- {(item) => No inline completion. Continue typing freely.
-
+
+
+
+ >
+ {(item) => No inline completion. Continue typing freely.
+
+
+
),
@@ -761,10 +776,14 @@ export const GroupedSuggestions: Story = {
+
+
+
+ {item.description}
+
+
+
-
-
-
- {item.description}
-
-
-