diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index d9c59864f..3b5812e30 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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) } diff --git a/pkg/agent/sanitize.go b/pkg/agent/sanitize.go index 66f2112eb..e5e6cf186 100644 --- a/pkg/agent/sanitize.go +++ b/pkg/agent/sanitize.go @@ -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 diff --git a/pkg/agent/sanitize_test.go b/pkg/agent/sanitize_test.go index 20d26cace..aae146e08 100644 --- a/pkg/agent/sanitize_test.go +++ b/pkg/agent/sanitize_test.go @@ -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) } }