From 99b189d3fb9090ef4dc031cdefd5f54ef7b07bba Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Thu, 19 Mar 2026 12:38:18 +0800 Subject: [PATCH] feat(subturn): implement token budget tracking for SubTurns --- pkg/agent/loop.go | 4 ++++ pkg/agent/subturn.go | 51 ++++++++++++++++++++++++++++++++++++++++- pkg/agent/turn_state.go | 45 ++++++++++++++++++++++++++++-------- pkg/tools/subagent.go | 2 ++ 4 files changed, 92 insertions(+), 10 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index e97fb14ff..6adaa423d 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1460,6 +1460,10 @@ func (al *AgentLoop) runLLMIteration( // Save finishReason to turnState for SubTurn truncation detection if ts := turnStateFromContext(ctx); ts != nil { ts.SetLastFinishReason(response.FinishReason) + // Save usage for token budget tracking + if response.Usage != nil { + ts.SetLastUsage(response.Usage) + } } go al.handleReasoning( diff --git a/pkg/agent/subturn.go b/pkg/agent/subturn.go index 78e55edc8..b8d986841 100644 --- a/pkg/agent/subturn.go +++ b/pkg/agent/subturn.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "strings" + "sync/atomic" "time" "github.com/sipeed/picoclaw/pkg/logger" @@ -127,6 +128,12 @@ type SubTurnConfig struct { // Used by evaluator-optimizer patterns to pass the full worker context across multiple iterations. InitialMessages []providers.Message + // InitialTokenBudget is a shared atomic counter for tracking remaining tokens. + // If set, the SubTurn will inherit this budget and deduct tokens after each LLM call. + // If nil, the SubTurn will inherit the parent's tokenBudget (if any). + // Used by team tool to enforce token limits across all team members. + InitialTokenBudget *atomic.Int64 + // Can be extended with temperature, topP, etc. } @@ -199,6 +206,7 @@ func (s *AgentLoopSpawner) SpawnSubTurn(ctx context.Context, cfg tools.SubTurnCo SystemPrompt: cfg.SystemPrompt, ActualSystemPrompt: cfg.ActualSystemPrompt, InitialMessages: cfg.InitialMessages, + InitialTokenBudget: cfg.InitialTokenBudget, MaxTokens: cfg.MaxTokens, Async: cfg.Async, Critical: cfg.Critical, @@ -292,6 +300,15 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S childTS.cancelFunc = cancel childTS.critical = cfg.Critical + // Token budget initialization/inheritance + // If InitialTokenBudget is explicitly provided (e.g., by team tool), use it. + // Otherwise, inherit from parent's tokenBudget (for nested SubTurns). + if cfg.InitialTokenBudget != nil { + childTS.tokenBudget = cfg.InitialTokenBudget + } else if parentTS.tokenBudget != nil { + childTS.tokenBudget = parentTS.tokenBudget + } + // IMPORTANT: Put childTS into childCtx so that code inside runTurn can retrieve it childCtx = withTurnState(childCtx, childTS) childCtx = WithAgentLoop(childCtx, al) // Propagate AgentLoop to child turn @@ -619,7 +636,39 @@ func runTurn(ctx context.Context, al *AgentLoop, ts *turnState, cfg SubTurnConfi continue // Retry with recovery prompt } - // 3. Success - return result with session history + // 3. Token budget enforcement (if configured) + // Check if budget is exhausted after this LLM call. If so, return gracefully + // with current result instead of continuing iterations. + if ts.tokenBudget != nil { + if usage := ts.GetLastUsage(); usage != nil { + newBudget := ts.tokenBudget.Add(-int64(usage.TotalTokens)) + + if newBudget <= 0 { + logger.WarnCF("subturn", "Token budget exhausted", + map[string]any{ + "turn_id": ts.turnID, + "deficit": -newBudget, + "tokens_used": usage.TotalTokens, + "final_budget": newBudget, + }) + + // Budget exhausted - return current result with marker + return &tools.ToolResult{ + ForLLM: finalContent + "\n\n[Token budget exhausted]", + Messages: childAgent.Sessions.GetHistory(ts.turnID), + }, nil + } + + logger.DebugCF("subturn", "Token budget updated", + map[string]any{ + "turn_id": ts.turnID, + "tokens_used": usage.TotalTokens, + "remaining_budget": newBudget, + }) + } + } + + // 4. Success - return result with session history return &tools.ToolResult{ ForLLM: finalContent, Messages: childAgent.Sessions.GetHistory(ts.turnID), diff --git a/pkg/agent/turn_state.go b/pkg/agent/turn_state.go index d5c98ff7f..1f7716ec7 100644 --- a/pkg/agent/turn_state.go +++ b/pkg/agent/turn_state.go @@ -67,6 +67,17 @@ type turnState struct { // Used by SubTurn to detect truncation and retry. // MUST be accessed under mu lock. lastFinishReason string + + // Token budget tracking + // tokenBudget is a shared atomic counter for tracking remaining tokens across team members. + // Inherited from parent or initialized from SubTurnConfig.InitialTokenBudget. + // Nil if no budget is set. + tokenBudget *atomic.Int64 + + // lastUsage stores the token usage from the last LLM call. + // Used by SubTurn to deduct from tokenBudget after each LLM iteration. + // MUST be accessed under mu lock. + lastUsage *providers.UsageInfo } // ====================== Public API ====================== @@ -134,7 +145,7 @@ func (al *AgentLoop) FormatTree(turnInfo *TurnInfo, prefix string, isLast bool) } var sb strings.Builder - + // Print current node marker := "├── " if isLast { @@ -154,7 +165,7 @@ func (al *AgentLoop) FormatTree(turnInfo *TurnInfo, prefix string, isLast bool) orphanMarker = " (Orphaned)" } - sb.WriteString(fmt.Sprintf("%s%s[%s] Depth:%d (%s)%s\n", prefix, marker, turnInfo.TurnID, turnInfo.Depth, status, orphanMarker)) + fmt.Fprintf(&sb, "%s%s[%s] Depth:%d (%s)%s\n", prefix, marker, turnInfo.TurnID, turnInfo.Depth, status, orphanMarker) // Prepare prefix for children childPrefix := prefix @@ -179,7 +190,7 @@ func (al *AgentLoop) FormatTree(turnInfo *TurnInfo, prefix string, isLast bool) if isLastChild { cMarker = "└── " } - sb.WriteString(fmt.Sprintf("%s%s[%s] (Completed/Cleaned Up)\n", childPrefix, cMarker, childID)) + fmt.Fprintf(&sb, "%s%s[%s] (Completed/Cleaned Up)\n", childPrefix, cMarker, childID) } } @@ -193,12 +204,12 @@ func newTurnState(ctx context.Context, id string, parent *turnState) *turnState // (spawnSubTurn) already creates one. The turnState stores the context and // cancelFunc provided by the caller to avoid redundant context wrapping. return &turnState{ - ctx: ctx, - cancelFunc: nil, // Will be set by the caller - turnID: id, - parentTurnID: parent.turnID, - depth: parent.depth + 1, - session: newEphemeralSession(parent.session), + ctx: ctx, + cancelFunc: nil, // Will be set by the caller + turnID: id, + parentTurnID: parent.turnID, + depth: parent.depth + 1, + session: newEphemeralSession(parent.session), parentTurnState: parent, // Store reference to parent for IsParentEnded() checks // NOTE: In this PoC, I use a fixed-size channel (16). // Under high concurrency or long-running sub-turns, this might fill up and cause @@ -233,6 +244,22 @@ func (ts *turnState) GetLastFinishReason() string { return ts.lastFinishReason } +// SetLastUsage stores the token usage from the last LLM call. +// This is used by SubTurn to track token consumption for budget enforcement. +func (ts *turnState) SetLastUsage(usage *providers.UsageInfo) { + ts.mu.Lock() + defer ts.mu.Unlock() + ts.lastUsage = usage +} + +// GetLastUsage retrieves the token usage from the last LLM call. +// Returns nil if no LLM call has been made yet. +func (ts *turnState) GetLastUsage() *providers.UsageInfo { + ts.mu.Lock() + defer ts.mu.Unlock() + return ts.lastUsage +} + // IsParentEnded is a convenience method to check if parent ended. // It returns the value of the parent's parentEnded atomic flag. diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 297fb13a5..39356cb1e 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "sync" + "sync/atomic" "time" "github.com/sipeed/picoclaw/pkg/providers" @@ -28,6 +29,7 @@ type SubTurnConfig struct { MaxContextRunes int // 0 = auto, -1 = no limit, >0 = explicit limit ActualSystemPrompt string InitialMessages []providers.Message + InitialTokenBudget *atomic.Int64 // Shared token budget for team members; nil if no budget } type SubagentTask struct {