fix(agent): refactor sanitizeToolPairs per review feedback

- Track retained tool_call IDs to eliminate secondary orphans: when an
  assistant message is dropped, its tool results are also dropped
- Split removed/modified counters for accurate logging
- Always call SetHistory after sanitization (not conditional on len)
- Add full field assertions to CompletePairs test
- Fix MultiToolCallPartialResults test: secondary orphans now cleaned
This commit is contained in:
Nikolas de Hor 2026-03-03 01:24:26 -03:00
parent 82a9962f6d
commit 52e510b56b
3 changed files with 81 additions and 52 deletions

View file

@ -1313,9 +1313,7 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
// Sanitize remaining history to fix orphaned tool pairs from truncation
remaining := agent.Sessions.GetHistory(sessionKey)
sanitized := sanitizeToolPairs(remaining)
if len(sanitized) != len(remaining) {
agent.Sessions.SetHistory(sessionKey, sanitized)
}
agent.Sessions.SetHistory(sessionKey, sanitized)
agent.Sessions.Save(sessionKey)
}

View file

@ -10,22 +10,16 @@ import (
// Orphaned messages are removed to prevent provider API errors (e.g.,
// Anthropic's "tool_use ids were provided that do not have a tool_use block").
//
// Uses a two-pass approach:
// 1. Forward pass: decide which assistant tool_call messages to retain,
// tracking which tool_call IDs survive. Also collect all tool_result IDs.
// 2. Forward pass: emit retained messages, dropping tool results whose
// tool_call was not retained.
//
// This is applied after history compression to fix pairs that were split
// when forceCompression() or summarizeSession() truncated the history.
func sanitizeToolPairs(messages []providers.Message) []providers.Message {
// Build set of tool_call IDs present in assistant messages
toolCallIDs := make(map[string]bool)
for _, m := range messages {
if m.Role == "assistant" && len(m.ToolCalls) > 0 {
for _, tc := range m.ToolCalls {
if tc.ID != "" {
toolCallIDs[tc.ID] = true
}
}
}
}
// Build set of tool_result IDs present
// Collect all tool_result IDs (needed to decide if assistant msgs survive)
toolResultIDs := make(map[string]bool)
for _, m := range messages {
if m.Role == "tool" && m.ToolCallID != "" {
@ -33,25 +27,17 @@ func sanitizeToolPairs(messages []providers.Message) []providers.Message {
}
}
// Filter: keep tool results only if their tool_call exists,
// and keep assistant tool_call messages only if all results exist
result := make([]providers.Message, 0, len(messages))
removed := 0
// Forward pass: decide which assistant tool_call messages to keep,
// tracking the set of retained tool_call IDs.
retainedCallIDs := make(map[string]bool)
type decision struct {
keep bool
modified bool // true if we strip tool_calls but keep text
}
assistantDecisions := make(map[int]decision) // index -> decision
for _, m := range messages {
switch {
case m.Role == "tool" && m.ToolCallID != "":
// Keep tool result only if its tool_call is present
if toolCallIDs[m.ToolCallID] {
result = append(result, m)
} else {
removed++
logger.DebugCF("agent", "sanitizeToolPairs: removing orphaned tool result",
map[string]interface{}{"tool_call_id": m.ToolCallID})
}
case m.Role == "assistant" && len(m.ToolCalls) > 0:
// Check if ALL tool_calls have matching results
for i, m := range messages {
if m.Role == "assistant" && len(m.ToolCalls) > 0 {
allHaveResults := true
for _, tc := range m.ToolCalls {
if tc.ID != "" && !toolResultIDs[tc.ID] {
@ -60,21 +46,52 @@ func sanitizeToolPairs(messages []providers.Message) []providers.Message {
}
}
if allHaveResults {
result = append(result, m)
assistantDecisions[i] = decision{keep: true}
for _, tc := range m.ToolCalls {
if tc.ID != "" {
retainedCallIDs[tc.ID] = true
}
}
} else if m.Content != "" {
// Keep the text content but strip the tool calls
assistantDecisions[i] = decision{keep: true, modified: true}
} else {
assistantDecisions[i] = decision{keep: false}
}
}
}
// Emit pass: build result using retained IDs
result := make([]providers.Message, 0, len(messages))
removed := 0
modified := 0
for i, m := range messages {
switch {
case m.Role == "tool" && m.ToolCallID != "":
if retainedCallIDs[m.ToolCallID] {
result = append(result, m)
} else {
removed++
logger.DebugCF("agent", "sanitizeToolPairs: stripping orphaned tool_calls from assistant message, keeping text content",
logger.DebugCF("agent", "sanitizeToolPairs: removing orphaned tool result",
map[string]interface{}{"tool_call_id": m.ToolCallID})
}
case m.Role == "assistant" && len(m.ToolCalls) > 0:
d := assistantDecisions[i]
if !d.keep {
removed++
logger.DebugCF("agent", "sanitizeToolPairs: removing orphaned assistant tool_call message",
map[string]interface{}{"tool_call_count": len(m.ToolCalls)})
} else if d.modified {
modified++
logger.DebugCF("agent", "sanitizeToolPairs: stripping orphaned tool_calls, keeping text",
map[string]interface{}{"tool_call_count": len(m.ToolCalls)})
result = append(result, providers.Message{
Role: "assistant",
Content: m.Content,
})
} else {
// No text content and missing results - drop entirely
removed++
logger.DebugCF("agent", "sanitizeToolPairs: removing orphaned assistant message with tool_calls",
map[string]interface{}{"tool_call_count": len(m.ToolCalls)})
result = append(result, m)
}
default:
@ -82,9 +99,9 @@ func sanitizeToolPairs(messages []providers.Message) []providers.Message {
}
}
if removed > 0 {
logger.WarnCF("agent", "sanitizeToolPairs: removed orphaned tool pair messages",
map[string]interface{}{"removed_count": removed})
if removed > 0 || modified > 0 {
logger.WarnCF("agent", "sanitizeToolPairs: cleaned orphaned tool pair messages",
map[string]interface{}{"removed": removed, "modified": modified})
}
return result

View file

@ -41,7 +41,19 @@ func TestSanitizeToolPairs_CompletePairs(t *testing.T) {
result := sanitizeToolPairs(messages)
if len(result) != 4 {
t.Errorf("expected 4 messages, got %d", len(result))
t.Fatalf("expected 4 messages, got %d", len(result))
}
if result[0].Role != "user" || result[0].Content != "What is the weather?" {
t.Errorf("msg[0]: expected user 'What is the weather?', got %s %q", result[0].Role, result[0].Content)
}
if result[1].Role != "assistant" || len(result[1].ToolCalls) != 1 || result[1].ToolCalls[0].ID != "call_1" {
t.Errorf("msg[1]: expected assistant with tool_call call_1, got role=%s calls=%d", result[1].Role, len(result[1].ToolCalls))
}
if result[2].Role != "tool" || result[2].ToolCallID != "call_1" {
t.Errorf("msg[2]: expected tool result for call_1, got role=%s id=%s", result[2].Role, result[2].ToolCallID)
}
if result[3].Role != "assistant" || result[3].Content != "The temperature is 72 degrees." {
t.Errorf("msg[3]: expected assistant final, got role=%s content=%q", result[3].Role, result[3].Content)
}
}
@ -148,14 +160,16 @@ func TestSanitizeToolPairs_MultiToolCallPartialResults(t *testing.T) {
result := sanitizeToolPairs(messages)
// The assistant msg is dropped (not all tool_calls have results).
// call_1's tool result stays because toolCallIDs includes call_1
// from the original scan. This is a secondary orphan handled by
// the existing Diegox-17 fix in BuildMessages (strips leading tool msgs).
if len(result) != 3 {
t.Errorf("expected 3 messages, got %d", len(result))
// call_1's tool result is also dropped because the assistant message
// that issued call_1 was removed — no secondary orphans left behind.
if len(result) != 2 {
t.Errorf("expected 2 messages (assistant + call_1 result both dropped), got %d", len(result))
}
if result[0].Role != "user" {
t.Errorf("expected first message to be user, got %s", result[0].Role)
if result[0].Role != "user" || result[0].Content != "Do two things" {
t.Errorf("expected first message to be user 'Do two things', got %s %q", result[0].Role, result[0].Content)
}
if result[1].Role != "user" || result[1].Content != "OK" {
t.Errorf("expected second message to be user 'OK', got %s %q", result[1].Role, result[1].Content)
}
}