Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
11 kB · 416 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417package usage
import ( "context" "net/http" "strings" "sync" "time"
log "github.com/sirupsen/logrus")
// DefaultServiceTier is retained for direct SDK and non-OpenAI usage callers.const DefaultServiceTier = "default"
// AutoServiceTier is the OpenAI request semantics when service_tier is omitted.// OpenAI HTTP handlers set it explicitly, without changing other providers'// historical direct-SDK default.const AutoServiceTier = "auto"
// Record contains the usage statistics captured for a single provider request.type Record struct { Provider string // ExecutorType stores the concrete executor type that handled the request. ExecutorType string Model string Alias string APIKey string SessionID string ParentSessionID string AuthID string AuthIndex string // AccessTokenSHA256 identifies the OAuth token version without exposing the token. AccessTokenSHA256 string AuthType string Source string // ReasoningEffort stores the translated upstream thinking level for request event logs. ReasoningEffort string // ServiceTier stores the client-requested service tier. ServiceTier string // RequestServiceTier is a deprecated input-only alias retained for existing // plugin callers. It is normalized into ServiceTier and never emitted. RequestServiceTier string // ResponseServiceTier stores the final tier reported by the upstream response. ResponseServiceTier string // Generate reports whether the client requested actual generation. // nil or true means generation is enabled; only an explicit false disables generation. // Use GenerateFlag to set the value and GenerateEnabled to read it with the default. Generate *bool // Stream reports whether the request was executed in streaming mode. Stream bool RequestedAt time.Time Latency time.Duration TTFT time.Duration Failed bool Fail Failure Detail Detail // ResponseHeaders stores a snapshot of upstream response headers for usage sinks. ResponseHeaders http.Header}
// Failure holds HTTP failure metadata for an upstream request attempt.type Failure struct { StatusCode int Body string}
// Detail holds the token usage breakdown.type Detail struct { InputTokens int64 OutputTokens int64 ReasoningTokens int64 CachedTokens int64 CacheReadTokens int64 CacheCreationTokens int64 TotalTokens int64 TokenBreakdown TokenBreakdown ResponseServiceTier string}
type requestedModelAliasContextKey struct{}type reasoningEffortContextKey struct{}type serviceTierContextKey struct{}type generateContextKey struct{}type streamContextKey struct{}
// WithRequestedModelAlias stores the client-requested model name for usage sinks.func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { if ctx == nil { ctx = context.Background() } alias = strings.TrimSpace(alias) if alias == "" { return ctx } return context.WithValue(ctx, requestedModelAliasContextKey{}, alias)}
// RequestedModelAliasFromContext returns the client-requested model name stored in ctx.func RequestedModelAliasFromContext(ctx context.Context) string { if ctx == nil { return "" } raw := ctx.Value(requestedModelAliasContextKey{}) switch value := raw.(type) { case string: return strings.TrimSpace(value) case []byte: return strings.TrimSpace(string(value)) default: return "" }}
// WithReasoningEffort stores the client-requested reasoning effort for usage sinks.func WithReasoningEffort(ctx context.Context, effort string) context.Context { if ctx == nil { ctx = context.Background() } effort = strings.TrimSpace(effort) if effort == "" { return ctx } return context.WithValue(ctx, reasoningEffortContextKey{}, effort)}
// ReasoningEffortFromContext returns the client-requested reasoning effort stored in ctx.func ReasoningEffortFromContext(ctx context.Context) string { if ctx == nil { return "" } raw := ctx.Value(reasoningEffortContextKey{}) switch value := raw.(type) { case string: return strings.TrimSpace(value) case []byte: return strings.TrimSpace(string(value)) default: return "" }}
// WithServiceTier stores the client-requested service tier for usage sinks.func WithServiceTier(ctx context.Context, tier string) context.Context { if ctx == nil { ctx = context.Background() } tier = strings.TrimSpace(tier) if tier == "" { tier = DefaultServiceTier } return context.WithValue(ctx, serviceTierContextKey{}, tier)}
// ServiceTierFromContext returns the client-requested service tier stored in ctx.func ServiceTierFromContext(ctx context.Context) string { if ctx == nil { return DefaultServiceTier } raw := ctx.Value(serviceTierContextKey{}) switch value := raw.(type) { case string: tier := strings.TrimSpace(value) if tier == "" { return DefaultServiceTier } return tier case []byte: tier := strings.TrimSpace(string(value)) if tier == "" { return DefaultServiceTier } return tier default: return DefaultServiceTier }}
// WithGenerate stores whether the client requested actual generation for usage sinks.// Missing context values default to true; only an explicit false disables generation.func WithGenerate(ctx context.Context, generate bool) context.Context { if ctx == nil { ctx = context.Background() } return context.WithValue(ctx, generateContextKey{}, generate)}
// GenerateFromContext returns whether the client requested actual generation.// Missing values default to true.func GenerateFromContext(ctx context.Context) bool { if ctx == nil { return true } raw := ctx.Value(generateContextKey{}) switch value := raw.(type) { case bool: return value default: return true }}
// WithStream stores whether the request was executed in streaming mode for usage sinks.func WithStream(ctx context.Context, stream bool) context.Context { if ctx == nil { ctx = context.Background() } return context.WithValue(ctx, streamContextKey{}, stream)}
// StreamFromContext returns whether the request was executed in streaming mode.// Missing values default to false.func StreamFromContext(ctx context.Context) bool { if ctx == nil { return false } raw := ctx.Value(streamContextKey{}) switch value := raw.(type) { case bool: return value default: return false }}
// GenerateFlag returns a pointer suitable for Record.Generate.func GenerateFlag(generate bool) *bool { return &generate}
// GenerateEnabled reports whether generation is enabled for the record field.// A nil value defaults to true so legacy callers that omit Generate keep the historical behavior.func GenerateEnabled(generate *bool) bool { if generate == nil { return true } return *generate}
// Plugin consumes usage records emitted by the proxy runtime.type Plugin interface { HandleUsage(ctx context.Context, record Record)}
type queueItem struct { ctx context.Context record Record}
// Manager maintains a queue of usage records and delivers them to registered plugins.type Manager struct { once sync.Once stopOnce sync.Once cancel context.CancelFunc
mu sync.Mutex cond *sync.Cond queue []queueItem closed bool
pluginsMu sync.RWMutex plugins []Plugin named map[string]int}
// NewManager constructs a manager with a buffered queue.func NewManager(buffer int) *Manager { m := &Manager{} m.cond = sync.NewCond(&m.mu) return m}
// Start launches the background dispatcher. Calling Start multiple times is safe.func (m *Manager) Start(ctx context.Context) { if m == nil { return } m.once.Do(func() { if ctx == nil { ctx = context.Background() } var workerCtx context.Context workerCtx, m.cancel = context.WithCancel(ctx) go m.run(workerCtx) })}
// Stop stops the dispatcher and drains the queue.func (m *Manager) Stop() { if m == nil { return } m.stopOnce.Do(func() { if m.cancel != nil { m.cancel() } m.mu.Lock() m.closed = true m.mu.Unlock() m.cond.Broadcast() })}
// Register appends a plugin to the delivery list.func (m *Manager) Register(plugin Plugin) { if m == nil || plugin == nil { return } m.pluginsMu.Lock() m.plugins = append(m.plugins, plugin) m.pluginsMu.Unlock()}
// RegisterNamed registers or replaces a plugin by name.func (m *Manager) RegisterNamed(name string, plugin Plugin) { if m == nil || plugin == nil { return } name = strings.TrimSpace(name) if name == "" { return }
m.pluginsMu.Lock() if m.named == nil { m.named = make(map[string]int) } if index, exists := m.named[name]; exists && index >= 0 && index < len(m.plugins) { m.plugins[index] = plugin m.pluginsMu.Unlock() return } m.named[name] = len(m.plugins) m.plugins = append(m.plugins, plugin) m.pluginsMu.Unlock()}
// Publish enqueues a usage record for processing. If no plugin is registered// the record will be discarded downstream.func (m *Manager) Publish(ctx context.Context, record Record) { if m == nil { return } // ensure worker is running even if Start was not called explicitly m.Start(context.Background()) m.mu.Lock() if m.closed { m.mu.Unlock() return } m.queue = append(m.queue, queueItem{ctx: ctx, record: record}) m.mu.Unlock() m.cond.Signal()}
func (m *Manager) run(ctx context.Context) { for { m.mu.Lock() for !m.closed && len(m.queue) == 0 { m.cond.Wait() } if len(m.queue) == 0 && m.closed { m.mu.Unlock() return } item := m.queue[0] m.queue = m.queue[1:] m.mu.Unlock() m.dispatch(item) }}
func (m *Manager) dispatch(item queueItem) { m.pluginsMu.RLock() plugins := make([]Plugin, len(m.plugins)) copy(plugins, m.plugins) m.pluginsMu.RUnlock() if len(plugins) == 0 { return } for _, plugin := range plugins { if plugin == nil { continue } safeInvoke(plugin, item.ctx, item.record) }}
func safeInvoke(plugin Plugin, ctx context.Context, record Record) { defer func() { if r := recover(); r != nil { log.Errorf("usage: plugin panic recovered: %v", r) } }() plugin.HandleUsage(ctx, record)}
var defaultManager = NewManager(512)
// DefaultManager returns the global usage manager instance.func DefaultManager() *Manager { return defaultManager }
// RegisterPlugin registers a plugin on the default manager.func RegisterPlugin(plugin Plugin) { DefaultManager().Register(plugin) }
// RegisterNamedPlugin registers or replaces a named plugin on the default manager.func RegisterNamedPlugin(name string, plugin Plugin) { DefaultManager().RegisterNamed(name, plugin) }
// PublishRecord publishes a record using the default manager.func PublishRecord(ctx context.Context, record Record) { DefaultManager().Publish(ctx, record) }
// StartDefault starts the default manager's dispatcher.func StartDefault(ctx context.Context) { DefaultManager().Start(ctx) }
// StopDefault stops the default manager's dispatcher.func StopDefault() { DefaultManager().Stop() }