opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
fix: add LLM rate limiting and fix heartbeat subagent lifecycle
This commit is contained in:
dj-oyu 2026-03-01 05:20:03 +09:00 • committed by GitHub
commit 0033e621ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 124 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
pkg/agent/instance.go
View file

@ -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

9
pkg/agent/loop.go
View file

@ -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,8 +943,13 @@ 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 {
if agent.SubagentMgr != nil {
agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout
}
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) {
@ -957,6 +963,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
cleanupCancel()
}
}
}
}()
// For background tasks (cron/heartbeat), generate a TaskID and send notification

22
pkg/providers/factory_provider.go
View file

@ -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 {

38
pkg/providers/openai_compat/provider.go
View file

@ -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)

46
pkg/tools/subagent.go
View file

@ -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
go sm.runTask(ctx, subagentTask, preset, callback)
// 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.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()

9
pkg/tools/subagent_reporter_test.go
View file

@ -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