wip: address comments

Author: 35C4n0r

cmd/server/server.go

@@ -140,8 +140,6 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		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)
@@ -183,7 +181,7 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		fmt.Println(srv.GetOpenAPI())
 		return nil
 	}
-	handleSignals(ctx, logger, srv, process)
+	handleSignals(ctx, logger, srv, process, pidFile)
 	logger.Info("Starting server on port", "port", port)
 	processExitCh := make(chan error, 1)
 	go func() {

cmd/server/signals.go

@@ -12,7 +12,7 @@ import (
 
 // performGracefulShutdown handles the common shutdown logic for all platforms.
 // It saves state, stops the HTTP server, closes the process, and exits.
-func performGracefulShutdown(sig os.Signal, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process) {
+func performGracefulShutdown(sig os.Signal, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process, pidFile string) {
 	logger.Info("Received shutdown signal, initiating graceful shutdown", "signal", sig)
 
 	// Save state
@@ -31,6 +31,11 @@ func performGracefulShutdown(sig os.Signal, logger *slog.Logger, srv *httpapi.Se
 		logger.Error("Failed to close process cleanly", "signal", sig, "error", err)
 	}
 
+	// Clean up PID file before exit
+	if pidFile != "" {
+		cleanupPIDFile(pidFile, logger)
+	}
+
 	// Exit cleanly
 	os.Exit(0)
 }

cmd/server/signals_unix.go

@@ -16,14 +16,14 @@ import (
 // handleSignals sets up signal handlers for:
 // - SIGTERM, SIGINT, SIGHUP: save conversation state, stop server, then close the process
 // - SIGUSR1: save conversation state without exiting
-func handleSignals(ctx context.Context, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process) {
+func handleSignals(ctx context.Context, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process, pidFile string) {
 	// 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)
+		performGracefulShutdown(sig, logger, srv, process, pidFile)
 	}()
 
 	// Handle SIGUSR1 for save without exit

cmd/server/signals_windows.go

@@ -15,13 +15,13 @@ import (
 
 // 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) {
+func handleSignals(ctx context.Context, logger *slog.Logger, srv *httpapi.Server, process *termexec.Process, pidFile string) {
 	// 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)
+		performGracefulShutdown(sig, logger, srv, process, pidFile)
 	}()
 }

lib/httpapi/events.go

@@ -137,7 +137,7 @@ func (e *EventEmitter) notifyChannels(eventType EventType, payload any) {
 	}
 }
 
-// UpdateMessagesAndEmitChanges assumes that only the last message can change or new messages can be added.
+// EmitMessages 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) EmitMessages(newMessages []st.ConversationMessage) {
 	e.mu.Lock()

lib/screentracker/pty_conversation.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"os"
 	"path/filepath"
@@ -193,7 +194,17 @@ func (c *PTYConversation) Start(ctx context.Context) {
 		}
 
 		if c.initialPromptReady && !c.loadStateSuccessful && c.cfg.StatePersistenceConfig.LoadState {
-			_ = c.loadState()
+			if err := c.loadStateLocked(); err != nil {
+				// Add error message to conversation so user is aware
+				errorMsg := fmt.Sprintf("AgentAPI state restoration failed, the conversation history may be missing: %v", err)
+				c.messages = append(c.messages, ConversationMessage{
+					Id:      len(c.messages),
+					Message: errorMsg,
+					Role:    ConversationRoleAgent,
+					Time:    c.cfg.Clock.Now(),
+				})
+				c.cfg.Logger.Error("Failed to load state", "error", err)
+			}
 			c.loadStateSuccessful = true
 		}
 
@@ -534,16 +545,14 @@ func (c *PTYConversation) Text() string {
 }
 
 func (c *PTYConversation) SaveState() error {
-	conversation := c.Messages()
-
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
 	stateFile := c.cfg.StatePersistenceConfig.StateFile
 	saveState := c.cfg.StatePersistenceConfig.SaveState
 
 	if !saveState {
-		c.cfg.Logger.Info("")
+		c.cfg.Logger.Info("State persistence is disabled")
 		return nil
 	}
 
@@ -553,6 +562,8 @@ func (c *PTYConversation) SaveState() error {
 		return nil
 	}
 
+	conversation := c.messagesLocked()
+
 	// Serialize initial prompt from message parts
 	var initialPromptStr string
 	if len(c.cfg.InitialPrompt) > 0 {
@@ -598,8 +609,8 @@ func (c *PTYConversation) SaveState() error {
 	return nil
 }
 
-// LoadState loads the state, this method assumes that caller holds the Lock
-func (c *PTYConversation) loadState() error {
+// loadStateLocked loads the state, this method assumes that caller holds the Lock
+func (c *PTYConversation) loadStateLocked() error {
 	stateFile := c.cfg.StatePersistenceConfig.StateFile
 	loadState := c.cfg.StatePersistenceConfig.LoadState
 
@@ -613,20 +624,25 @@ func (c *PTYConversation) loadState() error {
 		return nil
 	}
 
-	// Read state file
-	data, err := os.ReadFile(stateFile)
+	// Open state file
+	f, err := os.Open(stateFile)
 	if err != nil {
-		c.cfg.Logger.Warn("Failed to load state file", "path", stateFile, "err", err)
-		return xerrors.Errorf("failed to read state file: %w", err)
-	}
-
-	if len(data) == 0 {
-		c.cfg.Logger.Info("No previous state to load (file is empty)", "path", stateFile)
-		return nil
+		c.cfg.Logger.Warn("Failed to open state file", "path", stateFile, "err", err)
+		return xerrors.Errorf("failed to open state file: %w", err)
 	}
+	defer func() {
+		if closeErr := f.Close(); closeErr != nil {
+			c.cfg.Logger.Warn("Failed to close state file", "path", stateFile, "err", closeErr)
+		}
+	}()
 
 	var agentState AgentState
-	if err := json.Unmarshal(data, &agentState); err != nil {
+	decoder := json.NewDecoder(f)
+	if err := decoder.Decode(&agentState); err != nil {
+		if err == io.EOF {
+			c.cfg.Logger.Info("No previous state to load (file is empty)", "path", stateFile)
+			return nil
+		}
 		c.cfg.Logger.Warn("Failed to load state file (corrupted or invalid JSON)", "path", stateFile, "err", err)
 		return xerrors.Errorf("failed to unmarshal state (corrupted or invalid JSON): %w", err)
 	}
@@ -663,7 +679,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 == false {
+	if !c.userSentMessageAfterLoadState {
 		newScreen = "\n" + c.messages[len(c.messages)-1].Message
 	}