oauthmanager

package
v1.0.29 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Dec 9, 2025 License: MIT Imports: 9 Imported by: 0

README

OAuth Manager

Enterprise-grade OAuth credential management for AI providers.

Documentation

Complete documentation is available at:

Quick Start

import "github.com/cecil-the-coder/ai-provider-kit/pkg/oauthmanager"

// Create manager with multiple credentials
manager := oauthmanager.NewOAuthKeyManager("Gemini", credentials, refreshFunc)

// Use with automatic failover and token refresh
result, usage, err := manager.ExecuteWithFailover(ctx,
    func(ctx context.Context, cred *types.OAuthCredentialSet) (string, *types.Usage, error) {
        return makeAPICall(ctx, cred.AccessToken)
    })

See the full documentation for detailed usage, advanced features, and examples.

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

Examples

Constants

View Source
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

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

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL