fix: do not corrupt user message if late chunk arrives after error

Author: johnstcn

x/acpio/acp_conversation.go

@@ -179,9 +179,18 @@ func (c *ACPConversation) Text() string {
 // handleChunk is called for each streaming chunk from the agent.
 func (c *ACPConversation) handleChunk(chunk string) {
 	c.mu.Lock()
+	// Log and discard chunks that arrive after the prompt has completed or errored.
+	// This should not happen under normal operation — if it does, it indicates a
+	// bug in the ACP SDK or a race in the connection teardown.
+	if !c.prompting {
+		c.mu.Unlock()
+		c.logger.Error("received chunk while not prompting (late/unexpected chunk discarded)",
+			"chunkLen", len(chunk))
+		return
+	}
 	c.streamingResponse.WriteString(chunk)
-	// Update the last message (the streaming agent response)
-	if len(c.messages) > 0 {
+	// Only update the last message if it's the agent placeholder (defense-in-depth)
+	if len(c.messages) > 0 && c.messages[len(c.messages)-1].Role == st.ConversationRoleAgent {
 		c.messages[len(c.messages)-1].Message = c.streamingResponse.String()
 	}
 	messages := slices.Clone(c.messages)

x/acpio/acp_conversation_test.go

@@ -504,3 +504,41 @@ func Test_ErrorRemovesPartialMessage(t *testing.T) {
 	assert.Equal(t, screentracker.ConversationRoleUser, messages[0].Role)
 	assert.Equal(t, "test", messages[0].Message)
 }
+
+func Test_LateChunkAfterError_DoesNotCorruptUserMessage(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	mClock := quartz.NewMock(t)
+	mock := newMockAgentIO()
+	emitter := newMockEmitter()
+	started, done := mock.BlockWrite()
+
+	conv := acpio.NewACPConversation(ctx, mock, nil, nil, emitter, mClock)
+	conv.Start(ctx)
+
+	// Given: a send that fails with an error, removing the agent placeholder
+	err := conv.Send(screentracker.MessagePartText{Content: "hello"})
+	require.NoError(t, err)
+	<-started
+
+	mock.mu.Lock()
+	mock.writeErr = assert.AnError
+	mock.mu.Unlock()
+	close(done)
+
+	emitter.WaitForStatus(ctx, t, screentracker.ConversationStatusStable)
+
+	messages := conv.Messages()
+	require.Len(t, messages, 1, "agent placeholder should be removed after error")
+	assert.Equal(t, "hello", messages[0].Message)
+
+	// When: a late chunk arrives after the prompt has already errored
+	mock.SimulateChunks("late response data")
+
+	// Then: the user message is not corrupted
+	messages = conv.Messages()
+	require.Len(t, messages, 1, "no new messages should appear from a late chunk")
+	assert.Equal(t, "hello", messages[0].Message)
+	assert.Equal(t, screentracker.ConversationRoleUser, messages[0].Role)
+}