feat: implement state persistence (#177)

Co-authored-by: Cian Johnston <cian@coder.com>
Co-authored-by: Mathias Fredriksson <mafredri@gmail.com>

Author: 3 people

chat/src/app/layout.tsx

@@ -29,7 +29,7 @@ export default function RootLayout({
           disableTransitionOnChange
         >
           {children}
-          <Toaster richColors />
+          <Toaster richColors closeButton />
         </ThemeProvider>
       </body>
     </html>

chat/src/components/chat-provider.tsx

@@ -36,6 +36,12 @@ interface StatusChangeEvent {
   agent_type: string;
 }
 
+interface ErrorEventData {
+  message: string;
+  level: string;
+  time: string;
+}
+
 interface APIErrorDetail {
   location: string;
   message: string;
@@ -215,6 +221,25 @@ export function ChatProvider({ children }: PropsWithChildren) {
         setAgentType(data.agent_type === "" ? "unknown" : data.agent_type as AgentType);
       });
 
+      // Handle agent error events
+      eventSource.addEventListener("agent_error", (event) => {
+        const messageEvent = event as MessageEvent;
+        try {
+          const data: ErrorEventData = JSON.parse(messageEvent.data);
+
+          // Display error as toast notification that persists until manually dismissed
+          if (data.level === "error") {
+            toast.error(data.message, { duration: Infinity });
+          } else if (data.level === "warning") {
+            toast.warning(data.message, { duration: Infinity });
+          } else {
+            toast.info(data.message, { duration: Infinity });
+          }
+        } catch (e) {
+          console.error("Failed to parse agent_error event data:", e);
+        }
+      });
+
       // Handle connection open (server is online)
       eventSource.onopen = () => {
         // Connection is established, but we'll wait for status_change event

cmd/server/process_unix.go

@@ -0,0 +1,19 @@
+//go:build unix
+
+package server
+
+import (
+	"errors"
+	"os"
+	"syscall"
+)
+
+// isProcessRunning checks if a process with the given PID is running.
+func isProcessRunning(pid int) bool {
+	process, err := os.FindProcess(pid)
+	if err != nil {
+		return false
+	}
+	err = process.Signal(syscall.Signal(0))
+	return err == nil || errors.Is(err, syscall.EPERM)
+}

cmd/server/process_windows.go

@@ -0,0 +1,10 @@
+//go:build windows
+
+package server
+
+// isProcessRunning checks if a process with the given PID is running.
+// On Windows, Signal(0) is not supported, so this always returns false.
+// PID file liveness detection is best-effort on this platform.
+func isProcessRunning(_ int) bool {
+	return false
+}

cmd/server/server.go

@@ -8,9 +8,13 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sort"
+	"strconv"
 	"strings"
+	"time"
 
+	"github.com/coder/agentapi/lib/screentracker"
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -104,9 +108,51 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		}
 	}
 
-	printOpenAPI := viper.GetBool(FlagPrintOpenAPI)
+	// Get the variables related to state management
+	stateFile := viper.GetString(FlagStateFile)
+	loadState := false
+	saveState := false
+
+	// Validate state file configuration
+	if stateFile != "" {
+		if !viper.IsSet(FlagLoadState) {
+			loadState = true
+		} else {
+			loadState = viper.GetBool(FlagLoadState)
+		}
+
+		if !viper.IsSet(FlagSaveState) {
+			saveState = true
+		} else {
+			saveState = viper.GetBool(FlagSaveState)
+		}
+	} else {
+		if viper.IsSet(FlagLoadState) && viper.GetBool(FlagLoadState) {
+			return xerrors.Errorf("--load-state requires --state-file to be set")
+		}
+		if viper.IsSet(FlagSaveState) && viper.GetBool(FlagSaveState) {
+			return xerrors.Errorf("--save-state requires --state-file to be set")
+		}
+	}
+
 	experimentalACP := viper.GetBool(FlagExperimentalACP)
 
+	if experimentalACP && (saveState || loadState) {
+		return xerrors.Errorf("ACP mode doesn't support state persistence")
+	}
+
+	pidFile := viper.GetString(FlagPidFile)
+
+	// Write PID file if configured
+	if pidFile != "" {
+		if err := writePIDFile(pidFile, logger); err != nil {
+			return xerrors.Errorf("failed to write PID file: %w", err)
+		}
+		defer cleanupPIDFile(pidFile, logger)
+	}
+
+	printOpenAPI := viper.GetBool(FlagPrintOpenAPI)
+
 	if printOpenAPI && experimentalACP {
 		return xerrors.Errorf("flags --%s and --%s are mutually exclusive", FlagPrintOpenAPI, FlagExperimentalACP)
 	}
@@ -154,33 +200,45 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		AllowedHosts:   viper.GetStringSlice(FlagAllowedHosts),
 		AllowedOrigins: viper.GetStringSlice(FlagAllowedOrigins),
 		InitialPrompt:  initialPrompt,
+		StatePersistenceConfig: screentracker.StatePersistenceConfig{
+			StateFile: stateFile,
+			LoadState: loadState,
+			SaveState: saveState,
+		},
 	})
+
 	if err != nil {
 		return xerrors.Errorf("failed to create server: %w", err)
 	}
 	if printOpenAPI {
 		fmt.Println(srv.GetOpenAPI())
 		return nil
 	}
+
+	// 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)
+
 	logger.Info("Starting server on port", "port", port)
+
+	// Monitor process exit
 	processExitCh := make(chan error, 1)
-	// Wait for process exit in PTY mode
 	if process != nil {
 		go func() {
 			defer close(processExitCh)
+			defer gracefulCancel()
 			if err := process.Wait(); err != nil {
 				if errors.Is(err, termexec.ErrNonZeroExitCode) {
 					processExitCh <- xerrors.Errorf("========\n%s\n========\n: %w", strings.TrimSpace(process.ReadScreen()), err)
 				} else {
 					processExitCh <- xerrors.Errorf("failed to wait for process: %w", err)
 				}
 			}
-			if err := srv.Stop(ctx); err != nil {
-				logger.Error("Failed to stop server", "error", err)
-			}
 		}()
 	}
-	// Wait for process exit in ACP mode
 	if acpResult != nil {
 		go func() {
 			defer close(processExitCh)
@@ -193,13 +251,45 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 			}
 		}()
 	}
-	if err := srv.Start(); err != nil && err != context.Canceled && err != http.ErrServerClosed {
-		return xerrors.Errorf("failed to start server: %w", err)
+
+	// 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
+		}
+	}()
+
+	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)
 	}
+
 	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:
+		// Close the process
+		if err := process.Close(logger, 5*time.Second); err != nil {
+			logger.Error("Failed to close process cleanly", "error", err)
+		}
 	}
 	return nil
 }
@@ -213,6 +303,61 @@ var agentNames = (func() []string {
 	return names
 })()
 
+// writePIDFile writes the current process ID to the specified file
+func writePIDFile(pidFile string, logger *slog.Logger) error {
+	pid := os.Getpid()
+	pidContent := fmt.Sprintf("%d\n", pid)
+
+	// Create directory if it doesn't exist
+	dir := filepath.Dir(pidFile)
+	if err := os.MkdirAll(dir, 0o700); err != nil {
+		return xerrors.Errorf("failed to create PID file directory: %w", err)
+	}
+
+	// Check if PID file already exists
+	if existingPIDData, err := os.ReadFile(pidFile); err == nil {
+		existingPIDStr := strings.TrimSpace(string(existingPIDData))
+		if existingPID, err := strconv.Atoi(existingPIDStr); err == nil {
+			if isProcessRunning(existingPID) {
+				return xerrors.Errorf("another instance is already running with PID %d (PID file: %s)", existingPID, pidFile)
+			}
+			logger.Warn("Found stale PID file, will overwrite", "pidFile", pidFile, "stalePID", existingPID)
+		}
+	} else if !os.IsNotExist(err) {
+		return xerrors.Errorf("failed to read existing PID file: %w", err)
+	}
+
+	// Write PID file
+	if err := os.WriteFile(pidFile, []byte(pidContent), 0o600); err != nil {
+		return xerrors.Errorf("failed to write PID file: %w", err)
+	}
+
+	logger.Info("Wrote PID file", "pidFile", pidFile, "pid", pid)
+	return nil
+}
+
+// cleanupPIDFile removes the PID file if it was written by this process.
+func cleanupPIDFile(pidFile string, logger *slog.Logger) {
+	data, err := os.ReadFile(pidFile)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			logger.Error("Failed to read PID file for cleanup", "pidFile", pidFile, "error", err)
+		}
+		return
+	}
+	pidStr := strings.TrimSpace(string(data))
+	filePID, err := strconv.Atoi(pidStr)
+	if err != nil || filePID != os.Getpid() {
+		logger.Info("PID file belongs to another process, skipping cleanup", "pidFile", pidFile, "filePID", pidStr)
+		return
+	}
+	if err := os.Remove(pidFile); err != nil && !os.IsNotExist(err) {
+		logger.Error("Failed to remove PID file", "pidFile", pidFile, "error", err)
+	} else if err == nil {
+		logger.Info("Removed PID file", "pidFile", pidFile)
+	}
+}
+
 type flagSpec struct {
 	name         string
 	shorthand    string
@@ -232,6 +377,10 @@ const (
 	FlagAllowedOrigins  = "allowed-origins"
 	FlagExit            = "exit"
 	FlagInitialPrompt   = "initial-prompt"
+	FlagStateFile       = "state-file"
+	FlagLoadState       = "load-state"
+	FlagSaveState       = "save-state"
+	FlagPidFile         = "pid-file"
 	FlagExperimentalACP = "experimental-acp"
 )
 
@@ -271,6 +420,10 @@ func CreateServerCmd() *cobra.Command {
 		// localhost:3284 is the default origin when you open the chat interface in your browser. localhost:3000 and 3001 are used during development.
 		{FlagAllowedOrigins, "o", []string{"http://localhost:3284", "http://localhost:3000", "http://localhost:3001"}, "HTTP allowed origins. Use '*' for all, comma-separated list via flag, space-separated list via AGENTAPI_ALLOWED_ORIGINS env var", "stringSlice"},
 		{FlagInitialPrompt, "I", "", "Initial prompt for the agent. Recommended only if the agent doesn't support initial prompt in interaction mode. Will be read from stdin if piped (e.g., echo 'prompt' | agentapi server -- my-agent)", "string"},
+		{FlagStateFile, "s", "", "Path to file for saving/loading server state", "string"},
+		{FlagLoadState, "", false, "Load state from state-file on startup (defaults to true when state-file is set)", "bool"},
+		{FlagSaveState, "", false, "Save state to state-file on shutdown (defaults to true when state-file is set)", "bool"},
+		{FlagPidFile, "", "", "Path to file where the server process ID will be written for shutdown scripts", "string"},
 		{FlagExperimentalACP, "", false, "Use experimental ACP transport instead of PTY", "bool"},
 	}