feat: implement state persistence

cmd/server/server.go

@@ -8,6 +8,7 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 
@@ -103,6 +104,35 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		}
 	}
 
+	// Get the variables related to state management
+	stateFile := viper.GetString(StateFile)
+	loadState := true
+	saveState := true
+	if stateFile != "" {
+		if !viper.IsSet(LoadState) {
+			loadState = true
+		} else {
+			loadState = viper.GetBool(LoadState)
+		}
+
+		if !viper.IsSet(SaveState) {
+			saveState = true
+		} else {
+			saveState = viper.GetBool(SaveState)
+		}
+	}
+
+	pidFile := viper.GetString(PidFile)
+
+	// Write PID file if configured
+	if pidFile != "" {
+		if err := writePIDFile(pidFile, logger); err != nil {
+			return xerrors.Errorf("failed to write PID file: %w", err)
+		}
+		// Ensure PID file is cleaned up on exit
+		defer cleanupPIDFile(pidFile, logger)
+	}
+
 	printOpenAPI := viper.GetBool(FlagPrintOpenAPI)
 	var process *termexec.Process
 	if printOpenAPI {
@@ -128,7 +158,13 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		AllowedHosts:   viper.GetStringSlice(FlagAllowedHosts),
 		AllowedOrigins: viper.GetStringSlice(FlagAllowedOrigins),
 		InitialPrompt:  initialPrompt,
+		StatePersistenceConfig: httpapi.StatePersistenceConfig{
+			StateFile: stateFile,
+			LoadState: loadState,
+			SaveState: saveState,
+		},
 	})
+
 	if err != nil {
 		return xerrors.Errorf("failed to create server: %w", err)
 	}
@@ -137,6 +173,7 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		return nil
 	}
 	srv.StartSnapshotLoop(ctx)
+	srv.HandleSignals(ctx, process)
 	logger.Info("Starting server on port", "port", port)
 	processExitCh := make(chan error, 1)
 	go func() {
@@ -152,7 +189,7 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 			logger.Error("Failed to stop server", "error", err)
 		}
 	}()
-	if err := srv.Start(); err != nil && err != context.Canceled && err != http.ErrServerClosed {
+	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 {
@@ -172,6 +209,35 @@ 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, 0o755); err != nil {
+		return xerrors.Errorf("failed to create PID file directory: %w", err)
+	}
+
+	// Write PID file
+	if err := os.WriteFile(pidFile, []byte(pidContent), 0o644); 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 exists
+func cleanupPIDFile(pidFile string, logger *slog.Logger) {
+	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
@@ -191,6 +257,10 @@ const (
 	FlagAllowedOrigins = "allowed-origins"
 	FlagExit           = "exit"
 	FlagInitialPrompt  = "initial-prompt"
+	StateFile          = "state-file"
+	LoadState          = "load-state"
+	SaveState          = "save-state"
+	PidFile            = "pid-file"
 )
 
 func CreateServerCmd() *cobra.Command {
@@ -229,6 +299,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"},
+		{StateFile, "s", "", "Path to file for saving/loading server state", "string"},
+		{LoadState, "", false, "Load state from state-file on startup (defaults to true when state-file is set)", "bool"},
+		{SaveState, "", false, "Save state to state-file on shutdown (defaults to true when state-file is set)", "bool"},
+		{PidFile, "", "", "Path to file where the server process ID will be written for shutdown scripts", "string"},
 	}
 
 	for _, spec := range flagSpecs {

lib/httpapi/events.go

@@ -120,7 +120,7 @@ func (e *EventEmitter) notifyChannels(eventType EventType, payload any) {
 	}
 }
 
-// Assumes that only the last message can change or new messages can be added.
+// UpdateMessagesAndEmitChanges assumes that only the last message can change or new messages can be added.
 // If a new message is injected between existing messages (identified by Id), the behavior is undefined.
 func (e *EventEmitter) UpdateMessagesAndEmitChanges(newMessages []st.ConversationMessage) {
 	e.mu.Lock()

lib/httpapi/server.go

@@ -34,18 +34,20 @@ import (
 
 // Server represents the HTTP server
 type Server struct {
-	router       chi.Router
-	api          huma.API
-	port         int
-	srv          *http.Server
-	mu           sync.RWMutex
-	logger       *slog.Logger
-	conversation *st.Conversation
-	agentio      *termexec.Process
-	agentType    mf.AgentType
-	emitter      *EventEmitter
-	chatBasePath string
-	tempDir      string
+	router                 chi.Router
+	api                    huma.API
+	port                   int
+	srv                    *http.Server
+	mu                     sync.RWMutex
+	logger                 *slog.Logger
+	conversation           *st.PTYConversation
+	agentio                *termexec.Process
+	agentType              mf.AgentType
+	emitter                *EventEmitter
+	chatBasePath           string
+	tempDir                string
+	statePersistenceConfig StatePersistenceConfig
+	stateLoadComplete      bool
 }
 
 func (s *Server) NormalizeSchema(schema any) any {
@@ -94,14 +96,21 @@ func (s *Server) GetOpenAPI() string {
 // because the action of taking a snapshot takes time too.
 const snapshotInterval = 25 * time.Millisecond
 
+type StatePersistenceConfig struct {
+	StateFile string
+	LoadState bool
+	SaveState bool
+}
+
 type ServerConfig struct {
-	AgentType      mf.AgentType
-	Process        *termexec.Process
-	Port           int
-	ChatBasePath   string
-	AllowedHosts   []string
-	AllowedOrigins []string
-	InitialPrompt  string
+	AgentType              mf.AgentType
+	Process                *termexec.Process
+	Port                   int
+	ChatBasePath           string
+	AllowedHosts           []string
+	AllowedOrigins         []string
+	InitialPrompt          string
+	StatePersistenceConfig StatePersistenceConfig
 }
 
 // Validate allowed hosts don't contain whitespace, commas, schemes, or ports.
@@ -237,7 +246,7 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 		return mf.FormatToolCall(config.AgentType, message)
 	}
 
-	conversation := st.NewConversation(ctx, st.ConversationConfig{
+	conversation := st.NewPTY(ctx, st.PTYConversationConfig{
 		AgentType: config.AgentType,
 		AgentIO:   config.Process,
 		GetTime: func() time.Time {
@@ -260,16 +269,18 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 	logger.Info("Created temporary directory for uploads", "tempDir", tempDir)
 
 	s := &Server{
-		router:       router,
-		api:          api,
-		port:         config.Port,
-		conversation: conversation,
-		logger:       logger,
-		agentio:      config.Process,
-		agentType:    config.AgentType,
-		emitter:      emitter,
-		chatBasePath: strings.TrimSuffix(config.ChatBasePath, "/"),
-		tempDir:      tempDir,
+		router:                 router,
+		api:                    api,
+		port:                   config.Port,
+		conversation:           conversation,
+		logger:                 logger,
+		agentio:                config.Process,
+		agentType:              config.AgentType,
+		emitter:                emitter,
+		chatBasePath:           strings.TrimSuffix(config.ChatBasePath, "/"),
+		tempDir:                tempDir,
+		statePersistenceConfig: config.StatePersistenceConfig,
+		stateLoadComplete:      false,
 	}
 
 	// Register API routes
@@ -331,26 +342,32 @@ func sseMiddleware(ctx huma.Context, next func(huma.Context)) {
 }
 
 func (s *Server) StartSnapshotLoop(ctx context.Context) {
-	s.conversation.StartSnapshotLoop(ctx)
+	s.conversation.Start(ctx)
 	go func() {
 		for {
 			currentStatus := s.conversation.Status()
 
-			// Send initial prompt when agent becomes stable for the first time
-			if !s.conversation.InitialPromptSent && convertStatus(currentStatus) == AgentStatusStable {
+			// Send initial prompt & load state when agent becomes stable for the first time
+			if convertStatus(currentStatus) == AgentStatusStable {
 
-				if err := s.conversation.SendMessage(FormatMessage(s.agentType, s.conversation.InitialPrompt)...); err != nil {
-					s.logger.Error("Failed to send initial prompt", "error", err)
-				} else {
-					s.conversation.InitialPromptSent = true
-					s.conversation.ReadyForInitialPrompt = false
-					currentStatus = st.ConversationStatusChanging
-					s.logger.Info("Initial prompt sent successfully")
+				if !s.stateLoadComplete && s.statePersistenceConfig.LoadState {
+					_, _ = s.conversation.LoadState(s.statePersistenceConfig.StateFile)
+					s.stateLoadComplete = true
+				}
+				if !s.conversation.InitialPromptSent {
+					if err := s.conversation.Send(FormatMessage(s.agentType, s.conversation.InitialPrompt)...); err != nil {
+						s.logger.Error("Failed to send initial prompt", "error", err)
+					} else {
+						s.conversation.InitialPromptSent = true
+						s.conversation.ReadyForInitialPrompt = false
+						currentStatus = st.ConversationStatusChanging
+						s.logger.Info("Initial prompt sent successfully")
+					}
 				}
 			}
 			s.emitter.UpdateStatusAndEmitChanges(currentStatus, s.agentType)
 			s.emitter.UpdateMessagesAndEmitChanges(s.conversation.Messages())
-			s.emitter.UpdateScreenAndEmitChanges(s.conversation.Screen())
+			s.emitter.UpdateScreenAndEmitChanges(s.conversation.String())
 			time.Sleep(snapshotInterval)
 		}
 	}()
@@ -449,7 +466,7 @@ func (s *Server) createMessage(ctx context.Context, input *MessageRequest) (*Mes
 
 	switch input.Body.Type {
 	case MessageTypeUser:
-		if err := s.conversation.SendMessage(FormatMessage(s.agentType, input.Body.Content)...); err != nil {
+		if err := s.conversation.Send(FormatMessage(s.agentType, input.Body.Content)...); err != nil {
 			return nil, xerrors.Errorf("failed to send message: %w", err)
 		}
 	case MessageTypeRaw:
@@ -610,6 +627,30 @@ func (s *Server) cleanupTempDir() {
 	}
 }
 
+// saveAndCleanup saves the conversation state and cleans up before shutdown
+func (s *Server) saveAndCleanup(sig os.Signal, process *termexec.Process) {
+	// Save conversation state if configured (synchronously before closing process)
+	s.saveStateIfConfigured(sig.String())
+
+	// Now close the process
+	if err := process.Close(s.logger, 5*time.Second); err != nil {
+		s.logger.Error("Error closing process", "signal", sig, "error", err)
+	}
+}
+
+// saveStateIfConfigured saves the conversation state if configured
+func (s *Server) saveStateIfConfigured(source string) {
+	if s.statePersistenceConfig.SaveState && s.statePersistenceConfig.StateFile != "" {
+		if err := s.conversation.SaveState(s.conversation.Messages(), s.statePersistenceConfig.StateFile); err != nil {
+			s.logger.Error("Failed to save conversation state", "source", source, "error", err)
+		} else {
+			s.logger.Info("Saved conversation state", "source", source, "stateFile", s.statePersistenceConfig.StateFile)
+		}
+	} else {
+		s.logger.Warn("Save requested but state saving is not configured", "source", source)
+	}
+}
+
 // registerStaticFileRoutes sets up routes for serving static files
 func (s *Server) registerStaticFileRoutes() {
 	chatHandler := FileServerWithIndexFallback(s.chatBasePath)

lib/httpapi/server_signals_unix.go

@@ -0,0 +1,44 @@
+//go:build unix
+
+package httpapi
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/coder/agentapi/lib/termexec"
+)
+
+// HandleSignals sets up signal handlers for:
+// - SIGTERM, SIGINT, SIGHUP: save conversation state, then close the process
+// - SIGUSR1: save conversation state without exiting
+func (s *Server) HandleSignals(ctx context.Context, process *termexec.Process) {
+	// Handle shutdown signals (SIGTERM, SIGINT, SIGHUP)
+	shutdownCh := make(chan os.Signal, 1)
+	signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
+	go func() {
+		defer signal.Stop(shutdownCh)
+		sig := <-shutdownCh
+		s.logger.Info("Received shutdown signal, saving state before closing process", "signal", sig)
+
+		s.saveAndCleanup(sig, process)
+	}()
+
+	// Handle SIGUSR1 for save without exit
+	saveOnlyCh := make(chan os.Signal, 1)
+	signal.Notify(saveOnlyCh, syscall.SIGUSR1)
+	go func() {
+		defer signal.Stop(saveOnlyCh)
+		for {
+			select {
+			case <-saveOnlyCh:
+				s.logger.Info("Received SIGUSR1, saving state without exiting")
+				s.saveStateIfConfigured("SIGUSR1")
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+}

lib/httpapi/server_signals_windows.go

@@ -0,0 +1,27 @@
+//go:build windows
+
+package httpapi
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"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 (s *Server) HandleSignals(ctx context.Context, process *termexec.Process) {
+	// 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
+		s.logger.Info("Received shutdown signal, saving state before closing process", "signal", sig)
+
+		s.saveAndCleanup(sig, process)
+	}()
+}