db73cf5a39
Includes the current working directory changes. Historical commits and release binaries are excluded to limit Gitea storage. Source base: d7515973fbfd6dd2001fa156c483455b8081c7cb
375 lines
13 KiB
Go
375 lines
13 KiB
Go
package forwarder
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
usageFileName = "usage.json"
|
||
usageFileSchemaVersion = 2
|
||
usageRecentEventLimit = 500
|
||
|
||
usageEventKindProvider = "provider_call"
|
||
usageEventKindTurn = "turn_finalized"
|
||
usageTurnStatusDone = "completed"
|
||
)
|
||
|
||
type UsageFileStore struct {
|
||
path string
|
||
}
|
||
|
||
type usageFileDocument struct {
|
||
SchemaVersion int `json:"schema_version"`
|
||
UpdatedAt time.Time `json:"updated_at"`
|
||
Totals usageFileTotals `json:"totals"`
|
||
Daily []usageFileDaily `json:"daily"`
|
||
RecentEvents []usageFileEvent `json:"recent_events"`
|
||
EventIndex map[string]usageFileEvent `json:"event_index,omitempty"`
|
||
// eventIndexPartial 表示索引是从被裁剪的 RecentEvents 重建的（旧格式文件），
|
||
// 不参与序列化，仅用于加载后首批 UpsertEvent 的防双计判断。
|
||
eventIndexPartial bool
|
||
}
|
||
|
||
type usageFileTotals struct {
|
||
ProviderCalls int64 `json:"provider_calls"`
|
||
TurnsTotal int64 `json:"turns_total"`
|
||
ValidTurnsTotal int64 `json:"valid_turns_total"`
|
||
InvalidTurnsTotal int64 `json:"invalid_turns_total"`
|
||
InputTokens int64 `json:"input_tokens"`
|
||
OutputTokens int64 `json:"output_tokens"`
|
||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||
CacheWriteTokens int64 `json:"cache_write_tokens"`
|
||
TotalTokens int64 `json:"total_tokens"`
|
||
}
|
||
|
||
type usageFileDaily struct {
|
||
Date string `json:"date"`
|
||
ProviderCalls int64 `json:"provider_calls"`
|
||
TurnsTotal int64 `json:"turns_total"`
|
||
ValidTurnsTotal int64 `json:"valid_turns_total"`
|
||
InvalidTurnsTotal int64 `json:"invalid_turns_total"`
|
||
InputTokens int64 `json:"input_tokens"`
|
||
OutputTokens int64 `json:"output_tokens"`
|
||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||
CacheWriteTokens int64 `json:"cache_write_tokens"`
|
||
TotalTokens int64 `json:"total_tokens"`
|
||
}
|
||
|
||
type usageFileEvent struct {
|
||
EventID string `json:"event_id"`
|
||
Kind string `json:"kind,omitempty"`
|
||
Status string `json:"status,omitempty"`
|
||
At time.Time `json:"at"`
|
||
InputTokens int64 `json:"input_tokens"`
|
||
OutputTokens int64 `json:"output_tokens"`
|
||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||
CacheWriteTokens int64 `json:"cache_write_tokens"`
|
||
TotalTokens int64 `json:"total_tokens"`
|
||
UsagePresent bool `json:"usage_present"`
|
||
}
|
||
|
||
type usageFileDelta struct {
|
||
providerCalls int64
|
||
turnsTotal int64
|
||
validTurnsTotal int64
|
||
invalidTurnsTotal int64
|
||
inputTokens int64
|
||
outputTokens int64
|
||
cacheReadTokens int64
|
||
cacheWriteTokens int64
|
||
totalTokens int64
|
||
}
|
||
|
||
func NewUsageFileStore(historyRoot string) *UsageFileStore {
|
||
return &UsageFileStore{path: filepath.Join(strings.TrimSpace(historyRoot), usageFileName)}
|
||
}
|
||
|
||
func (store *UsageFileStore) UpsertEvent(event usageFileEvent) error {
|
||
if store == nil || strings.TrimSpace(store.path) == "" {
|
||
return nil
|
||
}
|
||
event.EventID = strings.TrimSpace(event.EventID)
|
||
if event.EventID == "" {
|
||
return nil
|
||
}
|
||
event.Kind = normalizeUsageEventKind(event.Kind)
|
||
event.Status = strings.TrimSpace(event.Status)
|
||
if event.At.IsZero() {
|
||
event.At = time.Now().UTC()
|
||
} else {
|
||
event.At = event.At.UTC()
|
||
}
|
||
event.InputTokens = nonNegativeInt64(event.InputTokens)
|
||
event.OutputTokens = nonNegativeInt64(event.OutputTokens)
|
||
event.CacheReadTokens = nonNegativeInt64(event.CacheReadTokens)
|
||
event.CacheWriteTokens = nonNegativeInt64(event.CacheWriteTokens)
|
||
event.TotalTokens = event.InputTokens + event.OutputTokens + event.CacheReadTokens + event.CacheWriteTokens
|
||
|
||
if err := os.MkdirAll(filepath.Dir(store.path), 0o755); err != nil {
|
||
return fmt.Errorf("create usage directory: %w", err)
|
||
}
|
||
release, err := acquireConversationLock(store.path + ".lock")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer release()
|
||
|
||
doc, err := readUsageFileDocument(store.path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if doc.EventIndex == nil {
|
||
doc.EventIndex = make(map[string]usageFileEvent)
|
||
}
|
||
oldEvent, found := doc.EventIndex[event.EventID]
|
||
if found {
|
||
applyUsageFileDelta(&doc, oldEvent.At, negateUsageFileDelta(usageFileEventDelta(oldEvent)))
|
||
} else if doc.eventIndexPartial && isOlderThanRetainedEvents(doc.RecentEvents, event.At) {
|
||
// 旧格式文件的部分索引中找不到该事件，且事件时间早于保留窗口内最早的事件，
|
||
// 说明它是历史事件（已计入 totals），拒绝重复累加。
|
||
doc.EventIndex[event.EventID] = event
|
||
doc.eventIndexPartial = false
|
||
doc.SchemaVersion = usageFileSchemaVersion
|
||
doc.UpdatedAt = time.Now().UTC()
|
||
return writeJSONFileAtomic(store.path, doc)
|
||
}
|
||
applyUsageFileDelta(&doc, event.At, usageFileEventDelta(event))
|
||
doc.RecentEvents = upsertRecentUsageEvent(doc.RecentEvents, event)
|
||
doc.RecentEvents = trimRecentUsageEvents(doc.RecentEvents, usageRecentEventLimit)
|
||
// EventIndex 必须持久化全量事件；若按 RecentEvents 重建，超出 500 条的旧事件
|
||
// 再次 upsert 时找不到旧值，旧 delta 不减、新 delta 全加，导致 totals 双计。
|
||
doc.EventIndex[event.EventID] = event
|
||
doc.SchemaVersion = usageFileSchemaVersion
|
||
doc.UpdatedAt = time.Now().UTC()
|
||
return writeJSONFileAtomic(store.path, doc)
|
||
}
|
||
|
||
func (store *UsageFileStore) LookupEvent(needle string) (usageFileEvent, bool, error) {
|
||
if store == nil || strings.TrimSpace(store.path) == "" {
|
||
return usageFileEvent{}, false, nil
|
||
}
|
||
doc, err := readUsageFileDocument(store.path)
|
||
if err != nil {
|
||
return usageFileEvent{}, false, err
|
||
}
|
||
trimmed := strings.TrimSpace(needle)
|
||
if trimmed == "" {
|
||
return usageFileEvent{}, false, nil
|
||
}
|
||
var aggregate usageFileEvent
|
||
found := false
|
||
events := doc.EventIndex
|
||
if len(events) == 0 {
|
||
events = make(map[string]usageFileEvent, len(doc.RecentEvents))
|
||
for _, event := range doc.RecentEvents {
|
||
if eventID := strings.TrimSpace(event.EventID); eventID != "" {
|
||
events[eventID] = event
|
||
}
|
||
}
|
||
}
|
||
for _, event := range events {
|
||
eventID := strings.TrimSpace(event.EventID)
|
||
if eventID != trimmed && !strings.HasPrefix(eventID, trimmed+"::") {
|
||
continue
|
||
}
|
||
if !found {
|
||
aggregate = usageFileEvent{EventID: trimmed, At: event.At}
|
||
found = true
|
||
}
|
||
if event.At.After(aggregate.At) {
|
||
aggregate.At = event.At
|
||
}
|
||
aggregate.InputTokens += nonNegativeInt64(event.InputTokens)
|
||
aggregate.OutputTokens += nonNegativeInt64(event.OutputTokens)
|
||
aggregate.CacheReadTokens += nonNegativeInt64(event.CacheReadTokens)
|
||
aggregate.CacheWriteTokens += nonNegativeInt64(event.CacheWriteTokens)
|
||
aggregate.TotalTokens += nonNegativeInt64(event.TotalTokens)
|
||
aggregate.UsagePresent = aggregate.UsagePresent || event.UsagePresent
|
||
}
|
||
if found {
|
||
return aggregate, true, nil
|
||
}
|
||
return usageFileEvent{}, false, nil
|
||
}
|
||
|
||
func readUsageFileDocument(path string) (usageFileDocument, error) {
|
||
body, err := os.ReadFile(path)
|
||
if err != nil {
|
||
if errors.Is(err, os.ErrNotExist) {
|
||
return usageFileDocument{
|
||
SchemaVersion: usageFileSchemaVersion,
|
||
Daily: make([]usageFileDaily, 0),
|
||
RecentEvents: make([]usageFileEvent, 0),
|
||
EventIndex: make(map[string]usageFileEvent),
|
||
}, nil
|
||
}
|
||
return usageFileDocument{}, fmt.Errorf("read usage file: %w", err)
|
||
}
|
||
var doc usageFileDocument
|
||
if err := json.Unmarshal(body, &doc); err != nil {
|
||
return usageFileDocument{}, fmt.Errorf("decode usage file: %w", err)
|
||
}
|
||
if doc.SchemaVersion == 0 {
|
||
doc.SchemaVersion = 1
|
||
}
|
||
doc.RecentEvents = trimRecentUsageEvents(doc.RecentEvents, usageRecentEventLimit)
|
||
if len(doc.EventIndex) == 0 {
|
||
// 旧格式文件没有全量索引，只能从被裁剪的 RecentEvents 重建。
|
||
// 标记为部分索引，UpsertEvent 据此防止历史事件被重复累加。
|
||
doc.EventIndex = buildUsageEventIndex(doc.RecentEvents)
|
||
doc.eventIndexPartial = true
|
||
}
|
||
return doc, nil
|
||
}
|
||
|
||
func upsertRecentUsageEvent(items []usageFileEvent, event usageFileEvent) []usageFileEvent {
|
||
event.EventID = strings.TrimSpace(event.EventID)
|
||
if event.EventID == "" {
|
||
return items
|
||
}
|
||
next := make([]usageFileEvent, 0, len(items)+1)
|
||
next = append(next, event)
|
||
for _, item := range items {
|
||
if strings.TrimSpace(item.EventID) == event.EventID {
|
||
continue
|
||
}
|
||
next = append(next, item)
|
||
}
|
||
return next
|
||
}
|
||
|
||
func trimRecentUsageEvents(items []usageFileEvent, limit int) []usageFileEvent {
|
||
if limit <= 0 || len(items) <= limit {
|
||
return items
|
||
}
|
||
return items[:limit]
|
||
}
|
||
|
||
func buildUsageEventIndex(items []usageFileEvent) map[string]usageFileEvent {
|
||
index := make(map[string]usageFileEvent, len(items))
|
||
for _, event := range items {
|
||
event.EventID = strings.TrimSpace(event.EventID)
|
||
if event.EventID == "" {
|
||
continue
|
||
}
|
||
event.Kind = normalizeUsageEventKind(event.Kind)
|
||
index[event.EventID] = event
|
||
}
|
||
return index
|
||
}
|
||
|
||
// isOlderThanRetainedEvents 判断事件时间是否早于保留窗口内最早的事件。
|
||
// 用于旧格式文件的部分索引场景：早于窗口的历史事件已计入 totals，不得重复累加。
|
||
func isOlderThanRetainedEvents(items []usageFileEvent, at time.Time) bool {
|
||
if len(items) == 0 {
|
||
return false
|
||
}
|
||
oldest := items[0].At
|
||
for _, item := range items[1:] {
|
||
if item.At.Before(oldest) {
|
||
oldest = item.At
|
||
}
|
||
}
|
||
return at.Before(oldest)
|
||
}
|
||
|
||
func normalizeUsageEventKind(kind string) string {
|
||
switch strings.TrimSpace(kind) {
|
||
case usageEventKindTurn:
|
||
return usageEventKindTurn
|
||
default:
|
||
return usageEventKindProvider
|
||
}
|
||
}
|
||
|
||
func usageFileEventDelta(event usageFileEvent) usageFileDelta {
|
||
switch normalizeUsageEventKind(event.Kind) {
|
||
case usageEventKindTurn:
|
||
delta := usageFileDelta{turnsTotal: 1}
|
||
if strings.TrimSpace(event.Status) == usageTurnStatusDone {
|
||
delta.validTurnsTotal = 1
|
||
} else {
|
||
delta.invalidTurnsTotal = 1
|
||
}
|
||
return delta
|
||
default:
|
||
return usageFileDelta{
|
||
providerCalls: 1,
|
||
inputTokens: nonNegativeInt64(event.InputTokens),
|
||
outputTokens: nonNegativeInt64(event.OutputTokens),
|
||
cacheReadTokens: nonNegativeInt64(event.CacheReadTokens),
|
||
cacheWriteTokens: nonNegativeInt64(event.CacheWriteTokens),
|
||
totalTokens: nonNegativeInt64(event.TotalTokens),
|
||
}
|
||
}
|
||
}
|
||
|
||
func negateUsageFileDelta(value usageFileDelta) usageFileDelta {
|
||
return usageFileDelta{
|
||
providerCalls: -value.providerCalls,
|
||
turnsTotal: -value.turnsTotal,
|
||
validTurnsTotal: -value.validTurnsTotal,
|
||
invalidTurnsTotal: -value.invalidTurnsTotal,
|
||
inputTokens: -value.inputTokens,
|
||
outputTokens: -value.outputTokens,
|
||
cacheReadTokens: -value.cacheReadTokens,
|
||
cacheWriteTokens: -value.cacheWriteTokens,
|
||
totalTokens: -value.totalTokens,
|
||
}
|
||
}
|
||
|
||
func applyUsageFileDelta(doc *usageFileDocument, at time.Time, delta usageFileDelta) {
|
||
if doc == nil {
|
||
return
|
||
}
|
||
doc.Totals.ProviderCalls = clampNonNegativeInt64(doc.Totals.ProviderCalls + delta.providerCalls)
|
||
doc.Totals.TurnsTotal = clampNonNegativeInt64(doc.Totals.TurnsTotal + delta.turnsTotal)
|
||
doc.Totals.ValidTurnsTotal = clampNonNegativeInt64(doc.Totals.ValidTurnsTotal + delta.validTurnsTotal)
|
||
doc.Totals.InvalidTurnsTotal = clampNonNegativeInt64(doc.Totals.InvalidTurnsTotal + delta.invalidTurnsTotal)
|
||
doc.Totals.InputTokens = clampNonNegativeInt64(doc.Totals.InputTokens + delta.inputTokens)
|
||
doc.Totals.OutputTokens = clampNonNegativeInt64(doc.Totals.OutputTokens + delta.outputTokens)
|
||
doc.Totals.CacheReadTokens = clampNonNegativeInt64(doc.Totals.CacheReadTokens + delta.cacheReadTokens)
|
||
doc.Totals.CacheWriteTokens = clampNonNegativeInt64(doc.Totals.CacheWriteTokens + delta.cacheWriteTokens)
|
||
doc.Totals.TotalTokens = clampNonNegativeInt64(doc.Totals.TotalTokens + delta.totalTokens)
|
||
|
||
date := at.UTC().Format("2006-01-02")
|
||
for index := range doc.Daily {
|
||
if doc.Daily[index].Date != date {
|
||
continue
|
||
}
|
||
applyUsageDailyDelta(&doc.Daily[index], delta)
|
||
return
|
||
}
|
||
item := usageFileDaily{Date: date}
|
||
applyUsageDailyDelta(&item, delta)
|
||
doc.Daily = append(doc.Daily, item)
|
||
}
|
||
|
||
func applyUsageDailyDelta(item *usageFileDaily, delta usageFileDelta) {
|
||
if item == nil {
|
||
return
|
||
}
|
||
item.ProviderCalls = clampNonNegativeInt64(item.ProviderCalls + delta.providerCalls)
|
||
item.TurnsTotal = clampNonNegativeInt64(item.TurnsTotal + delta.turnsTotal)
|
||
item.ValidTurnsTotal = clampNonNegativeInt64(item.ValidTurnsTotal + delta.validTurnsTotal)
|
||
item.InvalidTurnsTotal = clampNonNegativeInt64(item.InvalidTurnsTotal + delta.invalidTurnsTotal)
|
||
item.InputTokens = clampNonNegativeInt64(item.InputTokens + delta.inputTokens)
|
||
item.OutputTokens = clampNonNegativeInt64(item.OutputTokens + delta.outputTokens)
|
||
item.CacheReadTokens = clampNonNegativeInt64(item.CacheReadTokens + delta.cacheReadTokens)
|
||
item.CacheWriteTokens = clampNonNegativeInt64(item.CacheWriteTokens + delta.cacheWriteTokens)
|
||
item.TotalTokens = clampNonNegativeInt64(item.TotalTokens + delta.totalTokens)
|
||
}
|
||
|
||
func clampNonNegativeInt64(value int64) int64 {
|
||
if value < 0 {
|
||
return 0
|
||
}
|
||
return value
|
||
}
|