diff --git a/pkg/agent/context.go b/pkg/agent/context.go index cf5ce2913..7e1c026b5 100644 --- a/pkg/agent/context.go +++ b/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 - // --- INICIO DEL FIX --- - //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 --- + // Sanitize history to ensure valid turn ordering for all providers + history = sanitizeHistory(history) 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 +} diff --git a/pkg/agent/context_test.go b/pkg/agent/context_test.go new file mode 100644 index 000000000..bb431d479 --- /dev/null +++ b/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) + } +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index cd4276155..4585067b0 100644 --- a/pkg/agent/loop.go +++ b/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 { @@ -48,14 +51,15 @@ type AgentLoop struct { // processOptions configures how a message is processed type processOptions struct { - SessionKey string // Session identifier for history/context - Channel string // Target channel for tool execution - ChatID string // Target chat ID for tool execution - UserMessage string // User message content (may include prefix) - DefaultResponse string // Response when LLM returns empty - EnableSummary bool // Whether to trigger summarization - SendResponse bool // Whether to send response via bus - NoHistory bool // If true, don't load session history (for heartbeat) + SessionKey string // Session identifier for history/context + 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 + NoHistory bool // If true, don't load session history (for heartbeat) } // createToolRegistry creates a tool registry with common tools. @@ -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 - al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) - - // 4. Run LLM iteration loop - finalContent, iteration, err := al.runLLMIteration(ctx, messages, opts) + // 3. Run LLM iteration loop (no session saves until success) + historyOffset := 1 + len(history) // skip system prompt + existing history + finalContent, messages, 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. -func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.Message, opts processOptions) (string, int, 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, []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 - // We keep the current user message (opts.UserMessage) effectively + // Rebuild from session history only (user message already saved in step 3) 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 @@ -616,8 +570,9 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M // Build assistant message with tool calls assistantMsg := providers.Message{ - Role: "assistant", - Content: response.Content, + 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. diff --git a/pkg/providers/key_rotator.go b/pkg/providers/key_rotator.go new file mode 100644 index 000000000..7137193ec --- /dev/null +++ b/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 +} diff --git a/pkg/providers/types.go b/pkg/providers/types.go index 88b62e975..770c67b51 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -1,13 +1,17 @@ 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"` - Name string `json:"name,omitempty"` - Arguments map[string]interface{} `json:"arguments,omitempty"` + 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"` } type FunctionCall struct { @@ -16,10 +20,11 @@ type FunctionCall struct { } type LLMResponse struct { - Content string `json:"content"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - FinishReason string `json:"finish_reason"` - Usage *UsageInfo `json:"usage,omitempty"` + Content string `json:"content"` + 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"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - ToolCallID string `json:"tool_call_id,omitempty"` + 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 { diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 12bf33df0..cebc5b9a4 100644 --- a/pkg/session/manager.go +++ b/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() } diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 5ef5f4349..8e7c2279c 100644 --- a/pkg/session/manager_test.go +++ b/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)) + } +} diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index 1ca3fc35a..a0da52f81 100644 --- a/pkg/tools/shell.go +++ b/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`), diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index 1302079b4..6be575a95 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -97,8 +97,9 @@ func RunToolLoop(ctx context.Context, config ToolLoopConfig, messages []provider // 6. Build assistant message with tool calls assistantMsg := providers.Message{ - Role: "assistant", - Content: response.Content, + 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)