git.iko.soy/crush
on main help sign in
fix(session): preserve estimated usage marker

Keep estimated usage state in memory across session fetch-modify-save updates so unrelated saves do not clear the UI marker, and align the marker color with context percentages.

commit9595d1f073dde42e1749853688299a4bad19f496 authorGreg Slepak <contact@taoeffect.com> dateTue 12 May 2026, 18:10 (5 months ago) parent2736e48
3 files changed+125−7
compare with main
unifiedsplit
internal/session/session.go+43−6 internal/session/session_test.go+81 internal/ui/styles/quickstyle.go+1−1
internal/session/session.go+43−6
@@ -7,6 +7,7 @@ import (
77 "fmt"
88 "log/slog"
99 "strings"
10+ "sync"
1011
1112 "github.com/charmbracelet/crush/internal/db"
1213 "github.com/charmbracelet/crush/internal/event"
@@ -84,6 +85,12 @@ type service struct {
8485 *pubsub.Broker[Session]
8586 db *sql.DB
8687 q *db.Queries
88+
89+ // Estimated usage stays in memory so fetch-modify-save paths (e.g.,
90+ // updating todos or parent-session cost) do not rebuild a session from
91+ // SQLite and incorrectly clear the UI "~" marker.
92+ estimatedUsageMu sync.RWMutex
93+ estimatedUsage map[string]bool
8794}
8895
8996func (s *service) Create(ctx context.Context, title string) (Session, error) {
@@ -155,6 +162,7 @@ func (s *service) Delete(ctx context.Context, id string) error {
155162 }
156163
157164 session := s.fromDBItem(dbSession)
165+ s.clearEstimatedUsageState(dbSession.ID)
158166 s.Publish(pubsub.DeletedEvent, session)
159167 event.SessionDeleted()
160168 return nil
@@ -165,7 +173,9 @@ func (s *service) Get(ctx context.Context, id string) (Session, error) {
165173 if err != nil {
166174 return Session{}, err
167175 }
168− return s.fromDBItem(dbSession), nil
176+ session := s.fromDBItem(dbSession)
177+ s.applyEstimatedUsageState(&session)
178+ return session, nil
169179}
170180
171181func (s *service) GetLast(ctx context.Context) (Session, error) {
@@ -173,7 +183,9 @@ func (s *service) GetLast(ctx context.Context) (Session, error) {
173183 if err != nil {
174184 return Session{}, err
175185 }
176− return s.fromDBItem(dbSession), nil
186+ session := s.fromDBItem(dbSession)
187+ s.applyEstimatedUsageState(&session)
188+ return session, nil
177189}
178190
179191func (s *service) Save(ctx context.Context, session Session) (Session, error) {
@@ -201,6 +213,7 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) {
201213 return Session{}, err
202214 }
203215 estimatedUsage := session.EstimatedUsage
216+ s.setEstimatedUsageState(session.ID, estimatedUsage)
204217 session = s.fromDBItem(dbSession)
205218 session.EstimatedUsage = estimatedUsage
206219 s.Publish(pubsub.UpdatedEvent, session)
@@ -236,11 +249,34 @@ func (s *service) List(ctx context.Context) ([]Session, error) {
236249 sessions := make([]Session, len(dbSessions))
237250 for i, dbSession := range dbSessions {
238251 sessions[i] = s.fromDBItem(dbSession)
252+ s.applyEstimatedUsageState(&sessions[i])
239253 }
240254 return sessions, nil
241255}
242256
243−func (s service) fromDBItem(item db.Session) Session {
257+func (s *service) applyEstimatedUsageState(session *Session) {
258+ s.estimatedUsageMu.RLock()
259+ session.EstimatedUsage = s.estimatedUsage[session.ID]
260+ s.estimatedUsageMu.RUnlock()
261+}
262+
263+func (s *service) setEstimatedUsageState(sessionID string, estimatedUsage bool) {
264+ s.estimatedUsageMu.Lock()
265+ defer s.estimatedUsageMu.Unlock()
266+ if estimatedUsage {
267+ s.estimatedUsage[sessionID] = true
268+ return
269+ }
270+ delete(s.estimatedUsage, sessionID)
271+}
272+
273+func (s *service) clearEstimatedUsageState(sessionID string) {
274+ s.estimatedUsageMu.Lock()
275+ delete(s.estimatedUsage, sessionID)
276+ s.estimatedUsageMu.Unlock()
277+}
278+
279+func (s *service) fromDBItem(item db.Session) Session {
244280 todos, err := unmarshalTodos(item.Todos.String)
245281 if err != nil {
246282 slog.Error("Failed to unmarshal todos", "session_id", item.ID, "error", err)
@@ -285,9 +321,10 @@ func unmarshalTodos(data string) ([]Todo, error) {
285321func NewService(q *db.Queries, conn *sql.DB) Service {
286322 broker := pubsub.NewBroker[Session]()
287323 return &service{
288− Broker: broker,
289− db: conn,
290− q: q,
324+ Broker: broker,
325+ db: conn,
326+ q: q,
327+ estimatedUsage: make(map[string]bool),
291328 }
292329}
293330
internal/session/session_test.gonew file+81
@@ -0,0 +1,81 @@
1+package session
2+
3+import (
4+ "testing"
5+
6+ "github.com/charmbracelet/crush/internal/db"
7+ "github.com/stretchr/testify/require"
8+)
9+
10+func TestEstimatedUsageStateSurvivesFetchModifySave(t *testing.T) {
11+ dataDir := t.TempDir()
12+ t.Cleanup(func() {
13+ require.NoError(t, db.Release(dataDir))
14+ db.ResetPool()
15+ })
16+
17+ conn, err := db.Connect(t.Context(), dataDir)
18+ require.NoError(t, err)
19+
20+ sessions := NewService(db.New(conn), conn)
21+
22+ created, err := sessions.Create(t.Context(), "test")
23+ require.NoError(t, err)
24+ created.PromptTokens = 100
25+ created.CompletionTokens = 50
26+ created.EstimatedUsage = true
27+
28+ saved, err := sessions.Save(t.Context(), created)
29+ require.NoError(t, err)
30+ require.True(t, saved.EstimatedUsage)
31+
32+ fetched, err := sessions.Get(t.Context(), created.ID)
33+ require.NoError(t, err)
34+ require.True(t, fetched.EstimatedUsage)
35+
36+ fetched.Todos = []Todo{{
37+ Content: "Check estimate state",
38+ Status: TodoStatusInProgress,
39+ ActiveForm: "Checking estimate state",
40+ }}
41+
42+ updated, err := sessions.Save(t.Context(), fetched)
43+ require.NoError(t, err)
44+ require.True(t, updated.EstimatedUsage)
45+
46+ refetched, err := sessions.Get(t.Context(), created.ID)
47+ require.NoError(t, err)
48+ require.True(t, refetched.EstimatedUsage)
49+}
50+
51+func TestEstimatedUsageStateCanBeClearedByExplicitSave(t *testing.T) {
52+ dataDir := t.TempDir()
53+ t.Cleanup(func() {
54+ require.NoError(t, db.Release(dataDir))
55+ db.ResetPool()
56+ })
57+
58+ conn, err := db.Connect(t.Context(), dataDir)
59+ require.NoError(t, err)
60+
61+ sessions := NewService(db.New(conn), conn)
62+
63+ created, err := sessions.Create(t.Context(), "test")
64+ require.NoError(t, err)
65+ created.PromptTokens = 100
66+ created.CompletionTokens = 50
67+ created.EstimatedUsage = true
68+
69+ saved, err := sessions.Save(t.Context(), created)
70+ require.NoError(t, err)
71+ require.True(t, saved.EstimatedUsage)
72+
73+ saved.EstimatedUsage = false
74+ updated, err := sessions.Save(t.Context(), saved)
75+ require.NoError(t, err)
76+ require.False(t, updated.EstimatedUsage)
77+
78+ refetched, err := sessions.Get(t.Context(), created.ID)
79+ require.NoError(t, err)
80+ require.False(t, refetched.EstimatedUsage)
81+}
internal/ui/styles/quickstyle.go+1−1
@@ -763,7 +763,7 @@ func quickStyle(o quickStyleOpts) Styles {
763763 s.ModelInfo.Reasoning = lipgloss.NewStyle().Foreground(o.fgMostSubtle).PaddingLeft(2)
764764 s.ModelInfo.TokenCount = lipgloss.NewStyle().Foreground(o.fgMostSubtle)
765765 s.ModelInfo.TokenPercentage = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)
766− s.ModelInfo.EstimatedUsagePrefix = lipgloss.NewStyle().Foreground(o.fgBase)
766+ s.ModelInfo.EstimatedUsagePrefix = s.ModelInfo.TokenPercentage
767767 s.ModelInfo.Cost = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)
768768 s.ModelInfo.HypercreditIcon = lipgloss.NewStyle().Foreground(charmtone.Dolly)
769769 s.ModelInfo.HypercreditText = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)