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) ExecuteWithFailoverMessage(ctx context.Context, ...) (types.ChatMessage, *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) SetCredentialProvider(provider types.CredentialProvider)
- 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) ExecuteWithFailoverMessage ¶ added in v1.0.11
func (m *OAuthKeyManager) ExecuteWithFailoverMessage( ctx context.Context, operation func(context.Context, *types.OAuthCredentialSet) (types.ChatMessage, *types.Usage, error), ) (types.ChatMessage, *types.Usage, error)
ExecuteWithFailoverMessage attempts an operation with automatic failover (message variant) Returns full ChatMessage to preserve tool calls
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) SetCredentialProvider ¶ added in v1.0.25
func (m *OAuthKeyManager) SetCredentialProvider(provider types.CredentialProvider)
SetCredentialProvider sets a dynamic credential provider for this manager When set, credentials will be fetched from the provider on each request rather than using the cached credentials passed to NewOAuthKeyManager. This ensures that any external token refresh is immediately visible.
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