on main help sign in
fix: race condition (#1649)
commit4f6d0e9373ffb015cba1715a4985298e037d6d14 authorKujtim Hoxha <kujtimii.h@gmail.com> dateFri 19 Dec 2025, 13:19 (1 year ago) parent88b19ad
9 files changed+152−63
internal/agent/agent.go+25−8 internal/db/db.go+54−44 internal/db/querier.go+1 internal/db/sessions.sql.go+29 internal/db/sql/sessions.sql+9 internal/message/content.go+9 internal/message/message.go+9−3 internal/pubsub/broker.go+3−8 internal/session/session.go+13
internal/agent/agent.go+25−8
@@ -171,10 +171,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
171171 var wg sync.WaitGroup
172172 // Generate title if first message.
173173 if len(msgs) == 0 {
174+ titleCtx := ctx // Copy to avoid race with ctx reassignment below.
174175 wg.Go(func() {
175− sessionLock.Lock()
176− a.generateTitle(ctx, &currentSession, call.Prompt)
177− sessionLock.Unlock()
176+ a.generateTitle(titleCtx, call.SessionID, call.Prompt)
178177 })
179178 }
180179
@@ -723,7 +722,7 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
723722 return msgs, nil
724723}
725724
726−func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Session, prompt string) {
725+func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prompt string) {
727726 if prompt == "" {
728727 return
729728 }
@@ -768,8 +767,7 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
768767 return
769768 }
770769
771− session.Title = title
772−
770+ // Calculate usage and cost.
773771 var openrouterCost *float64
774772 for _, step := range resp.Steps {
775773 stepCost := a.openrouterCost(step.ProviderMetadata)
@@ -782,8 +780,27 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
782780 }
783781 }
784782
785− a.updateSessionUsage(a.smallModel, session, resp.TotalUsage, openrouterCost)
786− _, saveErr := a.sessions.Save(ctx, *session)
783+ modelConfig := a.smallModel.CatwalkCfg
784+ cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
785+ modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
786+ modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
787+ modelConfig.CostPer1MOut/1e6*float64(resp.TotalUsage.OutputTokens)
788+
789+ if a.isClaudeCode() {
790+ cost = 0
791+ }
792+
793+ // Use override cost if available (e.g., from OpenRouter).
794+ if openrouterCost != nil {
795+ cost = *openrouterCost
796+ }
797+
798+ promptTokens := resp.TotalUsage.InputTokens + resp.TotalUsage.CacheCreationTokens
799+ completionTokens := resp.TotalUsage.OutputTokens + resp.TotalUsage.CacheReadTokens
800+
801+ // Atomically update only title and usage fields to avoid overriding other
802+ // concurrent session updates.
803+ saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
787804 if saveErr != nil {
788805 slog.Error("failed to save session title & usage", "error", saveErr)
789806 return
internal/db/db.go+54−44
@@ -84,6 +84,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) {
8484 if q.updateSessionStmt, err = db.PrepareContext(ctx, updateSession); err != nil {
8585 return nil, fmt.Errorf("error preparing query UpdateSession: %w", err)
8686 }
87+ if q.updateSessionTitleAndUsageStmt, err = db.PrepareContext(ctx, updateSessionTitleAndUsage); err != nil {
88+ return nil, fmt.Errorf("error preparing query UpdateSessionTitleAndUsage: %w", err)
89+ }
8790 return &q, nil
8891}
8992
@@ -189,6 +192,11 @@ func (q *Queries) Close() error {
189192 err = fmt.Errorf("error closing updateSessionStmt: %w", cerr)
190193 }
191194 }
195+ if q.updateSessionTitleAndUsageStmt != nil {
196+ if cerr := q.updateSessionTitleAndUsageStmt.Close(); cerr != nil {
197+ err = fmt.Errorf("error closing updateSessionTitleAndUsageStmt: %w", cerr)
198+ }
199+ }
192200 return err
193201}
194202
@@ -226,53 +234,55 @@ func (q *Queries) queryRow(ctx context.Context, stmt *sql.Stmt, query string, ar
226234}
227235
228236type Queries struct {
229− db DBTX
230− tx *sql.Tx
231− createFileStmt *sql.Stmt
232− createMessageStmt *sql.Stmt
233− createSessionStmt *sql.Stmt
234− deleteFileStmt *sql.Stmt
235− deleteMessageStmt *sql.Stmt
236− deleteSessionStmt *sql.Stmt
237− deleteSessionFilesStmt *sql.Stmt
238− deleteSessionMessagesStmt *sql.Stmt
239− getFileStmt *sql.Stmt
240− getFileByPathAndSessionStmt *sql.Stmt
241− getMessageStmt *sql.Stmt
242− getSessionByIDStmt *sql.Stmt
243− listFilesByPathStmt *sql.Stmt
244− listFilesBySessionStmt *sql.Stmt
245− listLatestSessionFilesStmt *sql.Stmt
246− listMessagesBySessionStmt *sql.Stmt
247− listNewFilesStmt *sql.Stmt
248− listSessionsStmt *sql.Stmt
249− updateMessageStmt *sql.Stmt
250− updateSessionStmt *sql.Stmt
237+ db DBTX
238+ tx *sql.Tx
239+ createFileStmt *sql.Stmt
240+ createMessageStmt *sql.Stmt
241+ createSessionStmt *sql.Stmt
242+ deleteFileStmt *sql.Stmt
243+ deleteMessageStmt *sql.Stmt
244+ deleteSessionStmt *sql.Stmt
245+ deleteSessionFilesStmt *sql.Stmt
246+ deleteSessionMessagesStmt *sql.Stmt
247+ getFileStmt *sql.Stmt
248+ getFileByPathAndSessionStmt *sql.Stmt
249+ getMessageStmt *sql.Stmt
250+ getSessionByIDStmt *sql.Stmt
251+ listFilesByPathStmt *sql.Stmt
252+ listFilesBySessionStmt *sql.Stmt
253+ listLatestSessionFilesStmt *sql.Stmt
254+ listMessagesBySessionStmt *sql.Stmt
255+ listNewFilesStmt *sql.Stmt
256+ listSessionsStmt *sql.Stmt
257+ updateMessageStmt *sql.Stmt
258+ updateSessionStmt *sql.Stmt
259+ updateSessionTitleAndUsageStmt *sql.Stmt
251260}
252261
253262func (q *Queries) WithTx(tx *sql.Tx) *Queries {
254263 return &Queries{
255− db: tx,
256− tx: tx,
257− createFileStmt: q.createFileStmt,
258− createMessageStmt: q.createMessageStmt,
259− createSessionStmt: q.createSessionStmt,
260− deleteFileStmt: q.deleteFileStmt,
261− deleteMessageStmt: q.deleteMessageStmt,
262− deleteSessionStmt: q.deleteSessionStmt,
263− deleteSessionFilesStmt: q.deleteSessionFilesStmt,
264− deleteSessionMessagesStmt: q.deleteSessionMessagesStmt,
265− getFileStmt: q.getFileStmt,
266− getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,
267− getMessageStmt: q.getMessageStmt,
268− getSessionByIDStmt: q.getSessionByIDStmt,
269− listFilesByPathStmt: q.listFilesByPathStmt,
270− listFilesBySessionStmt: q.listFilesBySessionStmt,
271− listLatestSessionFilesStmt: q.listLatestSessionFilesStmt,
272− listMessagesBySessionStmt: q.listMessagesBySessionStmt,
273− listNewFilesStmt: q.listNewFilesStmt,
274− listSessionsStmt: q.listSessionsStmt,
275− updateMessageStmt: q.updateMessageStmt,
276− updateSessionStmt: q.updateSessionStmt,
264+ db: tx,
265+ tx: tx,
266+ createFileStmt: q.createFileStmt,
267+ createMessageStmt: q.createMessageStmt,
268+ createSessionStmt: q.createSessionStmt,
269+ deleteFileStmt: q.deleteFileStmt,
270+ deleteMessageStmt: q.deleteMessageStmt,
271+ deleteSessionStmt: q.deleteSessionStmt,
272+ deleteSessionFilesStmt: q.deleteSessionFilesStmt,
273+ deleteSessionMessagesStmt: q.deleteSessionMessagesStmt,
274+ getFileStmt: q.getFileStmt,
275+ getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,
276+ getMessageStmt: q.getMessageStmt,
277+ getSessionByIDStmt: q.getSessionByIDStmt,
278+ listFilesByPathStmt: q.listFilesByPathStmt,
279+ listFilesBySessionStmt: q.listFilesBySessionStmt,
280+ listLatestSessionFilesStmt: q.listLatestSessionFilesStmt,
281+ listMessagesBySessionStmt: q.listMessagesBySessionStmt,
282+ listNewFilesStmt: q.listNewFilesStmt,
283+ listSessionsStmt: q.listSessionsStmt,
284+ updateMessageStmt: q.updateMessageStmt,
285+ updateSessionStmt: q.updateSessionStmt,
286+ updateSessionTitleAndUsageStmt: q.updateSessionTitleAndUsageStmt,
277287 }
278288}
internal/db/querier.go+1
@@ -29,6 +29,7 @@ type Querier interface {
2929 ListSessions(ctx context.Context) ([]Session, error)
3030 UpdateMessage(ctx context.Context, arg UpdateMessageParams) error
3131 UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error)
32+ UpdateSessionTitleAndUsage(ctx context.Context, arg UpdateSessionTitleAndUsageParams) error
3233}
3334
3435var _ Querier = (*Queries)(nil)
internal/db/sessions.sql.go+29
@@ -199,3 +199,32 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S
199199 )
200200 return i, err
201201}
202+
203+const updateSessionTitleAndUsage = `-- name: UpdateSessionTitleAndUsage :exec
204+UPDATE sessions
205+SET
206+ title = ?,
207+ prompt_tokens = prompt_tokens + ?,
208+ completion_tokens = completion_tokens + ?,
209+ cost = cost + ?
210+WHERE id = ?
211+`
212+
213+type UpdateSessionTitleAndUsageParams struct {
214+ Title string `json:"title"`
215+ PromptTokens int64 `json:"prompt_tokens"`
216+ CompletionTokens int64 `json:"completion_tokens"`
217+ Cost float64 `json:"cost"`
218+ ID string `json:"id"`
219+}
220+
221+func (q *Queries) UpdateSessionTitleAndUsage(ctx context.Context, arg UpdateSessionTitleAndUsageParams) error {
222+ _, err := q.exec(ctx, q.updateSessionTitleAndUsageStmt, updateSessionTitleAndUsage,
223+ arg.Title,
224+ arg.PromptTokens,
225+ arg.CompletionTokens,
226+ arg.Cost,
227+ arg.ID,
228+ )
229+ return err
230+}
internal/db/sql/sessions.sql+9
@@ -46,6 +46,15 @@ SET
4646WHERE id = ?
4747RETURNING *;
4848
49+-- name: UpdateSessionTitleAndUsage :exec
50+UPDATE sessions
51+SET
52+ title = ?,
53+ prompt_tokens = prompt_tokens + ?,
54+ completion_tokens = completion_tokens + ?,
55+ cost = cost + ?
56+WHERE id = ?;
57+
4958
5059-- name: DeleteSession :exec
5160DELETE FROM sessions
internal/message/content.go+9
@@ -407,6 +407,15 @@ func (m *Message) SetToolResults(tr []ToolResult) {
407407 }
408408}
409409
410+// Clone returns a deep copy of the message with an independent Parts slice.
411+// This prevents race conditions when the message is modified concurrently.
412+func (m *Message) Clone() Message {
413+ clone := *m
414+ clone.Parts = make([]ContentPart, len(m.Parts))
415+ copy(clone.Parts, m.Parts)
416+ return clone
417+}
418+
410419func (m *Message) AddFinish(reason FinishReason, message, details string) {
411420 // remove any existing finish part
412421 for i, part := range m.Parts {
internal/message/message.go+9−3
@@ -51,7 +51,9 @@ func (s *service) Delete(ctx context.Context, id string) error {
5151 if err != nil {
5252 return err
5353 }
54− s.Publish(pubsub.DeletedEvent, message)
54+ // Clone the message before publishing to avoid race conditions with
55+ // concurrent modifications to the Parts slice.
56+ s.Publish(pubsub.DeletedEvent, message.Clone())
5557 return nil
5658}
5759
@@ -85,7 +87,9 @@ func (s *service) Create(ctx context.Context, sessionID string, params CreateMes
8587 if err != nil {
8688 return Message{}, err
8789 }
88− s.Publish(pubsub.CreatedEvent, message)
90+ // Clone the message before publishing to avoid race conditions with
91+ // concurrent modifications to the Parts slice.
92+ s.Publish(pubsub.CreatedEvent, message.Clone())
8993 return message, nil
9094}
9195
@@ -124,7 +128,9 @@ func (s *service) Update(ctx context.Context, message Message) error {
124128 return err
125129 }
126130 message.UpdatedAt = time.Now().Unix()
127− s.Publish(pubsub.UpdatedEvent, message)
131+ // Clone the message before publishing to avoid race conditions with
132+ // concurrent modifications to the Parts slice.
133+ s.Publish(pubsub.UpdatedEvent, message.Clone())
128134 return nil
129135}
130136
internal/pubsub/broker.go+3−8
@@ -92,22 +92,17 @@ func (b *Broker[T]) GetSubscriberCount() int {
9292
9393func (b *Broker[T]) Publish(t EventType, payload T) {
9494 b.mu.RLock()
95+ defer b.mu.RUnlock()
96+
9597 select {
9698 case <-b.done:
97− b.mu.RUnlock()
9899 return
99100 default:
100101 }
101102
102− subscribers := make([]chan Event[T], 0, len(b.subs))
103− for sub := range b.subs {
104− subscribers = append(subscribers, sub)
105− }
106− b.mu.RUnlock()
107−
108103 event := Event[T]{Type: t, Payload: payload}
109104
110− for _, sub := range subscribers {
105+ for sub := range b.subs {
111106 select {
112107 case sub <- event:
113108 default:
internal/session/session.go+13
@@ -50,6 +50,7 @@ type Service interface {
5050 Get(ctx context.Context, id string) (Session, error)
5151 List(ctx context.Context) ([]Session, error)
5252 Save(ctx context.Context, session Session) (Session, error)
53+ UpdateTitleAndUsage(ctx context.Context, sessionID, title string, promptTokens, completionTokens int64, cost float64) error
5354 Delete(ctx context.Context, id string) error
5455
5556 // Agent tool session management
@@ -156,6 +157,18 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) {
156157 return session, nil
157158}
158159
160+// UpdateTitleAndUsage updates only the title and usage fields atomically.
161+// This is safer than fetching, modifying, and saving the entire session.
162+func (s *service) UpdateTitleAndUsage(ctx context.Context, sessionID, title string, promptTokens, completionTokens int64, cost float64) error {
163+ return s.q.UpdateSessionTitleAndUsage(ctx, db.UpdateSessionTitleAndUsageParams{
164+ ID: sessionID,
165+ Title: title,
166+ PromptTokens: promptTokens,
167+ CompletionTokens: completionTokens,
168+ Cost: cost,
169+ })
170+}
171+
159172func (s *service) List(ctx context.Context) ([]Session, error) {
160173 dbSessions, err := s.q.ListSessions(ctx)
161174 if err != nil {