fix(adk): fix ack channel leak in concurrent preempt during TurnLoop shutdown

When drainAll runs during TurnLoop cleanup, a concurrent Push caller
that has already called holdRunLoop but not yet requestPreempt can
add an ack channel to pendingAckList after drainAll clears it. This
orphaned ack is never closed, causing the Push caller to hang.

Add a `drained` flag to preemptSignal. drainAll sets it, and
requestPreempt checks it — if drained, the ack is closed immediately
instead of being appended to pendingAckList.

Also fix the test to stop the loop concurrently with wg.Wait, since
the run loop may be blocked on buffer.Receive after processing all
preempts, and Stop is needed to unblock it and trigger drainAll.

Change-Id: I8d58b0e1478d2ae6f89e06d420b9252e43aa9cd6

Author: shentongmartin

adk/turn_loop.go

@@ -191,6 +191,7 @@ type preemptSignal struct {
 	agentCancelOpts  []AgentCancelOption
 	pendingAckList   []chan struct{}
 	notify           chan struct{}
+	drained          bool
 
 	currentTC     any
 	currentRunCtx context.Context
@@ -223,12 +224,13 @@ func (s *preemptSignal) holdAndGetTurn() (context.Context, any) {
 }
 
 // requestPreempt records a preempt request and wakes both waiters.
-// If holdCount is 0, no one is listening — close the ack immediately as a no-op.
+// If holdCount is 0 or the signal has been drained, no one is listening —
+// close the ack immediately as a no-op.
 func (s *preemptSignal) requestPreempt(ack chan struct{}, opts ...AgentCancelOption) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	if s.holdCount <= 0 {
+	if s.drained || s.holdCount <= 0 {
 		if ack != nil {
 			close(ack)
 		}
@@ -343,11 +345,13 @@ func (s *preemptSignal) endTurnAndUnhold() {
 // drainAll forcefully resets all preemptSignal state and closes any pending
 // ack channels. Called during TurnLoop cleanup to prevent ack channels from
 // leaking when the run loop exits (e.g. due to Stop) while a Push caller
-// still holds a reference.
+// still holds a reference. After drainAll, any subsequent holdRunLoop or
+// requestPreempt calls will be no-ops that close the ack immediately.
 func (s *preemptSignal) drainAll() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	s.drained = true
 	s.holdCount = 0
 	s.currentTC = nil
 	s.currentRunCtx = nil

adk/turn_loop_test.go

@@ -3122,17 +3122,22 @@ func TestTurnLoop_ConcurrentPreemptsDuringTurn(t *testing.T) {
 			if ok && ack != nil {
 				select {
 				case <-ack:
-				case <-time.After(30 * time.Second):
+				case <-time.After(5 * time.Second):
 					t.Error("ack channel not closed within timeout")
 				}
 			}
 		}(i)
 	}
 
-	wg.Wait()
-	time.Sleep(200 * time.Millisecond)
+	// Stop the loop concurrently. The run loop may be blocked on
+	// buffer.Receive after processing all preempts; Stop unblocks it
+	// and triggers drainAll which closes any orphaned ack channels.
+	go func() {
+		time.Sleep(500 * time.Millisecond)
+		loop.Stop(WithImmediate())
+	}()
 
-	loop.Stop(WithImmediate())
+	wg.Wait()
 	result := loop.Wait()
 	assert.NoError(t, result.ExitReason)
 	assert.True(t, atomic.LoadInt32(&genInputCount) >= 2, "should have had at least the initial turn + one preempted turn")