feat: implement state persistence

Author: 35C4n0r

cmd/server/server.go

@@ -103,6 +103,26 @@ 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)
+
 	printOpenAPI := viper.GetBool(FlagPrintOpenAPI)
 	var process *termexec.Process
 	if printOpenAPI {
@@ -128,7 +148,14 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		AllowedHosts:   viper.GetStringSlice(FlagAllowedHosts),
 		AllowedOrigins: viper.GetStringSlice(FlagAllowedOrigins),
 		InitialPrompt:  initialPrompt,
+		StatePersistenceCfg: httpapi.StatePersistenceCfg{
+			StateFile: stateFile,
+			LoadState: loadState,
+			SaveState: saveState,
+			PidFile:   pidFile,
+		},
 	})
+
 	if err != nil {
 		return xerrors.Errorf("failed to create server: %w", err)
 	}
@@ -137,6 +164,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 +180,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 {
@@ -191,6 +219,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 +261,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

@@ -11,11 +11,13 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"slices"
 	"sort"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 	"unicode"
 
@@ -34,18 +36,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.PTYConversation
-	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
+	statePersistenceCfg StatePersistenceCfg
+	stateLoadComplete   bool
 }
 
 func (s *Server) NormalizeSchema(schema any) any {
@@ -94,14 +98,22 @@ func (s *Server) GetOpenAPI() string {
 // because the action of taking a snapshot takes time too.
 const snapshotInterval = 25 * time.Millisecond
 
+type StatePersistenceCfg struct {
+	StateFile string
+	LoadState bool
+	SaveState bool
+	PidFile   string
+}
+
 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
+	StatePersistenceCfg StatePersistenceCfg
 }
 
 // Validate allowed hosts don't contain whitespace, commas, schemes, or ports.
@@ -260,16 +272,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,
+		statePersistenceCfg: config.StatePersistenceCfg,
+		stateLoadComplete:   false,
 	}
 
 	// Register API routes
@@ -337,15 +351,26 @@ func (s *Server) StartSnapshotLoop(ctx context.Context) {
 			currentStatus := s.conversation.Status()
 
 			// Send initial prompt when agent becomes stable for the first time
-			if !s.conversation.InitialPromptSent && convertStatus(currentStatus) == AgentStatusStable {
-
-				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")
+			if convertStatus(currentStatus) == AgentStatusStable {
+
+				if !s.stateLoadComplete && s.statePersistenceCfg.LoadState {
+					_, err := s.conversation.LoadState(s.statePersistenceCfg.StateFile)
+					if err != nil {
+						s.logger.Warn("Failed to load state file", "path", s.statePersistenceCfg.StateFile, "err", err)
+					} else {
+						s.logger.Info("Successfully loaded state", "path", s.statePersistenceCfg.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)
@@ -592,6 +617,15 @@ func (s *Server) Start() error {
 
 // Stop gracefully stops the HTTP server
 func (s *Server) Stop(ctx context.Context) error {
+	// Save conversation state if configured
+	if s.statePersistenceCfg.SaveState && s.statePersistenceCfg.StateFile != "" {
+		if err := s.conversation.SaveState(s.conversation.Messages(), s.statePersistenceCfg.StateFile); err != nil {
+			s.logger.Error("Failed to save conversation state", "error", err)
+		} else {
+			s.logger.Info("Saved conversation state", "stateFile", s.statePersistenceCfg.StateFile)
+		}
+	}
+
 	// Clean up temporary directory
 	s.cleanupTempDir()
 
@@ -610,6 +644,58 @@ func (s *Server) cleanupTempDir() {
 	}
 }
 
+// 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() {
+		sig := <-shutdownCh
+		s.logger.Info("Received shutdown signal, saving state before closing process", "signal", sig)
+
+		// Save conversation state if configured (synchronously before closing process)
+		if s.statePersistenceCfg.SaveState && s.statePersistenceCfg.StateFile != "" {
+			if err := s.conversation.SaveState(s.conversation.Messages(), s.statePersistenceCfg.StateFile); err != nil {
+				s.logger.Error("Failed to save conversation state on signal", "signal", sig, "error", err)
+			} else {
+				s.logger.Info("Saved conversation state on signal", "signal", sig, "stateFile", s.statePersistenceCfg.StateFile)
+			}
+		}
+
+		// 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)
+		}
+	}()
+
+	// Handle SIGUSR1 for save without exit
+	saveOnlyCh := make(chan os.Signal, 1)
+	signal.Notify(saveOnlyCh, syscall.SIGUSR1)
+	go func() {
+		for {
+			select {
+			case <-saveOnlyCh:
+				s.logger.Info("Received SIGUSR1, saving state without exiting")
+
+				// Save conversation state if configured
+				if s.statePersistenceCfg.SaveState && s.statePersistenceCfg.StateFile != "" {
+					if err := s.conversation.SaveState(s.conversation.Messages(), s.statePersistenceCfg.StateFile); err != nil {
+						s.logger.Error("Failed to save conversation state on SIGUSR1", "error", err)
+					} else {
+						s.logger.Info("Saved conversation state on SIGUSR1", "stateFile", s.statePersistenceCfg.StateFile)
+					}
+				} else {
+					s.logger.Warn("SIGUSR1 received but state saving is not configured")
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+}
+
 // registerStaticFileRoutes sets up routes for serving static files
 func (s *Server) registerStaticFileRoutes() {
 	chatHandler := FileServerWithIndexFallback(s.chatBasePath)

lib/httpapi/setup.go

@@ -4,10 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"os/signal"
 	"strings"
-	"syscall"
-	"time"
 
 	"github.com/coder/agentapi/lib/logctx"
 	mf "github.com/coder/agentapi/lib/msgfmt"
@@ -45,16 +42,5 @@ func SetupProcess(ctx context.Context, config SetupProcessConfig) (*termexec.Pro
 			return nil, err
 		}
 	}
-
-	// Handle SIGINT (Ctrl+C) and send it to the process
-	signalCh := make(chan os.Signal, 1)
-	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
-	go func() {
-		<-signalCh
-		if err := process.Close(logger, 5*time.Second); err != nil {
-			logger.Error("Error closing process", "error", err)
-		}
-	}()
-
 	return process, nil
 }

lib/screentracker/conversation.go

@@ -52,6 +52,8 @@ type MessagePart interface {
 // Conversation allows tracking of a conversation between a user and an agent.
 type Conversation interface {
 	Messages() []ConversationMessage
+	SaveState([]ConversationMessage, string) error
+	LoadState(string) ([]ConversationMessage, error)
 	Snapshot(string)
 	Start(context.Context)
 	Status() ConversationStatus
@@ -64,3 +66,10 @@ type ConversationMessage struct {
 	Role    ConversationRole
 	Time    time.Time
 }
+
+type AgentState struct {
+	Version           int                   `json:"version"`
+	Messages          []ConversationMessage `json:"messages"`
+	InitialPrompt     string                `json:"initial_prompt"`
+	InitialPromptSent bool                  `json:"initial_prompt_sent"`
+}