diff --git a/ui/packages/ui-plugin-bundler-kit/README.md b/ui/packages/ui-plugin-bundler-kit/README.md
index b5aaaf61a..d1817bedc 100644
--- a/ui/packages/ui-plugin-bundler-kit/README.md
+++ b/ui/packages/ui-plugin-bundler-kit/README.md
@@ -114,7 +114,7 @@ interface ViteUserConfig {
   /**
    * Custom Vite configuration
    */
-  vite: UserConfig;
+  vite: UserConfig | UserConfigFnObject;
 }
 ```
 
@@ -131,7 +131,7 @@ interface RsBuildUserConfig {
   /**
    * Custom Rsbuild configuration
    */
-  rsbuild: RsbuildConfig;
+  rsbuild: RsbuildConfig | ((env: ConfigParams) => RsbuildConfig);
 }
 ```
 
diff --git a/ui/packages/ui-plugin-bundler-kit/src/rsbuild.ts b/ui/packages/ui-plugin-bundler-kit/src/rsbuild.ts
index fa040e28f..08f6e24d2 100644
--- a/ui/packages/ui-plugin-bundler-kit/src/rsbuild.ts
+++ b/ui/packages/ui-plugin-bundler-kit/src/rsbuild.ts
@@ -3,6 +3,7 @@ import {
   mergeRsbuildConfig,
   type RsbuildConfig,
   type RsbuildMode,
+  type ConfigParams,
 } from "@rsbuild/core";
 import { getHaloPluginManifest } from "./utils/halo-plugin";
 import { DEFAULT_OUT_DIR_DEV, DEFAULT_OUT_DIR_PROD } from "./constants/build";
@@ -21,7 +22,7 @@ export interface RsBuildUserConfig {
   /**
    * Custom Rsbuild config.
    */
-  rsbuild: RsbuildConfig;
+  rsbuild: RsbuildConfig | ((env: ConfigParams) => RsbuildConfig);
 }
 
 function createRsbuildPresetsConfig(manifestPath: string) {
@@ -134,6 +135,10 @@ export function rsbuildConfig(config?: RsBuildUserConfig) {
   );
   return defineConfig((env) => {
     const presetsConfig = presetsConfigFn(env);
-    return mergeRsbuildConfig(presetsConfig, config?.rsbuild || {});
+    const userConfig =
+      typeof config?.rsbuild === "function"
+        ? config.rsbuild(env)
+        : config?.rsbuild || {};
+    return mergeRsbuildConfig(presetsConfig, userConfig);
   });
 }
diff --git a/ui/packages/ui-plugin-bundler-kit/src/vite.ts b/ui/packages/ui-plugin-bundler-kit/src/vite.ts
index 89fd73f1b..afc5b2815 100644
--- a/ui/packages/ui-plugin-bundler-kit/src/vite.ts
+++ b/ui/packages/ui-plugin-bundler-kit/src/vite.ts
@@ -1,4 +1,9 @@
-import { defineConfig, mergeConfig, UserConfig } from "vite";
+import {
+  defineConfig,
+  mergeConfig,
+  UserConfig,
+  UserConfigFnObject,
+} from "vite";
 import Vue from "@vitejs/plugin-vue";
 import { EXTERNALS, GLOBALS } from "./constants/externals";
 import { DEFAULT_OUT_DIR_DEV, DEFAULT_OUT_DIR_PROD } from "./constants/build";
@@ -16,7 +21,7 @@ export interface ViteUserConfig {
   /**
    * Custom Vite config.
    */
-  vite: UserConfig;
+  vite: UserConfig | UserConfigFnObject;
 }
 
 function createVitePresetsConfig(manifestPath: string) {
@@ -72,6 +77,10 @@ export function viteConfig(config?: ViteUserConfig) {
   );
   return defineConfig((env) => {
     const presetsConfig = presetsConfigFn(env);
-    return mergeConfig(presetsConfig, config?.vite || {});
+    const userConfig =
+      typeof config?.vite === "function"
+        ? config.vite(env)
+        : config?.vite || {};
+    return mergeConfig(presetsConfig, userConfig);
   });
 }