test: replace mock doubles in knowledge-fs core and agent tests with typed fakes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py

@@ -4,8 +4,8 @@
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass
 from types import SimpleNamespace
-from unittest.mock import Mock
 
 import pytest
 from dify_agent.layers.config import DifyConfigLayerConfig, DifyConfigSkillConfig
@@ -32,13 +32,21 @@ from core.app.apps.agent_app.runtime_request_builder import (
     AgentAppRuntimeRequestBuilder,
     AgentAppRuntimeRequestBuildError,
 )
-from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
+from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom
 from core.workflow.file_reference import build_file_reference
 from graphon.file import File, FileTransferMethod, FileType
 from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
 from graphon.model_runtime.entities.model_entities import ModelFeature
 from models.agent_config_entities import AgentSoulConfig
 from tests.unit_tests.config_override import apply_config_overrides
+from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema
+
+
+@dataclass(frozen=True)
+class _LayerDeps:
+    """The only part of the agent runtime Deps the user-prompt capability reads."""
+
+    layers: dict[str, dict[str, object]]
 
 
 @pytest.fixture(autouse=True)
@@ -178,7 +186,7 @@ def _ctx(
     files: tuple[File, ...] = (),
     image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
 ) -> AgentAppRuntimeBuildContext:
-    dify_context = SimpleNamespace(
+    dify_context = DifyRunContext(
         tenant_id="tenant-1",
         app_id="app-1",
         user_id="user-1",
@@ -187,7 +195,7 @@ def _ctx(
         trace_session_id="trace-session-1",
     )
     return AgentAppRuntimeBuildContext(
-        dify_context=dify_context,  # type: ignore[arg-type]
+        dify_context=dify_context,
         agent_id="agent-1",
         agent_config_snapshot_id="snap-1",
         agent_soul=soul,
@@ -363,10 +371,7 @@ class TestAgentAppRuntimeRequestBuilder:
             "core.app.apps.agent_app.runtime_request_builder.resolve_model_supports_vision",
             lambda **_kwargs: vision,
         )
-        model_factory = Mock()
-        model_factory.return_value.init_model_instance.return_value.get_model_schema.return_value = SimpleNamespace(
-            features=[ModelFeature.VISION] if vision else []
-        )
+        model_factory = RecordingModelFactory(schema=llm_schema(features=[ModelFeature.VISION] if vision else []))
         monkeypatch.setattr("core.app.knowledge_runtime_config.DifyModelFactory", model_factory)
 
         def image_content(file: File, **_kwargs: object) -> ImagePromptMessageContent:
@@ -421,7 +426,7 @@ class TestAgentAppRuntimeRequestBuilder:
         assert "knowledge_fs" in layers
         config = DifyUserPromptLayerConfig.model_validate(layers["agent_app_user_prompt"].config)
         # Render the serialized request through the actual runtime module, not just the API DTO.
-        deps = SimpleNamespace(
+        deps = _LayerDeps(
             layers={"agent_app_user_prompt": {"config": json.loads(config.model_dump_json()), "state": {}}}
         )
         prompts = user_prompt_layer.Capability("agent_app_user_prompt").build_user_content(deps)  # type: ignore[arg-type]

api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py

@@ -11,6 +11,7 @@ from flask import Flask, current_app
 from sqlalchemy import Engine, event, select
 from sqlalchemy.orm import Session, sessionmaker
 
+from core.app.app_config.entities import AppAdditionalFeatures
 from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueRetrieverResourcesEvent
 from core.app.entities.task_entities import MessageStreamResponse, StreamEvent, TaskStateMetadata, WorkflowTaskState
 from core.app.task_pipeline import message_cycle_manager as message_cycle_manager_module
@@ -25,6 +26,11 @@ from models.model import App, Conversation, MessageAnnotation, MessageFile
 from tests.unit_tests.model_factories import make_app, make_conversation
 
 
+@dataclass(frozen=True)
+class _AppConfigWithFeatures:
+    additional_features: AppAdditionalFeatures
+
+
 @dataclass(frozen=True)
 class _SQLiteDb:
     engine: Engine
@@ -547,10 +553,10 @@ class TestMessageCycleManagerOptimization:
     def test_knowledge_receipts_survive_message_metadata_without_collapsing_same_document(self, message_cycle_manager):
         from dify_agent.protocol.knowledge_fs import KnowledgeFsCitation
 
-        message_cycle_manager._application_generate_entity.app_config = SimpleNamespace(
-            additional_features=SimpleNamespace(show_retrieve_source=True)
+        message_cycle_manager._application_generate_entity.app_config = _AppConfigWithFeatures(
+            additional_features=AppAdditionalFeatures(show_retrieve_source=True)
         )
-        message_cycle_manager._task_state = SimpleNamespace(metadata=TaskStateMetadata())
+        message_cycle_manager._task_state = WorkflowTaskState(metadata=TaskStateMetadata())
         sources = [
             RetrievalSourceMetadata(
                 dataset_id="space",

api/tests/unit_tests/core/app/test_knowledge_runtime_config.py

@@ -1,14 +1,12 @@
-from types import SimpleNamespace
-from unittest.mock import Mock
-
 import pytest
 from dify_agent.protocol.knowledge_fs import KnowledgeFsError
 
 from core.app import knowledge_runtime_config as mapping
-from core.app.entities.app_invoke_entities import DifyRunContext
+from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom
 from graphon.model_runtime.entities.model_entities import ModelFeature
 from models.agent_config_entities import AgentSoulConfig
 from tests.unit_tests.config_override import apply_config_overrides
+from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema
 
 
 def soul() -> AgentSoulConfig:
@@ -22,23 +20,31 @@ def soul() -> AgentSoulConfig:
     )
 
 
+def run_context() -> DifyRunContext:
+    return DifyRunContext(
+        tenant_id="tenant-1",
+        app_id="app-1",
+        user_id="user-1",
+        user_from=UserFrom.ACCOUNT,
+        invoke_from=InvokeFrom.DEBUGGER,
+    )
+
+
 @pytest.mark.parametrize("vision", [True, False])
 def test_knowledge_vision_is_owned_by_agent_model_without_provider_env(
     monkeypatch: pytest.MonkeyPatch, vision: bool
 ) -> None:
     apply_config_overrides(monkeypatch, AGENT_SHELL_ENABLED=True)
-    factory = Mock()
-    factory.return_value.init_model_instance.return_value.get_model_schema.return_value = SimpleNamespace(
-        features=[ModelFeature.VISION] if vision else []
-    )
+    factory = RecordingModelFactory(schema=llm_schema(features=[ModelFeature.VISION] if vision else []))
     monkeypatch.setattr(mapping, "DifyModelFactory", factory)
-    context = Mock(spec=DifyRunContext)
+    context = run_context()
     config = mapping.build_knowledge_fs_layer_config(soul(), run_context=context)
     assert config is not None
     assert config.agent_supports_vision is vision
     assert config.spaces[0].name == "Docs"
-    factory.assert_called_once_with(run_context=context)
-    factory.return_value.init_model_instance.assert_called_once_with("provider", "model")
+    assert len(factory.run_contexts) == 1
+    assert factory.run_contexts[0] is context
+    assert factory.model_lookups == [("provider", "model")]
     assert "provider" not in config.model_dump_json()
 
 
@@ -46,10 +52,10 @@ def test_no_knowledge_has_no_model_or_shell_dependency_and_import_requires_rebin
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     apply_config_overrides(monkeypatch, AGENT_SHELL_ENABLED=False)
-    assert mapping.build_knowledge_fs_layer_config(AgentSoulConfig(), run_context=Mock(spec=DifyRunContext)) is None
+    assert mapping.build_knowledge_fs_layer_config(AgentSoulConfig(), run_context=run_context()) is None
     with pytest.raises(KnowledgeFsError, match="sandbox shell"):
-        mapping.build_knowledge_fs_layer_config(soul(), run_context=Mock(spec=DifyRunContext))
+        mapping.build_knowledge_fs_layer_config(soul(), run_context=run_context())
     imported = soul()
     imported.knowledge.spaces[0].is_missing = True
     with pytest.raises(KnowledgeFsError, match="Reselect"):
-        mapping.build_knowledge_fs_layer_config(imported, run_context=Mock(spec=DifyRunContext))
+        mapping.build_knowledge_fs_layer_config(imported, run_context=run_context())

api/tests/unit_tests/core/plugin/test_backwards_invocation_datasource.py

@@ -1,5 +1,7 @@
 import json
-from unittest.mock import MagicMock, patch
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from unittest.mock import patch
 
 import pytest
 from pydantic import ValidationError
@@ -16,6 +18,7 @@ from core.datasource.entities.datasource_entities import (
 from core.plugin.backwards_invocation.datasource import PluginDatasourceBackwardsInvocation
 from core.plugin.entities.request import RequestInvokeDatasource
 from models.account import Tenant
+from tests.unit_tests.knowledge_fs_core_fakes import Marker
 
 
 def _tenant() -> Tenant:
@@ -24,10 +27,67 @@ def _tenant() -> Tenant:
     return tenant
 
 
-def _tenant() -> Tenant:
-    tenant = Tenant(name="Test workspace")
-    tenant.id = "tenant-1"
-    return tenant
+@dataclass
+class _RuntimeState:
+    credentials: dict[str, object] = field(default_factory=dict)
+
+
+@dataclass
+class _FakeDatasource:
+    """Bound datasource plugin double that records each operation call and yields canned messages."""
+
+    provider_type: DatasourceProviderType
+    messages: list[object] = field(default_factory=list)
+    runtime: _RuntimeState = field(default_factory=_RuntimeState)
+    calls: list[tuple[str, dict[str, object]]] = field(default_factory=list)
+
+    def datasource_provider_type(self) -> DatasourceProviderType:
+        return self.provider_type
+
+    def _record(self, operation: str, kwargs: dict[str, object]) -> Iterator[object]:
+        self.calls.append((operation, kwargs))
+        return iter(self.messages)
+
+    def get_website_crawl(self, **kwargs: object) -> Iterator[object]:
+        return self._record("get_website_crawl", kwargs)
+
+    def get_online_document_pages(self, **kwargs: object) -> Iterator[object]:
+        return self._record("get_online_document_pages", kwargs)
+
+    def get_online_document_page_content(self, **kwargs: object) -> Iterator[object]:
+        return self._record("get_online_document_page_content", kwargs)
+
+    def online_drive_browse_files(self, **kwargs: object) -> Iterator[object]:
+        return self._record("online_drive_browse_files", kwargs)
+
+    def online_drive_download_file(self, **kwargs: object) -> Iterator[object]:
+        return self._record("online_drive_download_file", kwargs)
+
+
+@dataclass(frozen=True)
+class _ProviderEntity:
+    provider_type: DatasourceProviderType
+
+
+@dataclass
+class _FakeController:
+    entity: _ProviderEntity
+    datasource: _FakeDatasource | None = None
+    need_credentials: bool = True
+    requested_datasources: list[str] = field(default_factory=list)
+    validations: list[dict[str, object]] = field(default_factory=list)
+
+    def get_datasource(self, datasource_name: str) -> _FakeDatasource:
+        self.requested_datasources.append(datasource_name)
+        assert self.datasource is not None
+        return self.datasource
+
+    def validate_credentials(self, **kwargs: object) -> None:
+        self.validations.append(kwargs)
+
+
+def _controller(datasource: _FakeDatasource) -> _FakeController:
+    return _FakeController(entity=_ProviderEntity(datasource.provider_type), datasource=datasource)
 
 
 def datasource_payload(**overrides: object) -> RequestInvokeDatasource:
@@ -58,14 +118,8 @@ def test_request_rejects_operation_and_datasource_type_mismatch() -> None:
 def test_invoke_online_document_uses_dify_bound_runtime_and_credential() -> None:
     payload = datasource_payload()
     message = OnlineDocumentPagesMessage(result=[])
-    runtime = MagicMock()
-    runtime.runtime.credentials = dict[str, object]()
-    runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT
-    runtime.get_online_document_pages.return_value = iter([message])
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT
-    controller.need_credentials = True
-    controller.get_datasource.return_value = runtime
+    runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT, messages=[message])
+    controller = _controller(runtime)
 
     with (
         patch(
@@ -98,11 +152,17 @@ def test_invoke_online_document_uses_dify_bound_runtime_and_credential() -> None
         credential_id="credential-1",
     )
     assert runtime.runtime.credentials == {"token": "resolved-by-dify"}
-    runtime.get_online_document_pages.assert_called_once_with(
-        user_id="user-1",
-        datasource_parameters={"cursor": "next"},
-        provider_type=DatasourceProviderType.ONLINE_DOCUMENT,
-    )
+    assert controller.requested_datasources == ["notion_datasource"]
+    assert runtime.calls == [
+        (
+            "get_online_document_pages",
+            {
+                "user_id": "user-1",
+                "datasource_parameters": {"cursor": "next"},
+                "provider_type": DatasourceProviderType.ONLINE_DOCUMENT,
+            },
+        )
+    ]
 
 
 def test_validate_credentials_uses_resolved_provider_controller() -> None:
@@ -110,11 +170,8 @@ def test_validate_credentials_uses_resolved_provider_controller() -> None:
         datasource_parameters={},
         operation="validate_credentials",
     )
-    runtime = MagicMock()
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT
-    controller.need_credentials = True
-    controller.get_datasource.return_value = runtime
+    runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT)
+    controller = _controller(runtime)
 
     with (
         patch(
@@ -135,10 +192,8 @@ def test_validate_credentials_uses_resolved_provider_controller() -> None:
         )
 
     assert result == [{"result": True}]
-    controller.validate_credentials.assert_called_once_with(
-        user_id="user-1",
-        credentials={"token": "resolved-by-dify"},
-    )
+    assert controller.validations == [{"user_id": "user-1", "credentials": {"token": "resolved-by-dify"}}]
+    assert runtime.calls == []
 
 
 def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None:
@@ -149,15 +204,9 @@ def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None:
         operation="get_website_crawl",
         provider="langgenius/firecrawl_datasource/firecrawl",
     )
-    message = MagicMock()
-    runtime = MagicMock()
-    runtime.runtime.credentials = dict[str, object]()
-    runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL
-    runtime.get_website_crawl.return_value = iter([message])
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.WEBSITE_CRAWL
-    controller.need_credentials = True
-    controller.get_datasource.return_value = runtime
+    message = Marker("datasource-message")
+    runtime = _FakeDatasource(DatasourceProviderType.WEBSITE_CRAWL, messages=[message])
+    controller = _controller(runtime)
 
     with (
         patch(
@@ -179,11 +228,16 @@ def test_invoke_website_crawl_uses_the_bound_datasource_plugin() -> None:
 
     assert result == [message]
     assert runtime.runtime.credentials == {"api_key": "resolved-by-dify"}
-    runtime.get_website_crawl.assert_called_once_with(
-        user_id="user-1",
-        datasource_parameters={"url": "https://example.com"},
-        provider_type=DatasourceProviderType.WEBSITE_CRAWL,
-    )
+    assert runtime.calls == [
+        (
+            "get_website_crawl",
+            {
+                "user_id": "user-1",
+                "datasource_parameters": {"url": "https://example.com"},
+                "provider_type": DatasourceProviderType.WEBSITE_CRAWL,
+            },
+        )
+    ]
 
 
 def test_invoke_online_document_content_builds_the_dify_request_entity() -> None:
@@ -193,15 +247,9 @@ def test_invoke_online_document_content_builds_the_dify_request_entity() -> None
         operation="get_online_document_page_content",
         page=page,
     )
-    message = MagicMock()
-    runtime = MagicMock()
-    runtime.runtime.credentials = dict[str, object]()
-    runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT
-    runtime.get_online_document_page_content.return_value = iter([message])
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT
-    controller.need_credentials = True
-    controller.get_datasource.return_value = runtime
+    message = Marker("datasource-message")
+    runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT, messages=[message])
+    controller = _controller(runtime)
 
     with (
         patch(
@@ -222,11 +270,16 @@ def test_invoke_online_document_content_builds_the_dify_request_entity() -> None
         )
 
     assert result == [message]
-    runtime.get_online_document_page_content.assert_called_once_with(
-        user_id="user-1",
-        datasource_parameters=page,
-        provider_type=DatasourceProviderType.ONLINE_DOCUMENT,
-    )
+    assert runtime.calls == [
+        (
+            "get_online_document_page_content",
+            {
+                "user_id": "user-1",
+                "datasource_parameters": page,
+                "provider_type": DatasourceProviderType.ONLINE_DOCUMENT,
+            },
+        )
+    ]
 
 
 @pytest.mark.parametrize(
@@ -260,18 +313,9 @@ def test_invoke_online_drive_builds_typed_dify_requests(
         provider="langgenius/google_drive/google_drive",
         request=request_data,
     )
-    message = MagicMock()
-    runtime = MagicMock()
-    runtime.runtime.credentials = dict[str, object]()
-    runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE
-    {
-        "online_drive_browse_files": runtime.online_drive_browse_files,
-        "online_drive_download_file": runtime.online_drive_download_file,
-    }[method_name].return_value = iter([message])
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DRIVE
-    controller.need_credentials = True
-    controller.get_datasource.return_value = runtime
+    message = Marker("datasource-message")
+    runtime = _FakeDatasource(DatasourceProviderType.ONLINE_DRIVE, messages=[message])
+    controller = _controller(runtime)
 
     with (
         patch(
@@ -292,13 +336,12 @@ def test_invoke_online_drive_builds_typed_dify_requests(
         )
 
     assert result == [message]
-    call = {
-        "online_drive_browse_files": runtime.online_drive_browse_files,
-        "online_drive_download_file": runtime.online_drive_download_file,
-    }[method_name].call_args
-    assert call.kwargs["user_id"] == "user-1"
-    assert call.kwargs["provider_type"] == DatasourceProviderType.ONLINE_DRIVE
-    assert call.kwargs["request"] == request_type.model_validate(request_data)
+    assert len(runtime.calls) == 1
+    called_operation, call_kwargs = runtime.calls[0]
+    assert called_operation == method_name
+    assert call_kwargs["user_id"] == "user-1"
+    assert call_kwargs["provider_type"] == DatasourceProviderType.ONLINE_DRIVE
+    assert call_kwargs["request"] == request_type.model_validate(request_data)
 
 
 def test_online_drive_metadata_and_binary_survive_the_inner_event_stream() -> None:
@@ -397,10 +440,7 @@ def test_online_drive_metadata_and_binary_survive_the_inner_event_stream() -> No
 
 def test_invoke_rejects_missing_required_credential() -> None:
     payload = datasource_payload()
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DOCUMENT
-    controller.need_credentials = True
-    controller.get_datasource.return_value = MagicMock()
+    controller = _controller(_FakeDatasource(DatasourceProviderType.ONLINE_DOCUMENT))
 
     with (
         patch(
@@ -424,8 +464,7 @@ def test_invoke_rejects_missing_required_credential() -> None:
 
 def test_invoke_rejects_provider_declaration_type_mismatch() -> None:
     payload = datasource_payload()
-    controller = MagicMock()
-    controller.entity.provider_type = DatasourceProviderType.ONLINE_DRIVE
+    controller = _FakeController(entity=_ProviderEntity(DatasourceProviderType.ONLINE_DRIVE))
 
     with (
         patch(

api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py

@@ -1,6 +1,6 @@
 import json
-from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from dataclasses import dataclass, field
+from unittest.mock import patch
 
 import pytest
 
@@ -129,24 +129,111 @@ def test_invoke_summary_uses_same_user_scope_for_token_helpers():
     )
 
 
+def _tenant() -> Tenant:
+    tenant = Tenant(name="Test Workspace")
+    tenant.id = "tenant-1"
+    return tenant
+
+
+@dataclass
+class _RecordingModelInstance:
+    """Bound model instance double: returns a canned result and records invocation kwargs."""
+
+    result: str
+    calls: list[tuple[str, dict[str, object]]] = field(default_factory=list)
+
+    def invoke_text_embedding(self, **kwargs: object) -> str:
+        self.calls.append(("invoke_text_embedding", kwargs))
+        return self.result
+
+    def invoke_multimodal_embedding(self, **kwargs: object) -> str:
+        self.calls.append(("invoke_multimodal_embedding", kwargs))
+        return self.result
+
+    def invoke_rerank(self, **kwargs: object) -> str:
+        self.calls.append(("invoke_rerank", kwargs))
+        return self.result
+
+
+@dataclass(frozen=True)
+class _ModelProvider:
+    provider: str
+
+
+@dataclass(frozen=True)
+class _ActiveModel:
+    deprecated: bool
+    features: list[ModelFeature]
+    fetch_from: str
+    label: dict[str, str]
+    model: str
+    model_properties: dict[str, int]
+    model_type: ModelType
+    provider: _ModelProvider
+    status: str
+
+
+@dataclass
+class _ProviderConfigurations:
+    models: list[_ActiveModel]
+    get_models_calls: list[dict[str, object]] = field(default_factory=list)
+
+    def get_models(self, **kwargs: object) -> list[_ActiveModel]:
+        self.get_models_calls.append(kwargs)
+        return self.models
+
+
+@dataclass(frozen=True)
+class _ProviderManager:
+    configurations: _ProviderConfigurations
+
+    def get_configurations(self, _tenant_id: str) -> _ProviderConfigurations:
+        return self.configurations
+
+
+@dataclass(frozen=True)
+class _InstalledPlugin:
+    plugin_id: str
+    plugin_unique_identifier: str
+
+
+@dataclass(frozen=True)
+class _ParameterRule:
+    name: str
+    use_template: str | None
+    max: object
+    default: object = None
+
+
+@dataclass(frozen=True)
+class _ModelSchema:
+    model_properties: dict[str, int]
+    parameter_rules: list[_ParameterRule]
+
+
+@dataclass(frozen=True)
+class _SchemaModelInstance:
+    schema: _ModelSchema
+
+    def get_model_schema(self) -> _ModelSchema:
+        return self.schema
+
+
 def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance():
-    tenant = SimpleNamespace(id="tenant-1")
+    tenant = _tenant()
     payload = RequestInvokeTextEmbedding(
         provider="langgenius/openai/openai",
         model="text-embedding-3-small",
         texts=["query text"],
         input_type=EmbeddingInputType.QUERY,
     )
-    model_instance = SimpleNamespace(invoke_text_embedding=lambda **_: "unused")
+    model_instance = _RecordingModelInstance(result="embedding-result")
 
-    with (
-        patch.object(
-            PluginModelBackwardsInvocation,
-            "_get_bound_model_instance",
-            return_value=model_instance,
-        ) as mock_get_model,
-        patch.object(model_instance, "invoke_text_embedding", return_value="embedding-result") as mock_invoke,
-    ):
+    with patch.object(
+        PluginModelBackwardsInvocation,
+        "_get_bound_model_instance",
+        return_value=model_instance,
+    ) as mock_get_model:
         result = PluginModelBackwardsInvocation.invoke_text_embedding("user-1", tenant, payload)
 
     assert result == "embedding-result"
@@ -158,11 +245,13 @@ def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance():
         model="text-embedding-3-small",
         request_metadata={},
     )
-    mock_invoke.assert_called_once_with(texts=["query text"], input_type=EmbeddingInputType.QUERY)
+    assert model_instance.calls == [
+        ("invoke_text_embedding", {"texts": ["query text"], "input_type": EmbeddingInputType.QUERY})
+    ]
 
 
 def test_invoke_multimodal_embedding_uses_bound_model_instance():
-    tenant = SimpleNamespace(id="tenant-1")
+    tenant = _tenant()
     documents = [{"content": "AQID", "content_type": "image", "file_id": "asset.png"}]
     payload = RequestInvokeMultimodalEmbedding(
         provider="langgenius/clip/clip",
@@ -170,20 +259,13 @@ def test_invoke_multimodal_embedding_uses_bound_model_instance():
         documents=documents,
         input_type=EmbeddingInputType.DOCUMENT,
     )
-    model_instance = SimpleNamespace(invoke_multimodal_embedding=lambda **_: "unused")
+    model_instance = _RecordingModelInstance(result="embedding-result")
 
-    with (
-        patch.object(
-            PluginModelBackwardsInvocation,
-            "_get_bound_model_instance",
-            return_value=model_instance,
-        ) as mock_get_model,
-        patch.object(
-            model_instance,
-            "invoke_multimodal_embedding",
-            return_value="embedding-result",
-        ) as mock_invoke,
-    ):
+    with patch.object(
+        PluginModelBackwardsInvocation,
+        "_get_bound_model_instance",
+        return_value=model_instance,
+    ) as mock_get_model:
         result = PluginModelBackwardsInvocation.invoke_multimodal_embedding("user-1", tenant, payload)
 
     assert result == "embedding-result"
@@ -195,39 +277,38 @@ def test_invoke_multimodal_embedding_uses_bound_model_instance():
         model="clip-multimodal",
         request_metadata={},
     )
-    mock_invoke.assert_called_once_with(
-        multimodel_documents=documents,
-        input_type=EmbeddingInputType.DOCUMENT,
-    )
+    assert model_instance.calls == [
+        (
+            "invoke_multimodal_embedding",
+            {"multimodel_documents": documents, "input_type": EmbeddingInputType.DOCUMENT},
+        )
+    ]
 
 
 def test_invoke_rerank_accepts_omitted_threshold_and_top_n():
-    tenant = SimpleNamespace(id="tenant-1")
+    tenant = _tenant()
     payload = RequestInvokeRerank(
         provider="langgenius/cohere/cohere",
         model="rerank-v3.5",
         query="query",
         docs=["first", "second"],
     )
-    model_instance = SimpleNamespace(invoke_rerank=lambda **_: "unused")
+    model_instance = _RecordingModelInstance(result="rerank-result")
 
-    with (
-        patch.object(
-            PluginModelBackwardsInvocation,
-            "_get_bound_model_instance",
-            return_value=model_instance,
-        ),
-        patch.object(model_instance, "invoke_rerank", return_value="rerank-result") as mock_invoke,
+    with patch.object(
+        PluginModelBackwardsInvocation,
+        "_get_bound_model_instance",
+        return_value=model_instance,
     ):
         result = PluginModelBackwardsInvocation.invoke_rerank("user-1", tenant, payload)
 
     assert result == "rerank-result"
-    mock_invoke.assert_called_once_with(
-        query="query",
-        docs=["first", "second"],
-        score_threshold=None,
-        top_n=None,
-    )
+    assert model_instance.calls == [
+        (
+            "invoke_rerank",
+            {"query": "query", "docs": ["first", "second"], "score_threshold": None, "top_n": None},
+        )
+    ]
 
 
 @pytest.mark.parametrize(
@@ -237,8 +318,8 @@ def test_invoke_rerank_accepts_omitted_threshold_and_top_n():
 def test_list_models_returns_only_active_tenant_models_with_installed_identity(model_type, selected):
     model_name = "selected-model"
     payload = RequestListModels(model_type=model_type, model=model_name if selected else None, limit=10)
-    provider = SimpleNamespace(provider="langgenius/openai/openai")
-    active_model = SimpleNamespace(
+    provider = _ModelProvider(provider="langgenius/openai/openai")
+    active_model = _ActiveModel(
         deprecated=False,
         features=[ModelFeature.VISION],
         fetch_from="predefined-model",
@@ -249,9 +330,9 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m
         provider=provider,
         status="active",
     )
-    configurations = SimpleNamespace(get_models=MagicMock(return_value=[active_model]))
-    provider_manager = SimpleNamespace(get_configurations=lambda _tenant_id: configurations)
-    installed_plugin = SimpleNamespace(
+    configurations = _ProviderConfigurations(models=[active_model])
+    provider_manager = _ProviderManager(configurations=configurations)
+    installed_plugin = _InstalledPlugin(
         plugin_id="langgenius/openai",
         plugin_unique_identifier="langgenius/openai:1.0.0@sha256:installed",
     )
@@ -278,10 +359,7 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m
     assert result.items[0].provider == "openai"
     assert result.items[0].plugin_unique_identifier == installed_plugin.plugin_unique_identifier
     assert result.items[0].capabilities["features"] == ["vision"]
-    configurations.get_models.assert_called_once_with(
-        model_type=model_type,
-        only_active=True,
-    )
+    assert configurations.get_models_calls == [{"model_type": model_type, "only_active": True}]
     if selected:
         resolve_limits.assert_called_once_with(
             tenant_id="tenant-1", user_id="user-1", provider=provider.provider, model=model_name
@@ -298,11 +376,11 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m
 
 @pytest.mark.parametrize("parameter", ["max_tokens", "max_completion_tokens", "maxtoken", "custom_limit"])
 def test_model_token_limits_use_schema_maximum_not_default(parameter):
-    schema = SimpleNamespace(
+    schema = _ModelSchema(
         model_properties={"context_size": 131072},
-        parameter_rules=[SimpleNamespace(name=parameter, use_template="max_tokens", max=32768.0, default=1024)],
+        parameter_rules=[_ParameterRule(name=parameter, use_template="max_tokens", max=32768.0, default=1024)],
     )
-    instance = SimpleNamespace(get_model_schema=lambda: schema)
+    instance = _SchemaModelInstance(schema=schema)
     with patch.object(PluginModelBackwardsInvocation, "_get_bound_model_instance", return_value=instance) as bound:
         limits = PluginModelBackwardsInvocation._model_token_limits(
             tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model"
@@ -321,14 +399,14 @@ def test_model_token_limits_use_schema_maximum_not_default(parameter):
 
 @pytest.mark.parametrize("maximum", [None, 0, -1, True, float("inf"), float("nan"), 10**400, 12.5, "32768"])
 def test_model_token_limits_do_not_invent_unknown_output_capacity(maximum):
-    schema = SimpleNamespace(
+    schema = _ModelSchema(
         model_properties={"context_size": 8192},
-        parameter_rules=[SimpleNamespace(name="max_tokens", use_template=None, max=maximum, default=4096)],
+        parameter_rules=[_ParameterRule(name="max_tokens", use_template=None, max=maximum, default=4096)],
     )
     with patch.object(
         PluginModelBackwardsInvocation,
         "_get_bound_model_instance",
-        return_value=SimpleNamespace(get_model_schema=lambda: schema),
+        return_value=_SchemaModelInstance(schema=schema),
     ):
         limits = PluginModelBackwardsInvocation._model_token_limits(
             tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model"
@@ -338,17 +416,17 @@ def test_model_token_limits_do_not_invent_unknown_output_capacity(maximum):
 
 
 def test_model_token_limits_respect_the_tighter_declared_alias():
-    schema = SimpleNamespace(
+    schema = _ModelSchema(
         model_properties={},
         parameter_rules=[
-            SimpleNamespace(name="max_tokens", use_template=None, max=32768),
-            SimpleNamespace(name="max_output_tokens", use_template=None, max=8192),
+            _ParameterRule(name="max_tokens", use_template=None, max=32768),
+            _ParameterRule(name="max_output_tokens", use_template=None, max=8192),
         ],
     )
     with patch.object(
         PluginModelBackwardsInvocation,
         "_get_bound_model_instance",
-        return_value=SimpleNamespace(get_model_schema=lambda: schema),
+        return_value=_SchemaModelInstance(schema=schema),
     ):
         limits = PluginModelBackwardsInvocation._model_token_limits(
             tenant_id="tenant-1", user_id="user-1", provider="plugin/provider", model="selected-model"

api/tests/unit_tests/core/plugin/test_model_billing_attribution.py

@@ -1,13 +1,20 @@
 """Exercise backwards invocation through the real hosted quota owner, with no external I/O."""
 
 from collections.abc import Callable
-from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from dataclasses import dataclass
+from unittest.mock import patch
 
 import pytest
 
 from configs import dify_config
-from core.entities.provider_entities import ProviderQuotaType, QuotaUnit
+from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
+from core.entities.provider_entities import (
+    CustomConfiguration,
+    ProviderQuotaType,
+    QuotaConfiguration,
+    QuotaUnit,
+    SystemConfiguration,
+)
 from core.model_context import get_credit_usage_metadata
 from core.model_manager import ModelInstance
 from core.plugin.backwards_invocation.model import PluginModelBackwardsInvocation
@@ -17,12 +24,80 @@ from core.plugin.entities.request import (
     RequestInvokeRerank,
     RequestInvokeTextEmbedding,
 )
+from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime
 from enums import DeploymentEdition
+from graphon.model_runtime.entities.common_entities import I18nObject
 from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
 from graphon.model_runtime.entities.model_entities import ModelType
+from graphon.model_runtime.entities.provider_entities import ProviderEntity
+from graphon.model_runtime.model_providers.base.ai_model import AIModel
+from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
+from graphon.model_runtime.model_providers.base.rerank_model import RerankModel
+from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel
+from models.account import Tenant
 from models.provider import ProviderType
 from services.billing_service import BillingService
+from tests.unit_tests.knowledge_fs_core_fakes import Marker
+
+_MODEL_TYPE_CLASSES: dict[ModelType, type[AIModel]] = {
+    ModelType.LLM: LargeLanguageModel,
+    ModelType.TEXT_EMBEDDING: TextEmbeddingModel,
+    ModelType.RERANK: RerankModel,
+}
+
+
+@dataclass(frozen=True)
+class _ProviderManager:
+    bundle: ProviderModelBundle
+
+    def get_provider_model_bundle(self, **_kwargs: object) -> ProviderModelBundle:
+        return self.bundle
+
+    def get_configurations(self, _tenant_id: str) -> dict[str, ProviderConfiguration]:
+        return {self.bundle.configuration.provider.provider: self.bundle.configuration}
+
+
+def _tenant() -> Tenant:
+    tenant = Tenant(name="Test Workspace")
+    tenant.id = "tenant-1"
+    return tenant
+
+
+def _hosted_bundle(model_type: ModelType) -> ProviderModelBundle:
+    provider = ProviderEntity(
+        provider="langgenius/openai/openai",
+        label=I18nObject(en_US="OpenAI"),
+        supported_model_types=[model_type],
+        configurate_methods=[],
+    )
+    configuration = ProviderConfiguration(
+        tenant_id="tenant-1",
+        provider=provider,
+        preferred_provider_type=ProviderType.SYSTEM,
+        using_provider_type=ProviderType.SYSTEM,
+        system_configuration=SystemConfiguration(
+            enabled=True,
+            current_quota_type=ProviderQuotaType.TRIAL,
+            quota_configurations=[
+                QuotaConfiguration(
+                    quota_type=ProviderQuotaType.TRIAL,
+                    quota_unit=QuotaUnit.CREDITS,
+                    quota_limit=200,
+                    quota_used=0,
+                    is_valid=True,
+                    restrict_models=[],
+                )
+            ],
+        ),
+        custom_configuration=CustomConfiguration(provider=None),
+        model_settings=[],
+    )
+    model_type_instance = _MODEL_TYPE_CLASSES[model_type](
+        provider_schema=provider,
+        model_runtime=create_plugin_model_runtime(tenant_id="tenant-1", user_id="synthetic-end-user"),
+    )
+    return ProviderModelBundle(configuration=configuration, model_type_instance=model_type_instance)
 
 
 @pytest.mark.parametrize(
@@ -73,29 +148,7 @@ def test_model_invocation_preserves_billing_origin(
     config_overrides: Callable[..., None],
 ) -> None:
     config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD)
-    configuration = SimpleNamespace(
-        tenant_id="tenant-1",
-        provider=SimpleNamespace(provider="langgenius/openai/openai"),
-        model_settings=[],
-        using_provider_type=ProviderType.SYSTEM,
-        system_configuration=SimpleNamespace(
-            current_quota_type=ProviderQuotaType.TRIAL,
-            quota_configurations=[
-                SimpleNamespace(
-                    quota_type=ProviderQuotaType.TRIAL,
-                    quota_unit=QuotaUnit.CREDITS,
-                    quota_limit=200,
-                    restrict_models=[],
-                )
-            ],
-        ),
-        get_current_credentials=lambda **_: {"api_key": "local-test-only"},
-        get_provider_model=lambda **_: None,
-    )
-    bundle = SimpleNamespace(configuration=configuration, model_type_instance=MagicMock(model_type=model_type))
-    provider_manager = MagicMock()
-    provider_manager.get_provider_model_bundle.return_value = bundle
-    provider_manager.get_configurations.return_value = {"langgenius/openai/openai": configuration}
+    provider_manager = _ProviderManager(bundle=_hosted_bundle(model_type))
     metadata = {"app_type": "new_rag", "created_by": created_by} if created_by else {}
     payload = payload_type.model_validate(
         {"provider": "langgenius/openai/openai", "model": "selected-model", **body, **metadata}
@@ -117,15 +170,17 @@ def test_model_invocation_preserves_billing_origin(
             index=0, message=AssistantPromptMessage(content="OK"), usage=LLMUsage.empty_usage(), finish_reason="stop"
         ),
     )
-    provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else MagicMock()
+    provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else Marker("provider-result")
     with (
         patch("core.model_manager.create_plugin_provider_manager", return_value=provider_manager),
         patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager),
         patch.object(type(dify_config), "get_model_credits", return_value=3),
+        patch.object(ProviderConfiguration, "get_current_credentials", return_value={"api_key": "local-test-only"}),
+        patch.object(ProviderConfiguration, "get_provider_model", return_value=None),
         patch.object(ModelInstance, provider_method, return_value=provider_result),
         patch.object(BillingService, "_send_quota_request", side_effect=billing_transport),
     ):
-        response = invoke_method("synthetic-end-user", SimpleNamespace(id="tenant-1"), payload)
+        response = invoke_method("synthetic-end-user", _tenant(), payload)
         assert get_credit_usage_metadata() is None
         if model_type == ModelType.LLM:
             assert list(response)[0].delta.message.content == "OK"

api/tests/unit_tests/core/tools/builtin_tool/providers/knowledge_fs/test_create_research_task.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from types import SimpleNamespace
+from dataclasses import dataclass
 from unittest.mock import patch
 
 import pytest
@@ -41,6 +41,18 @@ class AppCapabilities:
         )
 
 
+class RevokedAppCapabilities:
+    def create_research_task(self, **_kwargs: object) -> KnowledgeFSResearchTaskResponse:
+        raise ToolInvokeError("binding revoked")
+
+
+@dataclass(frozen=True)
+class KnowledgeFSRuntimeStub:
+    """The only KnowledgeFS runtime component this tool reaches."""
+
+    app_capabilities: AppCapabilities | RevokedAppCapabilities
+
+
 def _run_context() -> DifyRunContext:
     return DifyRunContext(
         tenant_id="tenant-from-context",
@@ -103,7 +115,7 @@ def test_tool_is_a_real_agent_and_workflow_consumer_using_run_context_identity(
     caller_kind: KnowledgeFSAppSpaceJoinType,
 ) -> None:
     capabilities = AppCapabilities()
-    runtime = SimpleNamespace(app_capabilities=capabilities)
+    runtime = KnowledgeFSRuntimeStub(app_capabilities=capabilities)
     tool = _tool(tool_invoke_from=tool_invoke_from, run_context=_run_context())
     engine = create_engine("sqlite://")
 
@@ -212,7 +224,7 @@ def test_tool_rejects_cross_tenant_and_non_application_callers_before_runtime_io
 
 def test_tool_omits_unspecified_mode_and_maps_invalid_model_input_to_tool_error() -> None:
     capabilities = AppCapabilities()
-    runtime = SimpleNamespace(app_capabilities=capabilities)
+    runtime = KnowledgeFSRuntimeStub(app_capabilities=capabilities)
     tool = _tool(tool_invoke_from=ToolInvokeFrom.AGENT, run_context=_run_context())
     engine = create_engine("sqlite://")
 
@@ -256,11 +268,7 @@ def test_tool_omits_unspecified_mode_and_maps_invalid_model_input_to_tool_error(
 
 def test_tool_preserves_explicit_tool_errors_from_the_capability_boundary() -> None:
     tool = _tool(tool_invoke_from=ToolInvokeFrom.AGENT, run_context=_run_context())
-    runtime = SimpleNamespace(
-        app_capabilities=SimpleNamespace(
-            create_research_task=lambda **_: (_ for _ in ()).throw(ToolInvokeError("binding revoked"))
-        )
-    )
+    runtime = KnowledgeFSRuntimeStub(app_capabilities=RevokedAppCapabilities())
     engine = create_engine("sqlite://")
 
     with (

api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval_v2/test_automatic_metadata_filter.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import json
-from types import SimpleNamespace
+from dataclasses import dataclass, field
 
 import pytest
 
@@ -29,6 +29,14 @@ from graphon.nodes.llm.entities import ModelConfig
 from services.knowledge_fs.product_dto import KnowledgeFSMetadataFieldResponse
 
 
+@dataclass(frozen=True)
+class _ModelConfigEntity:
+    """The prompt and invocation settings the extractor reads from the resolved model config."""
+
+    stop: list[str]
+    parameters: dict[str, object] = field(default_factory=dict)
+
+
 def _field(name: str, field_type: str) -> KnowledgeFSMetadataFieldResponse:
     return KnowledgeFSMetadataFieldResponse.model_validate(
         {
@@ -61,7 +69,7 @@ def test_intersect_metadata_fields_keeps_only_fields_shared_with_identical_types
 
 
 def test_build_metadata_filter_prompt_mirrors_legacy_chat_and_completion_templates() -> None:
-    model_config = SimpleNamespace(stop=["</json>"])
+    model_config = _ModelConfigEntity(stop=["</json>"])
 
     chat_messages, chat_stop = build_metadata_filter_prompt(
         model_config=model_config,  # type: ignore[arg-type]
@@ -206,7 +214,7 @@ def test_extractor_invokes_the_configured_model_with_the_legacy_prompt_and_credi
         '{"metadata_map": [{"metadata_field_name": "department", "metadata_field_value": "finance", '
         '"comparison_operator": "="}]}'
     )
-    model_config_entity = SimpleNamespace(parameters={"temperature": 0.2}, stop=["END"])
+    model_config_entity = _ModelConfigEntity(parameters={"temperature": 0.2}, stop=["END"])
     fetched: list[dict[str, object]] = []
 
     def fake_fetch_model_config(**kwargs: object) -> tuple[object, object]:

api/tests/unit_tests/core/workflow/test_node_runtime.py

@@ -410,13 +410,24 @@ def test_dify_prepared_polling_llm_commits_successful_reservation() -> None:
     reservation.release.assert_not_called()
 
 
-def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result() -> None:
+def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     succeeded_result = LLMPollingResult.model_construct(status=LLMPollingStatus.SUCCEEDED, result=None)
-    polling_runtime = SimpleNamespace(
-        start_llm_polling=Mock(return_value=succeeded_result),
-        check_llm_polling=Mock(),
+    polling_runtime = PluginModelRuntime(
+        tenant_id="tenant-id",
+        user_id="user-id",
+        client=PluginModelClient(),
+        plugin_service=PluginService,
     )
-    reservation = MagicMock()
+    started_pollings: list[dict[str, object]] = []
+
+    def start_llm_polling(**kwargs: object) -> LLMPollingResult:
+        started_pollings.append(kwargs)
+        return succeeded_result
+
+    monkeypatch.setattr(polling_runtime, "start_llm_polling", start_llm_polling)
+    reservation = sentinel.quota_reservation
     model_instance = _QuotaManagedModelInstanceStub(
         model_schema=_build_model_schema(features=[ModelFeature.POLLING]),
         model_runtime=polling_runtime,
@@ -437,6 +448,7 @@ def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_resu
         )
 
     release_quota_safely.assert_called_once_with(reservation)
+    assert len(started_pollings) == 1
 
 
 def test_dify_prepared_polling_llm_releases_previous_reservation_on_restart() -> None:

api/tests/unit_tests/knowledge_fs_core_fakes.py

@@ -8,6 +8,8 @@ from __future__ import annotations
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
 
+from graphon.model_runtime.entities.common_entities import I18nObject
+from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType
 from services.knowledge_fs.observability import (
     KnowledgeFSBatchStatusMetric,
     KnowledgeFSCapabilityIssuanceMetric,
@@ -100,3 +102,39 @@ class RecordingOperationalMetrics:
 
     def register_control_space_state_gauge(self, read_counts: Callable[[], Mapping[str, int]]) -> None:
         self.gauge_readers.append(read_counts)
+
+
+def llm_schema(*, features: list[ModelFeature]) -> AIModelEntity:
+    return AIModelEntity(
+        model="model",
+        label=I18nObject(en_US="model"),
+        model_type=ModelType.LLM,
+        fetch_from=FetchFrom.PREDEFINED_MODEL,
+        model_properties={},
+        features=features,
+    )
+
+
+@dataclass
+class SchemaModelInstance:
+    schema: AIModelEntity
+
+    def get_model_schema(self) -> AIModelEntity:
+        return self.schema
+
+
+@dataclass
+class RecordingModelFactory:
+    """Stands in for DifyModelFactory and records the constructor context and model lookups."""
+
+    schema: AIModelEntity
+    run_contexts: list[object] = field(default_factory=list)
+    model_lookups: list[tuple[str, str]] = field(default_factory=list)
+
+    def __call__(self, *, run_context: object) -> RecordingModelFactory:
+        self.run_contexts.append(run_context)
+        return self
+
+    def init_model_instance(self, provider: str, model: str) -> SchemaModelInstance:
+        self.model_lookups.append((provider, model))
+        return SchemaModelInstance(self.schema)

api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from datetime import UTC, datetime
 
 import pytest
@@ -12,7 +13,6 @@ from models.knowledge_fs_cutover import (
 )
 from services.knowledge_fs.cutover_runtime_gate import SQLKnowledgeFSWorkspaceRuntimeGate
 from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError
-from tests.unit_tests.knowledge_fs_core_fakes import RecordingCallable
 
 _WATERMARK: KnowledgeFSCutoverRevisionWatermark = {
     "membership_epoch": 1,
@@ -34,39 +34,43 @@ def test_runtime_gate_fails_closed_without_a_cutover_ledger(sqlite_session: Sess
 
 
 class _Initializer:
-    def __init__(self, ensure_initialized: RecordingCallable) -> None:
-        self.ensure_initialized = ensure_initialized
+    def __init__(self, initialize: Callable[..., KnowledgeFSWorkspaceCutoverLedger]) -> None:
+        self._initialize = initialize
+        self.calls: list[str] = []
+
+    def ensure_initialized(self, *, tenant_id: str) -> KnowledgeFSWorkspaceCutoverLedger:
+        self.calls.append(tenant_id)
+        return self._initialize(tenant_id=tenant_id)
 
 
 @pytest.mark.parametrize("sqlite_session", [(KnowledgeFSWorkspaceCutoverLedger,)], indirect=True)
 def test_runtime_gate_initializes_a_missing_greenfield_workspace_once(sqlite_session: Session) -> None:
     session_maker = sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False)
 
-    def initialize(*, tenant_id: str) -> None:
+    def initialize(*, tenant_id: str) -> KnowledgeFSWorkspaceCutoverLedger:
+        ledger = KnowledgeFSWorkspaceCutoverLedger(
+            tenant_id=tenant_id,
+            source_revision_watermark=_WATERMARK,
+            applied_revision_watermark=_WATERMARK,
+            phase=KnowledgeFSWorkspaceCutoverPhase.CUTOVER,
+            freeze_at=datetime(2026, 7, 21, 11, 59, tzinfo=UTC),
+            cutover_at=datetime(2026, 7, 21, 12, 0, tzinfo=UTC),
+            product_routes_enabled=True,
+            capability_v2_enabled=True,
+            integrated_mode_enabled=True,
+            legacy_acl_read_only=True,
+        )
         with session_maker.begin() as session:
-            session.add(
-                KnowledgeFSWorkspaceCutoverLedger(
-                    tenant_id=tenant_id,
-                    source_revision_watermark=_WATERMARK,
-                    applied_revision_watermark=_WATERMARK,
-                    phase=KnowledgeFSWorkspaceCutoverPhase.CUTOVER,
-                    freeze_at=datetime(2026, 7, 21, 11, 59, tzinfo=UTC),
-                    cutover_at=datetime(2026, 7, 21, 12, 0, tzinfo=UTC),
-                    product_routes_enabled=True,
-                    capability_v2_enabled=True,
-                    integrated_mode_enabled=True,
-                    legacy_acl_read_only=True,
-                )
-            )
+            session.add(ledger)
+        return ledger
 
-    ensure_initialized = RecordingCallable(side_effect=initialize)
-    initializer = _Initializer(ensure_initialized)
+    initializer = _Initializer(initialize)
     gate = SQLKnowledgeFSWorkspaceRuntimeGate(session_maker, initializer=initializer)
 
     gate.require_product_routes(tenant_id="tenant-1")
     gate.require_capability_v2(tenant_id="tenant-1")
 
-    ensure_initialized.assert_called_once_with(tenant_id="tenant-1")
+    assert initializer.calls == ["tenant-1"]
 
 
 @pytest.mark.parametrize("sqlite_session", [(KnowledgeFSWorkspaceCutoverLedger,)], indirect=True)

api/tests/unit_tests/services/test_knowledge_fs_observability.py

@@ -33,12 +33,20 @@ class _FakeInstrument:
         self.log.append(call)
 
 
+@dataclass(frozen=True)
+class _GaugeRegistration:
+    name: str
+    callbacks: tuple[Callable[[object], Iterable[Observation]], ...]
+    description: str
+    unit: str
+
+
 @dataclass
 class _FakeMeter:
     log: list[RecordedCall] = field(default_factory=list)
     counters: dict[str, _FakeInstrument] = field(default_factory=dict)
     histograms: dict[str, _FakeInstrument] = field(default_factory=dict)
-    gauge_calls: list[RecordedCall] = field(default_factory=list)
+    gauge_calls: list[_GaugeRegistration] = field(default_factory=list)
 
     def create_counter(self, name: str, *, description: str, unit: str) -> _FakeInstrument:
         self.log.append(RecordedCall(args=(name,), kwargs={"description": description, "unit": unit}))
@@ -56,9 +64,10 @@ class _FakeMeter:
         description: str,
         unit: str,
     ) -> object:
-        call = RecordedCall(args=(name,), kwargs={"callbacks": callbacks, "description": description, "unit": unit})
-        self.gauge_calls.append(call)
-        self.log.append(call)
+        self.gauge_calls.append(_GaugeRegistration(name=name, callbacks=callbacks, description=description, unit=unit))
+        self.log.append(
+            RecordedCall(args=(name,), kwargs={"callbacks": callbacks, "description": description, "unit": unit})
+        )
         return object()
 
 
@@ -129,9 +138,9 @@ def test_control_space_current_state_gauge_uses_bounded_states_and_is_best_effor
 
     assert len(meter.gauge_calls) == 1
     gauge_call = meter.gauge_calls[0]
-    assert gauge_call.args == ("dify.knowledge_fs.control_spaces",)
-    assert gauge_call.kwargs["unit"] == "{space}"
-    callback = gauge_call.kwargs["callbacks"][0]
+    assert gauge_call.name == "dify.knowledge_fs.control_spaces"
+    assert gauge_call.unit == "{space}"
+    callback = gauge_call.callbacks[0]
     observations = tuple(callback(object()))
     assert [(observation.value, observation.attributes) for observation in observations] == [
         (3, {"aggregation_scope": "global_database_snapshot", "state": "provisioning"}),
@@ -139,7 +148,7 @@ def test_control_space_current_state_gauge_uses_bounded_states_and_is_best_effor
         (1, {"aggregation_scope": "global_database_snapshot", "state": "error"}),
     ]
 
-    assert "aggregate replica series with max, never sum" in gauge_call.kwargs["description"]
+    assert "aggregate replica series with max, never sum" in gauge_call.description
 
     counts_error.append(RuntimeError("database unavailable"))
     assert tuple(callback(object())) == ()

api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py

@@ -9,6 +9,7 @@ from services import knowledge_fs_capability
 from services.knowledge_fs import lifecycle_remote_http, remote_registry
 from tests.unit_tests.config_override import apply_config_overrides
 from tests.unit_tests.knowledge_fs_core_fakes import Marker
+from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote
 
 
 def test_configured_lifecycle_remote_assembles_capability_v2_http_client(
@@ -71,7 +72,7 @@ def test_lifecycle_remote_assembly_fails_closed_without_endpoint_or_capability_i
 def test_explicit_lifecycle_remote_factory_remains_available_for_isolated_process_tests(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    remote = Marker("remote")
+    remote = UnexpectedLifecycleRemote()
     monkeypatch.setattr(remote_registry, "_remote_factory", None)
 
     remote_registry.configure_knowledge_fs_lifecycle_remote_factory(lambda: remote)

api/tests/unit_tests/services/test_knowledge_fs_text_store.py

@@ -21,7 +21,7 @@ from services.knowledge_fs.text_store_native import (
     WeaviateTextStore,
     lexical_sparse_vector,
 )
-from services.knowledge_fs.vector_store import VectorStoreUnavailableError
+from services.knowledge_fs.vector_store import VectorResult, VectorStoreUnavailableError
 from services.knowledge_fs.vector_store_elasticsearch import ElasticsearchVectorStore
 from services.knowledge_fs.vector_store_weaviate import WeaviateVectorStore
 from tests.unit_tests.knowledge_fs_core_fakes import Marker
@@ -32,9 +32,10 @@ SCOPE = {"tenant_id": A, "knowledge_space_id": B, "index_version": "fts-v1"}
 
 @dataclass(frozen=True)
 class _StaticTextClient:
-    result: dict
+    result: VectorResult
 
-    def execute(self, _payload: TextPayload) -> dict:
+    def execute(self, payload: TextPayload) -> VectorResult:
+        assert isinstance(payload, TextPayload)
         return self.result
 
 

dify-agent/tests/local/dify_agent/agent_stub/server/test_knowledge.py

@@ -4,6 +4,7 @@ import asyncio
 import base64
 import io
 import json
+from dataclasses import dataclass
 from uuid import uuid4
 
 import httpx
@@ -14,6 +15,7 @@ from PIL import Image
 from pydantic import ValidationError
 from pydantic_ai.messages import BinaryContent, ModelRequest, ToolReturnPart, UserPromptPart
 
+from dify_agent.adapters.shell.protocols import CompleteShellCommandResult
 from dify_agent.agent_stub.server.knowledge import AgentStubKnowledgeHandler, _validate_thumbnail, project_result
 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 AgentStubPrincipal, AgentStubTokenCodec
@@ -34,6 +36,7 @@ from dify_agent.protocol.knowledge_fs import (
     KnowledgeFsPreparedRequest,
 )
 from dify_agent.storage.redis_keys import run_cancel_intent_key, run_record_key
+from dify_agent.runtime.context import Deps
 from dify_agent.storage.redis_knowledge_sessions import RedisKnowledgeSessionStore
 
 SPACE = "00000000-0000-4000-8000-000000000001"
@@ -383,12 +386,31 @@ def test_search_receipt_without_parse_id_reopens_and_pins_provenance():
         project_result(command("open", space="docs", receipt_id=RECEIPT), plan, raw, previous)
 
 
-def knowledge_capability(store, client, *, agent_supports_vision=False, **services):
-    from types import SimpleNamespace
+@dataclass(frozen=True)
+class _RunCtx:
+    """The run-context fields the KnowledgeFS capability reads around one run."""
 
+    deps: Deps
+    prompt: str
+
+
+def _remote_result(output: str) -> CompleteShellCommandResult:
+    return CompleteShellCommandResult(
+        job_id="job",
+        status="exited",
+        done=True,
+        exit_code=0,
+        output=output,
+        output_complete=True,
+        incomplete_reason=None,
+        offset=len(output),
+    )
+
+
+def knowledge_capability(store, client, *, agent_supports_vision=False, **services):
     from dify_agent.layers.knowledge_fs.layer import Capability, Config
     from dify_agent.layers.shell.layer import ShellSession
-    from dify_agent.runtime.context import Deps, Services
+    from dify_agent.runtime.context import Services
 
     deps = Deps(
         layers={
@@ -411,11 +433,10 @@ def knowledge_capability(store, client, *, agent_supports_vision=False, **servic
     )
     shell = ShellSession("shell", deps)
     deps.resources.shells["shell"] = shell
-    return Capability("knowledge_fs"), shell, SimpleNamespace(deps=deps, prompt="inspect evidence")
+    return Capability("knowledge_fs"), shell, _RunCtx(deps=deps, prompt="inspect evidence")
 
 
 def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch):
-    from types import SimpleNamespace
     from unittest.mock import AsyncMock
     from pydantic_ai import Agent
     from pydantic_ai.models.test import TestModel
@@ -425,8 +446,8 @@ def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch):
         redis, store, _ = await session_fixture()
         remote = AsyncMock(
             side_effect=[
-                SimpleNamespace(exit_code=0, output_complete=True, output="1\n"),
-                SimpleNamespace(exit_code=0, output_complete=True, output=delivery().result.model_dump_json()),
+                _remote_result("1\n"),
+                _remote_result(delivery().result.model_dump_json()),
             ]
         )
         monkeypatch.setattr(ShellSession, "run_remote_script", remote)
@@ -491,7 +512,6 @@ def test_real_layer_delivers_images_as_model_parts_and_cleans_up(monkeypatch):
 def test_layer_preflight_rejects_incompatible_deployments_and_closes_lease(
     monkeypatch, protocol, gateway_output, error
 ):
-    from types import SimpleNamespace
     from unittest.mock import AsyncMock
     from dify_agent.layers.shell.layer import ShellSession
 
@@ -504,8 +524,8 @@ def test_layer_preflight_rejects_incompatible_deployments_and_closes_lease(
             ).model_dump_json()
         remote = AsyncMock(
             side_effect=[
-                SimpleNamespace(exit_code=0, output_complete=True, output=protocol),
-                SimpleNamespace(exit_code=0, output_complete=True, output=output),
+                _remote_result(protocol),
+                _remote_result(output),
             ]
         )
         monkeypatch.setattr(ShellSession, "run_remote_script", remote)

dify-agent/tests/local/dify_agent/runtime/test_runner.py

@@ -504,11 +504,16 @@ async def test_empty_current_user_content_fails_before_model_io(monkeypatch, use
 @pytest.mark.anyio
 @pytest.mark.parametrize("failed", [False, True])
 async def test_quality_report_follows_final_run_event(monkeypatch, failed):
+    from dataclasses import dataclass
     from unittest.mock import AsyncMock
 
     from dify_agent.runtime.runner import RunSuccessOutcome
 
-    module = SimpleNamespace(publish_investigation=AsyncMock())
+    @dataclass(frozen=True)
+    class QualityReportModule:
+        publish_investigation: AsyncMock
+
+    module = QualityReportModule(publish_investigation=AsyncMock())
     async with httpx.AsyncClient() as client:
         services = Services(client, client)
         runner = AgentRunRunner(sink=InMemoryRunEventSink(), request=request(), run_id="quality-run", services=services)