1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
|
package main
// stats.go implements:
// - per-session cost/usage counters (audio in/out seconds, video frames,
// server usageMetadata true-ups) with USD cost from the active model rates.
// - an SSE hub that streams stats + audio on/off to the capture page, plus a
// short-lived token registry to correlate the separate SSE HTTP request with
// its capture WebSocket session.
//
// The server-side silence gate lives in gate.go.
import (
"encoding/json"
"math"
"sync"
"time"
)
// ---- per-session stats ----
// sessionStats accumulates cost-relevant counters for one capture session.
// Guarded by mu since the WS-receive goroutine (audio in), the live-receive
// goroutine (audio out, usageMetadata) and the SSE ticker all touch it.
type sessionStats struct {
mu sync.Mutex
started time.Time
audioInBytes int // PCM16 forwarded to Gemini (post-gate)
audioOutBytes int // PCM16 received from Gemini
videoFrames int
// server-authoritative token true-ups (from usageMetadata), by modality.
usageAudioIn int
usageImageIn int
usageTextIn int
usageAudioOut int
usageTextOut int
// additional authoritative token kinds the API reports. thoughts are billed
// at the output rate; toolUse (tool results fed back as input) and cached are
// billed as input. usageTotal is Gemini's own total_token_count.
usageThoughts int
usageToolUse int
usageCached int
usageTotal int
gateOn bool
}
func newSessionStats() *sessionStats { return &sessionStats{started: time.Now()} }
func (s *sessionStats) addAudioIn(nbytes int) {
s.mu.Lock()
s.audioInBytes += nbytes
s.mu.Unlock()
}
func (s *sessionStats) addAudioOut(nbytes int) {
s.mu.Lock()
s.audioOutBytes += nbytes
s.mu.Unlock()
}
func (s *sessionStats) addVideoFrame() {
s.mu.Lock()
s.videoFrames++
s.mu.Unlock()
}
func (s *sessionStats) setGate(on bool) {
s.mu.Lock()
s.gateOn = on
s.mu.Unlock()
}
// applyUsage folds a usageMetadata true-up into the per-modality usage counters.
// These are cumulative session totals from Gemini, so we overwrite (max) rather
// than add.
func (s *sessionStats) applyUsage(u *usageMetadata) {
if u == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
for _, d := range u.PromptTokensDetails {
switch d.Modality {
case "AUDIO":
s.usageAudioIn = maxInt(s.usageAudioIn, d.TokenCount)
case "IMAGE", "VIDEO":
s.usageImageIn = maxInt(s.usageImageIn, d.TokenCount)
case "TEXT":
s.usageTextIn = maxInt(s.usageTextIn, d.TokenCount)
}
}
for _, d := range u.ResponseTokensDetails {
switch d.Modality {
case "AUDIO":
s.usageAudioOut = maxInt(s.usageAudioOut, d.TokenCount)
case "TEXT":
s.usageTextOut = maxInt(s.usageTextOut, d.TokenCount)
}
}
// Extra token kinds. These arrive as their own cumulative counts (not always
// in the modality detail lists), so track them separately and price them:
// thoughts at the text-output rate, tool-use + cached at the text-input rate.
s.usageThoughts = maxInt(s.usageThoughts, u.ThoughtsTokenCount)
s.usageToolUse = maxInt(s.usageToolUse, u.ToolUsePromptTokenCount)
s.usageCached = maxInt(s.usageCached, u.CachedContentTokenCount)
s.usageTotal = maxInt(s.usageTotal, u.TotalTokenCount)
}
func maxInt(a, b int) int {
if a > b {
return a
}
return b
}
// sessionUsage is the resolved per-session token + cost accounting, ready to
// display or persist. Token counts are the authoritative usageMetadata values
// when present (audio falls back to a byte-based estimate until the first usage
// frame arrives). Cost is derived from the active model rates at capture time.
type sessionUsage struct {
// estimated audio seconds (from bytes), for the UI only.
AudioInSecs float64
AudioOutSecs float64
// token counts by kind.
AudioInTok int
AudioOutTok int
ImageInTok int
TextInTok int
TextOutTok int
ThoughtsTok int
ToolUseTok int
CachedTok int
TotalTok int
VideoFrames int
// per-kind USD cost and the summed total.
CostAudioIn float64
CostAudioOut float64
CostImageIn float64
CostTextIn float64
CostTextOut float64
CostThoughts float64
CostToolUse float64
CostCached float64
CostTotal float64
}
// computeUsageLocked resolves tokens and cost from the current counters. The
// caller must hold s.mu. Thoughts bill at the text-output rate ("output price
// includes thinking tokens"); tool-use and cached bill at the text-input rate
// (cached is priced at the full input rate rather than a discount, so we never
// undercount — caching is expected to be zero for this Live model).
func (s *sessionStats) computeUsageLocked() sessionUsage {
m := activeModel
// audio tokens: estimate from bytes; true-up if usage present.
audioInEstSecs := audioSecondsFromBytes(s.audioInBytes, m.AudioInSampleRate)
audioOutEstSecs := audioSecondsFromBytes(s.audioOutBytes, m.AudioOutSampleRate)
audioInTok := m.audioTokens(audioInEstSecs)
audioOutTok := m.audioTokens(audioOutEstSecs)
if s.usageAudioIn > 0 {
audioInTok = float64(s.usageAudioIn)
}
if s.usageAudioOut > 0 {
audioOutTok = float64(s.usageAudioOut)
}
imageInTok := float64(s.usageImageIn)
textInTok := float64(s.usageTextIn)
textOutTok := float64(s.usageTextOut)
thoughtsTok := float64(s.usageThoughts)
toolUseTok := float64(s.usageToolUse)
cachedTok := float64(s.usageCached)
u := sessionUsage{
AudioInSecs: audioInEstSecs,
AudioOutSecs: audioOutEstSecs,
AudioInTok: int(audioInTok),
AudioOutTok: int(audioOutTok),
ImageInTok: int(imageInTok),
TextInTok: int(textInTok),
TextOutTok: int(textOutTok),
ThoughtsTok: int(thoughtsTok),
ToolUseTok: int(toolUseTok),
CachedTok: int(cachedTok),
TotalTok: s.usageTotal,
VideoFrames: s.videoFrames,
}
u.CostAudioIn = costUSD(audioInTok, m.PriceAudioInPerMTok)
u.CostAudioOut = costUSD(audioOutTok, m.PriceAudioOutPerMTok)
u.CostImageIn = costUSD(imageInTok, m.PriceImageInPerMTok)
u.CostTextIn = costUSD(textInTok, m.PriceTextInPerMTok)
u.CostTextOut = costUSD(textOutTok, m.PriceTextOutPerMTok)
u.CostThoughts = costUSD(thoughtsTok, m.PriceTextOutPerMTok)
u.CostToolUse = costUSD(toolUseTok, m.PriceTextInPerMTok)
u.CostCached = costUSD(cachedTok, m.PriceTextInPerMTok)
u.CostTotal = u.CostAudioIn + u.CostAudioOut + u.CostImageIn + u.CostTextIn +
u.CostTextOut + u.CostThoughts + u.CostToolUse + u.CostCached
return u
}
// finalUsage returns the resolved usage/cost for persistence (thread-safe).
func (s *sessionStats) finalUsage() sessionUsage {
s.mu.Lock()
defer s.mu.Unlock()
return s.computeUsageLocked()
}
// snapshot builds the outward stats payload for the live UI. Token/cost values
// come from the shared computeUsageLocked so what the user sees matches what we
// persist.
func (s *sessionStats) snapshot() map[string]any {
s.mu.Lock()
defer s.mu.Unlock()
u := s.computeUsageLocked()
return map[string]any{
"elapsed_secs": int(time.Since(s.started).Seconds()),
"audio_in_secs": round2(u.AudioInSecs),
"audio_out_secs": round2(u.AudioOutSecs),
"audio_in_tokens": u.AudioInTok,
"audio_out_tokens": u.AudioOutTok,
"image_in_tokens": u.ImageInTok,
"video_frames": u.VideoFrames,
"gate_on": s.gateOn,
"cost_usd": round4(u.CostTotal),
"cost_breakdown": map[string]any{
"audio_in": round4(u.CostAudioIn),
"audio_out": round4(u.CostAudioOut),
"image_in": round4(u.CostImageIn),
"text_in": round4(u.CostTextIn),
"text_out": round4(u.CostTextOut),
"thoughts": round4(u.CostThoughts),
"tool_use": round4(u.CostToolUse),
},
}
}
func round2(f float64) float64 { return math.Round(f*100) / 100 }
func round4(f float64) float64 { return math.Round(f*10000) / 10000 }
// ---- SSE hub + token registry ----
// statsHub fans out SSE events (stats snapshots, audio_state) to the one SSE
// client for a session. Events are JSON-encoded envelopes.
type statsHub struct {
ch chan []byte
closed chan struct{}
once sync.Once
}
func newStatsHub() *statsHub {
return &statsHub{ch: make(chan []byte, 16), closed: make(chan struct{})}
}
// send queues an SSE event; drops if the buffer is full (stats are periodic, a
// dropped frame is harmless).
func (h *statsHub) send(event string, payload map[string]any) {
b, err := json.Marshal(payload)
if err != nil {
return
}
msg := append([]byte("event: "+event+"\ndata: "), b...)
msg = append(msg, '\n', '\n')
select {
case h.ch <- msg:
case <-h.closed:
default:
}
}
func (h *statsHub) close() {
h.once.Do(func() { close(h.closed) })
}
// statsRegistry correlates a capture session's stats hub with the separate SSE
// HTTP request, via a short-lived random token minted at WS-open time.
type statsRegistry struct {
mu sync.Mutex
hubs map[string]*statsHub
}
func newStatsRegistry() *statsRegistry { return &statsRegistry{hubs: map[string]*statsHub{}} }
func (r *statsRegistry) register(token string, h *statsHub) {
r.mu.Lock()
r.hubs[token] = h
r.mu.Unlock()
}
func (r *statsRegistry) get(token string) (*statsHub, bool) {
r.mu.Lock()
h, ok := r.hubs[token]
r.mu.Unlock()
return h, ok
}
func (r *statsRegistry) unregister(token string) {
r.mu.Lock()
if h, ok := r.hubs[token]; ok {
h.close()
delete(r.hubs, token)
}
r.mu.Unlock()
}
|