Canceling in-flight LLM streams with context.Context

A streaming LLM response can run for tens of seconds and emit thousands of tokens. Using go context cancel llm stream patterns in Go is the only clean way to abort that work when the caller closes the connection, hits a deadline, or your service receives a shutdown signal—without spawning orphaned goroutines that keep pulling tokens you’ll never use.

Step 1: Recognize the leak in a naive proxy

Most LLM streaming code looks like this: open a chat completion stream, loop on Recv(), write each token to the HTTP response writer. If the client disconnects, the server often keeps looping until the model finishes because nothing told the upstream request to stop.

The wasted tokens are real, but the bigger problem is resource buildup. Each hung stream holds a goroutine, a TCP connection, and a read buffer. Under load, a fleet of zombie streams will exhaust file descriptors before you notice.

The fix is to bind the upstream call to a context.Context that mirrors the lifecycle of the downstream request.

Step 2: Capture the inbound request context

Every http.Request in Go carries a context that the server cancels when the client connection closes or the server begins graceful shutdown. Use it directly.

func handleChat(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context() // cancels when client hangs up or server shuts down
    // ...
}

If you need a custom timeout on top of client disconnect, derive a new context:

ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
defer cancel()

This is the foundation of any go context cancel llm stream implementation: the context is the single signal that propagates everywhere.

Step 3: Open the upstream stream with that context

Use an OpenAI-compatible client. The go-openai library accepts a context on every stream call. Point the base URL at your provider—an OpenAI-compatible gateway such as n4n.ai lets you hit 240+ models with the same code and get automatic fallback if a provider degrades.

import (
    "github.com/sashabaranov/go-openai"
)

cfg := openai.DefaultConfig(os.Getenv("LLM_API_KEY"))
cfg.BaseURL = "https://api.n4n.ai/v1" // or api.openai.com/v1

client := openai.NewClientWithConfig(cfg)

req := openai.ChatCompletionRequest{
    Model:  "gpt-4o-mini",
    Stream: true,
    Messages: []openai.ChatMessage{
        {Role: openai.ChatMessageRoleUser, Content: "Explain Raft consensus"},
    },
}

stream, err := client.CreateChatCompletionStream(ctx, req)
if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
}
defer stream.Close()

Because ctx is the request context (or a derivative), cancelling it will tear down the underlying HTTP connection to the LLM provider. That is the core mechanism.

Step 4: Isolate the reader goroutine behind a channel

Never call stream.Recv() from the same goroutine that writes to the client. If the client disconnects, the write fails, but the reader would block on the next Recv(). Split the work:

type token struct {
    text string
    err  error
}

ch := make(chan token, 8)
go func() {
    defer close(ch)
    for {
        resp, err := stream.Recv()
        if err != nil {
            ch <- token{err: err}
            return
        }
        if len(resp.Choices) > 0 {
            ch <- token{text: resp.Choices[0].Delta.Content}
        }
    }
}()

The reader goroutine exits when Recv() returns an error. That error will be context.Canceled if the context was cancelled mid-stream. The bounded channel prevents the goroutine from blocking on a full queue after the writer stops consuming.

Step 5: Drive the response writer with a select on ctx.Done

Now the handler goroutine writes tokens until either the stream ends, the client disappears, or the timeout fires.

w.Header().Set("Content-Type", "text/event-stream")
flusher, ok := w.(http.Flusher)
if !ok {
    http.Error(w, "streaming unsupported", http.StatusInternalServerError)
    return
}

for {
    select {
    case <-ctx.Done():
        // Client gone or timeout: stop writing, return.
        // defer stream.Close() and cancel() handle cleanup.
        return
    case t, open := <-ch:
        if !open {
            return // stream finished naturally
        }
        if t.err != nil {
            if ctx.Err() == nil {
                // real upstream error, not cancellation
                log.Printf("upstream error: %v", t.err)
            }
            return
        }
        if _, err := w.Write([]byte(t.text)); err != nil {
            return // write failed => client disconnected
        }
        flusher.Flush()
    }
}

This select is where the go context cancel llm stream contract pays off: the moment ctx.Done() closes, the loop returns, the deferred stream.Close() runs, and the reader goroutine sees context.Canceled on its next Recv().

Step 6: Explicitly close the upstream stream

The defer stream.Close() in Step 3 is mandatory. The OpenAI streaming client holds an http.Response body; without Close(), the TCP connection stays open until the provider times out. In a proxy, add both:

defer cancel()       // cancels ctx, signals reader
defer stream.Close() // releases upstream connection

If you spawned the reader goroutine, it will observe the cancellation and exit. No goroutine leak.

Step 7: Verify cancellation end to end

Run the server locally and hit it with a slow prompt. Then kill the client mid-stream and confirm the server stops pulling tokens.

# terminal 1: start server
go run main.go

# terminal 2: start a curl, then Ctrl-C after 2 seconds
curl -N -X POST localhost:8080/chat \
  -d '{"prompt":"write a 5000 word essay"}' \
  -H 'Content-Type: application/json'
# press Ctrl-C

Add a log line in the reader goroutine when it receives context.Canceled:

if errors.Is(err, context.Canceled) {
    log.Println("stream reader exited: client cancelled")
}

You should see that log within a second of killing curl. If you instead see the server logging the full completion after the client is gone, your context is not wired to the upstream call.

For a stricter test, force a timeout:

ctx, cancel := context.WithTimeout(r.Context(), 500*time.Millisecond)

A model that takes longer than 500ms to first token should trigger the same cancellation path. Watch the process with goroutine profiles (pprof) to confirm the count drops back to baseline after the request ends.

Edge cases worth handling

Partial writes: If w.Write fails because the client disconnected, the subsequent ctx.Done() will fire; the select catches it on the next iteration. Don’t try to write an error frame after a failed write—the connection is already dead.

Provider that ignores close: Some upstreams keep sending TCP data even after you close the body. The go-openai client’s Close() drains a small amount, but your reader goroutine exiting is what matters; the OS reclaims the socket when the process drops references.

Multiple concurrent clients: Each request gets its own context, stream, and goroutine pair. The pattern scales linearly; there is no shared mutable state.

Using go context cancel llm stream design in your gateway or proxy keeps token spend bounded and your goroutine count honest. The code above is production-minimal—drop it into a handler and add your own auth and model routing.