Last updated 2026-05-17

Context & Cancellation

Every llmrouter call takes a context.Context, and that context is the master switch for everything downstream of the call. Cancel it and the entire pipeline — HTTP request, streaming response, internal producer goroutine — winds down deterministically. This page covers the contract, the propagation rules, and the patterns for gateways and deadline-driven workloads.

The contract

The signature is:

func (p *Provider) CompletionStream(
    ctx context.Context,
    req ChatRequest,
) (*Stream, error)

The ctx you pass in is wired into three places at once:

The outgoing HTTP request, via http.NewRequestWithContext. Cancellation cancels connect, TLS handshake, headers, and body read.
The internal producer goroutine, via a child context that llmrouter.NewStream derives.
The buffered channel send inside the producer, via a select { case s.chunks <- c: case <-ctx.Done(): }.

All three observe the same cancellation signal. That means there is exactly one place to "pull the plug" and it cleans up everything.

Cancellation propagation, end to end

What happens, in order, when you cancel the parent context:

The http.Client's in-flight read returns an error (typically wrapping context.Canceled).
The producer goroutine's SSE scanner loop notices either via ctx.Err() on its next iteration or via the underlying read failure.
The producer calls hooks.Finish(ctx.Err()), which stores the error and closes the chunks channel.
The consumer's range stream.Chunks() exits.
stream.Err() returns context.Canceled (or context.DeadlineExceeded).
The underlying TCP socket closes. No orphan goroutine, no orphan connection.

Or in code, from stream.go:

func (s *Stream) sendChunk(ctx context.Context, c Chunk) bool {
    select {
    case s.chunks <- c:
        return true
    case <-ctx.Done():
        return false
    }
}

When sendChunk returns false, the producer immediately calls hooks.Finish(ctx.Err()) and exits the goroutine. There's no "drain and retry" path — once the context is cancelled, the producer is done.

`Stream.Cancel()`

Stream.Cancel() cancels an internal child context that NewStream derived from the parent:

func NewStream(parent context.Context) (*Stream, context.Context, ProducerHooks) {
    ctx, cancel := context.WithCancel(parent)
    s := newStream(cancel)
    return s, ctx, ProducerHooks{ /* ... */ }
}

Calling stream.Cancel() is equivalent to cancelling the parent context, but only for this one stream. Use it when:

You want to abort one stream without affecting siblings sharing the same parent context (for example, in a fan-out pattern where you keep the fastest response and discard the others).
You want a clean "stop here" inside a consumer loop without needing access to the parent cancel closure.

The method is idempotent; multiple calls are safe.

Why this matters for gateways

If you're building a gateway, the most expensive bug you can ship is: a client disconnects, the gateway doesn't notice, the upstream HTTP call keeps streaming, and you pay for tokens nobody sees. Multiply by every disconnect on every long completion, every day.

The fix is one line: pass r.Context() into CompletionStream. Go's net/http cancels r.Context() when the client disconnects, and that cancellation propagates straight through:

func handler(p llmrouter.Provider) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // r.Context() is cancelled when the client TCP connection
        // closes. Wiring it into CompletionStream means the upstream
        // OpenAI request also closes — no orphan tokens billed.
        stream, err := p.CompletionStream(r.Context(), buildRequest(r))
        // ...
    }
}

The orphan-tokens test. Start a long completion, kill the client mid-stream (Ctrl-C the curl), and watch your upstream provider's billing dashboard. If usage keeps climbing, you're not wiring r.Context() through. If it stops within a couple of seconds, you are.

Timeouts

Use context.WithTimeout (or context.WithDeadline) to put a hard ceiling on a request:

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

stream, err := p.CompletionStream(ctx, req)
// If 30s elapse, ctx is cancelled with DeadlineExceeded.

Always defer cancel(). Without it, the context's timer goroutine doesn't release until the deadline fires, which matters if your function returns early.

See Configuration & Options for the relationship between context.WithTimeout and WithTimeout: the former is per-call; the latter is a per-provider HTTP-client setting and a defense-in-depth backstop. Use both.

Deadlines propagate naturally

A deadline in a parent context applies to every child. If you set a 10-second deadline at your HTTP handler entry and call multiple providers in sequence, all of them share that budget:

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

// All three calls compete for the same 10-second budget.
fast, _   := callProvider(ctx, primary, req)
medium, _ := callProvider(ctx, fallback, req)
slow, _   := callProvider(ctx, archive, req)

Use context.WithDeadline when you have an absolute time rather than a duration — e.g., a request-id-bound SLO:

deadline := requestStartedAt.Add(15 * time.Second)
ctx, cancel := context.WithDeadline(r.Context(), deadline)
defer cancel()

Example: client-disconnect-aware HTTP handler

A minimal SSE forwarder that respects client disconnect end-to-end. Note three things: r.Context() goes into CompletionStream; the loop checks r.Context().Done() for an aborted client; and stream.Err() is logged but not surfaced once we've started writing SSE (the headers are gone).

package main

import (
    "encoding/json"
    "errors"
    "fmt"
    "io"
    "log"
    "net/http"
    "os"

    "github.com/elloloop/llmrouter"
    "github.com/elloloop/llmrouter/providers/openai"
)

func main() {
    p, err := openai.New(llmrouter.WithAPIKey(os.Getenv("OPENAI_API_KEY")))
    if err != nil {
        log.Fatal(err)
    }
    http.HandleFunc("/chat", makeHandler(p))
    log.Fatal(http.ListenAndServe(":8080", nil))
}

func makeHandler(p llmrouter.Provider) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        body, err := io.ReadAll(r.Body)
        if err != nil {
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }
        var head struct{ Model string `json:"model"` }
        _ = json.Unmarshal(body, &head)

        // r.Context() is the master switch. If the client closes the
        // connection, the upstream HTTP call cancels automatically.
        stream, err := p.CompletionStream(r.Context(), llmrouter.ChatRequest{
            Model: head.Model,
            Raw:   body,
        })
        if err != nil {
            var ue *llmrouter.ErrUpstream
            if errors.As(err, &ue) {
                http.Error(w, ue.Body, ue.StatusCode)
                return
            }
            http.Error(w, err.Error(), http.StatusBadGateway)
            return
        }

        w.Header().Set("Content-Type", "text/event-stream")
        w.Header().Set("Cache-Control", "no-cache")
        flusher, _ := w.(http.Flusher)

        for chunk := range stream.Chunks() {
            // If the client is already gone, stop writing. The
            // upstream stream is already cancelling via r.Context();
            // this just avoids one extra Write to a dead socket.
            select {
            case <-r.Context().Done():
                log.Printf("client disconnected mid-stream")
                return
            default:
            }
            fmt.Fprintf(w, "data: %s\n\n", chunk.Raw)
            if flusher != nil {
                flusher.Flush()
            }
        }
        fmt.Fprint(w, "data: [DONE]\n\n")
        if flusher != nil {
            flusher.Flush()
        }

        if err := stream.Err(); err != nil {
            // Already wrote 200 OK headers; can't change status now.
            // Just log.
            if !errors.Is(err, r.Context().Err()) {
                log.Printf("stream error: %v", err)
            }
        }
    }
}

Example: a budget-aware streamer

Counting tokens (or characters as a stand-in) and cancelling when a per-request cap is exceeded. The pattern uses stream.Cancel() so the parent context stays usable for other work in the same handler:

package main

import (
    "context"
    "errors"
    "fmt"
    "io"
    "log"
    "os"
    "strings"
    "time"

    "github.com/elloloop/llmrouter"
    "github.com/elloloop/llmrouter/providers/openai"
)

// ErrBudgetExceeded is returned by streamWithBudget when the consumer
// stops the stream early because it ran out of budget.
var ErrBudgetExceeded = errors.New("budget exceeded")

// streamWithBudget writes chunks to w until the cumulative byte count
// exceeds maxBytes, then cancels the stream and returns
// ErrBudgetExceeded. Any text already written stays written.
func streamWithBudget(
    ctx context.Context,
    p llmrouter.Provider,
    req llmrouter.ChatRequest,
    w io.Writer,
    maxBytes int,
) error {
    stream, err := p.CompletionStream(ctx, req)
    if err != nil {
        return err
    }

    var (
        written int
        capped  bool
    )
    for chunk := range stream.Chunks() {
        for _, c := range chunk.Choices {
            if _, werr := io.WriteString(w, c.Delta.Content); werr != nil {
                stream.Cancel()
                // Keep draining so the goroutine exits cleanly.
                continue
            }
            written += len(c.Delta.Content)
        }
        if written >= maxBytes && !capped {
            capped = true
            stream.Cancel()
            // Don't break — drain until close so Err() unblocks.
        }
    }

    if serr := stream.Err(); serr != nil {
        if capped && errors.Is(serr, context.Canceled) {
            return ErrBudgetExceeded
        }
        return serr
    }
    return nil
}

func main() {
    p, err := openai.New(llmrouter.WithAPIKey(os.Getenv("OPENAI_API_KEY")))
    if err != nil {
        log.Fatal(err)
    }
    ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
    defer cancel()

    req := llmrouter.ChatRequest{
        Model: "gpt-4o-mini",
        Messages: []llmrouter.Message{
            llmrouter.TextMessage("user", "write a paragraph about goroutines"),
        },
    }
    var buf strings.Builder
    err = streamWithBudget(ctx, p, req, &buf, 200)
    fmt.Println(buf.String())
    if errors.Is(err, ErrBudgetExceeded) {
        fmt.Println("\n[truncated at budget]")
    } else if err != nil {
        log.Fatal(err)
    }
}

Drain after Cancel. The pattern above doesn't break out of the range loop after calling stream.Cancel(). The producer goroutine might be blocked on a buffered send; letting the loop continue until the channel closes guarantees the goroutine exits promptly and stream.Err() doesn't deadlock.

Common pitfalls

Using context.Background() in an HTTP handler. Always use r.Context(). Otherwise, client disconnects are invisible to your upstream call.
Forgetting defer cancel(). The cancel function returned by WithTimeout / WithCancel must be called or you leak the context's goroutine until the deadline fires.
Calling Err() before draining Chunks(). Err() blocks until the producer finishes. If you haven't been reading chunks, the producer is blocked on a buffered send and you'll deadlock.
Treating context.Canceled as an error. It's typically an intentional caller action — log at debug, not error.

Next steps

The Streaming Model — for the broader picture of how the producer goroutine interacts with the consumer.
Error Handling — how cancellation surfaces alongside upstream and network errors.
Byte Passthrough — the proxy example shows the recommended pattern for wiring r.Context() end-to-end.