clawgo/pkg/tools/subagent_runtime_context.go

package tools

import "context"

type SubagentExecutionStats struct {
	Iterations       int
	Attempts         int
	Restarts         int
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
	FailureCode      string
}

type subagentExecutionStatsKey struct{}
type subagentIterationBudgetKey struct{}

func WithSubagentExecutionStats(ctx context.Context, stats *SubagentExecutionStats) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, subagentExecutionStatsKey{}, stats)
}

func RecordSubagentExecutionStats(ctx context.Context, delta SubagentExecutionStats) {
	if ctx == nil {
		return
	}
	stats, _ := ctx.Value(subagentExecutionStatsKey{}).(*SubagentExecutionStats)
	if stats == nil {
		return
	}
	stats.Iterations += delta.Iterations
	stats.Attempts += delta.Attempts
	stats.Restarts += delta.Restarts
	stats.PromptTokens += delta.PromptTokens
	stats.CompletionTokens += delta.CompletionTokens
	stats.TotalTokens += delta.TotalTokens
	if delta.FailureCode != "" {
		stats.FailureCode = delta.FailureCode
	}
}

func WithSubagentIterationBudget(ctx context.Context, budget int) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, subagentIterationBudgetKey{}, budget)
}

func SubagentIterationBudget(ctx context.Context) (int, bool) {
	if ctx == nil {
		return 0, false
	}
	budget, ok := ctx.Value(subagentIterationBudgetKey{}).(int)
	return budget, ok && budget > 0
}