Files
2026-10-23 18:35:44 +08:00

71 lines
1.8 KiB
Go

package main
import (
"ai-router/internal/config"
"ai-router/internal/server"
"context"
"embed"
"flag"
"io/fs"
"log"
"net/http"
"os"
"os/signal"
"syscall"
"time"
)
//go:embed all:web/dist
var frontend embed.FS
func main() {
path := flag.String("config", "config.yaml", "YAML configuration path")
flag.Parse()
cfg, err := config.Load(*path)
if err != nil {
log.Fatal(err)
}
assets, err := fs.Sub(frontend, "web/dist")
if err != nil {
log.Fatal(err)
}
app, err := server.New(cfg.Storage.Path, cfg.Admin.Password, assets)
if err != nil {
log.Fatal(err)
}
defer func() {
if err := app.Close(); err != nil {
log.Print("AI Router shutdown: failed to persist all request logs or close database")
}
}()
if cfg.Server.PublicOrigin != "" {
if err := app.SetPublicOrigin(cfg.Server.PublicOrigin); err != nil {
log.Fatal(err)
}
}
if err := app.SetTrustedProxies(cfg.Server.TrustedProxies); err != nil {
log.Fatal(err)
}
if err := app.SetLimits(cfg.Limits.RequestsPerMinute, cfg.Limits.PerKeyConcurrency, cfg.Limits.GlobalConcurrency); err != nil {
log.Fatal(err)
}
app.SetLogRetention(cfg.Storage.LogRetentionDays)
srv := &http.Server{Addr: cfg.Server.Listen, Handler: app, ReadTimeout: 30 * time.Second, ReadHeaderTimeout: 10 * time.Second, IdleTimeout: 90 * time.Second, MaxHeaderBytes: 1 << 20}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
done := make(chan error, 1)
go func() { log.Printf("AI Router listening on %s", cfg.Server.Listen); done <- srv.ListenAndServe() }()
select {
case err := <-done:
if err != http.ErrServerClosed {
log.Printf("server: %v", err)
}
case <-ctx.Done():
shutdown, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
if err := srv.Shutdown(shutdown); err != nil {
_ = srv.Close()
}
}
}