CLIProxyAPI/internal/cmd/run.go

// Package cmd provides the main service execution functionality for the CLIProxyAPI.
// It contains the core logic for starting and managing the API proxy service,
// including authentication client management, server initialization, and graceful shutdown handling.
// The package handles loading authentication tokens, creating client pools, starting the API server,
// and monitoring configuration changes through file watchers.
package cmd

import (
	"context"
	"encoding/json"
	"github.com/luispater/CLIProxyAPI/internal/api"
	"github.com/luispater/CLIProxyAPI/internal/auth"
	"github.com/luispater/CLIProxyAPI/internal/client"
	"github.com/luispater/CLIProxyAPI/internal/config"
	"github.com/luispater/CLIProxyAPI/internal/util"
	"github.com/luispater/CLIProxyAPI/internal/watcher"
	log "github.com/sirupsen/logrus"
	"io/fs"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

// StartService initializes and starts the main API proxy service.
// It loads all available authentication tokens, creates a pool of clients,
// starts the API server, and handles graceful shutdown signals.
func StartService(cfg *config.Config, configPath string) {
	// Create a pool of API clients, one for each token file found.
	cliClients := make([]*client.Client, 0)
	err := filepath.Walk(cfg.AuthDir, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Process only JSON files in the auth directory.
		if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") {
			log.Debugf("Loading token from: %s", path)
			f, errOpen := os.Open(path)
			if errOpen != nil {
				return errOpen
			}
			defer func() {
				_ = f.Close()
			}()

			// Decode the token storage file.
			var ts auth.TokenStorage
			if err = json.NewDecoder(f).Decode(&ts); err == nil {
				// For each valid token, create an authenticated client.
				clientCtx := context.Background()
				log.Info("Initializing authentication for token...")
				httpClient, errGetClient := auth.GetAuthenticatedClient(clientCtx, &ts, cfg)
				if errGetClient != nil {
					// Log fatal will exit, but we return the error for completeness.
					log.Fatalf("failed to get authenticated client for token %s: %v", path, errGetClient)
					return errGetClient
				}
				log.Info("Authentication successful.")

				// Add the new client to the pool.
				cliClient := client.NewClient(httpClient, &ts, cfg)
				cliClients = append(cliClients, cliClient)
			}
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error walking auth directory: %v", err)
	}

	if len(cfg.GlAPIKey) > 0 {
		for i := 0; i < len(cfg.GlAPIKey); i++ {
			httpClient, errSetProxy := util.SetProxy(cfg, &http.Client{})
			if errSetProxy != nil {
				log.Fatalf("set proxy failed: %v", errSetProxy)
			}

			log.Debug("Initializing with Generative Language API key...")
			cliClient := client.NewClient(httpClient, nil, cfg, cfg.GlAPIKey[i])
			cliClients = append(cliClients, cliClient)
		}
	}

	// Create and start the API server with the pool of clients.
	apiServer := api.NewServer(cfg, cliClients)
	log.Infof("Starting API server on port %d", cfg.Port)

	// Start the API server in a goroutine so it doesn't block the main thread
	go func() {
		if err = apiServer.Start(); err != nil {
			log.Fatalf("API server failed to start: %v", err)
		}
	}()

	// Give the server a moment to start up
	time.Sleep(100 * time.Millisecond)
	log.Info("API server started successfully")

	// Setup file watcher for config and auth directory changes
	fileWatcher, errNewWatcher := watcher.NewWatcher(configPath, cfg.AuthDir, func(newClients []*client.Client, newCfg *config.Config) {
		// Update the API server with new clients and configuration
		apiServer.UpdateClients(newClients, newCfg)
	})
	if errNewWatcher != nil {
		log.Fatalf("failed to create file watcher: %v", errNewWatcher)
	}

	// Set initial state for the watcher
	fileWatcher.SetConfig(cfg)
	fileWatcher.SetClients(cliClients)

	// Start the file watcher
	watcherCtx, watcherCancel := context.WithCancel(context.Background())
	if errStartWatcher := fileWatcher.Start(watcherCtx); errStartWatcher != nil {
		log.Fatalf("failed to start file watcher: %v", errStartWatcher)
	}
	log.Info("file watcher started for config and auth directory changes")

	defer func() {
		watcherCancel()
		errStopWatcher := fileWatcher.Stop()
		if errStopWatcher != nil {
			log.Errorf("error stopping file watcher: %v", errStopWatcher)
		}
	}()

	// Set up a channel to listen for OS signals for graceful shutdown.
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Main loop to wait for shutdown signal.
	for {
		select {
		case <-sigChan:
			log.Debugf("Received shutdown signal. Cleaning up...")

			// Create a context with a timeout for the shutdown process.
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			_ = cancel

			// Stop the API server gracefully.
			if err = apiServer.Stop(ctx); err != nil {
				log.Debugf("Error stopping API server: %v", err)
			}

			log.Debugf("Cleanup completed. Exiting...")
			os.Exit(0)
		case <-time.After(5 * time.Second):
			// This case is currently empty and acts as a periodic check.
			// It could be used for periodic tasks in the future.
		}
	}
}