9595d1f · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(session): preserve estimated usage marker

9595d1fGreg Slepak5 months ago

internal/session/session.go+43−6

compare with main

unifiedsplit

@@ -7,6 +7,7 @@ import (

7 "fmt"

8 "log/slog"

9 "strings"

10+ "sync"

11 "github.com/charmbracelet/crush/internal/db"

12 "github.com/charmbracelet/crush/internal/db"

12 "github.com/charmbracelet/crush/internal/event"

13 "github.com/charmbracelet/crush/internal/event"

@@ -84,6 +85,12 @@ type service struct {

84 *pubsub.Broker[Session]

85 *pubsub.Broker[Session]

85 db *sql.DB

86 db *sql.DB

86 q *db.Queries

87 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

87}

94}

89func (s *service) Create(ctx context.Context, title string) (Session, error) {

96func (s *service) Create(ctx context.Context, title string) (Session, error) {

@@ -155,6 +162,7 @@ func (s *service) Delete(ctx context.Context, id string) error {

155 }

162 }

156

163

157 session := s.fromDBItem(dbSession)

164 session := s.fromDBItem(dbSession)

165+ s.clearEstimatedUsageState(dbSession.ID)

158 s.Publish(pubsub.DeletedEvent, session)

166 s.Publish(pubsub.DeletedEvent, session)

159 event.SessionDeleted()

167 event.SessionDeleted()

160 return nil

168 return nil

@@ -165,7 +173,9 @@ func (s *service) Get(ctx context.Context, id string) (Session, error) {

165 if err != nil {

173 if err != nil {

166 return Session{}, err

174 return Session{}, err

167 }

175 }

168− return s.fromDBItem(dbSession), nil

176+ session := s.fromDBItem(dbSession)

177+ s.applyEstimatedUsageState(&session)

178+ return session, nil

169}

179}

170

180

171func (s *service) GetLast(ctx context.Context) (Session, error) {

181func (s *service) GetLast(ctx context.Context) (Session, error) {

@@ -173,7 +183,9 @@ func (s *service) GetLast(ctx context.Context) (Session, error) {

173 if err != nil {

183 if err != nil {

174 return Session{}, err

184 return Session{}, err

175 }

185 }

176− return s.fromDBItem(dbSession), nil

186+ session := s.fromDBItem(dbSession)

187+ s.applyEstimatedUsageState(&session)

188+ return session, nil

177}

189}

178

190

179func (s *service) Save(ctx context.Context, session Session) (Session, error) {

191func (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) {

201 return Session{}, err

213 return Session{}, err

202 }

214 }

203 estimatedUsage := session.EstimatedUsage

215 estimatedUsage := session.EstimatedUsage

216+ s.setEstimatedUsageState(session.ID, estimatedUsage)

204 session = s.fromDBItem(dbSession)

217 session = s.fromDBItem(dbSession)

205 session.EstimatedUsage = estimatedUsage

218 session.EstimatedUsage = estimatedUsage

206 s.Publish(pubsub.UpdatedEvent, session)

219 s.Publish(pubsub.UpdatedEvent, session)

@@ -236,11 +249,34 @@ func (s *service) List(ctx context.Context) ([]Session, error) {

236 sessions := make([]Session, len(dbSessions))

249 sessions := make([]Session, len(dbSessions))

237 for i, dbSession := range dbSessions {

250 for i, dbSession := range dbSessions {

238 sessions[i] = s.fromDBItem(dbSession)

251 sessions[i] = s.fromDBItem(dbSession)

252+ s.applyEstimatedUsageState(&sessions[i])

239 }

253 }

240 return sessions, nil

254 return sessions, nil

241}

255}

242

256

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+

243−func (s service) fromDBItem(item db.Session) Session {

279+func (s *service) fromDBItem(item db.Session) Session {

244 todos, err := unmarshalTodos(item.Todos.String)

280 todos, err := unmarshalTodos(item.Todos.String)

245 if err != nil {

281 if err != nil {

246 slog.Error("Failed to unmarshal todos", "session_id", item.ID, "error", err)

282 slog.Error("Failed to unmarshal todos", "session_id", item.ID, "error", err)

@@ -285,9 +321,10 @@ func unmarshalTodos(data string) ([]Todo, error) {

285func NewService(q *db.Queries, conn *sql.DB) Service {

321func NewService(q *db.Queries, conn *sql.DB) Service {

286 broker := pubsub.NewBroker[Session]()

322 broker := pubsub.NewBroker[Session]()

287 return &service{

323 return &service{

288− Broker: broker,

324+ Broker: broker,

289− db: conn,

325+ db: conn,

290− q: q,

326+ q: q,

327+ estimatedUsage: make(map[string]bool),

291 }

328 }

292}

329}

293

330

internal/session/session_test.gonew file+81

@@ -0,0 +1,81 @@

1+package session

3+import (

4+ "testing"

6+ "github.com/charmbracelet/crush/internal/db"

7+ "github.com/stretchr/testify/require"

8+)

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 {

763 s.ModelInfo.Reasoning = lipgloss.NewStyle().Foreground(o.fgMostSubtle).PaddingLeft(2)

764 s.ModelInfo.TokenCount = lipgloss.NewStyle().Foreground(o.fgMostSubtle)

765 s.ModelInfo.TokenPercentage = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)

766− s.ModelInfo.EstimatedUsagePrefix = lipgloss.NewStyle().Foreground(o.fgBase)

766+ s.ModelInfo.EstimatedUsagePrefix = s.ModelInfo.TokenPercentage

767 s.ModelInfo.Cost = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)

768 s.ModelInfo.HypercreditIcon = lipgloss.NewStyle().Foreground(charmtone.Dolly)

769 s.ModelInfo.HypercreditText = lipgloss.NewStyle().Foreground(o.fgMoreSubtle)