diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index 30c3392939b63aa441551fd0437a5c69789d7b58..16c8e5955d9738283b11543755e82cb0e0584df4 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -371,9 +371,14 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
 			return a.messages.Update(genCtx, *currentAssistant)
 		},
 		StopWhen: []fantasy.StopCondition{
-			func(_ []fantasy.StepResult) bool {
+			func(steps []fantasy.StepResult) bool {
+				if len(steps) == 0 {
+					return false
+				}
+				lastStep := steps[len(steps)-1]
+				usage := lastStep.Usage
+				tokens := usage.InputTokens + usage.OutputTokens + usage.CacheCreationTokens + usage.CacheReadTokens
 				cw := int64(a.largeModel.CatwalkCfg.ContextWindow)
-				tokens := currentSession.CompletionTokens + currentSession.PromptTokens
 				remaining := cw - tokens
 				var threshold int64
 				if cw > 200_000 {