297 lines
9.0 KiB
Go
297 lines
9.0 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"cpa-ext/internal/collection"
|
|
"cpa-ext/internal/pricing"
|
|
"cpa-ext/internal/repository"
|
|
)
|
|
|
|
type App struct {
|
|
mu sync.RWMutex
|
|
config Config
|
|
usage *collection.Service
|
|
store *repository.SQLiteUsageRepository
|
|
prices map[string]pricing.Policy
|
|
host HostCaller
|
|
closed bool
|
|
seen atomic.Uint64
|
|
pending sync.Map
|
|
scopes sync.Map
|
|
observedUpstreams sync.Map
|
|
}
|
|
|
|
func NewApp() *App {
|
|
return &App{config: defaultConfig(), prices: make(map[string]pricing.Policy)}
|
|
}
|
|
|
|
func (a *App) HandleMethod(method string, request []byte) (response []byte, err error) {
|
|
defer func() {
|
|
if recovered := recover(); recovered != nil {
|
|
response = nil
|
|
err = fmt.Errorf("插件处理 %s 时发生异常: %v", method, recovered)
|
|
}
|
|
}()
|
|
|
|
switch method {
|
|
case MethodPluginRegister, MethodPluginReconfigure:
|
|
return a.configure(request)
|
|
case MethodUsageHandle:
|
|
return a.handleUsage(request)
|
|
case MethodFrontendIdentifier:
|
|
return a.frontendIdentifier()
|
|
case MethodFrontendAuthenticate:
|
|
return a.authenticate(request)
|
|
case MethodSchedulerPick:
|
|
return a.pickScheduler(request)
|
|
case MethodRequestBefore:
|
|
return a.interceptRequest(request, false)
|
|
case MethodRequestAfter:
|
|
return a.interceptRequest(request, true)
|
|
case MethodRequestComplete:
|
|
return a.handleRequestComplete(request)
|
|
case MethodManagementRegister:
|
|
return OKEnvelope(managementRegistration())
|
|
case MethodManagementHandle:
|
|
return a.handleManagement(request)
|
|
case MethodPluginShutdown:
|
|
a.Shutdown()
|
|
return OKEnvelope(struct{}{})
|
|
default:
|
|
return ErrorEnvelope("unknown_method", "不支持的插件方法: "+method, http.StatusNotFound), nil
|
|
}
|
|
}
|
|
|
|
func (a *App) configure(raw []byte) ([]byte, error) {
|
|
var req LifecycleRequest
|
|
if len(raw) > 0 {
|
|
if err := json.Unmarshal(raw, &req); err != nil {
|
|
return nil, fmt.Errorf("解析生命周期请求: %w", err)
|
|
}
|
|
}
|
|
if req.SchemaVersion == 0 {
|
|
req.SchemaVersion = 1
|
|
}
|
|
negotiated := min(req.SchemaVersion, SchemaVersion)
|
|
cfg, err := decodeConfig(req.ConfigYAML)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
usageRepository, err := repository.OpenSQLiteUsage(cfg.DatabasePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if _, err := usageRepository.BootstrapManagedKey(context.Background(), cfg.BootstrapName, cfg.BootstrapKey); err != nil {
|
|
_ = usageRepository.Close()
|
|
return nil, err
|
|
}
|
|
nextUsage := collection.NewService(usageRepository)
|
|
policies, err := usageRepository.ListPrices(context.Background())
|
|
if err != nil {
|
|
_ = nextUsage.Close()
|
|
return nil, err
|
|
}
|
|
nextPrices := make(map[string]pricing.Policy, len(policies))
|
|
for _, policy := range policies {
|
|
nextPrices[policy.Model] = policy
|
|
}
|
|
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.closed {
|
|
_ = nextUsage.Close()
|
|
return nil, fmt.Errorf("插件已经关闭")
|
|
}
|
|
previousUsage := a.usage
|
|
a.config = cfg
|
|
a.usage = nextUsage
|
|
a.store = usageRepository
|
|
a.prices = nextPrices
|
|
if previousUsage != nil {
|
|
_ = previousUsage.Close()
|
|
}
|
|
return OKEnvelope(registration(negotiated))
|
|
}
|
|
|
|
func registration(schemaVersion uint32) Registration {
|
|
return Registration{
|
|
SchemaVersion: schemaVersion,
|
|
Metadata: Metadata{
|
|
Name: PluginName,
|
|
Version: Version,
|
|
Author: "cpa-ext",
|
|
GitHubRepository: "https://git.pchuan.top/agent/cpa-plugin",
|
|
ConfigFields: []ConfigField{
|
|
{Name: "enabled", Type: "boolean", Description: "启用 CPA 扩展。"},
|
|
{Name: "codex_only", Type: "boolean", Description: "只接收 Codex/OpenAI 模型的用量事件。"},
|
|
{Name: "database_path", Type: "string", Description: "SQLite 数据库文件路径。"},
|
|
{Name: "bootstrap_name", Type: "string", Description: "首次启动时现有 Key 的名称。"},
|
|
{Name: "bootstrap_key", Type: "string", Description: "首次启动时导入的现有下游 Key。"},
|
|
},
|
|
},
|
|
Capabilities: Capabilities{
|
|
FrontendAuthProvider: true, FrontendAuthProviderExclusive: true,
|
|
Scheduler: true, RequestInterceptor: true,
|
|
RequestLifecyclePlugin: true, UsagePlugin: true, ManagementAPI: true,
|
|
},
|
|
}
|
|
}
|
|
|
|
func (a *App) handleUsage(raw []byte) ([]byte, error) {
|
|
var record UsageRecord
|
|
if err := json.Unmarshal(raw, &record); err != nil {
|
|
return nil, fmt.Errorf("解析用量事件: %w", err)
|
|
}
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
if a.closed {
|
|
return nil, fmt.Errorf("插件已经关闭")
|
|
}
|
|
if a.config.accepts(record) {
|
|
if a.usage == nil {
|
|
return nil, fmt.Errorf("用量数据库尚未初始化")
|
|
}
|
|
// CPA wire 类型只存在于适配层，采集模块接收与协议无关的观察值。
|
|
observed := collection.Record{
|
|
RequestID: record.RequestID,
|
|
ExecutionID: record.ExecutionID,
|
|
TraceID: record.TraceID,
|
|
RequestedAt: record.RequestedAt,
|
|
APIKey: record.APIKey,
|
|
AuthID: record.AuthID,
|
|
AuthIndex: record.AuthIndex,
|
|
AuthType: record.AuthType,
|
|
Model: record.Model,
|
|
ReasoningEffort: record.ReasoningEffort,
|
|
ServiceTier: record.ServiceTier,
|
|
Speed: record.Speed,
|
|
Failed: record.Failed,
|
|
ExecutorType: record.ExecutorType,
|
|
RequestType: usageRequestType(record.Endpoint),
|
|
Endpoint: record.Endpoint,
|
|
InputTokens: record.Detail.InputTokens,
|
|
OutputTokens: record.Detail.OutputTokens,
|
|
TotalTokens: record.Detail.TotalTokens,
|
|
CachedTokens: record.Detail.CachedTokens,
|
|
CacheReadTokens: record.Detail.CacheReadTokens,
|
|
CacheWriteTokens: record.Detail.CacheCreationTokens,
|
|
ReasoningTokens: record.Detail.ReasoningTokens,
|
|
TTFT: record.TTFT,
|
|
Latency: record.Latency,
|
|
ClientIP: record.ClientIP,
|
|
}
|
|
if key, keyErr := a.store.ManagedKeyByCredential(context.Background(), record.APIKey); keyErr == nil {
|
|
observed.ManagedKeyID = key.ID
|
|
observed.KeyAlias = key.Name
|
|
}
|
|
if policy, found := a.prices[strings.TrimSpace(record.Model)]; found {
|
|
cost, calculateErr := pricing.Calculate(policy, pricing.Usage{
|
|
InputTokens: record.Detail.InputTokens, CacheReadTokens: record.Detail.CacheReadTokens,
|
|
CacheWriteTokens: record.Detail.CacheCreationTokens, OutputTokens: record.Detail.OutputTokens,
|
|
ServiceTier: record.ServiceTier, Speed: record.Speed,
|
|
})
|
|
if calculateErr == nil {
|
|
observed.CostMicros = &cost.CostMicros
|
|
observed.PriceTier = cost.PriceTier
|
|
observed.FastRequested = cost.FastRequested
|
|
observed.FastPricingApplied = cost.FastApplied
|
|
observed.PriceMultiplierNumerator = cost.MultiplierNumerator
|
|
observed.PriceMultiplierDenominator = cost.MultiplierDenominator
|
|
}
|
|
}
|
|
if err := a.usage.Observe(context.Background(), observed); err != nil {
|
|
return nil, err
|
|
}
|
|
a.seen.Add(1)
|
|
}
|
|
return OKEnvelope(struct{}{})
|
|
}
|
|
|
|
func (a *App) handleRequestComplete(raw []byte) ([]byte, error) {
|
|
var completion RequestCompletion
|
|
if err := json.Unmarshal(raw, &completion); err != nil {
|
|
return nil, fmt.Errorf("解析请求终态: %w", err)
|
|
}
|
|
a.mu.RLock()
|
|
defer a.mu.RUnlock()
|
|
if a.closed {
|
|
return nil, fmt.Errorf("插件已经关闭")
|
|
}
|
|
if a.usage == nil {
|
|
return nil, fmt.Errorf("用量数据库尚未初始化")
|
|
}
|
|
endpoint, _ := completion.Metadata["request_path"].(string)
|
|
managedKeyID := metadataString(completion.Metadata, callerScopeMetadata)
|
|
if key, keyErr := a.store.ManagedKeyByReference(context.Background(), managedKeyID); keyErr == nil {
|
|
managedKeyID = key.ID
|
|
}
|
|
if pending, ok := a.pending.LoadAndDelete(completion.RequestID); managedKeyID == "" && ok {
|
|
managedKeyID, _ = pending.(string)
|
|
}
|
|
if err := a.usage.ObserveRequest(context.Background(), collection.RequestRecord{
|
|
ManagedKeyID: managedKeyID, RequestID: completion.RequestID, TraceID: completion.TraceID,
|
|
RequestedAt: completion.StartedAt, CompletedAt: completion.CompletedAt,
|
|
Model: completion.Model, RequestedModel: completion.RequestedModel,
|
|
SourceFormat: completion.SourceFormat, Stream: completion.Stream,
|
|
Outcome: string(completion.Outcome), StatusCode: completion.StatusCode,
|
|
Error: completion.Error, Endpoint: strings.TrimSpace(endpoint),
|
|
}); err != nil {
|
|
return nil, err
|
|
}
|
|
return OKEnvelope(struct{}{})
|
|
}
|
|
|
|
func usageRequestType(endpoint string) string {
|
|
if isCompactEndpoint(endpoint) {
|
|
return "JSON"
|
|
}
|
|
parts := strings.Fields(endpoint)
|
|
if len(parts) == 0 {
|
|
return ""
|
|
}
|
|
switch strings.ToUpper(parts[0]) {
|
|
case http.MethodPost:
|
|
return "SSE"
|
|
case http.MethodGet:
|
|
return "WS"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func isCompactEndpoint(endpoint string) bool {
|
|
path := strings.ToLower(strings.TrimSpace(endpoint))
|
|
if fields := strings.Fields(path); len(fields) > 1 {
|
|
path = fields[len(fields)-1]
|
|
}
|
|
if query := strings.IndexByte(path, '?'); query >= 0 {
|
|
path = path[:query]
|
|
}
|
|
return strings.HasSuffix(strings.TrimRight(path, "/"), "/responses/compact")
|
|
}
|
|
|
|
func (a *App) Seen() uint64 {
|
|
return a.seen.Load()
|
|
}
|
|
|
|
func (a *App) Shutdown() {
|
|
a.mu.Lock()
|
|
defer a.mu.Unlock()
|
|
if a.closed {
|
|
return
|
|
}
|
|
a.closed = true
|
|
if a.usage != nil {
|
|
_ = a.usage.Close()
|
|
a.usage = nil
|
|
a.store = nil
|
|
}
|
|
}
|