From 560a4456a72869ea2e69c295528cde033eaa2fd4 Mon Sep 17 00:00:00 2001 From: Rahul Bansal Date: Sat, 21 Feb 2026 11:56:31 +0530 Subject: [PATCH] fix: sanitizer drops 2nd+ tool_result in multi-call batches; add /doctor command The sanitizeHistoryForProvider check for tool messages only looked at the immediately preceding message, requiring it to be an assistant. But when an assistant makes 2+ tool calls, the 2nd tool_result follows another tool message. Fixed to look back past consecutive tool messages. Also adds /doctor slash command for in-session diagnosis and repair. --- pkg/agent/context.go | 17 +++++++- pkg/agent/loop.go | 64 ++++++++++++++++++++++++++++++ pkg/agent/sanitize_test.go | 80 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 159 insertions(+), 2 deletions(-) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 661fc0a04..5fee00a06 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -232,8 +232,21 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message logger.DebugCF("agent", "Dropping orphaned leading tool message", map[string]any{}) continue } - last := sanitized[len(sanitized)-1] - if last.Role != "assistant" || len(last.ToolCalls) == 0 { + // Look back past other tool messages to find the assistant that + // initiated these tool calls. An assistant with 2+ tool_calls + // produces consecutive tool results — the 2nd+ tool results have + // another tool message as their predecessor, not the assistant. + hasMatchingAssistant := false + for i := len(sanitized) - 1; i >= 0; i-- { + if sanitized[i].Role == "tool" { + continue // skip past earlier tool results in same batch + } + if sanitized[i].Role == "assistant" && len(sanitized[i].ToolCalls) > 0 { + hasMatchingAssistant = true + } + break // stop at first non-tool message + } + if !hasMatchingAssistant { logger.DebugCF("agent", "Dropping orphaned tool message", map[string]any{}) continue } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 0374f938d..8b38636d5 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1083,6 +1083,7 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage) /help Show this help message /new Start a new conversation /status Show current session info + /doctor Diagnose and repair current session /show model Show current model /show channel Show current channel /show agents Show registered agents @@ -1208,6 +1209,69 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage) Messages: %d in current session Max iterations: %d`, agent.Model, agent.ID, msg.Channel, len(history), agent.MaxIterations), true + case "/doctor": + // Diagnose and repair the current session in-place + route := al.registry.ResolveRoute(routing.RouteInput{ + Channel: msg.Channel, + AccountID: msg.Metadata["account_id"], + Peer: extractPeer(msg), + GuildID: msg.Metadata["guild_id"], + TeamID: msg.Metadata["team_id"], + }) + + agent, ok := al.registry.GetAgent(route.AgentID) + if !ok { + agent = al.registry.GetDefaultAgent() + } + if agent == nil { + return "No agent configured", true + } + + sessionKey := route.SessionKey + history := agent.Sessions.GetHistory(sessionKey) + if len(history) == 0 { + return "Session is empty — nothing to repair.", true + } + + repaired := repairOrphanedToolPairs(history) + + injected := len(repaired) - len(history) + if injected == 0 { + // Check if any messages were dropped (orphan tool_results removed) + // by comparing lengths — repairOrphanedToolPairs may drop AND inject + dropped := 0 + toolCallIDs := map[string]bool{} + for _, m := range history { + for _, tc := range m.ToolCalls { + if tc.ID != "" { + toolCallIDs[tc.ID] = true + } + } + } + for _, m := range history { + if m.Role == "tool" && m.ToolCallID != "" && !toolCallIDs[m.ToolCallID] { + dropped++ + } + } + if dropped == 0 { + return fmt.Sprintf("Session OK — %d messages, no issues found.", len(history)), true + } + } + + agent.Sessions.SetHistory(sessionKey, repaired) + agent.Sessions.Save(sessionKey) + + var parts []string + if len(repaired) > len(history) { + parts = append(parts, fmt.Sprintf("injected %d synthetic tool result(s)", len(repaired)-len(history))) + } + if len(repaired) < len(history) { + parts = append(parts, fmt.Sprintf("dropped %d orphaned tool result(s)", len(history)-len(repaired))) + } + + return fmt.Sprintf("Session repaired: %s. Messages: %d -> %d.", + strings.Join(parts, ", "), len(history), len(repaired)), true + case "/switch": if len(args) < 3 || args[1] != "to" { return "Usage: /switch [model|channel] to ", true diff --git a/pkg/agent/sanitize_test.go b/pkg/agent/sanitize_test.go index 6a848b7d9..e1ed2121c 100644 --- a/pkg/agent/sanitize_test.go +++ b/pkg/agent/sanitize_test.go @@ -97,6 +97,86 @@ func TestSanitizeHistoryForProvider_OrphanToolUseRepaired(t *testing.T) { } } +func TestSanitizeHistoryForProvider_MultiToolCallsPreserved(t *testing.T) { + // Reproduces the actual bug: assistant with 2 tool_calls, both have results, + // but the sanitizer was dropping the second tool_result because it only checked + // if the immediately preceding message was an assistant (not another tool). + history := []providers.Message{ + {Role: "user", Content: "do something"}, + { + Role: "assistant", Content: "checking two things", + ToolCalls: []providers.ToolCall{ + {ID: "call_A", Name: "read_file"}, + {ID: "call_B", Name: "read_file"}, + }, + }, + {Role: "tool", Content: "file contents A", ToolCallID: "call_A"}, + {Role: "tool", Content: "file contents B", ToolCallID: "call_B"}, + {Role: "assistant", Content: "here are the results"}, + } + sanitized := sanitizeHistoryForProvider(history) + + if len(sanitized) != 5 { + t.Fatalf("expected 5 messages (all preserved), got %d", len(sanitized)) + } + // Verify both tool results are present + if sanitized[2].Role != "tool" || sanitized[2].ToolCallID != "call_A" { + t.Errorf("message[2]: expected tool result for call_A, got role=%q id=%q", + sanitized[2].Role, sanitized[2].ToolCallID) + } + if sanitized[3].Role != "tool" || sanitized[3].ToolCallID != "call_B" { + t.Errorf("message[3]: expected tool result for call_B, got role=%q id=%q", + sanitized[3].Role, sanitized[3].ToolCallID) + } +} + +func TestSanitizeHistoryForProvider_RealSessionRegression(t *testing.T) { + // Reproduces the exact pattern from the user's corrupt session: + // assistant(2 calls) -> tool -> tool -> assistant -> user -> + // assistant(2 calls) -> tool -> tool -> assistant(1 call) -> tool -> ... + history := []providers.Message{ + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "tc1", Name: "read_file"}, {ID: "tc2", Name: "read_file"}, + }}, + {Role: "tool", Content: "file1", ToolCallID: "tc1"}, + {Role: "tool", Content: "file2", ToolCallID: "tc2"}, + {Role: "assistant", Content: "summary"}, + {Role: "user", Content: "do it"}, + {Role: "assistant", Content: "checking", ToolCalls: []providers.ToolCall{ + {ID: "tc3", Name: "list_dir"}, {ID: "tc4", Name: "read_file"}, + }}, + {Role: "tool", Content: "denied", ToolCallID: "tc3"}, + {Role: "tool", Content: "denied", ToolCallID: "tc4"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "tc5", Name: "exec"}, + }}, + {Role: "tool", Content: "output", ToolCallID: "tc5"}, + } + sanitized := sanitizeHistoryForProvider(history) + + // Count tool_use and tool_result IDs — must be balanced + toolCallIDs := map[string]bool{} + toolResultIDs := map[string]bool{} + for _, m := range sanitized { + for _, tc := range m.ToolCalls { + toolCallIDs[tc.ID] = true + } + if m.Role == "tool" && m.ToolCallID != "" { + toolResultIDs[m.ToolCallID] = true + } + } + for id := range toolCallIDs { + if !toolResultIDs[id] { + t.Errorf("orphaned tool_call %q — no matching tool_result after sanitize", id) + } + } + for id := range toolResultIDs { + if !toolCallIDs[id] { + t.Errorf("orphaned tool_result %q — no matching tool_call after sanitize", id) + } + } +} + func TestRepairOrphanedToolPairs_EmptyInput(t *testing.T) { repaired := repairOrphanedToolPairs(nil) if len(repaired) != 0 {