more self-review fixes

Author: johnstcn

cmd/server/server.go

@@ -114,20 +114,20 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 	var agentIO st.AgentIO
 	var transport = "pty"
 	var process *termexec.Process
-	var acpWait func() error
+	var acpResult *httpapi.SetupACPResult
 
 	if printOpenAPI {
 		agentIO = nil
 	} else if experimentalACP {
-		acpResult, err := httpapi.SetupACP(ctx, httpapi.SetupACPConfig{
+		var err error
+		acpResult, err = httpapi.SetupACP(ctx, httpapi.SetupACPConfig{
 			Program:     agent,
 			ProgramArgs: argsToPass[1:],
 		})
 		if err != nil {
 			return xerrors.Errorf("failed to setup ACP: %w", err)
 		}
 		acpIO := acpResult.AgentIO
-		acpWait = acpResult.Wait
 		agentIO = acpIO
 		transport = "acp"
 	} else {
@@ -181,10 +181,11 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		}()
 	}
 	// Wait for process exit in ACP mode
-	if acpWait != nil {
+	if acpResult != nil {
 		go func() {
 			defer close(processExitCh)
-			if err := acpWait(); err != nil {
+			defer close(acpResult.Done) // Signal cleanup goroutine to exit
+			if err := acpResult.Wait(); err != nil {
 				processExitCh <- xerrors.Errorf("ACP process exited: %w", err)
 			}
 			if err := srv.Stop(ctx); err != nil {

lib/httpapi/setup.go

@@ -69,7 +69,8 @@ type SetupACPConfig struct {
 // SetupACPResult contains the result of setting up an ACP process.
 type SetupACPResult struct {
 	AgentIO *acpio.ACPAgentIO
-	Wait    func() error // Calls cmd.Wait() and returns exit error
+	Wait    func() error  // Calls cmd.Wait() and returns exit error
+	Done    chan struct{} // Close this when Wait() returns to clean up goroutine
 }
 
 func SetupACP(ctx context.Context, config SetupACPConfig) (*SetupACPResult, error) {
@@ -93,23 +94,35 @@ func SetupACP(ctx context.Context, config SetupACPConfig) (*SetupACPResult, erro
 		return nil, fmt.Errorf("failed to start process: %w", err)
 	}
 
-	agentIO, err := acpio.NewWithPipes(ctx, stdin, stdout, logger)
+	agentIO, err := acpio.NewWithPipes(ctx, stdin, stdout, logger, os.Getwd)
 	if err != nil {
 		_ = cmd.Process.Kill()
 		return nil, fmt.Errorf("failed to initialize ACP connection: %w", err)
 	}
 
+	done := make(chan struct{})
 	go func() {
-		<-ctx.Done()
-		logger.Info("Context done, closing ACP agent")
-		_ = stdin.Close()
-		_ = stdout.Close()
-		_ = cmd.Process.Kill()
+		select {
+		case <-ctx.Done():
+			logger.Info("Context done, closing ACP agent")
+			_ = stdin.Close()
+			_ = stdout.Close()
+			// Try graceful shutdown first
+			_ = cmd.Process.Signal(syscall.SIGTERM)
+			// Force kill after timeout
+			time.AfterFunc(5*time.Second, func() {
+				_ = cmd.Process.Kill()
+			})
+		case <-done:
+			// Process exited normally, nothing to clean up
+			return
+		}
 	}()
 
 	return &SetupACPResult{
 		AgentIO: agentIO,
 		Wait:    cmd.Wait,
+		Done:    done,
 	}, nil
 }
 

x/acpio/acpio.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"strings"
 	"sync"
+	"time"
 
 	acp "github.com/coder/acp-go-sdk"
 	st "github.com/coder/agentapi/lib/screentracker"
@@ -15,6 +16,9 @@ import (
 // Compile-time assertion that ACPAgentIO implements st.AgentIO
 var _ st.AgentIO = (*ACPAgentIO)(nil)
 
+// DefaultPromptTimeout is the maximum time to wait for an agent response.
+const DefaultPromptTimeout = 5 * time.Minute
+
 // ACPAgentIO implements screentracker.AgentIO using the ACP protocol
 type ACPAgentIO struct {
 	ctx       context.Context
@@ -131,7 +135,7 @@ func (a *ACPAgentIO) SetOnChunk(fn func(chunk string)) {
 }
 
 // NewWithPipes creates an ACPAgentIO connected via the provided pipes
-func NewWithPipes(ctx context.Context, toAgent io.Writer, fromAgent io.Reader, logger *slog.Logger) (*ACPAgentIO, error) {
+func NewWithPipes(ctx context.Context, toAgent io.Writer, fromAgent io.Reader, logger *slog.Logger, getwd func() (string, error)) (*ACPAgentIO, error) {
 	if logger == nil {
 		logger = slog.Default()
 	}
@@ -155,8 +159,13 @@ func NewWithPipes(ctx context.Context, toAgent io.Writer, fromAgent io.Reader, l
 	logger.Debug("ACP initialized", "protocolVersion", initResp.ProtocolVersion)
 
 	// Create a session
+	cwd, err := getwd()
+	if err != nil {
+		logger.Error("Failed to get working directory", "error", err)
+		return nil, err
+	}
 	sessResp, err := conn.NewSession(ctx, acp.NewSessionRequest{
-		Cwd:        "/tmp",
+		Cwd:        cwd,
 		McpServers: []acp.McpServer{},
 	})
 	if err != nil {
@@ -199,7 +208,11 @@ func (a *ACPAgentIO) Write(data []byte) (int, error) {
 		"textLen", len(text),
 		"rawDataLen", len(data))
 
-	resp, err := a.conn.Prompt(a.ctx, acp.PromptRequest{
+	// Use a timeout to prevent hanging indefinitely
+	promptCtx, cancel := context.WithTimeout(a.ctx, DefaultPromptTimeout)
+	defer cancel()
+
+	resp, err := a.conn.Prompt(promptCtx, acp.PromptRequest{
 		SessionId: a.sessionID,
 		Prompt:    []acp.ContentBlock{acp.TextBlock(text)},
 	})