compaction

package
v0.3.4 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Mar 6, 2026 License: MPL-2.0 Imports: 9 Imported by: 0

Documentation

Overview

Package compaction provides context window management for AI agent conversations.

When conversations grow too long, they can exceed Claude's context window limits. This package implements strategies to compact conversation history while preserving essential context.

Strategies

The package supports two compaction strategies:

  • Summarization (StrategySummarization): Uses Claude to create a structured summary of older messages. The summary follows a 9-section format that preserves critical information about the conversation.

  • Hybrid (StrategyHybrid): A two-phase approach that first prunes tool outputs (replacing them with "[TOOL OUTPUT PRUNED]" placeholders), then summarizes if still needed. This is the default and most cost-effective strategy.

Usage

Create a Compactor with your configuration:

compactor := compaction.New(store, anthropicClient, &compaction.Config{
    Strategy:        compaction.StrategyHybrid,
    Trigger:         0.85,     // Trigger at 85% context usage
    TargetTokens:    80000,    // Target after compaction
    PreserveLastN:   10,       // Always keep last 10 messages
    ProtectedTokens: 40000,    // Never touch last 40K tokens
}, logger)

Check if compaction is needed and compact if necessary:

result, err := compactor.CompactIfNeeded(ctx, sessionID)
if err != nil {
    return err
}
if result != nil {
    log.Printf("Compacted: %d -> %d tokens", result.OriginalTokens, result.CompactedTokens)
}

Message Partitioning

Messages are partitioned into mutually exclusive categories:

  • Protected: Within the ProtectedTokens zone at the end of conversation.
  • Preserved: Messages marked with is_preserved=true.
  • Recent: Last PreserveLastN messages.
  • Summaries: Previous compaction summary messages (is_summary=true).
  • Compactable: Everything else, eligible for summarization.

Token Counting

Token counting uses Claude's API for accurate counts, with a character-based approximation fallback (~4 characters per token) if the API is unavailable.

Database Integration

The compactor uses the driver.Store interface for database operations:

  • GetMessages: Retrieve session messages
  • CreateMessage: Create summary messages
  • DeleteMessage: Remove archived messages
  • ArchiveMessage: Archive messages before deletion
  • CreateCompactionEvent: Record compaction operations
  • UpdateSession: Update compaction count

Thread Safety

The Compactor is safe for concurrent use. However, concurrent compactions on the same session should be avoided as they may lead to inconsistent results.

Index

Constants

View Source
const (
	DefaultStrategy            = StrategyHybrid
	DefaultTrigger             = 0.85  // 85% context usage
	DefaultTargetTokens        = 80000 // Target 80K tokens after compaction
	DefaultPreserveLastN       = 10    // Always keep last 10 messages
	DefaultProtectedTokens     = 40000 // Never touch last 40K tokens
	DefaultSummarizerModel     = "claude-3-5-haiku-20241022"
	DefaultMaxTokensForModel   = 200000 // Claude Sonnet 4.5 context window
	DefaultPreserveToolOutputs = false  // Prune tool outputs by default
	DefaultUseTokenCountingAPI = true   // Use Claude API for accurate counts
	DefaultSummarizerMaxTokens = 4096   // Max tokens for summarization response
)

Default configuration values based on production patterns.

View Source
const SummarizationSystemPrompt = `` /* 2262-byte string literal not displayed */

SummarizationSystemPrompt is the system prompt used for context summarization. This prompt instructs Claude to create a structured summary that preserves critical information from the conversation being compacted.

The 9-section structure is based on production patterns from AI coding assistants like Claude Code, Aider, and Cline.

Variables

View Source
var (
	// ErrInvalidConfig indicates invalid compaction configuration.
	ErrInvalidConfig = errors.New("invalid compaction configuration")

	// ErrNoMessagesToCompact indicates there are no messages eligible for compaction.
	ErrNoMessagesToCompact = errors.New("no messages to compact")

	// ErrCompactionInProgress indicates compaction is already running for this session.
	ErrCompactionInProgress = errors.New("compaction already in progress")

	// ErrSummarizationFailed indicates the summarization API call failed.
	ErrSummarizationFailed = errors.New("summarization failed")

	// ErrTokenCountingFailed indicates token counting failed.
	ErrTokenCountingFailed = errors.New("token counting failed")

	// ErrSessionNotFound indicates the session was not found.
	ErrSessionNotFound = errors.New("session not found")

	// ErrStorageError indicates a database operation failed.
	ErrStorageError = errors.New("storage operation failed")
)

Sentinel errors for compaction operations.

Functions

func BuildSummarizationUserPrompt added in v0.1.4

func BuildSummarizationUserPrompt(conversationText string) string

BuildSummarizationUserPrompt creates the user message for summarization. It includes the messages to be summarized and any additional context.

func BuildSummarizationUserPromptWithContext added in v0.1.4

func BuildSummarizationUserPromptWithContext(contextText, conversationText string) string

BuildSummarizationUserPromptWithContext creates the user message for summarization with additional context from preserved messages.

func EstimateTokensFromUsage added in v0.1.4

func EstimateTokensFromUsage(usage driver.Usage) int

EstimateTokensFromUsage calculates total tokens from a Usage struct.

func FormatMessagesAsText added in v0.1.4

func FormatMessagesAsText(messages []MessageForSummary) string

FormatMessagesAsText formats a slice of messages as readable text for summarization.

func WrapError added in v0.1.4

func WrapError(op string, err error) error

WrapError wraps an error with operation context. If err is nil, returns nil.

func WrapErrorWithSession added in v0.1.4

func WrapErrorWithSession(op string, sessionID uuid.UUID, err error) error

WrapErrorWithSession wraps an error with operation and session context.

Types

type CompactionContext added in v0.1.4

type CompactionContext struct {
	// SessionID is the session being compacted.
	SessionID uuid.UUID

	// Messages is the full list of messages in the session.
	Messages []*driver.Message

	// Partition is the result of partitioning the messages.
	Partition *MessagePartition

	// Config is the compaction configuration.
	Config *Config

	// TargetTokens is the target token count after compaction.
	TargetTokens int
}

CompactionContext contains all the context needed for a compaction operation.

func (*CompactionContext) ValidateForCompaction added in v0.1.4

func (c *CompactionContext) ValidateForCompaction() error

ValidateForCompaction checks if compaction should proceed.

type CompactionError added in v0.1.4

type CompactionError struct {
	// Op is the operation that failed (e.g., "Compact", "CountTokens", "Summarize")
	Op string

	// SessionID is the session ID if applicable
	SessionID uuid.UUID

	// Err is the underlying error
	Err error

	// Context holds additional key-value pairs for debugging
	Context map[string]any
}

CompactionError provides structured error context for compaction operations.

func NewCompactionError added in v0.1.4

func NewCompactionError(op string, err error) *CompactionError

NewCompactionError creates a new CompactionError with the given operation and underlying error.

func (*CompactionError) Error added in v0.1.4

func (e *CompactionError) Error() string

Error returns a formatted error message.

func (*CompactionError) Unwrap added in v0.1.4

func (e *CompactionError) Unwrap() error

Unwrap returns the underlying error for errors.Is/errors.As support.

func (*CompactionError) WithContext added in v0.1.4

func (e *CompactionError) WithContext(key string, value any) *CompactionError

WithContext adds a key-value pair to the error context and returns the error for chaining.

func (*CompactionError) WithSession added in v0.1.4

func (e *CompactionError) WithSession(sessionID uuid.UUID) *CompactionError

WithSession sets the session ID on the error and returns the error for chaining.

type Compactor added in v0.1.4

type Compactor[TTx any] struct {
	// contains filtered or unexported fields
}

Compactor provides context compaction for agent sessions. It uses the generic type parameter TTx to work with different database drivers.

func New added in v0.1.4

func New[TTx any](store driver.Store[TTx], client *anthropic.Client, config *Config, logger Logger) *Compactor[TTx]

New creates a new Compactor with the given configuration. If config is nil, default configuration is used.

func (*Compactor[TTx]) Compact added in v0.1.4

func (c *Compactor[TTx]) Compact(ctx context.Context, sessionID uuid.UUID) (*Result, error)

Compact performs compaction on the specified session.

func (*Compactor[TTx]) CompactIfNeeded added in v0.1.4

func (c *Compactor[TTx]) CompactIfNeeded(ctx context.Context, sessionID uuid.UUID) (*Result, error)

CompactIfNeeded performs compaction only if the session exceeds the trigger threshold. Returns nil result if compaction was not needed.

func (*Compactor[TTx]) Config added in v0.1.4

func (c *Compactor[TTx]) Config() *Config

Config returns the compactor's configuration.

func (*Compactor[TTx]) GetStats added in v0.1.4

func (c *Compactor[TTx]) GetStats(ctx context.Context, sessionID uuid.UUID) (*Stats, error)

GetStats returns statistics about a session's compaction state.

func (*Compactor[TTx]) NeedsCompaction added in v0.1.4

func (c *Compactor[TTx]) NeedsCompaction(ctx context.Context, sessionID uuid.UUID) (bool, error)

NeedsCompaction checks if a session needs compaction based on token usage.

type Config added in v0.1.4

type Config struct {
	// Strategy is the compaction strategy to use.
	// Default: StrategyHybrid
	Strategy Strategy

	// Trigger is the context usage threshold (0.0-1.0) that triggers compaction.
	// E.g., 0.85 means trigger at 85% context usage.
	// Default: 0.85
	Trigger float64

	// TargetTokens is the target token count after compaction.
	// The compactor will try to reduce context to this level.
	// Default: 80000
	TargetTokens int

	// PreserveLastN is the minimum number of recent messages to always preserve.
	// These messages are never summarized or removed.
	// Default: 10
	PreserveLastN int

	// ProtectedTokens is the token count at the end of context that is never summarized.
	// This protects the most recent conversation from compaction.
	// Default: 40000
	ProtectedTokens int

	// SummarizerModel is the Claude model to use for summarization.
	// Using a faster/cheaper model is recommended.
	// Default: "claude-3-5-haiku-20241022"
	SummarizerModel string

	// SummarizerMaxTokens is the maximum tokens for the summarization response.
	// Default: 4096
	SummarizerMaxTokens int

	// MaxTokensForModel is the maximum context window for the target model.
	// Used to calculate context usage percentage.
	// Default: 200000 (Sonnet 4.5 context window)
	MaxTokensForModel int

	// PreserveToolOutputs determines whether to keep tool outputs during hybrid compaction.
	// If false, tool outputs outside the protected zone are replaced with "[TOOL OUTPUT PRUNED]" placeholder.
	// Default: false
	PreserveToolOutputs bool

	// UseTokenCountingAPI determines whether to use Claude's token counting API.
	// If false or API fails, uses character-based approximation.
	// Default: true
	UseTokenCountingAPI bool
}

Config holds compaction configuration.

func DefaultConfig

func DefaultConfig() *Config

DefaultConfig returns a Config with sensible defaults based on production patterns.

func (*Config) ApplyDefaults added in v0.1.4

func (c *Config) ApplyDefaults()

ApplyDefaults fills in zero values with defaults.

func (*Config) TriggerThreshold added in v0.1.4

func (c *Config) TriggerThreshold() int

TriggerThreshold returns the absolute token count that triggers compaction.

func (*Config) Validate added in v0.1.4

func (c *Config) Validate() error

Validate validates the configuration and returns an error if invalid.

type HybridStrategy

type HybridStrategy struct {
	// contains filtered or unexported fields
}

HybridStrategy implements a two-phase compaction approach: 1. First, prune tool outputs from compactable messages (free, no API call) 2. If still over target, summarize remaining messages using the summarization strategy

func NewHybridStrategy

func NewHybridStrategy(summarizer *Summarizer, tokenCounter *TokenCounter, config *Config) *HybridStrategy

NewHybridStrategy creates a new hybrid compaction strategy.

func (*HybridStrategy) Execute added in v0.1.4

func (h *HybridStrategy) Execute(ctx context.Context, partition *MessagePartition) (*StrategyResult, error)

Execute performs the hybrid compaction strategy.

func (*HybridStrategy) Name

func (h *HybridStrategy) Name() Strategy

Name returns the strategy name.

type Logger added in v0.1.4

type Logger interface {
	Debug(msg string, args ...any)
	Info(msg string, args ...any)
	Warn(msg string, args ...any)
	Error(msg string, args ...any)
}

Logger interface for compaction logging.

type MessageForSummary added in v0.1.4

type MessageForSummary struct {
	Role    string
	Content string
}

MessageForSummary represents a simplified message for summarization.

type MessagePartition added in v0.1.4

type MessagePartition struct {
	// Protected messages are within the ProtectedTokens zone at the end
	// of the conversation. These are never touched during compaction.
	Protected []*driver.Message

	// Preserved messages have is_preserved=true and are never removed.
	// They may be from any position in the conversation.
	Preserved []*driver.Message

	// Recent messages are within the PreserveLastN count.
	// These are never summarized or removed.
	Recent []*driver.Message

	// Summaries are previous compaction summaries (is_summary=true).
	// These are kept as-is and included in new summarization context.
	Summaries []*driver.Message

	// Compactable messages are eligible for summarization and removal.
	// These are the messages that will be archived and replaced with a summary.
	Compactable []*driver.Message

	// Stats contains token counts for each partition.
	Stats PartitionStats
}

MessagePartition categorizes messages for compaction processing. Messages are partitioned into mutually exclusive categories based on their attributes and position in the conversation.

func (*MessagePartition) AllPreservedIDs added in v0.1.4

func (p *MessagePartition) AllPreservedIDs() []uuid.UUID

AllPreservedIDs returns the IDs of all messages that should be preserved (protected, preserved, recent, and summaries).

func (*MessagePartition) CanCompact added in v0.1.4

func (p *MessagePartition) CanCompact() bool

CanCompact returns true if there are messages eligible for compaction.

func (*MessagePartition) CompactableIDs added in v0.1.4

func (p *MessagePartition) CompactableIDs() []uuid.UUID

CompactableIDs returns the IDs of all compactable messages.

func (*MessagePartition) ContextMessages added in v0.1.4

func (p *MessagePartition) ContextMessages() []*driver.Message

ContextMessages returns messages that should be included in the summarization context but not summarized themselves. This includes previous summaries and preserved messages to maintain continuity.

func (*MessagePartition) MessagesForSummarization added in v0.1.4

func (p *MessagePartition) MessagesForSummarization() []*driver.Message

MessagesForSummarization returns all compactable messages in chronological order. These are the messages that will be summarized.

func (*MessagePartition) TokenReductionEstimate added in v0.1.4

func (p *MessagePartition) TokenReductionEstimate(summaryTokens int) int

TokenReductionEstimate estimates how many tokens could be saved by compaction. This is calculated as: compactable tokens - estimated summary size.

type PartitionStats added in v0.1.4

type PartitionStats struct {
	ProtectedTokens   int
	PreservedTokens   int
	RecentTokens      int
	SummaryTokens     int
	CompactableTokens int
	TotalTokens       int
}

PartitionStats contains token statistics for each partition.

type Partitioner

type Partitioner struct {
	// contains filtered or unexported fields
}

Partitioner handles the logic of partitioning messages for compaction.

func NewPartitioner

func NewPartitioner(tokenCounter *TokenCounter, config *Config) *Partitioner

NewPartitioner creates a new Partitioner with the given configuration.

func (*Partitioner) Partition

func (p *Partitioner) Partition(ctx context.Context, messages []*driver.Message) (*MessagePartition, error)

Partition categorizes messages into compaction groups. Messages are processed in reverse order (newest first) to correctly identify protected and recent messages.

type Result added in v0.1.4

type Result struct {
	// EventID is the ID of the compaction event record.
	EventID uuid.UUID

	// Strategy is the strategy that was used.
	Strategy Strategy

	// OriginalTokens is the token count before compaction.
	OriginalTokens int

	// CompactedTokens is the token count after compaction.
	CompactedTokens int

	// MessagesRemoved is the number of messages archived.
	MessagesRemoved int

	// PreservedMessageIDs is the list of message IDs that were preserved.
	PreservedMessageIDs []uuid.UUID

	// SummaryCreated indicates whether a summary message was created.
	SummaryCreated bool

	// Duration is how long the compaction took.
	Duration time.Duration
}

Result contains the outcome of a compaction operation.

type Stats added in v0.1.4

type Stats struct {
	// SessionID is the session being analyzed.
	SessionID uuid.UUID

	// TotalMessages is the number of messages in the session.
	TotalMessages int

	// TotalTokens is the estimated total token count.
	TotalTokens int

	// UsagePercent is the percentage of context window used.
	UsagePercent float64

	// CompactionCount is the number of times this session has been compacted.
	CompactionCount int

	// PreservedMessages is the count of preserved (non-compactable) messages.
	PreservedMessages int

	// SummaryMessages is the count of summary messages from previous compactions.
	SummaryMessages int

	// CompactableMessages is the count of messages eligible for compaction.
	CompactableMessages int

	// NeedsCompaction indicates if compaction should be triggered.
	NeedsCompaction bool
}

Stats contains statistics about a session's compaction state.

type Strategy

type Strategy string

Strategy represents a compaction strategy.

const (
	// StrategySummarization uses Claude to create a structured summary of older messages.
	StrategySummarization Strategy = "summarization"

	// StrategyHybrid prunes tool outputs first, then summarizes if still needed.
	// This is the recommended strategy as it's more cost-effective.
	StrategyHybrid Strategy = "hybrid"
)

type StrategyExecutor added in v0.1.4

type StrategyExecutor interface {
	// Name returns the strategy name (e.g., "summarization", "hybrid").
	Name() Strategy

	// Execute performs the compaction on the given messages.
	// It returns the result including the summary (if any) and the list of
	// message IDs to archive.
	Execute(ctx context.Context, partition *MessagePartition) (*StrategyResult, error)
}

StrategyExecutor defines the interface for compaction strategy implementations. Each strategy handles the actual compaction logic differently.

type StrategyFactory added in v0.1.4

type StrategyFactory struct {
	// contains filtered or unexported fields
}

StrategyFactory creates strategy executors based on configuration.

func NewStrategyFactory added in v0.1.4

func NewStrategyFactory(config *Config, tokenCounter *TokenCounter, summarizer *Summarizer) *StrategyFactory

NewStrategyFactory creates a new strategy factory.

func (*StrategyFactory) Create added in v0.1.4

func (f *StrategyFactory) Create() StrategyExecutor

Create returns the appropriate strategy executor for the configured strategy.

type StrategyResult added in v0.1.4

type StrategyResult struct {
	// SummaryText is the generated summary text (may be empty for prune-only).
	SummaryText string

	// SummaryTokens is the estimated token count of the summary.
	SummaryTokens int

	// ArchivedMessageIDs is the list of message IDs that should be archived.
	ArchivedMessageIDs []uuid.UUID

	// TokensRemoved is the estimated number of tokens removed.
	TokensRemoved int

	// TokensAfter is the estimated token count after compaction.
	TokensAfter int

	// Duration is how long the strategy execution took.
	Duration time.Duration
}

StrategyResult contains the result of executing a compaction strategy.

type SummarizationStrategy

type SummarizationStrategy struct {
	// contains filtered or unexported fields
}

SummarizationStrategy implements the StrategyExecutor interface using pure summarization.

func NewSummarizationStrategy

func NewSummarizationStrategy(summarizer *Summarizer, tokenCounter *TokenCounter) *SummarizationStrategy

NewSummarizationStrategy creates a new summarization strategy.

func (*SummarizationStrategy) Execute added in v0.1.4

Execute performs the summarization compaction.

func (*SummarizationStrategy) Name

func (s *SummarizationStrategy) Name() Strategy

Name returns the strategy name.

type Summarizer added in v0.1.4

type Summarizer struct {
	// contains filtered or unexported fields
}

Summarizer handles the creation of conversation summaries using Claude's streaming API.

func NewSummarizer added in v0.1.4

func NewSummarizer(client *anthropic.Client, model string, maxTokens int) *Summarizer

NewSummarizer creates a new Summarizer with the given Anthropic client and configuration.

func (*Summarizer) Summarize added in v0.1.4

func (s *Summarizer) Summarize(ctx context.Context, messages []*driver.Message) (string, error)

Summarize generates a summary of the given messages using Claude's streaming API. It returns the summary text and any error encountered.

func (*Summarizer) SummarizeWithContext added in v0.1.4

func (s *Summarizer) SummarizeWithContext(ctx context.Context, contextMsgs, toSummarize []*driver.Message) (string, error)

SummarizeWithContext generates a summary with additional context from previous summaries.

type TokenCountResult added in v0.1.4

type TokenCountResult struct {
	// TotalTokens is the total token count for all messages.
	TotalTokens int

	// UsedAPI indicates whether the Claude API was used (true) or the
	// character-based approximation fallback was used (false).
	UsedAPI bool

	// PerMessage contains the estimated token count per message.
	// Only populated when using the fallback approximation.
	PerMessage []int
}

TokenCountResult contains the result of a token count operation.

type TokenCounter

type TokenCounter struct {
	// contains filtered or unexported fields
}

TokenCounter provides token counting for messages using the Claude API with a character-based approximation fallback.

func NewTokenCounter

func NewTokenCounter(client *anthropic.Client, model string, useAPI bool) *TokenCounter

NewTokenCounter creates a new TokenCounter with the given Anthropic client. If useAPI is false, only character-based approximation will be used.

func (*TokenCounter) CountTokens

func (tc *TokenCounter) CountTokens(ctx context.Context, messages []*driver.Message) (*TokenCountResult, error)

CountTokens counts the tokens in the given messages. It first attempts to use the Claude API for accurate counting, falling back to character-based approximation if the API is unavailable or if useAPI is false.

func (*TokenCounter) CountTokensForContent added in v0.1.4

func (tc *TokenCounter) CountTokensForContent(ctx context.Context, text string) (int, error)

CountTokensForContent counts the tokens for a single text content.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL