from __future__ import annotations import asyncio from types import SimpleNamespace from unittest.mock import Mock import pytest from fastapi import HTTPException from langchain_core.messages import HumanMessage from langgraph.checkpoint.memory import InMemorySaver from langgraph.store.memory import InMemoryStore from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.gateway import analysis_projects, services from app.gateway import knowledge_research as run_policy from app.gateway.routers import llmwiki as llmwiki_router from app.gateway.routers.thread_runs import RunCreateRequest from deerflow.integrations.weknora.client import WeKnoraError from deerflow.persistence.base import Base from deerflow.persistence.knowledge_research import ( KnowledgeResearchSessionRow, KnowledgeResearchStore, ) from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore from deerflow.runtime.user_context import reset_current_user, set_current_user class _MappingStore: def __init__(self, mapping: dict) -> None: self.mapping = mapping self.calls: list[tuple[str, str, bool, bool]] = [] async def get_authorized( self, mapping_id: str, user_id: str, *, write: bool, is_admin: bool, ) -> dict | None: self.calls.append((mapping_id, user_id, write, is_admin)) if mapping_id != self.mapping["id"]: return None if is_admin or user_id == self.mapping["owner_user_id"]: return dict(self.mapping) if not write and self.mapping.get("publication_status") == "published": return dict(self.mapping) return None class _WikiClient: def __init__(self) -> None: self.created: list[tuple[str, dict]] = [] self.updated: list[tuple[str, str, dict]] = [] self.pages: dict[str, dict] = {} async def create_wiki_page(self, knowledge_base_id: str, payload: dict) -> dict: self.created.append((knowledge_base_id, dict(payload))) page = {"id": "page-one", **payload} self.pages[str(payload["slug"])] = page return dict(page) async def update_wiki_page( self, knowledge_base_id: str, slug: str, payload: dict, ) -> dict: self.updated.append((knowledge_base_id, slug, dict(payload))) page = {"id": "page-one", "slug": slug, **payload} self.pages[slug] = page return dict(page) async def get_wiki_page(self, knowledge_base_id: str, slug: str) -> dict: return dict(self.pages[slug]) def _request(**state_values): return SimpleNamespace( app=SimpleNamespace(state=SimpleNamespace(**state_values)) ) def test_knowledge_research_model_is_registered_for_startup_create_all() -> None: from deerflow.persistence import models assert models.KnowledgeResearchSessionRow is KnowledgeResearchSessionRow assert "knowledge_research_sessions" in Base.metadata.tables @pytest.mark.asyncio async def test_get_by_thread_missing_table_is_not_a_research_session() -> None: engine = create_async_engine("sqlite+aiosqlite:///:memory:") try: store = KnowledgeResearchStore( async_sessionmaker(engine, expire_on_commit=False) ) assert await store.get_by_thread("thread-one") is None finally: await engine.dispose() @pytest.mark.asyncio async def test_prepare_run_treats_missing_table_as_ordinary_chat() -> None: engine = create_async_engine("sqlite+aiosqlite:///:memory:") try: store = KnowledgeResearchStore( async_sessionmaker(engine, expire_on_commit=False) ) body = RunCreateRequest( input={"messages": [{"role": "user", "content": "普通问答"}]}, ) prepared, resolved = await run_policy.prepare_knowledge_research_run( body, "thread-one", _request(knowledge_research_store=store), ) assert prepared is body assert resolved is None finally: await engine.dispose() @pytest.mark.asyncio async def test_sql_store_keeps_question_title_separate_from_publication_title() -> None: engine = create_async_engine("sqlite+aiosqlite:///:memory:") try: async with engine.begin() as connection: await connection.run_sync(KnowledgeResearchSessionRow.__table__.create) store = KnowledgeResearchStore( async_sessionmaker(engine, expire_on_commit=False) ) created = await store.create( knowledge_base_id="kb-one", user_id="alice", title="原始问题", session_id="session-one", publication_id="publication-one", thread_id="thread-one", ) await store.update_owned( created["id"], "alice", draft_markdown="# 答案", ) saved = await store.save_publication( created["id"], "alice", title="可修改的文件名", markdown="# 答案", saved_hash="a" * 64, wiki_slug="knowledge-research/publication-one", wiki_page_id="page-one", wiki_folder_id="folder-one", ) assert saved is not None assert saved["title"] == "原始问题" assert saved["saved_title"] == "可修改的文件名" assert (await store.get_by_thread("thread-one"))["id"] == "session-one" assert (await store.get_publication("publication-one"))["id"] == "session-one" finally: await engine.dispose() @pytest.mark.asyncio async def test_memory_store_serializes_draft_update_with_publication() -> None: store = KnowledgeResearchStore(None) created = await store.create( knowledge_base_id="kb-one", user_id="alice", title="问题", session_id="session-one", publication_id="publication-one", thread_id="thread-one", ) await store.update_owned( created["id"], "alice", draft_markdown="# 旧草稿", ) async with store.locked_publication(created["id"], "alice") as publication: assert publication is not None pending_patch = asyncio.create_task( store.update_owned( created["id"], "alice", draft_markdown="# 新草稿", ) ) await asyncio.sleep(0) assert pending_patch.done() is False publication.save( title="文件名", markdown="# 旧草稿", saved_hash="a" * 64, wiki_slug="knowledge-research/publication-one", wiki_page_id="page-one", wiki_folder_id="folder-one", ) await pending_patch current = await store.get(created["id"]) assert current is not None assert current["draft_markdown"] == "# 新草稿" assert current["saved_markdown"] == "# 旧草稿" @pytest.mark.asyncio async def test_start_run_bypasses_global_fixed_answers_for_research( monkeypatch, ) -> None: """Research must reach scoped KB retrieval, never a global canned answer.""" async def no_project(body, _thread_id, _request): return body, None async def research_binding(body, _thread_id, _request): return body, {"id": "session-one", "mode": "knowledge_research"} async def forbidden_fixed_answer(**_kwargs): raise AssertionError("knowledge research must bypass fixed answers") class _ReachedModelResolution(RuntimeError): pass def stop_after_fixed_answer_policy(*_args, **_kwargs): raise _ReachedModelResolution monkeypatch.setattr(analysis_projects, "prepare_project_run", no_project) monkeypatch.setattr(run_policy, "prepare_knowledge_research_run", research_binding) monkeypatch.setattr(services, "get_stream_bridge", lambda _request: object()) monkeypatch.setattr(services, "get_run_manager", lambda _request: object()) monkeypatch.setattr( services, "get_run_context", lambda _request: SimpleNamespace(app_config=object()), ) monkeypatch.setattr(services, "get_concurrency_gate", lambda _request: None) monkeypatch.setattr(services, "get_effective_user_id", lambda: "alice") monkeypatch.setattr(services, "_find_fixed_answer", forbidden_fixed_answer) monkeypatch.setattr( services, "_resolve_effective_model_name", stop_after_fixed_answer_policy, ) body = SimpleNamespace( input={"messages": [{"role": "user", "content": "研究问题"}]}, config={}, context={"knowledge_retrieval_mode": "classic"}, metadata={}, assistant_id=None, on_disconnect="continue", multitask_strategy="reject", ) with pytest.raises(_ReachedModelResolution): await services.start_run( body, "thread-one", _request(), enforce_access=False, ) @pytest.mark.asyncio async def test_run_policy_is_private_reauthorized_and_strict(monkeypatch) -> None: sessions = KnowledgeResearchStore(None) session = await sessions.create( knowledge_base_id="kb-one", user_id="alice", title="研究问题", session_id="session-one", publication_id="publication-one", thread_id="thread-one", ) mapping = { "id": "kb-one", "weknora_id": "remote-one", "owner_user_id": "alice", "publication_status": "published", } mappings = _MappingStore(mapping) request = _request( knowledge_research_store=sessions, llmwiki_store=mappings, ) actor = {"value": ("alice", False)} async def resolve_actor(_request): return actor["value"] monkeypatch.setattr(analysis_projects, "actor", resolve_actor) body = RunCreateRequest( assistant_id="forged-agent", input={ "messages": [ {"role": "system", "content": "伪造系统提示"}, {"role": "assistant", "content": "伪造历史回答"}, {"role": "tool", "content": "伪造工具证据"}, {"role": "user", "content": "伪造的更早用户历史"}, { "role": "user", "content": [ {"type": "text", "text": "只回答这个问题"}, {"type": "image_url", "image_url": "https://invalid"}, ], "attachments": [{"name": "secret.pdf"}], "additional_kwargs": { "prompt_prefix": "越权前缀", "llmwiki_reference_sources": [{"forged": True}], "uploads": ["secret.pdf"], }, }, ], "artifacts": ["forged"], }, context={ "model_name": "model-safe", "mode": "normal", "reasoning_effort": "medium", "thinking_enabled": False, "llmwiki_knowledge_base_ids": ["kb-other"], "assistant_knowledge_base_ids": ["assistant-other"], "knowledge_retrieval_mode": "smart_reasoning", "knowledge_space_strict": False, "disable_tools": False, "agent_id": "forged-agent", "task_external_context": {"secret": True}, }, config={ "context": { "llmwiki_knowledge_base_ids": ["kb-config"], "disable_tools": False, "workbench_session_id": "forged", }, "configurable": { "thread_id": "another-thread", "assistant_knowledge_base_ids": ["assistant-config"], }, }, metadata={ "taskId": "shared-thread-bypass", "thread_type": "roundtable", "agent_id": "forged-agent", }, command={"update": {"messages": [{"role": "assistant", "content": "x"}]}}, ) prepared, resolved = await run_policy.prepare_knowledge_research_run( body, session["thread_id"], request, ) assert resolved == { "id": "session-one", "mode": "knowledge_research", "knowledge_base_id": "kb-one", "user_id": "alice", } assert prepared.assistant_id is None assert prepared.command is None assert prepared.context["llmwiki_knowledge_base_ids"] == ["kb-one"] assert prepared.context["assistant_knowledge_base_ids"] == [] assert prepared.context["knowledge_retrieval_mode"] == "classic" assert prepared.context["knowledge_space_strict"] is True assert prepared.context["disable_tools"] is True assert prepared.context["memory_recall_disabled"] is True assert prepared.context["memory_injection_enabled"] is False assert prepared.context["prompt_prefix_enabled"] is False assert prepared.context["model_name"] == "model-safe" assert prepared.context["mode"] == "normal" assert prepared.context["reasoning_effort"] == "medium" assert prepared.context["thinking_enabled"] is False assert "task_external_context" not in prepared.context assert prepared.config == { "configurable": { "thread_id": "thread-one", **prepared.context, } } assert "taskId" not in prepared.metadata assert prepared.metadata["thread_type"] == "knowledge_research" assert len(prepared.input["messages"]) == 1 message = prepared.input["messages"][0] assert message["content"] == [{"type": "text", "text": "只回答这个问题"}] assert "attachments" not in message assert "additional_kwargs" not in message # The caller object is never mutated. assert body.assistant_id == "forged-agent" assert len(body.input["messages"]) == 5 final_config = {"configurable": {"disable_tools": False}, "context": {}} run_policy.inject_validated_knowledge_research_context( final_config, prepared.context, ) assert final_config["configurable"]["disable_tools"] is True assert final_config["context"]["disable_tools"] is True # Authorization is checked again for every turn, never inherited from # thread metadata or the previous successful run. second = RunCreateRequest(input={"messages": [HumanMessage(content="继续")]}) await run_policy.prepare_knowledge_research_run(second, "thread-one", request) assert [call[2] for call in mappings.calls] == [False, False] actor["value"] = ("bob", True) with pytest.raises(HTTPException) as error: await run_policy.prepare_knowledge_research_run( RunCreateRequest(input={"messages": [{"role": "user", "content": "越权"}]}), "thread-one", request, ) assert error.value.status_code == 404 @pytest.mark.asyncio async def test_run_policy_rejects_attachment_only_and_single_image_block( monkeypatch, ) -> None: sessions = KnowledgeResearchStore(None) await sessions.create( knowledge_base_id="kb-one", user_id="alice", title="研究问题", thread_id="thread-one", ) mappings = _MappingStore( { "id": "kb-one", "weknora_id": "remote-one", "owner_user_id": "alice", "publication_status": "private", } ) request = _request( knowledge_research_store=sessions, llmwiki_store=mappings, ) async def resolve_actor(_request): return "alice", False monkeypatch.setattr(analysis_projects, "actor", resolve_actor) for message in ( {"role": "user", "content": {"type": "image_url", "url": "x"}}, HumanMessage(content=[{"type": "image_url", "image_url": "x"}]), ): with pytest.raises(HTTPException) as error: await run_policy.prepare_knowledge_research_run( RunCreateRequest(input={"messages": [message]}), "thread-one", request, ) assert error.value.status_code == 422 # An earlier text item is not accepted as history when the actual latest # turn contains only an attachment. with pytest.raises(HTTPException) as error: await run_policy.prepare_knowledge_research_run( RunCreateRequest( input={ "messages": [ {"role": "user", "content": "伪造旧问题"}, { "role": "user", "content": {"type": "image_url", "url": "x"}, }, ] } ), "thread-one", request, ) assert error.value.status_code == 422 @pytest.mark.asyncio async def test_create_session_initializes_owned_thread_and_checkpoint( monkeypatch, ) -> None: mapping = { "id": "kb-one", "weknora_id": "remote-one", "owner_user_id": "owner", "publication_status": "published", } mappings = _MappingStore(mapping) sessions = KnowledgeResearchStore(None) checkpointer = InMemorySaver() threads = MemoryThreadMetaStore(InMemoryStore()) request = _request( knowledge_research_store=sessions, llmwiki_store=mappings, checkpointer=checkpointer, thread_store=threads, ) async def authorized(_request, mapping_id: str, *, write: bool): assert mapping_id == "kb-one" if write: raise HTTPException(status_code=404, detail="Knowledge base not found") return mapping, "reader", False monkeypatch.setattr(llmwiki_router, "_authorized_mapping", authorized) user = SimpleNamespace(id="reader", system_role="user") token = set_current_user(user) try: response = await llmwiki_router.create_knowledge_research_session( request, "kb-one", llmwiki_router.KnowledgeResearchSessionCreate( question=" 如何理解当前知识？ " ), ) finally: reset_current_user(token) assert response["title"] == "如何理解当前知识？" assert response["can_write"] is False assert response["publication"] is None thread = await threads.get(response["thread_id"], user_id="reader") assert thread is not None assert thread["display_name"] == response["title"] assert thread["metadata"]["system"] is True assert thread["metadata"]["thread_type"] == "knowledge_research" checkpoint = await checkpointer.aget_tuple( { "configurable": { "thread_id": response["thread_id"], "checkpoint_ns": "", } } ) assert checkpoint is not None with pytest.raises(HTTPException) as publish_denied: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", response["id"], llmwiki_router.KnowledgeResearchPublicationSave(title="无权发布"), ) assert publish_denied.value.status_code == 404 @pytest.mark.asyncio async def test_publication_overwrites_stable_wiki_page_without_leaking_chat( monkeypatch, ) -> None: mapping = { "id": "kb-one", "weknora_id": "remote-one", "owner_user_id": "alice", "publication_status": "published", } mappings = _MappingStore(mapping) sessions = KnowledgeResearchStore(None) created = await sessions.create( knowledge_base_id="kb-one", user_id="alice", title="原始问题", session_id="session-one", publication_id="publication-one", thread_id="thread-one", ) await sessions.update_owned( created["id"], "alice", draft_markdown="# 初稿\n\n结论。", ) request = _request( knowledge_research_store=sessions, llmwiki_store=mappings, ) actor = {"value": "alice"} async def authorized(_request, mapping_id: str, *, write: bool): user_id = actor["value"] row = await mappings.get_authorized( mapping_id, user_id, write=write, is_admin=False, ) if row is None: raise HTTPException(status_code=404, detail="Knowledge base not found") return row, user_id, False client = _WikiClient() sync = Mock() monkeypatch.setattr(llmwiki_router, "_authorized_mapping", authorized) monkeypatch.setattr(llmwiki_router, "_client_or_503", lambda: client) async def ensure_directory(_client, remote_id: str, paths: list[list[str]]): assert remote_id == "remote-one" assert paths == [["深度研究"]] return {"深度研究": "folder-research"} monkeypatch.setattr(llmwiki_router, "ensure_directory_paths", ensure_directory) monkeypatch.setattr(llmwiki_router, "_schedule_local_wiki_sync", sync) first = await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave(title="研究报告"), ) first_hash = first["publication"]["saved_hash"] slug = first["publication"]["wiki_slug"] assert first["title"] == "原始问题" assert first["publication"]["title"] == "研究报告" assert slug == "knowledge-research/publication-one" assert client.created[0][1]["folder_id"] == "folder-research" assert client.created[0][1]["content"] == "# 初稿\n\n结论。" assert sync.call_count == 1 with pytest.raises(HTTPException) as missing_version: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave(title="缺少版本"), ) assert missing_version.value.status_code == 409 with pytest.raises(HTTPException) as stale: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave( title="仅改标题", expected_saved_hash="0" * 64, ), ) assert stale.value.status_code == 409 assert not client.updated renamed = await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave( title="仅改标题", expected_saved_hash=first_hash, ), ) assert renamed["publication"]["wiki_slug"] == slug assert renamed["title"] == "原始问题" assert renamed["publication"]["title"] == "仅改标题" assert renamed["publication"]["saved_hash"] != first_hash assert client.updated[0][1] == slug assert sync.call_count == 2 # If the previous response was lost, retrying the identical desired # publication with the old pre-save token is an idempotent no-op. retried = await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave( title="仅改标题", expected_saved_hash=first_hash, ), ) assert retried["publication"]["saved_hash"] == renamed["publication"]["saved_hash"] assert len(client.updated) == 1 assert sync.call_count == 2 entered_remote_update = asyncio.Event() release_remote_update = asyncio.Event() original_update = client.update_wiki_page async def blocking_update(*args, **kwargs): entered_remote_update.set() await release_remote_update.wait() return await original_update(*args, **kwargs) monkeypatch.setattr(client, "update_wiki_page", blocking_update) concurrent_a = asyncio.create_task( llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave( title="并发版本 A", expected_saved_hash=renamed["publication"]["saved_hash"], ), ) ) await entered_remote_update.wait() concurrent_b = asyncio.create_task( llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave( title="并发版本 B", expected_saved_hash=renamed["publication"]["saved_hash"], ), ) ) await asyncio.sleep(0) release_remote_update.set() concurrent_results = await asyncio.gather( concurrent_a, concurrent_b, return_exceptions=True, ) assert isinstance(concurrent_results[0], dict) assert isinstance(concurrent_results[1], HTTPException) assert concurrent_results[1].status_code == 409 assert len(client.updated) == 2 assert sync.call_count == 3 other = await sessions.create( knowledge_base_id="kb-one", user_id="bob", title="Bob 的私人会话", session_id="session-bob", publication_id="publication-bob", thread_id="thread-bob", ) await sessions.update_owned( other["id"], "bob", draft_markdown="# 私人内容", ) # KB owner has write permission but cannot read/publish another user's # private research session. with pytest.raises(HTTPException) as owner_cannot_publish_private_chat: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-bob", llmwiki_router.KnowledgeResearchPublicationSave(title="越权发布"), ) assert owner_cannot_publish_private_chat.value.status_code == 404 async def deposit_authorized(_request, mapping_id: str, *, write: bool): return { **mapping, "owner_user_id": "system", "name": llmwiki_router.CONVERSATION_DEPOSIT_KB_NAME, }, "alice", True monkeypatch.setattr( llmwiki_router, "_authorized_mapping", deposit_authorized, ) with pytest.raises(HTTPException) as managed_deposit: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave(title="禁止发布"), ) assert managed_deposit.value.status_code == 403 monkeypatch.setattr(llmwiki_router, "_authorized_mapping", authorized) # An unsaved draft must not alter the shared bookshelf snapshot. await sessions.update_owned( "session-one", "alice", draft_markdown="# 尚未保存的新内容", ) actor["value"] = "bob" listing = await llmwiki_router.list_knowledge_research_sessions( request, "kb-one", ) assert [row["id"] for row in listing["sessions"]] == ["session-bob"] assert all(row["id"] != "session-one" for row in listing["sessions"]) assert len(listing["publications"]) == 1 public_book = listing["publications"][0] assert public_book["id"] == "publication-one" assert public_book["can_continue"] is False assert public_book["can_write"] is False assert "session_id" not in public_book assert "thread_id" not in public_book assert "content_markdown" not in public_book detail = await llmwiki_router.get_knowledge_research_publication( request, "kb-one", "publication-one", ) assert detail["content_markdown"] == "# 初稿\n\n结论。" assert "session_id" not in detail assert "thread_id" not in detail assert "draft_markdown" not in detail with pytest.raises(HTTPException) as private_chat: await llmwiki_router.get_knowledge_research_session( request, "kb-one", "session-one", ) assert private_chat.value.status_code == 404 @pytest.mark.asyncio @pytest.mark.parametrize(("upstream_status", "expected_status"), [(400, 422), (503, 502)]) async def test_publication_translates_upstream_directory_errors( monkeypatch, upstream_status: int, expected_status: int, ) -> None: mapping = { "id": "kb-one", "weknora_id": "remote-one", "owner_user_id": "alice", "publication_status": "private", } sessions = KnowledgeResearchStore(None) created = await sessions.create( knowledge_base_id="kb-one", user_id="alice", title="问题", session_id="session-one", publication_id="publication-one", thread_id="thread-one", ) await sessions.update_owned( created["id"], "alice", draft_markdown="# 正文", ) request = _request(knowledge_research_store=sessions) async def authorized(_request, mapping_id: str, *, write: bool): assert mapping_id == "kb-one" assert write is True return mapping, "alice", False async def failed_directory(*_args, **_kwargs): raise WeKnoraError("上游目录失败", status_code=upstream_status) sync = Mock() monkeypatch.setattr(llmwiki_router, "_authorized_mapping", authorized) monkeypatch.setattr(llmwiki_router, "_client_or_503", _WikiClient) monkeypatch.setattr(llmwiki_router, "ensure_directory_paths", failed_directory) monkeypatch.setattr(llmwiki_router, "_schedule_local_wiki_sync", sync) with pytest.raises(HTTPException) as error: await llmwiki_router.save_knowledge_research_publication( request, "kb-one", "session-one", llmwiki_router.KnowledgeResearchPublicationSave(title="报告"), ) assert error.value.status_code == expected_status assert "上游目录失败" in str(error.value.detail) assert (await sessions.get("session-one"))["saved_hash"] is None sync.assert_not_called()