A lot of contracts use the phrase 'courts in X', which certainly raises an ambiguity. What most concerns contract drafters is that the court case be heard on their home turf, without considering the federal/state aspect.

To my mind the phrase 'courts of X' is better and as you say specifically alluding to the federal/state aspect makes it clearer still.