diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 74e9feda2e38e..d51aed44de80b 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -28,6 +28,7 @@ from __future__ import annotations +import json import logging import os import tempfile @@ -38,6 +39,7 @@ from typing import Any, Optional, Tuple from agent.model_metadata import estimate_request_tokens_rough +from agent.tool_dispatch_helpers import make_tool_result_message logger = logging.getLogger(__name__) @@ -52,6 +54,86 @@ ) +def _todo_state_messages_for_compression(agent: Any) -> list[dict]: + """Return provider-valid tool history that preserves active todos. + + This state is internal continuity, not user intent. It must therefore be + represented as canonical ``todo`` tool history rather than as a synthetic + ``role=user`` message that can become the next active instruction after a + compaction boundary. + + The serialized summary describes only the carried active payload. Completed + and cancelled items are deliberately omitted so finished work does not + re-enter the model context after compression. + """ + todo_store = getattr(agent, "_todo_store", None) + if todo_store is None: + return [] + + active_items_fn = getattr(todo_store, "active_items", None) + if not callable(active_items_fn): + return [] + active_items = active_items_fn() + if not isinstance(active_items, list): + return [] + if not active_items: + return [] + + pending = sum(1 for item in active_items if item["status"] == "pending") + in_progress = sum(1 for item in active_items if item["status"] == "in_progress") + content = json.dumps( + { + "todos": active_items, + "summary": { + "total": len(active_items), + "pending": pending, + "in_progress": in_progress, + "completed": 0, + "cancelled": 0, + }, + }, + ensure_ascii=False, + ) + tool_call_id = f"context-compression-todo-{uuid.uuid4().hex[:8]}" + assistant_msg = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": {"name": "todo", "arguments": "{}"}, + } + ], + } + return [assistant_msg, make_tool_result_message("todo", content, tool_call_id)] + + +def _insert_todo_state_messages( + compressed: list[dict], + todo_state_messages: list[dict], +) -> list[dict]: + """Insert compressed todo state at a provider-valid continuation point. + + The todo state is synthetic assistant/tool history, not user intent. Place + it after the latest real user turn so it cannot become a leading assistant + message on providers that require user-first histories, while preserving + adjacency between the assistant tool call and its tool result. + """ + if not todo_state_messages: + return compressed + + for idx in range(len(compressed) - 1, -1, -1): + if compressed[idx].get("role") == "user": + return [ + *compressed[: idx + 1], + *todo_state_messages, + *compressed[idx + 1 :], + ] + + return compressed + + def _compression_lock_holder(agent: Any) -> str: """Build a unique holder id for the lock: pid:tid:agent-instance:uuid. @@ -644,9 +726,9 @@ def _release_lock() -> None: "check auxiliary.compression.model in config.yaml." ) - todo_snapshot = agent._todo_store.format_for_injection() - if todo_snapshot: - compressed.append({"role": "user", "content": todo_snapshot}) + todo_state_messages = _todo_state_messages_for_compression(agent) + if todo_state_messages: + compressed = _insert_todo_state_messages(compressed, todo_state_messages) agent._invalidate_system_prompt() new_system_prompt = agent._build_system_prompt(system_message) diff --git a/tests/agent/test_context_compression_todo_leak.py b/tests/agent/test_context_compression_todo_leak.py new file mode 100644 index 0000000000000..462d594fe59c3 --- /dev/null +++ b/tests/agent/test_context_compression_todo_leak.py @@ -0,0 +1,208 @@ +"""Regression tests for preserving todo state across context compression.""" + +from typing import Any, cast +from unittest.mock import MagicMock, patch + +from run_agent import AIAgent +from tools.todo_tool import TodoStore + + +def _agent_with_fake_compressor(): + agent = cast(Any, object.__new__(AIAgent)) + agent.session_id = "todo-leak-session" + agent.model = "test/model" + agent.tools = [] + agent.platform = "cli" + agent.log_prefix = "" + agent.quiet_mode = True + agent.compression_in_place = False + agent._compression_feasibility_checked = True + agent._session_db = None + agent._memory_manager = None + agent._cached_system_prompt = None + agent._last_flushed_db_idx = 0 + agent._flushed_db_message_ids = set() + agent._todo_store = TodoStore() + agent._todo_store.write( + [ + { + "id": "stale", + "content": "STALE TODO SHOULD NOT BE A USER MESSAGE", + "status": "in_progress", + }, + {"id": "done", "content": "done item", "status": "completed"}, + ] + ) + agent._emit_status = lambda *args, **kwargs: None + agent._emit_warning = lambda *args, **kwargs: None + agent._vprint = lambda *args, **kwargs: None + agent._invalidate_system_prompt = lambda *args, **kwargs: None + agent._build_system_prompt = lambda *args, **kwargs: "rebuilt-system-prompt" + agent.commit_memory_session = lambda *args, **kwargs: None + + compressor = MagicMock() + compressor.compress.return_value = [ + { + "role": "assistant", + "content": "[CONTEXT COMPACTION — REFERENCE ONLY]\nsummary", + }, + {"role": "user", "content": "latest real user request"}, + ] + compressor._last_compress_aborted = False + compressor._last_summary_error = None + compressor._last_aux_model_failure_model = None + compressor._last_aux_model_failure_error = None + compressor.compression_count = 1 + agent.context_compressor = compressor + return agent + + +def test_compress_context_does_not_append_todo_snapshot_as_user_message(): + from agent.conversation_compression import compress_context + + agent = _agent_with_fake_compressor() + compressed, _ = compress_context( + agent, + [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "reply"}, + ], + "system", + approx_tokens=100, + ) + + leaked = [ + msg + for msg in compressed + if msg.get("role") == "user" + and isinstance(msg.get("content"), str) + and msg["content"].startswith( + "[Your active task list was preserved across context compression]" + ) + ] + assert leaked == [] + user_messages = [msg for msg in compressed if msg.get("role") == "user"] + assert user_messages[-1].get("content") == "latest real user request" + + +def test_compress_context_preserves_active_todos_as_paired_tool_state(): + from agent.conversation_compression import compress_context + + agent = _agent_with_fake_compressor() + compressed, _ = compress_context( + agent, + [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "reply"}, + ], + "system", + approx_tokens=100, + ) + + tool_indexes = [ + idx for idx, msg in enumerate(compressed) if msg.get("role") == "tool" + ] + assert len(tool_indexes) == 1 + tool_msg = compressed[tool_indexes[0]] + assert tool_msg.get("name") == "todo" + assert tool_msg.get("tool_name") == "todo" + assert tool_msg.get("tool_call_id") + assert AIAgent._tool_response_matches_todo_call(compressed, tool_indexes[0]) is True + assert "STALE TODO SHOULD NOT BE A USER MESSAGE" in tool_msg.get("content", "") + assert "done item" not in tool_msg.get("content", "") + + +def test_compress_context_inserts_todo_tool_state_after_latest_real_user_turn(): + from agent.conversation_compression import compress_context + + agent = _agent_with_fake_compressor() + compressed, _ = compress_context( + agent, + [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "reply"}, + ], + "system", + approx_tokens=100, + ) + + latest_user_index = max( + idx + for idx, msg in enumerate(compressed) + if msg.get("role") == "user" and msg.get("content") == "latest real user request" + ) + tool_index = next(idx for idx, msg in enumerate(compressed) if msg.get("role") == "tool") + assistant_index = tool_index - 1 + + assert assistant_index > latest_user_index + assert compressed[assistant_index].get("role") == "assistant" + assert compressed[assistant_index].get("tool_calls") + assert compressed[tool_index].get("tool_call_id") == compressed[assistant_index]["tool_calls"][0]["id"] + + +def test_insert_todo_state_messages_preserves_tail_after_latest_user(): + from agent.conversation_compression import _insert_todo_state_messages + + compressed = [ + {"role": "user", "content": "first request"}, + {"role": "assistant", "content": "first reply"}, + {"role": "user", "content": "latest request"}, + {"role": "assistant", "content": "tail reply"}, + ] + todo_state_messages = [ + {"role": "assistant", "content": None, "tool_calls": [{"id": "todo-1"}]}, + {"role": "tool", "tool_call_id": "todo-1", "content": "{}"}, + ] + + result = _insert_todo_state_messages(compressed, todo_state_messages) + + assert result == [ + {"role": "user", "content": "first request"}, + {"role": "assistant", "content": "first reply"}, + {"role": "user", "content": "latest request"}, + *todo_state_messages, + {"role": "assistant", "content": "tail reply"}, + ] + + +def test_insert_todo_state_messages_without_user_turn_leaves_history_unchanged(): + from agent.conversation_compression import _insert_todo_state_messages + + compressed = [{"role": "assistant", "content": "summary only"}] + todo_state_messages = [ + {"role": "assistant", "content": None, "tool_calls": [{"id": "todo-1"}]}, + {"role": "tool", "tool_call_id": "todo-1", "content": "{}"}, + ] + + assert _insert_todo_state_messages(compressed, todo_state_messages) == compressed + + +def test_compressed_todo_tool_state_hydrates_fresh_agent_store(): + from agent.conversation_compression import compress_context + + agent = _agent_with_fake_compressor() + compressed, _ = compress_context( + agent, + [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "reply"}, + ], + "system", + approx_tokens=100, + ) + + fresh = cast(Any, object.__new__(AIAgent)) + fresh.session_id = "todo-leak-session" + fresh.quiet_mode = True + fresh.log_prefix = "" + fresh._todo_store = TodoStore() + with patch("run_agent._set_interrupt"): + fresh._hydrate_todo_store(compressed) + + assert fresh._todo_store.read() == [ + { + "id": "stale", + "content": "STALE TODO SHOULD NOT BE A USER MESSAGE", + "status": "in_progress", + } + ] diff --git a/tests/tools/test_todo_tool.py b/tests/tools/test_todo_tool.py index dbb64e80ee6ad..26d7f2cb57cda 100644 --- a/tests/tools/test_todo_tool.py +++ b/tests/tools/test_todo_tool.py @@ -48,6 +48,31 @@ def test_non_empty_store(self): assert store.has_items() is True +class TestActiveItems: + def test_returns_only_pending_and_in_progress_items(self): + store = TodoStore() + store.write([ + {"id": "1", "content": "Done", "status": "completed"}, + {"id": "2", "content": "Next", "status": "pending"}, + {"id": "3", "content": "Working", "status": "in_progress"}, + {"id": "4", "content": "Skipped", "status": "cancelled"}, + ]) + + assert store.active_items() == [ + {"id": "2", "content": "Next", "status": "pending"}, + {"id": "3", "content": "Working", "status": "in_progress"}, + ] + + def test_returns_copy(self): + store = TodoStore() + store.write([{"id": "1", "content": "Task", "status": "pending"}]) + + active = store.active_items() + active[0]["content"] = "MUTATED" + + assert store.read()[0]["content"] == "Task" + + class TestFormatForInjection: def test_empty_returns_none(self): store = TodoStore() @@ -122,10 +147,10 @@ def test_no_store_returns_error(self): class TestTodoStoreBounds: """Bounds on persisted todo state (GHSA-5g4g-6jrg-mw3g hardening). - The todo list is re-injected into context after every compression event, - so an unbounded item — whether authored by the model or replayed from - caller-supplied history on the API server's _hydrate_todo_store path — - would defeat the compression it rides through. These pin the caps. + Active todos are carried across compression, so an unbounded item — whether + authored by the model or replayed from caller-supplied history on the API + server's _hydrate_todo_store path — would defeat the compression it rides + through. These pin the caps. Not a security boundary (the API surface is authenticated and the caller supplies their own history); this is footgun containment / parity. """ diff --git a/tools/todo_tool.py b/tools/todo_tool.py index 3c657c034d6d1..5320e2a34a83f 100644 --- a/tools/todo_tool.py +++ b/tools/todo_tool.py @@ -4,8 +4,8 @@ Provides an in-memory task list the agent uses to decompose complex tasks, track progress, and maintain focus across long conversations. The state -lives on the AIAgent instance (one per session) and is re-injected into -the conversation after context compression events. +lives on the AIAgent instance (one per session) and is preserved across +context compression as canonical todo tool history. Design: - Single `todo` tool: provide `todos` param to write, omit to read @@ -22,18 +22,20 @@ VALID_STATUSES = {"pending", "in_progress", "completed", "cancelled"} # Bounds on persisted todo state. The todo list is a planning aid the model -# re-reads after every context-compression event (see format_for_injection), -# so unbounded item content or count defeats the compression it rides through. +# re-reads after context compression via canonical todo tool history, so +# unbounded item content or count defeats the compression it rides through. # These caps keep a single oversized item (whether authored by the model or # replayed from caller-supplied history on the API server) from inflating the -# re-injection block. Generous relative to real plans — a todo item is a short -# task description, and active lists are a handful of items, not hundreds. +# preserved tool-result payload. Generous relative to real plans — a todo item +# is a short task description, and active lists are a handful of items, not +# hundreds. MAX_TODO_CONTENT_CHARS = 4000 MAX_TODO_ITEMS = 256 # Upper bound on a single todo tool-result payload accepted during history # hydration. The gateway/API server replays caller-supplied conversation # history to rebuild the store, so an oversized forged result is dropped -# before it is parsed and re-injected (see AIAgent._hydrate_todo_store). +# before it is parsed into the in-memory TodoStore (see +# AIAgent._hydrate_todo_store). MAX_TODO_RESULT_CHARS = 512_000 _TRUNCATION_MARKER = "… [truncated]" @@ -108,12 +110,21 @@ def has_items(self) -> bool: """Check if there are any items in the list.""" return bool(self._items) + def active_items(self) -> List[Dict[str, str]]: + """Return a copy of todos that should remain active after compression.""" + return [ + item.copy() + for item in self._items + if item["status"] in {"pending", "in_progress"} + ] + def format_for_injection(self) -> Optional[str]: """ - Render the todo list for post-compression injection. + Render active todos for legacy display contexts. - Returns a human-readable string to append to the compressed - message history, or None if the list is empty. + Returns a human-readable string, or None if the list has no active + pending/in-progress items. Context compression preserves todos as + canonical tool history instead of appending this text as a user message. """ if not self._items: return None @@ -126,12 +137,9 @@ def format_for_injection(self) -> Optional[str]: "cancelled": "[~]", } - # Only inject pending/in_progress items — completed/cancelled ones - # cause the model to re-do finished work after compression. - active_items = [ - item for item in self._items - if item["status"] in {"pending", "in_progress"} - ] + # Only render pending/in_progress items — completed/cancelled ones + # cause the model to re-do finished work if carried across compression. + active_items = self.active_items() if not active_items: return None