Skip to content

Commit 170d4d7

Browse files
authored
Merge pull request #2377 from rumpl/sanitize-tool-calls
Sanitize message history
2 parents 15a3405 + d1d820b commit 170d4d7

File tree

6 files changed

+252
-335
lines changed

6 files changed

+252
-335
lines changed

pkg/model/provider/anthropic/beta_client.go

Lines changed: 0 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,6 @@ func (c *Client) createBetaStream(
4444
slog.Error("Failed to convert messages for Anthropic Beta request", "error", err)
4545
return nil, err
4646
}
47-
if err := validateAnthropicSequencingBeta(converted); err != nil {
48-
slog.Warn("Invalid message sequencing for Anthropic Beta API detected, attempting self-repair", "error", err)
49-
converted = repairAnthropicSequencingBeta(converted)
50-
if err2 := validateAnthropicSequencingBeta(converted); err2 != nil {
51-
slog.Error("Failed to self-repair Anthropic Beta sequencing", "error", err2)
52-
return nil, err
53-
}
54-
}
5547
if len(converted) == 0 {
5648
return nil, errors.New("no messages to send after conversion: all messages were filtered out")
5749
}
@@ -148,35 +140,6 @@ func (c *Client) createBetaStream(
148140
return ad, nil
149141
}
150142

151-
// validateAnthropicSequencingBeta performs the same validation as standard API but for Beta payloads
152-
func validateAnthropicSequencingBeta(msgs []anthropic.BetaMessageParam) error {
153-
return validateSequencing(msgs)
154-
}
155-
156-
// repairAnthropicSequencingBeta inserts a synthetic user message with tool_result blocks
157-
// for any assistant tool_use blocks that don't have corresponding tool_result blocks
158-
// in the immediate next user message.
159-
func repairAnthropicSequencingBeta(msgs []anthropic.BetaMessageParam) []anthropic.BetaMessageParam {
160-
return repairSequencing(msgs, func(toolUseIDs map[string]struct{}) anthropic.BetaMessageParam {
161-
blocks := make([]anthropic.BetaContentBlockParamUnion, 0, len(toolUseIDs))
162-
for id := range toolUseIDs {
163-
slog.Debug("Creating synthetic tool_result", "tool_use_id", id)
164-
blocks = append(blocks, anthropic.BetaContentBlockParamUnion{
165-
OfToolResult: &anthropic.BetaToolResultBlockParam{
166-
ToolUseID: id,
167-
Content: []anthropic.BetaToolResultBlockParamContentUnion{
168-
{OfText: &anthropic.BetaTextBlockParam{Text: "(tool execution failed)"}},
169-
},
170-
},
171-
})
172-
}
173-
return anthropic.BetaMessageParam{
174-
Role: anthropic.BetaMessageParamRoleUser,
175-
Content: blocks,
176-
}
177-
})
178-
}
179-
180143
// countAnthropicTokensBeta calls Anthropic's Count Tokens API for the provided Beta API payload
181144
// and returns the number of input tokens.
182145
func countAnthropicTokensBeta(

pkg/model/provider/anthropic/beta_converter_test.go

Lines changed: 0 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -78,14 +78,6 @@ func TestConvertBetaMessages_MergesConsecutiveToolMessages(t *testing.T) {
7878
require.True(t, ok)
7979
content := contentArray(userMsg2Map)
8080
require.Len(t, content, 2, "User message should have 2 tool_result blocks")
81-
82-
toolResultIDs := collectToolResultIDs(content)
83-
assert.Contains(t, toolResultIDs, "tool_call_1")
84-
assert.Contains(t, toolResultIDs, "tool_call_2")
85-
86-
// Most importantly: validate that the sequence is valid for Anthropic API
87-
err = validateAnthropicSequencingBeta(betaMessages)
88-
require.NoError(t, err, "Converted messages should pass Anthropic sequencing validation")
8981
}
9082

9183
func TestConvertBetaMessages_SingleToolMessage(t *testing.T) {
@@ -123,68 +115,4 @@ func TestConvertBetaMessages_SingleToolMessage(t *testing.T) {
123115
betaMessages, err := testClient().convertBetaMessages(t.Context(), messages)
124116
require.NoError(t, err)
125117
require.Len(t, betaMessages, 4)
126-
127-
// Validate sequence
128-
err = validateAnthropicSequencingBeta(betaMessages)
129-
require.NoError(t, err)
130-
}
131-
132-
func TestConvertBetaMessages_NonConsecutiveToolMessages(t *testing.T) {
133-
// When tool messages are separated by other messages (edge case)
134-
// Each tool message group should be handled independently
135-
messages := []chat.Message{
136-
{
137-
Role: chat.MessageRoleUser,
138-
Content: "First request",
139-
},
140-
{
141-
Role: chat.MessageRoleAssistant,
142-
Content: "",
143-
ToolCalls: []tools.ToolCall{
144-
{
145-
ID: "tool_1",
146-
Type: "function",
147-
Function: tools.FunctionCall{
148-
Name: "test_tool",
149-
Arguments: `{}`,
150-
},
151-
},
152-
},
153-
},
154-
{
155-
Role: chat.MessageRoleTool,
156-
Content: "Tool result 1",
157-
ToolCallID: "tool_1",
158-
},
159-
{
160-
Role: chat.MessageRoleAssistant,
161-
Content: "Intermediate response",
162-
ToolCalls: []tools.ToolCall{
163-
{
164-
ID: "tool_2",
165-
Type: "function",
166-
Function: tools.FunctionCall{
167-
Name: "test_tool",
168-
Arguments: `{}`,
169-
},
170-
},
171-
},
172-
},
173-
{
174-
Role: chat.MessageRoleTool,
175-
Content: "Tool result 2",
176-
ToolCallID: "tool_2",
177-
},
178-
{
179-
Role: chat.MessageRoleAssistant,
180-
Content: "Final response",
181-
},
182-
}
183-
184-
betaMessages, err := testClient().convertBetaMessages(t.Context(), messages)
185-
require.NoError(t, err)
186-
187-
// Validate the entire sequence
188-
err = validateAnthropicSequencingBeta(betaMessages)
189-
require.NoError(t, err, "Messages with non-consecutive tool calls should still validate")
190118
}

pkg/model/provider/anthropic/client.go

Lines changed: 0 additions & 149 deletions
Original file line numberDiff line numberDiff line change
@@ -271,15 +271,6 @@ func (c *Client) CreateChatCompletionStream(
271271
slog.Error("Failed to convert messages for Anthropic request", "error", err)
272272
return nil, err
273273
}
274-
// Preflight validation to ensure tool_use/tool_result sequencing is valid
275-
if err := validateAnthropicSequencing(converted); err != nil {
276-
slog.Warn("Invalid message sequencing for Anthropic detected, attempting self-repair", "error", err)
277-
converted = repairAnthropicSequencing(converted)
278-
if err2 := validateAnthropicSequencing(converted); err2 != nil {
279-
slog.Error("Failed to self-repair Anthropic sequencing", "error", err2)
280-
return nil, err
281-
}
282-
}
283274
if len(converted) == 0 {
284275
return nil, errors.New("no messages to send after conversion: all messages were filtered out")
285276
}
@@ -735,27 +726,6 @@ func (c *Client) FileManager() *FileManager {
735726
return c.fileManager
736727
}
737728

738-
// validateAnthropicSequencing verifies that for every assistant message that includes
739-
// one or more tool_use blocks, the immediately following message is a user message
740-
// that includes tool_result blocks for all those tool_use IDs (grouped into that single message).
741-
func validateAnthropicSequencing(msgs []anthropic.MessageParam) error {
742-
return validateSequencing(msgs)
743-
}
744-
745-
// repairAnthropicSequencing inserts a synthetic user message containing tool_result blocks
746-
// immediately after any assistant message that has tool_use blocks missing a corresponding
747-
// tool_result in the next user message. This is a best-effort local repair to keep the
748-
// conversation valid for Anthropic while preserving original messages, to keep the agent loop running.
749-
func repairAnthropicSequencing(msgs []anthropic.MessageParam) []anthropic.MessageParam {
750-
return repairSequencing(msgs, func(toolUseIDs map[string]struct{}) anthropic.MessageParam {
751-
blocks := make([]anthropic.ContentBlockParamUnion, 0, len(toolUseIDs))
752-
for id := range toolUseIDs {
753-
blocks = append(blocks, anthropic.NewToolResultBlock(id, "(tool execution failed)", false))
754-
}
755-
return anthropic.NewUserMessage(blocks...)
756-
})
757-
}
758-
759729
// marshalToMap is a helper that converts any value to a map[string]any via JSON marshaling.
760730
// This is used to inspect SDK union types without depending on their internal structure.
761731
// It's shared by both standard and Beta API validation/repair code.
@@ -780,125 +750,6 @@ func contentArray(m map[string]any) []any {
780750
return nil
781751
}
782752

783-
// validateSequencing generically validates that every assistant message with tool_use blocks
784-
// is immediately followed by a user message with corresponding tool_result blocks.
785-
// It works on both standard (MessageParam) and Beta (BetaMessageParam) types by
786-
// marshaling to map[string]any for inspection.
787-
func validateSequencing[T any](msgs []T) error {
788-
for i := range msgs {
789-
m, ok := marshalToMap(msgs[i])
790-
if !ok || m["role"] != "assistant" {
791-
continue
792-
}
793-
794-
toolUseIDs := collectToolUseIDs(contentArray(m))
795-
if len(toolUseIDs) == 0 {
796-
continue
797-
}
798-
799-
if i+1 >= len(msgs) {
800-
slog.Warn("Anthropic sequencing invalid: assistant tool_use present but no next user tool_result message", "assistant_index", i)
801-
return errors.New("assistant tool_use present but no subsequent user message with tool_result blocks")
802-
}
803-
804-
next, ok := marshalToMap(msgs[i+1])
805-
if !ok || next["role"] != "user" {
806-
slog.Warn("Anthropic sequencing invalid: next message after assistant tool_use is not user", "assistant_index", i, "next_role", next["role"])
807-
return errors.New("assistant tool_use must be followed by a user message containing corresponding tool_result blocks")
808-
}
809-
810-
toolResultIDs := collectToolResultIDs(contentArray(next))
811-
missing := differenceIDs(toolUseIDs, toolResultIDs)
812-
if len(missing) > 0 {
813-
slog.Warn("Anthropic sequencing invalid: missing tool_result for tool_use id in next user message", "assistant_index", i, "tool_use_id", missing[0], "missing_count", len(missing))
814-
return fmt.Errorf("missing tool_result for tool_use id %s in the next user message", missing[0])
815-
}
816-
}
817-
return nil
818-
}
819-
820-
// repairSequencing generically inserts a synthetic user message after any assistant
821-
// tool_use message that is missing corresponding tool_result blocks. The makeSynthetic
822-
// callback builds the appropriate user message type for the remaining tool_use IDs.
823-
func repairSequencing[T any](msgs []T, makeSynthetic func(toolUseIDs map[string]struct{}) T) []T {
824-
if len(msgs) == 0 {
825-
return msgs
826-
}
827-
repaired := make([]T, 0, len(msgs)+2)
828-
for i := range msgs {
829-
repaired = append(repaired, msgs[i])
830-
831-
m, ok := marshalToMap(msgs[i])
832-
if !ok || m["role"] != "assistant" {
833-
continue
834-
}
835-
836-
toolUseIDs := collectToolUseIDs(contentArray(m))
837-
if len(toolUseIDs) == 0 {
838-
continue
839-
}
840-
841-
// Remove any IDs that already have results in the next user message
842-
if i+1 < len(msgs) {
843-
if next, ok := marshalToMap(msgs[i+1]); ok && next["role"] == "user" {
844-
toolResultIDs := collectToolResultIDs(contentArray(next))
845-
for id := range toolResultIDs {
846-
delete(toolUseIDs, id)
847-
}
848-
}
849-
}
850-
851-
if len(toolUseIDs) > 0 {
852-
slog.Debug("Inserting synthetic user message for missing tool_results",
853-
"assistant_index", i,
854-
"missing_count", len(toolUseIDs))
855-
repaired = append(repaired, makeSynthetic(toolUseIDs))
856-
}
857-
}
858-
return repaired
859-
}
860-
861-
func collectToolUseIDs(content []any) map[string]struct{} {
862-
ids := make(map[string]struct{})
863-
for _, c := range content {
864-
if cb, ok := c.(map[string]any); ok {
865-
if t, _ := cb["type"].(string); t == "tool_use" {
866-
if id, _ := cb["id"].(string); id != "" {
867-
ids[id] = struct{}{}
868-
}
869-
}
870-
}
871-
}
872-
return ids
873-
}
874-
875-
func collectToolResultIDs(content []any) map[string]struct{} {
876-
ids := make(map[string]struct{})
877-
for _, c := range content {
878-
if cb, ok := c.(map[string]any); ok {
879-
if t, _ := cb["type"].(string); t == "tool_result" {
880-
if id, _ := cb["tool_use_id"].(string); id != "" {
881-
ids[id] = struct{}{}
882-
}
883-
}
884-
}
885-
}
886-
return ids
887-
}
888-
889-
func differenceIDs(a, b map[string]struct{}) []string {
890-
if len(a) == 0 {
891-
return nil
892-
}
893-
var missing []string
894-
for id := range a {
895-
if _, ok := b[id]; !ok {
896-
missing = append(missing, id)
897-
}
898-
}
899-
return missing
900-
}
901-
902753
// validThinkingTokens validates that the token budget is within the
903754
// acceptable range for Anthropic (>= 1024 and < maxTokens).
904755
// Returns (tokens, true) if valid, or (0, false) with a warning log if not.

pkg/model/provider/anthropic/client_test.go

Lines changed: 0 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -235,80 +235,6 @@ func TestSystemMessages_InterspersedExtractedAndExcluded(t *testing.T) {
235235
}
236236
}
237237

238-
func TestSequencingRepair_Standard(t *testing.T) {
239-
msgs := []chat.Message{
240-
{Role: chat.MessageRoleUser, Content: "start"},
241-
{
242-
Role: chat.MessageRoleAssistant,
243-
ToolCalls: []tools.ToolCall{
244-
{ID: "tool-1", Function: tools.FunctionCall{Name: "do_thing", Arguments: "{}"}},
245-
},
246-
},
247-
// Intentionally missing the user/tool_result message here
248-
{Role: chat.MessageRoleUser, Content: "continue"},
249-
}
250-
251-
converted, err := testClient().convertMessages(t.Context(), msgs)
252-
require.NoError(t, err)
253-
err = validateAnthropicSequencing(converted)
254-
require.Error(t, err)
255-
256-
repaired := repairAnthropicSequencing(converted)
257-
err = validateAnthropicSequencing(repaired)
258-
require.NoError(t, err)
259-
}
260-
261-
func TestSequencingRepair_Beta(t *testing.T) {
262-
msgs := []chat.Message{
263-
{Role: chat.MessageRoleUser, Content: "start"},
264-
{
265-
Role: chat.MessageRoleAssistant,
266-
ToolCalls: []tools.ToolCall{
267-
{ID: "tool-1", Function: tools.FunctionCall{Name: "do_thing", Arguments: "{}"}},
268-
},
269-
},
270-
// Intentionally missing the user/tool_result message here
271-
{Role: chat.MessageRoleUser, Content: "continue"},
272-
}
273-
274-
converted, err := testClient().convertBetaMessages(t.Context(), msgs)
275-
require.NoError(t, err)
276-
err = validateAnthropicSequencingBeta(converted)
277-
require.Error(t, err)
278-
279-
repaired := repairAnthropicSequencingBeta(converted)
280-
err = validateAnthropicSequencingBeta(repaired)
281-
require.NoError(t, err)
282-
}
283-
284-
func TestConvertMessages_DropOrphanToolResults_NoPrecedingToolUse(t *testing.T) {
285-
msgs := []chat.Message{
286-
{Role: chat.MessageRoleUser, Content: "start"},
287-
// Orphan tool result (no assistant tool_use immediately before)
288-
{Role: chat.MessageRoleTool, ToolCallID: "tool-1", Content: "result-1"},
289-
{Role: chat.MessageRoleUser, Content: "continue"},
290-
}
291-
292-
converted, err := testClient().convertMessages(t.Context(), msgs)
293-
require.NoError(t, err)
294-
// Expect only the two user text messages to appear
295-
require.Len(t, converted, 2)
296-
297-
// Ensure none of the converted messages contain tool_result blocks
298-
for i := range converted {
299-
b, err := json.Marshal(converted[i])
300-
require.NoError(t, err)
301-
var m map[string]any
302-
require.NoError(t, json.Unmarshal(b, &m))
303-
content, _ := m["content"].([]any)
304-
for _, c := range content {
305-
if cb, ok := c.(map[string]any); ok {
306-
assert.NotEqual(t, "tool_result", cb["type"], "unexpected orphan tool_result included in payload")
307-
}
308-
}
309-
}
310-
}
311-
312238
func TestConvertMessages_GroupToolResults_AfterAssistantToolUse(t *testing.T) {
313239
msgs := []chat.Message{
314240
{Role: chat.MessageRoleUser, Content: "start"},
@@ -329,9 +255,6 @@ func TestConvertMessages_GroupToolResults_AfterAssistantToolUse(t *testing.T) {
329255
// Expect: user(start), assistant(tool_use), user(grouped tool_result), user(ok)
330256
require.Len(t, converted, 4)
331257

332-
// Validate sequencing is acceptable to Anthropic
333-
require.NoError(t, validateAnthropicSequencing(converted))
334-
335258
b, err := json.Marshal(converted[2])
336259
require.NoError(t, err)
337260
var m map[string]any

0 commit comments

Comments
 (0)