From 57b08faa00f18fb21e4c12b8341e4fdb25a711a7 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Thu, 16 Jul 2026 11:52:20 +0530 Subject: [PATCH] Cancel AG-UI SSE request when the consumer stops early MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The aguiprovider streaming run passed the caller's context straight to the SSE client's Stream. On any early return from the run closure — the consumer stopping mid-stream, or a decode/event error — the request was never cancelled, so the client's background reader goroutine (which owns the HTTP response body) leaked until the client's ReadTimeout (5 minutes by default, or forever when ReadTimeout is 0). agent.Run does not supply a cancellable context, so nothing else tore the request down. Derive a cancellable context in the run closure and defer its cancel, so every early return releases the reader goroutine and response body. The sibling a2a and copilot providers already tear down on early return, so this brings aguiprovider in line. Adds TestAGUIAgentRun_EarlyStreamStop_CancelsUpstreamRequest: the server keeps the stream open after one event and waits for the client to disconnect; the test fails (times out on disconnect) before this change and passes after. --- provider/aguiprovider/agui.go | 11 ++++- provider/aguiprovider/agui_leak_test.go | 63 +++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 provider/aguiprovider/agui_leak_test.go diff --git a/provider/aguiprovider/agui.go b/provider/aguiprovider/agui.go index 4311638d..0b8ee0a4 100644 --- a/provider/aguiprovider/agui.go +++ b/provider/aguiprovider/agui.go @@ -95,7 +95,16 @@ func (p *provider) run(ctx context.Context, messages []*message.Message, options ForwardedProps: map[string]any{}, } - frames, errs, err := p.client.Stream(aguiSSEClient.StreamOptions{Context: ctx, Payload: payload}) + // Derive a cancellable context so that any early return from this + // closure — the consumer stopping mid-stream, or a decode/event error — + // tears down the SSE request and releases the client's reader goroutine + // and HTTP response body. agent.Run does not supply a cancellable + // context, so without this the reader goroutine leaks until the client's + // ReadTimeout (or indefinitely when ReadTimeout is 0). + streamCtx, cancel := context.WithCancel(ctx) + defer cancel() + + frames, errs, err := p.client.Stream(aguiSSEClient.StreamOptions{Context: streamCtx, Payload: payload}) if err != nil { yield(nil, err) return diff --git a/provider/aguiprovider/agui_leak_test.go b/provider/aguiprovider/agui_leak_test.go new file mode 100644 index 00000000..86d1773b --- /dev/null +++ b/provider/aguiprovider/agui_leak_test.go @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft. All rights reserved. + +package aguiprovider_test + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + aguiEvents "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/core/events" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/provider/aguiprovider" +) + +// TestAGUIAgentRun_EarlyStreamStop_CancelsUpstreamRequest verifies that when a +// consumer stops ranging over the streaming response early, the provider tears +// down the SSE request rather than leaking the client's reader goroutine and +// HTTP response body. +// +// The server sends one event then keeps the stream open and waits for the +// client to disconnect. With the fix, the early stop cancels the request and +// the server observes the disconnect promptly; without it the request lingers +// until the client's ReadTimeout and the disconnect never arrives in time. +func TestAGUIAgentRun_EarlyStreamStop_CancelsUpstreamRequest(t *testing.T) { + disconnected := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + writeSSE(t, w, aguiEvents.NewRunStartedEvent("thread-1", "run-1")) + writeSSE(t, w, aguiEvents.NewTextMessageStartEvent("msg-1", aguiEvents.WithRole("assistant"))) + writeSSE(t, w, aguiEvents.NewTextMessageContentEvent("msg-1", "Hello")) + // Deliberately do not finish the run: keep the stream open and wait for + // the client to disconnect. The safety timeout lets the handler return + // (so the test server can close) if the disconnect never comes. + select { + case <-r.Context().Done(): + close(disconnected) + case <-time.After(10 * time.Second): + } + })) + defer server.Close() + + a := aguiprovider.NewAgent(newTestClient(server.URL), aguiprovider.AgentConfig{}) + + // Consume a single update, then stop. Breaking out of the range makes the + // downstream yield return false, which must propagate to the provider and + // cancel the upstream SSE request. + for _, err := range a.RunText(context.Background(), "hi", agent.Stream(true)) { + if err != nil { + t.Fatalf("unexpected error before first update: %v", err) + } + break + } + + select { + case <-disconnected: + // Fixed: the early stop cancelled the SSE request. + case <-time.After(3 * time.Second): + t.Fatal("server did not observe client disconnect after early stream termination: " + + "the SSE reader goroutine and HTTP response body leaked") + } +}