From 4efe582e42a16c7e0dd653fec2ca3372bfee4a61 Mon Sep 17 00:00:00 2001 From: mrigankad Date: Fri, 20 Feb 2026 18:14:50 +0530 Subject: [PATCH] refactor(agent): address code review on sanitize functions sanitizeHistoryForProvider (context.go): - Replace O(n^2) interveningTurn scan with an O(1) hasInterveningTurn bool maintained incrementally: reset when a tool-call assistant is accepted, set when any user or non-tool-call assistant is appended after it. - Add comment explaining why an empty ToolCallID skips the ID-based check. sanitizeToolPairs (loop_test.go): - Add table-driven tests covering all cases flagged in review: empty input, all matched (no-op), orphaned tool result, orphaned tool call with/without text content, partial match. Co-Authored-By: Claude Sonnet 4.6 --- pkg/agent/context.go | 26 ++++++---- pkg/agent/loop_test.go | 115 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 9 deletions(-) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 530b9ec86..fd579ed86 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -217,6 +217,10 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message // Index (into sanitized) of the last assistant message that carried ToolCalls. // -1 means no such message has been seen yet. lastCallAssistantIdx := -1 + // hasInterveningTurn is true when a user or non-tool-call assistant message + // has been appended since lastCallAssistantIdx was last set. Maintained + // incrementally (O(1) per message) to avoid an O(n) scan for each tool result. + hasInterveningTurn := false for _, msg := range history { switch msg.Role { @@ -229,14 +233,7 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message // Any user message or any assistant message between the tool-call // assistant and the current position means this result belongs to a // now-lost turn and should be dropped. - interveningTurn := false - for i := lastCallAssistantIdx + 1; i < len(sanitized); i++ { - if r := sanitized[i].Role; r == "user" || r == "assistant" { - interveningTurn = true - break - } - } - if interveningTurn { + if hasInterveningTurn { logger.DebugCF("agent", "Dropping orphaned tool message: intervening turn", map[string]interface{}{"tool_call_id": msg.ToolCallID}) continue @@ -245,6 +242,9 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message // the associated assistant message. This catches the case where a // mid-compression cut left a result whose parent call was dropped but // another assistant-with-calls happened to precede it in the kept half. + // An empty ToolCallID skips this check — some providers omit the ID for + // synthetic or legacy tool messages; the intervening-turn check above is + // sufficient in those cases. if msg.ToolCallID != "" { found := false for _, tc := range sanitized[lastCallAssistantIdx].ToolCalls { @@ -277,10 +277,18 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message // Record where this assistant lands so tool-result validation can // reference it. len(sanitized) is the index after append. lastCallAssistantIdx = len(sanitized) + hasInterveningTurn = false + } else if lastCallAssistantIdx >= 0 { + // An assistant without tool calls is an intervening turn relative + // to the last tool-call assistant. + hasInterveningTurn = true } sanitized = append(sanitized, msg) - default: + default: // user (and any future roles) + if lastCallAssistantIdx >= 0 { + hasInterveningTurn = true + } sanitized = append(sanitized, msg) } } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 360685eca..f68d10b25 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -613,3 +613,118 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { t.Errorf("Expected history to be compressed (len < 8), got %d", len(finalHistory)) } } + +func tc(id string) providers.ToolCall { return providers.ToolCall{ID: id} } + +func TestSanitizeToolPairs(t *testing.T) { + tests := []struct { + name string + in []providers.Message + want []providers.Message + }{ + { + name: "empty input", + in: []providers.Message{}, + want: []providers.Message{}, + }, + { + name: "all tool calls matched - no-op", + in: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", ToolCalls: []providers.ToolCall{tc("a"), tc("b")}}, + {Role: "tool", ToolCallID: "a"}, + {Role: "tool", ToolCallID: "b"}, + }, + want: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", ToolCalls: []providers.ToolCall{tc("a"), tc("b")}}, + {Role: "tool", ToolCallID: "a"}, + {Role: "tool", ToolCallID: "b"}, + }, + }, + { + name: "orphaned tool result is dropped", + in: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "tool", ToolCallID: "orphan"}, + {Role: "assistant", Content: "ok"}, + }, + want: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "ok"}, + }, + }, + { + name: "orphaned tool call with text content - call stripped, text preserved", + in: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "thinking...", ToolCalls: []providers.ToolCall{tc("x")}}, + }, + want: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "thinking..."}, + }, + }, + { + name: "orphaned tool call with no content - message dropped", + in: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", ToolCalls: []providers.ToolCall{tc("x")}}, + }, + want: []providers.Message{ + {Role: "user", Content: "hi"}, + }, + }, + { + name: "partial match - unmatched calls stripped, matched calls and text preserved", + in: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "let me check", ToolCalls: []providers.ToolCall{tc("a"), tc("b"), tc("c")}}, + {Role: "tool", ToolCallID: "a"}, + {Role: "tool", ToolCallID: "c"}, + }, + want: []providers.Message{ + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "let me check", ToolCalls: []providers.ToolCall{tc("a"), tc("c")}}, + {Role: "tool", ToolCallID: "a"}, + {Role: "tool", ToolCallID: "c"}, + }, + }, + { + name: "user messages pass through unchanged", + in: []providers.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "world"}, + }, + want: []providers.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "world"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := sanitizeToolPairs(tt.in) + if len(got) != len(tt.want) { + t.Fatalf("len(got) = %d, len(want) = %d\ngot: %+v\nwant: %+v", len(got), len(tt.want), got, tt.want) + } + for i := range got { + g, w := got[i], tt.want[i] + if g.Role != w.Role || g.Content != w.Content || g.ToolCallID != w.ToolCallID { + t.Errorf("msg[%d]: got {Role:%q Content:%q ToolCallID:%q}, want {Role:%q Content:%q ToolCallID:%q}", + i, g.Role, g.Content, g.ToolCallID, w.Role, w.Content, w.ToolCallID) + } + if len(g.ToolCalls) != len(w.ToolCalls) { + t.Errorf("msg[%d] ToolCalls len: got %d, want %d", i, len(g.ToolCalls), len(w.ToolCalls)) + continue + } + for j := range g.ToolCalls { + if g.ToolCalls[j].ID != w.ToolCalls[j].ID { + t.Errorf("msg[%d].ToolCalls[%d].ID: got %q, want %q", i, j, g.ToolCalls[j].ID, w.ToolCalls[j].ID) + } + } + } + }) + } +}