Merge pull request #1 from dimonb/codex/fix-tasktool-plan-ordering

fix(tasktool): serialize sequential tool calls per turn
This commit is contained in:
DimonB 2026-03-08 15:00:47 +02:00 committed by GitHub
commit 1d5c4e93ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 351 additions and 91 deletions

View file

@ -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 {

View file

@ -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

View file

@ -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",

View file

@ -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)
}

View file

@ -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{

View file

@ -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" +

View file

@ -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 {