From 911a3a7ad4646447b762a84c3ba20fac6d14d0d1 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 2 Oct 2026 08:18:01 +0000 Subject: [PATCH] fix(agent): make follow-up resumes asynchronous --- README.md | 5 +- agent.go | 122 +++++++++++++++++++------- agent_resume_attribution_test.go | 9 ++ agent_resume_test.go | 146 +++++++++++++++++++++++++++++-- agent_test.go | 18 ++++ tools.go | 2 +- 6 files changed, 260 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index f66616f..98a10a0 100644 --- a/README.md +++ b/README.md @@ -680,10 +680,11 @@ repair resume semantics. Excluding `send_agent_message` removes that entire tool not just its resume behavior. -When enabled, three built-in tools are injected: +When enabled, four built-in tools are injected: - **`spawn_agent`** — Spawns a sub-agent with a task. Set `background: true` for non-blocking execution. - **`check_agent`** — Checks the status of a background agent by ID. -- **`get_agent_result`** — Retrieves the result of a background agent. Set `wait: true` to block until done. +- **`get_agent_result`** — Retrieves the result of a background agent. Use `wait: false` for a non-blocking snapshot. Results are delivered automatically; `wait: true` blocks for at most 30 seconds and should be used only when no other work can proceed. +- **`send_agent_message`** — Sends a follow-up without blocking. It injects into a running agent or resumes a finished agent in the background; the completion result is delivered automatically. **Foreground Agents (Blocking):** diff --git a/agent.go b/agent.go index 86abd9a..7ddf836 100644 --- a/agent.go +++ b/agent.go @@ -7,6 +7,7 @@ import ( "slices" "strings" "sync" + "time" "github.com/google/uuid" "github.com/sashabaranov/go-openai" @@ -113,9 +114,11 @@ type CheckAgentArgs struct { // GetAgentResultArgs are the arguments for retrieving a background agent's result. type GetAgentResultArgs struct { AgentID string `json:"agent_id" description:"The ID of the background agent"` - Wait bool `json:"wait" description:"If true, blocks until the agent finishes. If false, returns immediately with current status."` + Wait bool `json:"wait" description:"If true, waits briefly for completion. Prefer false: background completions are delivered automatically, and waiting blocks the current tool call."` } +const defaultAgentResultWaitTimeout = 30 * time.Second + // AgentState tracks the lifecycle of a single sub-agent. type AgentState struct { ID string @@ -323,12 +326,13 @@ func (r *CheckAgentRunnerForTest) Run(args CheckAgentArgs) (string, any, error) // GetAgentResultRunnerForTest exposes the getAgentResultRunner for testing. type GetAgentResultRunnerForTest struct { - Manager *AgentManager - Ctx context.Context + Manager *AgentManager + Ctx context.Context + WaitTimeout time.Duration } func (r *GetAgentResultRunnerForTest) Run(args GetAgentResultArgs) (string, any, error) { - inner := &getAgentResultRunner{manager: r.Manager, ctx: r.Ctx} + inner := &getAgentResultRunner{manager: r.Manager, ctx: r.Ctx, waitTimeout: r.WaitTimeout} return inner.Run(args) } @@ -608,7 +612,6 @@ func (r *spawnAgentRunner) Run(args SpawnAgentArgs) (string, any, error) { // the parent loop. Shared by the foreground (detachable) and background spawn // branches so the lifecycle bookkeeping lives in one place. func (r *spawnAgentRunner) runAgent(agent *AgentState, llm LLM, frag Fragment, opts []Option, spec AgentRunSpec, ctx context.Context, cancel context.CancelFunc) { - defer close(agent.done) defer cancel() var ( @@ -638,6 +641,10 @@ func (r *spawnAgentRunner) runAgent(agent *AgentState, llm LLM, frag Fragment, o } r.manager.mu.Unlock() + // Publish terminal state before external callbacks. A slow callback must not + // keep Wait/get_agent_result blocked after the agent has finished. + close(agent.done) + // Fire completion callback. if r.agentCompletionCallback != nil { r.agentCompletionCallback(agent) @@ -713,8 +720,9 @@ func (r *checkAgentRunner) Run(args CheckAgentArgs) (string, any, error) { // getAgentResultRunner implements Tool[GetAgentResultArgs]. type getAgentResultRunner struct { - manager *AgentManager - ctx context.Context + manager *AgentManager + ctx context.Context + waitTimeout time.Duration } func (r *getAgentResultRunner) Run(args GetAgentResultArgs) (string, any, error) { @@ -725,13 +733,20 @@ func (r *getAgentResultRunner) Run(args GetAgentResultArgs) (string, any, error) if agent.Status == AgentStatusRunning { if !args.Wait { - return fmt.Sprintf("Agent %s is still running. Use wait=true to block until completion.", args.AgentID), nil, nil + return fmt.Sprintf("Agent %s is still running. Its result will be delivered automatically; use wait=true only when no other work can proceed.", args.AgentID), nil, nil } - // Block until done or context cancelled + waitTimeout := r.waitTimeout + if waitTimeout <= 0 { + waitTimeout = defaultAgentResultWaitTimeout + } + timer := time.NewTimer(waitTimeout) + defer timer.Stop() select { case <-agent.done: + case <-timer.C: + return fmt.Sprintf("Agent %s is still running after waiting %s. Continue other work or check again later; completion will be delivered automatically.", args.AgentID, waitTimeout), agent.Status, nil case <-r.ctx.Done(): - return fmt.Sprintf("Timed out waiting for agent %s", args.AgentID), nil, r.ctx.Err() + return fmt.Sprintf("Stopped waiting for agent %s", args.AgentID), nil, r.ctx.Err() } } @@ -817,7 +832,7 @@ func newGetAgentResultTool(manager *AgentManager, ctx context.Context) ToolDefin &getAgentResultRunner{manager: manager, ctx: ctx}, GetAgentResultArgs{}, "get_agent_result", - "Get the result of a background sub-agent. Set wait=true to block until the agent finishes.", + "Get the result of a background sub-agent. Use wait=false for a non-blocking snapshot. Results are delivered automatically; wait=true blocks briefly and should be used only when no other work can proceed.", ) } @@ -831,10 +846,13 @@ type SendAgentMessageArgs struct { // a live message into a running agent or re-runs a finished agent from its prior // context with the new message appended. type sendAgentMessageRunner struct { - manager *AgentManager - ctx context.Context - llm LLM - subOpts []Option + manager *AgentManager + ctx context.Context + llm LLM + subOpts []Option + messageInjectionChan chan openai.ChatCompletionMessage + completionCB func(*AgentState) + completionFormatter func(*AgentState) string } func (r *sendAgentMessageRunner) Run(args SendAgentMessageArgs) (string, any, error) { @@ -850,32 +868,74 @@ func (r *sendAgentMessageRunner) Run(args SendAgentMessageArgs) (string, any, er return fmt.Sprintf("Message delivered to running agent %s.", args.AgentID), nil, nil } - // Completed/failed: resume by appending the message to the stored fragment and re-running. + // Completed/failed: atomically transition back to running, then resume in + // the background. send_agent_message is fire-and-forget for both live and + // completed agents; completion arrives through the normal callback/notice. + r.manager.mu.Lock() + if agent.Status == AgentStatusRunning { + r.manager.mu.Unlock() + if err := r.manager.Inject(args.AgentID, args.Message); err != nil { + return fmt.Sprintf("Could not message agent %s: %v", args.AgentID, err), nil, nil + } + return fmt.Sprintf("Message delivered to running agent %s.", args.AgentID), nil, nil + } if agent.Fragment == nil { + r.manager.mu.Unlock() return fmt.Sprintf("Agent %s has no stored context to resume", args.AgentID), nil, nil } resumed := agent.Fragment.AddMessage(UserMessageRole, args.Message) - opts := append([]Option{WithContext(r.ctx)}, r.subOpts...) - // Match normal spawn: inherited approval callbacks belong to this child. - opts = append(opts, withAgentIDStamp(args.AgentID)) - result, err := ExecuteTools(r.llm, resumed, opts...) - if err != nil { - return fmt.Sprintf("Resume of agent %s failed: %v", args.AgentID, err), nil, nil - } - r.manager.mu.Lock() - agent.Status = AgentStatusCompleted - agent.Result = result.LastMessage().Content - agent.Fragment = &result + runCtx, cancel := context.WithCancel(context.WithoutCancel(r.ctx)) + agent.Status = AgentStatusRunning + agent.Error = nil + agent.Cancel = cancel + agent.done = make(chan struct{}) + agent.inject = make(chan openai.ChatCompletionMessage, 8) + done := agent.done r.manager.mu.Unlock() - return agent.Result, result, nil + + opts := append([]Option{WithContext(runCtx)}, r.subOpts...) + // Match normal spawn: inherited approval callbacks belong to this child. + opts = append(opts, withAgentIDStamp(args.AgentID), WithMessageInjectionChan(agent.inject)) + go func() { + defer cancel() + result, err := ExecuteTools(r.llm, resumed, opts...) + r.manager.mu.Lock() + if err != nil { + agent.Status = AgentStatusFailed + agent.Error = err + agent.Result = fmt.Sprintf("Failed: %v", err) + } else { + agent.Status = AgentStatusCompleted + agent.Result = result.LastMessage().Content + agent.Fragment = &result + } + r.manager.mu.Unlock() + close(done) + if r.completionCB != nil { + r.completionCB(agent) + } + if r.messageInjectionChan != nil { + notice := openai.ChatCompletionMessage{Role: "user", Content: formatAgentCompletion(agent, r.completionFormatter)} + select { + case r.messageInjectionChan <- notice: + default: + } + } + }() + + return fmt.Sprintf("Agent %s resumed in the background. Its result will be delivered automatically.", args.AgentID), args.AgentID, nil } // newSendAgentMessageTool creates the send_agent_message tool definition. -func newSendAgentMessageTool(manager *AgentManager, ctx context.Context, llm LLM, subOpts []Option) ToolDefinitionInterface { +func newSendAgentMessageTool(manager *AgentManager, ctx context.Context, llm LLM, subOpts []Option, injectionChan chan openai.ChatCompletionMessage, completionCB func(*AgentState), completionFormatter func(*AgentState) string) ToolDefinitionInterface { return NewToolDefinition( - &sendAgentMessageRunner{manager: manager, ctx: ctx, llm: llm, subOpts: subOpts}, + &sendAgentMessageRunner{ + manager: manager, ctx: ctx, llm: llm, subOpts: subOpts, + messageInjectionChan: injectionChan, completionCB: completionCB, + completionFormatter: completionFormatter, + }, SendAgentMessageArgs{}, "send_agent_message", - "Send a follow-up message to a sub-agent. If it is still running the message is injected live; if it has finished, the agent resumes from its prior context.", + "Send a follow-up message to a sub-agent without blocking. Running agents receive it immediately; finished agents resume in the background and deliver their result automatically.", ) } diff --git a/agent_resume_attribution_test.go b/agent_resume_attribution_test.go index b1adfb3..e6eefda 100644 --- a/agent_resume_attribution_test.go +++ b/agent_resume_attribution_test.go @@ -25,6 +25,15 @@ func TestFinishedAgentRepeatedResumePreservesAttribution(t *testing.T) { if _, _, err := runner.Run(SendAgentMessageArgs{AgentID: "child", Message: "continue"}); err != nil { t.Fatal(err) } + agent, ok := manager.Get("child") + if !ok { + t.Fatal("resumed agent missing") + } + select { + case <-agent.done: + case <-ctx.Done(): + t.Fatal("resumed agent did not finish") + } } if count != 2 || len(ids) != 2 { t.Fatalf("executions=%d callbacks=%d", count, len(ids)) diff --git a/agent_resume_test.go b/agent_resume_test.go index 2579f4c..21af553 100644 --- a/agent_resume_test.go +++ b/agent_resume_test.go @@ -36,6 +36,44 @@ func (m *replyLLM) CreateChatCompletion(_ context.Context, _ openai.ChatCompleti }}, LLMUsage{}, nil } +type blockingReplyLLM struct { + started chan struct{} + release chan struct{} + reply string +} + +func (m *blockingReplyLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { + select { + case <-m.started: + default: + close(m.started) + } + select { + case <-m.release: + return f.AddMessage(AssistantMessageRole, m.reply), nil + case <-ctx.Done(): + return f, ctx.Err() + } +} + +func (m *blockingReplyLLM) CreateChatCompletion(ctx context.Context, _ openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) { + select { + case <-m.started: + default: + close(m.started) + } + select { + case <-m.release: + return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{ + Choices: []openai.ChatCompletionChoice{{ + Message: openai.ChatCompletionMessage{Role: AssistantMessageRole.String(), Content: m.reply}, + }}, + }}, LLMUsage{}, nil + case <-ctx.Done(): + return LLMReply{}, LLMUsage{}, ctx.Err() + } +} + func closedChan() chan struct{} { c := make(chan struct{}); close(c); return c } func TestInjectDeliversToRunningAgent(t *testing.T) { @@ -74,7 +112,7 @@ func TestInjectUnknownAgentErrors(t *testing.T) { } } -func TestSendAgentMessageResumesCompletedAgent(t *testing.T) { +func TestSendAgentMessageResumesCompletedAgentWithoutBlocking(t *testing.T) { m := NewAgentManager() frag := NewFragment(openai.ChatCompletionMessage{Role: "user", Content: "first task"}) agent := &AgentState{ @@ -84,14 +122,75 @@ func TestSendAgentMessageResumesCompletedAgent(t *testing.T) { } m.Register(agent) - llm := newReplyLLM("second result") - runner := &sendAgentMessageRunner{manager: m, ctx: context.Background(), llm: llm} - out, _, err := runner.Run(SendAgentMessageArgs{AgentID: "done1", Message: "now do more"}) - if err != nil { - t.Fatalf("resume errored: %v", err) + llm := &blockingReplyLLM{ + started: make(chan struct{}), + release: make(chan struct{}), + reply: "second result", + } + completed := make(chan *AgentState, 1) + notices := make(chan openai.ChatCompletionMessage, 1) + runner := &sendAgentMessageRunner{ + manager: m, + ctx: context.Background(), + llm: llm, + completionCB: func(a *AgentState) { + completed <- a + }, + messageInjectionChan: notices, + } + + returned := make(chan string, 1) + go func() { + out, _, err := runner.Run(SendAgentMessageArgs{AgentID: "done1", Message: "now do more"}) + if err != nil { + returned <- "error: " + err.Error() + return + } + returned <- out + }() + + select { + case out := <-returned: + if !strings.Contains(out, "resumed") || !strings.Contains(out, "done1") { + t.Fatalf("expected immediate resume acknowledgement, got %q", out) + } + case <-time.After(200 * time.Millisecond): + t.Fatal("send_agent_message blocked on the resumed agent") + } + + select { + case <-llm.started: + case <-time.After(time.Second): + t.Fatal("resumed agent did not start") + } + if got := agent.Status; got != AgentStatusRunning { + t.Fatalf("status while resumed = %q, want %q", got, AgentStatusRunning) + } + + close(llm.release) + select { + case <-agent.done: + case <-time.After(time.Second): + t.Fatal("resumed agent did not finish") } - if !strings.Contains(out, "second result") { - t.Fatalf("expected re-run result, got %q", out) + if agent.Status != AgentStatusCompleted || agent.Result != "second result" { + t.Fatalf("resumed state = (%q, %q), want completed second result", agent.Status, agent.Result) + } + select { + case got := <-completed: + if got != agent { + t.Fatal("completion callback received another agent") + } + case <-time.After(time.Second): + t.Fatal("resume completion callback did not fire") + } + select { + case notice := <-notices: + if !strings.Contains(notice.Content, "done1") || !strings.Contains(notice.Content, "second result") { + t.Fatalf("unexpected completion notice %q", notice.Content) + } + case <-time.After(time.Second): + t.Fatal("resume completion notice was not injected") } } @@ -125,6 +224,37 @@ func TestSendAgentMessageUnknownAgent(t *testing.T) { } } +func TestAgentDoneClosesBeforeCompletionCallbackReturns(t *testing.T) { + m := NewAgentManager() + callbackStarted := make(chan struct{}) + callbackRelease := make(chan struct{}) + runner := &spawnAgentRunner{ + llm: newReplyLLM("finished"), + manager: m, + ctx: context.Background(), + agentCompletionCallback: func(*AgentState) { + close(callbackStarted) + <-callbackRelease + }, + } + ctx, cancel := context.WithCancel(context.Background()) + agent := &AgentState{ID: "callback-order", Status: AgentStatusRunning, done: make(chan struct{})} + m.Register(agent) + go runner.runAgent(agent, runner.llm, NewFragment(openai.ChatCompletionMessage{Role: "user", Content: "work"}), nil, AgentRunSpec{}, ctx, cancel) + + select { + case <-callbackStarted: + case <-time.After(time.Second): + t.Fatal("completion callback did not start") + } + select { + case <-agent.done: + case <-time.After(200 * time.Millisecond): + t.Fatal("agent completion remained blocked by completion callback") + } + close(callbackRelease) +} + // Inject must not block the caller: an agent that finished never reads its // channel again, so a blocking send could hang a UI forever. func TestInjectRejectsFinishedAgent(t *testing.T) { diff --git a/agent_test.go b/agent_test.go index 77ca652..3f0dd7e 100644 --- a/agent_test.go +++ b/agent_test.go @@ -302,6 +302,24 @@ var _ = Describe("Sub-Agent Spawning", func() { Expect(result).To(Equal("waited result")) }) + It("should stop waiting after the bounded wait window", func() { + manager := NewAgentManager() + done := make(chan struct{}) + agent := &AgentState{ID: "bounded-wait", Task: "waiting", Status: AgentStatusRunning} + SetAgentDone(agent, done) + manager.Register(agent) + + runner := &GetAgentResultRunnerForTest{ + Manager: manager, Ctx: context.Background(), WaitTimeout: 20 * time.Millisecond, + } + started := time.Now() + result, data, err := runner.Run(GetAgentResultArgs{AgentID: "bounded-wait", Wait: true}) + Expect(err).ToNot(HaveOccurred()) + Expect(time.Since(started)).To(BeNumerically("<", 500*time.Millisecond)) + Expect(result).To(ContainSubstring("still running")) + Expect(data).To(Equal(AgentStatusRunning)) + }) + It("should return status when not waiting for running agent", func() { manager := NewAgentManager() done := make(chan struct{}) diff --git a/tools.go b/tools.go index 1ed9bd7..fd1bb7d 100644 --- a/tools.go +++ b/tools.go @@ -1267,7 +1267,7 @@ func prepareAgentTools(o *Options, llm LLM) []ToolDefinitionInterface { newSpawnAgentTool(agentLLM, o.tools, o.agentManager, o.context, subAgentOpts, o.streamCallback, o.messageInjectionChan, o.agentCompletionCallback, o.agentSpawnCallback, o.agentCompletionFormatter, o.agentDefinitions, o.agentLLMFactory, o.agentDispatcher), newCheckAgentTool(o.agentManager), newGetAgentResultTool(o.agentManager, o.context), - newSendAgentMessageTool(o.agentManager, o.context, agentLLM, subAgentOpts), + newSendAgentMessageTool(o.agentManager, o.context, agentLLM, subAgentOpts, o.messageInjectionChan, o.agentCompletionCallback, o.agentCompletionFormatter), }, o.agentTools) // Explicit selections own the bundled names: replace inherited definitions // rather than registering duplicates with stale runners or managers. The