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.
This commit is contained in:
pkonowrocki
parent
7f60392d88
commit
0fa8e78424
9 changed files with 501 additions and 62 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
151
pkg/agent/context_test.go
Normal file
151
pkg/agent/context_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -25,7 +25,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 {
|
||||
|
|
@ -45,14 +48,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.
|
||||
|
|
@ -258,6 +262,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,
|
||||
|
|
@ -319,6 +324,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)
|
||||
|
|
@ -344,16 +357,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
|
||||
}
|
||||
|
|
@ -361,21 +372,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)
|
||||
}
|
||||
|
||||
// 8. Optional: send response via bus
|
||||
// 7. Optional: send response via bus
|
||||
if opts.SendResponse {
|
||||
al.bus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
|
|
@ -384,7 +399,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{}{
|
||||
|
|
@ -397,8 +412,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
|
||||
|
||||
|
|
@ -435,10 +450,18 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
|
|||
})
|
||||
|
||||
// Call LLM
|
||||
response, err := al.provider.Chat(ctx, messages, providerToolDefs, al.model, map[string]interface{}{
|
||||
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)),
|
||||
))
|
||||
response, err := al.provider.Chat(llmCtx, messages, providerToolDefs, al.model, map[string]interface{}{
|
||||
"max_tokens": 8192,
|
||||
"temperature": 0.7,
|
||||
})
|
||||
llmSpan.End()
|
||||
|
||||
if err != nil {
|
||||
logger.ErrorCF("agent", "LLM call failed",
|
||||
|
|
@ -446,7 +469,7 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
|
|||
"iteration": iteration,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return "", iteration, fmt.Errorf("LLM call failed: %w", err)
|
||||
return "", messages, iteration, fmt.Errorf("LLM call failed: %w", err)
|
||||
}
|
||||
|
||||
// Check if no tool calls - we're done
|
||||
|
|
@ -474,8 +497,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)
|
||||
|
|
@ -486,13 +510,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
|
||||
|
|
@ -520,7 +542,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 {
|
||||
|
|
@ -548,13 +575,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.
|
||||
|
|
|
|||
33
pkg/providers/key_rotator.go
Normal file
33
pkg/providers/key_rotator.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
||||
|
|
|
|||
93
pkg/session/manager_test.go
Normal file
93
pkg/session/manager_test.go
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
package session
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
|
@ -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`),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue