Merge pull request #8 from dj-oyu/fix/rate-limit-and-heartbeat-lifecycle

fix: add LLM rate limiting and fix heartbeat subagent lifecycle

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

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) {
+		if opts.Background {
-			commitMsg := "heartbeat: auto-save"
+			if agent.SubagentMgr != nil {
-			wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false)
+				agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout
-			if wtResult != nil && wtResult.CommitsAhead > 0 && !constants.IsInternalChannel(opts.Channel) {
+			}
-				cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second)
+			if agent.IsInWorktree(opts.SessionKey) {
-				_ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{
+				commitMsg := "heartbeat: auto-save"
-					Channel: opts.Channel,
+				wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false)
-					ChatID:  opts.ChatID,
+				if wtResult != nil && wtResult.CommitsAhead > 0 && !constants.IsInternalChannel(opts.Channel) {
-					Content: fmt.Sprintf("Heartbeat made code changes on branch `%s` (%d commits).",
+					cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second)
-						wtResult.Branch, wtResult.CommitsAhead),
+					_ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{
-				})
+						Channel: opts.Channel,
-				cleanupCancel()
+						ChatID:  opts.ChatID,
+						Content: fmt.Sprintf("Heartbeat made code changes on branch `%s` (%d commits).",
+							wtResult.Branch, wtResult.CommitsAhead),
+					})
+					cleanupCancel()
+				}
 			}
 		}
 	}()

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(
+		return NewHTTPProviderWithOptions(cfg.APIKey, apiBase, cfg.Proxy,
-			cfg.APIKey,
+			openai_compat.WithMaxTokensField(cfg.MaxTokensField),
-			apiBase,
+			openai_compat.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second),
-			cfg.Proxy,
+			openai_compat.WithMinInterval(rpmToMinInterval(cfg.RPM)),
-			cfg.MaxTokensField,
-			cfg.RequestTimeout,
 		), 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 {

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)

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
@@ -28,11 +32,13 @@ type SubagentTask struct {
 	Iterations    int            `json:"-"`
 	ToolCalls     int            `json:"-"`
 	ToolStats     map[string]int `json:"-"`
+	cancel        context.CancelFunc
 }
 
 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 +128,17 @@ func (sm *SubagentManager) Spawn(
 
 	sm.reporter.ReportSpawn(taskID, label, task)
 
-	// Start task in background with context cancellation support
+	// Start task in background with a detached context that has a hard timeout.
-	go sm.runTask(ctx, subagentTask, preset, callback)
+	// 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.WithTimeout(context.Background(), spawnTimeout)
+	subagentTask.cancel = spawnCancel
+	sm.wg.Add(1)
+	go func() {
+		defer sm.wg.Done()
+		sm.runTask(spawnCtx, subagentTask, preset, callback)
+	}()
 
 	if label != "" {
 		return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil
@@ -366,6 +381,33 @@ func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string)
 	return registry
 }
 
+// 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.
+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()

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())
+	_, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil)
-	defer cancel()
-
-	_, err := mgr.Spawn(ctx, "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 via CancelTask — the spawned goroutine's detached context is canceled.
-	cancel()
+	mgr.CancelTask("subagent-1")
 
 	// Collect events until agent_gc.
 	var events []orch.Event