goclaw/internal/oauth/token.go

package oauth

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/store"
)

const (
	// DefaultProviderName is the provider name for ChatGPT OAuth.
	DefaultProviderName = "openai-codex"

	// refreshTokenSecretKey is the config_secrets key for the refresh token.
	refreshTokenSecretKey = "oauth.openai-codex.refresh_token"

	// refreshMargin is how early before expiry we refresh the token.
	refreshMargin = 5 * time.Minute
)

// OAuthSettings is stored in llm_providers.settings JSONB (non-sensitive metadata).
type OAuthSettings struct {
	ExpiresAt int64  `json:"expires_at"` // unix timestamp
	Scopes    string `json:"scopes,omitempty"`
}

// DBTokenSource provides a valid access token backed by the llm_providers + config_secrets tables.
// Implements providers.TokenSource.
type DBTokenSource struct {
	providerStore store.ProviderStore
	secretsStore  store.ConfigSecretsStore
	providerName  string

	mu          sync.Mutex
	cachedToken string
	expiresAt   time.Time
}

// NewDBTokenSource creates a DB-backed token source.
func NewDBTokenSource(provStore store.ProviderStore, secretsStore store.ConfigSecretsStore, providerName string) *DBTokenSource {
	return &DBTokenSource{
		providerStore: provStore,
		secretsStore:  secretsStore,
		providerName:  providerName,
	}
}

// Token returns a valid access token, refreshing if expired or about to expire.
func (ts *DBTokenSource) Token() (string, error) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	// Use cached token if still valid
	if ts.cachedToken != "" && time.Until(ts.expiresAt) > refreshMargin {
		return ts.cachedToken, nil
	}

	ctx := context.Background()

	// Load from DB if not cached
	if ts.cachedToken == "" {
		p, err := ts.providerStore.GetProviderByName(ctx, ts.providerName)
		if err != nil {
			return "", fmt.Errorf("load oauth provider %q: %w", ts.providerName, err)
		}
		ts.cachedToken = p.APIKey

		var settings OAuthSettings
		if len(p.Settings) > 0 {
			_ = json.Unmarshal(p.Settings, &settings)
		}
		if settings.ExpiresAt > 0 {
			ts.expiresAt = time.Unix(settings.ExpiresAt, 0)
		}
	}

	// Refresh if expired or expiring soon
	if time.Until(ts.expiresAt) < refreshMargin {
		if err := ts.refresh(ctx); err != nil {
			// If refresh fails but we still have a token, return it (might still work)
			if ts.cachedToken != "" {
				slog.Warn("oauth token refresh failed, using existing token", "error", err)
				return ts.cachedToken, nil
			}
			return "", fmt.Errorf("refresh oauth token: %w", err)
		}
	}

	return ts.cachedToken, nil
}

// refresh gets the refresh token from config_secrets, calls RefreshOpenAIToken, and updates DB.
func (ts *DBTokenSource) refresh(ctx context.Context) error {
	refreshToken, err := ts.secretsStore.Get(ctx, refreshTokenSecretKey)
	if err != nil {
		return fmt.Errorf("get refresh token: %w", err)
	}

	slog.Info("refreshing OpenAI OAuth token")
	newToken, err := RefreshOpenAIToken(refreshToken)
	if err != nil {
		return err
	}

	// Update cached values
	ts.cachedToken = newToken.AccessToken
	ts.expiresAt = time.Now().Add(time.Duration(newToken.ExpiresIn) * time.Second)

	// Update provider api_key (access token) in DB
	p, err := ts.providerStore.GetProviderByName(ctx, ts.providerName)
	if err != nil {
		return fmt.Errorf("get provider for update: %w", err)
	}

	settings := OAuthSettings{
		ExpiresAt: ts.expiresAt.Unix(),
	}
	settingsJSON, _ := json.Marshal(settings)

	if err := ts.providerStore.UpdateProvider(ctx, p.ID, map[string]any{
		"api_key":  newToken.AccessToken,
		"settings": json.RawMessage(settingsJSON),
	}); err != nil {
		slog.Warn("failed to persist refreshed access token", "error", err)
	}

	// Update refresh token if a new one was issued
	if newToken.RefreshToken != "" {
		if err := ts.secretsStore.Set(ctx, refreshTokenSecretKey, newToken.RefreshToken); err != nil {
			slog.Warn("failed to persist new refresh token", "error", err)
		}
	}

	return nil
}

// SaveOAuthResult persists OAuth tokens after a successful exchange.
// Creates or updates the provider in llm_providers and stores refresh token in config_secrets.
// Returns the provider ID.
func (ts *DBTokenSource) SaveOAuthResult(ctx context.Context, tokenResp *OpenAITokenResponse) (uuid.UUID, error) {
	expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
	settings := OAuthSettings{
		ExpiresAt: expiresAt.Unix(),
		Scopes:    tokenResp.Scope,
	}
	settingsJSON, _ := json.Marshal(settings)

	// Update cache
	ts.mu.Lock()
	ts.cachedToken = tokenResp.AccessToken
	ts.expiresAt = expiresAt
	ts.mu.Unlock()

	// Check if provider already exists
	existing, err := ts.providerStore.GetProviderByName(ctx, ts.providerName)
	if err == nil {
		// Update existing provider
		if err := ts.providerStore.UpdateProvider(ctx, existing.ID, map[string]any{
			"api_key":  tokenResp.AccessToken,
			"settings": json.RawMessage(settingsJSON),
			"enabled":  true,
		}); err != nil {
			return uuid.Nil, fmt.Errorf("update provider: %w", err)
		}

		// Save refresh token
		if tokenResp.RefreshToken != "" {
			if err := ts.secretsStore.Set(ctx, refreshTokenSecretKey, tokenResp.RefreshToken); err != nil {
				return uuid.Nil, fmt.Errorf("save refresh token: %w", err)
			}
		}

		return existing.ID, nil
	}

	// Create new provider
	p := &store.LLMProviderData{
		Name:         ts.providerName,
		DisplayName:  "ChatGPT (OAuth)",
		ProviderType: store.ProviderChatGPTOAuth,
		APIBase:      "https://chatgpt.com/backend-api",
		APIKey:       tokenResp.AccessToken,
		Enabled:      true,
		Settings:     settingsJSON,
	}
	if err := ts.providerStore.CreateProvider(ctx, p); err != nil {
		return uuid.Nil, fmt.Errorf("create provider: %w", err)
	}

	// Save refresh token
	if tokenResp.RefreshToken != "" {
		if err := ts.secretsStore.Set(ctx, refreshTokenSecretKey, tokenResp.RefreshToken); err != nil {
			return uuid.Nil, fmt.Errorf("save refresh token: %w", err)
		}
	}

	return p.ID, nil
}

// Delete removes the OAuth provider from DB and its refresh token from config_secrets.
func (ts *DBTokenSource) Delete(ctx context.Context) error {
	ts.mu.Lock()
	ts.cachedToken = ""
	ts.expiresAt = time.Time{}
	ts.mu.Unlock()

	// Delete refresh token from config_secrets
	_ = ts.secretsStore.Delete(ctx, refreshTokenSecretKey)

	// Delete provider from llm_providers
	p, err := ts.providerStore.GetProviderByName(ctx, ts.providerName)
	if err != nil {
		return nil // already gone
	}
	return ts.providerStore.DeleteProvider(ctx, p.ID)
}

// Exists checks if an OAuth provider exists and has a valid token.
func (ts *DBTokenSource) Exists(ctx context.Context) bool {
	p, err := ts.providerStore.GetProviderByName(ctx, ts.providerName)
	return err == nil && p.APIKey != ""
}