fix(auth,client): cover token refresh with tests and repair the MCP 401 retry

cmd/cli/cmd/mcp.go

@@ -59,7 +59,10 @@ var mcpServeCmd = &cobra.Command{
 		tokenProvider := func() (string, error) {
 			return lib.GetNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken)
 		}
-		httpClient, clientErr := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider)
+		refreshProvider := func() (string, error) {
+			return lib.RefreshNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken)
+		}
+		httpClient, clientErr := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider, refreshProvider)
 		if clientErr != nil {
 			return fmt.Errorf("failed to create client: %w", clientErr)
 		}

internal/auth/login.go

@@ -3,6 +3,7 @@ package auth
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -18,6 +19,10 @@ import (
 
 var httpClient = &http.Client{Timeout: 30 * time.Second}
 
+// authBaseURL builds the API origin for the auth endpoints. It is a var so
+// tests can point the auth calls at an httptest server.
+var authBaseURL = func(host string) string { return "https://" + apiHost(host) }
+
 type cliSessionResponse struct {
 	SessionID string `json:"session_id"`
 	AuthURL   string `json:"auth_url"`
@@ -198,8 +203,35 @@ func GetValidToken(ctx context.Context, host string) (string, error) {
 	return hostCreds.AccessToken, nil
 }
 
+// ErrNotRefreshable reports that the stored credentials for a host carry no
+// refresh token, so no amount of retrying will produce a new access token.
+var ErrNotRefreshable = errors.New("no refresh token stored")
+
+// ForceRefreshToken exchanges the stored refresh token for a new access token
+// regardless of the recorded expiry. GetValidToken only refreshes once
+// ExpiresAt has elapsed, which leaves callers no way to recover from a token
+// the server rejects while the CLI still believes it is valid -- revocation, a
+// killed session, or a local clock running behind the server's.
+func ForceRefreshToken(ctx context.Context, host string) (string, error) {
+	creds, err := LoadCredentials()
+	if err != nil {
+		return "", fmt.Errorf("not authenticated - run 'nullify auth login'")
+	}
+
+	hostCreds, ok := creds[CredentialKey(host)]
+	if !ok {
+		return "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host)
+	}
+	if hostCreds.RefreshToken == "" {
+		return "", ErrNotRefreshable
+	}
+
+	logger.L(ctx).Debug("forcing access token refresh")
+	return refreshToken(ctx, host, hostCreds.RefreshToken)
+}
+
 func createCLISession(ctx context.Context, host string, port int) (*cliSessionResponse, error) {
-	url := fmt.Sprintf("https://%s/auth/cli/session", apiHost(host))
+	url := authBaseURL(host) + "/auth/cli/session"
 
 	bodyData, err := json.Marshal(map[string]int{"port": port})
 	if err != nil {
@@ -233,7 +265,7 @@ func createCLISession(ctx context.Context, host string, port int) (*cliSessionRe
 }
 
 func fetchCLIToken(ctx context.Context, host string, sessionID string) (*cliTokenResponse, error) {
-	url := fmt.Sprintf("https://%s/auth/cli/token", apiHost(host))
+	url := authBaseURL(host) + "/auth/cli/token"
 
 	bodyData, err := json.Marshal(map[string]string{"session_id": sessionID})
 	if err != nil {
@@ -267,7 +299,7 @@ func fetchCLIToken(ctx context.Context, host string, sessionID string) (*cliToke
 }
 
 func refreshToken(ctx context.Context, host string, refreshTok string) (string, error) {
-	refreshURL := fmt.Sprintf("https://%s/auth/refresh_token", apiHost(host))
+	refreshURL := authBaseURL(host) + "/auth/refresh_token"
 
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, refreshURL, nil)
 	if err != nil {
@@ -296,17 +328,21 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string,
 		return "", err
 	}
 
-	// The endpoint delivers the new access token as a Set-Cookie header
-	// rather than in the JSON body; fall back to it when the body has none.
+	// GET /auth/refresh_token carries no token in its JSON body -- see
+	// RefreshTokenOutput in auth/pkg/endpoints/refresh_token.go, which has no
+	// token field at all. The access token arrives only as a Set-Cookie header,
+	// so this is the sole source on every refresh, not a fallback. The body
+	// check is kept in case the endpoint ever starts returning one.
 	if result.AccessToken == "" {
 		for _, c := range resp.Cookies() {
-			if c.Name == "access_token" {
-				result.AccessToken = c.Value
-				if result.ExpiresIn == 0 && c.MaxAge > 0 {
-					result.ExpiresIn = c.MaxAge
-				}
-				break
+			if c.Name != "access_token" || c.Value == "" {
+				continue
+			}
+			result.AccessToken = c.Value
+			if result.ExpiresIn == 0 {
+				result.ExpiresIn = cookieLifetime(c)
 			}
+			break
 		}
 	}
 
@@ -329,6 +365,24 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string,
 	return result.AccessToken, nil
 }
 
+// cookieLifetime returns a cookie's remaining lifetime in seconds, preferring
+// Max-Age and falling back to Expires. Returns 0 when the cookie carries
+// neither, which refreshToken stores as an already-elapsed ExpiresAt: an
+// unknown lifetime forces a refresh on next use rather than being trusted.
+// Storing 0 instead would read as "never expires" in GetValidToken, which
+// only checks an expiry it was actually given.
+func cookieLifetime(c *http.Cookie) int {
+	if c.MaxAge > 0 {
+		return c.MaxAge
+	}
+	if !c.Expires.IsZero() {
+		if secs := int(time.Until(c.Expires).Seconds()); secs > 0 {
+			return secs
+		}
+	}
+	return 0
+}
+
 // apiHost returns the API hostname, prepending "api." if not already present.
 func apiHost(host string) string {
 	if strings.HasPrefix(host, "api.") {

internal/auth/login_test.go

@@ -0,0 +1,222 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+// newAuthServer points the package's auth calls at h and isolates the
+// credential store, which SaveHostCredentials writes under os.UserHomeDir().
+func newAuthServer(t *testing.T, h http.HandlerFunc) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(h)
+	oldBase, oldClient := authBaseURL, httpClient
+	authBaseURL = func(string) string { return srv.URL }
+	httpClient = srv.Client()
+	t.Cleanup(func() {
+		authBaseURL, httpClient = oldBase, oldClient
+		srv.Close()
+	})
+	t.Setenv("HOME", t.TempDir())
+	return srv
+}
+
+// refreshHandler mirrors what GET /auth/refresh_token actually returns: the
+// access token only as a Set-Cookie header, with the metadata in the body.
+// See setCookies and RefreshTokenOutput in auth/pkg/endpoints.
+func refreshHandler(w http.ResponseWriter, _ *http.Request) {
+	http.SetCookie(w, &http.Cookie{
+		Name:     "access_token",
+		Value:    "new-id-token",
+		MaxAge:   3600,
+		Path:     "/",
+		Secure:   true,
+		HttpOnly: true,
+		SameSite: http.SameSiteNoneMode,
+	})
+	w.Header().Set("Content-Type", "application/json")
+	_, _ = w.Write([]byte(`{"expiresIn":3600,"installationIds":[],"expiresAt":1,` +
+		`"queryParameters":{"githubOwnerId":"42"},"isTriaged":false,"isNullifyDev":false}`))
+}
+
+func TestRefreshTokenReadsSetCookie(t *testing.T) {
+	newAuthServer(t, refreshHandler)
+
+	got, err := refreshToken(context.Background(), "example.nullify.ai", "stored-refresh")
+	if err != nil {
+		t.Fatalf("refreshToken: %v", err)
+	}
+	if got != "new-id-token" {
+		t.Fatalf("token = %q, want new-id-token", got)
+	}
+
+	creds, err := LoadCredentials()
+	if err != nil {
+		t.Fatalf("LoadCredentials: %v", err)
+	}
+	saved := creds[CredentialKey("example.nullify.ai")]
+	if saved.AccessToken != "new-id-token" {
+		t.Errorf("saved AccessToken = %q, want new-id-token", saved.AccessToken)
+	}
+	// The endpoint does not rotate the refresh cookie, so the stored one must
+	// survive or every later refresh fails.
+	if saved.RefreshToken != "stored-refresh" {
+		t.Errorf("saved RefreshToken = %q, want stored-refresh", saved.RefreshToken)
+	}
+	if d := saved.ExpiresAt - time.Now().Unix(); d < 3500 || d > 3600 {
+		t.Errorf("ExpiresAt delta = %ds, want ~3600", d)
+	}
+	if saved.QueryParameters["githubOwnerId"] != "42" {
+		t.Errorf("saved QueryParameters = %v, want githubOwnerId=42", saved.QueryParameters)
+	}
+}
+
+func TestRefreshTokenBodyTokenWins(t *testing.T) {
+	newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
+		http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "cookie-token", MaxAge: 3600})
+		_, _ = w.Write([]byte(`{"accessToken":"body-token","expiresIn":60}`))
+	})
+
+	got, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
+	if err != nil {
+		t.Fatalf("refreshToken: %v", err)
+	}
+	if got != "body-token" {
+		t.Errorf("token = %q, want body-token to take precedence", got)
+	}
+}
+
+func TestRefreshTokenIgnoresEmptyCookie(t *testing.T) {
+	newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
+		// A cleared cookie ahead of the live one must not win.
+		http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "", MaxAge: -1})
+		http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "real-token", MaxAge: 3600})
+		_, _ = w.Write([]byte(`{"expiresIn":3600}`))
+	})
+
+	got, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
+	if err != nil {
+		t.Fatalf("refreshToken: %v", err)
+	}
+	if got != "real-token" {
+		t.Errorf("token = %q, want real-token", got)
+	}
+}
+
+func TestRefreshTokenExpiresFallback(t *testing.T) {
+	newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
+		// Max-Age absent: the lifetime must come from Expires rather than
+		// collapsing to an already-expired credential.
+		http.SetCookie(w, &http.Cookie{
+			Name:    "access_token",
+			Value:   "expires-token",
+			Expires: time.Now().Add(time.Hour),
+		})
+		_, _ = w.Write([]byte(`{}`))
+	})
+
+	if _, err := refreshToken(context.Background(), "example.nullify.ai", "rt"); err != nil {
+		t.Fatalf("refreshToken: %v", err)
+	}
+
+	creds, err := LoadCredentials()
+	if err != nil {
+		t.Fatalf("LoadCredentials: %v", err)
+	}
+	if d := creds[CredentialKey("example.nullify.ai")].ExpiresAt - time.Now().Unix(); d < 3500 {
+		t.Errorf("ExpiresAt delta = %ds, want ~3600 from the Expires attribute", d)
+	}
+}
+
+func TestRefreshTokenNoTokenAnywhere(t *testing.T) {
+	newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write([]byte(`{"expiresIn":3600}`))
+	})
+
+	_, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
+	if err == nil {
+		t.Fatal("want an error when neither body nor cookie carries a token")
+	}
+}
+
+func TestRefreshTokenNon200(t *testing.T) {
+	newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusUnauthorized)
+	})
+
+	_, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
+	if err == nil {
+		t.Fatal("want an error on a non-200 refresh")
+	}
+}
+
+func TestGetValidTokenReturnsStoredTokenWhenUnexpired(t *testing.T) {
+	newAuthServer(t, func(http.ResponseWriter, *http.Request) {
+		t.Error("refresh must not be attempted while the token is still valid")
+	})
+	seed(t, "example.nullify.ai", "live-token", "rt", time.Now().Add(time.Hour).Unix())
+
+	got, err := GetValidToken(context.Background(), "example.nullify.ai")
+	if err != nil {
+		t.Fatalf("GetValidToken: %v", err)
+	}
+	if got != "live-token" {
+		t.Errorf("token = %q, want live-token", got)
+	}
+}
+
+func TestGetValidTokenRefreshesWhenExpired(t *testing.T) {
+	newAuthServer(t, refreshHandler)
+	seed(t, "example.nullify.ai", "old-token", "rt", time.Now().Add(-time.Hour).Unix())
+
+	got, err := GetValidToken(context.Background(), "example.nullify.ai")
+	if err != nil {
+		t.Fatalf("GetValidToken: %v", err)
+	}
+	if got != "new-id-token" {
+		t.Errorf("token = %q, want the refreshed new-id-token", got)
+	}
+}
+
+// ForceRefreshToken exists so a caller holding a token the server rejects can
+// get a new one even though the stored expiry still looks valid.
+func TestForceRefreshTokenIgnoresUnexpiredCredential(t *testing.T) {
+	newAuthServer(t, refreshHandler)
+	seed(t, "example.nullify.ai", "revoked-token", "rt", time.Now().Add(time.Hour).Unix())
+
+	got, err := ForceRefreshToken(context.Background(), "example.nullify.ai")
+	if err != nil {
+		t.Fatalf("ForceRefreshToken: %v", err)
+	}
+	if got != "new-id-token" {
+		t.Errorf("token = %q, want new-id-token", got)
+	}
+}
+
+func TestForceRefreshTokenWithoutRefreshToken(t *testing.T) {
+	newAuthServer(t, func(http.ResponseWriter, *http.Request) {
+		t.Error("refresh must not be attempted without a stored refresh token")
+	})
+	seed(t, "example.nullify.ai", "tok", "", time.Now().Add(time.Hour).Unix())
+
+	_, err := ForceRefreshToken(context.Background(), "example.nullify.ai")
+	if !errors.Is(err, ErrNotRefreshable) {
+		t.Fatalf("err = %v, want ErrNotRefreshable", err)
+	}
+}
+
+func seed(t *testing.T, host, access, refresh string, expiresAt int64) {
+	t.Helper()
+	err := SaveHostCredentials(host, HostCredentials{
+		AccessToken:  access,
+		RefreshToken: refresh,
+		ExpiresAt:    expiresAt,
+	})
+	if err != nil {
+		t.Fatalf("seed credentials: %v", err)
+	}
+}