More correct price calculation

dgageot · dgageot · commit 82bfbe39cb1b · 2025-09-02T11:41:22.000+02:00
Signed-off-by: David Gageot &lt;david.gageot@docker.com&gt;
diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
@@ -318,12 +318,15 @@ func (r *Runtime) handleStream(stream chat.MessageStream, a *agent.Agent, sess *
 		}
 
 		if response.Usage != nil {
-			sess.InputTokens = response.Usage.InputTokens + response.Usage.CachedInputTokens
-			sess.OutputTokens = response.Usage.OutputTokens + response.Usage.CachedOutputTokens
 			if m != nil {
-				cost := m.Cost.Input/1e6*float64(sess.InputTokens) + m.Cost.Output/1e6*float64(sess.OutputTokens) + m.Cost.CacheRead/1e6*float64(response.Usage.CachedInputTokens) + m.Cost.CacheWrite/1e6*float64(response.Usage.CachedOutputTokens)
-				sess.Cost += cost
+				sess.Cost += (float64(response.Usage.InputTokens)*m.Cost.Input +
+					float64(response.Usage.OutputTokens)*m.Cost.Output +
+					float64(response.Usage.CachedInputTokens)*m.Cost.CacheRead +
+					float64(response.Usage.CachedOutputTokens)*m.Cost.CacheWrite) / 1e6
 			}
+
+			sess.InputTokens = response.Usage.InputTokens + response.Usage.CachedInputTokens
+			sess.OutputTokens = response.Usage.OutputTokens + response.Usage.CachedOutputTokens
 		}
 
 		if len(response.Choices) == 0 {

Original file line number	Diff line number	Diff line change
`@@ -318,12 +318,15 @@ func (r Runtime) handleStream(stream chat.MessageStream, a agent.Agent, sess *`
`318`	`318`	`}`
`319`	`319`
`320`	`320`	`if response.Usage != nil {`
`321`		`- sess.InputTokens = response.Usage.InputTokens + response.Usage.CachedInputTokens`
`322`		`- sess.OutputTokens = response.Usage.OutputTokens + response.Usage.CachedOutputTokens`
`323`	`321`	`if m != nil {`
`324`		`- cost := m.Cost.Input/1e6float64(sess.InputTokens) + m.Cost.Output/1e6float64(sess.OutputTokens) + m.Cost.CacheRead/1e6float64(response.Usage.CachedInputTokens) + m.Cost.CacheWrite/1e6float64(response.Usage.CachedOutputTokens)`
`325`		`- sess.Cost += cost`
	`322`	`+ sess.Cost += (float64(response.Usage.InputTokens)*m.Cost.Input +`
	`323`	`+ float64(response.Usage.OutputTokens)*m.Cost.Output +`
	`324`	`+ float64(response.Usage.CachedInputTokens)*m.Cost.CacheRead +`
	`325`	`+ float64(response.Usage.CachedOutputTokens)*m.Cost.CacheWrite) / 1e6`
`326`	`326`	`}`
	`327`	`+`
	`328`	`+ sess.InputTokens = response.Usage.InputTokens + response.Usage.CachedInputTokens`
	`329`	`+ sess.OutputTokens = response.Usage.OutputTokens + response.Usage.CachedOutputTokens`
`327`	`330`	`}`
`328`	`331`
`329`	`332`	`if len(response.Choices) == 0 {`