feat: remove anti-pattern for graceful shutdown

Author: 35C4n0r

cmd/server/server.go

@@ -181,8 +181,22 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		fmt.Println(srv.GetOpenAPI())
 		return nil
 	}
-	handleSignals(ctx, logger, srv, process, pidFile)
+
+	// Create a context for graceful shutdown
+	gracefulCtx, gracefulCancel := context.WithCancel(ctx)
+	defer gracefulCancel()
+
+	// Setup signal handlers (they will call gracefulCancel)
+	handleSignals(gracefulCtx, gracefulCancel, logger, srv)
+
+	// Setup PID file cleanup
+	if pidFile != "" {
+		defer cleanupPIDFile(pidFile, logger)
+	}
+
 	logger.Info("Starting server on port", "port", port)
+
+	// Monitor process exit
 	processExitCh := make(chan error, 1)
 	go func() {
 		defer close(processExitCh)
@@ -193,18 +207,52 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 				processExitCh <- xerrors.Errorf("failed to wait for process: %w", err)
 			}
 		}
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		if err := srv.Stop(shutdownCtx); err != nil {
-			logger.Error("Failed to stop server after process exit", "error", err)
+
+		select {
+		case <-gracefulCtx.Done():
+		default:
+			gracefulCancel()
+		}
+	}()
+
+	// Start the server
+	serverErrCh := make(chan error, 1)
+	go func() {
+		defer close(serverErrCh)
+		if err := srv.Start(); err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, http.ErrServerClosed) {
+			serverErrCh <- err
 		}
 	}()
-	if err := srv.Start(); err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, http.ErrServerClosed) {
-		return xerrors.Errorf("failed to start server: %w", err)
+
+	select {
+	case err := <-serverErrCh:
+		if err != nil {
+			return xerrors.Errorf("failed to start server: %w", err)
+		}
+	case <-gracefulCtx.Done():
 	}
+
+	if err := srv.SaveState("shutdown"); err != nil {
+		logger.Error("Failed to save state during shutdown", "error", err)
+	}
+
+	// Stop the HTTP server
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := srv.Stop(shutdownCtx); err != nil {
+		logger.Error("Failed to stop HTTP server", "error", err)
+	}
+
+	// Close the process
+	if err := process.Close(logger, 5*time.Second); err != nil {
+		logger.Error("Failed to close process cleanly", "error", err)
+	}
+
 	select {
 	case err := <-processExitCh:
-		return xerrors.Errorf("agent exited with error: %w", err)
+		if err != nil {
+			return xerrors.Errorf("agent exited with error: %w", err)
+		}
 	default:
 	}
 	return nil

cmd/server/signals.go

@@ -10,20 +10,20 @@ import (
 	"syscall"
 
 	"github.com/coder/agentapi/lib/httpapi"
-	"github.com/coder/agentapi/lib/termexec"
 )
 
 // handleSignals sets up signal handlers for:
-// - SIGTERM, SIGINT, SIGHUP: save conversation state, stop server, then close the process
+// - SIGTERM, SIGINT, SIGHUP: trigger graceful shutdown by canceling the context
 // - SIGUSR1: save conversation state without exiting
-func handleSignals(ctx context.Context, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process, pidFile string) {
+func handleSignals(ctx context.Context, cancel context.CancelFunc, logger *slog.Logger, srv *httpapi.Server) {
 	// Handle shutdown signals (SIGTERM, SIGINT, SIGHUP)
 	shutdownCh := make(chan os.Signal, 1)
 	signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGINT)
 	go func() {
 		defer signal.Stop(shutdownCh)
 		sig := <-shutdownCh
-		performGracefulShutdown(sig, logger, srv, process, pidFile)
+		logger.Info("Received shutdown signal", "signal", sig)
+		cancel()
 	}()
 
 	// Handle SIGUSR1 for save without exit

cmd/server/signals_unix.go

@@ -10,18 +10,18 @@ import (
 	"syscall"
 
 	"github.com/coder/agentapi/lib/httpapi"
-	"github.com/coder/agentapi/lib/termexec"
 )
 
 // handleSignals sets up signal handlers for Windows.
 // Only handles SIGTERM and SIGINT (SIGHUP and SIGUSR1 don't exist on Windows).
-func handleSignals(ctx context.Context, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process, pidFile string) {
+func handleSignals(ctx context.Context, cancel context.CancelFunc, logger *slog.Logger, srv *httpapi.Server) {
 	// Handle shutdown signals (SIGTERM, SIGINT only on Windows)
 	shutdownCh := make(chan os.Signal, 1)
 	signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		defer signal.Stop(shutdownCh)
 		sig := <-shutdownCh
-		performGracefulShutdown(sig, logger, srv, process, pidFile)
+		logger.Info("Received shutdown signal", "signal", sig)
+		cancel()
 	}()
 }

cmd/server/signals_windows.go

@@ -49,6 +49,8 @@ type Server struct {
 	chatBasePath string
 	tempDir      string
 	clock        quartz.Clock
+	shutdownCtx  context.Context
+	shutdown     context.CancelFunc
 }
 
 func (s *Server) NormalizeSchema(schema any) any {
@@ -275,6 +277,8 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 	}
 	logger.Info("Created temporary directory for uploads", "tempDir", tempDir)
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	s := &Server{
 		router:       router,
 		api:          api,
@@ -287,6 +291,8 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 		chatBasePath: strings.TrimSuffix(config.ChatBasePath, "/"),
 		tempDir:      tempDir,
 		clock:        config.Clock,
+		shutdownCtx:  ctx,
+		shutdown:     cancel,
 	}
 
 	// Register API routes
@@ -514,6 +520,7 @@ func (s *Server) uploadFiles(ctx context.Context, input *struct {
 func (s *Server) subscribeEvents(ctx context.Context, input *struct{}, send sse.Sender) {
 	subscriberId, ch, stateEvents := s.emitter.Subscribe()
 	defer s.emitter.Unsubscribe(subscriberId)
+
 	s.logger.Info("New subscriber", "subscriberId", subscriberId)
 	for _, event := range stateEvents {
 		if event.Type == EventTypeScreenUpdate {
@@ -539,6 +546,9 @@ func (s *Server) subscribeEvents(ctx context.Context, input *struct{}, send sse.
 				s.logger.Error("Failed to send event", "subscriberId", subscriberId, "error", err)
 				return
 			}
+		case <-s.shutdownCtx.Done():
+			s.logger.Info("Server stop initiated, unsubscribing.", "subscriberId", subscriberId)
+			return
 		case <-ctx.Done():
 			s.logger.Info("Context done", "subscriberId", subscriberId)
 			return
@@ -573,6 +583,9 @@ func (s *Server) subscribeScreen(ctx context.Context, input *struct{}, send sse.
 				s.logger.Error("Failed to send screen event", "subscriberId", subscriberId, "error", err)
 				return
 			}
+		case <-s.shutdownCtx.Done():
+			s.logger.Info("Server stop initiated, unsubscribing.", "subscriberId", subscriberId)
+			return
 		case <-ctx.Done():
 			s.logger.Info("Screen context done", "subscriberId", subscriberId)
 			return
@@ -595,6 +608,8 @@ func (s *Server) Start() error {
 func (s *Server) Stop(ctx context.Context) error {
 	var err error
 	s.stopOnce.Do(func() {
+		s.shutdown()
+
 		// Clean up temporary directory
 		s.cleanupTempDir()