fix cancell

4 files changed, 42 insertions, 1 deletions

diff --git a/internal/llm/provider/anthropic.go b/internal/llm/provider/anthropic.go
index 88b874806..93c4308ad 100644
--- a/internal/llm/provider/anthropic.go
+++ b/internal/llm/provider/anthropic.go
@@ -92,6 +92,7 @@ func NewAnthropicProvider(opts ...AnthropicOption) (Provider, error) {
 }
 
 func (a *anthropicProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
+	messages = cleanupMessages(messages)
 	anthropicMessages := a.convertToAnthropicMessages(messages)
 	anthropicTools := a.convertToAnthropicTools(tools)
 
@@ -135,6 +136,7 @@ func (a *anthropicProvider) SendMessages(ctx context.Context, messages []message
 }
 
 func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error) {
+	messages = cleanupMessages(messages)
 	anthropicMessages := a.convertToAnthropicMessages(messages)
 	anthropicTools := a.convertToAnthropicTools(tools)
 
diff --git a/internal/llm/provider/gemini.go b/internal/llm/provider/gemini.go
index 34bfc5283..2d1db2b64 100644
--- a/internal/llm/provider/gemini.go
+++ b/internal/llm/provider/gemini.go
@@ -154,6 +154,7 @@ func (p *geminiProvider) extractTokenUsage(resp *genai.GenerateContentResponse)
 }
 
 func (p *geminiProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
+	messages = cleanupMessages(messages)
 	model := p.client.GenerativeModel(p.model.APIModel)
 	model.SetMaxOutputTokens(p.maxTokens)
 
@@ -206,6 +207,7 @@ func (p *geminiProvider) SendMessages(ctx context.Context, messages []message.Me
 }
 
 func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error) {
+	messages = cleanupMessages(messages)
 	model := p.client.GenerativeModel(p.model.APIModel)
 	model.SetMaxOutputTokens(p.maxTokens)
 
diff --git a/internal/llm/provider/openai.go b/internal/llm/provider/openai.go
index c8e04d5ee..dbfde3fa8 100644
--- a/internal/llm/provider/openai.go
+++ b/internal/llm/provider/openai.go
@@ -163,6 +163,7 @@ func (p *openaiProvider) extractTokenUsage(usage openai.CompletionUsage) TokenUs
 }
 
 func (p *openaiProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
+	messages = cleanupMessages(messages)
 	chatMessages := p.convertToOpenAIMessages(messages)
 	openaiTools := p.convertToOpenAITools(tools)
 
@@ -206,6 +207,7 @@ func (p *openaiProvider) SendMessages(ctx context.Context, messages []message.Me
 }
 
 func (p *openaiProvider) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error) {
+	messages = cleanupMessages(messages)
 	chatMessages := p.convertToOpenAIMessages(messages)
 	openaiTools := p.convertToOpenAITools(tools)
 
@@ -276,4 +278,3 @@ func (p *openaiProvider) StreamResponse(ctx context.Context, messages []message.
 
 	return eventChan, nil
 }
-
diff --git a/internal/llm/provider/provider.go b/internal/llm/provider/provider.go
index cd3a7aa05..938a8c0ad 100644
--- a/internal/llm/provider/provider.go
+++ b/internal/llm/provider/provider.go
@@ -52,3 +52,39 @@ type Provider interface {
 
 	StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error)
 }
+
+func cleanupMessages(messages []message.Message) []message.Message {
+	// First pass: filter out canceled messages
+	var cleanedMessages []message.Message
+	for _, msg := range messages {
+		if msg.FinishReason() != "canceled" {
+			cleanedMessages = append(cleanedMessages, msg)
+		}
+	}
+
+	// Second pass: filter out tool messages without a corresponding tool call
+	var result []message.Message
+	toolMessageIDs := make(map[string]bool)
+
+	for _, msg := range cleanedMessages {
+		if msg.Role == message.Assistant {
+			for _, toolCall := range msg.ToolCalls() {
+				toolMessageIDs[toolCall.ID] = true // Mark as referenced
+			}
+		}
+	}
+
+	// Keep only messages that aren't unreferenced tool messages
+	for _, msg := range cleanedMessages {
+		if msg.Role == message.Tool {
+			for _, toolCall := range msg.ToolResults() {
+				if referenced, exists := toolMessageIDs[toolCall.ToolCallID]; exists && referenced {
+					result = append(result, msg)
+				}
+			}
+		} else {
+			result = append(result, msg)
+		}
+	}
+	return result
+}