Fix compacted tool history for resumed sessions

This commit is contained in:
wuyang6
2026-06-24 14:18:15 +08:00
parent 45792c8fd5
commit 3637e2c73a
4 changed files with 172 additions and 7 deletions
+7 -5
View File
@@ -37,7 +37,7 @@ from .agent_prompting import (
build_system_prompt_parts, build_system_prompt_parts,
render_system_prompt, render_system_prompt,
) )
from .agent_session import AgentSessionState from .agent_session import AgentSessionState, sanitize_model_message_sequence
from .agent_slash_commands import preprocess_slash_command from .agent_slash_commands import preprocess_slash_command
from .agent_tools import ( from .agent_tools import (
AgentTool, AgentTool,
@@ -2468,7 +2468,9 @@ class LocalCodingAgent:
if compact_end <= prefix_count: if compact_end <= prefix_count:
return False return False
candidates = session.messages[prefix_count:compact_end] candidates = session.messages[prefix_count:compact_end]
preserved_tail = list(session.messages[compact_end:]) preserved_tail = sanitize_model_message_sequence(
list(session.messages[compact_end:])
)
if not candidates: if not candidates:
return False return False
compacted_tokens = sum( compacted_tokens = sum(
@@ -2492,13 +2494,13 @@ class LocalCodingAgent:
turn_index=turn_index, turn_index=turn_index,
estimated_tokens_before=usage_total, estimated_tokens_before=usage_total,
estimated_tokens_removed=compacted_tokens, estimated_tokens_removed=compacted_tokens,
preserved_tail_count=tail_count, preserved_tail_count=len(preserved_tail),
preserved_tail=preserved_tail, preserved_tail=preserved_tail,
) )
session.messages = ( session.messages = (
session.messages[:prefix_count] session.messages[:prefix_count]
+ [compact_message] + [compact_message]
+ session.messages[compact_end:] + preserved_tail
) )
stream_events.append( stream_events.append(
{ {
@@ -2507,7 +2509,7 @@ class LocalCodingAgent:
'compacted_message_count': len(candidates), 'compacted_message_count': len(candidates),
'estimated_tokens_before': usage_total, 'estimated_tokens_before': usage_total,
'estimated_tokens_removed': compacted_tokens, 'estimated_tokens_removed': compacted_tokens,
'preserved_tail_count': tail_count, 'preserved_tail_count': len(preserved_tail),
'preserved_tail_ids': [ 'preserved_tail_ids': [
message.message_id for message in preserved_tail if message.message_id message.message_id for message in preserved_tail if message.message_id
], ],
+79
View File
@@ -96,6 +96,84 @@ class AgentMessage:
) )
def _tool_call_id(tool_call: JSONDict) -> str | None:
raw_id = tool_call.get('id')
return raw_id if isinstance(raw_id, str) and raw_id else None
def _assistant_tool_call_ids(message: AgentMessage) -> list[str]:
if message.role != 'assistant' or not message.tool_calls:
return []
ids: list[str] = []
for tool_call in message.tool_calls:
call_id = _tool_call_id(tool_call)
if call_id:
ids.append(call_id)
return ids
def sanitize_model_message_sequence(
messages: list[AgentMessage],
) -> list[AgentMessage]:
"""Remove invalid tool-call fragments from model-facing history.
Chat backends require a tool result to appear immediately after the
assistant message that requested it. Compaction can otherwise preserve a
tail starting with a ``tool`` message after the matching assistant tool call
has been summarized away. Display history is append-only and should not use
this helper; it is only for model-facing context.
"""
cleaned: list[AgentMessage] = []
index = 0
while index < len(messages):
message = messages[index]
if message.role == 'tool':
index += 1
continue
tool_call_ids = _assistant_tool_call_ids(message)
if not tool_call_ids:
cleaned.append(message)
index += 1
continue
group: list[AgentMessage] = []
seen: set[str] = set()
cursor = index + 1
while cursor < len(messages) and messages[cursor].role == 'tool':
tool_message = messages[cursor]
call_id = tool_message.tool_call_id
if (
call_id is None
or call_id not in tool_call_ids
or call_id in seen
):
break
group.append(tool_message)
seen.add(call_id)
cursor += 1
if len(seen) == len(tool_call_ids):
cleaned.append(message)
cleaned.extend(group)
index = cursor
continue
# Drop the incomplete assistant tool-call request and any immediately
# adjacent partial tool results tied to it. Keeping either side would
# make the next model request invalid.
index += 1
while (
index < len(messages)
and messages[index].role == 'tool'
and messages[index].tool_call_id in set(tool_call_ids)
):
index += 1
return cleaned
@dataclass @dataclass
class AgentSessionState: class AgentSessionState:
system_prompt_parts: tuple[str, ...] system_prompt_parts: tuple[str, ...]
@@ -592,6 +670,7 @@ class AgentSessionState:
for message in messages for message in messages
if isinstance(message, dict) if isinstance(message, dict)
] ]
model_messages = sanitize_model_message_sequence(model_messages)
display_source = messages if display_messages is None else display_messages display_source = messages if display_messages is None else display_messages
visible_messages = [ visible_messages = [
AgentMessage.from_openai_message(message) AgentMessage.from_openai_message(message)
+2 -2
View File
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any
from .agent_context_usage import estimate_tokens from .agent_context_usage import estimate_tokens
from .agent_types import UsageStats from .agent_types import UsageStats
from .agent_session import AgentMessage from .agent_session import AgentMessage, sanitize_model_message_sequence
if TYPE_CHECKING: if TYPE_CHECKING:
from .agent_runtime import LocalCodingAgent from .agent_runtime import LocalCodingAgent
@@ -522,7 +522,7 @@ def compact_conversation(
) )
candidates = list(session.messages[prefix_count:compact_end]) candidates = list(session.messages[prefix_count:compact_end])
preserved_tail = list(session.messages[compact_end:]) preserved_tail = sanitize_model_message_sequence(list(session.messages[compact_end:]))
if not candidates: if not candidates:
return CompactionResult( return CompactionResult(
+84
View File
@@ -0,0 +1,84 @@
from __future__ import annotations
import unittest
from src.agent_session import AgentMessage, sanitize_model_message_sequence
class TestSanitizeModelMessageSequence(unittest.TestCase):
def test_removes_orphan_tool_message(self) -> None:
messages = [
AgentMessage(role='user', content='summary', message_id='summary'),
AgentMessage(
role='tool',
content='orphan result',
tool_call_id='call_1',
message_id='tool_1',
),
AgentMessage(role='assistant', content='review', message_id='assistant_1'),
AgentMessage(role='user', content='continue', message_id='user_1'),
]
cleaned = sanitize_model_message_sequence(messages)
self.assertEqual(
[message.message_id for message in cleaned],
['summary', 'assistant_1', 'user_1'],
)
def test_preserves_complete_assistant_tool_group(self) -> None:
assistant = AgentMessage(
role='assistant',
content='',
tool_calls=(
{
'id': 'call_1',
'type': 'function',
'function': {'name': 'read_file', 'arguments': '{}'},
},
),
message_id='assistant_1',
)
tool = AgentMessage(
role='tool',
content='ok',
tool_call_id='call_1',
message_id='tool_1',
)
final = AgentMessage(role='assistant', content='done', message_id='assistant_2')
cleaned = sanitize_model_message_sequence([assistant, tool, final])
self.assertEqual(
[message.message_id for message in cleaned],
['assistant_1', 'tool_1', 'assistant_2'],
)
def test_removes_incomplete_assistant_tool_call_group(self) -> None:
messages = [
AgentMessage(role='user', content='question', message_id='user_1'),
AgentMessage(
role='assistant',
content='',
tool_calls=(
{
'id': 'call_1',
'type': 'function',
'function': {'name': 'read_file', 'arguments': '{}'},
},
),
message_id='assistant_1',
),
AgentMessage(role='user', content='next', message_id='user_2'),
]
cleaned = sanitize_model_message_sequence(messages)
self.assertEqual(
[message.message_id for message in cleaned],
['user_1', 'user_2'],
)
if __name__ == '__main__':
unittest.main()