Return 401+WWW-Authenticate when vMCP upstream token is unrefreshable

pkg/auth/token.go

@@ -1152,7 +1152,7 @@ type RFC6750Error struct {
 }
 
 // writeOAuthError writes an RFC 6750 compliant JSON error response.
-func writeOAuthError(w http.ResponseWriter, errorCode, description string, status int) {
+func writeOAuthError(w http.ResponseWriter, errorCode, description string) {
 	body, err := json.Marshal(RFC6750Error{
 		Error:            errorCode,
 		ErrorDescription: description,
@@ -1162,26 +1162,28 @@ func writeOAuthError(w http.ResponseWriter, errorCode, description string, statu
 		return
 	}
 	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(status)
+	w.WriteHeader(http.StatusUnauthorized)
 	_, _ = w.Write(body)
 }
 
 // loadUpstreamTokens extracts the token session ID from claims and loads
-// all upstream provider tokens for that session. Returns (nil, nil) if no
+// all upstream provider tokens for that session. Returns (nil, nil, nil) if no
 // tsid claim exists. Returns a non-nil error when a tsid claim is present
-// but token loading fails (infrastructure error).
-func (v *TokenValidator) loadUpstreamTokens(ctx context.Context, claims jwt.MapClaims) (map[string]string, error) {
+// but token loading fails due to an infrastructure error (storage unavailable).
+// A non-nil failed slice means one or more providers' tokens could not be
+// refreshed; the caller should return HTTP 401 to trigger re-authentication.
+func (v *TokenValidator) loadUpstreamTokens(ctx context.Context, claims jwt.MapClaims) (map[string]string, []string, error) {
 	tsid, ok := claims[upstreamtoken.TokenSessionIDClaimKey].(string)
 	if !ok || tsid == "" {
-		return nil, nil
+		return nil, nil, nil
 	}
 
-	tokens, err := v.upstreamTokenReader.GetAllValidTokens(ctx, tsid)
+	tokens, failed, err := v.upstreamTokenReader.GetAllValidTokens(ctx, tsid)
 	if err != nil {
-		return nil, fmt.Errorf("load upstream tokens for session %s: %w", tsid, err)
+		return nil, nil, fmt.Errorf("load upstream tokens for session %s: %w", tsid, err)
 	}
 
-	return tokens, nil
+	return tokens, failed, nil
 }
 
 // Middleware creates an HTTP middleware that validates JWT tokens and creates Identity.
@@ -1191,15 +1193,15 @@ func (v *TokenValidator) Middleware(next http.Handler) http.Handler {
 		tokenString, err := ExtractBearerToken(r)
 		if err != nil {
 			w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidRequest, err.Error()))
-			writeOAuthError(w, OAuthErrInvalidRequest, err.Error(), http.StatusUnauthorized)
+			writeOAuthError(w, OAuthErrInvalidRequest, err.Error())
 			return
 		}
 
 		// Validate the token
 		claims, err := v.ValidateToken(r.Context(), tokenString)
 		if err != nil {
 			w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidToken, err.Error()))
-			writeOAuthError(w, OAuthErrInvalidToken, fmt.Sprintf("Invalid token: %v", err), http.StatusUnauthorized)
+			writeOAuthError(w, OAuthErrInvalidToken, fmt.Sprintf("Invalid token: %v", err))
 			return
 		}
 
@@ -1208,7 +1210,7 @@ func (v *TokenValidator) Middleware(next http.Handler) http.Handler {
 		if err != nil {
 			slog.Error("failed to convert claims to identity", "error", err)
 			w.Header().Set("WWW-Authenticate", v.buildWWWAuthenticate(OAuthErrInvalidToken, err.Error()))
-			writeOAuthError(w, OAuthErrInvalidToken, "Invalid authentication claims", http.StatusUnauthorized)
+			writeOAuthError(w, OAuthErrInvalidToken, "Invalid authentication claims")
 			return
 		}
 
@@ -1223,7 +1225,7 @@ func (v *TokenValidator) Middleware(next http.Handler) http.Handler {
 		// never mutated afterwards (see the UpstreamTokens doc comment in identity.go).
 		if v.upstreamTokenReader != nil {
 			loadCtx := WithIdentity(r.Context(), identity)
-			tokens, loadErr := v.loadUpstreamTokens(loadCtx, claims)
+			tokens, failed, loadErr := v.loadUpstreamTokens(loadCtx, claims)
 			if loadErr != nil {
 				slog.WarnContext(loadCtx, "upstream token storage unavailable",
 					"error", loadErr,
@@ -1232,6 +1234,20 @@ func (v *TokenValidator) Middleware(next http.Handler) http.Handler {
 				return
 			}
 			identity.UpstreamTokens = tokens
+			// Conservative policy: any provider refresh failure rejects the entire
+			// request. In vMCP multi-provider sessions this may block requests that
+			// could partially succeed if only some backends need the failing provider.
+			// A per-backend provider check would require routing context not available
+			// at this layer; this is the correct tradeoff for now.
+			if len(failed) > 0 {
+				slog.WarnContext(loadCtx, "upstream token refresh failed; returning 401",
+					"providers", failed)
+				w.Header().Set("WWW-Authenticate",
+					v.buildWWWAuthenticate(OAuthErrInvalidToken, "upstream token is no longer valid; re-authentication required"))
+				writeOAuthError(w, OAuthErrInvalidToken,
+					"upstream token is no longer valid; re-authentication required")
+				return
+			}
 		}
 
 		// Add the fully-populated Identity to the request context and serve.

pkg/auth/token_test.go

@@ -2236,14 +2236,15 @@ func TestLoadUpstreamTokens(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-abc").
-			Return(map[string]string{"github": "gh-token", "atlassian": "atl-token"}, nil)
+			Return(map[string]string{"github": "gh-token", "atlassian": "atl-token"}, []string(nil), nil)
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
 			"sub":                                "user123",
 			upstreamtoken.TokenSessionIDClaimKey: "session-abc",
 		})
 		require.NoError(t, err)
+		require.Nil(t, failed)
 		require.Equal(t, map[string]string{"github": "gh-token", "atlassian": "atl-token"}, result)
 	})
 
@@ -2254,9 +2255,10 @@ func TestLoadUpstreamTokens(t *testing.T) {
 		// No EXPECT — reader should not be called
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{"sub": "user123"})
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{"sub": "user123"})
 		require.NoError(t, err)
 		require.Nil(t, result)
+		require.Nil(t, failed)
 	})
 
 	t.Run("returns nil when tsid is empty string", func(t *testing.T) {
@@ -2265,12 +2267,13 @@ func TestLoadUpstreamTokens(t *testing.T) {
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
 			"sub":                                "user123",
 			upstreamtoken.TokenSessionIDClaimKey: "",
 		})
 		require.NoError(t, err)
 		require.Nil(t, result)
+		require.Nil(t, failed)
 	})
 
 	t.Run("returns nil when tsid is non-string type", func(t *testing.T) {
@@ -2279,44 +2282,64 @@ func TestLoadUpstreamTokens(t *testing.T) {
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
 			"sub":                                "user123",
 			upstreamtoken.TokenSessionIDClaimKey: 12345,
 		})
 		require.NoError(t, err)
 		require.Nil(t, result)
+		require.Nil(t, failed)
 	})
 
 	t.Run("returns error when reader fails", func(t *testing.T) {
 		t.Parallel()
 		ctrl := gomock.NewController(t)
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-abc").
-			Return(nil, errors.New("storage unavailable"))
+			Return(nil, nil, errors.New("storage unavailable"))
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
 			"sub":                                "user123",
 			upstreamtoken.TokenSessionIDClaimKey: "session-abc",
 		})
 		require.Error(t, err)
 		require.Nil(t, result)
+		require.Nil(t, failed)
+	})
+
+	t.Run("returns failed providers from reader", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
+		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-abc").
+			Return(map[string]string{}, []string{"github"}, nil)
+
+		v := &TokenValidator{upstreamTokenReader: reader}
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+			"sub":                                "user123",
+			upstreamtoken.TokenSessionIDClaimKey: "session-abc",
+		})
+		require.NoError(t, err)
+		require.Equal(t, map[string]string{}, result)
+		require.Equal(t, []string{"github"}, failed)
 	})
 
 	t.Run("returns empty map from reader", func(t *testing.T) {
 		t.Parallel()
 		ctrl := gomock.NewController(t)
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-abc").
-			Return(map[string]string{}, nil)
+			Return(map[string]string{}, []string(nil), nil)
 
 		v := &TokenValidator{upstreamTokenReader: reader}
-		result, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
+		result, failed, err := v.loadUpstreamTokens(context.Background(), jwt.MapClaims{
 			"sub":                                "user123",
 			upstreamtoken.TokenSessionIDClaimKey: "session-abc",
 		})
 		require.NoError(t, err)
 		require.Equal(t, map[string]string{}, result)
+		require.Nil(t, failed)
 	})
 }
 
@@ -2386,7 +2409,7 @@ func TestMiddleware_UpstreamTokenEnrichment(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-xyz").
-			Return(map[string]string{"github": "gh-tok"}, nil)
+			Return(map[string]string{"github": "gh-tok"}, []string(nil), nil)
 		v := makeValidator(t, WithUpstreamTokenReader(reader))
 
 		var captured *Identity
@@ -2417,10 +2440,10 @@ func TestMiddleware_UpstreamTokenEnrichment(t *testing.T) {
 		var loadCtxCanonicalUser string
 		var loadCtxHadCanonicalUser bool
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-xyz").
-			DoAndReturn(func(ctx context.Context, _ string) (map[string]string, error) {
+			DoAndReturn(func(ctx context.Context, _ string) (map[string]string, []string, error) {
 				loadCtxIdentity, loadCtxHadIdentity = IdentityFromContext(ctx)
 				loadCtxCanonicalUser, loadCtxHadCanonicalUser = CanonicalUserFromContext(ctx)
-				return map[string]string{"github": "gh-tok"}, nil
+				return map[string]string{"github": "gh-tok"}, nil, nil
 			})
 		v := makeValidator(t, WithUpstreamTokenReader(reader))
 		handler := v.Middleware(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}))
@@ -2447,7 +2470,7 @@ func TestMiddleware_UpstreamTokenEnrichment(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
 		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-xyz").
-			Return(nil, errors.New("redis down"))
+			Return(nil, nil, errors.New("redis down"))
 		v := makeValidator(t, WithUpstreamTokenReader(reader))
 
 		nextCalled := false
@@ -2464,6 +2487,54 @@ func TestMiddleware_UpstreamTokenEnrichment(t *testing.T) {
 		require.False(t, nextCalled)
 	})
 
+	t.Run("returns 401 with WWW-Authenticate when provider refresh failed", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
+		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-xyz").
+			Return(map[string]string{}, []string{"github"}, nil)
+		v := makeValidator(t, WithUpstreamTokenReader(reader))
+
+		nextCalled := false
+		handler := v.Middleware(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+			nextCalled = true
+		}))
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Authorization", "Bearer "+signToken(claimsWithTsid))
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		require.Equal(t, http.StatusUnauthorized, rr.Code)
+		require.False(t, nextCalled)
+		wwwAuth := rr.Header().Get("WWW-Authenticate")
+		require.Contains(t, wwwAuth, `error="invalid_token"`)
+	})
+
+	t.Run("returns 401 when one of multiple providers fails refresh", func(t *testing.T) {
+		t.Parallel()
+		ctrl := gomock.NewController(t)
+		reader := upstreamtokenmocks.NewMockTokenReader(ctrl)
+		// atlassian succeeded, github failed — the middleware must still reject the request
+		reader.EXPECT().GetAllValidTokens(gomock.Any(), "session-xyz").
+			Return(map[string]string{"atlassian": "atl-tok"}, []string{"github"}, nil)
+		v := makeValidator(t, WithUpstreamTokenReader(reader))
+
+		nextCalled := false
+		handler := v.Middleware(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+			nextCalled = true
+		}))
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Authorization", "Bearer "+signToken(claimsWithTsid))
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		require.Equal(t, http.StatusUnauthorized, rr.Code)
+		require.False(t, nextCalled, "next must not be called when any provider fails")
+		require.Contains(t, rr.Header().Get("WWW-Authenticate"), `error="invalid_token"`)
+	})
+
 	t.Run("no enrichment without tsid", func(t *testing.T) {
 		t.Parallel()
 		ctrl := gomock.NewController(t)

pkg/auth/upstreamtoken/service.go

@@ -75,19 +75,20 @@ func (s *InProcessService) GetValidTokens(ctx context.Context, sessionID, provid
 }
 
 // GetAllValidTokens returns access tokens for all upstream providers in a session.
-// Expired tokens are refreshed transparently; if refresh fails, the provider is
-// omitted from the result so downstream middleware can return a clean 401.
-func (s *InProcessService) GetAllValidTokens(ctx context.Context, sessionID string) (map[string]string, error) {
+// Expired tokens are refreshed transparently; if refresh fails, the provider name
+// is included in the returned failed slice so the caller can return HTTP 401.
+func (s *InProcessService) GetAllValidTokens(ctx context.Context, sessionID string) (map[string]string, []string, error) {
 	allTokens, err := s.storage.GetAllUpstreamTokens(ctx, sessionID)
 	if err != nil {
-		return nil, fmt.Errorf("bulk read upstream tokens: %w", err)
+		return nil, nil, fmt.Errorf("bulk read upstream tokens: %w", err)
 	}
 
 	if len(allTokens) == 0 {
-		return map[string]string{}, nil
+		return map[string]string{}, nil, nil
 	}
 
 	result := make(map[string]string, len(allTokens))
+	var failed []string
 	// TODO(auth): Refresh providers in parallel using errgroup to avoid
 	// worst-case latency of N refreshes when multiple providers need refresh.
 	for providerName, tokens := range allTokens {
@@ -104,18 +105,18 @@ func (s *InProcessService) GetAllValidTokens(ctx context.Context, sessionID stri
 		// Token is expired — attempt refresh.
 		refreshed, refreshErr := s.refreshOrFail(ctx, sessionID, providerName, tokens)
 		if refreshErr != nil {
-			// Refresh failed — omit provider so downstream middleware returns 401.
-			slog.WarnContext(ctx, "omitting provider with unrefreshable expired token",
+			slog.WarnContext(ctx, "upstream token refresh failed; provider will require re-authentication",
 				"session_id", sessionID,
 				"provider", providerName,
 				"error", refreshErr,
 			)
+			failed = append(failed, providerName)
 			continue
 		}
 		result[providerName] = refreshed.AccessToken
 	}
 
-	return result, nil
+	return result, failed, nil
 }
 
 // refreshOrFail attempts a refresh via the shared refresher and maps errors to