11package main
22
33import (
4+ "context"
45 "flag"
56 "fmt"
67 "log"
78 "net/http"
89 "os"
10+ "os/signal"
911 "path/filepath"
12+ "syscall"
1013 "time"
1114
1215 "github.com/awsl-project/maxx/internal/adapter/client"
@@ -18,9 +21,9 @@ import (
1821 "github.com/awsl-project/maxx/internal/handler"
1922 "github.com/awsl-project/maxx/internal/repository/cached"
2023 "github.com/awsl-project/maxx/internal/repository/sqlite"
21- "github.com/awsl-project/maxx/internal/stats"
2224 "github.com/awsl-project/maxx/internal/router"
2325 "github.com/awsl-project/maxx/internal/service"
26+ "github.com/awsl-project/maxx/internal/stats"
2427 "github.com/awsl-project/maxx/internal/version"
2528 "github.com/awsl-project/maxx/internal/waiter"
2629)
@@ -118,6 +121,23 @@ func main() {
118121 } else if count > 0 {
119122 log .Printf ("Marked %d stale requests as failed" , count )
120123 }
124+ // Also mark stale upstream attempts as failed
125+ if count , err := attemptRepo .MarkStaleAttemptsFailed (); err != nil {
126+ log .Printf ("Warning: Failed to mark stale attempts: %v" , err )
127+ } else if count > 0 {
128+ log .Printf ("Marked %d stale upstream attempts as failed" , count )
129+ }
130+ // Fix legacy failed requests/attempts without end_time
131+ if count , err := proxyRequestRepo .FixFailedRequestsWithoutEndTime (); err != nil {
132+ log .Printf ("Warning: Failed to fix failed requests without end_time: %v" , err )
133+ } else if count > 0 {
134+ log .Printf ("Fixed %d failed requests without end_time" , count )
135+ }
136+ if count , err := attemptRepo .FixFailedAttemptsWithoutEndTime (); err != nil {
137+ log .Printf ("Warning: Failed to fix failed attempts without end_time: %v" , err )
138+ } else if count > 0 {
139+ log .Printf ("Fixed %d failed attempts without end_time" , count )
140+ }
121141
122142 // Create cached repositories
123143 cachedProviderRepo := cached .NewProviderRepository (providerRepo )
@@ -228,6 +248,7 @@ func main() {
228248 responseModelRepo ,
229249 * addr ,
230250 r , // Router implements ProviderAdapterRefresher interface
251+ wsHub ,
231252 )
232253
233254 // Create backup service
@@ -257,8 +278,12 @@ func main() {
257278 log .Println ("Proxy token authentication is enabled" )
258279 }
259280
281+ // Create request tracker for graceful shutdown
282+ requestTracker := core .NewRequestTracker ()
283+
260284 // Create handlers
261285 proxyHandler := handler .NewProxyHandler (clientAdapter , exec , cachedSessionRepo , tokenAuthMiddleware )
286+ proxyHandler .SetRequestTracker (requestTracker )
262287 adminHandler := handler .NewAdminHandler (adminService , backupService , logPath )
263288 authHandler := handler .NewAuthHandler (authMiddleware )
264289 antigravityHandler := handler .NewAntigravityHandler (adminService , antigravityQuotaRepo , wsHub )
@@ -309,22 +334,55 @@ func main() {
309334 // Wrap with logging middleware
310335 loggedMux := handler .LoggingMiddleware (mux )
311336
312- // Start server
337+ // Create HTTP server
338+ server := & http.Server {
339+ Addr : * addr ,
340+ Handler : loggedMux ,
341+ }
342+
343+ // Start server in goroutine
313344 log .Printf ("Starting Maxx server %s on %s" , version .Info (), * addr )
314345 log .Printf ("Data directory: %s" , dataDirPath )
315- log .Printf (" Database: %s" , dbPath )
316- log .Printf (" Log file: %s" , logPath )
317- log .Printf ("Admin API: http://localhost%s/api/admin/" , * addr )
318- log .Printf ("WebSocket: ws://localhost%s/ws" , * addr )
319- log .Printf ("Proxy endpoints:" )
320- log .Printf (" Claude: http://localhost%s/v1/messages" , * addr )
321- log .Printf (" OpenAI: http://localhost%s/v1/chat/completions" , * addr )
322- log .Printf (" Codex: http://localhost%s/v1/responses" , * addr )
323- log .Printf (" Gemini: http://localhost%s/v1beta/models/{model}:generateContent" , * addr )
324- log .Printf ("Project proxy: http://localhost%s/{project-slug}/v1/messages (etc.)" , * addr )
325-
326- if err := http .ListenAndServe (* addr , loggedMux ); err != nil {
327- log .Printf ("Server error: %v" , err )
328- os .Exit (1 )
346+
347+ go func () {
348+ if err := server .ListenAndServe (); err != nil && err != http .ErrServerClosed {
349+ log .Printf ("Server error: %v" , err )
350+ os .Exit (1 )
351+ }
352+ }()
353+
354+ // Wait for interrupt signal (SIGINT or SIGTERM)
355+ sigCh := make (chan os.Signal , 1 )
356+ signal .Notify (sigCh , syscall .SIGINT , syscall .SIGTERM )
357+ sig := <- sigCh
358+ log .Printf ("Received signal %v, initiating graceful shutdown..." , sig )
359+
360+ // Step 1: Wait for active proxy requests to complete
361+ activeCount := requestTracker .ActiveCount ()
362+ if activeCount > 0 {
363+ log .Printf ("Waiting for %d active proxy requests to complete..." , activeCount )
364+ completed := requestTracker .GracefulShutdown (core .GracefulShutdownTimeout )
365+ if ! completed {
366+ log .Printf ("Graceful shutdown timeout, some requests may be interrupted" )
367+ } else {
368+ log .Printf ("All proxy requests completed successfully" )
369+ }
370+ } else {
371+ // Mark as shutting down to reject new requests
372+ requestTracker .GracefulShutdown (0 )
373+ log .Printf ("No active proxy requests" )
374+ }
375+
376+ // Step 2: Shutdown HTTP server
377+ shutdownCtx , cancel := context .WithTimeout (context .Background (), core .HTTPShutdownTimeout )
378+ defer cancel ()
379+
380+ if err := server .Shutdown (shutdownCtx ); err != nil {
381+ log .Printf ("HTTP server graceful shutdown failed: %v, forcing close" , err )
382+ if closeErr := server .Close (); closeErr != nil {
383+ log .Printf ("Force close error: %v" , closeErr )
384+ }
329385 }
386+
387+ log .Printf ("Server stopped" )
330388}
0 commit comments