From 3e2aba45ed1b1cf980534c6629e8d285a1cfff18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Tue, 2 Jun 2026 22:58:31 +0100 Subject: [PATCH] fix(agentctl): sticky-max context window size with reset on model switch (#1254) --- .../adapter/transport/acp/adapter_session.go | 1 + .../adapter/transport/acp/adapter_updates.go | 47 +++++++--- .../adapter/transport/acp/conversion_test.go | 92 +++++++++++++++++++ .../adapter/transport/acp/usage_test.go | 34 ++++--- 4 files changed, 149 insertions(+), 25 deletions(-) diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_session.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_session.go index 92dc4225ac..1556a94036 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_session.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_session.go @@ -525,6 +525,7 @@ func (a *Adapter) SetModel(ctx context.Context, modelID string) error { if err != nil { return fmt.Errorf("set session model failed: %w", err) } + a.resetContextWindowMaxSize(sessionID) a.emitSetModelEvent(sessionID, modelID, available, cachedConfig) return nil } diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go index 7664376386..fc56649064 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go @@ -270,27 +270,50 @@ type acpUsageUpdate struct { type usageTracker struct { lastUsed int64 lastCostSubcents int64 + // maxSize is the largest context-window size reported for this session. + // claude-acp's default model emits a 200K turn-start frame and a 1M + // turn-end frame; sticky-max prevents the stale 200K from shrinking + // the displayed window after the authoritative 1M frame lands. + maxSize int64 } -// recordUsageDelta updates the per-session tracker. Returns the -// integer delta against the previous `used` snapshot — the next prompt -// complete call consumes this via consumeUsageDelta. Cost is forwarded -// as the latest value (claude-acp emits a per-turn cumulative cost; we -// treat the most recent value as the turn's cost). -func (a *Adapter) recordUsageDelta(sessionID string, used int64, costSubcents int64) { - a.mu.Lock() - defer a.mu.Unlock() +// ensureUsageTracker returns the per-session usage tracker, creating it if +// needed. Callers must hold a.mu (write lock). +func (a *Adapter) ensureUsageTracker(sessionID string) *usageTracker { tr := a.usageBySession[sessionID] if tr == nil { tr = &usageTracker{} a.usageBySession[sessionID] = tr } + return tr +} + +// recordUsageAndMaxSize updates per-session usage/cost tracking and returns +// the sticky-max context window size for the session. +func (a *Adapter) recordUsageAndMaxSize(sessionID string, size, used, costSubcents int64) int64 { + a.mu.Lock() + defer a.mu.Unlock() + tr := a.ensureUsageTracker(sessionID) + if size > tr.maxSize { + tr.maxSize = size + } if used > tr.lastUsed { tr.lastUsed = used } if costSubcents > 0 { tr.lastCostSubcents = costSubcents } + return tr.maxSize +} + +// resetContextWindowMaxSize clears the sticky-max window for a session so the +// next usage_update can establish the new model's window after SetModel. +func (a *Adapter) resetContextWindowMaxSize(sessionID string) { + a.mu.Lock() + defer a.mu.Unlock() + if tr := a.usageBySession[sessionID]; tr != nil { + tr.maxSize = 0 + } } // consumeUsageDelta returns the delta of cumulative `used` since the @@ -339,16 +362,16 @@ func (a *Adapter) tryConvertUntypedUpdate(rawNotification []byte, sessionID stri if usage.Cost != nil { costSubcents = int64(usage.Cost.Amount * 10000) } - a.recordUsageDelta(sessionID, usage.Used, costSubcents) + effectiveSize := a.recordUsageAndMaxSize(sessionID, usage.Size, usage.Used, costSubcents) - remaining := max(usage.Size-usage.Used, 0) + remaining := max(effectiveSize-usage.Used, 0) return &AgentEvent{ Type: streams.EventTypeContextWindow, SessionID: sessionID, - ContextWindowSize: usage.Size, + ContextWindowSize: effectiveSize, ContextWindowUsed: usage.Used, ContextWindowRemaining: remaining, - ContextEfficiency: float64(usage.Used) / float64(usage.Size) * 100, + ContextEfficiency: float64(usage.Used) / float64(effectiveSize) * 100, } } diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/conversion_test.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/conversion_test.go index 22a4ac1cb8..09bc6d2144 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/conversion_test.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/conversion_test.go @@ -1,6 +1,7 @@ package acp import ( + "fmt" "testing" "github.com/coder/acp-go-sdk" @@ -891,6 +892,97 @@ func TestTryConvertUntypedUpdate_ZeroSize(t *testing.T) { } } +func usageUpdateRaw(size, used int64) []byte { + return []byte(fmt.Sprintf( + `{"sessionId":"s1","update":{"sessionUpdate":"usage_update","size":%d,"used":%d}}`, + size, used, + )) +} + +func TestTryConvertUntypedUpdate_StickyMaxSize(t *testing.T) { + t.Run("default turn raises max from 200K to 1M", func(t *testing.T) { + a := newTestAdapter() + + first := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 5_000), "s1") + if first == nil || first.ContextWindowSize != 200_000 { + t.Fatalf("first size = %d, want 200000", sizeOrZero(first)) + } + + second := a.tryConvertUntypedUpdate(usageUpdateRaw(1_000_000, 5_000), "s1") + if second == nil || second.ContextWindowSize != 1_000_000 { + t.Fatalf("second size = %d, want 1000000", sizeOrZero(second)) + } + }) + + t.Run("stale 200K start-frame cannot shrink after 1M end-frame", func(t *testing.T) { + a := newTestAdapter() + _, _ = a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 5_000), "s1"), a.tryConvertUntypedUpdate(usageUpdateRaw(1_000_000, 5_000), "s1") + + stale := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 233_900), "s1") + if stale == nil { + t.Fatal("expected non-nil stale frame result") + } + if stale.ContextWindowSize != 1_000_000 { + t.Errorf("ContextWindowSize = %d, want 1000000", stale.ContextWindowSize) + } + expectedEff := float64(233_900) / float64(1_000_000) * 100 + if stale.ContextEfficiency != expectedEff { + t.Errorf("ContextEfficiency = %f, want %f", stale.ContextEfficiency, expectedEff) + } + if stale.ContextWindowRemaining != 766_100 { + t.Errorf("ContextWindowRemaining = %d, want 766100", stale.ContextWindowRemaining) + } + }) + + t.Run("sonnet stays at 200K", func(t *testing.T) { + a := newTestAdapter() + first := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 5_000), "s1") + second := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 7_000), "s1") + if first.ContextWindowSize != 200_000 || second.ContextWindowSize != 200_000 { + t.Fatalf("sizes = %d, %d; want 200000 both", first.ContextWindowSize, second.ContextWindowSize) + } + }) + + t.Run("sonnet[1m] is 1M from first frame", func(t *testing.T) { + a := newTestAdapter() + result := a.tryConvertUntypedUpdate(usageUpdateRaw(1_000_000, 5_000), "s1") + if result == nil || result.ContextWindowSize != 1_000_000 { + t.Fatalf("size = %d, want 1000000", sizeOrZero(result)) + } + }) + + t.Run("sessions track max independently", func(t *testing.T) { + a := newTestAdapter() + _, _ = a.tryConvertUntypedUpdate(usageUpdateRaw(1_000_000, 5_000), "s1"), a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 5_000), "s2") + s1 := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 10_000), "s1") + s2 := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 10_000), "s2") + if s1.ContextWindowSize != 1_000_000 { + t.Errorf("s1 size = %d, want 1000000", s1.ContextWindowSize) + } + if s2.ContextWindowSize != 200_000 { + t.Errorf("s2 size = %d, want 200000", s2.ContextWindowSize) + } + }) + + t.Run("reset after model switch allows downshift to 200K", func(t *testing.T) { + a := newTestAdapter() + _, _ = a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 5_000), "s1"), a.tryConvertUntypedUpdate(usageUpdateRaw(1_000_000, 5_000), "s1") + + a.resetContextWindowMaxSize("s1") + afterSwitch := a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 26_000), "s1") + if afterSwitch == nil || afterSwitch.ContextWindowSize != 200_000 { + t.Fatalf("after switch size = %d, want 200000", sizeOrZero(afterSwitch)) + } + }) +} + +func sizeOrZero(ev *AgentEvent) int64 { + if ev == nil { + return 0 + } + return ev.ContextWindowSize +} + // --- emitInitialModeState --- func TestEmitInitialModeState(t *testing.T) { diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go index d913e52e37..2f365e1b91 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go @@ -1,6 +1,7 @@ package acp import ( + "fmt" "testing" "github.com/coder/acp-go-sdk" @@ -156,16 +157,17 @@ func TestExtractUsage(t *testing.T) { } // TestUsageTracker_CumulativeDelta asserts the codex-acp fallback path: -// usage_update updates the cumulative used counter; consumeUsageDelta -// returns the running total and resets to zero so the next turn starts -// fresh. Cost (when present) is forwarded as the latest snapshot. +// usage_update updates the cumulative used counter via tryConvertUntypedUpdate; +// consumeUsageDelta returns the running total and resets to zero so the next +// turn starts fresh. func TestUsageTracker_CumulativeDelta(t *testing.T) { a := newTestAdapter() + const sess = "sess-codex" - a.recordUsageDelta("sess-codex", 100, 0) - a.recordUsageDelta("sess-codex", 350, 0) + a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 100), sess) + a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 350), sess) - delta, cost := a.consumeUsageDelta("sess-codex") + delta, cost := a.consumeUsageDelta(sess) if delta != 350 { t.Errorf("first consume delta = %d, want 350", delta) } @@ -173,9 +175,8 @@ func TestUsageTracker_CumulativeDelta(t *testing.T) { t.Errorf("first consume cost = %d, want 0 (no cost reported)", cost) } - // After consume the tracker resets; new turn starts from 0. - a.recordUsageDelta("sess-codex", 200, 0) - delta, _ = a.consumeUsageDelta("sess-codex") + a.tryConvertUntypedUpdate(usageUpdateRaw(200_000, 200), sess) + delta, _ = a.consumeUsageDelta(sess) if delta != 200 { t.Errorf("second consume delta = %d, want 200 (reset baseline)", delta) } @@ -186,20 +187,27 @@ func TestUsageTracker_CumulativeDelta(t *testing.T) { // recent value; consume returns it and resets. func TestUsageTracker_ForwardsCost(t *testing.T) { a := newTestAdapter() + const sess = "sess-claude" - a.recordUsageDelta("sess-claude", 25068, 615) // 0.06156125 USD -> 616 (rounded) - _, cost := a.consumeUsageDelta("sess-claude") + a.tryConvertUntypedUpdate(usageUpdateCostRaw(200_000, 25_068, 0.06156125), sess) + _, cost := a.consumeUsageDelta(sess) if cost != 615 { t.Errorf("cost = %d, want 615 (subcents)", cost) } - // Second consume after no updates returns zero (reset). - _, cost = a.consumeUsageDelta("sess-claude") + _, cost = a.consumeUsageDelta(sess) if cost != 0 { t.Errorf("second consume cost = %d, want 0", cost) } } +func usageUpdateCostRaw(size, used int64, costUSD float64) []byte { + return []byte(fmt.Sprintf( + `{"sessionId":"s1","update":{"sessionUpdate":"usage_update","size":%d,"used":%d,"cost":{"amount":%g,"currency":"USD"}}}`, + size, used, costUSD, + )) +} + // TestConsumeUsageDelta_UnknownSession returns zero for sessions that // never recorded usage (e.g. claude-acp where the typed resp.Usage is // the authoritative path and the tracker was never touched).