Merge pull request #1469 from thotam/fix/session-context-usage-tokens

fix(sessions): store last-call context tokens instead of run-cumulative total

Author: clark-cant

internal/agent/loop_pipeline_callbacks.go

@@ -93,7 +93,7 @@ type pipelineCallbackSet struct {
 	checkReadOnly      func(state *pipeline.RunState) (*providers.Message, bool)
 	sanitizeContent    func(string) string
 	flushMessages      func(ctx context.Context, sessionKey string, msgs []providers.Message) error
-	updateMetadata     func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error
+	updateMetadata     func(ctx context.Context, sessionKey string, usage, lastUsage providers.Usage, msgCount int) error
 	bootstrapCleanup   func(ctx context.Context, state *pipeline.RunState) error
 	maybeSummarize     func(ctx context.Context, sessionKey string)
 }
@@ -713,14 +713,21 @@ func (l *Loop) makeFlushMessages(req *RunRequest) func(ctx context.Context, sess
 	}
 }
 
-func (l *Loop) makeUpdateMetadata(req *RunRequest) func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error {
-	return func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error {
+func (l *Loop) makeUpdateMetadata(req *RunRequest) func(ctx context.Context, sessionKey string, usage, lastUsage providers.Usage, msgCount int) error {
+	return func(ctx context.Context, sessionKey string, usage, lastUsage providers.Usage, msgCount int) error {
 		l.sessions.UpdateMetadata(ctx, sessionKey, l.model, l.provider.Name(), req.Channel)
 		l.sessions.AccumulateTokens(ctx, sessionKey, int64(usage.PromptTokens), int64(usage.CompletionTokens))
 		// Persist session to DB (matching v2 finalizeRun behavior).
 		// FlushMessages already ran, so all pending messages are in the cache.
-		if usage.PromptTokens > 0 {
-			l.sessions.SetLastPromptTokens(ctx, sessionKey, usage.PromptTokens, msgCount)
+		// Calibration uses the FINAL iteration's context size, NOT the
+		// run-cumulative total — the total sums every think→act→observe
+		// iteration and inflated the sessions "context used" display and
+		// compaction decisions by the iteration count.
+		// Current context = final prompt (incl. cached segments) + final output:
+		// the last reply joins history, so it occupies the next request's prompt.
+		// This matches msgCount, which already counts the flushed reply.
+		if lastCtx := lastUsage.ContextTokens(); lastCtx > 0 {
+			l.sessions.SetLastPromptTokens(ctx, sessionKey, lastCtx+lastUsage.CompletionTokens, msgCount)
 		}
 		l.sessions.Save(ctx, sessionKey)
 		return nil

internal/pipeline/deps.go

@@ -121,9 +121,13 @@ type PipelineDeps struct {
 	DeduplicateMediaSuffix func(content, suffix string) string
 	IsSilentReply          func(content string) bool
 	EmitSessionCompleted   func(ctx context.Context, sessionKey string, msgCount, tokensUsed, compactionCount int)
-	UpdateMetadata         func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error
-	BootstrapCleanup       func(ctx context.Context, state *RunState) error
-	MaybeSummarize         func(ctx context.Context, sessionKey string)
+	// UpdateMetadata persists run token accounting: usage is the run-cumulative
+	// total (billing/AccumulateTokens); lastUsage is the final iteration's own
+	// usage — its ContextTokens() is the session's current context size
+	// (SetLastPromptTokens → sessions context display + compaction calibration).
+	UpdateMetadata   func(ctx context.Context, sessionKey string, usage, lastUsage providers.Usage, msgCount int) error
+	BootstrapCleanup func(ctx context.Context, state *RunState) error
+	MaybeSummarize   func(ctx context.Context, sessionKey string)
 }
 
 // FireHook is nil-safe. Returns FireResult{Decision: DecisionAllow} when no

internal/pipeline/finalize_stage.go

@@ -133,7 +133,7 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error {
 	// 5. Update session metadata (token usage)
 	if s.deps.UpdateMetadata != nil {
 		msgCount := historyCountBeforeFlush + len(persistablePending)
-		if err := s.deps.UpdateMetadata(ctx, state.Input.SessionKey, state.Think.TotalUsage, msgCount); err != nil {
+		if err := s.deps.UpdateMetadata(ctx, state.Input.SessionKey, state.Think.TotalUsage, state.Think.LastUsage, msgCount); err != nil {
 			slog.Warn("finalize metadata update failed", "err", err)
 		}
 	}

internal/pipeline/finalize_stage_image_delivery_test.go

@@ -23,7 +23,7 @@ func TestFinalizeStage_GeneratedImageReachesMediaResults(t *testing.T) {
 
 	deps := &PipelineDeps{
 		FlushMessages:  func(context.Context, string, []providers.Message) error { return nil },
-		UpdateMetadata: func(context.Context, string, providers.Usage, int) error { return nil },
+		UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil },
 		// Simulates persistAssistantImages: real implementation decodes base64,
 		// writes to workspace/media/, and appends MediaRefs to msg.
 		PersistAssistantImages: func(msg *providers.Message, _ string) {
@@ -73,7 +73,7 @@ func TestFinalizeStage_NoGeneratedImagesNoMediaResults(t *testing.T) {
 
 	deps := &PipelineDeps{
 		FlushMessages:  func(context.Context, string, []providers.Message) error { return nil },
-		UpdateMetadata: func(context.Context, string, providers.Usage, int) error { return nil },
+		UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil },
 	}
 
 	stage := NewFinalizeStage(deps)

internal/pipeline/finalize_stage_metadata_test.go

@@ -11,13 +11,15 @@ func TestFinalizeStage_UpdateMetadataReceivesPersistedMessageCount(t *testing.T)
 	t.Parallel()
 
 	var gotUsage providers.Usage
+	var gotLastUsage providers.Usage
 	var gotMsgCount int
 	deps := &PipelineDeps{
 		FlushMessages: func(_ context.Context, _ string, _ []providers.Message) error {
 			return nil
 		},
-		UpdateMetadata: func(_ context.Context, _ string, usage providers.Usage, msgCount int) error {
+		UpdateMetadata: func(_ context.Context, _ string, usage, lastUsage providers.Usage, msgCount int) error {
 			gotUsage = usage
+			gotLastUsage = lastUsage
 			gotMsgCount = msgCount
 			return nil
 		},
@@ -33,13 +35,21 @@ func TestFinalizeStage_UpdateMetadataReceivesPersistedMessageCount(t *testing.T)
 	state.Messages.AppendPending(providers.Message{Role: "assistant", Content: "transient", Transient: true})
 	state.Observe.FinalContent = "final answer"
 	state.Think.TotalUsage = providers.Usage{PromptTokens: 1234, CompletionTokens: 56}
+	state.Think.LastUsage = providers.Usage{PromptTokens: 700, CompletionTokens: 30}
 
 	if err := stage.Execute(context.Background(), state); err != nil {
 		t.Fatalf("Execute() error: %v", err)
 	}
 	if gotUsage.PromptTokens != 1234 || gotUsage.CompletionTokens != 56 {
 		t.Fatalf("UpdateMetadata usage = %+v, want prompt=1234 completion=56", gotUsage)
 	}
+	// Regression guard: the last-call usage (current context size) must reach
+	// UpdateMetadata separately from the run-cumulative total — the sessions
+	// "context used" bar and compaction calibration read it via
+	// SetLastPromptTokens. Passing the total inflated it by the iteration count.
+	if gotLastUsage.PromptTokens != 700 {
+		t.Fatalf("UpdateMetadata lastUsage.PromptTokens = %d, want 700 (final iteration only)", gotLastUsage.PromptTokens)
+	}
 	if gotMsgCount != 4 {
 		t.Fatalf("UpdateMetadata msgCount = %d, want 4 persisted messages", gotMsgCount)
 	}

internal/pipeline/substates.go

@@ -29,8 +29,14 @@ type ContextState struct {
 
 // ThinkState: owned by ThinkStage.
 type ThinkState struct {
-	LastResponse    *providers.ChatResponse
-	TotalUsage      providers.Usage
+	LastResponse *providers.ChatResponse
+	TotalUsage   providers.Usage
+	// LastUsage snapshots the most recent iteration that reported prompt tokens.
+	// Unlike TotalUsage (run-cumulative), it reflects the actual size of the last
+	// prompt sent to the model — the session's current context. Consumed by
+	// FinalizeStage → UpdateMetadata → SetLastPromptTokens for the sessions
+	// context-usage display and compaction calibration.
+	LastUsage       providers.Usage
 	TruncRetries    int  // consecutive truncation retries (max 3)
 	OverflowRetries int  // context overflow compact+retry attempts (max 1)
 	StreamingActive bool // true during active stream

internal/pipeline/think_stage.go

@@ -95,6 +95,12 @@ func (s *ThinkStage) Execute(ctx context.Context, state *RunState) error {
 		if resp.Usage.PromptTokensIncludeCachedSegments {
 			state.Think.TotalUsage.PromptTokensIncludeCachedSegments = true
 		}
+		// Snapshot per-call usage: the LAST iteration's prompt size IS the
+		// session's current context. Keep the previous snapshot when a response
+		// carries no prompt tokens (e.g. providers that omit usage on some turns).
+		if resp.Usage.PromptTokens > 0 {
+			state.Think.LastUsage = *resp.Usage
+		}
 	}
 
 	if isEmptyLengthResponse(resp) {

internal/pipeline/think_stage_last_usage_test.go

@@ -0,0 +1,88 @@
+package pipeline
+
+import (
+	"context"
+	"testing"
+
+	"github.com/nextlevelbuilder/goclaw/internal/providers"
+)
+
+// TestThinkStage_TracksLastUsagePerIteration guards the per-iteration usage
+// snapshot used for the sessions "context used" display and compaction
+// calibration. Regression: SetLastPromptTokens received the RUN-CUMULATIVE
+// TotalUsage.PromptTokens (sum over all think→act→observe iterations), so a
+// tool-heavy run reported e.g. 901K "context" against a 258K window and
+// over-triggered compaction. LastUsage must hold only the FINAL iteration's
+// own usage — the actual size of the last prompt sent to the model.
+func TestThinkStage_TracksLastUsagePerIteration(t *testing.T) {
+	t.Parallel()
+
+	usages := []providers.Usage{
+		{PromptTokens: 100000, CompletionTokens: 500, TotalTokens: 100500},
+		{PromptTokens: 120000, CompletionTokens: 300, TotalTokens: 120300, CacheReadTokens: 90000},
+	}
+	call := 0
+	deps := &PipelineDeps{
+		Config: PipelineConfig{MaxIterations: 10, MaxTokens: 1000},
+		CallLLM: func(_ context.Context, _ *RunState, _ providers.ChatRequest) (*providers.ChatResponse, error) {
+			u := usages[call]
+			call++
+			return &providers.ChatResponse{FinishReason: "stop", Content: "done", Usage: &u}, nil
+		},
+	}
+	stage := NewThinkStage(deps)
+	state := defaultState()
+
+	if err := stage.Execute(context.Background(), state); err != nil {
+		t.Fatalf("Execute() 1: %v", err)
+	}
+	if err := stage.Execute(context.Background(), state); err != nil {
+		t.Fatalf("Execute() 2: %v", err)
+	}
+
+	// TotalUsage accumulates (unchanged behavior).
+	if state.Think.TotalUsage.PromptTokens != 220000 {
+		t.Errorf("TotalUsage.PromptTokens = %d, want 220000", state.Think.TotalUsage.PromptTokens)
+	}
+	// LastUsage holds only the final iteration's usage — NOT the sum.
+	if state.Think.LastUsage.PromptTokens != 120000 {
+		t.Errorf("LastUsage.PromptTokens = %d, want 120000 (final iteration only)", state.Think.LastUsage.PromptTokens)
+	}
+	if state.Think.LastUsage.CacheReadTokens != 90000 {
+		t.Errorf("LastUsage.CacheReadTokens = %d, want 90000", state.Think.LastUsage.CacheReadTokens)
+	}
+}
+
+// TestThinkStage_LastUsageKeptWhenFinalResponseLacksUsage: a final iteration
+// whose response carries no usage (nil or zero prompt tokens) must not wipe the
+// last usable snapshot.
+func TestThinkStage_LastUsageKeptWhenFinalResponseLacksUsage(t *testing.T) {
+	t.Parallel()
+
+	responses := []*providers.ChatResponse{
+		{FinishReason: "stop", Content: "a", Usage: &providers.Usage{PromptTokens: 50000, CompletionTokens: 10}},
+		{FinishReason: "stop", Content: "b", Usage: nil},
+	}
+	call := 0
+	deps := &PipelineDeps{
+		Config: PipelineConfig{MaxIterations: 10, MaxTokens: 1000},
+		CallLLM: func(_ context.Context, _ *RunState, _ providers.ChatRequest) (*providers.ChatResponse, error) {
+			r := responses[call]
+			call++
+			return r, nil
+		},
+	}
+	stage := NewThinkStage(deps)
+	state := defaultState()
+
+	if err := stage.Execute(context.Background(), state); err != nil {
+		t.Fatalf("Execute() 1: %v", err)
+	}
+	if err := stage.Execute(context.Background(), state); err != nil {
+		t.Fatalf("Execute() 2: %v", err)
+	}
+
+	if state.Think.LastUsage.PromptTokens != 50000 {
+		t.Errorf("LastUsage.PromptTokens = %d, want 50000 (kept from iteration 1)", state.Think.LastUsage.PromptTokens)
+	}
+}

internal/providers/call_usage.go

@@ -32,6 +32,19 @@ func SumCallUsage(calls []CallUsage) Usage {
 	return total
 }
 
+// ContextTokens returns the full prompt-side size of a single LLM call —
+// the actual context the model saw — normalizing provider cache accounting.
+// Anthropic-style usage reports input_tokens EXCLUDING cached segments, so
+// cache read/creation tokens must be added back; OpenAI-style usage
+// (PromptTokensIncludeCachedSegments=true) already includes them.
+// Only meaningful for per-call usage, not run-cumulative totals.
+func (u Usage) ContextTokens() int {
+	if u.PromptTokensIncludeCachedSegments {
+		return u.PromptTokens
+	}
+	return u.PromptTokens + u.CacheReadTokens + u.CacheCreationTokens
+}
+
 // SumCallCost totals the per-call USD cost (0 when pricing was unavailable).
 func SumCallCost(calls []CallUsage) float64 {
 	var total float64

internal/providers/usage_context_tokens_test.go

@@ -0,0 +1,49 @@
+package providers
+
+import "testing"
+
+// TestUsage_ContextTokens guards the "current context size" computation used by
+// the sessions context-usage display and compaction calibration. Anthropic-style
+// accounting reports input_tokens EXCLUDING cached segments, so the real context
+// of a call is prompt + cache_read + cache_creation. OpenAI-style accounting
+// (PromptTokensIncludeCachedSegments=true) already includes cached tokens.
+func TestUsage_ContextTokens(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name  string
+		usage Usage
+		want  int
+	}{
+		{
+			name: "anthropic-style excludes cached segments — add them back",
+			usage: Usage{
+				PromptTokens:        1000,
+				CacheReadTokens:     80000,
+				CacheCreationTokens: 5000,
+			},
+			want: 86000,
+		},
+		{
+			name: "openai-style already includes cached segments",
+			usage: Usage{
+				PromptTokens:                      86000,
+				CacheReadTokens:                   80000,
+				PromptTokensIncludeCachedSegments: true,
+			},
+			want: 86000,
+		},
+		{
+			name:  "zero value",
+			usage: Usage{},
+			want:  0,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.usage.ContextTokens(); got != tt.want {
+				t.Errorf("ContextTokens() = %d, want %d", got, tt.want)
+			}
+		})
+	}
+}