refactor(sessions): import driver via vite plugin

packages/astro/dev-only.d.ts

@@ -44,3 +44,8 @@ declare module 'virtual:astro:middleware' {
 	const middleware: AstroMiddlewareInstance;
 	export default middleware;
 }
+
+declare module 'virtual:astro:session-driver' {
+	import type { Driver } from 'unstorage';
+	export const driver: Driver;
+}

packages/astro/src/core/app/entrypoint.ts

@@ -11,6 +11,7 @@ const manifest: SSRManifest = Object.assign(serializedManifest, {
 	renderers,
 	actions: () => import('virtual:astro:actions/entrypoint'),
 	middleware: () => import('virtual:astro:middleware'),
+	sessionDriver: () => import('virtual:astro:session-driver'),
 	routes,
 });
 

packages/astro/src/core/app/types.ts

@@ -17,6 +17,7 @@ import type {
 } from '../../types/public/internal.js';
 import type { SinglePageBuiltModule } from '../build/types.js';
 import type { CspDirective } from '../csp/config.js';
+import type { SessionDriver } from '../session.js';
 import type { RoutingStrategies } from './common.js';
 
 type ComponentPath = string;
@@ -84,6 +85,7 @@ export type SSRManifest = {
 	i18n: SSRManifestI18n | undefined;
 	middleware?: () => Promise<AstroMiddlewareInstance> | AstroMiddlewareInstance;
 	actions?: () => Promise<SSRActions> | SSRActions;
+	sessionDriver?: () => Promise<{ default: SessionDriver | null }>;
 	checkOrigin: boolean;
 	allowedDomains?: Partial<RemotePattern>[];
 	sessionConfig?: ResolvedSessionConfig<any>;

packages/astro/src/core/base-pipeline.ts

@@ -22,6 +22,7 @@ import { sequence } from './middleware/sequence.js';
 import { RedirectSinglePageBuiltModule } from './redirects/index.js';
 import { RouteCache } from './render/route-cache.js';
 import { createDefaultRoutes } from './routing/default.js';
+import type { SessionDriver } from './session.js';
 
 /**
  * The `Pipeline` represents the static parts of rendering that do not change between requests.
@@ -33,6 +34,7 @@ export abstract class Pipeline {
 	readonly internalMiddleware: MiddlewareHandler[];
 	resolvedMiddleware: MiddlewareHandler | undefined = undefined;
 	resolvedActions: SSRActions | undefined = undefined;
+	resolvedSessionDriver: SessionDriver | null | undefined = undefined;
 
 	constructor(
 		readonly logger: Logger,
@@ -69,6 +71,7 @@ export abstract class Pipeline {
 		readonly defaultRoutes = createDefaultRoutes(manifest),
 
 		readonly actions = manifest.actions,
+		readonly sessionDriver = manifest.sessionDriver,
 	) {
 		this.internalMiddleware = [];
 		// We do use our middleware only if the user isn't using the manual setup
@@ -140,6 +143,24 @@ export abstract class Pipeline {
 		return NOOP_ACTIONS_MOD;
 	}
 
+	async getSessionDriver(): Promise<SessionDriver | null> {
+		// Return cached value if already resolved (including null)
+		if (this.resolvedSessionDriver !== undefined) {
+			return this.resolvedSessionDriver;
+		}
+
+		// Try to load the driver from the manifest
+		if (this.sessionDriver) {
+			const driverModule = await this.sessionDriver();
+			this.resolvedSessionDriver = driverModule?.default || null;
+			return this.resolvedSessionDriver;
+		}
+
+		// No driver configured
+		this.resolvedSessionDriver = null;
+		return null;
+	}
+
 	async getAction(path: string): Promise<ActionClient<unknown, ActionAccept, ZodType>> {
 		const pathKeys = path.split('.').map((key) => decodeURIComponent(key));
 		let { server } = await this.getActions();

packages/astro/src/core/build/plugins/plugin-manifest.ts

@@ -1,7 +1,6 @@
 import { fileURLToPath } from 'node:url';
 import type { OutputChunk } from 'rollup';
 import { glob } from 'tinyglobby';
-import { type BuiltinDriverName, builtinDrivers } from 'unstorage';
 import type { Plugin as VitePlugin } from 'vite';
 import { getAssetsPrefix } from '../../../assets/utils/getAssetsPrefix.js';
 import { normalizeTheLocale } from '../../../i18n/index.js';
@@ -43,25 +42,7 @@ const replaceExp = new RegExp(`['"]${manifestReplace}['"]`, 'g');
 export const SSR_MANIFEST_VIRTUAL_MODULE_ID = '@astrojs-manifest';
 export const RESOLVED_SSR_MANIFEST_VIRTUAL_MODULE_ID = '\0' + SSR_MANIFEST_VIRTUAL_MODULE_ID;
 
-function resolveSessionDriver(driver: string | undefined): string | null {
-	if (!driver) {
-		return null;
-	}
-	try {
-		if (driver === 'fs') {
-			return import.meta.resolve(builtinDrivers.fsLite, import.meta.url);
-		}
-		if (driver in builtinDrivers) {
-			return import.meta.resolve(builtinDrivers[driver as BuiltinDriverName], import.meta.url);
-		}
-	} catch {
-		return null;
-	}
-
-	return driver;
-}
-
-function vitePluginManifest(options: StaticBuildOptions, internals: BuildInternals): VitePlugin {
+function vitePluginManifest(internals: BuildInternals): VitePlugin {
 	return {
 		name: '@astro/plugin-build-manifest',
 		enforce: 'post',
@@ -85,11 +66,8 @@ function vitePluginManifest(options: StaticBuildOptions, internals: BuildInterna
 					`import { _privateSetManifestDontUseThis } from 'astro:ssr-manifest'`,
 				];
 
-				const resolvedDriver = resolveSessionDriver(options.settings.config.session?.driver);
-
 				const contents = [
 					`const manifest = _deserializeManifest('${manifestReplace}');`,
-					`if (manifest.sessionConfig) manifest.sessionConfig.driverModule = ${resolvedDriver ? `() => import(${JSON.stringify(resolvedDriver)})` : 'null'};`,
 					`_privateSetManifestDontUseThis(manifest);`,
 				];
 				const exports = [`export { manifest }`];
@@ -124,7 +102,7 @@ export function pluginManifest(
 		hooks: {
 			'build:before': () => {
 				return {
-					vitePlugin: vitePluginManifest(options, internals),
+					vitePlugin: vitePluginManifest(internals),
 				};
 			},
 

packages/astro/src/core/build/plugins/plugin-ssr.ts

@@ -5,6 +5,7 @@ import { ASTRO_RENDERERS_MODULE_ID } from '../../../vite-plugin-renderers/index.
 import { MIDDLEWARE_MODULE_ID } from '../../middleware/vite-plugin.js';
 import { routeIsRedirect } from '../../redirects/index.js';
 import { VIRTUAL_ISLAND_MAP_ID } from '../../server-islands/vite-plugin-server-islands.js';
+import { VIRTUAL_SESSION_DRIVER_ID } from '../../session/vite-plugin.js';
 import { addRollupInput } from '../add-rollup-input.js';
 import type { BuildInternals } from '../internal.js';
 import type { AstroBuildPlugin } from '../plugin.js';
@@ -180,7 +181,8 @@ function generateSSRCode(adapter: AstroAdapter) {
 		`    serverIslandMap,`,
 		`    renderers,`,
 		`    actions: () => import("${ENTRYPOINT_VIRTUAL_MODULE_ID}"),`,
-		`    middleware: ${edgeMiddleware ? 'undefined' : `() => import("${MIDDLEWARE_MODULE_ID}")`}`,
+		`    middleware: ${edgeMiddleware ? 'undefined' : `() => import("${MIDDLEWARE_MODULE_ID}")`},`,
+		`    sessionDriver: () => import("${VIRTUAL_SESSION_DRIVER_ID}")`,
 		`});`,
 		`const _args = ${adapter.args ? JSON.stringify(adapter.args, null, 4) : 'undefined'};`,
 		adapter.exports

packages/astro/src/core/create-vite.ts

@@ -43,6 +43,7 @@ import { createViteLogger } from './logger/vite.js';
 import { vitePluginMiddleware } from './middleware/vite-plugin.js';
 import { joinPaths } from './path.js';
 import { vitePluginServerIslands } from './server-islands/vite-plugin-server-islands.js';
+import { vitePluginSessionDriver } from './session/vite-plugin.js';
 import { isObject } from './util.js';
 
 type CreateViteOptions = {
@@ -174,6 +175,7 @@ export async function createVite(
 			astroInternationalization({ settings }),
 			vitePluginActions({ fs, settings }),
 			vitePluginServerIslands({ settings, logger }),
+			vitePluginSessionDriver({ settings }),
 			astroContainer(),
 			astroHmrReloadPlugin(),
 		],

packages/astro/src/core/render-context.ts

@@ -77,9 +77,7 @@ export class RenderContext {
 		public props: Props = {},
 		public partial: undefined | boolean = undefined,
 		public shouldInjectCspMetaTags = !!pipeline.manifest.csp,
-		public session: AstroSession | undefined = pipeline.manifest.sessionConfig
-			? new AstroSession(cookies, pipeline.manifest.sessionConfig, pipeline.runtimeMode)
-			: undefined,
+		public session: AstroSession | undefined = undefined,
 	) {}
 
 	/**
@@ -109,12 +107,23 @@ export class RenderContext {
 	}: CreateRenderContext): Promise<RenderContext> {
 		const pipelineMiddleware = await pipeline.getMiddleware();
 		const pipelineActions = actions ?? (await pipeline.getActions());
+		const pipelineSessionDriver = await pipeline.getSessionDriver();
 		setOriginPathname(
 			request,
 			pathname,
 			pipeline.manifest.trailingSlash,
 			pipeline.manifest.buildFormat,
 		);
+		const cookies = new AstroCookies(request);
+		const session =
+			pipeline.manifest.sessionConfig && pipelineSessionDriver
+				? new AstroSession(
+						cookies,
+						pipeline.manifest.sessionConfig,
+						pipeline.runtimeMode,
+						pipelineSessionDriver,
+					)
+				: undefined;
 		return new RenderContext(
 			pipeline,
 			locals,
@@ -125,12 +134,13 @@ export class RenderContext {
 			routeData,
 			status,
 			clientAddress,
-			undefined,
+			cookies,
 			undefined,
 			undefined,
 			props,
 			partial,
 			shouldInjectCspMetaTags ?? !!pipeline.manifest.csp,
+			session,
 		);
 	}
 	/**

packages/astro/src/core/session.ts

@@ -1,13 +1,6 @@
 import { stringify as rawStringify, unflatten as rawUnflatten } from 'devalue';
 
-import {
-	type BuiltinDriverName,
-	type BuiltinDriverOptions,
-	builtinDrivers,
-	createStorage,
-	type Driver,
-	type Storage,
-} from 'unstorage';
+import { type BuiltinDriverOptions, createStorage, type Driver, type Storage } from 'unstorage';
 import type {
 	ResolvedSessionConfig,
 	RuntimeMode,
@@ -23,6 +16,10 @@ export const PERSIST_SYMBOL = Symbol();
 const DEFAULT_COOKIE_NAME = 'astro-session';
 const VALID_COOKIE_REGEX = /^[\w-]+$/;
 
+export type SessionDriver<TDriver extends SessionDriverName = any> = (
+	config: SessionConfig<TDriver>['options'],
+) => import('unstorage').Driver;
+
 interface SessionEntry {
 	data: any;
 	expires?: number;
@@ -72,6 +69,8 @@ export class AstroSession<TDriver extends SessionDriverName = any> {
 	// When we load the data from storage, we need to merge it with the local partial data,
 	// preserving in-memory changes and deletions.
 	#partial = true;
+	// The driver factory function provided by the pipeline
+	#driverFactory: SessionDriver | null | undefined;
 
 	static #sharedStorage = new Map<string, Storage>();
 
@@ -82,6 +81,7 @@ export class AstroSession<TDriver extends SessionDriverName = any> {
 			...config
 		}: NonNullable<ResolvedSessionConfig<TDriver>>,
 		runtimeMode?: RuntimeMode,
+		driverFactory?: ((config: SessionConfig<TDriver>['options']) => Driver) | null,
 	) {
 		const { driver } = config;
 		if (!driver) {
@@ -93,6 +93,7 @@ export class AstroSession<TDriver extends SessionDriverName = any> {
 			});
 		}
 		this.#cookies = cookies;
+		this.#driverFactory = driverFactory;
 		let cookieConfigObject: AstroCookieSetOptions | undefined;
 		if (typeof cookieConfig === 'object') {
 			const { name = DEFAULT_COOKIE_NAME, ...rest } = cookieConfig;
@@ -441,46 +442,19 @@ export class AstroSession<TDriver extends SessionDriverName = any> {
 			(this.#config.options as BuiltinDriverOptions['fs-lite']).base ??= '.astro/session';
 		}
 
-		let driver: ((config: SessionConfig<TDriver>['options']) => Driver) | null = null;
-
-		try {
-			if (this.#config.driverModule) {
-				driver = (await this.#config.driverModule()).default;
-			} else if (this.#config.driver) {
-				const driverName = resolveSessionDriverName(this.#config.driver);
-				if (driverName) {
-					driver = (await import(/* @vite-ignore */ driverName)).default;
-				}
-			}
-		} catch (err: any) {
-			// If the driver failed to load, throw an error.
-			if (err.code === 'ERR_MODULE_NOT_FOUND') {
-				throw new AstroError(
-					{
-						...SessionStorageInitError,
-						message: SessionStorageInitError.message(
-							err.message.includes(`Cannot find package`)
-								? 'The driver module could not be found.'
-								: err.message,
-							this.#config.driver,
-						),
-					},
-					{ cause: err },
-				);
-			}
-			throw err;
-		}
-
-		if (!driver) {
+		// Get the driver factory from the pipeline
+		if (!this.#driverFactory) {
 			throw new AstroError({
 				...SessionStorageInitError,
 				message: SessionStorageInitError.message(
-					'The module did not export a driver.',
+					'Astro could not load the driver correctly. Does it exist?',
 					this.#config.driver,
 				),
 			});
 		}
 
+		const driver = this.#driverFactory;
+
 		try {
 			this.#storage = createStorage({
 				driver: driver(this.#config.options),
@@ -498,21 +472,3 @@ export class AstroSession<TDriver extends SessionDriverName = any> {
 		}
 	}
 }
-
-function resolveSessionDriverName(driver: string | undefined): string | null {
-	if (!driver) {
-		return null;
-	}
-	try {
-		if (driver === 'fs') {
-			return builtinDrivers.fsLite;
-		}
-		if (driver in builtinDrivers) {
-			return builtinDrivers[driver as BuiltinDriverName];
-		}
-	} catch {
-		return null;
-	}
-
-	return driver;
-}

packages/astro/src/core/session/vite-plugin.ts

@@ -0,0 +1,36 @@
+import { type BuiltinDriverName, builtinDrivers } from 'unstorage';
+import type { Plugin as VitePlugin } from 'vite';
+import type { AstroSettings } from '../../types/astro.js';
+
+export const VIRTUAL_SESSION_DRIVER_ID = 'virtual:astro:session-driver';
+const RESOLVED_VIRTUAL_SESSION_DRIVER_ID = '\0' + VIRTUAL_SESSION_DRIVER_ID;
+
+export function vitePluginSessionDriver({ settings }: { settings: AstroSettings }): VitePlugin {
+	return {
+		name: VIRTUAL_SESSION_DRIVER_ID,
+		enforce: 'pre',
+
+		async resolveId(id) {
+			if (id === VIRTUAL_SESSION_DRIVER_ID) {
+				if (settings.config.session) {
+					if (settings.config.session.driver === 'fs') {
+						return await this.resolve(builtinDrivers.fsLite);
+					}
+					if (settings.config.session.driver && settings.config.session.driver in builtinDrivers) {
+						return await this.resolve(
+							builtinDrivers[settings.config.session.driver as BuiltinDriverName],
+						);
+					}
+				} else {
+					return RESOLVED_VIRTUAL_SESSION_DRIVER_ID;
+				}
+			}
+		},
+
+		async load(id) {
+			if (id === RESOLVED_VIRTUAL_SESSION_DRIVER_ID) {
+				return { code: 'export default null;' };
+			}
+		},
+	};
+}