From ec060ad7e1e2d2a8c38e40089bf5700c65db2d89 Mon Sep 17 00:00:00 2001 From: viettranx Date: Mon, 30 Mar 2026 14:12:00 +0700 Subject: [PATCH] fix(agent): unblock stuck agent on /stop, auto-complete nil-result tasks, graceful shutdown (#527, #504, #39) - loop.go: replace blocking range with select+ctx.Done() for parallel tool collection; document finalization trade-off on early cancel - loop_cancel_test.go: tests for context cancellation and normal completion paths - gateway_consumer_post_turn.go: always auto-complete team task in default case, fallback message when outcome.Result is nil - gateway.go: call sched.Stop() with 5s drain before context cancel on shutdown --- cmd/gateway_consumer_post_turn.go | 40 ++++++++++-------- internal/agent/loop.go | 36 +++++++++++----- internal/agent/loop_cancel_test.go | 68 ++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 29 deletions(-) create mode 100644 internal/agent/loop_cancel_test.go diff --git a/cmd/gateway_consumer_post_turn.go b/cmd/gateway_consumer_post_turn.go index 33bd24f9..1ef590f9 100644 --- a/cmd/gateway_consumer_post_turn.go +++ b/cmd/gateway_consumer_post_turn.go @@ -171,28 +171,32 @@ func resolveTeamTaskOutcome( default: // Agent turn ended without terminal action — auto-complete. // Covers: Progressed, Commented, Claimed, or no flags at all. + result := "" if outcome.Result != nil { - result := outcome.Result.Content + result = outcome.Result.Content if len(outcome.Result.Deliverables) > 0 { result = strings.Join(outcome.Result.Deliverables, "\n\n---\n\n") } - if len(result) > 100_000 { - result = result[:100_000] + "\n[truncated]" - } - if err := deps.TeamStore.CompleteTask(ctx, meta.TaskID, meta.TeamID, result); err != nil { - slog.Warn("auto-complete: CompleteTask error", "task_id", meta.TaskID, "error", err) - } else { - bus.BroadcastForTenant(deps.MsgBus, protocol.EventTeamTaskCompleted, store.TenantIDFromContext(ctx), tools.BuildTaskEventPayload( - meta.TeamID.String(), meta.TaskID.String(), - store.TeamTaskStatusCompleted, - "agent", toAgent, - tools.WithTaskInfo(taskNumber, taskSubject), - tools.WithOwnerAgentKey(toAgent), - tools.WithChannel(taskChannel), - tools.WithChatID(taskChatID), - tools.WithTimestamp(now), - )) - } + } + if result == "" { + result = "Agent run ended without explicit result" + } + if len(result) > 100_000 { + result = result[:100_000] + "\n[truncated]" + } + if err := deps.TeamStore.CompleteTask(ctx, meta.TaskID, meta.TeamID, result); err != nil { + slog.Warn("auto-complete: CompleteTask error", "task_id", meta.TaskID, "error", err) + } else { + bus.BroadcastForTenant(deps.MsgBus, protocol.EventTeamTaskCompleted, store.TenantIDFromContext(ctx), tools.BuildTaskEventPayload( + meta.TeamID.String(), meta.TaskID.String(), + store.TeamTaskStatusCompleted, + "agent", toAgent, + tools.WithTaskInfo(taskNumber, taskSubject), + tools.WithOwnerAgentKey(toAgent), + tools.WithChannel(taskChannel), + tools.WithChatID(taskChatID), + tools.WithTimestamp(now), + )) } } diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 32037b6d..2bb10e2d 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -20,6 +20,17 @@ import ( "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) +// indexedResult holds the output of a single parallel tool execution, preserving +// the original call index so results can be sorted back into deterministic order. +type indexedResult struct { + idx int + tc providers.ToolCall + registryName string + result *tools.Result + argsJSON string + spanStart time.Time +} + func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) { // Per-run emit wrapper: enriches every AgentEvent with delegation + routing context. emitRun := func(event AgentEvent) { @@ -551,14 +562,6 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Each goroutine performs lazy MCP activation + policy checking independently. // Tool instances are immutable (context-based) so concurrent access is safe. // Results are collected then processed sequentially for deterministic ordering. - type indexedResult struct { - idx int - tc providers.ToolCall - registryName string - result *tools.Result - argsJSON string - spanStart time.Time - } // 1. Emit all tool.call events upfront (client sees all calls starting) for _, tc := range resp.ToolCalls { @@ -625,10 +628,21 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Close channel after all goroutines complete (run in separate goroutine to avoid deadlock) go func() { wg.Wait(); close(resultCh) }() - // 3. Collect results + // 3. Collect results (respect context cancellation to allow /stop) collected := make([]indexedResult, 0, len(resp.ToolCalls)) - for r := range resultCh { - collected = append(collected, r) + collectLoop: + for range resp.ToolCalls { + select { + case r, ok := <-resultCh: + if !ok { + break collectLoop + } + collected = append(collected, r) + case <-ctx.Done(): + // Trade-off: responsive /stop cancellation skips finalizeRun() cleanup + // (bootstrap cleanup, message flush). Stuck agent is worse than lost finalization. + return nil, ctx.Err() + } } // 4. Sort by original index → deterministic message ordering diff --git a/internal/agent/loop_cancel_test.go b/internal/agent/loop_cancel_test.go new file mode 100644 index 00000000..e31ad470 --- /dev/null +++ b/internal/agent/loop_cancel_test.go @@ -0,0 +1,68 @@ +package agent + +import ( + "context" + "testing" +) + +func TestParallelToolCollection_ContextCancel(t *testing.T) { + // Simulate: 3 tool calls, but ctx cancelled before all arrive + ctx, cancel := context.WithCancel(context.Background()) + resultCh := make(chan indexedResult, 3) + + // Send only 1 result, then cancel + resultCh <- indexedResult{idx: 0} + cancel() + + collected := make([]indexedResult, 0, 3) + var err error + +collectLoop: + for i := 0; i < 3; i++ { + select { + case r, ok := <-resultCh: + if !ok { + break collectLoop + } + collected = append(collected, r) + case <-ctx.Done(): + err = ctx.Err() + break collectLoop + } + } + + if err == nil { + t.Fatal("expected context.Canceled error") + } + // We should have collected 1 result before cancellation (or 0 — timing dependent) + if len(collected) > 1 { + t.Errorf("expected at most 1 collected, got %d", len(collected)) + } +} + +func TestParallelToolCollection_AllComplete(t *testing.T) { + resultCh := make(chan indexedResult, 3) + resultCh <- indexedResult{idx: 0} + resultCh <- indexedResult{idx: 1} + resultCh <- indexedResult{idx: 2} + + ctx := context.Background() + collected := make([]indexedResult, 0, 3) + +collectLoop: + for i := 0; i < 3; i++ { + select { + case r, ok := <-resultCh: + if !ok { + break collectLoop + } + collected = append(collected, r) + case <-ctx.Done(): + t.Fatal("context should not be cancelled") + } + } + + if len(collected) != 3 { + t.Fatalf("expected 3, got %d", len(collected)) + } +}