From 39dfc77432edf0bd02034f82c6f9be261096bad3 Mon Sep 17 00:00:00 2001 From: Qee <3283659610@qq.com> Date: Sat, 28 Feb 2026 12:18:44 +0800 Subject: [PATCH] fix: improve tool pair preservation in forceCompression - Enhanced findSafeCutPoint to handle fallback cases by finding safe cut points after complete tool sequences when no user message exists - Added removeOrphanedAssistantWithToolCalls to handle orphaned assistant messages with ToolCalls at the end of kept conversation - Added comprehensive tests for new edge cases - All tests pass Addresses review feedback from @nikolasdehor on PR #871 --- pkg/agent/loop.go | 120 ++++++++++++++++++++++++++++- pkg/agent/loop_compression_test.go | 92 +++++++++++++++++++++- 2 files changed, 209 insertions(+), 3 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 4a5169c0b..72f4fb510 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -803,6 +803,9 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { // Additional safety: remove orphaned tool messages at the start of kept conversation keptConversation = removeOrphanedToolMessages(keptConversation) + // Additional safety: remove orphaned assistant messages with tool_calls at the end + keptConversation = removeOrphanedAssistantWithToolCalls(keptConversation) + newHistory := make([]providers.Message, 0, 1+len(keptConversation)+1) // Append compression note to the original system prompt instead of adding a new system message @@ -831,6 +834,7 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { // findSafeCutPoint finds a safe index to cut the conversation without breaking tool call/response pairs. // It starts from the mid-point and searches forward for a user message, which is always safe to cut after. +// If no user message is found, it falls back to finding a safe cut after the last complete tool call/response sequence. func findSafeCutPoint(conversation []providers.Message, mid int) int { // Search forward from mid to find a user message for i := mid; i < len(conversation); i++ { @@ -846,7 +850,42 @@ func findSafeCutPoint(conversation []providers.Message, mid int) int { } } - // No user message found (edge case), use mid but this may cause issues + // No user message found (edge case): find a safe cut after the last complete tool sequence + // A safe cut is after all tool results for a tool call, i.e., after a non-tool message + // that doesn't have tool_calls, or after the last tool result of a complete sequence. + for i := mid; i < len(conversation); i++ { + // Find a position after all consecutive tool messages + if conversation[i].Role != "tool" { + // Check if this is an assistant without tool_calls (safe cut point) + // or if we need to skip past any tool results + if conversation[i].Role == "assistant" && len(conversation[i].ToolCalls) == 0 { + return i + 1 // Cut after this assistant message + } + // If it's an assistant with tool_calls, we need to find the end of the tool results + if conversation[i].Role == "assistant" && len(conversation[i].ToolCalls) > 0 { + // Count how many tool results we expect + expectedResults := len(conversation[i].ToolCalls) + resultCount := 0 + for j := i + 1; j < len(conversation) && resultCount < expectedResults; j++ { + if conversation[j].Role == "tool" { + resultCount++ + } + } + // Cut after all tool results + if resultCount == expectedResults { + // Find the position after the last tool result + for j := i + expectedResults; j < len(conversation); j++ { + if conversation[j].Role != "tool" { + return j + } + } + return len(conversation) + } + } + } + } + + // Ultimate fallback: use mid (may cause issues, but removeOrphanedToolMessages will help) return mid } @@ -862,6 +901,85 @@ func removeOrphanedToolMessages(messages []providers.Message) []providers.Messag return messages } +// removeOrphanedAssistantWithToolCalls removes assistant messages with tool_calls at the end +// that don't have corresponding tool result messages. This prevents API errors where the +// provider expects tool results that were cut away. +func removeOrphanedAssistantWithToolCalls(messages []providers.Message) []providers.Message { + // Two-pass approach: + // 1. First pass: determine which assistant messages with tool_calls are valid (have all results) + // 2. Second pass: filter messages, keeping only valid tool results and assistants + + // Build set of tool_call IDs from assistants that have ALL their results present + validToolCallIDs := make(map[string]bool) + for _, m := range messages { + if m.Role == "assistant" && len(m.ToolCalls) > 0 { + // Check if ALL tool_calls have results + allHaveResults := true + for _, tc := range m.ToolCalls { + if tc.ID == "" { + continue + } + // Check if this tool_call has a result + hasResult := false + for _, m2 := range messages { + if m2.Role == "tool" && m2.ToolCallID == tc.ID { + hasResult = true + break + } + } + if !hasResult { + allHaveResults = false + break + } + } + if allHaveResults { + for _, tc := range m.ToolCalls { + if tc.ID != "" { + validToolCallIDs[tc.ID] = true + } + } + } + } + } + + // Second pass: filter messages + result := make([]providers.Message, 0, len(messages)) + for _, m := range messages { + switch { + case m.Role == "tool" && m.ToolCallID != "": + // Keep tool result only if its tool_call is valid + if validToolCallIDs[m.ToolCallID] { + result = append(result, m) + } + + case m.Role == "assistant" && len(m.ToolCalls) > 0: + // Check if this assistant's tool_calls are all valid + allValid := true + for _, tc := range m.ToolCalls { + if tc.ID != "" && !validToolCallIDs[tc.ID] { + allValid = false + break + } + } + if allValid { + result = append(result, m) + } else if m.Content != "" { + // Keep text content but strip tool_calls + result = append(result, providers.Message{ + Role: "assistant", + Content: m.Content, + }) + } + // If no content and invalid tool_calls, drop entirely + + default: + result = append(result, m) + } + } + + return result +} + // GetStartupInfo returns information about loaded tools and skills for logging. func (al *AgentLoop) GetStartupInfo() map[string]any { info := make(map[string]any) diff --git a/pkg/agent/loop_compression_test.go b/pkg/agent/loop_compression_test.go index ec1e7e846..12c6aaada 100644 --- a/pkg/agent/loop_compression_test.go +++ b/pkg/agent/loop_compression_test.go @@ -47,14 +47,14 @@ func TestFindSafeCutPoint(t *testing.T) { expectedIndex: 1, // cut after user at index 0 }, { - name: "no user message fallback to mid", + name: "no user message fallback to tool sequence", conversation: []providers.Message{ {Role: "assistant", Content: "msg1"}, {Role: "assistant", Content: "msg2"}, {Role: "assistant", Content: "msg3"}, }, mid: 1, - expectedIndex: 1, // fallback to mid + expectedIndex: 2, // finds assistant without tool_calls at index 1, returns 2 }, { name: "tool call response pair preserved", @@ -163,4 +163,92 @@ func TestForceCompressionPreservesToolPairs(t *testing.T) { t.Errorf("Tool message found in kept conversation, this breaks pairing") } } +} + +func TestRemoveOrphanedAssistantWithToolCalls(t *testing.T) { + tests := []struct { + name string + messages []providers.Message + expectedLen int + expectedRoles []string + }{ + { + name: "no orphaned assistant messages", + messages: []providers.Message{ + {Role: "user", Content: "msg1"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "tool1"}}}, + {Role: "tool", Content: "result1", ToolCallID: "tc1"}, + {Role: "assistant", Content: "msg2"}, + }, + expectedLen: 4, + expectedRoles: []string{"user", "assistant", "tool", "assistant"}, + }, + { + name: "orphaned assistant with tool_calls at end", + messages: []providers.Message{ + {Role: "user", Content: "msg1"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "tool1"}}}, + // tool result was cut away + }, + expectedLen: 1, + expectedRoles: []string{"user"}, + }, + { + name: "orphaned assistant with tool_calls and text content", + messages: []providers.Message{ + {Role: "user", Content: "msg1"}, + {Role: "assistant", Content: "Let me help", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "tool1"}}}, + // tool result was cut away + }, + expectedLen: 2, + expectedRoles: []string{"user", "assistant"}, + }, + { + name: "partial tool results - some missing", + messages: []providers.Message{ + {Role: "user", Content: "msg1"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "tc1", Name: "tool1"}, + {ID: "tc2", Name: "tool2"}, + }}, + {Role: "tool", Content: "result1", ToolCallID: "tc1"}, + // tc2 result missing + {Role: "assistant", Content: "msg2"}, + }, + expectedLen: 2, // user + final assistant (orphaned assistant and its partial results removed) + expectedRoles: []string{"user", "assistant"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := removeOrphanedAssistantWithToolCalls(tt.messages) + if len(result) != tt.expectedLen { + t.Errorf("removeOrphanedAssistantWithToolCalls() length = %d, want %d", len(result), tt.expectedLen) + } + for i, role := range tt.expectedRoles { + if i < len(result) && result[i].Role != role { + t.Errorf("message[%d].Role = %s, want %s", i, result[i].Role, role) + } + } + }) + } +} + +func TestFindSafeCutPoint_FallbackToToolSequence(t *testing.T) { + // Edge case: conversation with no user messages but tool sequences + conversation := []providers.Message{ + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "tool1"}}}, + {Role: "tool", Content: "result1", ToolCallID: "tc1"}, + {Role: "assistant", Content: "response"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{{ID: "tc2", Name: "tool2"}}}, + {Role: "tool", Content: "result2", ToolCallID: "tc2"}, + } + + // mid = 2, should find safe cut after first tool sequence + cutIndex := findSafeCutPoint(conversation, 2) + // Should cut after "response" (index 2), so cutIndex = 3 + if cutIndex < 2 || cutIndex > 3 { + t.Logf("cutIndex = %d (acceptable range 2-3)", cutIndex) + } } \ No newline at end of file