fix(todos): keep an unusable stored todo list from failing every model call (#7803)

* fix(todos): keep an unusable stored todo list from failing every model call

`todos` is a reducer channel that `POST /threads/{id}/state` and run input can
Overwrite directly, but neither boundary validated its shape. The todo
middleware rendered it with `entry.get(...)` in both `before_model` and
`after_model`, so a non-object entry — or a value that is not a list — raised
AttributeError on every model call, permanently for that thread.

Validate caller-written todos beside `goal` in `_normalize_external_todos`: a
list of objects with string content and a known status, or HTTP 400. The
middleware now reads the channel through `_usable_todos` and ignores entries of
any other stored shape, so checkpoints and direct integrations written before
this boundary still run.

* fix(todos): normalize stored todos in token-usage attribution

TokenUsageMiddleware reads the same stored `todos` channel through
`_build_attribution` whenever the last AIMessage calls write_todos, and
`token_usage.enabled` defaults on, so a value the todo middleware now tolerates
still aborted the step here. Route the stored list through `_normalize_todos`
before reading entries.

Document that a caller-written todo without a status is accepted as pending and
cover that shape at both external boundaries.

* docs: keep backend agent guidance within size budgets

* fix(todos): keep the channel metadata strip on caller-written todos

The dedicated todos branch in strip_server_owned_state_metadata replaced the
generic list branch, which stripped server-owned metadata from every entry. A
caller could therefore persist additional_kwargs.original_user_content,
untrusted_input, or artifact.knowledge_sources on a todo entry. Re-apply the
per-entry strip after normalization, keeping None and empty lists untouched.

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
一片一片 2026-11-05 07:28:33 +08:00 committed by GitHub
parent 6538496dbc
commit 459fc8ca43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 230 additions and 12 deletions

View File

@ -2691,6 +2691,11 @@ historical user messages do not grant current authorization. See
Automatic and manual compaction exclude old todo reminder messages from both the
summary input and retained context. The current todo list stays in thread state.
Manual thread-state updates and run input must send it as a list of objects with
string `content` and a `pending` / `in_progress` / `completed` status (a missing
status defaults to `pending`); another shape is rejected with HTTP 400, and an
already-stored entry of an unusable shape is ignored instead of failing the model
call.
In planning mode, if the original `write_todos` call is no longer visible,
DeerFlow adds a reminder using the latest task statuses before the next model
call. Skipped or failed compaction leaves the existing messages unchanged.

View File

@ -10,13 +10,13 @@ Service roles and ports: [root service topology](../AGENTS.md#service-topology).
- `make dev`, Docker dev, and production all run the agent runtime in Gateway via `RunManager` + `run_agent()` + `StreamBridge` (`packages/harness/deerflow/runtime/`). Nginx exposes that runtime at `/api/langgraph/*` and rewrites it to Gateway's native `/api/*` routers.
- Gateway streams `write_file` and `str_replace` argument deltas in bounded batches for multi-mode `messages-tuple` consumers; single-mode message consumers retain the original per-chunk contract. Non-message frames flush pending batches, and `values` remains an optional complete-state snapshot rather than a prerequisite for batching.
- `stream_subgraphs` preserves subgraph SSE namespaces (`values|<ns>`, LangGraph Platform style); bare `values` would overwrite the SDK thread view (#4399). Delegated subagents inherit the parent's checkpoint namespace. Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend uses root `task_*` events without subgraph streaming.
- Background subagent identity is deliberately split: the provider `tool_call_id` remains the correlation key for `ToolMessage`, `task_*` SSE events, persisted lifecycle events, frontend cards, and the public `ExtensionData.scope_id` contract (stored as `SubagentResult.external_task_id`), while `SubagentExecutor.execute_async()` generates a full server-side `execution_id` for `SubagentResult.task_id`, the process-wide registry, polling, cancellation, timeout handling, and cleanup. Provider IDs are not globally unique across parent runs, so they must never become registry ownership keys; scheduler closures retain their own `SubagentResult` rather than resolving ownership again through the mutable registry. Terminal subagent token usage travels in the current run's `ToolMessage.additional_kwargs` and is attributed from message state, never through a process-global provider-ID cache.
- Subagent IDs: provider `tool_call_id` correlates `ToolMessage`, `task_*` SSE/persisted events, frontend cards and `ExtensionData.scope_id` (`SubagentResult.external_task_id`). Unique server `execution_id` owns `SubagentResult.task_id`, registry, polling, cancellation, timeout and cleanup. Provider IDs can collide across parent runs: scheduler closures retain their own result, never recover ownership through the mutable registry. Attribute terminal usage from the current run's `ToolMessage.additional_kwargs` and message state, never a global provider-ID cache.
- Scheduled runs use Gateway's normal run path. `launch_scheduled_thread_run` reads `get_app_config()` at dispatch: `scheduler.recursion_limit` defaults to 1000 (web UI), capped by `max_recursion_limit`; YAML changes apply next run without restart. IM enqueue checks recipients before locking preferences in the outcome transaction.
- Run-history `status` filters are occurrence states, not task states. `ScheduledTaskRunStatus` in `persistence/scheduled_tasks/model.py` is the shared API/repository vocabulary and must match the active and terminal occurrence-status sets. Keep owner lookup before reading history, and apply SQL task/status predicates before pagination; omitted status preserves the existing response.
- The background scheduler is single-instance by default. `scheduler.multi_instance=true` opts into lease-aware recovery across Gateway instances and requires shared Postgres, `run_ownership.heartbeat_enabled=true`, and `run_events.backend=db`; otherwise startup rejects the configuration. Live scheduled runs are preserved when a peer starts; expired launch claims return to the durable queue, expired run leases are atomically taken over, stale launch writes are fenced by lease ownership, and the Postgres advisory-locked budget makes `max_concurrent_runs` a shared global cap for `launching`/`running` rows.
- The multi-process startup gate also fires on `deployment.multi_instance: true`; contract in `docs/CONFIGURATION.md`.
- Long-running MCP work uses a separate durable task runtime (`McpTaskService` + `mcp_tasks`, lease-based recovery) rather than keeping remote task IDs or status polling inside the Agent loop; only submit remains Agent-visible, the database is the source of truth, and `ThreadState` receives only a bounded current-thread projection. Full contract — including notification retries, dead-lettering, and the cancel endpoint's worker-stopped 503 (leases, cancellation fencing, delivery idempotency, management-tool exposure): [packages/harness/deerflow/mcp/AGENTS.md](packages/harness/deerflow/mcp/AGENTS.md).
- Scheduled-task dispatch permits one active occurrence per task via `uq_scheduled_task_run_active` (`task_id WHERE status IN ('queued','launching','running')`). Durable `queued` rows survive restarts; only lease-fenced `launching` may call Gateway launch; `running` references the durable run. Stable admission idempotency keys reuse that run after recovery. Reused-thread `ConflictError` returns `launching` to `queued`; other launch errors become `failed`. Atomic queue claims enforce `max_concurrent_runs` and the per-owner cap, excluding waiting rows; the budget count and its UPDATE are separate statements, so writers must serialize before the count (Postgres advisory lock; SQLite `BEGIN IMMEDIATE`, whose deferred transaction otherwise reserves the writer only at the UPDATE) or claims on distinct rows overshoot the cap. Repeated triggers coalesce; same-thread FIFO blocks behind older active rows. Queue admission, PATCH/resume, pause and delete lock the parent before the occurrence, freezing active task definitions. Pause/delete atomically cancel `queued` work but reject `launching`/`running`; PATCH/resume reject all active states. Only queued conflicts offer pause cancellation. Manual triggers may queue/run while paused. Manual-trial run origins reuse the admitted `scheduled_for`, so an idempotent retry submits the same request instead of colliding with the first attempt.
- Scheduled-task dispatch: `uq_scheduled_task_run_active` permits one `queued`/`launching`/`running` occurrence per task. Queues survive restarts; only lease-fenced `launching` calls Gateway; `running` references the durable run. Stable admission keys reuse it after recovery. Reused-thread `ConflictError` requeues; other launch errors fail. Claims exclude waiting rows and enforce `max_concurrent_runs`/per-owner caps: serialize count+UPDATE with the Postgres advisory lock or SQLite `BEGIN IMMEDIATE`. Coalesce triggers; preserve same-thread FIFO. Admission, PATCH/resume, pause/delete lock parent before occurrence, freezing active definitions. Pause/delete atomically cancel queued work and reject launching/running; PATCH/resume reject all active states. Only queued conflicts offer pause cancellation. Manual triggers run/queue while paused; trial origins reuse admitted `scheduled_for` for idempotent retries.
- Scheduled-task lifecycle, recovery and queue timeout: [persistence guide](packages/harness/deerflow/persistence/AGENTS.md#scheduled-task-lifecycle).
- `POST /api/scheduled-tasks/preview-cron` requires authenticated `threads:read`. Bounded cron previews call the shared scheduler calculator in `asyncio.to_thread`, preserving its DST semantics. Capture the optional aware reference once; return UTC and offset-bearing local occurrences without acquiring task/thread/run stores or dispatching work. This advisory API does not reserve execution.
- Gateway MCP/skill updates write `extensions_config.json` at runtime. Production Compose mounts it read-write (`config.yaml` stays `:ro`); Helm seeds a writable home-volume copy. Hold both `extensions_config_write_lock` and the sidecar advisory `extensions_config_file_lock` during read-modify-write to prevent lost updates across workers. `atomic_write_extensions_config` uses temp-file replacement; only mount-point `EBUSY` permits a non-atomic in-place overwrite (crashes can truncate). Warn about this fallback once per target; propagate other errors. `get_extensions_config()` revalidates reads by file signature; see [config caching](packages/harness/deerflow/config/AGENTS.md). Tests: `tests/test_compose_extensions_config_writable.py`, `tests/test_extensions_config_atomic_write.py`, `tests/test_helm_extensions_config_writable.py`.
@ -192,6 +192,10 @@ and an explicit invalid/null top-level value clears a legacy preference. Treat
`disable_clarification` like `non_interactive`. Before run/state writes,
`_normalize_input_messages` rejects canonical external system/developer roles;
only `AUTH_SOURCE_INTERNAL` run input may retain them.
`_normalize_external_todos` rejects invalid caller todos with HTTP 400: entries
must be objects with string content and a known status (default `pending`). Todo
and token-usage readers normalize malformed stored/direct-integration values so
model calls remain usable.
`contracts/ui_locales.v1.json` pins the UI/Gateway locale allowlists in tests;
DeerMem's portable map stays separate. Scheduled/MCP notification runs omit

View File

@ -642,6 +642,36 @@ def _normalize_input_messages(
#: record that clearing a satisfied goal writes; a caller copy would forge one.
SERVER_OWNED_STATE_CHANNELS = frozenset({"sandbox", "thread_data", "viewed_images", "goal_outcome", "applied_run_inputs"})
#: The ``write_todos`` status contract the todo middleware renders.
_TODO_STATUSES = frozenset({"pending", "in_progress", "completed"})
def _normalize_external_todos(value: Any, *, location: str) -> list[dict[str, Any]] | None:
"""Reject todo lists the middleware cannot render.
``todos`` is a reducer channel a caller may Overwrite directly, and the todo
middleware reads it on every model call. A value that is not a list of
objects with string content would fail those calls on every turn instead of
being refused here, so validate the shape at the boundary the same way
``goal`` is.
"""
if value is None:
return None
if not isinstance(value, list):
raise HTTPException(status_code=400, detail=f"{location} must be a list of todo objects")
normalized: list[dict[str, Any]] = []
for index, entry in enumerate(value):
if not isinstance(entry, Mapping):
raise HTTPException(status_code=400, detail=f"{location}[{index}] must be an object")
content = entry.get("content")
if not isinstance(content, str):
raise HTTPException(status_code=400, detail=f"{location}[{index}].content must be a string")
status = entry.get("status", "pending")
if not isinstance(status, str) or status not in _TODO_STATUSES:
raise HTTPException(status_code=400, detail=f"{location}[{index}].status must be pending, in_progress, or completed")
normalized.append({**entry, "content": content, "status": status})
return normalized
def _normalize_external_goal(value: Any, *, location: str) -> dict[str, Any] | None:
"""Require usable goal text without rebuilding caller-supplied goal state."""
@ -687,6 +717,12 @@ def strip_server_owned_state_metadata(values: Mapping[str, Any]) -> dict[str, An
for channel, value in values.items():
if channel == "goal":
stripped[channel] = _strip_external_metadata_from_message_like(_normalize_external_goal(value, location="values.goal"))
elif channel == "todos":
normalized_todos = _normalize_external_todos(value, location="values.todos")
# Todos keep their extra fields, but they are still caller-supplied
# channel values: a dedicated validator must not bypass the
# per-entry metadata strip the generic list branch applies.
stripped[channel] = None if normalized_todos is None else [_strip_external_metadata_from_message_like(item) for item in normalized_todos]
elif channel == "messages" and value is not None:
stripped[channel] = _normalize_input_messages(value, location="values.messages")
elif channel == "delegations" and isinstance(value, list):
@ -762,6 +798,8 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool
# merge_goal ignores None, so only a goal value replaces the goal.
if result.get("goal") is not None:
result = {**result, "goal": _normalize_external_goal(result["goal"], location="input.goal"), "goal_outcome": None}
if "todos" in result:
result = {**result, "todos": _normalize_external_todos(result["todos"], location="input.todos")}
return result

View File

@ -153,9 +153,9 @@ Before changing a later authorization phase, read the [authorization RFC](../../
18. **SkillActivationMiddleware** - Detects strict `/skill-name task` syntax on the latest real user message, resolves only enabled and runtime-allowed skills, injects the `SKILL.md` body as hidden current-turn context, and records a `middleware:skill_activation` audit event; a retried activation call replays the recorded reminder without re-auditing. Skill `allowed-tools` is compatibility metadata only: activation never filters, grants, or denies tools — Agent assembly/execution authorization owns the permission ceiling.
19. **DurableContextMiddleware** - Before compaction, captures task dispatch/result summaries in `ThreadState.delegations` and loaded skill name/path/description (never bodies) in `ThreadState.skill_context`; projects them per request. Static authority uses `SystemMessage`; untrusted `summary_text`, delegation results and skill descriptions use a separate hidden `HumanMessage` data block, never persisted as messages or promoted to system instructions. Subagents attach it immediately before summarization: summary precedes preserved AI/tool history for strict providers. See [tool artifacts](TOOL_ARTIFACTS.md) for projection, retry limits and delegation scope.
20. **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.
21. **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.
21. **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. Both hooks read stored todos through `_usable_todos`: thread-state updates and direct integrations can Overwrite that channel, and a non-mapping entry would otherwise raise on every model call.
22. **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.
TODO attribution skips non-string statuses before membership; preserve args/TODO state for normal tool validation. Attribution normalizes the stored `todos` channel with `_normalize_todos` before reading entries, the same tolerance the todo middleware applies.
23. **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.
24. **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
25. **ViewImageMiddleware** - *(vision models only)* Enforces the shared vision cap on remote rereads (see [download contract](../../sandbox/AGENTS.md)). Rebuilds hidden base64 context in `ModelRequest.messages`, keeping lightweight `viewed_images` in checkpoints. Caps each request at four recent images and 28 MiB base64; constructor limits may tighten caps. Skips over-budget reads. Re-views refresh reducer recency. Validated blob refs precede compatibility paths. `_image_injection_plan` removes stale server-marked context, including old `before_model`/`after_model` payloads, then shares eligibility for sync/async injection. Removal requires both the reserved ID prefix and marker, preserving user and unmarked legacy messages. Both paths recheck `sandbox:execute` before I/O. Host copies require `actual_path` to match the current user/thread's canonical virtual path, even for old checkpoints or replacement sandboxes; absent/mismatched identity skips reads. Only the async wrapper passes `authorization_checked=True` after its gate to avoid duplicate checks in its worker.

View File

@ -20,11 +20,10 @@ re-engaging would only re-emit the same oversized tool call into the same cap.
from __future__ import annotations
import threading
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Iterable, Mapping
from typing import Any, override
from langchain.agents.middleware import TodoListMiddleware
from langchain.agents.middleware.todo import Todo
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse, hook_config
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.runtime import Runtime
@ -52,7 +51,22 @@ def _reminder_in_messages(messages: list[Any]) -> bool:
return False
def _format_todos(todos: list[Todo]) -> str:
def _usable_todos(state: ThreadState) -> list[Mapping[str, Any]]:
"""Return the stored todo entries this middleware can render.
``todos`` is a reducer channel that checkpoint writes and direct integrations
can Overwrite with any shape, so a stored value is not guaranteed to be the
list of mappings ``write_todos`` produces. Entries of any other shape are
dropped: one of them would otherwise raise on every model call of the thread,
because both hooks read this channel before the model runs.
"""
todos = state.get("todos")
if not isinstance(todos, list):
return []
return [todo for todo in todos if isinstance(todo, Mapping)]
def _format_todos(todos: Iterable[Mapping[str, Any]]) -> str:
"""Format a list of Todo items into a human-readable string."""
lines: list[str] = []
for todo in todos:
@ -62,7 +76,7 @@ def _format_todos(todos: list[Todo]) -> str:
return "\n".join(lines)
def _format_completion_reminder(todos: list[Todo]) -> str:
def _format_completion_reminder(todos: Iterable[Mapping[str, Any]]) -> str:
"""Format a completion reminder for incomplete todo items."""
incomplete = [t for t in todos if t.get("status") != "completed"]
incomplete_text = "\n".join(f"- [{t.get('status', 'pending')}] {t.get('content', '')}" for t in incomplete)
@ -153,7 +167,7 @@ class TodoMiddleware(TodoListMiddleware):
# write_todos was denied for this build; context-loss detection only
# exists to protect todo state the model cannot update anyway.
return None
todos: list[Todo] = state.get("todos") or [] # type: ignore[assignment]
todos = _usable_todos(state)
if not todos:
return None
@ -357,7 +371,7 @@ class TodoMiddleware(TodoListMiddleware):
return None
# 3. Allow exit when all todos are completed or there are no todos.
todos: list[Todo] = state.get("todos") or [] # type: ignore[assignment]
todos = _usable_todos(state)
if not todos or all(t.get("status") == "completed" for t in todos):
return None

View File

@ -241,10 +241,13 @@ def _subagent_usage_from_tool_message(message: ToolMessage) -> dict[str, int] |
return normalize_token_usage(additional_kwargs.get(SUBAGENT_TOKEN_USAGE_KEY))
def _build_attribution(message: AIMessage, todos: list[Todo]) -> dict[str, Any]:
def _build_attribution(message: AIMessage, todos: Any) -> dict[str, Any]:
tool_calls = getattr(message, "tool_calls", None) or []
actions: list[dict[str, Any]] = []
current_todos = list(todos)
# ``todos`` is the stored channel value: checkpoint writes and direct
# integrations can Overwrite it with any shape, so normalize it before
# reading entries instead of trusting the write_todos shape.
current_todos = _normalize_todos(todos)
for raw_tool_call in tool_calls:
if not isinstance(raw_tool_call, dict):

View File

@ -0,0 +1,82 @@
"""Todo writes are validated before they can break every model call on a thread."""
from __future__ import annotations
import pytest
from fastapi import HTTPException
from app.gateway import services
BOUNDARIES = [services.normalize_input, services.strip_server_owned_state_metadata]
@pytest.mark.parametrize("boundary", BOUNDARIES)
@pytest.mark.parametrize(
"todos",
[
pytest.param("not a list", id="string"),
pytest.param({"content": "x"}, id="mapping"),
pytest.param(["oops"], id="string-entry"),
pytest.param([["nested"]], id="list-entry"),
pytest.param([None], id="null-entry"),
pytest.param([{"content": 5}], id="non-string-content"),
pytest.param([{"status": "pending"}], id="missing-content"),
pytest.param([{"content": "x", "status": "done"}], id="unknown-status"),
pytest.param([{"content": "x", "status": None}], id="null-status"),
pytest.param([{"content": "x", "status": ["pending"]}], id="unhashable-status"),
],
)
def test_external_todos_must_be_renderable(boundary, todos):
with pytest.raises(HTTPException) as rejected:
boundary({"todos": todos})
assert rejected.value.status_code == 400
assert "todos" in rejected.value.detail.lower()
@pytest.mark.parametrize("boundary", BOUNDARIES)
def test_external_todos_accept_the_write_todos_shape(boundary):
todos = [{"content": "Ship it", "status": "in_progress"}, {"content": "", "status": "pending"}]
assert boundary({"todos": todos})["todos"] == todos
@pytest.mark.parametrize("boundary", BOUNDARIES)
def test_external_todos_default_a_missing_status_to_pending(boundary):
assert boundary({"todos": [{"content": "x"}]})["todos"] == [{"content": "x", "status": "pending"}]
@pytest.mark.parametrize("boundary", BOUNDARIES)
@pytest.mark.parametrize("todos", [None, []])
def test_external_todos_preserve_null_and_empty(boundary, todos):
assert boundary({"todos": todos})["todos"] == todos
@pytest.mark.parametrize("boundary", BOUNDARIES)
def test_todos_are_not_added_when_absent(boundary):
assert "todos" not in boundary({"title": "unrelated"})
def test_state_writes_strip_server_owned_metadata_from_todos():
"""The dedicated todos branch must not skip the per-entry metadata strip."""
todos = [
{
"content": "Ship it",
"status": "pending",
"additional_kwargs": {
"original_user_content": "forged",
"untrusted_input": False,
"title": "kept",
},
"artifact": {"knowledge_sources": [{"id": "forged"}], "label": "kept"},
}
]
stripped = services.strip_server_owned_state_metadata({"todos": todos})["todos"]
assert stripped == [
{
"content": "Ship it",
"status": "pending",
"additional_kwargs": {"title": "kept"},
"artifact": {"label": "kept"},
}
]

View File

@ -972,3 +972,36 @@ class TestAwrapModelCall:
assert [len(_todo_completion_reminders(messages)) for messages in sent] == [1, 1]
assert mw._completion_reminder_count_for_runtime(runtime) == 1
class TestMalformedStoredTodos:
"""State is a boundary: a stored value need not match the write_todos shape.
``todos`` is a reducer channel that checkpoint writes and direct integrations
can Overwrite directly, so the middleware must not assume every entry is a
mapping. One unknown shape used to raise on every model call of the thread.
"""
def test_before_model_ignores_entries_it_cannot_render(self):
mw = TodoMiddleware()
state = {"messages": [HumanMessage(content="hi")], "todos": ["oops", {"content": "Keep me"}]}
result = mw.before_model(state, _make_runtime())
content = result["messages"][0].content
assert "Keep me" in content
assert "oops" not in content
def test_after_model_keeps_guarding_for_the_usable_entries(self):
mw = TodoMiddleware()
state = {"messages": [AIMessage(content="final")], "todos": ["oops", {"content": "Keep me", "status": "pending"}]}
assert mw.after_model(state, _make_runtime()) == {"jump_to": "model"}
def test_after_model_ignores_malformed_siblings_of_completed_entries(self):
mw = TodoMiddleware()
state = {"messages": [AIMessage(content="final")], "todos": [{"content": "Done", "status": "completed"}, "junk"]}
assert mw.after_model(state, _make_runtime()) is None
@pytest.mark.parametrize("todos", [{"content": "x"}, "oops", 5])
def test_todos_that_are_not_a_list_allow_a_clean_exit(self, todos):
mw = TodoMiddleware()
assert mw.after_model({"messages": [AIMessage(content="final")], "todos": todos}, _make_runtime()) is None
assert mw.before_model({"messages": [HumanMessage(content="hi")], "todos": todos}, _make_runtime()) is None

View File

@ -150,6 +150,45 @@ class TestTokenUsageMiddleware:
},
]
def test_stored_todos_with_unreadable_entries_still_get_attribution(self):
middleware = TokenUsageMiddleware()
message = AIMessage(
content="",
tool_calls=[
{
"id": "write_todos:1",
"name": "write_todos",
"args": {"todos": [{"content": "Ship it", "status": "in_progress"}]},
}
],
usage_metadata={"input_tokens": 100, "output_tokens": 20, "total_tokens": 120},
)
state = {"messages": [message], "todos": ["oops", None, {"content": "Ship it", "status": "pending"}]}
result = middleware.after_model(state, _make_runtime())
attribution = result["messages"][0].additional_kwargs[TOKEN_USAGE_ATTRIBUTION_KEY]
assert attribution["actions"] == [{"kind": "todo_start", "content": "Ship it", "tool_call_id": "write_todos:1"}]
def test_stored_todos_that_are_not_a_list_still_get_attribution(self):
middleware = TokenUsageMiddleware()
message = AIMessage(
content="",
tool_calls=[
{
"id": "write_todos:1",
"name": "write_todos",
"args": {"todos": [{"content": "Ship it", "status": "in_progress"}]},
}
],
)
state = {"messages": [message], "todos": {"content": "Ship it", "status": "pending"}}
result = middleware.after_model(state, _make_runtime())
attribution = result["messages"][0].additional_kwargs[TOKEN_USAGE_ATTRIBUTION_KEY]
assert attribution["actions"] == [{"kind": "todo_start", "content": "Ship it", "tool_call_id": "write_todos:1"}]
def test_annotates_subagent_and_search_steps(self):
middleware = TokenUsageMiddleware()
message = AIMessage(