git.iko.soy/crush
on main help sign in
fix: token calculation (#2004)
5ddf9d1Kujtim Hoxha
internal/agent/agent.go+6−8
compare with main
unifiedsplit
@@ -372,20 +372,18 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
372 }
372 }
373 currentAssistant.AddFinish(finishReason, "", "")
373 currentAssistant.AddFinish(finishReason, "", "")
374 sessionLock.Lock()
374 sessionLock.Lock()
375+ defer sessionLock.Unlock()
376+
375− updatedSession, getSessionErr := a.sessions.Get(genCtx, call.SessionID)
377+ updatedSession, getSessionErr := a.sessions.Get(ctx, call.SessionID)
376 if getSessionErr != nil {
378 if getSessionErr != nil {
377− sessionLock.Unlock()
378 return getSessionErr
379 return getSessionErr
379 }
380 }
380 a.updateSessionUsage(largeModel, &updatedSession, stepResult.Usage, a.openrouterCost(stepResult.ProviderMetadata))
381 a.updateSessionUsage(largeModel, &updatedSession, stepResult.Usage, a.openrouterCost(stepResult.ProviderMetadata))
381− _, sessionErr := a.sessions.Save(genCtx, updatedSession)
382+ _, sessionErr := a.sessions.Save(ctx, updatedSession)
382− if sessionErr == nil {
383− currentSession = updatedSession
384− }
385− sessionLock.Unlock()
386 if sessionErr != nil {
383 if sessionErr != nil {
387 return sessionErr
384 return sessionErr
388 }
385 }
386+ currentSession = updatedSession
389 return a.messages.Update(genCtx, *currentAssistant)
387 return a.messages.Update(genCtx, *currentAssistant)
390 },
388 },
391 StopWhen: []fantasy.StopCondition{
389 StopWhen: []fantasy.StopCondition{
@@ -898,7 +896,7 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
898 }
896 }
899
897
900 session.CompletionTokens = usage.OutputTokens
898 session.CompletionTokens = usage.OutputTokens
901− session.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
899+ session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
902}
900}
903
901
904func (a *sessionAgent) Cancel(sessionID string) {
902func (a *sessionAgent) Cancel(sessionID string) {