@@ -124,6 +124,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
124124 ai.WithTools(a.tools...),
125125 )
126126
127+ sessionLock := sync.Mutex{}
127128 currentSession, err := a.sessions.Get(ctx, call.SessionID)
128129 if err != nil {
129130 return nil, fmt.Errorf("failed to get session: %w", err)
@@ -138,7 +139,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
138139 // Generate title if first message
139140 if len(msgs) == 0 {
140141 wg.Go(func() {
141− a.generateTitle(ctx, currentSession, call.Prompt)
142+ sessionLock.Lock()
143+ a.generateTitle(ctx, ¤tSession, call.Prompt)
144+ sessionLock.Unlock()
142145 })
143146 }
144147
@@ -304,8 +307,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
304307 case ai.FinishReasonToolCalls:
305308 finishReason = message.FinishReasonToolUse
306309 }
310+ slog.Info("OnStepFinish", "reason", stepResult.FinishReason)
307311 currentAssistant.AddFinish(finishReason, "", "")
308312 a.updateSessionUsage(a.largeModel, ¤tSession, stepResult.Usage)
313+ sessionLock.Lock()
314+ _, sessionErr := a.sessions.Save(genCtx, currentSession)
315+ sessionLock.Unlock()
316+ if sessionErr != nil {
317+ return sessionErr
318+ }
309319 return a.messages.Update(genCtx, *currentAssistant)
310320 },
311321 })
@@ -525,7 +535,7 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
525535 return msgs, nil
526536}
527537
528−func (a *sessionAgent) generateTitle(ctx context.Context, session session.Session, prompt string) {
538+func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Session, prompt string) {
529539 if prompt == "" {
530540 return
531541 }
@@ -559,8 +569,8 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session session.Sessio
559569 }
560570
561571 session.Title = title
562− a.updateSessionUsage(a.smallModel, &session, resp.TotalUsage)
563− _, saveErr := a.sessions.Save(ctx, session)
572+ a.updateSessionUsage(a.smallModel, session, resp.TotalUsage)
573+ _, saveErr := a.sessions.Save(ctx, *session)
564574 if saveErr != nil {
565575 slog.Error("failed to save session title & usage", "error", saveErr)
566576 return