feat(adk): PreserveSkillsConfig support MaxTokensPerSkill and SkillsTokenBudget (#967)

Author: mrh997

adk/middlewares/summarization/finalizer_builder.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"unicode/utf8"
 
 	"github.com/bytedance/sonic"
 
@@ -97,16 +98,30 @@ func (b *FinalizerBuilder) Build() (FinalizeFunc, error) {
 }
 
 type PreserveSkillsConfig struct {
+	// SkillToolName is the tool name used for loading skills.
+	// Must match the tool name configured in the ADK skill middleware.
+	// Optional. Defaults to "skill".
+	SkillToolName string
+
 	// MaxSkills limits the maximum number of skills to preserve.
 	// = 0 means do not preserve any skills (disabled).
 	// > 0 means preserve up to this many most recent skills.
 	// Optional. Defaults to 5.
 	MaxSkills *int
 
-	// SkillToolName is the tool name used for loading skills.
-	// Must match the tool name configured in the ADK skill middleware.
-	// Optional. Defaults to "skill".
-	SkillToolName string
+	// MaxTokensPerSkill limits the maximum token count for a single preserved skill.
+	// Skills exceeding this limit are truncated, with the truncated portion replaced
+	// by a short marker text (e.g. "[... skill content truncated ...]").
+	// Note: if this value is set smaller than the token count of the marker text itself,
+	// the skill will contain only the marker text with no original content preserved.
+	// Optional. Defaults to 5000.
+	MaxTokensPerSkill *int
+
+	// SkillsTokenBudget limits the total token count for all preserved skills combined.
+	// Skills are preserved from most recent to oldest; once the budget is exhausted,
+	// remaining skills are dropped.
+	// Optional. Defaults to 25000.
+	SkillsTokenBudget *int
 }
 
 // PreserveSkills extracts skill contents loaded by the ADK skill middleware
@@ -149,6 +164,12 @@ func (c *PreserveSkillsConfig) check() error {
 	if c.MaxSkills != nil && *c.MaxSkills < 0 {
 		return fmt.Errorf("MaxSkills must be non-negative")
 	}
+	if c.MaxTokensPerSkill != nil && *c.MaxTokensPerSkill < 0 {
+		return fmt.Errorf("MaxTokensPerSkill must be non-negative")
+	}
+	if c.SkillsTokenBudget != nil && *c.SkillsTokenBudget < 0 {
+		return fmt.Errorf("SkillsTokenBudget must be non-negative")
+	}
 	return nil
 }
 
@@ -158,7 +179,11 @@ type skillInfo struct {
 }
 
 func buildPreservedSkillsText(_ context.Context, messages []adk.Message, config *PreserveSkillsConfig) (string, error) {
-	const defaultSkillTool = "skill"
+	const (
+		defaultSkillTool         = "skill"
+		defaultMaxTokensPerSkill = 5000
+		defaultSkillsTokenBudget = 25000
+	)
 
 	if config == nil {
 		config = &PreserveSkillsConfig{}
@@ -177,6 +202,16 @@ func buildPreservedSkillsText(_ context.Context, messages []adk.Message, config
 		skillTool = config.SkillToolName
 	}
 
+	maxTokensPerSkill := defaultMaxTokensPerSkill
+	if config.MaxTokensPerSkill != nil {
+		maxTokensPerSkill = *config.MaxTokensPerSkill
+	}
+
+	skillsTokenBudget := defaultSkillsTokenBudget
+	if config.SkillsTokenBudget != nil {
+		skillsTokenBudget = *config.SkillsTokenBudget
+	}
+
 	var skills []*skillInfo
 	argsMap := make(map[string]string)
 
@@ -188,6 +223,7 @@ func buildPreservedSkillsText(_ context.Context, messages []adk.Message, config
 					argsMap[tc.ID] = tc.Function.Arguments
 				}
 			}
+
 		case schema.Tool:
 			arguments, ok := argsMap[msg.ToolCallID]
 			if !ok {
@@ -231,8 +267,39 @@ func buildPreservedSkillsText(_ context.Context, messages []adk.Message, config
 		skills = skills[len(skills)-maxSkills:]
 	}
 
+	totalTokens := 0
+	var budgetedSkills []*skillInfo
+	for i := len(skills) - 1; i >= 0; i-- {
+		skill := skills[i]
+		tokens := estimateTokenCount(skill.Content)
+
+		if tokens > maxTokensPerSkill {
+			skill = &skillInfo{
+				Name:    skill.Name,
+				Content: truncateSkillContent(skill.Content, maxTokensPerSkill),
+			}
+			tokens = maxTokensPerSkill
+		}
+
+		if totalTokens+tokens > skillsTokenBudget {
+			break
+		}
+
+		totalTokens += tokens
+		budgetedSkills = append(budgetedSkills, skill)
+	}
+
+	if len(budgetedSkills) == 0 {
+		return "", nil
+	}
+
+	// Reverse to restore chronological order.
+	for i, j := 0, len(budgetedSkills)-1; i < j; i, j = i+1, j-1 {
+		budgetedSkills[i], budgetedSkills[j] = budgetedSkills[j], budgetedSkills[i]
+	}
+
 	var parts []string
-	for _, skill := range skills {
+	for _, skill := range budgetedSkills {
 		parts = append(parts, fmt.Sprintf(skillSectionFormat, skill.Name, skill.Content))
 	}
 
@@ -242,3 +309,33 @@ func buildPreservedSkillsText(_ context.Context, messages []adk.Message, config
 
 	return skillsText, nil
 }
+
+// truncateSkillContent truncates skill content to fit within maxTokens.
+// It keeps the first portion of the content and appends a truncation marker
+// (e.g. "[... skill content truncated ...]") to indicate the omission.
+// If maxTokens is smaller than the marker itself, only the marker is returned.
+func truncateSkillContent(content string, maxTokens int) string {
+	if len(content) == 0 {
+		return content
+	}
+
+	if estimateTokenCount(content) <= maxTokens {
+		return content
+	}
+
+	marker := getSkillTruncationMarker()
+	targetBytes := estimateTokenBytes(maxTokens) - len(marker)
+	if targetBytes < 0 {
+		targetBytes = 0
+	}
+	if targetBytes > len(content) {
+		targetBytes = len(content)
+	}
+
+	// Back up to a valid UTF-8 rune boundary.
+	for targetBytes > 0 && targetBytes < len(content) && !utf8.RuneStart(content[targetBytes]) {
+		targetBytes--
+	}
+
+	return content[:targetBytes] + marker
+}

adk/middlewares/summarization/finalizer_builder_test.go

@@ -19,7 +19,9 @@ package summarization
 import (
 	"context"
 	"errors"
+	"strings"
 	"testing"
+	"unicode/utf8"
 
 	"github.com/stretchr/testify/assert"
 
@@ -212,6 +214,24 @@ func TestPreserveSkillsConfigCheck(t *testing.T) {
 		err := c.check()
 		assert.NoError(t, err)
 	})
+
+	t.Run("negative max tokens per skill", func(t *testing.T) {
+		c := &PreserveSkillsConfig{
+			MaxTokensPerSkill: ptr(-1),
+		}
+		err := c.check()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "MaxTokensPerSkill must be non-negative")
+	})
+
+	t.Run("negative skills token budget", func(t *testing.T) {
+		c := &PreserveSkillsConfig{
+			SkillsTokenBudget: ptr(-1),
+		}
+		err := c.check()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "SkillsTokenBudget must be non-negative")
+	})
 }
 
 func TestPreserveSkillsViaBuilder(t *testing.T) {
@@ -434,4 +454,128 @@ func TestBuildPreservedSkillsText(t *testing.T) {
 		assert.NotContains(t, text, "skill2")
 		assert.NotContains(t, text, "c2")
 	})
+
+	t.Run("per skill token limit truncates large skills", func(t *testing.T) {
+		// estimateTokenCount = (len+3)/4
+		// "short" = 5 chars → 2 tokens
+		// strings.Repeat("x", 100) = 100 chars → 25 tokens
+		largeContent := strings.Repeat("x", 100)
+		messages := []adk.Message{
+			{
+				Role: schema.Assistant,
+				ToolCalls: []schema.ToolCall{
+					{ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "small"}`}},
+					{ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "large"}`}},
+				},
+			},
+			{Role: schema.Tool, ToolCallID: "call_1", Content: "short"},
+			{Role: schema.Tool, ToolCallID: "call_2", Content: largeContent},
+		}
+
+		// MaxTokensPerSkill=10: "short"→2 tokens (ok), largeContent→25 tokens (truncated)
+		text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{
+			MaxSkills:         ptr(10),
+			MaxTokensPerSkill: ptr(10),
+			SkillToolName:     "load_skill",
+		})
+		assert.NoError(t, err)
+		// small skill preserved as-is
+		assert.Contains(t, text, "small")
+		assert.Contains(t, text, "short")
+		// large skill is truncated, not dropped — name still present, full content gone
+		assert.Contains(t, text, "large")
+		assert.NotContains(t, text, largeContent)
+		assert.Contains(t, text, "skill content truncated for compaction")
+	})
+
+	t.Run("total token budget drops excess skills", func(t *testing.T) {
+		// Each content is 40 chars → (40+3)/4 = 10 tokens
+		content := strings.Repeat("a", 40)
+		messages := []adk.Message{
+			{
+				Role: schema.Assistant,
+				ToolCalls: []schema.ToolCall{
+					{ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill1"}`}},
+					{ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill2"}`}},
+					{ID: "call_3", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill3"}`}},
+				},
+			},
+			{Role: schema.Tool, ToolCallID: "call_1", Content: content},
+			{Role: schema.Tool, ToolCallID: "call_2", Content: content},
+			{Role: schema.Tool, ToolCallID: "call_3", Content: content},
+		}
+
+		// Budget=15: skill3=10 tokens fits, skill2=10 tokens → 10+10=20 > 15, stop.
+		text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{
+			MaxSkills:         ptr(10),
+			SkillsTokenBudget: ptr(15),
+			SkillToolName:     "load_skill",
+		})
+		assert.NoError(t, err)
+		assert.Contains(t, text, "skill3")
+		assert.NotContains(t, text, "skill1")
+		assert.NotContains(t, text, "skill2")
+	})
+
+	t.Run("token budget and per-skill limit combined", func(t *testing.T) {
+		// s1: 16 chars → 4 tokens
+		// s2: 200 chars → 50 tokens (exceeds per-skill limit of 20, gets truncated)
+		// s3: 24 chars → 6 tokens
+		// s4: 24 chars → 6 tokens
+		messages := []adk.Message{
+			{
+				Role: schema.Assistant,
+				ToolCalls: []schema.ToolCall{
+					{ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s1"}`}},
+					{ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s2"}`}},
+					{ID: "call_3", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s3"}`}},
+					{ID: "call_4", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s4"}`}},
+				},
+			},
+			{Role: schema.Tool, ToolCallID: "call_1", Content: strings.Repeat("a", 16)},
+			{Role: schema.Tool, ToolCallID: "call_2", Content: strings.Repeat("b", 200)},
+			{Role: schema.Tool, ToolCallID: "call_3", Content: strings.Repeat("c", 24)},
+			{Role: schema.Tool, ToolCallID: "call_4", Content: strings.Repeat("d", 24)},
+		}
+
+		// Per-skill limit: 20 (s2 with 50 tokens is truncated to 20)
+		// Budget: 30 (from most recent: s4=6, s3=6, s2=20, total=32 > 30, so s2 cannot fit)
+		// Result: s4 and s3 preserved
+		text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{
+			MaxSkills:         ptr(10),
+			MaxTokensPerSkill: ptr(20),
+			SkillsTokenBudget: ptr(30),
+			SkillToolName:     "load_skill",
+		})
+		assert.NoError(t, err)
+		assert.Contains(t, text, "s3")
+		assert.Contains(t, text, "s4")
+		assert.NotContains(t, text, "\"s1\"")
+		assert.NotContains(t, text, "\"s2\"")
+	})
+
+	t.Run("truncated skill content preserves only prefix", func(t *testing.T) {
+		// Use a long content and generous maxTokens so the prefix is clearly visible.
+		content := strings.Repeat("abcdefghij", 100) // 1000 bytes → 250 tokens
+		// maxTokens=125 → targetBytes = 500, minus ~101 marker bytes → ~399 prefix bytes
+		truncated := truncateSkillContent(content, 125)
+		assert.True(t, strings.HasPrefix(truncated, "abcdefghij")) // prefix preserved
+		assert.Contains(t, truncated, "skill content truncated for compaction")
+		assert.NotEqual(t, content, truncated)
+		// Ends with marker, not with original content suffix
+		assert.True(t, strings.HasSuffix(truncated, "]"))
+		// No suffix from original content
+		assert.False(t, strings.HasSuffix(truncated, "abcdefghij]"))
+	})
+
+	t.Run("truncated multibyte content does not produce invalid utf8", func(t *testing.T) {
+		// Each Chinese char is 3 bytes. 334 chars = 1002 bytes → 251 tokens
+		content := strings.Repeat("中", 334)
+		// maxTokens=125 → targetBytes=500, minus marker ~101 bytes → ~399 bytes
+		// 399 / 3 = 133 full Chinese chars, no partial rune
+		truncated := truncateSkillContent(content, 125)
+		assert.True(t, utf8.ValidString(truncated))
+		assert.True(t, strings.HasPrefix(truncated, "中中中"))
+		assert.Contains(t, truncated, "skill content truncated for compaction")
+	})
 }

adk/middlewares/summarization/prompt.go

@@ -327,3 +327,14 @@ func getSkillPreamble() string {
 		Chinese: skillPreambleZh,
 	})
 }
+
+const skillTruncationMarker = "\n\n[... skill content truncated for compaction; use Read on the skill path if you need the full text]"
+
+const skillTruncationMarkerZh = "\n\n[... skill 内容已在压缩时截断，如需完整内容请通过 Read 读取 skill 对应的文件路径]"
+
+func getSkillTruncationMarker() string {
+	return internal.SelectPrompt(internal.I18nPrompts{
+		English: skillTruncationMarker,
+		Chinese: skillTruncationMarkerZh,
+	})
+}

adk/middlewares/summarization/summarization.go

@@ -296,9 +296,9 @@ func SummarizeMessages(ctx context.Context, cfg *Config, messages []adk.Message)
 		return nil, err
 	}
 
-	ctx = context.WithValue(ctx, ctxKeyModelInput{}, modelInput)
+	finalizeCtx := context.WithValue(ctx, ctxKeyModelInput{}, modelInput)
 
-	_, finalMsgs, err := m.finalizeSummary(ctx, messages, rawSummary)
+	_, finalMsgs, err := m.finalizeSummary(finalizeCtx, messages, rawSummary)
 	if err != nil {
 		return nil, err
 	}
@@ -355,10 +355,10 @@ func (m *middleware) BeforeModelRewriteState(ctx context.Context, state *adk.Cha
 		return nil, nil, err
 	}
 
-	ctx = context.WithValue(ctx, ctxKeyModelInput{}, modelInput)
+	finalizeCtx := context.WithValue(ctx, ctxKeyModelInput{}, modelInput)
 
 	var finalMsgs []adk.Message
-	ctx, finalMsgs, err = m.finalizeSummary(ctx, beforeState.Messages, rawSummary)
+	_, finalMsgs, err = m.finalizeSummary(finalizeCtx, beforeState.Messages, rawSummary)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -503,6 +503,10 @@ func estimateTokenCount(text string) int {
 	return (len(text) + 3) / 4
 }
 
+func estimateTokenBytes(tokens int) int {
+	return tokens * 4
+}
+
 func (m *middleware) summarize(ctx context.Context, originalMsgs []adk.Message) (adk.Message, []adk.Message, error) {
 	_, contextMsgs := m.splitSystemAndContextMsgs(originalMsgs)