diff --git a/frontend/src/api.ts b/frontend/src/api.ts index 39999ba..75565d2 100644 --- a/frontend/src/api.ts +++ b/frontend/src/api.ts @@ -50,15 +50,28 @@ export async function getStatus(): Promise { } // PRD-43: all API calls carry HttpOnly auth cookies via credentials:"include". -// On a 401 response, fetchJSON attempts exactly one silent refresh; if that -// succeeds, the original request is retried once. If the refresh returns 401 -// too, we redirect the user to /login (except for the auth endpoints -// themselves, which handle 401 as a "bad credentials" UX state). +// On a genuine auth 401, fetchJSON attempts exactly one silent refresh; if +// that succeeds, the original request is retried once. If the refresh fails, +// we redirect the user to /login (except for the auth endpoints themselves, +// which handle 401 as a "bad credentials" UX state). +// +// A 401 is only treated as an auth failure when the response carries a +// `WWW-Authenticate` header, which the auth middleware sets (see +// internal/auth/middleware.go). This distinguishes a real expired-session +// 401 from a 401 that merely originated upstream — e.g. a gated-model +// response relayed from HuggingFace by /recommend, /estimate, or +// /memory-breakdown. Without this guard, selecting a gated model bounced +// the user to /login (and, with auth disabled, /auth/refresh returns 503 +// so the redirect fired every time). +function isAuthChallenge(res: Response): boolean { + return res.status === 401 && res.headers.has("WWW-Authenticate"); +} + async function fetchJSON(url: string, init?: RequestInit): Promise { const withCreds: RequestInit = { credentials: "include", ...(init ?? {}) }; let res = await fetch(url, withCreds); - if (res.status === 401 && !isAuthEndpoint(url)) { + if (isAuthChallenge(res) && !isAuthEndpoint(url)) { const refreshed = await trySilentRefresh(); if (refreshed) { res = await fetch(url, withCreds); diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 3561a11..f9a3dc7 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -956,12 +956,7 @@ func (s *Server) handleRecommend(w http.ResponseWriter, r *http.Request) { // Fetch model config (from S3 cache if available, else HuggingFace). modelCfg, err := s.FetchModelConfig(r.Context(), modelID, hfToken) if err != nil { - var hfErr *recommend.HFError - if errors.As(err, &hfErr) { - writeError(w, hfErr.StatusCode, hfErr.Message) - return - } - writeError(w, http.StatusBadGateway, "failed to fetch model metadata from HuggingFace") + writeHFError(w, err) return } @@ -1149,6 +1144,28 @@ func writeError(w http.ResponseWriter, code int, msg string) { writeJSON(w, code, map[string]string{"error": msg}) } +// writeHFError translates an error from FetchModelConfig into an HTTP +// response. It exists so a gated/unauthorized HuggingFace model never +// surfaces to the browser as the API's own 401/403: the frontend's +// fetchJSON treats any 401 as a session expiry and bounces the user to +// /login (see frontend/src/api.ts). HuggingFace 401/403 (gated model, +// missing/expired platform token, or HF rate-limiting) is a +// request-level problem, not an auth failure, so we remap it to 422 +// Unprocessable Entity. The original message ("model is gated — provide +// an HF token …") is preserved. Non-HFError failures become 502. +func writeHFError(w http.ResponseWriter, err error) { + var hfErr *recommend.HFError + if errors.As(err, &hfErr) { + code := hfErr.StatusCode + if code == http.StatusUnauthorized || code == http.StatusForbidden { + code = http.StatusUnprocessableEntity + } + writeError(w, code, hfErr.Message) + return + } + writeError(w, http.StatusBadGateway, "failed to fetch model metadata from HuggingFace") +} + // handleListScenarios returns all available benchmark scenarios. func (s *Server) handleListScenarios(w http.ResponseWriter, r *http.Request) { const cacheKey = "scenarios" diff --git a/internal/api/handlers_estimate.go b/internal/api/handlers_estimate.go index 552faf2..5495976 100644 --- a/internal/api/handlers_estimate.go +++ b/internal/api/handlers_estimate.go @@ -1,7 +1,6 @@ package api import ( - "errors" "net/http" "sort" "strconv" @@ -121,12 +120,7 @@ func (s *Server) handleEstimate(w http.ResponseWriter, r *http.Request) { // Fetch model config (from S3 cache if available, else HuggingFace). modelCfg, err := s.FetchModelConfig(ctx, modelID, hfToken) if err != nil { - var hfErr *recommend.HFError - if errors.As(err, &hfErr) { - writeError(w, hfErr.StatusCode, hfErr.Message) - return - } - writeError(w, http.StatusBadGateway, "failed to fetch model metadata from HuggingFace") + writeHFError(w, err) return } diff --git a/internal/api/handlers_hferror_test.go b/internal/api/handlers_hferror_test.go new file mode 100644 index 0000000..066693d --- /dev/null +++ b/internal/api/handlers_hferror_test.go @@ -0,0 +1,106 @@ +package api + +// Regression tests for the gated-model login-bounce bug. +// +// When a user selected a gated HuggingFace model on the New Benchmark +// (or Estimate) page, /recommend, /estimate, and /memory-breakdown +// fetched the model config from HuggingFace, which returned 401. The +// handlers relayed that 401 verbatim, and the frontend's fetchJSON +// treats any 401 as an expired session — bouncing the user to /login. +// The symptom only appeared with auth disabled, because /auth/refresh +// returns 503 there so the silent-refresh retry always failed. +// +// The fix: an HF 401/403 (gated model, missing/expired platform token, +// or HF rate-limiting) is remapped to 422 so it never collides with the +// app's own auth 401. See writeHFError in handlers.go. + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/accelbench/accelbench/internal/recommend" + + "k8s.io/client-go/kubernetes/fake" +) + +// hfErrorServer builds a server whose HF client always fails with the +// given HFError, mirroring a gated-model response. Auth is disabled +// (NewServerWithHFClient default), matching the reported environment. +func hfErrorServer(hfErr error) *http.ServeMux { + repo := seedRepo() + client := fake.NewSimpleClientset() + hf := &recommend.MockHFClient{ + FetchModelConfigFunc: func(modelID, hfToken string) (*recommend.ModelConfig, error) { + return nil, hfErr + }, + } + srv := NewServerWithHFClient(repo, client, hf, "test-pod") + mux := http.NewServeMux() + srv.RegisterRoutes(mux) + return mux +} + +func TestHFGatedError_NotRelayedAs401(t *testing.T) { + // The three endpoints that fetch model config from HuggingFace. + // meta-llama/Llama-3.1-8B is seeded by seedRepo(); the mock HF + // client rejects it as gated regardless. + endpoints := []string{ + "/api/v1/recommend?model=meta-llama/Llama-3.1-8B&instance_type=g5.xlarge", + "/api/v1/estimate?model=meta-llama/Llama-3.1-8B", + "/api/v1/memory-breakdown?model=meta-llama/Llama-3.1-8B&instance_type=g5.xlarge", + } + + // Both statuses HuggingFace uses for gated/unauthorized models. + for _, hfStatus := range []int{http.StatusUnauthorized, http.StatusForbidden} { + mux := hfErrorServer(&recommend.HFError{ + StatusCode: hfStatus, + Message: "model is gated — provide an HF token with access", + }) + for _, ep := range endpoints { + t.Run(http.StatusText(hfStatus)+" "+ep, func(t *testing.T) { + req := httptest.NewRequest("GET", ep, nil) + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + if w.Code == http.StatusUnauthorized || w.Code == http.StatusForbidden { + t.Fatalf("status = %d; a gated-model HF %d must not surface as the app's own 401/403 (that bounces the user to /login). body = %s", + w.Code, hfStatus, w.Body.String()) + } + if w.Code != http.StatusUnprocessableEntity { + t.Errorf("status = %d, want 422; body = %s", w.Code, w.Body.String()) + } + // The original, actionable message must be preserved. + if body := w.Body.String(); !contains(body, "gated") { + t.Errorf("body = %q, want it to preserve the gated-model message", body) + } + }) + } + } +} + +// A non-HFError failure from the HF fetch (e.g. network error) should be +// a 502, never a 401. +func TestHFGenericError_Is502(t *testing.T) { + mux := hfErrorServer(errString("connection refused")) + req := httptest.NewRequest("GET", "/api/v1/recommend?model=meta-llama/Llama-3.1-8B&instance_type=g5.xlarge", nil) + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + if w.Code != http.StatusBadGateway { + t.Errorf("status = %d, want 502; body = %s", w.Code, w.Body.String()) + } +} + +type errString string + +func (e errString) Error() string { return string(e) } + +func contains(s, sub string) bool { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/internal/api/handlers_recommend.go b/internal/api/handlers_recommend.go index 243af47..8968dc2 100644 --- a/internal/api/handlers_recommend.go +++ b/internal/api/handlers_recommend.go @@ -1,7 +1,6 @@ package api import ( - "errors" "fmt" "net/http" @@ -64,12 +63,7 @@ func (s *Server) handleMemoryBreakdown(w http.ResponseWriter, r *http.Request) { // Fetch model config (from S3 cache if available, else HuggingFace). modelCfg, err := s.FetchModelConfig(r.Context(), modelID, hfToken) if err != nil { - var hfErr *recommend.HFError - if errors.As(err, &hfErr) { - writeError(w, hfErr.StatusCode, hfErr.Message) - return - } - writeError(w, http.StatusBadGateway, "failed to fetch model metadata") + writeHFError(w, err) return } diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go index 1cca831..3f22af8 100644 --- a/internal/auth/auth_test.go +++ b/internal/auth/auth_test.go @@ -130,6 +130,13 @@ func TestMiddleware_MissingCookie_Returns401(t *testing.T) { if rec.Code != http.StatusUnauthorized { t.Errorf("status = %d, want 401", rec.Code) } + // A genuine auth 401 must carry WWW-Authenticate so the frontend can + // tell it apart from a 401 relayed from an upstream (e.g. a gated + // HuggingFace model). Without this marker fetchJSON would not run its + // refresh + redirect-to-login flow. + if got := rec.Header().Get("WWW-Authenticate"); got == "" { + t.Errorf("WWW-Authenticate header missing on auth 401") + } } func TestMiddleware_ExpiredToken_Returns401(t *testing.T) { diff --git a/internal/auth/middleware.go b/internal/auth/middleware.go index 3208c08..8e8d858 100644 --- a/internal/auth/middleware.go +++ b/internal/auth/middleware.go @@ -90,6 +90,12 @@ func Middleware(cfg Config, verifier *Verifier) func(http.Handler) http.Handler } func writeUnauthorized(w http.ResponseWriter, code string) { + // WWW-Authenticate marks this as a genuine authentication failure so + // the frontend can distinguish it from a 401 that merely originated + // from an upstream (e.g. a gated-model 401 relayed from HuggingFace). + // Only on this header does fetchJSON run its silent-refresh + + // redirect-to-login flow. See frontend/src/api.ts. + w.Header().Set("WWW-Authenticate", "Bearer") w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusUnauthorized) _ = json.NewEncoder(w).Encode(map[string]string{"error": code})