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 <noreply@anthropic.com>
This commit is contained in:
parent
453f6ca81e
commit
fb9e8a3618
5 changed files with 92 additions and 20 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue