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)