from types import SimpleNamespace import pytest from app.gateway.services import ( apply_knowledge_space_strict_scope, enforce_knowledge_retrieval_mode_run_config, enforce_knowledge_space_strict_run_config, knowledge_space_strict_excluded_tools, ) def test_strict_scope_preserves_existing_and_blocks_configured_web_tools() -> None: config = SimpleNamespace( tools=[ SimpleNamespace(name="company_search", group="web"), SimpleNamespace(name="calculator", group="utility"), ] ) excluded = knowledge_space_strict_excluded_tools(["custom_disabled"], app_config=config) assert "custom_disabled" in excluded assert "company_search" in excluded assert "web_search" in excluded assert "browser_fetch_page" in excluded assert "calculator" not in excluded assert "llmwiki_search" not in excluded def test_harness_strict_scope_keeps_skills_and_report_delivery_available() -> None: from deerflow.agents.lead_agent.agent import _knowledge_space_strict_exclusions excluded = set(_knowledge_space_strict_exclusions([])) assert "web_search" in excluded assert "task" in excluded assert "skill_view" not in excluded assert "skill_list" not in excluded assert "search_skills" not in excluded assert "read_file" not in excluded assert "deep_research_report" not in excluded def test_selected_knowledge_space_activates_fail_closed_external_retrieval_block() -> None: context = {"excluded_tools": ["custom_disabled"]} config = SimpleNamespace( tools=[ SimpleNamespace(name="company_search", group="web"), SimpleNamespace(name="calculator", group="utility"), ] ) activated = apply_knowledge_space_strict_scope( context, selection_requested=True, knowledge_base_ids=[], assistant_knowledge_base_ids=[], app_config=config, ) assert activated is True assert context["knowledge_space_strict"] is True assert "web_search" in context["excluded_tools"] assert "browser_fetch_page" in context["excluded_tools"] assert "company_search" in context["excluded_tools"] assert "llmwiki_search" not in context["excluded_tools"] assert "calculator" not in context["excluded_tools"] def test_selected_scope_overrides_forged_client_runtime_config() -> None: config = { "configurable": { "knowledge_space_strict": False, "excluded_tools": [], "llmwiki_knowledge_base_ids": ["forged"], "assistant_knowledge_base_ids": ["forged-assistant"], }, "context": { "knowledge_space_strict": False, "excluded_tools": ["client-disabled"], "llmwiki_knowledge_base_ids": ["forged"], "assistant_knowledge_base_ids": ["forged-assistant"], }, } enforce_knowledge_space_strict_run_config( config, knowledge_base_ids=["wiki-authorized"], assistant_knowledge_base_ids=["assistant-authorized"], excluded_tools=["server-disabled"], app_config=SimpleNamespace(tools=[]), ) for scope in ("configurable", "context"): assert config[scope]["knowledge_space_strict"] is True assert config[scope]["llmwiki_knowledge_base_ids"] == [ "wiki-authorized" ] assert config[scope]["assistant_knowledge_base_ids"] == [ "assistant-authorized" ] assert "web_search" in config[scope]["excluded_tools"] assert "server-disabled" in config[scope]["excluded_tools"] def test_retrieval_modes_mount_mutually_exclusive_knowledge_tool_surfaces() -> None: classic = { "configurable": { "llmwiki_knowledge_base_ids": ["kb-1"], "assistant_knowledge_base_ids": [], "excluded_tools": [ "llmwiki_search", "llmwiki_read_wiki_page", "ordinary_knowledge_search", ], }, "context": { "llmwiki_knowledge_base_ids": ["kb-1"], "assistant_knowledge_base_ids": [], "excluded_tools": [], }, } enforce_knowledge_retrieval_mode_run_config(classic, "classic") for scope in ("configurable", "context"): excluded = set(classic[scope]["excluded_tools"]) assert "llmwiki_search" not in excluded assert "llmwiki_read_wiki_page" not in excluded assert "llmwiki_browse_wiki_index" not in excluded assert "ordinary_knowledge_search" in excluded assert "ordinary_wiki_search" in excluded @pytest.mark.parametrize("delivery", [False, True]) def test_classic_strict_knowledge_scope_builds_lead_agent(monkeypatch, delivery) -> None: from deerflow.agents.deep_research import retrieval from deerflow.agents.lead_agent import agent as lead_agent captured_tools: dict = {} app_config = SimpleNamespace( get_model_config=lambda _name: SimpleNamespace(supports_thinking=True), ) monkeypatch.setattr(lead_agent, "_resolve_model_name", lambda *_args, **_kwargs: "test-model") monkeypatch.setattr(lead_agent, "validate_agent_id", lambda _value: "roundtable-coordinator") monkeypatch.setattr(lead_agent, "create_chat_model", lambda **_kwargs: object()) monkeypatch.setattr(lead_agent, "create_agent", lambda **kwargs: kwargs) monkeypatch.setattr(lead_agent, "load_agent_soul", lambda _agent: "") monkeypatch.setattr(lead_agent, "_log_es_query_routing_registration", lambda **_kwargs: None) monkeypatch.setattr(lead_agent, "_build_middlewares", lambda *_args, **_kwargs: []) monkeypatch.setattr(lead_agent, "apply_prompt_template", lambda **_kwargs: "test prompt") monkeypatch.setattr(lead_agent, "_filter_available_skills_for_excluded_tools", lambda skills, _excluded: skills) monkeypatch.setattr(retrieval, "resolve_collector_skill_visibility", lambda _agent, skills, _extra: skills) monkeypatch.setattr(retrieval, "runtime_extra_skills", lambda _config: []) import deerflow.tools def fake_get_available_tools(**kwargs): captured_tools.update(kwargs) return [] monkeypatch.setattr(deerflow.tools, "get_available_tools", fake_get_available_tools) result = lead_agent._make_lead_agent( { "configurable": { "is_bootstrap": False, "writing_mode": True, "workbench_write_now": delivery, "workbench_session_id": "test-session" if delivery else None, "knowledge_space_strict": True, "knowledge_retrieval_mode": "classic", } }, app_config=app_config, ) assert result["system_prompt"].startswith("test prompt") assert captured_tools["include_mcp"] is False assert "web_search" in captured_tools["excluded_tools"] assert ("deep_research_report" in {tool.name for tool in result["tools"]}) is delivery smart = { "configurable": { "llmwiki_knowledge_base_ids": ["kb-1"], "assistant_knowledge_base_ids": [], "excluded_tools": ["ordinary_knowledge_search"], }, "context": { "llmwiki_knowledge_base_ids": ["kb-1"], "assistant_knowledge_base_ids": [], "excluded_tools": [], }, } enforce_knowledge_retrieval_mode_run_config(smart, "smart_reasoning") for scope in ("configurable", "context"): excluded = set(smart[scope]["excluded_tools"]) assert "ordinary_knowledge_search" not in excluded assert "ordinary_wiki_search" not in excluded assert "llmwiki_search" in excluded assert "llmwiki_read_wiki_page" in excluded assert "llmwiki_browse_wiki_index" in excluded def test_classic_without_an_ordinary_space_hides_page_and_index_tools() -> None: config = { "configurable": { "llmwiki_knowledge_base_ids": [], "assistant_knowledge_base_ids": ["assistant-1"], }, "context": { "llmwiki_knowledge_base_ids": [], "assistant_knowledge_base_ids": ["assistant-1"], }, } enforce_knowledge_retrieval_mode_run_config(config, "classic") for scope in ("configurable", "context"): excluded = set(config[scope]["excluded_tools"]) assert "llmwiki_search" not in excluded assert "llmwiki_read_wiki_page" in excluded assert "llmwiki_browse_wiki_index" in excluded