From 5ddf9d1068cc01d2cea20500cecd6fabf9370fb4 Mon Sep 17 00:00:00 2001
From: Kujtim Hoxha <kujtimii.h@gmail.com>
Date: Tue, 27 Jan 2026 14:11:37 +0100
Subject: [PATCH] fix: token calculation (#2004)

---
 internal/agent/agent.go | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index d5b0ed5f1cf4333922ef65c6e0a0fafcf92710a2..815ba2fa8f3c78db8de593849a83ed161e1ee008 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -372,20 +372,18 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
 			}
 			currentAssistant.AddFinish(finishReason, "", "")
 			sessionLock.Lock()
-			updatedSession, getSessionErr := a.sessions.Get(genCtx, call.SessionID)
+			defer sessionLock.Unlock()
+
+			updatedSession, getSessionErr := a.sessions.Get(ctx, call.SessionID)
 			if getSessionErr != nil {
-				sessionLock.Unlock()
 				return getSessionErr
 			}
 			a.updateSessionUsage(largeModel, &updatedSession, stepResult.Usage, a.openrouterCost(stepResult.ProviderMetadata))
-			_, sessionErr := a.sessions.Save(genCtx, updatedSession)
-			if sessionErr == nil {
-				currentSession = updatedSession
-			}
-			sessionLock.Unlock()
+			_, sessionErr := a.sessions.Save(ctx, updatedSession)
 			if sessionErr != nil {
 				return sessionErr
 			}
+			currentSession = updatedSession
 			return a.messages.Update(genCtx, *currentAssistant)
 		},
 		StopWhen: []fantasy.StopCondition{
@@ -898,7 +896,7 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
 	}
 
 	session.CompletionTokens = usage.OutputTokens
-	session.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
+	session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
 }
 
 func (a *sessionAgent) Cancel(sessionID string) {