anthropic-proxy/main.go

package main

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/fujin/anthropic-proxy/internal/auth"
	"github.com/fujin/anthropic-proxy/internal/config"
	"github.com/fujin/anthropic-proxy/internal/embedded"
	"github.com/fujin/anthropic-proxy/internal/logging"
	"github.com/fujin/anthropic-proxy/internal/proxy"
	"github.com/fujin/anthropic-proxy/internal/ratelimit"
	"github.com/fujin/anthropic-proxy/internal/server"
	"github.com/fujin/anthropic-proxy/internal/telemetry"
	"github.com/rs/zerolog/log"
)

func initCredential() (*auth.Credential, error) {
	creds, err := auth.LoadDefaultCredentials()
	if err != nil {
		return nil, fmt.Errorf("load credentials: %w", err)
	}

	var cred *auth.Credential
	if len(creds) > 0 {
		cred = creds[0]
		// If token is expired, try refresh first
		if !cred.ExpiresAt.IsZero() && time.Now().After(cred.ExpiresAt) {
			log.Info().Msg("token expired, attempting refresh")
			refreshCtx, refreshCancel := context.WithTimeout(context.Background(), 15*time.Second)
			refreshErr := auth.RefreshToken(refreshCtx, cred)
			refreshCancel()
			if refreshErr != nil {
				log.Warn().Err(refreshErr).Msg("refresh failed, initiating login")
				cred = nil // fall through to login
			} else {
				log.Info().Msg("token refreshed")
			}
		}
	}

	if cred == nil {
		fi, statErr := os.Stdin.Stat()
		if statErr == nil && (fi.Mode()&os.ModeCharDevice) == 0 {
			return nil, fmt.Errorf("no valid credentials found; run the proxy interactively for initial login")
		}
		log.Info().Msg("no credentials found, starting OAuth login")
		cred, err = auth.Login(context.Background())
		if err != nil {
			return nil, fmt.Errorf("login failed: %w", err)
		}
	}

	log.Info().Str("credential", cred.Email).Msg("credential loaded")
	return cred, nil
}

func initEmbedded(cfg *config.Config) (cleanup func(), err error) {
	if !cfg.Telemetry.Embedded.Enabled {
		return func() {}, nil
	}

	var cleanups []func()

	vm := embedded.NewVM(cfg.Telemetry.Embedded, cfg.Port)
	if err := vm.Start(); err != nil {
		log.Error().Err(err).Msg("failed to start victoria-metrics")
	} else {
		cleanups = append(cleanups, vm.Stop)
	}

	perses := embedded.NewPerses(cfg.Telemetry.Embedded, cfg.Port)
	if err := perses.Start(); err != nil {
		log.Error().Err(err).Msg("failed to start perses")
	} else {
		cleanups = append(cleanups, perses.Stop)
	}

	return func() {
		for i := len(cleanups) - 1; i >= 0; i-- {
			cleanups[i]()
		}
	}, nil
}

func run() error {
	cfg, err := config.Load("config.yaml")
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}

	// Create usage tracker (started later once credential is loaded)
	var credForTracker *auth.Credential
	tracker := ratelimit.NewTracker(func() string {
		if credForTracker == nil {
			return ""
		}
		return credForTracker.Token()
	})

	// Initialize telemetry (metrics always active; OTLP export when endpoint set)
	telemetryShutdown, logBridge, metricsHandler, err := telemetry.Setup(context.Background(), cfg.Telemetry, tracker)
	if err != nil {
		return fmt.Errorf("telemetry setup: %w", err)
	}
	defer telemetryShutdown(context.Background())

	var extraWriters []io.Writer
	if logBridge != nil {
		extraWriters = append(extraWriters, logBridge)
	}

	logging.Setup(cfg.Logging, extraWriters...)

	cred, err := initCredential()
	if err != nil {
		return err
	}

	credForTracker = cred

	pool := auth.NewPool([]*auth.Credential{cred})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	pool.RefreshExpiring(context.Background())
	auth.StartBackgroundRefresh(ctx, pool)
	tracker.Start(ctx)

	var profile *proxy.SniffedProfile
	if cfg.ClaudeBinary != "" {
		log.Info().Str("binary", cfg.ClaudeBinary).Msg("sniffing claude-code")
		profile, err = proxy.SniffClaudeCode(cfg.ClaudeBinary)
		if err != nil {
			log.Warn().Err(err).Msg("sniff failed, using defaults")
		}
	}

	embeddedCleanup, err := initEmbedded(cfg)
	if err != nil {
		return err
	}
	defer embeddedCleanup()

	log.Info().Int("port", cfg.Port).Msg("starting server")
	srv := server.New(cfg, pool, profile, tracker, metricsHandler)

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-quit
		log.Info().Msg("shutting down")
		cancel()

		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer shutdownCancel()

		if err := srv.Shutdown(shutdownCtx); err != nil {
			log.Error().Err(err).Msg("shutdown error")
		}
	}()

	if err := srv.Start(); err != nil && err != http.ErrServerClosed {
		return err
	}

	return nil
}

func main() {
	if err := run(); err != nil {
		log.Error().Err(err).Msg("fatal error")
		os.Exit(1)
	}
}