loop: preserve cumulative usage across conversation compaction
conversation: add optional CumulativeUsage parameter to New function
Update CompactConversation to preserve cumulative usage statistics when
resetting the conversation, preventing usage numbers from being lost during
compaction.
Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: s11dcb84790847494k
diff --git a/dockerimg/createdockerfile.go b/dockerimg/createdockerfile.go
index a1ea266..8343005 100644
--- a/dockerimg/createdockerfile.go
+++ b/dockerimg/createdockerfile.go
@@ -165,7 +165,7 @@
return llm.TextContent("OK"), nil
}
- convo := conversation.New(ctx, srv)
+ convo := conversation.New(ctx, srv, nil)
convo.Tools = []*llm.Tool{{
Name: "dockerfile",
diff --git a/llm/conversation/convo.go b/llm/conversation/convo.go
index 95e4fba..f3161ed 100644
--- a/llm/conversation/convo.go
+++ b/llm/conversation/convo.go
@@ -117,13 +117,16 @@
// New creates a new conversation with Claude with sensible defaults.
// ctx is the context for the entire conversation.
-func New(ctx context.Context, srv llm.Service) *Convo {
+func New(ctx context.Context, srv llm.Service, usage *CumulativeUsage) *Convo {
id := newConvoID()
+ if usage == nil {
+ usage = newUsage()
+ }
return &Convo{
Ctx: skribe.ContextWithAttr(ctx, slog.String("convo_id", id)),
Service: srv,
PromptCaching: true,
- usage: newUsage(),
+ usage: usage,
Listener: &NoopListener{},
ID: id,
toolUseCancel: map[string]context.CancelCauseFunc{},
diff --git a/llm/conversation/convo_test.go b/llm/conversation/convo_test.go
index 8794468..694b977 100644
--- a/llm/conversation/convo_test.go
+++ b/llm/conversation/convo_test.go
@@ -30,7 +30,7 @@
APIKey: apiKey,
HTTPC: rr.Client(),
}
- convo := New(ctx, srv)
+ convo := New(ctx, srv, nil)
const name = "Cornelius"
res, err := convo.SendUserTextMessage("Hi, my name is " + name)
@@ -92,7 +92,7 @@
srv := &ant.Service{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
- convo := New(context.Background(), srv)
+ convo := New(context.Background(), srv, nil)
var cancelCalled bool
var cancelledWithErr error
@@ -248,7 +248,7 @@
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
srv := &ant.Service{}
- convo := New(context.Background(), srv)
+ convo := New(context.Background(), srv, nil)
// Create request with messages
req := &llm.Request{
diff --git a/loop/agent.go b/loop/agent.go
index 092181b..5c36bcf 100644
--- a/loop/agent.go
+++ b/loop/agent.go
@@ -622,9 +622,12 @@
contextWindow := a.config.Service.TokenContextWindow()
currentContextSize := lastUsage.InputTokens + lastUsage.CacheReadInputTokens + lastUsage.CacheCreationInputTokens
+ // Preserve cumulative usage across compaction
+ cumulativeUsage := a.convo.CumulativeUsage()
+
// Reset conversation state but keep all other state (git, working dir, etc.)
a.firstMessageIndex = len(a.history)
- a.convo = a.initConvo()
+ a.convo = a.initConvoWithUsage(&cumulativeUsage)
a.mu.Unlock()
@@ -1205,8 +1208,13 @@
// It must not be called until all agent fields are initialized,
// particularly workingDir and git.
func (a *Agent) initConvo() *conversation.Convo {
+ return a.initConvoWithUsage(nil)
+}
+
+// initConvoWithUsage initializes the conversation with optional preserved usage.
+func (a *Agent) initConvoWithUsage(usage *conversation.CumulativeUsage) *conversation.Convo {
ctx := a.config.Context
- convo := conversation.New(ctx, a.config.Service)
+ convo := conversation.New(ctx, a.config.Service, usage)
convo.PromptCaching = true
convo.Budget = a.config.Budget
convo.SystemPrompt = a.renderSystemPrompt()