From 61c93786da95cb146adfd0a1a72dfd02ffaa684c Mon Sep 17 00:00:00 2001 From: Dmitrii Balabanov Date: Sun, 8 Mar 2026 14:58:29 +0200 Subject: [PATCH] fix(tasktool): serialize sequential tool calls per turn --- pkg/agent/loop.go | 162 +++++++++++++++++++++---------------- pkg/agent/loop_test.go | 159 ++++++++++++++++++++++++++++++++++++ pkg/tools/base.go | 9 +++ pkg/tools/registry.go | 12 +++ pkg/tools/registry_test.go | 28 +++++++ pkg/tools/tasktool.go | 4 + pkg/tools/toolloop.go | 68 +++++++++++----- 7 files changed, 351 insertions(+), 91 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 986b86b88..9d304271f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1158,86 +1158,110 @@ func (al *AgentLoop) runLLMIteration( // Save assistant message with tool calls to session agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) - // Execute tool calls in parallel + // Execute tool calls, preserving model order for tools that require it. type indexedAgentResult struct { result *tools.ToolResult tc providers.ToolCall } agentResults := make([]indexedAgentResult, len(normalizedToolCalls)) - var wg sync.WaitGroup + executeToolCall := func(idx int, tc providers.ToolCall) { + argsJSON, _ := json.Marshal(tc.Arguments) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), + map[string]any{ + "agent_id": agent.ID, + "tool": tc.Name, + "iteration": iteration, + }) - for i, tc := range normalizedToolCalls { - agentResults[i].tc = tc - - wg.Add(1) - go func(idx int, tc providers.ToolCall) { - defer wg.Done() - - argsJSON, _ := json.Marshal(tc.Arguments) - argsPreview := utils.Truncate(string(argsJSON), 200) - logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), - map[string]any{ - "agent_id": agent.ID, - "tool": tc.Name, - "iteration": iteration, - }) - - // Create async callback for tools that implement AsyncExecutor. - // When the background work completes, this publishes the result - // as an inbound system message so processSystemMessage routes it - // back to the user via the normal agent loop. - asyncCallback := func(_ context.Context, result *tools.ToolResult) { - // Send ForUser content directly to the user (immediate feedback), - // mirroring the synchronous tool execution path. - if !result.Silent && result.ForUser != "" { - outCtx, outCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer outCancel() - _ = al.bus.PublishOutbound(outCtx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: result.ForUser, - }) - } - - // Determine content for the agent loop (ForLLM or error). - content := result.ForLLM - if content == "" && result.Err != nil { - content = result.Err.Error() - } - if content == "" { - return - } - - logger.InfoCF("agent", "Async tool completed, publishing result", - map[string]any{ - "tool": tc.Name, - "content_len": len(content), - "channel": opts.Channel, - }) - - pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer pubCancel() - _ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{ - Channel: "system", - SenderID: fmt.Sprintf("async:%s", tc.Name), - ChatID: fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID), - Content: content, + // Create async callback for tools that implement AsyncExecutor. + // When the background work completes, this publishes the result + // as an inbound system message so processSystemMessage routes it + // back to the user via the normal agent loop. + asyncCallback := func(_ context.Context, result *tools.ToolResult) { + // Send ForUser content directly to the user (immediate feedback), + // mirroring the synchronous tool execution path. + if !result.Silent && result.ForUser != "" { + outCtx, outCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer outCancel() + _ = al.bus.PublishOutbound(outCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: result.ForUser, }) } - toolResult := agent.Tools.ExecuteWithContext( - ctx, - tc.Name, - tc.Arguments, - opts.Channel, - opts.ChatID, - asyncCallback, - ) - agentResults[idx].result = toolResult - }(i, tc) + // Determine content for the agent loop (ForLLM or error). + content := result.ForLLM + if content == "" && result.Err != nil { + content = result.Err.Error() + } + if content == "" { + return + } + + logger.InfoCF("agent", "Async tool completed, publishing result", + map[string]any{ + "tool": tc.Name, + "content_len": len(content), + "channel": opts.Channel, + }) + + pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + _ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{ + Channel: "system", + SenderID: fmt.Sprintf("async:%s", tc.Name), + ChatID: fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID), + Content: content, + }) + } + + toolResult := agent.Tools.ExecuteWithContext( + ctx, + tc.Name, + tc.Arguments, + opts.Channel, + opts.ChatID, + asyncCallback, + ) + agentResults[idx].result = toolResult + } + + executeParallelBatch := func(start, end int) { + var wg sync.WaitGroup + for i := start; i < end; i++ { + tc := normalizedToolCalls[i] + wg.Add(1) + go func(idx int, tc providers.ToolCall) { + defer wg.Done() + executeToolCall(idx, tc) + }(i, tc) + } + wg.Wait() + } + + batchStart := -1 + for i, tc := range normalizedToolCalls { + agentResults[i].tc = tc + + if agent.Tools.ExecutesSequentially(tc.Name) { + if batchStart != -1 { + executeParallelBatch(batchStart, i) + batchStart = -1 + } + executeToolCall(i, tc) + continue + } + + if batchStart == -1 { + batchStart = i + } + } + if batchStart != -1 { + executeParallelBatch(batchStart, len(normalizedToolCalls)) } - wg.Wait() // Process results in original order (send to user, save to session) for _, r := range agentResults { diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 5625dafba..acd81bef1 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -7,6 +7,7 @@ import ( "path/filepath" "slices" "strings" + "sync" "testing" "time" @@ -16,6 +17,7 @@ import ( "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/routing" + "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -384,6 +386,110 @@ func (m *taskToolPlanMockProvider) GetDefaultModel() string { return "tasktool-mock-model" } +type taskToolRaceMockProvider struct { + calls int +} + +func (m *taskToolRaceMockProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + if m.calls == 1 { + return &providers.LLMResponse{ + ToolCalls: []providers.ToolCall{ + { + ID: "call_tasktool_create", + Name: "tasktool", + Arguments: map[string]any{ + "action": "create_plan", + "tasks": []any{ + map[string]any{ + "id": "step_1", + "description": "Create the plan", + }, + }, + }, + }, + { + ID: "call_tasktool_update", + Name: "tasktool", + Arguments: map[string]any{ + "action": "update_task", + "task_id": "step_1", + "status": string(session.TaskStatusCompleted), + "result": "done", + }, + }, + }, + }, nil + } + + return &providers.LLMResponse{ + Content: "", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (m *taskToolRaceMockProvider) GetDefaultModel() string { + return "tasktool-race-mock-model" +} + +type blockingSequentialTaskTool struct { + inner *tools.TaskTool + createOnce sync.Once + updateOnce sync.Once + createStarted chan struct{} + updateStarted chan struct{} +} + +func newBlockingSequentialTaskTool(inner *tools.TaskTool) *blockingSequentialTaskTool { + return &blockingSequentialTaskTool{ + inner: inner, + createStarted: make(chan struct{}), + updateStarted: make(chan struct{}), + } +} + +func (t *blockingSequentialTaskTool) Name() string { + return t.inner.Name() +} + +func (t *blockingSequentialTaskTool) Description() string { + return t.inner.Description() +} + +func (t *blockingSequentialTaskTool) Parameters() map[string]any { + return t.inner.Parameters() +} + +func (t *blockingSequentialTaskTool) ExecuteSequentially() bool { + return true +} + +func (t *blockingSequentialTaskTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + action, _ := args["action"].(string) + switch action { + case "create_plan": + t.createOnce.Do(func() { close(t.createStarted) }) + + select { + case <-t.updateStarted: + // If sibling tool calls are still fanned out in parallel, allow the + // update path to reach TaskManager.UpdateTask before the plan exists. + time.Sleep(10 * time.Millisecond) + case <-time.After(50 * time.Millisecond): + } + case "update_task": + t.updateOnce.Do(func() { close(t.updateStarted) }) + } + + return t.inner.Execute(ctx, args) +} + // mockCustomTool is a simple mock tool for registration testing type mockCustomTool struct{} @@ -731,6 +837,59 @@ func TestTaskTool_DirectModeWithoutChannelManagerReturnsPlan(t *testing.T) { } } +func TestTaskTool_CreatePlanAndUpdateTaskSameTurnRunsInOrder(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = tmpDir + cfg.Agents.Defaults.Model = "test-model" + cfg.Agents.Defaults.MaxTokens = 4096 + cfg.Agents.Defaults.MaxToolIterations = 4 + + msgBus := bus.NewMessageBus() + provider := &taskToolRaceMockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { + t.Fatal("No default agent found") + } + + taskTool, ok := defaultAgent.Tools.Get("tasktool") + if !ok { + t.Fatal("tasktool is not registered") + } + + inner, ok := taskTool.(*tools.TaskTool) + if !ok { + t.Fatalf("tasktool has unexpected type %T", taskTool) + } + + defaultAgent.Tools.Register(newBlockingSequentialTaskTool(inner)) + + if _, err := al.ProcessDirect(context.Background(), "make a plan and complete it", "cli:race"); err != nil { + t.Fatalf("ProcessDirect failed: %v", err) + } + + st := al.taskManager.Get(routing.BuildAgentMainSessionKey(defaultAgent.ID)) + if st == nil { + t.Fatal("expected task plan to be stored") + } + if len(st.Tasks) != 1 { + t.Fatalf("expected 1 task, got %d", len(st.Tasks)) + } + if st.Tasks[0].Status != session.TaskStatusCompleted { + t.Fatalf("expected task status %q, got %q", session.TaskStatusCompleted, st.Tasks[0].Status) + } + if st.Tasks[0].Result != "done" { + t.Fatalf("expected task result %q, got %q", "done", st.Tasks[0].Result) + } +} + // failFirstMockProvider fails on the first N calls with a specific error type failFirstMockProvider struct { failures int diff --git a/pkg/tools/base.go b/pkg/tools/base.go index f930742eb..f61316667 100644 --- a/pkg/tools/base.go +++ b/pkg/tools/base.go @@ -93,6 +93,15 @@ type AsyncExecutor interface { ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult } +// SequentialTool marks tools that must execute in model order within a single +// LLM turn, instead of being fanned out in parallel with sibling tool calls. +// This is intended for tools whose calls mutate shared state and can depend on +// earlier calls from the same assistant message. +type SequentialTool interface { + Tool + ExecuteSequentially() bool +} + func ToolToSchema(tool Tool) map[string]any { return map[string]any{ "type": "function", diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index ca8436c67..1b77d03ef 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -40,6 +40,18 @@ func (r *ToolRegistry) Get(name string) (Tool, bool) { return tool, ok } +// ExecutesSequentially reports whether the named tool must preserve model order +// within a single LLM turn instead of being fanned out with sibling calls. +func (r *ToolRegistry) ExecutesSequentially(name string) bool { + tool, ok := r.Get(name) + if !ok { + return false + } + + sequential, ok := tool.(SequentialTool) + return ok && sequential.ExecuteSequentially() +} + func (r *ToolRegistry) Execute(ctx context.Context, name string, args map[string]any) *ToolResult { return r.ExecuteWithContext(ctx, name, args, "", "", nil) } diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index 92d7d5abd..243c1bab5 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -45,6 +45,15 @@ func (m *mockAsyncRegistryTool) ExecuteAsync(_ context.Context, args map[string] return m.result } +type mockSequentialRegistryTool struct { + mockRegistryTool + sequential bool +} + +func (m *mockSequentialRegistryTool) ExecuteSequentially() bool { + return m.sequential +} + // --- helpers --- func newMockTool(name, desc string) *mockRegistryTool { @@ -104,6 +113,25 @@ func TestToolRegistry_RegisterOverwrite(t *testing.T) { } } +func TestToolRegistry_ExecutesSequentially(t *testing.T) { + r := NewToolRegistry() + r.Register(&mockSequentialRegistryTool{ + mockRegistryTool: *newMockTool("seq", "ordered"), + sequential: true, + }) + r.Register(newMockTool("plain", "parallel")) + + if !r.ExecutesSequentially("seq") { + t.Fatal("expected sequential tool to be detected") + } + if r.ExecutesSequentially("plain") { + t.Fatal("expected non-sequential tool to remain parallel") + } + if r.ExecutesSequentially("missing") { + t.Fatal("expected missing tool to report false") + } +} + func TestToolRegistry_Execute_Success(t *testing.T) { r := NewToolRegistry() r.Register(&mockRegistryTool{ diff --git a/pkg/tools/tasktool.go b/pkg/tools/tasktool.go index 13cd240e6..738007d31 100644 --- a/pkg/tools/tasktool.go +++ b/pkg/tools/tasktool.go @@ -30,6 +30,10 @@ func (t *TaskTool) Name() string { return "tasktool" } +func (t *TaskTool) ExecuteSequentially() bool { + return true +} + func (t *TaskTool) Description() string { return "Manage planning mode tasks. Use action='create_plan' to start a new plan with a list of tasks. Use action='update_task' to update the status of an existing task and return the current plan state.\n\n" + "CRITICAL INSTRUCTIONS:\n" + diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index 244f0d4a2..12a4bc7fd 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -122,40 +122,64 @@ func RunToolLoop( } messages = append(messages, assistantMsg) - // 7. Execute tool calls in parallel + // 7. Execute tool calls, preserving model order for tools that require it. type indexedResult struct { result *ToolResult tc providers.ToolCall } results := make([]indexedResult, len(normalizedToolCalls)) - var wg sync.WaitGroup + executeToolCall := func(idx int, tc providers.ToolCall) { + argsJSON, _ := json.Marshal(tc.Arguments) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), + map[string]any{ + "tool": tc.Name, + "iteration": iteration, + }) + var toolResult *ToolResult + if config.Tools != nil { + toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil) + } else { + toolResult = ErrorResult("No tools available") + } + results[idx].result = toolResult + } + + executeParallelBatch := func(start, end int) { + var wg sync.WaitGroup + for i := start; i < end; i++ { + tc := normalizedToolCalls[i] + wg.Add(1) + go func(idx int, tc providers.ToolCall) { + defer wg.Done() + executeToolCall(idx, tc) + }(i, tc) + } + wg.Wait() + } + + batchStart := -1 for i, tc := range normalizedToolCalls { results[i].tc = tc - wg.Add(1) - go func(idx int, tc providers.ToolCall) { - defer wg.Done() - - argsJSON, _ := json.Marshal(tc.Arguments) - argsPreview := utils.Truncate(string(argsJSON), 200) - logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), - map[string]any{ - "tool": tc.Name, - "iteration": iteration, - }) - - var toolResult *ToolResult - if config.Tools != nil { - toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil) - } else { - toolResult = ErrorResult("No tools available") + if config.Tools != nil && config.Tools.ExecutesSequentially(tc.Name) { + if batchStart != -1 { + executeParallelBatch(batchStart, i) + batchStart = -1 } - results[idx].result = toolResult - }(i, tc) + executeToolCall(i, tc) + continue + } + + if batchStart == -1 { + batchStart = i + } + } + if batchStart != -1 { + executeParallelBatch(batchStart, len(normalizedToolCalls)) } - wg.Wait() // Append results in original order for _, r := range results {