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
This commit is contained in:
Qee
parent
997f31d10a
commit
39dfc77432
2 changed files with 209 additions and 3 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue