Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 55 additions & 11 deletions agent/compaction/compaction_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ func TestTruncationStrategy_ExcludesOldestGroups(t *testing.T) {
index := compaction.CreateMessageIndex(turnMessages(3), nil)
strategy := &compaction.TruncationStrategy{
Trigger: compaction.GroupsExceed(2),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand All @@ -142,7 +142,7 @@ func TestTruncationStrategy_SkipsPreExcludedAndSystemGroups(t *testing.T) {
}, nil)
index.Groups[1].IsExcluded = true
strategy := &compaction.TruncationStrategy{
MinimumPreservedGroups: 1,
MinimumPreservedGroups: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand Down Expand Up @@ -187,7 +187,7 @@ func TestSlidingWindowStrategy_ExcludesOldestTurns(t *testing.T) {
index := compaction.CreateMessageIndex(turnMessages(3), nil)
strategy := &compaction.SlidingWindowStrategy{
Trigger: compaction.TurnsExceed(1),
MinimumPreservedTurns: 1,
MinimumPreservedTurns: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand Down Expand Up @@ -215,7 +215,7 @@ func TestSlidingWindowStrategy_PreservesTurnZeroGroups(t *testing.T) {
}, nil)
strategy := &compaction.SlidingWindowStrategy{
Trigger: compaction.TurnsExceed(1),
MinimumPreservedTurns: 1,
MinimumPreservedTurns: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand All @@ -233,6 +233,48 @@ func TestSlidingWindowStrategy_PreservesTurnZeroGroups(t *testing.T) {
}
}

func TestTruncationStrategy_ExplicitZeroPreservesNone(t *testing.T) {
index := compaction.CreateMessageIndex([]*message.Message{
textMessage(message.RoleSystem, "system"),
textMessage(message.RoleAssistant, "g1"),
textMessage(message.RoleAssistant, "g2"),
textMessage(message.RoleAssistant, "g3"),
textMessage(message.RoleAssistant, "g4"),
textMessage(message.RoleAssistant, "g5"),
}, nil)
strategy := &compaction.TruncationStrategy{MinimumPreservedGroups: ptr(0)}

compacted, err := strategy.Compact(t.Context(), index)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !compacted {
t.Fatal("expected compaction with an explicit zero floor")
}
if got := index.IncludedNonSystemGroupCount(); got != 0 {
t.Fatalf("expected all removable non-system groups excluded, got %d preserved", got)
}
if index.Groups[0].IsExcluded {
t.Fatal("expected system group to be preserved")
}
}

func TestSlidingWindowStrategy_NegativeMinimumClampsToZero(t *testing.T) {
index := compaction.CreateMessageIndex(turnMessages(3), nil)
strategy := &compaction.SlidingWindowStrategy{MinimumPreservedTurns: ptr(-5)}

compacted, err := strategy.Compact(t.Context(), index)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !compacted {
t.Fatal("expected compaction when the floor clamps to zero")
}
if got := index.IncludedTurnCount(); got != 0 {
t.Fatalf("expected a negative floor to clamp to zero and exclude all turns, got %d turns preserved", got)
}
}

func TestSlidingWindowStrategy_ZeroValueUsesDefaults(t *testing.T) {
index := compaction.CreateMessageIndex(turnMessages(3), nil)
strategy := &compaction.SlidingWindowStrategy{}
Expand Down Expand Up @@ -273,7 +315,7 @@ func TestToolResultStrategy_CollapsesOldToolGroups(t *testing.T) {
index := compaction.CreateMessageIndex(messages, nil)
strategy := &compaction.ToolResultStrategy{
Trigger: compaction.HasToolCalls(),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand Down Expand Up @@ -345,7 +387,7 @@ func TestSummarizationStrategy_InsertsSummaryAndPreservesRecentGroups(t *testing
strategy := &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: summarizer,
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
SummarizationPrompt: "summarize",
}

Expand Down Expand Up @@ -417,7 +459,7 @@ func TestSummarizationStrategy_RestoresGroupsWhenSummarizerFails(t *testing.T) {
strategy := &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: compaction.SummarizerFunc(func(context.Context, []*message.Message) (string, error) { return "", expected }),
MinimumPreservedGroups: 1,
MinimumPreservedGroups: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand All @@ -437,7 +479,7 @@ func TestSummarizationStrategy_PropagatesCancellation(t *testing.T) {
strategy := &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: compaction.SummarizerFunc(func(context.Context, []*message.Message) (string, error) { return "", context.Canceled }),
MinimumPreservedGroups: 1,
MinimumPreservedGroups: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand All @@ -457,7 +499,7 @@ func TestNewProvider_CompactsAndPersistsIndex(t *testing.T) {
provider := compaction.NewContextProvider(compaction.ContextProviderConfig{
Strategy: &compaction.TruncationStrategy{
Trigger: compaction.GroupsExceed(2),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
},
SourceID: "compaction-test",
})
Expand Down Expand Up @@ -494,7 +536,7 @@ func TestNewProvider_SourceStampsGeneratedMessages(t *testing.T) {
Strategy: &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: compaction.SummarizerFunc(func(context.Context, []*message.Message) (string, error) { return "older context", nil }),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
},
SourceID: "compaction-test",
})
Expand All @@ -518,7 +560,7 @@ func TestNewProvider_CompactsWithoutSession(t *testing.T) {
provider := compaction.NewContextProvider(compaction.ContextProviderConfig{
Strategy: &compaction.TruncationStrategy{
Trigger: compaction.GroupsExceed(2),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
},
})

Expand Down Expand Up @@ -583,3 +625,5 @@ func messageTexts(messages []*message.Message) []string {
}
return texts
}

func ptr[T any](value T) *T { return &value }
5 changes: 3 additions & 2 deletions agent/compaction/contextwindow.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,15 +78,16 @@ func (s *ContextWindowStrategy) Compact(ctx context.Context, index *MessageIndex
toolEvictionTokens := int(float64(inputBudget) * toolEviction)
truncationTokens := int(float64(inputBudget) * truncation)

minimumPreservedGroups := 2
pipeline := &PipelineStrategy{
Strategies: []Strategy{
&ToolResultStrategy{
Trigger: TokensExceed(toolEvictionTokens),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: &minimumPreservedGroups,
},
&TruncationStrategy{
Trigger: TokensExceed(truncationTokens),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: &minimumPreservedGroups,
},
},
}
Expand Down
4 changes: 2 additions & 2 deletions agent/compaction/index_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func TestMessageIndex_ClassifiesStrategySummaryMessages(t *testing.T) {
strategy := &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: compaction.SummarizerFunc(func(context.Context, []*message.Message) (string, error) { return "older context", nil }),
MinimumPreservedGroups: 1,
MinimumPreservedGroups: ptr(1),
}

compacted, err := strategy.Compact(t.Context(), index)
Expand Down Expand Up @@ -183,7 +183,7 @@ func TestMessageIndex_UpdatePreservesStateFromCompactedProjection(t *testing.T)
strategy := &compaction.SummarizationStrategy{
Trigger: compaction.GroupsExceed(2),
Summarizer: compaction.SummarizerFunc(func(context.Context, []*message.Message) (string, error) { return "older context", nil }),
MinimumPreservedGroups: 2,
MinimumPreservedGroups: ptr(2),
SummarizationPrompt: "summarize",
}

Expand Down
11 changes: 8 additions & 3 deletions agent/compaction/slidingwindow.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
package compaction

import (
"cmp"
"context"
"slices"
)
Expand All @@ -25,7 +24,10 @@ type SlidingWindowStrategy struct {

// MinimumPreservedTurns is the minimum number of most-recent user turns to preserve.
// Groups with nil or non-positive turn indexes are preserved independently of this value.
MinimumPreservedTurns int
//
// When nil, a default floor is used. An explicit value is honored as-is, so a pointer to 0
// disables the floor entirely; a negative value is clamped to 0.
MinimumPreservedTurns *int
}

// Compact compacts index in place.
Expand All @@ -35,7 +37,10 @@ func (strategy *SlidingWindowStrategy) Compact(_ context.Context, index *Message
return false, nil
}

minimumPreservedTurns := cmp.Or(max(strategy.MinimumPreservedTurns, 0), defaultMinimumPreservedSlidingWindowTurns)
minimumPreservedTurns := defaultMinimumPreservedSlidingWindowTurns
if strategy.MinimumPreservedTurns != nil {
minimumPreservedTurns = max(*strategy.MinimumPreservedTurns, 0)
}
turnGroups := make(map[int][]int)
var turnOrder []int
for i, group := range index.Groups {
Expand Down
10 changes: 8 additions & 2 deletions agent/compaction/summarization.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,10 @@ type SummarizationStrategy struct {

// MinimumPreservedGroups is the minimum number of most-recent non-system groups to preserve.
// This is a hard floor; summarization will not summarize groups within this protected window.
MinimumPreservedGroups int
//
// When nil, a default floor is used. An explicit value is honored as-is, so a pointer to 0
// disables the floor entirely; a negative value is clamped to 0.
MinimumPreservedGroups *int

// SummarizationPrompt is the system prompt prepended to messages sent to Summarizer.
// When empty, a default prompt is used.
Expand All @@ -75,7 +78,10 @@ func (strategy *SummarizationStrategy) Compact(ctx context.Context, index *Messa
if strategy.Summarizer == nil {
return false, nil
}
minimumPreservedGroups := cmp.Or(max(strategy.MinimumPreservedGroups, 0), defaultMinimumPreservedSummarizationGroups)
minimumPreservedGroups := defaultMinimumPreservedSummarizationGroups
if strategy.MinimumPreservedGroups != nil {
minimumPreservedGroups = max(*strategy.MinimumPreservedGroups, 0)
}
summarizationPrompt := cmp.Or(strategy.SummarizationPrompt, defaultSummarizationPrompt)
summaryUnavailableMessage := cmp.Or(strategy.SummaryUnavailableMessage, "[Summary unavailable]")

Expand Down
11 changes: 8 additions & 3 deletions agent/compaction/toolresult.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
package compaction

import (
"cmp"
"context"
"fmt"
"slices"
Expand All @@ -29,7 +28,10 @@ type ToolResultStrategy struct {

// MinimumPreservedGroups is the minimum number of most-recent non-system groups to preserve.
// This is a hard floor; tool-call groups within this protected window are not collapsed.
MinimumPreservedGroups int
//
// When nil, a default floor is used. An explicit value is honored as-is, so a pointer to 0
// disables the floor entirely; a negative value is clamped to 0.
MinimumPreservedGroups *int

// ToolCallFormatter formats a tool-call group as a compact summary string.
// When nil, DefaultToolCallFormatter is used, which produces a YAML-like block listing
Expand All @@ -44,7 +46,10 @@ func (strategy *ToolResultStrategy) Compact(_ context.Context, index *MessageInd
return false, nil
}

minimumPreservedGroups := cmp.Or(max(strategy.MinimumPreservedGroups, 0), defaultMinimumPreservedToolResultGroups)
minimumPreservedGroups := defaultMinimumPreservedToolResultGroups
if strategy.MinimumPreservedGroups != nil {
minimumPreservedGroups = max(*strategy.MinimumPreservedGroups, 0)
}
var nonSystemIncludedIndices []int
for i, group := range index.Groups {
if !group.IsExcluded && group.Kind != GroupKindSystem {
Expand Down
11 changes: 8 additions & 3 deletions agent/compaction/truncation.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
package compaction

import (
"cmp"
"context"
)

Expand All @@ -24,7 +23,10 @@ type TruncationStrategy struct {

// MinimumPreservedGroups is the minimum number of most-recent non-system groups to preserve.
// This is a hard floor; truncation will not remove groups beyond this limit.
MinimumPreservedGroups int
//
// When nil, a default floor is used. An explicit value is honored as-is, so a pointer to 0
// disables the floor entirely; a negative value is clamped to 0.
MinimumPreservedGroups *int
}

// Compact compacts index in place.
Expand All @@ -34,7 +36,10 @@ func (strategy *TruncationStrategy) Compact(_ context.Context, index *MessageInd
return false, nil
}

minimumPreservedGroups := cmp.Or(max(strategy.MinimumPreservedGroups, 0), defaultMinimumPreservedTruncationGroups)
minimumPreservedGroups := defaultMinimumPreservedTruncationGroups
if strategy.MinimumPreservedGroups != nil {
minimumPreservedGroups = max(*strategy.MinimumPreservedGroups, 0)
}
removableCount := index.IncludedNonSystemGroupCount()
maxRemovable := removableCount - minimumPreservedGroups
if maxRemovable <= 0 {
Expand Down
10 changes: 6 additions & 4 deletions examples/02-agents/agents/step18_compaction_pipeline/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,23 +58,23 @@ func main() {
// 1. Gentle: collapse old tool-call groups into short summaries.
&compaction.ToolResultStrategy{
Trigger: compaction.MessagesExceed(7),
MinimumPreservedGroups: 4,
MinimumPreservedGroups: ptr(4),
},
// 2. Moderate: use an LLM to summarize older conversation spans into a concise message.
&compaction.SummarizationStrategy{
Trigger: compaction.TokensExceed(0x500),
Summarizer: summarizer,
MinimumPreservedGroups: 4,
MinimumPreservedGroups: ptr(4),
},
// 3. Aggressive: keep only the last N user turns and their responses.
&compaction.SlidingWindowStrategy{
Trigger: compaction.TurnsExceed(4),
MinimumPreservedTurns: 4,
MinimumPreservedTurns: ptr(4),
},
// 4. Emergency: drop oldest groups until under the token budget.
&compaction.TruncationStrategy{
Trigger: compaction.TokensExceed(0x8000),
MinimumPreservedGroups: 8,
MinimumPreservedGroups: ptr(8),
},
},
}
Expand Down Expand Up @@ -126,6 +126,8 @@ When responding, be extra descriptive and use as many words as possible without
}
}

func ptr[T any](value T) *T { return &value }

func lookupPrice(_ context.Context, productName string) (string, error) {
switch strings.ToUpper(productName) {
case "LAPTOP":
Expand Down
Loading