From 49dcd49f9163c067cf0fa5431b35063bd8abcd9c Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Sun, 4 Oct 2026 21:21:39 -0700 Subject: [PATCH] fix(agents): keep malformed todo statuses recoverable (#6299) * fix(agents): keep malformed todo statuses recoverable * docs(agents): fit todo guidance within inherited size limit * docs(agents): clarify todo attribution guidance --------- Co-authored-by: Willem Jiang --- README.md | 2 + .../deerflow/agents/middlewares/AGENTS.md | 1 + .../middlewares/token_usage_middleware.py | 2 +- backend/tests/test_token_usage_middleware.py | 58 ++++++++++++++++++- 4 files changed, 61 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index b9c0f90697..16f2a47633 100644 --- a/README.md +++ b/README.md @@ -688,6 +688,8 @@ Text results retain their original representation, including large numeric IDs a Copy actions copy only the displayed preview. This is a frontend view of data already received by the browser, without an additional secret-redaction layer. +In plan mode, malformed TODO statuses return normal tool-validation errors without aborting token attribution, so the agent can correct the call. + Tool-produced paths and URLs can be retained as short artifact handles across context compaction (`tool_artifacts` in `config.yaml`). Handles distinguish separate tool-result occurrences, even when a provider reuses call IDs. Detected file URLs preserve their query strings and fragments. When PII redaction is enabled, model-visible artifact labels follow that policy; internal references stay intact for tool argument resolution. The configured registry limit retains the newest artifacts, while checkpointed processing identities prevent evicted results from being recaptured after restart. Resolution runs before authorization and write-safety checks; unknown or expired handles return an error without executing the tool. Small unknown structured results may be retained as complete JSON up to 4096 UTF-8 bytes; empty or oversized payloads are skipped. Handles are agent-local: task arguments resolve parent handles to concrete references, and delegated reports must return concrete references rather than child-local handles. A truncated model projection reports how many handles are omitted. DeerFlow supports configurable MCP servers and skills to extend its capabilities. diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 6fa8619705..c99ed00cb7 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -100,6 +100,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../ 21. **SummarizationMiddleware** - *(optional)* Compacts near token limits. Fraction trigger/keep profiles belong to the active run model; an explicit summary model only generates. Memory flush/manual compaction follow runtime/checkpoint agent policy. Preserve the latest user request and DynamicContext reminders; a backward cutoff may retain old AI/tool turns or make first-turn compaction a no-op. 22. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with `write_todos`. Skips reminders on `model_length_termination` so capped turns end cleanly. A raising model call requeues its reminder, uncounted, for the retry. 23. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable. + TODO attribution skips non-string statuses before membership; preserve args/TODO state for normal tool validation. 24. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint. 25. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket 26. **ViewImageMiddleware** - *(vision models only)* Rebuilds hidden base64 image context in `ModelRequest.messages`, never in new checkpoint state; state retains lightweight `viewed_images` metadata. `_image_injection_plan` first sweeps stale server-marked context (including payloads stranded by the old `before_model`/`after_model` lifecycle), then shares one read-eligibility predicate between sync injection and async authorization. The sweep requires both the reserved ID prefix and server-owned marker, preserving user messages and unmarked legacy payloads. Before reading images, both paths recheck `sandbox:execute`, including restored views after role changes; denial skips I/O and removes marked stale context. Host-copy reads additionally require the saved `actual_path` to match the current user/thread's canonical virtual image path, including for old checkpoints and replacement sandboxes; a missing or mismatched thread identity skips that host read. Only the async wrapper passes `authorization_checked=True` after its gate, avoiding a duplicate synchronous provider call in its worker. diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py index 629387cffe..b72204db8f 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py @@ -42,7 +42,7 @@ def _normalize_todos(value: Any) -> list[Todo]: if content is not None: todo["content"] = content - if status in {"pending", "in_progress", "completed"}: + if isinstance(status, str) and status in {"pending", "in_progress", "completed"}: todo["status"] = status normalized.append(todo) diff --git a/backend/tests/test_token_usage_middleware.py b/backend/tests/test_token_usage_middleware.py index 545c57c5f7..dd6ef62178 100644 --- a/backend/tests/test_token_usage_middleware.py +++ b/backend/tests/test_token_usage_middleware.py @@ -3,13 +3,18 @@ import logging from unittest.mock import MagicMock -from langchain_core.messages import AIMessage, ToolMessage +import pytest +from langchain.agents import create_agent +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langgraph.graph.message import add_messages +from deerflow.agents.middlewares.todo_middleware import TodoMiddleware from deerflow.agents.middlewares.token_usage_middleware import ( TOKEN_USAGE_ATTRIBUTION_KEY, TokenUsageMiddleware, _build_todo_actions, + _normalize_todos, ) from deerflow.subagents.status_contract import SUBAGENT_TOKEN_USAGE_KEY @@ -367,3 +372,54 @@ class TestBuildTodoActions: ] actions = _build_todo_actions(previous, next_todos) assert any(a.get("kind") == "todo_remove" and a.get("content") == "B" for a in actions), f"Expected todo_remove for B but got: {actions}" + + +@pytest.mark.parametrize( + "status, expected_status", + [("pending", "pending"), ("in_progress", "in_progress"), ("completed", "completed"), ("unknown", None), (None, None), (True, None), (1, None), ([], None), ({}, None)], +) +def test_normalize_todos_only_keeps_valid_string_statuses(status, expected_status): + todo = {"content": "Example task", "status": status} + expected = {"content": "Example task"} + if expected_status is not None: + expected["status"] = expected_status + + assert _normalize_todos([todo]) == [expected] + + +class _ToolCallingModel(FakeMessagesListChatModel): + def bind_tools(self, tools, *, tool_choice=None, **kwargs): + return self + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("status", [[], {}, "unknown", "completed"], ids=["list", "dict", "unknown-string", "valid"]) +async def test_todo_validation_remains_recoverable_with_token_attribution(status, use_async): + """Attribution runs before ToolNode validates model-generated TODO arguments.""" + todo = {"content": "Example task", "status": status} + model = _ToolCallingModel( + responses=[ + AIMessage(content="", tool_calls=[{"name": "write_todos", "id": "todo-call", "args": {"todos": [todo]}}]), + AIMessage(content="Final answer"), + ] + ) + graph = create_agent(model=model, tools=[], middleware=[TodoMiddleware(), TokenUsageMiddleware()]) + state = {"messages": [HumanMessage(content="Plan the example task")]} + + result = await graph.ainvoke(state) if use_async else graph.invoke(state) + + assert result["messages"][-1].content == "Final answer" + tool_messages = [message for message in result["messages"] if isinstance(message, ToolMessage)] + assert len(tool_messages) == 1 + assert tool_messages[0].tool_call_id == "todo-call" + if status == "completed": + assert tool_messages[0].status == "success" + assert result["todos"] == [todo] + else: + assert tool_messages[0].status == "error" + assert "todos.0.status" in tool_messages[0].content + assert not result.get("todos") + dispatch = next(message for message in result["messages"] if isinstance(message, AIMessage) and message.tool_calls) + assert dispatch.tool_calls[0]["args"] == {"todos": [todo]} + assert dispatch.additional_kwargs[TOKEN_USAGE_ATTRIBUTION_KEY]["tool_call_ids"] == ["todo-call"]