wip: apply suggestions from cian

Author: 35C4n0r

cmd/server/server.go

@@ -107,33 +107,33 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 	}
 
 	// Get the variables related to state management
-	stateFile := viper.GetString(StateFile)
+	stateFile := viper.GetString(FlagStateFile)
 	loadState := false
 	saveState := false
 
 	// Validate state file configuration
 	if stateFile != "" {
-		if !viper.IsSet(LoadState) {
+		if !viper.IsSet(FlagLoadState) {
 			loadState = true
 		} else {
-			loadState = viper.GetBool(LoadState)
+			loadState = viper.GetBool(FlagLoadState)
 		}
 
-		if !viper.IsSet(SaveState) {
+		if !viper.IsSet(FlagSaveState) {
 			saveState = true
 		} else {
-			saveState = viper.GetBool(SaveState)
+			saveState = viper.GetBool(FlagSaveState)
 		}
 	} else {
-		if viper.IsSet(LoadState) && viper.GetBool(LoadState) {
+		if viper.IsSet(FlagLoadState) && viper.GetBool(FlagLoadState) {
 			return xerrors.Errorf("--load-state requires --state-file to be set")
 		}
-		if viper.IsSet(SaveState) && viper.GetBool(SaveState) {
+		if viper.IsSet(FlagSaveState) && viper.GetBool(FlagSaveState) {
 			return xerrors.Errorf("--save-state requires --state-file to be set")
 		}
 	}
 
-	pidFile := viper.GetString(PidFile)
+	pidFile := viper.GetString(FlagPidFile)
 
 	// Write PID file if configured
 	if pidFile != "" {
@@ -315,10 +315,10 @@ const (
 	FlagAllowedOrigins = "allowed-origins"
 	FlagExit           = "exit"
 	FlagInitialPrompt  = "initial-prompt"
-	StateFile          = "state-file"
-	LoadState          = "load-state"
-	SaveState          = "save-state"
-	PidFile            = "pid-file"
+	FlagStateFile      = "state-file"
+	FlagLoadState      = "load-state"
+	FlagSaveState      = "save-state"
+	FlagPidFile        = "pid-file"
 )
 
 func CreateServerCmd() *cobra.Command {
@@ -357,10 +357,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"},
+		{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"},
 	}
 
 	for _, spec := range flagSpecs {

lib/screentracker/pty_conversation.go

@@ -596,7 +596,7 @@ func (c *PTYConversation) SaveState() error {
 	// Clear dirty flag after successful save
 	c.dirty = false
 
-	c.cfg.Logger.Info(fmt.Sprintf("State saved successfully to: %s", stateFile))
+	c.cfg.Logger.Info("State saved successfully", "path", stateFile)
 
 	return nil
 }
@@ -606,7 +606,7 @@ func (c *PTYConversation) loadStateLocked() error {
 	stateFile := c.cfg.StatePersistenceConfig.StateFile
 	loadState := c.cfg.StatePersistenceConfig.LoadState
 
-	if !loadState {
+	if !loadState || c.loadStateSuccessful {
 		return nil
 	}
 
@@ -647,7 +647,7 @@ func (c *PTYConversation) loadStateLocked() error {
 
 	// Store the first stable snapshot for filtering later
 	snapshots := c.snapshotBuffer.GetAll()
-	if len(snapshots) > 0 {
+	if len(snapshots) > 0 && c.cfg.FormatMessage != nil {
 		c.firstStableSnapshot = c.cfg.FormatMessage(strings.TrimSpace(snapshots[len(snapshots)-1].screen), "")
 	}
 
@@ -669,7 +669,7 @@ func (c *PTYConversation) adjustScreenAfterStateLoad(screen string) string {
 	// Before the first user message after loading state, return the last message from the loaded state.
 	// This prevents computing incorrect diffs from the restored screen, as the agent's message should
 	// remain stable until the user continues the conversation.
-	if !c.userSentMessageAfterLoadState {
+	if !c.userSentMessageAfterLoadState && len(c.messages) > 0 {
 		newScreen = "\n" + c.messages[len(c.messages)-1].Message
 	}