Merge pull request #108 from 73ai/background-agent-refactor

Background agent improvements: retry, env scrub, tasks, budget, reactive triggers

Author: priyanshujain

agent/agent.go

@@ -40,6 +40,7 @@ type Agent struct {
 	summarizer          Summarizer
 	rateLimiter         *provider.RateLimiter
 	usageRecorder       UsageRecorder
+	budgetChecker       BudgetChecker
 }
 
 // Option configures an Agent.
@@ -95,6 +96,11 @@ func WithSummarizer(s Summarizer) Option {
 	return func(a *Agent) { a.summarizer = s }
 }
 
+// WithBudgetChecker sets a budget checker that is called after each LLM call.
+func WithBudgetChecker(bc BudgetChecker) Option {
+	return func(a *Agent) { a.budgetChecker = bc }
+}
+
 // New creates a new Agent.
 func New(p provider.Provider, model string, executor ToolExecutor, opts ...Option) *Agent {
 	a := &Agent{
@@ -143,6 +149,16 @@ func (a *Agent) Run(ctx context.Context, input string) (string, error) {
 		if a.usageRecorder != nil {
 			a.usageRecorder.RecordUsage(a.model, resp.Usage)
 		}
+		if a.budgetChecker != nil {
+			if err := a.budgetChecker.CheckBudget(); err != nil {
+				// Return partial text on budget exceeded (graceful).
+				text := resp.TextContent()
+				if text == "" {
+					text = "(budget exceeded before completion)"
+				}
+				return text, err
+			}
+		}
 
 		// Append assistant response to history.
 		a.history = append(a.history, provider.Message{

agent/agent_test.go

@@ -582,6 +582,55 @@ func TestLoop_TokenBasedCompaction(t *testing.T) {
 	}
 }
 
+func TestLoop_BudgetExceeded(t *testing.T) {
+	mp := &mockProvider{
+		responses: []*provider.ChatResponse{
+			{
+				Content:    []provider.ContentBlock{{Type: provider.ContentText, Text: "Partial result"}},
+				StopReason: provider.StopEndTurn,
+				Usage:      provider.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000},
+			},
+		},
+	}
+	exec := &mockExecutor{results: map[string]string{}}
+	bt := NewBudgetTracker(0.001, nil) // $0.001 budget — 1M tokens of sonnet costs $18
+	a := New(mp, "claude-sonnet-4-6", exec, WithUsageRecorder(bt), WithBudgetChecker(bt))
+
+	result, err := a.Run(context.Background(), "test")
+	if err == nil {
+		t.Fatal("expected budget exceeded error")
+	}
+	if !strings.Contains(err.Error(), "budget exceeded") {
+		t.Errorf("error = %q, want budget exceeded", err.Error())
+	}
+	if result != "Partial result" {
+		t.Errorf("result = %q, want partial text returned", result)
+	}
+}
+
+func TestLoop_BudgetNotExceeded(t *testing.T) {
+	mp := &mockProvider{
+		responses: []*provider.ChatResponse{
+			{
+				Content:    []provider.ContentBlock{{Type: provider.ContentText, Text: "Full result"}},
+				StopReason: provider.StopEndTurn,
+				Usage:      provider.Usage{InputTokens: 100, OutputTokens: 100},
+			},
+		},
+	}
+	exec := &mockExecutor{results: map[string]string{}}
+	bt := NewBudgetTracker(10.0, nil) // $10 budget — more than enough
+	a := New(mp, "claude-sonnet-4-6", exec, WithUsageRecorder(bt), WithBudgetChecker(bt))
+
+	result, err := a.Run(context.Background(), "test")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result != "Full result" {
+		t.Errorf("result = %q, want Full result", result)
+	}
+}
+
 func TestLoop_TracksLastInputTokens(t *testing.T) {
 	mp := &mockProvider{
 		responses: []*provider.ChatResponse{

agent/budget.go

@@ -0,0 +1,58 @@
+package agent
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/73ai/openbotkit/provider"
+)
+
+// BudgetChecker checks whether the accumulated cost has exceeded a budget.
+type BudgetChecker interface {
+	CheckBudget() error
+}
+
+// BudgetTracker wraps a UsageRecorder and accumulates cost per call.
+// It implements both UsageRecorder and BudgetChecker.
+type BudgetTracker struct {
+	maxBudget float64
+	inner     UsageRecorder
+	mu        sync.Mutex
+	total     float64
+}
+
+// NewBudgetTracker creates a tracker that enforces a cost budget.
+// If maxBudget is 0, budget checking is disabled (unlimited).
+func NewBudgetTracker(maxBudget float64, inner UsageRecorder) *BudgetTracker {
+	return &BudgetTracker{maxBudget: maxBudget, inner: inner}
+}
+
+func (bt *BudgetTracker) RecordUsage(model string, usage provider.Usage) {
+	cost := provider.EstimateCost(model, usage)
+	bt.mu.Lock()
+	bt.total += cost
+	bt.mu.Unlock()
+	if bt.inner != nil {
+		bt.inner.RecordUsage(model, usage)
+	}
+}
+
+func (bt *BudgetTracker) CheckBudget() error {
+	if bt.maxBudget <= 0 {
+		return nil
+	}
+	bt.mu.Lock()
+	total := bt.total
+	bt.mu.Unlock()
+	if total >= bt.maxBudget {
+		return fmt.Errorf("budget exceeded: $%.4f spent of $%.4f limit", total, bt.maxBudget)
+	}
+	return nil
+}
+
+// Total returns the accumulated cost so far.
+func (bt *BudgetTracker) Total() float64 {
+	bt.mu.Lock()
+	defer bt.mu.Unlock()
+	return bt.total
+}

agent/budget_test.go

@@ -0,0 +1,61 @@
+package agent
+
+import (
+	"testing"
+
+	"github.com/73ai/openbotkit/provider"
+)
+
+type mockRecorder struct {
+	calls int
+}
+
+func (m *mockRecorder) RecordUsage(_ string, _ provider.Usage) {
+	m.calls++
+}
+
+func TestBudgetTracker_UnderBudget(t *testing.T) {
+	bt := NewBudgetTracker(1.0, nil)
+	bt.RecordUsage("claude-sonnet-4-6", provider.Usage{InputTokens: 1000, OutputTokens: 1000})
+	if err := bt.CheckBudget(); err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestBudgetTracker_ExceedsBudget(t *testing.T) {
+	bt := NewBudgetTracker(0.01, nil)
+	// 1M tokens of sonnet = $18, well over $0.01
+	bt.RecordUsage("claude-sonnet-4-6", provider.Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000})
+	if err := bt.CheckBudget(); err == nil {
+		t.Error("expected budget exceeded error")
+	}
+}
+
+func TestBudgetTracker_Unlimited(t *testing.T) {
+	bt := NewBudgetTracker(0, nil)
+	bt.RecordUsage("claude-opus-4-6", provider.Usage{InputTokens: 10_000_000, OutputTokens: 10_000_000})
+	if err := bt.CheckBudget(); err != nil {
+		t.Errorf("unlimited budget should never error: %v", err)
+	}
+}
+
+func TestBudgetTracker_ChainsInnerRecorder(t *testing.T) {
+	inner := &mockRecorder{}
+	bt := NewBudgetTracker(1.0, inner)
+	bt.RecordUsage("claude-sonnet-4-6", provider.Usage{InputTokens: 100})
+	bt.RecordUsage("claude-sonnet-4-6", provider.Usage{InputTokens: 200})
+	if inner.calls != 2 {
+		t.Errorf("inner.calls = %d, want 2", inner.calls)
+	}
+}
+
+func TestBudgetTracker_Total(t *testing.T) {
+	bt := NewBudgetTracker(1.0, nil)
+	if bt.Total() != 0 {
+		t.Errorf("initial total = %f, want 0", bt.Total())
+	}
+	bt.RecordUsage("claude-sonnet-4-6", provider.Usage{InputTokens: 1000, OutputTokens: 1000})
+	if bt.Total() <= 0 {
+		t.Error("expected positive total after recording usage")
+	}
+}

agent/tools/agent_runner.go

@@ -119,7 +119,7 @@ func (r *AgentRunner) buildArgs(opts runOptions) []string {
 }
 
 func (r *AgentRunner) buildEnv() []string {
-	env := os.Environ()
+	env := scrubEnv(os.Environ())
 	if r.info.Kind == AgentClaude {
 		return filterEnv(env, "CLAUDECODE")
 	}
@@ -137,3 +137,74 @@ func filterEnv(env []string, key string) []string {
 	}
 	return filtered
 }
+
+// scrubEnv removes environment variables that contain sensitive values
+// (API keys, tokens, secrets, passwords) from the given list.
+func scrubEnv(env []string) []string {
+	filtered := make([]string, 0, len(env))
+	for _, e := range env {
+		eqIdx := strings.IndexByte(e, '=')
+		if eqIdx < 0 {
+			filtered = append(filtered, e)
+			continue
+		}
+		key := e[:eqIdx]
+		if !isSensitiveKey(key) {
+			filtered = append(filtered, e)
+		}
+	}
+	return filtered
+}
+
+var safeKeys = map[string]bool{
+	"PATH": true, "HOME": true, "USER": true, "SHELL": true,
+	"TERM": true, "LANG": true, "TMPDIR": true, "PWD": true,
+	"GOPATH": true, "GOROOT": true, "GOBIN": true,
+	"EDITOR": true, "VISUAL": true, "PAGER": true,
+	"HOSTNAME": true, "LOGNAME": true, "DISPLAY": true,
+	"COLORTERM": true, "TERM_PROGRAM": true, "SHLVL": true,
+}
+
+var safePrefixes = []string{"LC_", "XDG_"}
+
+var sensitivePrefixes = []string{
+	"ANTHROPIC_", "OPENAI_", "GOOGLE_API_", "GEMINI_",
+	"GROQ_", "OPENROUTER_", "AWS_SECRET_", "CLAUDECODE",
+}
+
+var sensitiveSuffixes = []string{
+	"_KEY", "_SECRET", "_TOKEN", "_PASSWORD", "_CREDENTIAL", "_AUTH",
+	"_PRIVATE_KEY", "_DSN",
+}
+
+var sensitiveExact = map[string]bool{
+	"GITHUB_TOKEN": true, "GH_TOKEN": true,
+	"DATABASE_URL": true, "REDIS_URL": true, "MONGODB_URI": true,
+	"AMQP_URL": true, "ELASTICSEARCH_URL": true,
+}
+
+func isSensitiveKey(key string) bool {
+	if safeKeys[key] {
+		return false
+	}
+	for _, p := range safePrefixes {
+		if strings.HasPrefix(key, p) {
+			return false
+		}
+	}
+	if sensitiveExact[key] {
+		return true
+	}
+	upper := strings.ToUpper(key)
+	for _, p := range sensitivePrefixes {
+		if strings.HasPrefix(upper, p) {
+			return true
+		}
+	}
+	for _, s := range sensitiveSuffixes {
+		if strings.HasSuffix(upper, s) {
+			return true
+		}
+	}
+	return false
+}