From 1965cd9281162e886f7a3d6a50b04dde16de9798 Mon Sep 17 00:00:00 2001
From: jlenon7 <lenonSec7@gmail.com>
Date: Thu, 2 Jan 2025 16:34:19 -0300
Subject: [PATCH] chore(edge): add global attributes to style/script tag

---
 package-lock.json          |  4 +--
 package.json               |  2 +-
 src/fastify/FastifyVite.ts |  7 ++---
 src/vite/Vite.ts           | 52 +++++++++++++++++++++++++++++++++-----
 4 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index d393c17..af4f34e 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@athenna/vite",
-  "version": "5.1.0",
+  "version": "5.2.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@athenna/vite",
-      "version": "5.1.0",
+      "version": "5.2.0",
       "license": "MIT",
       "dependencies": {
         "@fastify/middie": "^9.0.2",
diff --git a/package.json b/package.json
index fae095f..f458526 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@athenna/vite",
-  "version": "5.1.0",
+  "version": "5.2.0",
   "description": "Vite plugin for Athenna Framework.",
   "license": "MIT",
   "author": "João Lenon <lenon@athenna.io>",
diff --git a/src/fastify/FastifyVite.ts b/src/fastify/FastifyVite.ts
index 42eab64..81be778 100644
--- a/src/fastify/FastifyVite.ts
+++ b/src/fastify/FastifyVite.ts
@@ -61,9 +61,6 @@ export class FastifyVite {
      * Don't start vite server under production.
      */
     if (this.options.dev) {
-      debug('vite server creation bypassed because app is in production mode.')
-      return
-    } else {
       const { createServer, loadConfigFromFile } = await import('vite')
 
       const defaultConfig = {
@@ -90,6 +87,10 @@ export class FastifyVite {
       await this.scope.register(fastifyMiddie, { hook: 'onRequest' })
 
       this.scope.use(this.devServer.middlewares)
+
+      this.scope.addHook('onClose', () => this.devServer.close())
+    } else {
+      debug('vite server creation bypassed because app is in production mode.')
     }
 
     const vite = new Vite(this.options, this.devServer)
diff --git a/src/vite/Vite.ts b/src/vite/Vite.ts
index 6a1075e..9a81856 100644
--- a/src/vite/Vite.ts
+++ b/src/vite/Vite.ts
@@ -15,7 +15,7 @@
 import path from 'node:path'
 
 import { File } from '@athenna/common'
-import type { FastifyViteOptions } from '#src/types'
+import type { FastifyViteOptions, SetAttributes } from '#src/types'
 import type { Manifest, ModuleNode, ViteDevServer } from 'vite'
 
 const styleFileRegex = /\.(css|less|sass|scss|styl|stylus|pcss|postcss)($|\?)/
@@ -148,23 +148,61 @@ export class Vite {
     return false
   }
 
+  /**
+   * Unwrap attributes from the user defined function or return
+   * the attributes as it is.
+   */
+  public unwrapAttributes(
+    src: string,
+    url: string,
+    attributes?: SetAttributes
+  ) {
+    if (typeof attributes === 'function') {
+      return attributes({ src, url })
+    }
+
+    return attributes
+  }
+
   /**
    * Create a style tag for the given path.
    */
-  public makeStyleTag(url: string, attributes?: any) {
+  public makeStyleTag(src: string, url: string, attributes?: any) {
+    const customAttributes = this.unwrapAttributes(
+      src,
+      url,
+      this.options?.styleAttributes
+    )
+
     return this.generateElement({
       tag: 'link',
-      attributes: { rel: 'stylesheet', href: url, ...attributes }
+      attributes: {
+        rel: 'stylesheet',
+        ...customAttributes,
+        ...attributes,
+        href: url
+      }
     })
   }
 
   /**
    * Create a script tag for the given path
    */
-  public makeScriptTag(url: string, attributes?: any) {
+  public makeScriptTag(src: string, url: string, attributes?: any) {
+    const customAttributes = this.unwrapAttributes(
+      src,
+      url,
+      this.options?.scriptAttributes
+    )
+
     return this.generateElement({
       tag: 'script',
-      attributes: { type: 'module', src: url, ...attributes },
+      attributes: {
+        type: 'module',
+        ...customAttributes,
+        ...attributes,
+        src: url
+      },
       children: []
     })
   }
@@ -182,10 +220,10 @@ export class Vite {
     }
 
     if (this.isCssPath(asset)) {
-      return this.makeStyleTag(url, attributes)
+      return this.makeStyleTag(asset, url, attributes)
     }
 
-    return this.makeScriptTag(url, attributes)
+    return this.makeScriptTag(asset, url, attributes)
   }
 
   /**