4f6d0e9 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: race condition (#1649)

4f6d0e9Kujtim Hoxha1 year ago

internal/agent/agent.go+25−8

compare with main

unifiedsplit

@@ -171,10 +171,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

171 var wg sync.WaitGroup

172 // Generate title if first message.

173 if len(msgs) == 0 {

174+ titleCtx := ctx // Copy to avoid race with ctx reassignment below.

174 wg.Go(func() {

175 wg.Go(func() {

175− sessionLock.Lock()

176− a.generateTitle(ctx, &currentSession, call.Prompt)

176+ a.generateTitle(titleCtx, call.SessionID, call.Prompt)

177− sessionLock.Unlock()

178 })

177 })

179 }

178 }

180

179

@@ -723,7 +722,7 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S

723 return msgs, nil

722 return msgs, nil

724}

723}

725

724

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) {

727 if prompt == "" {

726 if prompt == "" {

728 return

727 return

729 }

728 }

@@ -768,8 +767,7 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi

768 return

767 return

769 }

768 }

770

769

771− session.Title = title

770+ // Calculate usage and cost.

772−

773 var openrouterCost *float64

771 var openrouterCost *float64

774 for _, step := range resp.Steps {

772 for _, step := range resp.Steps {

775 stepCost := a.openrouterCost(step.ProviderMetadata)

773 stepCost := a.openrouterCost(step.ProviderMetadata)

@@ -782,8 +780,27 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi

782 }

780 }

783 }

781 }

784

782

785− a.updateSessionUsage(a.smallModel, session, resp.TotalUsage, openrouterCost)

783+ modelConfig := a.smallModel.CatwalkCfg

786− _, saveErr := a.sessions.Save(ctx, *session)

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)

787 if saveErr != nil {

804 if saveErr != nil {

788 slog.Error("failed to save session title & usage", "error", saveErr)

805 slog.Error("failed to save session title & usage", "error", saveErr)

789 return

806 return

internal/db/db.go+54−44

@@ -84,6 +84,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) {

84 if q.updateSessionStmt, err = db.PrepareContext(ctx, updateSession); err != nil {

85 return nil, fmt.Errorf("error preparing query UpdateSession: %w", err)

86 }

87+ if q.updateSessionTitleAndUsageStmt, err = db.PrepareContext(ctx, updateSessionTitleAndUsage); err != nil {

88+ return nil, fmt.Errorf("error preparing query UpdateSessionTitleAndUsage: %w", err)

89+ }

87 return &q, nil

90 return &q, nil

88}

91}

@@ -189,6 +192,11 @@ func (q *Queries) Close() error {

189 err = fmt.Errorf("error closing updateSessionStmt: %w", cerr)

192 err = fmt.Errorf("error closing updateSessionStmt: %w", cerr)

190 }

193 }

191 }

194 }

195+ if q.updateSessionTitleAndUsageStmt != nil {

196+ if cerr := q.updateSessionTitleAndUsageStmt.Close(); cerr != nil {

197+ err = fmt.Errorf("error closing updateSessionTitleAndUsageStmt: %w", cerr)

198+ }

199+ }

192 return err

200 return err

193}

201}

194

202

@@ -226,53 +234,55 @@ func (q *Queries) queryRow(ctx context.Context, stmt *sql.Stmt, query string, ar

226}

234}

227

235

228type Queries struct {

236type Queries struct {

229− db DBTX

237+ db DBTX

230− tx *sql.Tx

238+ tx *sql.Tx

231− createFileStmt *sql.Stmt

239+ createFileStmt *sql.Stmt

232− createMessageStmt *sql.Stmt

240+ createMessageStmt *sql.Stmt

233− createSessionStmt *sql.Stmt

241+ createSessionStmt *sql.Stmt

234− deleteFileStmt *sql.Stmt

242+ deleteFileStmt *sql.Stmt

235− deleteMessageStmt *sql.Stmt

243+ deleteMessageStmt *sql.Stmt

236− deleteSessionStmt *sql.Stmt

244+ deleteSessionStmt *sql.Stmt

237− deleteSessionFilesStmt *sql.Stmt

245+ deleteSessionFilesStmt *sql.Stmt

238− deleteSessionMessagesStmt *sql.Stmt

246+ deleteSessionMessagesStmt *sql.Stmt

239− getFileStmt *sql.Stmt

247+ getFileStmt *sql.Stmt

240− getFileByPathAndSessionStmt *sql.Stmt

248+ getFileByPathAndSessionStmt *sql.Stmt

241− getMessageStmt *sql.Stmt

249+ getMessageStmt *sql.Stmt

242− getSessionByIDStmt *sql.Stmt

250+ getSessionByIDStmt *sql.Stmt

243− listFilesByPathStmt *sql.Stmt

251+ listFilesByPathStmt *sql.Stmt

244− listFilesBySessionStmt *sql.Stmt

252+ listFilesBySessionStmt *sql.Stmt

245− listLatestSessionFilesStmt *sql.Stmt

253+ listLatestSessionFilesStmt *sql.Stmt

246− listMessagesBySessionStmt *sql.Stmt

254+ listMessagesBySessionStmt *sql.Stmt

247− listNewFilesStmt *sql.Stmt

255+ listNewFilesStmt *sql.Stmt

248− listSessionsStmt *sql.Stmt

256+ listSessionsStmt *sql.Stmt

249− updateMessageStmt *sql.Stmt

257+ updateMessageStmt *sql.Stmt

250− updateSessionStmt *sql.Stmt

258+ updateSessionStmt *sql.Stmt

259+ updateSessionTitleAndUsageStmt *sql.Stmt

251}

260}

252

261

253func (q *Queries) WithTx(tx *sql.Tx) *Queries {

262func (q *Queries) WithTx(tx *sql.Tx) *Queries {

254 return &Queries{

263 return &Queries{

255− db: tx,

264+ db: tx,

256− tx: tx,

265+ tx: tx,

257− createFileStmt: q.createFileStmt,

266+ createFileStmt: q.createFileStmt,

258− createMessageStmt: q.createMessageStmt,

267+ createMessageStmt: q.createMessageStmt,

259− createSessionStmt: q.createSessionStmt,

268+ createSessionStmt: q.createSessionStmt,

260− deleteFileStmt: q.deleteFileStmt,

269+ deleteFileStmt: q.deleteFileStmt,

261− deleteMessageStmt: q.deleteMessageStmt,

270+ deleteMessageStmt: q.deleteMessageStmt,

262− deleteSessionStmt: q.deleteSessionStmt,

271+ deleteSessionStmt: q.deleteSessionStmt,

263− deleteSessionFilesStmt: q.deleteSessionFilesStmt,

272+ deleteSessionFilesStmt: q.deleteSessionFilesStmt,

264− deleteSessionMessagesStmt: q.deleteSessionMessagesStmt,

273+ deleteSessionMessagesStmt: q.deleteSessionMessagesStmt,

265− getFileStmt: q.getFileStmt,

274+ getFileStmt: q.getFileStmt,

266− getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,

275+ getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,

267− getMessageStmt: q.getMessageStmt,

276+ getMessageStmt: q.getMessageStmt,

268− getSessionByIDStmt: q.getSessionByIDStmt,

277+ getSessionByIDStmt: q.getSessionByIDStmt,

269− listFilesByPathStmt: q.listFilesByPathStmt,

278+ listFilesByPathStmt: q.listFilesByPathStmt,

270− listFilesBySessionStmt: q.listFilesBySessionStmt,

279+ listFilesBySessionStmt: q.listFilesBySessionStmt,

271− listLatestSessionFilesStmt: q.listLatestSessionFilesStmt,

280+ listLatestSessionFilesStmt: q.listLatestSessionFilesStmt,

272− listMessagesBySessionStmt: q.listMessagesBySessionStmt,

281+ listMessagesBySessionStmt: q.listMessagesBySessionStmt,

273− listNewFilesStmt: q.listNewFilesStmt,

282+ listNewFilesStmt: q.listNewFilesStmt,

274− listSessionsStmt: q.listSessionsStmt,

283+ listSessionsStmt: q.listSessionsStmt,

275− updateMessageStmt: q.updateMessageStmt,

284+ updateMessageStmt: q.updateMessageStmt,

276− updateSessionStmt: q.updateSessionStmt,

285+ updateSessionStmt: q.updateSessionStmt,

286+ updateSessionTitleAndUsageStmt: q.updateSessionTitleAndUsageStmt,

277 }

287 }

278}

288}

internal/db/querier.go+1

@@ -29,6 +29,7 @@ type Querier interface {

29 ListSessions(ctx context.Context) ([]Session, error)

30 UpdateMessage(ctx context.Context, arg UpdateMessageParams) error

31 UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error)

32+ UpdateSessionTitleAndUsage(ctx context.Context, arg UpdateSessionTitleAndUsageParams) error

32}

33}

34var _ Querier = (*Queries)(nil)

35var _ Querier = (*Queries)(nil)

internal/db/sessions.sql.go+29

@@ -199,3 +199,32 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S

199 )

200 return i, err

201}

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

46WHERE id = ?

47RETURNING *;

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+

50-- name: DeleteSession :exec

59-- name: DeleteSession :exec

51DELETE FROM sessions

60DELETE FROM sessions

internal/message/content.go+9

@@ -407,6 +407,15 @@ func (m *Message) SetToolResults(tr []ToolResult) {

407 }

408}

409

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+

410func (m *Message) AddFinish(reason FinishReason, message, details string) {

419func (m *Message) AddFinish(reason FinishReason, message, details string) {

411 // remove any existing finish part

420 // remove any existing finish part

412 for i, part := range m.Parts {

421 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 {

51 if err != nil {

52 return err

53 }

54+ // Clone the message before publishing to avoid race conditions with

55+ // concurrent modifications to the Parts slice.

54− s.Publish(pubsub.DeletedEvent, message)

56+ s.Publish(pubsub.DeletedEvent, message.Clone())

55 return nil

57 return nil

56}

58}

@@ -85,7 +87,9 @@ func (s *service) Create(ctx context.Context, sessionID string, params CreateMes

85 if err != nil {

87 if err != nil {

86 return Message{}, err

88 return Message{}, err

87 }

89 }

90+ // Clone the message before publishing to avoid race conditions with

91+ // concurrent modifications to the Parts slice.

88− s.Publish(pubsub.CreatedEvent, message)

92+ s.Publish(pubsub.CreatedEvent, message.Clone())

89 return message, nil

93 return message, nil

90}

94}

@@ -124,7 +128,9 @@ func (s *service) Update(ctx context.Context, message Message) error {

124 return err

128 return err

125 }

129 }

126 message.UpdatedAt = time.Now().Unix()

130 message.UpdatedAt = time.Now().Unix()

131+ // Clone the message before publishing to avoid race conditions with

132+ // concurrent modifications to the Parts slice.

127− s.Publish(pubsub.UpdatedEvent, message)

133+ s.Publish(pubsub.UpdatedEvent, message.Clone())

128 return nil

134 return nil

129}

135}

130

136

internal/pubsub/broker.go+3−8

@@ -92,22 +92,17 @@ func (b *Broker[T]) GetSubscriberCount() int {

93func (b *Broker[T]) Publish(t EventType, payload T) {

94 b.mu.RLock()

95+ defer b.mu.RUnlock()

96+

95 select {

97 select {

96 case <-b.done:

98 case <-b.done:

97− b.mu.RUnlock()

98 return

99 return

99 default:

100 default:

100 }

101 }

101

102

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−

108 event := Event[T]{Type: t, Payload: payload}

103 event := Event[T]{Type: t, Payload: payload}

109

104

110− for _, sub := range subscribers {

105+ for sub := range b.subs {

111 select {

106 select {

112 case sub <- event:

107 case sub <- event:

113 default:

108 default:

internal/session/session.go+13

@@ -50,6 +50,7 @@ type Service interface {

50 Get(ctx context.Context, id string) (Session, error)

51 List(ctx context.Context) ([]Session, error)

52 Save(ctx context.Context, session Session) (Session, error)

53+ UpdateTitleAndUsage(ctx context.Context, sessionID, title string, promptTokens, completionTokens int64, cost float64) error

53 Delete(ctx context.Context, id string) error

54 Delete(ctx context.Context, id string) error

55 // Agent tool session management

56 // Agent tool session management

@@ -156,6 +157,18 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) {

156 return session, nil

157 return session, nil

157}

158}

158

159

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+

159func (s *service) List(ctx context.Context) ([]Session, error) {

172func (s *service) List(ctx context.Context) ([]Session, error) {

160 dbSessions, err := s.q.ListSessions(ctx)

173 dbSessions, err := s.q.ListSessions(ctx)

161 if err != nil {

174 if err != nil {