From 46a5dde93abf6fe768bd9fdae4eb08a143f3d053 Mon Sep 17 00:00:00 2001 From: Rahul Bansal Date: Sat, 21 Feb 2026 11:30:41 +0530 Subject: [PATCH] feat: add repairOrphanedToolPairs utility for fixing orphaned tool_use/tool_result pairs --- pkg/agent/sanitize.go | 95 ++++++++++++++++++++++++++++++++++++++ pkg/agent/sanitize_test.go | 90 ++++++++++++++++++++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 pkg/agent/sanitize.go create mode 100644 pkg/agent/sanitize_test.go diff --git a/pkg/agent/sanitize.go b/pkg/agent/sanitize.go new file mode 100644 index 000000000..fdcdc9f6b --- /dev/null +++ b/pkg/agent/sanitize.go @@ -0,0 +1,95 @@ +package agent + +import ( + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/providers" +) + +// repairOrphanedToolPairs scans messages and: +// 1. Injects synthetic tool_result messages for any tool_use that lacks a matching result +// 2. Drops tool_result messages that lack a preceding tool_use +func repairOrphanedToolPairs(msgs []providers.Message) []providers.Message { + if len(msgs) == 0 { + return msgs + } + + // 1. Collect all tool_call IDs from assistant messages + toolCallIDs := map[string]bool{} + for _, m := range msgs { + if m.Role == "assistant" { + for _, tc := range m.ToolCalls { + if tc.ID != "" { + toolCallIDs[tc.ID] = true + } + } + } + } + + // 2. Drop orphaned tool_results (no matching tool_call) + filtered := make([]providers.Message, 0, len(msgs)) + for i, m := range msgs { + if m.Role == "tool" && m.ToolCallID != "" && !toolCallIDs[m.ToolCallID] { + logger.DebugCF("agent", "Dropping orphaned tool_result", map[string]any{ + "tool_call_id": m.ToolCallID, + "index": i, + }) + continue + } + filtered = append(filtered, m) + } + + // 3. Collect existing tool_result IDs + resultIDs := map[string]bool{} + for _, m := range filtered { + if m.Role == "tool" && m.ToolCallID != "" { + resultIDs[m.ToolCallID] = true + } + } + + // 4. Build repaired slice, injecting synthetic results for orphaned tool_calls. + // Track pending tool_calls from each assistant message and flush missing + // results after the last consecutive tool result (or at end of input). + repaired := make([]providers.Message, 0, len(filtered)) + var pendingCalls []providers.ToolCall + + flushPending := func() { + for _, tc := range pendingCalls { + if tc.ID == "" || resultIDs[tc.ID] { + continue + } + name := tc.Name + if tc.Function != nil { + name = tc.Function.Name + } + logger.DebugCF("agent", "Injecting synthetic tool_result for orphaned tool_call", map[string]any{ + "tool_call_id": tc.ID, + "tool_name": name, + }) + repaired = append(repaired, providers.Message{ + Role: "tool", + ToolCallID: tc.ID, + Content: "[tool result unavailable — session history was compressed]", + }) + resultIDs[tc.ID] = true + } + pendingCalls = nil + } + + for _, m := range filtered { + // When we hit a non-tool message and have pending calls, flush synthetics + if m.Role != "tool" && pendingCalls != nil { + flushPending() + } + + repaired = append(repaired, m) + + if m.Role == "assistant" && len(m.ToolCalls) > 0 { + pendingCalls = m.ToolCalls + } + } + + // Flush any remaining pending calls at the end of input + flushPending() + + return repaired +} diff --git a/pkg/agent/sanitize_test.go b/pkg/agent/sanitize_test.go new file mode 100644 index 000000000..151c35936 --- /dev/null +++ b/pkg/agent/sanitize_test.go @@ -0,0 +1,90 @@ +package agent + +import ( + "testing" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +func TestRepairOrphanedToolPairs_NoOrphans(t *testing.T) { + msgs := []providers.Message{ + {Role: "user", Content: "hello"}, + { + Role: "assistant", Content: "let me check", + ToolCalls: []providers.ToolCall{{ID: "call_1", Name: "exec"}}, + }, + {Role: "tool", Content: "output", ToolCallID: "call_1"}, + {Role: "assistant", Content: "done"}, + } + repaired := repairOrphanedToolPairs(msgs) + if len(repaired) != 4 { + t.Errorf("expected 4 messages, got %d", len(repaired)) + } +} + +func TestRepairOrphanedToolPairs_OrphanToolUseAtEnd(t *testing.T) { + msgs := []providers.Message{ + {Role: "user", Content: "hello"}, + { + Role: "assistant", Content: "let me check", + ToolCalls: []providers.ToolCall{{ID: "call_1", Name: "exec"}}, + }, + } + repaired := repairOrphanedToolPairs(msgs) + if len(repaired) != 3 { + t.Fatalf("expected 3 messages, got %d", len(repaired)) + } + if repaired[2].Role != "tool" { + t.Errorf("expected injected tool message, got role=%q", repaired[2].Role) + } + if repaired[2].ToolCallID != "call_1" { + t.Errorf("expected ToolCallID=call_1, got %q", repaired[2].ToolCallID) + } +} + +func TestRepairOrphanedToolPairs_MultipleToolCallsPartialResults(t *testing.T) { + msgs := []providers.Message{ + {Role: "user", Content: "hello"}, + { + Role: "assistant", Content: "checking two things", + ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "exec"}, + {ID: "call_2", Name: "web_fetch"}, + }, + }, + {Role: "tool", Content: "output1", ToolCallID: "call_1"}, + } + repaired := repairOrphanedToolPairs(msgs) + if len(repaired) != 4 { + t.Fatalf("expected 4 messages, got %d", len(repaired)) + } + if repaired[3].ToolCallID != "call_2" { + t.Errorf("expected injected result for call_2, got ToolCallID=%q", repaired[3].ToolCallID) + } +} + +func TestRepairOrphanedToolPairs_OrphanToolResultDropped(t *testing.T) { + msgs := []providers.Message{ + {Role: "user", Content: "hello"}, + {Role: "tool", Content: "orphaned output", ToolCallID: "call_orphan"}, + {Role: "assistant", Content: "hi"}, + } + repaired := repairOrphanedToolPairs(msgs) + if len(repaired) != 2 { + t.Fatalf("expected 2 messages (user + assistant), got %d", len(repaired)) + } + if repaired[0].Role != "user" || repaired[1].Role != "assistant" { + t.Errorf("unexpected message roles: %q, %q", repaired[0].Role, repaired[1].Role) + } +} + +func TestRepairOrphanedToolPairs_EmptyInput(t *testing.T) { + repaired := repairOrphanedToolPairs(nil) + if len(repaired) != 0 { + t.Errorf("expected 0 messages for nil input, got %d", len(repaired)) + } + repaired = repairOrphanedToolPairs([]providers.Message{}) + if len(repaired) != 0 { + t.Errorf("expected 0 messages for empty input, got %d", len(repaired)) + } +}