From fb9e8a3618e67bf64d08fb66ec3249d5eb37fe21 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 1 Mar 2026 05:09:02 +0900 Subject: [PATCH 1/3] fix: add LLM rate limiting and fix heartbeat subagent lifecycle Two issues addressed: 1. Per-provider rate limiting via RPM config field. When multiple subagents share the same provider instance, concurrent API calls can trigger rate limit errors. WithMinInterval() enforces a minimum wait between requests, derived from the existing model_list rpm setting. 2. Spawned subagents now outlive their parent session. Previously, heartbeat runAgentLoop would cancel its child context on exit, killing any spawned subagents mid-flight and cleaning up worktrees they were still using. Fix: spawn goroutines use context.Background() and the heartbeat cleanup defer waits for SubagentManager.WaitAll() before deactivating worktrees. Co-Authored-By: Claude Opus 4.6 --- pkg/agent/instance.go | 4 +++ pkg/agent/loop.go | 31 ++++++++++++-------- pkg/providers/factory_provider.go | 22 ++++++++++---- pkg/providers/openai_compat/provider.go | 38 +++++++++++++++++++++++++ pkg/tools/subagent.go | 17 +++++++++-- 5 files changed, 92 insertions(+), 20 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 5403c190e..1ba2faea1 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -40,6 +40,10 @@ type AgentInstance struct { PlanFallbacks []string PlanCandidates []providers.FallbackCandidate + // SubagentMgr is set during registerSharedTools when orchestration is enabled. + // Used by runAgentLoop to wait for spawned subagents before worktree cleanup. + SubagentMgr *tools.SubagentManager + // Interview staleness tracking: consecutive turns where MEMORY.md was not updated. interviewStaleCount int interviewMemoryLen int diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 066a53747..bc643b6b1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -305,6 +305,7 @@ func registerSharedTools( webSearchOpts, ) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + agent.SubagentMgr = subagentManager spawnTool := tools.NewSpawnTool(subagentManager) currentAgentID := agentID spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { @@ -942,19 +943,25 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt } // Guarantee heartbeat worktree cleanup on ALL exit paths (error, panic, normal). + // Wait for spawned subagents first so they aren't killed mid-flight. defer func() { - if opts.Background && agent.IsInWorktree(opts.SessionKey) { - commitMsg := "heartbeat: auto-save" - wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false) - if wtResult != nil && wtResult.CommitsAhead > 0 && !constants.IsInternalChannel(opts.Channel) { - cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) - _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: fmt.Sprintf("Heartbeat made code changes on branch `%s` (%d commits).", - wtResult.Branch, wtResult.CommitsAhead), - }) - cleanupCancel() + if opts.Background { + if agent.SubagentMgr != nil { + agent.SubagentMgr.WaitAll() + } + if agent.IsInWorktree(opts.SessionKey) { + commitMsg := "heartbeat: auto-save" + wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false) + if wtResult != nil && wtResult.CommitsAhead > 0 && !constants.IsInternalChannel(opts.Channel) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: fmt.Sprintf("Heartbeat made code changes on branch `%s` (%d commits).", + wtResult.Branch, wtResult.CommitsAhead), + }) + cleanupCancel() + } } } }() diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index cc5d905ad..8c6b5b692 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -90,6 +90,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err openai_compat.WithMaxTokensField(cfg.MaxTokensField), openai_compat.WithStream(boolDefault(cfg.Stream, false)), openai_compat.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second), + openai_compat.WithMinInterval(rpmToMinInterval(cfg.RPM)), ), modelID, nil case "minimax": @@ -106,6 +107,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err openai_compat.WithMaxTokensField(cfg.MaxTokensField), openai_compat.WithStream(boolDefault(cfg.Stream, true)), openai_compat.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second), + openai_compat.WithMinInterval(rpmToMinInterval(cfg.RPM)), ), modelID, nil case "openrouter", "groq", "zhipu", "gemini", "nvidia", @@ -123,6 +125,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err openai_compat.WithMaxTokensField(cfg.MaxTokensField), openai_compat.WithStream(boolDefault(cfg.Stream, false)), openai_compat.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second), + openai_compat.WithMinInterval(rpmToMinInterval(cfg.RPM)), ), modelID, nil case "anthropic": @@ -142,12 +145,10 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err if cfg.APIKey == "" { return nil, "", fmt.Errorf("api_key is required for anthropic protocol (model: %s)", cfg.Model) } - return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( - cfg.APIKey, - apiBase, - cfg.Proxy, - cfg.MaxTokensField, - cfg.RequestTimeout, + return NewHTTPProviderWithOptions(cfg.APIKey, apiBase, cfg.Proxy, + openai_compat.WithMaxTokensField(cfg.MaxTokensField), + openai_compat.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second), + openai_compat.WithMinInterval(rpmToMinInterval(cfg.RPM)), ), modelID, nil case "antigravity": @@ -227,6 +228,15 @@ func getDefaultAPIBase(protocol string) string { } } +// rpmToMinInterval converts a requests-per-minute limit to a minimum interval +// between consecutive requests. Returns 0 (no throttle) when rpm <= 0. +func rpmToMinInterval(rpm int) time.Duration { + if rpm <= 0 { + return 0 + } + return time.Minute / time.Duration(rpm) +} + // boolDefault dereferences a *bool, returning def when nil. func boolDefault(p *bool, def bool) bool { if p != nil { diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 2ed956991..453df1daf 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -11,6 +11,7 @@ import ( "net/http" "net/url" "strings" + "sync" "time" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" @@ -36,6 +37,12 @@ type Provider struct { maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models) stream bool // Use SSE streaming internally (accumulates into a single LLMResponse) httpClient *http.Client + + // Rate limiting: minimum interval between consecutive API requests. + // Shared across all goroutines using this provider instance. + mu sync.Mutex + lastRequestAt time.Time + minInterval time.Duration } // Option is a functional option for configuring a Provider. @@ -69,6 +76,14 @@ func WithStream(stream bool) Option { } } +// WithMinInterval sets the minimum interval between consecutive API requests. +// This prevents rate limit errors when many subagents share the same provider. +func WithMinInterval(d time.Duration) Option { + return func(p *Provider) { + p.minInterval = d + } +} + // WithEndpointPath sets the API path appended to apiBase (default: "/chat/completions"). func WithEndpointPath(path string) Option { return func(p *Provider) { @@ -211,6 +226,25 @@ func (p *Provider) buildHTTPRequest( return req, nil } +// waitForInterval enforces the minimum interval between consecutive API requests. +// It sleeps if needed, then records the current time as the last request time. +func (p *Provider) waitForInterval() { + if p.minInterval <= 0 { + return + } + p.mu.Lock() + if !p.lastRequestAt.IsZero() { + elapsed := time.Since(p.lastRequestAt) + if wait := p.minInterval - elapsed; wait > 0 { + p.mu.Unlock() + time.Sleep(wait) + p.mu.Lock() + } + } + p.lastRequestAt = time.Now() + p.mu.Unlock() +} + func (p *Provider) Chat( ctx context.Context, messages []Message, @@ -233,6 +267,8 @@ func (p *Provider) Chat( return nil, err } + p.waitForInterval() + resp, err := p.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("failed to send request: %w", err) @@ -272,6 +308,8 @@ func (p *Provider) ChatStream( return nil, err } + p.waitForInterval() + resp, err := p.httpClient.Do(req) //nolint:bodyclose // closed in goroutine or error path below if err != nil { return nil, fmt.Errorf("failed to send request: %w", err) diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 5dd9853ee..f4acb4a4f 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -33,6 +33,7 @@ type SubagentTask struct { type SubagentManager struct { tasks map[string]*SubagentTask mu sync.RWMutex + wg sync.WaitGroup // tracks running spawn goroutines provider providers.LLMProvider defaultModel string bus *bus.MessageBus @@ -122,8 +123,14 @@ func (sm *SubagentManager) Spawn( sm.reporter.ReportSpawn(taskID, label, task) - // Start task in background with context cancellation support - go sm.runTask(ctx, subagentTask, preset, callback) + // Start task in background with a detached context. + // The spawned goroutine must outlive the parent (e.g. heartbeat session) + // which may finish before the subagent completes. + sm.wg.Add(1) + go func() { + defer sm.wg.Done() + sm.runTask(context.Background(), subagentTask, preset, callback) + }() if label != "" { return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil @@ -366,6 +373,12 @@ func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string) return registry } +// WaitAll blocks until all spawned subagent goroutines have finished. +// Used by heartbeat cleanup to avoid destroying worktrees while subagents are still running. +func (sm *SubagentManager) WaitAll() { + sm.wg.Wait() +} + func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) { sm.mu.RLock() defer sm.mu.RUnlock() From b7ed441ef4c742d81591bf66933fefc0e770ba3e Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 1 Mar 2026 05:14:51 +0900 Subject: [PATCH 2/3] fix: update cancel test for detached spawn context Spawn now uses a detached context. Add CancelTask() method with per-task cancel func so subagents can still be explicitly stopped. Update TestSubagentManager_Spawn_CancelledDuringExecution to use CancelTask() instead of parent context cancellation. Co-Authored-By: Claude Opus 4.6 --- pkg/tools/subagent.go | 16 +++++++++++++++- pkg/tools/subagent_reporter_test.go | 9 +++------ 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index f4acb4a4f..79b6300e2 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -28,6 +28,7 @@ type SubagentTask struct { Iterations int `json:"-"` ToolCalls int `json:"-"` ToolStats map[string]int `json:"-"` + cancel context.CancelFunc } type SubagentManager struct { @@ -126,10 +127,13 @@ func (sm *SubagentManager) Spawn( // Start task in background with a detached context. // The spawned goroutine must outlive the parent (e.g. heartbeat session) // which may finish before the subagent completes. + // The cancel func is stored on the task so CancelTask() can stop it. + spawnCtx, spawnCancel := context.WithCancel(context.Background()) + subagentTask.cancel = spawnCancel sm.wg.Add(1) go func() { defer sm.wg.Done() - sm.runTask(context.Background(), subagentTask, preset, callback) + sm.runTask(spawnCtx, subagentTask, preset, callback) }() if label != "" { @@ -379,6 +383,16 @@ func (sm *SubagentManager) WaitAll() { sm.wg.Wait() } +// CancelTask cancels the context for a running subagent task. +func (sm *SubagentManager) CancelTask(taskID string) { + sm.mu.RLock() + task, ok := sm.tasks[taskID] + sm.mu.RUnlock() + if ok && task.cancel != nil { + task.cancel() + } +} + func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) { sm.mu.RLock() defer sm.mu.RUnlock() diff --git a/pkg/tools/subagent_reporter_test.go b/pkg/tools/subagent_reporter_test.go index f378b3060..0dcc2d6ac 100644 --- a/pkg/tools/subagent_reporter_test.go +++ b/pkg/tools/subagent_reporter_test.go @@ -196,10 +196,7 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { bp := newBlockingProvider() mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - _, err := mgr.Spawn(ctx, "long task", "cancel-me", "", "cli", "direct", "", nil) + _, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil) if err != nil { t.Fatalf("Spawn() error: %v", err) } @@ -211,8 +208,8 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { t.Fatal("timed out waiting for blockingProvider to enter Chat") } - // Now cancel — the LLM call unblocks with ctx.Err(). - cancel() + // Cancel via CancelTask — the spawned goroutine's detached context is canceled. + mgr.CancelTask("subagent-1") // Collect events until agent_gc. var events []orch.Event From 0b603d74039e327d649b5a06e3ca31ba9c9aca4d Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 1 Mar 2026 05:17:06 +0900 Subject: [PATCH 3/3] fix: add spawn timeout and WaitAll deadline to prevent infinite hangs - Spawn goroutines now have a 30-minute hard timeout (spawnTimeout) - WaitAll accepts a timeout parameter so heartbeat cleanup doesn't block forever if a subagent gets stuck - CancelTask stores per-task cancel func for explicit cancellation Co-Authored-By: Claude Opus 4.6 --- pkg/agent/loop.go | 2 +- pkg/tools/subagent.go | 27 +++++++++++++++++++++------ 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index bc643b6b1..218ef14c5 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -947,7 +947,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt defer func() { if opts.Background { if agent.SubagentMgr != nil { - agent.SubagentMgr.WaitAll() + agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout } if agent.IsInWorktree(opts.SessionKey) { commitMsg := "heartbeat: auto-save" diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 79b6300e2..beecb20b4 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -14,6 +14,10 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" ) +// spawnTimeout is the hard upper bound for a single spawn goroutine. +// MaxIterations × HTTP timeout provides the soft limit; this is a safety net. +const spawnTimeout = 30 * time.Minute + type SubagentTask struct { ID string Task string @@ -124,11 +128,11 @@ func (sm *SubagentManager) Spawn( sm.reporter.ReportSpawn(taskID, label, task) - // Start task in background with a detached context. + // Start task in background with a detached context that has a hard timeout. // The spawned goroutine must outlive the parent (e.g. heartbeat session) // which may finish before the subagent completes. // The cancel func is stored on the task so CancelTask() can stop it. - spawnCtx, spawnCancel := context.WithCancel(context.Background()) + spawnCtx, spawnCancel := context.WithTimeout(context.Background(), spawnTimeout) subagentTask.cancel = spawnCancel sm.wg.Add(1) go func() { @@ -377,10 +381,21 @@ func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string) return registry } -// WaitAll blocks until all spawned subagent goroutines have finished. -// Used by heartbeat cleanup to avoid destroying worktrees while subagents are still running. -func (sm *SubagentManager) WaitAll() { - sm.wg.Wait() +// WaitAll blocks until all spawned subagent goroutines have finished +// or the timeout expires. Returns true if all goroutines finished, +// false on timeout. +func (sm *SubagentManager) WaitAll(timeout time.Duration) bool { + done := make(chan struct{}) + go func() { + sm.wg.Wait() + close(done) + }() + select { + case <-done: + return true + case <-time.After(timeout): + return false + } } // CancelTask cancels the context for a running subagent task.