test(oauth): pin the callback server's header bound and its join on Serve

Two of the three things CallbackServer does had no test that noticed their
removal, the same as when #1051 introduced them on LoopbackListener: the
five-second header read bound, and Close waiting for the Serve goroutine.

The header test stalls a client mid-header and never calls Close, so it can
only pass on ReadHeaderTimeout; the bound is a package var so the test can
shorten it. It synchronises on StateNew. Go reports StateActive only after a
request read returns, and a stalled header read returns only when the header
timeout fires, so a helper waiting for StateActive depended on the bound under
test and turned its removal into a SETUP INVALID instead of a failure.

The join test uses the case where the join matters: with no connections open,
Shutdown closes the listener and returns at once, so without the join Close can
return before Serve does. Without it the test failed 10 runs out of 10.

internal/mcp/oauth.go

@@ -494,31 +494,28 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
 		err  error
 	}
 	resultChan := make(chan callbackResult, 1)
-	server := &http.Server{
-		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if r.URL.Path != "/callback" {
-				http.NotFound(w, r)
-				return
-			}
-			code, parseErr := flow.parseCallback(r.URL.Query())
-			if parseErr != nil {
-				w.WriteHeader(http.StatusBadRequest)
-				_, _ = io.WriteString(w, "Authorization failed. You may close this window.")
-			} else {
-				_, _ = io.WriteString(w, "Authorization complete. You may close this window.")
-			}
-			select {
-			case resultChan <- callbackResult{code: code, err: parseErr}:
-			default:
-			}
-		}),
-	}
-	go func() { _ = server.Serve(listener) }()
-	defer func() {
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
-		defer cancel()
-		_ = server.Shutdown(shutdownCtx)
-	}()
+	// The shared callback server bounds header reads, force-closes a client
+	// still connected when shutdown runs out of time, and waits for Serve to
+	// return. The bare http.Server this used to build did none of that, so a
+	// client stalled mid-header outlived Login (#1028).
+	server := oauth.StartCallbackServer(listener, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/callback" {
+			http.NotFound(w, r)
+			return
+		}
+		code, parseErr := flow.parseCallback(r.URL.Query())
+		if parseErr != nil {
+			w.WriteHeader(http.StatusBadRequest)
+			_, _ = io.WriteString(w, "Authorization failed. You may close this window.")
+		} else {
+			_, _ = io.WriteString(w, "Authorization complete. You may close this window.")
+		}
+		select {
+		case resultChan <- callbackResult{code: code, err: parseErr}:
+		default:
+		}
+	}))
+	defer server.Close()
 
 	open := options.OpenBrowser
 	if open == nil {

internal/mcp/oauth_test.go

@@ -6,6 +6,8 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"fmt"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -668,3 +670,86 @@ func TestFullFlowUsesProtectedResourceDiscovery(t *testing.T) {
 		t.Fatalf("flow hits = probe:%d resource:%d authorization:%d token:%d", resourceProbeHits.Load(), resourceMetadataHits.Load(), authorizationMetadataHits.Load(), tokenHits.Load())
 	}
 }
+
+// A CLIENT STALLED MID-HEADER MUST NOT OUTLIVE LOGIN (#1028).
+//
+// Login built a bare http.Server and ended it with Shutdown under a one-second
+// deadline, and Shutdown only closes idle connections. A client that has sent
+// part of a request header is not idle, so it survived the deadline and kept
+// its goroutine, and the Serve goroutine, alive after Login returned.
+//
+// No sleep is needed to get the stalled connection accepted before Login ends:
+// it is dialled before the real callback's connection, Serve accepts in order,
+// and Login cannot finish until the real callback has been served.
+func TestLoginDoesNotLeaveAStalledCallbackClientOpen(t *testing.T) {
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"access_token": "access-final", "token_type": "Bearer", "expires_in": 3600,
+		})
+	}))
+	defer tokenServer.Close()
+
+	var stalled net.Conn
+	opener := func(authURL string) error {
+		parsed, err := url.Parse(authURL)
+		if err != nil {
+			return err
+		}
+		redirect, err := url.Parse(parsed.Query().Get("redirect_uri"))
+		if err != nil {
+			return err
+		}
+		stalled, err = net.Dial("tcp", redirect.Host)
+		if err != nil {
+			return err
+		}
+		if _, err := fmt.Fprintf(stalled, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", redirect.Host); err != nil {
+			return err
+		}
+		cbURL := redirect.String() + "?code=auth-code&state=" + url.QueryEscape(parsed.Query().Get("state"))
+		go func() {
+			if resp, err := http.Get(cbURL); err == nil {
+				resp.Body.Close()
+			}
+		}()
+		return nil
+	}
+
+	_, err := Login(context.Background(), LoginOptions{
+		ServerName: "demo",
+		ServerURL:  "https://issuer.invalid",
+		Config: OAuthConfig{
+			ClientID:              "client-123",
+			AuthorizationEndpoint: "https://issuer.invalid/authorize",
+			TokenEndpoint:         tokenServer.URL,
+		},
+		HTTPClient:  http.DefaultClient,
+		OpenBrowser: opener,
+		Timeout:     5 * time.Second,
+		Now:         time.Now,
+	})
+	if err != nil {
+		t.Fatalf("Login() error = %v", err)
+	}
+	if stalled == nil {
+		t.Fatal("SETUP INVALID: the stalled client never connected, so this proves nothing")
+	}
+	defer stalled.Close()
+	requireCallbackClientClosed(t, stalled)
+}
+
+// requireCallbackClientClosed fails unless the server end of conn is already
+// closed: a read has to fail with something other than a timeout.
+func requireCallbackClientClosed(t *testing.T, conn net.Conn) {
+	t.Helper()
+	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
+	var one [1]byte
+	_, err := conn.Read(one[:])
+	var netErr net.Error
+	if errors.As(err, &netErr) && netErr.Timeout() {
+		t.Fatal("the stalled callback client was still connected after Login returned")
+	}
+	if err == nil {
+		t.Fatal("the stalled callback client received data instead of being closed")
+	}
+}

internal/oauth/callback_server.go

@@ -0,0 +1,70 @@
+package oauth
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"time"
+)
+
+// callbackReadHeaderTimeout bounds how long an accepted connection may take to
+// send its request header. A single-use loopback callback server has no reason
+// to wait longer, and without a bound a client that stalls mid-header holds the
+// connection, and the server goroutine, open indefinitely. A var so the
+// package tests can shorten it.
+var callbackReadHeaderTimeout = 5 * time.Second
+
+// callbackShutdownBudget is how long Close waits for in-flight requests before
+// it force-closes whatever is still open.
+const callbackShutdownBudget = time.Second
+
+// CallbackServer is the lifecycle every loopback OAuth callback server needs,
+// in one place.
+//
+// It exists because the same server was written out three times (the shared
+// LoopbackListener, MCP's Login and OpenRouter's login) and a fix to one of
+// them did not reach the others. http.Server.Shutdown only closes idle
+// connections, so a client that has sent part of a request header is active,
+// survives the shutdown deadline, and keeps its goroutine and the Serve
+// goroutine alive past the end of the flow (#1028). #1051 closed that for
+// LoopbackListener; the other two kept the original construction.
+type CallbackServer struct {
+	server *http.Server
+	done   chan struct{}
+}
+
+// StartCallbackServer serves handler on listener until Close.
+func StartCallbackServer(listener net.Listener, handler http.Handler) *CallbackServer {
+	return startCallbackServer(listener, handler, nil)
+}
+
+// startCallbackServer is StartCallbackServer with a ConnState hook, which the
+// package's own tests use to see a connection reach a given state.
+func startCallbackServer(listener net.Listener, handler http.Handler, connState func(net.Conn, http.ConnState)) *CallbackServer {
+	s := &CallbackServer{
+		server: &http.Server{
+			Handler:           handler,
+			ReadHeaderTimeout: callbackReadHeaderTimeout,
+			ConnState:         connState,
+		},
+		done: make(chan struct{}),
+	}
+	go func() {
+		_ = s.server.Serve(listener)
+		close(s.done)
+	}()
+	return s
+}
+
+// Close shuts the server down within callbackShutdownBudget. A connection still
+// open when the budget runs out (a stalled request header, say) is force-closed
+// rather than leaked, and Close returns only after the Serve goroutine has
+// finished, so nothing this server started outlives the call. Idempotent.
+func (s *CallbackServer) Close() {
+	ctx, cancel := context.WithTimeout(context.Background(), callbackShutdownBudget)
+	defer cancel()
+	if err := s.server.Shutdown(ctx); err != nil {
+		_ = s.server.Close()
+	}
+	<-s.done
+}

internal/oauth/callback_server_test.go

@@ -0,0 +1,100 @@
+package oauth
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+)
+
+// startStalledCallbackClient starts a CallbackServer and connects a client that
+// sends part of a request header and then stops, which is the shape #1028 is
+// about. It returns once the server has ACCEPTED the connection, using the
+// ConnState hook rather than a sleep. StateNew, not StateActive: Go reports
+// StateActive only after a request read returns, and a stalled header read
+// returns only when the header timeout fires, so waiting for StateActive would
+// make the setup depend on the very bound the test is about.
+func startStalledCallbackClient(t *testing.T) (*CallbackServer, net.Conn) {
+	t.Helper()
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	accepted := make(chan struct{}, 1)
+	server := startCallbackServer(listener, http.NotFoundHandler(), func(_ net.Conn, state http.ConnState) {
+		if state == http.StateNew {
+			select {
+			case accepted <- struct{}{}:
+			default:
+			}
+		}
+	})
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		server.Close()
+		t.Fatal(err)
+	}
+	if _, err := fmt.Fprintf(conn, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", listener.Addr()); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("SETUP INVALID: the server never accepted the stalled client")
+	}
+	return server, conn
+}
+
+// connClosedWithin reports whether the server end of conn closes before the
+// deadline: a read must fail with something other than a timeout.
+func connClosedWithin(conn net.Conn, deadline time.Duration) bool {
+	_ = conn.SetReadDeadline(time.Now().Add(deadline))
+	var one [1]byte
+	_, err := conn.Read(one[:])
+	var netErr net.Error
+	if errors.As(err, &netErr) && netErr.Timeout() {
+		return false
+	}
+	return err != nil
+}
+
+// THE HEADER READ IS BOUNDED EVEN BEFORE CLOSE. The force-close in Close ends a
+// stalled client when the flow ends; the header timeout ends it while the flow
+// is still waiting for its callback, which can be minutes. Close is not called
+// until after the assertion, so this can only pass on ReadHeaderTimeout.
+func TestCallbackServerBoundsAStalledHeaderWithoutClose(t *testing.T) {
+	previous := callbackReadHeaderTimeout
+	callbackReadHeaderTimeout = 200 * time.Millisecond
+	defer func() { callbackReadHeaderTimeout = previous }()
+
+	server, conn := startStalledCallbackClient(t)
+	defer server.Close()
+	defer conn.Close()
+
+	if !connClosedWithin(conn, 3*time.Second) {
+		t.Fatal("a client stalled mid-header was still connected with Close never called; the header read is unbounded")
+	}
+}
+
+// CLOSE RETURNS ONLY AFTER SERVE HAS. Its doc comment promises that nothing the
+// server started outlives the call. With no connections open, Shutdown closes
+// the listener and returns at once, so this is the case where Close could beat
+// the Serve goroutine back; the goroutine closing done is the last thing it
+// does, so done has to be closed by the time Close returns.
+func TestCallbackServerCloseWaitsForServe(t *testing.T) {
+	for attempt := 0; attempt < 50; attempt++ {
+		listener, err := net.Listen("tcp", "127.0.0.1:0")
+		if err != nil {
+			t.Fatal(err)
+		}
+		server := StartCallbackServer(listener, http.NotFoundHandler())
+		server.Close()
+		select {
+		case <-server.done:
+		default:
+			t.Fatalf("attempt %d: Close returned while the Serve goroutine was still running", attempt)
+		}
+	}
+}

internal/oauth/loopback.go

@@ -9,7 +9,6 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
-	"time"
 )
 
 // LoopbackListener is a single-use loopback HTTP server that captures an OAuth
@@ -19,8 +18,7 @@ type LoopbackListener struct {
 	listener net.Listener
 	state    string
 	result   chan callbackResult
-	server   *http.Server
-	done     chan struct{}
+	server   *CallbackServer
 }
 
 type callbackResult struct {
@@ -58,20 +56,8 @@ func newLoopbackListener(state string, port int, connState func(net.Conn, http.C
 		listener: ln,
 		state:    state,
 		result:   make(chan callbackResult, 1),
-		done:     make(chan struct{}),
 	}
-	l.server = &http.Server{
-		Handler: http.HandlerFunc(l.handle),
-		// This is a single-use, loopback-only callback server: a client that
-		// never finishes sending its request header must not be able to hold
-		// the accepted connection (and Close, below) open indefinitely.
-		ReadHeaderTimeout: 5 * time.Second,
-		ConnState:         connState,
-	}
-	go func() {
-		_ = l.server.Serve(ln)
-		close(l.done)
-	}()
+	l.server = startCallbackServer(ln, http.HandlerFunc(l.handle), connState)
 	return l, nil
 }
 
@@ -122,14 +108,10 @@ func (l *LoopbackListener) Wait(ctx context.Context) (string, error) {
 // Close shuts the listener down (bounded), idempotent. If a client leaves a
 // connection open past the shutdown deadline (e.g. a stalled request header),
 // Close forcibly closes it rather than leaking it past the end of the flow,
-// and waits for the Serve goroutine to finish before returning.
+// and waits for the Serve goroutine to finish before returning. See
+// CallbackServer.Close.
 func (l *LoopbackListener) Close() {
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-	if err := l.server.Shutdown(shutdownCtx); err != nil {
-		_ = l.server.Close()
-	}
-	<-l.done
+	l.server.Close()
 }
 
 // parseCallback validates the redirect query and returns the authorization code,

internal/provideroauth/openrouter.go

@@ -79,7 +79,11 @@ func OpenRouterLogin(ctx context.Context, opts OpenRouterOptions) (string, error
 
 	codeCh := make(chan string, 1)
 	errCh := make(chan error, 1)
-	server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	// The shared callback server bounds header reads, force-closes a client
+	// still connected when shutdown runs out of time, and waits for Serve to
+	// return. The bare http.Server this used to build did none of that, so a
+	// client stalled mid-header outlived the login (#1028).
+	server := oauth.StartCallbackServer(listener, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != "/callback" {
 			http.NotFound(w, r)
 			return
@@ -98,13 +102,8 @@ func OpenRouterLogin(ctx context.Context, opts OpenRouterOptions) (string, error
 		case errCh <- errors.New("provideroauth: callback missing authorization code"):
 		default:
 		}
-	})}
-	go func() { _ = server.Serve(listener) }()
-	defer func() {
-		shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), time.Second)
-		defer cancelShutdown()
-		_ = server.Shutdown(shutdownCtx)
-	}()
+	}))
+	defer server.Close()
 
 	authURL := base + openRouterAuthPath + "?" + url.Values{
 		"callback_url":          {callbackURL},

internal/provideroauth/openrouter_test.go

@@ -4,12 +4,15 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 )
 
 // simulateBrowser returns an OpenBrowser func that authorizes by GETting the
@@ -129,3 +132,62 @@ func TestOpenRouterLoginBrowserErrorPropagates(t *testing.T) {
 		t.Fatal("a browser-open failure should propagate")
 	}
 }
+
+// A CLIENT STALLED MID-HEADER MUST NOT OUTLIVE THE LOGIN (#1028). Same defect
+// as MCP's Login: a bare http.Server ended by Shutdown alone, which leaves a
+// connection that has sent part of a request header open after the flow ends.
+// The stalled client is dialled before the real callback, so Serve has
+// accepted it by the time the login can complete; no sleep is involved.
+func TestOpenRouterLoginDoesNotLeaveAStalledCallbackClientOpen(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]string{"key": "sk-or-mock-123"})
+	}))
+	defer srv.Close()
+
+	var stalled net.Conn
+	open := func(authURL string) error {
+		u, err := url.Parse(authURL)
+		if err != nil {
+			return err
+		}
+		cb, err := url.Parse(u.Query().Get("callback_url"))
+		if err != nil {
+			return err
+		}
+		stalled, err = net.Dial("tcp", cb.Host)
+		if err != nil {
+			return err
+		}
+		if _, err := fmt.Fprintf(stalled, "GET /callback HTTP/1.1\r\nHost: %s\r\nX-Stall:", cb.Host); err != nil {
+			return err
+		}
+		resp, err := http.Get(cb.String() + "?code=TESTCODE") //nolint:noctx // test loopback
+		if err != nil {
+			return err
+		}
+		_ = resp.Body.Close()
+		return nil
+	}
+
+	if _, err := OpenRouterLogin(context.Background(), OpenRouterOptions{
+		BaseURL: srv.URL, HTTPClient: srv.Client(), OpenBrowser: open,
+	}); err != nil {
+		t.Fatalf("OpenRouterLogin: %v", err)
+	}
+	if stalled == nil {
+		t.Fatal("SETUP INVALID: the stalled client never connected, so this proves nothing")
+	}
+	defer stalled.Close()
+
+	_ = stalled.SetReadDeadline(time.Now().Add(2 * time.Second))
+	var one [1]byte
+	_, err := stalled.Read(one[:])
+	var netErr net.Error
+	if errors.As(err, &netErr) && netErr.Timeout() {
+		t.Fatal("the stalled callback client was still connected after OpenRouterLogin returned")
+	}
+	if err == nil {
+		t.Fatal("the stalled callback client received data instead of being closed")
+	}
+}