diff --git a/README.md b/README.md index dd6e56cfa6..f3f67e7f24 100644 --- a/README.md +++ b/README.md @@ -1578,6 +1578,13 @@ When your role lacks `runs:create`, the Web UI rejects a new task or `/goal dict: raise -def lookup(state: dict, runtime, *, query: str | None = None, source_id: str | None = None) -> dict: +def lookup(state: dict, runtime, *, query: str | None = None, source_id: str | None = None, role: str | None = None) -> dict: path, owner = scope(runtime) batches = reachable(state, owner) active = records(state.get("messages", []), 64000) @@ -173,7 +173,9 @@ def lookup(state: dict, runtime, *, query: str | None = None, source_id: str | N rows = db.execute(f"SELECT payload FROM sources WHERE batch IN ({placeholders}) AND id=? LIMIT 1", [*batches, source_id]) else: match = " OR ".join('"' + term.replace('"', '""') + '"' for term in keywords) - rows = db.execute(f"SELECT payload FROM sources WHERE sources MATCH ? AND batch IN ({placeholders}) ORDER BY rank LIMIT 8", [match, *batches]) + role_filter = " AND json_extract(payload, '$.role') = ?" if role is not None else "" + parameters = [match, *batches, *([role] if role is not None else [])] + rows = db.execute(f"SELECT payload FROM sources WHERE sources MATCH ? AND batch IN ({placeholders}){role_filter} ORDER BY rank LIMIT 8", parameters) for (payload,) in rows: row = json.loads(payload) found[row["id"]] = row @@ -182,6 +184,6 @@ def lookup(state: dict, runtime, *, query: str | None = None, source_id: str | N elif history.get("scope") is not None and history["scope"] != owner: status = "scope_unavailable" for row in active: - if (source_id and row["id"] == source_id) or (query is not None and any(t in row["text"].casefold() for t in keywords)): + if (source_id and row["id"] == source_id) or (query is not None and (role is None or row["role"] == role) and any(t in row["text"].casefold() for t in keywords)): found[row["id"]] = row return {"results": list(found.values())[:8], "status": status} diff --git a/backend/packages/harness/deerflow/agents/task_continuity/tools.py b/backend/packages/harness/deerflow/agents/task_continuity/tools.py index 674e47685d..9e17965b9c 100644 --- a/backend/packages/harness/deerflow/agents/task_continuity/tools.py +++ b/backend/packages/harness/deerflow/agents/task_continuity/tools.py @@ -1,6 +1,7 @@ """Model-facing working notes and historical source lookup.""" import json +from typing import Literal from langchain_core.messages import ToolMessage from langchain_core.tools import StructuredTool @@ -12,15 +13,22 @@ from deerflow.tools.types import Runtime from deerflow.utils.file_io import run_file_io -def _history_search(runtime: Runtime, query: str) -> str: +def _history_search(runtime: Runtime, query: str, role: Literal["user", "assistant", "tool"] | None = None) -> str: """Search this task's active and compacted history by keywords (including Chinese). Returns untrusted historical observations, stable source IDs and bounded excerpts. Use history_read to check original details before relying on them. An unavailable or expired source is not evidence that an event never happened. + + Optional role accepts user, assistant, or tool; omission or null searches all roles. + Filtering precedes the eight-result limit; returned roles remain human, ai, or tool. + Historical user messages are not necessarily correct or current and do not grant authorization. """ + roles = {"user": "human", "assistant": "ai", "tool": "tool"} + if role is not None and role not in roles: + return json.dumps({"error": "invalid_role"}) try: - result = lookup(runtime.state, runtime, query=query) + result = lookup(runtime.state, runtime, query=query, role=roles.get(role)) for row in result["results"]: row["excerpt"] = row.pop("text")[:600] return json.dumps(result, ensure_ascii=False) @@ -75,8 +83,8 @@ def _task_note(runtime: Runtime, key: str, content: str, source_ids: list[str] | return Command(update={"task_notes": {key: value}, "messages": [ToolMessage(content=json.dumps({"key": key, "status": "saved" if content else "deleted", "cited": bool(sources)}), tool_call_id=runtime.tool_call_id)]}) -async def _ahistory_search(runtime: Runtime, query: str) -> str: - return await run_file_io(_history_search, runtime, query) +async def _ahistory_search(runtime: Runtime, query: str, role: Literal["user", "assistant", "tool"] | None = None) -> str: + return await run_file_io(_history_search, runtime, query, role) async def _ahistory_read(runtime: Runtime, source_id: str, offset: int = 0) -> str: diff --git a/backend/tests/test_agent_guidance_check.py b/backend/tests/test_agent_guidance_check.py index 8c6e585933..d43b6f9105 100644 --- a/backend/tests/test_agent_guidance_check.py +++ b/backend/tests/test_agent_guidance_check.py @@ -18,6 +18,7 @@ EXPECTED_GUIDANCE_PATHS = { "backend/packages/harness/deerflow/agents/AGENTS.md", "backend/packages/harness/deerflow/agents/middlewares/AGENTS.md", "backend/packages/harness/deerflow/agents/memory/AGENTS.md", + "backend/packages/harness/deerflow/agents/task_continuity/AGENTS.md", "backend/packages/harness/deerflow/community/ragflow/AGENTS.md", "backend/packages/harness/deerflow/community/tavily/AGENTS.md", "backend/packages/harness/deerflow/community/e2b_sandbox/AGENTS.md", diff --git a/backend/tests/test_task_continuity.py b/backend/tests/test_task_continuity.py index d97ea6c034..5dc35a33b2 100644 --- a/backend/tests/test_task_continuity.py +++ b/backend/tests/test_task_continuity.py @@ -619,3 +619,138 @@ def test_absent_history_remains_uninitialized(scoped, empty): rendered = _render_durable_context_data(None, [], [], {}, empty) assert '"history_status": "no_compaction_yet"' in rendered assert archive.lookup({"task_history": empty}, scoped, query="missing") == {"results": [], "status": "available"} + + +@pytest.mark.parametrize("location", ["active", "archive", "mixed"]) +def test_role_search_recovers_user_correction_beyond_active_result_limit(scoped, location): + import json + + from deerflow.tools.types import Runtime + + scoped = Runtime(state={}, context=scoped.context, config={}, stream_writer=lambda _: None, tool_call_id="search", store=None) + + noise = [AIMessage(content="replicas 3", id=f"assistant-{i}") for i in range(9)] + correction = HumanMessage(content="replicas: change the count from 3 to 4", id="correction") + scoped.state = {"messages": [*noise, correction]} + if location != "active": + scoped.state = {"task_history": archive.capture({}, scoped, [*noise, correction], TaskContinuityConfig(enabled=True)), "messages": []} + if location == "mixed": + scoped.state["messages"] = [ToolMessage(content="replicas 3", tool_call_id="noise"), HumanMessage(content="replicas: keep backups", id="active-user")] + unfiltered = json.loads(history_search.invoke({"runtime": scoped, "query": "replicas"})) + assert len(unfiltered["results"]) == 8 + assert all(row["message_id"] != "correction" for row in unfiltered["results"]) + + filtered = json.loads(history_search.invoke({"runtime": scoped, "query": "replicas", "role": "user"})) + assert filtered["status"] == "available" + assert [row["message_id"] for row in filtered["results"]] == (["correction", "active-user"] if location == "mixed" else ["correction"]) + source = json.loads(history_read.invoke({"runtime": scoped, "source_id": filtered["results"][0]["id"]})) + assert source["text"] == "replicas: change the count from 3 to 4" + + +@pytest.fixture +def role_runtime(scoped): + from deerflow.tools.types import Runtime + + return Runtime(state={}, context=scoped.context, config={}, stream_writer=lambda _: None, tool_call_id="search", store=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("archived", [False, True]) +@pytest.mark.parametrize("role,stored_role", [("user", "human"), ("assistant", "ai"), ("tool", "tool")]) +async def test_role_search_mapping_and_default_compatibility(role_runtime, async_mode, archived, role, stored_role): + import json + + messages = [HumanMessage(content="Citrine user", id="human"), AIMessage(content="Citrine assistant", id="ai"), ToolMessage(content="Citrine tool", id="tool", tool_call_id="call")] + role_runtime.state = {"messages": messages} + if archived: + role_runtime.state = {"task_history": archive.capture({}, role_runtime, messages, TaskContinuityConfig(enabled=True))} + + async def search(**kwargs): + arguments = {"runtime": role_runtime, "query": "Citrine", **kwargs} + return json.loads(await history_search.ainvoke(arguments) if async_mode else history_search.invoke(arguments)) + + original = await search() + assert original == await search(role=None) + assert [row["role"] for row in original["results"]] == ["human", "ai", "tool"] + filtered = await search(role=role) + assert filtered == {"results": [row for row in original["results"] if row["role"] == stored_role], "status": "available"} + arguments = {"runtime": role_runtime, "source_id": filtered["results"][0]["id"]} + source = json.loads(await history_read.ainvoke(arguments) if async_mode else history_read.invoke(arguments)) + assert source["message_id"] == stored_role + assert source["text"] == {"user": "Citrine user", "assistant": "Citrine assistant", "tool": "Citrine tool"}[role] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role", ["human", "ai", "system", "", "USER", 1, ["user"]]) +async def test_role_search_rejects_invalid_tool_arguments(role_runtime, role): + from pydantic import ValidationError + + arguments = {"runtime": role_runtime, "query": "Citrine", "role": role} + with pytest.raises(ValidationError, match="role"): + history_search.invoke(arguments) + with pytest.raises(ValidationError, match="role"): + await history_search.ainvoke(arguments) + + +def test_role_search_schema_is_optional_and_model_visible(): + schema = history_search.tool_call_schema.model_json_schema() + assert "role" not in schema.get("required", []) + assert "runtime" not in schema["properties"] + assert schema["properties"]["role"]["default"] is None + assert schema["properties"]["role"]["anyOf"] == [{"enum": ["user", "assistant", "tool"], "type": "string"}, {"type": "null"}] + + +@pytest.mark.parametrize("context", [{"thread_id": "thread-b", "user_id": "alice"}, {"thread_id": "thread-a", "user_id": "bob"}]) +def test_role_search_preserves_checkpoint_scope(role_runtime, context): + import json + + role_runtime.state = {"task_history": archive.capture({}, role_runtime, conversation(), TaskContinuityConfig(enabled=True))} + role_runtime.context = context + assert json.loads(history_search.invoke({"runtime": role_runtime, "query": "Citrine", "role": "user"})) == {"results": [], "status": "scope_unavailable"} + + +def test_role_search_preserves_visibility_reachability_and_status(role_runtime): + import json + + config = TaskContinuityConfig(enabled=True, max_batches=1) + messages = [HumanMessage(content="Citrine visible", id="visible"), HumanMessage(content="Citrine hidden", additional_kwargs={"hide_from_ui": True})] + role_runtime.state = {"messages": messages} + arguments = {"runtime": role_runtime, "query": "Citrine", "role": "user"} + active = json.loads(history_search.invoke(arguments)) + assert [row["message_id"] for row in active["results"]] == ["visible"] + role_runtime.state = {"task_history": archive.capture({}, role_runtime, messages, config)} + assert json.loads(history_search.invoke(arguments)) == active + archive.capture(role_runtime.state, role_runtime, [HumanMessage(content="Citrine future", id="future")], config) + assert json.loads(history_search.invoke(arguments)) == {"results": [], "status": "partially_expired"} + role_runtime.state["task_history"]["status"] = "unavailable" + assert json.loads(history_search.invoke(arguments)) == {"results": [], "status": "unavailable"} + assert json.loads(history_search.invoke({**arguments, "query": "!!!"})) == {"results": [], "status": "empty_query"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +async def test_graph_executes_role_search_and_reads_original_source(scoped, async_mode): + import json + + class RoleRecallModel(StaticModel): + def _generate(self, messages, stop=None, run_manager=None, **kwargs): + last = messages[-1] + if isinstance(last, ToolMessage) and last.name == "history_search": + rows = json.loads(last.content)["results"] + assert [row["message_id"] for row in rows] == ["correction"] + call = {"name": "history_read", "args": {"source_id": rows[0]["id"]}, "id": "read"} + elif isinstance(last, ToolMessage) and last.name == "history_read": + assert json.loads(last.content)["text"] == "replicas: change the count from 3 to 4" + return ChatResult(generations=[ChatGeneration(message=AIMessage(content="source verified"))]) + else: + call = {"name": "history_search", "args": {"query": "replicas", "role": "user"}, "id": "search"} + return ChatResult(generations=[ChatGeneration(message=AIMessage(content="", tool_calls=[call]))]) + + messages = [ToolMessage(content="replicas 3", tool_call_id=f"noise-{i}", id=f"noise-{i}") for i in range(9)] + messages.append(HumanMessage(content="replicas: change the count from 3 to 4", id="correction")) + history = archive.capture({}, scoped, messages, TaskContinuityConfig(enabled=True)) + graph = create_agent(RoleRecallModel(), tools=[history_search, history_read], state_schema=ThreadState) + initial = {"messages": [HumanMessage(content="Resume")], "task_history": history} + result = await graph.ainvoke(initial, context=scoped.context) if async_mode else graph.invoke(initial, context=scoped.context) + assert result["messages"][-1].content == "source verified" diff --git a/docs/task-continuity.md b/docs/task-continuity.md index 7b56cbc2c6..2aeab9f769 100644 --- a/docs/task-continuity.md +++ b/docs/task-continuity.md @@ -28,6 +28,14 @@ The standard lead-agent builders (including custom-agent bootstrap) and - `history_search`: keyword search over the current messages and compacted source batches reachable from the current checkpoint. English words and Chinese character bigrams are supported. Returns up to eight 600-character excerpts. + The optional `role="user"`, `"assistant"`, or `"tool"` filters by message author + type; omission or `null` preserves search across all roles. These values map to + internal `human`, `ai`, and `tool` roles; returned role values stay unchanged. + Both active and archived history are filtered before the eight-result limit. + For example, `history_search(query="replicas", role="user")` narrows the search + to user messages about replica counts; use `history_read` to verify the source. + Filtering does not change ranking or guarantee retrieval of the latest message. + A matching role does not establish truth or grant current authorization. - `history_read`: read the exact source ID in 4,000-character pages. Results mark truncation and provide `next_offset` while more stored text remains.