claude-code-proxy/proxy/internal/service/storage_analytics.go

package service

import (
	"database/sql"
	"encoding/json"

	"github.com/seifghazi/claude-code-monitor/internal/model"
)

type responseBodySummary struct {
	Usage      *model.AnthropicUsage `json:"usage"`
	StopReason string                `json:"stop_reason"`
}

func decodeStoredResponse(responseJSON sql.NullString) (*model.ResponseLog, bool) {
	if !responseJSON.Valid || responseJSON.String == "" {
		return nil, false
	}

	var resp model.ResponseLog
	if err := json.Unmarshal([]byte(responseJSON.String), &resp); err != nil {
		return nil, false
	}

	return &resp, true
}

func decodeResponseBodySummary(body json.RawMessage) (*responseBodySummary, bool) {
	if len(body) == 0 {
		return nil, false
	}

	var summary responseBodySummary
	if err := json.Unmarshal(body, &summary); err != nil {
		return nil, false
	}

	return &summary, true
}

func totalTokensFromUsage(usage *model.AnthropicUsage) int64 {
	if usage == nil {
		return 0
	}

	return int64(
		usage.InputTokens +
			usage.OutputTokens +
			usage.CacheReadInputTokens +
			usage.CacheCreationInputTokens)
}

func totalTokensFromStoredResponse(responseJSON sql.NullString) int64 {
	input, output, cache, ok := usageCountsFromStoredResponse(responseJSON)
	if !ok {
		return 0
	}

	return input + output + cache
}

func responseTimeFromStoredResponse(responseJSON sql.NullString) int64 {
	resp, ok := decodeStoredResponse(responseJSON)
	if !ok {
		return 0
	}

	return resp.ResponseTime
}

func applyStoredResponseToSummary(summary *model.RequestSummary, responseJSON sql.NullString) {
	resp, ok := decodeStoredResponse(responseJSON)
	if !ok {
		return
	}

	summary.StatusCode = resp.StatusCode
	summary.ResponseTime = resp.ResponseTime

	bodySummary, ok := decodeResponseBodySummary(resp.Body)
	if !ok {
		return
	}

	summary.Usage = bodySummary.Usage
	summary.StopReason = bodySummary.StopReason
}

func usageCountsFromStoredResponse(responseJSON sql.NullString) (input, output, cache int64, ok bool) {
	resp, ok := decodeStoredResponse(responseJSON)
	if !ok {
		return 0, 0, 0, false
	}

	bodySummary, ok := decodeResponseBodySummary(resp.Body)
	if !ok || bodySummary.Usage == nil {
		return 0, 0, 0, false
	}

	return int64(bodySummary.Usage.InputTokens),
		int64(bodySummary.Usage.OutputTokens),
		int64(bodySummary.Usage.CacheCreationInputTokens + bodySummary.Usage.CacheReadInputTokens),
		true
}

func addDailyTokens(dailyMap map[string]*model.DailyTokens, date, modelName string, tokens int64) {
	if daily, ok := dailyMap[date]; ok {
		daily.Tokens += tokens
		daily.Requests++
		daily.Models = addDailyModelStat(daily.Models, modelName, tokens)
		return
	}

	dailyMap[date] = &model.DailyTokens{
		Date:     date,
		Tokens:   tokens,
		Requests: 1,
		Models:   addDailyModelStat(nil, modelName, tokens),
	}
}

func addHourlyTokens(bucketMap map[string]*model.HourlyTokens, bucketKey, bucketLabel, modelName string, tokens int64) {
	if bucket, ok := bucketMap[bucketKey]; ok {
		bucket.Tokens += tokens
		bucket.Requests++
		bucket.Models = addDailyModelStat(bucket.Models, modelName, tokens)
		return
	}

	bucketMap[bucketKey] = &model.HourlyTokens{
		Hour:     0,
		Label:    bucketLabel,
		Tokens:   tokens,
		Requests: 1,
		Models:   addDailyModelStat(nil, modelName, tokens),
	}
}

func addModelTokens(modelMap map[string]*model.ModelTokens, modelName string, tokens int64) {
	if modelStat, ok := modelMap[modelName]; ok {
		modelStat.Tokens += tokens
		modelStat.Requests++
		return
	}

	modelMap[modelName] = &model.ModelTokens{
		Model:    modelName,
		Tokens:   tokens,
		Requests: 1,
	}
}

func addDailyModelStat(models map[string]model.DailyModelStat, modelName string, tokens int64) map[string]model.DailyModelStat {
	if models == nil {
		models = make(map[string]model.DailyModelStat)
	}

	modelStat := models[modelName]
	modelStat.Tokens += tokens
	modelStat.Requests++
	models[modelName] = modelStat
	return models
}