git.iko.soy/crush
on main help sign in
fix: add title lock
5401adfkujtimiihoxha
internal/agent/agent.go+14−4
compare with main
unifiedsplit
@@ -124,6 +124,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
124 ai.WithTools(a.tools...),
124 ai.WithTools(a.tools...),
125 )
125 )
126
126
127+ sessionLock := sync.Mutex{}
127 currentSession, err := a.sessions.Get(ctx, call.SessionID)
128 currentSession, err := a.sessions.Get(ctx, call.SessionID)
128 if err != nil {
129 if err != nil {
129 return nil, fmt.Errorf("failed to get session: %w", err)
130 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
138 // Generate title if first message
139 // Generate title if first message
139 if len(msgs) == 0 {
140 if len(msgs) == 0 {
140 wg.Go(func() {
141 wg.Go(func() {
142+ sessionLock.Lock()
141− a.generateTitle(ctx, currentSession, call.Prompt)
143+ a.generateTitle(ctx, &currentSession, call.Prompt)
144+ sessionLock.Unlock()
142 })
145 })
143 }
146 }
144
147
@@ -304,8 +307,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
304 case ai.FinishReasonToolCalls:
307 case ai.FinishReasonToolCalls:
305 finishReason = message.FinishReasonToolUse
308 finishReason = message.FinishReasonToolUse
306 }
309 }
310+ slog.Info("OnStepFinish", "reason", stepResult.FinishReason)
307 currentAssistant.AddFinish(finishReason, "", "")
311 currentAssistant.AddFinish(finishReason, "", "")
308 a.updateSessionUsage(a.largeModel, &currentSession, stepResult.Usage)
312 a.updateSessionUsage(a.largeModel, &currentSession, stepResult.Usage)
313+ sessionLock.Lock()
314+ _, sessionErr := a.sessions.Save(genCtx, currentSession)
315+ sessionLock.Unlock()
316+ if sessionErr != nil {
317+ return sessionErr
318+ }
309 return a.messages.Update(genCtx, *currentAssistant)
319 return a.messages.Update(genCtx, *currentAssistant)
310 },
320 },
311 })
321 })
@@ -525,7 +535,7 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
525 return msgs, nil
535 return msgs, nil
526}
536}
527
537
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) {
529 if prompt == "" {
539 if prompt == "" {
530 return
540 return
531 }
541 }
@@ -559,8 +569,8 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session session.Sessio
559 }
569 }
560
570
561 session.Title = title
571 session.Title = title
562− a.updateSessionUsage(a.smallModel, &session, resp.TotalUsage)
572+ a.updateSessionUsage(a.smallModel, session, resp.TotalUsage)
563− _, saveErr := a.sessions.Save(ctx, session)
573+ _, saveErr := a.sessions.Save(ctx, *session)
564 if saveErr != nil {
574 if saveErr != nil {
565 slog.Error("failed to save session title & usage", "error", saveErr)
575 slog.Error("failed to save session title & usage", "error", saveErr)
566 return
576 return