fix: fallback to Request.Context in Context

.Stream when CloseNotify is nil
This commit is contained in:
jassus23 2026-05-04 11:13:50 +03:00
parent 1f5552f6d2
commit 3317f91a1b
2 changed files with 25 additions and 0 deletions

View File

@ -1328,10 +1328,16 @@ func (c *Context) SSEvent(name string, message any) {
func (c *Context) Stream(step func(w io.Writer) bool) bool {
w := c.Writer
clientGone := w.CloseNotify()
var requestGone <-chan struct{}
if c.Request != nil {
requestGone = c.Request.Context().Done()
}
for {
select {
case <-clientGone:
return true
case <-requestGone:
return true
default:
keepOpen := step(w)
w.Flush()

View File

@ -3078,6 +3078,25 @@ func TestContextStreamWithClientGone(t *testing.T) {
assert.Equal(t, "test", w.Body.String())
}
func TestContextStreamWithRequestContextDone(t *testing.T) {
w := httptest.NewRecorder()
c, _ := CreateTestContext(w)
reqCtx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodGet, "/", nil).WithContext(reqCtx)
disconnected := c.Stream(func(writer io.Writer) bool {
_, err := writer.Write([]byte("test"))
require.NoError(t, err)
return true
})
assert.True(t, disconnected)
assert.Empty(t, w.Body.String())
}
func TestContextResetInHandler(t *testing.T) {
w := CreateTestResponseRecorder()
c, _ := CreateTestContext(w)