@agentuity/runtime 1.0.47 → 2.0.0-beta.0

package/src/bootstrap.ts

@@ -0,0 +1,316 @@
+/**
+ * Server lifecycle helpers.
+ *
+ * These functions are called by createApp() to set up routes, middleware,
+ * and the Bun HTTP server. They're kept separate to keep createApp() focused
+ * on orchestration while these handle the details.
+ */
+
+import type { Context } from 'hono';
+import { websocket, serveStatic } from 'hono/bun';
+import { readFileSync, existsSync } from 'node:fs';
+import { join } from 'node:path';
+import { mimeTypes } from '@agentuity/server';
+
+import { runShutdown } from './app';
+import type { AnalyticsOptions, WorkbenchOptions } from './app';
+import { createRouter } from './router';
+import { createWebSessionMiddleware } from './middleware';
+import { enableProcessExitProtection } from './_process-protection';
+import { hasWaitUntilPending } from './_waituntil';
+import { getOrganizationId, getProjectId, isDevMode as runtimeIsDevMode } from './_config';
+import { BEACON_SCRIPT } from '@agentuity/frontend';
+
+// ============================================================================
+// Mode detection
+// ============================================================================
+
+/**
+ * Runtime mode detection.
+ * Dynamic string concatenation prevents Bun.build from inlining NODE_ENV.
+ * @see https://github.com/oven-sh/bun/issues/20183
+ */
+const getEnv = (key: string) => process.env[key];
+/**
+ * Check if running in development mode.
+ *
+ * The CLI dev server explicitly sets NODE_ENV='development'. In production
+ * (cloud deployment, CI integration test running a built app.js), NODE_ENV
+ * may be 'production' or unset entirely. When unset, we assume production
+ * — the dev server always sets it, so absence means production.
+ */
+export const isDevelopment = () => getEnv('NODE' + '_' + 'ENV') === 'development';
+
+// ============================================================================
+// Analytics helpers
+// ============================================================================
+
+/** Resolve analytics config with defaults */
+export function resolveAnalyticsConfig(
+	analytics: boolean | AnalyticsOptions | undefined
+): AnalyticsOptions & { enabled: boolean } {
+	if (analytics === false) {
+		return { enabled: false };
+	}
+	const opts = typeof analytics === 'object' ? analytics : {};
+	return {
+		enabled: opts.enabled !== false,
+		requireConsent: opts.requireConsent ?? false,
+		trackClicks: opts.trackClicks ?? true,
+		trackScroll: opts.trackScroll ?? true,
+		trackOutboundLinks: opts.trackOutboundLinks ?? true,
+		trackForms: opts.trackForms ?? false,
+		trackWebVitals: opts.trackWebVitals ?? true,
+		trackErrors: opts.trackErrors ?? true,
+		trackSPANavigation: opts.trackSPANavigation ?? true,
+		sampleRate: opts.sampleRate ?? 1,
+		excludePatterns: opts.excludePatterns ?? [],
+		globalProperties: opts.globalProperties ?? {},
+	};
+}
+
+/** Resolve workbench config */
+export function resolveWorkbenchConfig(
+	workbench: boolean | string | WorkbenchOptions | undefined
+): { enabled: boolean; route: string; headers: Record<string, string> } {
+	if (!workbench) {
+		return { enabled: false, route: '/workbench', headers: {} };
+	}
+	if (workbench === true) {
+		return { enabled: true, route: '/workbench', headers: {} };
+	}
+	if (typeof workbench === 'string') {
+		return { enabled: true, route: workbench, headers: {} };
+	}
+	return {
+		enabled: true,
+		route: workbench.route ?? '/workbench',
+		headers: workbench.headers ?? {},
+	};
+}
+
+/** Inject analytics scripts into HTML */
+function injectAnalytics(
+	html: string,
+	analyticsConfig: AnalyticsOptions & { enabled: boolean }
+): string {
+	if (!analyticsConfig.enabled) return html;
+
+	const orgId = getOrganizationId() || '';
+	const projectId = getProjectId() || '';
+	const isDevmode = runtimeIsDevMode();
+
+	const pageConfig = {
+		...analyticsConfig,
+		orgId,
+		projectId,
+		isDevmode,
+	};
+
+	const configScript = `<script>window.__AGENTUITY_ANALYTICS__=${JSON.stringify(pageConfig)};</script>`;
+	const sessionScript = '<script src="/_agentuity/webanalytics/session.js" async></script>';
+
+	// In production, the beacon is already in HTML as a CDN asset (data-agentuity-beacon marker)
+	const beaconMarker = '<script data-agentuity-beacon';
+	if (html.includes(beaconMarker)) {
+		const injection = configScript + sessionScript;
+		return html.replace(beaconMarker, injection + beaconMarker);
+	}
+
+	// Development: beacon served from local route
+	const beaconScript = '<script src="/_agentuity/webanalytics/analytics.js"></script>';
+	const injection = configScript + sessionScript + beaconScript;
+
+	if (html.includes('</head>')) {
+		return html.replace('</head>', injection + '</head>');
+	}
+	if (html.includes('<body')) {
+		return html.replace(/<body([^>]*)>/, `<body$1>${injection}`);
+	}
+	return injection + html;
+}
+
+/** Register analytics routes on the app */
+export function registerAnalyticsRoutes(
+	app: ReturnType<typeof createRouter>,
+	_analyticsConfig: AnalyticsOptions & { enabled: boolean }
+): void {
+	app.get(
+		'/_agentuity/webanalytics/session.js',
+		createWebSessionMiddleware(),
+		async (c: Context) => {
+			const threadId = c.get('_webThreadId') || '';
+			const sessionData = JSON.stringify({ threadId });
+			const sessionScript = `window.__AGENTUITY_SESSION__=${sessionData};`;
+			return new Response(sessionScript, {
+				headers: {
+					'Content-Type': 'application/javascript; charset=utf-8',
+					'Cache-Control': 'no-store, no-cache, must-revalidate',
+				},
+			});
+		}
+	);
+
+	if (isDevelopment()) {
+		app.get('/_agentuity/webanalytics/analytics.js', async () => {
+			return new Response(BEACON_SCRIPT, {
+				headers: {
+					'Content-Type': 'application/javascript; charset=utf-8',
+					'Cache-Control': 'no-store, no-cache, must-revalidate',
+				},
+			});
+		});
+	}
+}
+
+// ============================================================================
+// Health routes
+// ============================================================================
+
+export function registerHealthRoutes(app: ReturnType<typeof createRouter>): void {
+	if (!isDevelopment()) {
+		const healthHandler = (c: Context) => {
+			return c.text('OK', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+		};
+		const idleHandler = (c: Context) => {
+			const server = globalThis.__AGENTUITY_SERVER__;
+			if (!server) return c.text('NO', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+			if (hasWaitUntilPending())
+				return c.text('NO', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+			if (server.pendingRequests > 1)
+				return c.text('NO', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+			if (server.pendingWebSockets > 0)
+				return c.text('NO', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+			return c.text('OK', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+		};
+		app.get('/_agentuity/health', healthHandler);
+		app.get('/_health', healthHandler);
+		app.get('/_agentuity/idle', idleHandler);
+		app.get('/_idle', idleHandler);
+	}
+
+	if (isDevelopment()) {
+		app.get('/_agentuity/ready', (c: Context) => {
+			return c.text('OK', 200, { 'Content-Type': 'text/plain; charset=utf-8' });
+		});
+	}
+}
+
+// ============================================================================
+// Web routes (production static serving)
+// ============================================================================
+
+export function registerWebRoutes(
+	app: ReturnType<typeof createRouter>,
+	analyticsConfig: AnalyticsOptions & { enabled: boolean }
+): void {
+	if (isDevelopment()) return;
+
+	// Resolve client dir relative to the built app.js location (import.meta.dir)
+	// rather than process.cwd(), since the app may be run from inside .agentuity/
+	const appDir = typeof import.meta.dir === 'string' ? import.meta.dir : process.cwd();
+	const clientDir = join(appDir, 'client');
+	// Fallback: try process.cwd()/.agentuity/client if appDir/client doesn't exist
+	const resolvedClientDir = existsSync(join(clientDir, 'index.html'))
+		? clientDir
+		: join(process.cwd(), '.agentuity', 'client');
+	const indexHtmlPath = join(resolvedClientDir, 'index.html');
+	const baseIndexHtml = existsSync(indexHtmlPath) ? readFileSync(indexHtmlPath, 'utf-8') : '';
+
+	if (!baseIndexHtml) return;
+
+	const prodHtmlHandler = (c: Context) => {
+		if (analyticsConfig.enabled) {
+			const html = injectAnalytics(baseIndexHtml, analyticsConfig);
+			return c.html(html);
+		}
+		return c.html(baseIndexHtml);
+	};
+
+	app.get('/', prodHtmlHandler);
+	app.use('/assets/*', serveStatic({ root: resolvedClientDir, mimes: mimeTypes }));
+	app.use(
+		'/*',
+		serveStatic({
+			root: resolvedClientDir,
+			rewriteRequestPath: (path: string) => path,
+			mimes: mimeTypes,
+		})
+	);
+
+	app.all('/_agentuity/*', (c: Context) => c.notFound());
+	app.all('/api/*', (c: Context) => c.notFound());
+
+	app.get('*', (c: Context) => {
+		const path = c.req.path;
+		if (/\.[a-zA-Z0-9]+$/.test(path)) {
+			return c.notFound();
+		}
+		return prodHtmlHandler(c);
+	});
+}
+
+// ============================================================================
+// Workbench UI route (dev only)
+// ============================================================================
+
+export function registerWorkbenchUI(
+	app: ReturnType<typeof createRouter>,
+	workbenchConfig: { enabled: boolean; route: string }
+): void {
+	if (!workbenchConfig.enabled || !isDevelopment()) return;
+
+	const workbenchSrcDir = join(process.cwd(), '.agentuity', 'workbench-src');
+	const workbenchIndexPath = join(workbenchSrcDir, 'index.html');
+
+	app.get(workbenchConfig.route, async (c: Context) => {
+		const html = await Bun.file(workbenchIndexPath).text();
+		const withVite = html
+			.replace('src="./main.tsx"', `src="/@fs${workbenchSrcDir}/main.tsx"`)
+			.replace('href="./styles.css"', `href="/@fs${workbenchSrcDir}/styles.css"`);
+		return c.html(withVite);
+	});
+}
+
+// ============================================================================
+// Server startup
+// ============================================================================
+
+export function startServer(
+	app: ReturnType<typeof createRouter>,
+	options?: { requestTimeout?: number }
+): void {
+	if (typeof Bun === 'undefined') return;
+
+	enableProcessExitProtection();
+
+	const port = parseInt(process.env.PORT || '3500', 10);
+	const requestTimeout = options?.requestTimeout ?? 0;
+
+	const server = Bun.serve({
+		fetch: (req, server) => {
+			server.timeout(req, requestTimeout);
+			return app.fetch(req, server);
+		},
+		websocket,
+		port,
+		hostname: '127.0.0.1',
+		development: isDevelopment(),
+	});
+
+	globalThis.__AGENTUITY_SERVER__ = server;
+
+	if (!isDevelopment()) {
+		const handleShutdown = async (_signal: string) => {
+			try {
+				await runShutdown();
+			} catch {
+				// Ignore shutdown errors
+			}
+			process.exit(0);
+		};
+
+		process.once('SIGTERM', () => handleShutdown('SIGTERM'));
+		process.once('SIGINT', () => handleShutdown('SIGINT'));
+	}
+}

package/src/dev-patches/aisdk.ts

@@ -12,19 +12,16 @@
 function warnMissingKey(envKey: string): void {
 	const isDev =
 		process.env.AGENTUITY_ENVIRONMENT === 'development' || process.env.NODE_ENV !== 'production';
+
+	// Align no-bundle runtime behavior with historical build-time expectations:
+	// in development, don't emit eager startup errors for every provider.
 	if (isDev) {
-		console.error('[ERROR] No credentials found for this AI provider. To fix this, either:');
-		console.error(
-			'  1. Login to Agentuity Cloud (agentuity auth login) to use the AI Gateway (recommended)'
-		);
-		console.error(`  2. Set ${envKey} in your .env file to use the provider directly`);
-	} else {
-		console.error(`[ERROR] The environment variable ${envKey} is required. Either:`);
-		console.error(
-			'  1. Use Agentuity Cloud AI Gateway by ensuring AGENTUITY_SDK_KEY is configured'
-		);
-		console.error(`  2. Set ${envKey} using "agentuity env set ${envKey}" and redeploy`);
+		return;
 	}
+
+	console.error(`[ERROR] The environment variable ${envKey} is required. Either:`);
+	console.error('  1. Use Agentuity Cloud AI Gateway by ensuring AGENTUITY_SDK_KEY is configured');
+	console.error(`  2. Set ${envKey} using "agentuity env set ${envKey}" and redeploy`);
 }
 
 /**

package/src/dev-patches/gateway.ts

@@ -21,19 +21,17 @@ const GATEWAY_CONFIGS: GatewayConfig[] = [
 function warnMissingKey(envKey: string): void {
 	const isDev =
 		process.env.AGENTUITY_ENVIRONMENT === 'development' || process.env.NODE_ENV !== 'production';
+
+	// Align no-bundle runtime behavior with historical build-time expectations:
+	// in development, don't emit eager startup errors for every provider.
+	// Missing credentials should surface when/if that provider is actually used.
 	if (isDev) {
-		console.error('[ERROR] No credentials found for this AI provider. To fix this, either:');
-		console.error(
-			'  1. Login to Agentuity Cloud (agentuity auth login) to use the AI Gateway (recommended)'
-		);
-		console.error(`  2. Set ${envKey} in your .env file to use the provider directly`);
-	} else {
-		console.error(`[ERROR] The environment variable ${envKey} is required. Either:`);
-		console.error(
-			'  1. Use Agentuity Cloud AI Gateway by ensuring AGENTUITY_SDK_KEY is configured'
-		);
-		console.error(`  2. Set ${envKey} using "agentuity env set ${envKey}" and redeploy`);
+		return;
 	}
+
+	console.error(`[ERROR] The environment variable ${envKey} is required. Either:`);
+	console.error('  1. Use Agentuity Cloud AI Gateway by ensuring AGENTUITY_SDK_KEY is configured');
+	console.error(`  2. Set ${envKey} using "agentuity env set ${envKey}" and redeploy`);
 }
 
 /**

package/src/globals.d.ts

@@ -0,0 +1,28 @@
+/**
+ * Type-safe globalThis declarations for Agentuity runtime.
+ *
+ * String-keyed globals that persist across bun --hot reloads.
+ * Symbol-keyed globals are accessed via the typed helpers in _globals.ts.
+ */
+
+declare global {
+	// eslint-disable-next-line no-var
+	var __AGENTUITY_SERVER__:
+		| {
+				stop: (closeActiveConnections?: boolean) => void;
+				port: number | undefined;
+				pendingRequests: number;
+				pendingWebSockets: number;
+		  }
+		| undefined;
+
+	// eslint-disable-next-line no-var
+	var __AGENTUITY_BUN_SUBPROCESS__:
+		| {
+				kill: (signal?: number | NodeJS.Signals) => void;
+				exitCode: number | null;
+		  }
+		| undefined;
+}
+
+export {};

package/src/handlers/_route-meta.ts

@@ -0,0 +1,31 @@
+/**
+ * Route metadata tagging for build-time route type discovery.
+ *
+ * Handler wrappers (websocket, sse, stream, cron) stamp this symbol on
+ * the returned middleware so the build tool can detect the route type
+ * from `router.routes` without AST parsing.
+ */
+
+export const ROUTE_META = Symbol.for('agentuity:route-meta');
+
+export interface RouteMeta {
+	type: 'websocket' | 'sse' | 'stream' | 'cron';
+}
+
+/**
+ * Tag a handler/middleware with route metadata.
+ */
+export function tagRoute<T extends (...args: any[]) => any>(handler: T, meta: RouteMeta): T {
+	(handler as any)[ROUTE_META] = meta;
+	return handler;
+}
+
+/**
+ * Read route metadata from a handler/middleware.
+ */
+export function getRouteMeta(handler: unknown): RouteMeta | undefined {
+	if (typeof handler === 'function') {
+		return (handler as any)[ROUTE_META];
+	}
+	return undefined;
+}

package/src/handlers/cron.ts

@@ -2,6 +2,7 @@ import type { Context, Handler } from 'hono';
 import { returnResponse } from '../_util';
 import type { Env } from '../app';
 import { verifySignature } from '../signature';
+import { tagRoute } from './_route-meta';
 
 /**
  * Handler function for cron jobs.
@@ -92,7 +93,7 @@ export function cron<E extends Env = Env>(
 		handler = maybeHandler!;
 	}
 
-	return async (c: Context<E>) => {
+	const cronHandler: Handler<E> = async (c: Context<E>) => {
 		if (c.req.method !== 'POST') {
 			throw new Error(
 				`Cron endpoint must use POST method, but received ${c.req.method}. ` +
@@ -127,6 +128,8 @@ export function cron<E extends Env = Env>(
 
 		return returnResponse(c, result);
 	};
+
+	return tagRoute(cronHandler, { type: 'cron' });
 }
 
 /**

package/src/handlers/sse.ts

@@ -3,8 +3,8 @@ import { stream as honoStream } from 'hono/streaming';
 import { context as otelContext, ROOT_CONTEXT } from '@opentelemetry/api';
 import { StructuredError } from '@agentuity/core';
 import type { Schema } from '@agentuity/schema';
-import { getAgentAsyncLocalStorage } from '../_context';
 import type { Env } from '../app';
+import { tagRoute } from './_route-meta';
 
 /**
  * Error thrown when sse() is called without a handler function.
@@ -83,7 +83,7 @@ export interface SSEOptions<TOutput = unknown> {
 	 *
 	 * This schema is used for:
 	 * - Type inference in generated `routes.ts` registry
-	 * - Automatic typing of `useEventStream` hook's `data` property
+	 * - Automatic typing of `EventSource/EventStreamManager` hook's `data` property
 	 *
 	 * The schema is NOT used for runtime validation - SSE messages are sent
 	 * as-is through the stream. Use this for TypeScript type safety only.
@@ -187,8 +187,8 @@ function formatSSEMessage(message: SSEMessage): string {
  *   stream.close();
  * }));
  *
- * // On the frontend, useEventStream will now have typed data:
- * // const { data } = useEventStream('/api/stream');
+ * // On the frontend, EventSource/EventStreamManager will now have typed data:
+ * // const { data } = EventSource/EventStreamManager('/api/stream');
  * // data.type is 'token' | 'complete' | 'error'
  * ```
  *
@@ -225,10 +225,7 @@ export function sse<E extends Env = Env, TOutput = unknown>(
 
 	// Note: options.output is captured for type inference but not used at runtime
 	// The CLI extracts this during build to generate typed route registries
-	return (c: Context<E>) => {
-		const asyncLocalStorage = getAgentAsyncLocalStorage();
-		const capturedContext = asyncLocalStorage.getStore();
-
+	const sseHandler: Handler<E> = (c: Context<E>) => {
 		// Track stream completion for deferred session/thread saving
 		// This promise resolves when the stream closes (normally or via abort)
 		let resolveDone: (() => void) | undefined;
@@ -345,10 +342,6 @@ export function sse<E extends Env = Env, TOutput = unknown>(
 				}
 			};
 
-			// Run handler with AsyncLocalStorage context propagation.
-			// honoStream already uses a fire-and-forget pattern internally,
-			// so we can safely await here - the response is already being sent.
-			//
 			// IMPORTANT: We run in ROOT_CONTEXT (no active OTEL span) to avoid a Bun bug
 			// where OTEL-instrumented fetch conflicts with streaming responses.
 			// This causes "ReadableStream has already been used" errors when AI SDK's
@@ -357,13 +350,9 @@ export function sse<E extends Env = Env, TOutput = unknown>(
 			// our OTEL fetch wrapper use the original unpatched fetch.
 			// See: https://github.com/agentuity/sdk/issues/471
 			// See: https://github.com/oven-sh/bun/issues/24766
-			await otelContext.with(ROOT_CONTEXT, async () => {
-				if (capturedContext) {
-					await asyncLocalStorage.run(capturedContext, runInContext);
-				} else {
-					await runInContext();
-				}
-			});
+			await otelContext.with(ROOT_CONTEXT, runInContext);
 		});
 	};
+
+	return tagRoute(sseHandler, { type: 'sse' });
 }

package/src/handlers/stream.ts

@@ -1,9 +1,9 @@
 import type { Context, Handler } from 'hono';
 import { stream as honoStream } from 'hono/streaming';
 import { context as otelContext, ROOT_CONTEXT } from '@opentelemetry/api';
-import { getAgentAsyncLocalStorage } from '../_context';
 import type { Env } from '../app';
 import { STREAM_DONE_PROMISE_KEY, IS_STREAMING_RESPONSE_KEY } from './sse';
+import { tagRoute } from './_route-meta';
 
 /**
  * Handler function for streaming responses.
@@ -57,10 +57,7 @@ export type StreamHandler<E extends Env = Env> = (
  * @returns Hono handler for streaming response
  */
 export function stream<E extends Env = Env>(handler: StreamHandler<E>): Handler<E> {
-	return (c: Context<E>) => {
-		const asyncLocalStorage = getAgentAsyncLocalStorage();
-		const capturedContext = asyncLocalStorage.getStore();
-
+	const streamHandler: Handler<E> = (c: Context<E>) => {
 		// Track stream completion for deferred session/thread saving
 		// This promise resolves when the stream completes (pipe finishes or errors)
 		let resolveDone: (() => void) | undefined;
@@ -116,13 +113,9 @@ export function stream<E extends Env = Env>(handler: StreamHandler<E>): Handler<
 			// our OTEL fetch wrapper use the original unpatched fetch.
 			// See: https://github.com/agentuity/sdk/issues/471
 			// See: https://github.com/oven-sh/bun/issues/24766
-			await otelContext.with(ROOT_CONTEXT, async () => {
-				if (capturedContext) {
-					await asyncLocalStorage.run(capturedContext, runInContext);
-				} else {
-					await runInContext();
-				}
-			});
+			await otelContext.with(ROOT_CONTEXT, runInContext);
 		});
 	};
+
+	return tagRoute(streamHandler, { type: 'stream' });
 }

package/src/handlers/websocket.ts

@@ -1,8 +1,8 @@
 import type { Context, MiddlewareHandler } from 'hono';
 import { upgradeWebSocket } from 'hono/bun';
 import { context as otelContext, ROOT_CONTEXT } from '@opentelemetry/api';
-import { getAgentAsyncLocalStorage } from '../_context';
 import type { Env } from '../app';
+import { tagRoute } from './_route-meta';
 
 /**
  * Context key for WebSocket close promise.
@@ -92,16 +92,15 @@ export type WebSocketHandler<E extends Env = Env> = (
  * @param handler - Synchronous handler function receiving context and WebSocket connection
  * @returns Hono middleware handler for WebSocket upgrade
  */
-export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): MiddlewareHandler<E> {
+export function websocket<E extends Env = Env>(
+	handler: WebSocketHandler<E>
+): MiddlewareHandler<E, string, { outputFormat: 'ws' }> {
 	const wsHandler = upgradeWebSocket((c: Context<E>) => {
 		let openHandler: ((event: Event) => void | Promise<void>) | undefined;
 		let messageHandler: ((event: MessageEvent) => void | Promise<void>) | undefined;
 		let closeHandler: ((event: CloseEvent) => void | Promise<void>) | undefined;
 		let initialized = false;
 
-		const asyncLocalStorage = getAgentAsyncLocalStorage();
-		const capturedContext = asyncLocalStorage.getStore();
-
 		// Create done promise for session lifecycle deferral, but ONLY for actual
 		// WebSocket upgrade requests. The factory runs unconditionally for every
 		// request hitting this route (Hono calls createEvents before attempting
@@ -144,11 +143,7 @@ export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): Mi
 		// See: https://github.com/oven-sh/bun/issues/24766
 		const runHandler = () => {
 			otelContext.with(ROOT_CONTEXT, () => {
-				if (capturedContext) {
-					asyncLocalStorage.run(capturedContext, () => handler(c, wsConnection));
-				} else {
-					handler(c, wsConnection);
-				}
+				handler(c, wsConnection);
 			});
 			initialized = true;
 		};
@@ -162,14 +157,7 @@ export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): Mi
 					wsConnection.send = (data) => ws.send(data);
 
 					if (openHandler) {
-						const h = openHandler;
-						await otelContext.with(ROOT_CONTEXT, async () => {
-							if (capturedContext) {
-								await asyncLocalStorage.run(capturedContext, () => h(event));
-							} else {
-								await h(event);
-							}
-						});
+						await otelContext.with(ROOT_CONTEXT, () => openHandler!(event));
 					}
 				} catch (err) {
 					c.var.logger?.error('WebSocket onOpen error:', err);
@@ -184,14 +172,7 @@ export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): Mi
 						runHandler();
 					}
 					if (messageHandler) {
-						const h = messageHandler;
-						await otelContext.with(ROOT_CONTEXT, async () => {
-							if (capturedContext) {
-								await asyncLocalStorage.run(capturedContext, () => h(event));
-							} else {
-								await h(event);
-							}
-						});
+						await otelContext.with(ROOT_CONTEXT, () => messageHandler!(event));
 					}
 				} catch (err) {
 					c.var.logger?.error('WebSocket onMessage error:', err);
@@ -202,14 +183,7 @@ export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): Mi
 			onClose: async (event: CloseEvent, _ws: any) => {
 				try {
 					if (closeHandler) {
-						const h = closeHandler;
-						await otelContext.with(ROOT_CONTEXT, async () => {
-							if (capturedContext) {
-								await asyncLocalStorage.run(capturedContext, () => h(event));
-							} else {
-								await h(event);
-							}
-						});
+						await otelContext.with(ROOT_CONTEXT, () => closeHandler!(event));
 					}
 				} catch (err) {
 					c.var.logger?.error('WebSocket onClose error:', err);
@@ -222,8 +196,8 @@ export function websocket<E extends Env = Env>(handler: WebSocketHandler<E>): Mi
 		};
 	});
 
-	const middleware: MiddlewareHandler<E> = (c, next) =>
-		(wsHandler as unknown as MiddlewareHandler<E>)(c, next);
+	const middleware: MiddlewareHandler<E, string, { outputFormat: 'ws' }> = (c, next) =>
+		(wsHandler as unknown as MiddlewareHandler<E, string, { outputFormat: 'ws' }>)(c, next);
 
-	return middleware;
+	return tagRoute(middleware, { type: 'websocket' });
 }