Documentation
¶
Overview ¶
Package oauthmanager provides OAuth credential management and rotation for AI providers. It includes health tracking, automatic refresh, failover, and monitoring for OAuth tokens.
Index ¶
- Constants
- func Clone(c *types.OAuthCredentialSet) *types.OAuthCredentialSet
- func IsExpired(c *types.OAuthCredentialSet) bool
- func NeedsRefresh(cred *types.OAuthCredentialSet) bool
- func NoOpRefreshFunc(ctx context.Context, cred *types.OAuthCredentialSet) (*types.OAuthCredentialSet, error)
- type CredentialHealthInfo
- type CredentialMetrics
- type HealthSummary
- type MonitoringConfig
- type OAuthKeyManager
- func (m *OAuthKeyManager) AutoDecommissionExpired() ([]string, error)
- func (m *OAuthKeyManager) CheckAlerts() []WebhookEvent
- func (m *OAuthKeyManager) CheckRotationNeeded() []string
- func (m *OAuthKeyManager) CompleteRotation(credentialID string) error
- func (m *OAuthKeyManager) ExecuteWithFailover(ctx context.Context, ...) (string, *types.Usage, error)
- func (m *OAuthKeyManager) ExportJSON() ([]byte, error)
- func (m *OAuthKeyManager) ExportPrometheus() *PrometheusMetrics
- func (m *OAuthKeyManager) GetAllMetrics() map[string]*CredentialMetrics
- func (m *OAuthKeyManager) GetCredentialHealth(credentialID string) *credentialHealth
- func (m *OAuthKeyManager) GetCredentialHealthInfo(credentialID string) *CredentialHealthInfo
- func (m *OAuthKeyManager) GetCredentialMetrics(credentialID string) *CredentialMetrics
- func (m *OAuthKeyManager) GetCredentials() []*types.OAuthCredentialSet
- func (m *OAuthKeyManager) GetHealthSummary() map[string]interface{}
- func (m *OAuthKeyManager) GetMonitoringConfig() *MonitoringConfig
- func (m *OAuthKeyManager) GetNextCredential(ctx context.Context) (*types.OAuthCredentialSet, error)
- func (m *OAuthKeyManager) GetRefreshStrategy() *RefreshStrategy
- func (m *OAuthKeyManager) GetRotationPolicy() *RotationPolicy
- func (m *OAuthKeyManager) GetRotationState(credentialID string) *credentialRotationState
- func (m *OAuthKeyManager) MarkForRotation(credentialID string, newCredential *types.OAuthCredentialSet) error
- func (m *OAuthKeyManager) RecordRequest(credentialID string, tokensUsed int64, latency time.Duration, success bool)
- func (m *OAuthKeyManager) ReportFailure(credentialID string, err error)
- func (m *OAuthKeyManager) ReportSuccess(credentialID string)
- func (m *OAuthKeyManager) SetMonitoringConfig(config *MonitoringConfig)
- func (m *OAuthKeyManager) SetRefreshStrategy(strategy *RefreshStrategy)
- func (m *OAuthKeyManager) SetRotationPolicy(policy *RotationPolicy)
- type PrometheusMetrics
- type RefreshFunc
- type RefreshStrategy
- type RotationPolicy
- type WebhookEvent
Examples ¶
Constants ¶
const RefreshBufferTime = 5 * time.Minute
RefreshBufferTime is the time before expiry when refresh should occur Tokens expiring within this window should be refreshed proactively
Variables ¶
This section is empty.
Functions ¶
func Clone ¶
func Clone(c *types.OAuthCredentialSet) *types.OAuthCredentialSet
Clone creates a deep copy of the credential set Used when updating credentials to avoid race conditions
func IsExpired ¶
func IsExpired(c *types.OAuthCredentialSet) bool
IsExpired checks if the access token is expired or will expire soon Uses a 5-minute buffer to avoid edge cases
func NeedsRefresh ¶
func NeedsRefresh(cred *types.OAuthCredentialSet) bool
NeedsRefresh checks if a credential needs to be refreshed Returns true if the token expires within RefreshBufferTime or is already expired Returns false if: - Credential is nil - ExpiresAt is not set (zero time) - Token has sufficient time remaining (> RefreshBufferTime)
func NoOpRefreshFunc ¶
func NoOpRefreshFunc(ctx context.Context, cred *types.OAuthCredentialSet) (*types.OAuthCredentialSet, error)
NoOpRefreshFunc is a refresh function that always returns an error Used when token refresh is not supported or configured for a provider This allows the OAuthKeyManager to be created even when refresh is not available
Types ¶
type CredentialHealthInfo ¶
type CredentialHealthInfo struct {
ID string `json:"id"`
IsHealthy bool `json:"is_healthy"`
IsAvailable bool `json:"is_available"`
FailureCount int `json:"failure_count"`
SuccessRate float64 `json:"success_rate"`
RequestCount int64 `json:"request_count"`
TokensUsed int64 `json:"tokens_used"`
AverageLatency time.Duration `json:"average_latency"`
LastUsed time.Time `json:"last_used"`
LastSuccess time.Time `json:"last_success"`
LastFailure time.Time `json:"last_failure"`
BackoffUntil time.Time `json:"backoff_until,omitempty"`
RefreshCount int `json:"refresh_count"`
LastRefresh time.Time `json:"last_refresh"`
RefreshFailCount int `json:"refresh_fail_count"`
ExpiresAt time.Time `json:"expires_at,omitempty"`
MarkedForRotation bool `json:"marked_for_rotation"`
DecommissionAt time.Time `json:"decommission_at,omitempty"`
}
CredentialHealthInfo provides detailed health info for a single credential
type CredentialMetrics ¶
type CredentialMetrics struct {
// Usage tracking
RequestCount int64 // Total requests using this credential
SuccessCount int64 // Successful requests
ErrorCount int64 // Failed requests
// Token usage
TokensUsed int64 // Cumulative tokens consumed
// Performance
TotalLatency time.Duration // Cumulative latency
AverageLatency time.Duration // Calculated average
// Timing
FirstUsed time.Time // When credential was first used
LastUsed time.Time // Most recent usage
// Refresh tracking
RefreshCount int // Number of token refreshes
LastRefreshTime time.Time // Last successful refresh
// contains filtered or unexported fields
}
CredentialMetrics tracks usage and performance metrics for a single credential
func NewCredentialMetrics ¶
func NewCredentialMetrics() *CredentialMetrics
NewCredentialMetrics creates a new metrics tracker for a credential
func (*CredentialMetrics) GetRequestsPerHour ¶
func (m *CredentialMetrics) GetRequestsPerHour() float64
GetRequestsPerHour returns the average requests per hour based on usage history
func (*CredentialMetrics) GetSnapshot ¶
func (m *CredentialMetrics) GetSnapshot() *CredentialMetrics
GetSnapshot returns a thread-safe copy of the current metrics
func (*CredentialMetrics) GetSuccessRate ¶
func (m *CredentialMetrics) GetSuccessRate() float64
GetSuccessRate returns the success rate as a percentage (0.0 to 1.0)
type HealthSummary ¶
type HealthSummary struct {
TotalCredentials int `json:"total_credentials"`
HealthyCredentials int `json:"healthy_credentials"`
UnhealthyCredentials int `json:"unhealthy_credentials"`
CredentialsInBackoff int `json:"credentials_in_backoff"`
TotalRequests int64 `json:"total_requests"`
SuccessRate float64 `json:"success_rate"`
Credentials map[string]CredentialHealthInfo `json:"credentials"`
}
HealthSummary provides a high-level health overview
type MonitoringConfig ¶
type MonitoringConfig struct {
// Webhooks
WebhookURL string // URL for webhook notifications
WebhookEvents []string // Events to notify: "refresh", "failure", "rotation", "expiry_warning"
// Alerting
AlertOnHighFailureRate bool // Alert if error rate > threshold
FailureRateThreshold float64 // Default: 0.25 (25%)
AlertOnRefreshFailure bool // Alert on refresh failures
AlertOnExpirySoon bool // Alert when tokens expire soon
ExpiryWarningTime time.Duration // Default: 24 hours
}
MonitoringConfig defines monitoring and alerting configuration
func DefaultMonitoringConfig ¶
func DefaultMonitoringConfig() *MonitoringConfig
DefaultMonitoringConfig returns a default monitoring configuration
type OAuthKeyManager ¶
type OAuthKeyManager struct {
// contains filtered or unexported fields
}
OAuthKeyManager manages multiple OAuth credentials with load balancing and failover
Example ¶
Example test demonstrating usage
// Create credentials
credentials := []*types.OAuthCredentialSet{
{
ID: "team-account",
ClientID: "client-id-1",
ClientSecret: "client-secret-1",
AccessToken: "access-token-1",
RefreshToken: "refresh-token-1",
ExpiresAt: time.Now().Add(1 * time.Hour),
},
{
ID: "personal-account",
ClientID: "client-id-2",
ClientSecret: "client-secret-2",
AccessToken: "access-token-2",
RefreshToken: "refresh-token-2",
ExpiresAt: time.Now().Add(2 * time.Hour),
},
}
// Create manager
manager := NewOAuthKeyManager("Gemini", credentials, nil)
// Use in operation with automatic failover
ctx := context.Background()
result, usage, err := manager.ExecuteWithFailover(ctx,
func(ctx context.Context, cred *types.OAuthCredentialSet) (string, *types.Usage, error) {
// Make API call with cred.AccessToken
return fmt.Sprintf("API response using %s", cred.ID), &types.Usage{TotalTokens: 100}, nil
})
if err != nil {
fmt.Printf("Error: %v\n", err)
return
}
fmt.Printf("Result: %s, Tokens: %d\n", result, usage.TotalTokens)
func NewOAuthKeyManager ¶
func NewOAuthKeyManager(providerName string, credentials []*types.OAuthCredentialSet, refreshFunc RefreshFunc) *OAuthKeyManager
NewOAuthKeyManager creates a new OAuth key manager with multiple credentials The refreshFunc parameter is a provider-specific function that refreshes OAuth tokens If no refresh is needed, pass NoOpRefreshFunc
func (*OAuthKeyManager) AutoDecommissionExpired ¶
func (m *OAuthKeyManager) AutoDecommissionExpired() ([]string, error)
AutoDecommissionExpired automatically decommissions credentials whose grace period has expired Returns the list of decommissioned credential IDs
func (*OAuthKeyManager) CheckAlerts ¶
func (m *OAuthKeyManager) CheckAlerts() []WebhookEvent
CheckAlerts checks for conditions that should trigger alerts Returns a list of alerts that were generated
func (*OAuthKeyManager) CheckRotationNeeded ¶
func (m *OAuthKeyManager) CheckRotationNeeded() []string
CheckRotationNeeded returns a list of credential IDs that need rotation
func (*OAuthKeyManager) CompleteRotation ¶
func (m *OAuthKeyManager) CompleteRotation(credentialID string) error
CompleteRotation removes a credential that has completed its rotation grace period
func (*OAuthKeyManager) ExecuteWithFailover ¶
func (m *OAuthKeyManager) ExecuteWithFailover( ctx context.Context, operation func(context.Context, *types.OAuthCredentialSet) (string, *types.Usage, error), ) (string, *types.Usage, error)
ExecuteWithFailover attempts an operation with automatic failover to next credential on failure Automatically refreshes tokens if they are expired or expiring soon
func (*OAuthKeyManager) ExportJSON ¶
func (m *OAuthKeyManager) ExportJSON() ([]byte, error)
ExportJSON exports all metrics and health data as JSON
func (*OAuthKeyManager) ExportPrometheus ¶
func (m *OAuthKeyManager) ExportPrometheus() *PrometheusMetrics
ExportPrometheus exports metrics in Prometheus-compatible format
func (*OAuthKeyManager) GetAllMetrics ¶
func (m *OAuthKeyManager) GetAllMetrics() map[string]*CredentialMetrics
GetAllMetrics returns a map of all credential metrics Returns a map of credential ID to metrics snapshot
func (*OAuthKeyManager) GetCredentialHealth ¶
func (m *OAuthKeyManager) GetCredentialHealth(credentialID string) *credentialHealth
GetCredentialHealth returns the health status of a specific credential Returns nil if credential ID not found
func (*OAuthKeyManager) GetCredentialHealthInfo ¶
func (m *OAuthKeyManager) GetCredentialHealthInfo(credentialID string) *CredentialHealthInfo
GetCredentialHealthInfo returns detailed health info for a specific credential
func (*OAuthKeyManager) GetCredentialMetrics ¶
func (m *OAuthKeyManager) GetCredentialMetrics(credentialID string) *CredentialMetrics
GetCredentialMetrics returns the metrics for a specific credential Returns nil if the credential ID is not found
func (*OAuthKeyManager) GetCredentials ¶
func (m *OAuthKeyManager) GetCredentials() []*types.OAuthCredentialSet
GetCredentials returns a copy of the credentials slice Returns clones to prevent external modification
func (*OAuthKeyManager) GetHealthSummary ¶
func (m *OAuthKeyManager) GetHealthSummary() map[string]interface{}
GetHealthSummary returns a comprehensive health summary for all credentials
func (*OAuthKeyManager) GetMonitoringConfig ¶
func (m *OAuthKeyManager) GetMonitoringConfig() *MonitoringConfig
GetMonitoringConfig returns the current monitoring configuration
func (*OAuthKeyManager) GetNextCredential ¶
func (m *OAuthKeyManager) GetNextCredential(ctx context.Context) (*types.OAuthCredentialSet, error)
GetNextCredential returns the next available OAuth credential using round-robin load balancing
func (*OAuthKeyManager) GetRefreshStrategy ¶
func (m *OAuthKeyManager) GetRefreshStrategy() *RefreshStrategy
GetRefreshStrategy returns the current refresh strategy
func (*OAuthKeyManager) GetRotationPolicy ¶
func (m *OAuthKeyManager) GetRotationPolicy() *RotationPolicy
GetRotationPolicy returns the current rotation policy
func (*OAuthKeyManager) GetRotationState ¶
func (m *OAuthKeyManager) GetRotationState(credentialID string) *credentialRotationState
GetRotationState returns the rotation state for a specific credential
func (*OAuthKeyManager) MarkForRotation ¶
func (m *OAuthKeyManager) MarkForRotation(credentialID string, newCredential *types.OAuthCredentialSet) error
MarkForRotation marks a credential for rotation and adds the new credential
func (*OAuthKeyManager) RecordRequest ¶
func (m *OAuthKeyManager) RecordRequest(credentialID string, tokensUsed int64, latency time.Duration, success bool)
RecordRequest records metrics for a single request on the manager
func (*OAuthKeyManager) ReportFailure ¶
func (m *OAuthKeyManager) ReportFailure(credentialID string, err error)
ReportFailure reports that an API call failed with this credential
func (*OAuthKeyManager) ReportSuccess ¶
func (m *OAuthKeyManager) ReportSuccess(credentialID string)
ReportSuccess reports that an API call succeeded with this credential
func (*OAuthKeyManager) SetMonitoringConfig ¶
func (m *OAuthKeyManager) SetMonitoringConfig(config *MonitoringConfig)
SetMonitoringConfig sets the monitoring configuration for the manager
func (*OAuthKeyManager) SetRefreshStrategy ¶
func (m *OAuthKeyManager) SetRefreshStrategy(strategy *RefreshStrategy)
SetRefreshStrategy sets the refresh strategy for the manager
func (*OAuthKeyManager) SetRotationPolicy ¶
func (m *OAuthKeyManager) SetRotationPolicy(policy *RotationPolicy)
SetRotationPolicy sets the rotation policy for the manager
type PrometheusMetrics ¶
type PrometheusMetrics struct {
RequestsTotal map[string]int64 // Counter per credential
SuccessTotal map[string]int64 // Success counter per credential
ErrorsTotal map[string]int64 // Error counter per credential
TokensUsedTotal map[string]int64 // Token usage counter per credential
RefreshesTotal map[string]int64 // Refresh counter per credential
LatencyHistogram map[string][]time.Duration // Latency samples per credential
}
PrometheusMetrics represents metrics in Prometheus-compatible format
type RefreshFunc ¶
type RefreshFunc func(ctx context.Context, cred *types.OAuthCredentialSet) (*types.OAuthCredentialSet, error)
RefreshFunc is a provider-specific function that refreshes OAuth tokens It receives the current credential set and returns an updated credential set with new tokens The provider is responsible for: - Making the token refresh request to the OAuth server - Parsing the response and creating updated credential set - Updating LastRefresh and RefreshCount fields
type RefreshStrategy ¶
type RefreshStrategy struct {
// Buffer configuration
BufferTime time.Duration // Time before expiry to refresh (default: 5 min)
// Adaptive refresh
AdaptiveBuffer bool // Adjust buffer based on usage patterns
MinBuffer time.Duration // Minimum buffer (default: 1 min)
MaxBuffer time.Duration // Maximum buffer (default: 15 min)
// Preemptive refresh
PreemptiveRefresh bool // Refresh before buffer for high-traffic creds
HighTrafficThreshold int64 // Requests/hour to trigger preemptive (default: 100)
}
RefreshStrategy defines how and when OAuth tokens should be refreshed
func AdaptiveRefreshStrategy ¶
func AdaptiveRefreshStrategy() *RefreshStrategy
AdaptiveRefreshStrategy returns a refresh strategy with adaptive buffer enabled
func ConservativeRefreshStrategy ¶
func ConservativeRefreshStrategy() *RefreshStrategy
ConservativeRefreshStrategy returns a strategy that refreshes tokens very early
func DefaultRefreshStrategy ¶
func DefaultRefreshStrategy() *RefreshStrategy
DefaultRefreshStrategy returns the default refresh strategy with sensible defaults
func (*RefreshStrategy) CalculateBufferTime ¶
func (s *RefreshStrategy) CalculateBufferTime(metrics *CredentialMetrics) time.Duration
CalculateBufferTime calculates the appropriate buffer time based on usage metrics
func (*RefreshStrategy) ShouldRefresh ¶
func (s *RefreshStrategy) ShouldRefresh(cred *types.OAuthCredentialSet, metrics *CredentialMetrics) bool
ShouldRefresh determines if a credential should be refreshed based on the strategy
type RotationPolicy ¶
type RotationPolicy struct {
Enabled bool // Enable automatic rotation
RotationInterval time.Duration // How often to rotate (default: 30 days)
GracePeriod time.Duration // Overlap period (default: 7 days)
AutoDecommission bool // Automatically remove old credentials
// Callbacks
OnRotationNeeded func(credentialID string) error // Called when rotation is due
OnDecommission func(credentialID string) error // Called before removal
}
RotationPolicy defines how and when credentials should be rotated
func DefaultRotationPolicy ¶
func DefaultRotationPolicy() *RotationPolicy
DefaultRotationPolicy returns a default rotation policy
type WebhookEvent ¶
type WebhookEvent struct {
Type string `json:"type"` // "refresh", "failure", "rotation", "expiry_warning"
CredentialID string `json:"credential_id"`
Timestamp time.Time `json:"timestamp"`
Message string `json:"message"`
Details map[string]interface{} `json:"details,omitempty"`
}
WebhookEvent represents an event to be sent via webhook