Cancel AG-UI SSE request when the consumer stops early

[PratikDhanave]

Summary

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 (yield returns false), 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 indefinitely when ReadTimeout is 0). agent.Run wraps the context only with context.WithValue, so nothing else tore the request down.

Fix

Derive a cancellable context in the run closure and defer cancel(), so every early return releases the reader goroutine and response body:

streamCtx, cancel := context.WithCancel(ctx)
defer cancel()
frames, errs, err := p.client.Stream(aguiSSEClient.StreamOptions{Context: streamCtx, Payload: payload})

The sibling a2a (pull-based iterator with defer body.Close()) and copilot (deferred unsubscribe/Disconnect) providers already tear down on early return; this brings aguiprovider in line.

Public API

No exported symbols change.

Tests

Adds TestAGUIAgentRun_EarlyStreamStop_CancelsUpstreamRequest: the test server sends one event, keeps the stream open, and waits for the client to disconnect. It fails before this change (the disconnect never arrives; the handler's safety timeout expires) and passes after (early stop cancels the request, server observes the disconnect promptly).

go test -race ./provider/aguiprovider passes; go build ./... and go vet clean; gofumpt clean.

[Copilot]

Pull request overview

Fixes a resource-leak scenario in the AG-UI SSE streaming provider by ensuring upstream SSE requests are cancelled when the downstream consumer stops early (or when decode/event errors cause an early return), preventing the SSE client’s background reader goroutine and HTTP response body from lingering until timeout.

Changes:

• Wrap the provider’s streaming request context with context.WithCancel and defer cancel() to guarantee teardown on any early return.
• Add a regression test that verifies the server observes a client disconnect promptly when the consumer breaks out of the stream early.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
provider/aguiprovider/agui.go	Ensures SSE request cancellation on early stream termination to avoid leaking the client reader goroutine/response body.
provider/aguiprovider/agui_leak_test.go	Adds a targeted regression test confirming early consumer stop triggers upstream disconnect/cancellation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.