feat: address maf's review

Author: 35C4n0r

cmd/server/server.go

@@ -140,6 +140,7 @@ 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)
 		}
+		defer cleanupPIDFile(pidFile, logger)
 	}
 
 	printOpenAPI := viper.GetBool(FlagPrintOpenAPI)
@@ -189,30 +190,20 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 	// Setup signal handlers (they will call gracefulCancel)
 	handleSignals(gracefulCtx, gracefulCancel, logger, srv)
 
-	// Setup PID file cleanup
-	if pidFile != "" {
-		defer cleanupPIDFile(pidFile, logger)
-	}
-
 	logger.Info("Starting server on port", "port", port)
 
 	// Monitor process exit
 	processExitCh := make(chan error, 1)
 	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)
 			}
 		}
-
-		select {
-		case <-gracefulCtx.Done():
-		default:
-			gracefulCancel()
-		}
 	}()
 
 	// Start the server
@@ -243,17 +234,16 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		logger.Error("Failed to stop HTTP server", "error", err)
 	}
 
-	// Close the process
-	if err := process.Close(logger, 5*time.Second); err != nil {
-		logger.Error("Failed to close process cleanly", "error", err)
-	}
-
 	select {
 	case err := <-processExitCh:
 		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
 }

e2e/echo_test.go

@@ -138,9 +138,6 @@ func TestE2E(t *testing.T) {
 		// Step 2: Stop server (triggers state save)
 		cleanup()
 
-		// Give filesystem a moment to sync
-		time.Sleep(100 * time.Millisecond)
-
 		// Verify state file was created
 		require.FileExists(t, stateFile, "State file should exist after shutdown")
 
@@ -194,7 +191,6 @@ func TestE2E(t *testing.T) {
 
 		// Step 2: Close server
 		cleanup()
-		time.Sleep(100 * time.Millisecond)
 		require.FileExists(t, stateFile, "State file should exist after shutdown")
 
 		// Step 3: Restart WITHOUT an initial prompt
@@ -213,7 +209,6 @@ func TestE2E(t *testing.T) {
 
 		// Step 5: Close server
 		cleanup2()
-		time.Sleep(100 * time.Millisecond)
 
 		// Step 6: Restart with same initial prompt
 		_, apiClient3, cleanup3 := setup(ctx, t, &params{
@@ -256,7 +251,6 @@ func TestE2E(t *testing.T) {
 
 		// Step 2: Close server
 		cleanup()
-		time.Sleep(100 * time.Millisecond)
 		require.FileExists(t, stateFile, "State file should exist after shutdown")
 
 		// Step 3: Restart with DIFFERENT initial prompt using a different script
@@ -384,7 +378,7 @@ func setup(ctx context.Context, t testing.TB, p *params, waitForStable bool) ([]
 	// Create cleanup function
 	cleanup := func() {
 		if cmd.Process != nil {
-			// Send SIGTERM to allow graceful shutdown and state save
+			// Send SIGINT to allow graceful shutdown and state save
 			_ = cmd.Process.Signal(os.Interrupt)
 			// Wait for process to exit gracefully (with timeout)
 			done := make(chan error, 1)
@@ -394,7 +388,7 @@ func setup(ctx context.Context, t testing.TB, p *params, waitForStable bool) ([]
 			select {
 			case <-done:
 				// Process exited gracefully
-			case <-time.After(5 * time.Second):
+			case <-time.After(10 * time.Second):
 				// Timeout, force kill
 				_ = cmd.Process.Kill()
 				<-done

lib/httpapi/server.go

@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -615,7 +616,9 @@ func (s *Server) Stop(ctx context.Context) error {
 		s.cleanupTempDir()
 
 		if s.srv != nil {
-			err = s.srv.Shutdown(ctx)
+			if err = s.srv.Shutdown(ctx); errors.Is(err, http.ErrServerClosed) {
+				err = nil
+			}
 		}
 	})
 	return err

lib/screentracker/pty_conversation.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log/slog"
 	"os"
 	"path/filepath"
@@ -215,7 +214,6 @@ func (c *PTYConversation) Start(ctx context.Context) {
 			}
 		}
 
-		// Enqueue initial prompt once after agent is ready (and after state is potentially loaded)
 		if c.initialPromptReady && len(c.cfg.InitialPrompt) > 0 && !c.initialPromptSent {
 			c.outboundQueue <- outboundMessage{parts: c.cfg.InitialPrompt, errCh: nil}
 			c.initialPromptSent = true
@@ -576,27 +574,34 @@ func (c *PTYConversation) SaveState() error {
 		initialPromptStr = buildStringFromMessageParts(c.cfg.InitialPrompt)
 	}
 
-	// Use atomic write: write to temp file, then rename to target path
-	data, err := json.MarshalIndent(AgentState{
-		Version:           1,
-		Messages:          conversation,
-		InitialPrompt:     initialPromptStr,
-		InitialPromptSent: c.initialPromptSent,
-	}, "", " ")
-	if err != nil {
-		return xerrors.Errorf("failed to marshal state: %w", err)
-	}
-
 	// Create directory if it doesn't exist
 	dir := filepath.Dir(stateFile)
 	if err := os.MkdirAll(dir, 0o700); err != nil {
 		return xerrors.Errorf("failed to create state directory: %w", err)
 	}
 
-	// Write to temp file
+	// Use atomic write: write to temp file, then rename to target path
 	tempFile := stateFile + ".tmp"
-	if err := os.WriteFile(tempFile, data, 0o600); err != nil {
-		return xerrors.Errorf("failed to write temp state file: %w", err)
+	f, err := os.OpenFile(tempFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
+	if err != nil {
+		return xerrors.Errorf("failed to create temp state file: %w", err)
+	}
+
+	// Encode directly to file to avoid loading entire JSON into memory
+	encoder := json.NewEncoder(f)
+	encoder.SetIndent("", " ")
+	if err := encoder.Encode(AgentState{
+		Version:           1,
+		Messages:          conversation,
+		InitialPrompt:     initialPromptStr,
+		InitialPromptSent: c.initialPromptSent,
+	}); err != nil {
+		return xerrors.Errorf("failed to encode state: %w", err)
+	}
+
+	// Close file before rename
+	if err := f.Close(); err != nil {
+		return xerrors.Errorf("failed to close temp state file: %w", err)
 	}
 
 	// Atomic rename
@@ -641,10 +646,6 @@ func (c *PTYConversation) loadStateLocked() error {
 	var agentState AgentState
 	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
-		}
 		return xerrors.Errorf("failed to unmarshal state (corrupted or invalid JSON): %w", err)
 	}
 

lib/termexec/termexec.go

@@ -163,7 +163,7 @@ func (p *Process) Close(logger *slog.Logger, timeout time.Duration) error {
 	case err := <-exited:
 		var pathErr *os.SyscallError
 		// ECHILD is expected if the process has already exited
-		if err != nil && !(errors.As(err, &pathErr) && pathErr.Err == syscall.ECHILD) {
+		if err != nil && !(errors.As(err, &pathErr) && errors.Is(pathErr.Err, syscall.ECHILD)) {
 			exitErr = xerrors.Errorf("process exited with error: %w", err)
 		}
 	}