fix(mcp): use session-level context for login polling, not per-request

The per-request ctx passed to MCP tool handlers is cancelled when the
handler returns. The previous commit selected on that ctx in the login
polling goroutine, which killed the poll immediately after returning
the browser URL — breaking in-MCP authentication.

Fix: add a sessionCtx field to Server, set it in Run() (called by
RunStdio and tests), and select on that instead. The session context
is only cancelled when the MCP transport closes (stdin EOF), which is
the correct signal to abandon login polling.

Also adds TestLoginTool_PollSurvivesAcrossToolCalls: a regression test
that starts a login flow, lets the mock auth complete between tool
calls, and verifies the client is authenticated on the second call.
This would have caught the per-request ctx bug.

https://claude.ai/code/session_01EpXZqTmgybtjgmSALukH8d

Author: claude

pkg/gateway/mcp/server.go

@@ -18,6 +18,12 @@ type Server struct {
 	client    *hookdeck.Client
 	cfg       *config.Config
 	mcpServer *mcpsdk.Server
+
+	// sessionCtx is the context passed to RunStdio. It is cancelled when the
+	// MCP transport closes (stdin EOF). Background goroutines (e.g. login
+	// polling) should select on this — NOT on the per-request ctx passed to
+	// tool handlers, which is cancelled when the handler returns.
+	sessionCtx context.Context
 }
 
 // NewServer creates an MCP server with all Hookdeck tools registered.
@@ -56,7 +62,7 @@ func (s *Server) registerTools() {
 				"reauth": {Type: "boolean", Desc: "If true, clear stored credentials and start a new browser login. Use when project listing fails — complete login in the browser, then retry hookdeck_projects."},
 			}),
 		},
-		s.wrapWithTelemetry("hookdeck_login", handleLogin(s.client, s.cfg)),
+		s.wrapWithTelemetry("hookdeck_login", handleLogin(s)),
 	)
 }
 
@@ -126,5 +132,13 @@ func extractAction(req *mcpsdk.CallToolRequest) string {
 // RunStdio starts the MCP server on stdin/stdout and blocks until the
 // connection is closed (i.e. stdin reaches EOF).
 func (s *Server) RunStdio(ctx context.Context) error {
-	return s.mcpServer.Run(ctx, &mcpsdk.StdioTransport{})
+	return s.Run(ctx, &mcpsdk.StdioTransport{})
+}
+
+// Run starts the MCP server on the given transport. It stores ctx as the
+// session-level context so background goroutines (e.g. login polling) can
+// detect when the session ends.
+func (s *Server) Run(ctx context.Context, transport mcpsdk.Transport) error {
+	s.sessionCtx = ctx
+	return s.mcpServer.Run(ctx, transport)
 }

pkg/gateway/mcp/server_test.go

@@ -9,6 +9,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 
 	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/stretchr/testify/assert"
@@ -45,7 +46,7 @@ func connectInMemory(t *testing.T, client *hookdeck.Client) *mcpsdk.ClientSessio
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 	go func() {
-		_ = srv.mcpServer.Run(ctx, serverTransport)
+		_ = srv.Run(ctx, serverTransport)
 	}()
 
 	mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{
@@ -1143,7 +1144,7 @@ func TestLoginTool_ReauthStartsFreshLogin(t *testing.T) {
 	serverTransport, clientTransport := mcpsdk.NewInMemoryTransports()
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
-	go func() { _ = srv.mcpServer.Run(ctx, serverTransport) }()
+	go func() { _ = srv.Run(ctx, serverTransport) }()
 
 	mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "test", Version: "0.0.1"}, nil)
 	session, err := mcpClient.Connect(ctx, clientTransport, nil)
@@ -1182,7 +1183,7 @@ func TestLoginTool_ReturnsURLImmediately(t *testing.T) {
 	serverTransport, clientTransport := mcpsdk.NewInMemoryTransports()
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
-	go func() { _ = srv.mcpServer.Run(ctx, serverTransport) }()
+	go func() { _ = srv.Run(ctx, serverTransport) }()
 
 	mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "test", Version: "0.0.1"}, nil)
 	session, err := mcpClient.Connect(ctx, clientTransport, nil)
@@ -1218,7 +1219,7 @@ func TestLoginTool_InProgressShowsURL(t *testing.T) {
 	serverTransport, clientTransport := mcpsdk.NewInMemoryTransports()
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
-	go func() { _ = srv.mcpServer.Run(ctx, serverTransport) }()
+	go func() { _ = srv.Run(ctx, serverTransport) }()
 
 	mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "test", Version: "0.0.1"}, nil)
 	session, err := mcpClient.Connect(ctx, clientTransport, nil)
@@ -1236,6 +1237,70 @@ func TestLoginTool_InProgressShowsURL(t *testing.T) {
 	assert.Contains(t, text, "https://hookdeck.com/auth?code=xyz")
 }
 
+func TestLoginTool_PollSurvivesAcrossToolCalls(t *testing.T) {
+	// Regression: the login polling goroutine must use the session-level
+	// context, not the per-request ctx (which is cancelled when the handler
+	// returns). If the goroutine selected on per-request ctx, it would be
+	// cancelled immediately and the second hookdeck_login call would see a
+	// "login cancelled" error instead of "Already authenticated".
+	pollCount := 0
+	api := mockAPI(t, map[string]http.HandlerFunc{
+		"/2025-07-01/cli-auth": func(w http.ResponseWriter, r *http.Request) {
+			json.NewEncoder(w).Encode(map[string]any{
+				"browser_url": "https://hookdeck.com/auth?code=survive",
+				"poll_url":    "http://" + r.Host + "/2025-07-01/cli-auth/poll?key=survive",
+			})
+		},
+		"/2025-07-01/cli-auth/poll": func(w http.ResponseWriter, r *http.Request) {
+			pollCount++
+			if pollCount >= 2 {
+				// Simulate user completing browser auth on 2nd poll.
+				json.NewEncoder(w).Encode(map[string]any{
+					"claimed":           true,
+					"key":               "sk_test_survive12345",
+					"team_id":           "proj_survive",
+					"team_name":         "Survive Project",
+					"team_mode":         "console",
+					"user_name":         "test-user",
+					"organization_name": "test-org",
+				})
+				return
+			}
+			json.NewEncoder(w).Encode(map[string]any{"claimed": false})
+		},
+	})
+
+	unauthClient := newTestClient(api.URL, "")
+	cfg := &config.Config{APIBaseURL: api.URL}
+	srv := NewServer(unauthClient, cfg)
+
+	serverTransport, clientTransport := mcpsdk.NewInMemoryTransports()
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	go func() { _ = srv.Run(ctx, serverTransport) }()
+
+	mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "test", Version: "0.0.1"}, nil)
+	session, err := mcpClient.Connect(ctx, clientTransport, nil)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = session.Close() })
+
+	// First call initiates the flow — handler returns immediately.
+	result := callTool(t, session, "hookdeck_login", map[string]any{})
+	assert.False(t, result.IsError)
+	assert.Contains(t, textContent(t, result), "https://hookdeck.com/auth?code=survive")
+
+	// Wait briefly for the polling goroutine to complete (poll interval is 2s
+	// in production, but the mock returns instantly so it completes quickly).
+	time.Sleep(500 * time.Millisecond)
+
+	// Second call — if the goroutine survived, the client is now authenticated.
+	result2 := callTool(t, session, "hookdeck_login", map[string]any{})
+	assert.False(t, result2.IsError)
+	text := textContent(t, result2)
+	assert.Contains(t, text, "Already authenticated")
+	assert.Equal(t, "sk_test_survive12345", unauthClient.APIKey)
+}
+
 // ---------------------------------------------------------------------------
 // API error scenarios (shared across tools)
 // ---------------------------------------------------------------------------

pkg/gateway/mcp/tool_login.go

@@ -36,7 +36,9 @@ type loginState struct {
 	err        error         // non-nil if polling failed
 }
 
-func handleLogin(client *hookdeck.Client, cfg *config.Config) mcpsdk.ToolHandler {
+func handleLogin(srv *Server) mcpsdk.ToolHandler {
+	client := srv.client
+	cfg := srv.cfg
 	var stateMu sync.Mutex
 	var state *loginState
 
@@ -121,9 +123,11 @@ func handleLogin(client *hookdeck.Client, cfg *config.Config) mcpsdk.ToolHandler
 		}
 
 		// Poll in the background so we return the URL to the agent immediately.
-		// WaitForAPIKey blocks with time.Sleep; run it in a goroutine and
-		// select on ctx so we abandon the attempt when the session closes.
-		go func(s *loginState, ctx context.Context) {
+		// WaitForAPIKey blocks with time.Sleep internally, so we run it in an
+		// inner goroutine and select on the session-level context (not the
+		// per-request ctx, which is cancelled when this handler returns).
+		sessionCtx := srv.sessionCtx
+		go func(s *loginState) {
 			defer close(s.done)
 
 			type pollResult struct {
@@ -138,7 +142,7 @@ func handleLogin(client *hookdeck.Client, cfg *config.Config) mcpsdk.ToolHandler
 
 			var response *hookdeck.PollAPIKeyResponse
 			select {
-			case <-ctx.Done():
+			case <-sessionCtx.Done():
 				s.err = fmt.Errorf("login cancelled: MCP session closed")
 				log.Debug("Login polling cancelled — MCP session closed")
 				return
@@ -187,7 +191,7 @@ func handleLogin(client *hookdeck.Client, cfg *config.Config) mcpsdk.ToolHandler
 				"user":    response.UserName,
 				"project": response.ProjectName,
 			}).Info("MCP login completed successfully")
-		}(state, ctx)
+		}(state)
 
 		// Return the URL immediately so the agent can show it to the user.
 		return TextResult(fmt.Sprintf(