fix(coding/acp): deliver the controller cancel before a cancelled turn ends

internal/coding/acp/cancel.go

@@ -0,0 +1,30 @@
+package acp
+
+import "context"
+
+// deliverCancelOnDone arranges for deliver to be called once when ctx is done and
+// returns the function that disarms it. The returned function guarantees the
+// delivery has happened by the time it returns whenever ctx is already done:
+// context.AfterFunc runs its callback on its own goroutine, so a prompt that
+// returns on its own cancellation can otherwise prevent the delivery and leave the
+// controller running with no cancel.
+func deliverCancelOnDone(ctx context.Context, deliver func()) func() {
+	delivered := make(chan struct{})
+	stop := context.AfterFunc(ctx, func() {
+		defer close(delivered)
+
+		deliver()
+	})
+
+	return func() {
+		switch {
+		case !stop():
+			// The callback is on its goroutine: let it finish before the caller moves on.
+			<-delivered
+		case ctx.Err() != nil:
+			// The callback was prevented and the context is done anyway, which is what
+			// happens when the prompt returns on its own cancellation first.
+			deliver()
+		}
+	}
+}

internal/coding/acp/cancel_test.go

@@ -0,0 +1,63 @@
+package acp
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestDeliverCancelOnDone pins the guarantee a cancelled turn relies on: the
+// delivery has happened by the time the disarm returns, whichever goroutine the
+// scheduler ran first, and exactly once.
+func TestDeliverCancelOnDone(t *testing.T) {
+	t.Parallel()
+
+	t.Run("cancelled before the disarm", func(t *testing.T) {
+		t.Parallel()
+
+		// The loop is the point: the callback may have run before the disarm or be
+		// prevented by it, and both orderings have to deliver exactly once.
+		for range 100 {
+			ctx, cancel := context.WithCancel(t.Context())
+
+			var deliveries atomic.Int64
+
+			stop := deliverCancelOnDone(ctx, func() { deliveries.Add(1) })
+
+			cancel()
+			stop()
+
+			require.EqualValues(t, 1, deliveries.Load())
+		}
+	})
+
+	t.Run("disarmed before the cancel delivers nothing", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(t.Context())
+
+		var deliveries atomic.Int64
+
+		stop := deliverCancelOnDone(ctx, func() { deliveries.Add(1) })
+
+		stop()
+		cancel()
+
+		assert.Zero(t, deliveries.Load(), "a turn that ends on its own must not cancel the controller")
+	})
+
+	t.Run("an open context delivers nothing", func(t *testing.T) {
+		t.Parallel()
+
+		var deliveries atomic.Int64
+
+		stop := deliverCancelOnDone(t.Context(), func() { deliveries.Add(1) })
+
+		stop()
+
+		assert.Zero(t, deliveries.Load())
+	})
+}

internal/coding/acp/session.go

@@ -421,7 +421,7 @@ func (s *session) beginTurn(
 	s.idle = make(chan struct{})
 	s.mu.Unlock()
 
-	stopCancel := context.AfterFunc(turnCtx, func() {
+	stopCancel := deliverCancelOnDone(turnCtx, func() {
 		_ = s.controller.Cancel()
 	})
 	finish := func() {

internal/coding/acp/session_test.go

@@ -177,12 +177,13 @@ func TestSessionRequestContextCancellationCancelsRuntime(t *testing.T) {
 	prompted := <-result
 	require.NoError(t, prompted.err)
 	assert.Equal(t, acpsdk.StopReasonCancelled, prompted.stop)
-	assert.Eventually(t, func() bool {
-		controller.mu.Lock()
-		defer controller.mu.Unlock()
+	// The turn delivers the cancellation before it finishes, so this needs no wait: it
+	// used to be an Eventually that ran its full thirty seconds whenever the prompt
+	// returned before the AfterFunc callback had been scheduled.
+	controller.mu.Lock()
+	defer controller.mu.Unlock()
 
-		return controller.cancelCount > 0
-	}, 30*time.Second, 10*time.Millisecond)
+	assert.Equal(t, 1, controller.cancelCount)
 }
 
 func TestSessionModeControlSerializesAfterActivePrompt(t *testing.T) {