fix race in ACPConversation: wait for chunk after Prompt() returns

Author: johnstcn

x/acpio/acp_conversation.go

@@ -6,6 +6,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"time"
 
 	st "github.com/coder/agentapi/lib/screentracker"
 	"github.com/coder/quartz"
@@ -31,7 +32,8 @@ type ACPConversation struct {
 	agentIO           ChunkableAgentIO
 	messages          []st.ConversationMessage
 	nextID            int  // monotonically increasing message ID
-	prompting         bool // true while agent is processing
+	prompting         bool          // true while agent is processing
+	chunkReceived     chan struct{}  // signals that handleChunk has accumulated a chunk
 	streamingResponse strings.Builder
 	logger            *slog.Logger
 	emitter           st.Emitter
@@ -68,6 +70,7 @@ func NewACPConversation(ctx context.Context, agentIO ChunkableAgentIO, logger *s
 		initialPrompt: initialPrompt,
 		emitter:       emitter,
 		clock:         clock,
+		chunkReceived: make(chan struct{}, 1),
 	}
 	return c
 }
@@ -202,13 +205,25 @@ func (c *ACPConversation) handleChunk(chunk string) {
 	screen := c.streamingResponse.String()
 	c.mu.Unlock()
 
+	// Signal that a chunk has been received (non-blocking; a pending signal is sufficient).
+	select {
+	case c.chunkReceived <- struct{}{}:
+	default:
+	}
+
 	c.emitter.EmitMessages(messages)
 	c.emitter.EmitStatus(status)
 	c.emitter.EmitScreen(screen)
 }
 
 // executePrompt runs the actual agent request and returns any error.
 func (c *ACPConversation) executePrompt(messageParts []st.MessagePart) error {
+	// Drain any stale signal before sending the prompt.
+	select {
+	case <-c.chunkReceived:
+	default:
+	}
+
 	var err error
 	for _, part := range messageParts {
 		if c.ctx.Err() != nil {
@@ -221,6 +236,13 @@ func (c *ACPConversation) executePrompt(messageParts []st.MessagePart) error {
 		}
 	}
 
+	// The ACP SDK dispatches SessionUpdate notifications as goroutines, so
+	// the chunk may arrive after conn.Prompt() returns. Wait up to 100ms.
+	select {
+	case <-c.chunkReceived:
+	case <-time.After(100 * time.Millisecond):
+	}
+
 	c.mu.Lock()
 	c.prompting = false