feat(subturn): implement token budget tracking for SubTurns

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(

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

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.
 

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 {