487 lines
18 KiB
Go
487 lines
18 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
"unicode"
|
|
|
|
managedaccess "billing/internal/access"
|
|
"billing/internal/collection"
|
|
"billing/internal/repository"
|
|
)
|
|
|
|
type createManagedKeyRequest struct {
|
|
Name string `json:"name"`
|
|
Secret string `json:"secret,omitempty"`
|
|
RouteMode string `json:"route_mode,omitempty"`
|
|
UpstreamAccountID string `json:"upstream_account_id,omitempty"`
|
|
AllModels *bool `json:"all_models,omitempty"`
|
|
Models *[]string `json:"models,omitempty"`
|
|
Billing *billingSettingsRequest `json:"billing,omitempty"`
|
|
}
|
|
|
|
type updateManagedKeyRequest struct {
|
|
ID string `json:"id"`
|
|
Name *string `json:"name,omitempty"`
|
|
Status *string `json:"status,omitempty"`
|
|
RouteMode *string `json:"route_mode,omitempty"`
|
|
UpstreamAccountID *string `json:"upstream_account_id,omitempty"`
|
|
AllModels *bool `json:"all_models,omitempty"`
|
|
Models *[]string `json:"models,omitempty"`
|
|
Billing *billingSettingsRequest `json:"billing,omitempty"`
|
|
}
|
|
|
|
type billingSettingsRequest struct {
|
|
QuotaUSD string `json:"quota_usd"`
|
|
ResetPeriod string `json:"reset_period"`
|
|
NextResetAt *time.Time `json:"next_reset_at,omitempty"`
|
|
MaxConcurrency int `json:"max_concurrency"`
|
|
}
|
|
|
|
type billingStateDTO struct {
|
|
QuotaUSD string `json:"quota_usd"`
|
|
SpentUSD string `json:"spent_usd"`
|
|
LifetimeSpentUSD string `json:"lifetime_spent_usd"`
|
|
BalanceUSD string `json:"balance_usd"`
|
|
ResetPeriod string `json:"reset_period"`
|
|
NextResetAt *time.Time `json:"next_reset_at,omitempty"`
|
|
MaxConcurrency int `json:"max_concurrency"`
|
|
ActiveRequests int `json:"active_requests"`
|
|
CycleStartedAt time.Time `json:"cycle_started_at"`
|
|
}
|
|
|
|
type managedKeyDTO struct {
|
|
managedaccess.ManagedKey
|
|
Billing billingStateDTO `json:"billing"`
|
|
}
|
|
|
|
type ledgerEntryDTO struct {
|
|
ID int64 `json:"id"`
|
|
Kind string `json:"kind"`
|
|
AmountUSD string `json:"amount_usd"`
|
|
BalanceAfterUSD string `json:"balance_after_usd"`
|
|
RequestID string `json:"request_id,omitempty"`
|
|
ExecutionID string `json:"execution_id,omitempty"`
|
|
Model string `json:"model,omitempty"`
|
|
OccurredAt time.Time `json:"occurred_at"`
|
|
}
|
|
|
|
type archiveManagedKeyRequest struct {
|
|
ID string `json:"id"`
|
|
}
|
|
|
|
type managedKeyStatsResponse struct {
|
|
Stats managedaccess.KeyStats `json:"stats"`
|
|
Recent []usageListItem `json:"recent"`
|
|
}
|
|
|
|
func (a *App) listManagedKeys(includeArchived bool) ManagementResponse {
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
|
|
}
|
|
keys, err := store.ListManagedKeys(context.Background(), includeArchived)
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
items := make([]managedKeyDTO, 0, len(keys))
|
|
for _, key := range keys {
|
|
state, stateErr := store.BillingState(context.Background(), key.ID, time.Now().UTC())
|
|
if stateErr != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", stateErr.Error())
|
|
}
|
|
items = append(items, managedKeyResponse(key, state))
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"keys": items})
|
|
}
|
|
|
|
func (a *App) createManagedKey(body []byte) ManagementResponse {
|
|
var req createManagedKeyRequest
|
|
if err := json.Unmarshal(body, &req); err != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_request", "无法解析 Key 请求")
|
|
}
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
|
|
}
|
|
name := strings.TrimSpace(req.Name)
|
|
if name == "" || len(name) > 64 {
|
|
return managementError(http.StatusBadRequest, "invalid_name", "Key 名称必须为 1-64 个字符")
|
|
}
|
|
secret := strings.TrimSpace(req.Secret)
|
|
if secret == "" {
|
|
var generateErr error
|
|
secret, generateErr = generatedToken("cpa_")
|
|
if generateErr != nil {
|
|
return managementError(http.StatusInternalServerError, "random_unavailable", "无法生成 Key")
|
|
}
|
|
}
|
|
if !validManagedSecret(secret) {
|
|
return managementError(http.StatusBadRequest, "invalid_key", "Key 必须为 6-256 个不含空白或控制字符的字符")
|
|
}
|
|
keyID, err := generatedToken("key_")
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "random_unavailable", "无法生成 Key ID")
|
|
}
|
|
key := managedaccess.ManagedKey{
|
|
ID: keyID, Name: name, Secret: secret,
|
|
Status: managedaccess.StatusActive, RouteMode: managedaccess.RouteAuto, AllModels: true,
|
|
}
|
|
if template, err := store.ManagedKeyByID(context.Background(), "key_default"); err == nil {
|
|
key.RouteMode = template.RouteMode
|
|
key.UpstreamAccountID = template.UpstreamAccountID
|
|
key.AllModels = template.AllModels
|
|
key.Models = append([]string(nil), template.Models...)
|
|
}
|
|
if req.RouteMode != "" {
|
|
key.RouteMode = req.RouteMode
|
|
}
|
|
if req.UpstreamAccountID != "" {
|
|
key.UpstreamAccountID = strings.TrimSpace(req.UpstreamAccountID)
|
|
}
|
|
if req.AllModels != nil {
|
|
key.AllModels = *req.AllModels
|
|
}
|
|
if req.Models != nil {
|
|
key.Models = normalizeModels(*req.Models)
|
|
}
|
|
if key.RouteMode == managedaccess.RouteAuto {
|
|
key.UpstreamAccountID = ""
|
|
}
|
|
if key.AllModels {
|
|
key.Models = nil
|
|
}
|
|
if err := validateManagedKeyRule(store, key); err != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_rule", err.Error())
|
|
}
|
|
var requestedBilling *managedaccess.BillingSettings
|
|
if req.Billing != nil {
|
|
settings, settingsErr := billingSettings(*req.Billing)
|
|
if settingsErr != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error())
|
|
}
|
|
requestedBilling = &settings
|
|
}
|
|
if err := store.CreateManagedKey(context.Background(), key); err != nil {
|
|
return managementError(http.StatusConflict, "key_conflict", err.Error())
|
|
}
|
|
if requestedBilling != nil {
|
|
if _, settingsErr := store.UpdateBilling(context.Background(), key.ID, *requestedBilling, time.Now().UTC()); settingsErr != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error())
|
|
}
|
|
}
|
|
created, _ := store.ManagedKeyByID(context.Background(), key.ID)
|
|
state, _ := store.BillingState(context.Background(), key.ID, time.Now().UTC())
|
|
return jsonManagementResponse(http.StatusCreated, managedKeyResponse(created, state))
|
|
}
|
|
|
|
func (a *App) updateManagedKey(body []byte) ManagementResponse {
|
|
var req updateManagedKeyRequest
|
|
if err := json.Unmarshal(body, &req); err != nil || strings.TrimSpace(req.ID) == "" {
|
|
return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
|
|
}
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
|
|
}
|
|
key, err := store.ManagedKeyByID(context.Background(), req.ID)
|
|
if err != nil {
|
|
return managementError(http.StatusNotFound, "key_not_found", "Key 不存在")
|
|
}
|
|
if req.Name != nil {
|
|
key.Name = strings.TrimSpace(*req.Name)
|
|
}
|
|
if req.Status != nil {
|
|
key.Status = strings.TrimSpace(*req.Status)
|
|
}
|
|
if req.RouteMode != nil {
|
|
key.RouteMode = strings.TrimSpace(*req.RouteMode)
|
|
}
|
|
if req.UpstreamAccountID != nil {
|
|
key.UpstreamAccountID = strings.TrimSpace(*req.UpstreamAccountID)
|
|
}
|
|
if req.AllModels != nil {
|
|
key.AllModels = *req.AllModels
|
|
}
|
|
if req.Models != nil {
|
|
key.Models = normalizeModels(*req.Models)
|
|
}
|
|
if key.RouteMode == managedaccess.RouteAuto {
|
|
key.UpstreamAccountID = ""
|
|
}
|
|
if key.AllModels {
|
|
key.Models = nil
|
|
}
|
|
if key.Name == "" || len(key.Name) > 64 {
|
|
return managementError(http.StatusBadRequest, "invalid_name", "Key 名称必须为 1-64 个字符")
|
|
}
|
|
if key.Status != managedaccess.StatusActive && key.Status != managedaccess.StatusDisabled {
|
|
return managementError(http.StatusBadRequest, "invalid_status", "状态只能是 active 或 disabled")
|
|
}
|
|
if err := validateManagedKeyRule(store, key); err != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_rule", err.Error())
|
|
}
|
|
var requestedBilling *managedaccess.BillingSettings
|
|
if req.Billing != nil {
|
|
settings, settingsErr := billingSettings(*req.Billing)
|
|
if settingsErr != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error())
|
|
}
|
|
requestedBilling = &settings
|
|
}
|
|
if err := store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
status := http.StatusConflict
|
|
if errors.Is(err, repository.ErrManagedKeyNotFound) {
|
|
status = http.StatusNotFound
|
|
}
|
|
return managementError(status, "update_failed", err.Error())
|
|
}
|
|
if requestedBilling != nil {
|
|
if _, settingsErr := store.UpdateBilling(context.Background(), key.ID, *requestedBilling, time.Now().UTC()); settingsErr != nil {
|
|
return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error())
|
|
}
|
|
}
|
|
updated, _ := store.ManagedKeyByID(context.Background(), key.ID)
|
|
state, _ := store.BillingState(context.Background(), key.ID, time.Now().UTC())
|
|
return jsonManagementResponse(http.StatusOK, managedKeyResponse(updated, state))
|
|
}
|
|
|
|
func billingSettings(request billingSettingsRequest) (managedaccess.BillingSettings, error) {
|
|
quota, err := parseDecimalMicros(request.QuotaUSD)
|
|
if err != nil {
|
|
return managedaccess.BillingSettings{}, errors.New("额度必须是非负美元金额，最多六位小数")
|
|
}
|
|
period := strings.TrimSpace(request.ResetPeriod)
|
|
if period == "" {
|
|
period = managedaccess.ResetNone
|
|
}
|
|
if period != managedaccess.ResetNone && period != managedaccess.ResetDaily && period != managedaccess.ResetWeekly && period != managedaccess.ResetMonthly {
|
|
return managedaccess.BillingSettings{}, errors.New("重置周期只能是 none、daily、weekly 或 monthly")
|
|
}
|
|
if period != managedaccess.ResetNone && request.NextResetAt != nil && !request.NextResetAt.After(time.Now()) {
|
|
return managedaccess.BillingSettings{}, errors.New("下次重置时间必须晚于当前时间")
|
|
}
|
|
maxConcurrency := request.MaxConcurrency
|
|
if maxConcurrency == 0 {
|
|
maxConcurrency = 4
|
|
}
|
|
if maxConcurrency < 1 || maxConcurrency > 64 {
|
|
return managedaccess.BillingSettings{}, errors.New("并发上限必须为 1-64")
|
|
}
|
|
return managedaccess.BillingSettings{QuotaMicros: quota, ResetPeriod: period,
|
|
NextResetAt: request.NextResetAt, MaxConcurrency: maxConcurrency}, nil
|
|
}
|
|
|
|
func managedKeyResponse(key managedaccess.ManagedKey, state managedaccess.BillingState) managedKeyDTO {
|
|
return managedKeyDTO{ManagedKey: key, Billing: billingStateResponse(state)}
|
|
}
|
|
|
|
func (a *App) resetManagedKeyBilling(body []byte) ManagementResponse {
|
|
var request struct {
|
|
ID string `json:"id"`
|
|
All bool `json:"all"`
|
|
}
|
|
if err := json.Unmarshal(body, &request); err != nil || !request.All && strings.TrimSpace(request.ID) == "" {
|
|
return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
|
|
}
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "额度数据库尚未初始化")
|
|
}
|
|
if request.All {
|
|
count, err := store.ResetAllBilling(context.Background(), time.Now().UTC())
|
|
if err != nil {
|
|
return managementError(http.StatusBadRequest, "reset_failed", err.Error())
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]int{"reset_count": count})
|
|
}
|
|
state, err := store.ResetBilling(context.Background(), request.ID, time.Now().UTC())
|
|
if err != nil {
|
|
return managementError(http.StatusBadRequest, "reset_failed", err.Error())
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, billingStateDTO{
|
|
QuotaUSD: formatMicros(state.QuotaMicros), SpentUSD: formatMicros(state.SpentMicros),
|
|
LifetimeSpentUSD: formatMicros(state.LifetimeSpentMicros),
|
|
BalanceUSD: formatMicros(state.BalanceMicros), ResetPeriod: state.ResetPeriod,
|
|
NextResetAt: state.NextResetAt, MaxConcurrency: state.MaxConcurrency,
|
|
ActiveRequests: state.ActiveRequests, CycleStartedAt: state.CycleStartedAt,
|
|
})
|
|
}
|
|
|
|
func (a *App) managedKeyLedger(query url.Values) ManagementResponse {
|
|
keyID := strings.TrimSpace(query.Get("id"))
|
|
if keyID == "" {
|
|
return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
|
|
}
|
|
before, _ := strconv.ParseInt(query.Get("before"), 10, 64)
|
|
limit, _ := strconv.Atoi(query.Get("limit"))
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "额度数据库尚未初始化")
|
|
}
|
|
entries, err := store.ListBillingLedger(context.Background(), keyID, before, limit)
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
items := make([]ledgerEntryDTO, 0, len(entries))
|
|
for _, entry := range entries {
|
|
items = append(items, ledgerEntryDTO{ID: entry.ID, Kind: entry.Kind,
|
|
AmountUSD: formatMicros(entry.AmountMicros), BalanceAfterUSD: formatMicros(entry.BalanceAfterMicros),
|
|
RequestID: entry.RequestID, ExecutionID: entry.ExecutionID, Model: entry.Model, OccurredAt: entry.OccurredAt})
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"entries": items})
|
|
}
|
|
|
|
func (a *App) archiveManagedKey(body []byte) ManagementResponse {
|
|
var req archiveManagedKeyRequest
|
|
if err := json.Unmarshal(body, &req); err != nil || strings.TrimSpace(req.ID) == "" {
|
|
return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID")
|
|
}
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
|
|
}
|
|
if err := store.ArchiveManagedKey(context.Background(), req.ID); err != nil {
|
|
return managementError(http.StatusNotFound, "key_not_found", "Key 不存在或已经归档")
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]bool{"archived": true})
|
|
}
|
|
|
|
func (a *App) managedKeyStats(id string) ManagementResponse {
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "Key 数据库尚未初始化")
|
|
}
|
|
if _, err := store.ManagedKeyByID(context.Background(), id); err != nil {
|
|
return managementError(http.StatusNotFound, "key_not_found", "Key 不存在")
|
|
}
|
|
location, _ := time.LoadLocation("Asia/Shanghai")
|
|
now := time.Now().In(location)
|
|
today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, location)
|
|
stats, err := store.KeyStats(context.Background(), id, today)
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
records, err := store.QueryUsage(context.Background(), collection.UsageQuery{Page: 1, PageSize: 50, KeyID: id})
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
recent := make([]usageListItem, 0, len(records.Records))
|
|
for _, record := range records.Records {
|
|
recent = append(recent, usageItem(record))
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, managedKeyStatsResponse{Stats: stats, Recent: recent})
|
|
}
|
|
|
|
func (a *App) upstreamAccounts() ManagementResponse {
|
|
syncErr := a.syncUpstreamAccounts(context.Background())
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "上游账号数据库尚未初始化")
|
|
}
|
|
if observedErr := a.flushObservedUpstreams(context.Background(), store); observedErr != nil && syncErr == nil {
|
|
syncErr = observedErr
|
|
}
|
|
accounts, err := store.ListUpstreamAccounts(context.Background())
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
response := map[string]any{"accounts": accounts}
|
|
if syncErr != nil {
|
|
response["warning"] = syncErr.Error()
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, response)
|
|
}
|
|
|
|
func (a *App) modelSuggestions() ManagementResponse {
|
|
store, ok := a.currentStore()
|
|
if !ok {
|
|
return managementError(http.StatusServiceUnavailable, "database_unavailable", "数据库尚未初始化")
|
|
}
|
|
models, err := store.ModelSuggestions(context.Background())
|
|
if err != nil {
|
|
return managementError(http.StatusInternalServerError, "database_error", err.Error())
|
|
}
|
|
return jsonManagementResponse(http.StatusOK, map[string]any{"models": models})
|
|
}
|
|
|
|
func (a *App) currentStore() (*repository.SQLiteUsageRepository, bool) {
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
return a.store, !a.closed && a.store != nil
|
|
}
|
|
|
|
func validateManagedKeyRule(store *repository.SQLiteUsageRepository, key managedaccess.ManagedKey) error {
|
|
if key.RouteMode != managedaccess.RouteAuto && key.RouteMode != managedaccess.RouteStrict {
|
|
return errors.New("路由模式只能是 auto 或 strict")
|
|
}
|
|
if key.RouteMode == managedaccess.RouteAuto {
|
|
key.UpstreamAccountID = ""
|
|
} else {
|
|
if strings.TrimSpace(key.UpstreamAccountID) == "" {
|
|
return errors.New("严格路由必须选择上游账号")
|
|
}
|
|
if _, err := store.UpstreamAccountByID(context.Background(), key.UpstreamAccountID); err != nil {
|
|
return errors.New("选择的上游账号不存在")
|
|
}
|
|
}
|
|
if !key.AllModels && len(normalizeModels(key.Models)) == 0 {
|
|
return errors.New("未允许全部模型时至少配置一个模型")
|
|
}
|
|
for _, pattern := range key.Models {
|
|
if len(pattern) > 128 || strings.IndexFunc(pattern, func(value rune) bool { return unicode.IsSpace(value) || unicode.IsControl(value) }) >= 0 {
|
|
return errors.New("模型规则不能包含空白或控制字符，且最长 128 个字符")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func normalizeModels(models []string) []string {
|
|
seen := make(map[string]struct{})
|
|
result := make([]string, 0, len(models))
|
|
for _, raw := range models {
|
|
model := strings.ToLower(strings.TrimSpace(raw))
|
|
if model == "" {
|
|
continue
|
|
}
|
|
if _, exists := seen[model]; exists {
|
|
continue
|
|
}
|
|
seen[model] = struct{}{}
|
|
result = append(result, model)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func generatedToken(prefix string) (string, error) {
|
|
raw := make([]byte, 24)
|
|
if _, err := rand.Read(raw); err != nil {
|
|
return "", err
|
|
}
|
|
return prefix + base64.RawURLEncoding.EncodeToString(raw), nil
|
|
}
|
|
|
|
func validManagedSecret(secret string) bool {
|
|
if len(secret) < 6 || len(secret) > 256 {
|
|
return false
|
|
}
|
|
for _, value := range secret {
|
|
if unicode.IsSpace(value) || unicode.IsControl(value) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func managementError(status int, code, message string) ManagementResponse {
|
|
return jsonManagementResponse(status, map[string]any{"error": map[string]string{"code": code, "message": message}})
|
|
}
|