Fix compacted tool history for resumed sessions

2026-06-24 14:18:15 +08:00
parent 45792c8fd5
commit 3637e2c73a
4 changed files with 172 additions and 7 deletions
@@ -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()