fix: resolve Gemini turn-ordering 400 error from corrupted session history

Addresses two bugs causing invalid message ordering:

1. Atomic message saving: moves all session saves to after successful LLM
   completion, preventing orphaned user messages when LLM calls fail.

2. Comprehensive history sanitization: replaces partial tool-message stripping
   with sanitizeHistory() that handles leading non-user messages, consecutive
   user messages, orphaned tool results, and trailing incomplete tool-call
   sequences.

3. Smart truncation: TruncateHistory now scans forward to find the nearest
   user message boundary, preventing summarization from creating invalid
   mid-sequence starting points.

pkg/agent/context.go

@@ -189,16 +189,8 @@ func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary str
 		systemPrompt += "\n\n## Summary of Previous Conversation\n\n" + summary
 	}
 
-	//This fix prevents the session memory from LLM failure due to elimination of toolu_IDs required from LLM
+	// Sanitize history to ensure valid turn ordering for all providers
-	// --- INICIO DEL FIX ---
+	history = sanitizeHistory(history)
-	//Diegox-17
-	for len(history) > 0 && (history[0].Role == "tool") {
-		logger.DebugCF("agent", "Removing orphaned tool message from history to prevent LLM error",
-			map[string]interface{}{"role": history[0].Role})
-		history = history[1:]
-	}
-	//Diegox-17
-	// --- FIN DEL FIX ---
 
 	messages = append(messages, providers.Message{
 		Role:    "system",
@@ -207,10 +199,23 @@ func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary str
 
 	messages = append(messages, history...)
 
-	messages = append(messages, providers.Message{
+	userMsg := providers.Message{
 		Role:    "user",
 		Content: currentMessage,
+	}
+	if len(media) > 0 {
+		parts := []providers.ContentPart{
+			{Type: "text", Text: currentMessage},
+		}
+		for _, url := range media {
+			parts = append(parts, providers.ContentPart{
+				Type:     "image_url",
+				ImageURL: &providers.ImageURL{URL: url},
 			})
+		}
+		userMsg.ContentParts = parts
+	}
+	messages = append(messages, userMsg)
 
 	return messages
 }
@@ -266,3 +271,106 @@ func (cb *ContextBuilder) GetSkillsInfo() map[string]interface{} {
 		"names":     skillNames,
 	}
 }
+
+// sanitizeHistory ensures valid turn ordering for all LLM providers.
+// It handles corruption from truncation, failed LLM calls, or race conditions.
+func sanitizeHistory(history []providers.Message) []providers.Message {
+	if len(history) == 0 {
+		return history
+	}
+
+	// 1. Remove leading non-user messages (tool results, assistant with tool_calls)
+	for len(history) > 0 && history[0].Role != "user" {
+		logger.DebugCF("agent", "Removing leading non-user message from history",
+			map[string]interface{}{"role": history[0].Role})
+		history = history[1:]
+	}
+
+	if len(history) == 0 {
+		return history
+	}
+
+	// 2. Walk through and build a valid sequence
+	sanitized := make([]providers.Message, 0, len(history))
+	for i := 0; i < len(history); i++ {
+		msg := history[i]
+
+		// Skip consecutive user messages (keep the last one in a run)
+		if msg.Role == "user" && i+1 < len(history) && history[i+1].Role == "user" {
+			logger.DebugCF("agent", "Removing duplicate consecutive user message",
+				map[string]interface{}{"index": i})
+			continue
+		}
+
+		// Skip tool messages that don't follow an assistant message with tool_calls
+		if msg.Role == "tool" {
+			if len(sanitized) == 0 || sanitized[len(sanitized)-1].Role != "assistant" || len(sanitized[len(sanitized)-1].ToolCalls) == 0 {
+				// Check if the preceding message (allowing for other tool messages) was an assistant with tool_calls
+				hasMatchingAssistant := false
+				for j := len(sanitized) - 1; j >= 0; j-- {
+					if sanitized[j].Role == "tool" {
+						continue
+					}
+					if sanitized[j].Role == "assistant" && len(sanitized[j].ToolCalls) > 0 {
+						hasMatchingAssistant = true
+					}
+					break
+				}
+				if !hasMatchingAssistant {
+					logger.DebugCF("agent", "Removing orphaned tool message from history",
+						map[string]interface{}{"index": i, "tool_call_id": msg.ToolCallID})
+					continue
+				}
+			}
+		}
+
+		sanitized = append(sanitized, msg)
+	}
+
+	// 3. Remove trailing incomplete tool-call sequences
+	// (assistant with tool_calls at the end without all corresponding tool results)
+	for len(sanitized) > 0 {
+		last := sanitized[len(sanitized)-1]
+		if last.Role == "assistant" && len(last.ToolCalls) > 0 {
+			logger.DebugCF("agent", "Removing trailing assistant with unanswered tool_calls",
+				map[string]interface{}{"tool_calls": len(last.ToolCalls)})
+			sanitized = sanitized[:len(sanitized)-1]
+			continue
+		}
+		// Also check if we end with tool results but the preceding assistant
+		// doesn't have all its tool_calls answered
+		if last.Role == "tool" {
+			// Find the preceding assistant message
+			assistantIdx := -1
+			for j := len(sanitized) - 2; j >= 0; j-- {
+				if sanitized[j].Role == "assistant" && len(sanitized[j].ToolCalls) > 0 {
+					assistantIdx = j
+					break
+				}
+				if sanitized[j].Role != "tool" {
+					break
+				}
+			}
+			if assistantIdx >= 0 {
+				// Count tool results after the assistant
+				expectedCount := len(sanitized[assistantIdx].ToolCalls)
+				actualCount := 0
+				for j := assistantIdx + 1; j < len(sanitized); j++ {
+					if sanitized[j].Role == "tool" {
+						actualCount++
+					}
+				}
+				if actualCount < expectedCount {
+					// Incomplete sequence — remove the assistant and all its tool results
+					logger.DebugCF("agent", "Removing trailing incomplete tool-call sequence",
+						map[string]interface{}{"expected": expectedCount, "actual": actualCount})
+					sanitized = sanitized[:assistantIdx]
+					continue
+				}
+			}
+		}
+		break
+	}
+
+	return sanitized
+}

pkg/agent/context_test.go

@@ -0,0 +1,151 @@
+package agent
+
+import (
+	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/providers"
+)
+
+func TestSanitizeHistory_LeadingToolMessages(t *testing.T) {
+	history := []providers.Message{
+		{Role: "tool", Content: "orphaned result", ToolCallID: "call_1"},
+		{Role: "tool", Content: "orphaned result 2", ToolCallID: "call_2"},
+		{Role: "user", Content: "hello"},
+		{Role: "assistant", Content: "hi"},
+	}
+
+	result := sanitizeHistory(history)
+
+	if len(result) != 2 {
+		t.Fatalf("expected 2 messages, got %d", len(result))
+	}
+	if result[0].Role != "user" {
+		t.Errorf("expected first message to be user, got %s", result[0].Role)
+	}
+}
+
+func TestSanitizeHistory_LeadingAssistantWithToolCalls(t *testing.T) {
+	history := []providers.Message{
+		{
+			Role: "assistant",
+			ToolCalls: []providers.ToolCall{
+				{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "test"}},
+			},
+		},
+		{Role: "tool", Content: "result", ToolCallID: "call_1"},
+		{Role: "user", Content: "hello"},
+		{Role: "assistant", Content: "hi"},
+	}
+
+	result := sanitizeHistory(history)
+
+	if len(result) != 2 {
+		t.Fatalf("expected 2 messages, got %d", len(result))
+	}
+	if result[0].Role != "user" {
+		t.Errorf("expected first message to be user, got %s", result[0].Role)
+	}
+}
+
+func TestSanitizeHistory_ConsecutiveUsers(t *testing.T) {
+	history := []providers.Message{
+		{Role: "user", Content: "first"},
+		{Role: "user", Content: "second"},
+		{Role: "user", Content: "third"},
+		{Role: "assistant", Content: "response"},
+	}
+
+	result := sanitizeHistory(history)
+
+	if len(result) != 2 {
+		t.Fatalf("expected 2 messages, got %d", len(result))
+	}
+	if result[0].Content != "third" {
+		t.Errorf("expected last user message 'third', got %q", result[0].Content)
+	}
+}
+
+func TestSanitizeHistory_ValidHistory(t *testing.T) {
+	history := []providers.Message{
+		{Role: "user", Content: "hello"},
+		{
+			Role: "assistant",
+			ToolCalls: []providers.ToolCall{
+				{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "test"}},
+			},
+		},
+		{Role: "tool", Content: "result", ToolCallID: "call_1"},
+		{Role: "assistant", Content: "done"},
+		{Role: "user", Content: "thanks"},
+		{Role: "assistant", Content: "welcome"},
+	}
+
+	result := sanitizeHistory(history)
+
+	if len(result) != len(history) {
+		t.Fatalf("expected %d messages (unchanged), got %d", len(history), len(result))
+	}
+	for i := range result {
+		if result[i].Role != history[i].Role {
+			t.Errorf("message %d: expected role %s, got %s", i, history[i].Role, result[i].Role)
+		}
+	}
+}
+
+func TestSanitizeHistory_TrailingOrphanedToolCalls(t *testing.T) {
+	history := []providers.Message{
+		{Role: "user", Content: "hello"},
+		{Role: "assistant", Content: "hi"},
+		{Role: "user", Content: "run tools"},
+		{
+			Role: "assistant",
+			ToolCalls: []providers.ToolCall{
+				{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "tool1"}},
+				{ID: "call_2", Type: "function", Function: &providers.FunctionCall{Name: "tool2"}},
+			},
+		},
+		// Only one tool result for two tool calls — incomplete
+		{Role: "tool", Content: "result1", ToolCallID: "call_1"},
+	}
+
+	result := sanitizeHistory(history)
+
+	// Should remove the incomplete tool-call sequence (assistant + partial tool results)
+	if len(result) != 3 {
+		t.Fatalf("expected 3 messages, got %d: %+v", len(result), result)
+	}
+	if result[2].Role != "user" {
+		t.Errorf("expected last message to be user, got %s", result[2].Role)
+	}
+}
+
+func TestSanitizeHistory_Empty(t *testing.T) {
+	result := sanitizeHistory(nil)
+	if len(result) != 0 {
+		t.Fatalf("expected empty result, got %d messages", len(result))
+	}
+}
+
+func TestSanitizeHistory_TrailingAssistantWithToolCallsNoResults(t *testing.T) {
+	history := []providers.Message{
+		{Role: "user", Content: "hello"},
+		{Role: "assistant", Content: "hi"},
+		{Role: "user", Content: "do something"},
+		{
+			Role: "assistant",
+			ToolCalls: []providers.ToolCall{
+				{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "tool1"}},
+			},
+		},
+	}
+
+	result := sanitizeHistory(history)
+
+	// Should remove trailing assistant with unanswered tool_calls
+	if len(result) != 3 {
+		t.Fatalf("expected 3 messages, got %d", len(result))
+	}
+	if result[2].Role != "user" {
+		t.Errorf("expected last message to be user, got %s", result[2].Role)
+	}
+}

pkg/agent/loop.go

@@ -27,7 +27,10 @@ import (
 	"github.com/sipeed/picoclaw/pkg/session"
 	"github.com/sipeed/picoclaw/pkg/state"
 	"github.com/sipeed/picoclaw/pkg/tools"
+	"github.com/sipeed/picoclaw/pkg/tracing"
 	"github.com/sipeed/picoclaw/pkg/utils"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 )
 
 type AgentLoop struct {
@@ -52,6 +55,7 @@ type processOptions struct {
 	Channel         string   // Target channel for tool execution
 	ChatID          string   // Target chat ID for tool execution
 	UserMessage     string   // User message content (may include prefix)
+	Media           []string // Media URLs (images) attached to the message
 	DefaultResponse string   // Response when LLM returns empty
 	EnableSummary   bool     // Whether to trigger summarization
 	SendResponse    bool     // Whether to send response via bus
@@ -280,6 +284,7 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
 		Channel:         msg.Channel,
 		ChatID:          msg.ChatID,
 		UserMessage:     msg.Content,
+		Media:           msg.Media,
 		DefaultResponse: "I've completed processing but have no response to give.",
 		EnableSummary:   true,
 		SendResponse:    false,
@@ -341,6 +346,14 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe
 // runAgentLoop is the core message processing logic.
 // It handles context building, LLM calls, tool execution, and response handling.
 func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (string, error) {
+	ctx, span := tracing.Tracer("agent").Start(ctx, "agent.processMessage",
+		trace.WithAttributes(
+			attribute.String("session_key", opts.SessionKey),
+			attribute.String("channel", opts.Channel),
+			attribute.String("chat_id", opts.ChatID),
+		))
+	defer span.End()
+
 	// 0. Record last channel for heartbeat notifications (skip internal channels)
 	if opts.Channel != "" && opts.ChatID != "" {
 		// Don't record internal channels (cli, system, subagent)
@@ -366,16 +379,14 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
 		history,
 		summary,
 		opts.UserMessage,
-		nil,
+		opts.Media,
 		opts.Channel,
 		opts.ChatID,
 	)
 
-	// 3. Save user message to session
+	// 3. Run LLM iteration loop (no session saves until success)
-	al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage)
+	historyOffset := 1 + len(history) // skip system prompt + existing history
-
+	finalContent, messages, iteration, err := al.runLLMIteration(ctx, messages, opts)
-	// 4. Run LLM iteration loop
-	finalContent, iteration, err := al.runLLMIteration(ctx, messages, opts)
 	if err != nil {
 		return "", err
 	}
@@ -383,21 +394,25 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
 	// If last tool had ForUser content and we already sent it, we might not need to send final response
 	// This is controlled by the tool's Silent flag and ForUser content
 
-	// 5. Handle empty response
+	// 4. Handle empty response
 	if finalContent == "" {
 		finalContent = opts.DefaultResponse
 	}
 
-	// 6. Save final assistant message to session
+	// 5. Atomically save all new messages (user + intermediate tool calls + final assistant) to session
+	// This prevents orphaned messages if the LLM call fails mid-way
+	for _, msg := range messages[historyOffset:] {
+		al.sessions.AddFullMessage(opts.SessionKey, msg)
+	}
 	al.sessions.AddMessage(opts.SessionKey, "assistant", finalContent)
 	al.sessions.Save(opts.SessionKey)
 
-	// 7. Optional: summarization
+	// 6. Optional: summarization
 	if opts.EnableSummary {
 		al.maybeSummarize(opts.SessionKey, opts.Channel, opts.ChatID)
 	}
 
-	// 8. Optional: send response via bus
+	// 7. Optional: send response via bus
 	if opts.SendResponse {
 		al.bus.PublishOutbound(bus.OutboundMessage{
 			Channel: opts.Channel,
@@ -406,7 +421,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
 		})
 	}
 
-	// 9. Log response
+	// 8. Log response
 	responsePreview := utils.Truncate(finalContent, 120)
 	logger.InfoCF("agent", fmt.Sprintf("Response: %s", responsePreview),
 		map[string]interface{}{
@@ -419,8 +434,8 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
 }
 
 // runLLMIteration executes the LLM call loop with tool handling.
-// Returns the final content, iteration count, and any error.
+// Returns the final content, updated messages slice, iteration count, and any error.
-func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.Message, opts processOptions) (string, int, error) {
+func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.Message, opts processOptions) (string, []providers.Message, int, error) {
 	iteration := 0
 	var finalContent string
 
@@ -456,23 +471,30 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 				"tools_json":    formatToolsForLog(providerToolDefs),
 			})
 
+		// Call LLM with retry for context window errors
+		llmCtx, llmSpan := tracing.Tracer("agent").Start(ctx, "agent.llm.call",
+			trace.WithAttributes(
+				attribute.Int("iteration", iteration),
+				attribute.String("model", al.model),
+				attribute.Int("messages_count", len(messages)),
+				attribute.Int("tools_count", len(providerToolDefs)),
+			))
+
 		var response *providers.LLMResponse
 		var err error
 
-		// Retry loop for context/token errors
 		maxRetries := 2
 		for retry := 0; retry <= maxRetries; retry++ {
-			response, err = al.provider.Chat(ctx, messages, providerToolDefs, al.model, map[string]interface{}{
+			response, err = al.provider.Chat(llmCtx, messages, providerToolDefs, al.model, map[string]interface{}{
 				"max_tokens":  8192,
 				"temperature": 0.7,
 			})
 
 			if err == nil {
-				break // Success
+				break
 			}
 
 			errMsg := strings.ToLower(err.Error())
-			// Check for context window errors (provider specific, but usually contain "token" or "invalid")
 			isContextError := strings.Contains(errMsg, "token") ||
 				strings.Contains(errMsg, "context") ||
 				strings.Contains(errMsg, "invalidparameter") ||
@@ -484,7 +506,6 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 					"retry": retry,
 				})
 
-				// Notify user on first retry only
 				if retry == 0 && !constants.IsInternalChannel(opts.Channel) && opts.SendResponse {
 					al.bus.PublishOutbound(bus.OutboundMessage{
 						Channel: opts.Channel,
@@ -493,83 +514,16 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 					})
 				}
 
-				// Force compression
 				al.forceCompression(opts.SessionKey)
 
-				// Rebuild messages with compressed history
-				// Note: We need to reload history from session manager because forceCompression changed it
 				newHistory := al.sessions.GetHistory(opts.SessionKey)
 				newSummary := al.sessions.GetSummary(opts.SessionKey)
 
-				// Re-create messages for the next attempt
+				// Rebuild from session history only (user message already saved in step 3)
-				// We keep the current user message (opts.UserMessage) effectively
 				messages = al.contextBuilder.BuildMessages(
 					newHistory,
 					newSummary,
-					opts.UserMessage,
+					"",
-					nil,
-					opts.Channel,
-					opts.ChatID,
-				)
-
-				// Important: If we are in the middle of a tool loop (iteration > 1),
-				// rebuilding messages from session history might duplicate the flow or miss context
-				// if intermediate steps weren't saved correctly.
-				// However, al.sessions.AddFullMessage is called after every tool execution,
-				// so GetHistory should reflect the current state including partial tool execution.
-				// But we need to ensure we don't duplicate the user message which is appended in BuildMessages.
-				// BuildMessages(history...) takes the stored history and appends the *current* user message.
-				// If iteration > 1, the "current user message" was already added to history in step 3 of runAgentLoop.
-				// So if we pass opts.UserMessage again, we might duplicate it?
-				// Actually, step 3 is: al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage)
-				// So GetHistory ALREADY contains the user message!
-
-				// CORRECTION:
-				// BuildMessages combines: [System] + [History] + [CurrentMessage]
-				// But Step 3 added CurrentMessage to History.
-				// So if we use GetHistory now, it has the user message.
-				// If we pass opts.UserMessage to BuildMessages, it adds it AGAIN.
-
-				// For retry in the middle of a loop, we should rely on what's in the session.
-				// BUT checking BuildMessages implementation:
-				// It appends history... then appends currentMessage.
-
-				// Logic fix for retry:
-				// If iteration == 1, opts.UserMessage corresponds to the user input.
-				// If iteration > 1, we are processing tool results. The "messages" passed to Chat
-				// already accumulated tool outputs.
-				// Rebuilding from session history is safest because it persists state.
-				// Start fresh with rebuilt history.
-
-				// Special case: standard BuildMessages appends "currentMessage".
-				// If we are strictly retrying the *LLM call*, we want the exact same state as before but compressed.
-				// However, the "messages" argument passed to runLLMIteration is constructed by the caller.
-				// If we rebuild from Session, we need to know if "currentMessage" should be appended or is already in history.
-
-				// In runAgentLoop:
-				// 3. sessions.AddMessage(userMsg)
-				// 4. runLLMIteration(..., UserMessage)
-
-				// So History contains the user message.
-				// BuildMessages typically appends the user message as a *new* pending message.
-				// Wait, standard BuildMessages usage in runAgentLoop:
-				// messages := BuildMessages(history (has old), UserMessage)
-				// THEN AddMessage(UserMessage).
-				// So "history" passed to BuildMessages does NOT contain the current UserMessage yet.
-
-				// But here, inside the loop, we have already saved it.
-				// So GetHistory() includes the current user message.
-				// If we call BuildMessages(GetHistory(), UserMessage), we get duplicates.
-
-				// Hack/Fix:
-				// If we are retrying, we rebuild from Session History ONLY.
-				// We pass empty string as "currentMessage" to BuildMessages
-				// because the "current message" is already saved in history (step 3).
-
-				messages = al.contextBuilder.BuildMessages(
-					newHistory,
-					newSummary,
-					"", // Empty because history already contains the relevant messages
 					nil,
 					opts.Channel,
 					opts.ChatID,
@@ -578,9 +532,9 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 				continue
 			}
 
-			// Real error or success, break loop
 			break
 		}
+		llmSpan.End()
 
 		if err != nil {
 			logger.ErrorCF("agent", "LLM call failed",
@@ -588,7 +542,7 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 					"iteration": iteration,
 					"error":     err.Error(),
 				})
-			return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err)
+			return "", messages, iteration, fmt.Errorf("LLM call failed: %w", err)
 		}
 
 		// Check if no tool calls - we're done
@@ -618,6 +572,7 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 		assistantMsg := providers.Message{
 			Role:          "assistant",
 			Content:       response.Content,
+			RawAPIMessage: response.RawAssistantMessage,
 		}
 		for _, tc := range response.ToolCalls {
 			argumentsJSON, _ := json.Marshal(tc.Arguments)
@@ -628,13 +583,11 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 					Name:      tc.Name,
 					Arguments: string(argumentsJSON),
 				},
+				ExtraContent: tc.ExtraContent,
 			})
 		}
 		messages = append(messages, assistantMsg)
 
-		// Save assistant message with tool calls to session
-		al.sessions.AddFullMessage(opts.SessionKey, assistantMsg)
-
 		// Execute tool calls
 		for _, tc := range response.ToolCalls {
 			// Log tool call with arguments preview
@@ -662,7 +615,12 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 				}
 			}
 
+			_, toolSpan := tracing.Tracer("agent").Start(ctx, "agent.tool.execute",
+				trace.WithAttributes(
+					attribute.String("tool_name", tc.Name),
+				))
 			toolResult := al.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, opts.Channel, opts.ChatID, asyncCallback)
+			toolSpan.End()
 
 			// Send ForUser content to user immediately if not Silent
 			if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse {
@@ -690,13 +648,10 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
 				ToolCallID: tc.ID,
 			}
 			messages = append(messages, toolResultMsg)
-
-			// Save tool result message to session
-			al.sessions.AddFullMessage(opts.SessionKey, toolResultMsg)
 		}
 	}
 
-	return finalContent, iteration, nil
+	return finalContent, messages, iteration, nil
 }
 
 // updateToolContexts updates the context for tools that need channel/chatID info.

pkg/providers/key_rotator.go

@@ -0,0 +1,33 @@
+package providers
+
+import (
+	"sync"
+)
+
+// KeyRotator provides thread-safe round-robin API key selection.
+type KeyRotator struct {
+	keys  []string
+	index uint64
+	mu    sync.Mutex
+}
+
+// NewKeyRotator creates a new KeyRotator with the given keys.
+func NewKeyRotator(keys []string) *KeyRotator {
+	return &KeyRotator{
+		keys: keys,
+	}
+}
+
+// Len returns the number of API keys.
+func (kr *KeyRotator) Len() int {
+	return len(kr.keys)
+}
+
+// Next returns the next API key in round-robin order.
+func (kr *KeyRotator) Next() string {
+	kr.mu.Lock()
+	defer kr.mu.Unlock()
+	key := kr.keys[kr.index%uint64(len(kr.keys))]
+	kr.index++
+	return key
+}

pkg/providers/types.go

@@ -1,11 +1,15 @@
 package providers
 
-import "context"
+import (
+	"context"
+	"encoding/json"
+)
 
 type ToolCall struct {
 	ID           string                 `json:"id"`
 	Type         string                 `json:"type,omitempty"`
 	Function     *FunctionCall          `json:"function,omitempty"`
+	ExtraContent map[string]interface{} `json:"extra_content,omitempty"`
 	Name         string                 `json:"name,omitempty"`
 	Arguments    map[string]interface{} `json:"arguments,omitempty"`
 }
@@ -20,6 +24,7 @@ type LLMResponse struct {
 	ToolCalls           []ToolCall      `json:"tool_calls,omitempty"`
 	FinishReason        string          `json:"finish_reason"`
 	Usage               *UsageInfo      `json:"usage,omitempty"`
+	RawAssistantMessage json.RawMessage `json:"-"`
 }
 
 type UsageInfo struct {
@@ -28,11 +33,23 @@ type UsageInfo struct {
 	TotalTokens      int `json:"total_tokens"`
 }
 
+type ContentPart struct {
+	Type     string    `json:"type"`
+	Text     string    `json:"text,omitempty"`
+	ImageURL *ImageURL `json:"image_url,omitempty"`
+}
+
+type ImageURL struct {
+	URL string `json:"url"`
+}
+
 type Message struct {
 	Role         string          `json:"role"`
 	Content      string          `json:"content"`
+	ContentParts []ContentPart   `json:"content_parts,omitempty"`
 	ToolCalls    []ToolCall      `json:"tool_calls,omitempty"`
 	ToolCallID   string          `json:"tool_call_id,omitempty"`
+	RawAPIMessage json.RawMessage `json:"raw_api_message,omitempty"`
 }
 
 type LLMProvider interface {

pkg/session/manager.go

@@ -141,7 +141,18 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
 		return
 	}
 
-	session.Messages = session.Messages[len(session.Messages)-keepLast:]
+	startIdx := len(session.Messages) - keepLast
+	// Adjust to start at a user message boundary for valid turn ordering.
+	// This prevents truncation from landing mid-sequence (e.g., starting
+	// with assistant+tool_calls or tool results).
+	for startIdx < len(session.Messages) && session.Messages[startIdx].Role != "user" {
+		startIdx++
+	}
+	if startIdx >= len(session.Messages) {
+		// No user message found — keep all messages rather than corrupt history
+		return
+	}
+	session.Messages = session.Messages[startIdx:]
 	session.Updated = time.Now()
 }
 

pkg/session/manager_test.go

@@ -4,6 +4,8 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/providers"
 )
 
 func TestSanitizeFilename(t *testing.T) {
@@ -72,3 +74,89 @@ func TestSave_RejectsPathTraversal(t *testing.T) {
 		}
 	}
 }
+
+func TestTruncateHistory_StartsAtUserBoundary(t *testing.T) {
+	sm := NewSessionManager("")
+	key := "test-session"
+
+	// Build a session with: user, assistant(tc), tool, user, assistant
+	sm.AddMessage(key, "user", "first")
+	sm.AddFullMessage(key, providers.Message{
+		Role: "assistant",
+		ToolCalls: []providers.ToolCall{
+			{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "test"}},
+		},
+	})
+	sm.AddFullMessage(key, providers.Message{
+		Role:       "tool",
+		Content:    "result",
+		ToolCallID: "call_1",
+	})
+	sm.AddMessage(key, "user", "second")
+	sm.AddMessage(key, "assistant", "response")
+
+	// keepLast=3 would normally start at index 2 (tool message)
+	// Smart truncation should advance to index 3 (user message)
+	sm.TruncateHistory(key, 3)
+
+	history := sm.GetHistory(key)
+	if len(history) != 2 {
+		t.Fatalf("expected 2 messages after truncation, got %d: %+v", len(history), history)
+	}
+	if history[0].Role != "user" {
+		t.Errorf("expected first message to be user, got %s", history[0].Role)
+	}
+	if history[0].Content != "second" {
+		t.Errorf("expected first message content 'second', got %q", history[0].Content)
+	}
+}
+
+func TestTruncateHistory_AlreadyAtUserBoundary(t *testing.T) {
+	sm := NewSessionManager("")
+	key := "test-session"
+
+	sm.AddMessage(key, "user", "hello")
+	sm.AddMessage(key, "assistant", "hi")
+	sm.AddMessage(key, "user", "bye")
+	sm.AddMessage(key, "assistant", "goodbye")
+
+	sm.TruncateHistory(key, 2)
+
+	history := sm.GetHistory(key)
+	if len(history) != 2 {
+		t.Fatalf("expected 2 messages, got %d", len(history))
+	}
+	if history[0].Role != "user" {
+		t.Errorf("expected first message to be user, got %s", history[0].Role)
+	}
+}
+
+func TestTruncateHistory_NoUserMessage(t *testing.T) {
+	sm := NewSessionManager("")
+	key := "test-session"
+
+	// Only non-user messages
+	sm.AddMessage(key, "assistant", "hello")
+	sm.AddFullMessage(key, providers.Message{
+		Role: "assistant",
+		ToolCalls: []providers.ToolCall{
+			{ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "test"}},
+		},
+	})
+	sm.AddFullMessage(key, providers.Message{
+		Role:       "tool",
+		Content:    "result",
+		ToolCallID: "call_1",
+	})
+
+	original := sm.GetHistory(key)
+	origLen := len(original)
+
+	// Should keep all messages since no user boundary exists
+	sm.TruncateHistory(key, 1)
+
+	history := sm.GetHistory(key)
+	if len(history) != origLen {
+		t.Fatalf("expected %d messages (unchanged), got %d", origLen, len(history))
+	}
+}

pkg/tools/shell.go

@@ -26,7 +26,7 @@ func NewExecTool(workingDir string, restrict bool) *ExecTool {
 		regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`),
 		regexp.MustCompile(`\bdel\s+/[fq]\b`),
 		regexp.MustCompile(`\brmdir\s+/s\b`),
-		regexp.MustCompile(`\b(format|mkfs|diskpart)\b\s`), // Match disk wiping commands (must be followed by space/args)
+		regexp.MustCompile(`(^|\s)(format|mkfs|diskpart)\s`), // Match disk wiping commands as standalone (not --format flags)
 		regexp.MustCompile(`\bdd\s+if=`),
 		regexp.MustCompile(`>\s*/dev/sd[a-z]\b`), // Block writes to disk devices (but allow /dev/null)
 		regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`),

pkg/tools/toolloop.go

@@ -99,6 +99,7 @@ func RunToolLoop(ctx context.Context, config ToolLoopConfig, messages []provider
 		assistantMsg := providers.Message{
 			Role:          "assistant",
 			Content:       response.Content,
+			RawAPIMessage: response.RawAssistantMessage,
 		}
 		for _, tc := range response.ToolCalls {
 			argumentsJSON, _ := json.Marshal(tc.Arguments)
@@ -109,6 +110,7 @@ func RunToolLoop(ctx context.Context, config ToolLoopConfig, messages []provider
 					Name:      tc.Name,
 					Arguments: string(argumentsJSON),
 				},
+				ExtraContent: tc.ExtraContent,
 			})
 		}
 		messages = append(messages, assistantMsg)