From a014ba0e49596c81a93be0fc3cd5c34fb3fbcdbf Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Fri, 23 Oct 2026 01:14:55 +0800
Subject: [PATCH] test(pluginhost,handlers): fix flaky races in plugin load and
 stream cancel tests

- Add a loadCleanupDoneHook callback in pluginhost.Host so canceled plugin load tests wait until the in-flight load token is cleared after client shutdown.
- Drain dataChan after canceling the context in TestHandlerPluginExecutorPublishesStreamCancellation before closing the chunk channel to avoid racing ctx.Done() against normal stream EOF.

Closes: #6885
---
 internal/pluginhost/host.go                    |  5 +++++
 internal/pluginhost/host_test.go               | 18 ++++++++++--------
 .../handlers_plugin_executor_usage_test.go     |  6 +++++-
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go
index 17456e7243..f23854cfde 100644
--- a/internal/pluginhost/host.go
+++ b/internal/pluginhost/host.go
@@ -92,6 +92,8 @@ type Host struct {
 	modelStreams           *modelStreamBridge
 	callbackContexts       *callbackContextRegistry
 	snapshot               atomic.Value
+	// loadCleanupDoneHook, set by tests, runs after a late load cleanup dropped its load token.
+	loadCleanupDoneHook func()
 }
 
 func New() *Host {
@@ -578,6 +580,9 @@ func (h *Host) finishPluginLoadCleanup(id string, request *pluginLoadRequest, lo
 	go func() {
 		h.discardLoadedPlugin(loaded)
 		h.clearLoadingRequest(id, request)
+		if h.loadCleanupDoneHook != nil {
+			h.loadCleanupDoneHook()
+		}
 	}()
 }
 
diff --git a/internal/pluginhost/host_test.go b/internal/pluginhost/host_test.go
index 9531687b59..1141893345 100644
--- a/internal/pluginhost/host_test.go
+++ b/internal/pluginhost/host_test.go
@@ -1695,6 +1695,8 @@ func TestHostCanceledInitializationDiscardsBlockedClient(t *testing.T) {
 		registration: validTestPlugin("alpha"),
 	}
 	h := NewForTest(&blockingHostCallLoader{client: client})
+	cleanupDone := make(chan struct{})
+	h.loadCleanupDoneHook = func() { close(cleanupDone) }
 	cfg := &config.Config{Plugins: config.PluginsConfig{
 		Enabled: true,
 		Dir:     makePluginDir(t, "alpha"),
@@ -1714,10 +1716,9 @@ func TestHostCanceledInitializationDiscardsBlockedClient(t *testing.T) {
 	}
 
 	close(client.release)
-	deadline := time.Now().Add(time.Second)
-	for client.shutdown.Load() == 0 && time.Now().Before(deadline) {
-		time.Sleep(time.Millisecond)
-	}
+	// Late cleanup shuts the client down first and drops the load token after;
+	// wait for the end of the cleanup, not for the shutdown alone.
+	waitForHostTestSignal(t, cleanupDone, "late load cleanup")
 	if got := client.shutdown.Load(); got != 1 {
 		t.Fatalf("blocked initialization client shutdown calls = %d, want 1", got)
 	}
@@ -1734,6 +1735,8 @@ func TestHostCancellationUnderMutationLockDoesNotInsertLoadedPlugin(t *testing.T
 		registration: validTestPlugin("alpha"),
 	}
 	h := NewForTest(&blockingHostCallLoader{client: client})
+	cleanupDone := make(chan struct{})
+	h.loadCleanupDoneHook = func() { close(cleanupDone) }
 	cfg := &config.Config{Plugins: config.PluginsConfig{
 		Enabled: true,
 		Dir:     makePluginDir(t, "alpha"),
@@ -1754,10 +1757,9 @@ func TestHostCancellationUnderMutationLockDoesNotInsertLoadedPlugin(t *testing.T
 	h.mu.Unlock()
 	waitForHostTestSignal(t, applyDone, "canceled plugin apply")
 
-	deadline := time.Now().Add(time.Second)
-	for client.shutdown.Load() == 0 && time.Now().Before(deadline) {
-		time.Sleep(time.Millisecond)
-	}
+	// Late cleanup shuts the client down first and drops the load token after;
+	// wait for the end of the cleanup, not for the shutdown alone.
+	waitForHostTestSignal(t, cleanupDone, "late load cleanup")
 	if got := client.shutdown.Load(); got != 1 {
 		t.Fatalf("late client shutdown calls = %d, want 1", got)
 	}
diff --git a/sdk/api/handlers/handlers_plugin_executor_usage_test.go b/sdk/api/handlers/handlers_plugin_executor_usage_test.go
index c83850beaa..211f3c6dd4 100644
--- a/sdk/api/handlers/handlers_plugin_executor_usage_test.go
+++ b/sdk/api/handlers/handlers_plugin_executor_usage_test.go
@@ -550,9 +550,13 @@ func TestHandlerPluginExecutorPublishesStreamCancellation(t *testing.T) {
 	handler.SetModelRouterHost(mockHost)
 
 	dataChan, _, _ := handler.ExecuteStreamWithAuthManager(ctx, "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "")
-	// Receive first chunk then cancel context
+	// Receive the first chunk, then cancel. The chunk channel stays open until the stream
+	// goroutine has exited: closed early, it races ctx.Done() in nextStreamChunk's select and the
+	// stream can end as a normal completion (about 1 in 1000 runs under -race).
 	<-dataChan
 	cancel()
+	for range dataChan {
+	}
 	close(chunks)
 
 	record := plugin.waitRecord(t)