From d37c44db901ef9f8f5a425bd7a16198e8d5909cb Mon Sep 17 00:00:00 2001 From: Anton Bogdanovich <27antonb@gmail.com> Date: Fri, 8 May 2026 20:11:47 -0700 Subject: [PATCH] fix(spawn): add async delivery policy for specialist results --- pkg/agent/eventbus_test.go | 88 +++++++++++++++++++++++++++++- pkg/agent/pipeline_execute.go | 42 +++++++++++++- pkg/tools/result_test.go | 11 ++++ pkg/tools/shared/result.go | 24 ++++++++ pkg/tools/shared_facade.go | 6 ++ pkg/tools/spawn.go | 100 ++++++++++++++++++++-------------- pkg/tools/spawn_test.go | 82 ++++++++++++++++++++++++++++ 7 files changed, 308 insertions(+), 45 deletions(-) diff --git a/pkg/agent/eventbus_test.go b/pkg/agent/eventbus_test.go index 86d7f4afa..f0e913038 100644 --- a/pkg/agent/eventbus_test.go +++ b/pkg/agent/eventbus_test.go @@ -699,6 +699,7 @@ type asyncFollowUpTool struct { name string followUpText string completionSig chan struct{} + deliveryMode tools.AsyncDeliveryMode } func (t *asyncFollowUpTool) Name() string { @@ -726,7 +727,11 @@ func (t *asyncFollowUpTool) ExecuteAsync( cb tools.AsyncCallback, ) *tools.ToolResult { go func() { - cb(ctx, &tools.ToolResult{ForLLM: t.followUpText}) + res := &tools.ToolResult{ForLLM: t.followUpText} + if t.deliveryMode != "" { + res.WithAsyncDelivery(t.deliveryMode) + } + cb(ctx, res) if t.completionSig != nil { close(t.completionSig) } @@ -738,3 +743,84 @@ var ( _ tools.Tool = (*mockCustomTool)(nil) _ tools.AsyncExecutor = (*asyncFollowUpTool)(nil) ) + +func TestAgentLoop_AsyncToolUserOnly_DoesNotEmitFollowUpQueued(t *testing.T) { + tmpDir := t.TempDir() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + provider := &toolCallProvider{ + toolCalls: []providers.ToolCall{ + { + ID: "call_async_1", + Type: "function", + Name: "async_followup_user_only", + Function: &providers.FunctionCall{ + Name: "async_followup_user_only", + Arguments: "{}", + }, + Arguments: map[string]any{}, + }, + }, + finalResp: "async launched", + } + + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + doneCh := make(chan struct{}) + al.RegisterTool(&asyncFollowUpTool{ + name: "async_followup_user_only", + followUpText: "background result", + completionSig: doneCh, + deliveryMode: tools.AsyncDeliveryUserOnly, + }) + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { + t.Fatal("expected default agent") + } + + runtimeCh, closeRuntimeEvents := subscribeRuntimeEventsForTest( + t, + al, + 8, + runtimeevents.KindAgentFollowUpQueued, + ) + defer closeRuntimeEvents() + + resp, err := al.runAgentLoop(context.Background(), defaultAgent, processOptions{ + SessionKey: "session-1", + Channel: "cli", + ChatID: "direct", + UserMessage: "run async tool", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + if resp != "async launched" { + t.Fatalf("expected final response 'async launched', got %q", resp) + } + + select { + case <-doneCh: + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for async tool completion") + } + + select { + case evt := <-runtimeCh: + t.Fatalf("unexpected follow-up queued event: %+v", evt) + case <-time.After(200 * time.Millisecond): + } +} diff --git a/pkg/agent/pipeline_execute.go b/pkg/agent/pipeline_execute.go index 567e56d17..05c3a4285 100644 --- a/pkg/agent/pipeline_execute.go +++ b/pkg/agent/pipeline_execute.go @@ -103,6 +103,42 @@ func inferSkillNamesFromToolCall(ts *turnState, toolName string, toolArgs map[st return names } +func effectiveAsyncToolResultDelivery(result *tools.ToolResult) tools.AsyncDeliveryMode { + if result == nil || result.AsyncDelivery == "" { + return tools.AsyncDeliveryUserAndParent + } + return result.AsyncDelivery +} + +func shouldPublishAsyncToolResultToUser(result *tools.ToolResult) bool { + if result == nil { + return false + } + switch effectiveAsyncToolResultDelivery(result) { + case tools.AsyncDeliveryParentOnly: + return false + default: + return !result.Silent && result.ForUser != "" + } +} + +func shouldQueueAsyncToolResultForParent(result *tools.ToolResult) bool { + if result == nil { + return false + } + content := result.ContentForLLM() + if content == "" { + return false + } + switch effectiveAsyncToolResultDelivery(result) { + case tools.AsyncDeliveryUserOnly: + return false + case tools.AsyncDeliveryParentOnly, tools.AsyncDeliveryUserAndParent: + return true + } + return true +} + // ExecuteTools executes the tool loop, handling BeforeTool/ApproveTool/AfterTool hooks, // tool execution with async callbacks, media delivery, and steering injection. // Returns ToolControl indicating what the coordinator should do next: @@ -474,17 +510,17 @@ toolLoop: toolCallID := tc.ID asyncToolName := toolName asyncCallback := func(_ context.Context, result *tools.ToolResult) { - if !result.Silent && result.ForUser != "" { + if shouldPublishAsyncToolResultToUser(result) { outCtx, outCancel := context.WithTimeout(context.Background(), 5*time.Second) defer outCancel() _ = al.bus.PublishOutbound(outCtx, outboundMessageForTurn(ts, result.ForUser)) } - content := result.ContentForLLM() - if content == "" { + if !shouldQueueAsyncToolResultForParent(result) { return } + content := result.ContentForLLM() content = al.cfg.FilterSensitiveData(content) logger.InfoCF("agent", "Async tool completed, publishing result", diff --git a/pkg/tools/result_test.go b/pkg/tools/result_test.go index 5f08cb4fa..edbd91c01 100644 --- a/pkg/tools/result_test.go +++ b/pkg/tools/result_test.go @@ -56,6 +56,17 @@ func TestAsyncResult(t *testing.T) { if !result.Async { t.Error("Expected Async to be true") } + if result.AsyncDelivery != "" { + t.Errorf("Expected empty AsyncDelivery by default, got %q", result.AsyncDelivery) + } +} + +func TestToolResultWithAsyncDelivery(t *testing.T) { + result := AsyncResult("async task started").WithAsyncDelivery(AsyncDeliveryUserOnly) + + if result.AsyncDelivery != AsyncDeliveryUserOnly { + t.Fatalf("AsyncDelivery = %q, want %q", result.AsyncDelivery, AsyncDeliveryUserOnly) + } } func TestErrorResult(t *testing.T) { diff --git a/pkg/tools/shared/result.go b/pkg/tools/shared/result.go index e4b16f7b3..b136d6275 100644 --- a/pkg/tools/shared/result.go +++ b/pkg/tools/shared/result.go @@ -12,6 +12,14 @@ const ( ArtifactPathsLLMNote = "Use `send_file` with one of these paths to send it to the user, or use file/exec tools to save it inside the workspace if requested." ) +type AsyncDeliveryMode string + +const ( + AsyncDeliveryUserOnly AsyncDeliveryMode = "user_only" + AsyncDeliveryParentOnly AsyncDeliveryMode = "parent_only" + AsyncDeliveryUserAndParent AsyncDeliveryMode = "user_and_parent" +) + // ToolResult represents the structured return value from tool execution. // It provides clear semantics for different types of results and supports // async operations, user-facing messages, and error handling. @@ -37,6 +45,16 @@ type ToolResult struct { // When true, the tool will complete later and notify via callback. Async bool `json:"async"` + // AsyncDelivery controls how the final async result should be routed when + // the background work completes. + // + // Empty means "use runtime default behavior". + // Supported values: + // - user_only + // - parent_only + // - user_and_parent + AsyncDelivery AsyncDeliveryMode `json:"async_delivery,omitempty"` + // Err is the underlying error (not JSON serialized). // Used for internal error handling and logging. Err error `json:"-"` @@ -221,3 +239,9 @@ func (tr *ToolResult) WithResponseHandled() *ToolResult { tr.ResponseHandled = true return tr } + +// WithAsyncDelivery sets the async delivery policy for this tool result. +func (tr *ToolResult) WithAsyncDelivery(mode AsyncDeliveryMode) *ToolResult { + tr.AsyncDelivery = mode + return tr +} diff --git a/pkg/tools/shared_facade.go b/pkg/tools/shared_facade.go index 8409ea060..042b05145 100644 --- a/pkg/tools/shared_facade.go +++ b/pkg/tools/shared_facade.go @@ -36,6 +36,10 @@ const ( ToolPromptSlotMCP = toolshared.ToolPromptSlotMCP ToolPromptSourceRegistry = toolshared.ToolPromptSourceRegistry ToolPromptSourceDiscovery = toolshared.ToolPromptSourceDiscovery + + AsyncDeliveryUserOnly = toolshared.AsyncDeliveryUserOnly + AsyncDeliveryParentOnly = toolshared.AsyncDeliveryParentOnly + AsyncDeliveryUserAndParent = toolshared.AsyncDeliveryUserAndParent ) func WithToolContext(ctx context.Context, channel, chatID string) context.Context { @@ -101,6 +105,8 @@ func SilentResult(forLLM string) *ToolResult { return toolshared.SilentResult(forLLM) } +type AsyncDeliveryMode = toolshared.AsyncDeliveryMode + func AsyncResult(forLLM string) *ToolResult { return toolshared.AsyncResult(forLLM) } diff --git a/pkg/tools/spawn.go b/pkg/tools/spawn.go index a9a373856..ed32b6335 100644 --- a/pkg/tools/spawn.go +++ b/pkg/tools/spawn.go @@ -7,6 +7,7 @@ import ( ) type SpawnTool struct { + manager *SubagentManager spawner SubTurnSpawner defaultModel string maxTokens int @@ -22,6 +23,7 @@ func NewSpawnTool(manager *SubagentManager) *SpawnTool { return &SpawnTool{} } return &SpawnTool{ + manager: manager, defaultModel: manager.defaultModel, maxTokens: manager.maxTokens, temperature: manager.temperature, @@ -31,6 +33,27 @@ func NewSpawnTool(manager *SubagentManager) *SpawnTool { // SetSpawner sets the SubTurnSpawner for direct sub-turn execution. func (t *SpawnTool) SetSpawner(spawner SubTurnSpawner) { t.spawner = spawner + if t.manager != nil && spawner != nil { + t.manager.SetSpawner(func( + ctx context.Context, + task, label, agentID string, + tools *ToolRegistry, + maxTokens int, + temperature float64, + hasMaxTokens, hasTemperature bool, + ) (*ToolResult, error) { + return spawner.SpawnSubTurn(ctx, SubTurnConfig{ + TargetAgentID: strings.TrimSpace(agentID), + Model: t.defaultModel, + Tools: nil, + SystemPrompt: buildSpawnSystemPrompt(task, label), + MaxTokens: maxTokens, + Temperature: temperature, + Async: false, + Critical: true, + }) + }) + } } func (t *SpawnTool) Name() string { @@ -101,16 +124,40 @@ func (t *SpawnTool) execute( } } - // Build system prompt for spawned subagent - systemPrompt := fmt.Sprintf( - `You are a spawned subagent running in the background. Complete the given task independently and report back when done. + // Preferred path: route through SubagentManager so spawn_status and + // background execution share the same task registry. + if t.manager != nil { + wrappedCallback := cb + if cb != nil { + wrappedCallback = func(cbCtx context.Context, res *ToolResult) { + if res != nil { + res.WithAsyncDelivery(AsyncDeliveryUserOnly) + } + cb(cbCtx, res) + } + } + ack, err := t.manager.Spawn( + ctx, + task, + label, + strings.TrimSpace(agentID), + ToolChannel(ctx), + ToolChatID(ctx), + wrappedCallback, + ) + if err != nil { + return ErrorResult(fmt.Sprintf("Spawn failed: %v", err)).WithError(err) + } + return AsyncResult(ack) + } -Task: %s`, - task, - ) + // Fallback: manager not configured + return ErrorResult("Subagent manager not configured") +} +func buildSpawnSystemPrompt(task, label string) string { if label != "" { - systemPrompt = fmt.Sprintf( + return fmt.Sprintf( `You are a spawned subagent labeled "%s" running in the background. Complete the given task independently and report back when done. Task: %s`, @@ -118,38 +165,9 @@ Task: %s`, task, ) } - - // Use spawner if available (direct SpawnSubTurn call) - if t.spawner != nil { - // Launch async sub-turn in goroutine - go func() { - result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{ - Model: t.defaultModel, - Tools: nil, // Will inherit from parent via context - SystemPrompt: systemPrompt, - MaxTokens: t.maxTokens, - Temperature: t.temperature, - Async: true, // Async execution - Critical: true, // Background spawn should survive parent turn completion - TargetAgentID: targetAgentID, - }) - if err != nil { - result = ErrorResult(fmt.Sprintf("Spawn failed: %v", err)).WithError(err) - } - - // Call callback if provided - if cb != nil { - cb(ctx, result) - } - }() - - // Return immediate acknowledgment - if label != "" { - return AsyncResult(fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task)) - } - return AsyncResult(fmt.Sprintf("Spawned subagent for task: %s", task)) - } - - // Fallback: spawner not configured - return ErrorResult("Subagent manager not configured") + return fmt.Sprintf( + `You are a spawned subagent running in the background. Complete the given task independently and report back when done. +Task: %s`, + task, + ) } diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index c91c79578..5c366769e 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -4,6 +4,7 @@ import ( "context" "strings" "testing" + "time" ) // mockSpawner implements SubTurnSpawner for testing. @@ -113,3 +114,84 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) { t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) } } + +func TestSpawnTool_SpawnStatusSeesSpawnedTask(t *testing.T) { + provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test") + spawnTool := NewSpawnTool(manager) + spawner := &mockSpawner{done: make(chan struct{})} + spawnTool.SetSpawner(spawner) + statusTool := NewSpawnStatusTool(manager) + + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + args := map[string]any{ + "task": "Write a haiku about coding", + "label": "haiku-task", + "agent_id": "deep-research", + } + + result := spawnTool.Execute(ctx, args) + if result == nil { + t.Fatal("Result should not be nil") + } + if result.IsError { + t.Fatalf("Expected success for valid task, got error: %s", result.ForLLM) + } + if !result.Async { + t.Fatal("SpawnTool should return async result") + } + + deadline := time.Now().Add(2 * time.Second) + for { + status := statusTool.Execute(ctx, map[string]any{}) + if status == nil { + t.Fatal("status result should not be nil") + } + if status.IsError { + t.Fatalf("spawn_status returned error: %s", status.ForLLM) + } + if strings.Contains(status.ForLLM, "subagent-1") { + if !strings.Contains(status.ForLLM, "haiku-task") { + t.Fatalf("expected label in status output, got: %s", status.ForLLM) + } + break + } + if time.Now().After(deadline) { + t.Fatalf("spawn_status never observed spawned task; last output: %s", status.ForLLM) + } + time.Sleep(10 * time.Millisecond) + } + + <-spawner.done +} + +func TestSpawnTool_ExecuteAsync_MarksCallbackResultUserOnly(t *testing.T) { + provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test") + tool := NewSpawnTool(manager) + spawner := &mockSpawner{} + tool.SetSpawner(spawner) + + done := make(chan *ToolResult, 1) + result := tool.ExecuteAsync(context.Background(), map[string]any{ + "task": "Write a haiku about coding", + }, func(_ context.Context, res *ToolResult) { + done <- res + }) + + if result == nil || !result.Async { + t.Fatal("expected async acknowledgment result") + } + + select { + case cbResult := <-done: + if cbResult == nil { + t.Fatal("expected callback result") + } + if cbResult.AsyncDelivery != AsyncDeliveryUserOnly { + t.Fatalf("AsyncDelivery = %q, want %q", cbResult.AsyncDelivery, AsyncDeliveryUserOnly) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for spawn callback result") + } +}