@m6d/cortex-server 1.2.0 → 1.4.0

package/src/ai/tools/execute-code.tool.ts

@@ -2,13 +2,7 @@ import { tool } from "ai";
 import z from "zod";
 import type { ResolvedCortexAgentConfig } from "../../config.ts";
 import { fetchBackend } from "../fetch.ts";
-
-type ApiHelper = {
-    get: (path: string, queryParams?: Record<string, unknown>) => Promise<unknown>;
-    post: (path: string, body?: unknown) => Promise<unknown>;
-    put: (path: string, body?: unknown) => Promise<unknown>;
-    del: (path: string) => Promise<unknown>;
-};
+import { getQuickJS, Scope, shouldInterruptAfterDeadline } from "quickjs-emscripten";
 
 function buildQueryString(params: Record<string, unknown>) {
     const searchParams = new URLSearchParams();
@@ -26,7 +20,8 @@ function buildQueryString(params: Record<string, unknown>) {
     return qs ? `?${qs}` : "";
 }
 
-function createApiHelper(
+async function runInSandbox(
+    code: string,
     backendFetch: NonNullable<ResolvedCortexAgentConfig["backendFetch"]>,
     token: string,
 ) {
@@ -53,22 +48,82 @@ function createApiHelper(
         return response.json();
     }
 
-    return {
-        get: (path, queryParams) => {
-            const fullPath = queryParams ? path + buildQueryString(queryParams) : path;
-            return request("GET", fullPath);
-        },
-        post: (path, body) => request("POST", path, body),
-        put: (path, body) => request("PUT", path, body),
-        del: (path) => request("DELETE", path),
-    } satisfies ApiHelper;
-}
+    return await Scope.withScopeAsync(async (scope) => {
+        const QuickJS = await getQuickJS();
+
+        const runtime = scope.manage(QuickJS.newRuntime());
+        runtime.setMemoryLimit(1024 * 1024 * 20); // 20MB
+        runtime.setInterruptHandler(shouldInterruptAfterDeadline(Date.now() + 10_000)); // 60 seconds;
+
+        const vm = scope.manage(runtime.newContext());
+
+        function drainJobs() {
+            vm.runtime.executePendingJobs();
+        }
+
+        function makeApiFunction(fn: (...args: unknown[]) => Promise<unknown>) {
+            return scope.manage(
+                vm.newFunction(fn.name, (...argHandles) => {
+                    const args = argHandles.map(vm.dump.bind(vm));
+                    const promise = scope.manage(
+                        vm.newPromise(
+                            fn(...args).then((result) =>
+                                scope.manage(
+                                    vm.unwrapResult(vm.evalCode(`(${JSON.stringify(result)})`)),
+                                ),
+                            ),
+                        ),
+                    );
 
-const AsyncFunction: new (
-    ...args: [...paramNames: string[], body: string]
-) => (...args: unknown[]) => Promise<unknown> = Object.getPrototypeOf(
-    async function () {},
-).constructor;
+                    promise.settled.finally(() => {
+                        drainJobs();
+                    });
+
+                    return promise.handle;
+                }),
+            );
+        }
+
+        const apiHandle = scope.manage(vm.newObject());
+        vm.setProp(
+            apiHandle,
+            "get",
+            makeApiFunction(async (path, queryParams?) => {
+                const fullPath = queryParams
+                    ? path + buildQueryString(queryParams as Record<string, unknown>)
+                    : path;
+                return await request("GET", fullPath as string);
+            }),
+        );
+        vm.setProp(
+            apiHandle,
+            "post",
+            makeApiFunction(async (path, body) => await request("POST", path as string, body)),
+        );
+        vm.setProp(
+            apiHandle,
+            "put",
+            makeApiFunction(async (path, body) => await request("PUT", path as string, body)),
+        );
+        vm.setProp(
+            apiHandle,
+            "del",
+            makeApiFunction(async (path) => await request("DELETE", path as string)),
+        );
+
+        vm.setProp(vm.global, "api", apiHandle);
+
+        const runResult = vm.evalCode(`(async () => {
+                ${code}
+            })()`);
+        const resultHandle = scope.manage(vm.unwrapResult(runResult));
+        const pending = vm.resolvePromise(resultHandle);
+        drainJobs();
+        const resolvedResult = await pending;
+        const result = scope.manage(vm.unwrapResult(resolvedResult));
+        return vm.dump(result);
+    });
+}
 
 export function createExecuteCodeTool(
     backendFetch: NonNullable<ResolvedCortexAgentConfig["backendFetch"]>,
@@ -86,11 +141,8 @@ export function createExecuteCodeTool(
                 ),
         }),
         execute: async ({ code }) => {
-            const apiHelper = createApiHelper(backendFetch, token);
-
             try {
-                const fn = new AsyncFunction("api", code);
-                const result = await fn(apiHelper);
+                const result = await runInSandbox(code, backendFetch, token);
                 return JSON.stringify(result);
             } catch (e) {
                 const message = e instanceof Error ? e.message : String(e);

package/src/ai/tools/query-graph.tool.ts

@@ -24,7 +24,7 @@ ORDER BY score DESC;
                 .string()
                 .optional()
                 .describe(
-                    'Optional JSON-encoded string of query parameters. Example: `{"name": "LeaveBalance"}` if you know the exact name; for parameters that need to be embedded first prepend the name with `#`, e.g., `{"#paramName": "Text to be embedded before passed to query"}`',
+                    'Optional JSON-encoded string of query parameters. Example: `{"name": "LeaveBalance"}` if you know the exact name; for parameters that need to be embedded first prepend the name with `#`, e.g., `{"#embedding": "Text to be embedded before passed to query"}`',
                 ),
         }),
         execute: async ({ query, parameters }) => {

package/src/cli/extract-endpoints.ts

@@ -39,7 +39,7 @@ const AUTO_START = "// @auto-generated-start";
 const AUTO_END = "// @auto-generated-end";
 const MAX_DEPTH = 8;
 
-const toCamelCase = (s: string): string =>
+const toCamelCase = (s: string) =>
     s
         .split(".")
         .map((seg) => seg.replace(/^[A-Z]/, (c) => c.toLowerCase()))
@@ -54,13 +54,13 @@ const dedupe = (items: Prop[]): Prop[] => {
     return items.filter((x) => (seen.has(x.name) ? false : (seen.add(x.name), true)));
 };
 
-function normalizePath(raw: string): string {
+function normalizePath(raw: string) {
     let p = raw.startsWith("/") ? raw : `/${raw}`;
     p = p.replace(/\{([^}:?]+)(?::[^}]+)?\??\}/g, "{$1}").replace(/^\/api(?=\/|$)/i, "") || "/";
     return p.length > 1 ? p.replace(/\/+$/, "") : p;
 }
 
-function walkFiles(dir: string, accept: (filePath: string) => boolean): string[] {
+function walkFiles(dir: string, accept: (filePath: string) => boolean) {
     if (!fs.existsSync(dir)) return [];
     const out: string[] = [];
     function walk(d: string) {
@@ -74,7 +74,7 @@ function walkFiles(dir: string, accept: (filePath: string) => boolean): string[]
     return out;
 }
 
-function endpointFiles(root: string): Map<string, string> {
+function endpointFiles(root: string) {
     const out = new Map<string, string>();
     for (const filePath of walkFiles(root, (p) => p.endsWith(".endpoint.ts"))) {
         const c = fs.readFileSync(filePath, "utf8");
@@ -85,7 +85,7 @@ function endpointFiles(root: string): Map<string, string> {
     return out;
 }
 
-function resolvePointer(doc: unknown, ref: string): unknown {
+function resolvePointer(doc: unknown, ref: string) {
     if (!ref.startsWith("#/")) return null;
     let cur: unknown = doc;
     for (const part of ref.slice(2).split("/")) {
@@ -96,7 +96,7 @@ function resolvePointer(doc: unknown, ref: string): unknown {
     return cur;
 }
 
-function deref(doc: unknown, v: unknown): Obj {
+function deref(doc: unknown, v: unknown) {
     if (!isObj(v)) return {};
     let x = v;
     const seen = new Set<string>();
@@ -108,7 +108,7 @@ function deref(doc: unknown, v: unknown): Obj {
     return x;
 }
 
-function schemaType(s: Obj): string {
+function schemaType(s: Obj) {
     const t = Array.isArray(s.type)
         ? s.type.find((x) => typeof x === "string" && x !== "null")
         : s.type;
@@ -133,7 +133,7 @@ function isJsonNodeSchema(s: Obj) {
     );
 }
 
-function scalarType(s: Obj): string {
+function scalarType(s: Obj) {
     if (Array.isArray(s.enum) && s.enum.every((v) => typeof v === "string")) {
         return s.enum.map((v) => `'${String(v).replace(/'/g, "\\'")}'`).join(" | ");
     }
@@ -149,7 +149,7 @@ function scalarType(s: Obj): string {
     return "unknown";
 }
 
-function toProps(doc: unknown, schemaIn: unknown, depth = 0): Prop[] {
+function toProps(doc: unknown, schemaIn: unknown, depth = 0) {
     if (depth > MAX_DEPTH) return [];
     const s = deref(doc, schemaIn);
     if (isJsonNodeRef(schemaIn) || isJsonNodeSchema(s)) return [];
@@ -198,7 +198,7 @@ function toProp(doc: unknown, name: string, schemaIn: unknown, required: boolean
     return { name, required: req, type: scalarType(s) };
 }
 
-function pickContent(contentIn: unknown): { schema: unknown; mimes: string[] } | null {
+function pickContent(contentIn: unknown) {
     const content = obj(contentIn);
     const mimes = Object.keys(content);
     if (!mimes.length) return null;
@@ -319,7 +319,7 @@ function extractEndpoint(
     method: HttpMethod,
     pathItemIn: unknown,
     operationIn: unknown,
-): Endpoint {
+) {
     const pathItem = obj(pathItemIn);
     const operation = obj(operationIn);
 
@@ -393,10 +393,10 @@ function extractEndpoint(
         responseKind: parsed.responseKind,
         successStatus: parsed.successStatus,
         errorStatuses: parsed.errorStatuses,
-    };
+    } satisfies Endpoint;
 }
 
-function parseSwaggerEndpoints(swagger: unknown): Map<string, Endpoint> {
+function parseSwaggerEndpoints(swagger: unknown) {
     const out = new Map<string, Endpoint>();
     const doc = obj(swagger);
     for (const [route, pathItem] of Object.entries(obj(doc.paths))) {
@@ -410,7 +410,7 @@ function parseSwaggerEndpoints(swagger: unknown): Map<string, Endpoint> {
     return out;
 }
 
-function serializeProps(props: Prop[], depth: number): string {
+function serializeProps(props: Prop[], depth: number) {
     return props
         .map((p) => {
             const indent = "    ".repeat(depth);
@@ -431,7 +431,7 @@ function serializeProps(props: Prop[], depth: number): string {
         .join("\n");
 }
 
-function blockFor(endpoint: Endpoint): string {
+function blockFor(endpoint: Endpoint) {
     return `    autoGenerated: {
         params: [
 ${serializeProps(endpoint.params, 2)}
@@ -461,7 +461,7 @@ function formatGeneratedFiles(cwd: string, filePaths: string[]) {
     }
 }
 
-async function fetchJson(url: string): Promise<unknown> {
+async function fetchJson(url: string) {
     const response = await fetch(url, {
         headers: { Accept: "application/json" },
     });
@@ -471,7 +471,7 @@ async function fetchJson(url: string): Promise<unknown> {
     return response.json();
 }
 
-function resolveEndpoint(key: string, extracted: Map<string, Endpoint>): Endpoint | undefined {
+function resolveEndpoint(key: string, extracted: Map<string, Endpoint>) {
     const [method, route] = key.split(":");
     if (!method || !route) return undefined;
     return (
@@ -483,7 +483,7 @@ function resolveEndpoint(key: string, extracted: Map<string, Endpoint>): Endpoin
     );
 }
 
-export async function extractEndpoints(options: ExtractEndpointsOptions): Promise<void> {
+export async function extractEndpoints(options: ExtractEndpointsOptions) {
     const { swaggerUrl, domainsDir, write = false } = options;
 
     console.log(`Fetching Swagger from ${swaggerUrl}`);

package/src/config.ts

@@ -3,6 +3,7 @@ import type { DatabaseAdapter } from "./adapters/database";
 import type { StorageAdapter } from "./adapters/storage";
 import type { DomainDef } from "./graph/types.ts";
 import type { RequestInterceptorOptions } from "./ai/interceptors/request-interceptor.ts";
+import type { ContextConfig } from "./ai/context/types.ts";
 
 export type KnowledgeConfig = {
     swagger?: { url: string };
@@ -64,6 +65,7 @@ export type CortexAgentDefinition = {
         url: string;
         apiKey: string;
     };
+    context?: Partial<ContextConfig>;
     knowledge?: KnowledgeConfig | null;
 };
 
@@ -110,6 +112,7 @@ export type ResolvedCortexAgentConfig = {
         args: Record<string, unknown>;
     }) => void;
     onStreamFinish?: (result: { messages: UIMessage[]; isAborted: boolean }) => void;
+    context: ContextConfig;
     knowledge?: KnowledgeConfig;
 };
 
@@ -144,6 +147,7 @@ export type CortexConfig = {
         url: string;
         apiKey: string;
     };
+    context?: Partial<ContextConfig>;
     knowledge?: KnowledgeConfig;
     agents: Record<string, CortexAgentDefinition>;
 };

package/src/db/migrations/20260315000000_add_context_meta/migration.sql

@@ -0,0 +1 @@
+ALTER TABLE [ai].[threads] ADD [context_meta] nvarchar(max);

package/src/db/schema.ts

@@ -1,6 +1,8 @@
 import type { ToolUIPart, UIMessage } from "ai";
 import { sql } from "drizzle-orm";
 import { customType, datetime2, index, mssqlSchema, nvarchar } from "drizzle-orm/mssql-core";
+import type { MessageMetadata } from "src/types";
+import type { ThreadContextMeta } from "../ai/context/types.ts";
 
 const uniqueIdentifier = customType<{ data: string }>({
     dataType() {
@@ -28,6 +30,7 @@ export const threads = aiSchema.table("threads", {
     agentId: nvarchar({ length: 128 }).notNull().default("default"),
     title: nvarchar({ length: 256 }),
     session: nvarchar({ mode: "json", length: "max" }).$type<Record<string, unknown>>(),
+    contextMeta: nvarchar({ mode: "json", length: "max" }).$type<ThreadContextMeta>(),
     ...auditColumns,
 });
 
@@ -37,7 +40,9 @@ export const messages = aiSchema.table("messages", {
         .references(() => threads.id, { onDelete: "cascade" })
         .notNull(),
     text: nvarchar({ length: "max" }),
-    content: nvarchar({ mode: "json", length: "max" }).notNull().$type<UIMessage>(),
+    content: nvarchar({ mode: "json", length: "max" })
+        .notNull()
+        .$type<UIMessage<MessageMetadata>>(),
     role: nvarchar({
         enum: ["system", "user", "assistant", "tool"],
         length: 16,

package/src/factory.ts

@@ -2,6 +2,8 @@ import { Hono } from "hono";
 import { websocket } from "hono/bun";
 import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
 import type { CortexConfig, ResolvedCortexAgentConfig } from "./config.ts";
+import { DEFAULT_CONTEXT_CONFIG } from "./ai/context/types.ts";
+import type { ContextConfig } from "./ai/context/types.ts";
 import type { AppEnv, CortexAppEnv } from "./types.ts";
 import type { DomainDef } from "./graph/types.ts";
 import { createUserLoaderMiddleware } from "./auth/middleware.ts";
@@ -61,7 +63,7 @@ export function createCortex(config: CortexConfig) {
             const userLoader = createUserLoaderMiddleware(config.auth);
             app.use("*", userLoader);
 
-            // WebSocket route (global, not agent-scoped)
+            // Compatibility WebSocket route
             app.route("/", createWsRoute());
 
             // Agent-scoped routes
@@ -72,6 +74,12 @@ export function createCortex(config: CortexConfig) {
                 const agentDef = config.agents[agentId];
                 if (!agentDef) return c.json({ error: "Agent not found" }, 404);
 
+                const resolvedContext: ContextConfig = {
+                    ...DEFAULT_CONTEXT_CONFIG,
+                    ...config.context,
+                    ...agentDef.context,
+                };
+
                 const resolvedConfig: ResolvedCortexAgentConfig = {
                     db,
                     storage,
@@ -89,6 +97,7 @@ export function createCortex(config: CortexConfig) {
                     loadSessionData: agentDef.loadSessionData,
                     onToolCall: agentDef.onToolCall,
                     onStreamFinish: agentDef.onStreamFinish,
+                    context: resolvedContext,
                 };
 
                 c.set("agentConfig", resolvedConfig);
@@ -99,6 +108,7 @@ export function createCortex(config: CortexConfig) {
             agentApp.route("/", createThreadRoutes());
             agentApp.route("/", createChatRoutes());
             agentApp.route("/", createFileRoutes());
+            agentApp.route("/", createWsRoute({ useAgentParam: true }));
 
             app.route("/agents/:agentId", agentApp);
 

package/src/index.ts

@@ -7,6 +7,8 @@ export type {
     StorageConfig,
 } from "./config";
 
+export type { ContextConfig, ThreadContextMeta } from "./ai/context/types";
+
 // Types consumers may need
 export type { Thread, AppEnv } from "./types";
 

package/src/routes/chat.ts

@@ -5,6 +5,9 @@ import { HTTPException } from "hono/http-exception";
 import type { CortexAppEnv } from "../types.ts";
 import { requireAuth } from "../auth/middleware.ts";
 import { stream } from "../ai/index.ts";
+import { subscribe, abortStream } from "../ai/active-streams.ts";
+import { notify } from "../ws/index.ts";
+import { toThreadSummary } from "../types.ts";
 
 export function createChatRoutes() {
     const app = new Hono<CortexAppEnv>();
@@ -21,12 +24,13 @@ export function createChatRoutes() {
         ),
         async function (c) {
             const config = c.get("agentConfig");
+            const agentId = c.get("agentId");
             const { id: userId, token } = c.get("user");
             const { messages, id: threadId } = c.req.valid("json");
 
             const thread = await config.db.threads.getById(userId, threadId);
 
-            if (!thread) {
+            if (!thread || thread.agentId !== agentId) {
                 throw new HTTPException(404, { message: "Not found" });
             }
 
@@ -34,5 +38,46 @@ export function createChatRoutes() {
         },
     );
 
+    app.get("/chat/:chatId/stream", requireAuth, async function (c) {
+        const config = c.get("agentConfig");
+        const agentId = c.get("agentId");
+        const chatId = c.req.param("chatId");
+        const thread = await config.db.threads.getById(c.get("user").id, chatId);
+        if (!thread || thread.agentId !== agentId) {
+            throw new HTTPException(404, { message: "Not found" });
+        }
+
+        const subscriberStream = subscribe(chatId);
+        if (!subscriberStream) {
+            return c.body(null, 204);
+        }
+
+        return new Response(subscriberStream.pipeThrough(new TextEncoderStream()), {
+            headers: {
+                "Content-Type": "text/event-stream",
+                "Cache-Control": "no-cache",
+                Connection: "keep-alive",
+            },
+        });
+    });
+
+    app.post("/chat/:chatId/abort", requireAuth, async function (c) {
+        const config = c.get("agentConfig");
+        const agentId = c.get("agentId");
+        const userId = c.get("user").id;
+        const chatId = c.req.param("chatId");
+        const thread = await config.db.threads.getById(userId, chatId);
+        if (!thread || thread.agentId !== agentId) {
+            throw new HTTPException(404, { message: "Not found" });
+        }
+
+        abortStream(chatId);
+        notify(userId, agentId, {
+            type: "thread:run-finished",
+            payload: { thread: toThreadSummary(thread, false) },
+        });
+        return c.body(null, 204);
+    });
+
     return app;
 }

package/src/routes/threads.ts

@@ -4,6 +4,9 @@ import z from "zod";
 import type { CortexAppEnv } from "../types.ts";
 import { requireAuth } from "../auth/middleware.ts";
 import { generateTitle } from "../ai/index.ts";
+import { abortStream, isStreamRunning } from "../ai/active-streams.ts";
+import { notify } from "../ws/index.ts";
+import { toThreadSummary } from "../types.ts";
 
 export function createThreadRoutes() {
     const app = new Hono<CortexAppEnv>();
@@ -13,7 +16,7 @@ export function createThreadRoutes() {
         const agentId = c.get("agentId");
         const threads = await config.db.threads
             .list(c.get("user").id, agentId)
-            .then((x) => x.map((y) => ({ id: y.id, title: y.title, createdAt: y.createdAt })));
+            .then((x) => x.map((y) => toThreadSummary(y, isStreamRunning(y.id))));
         return c.json(threads);
     });
 
@@ -24,28 +27,55 @@ export function createThreadRoutes() {
         async function (c) {
             const config = c.get("agentConfig");
             const agentId = c.get("agentId");
+            const userId = c.get("user").id;
             const { prompt } = c.req.valid("json");
-            const { id, createdAt, title } = await config.db.threads.create(
-                c.get("user").id,
-                agentId,
-            );
-            generateTitle(id, prompt, c.get("user").id, config);
-            return c.json({ id, title, createdAt });
+            const thread = await config.db.threads.create(userId, agentId);
+
+            notify(userId, agentId, {
+                type: "thread:created",
+                payload: { thread: toThreadSummary(thread, false) },
+            });
+
+            generateTitle(thread.id, prompt, userId, config);
+
+            return c.json(toThreadSummary(thread, false));
         },
     );
 
     app.delete("/threads/:threadId", requireAuth, async function (c) {
         const config = c.get("agentConfig");
+        const agentId = c.get("agentId");
+        const userId = c.get("user").id;
         const threadId = c.req.param("threadId");
-        await config.db.threads.delete(c.get("user").id, threadId);
+
+        const thread = await config.db.threads.getById(userId, threadId);
+        if (!thread || thread.agentId !== agentId) {
+            return c.body(null, 404);
+        }
+
+        abortStream(threadId);
+        await config.db.threads.delete(userId, threadId);
+
+        notify(userId, agentId, {
+            type: "thread:deleted",
+            payload: { threadId },
+        });
+
         return c.body(null, 204);
     });
 
     app.get("/threads/:threadId/messages", requireAuth, async function (c) {
         const config = c.get("agentConfig");
+        const agentId = c.get("agentId");
+        const userId = c.get("user").id;
         const threadId = c.req.param("threadId");
+        const thread = await config.db.threads.getById(userId, threadId);
+        if (!thread || thread.agentId !== agentId) {
+            return c.body(null, 404);
+        }
+
         const messages = await config.db.messages
-            .list(c.get("user").id, threadId)
+            .list(userId, threadId)
             .then((x) => x.map((y) => y.content));
         return c.json(messages);
     });
@@ -56,8 +86,15 @@ export function createThreadRoutes() {
         zValidator("json", z.object({ session: z.record(z.unknown()) })),
         async function (c) {
             const config = c.get("agentConfig");
+            const agentId = c.get("agentId");
+            const userId = c.get("user").id;
             const threadId = c.req.param("threadId");
             const { session } = c.req.valid("json");
+            const thread = await config.db.threads.getById(userId, threadId);
+            if (!thread || thread.agentId !== agentId) {
+                return c.body(null, 404);
+            }
+
             await config.db.threads.updateSession(threadId, session);
             return c.body(null, 204);
         },

package/src/routes/ws.ts

@@ -3,31 +3,45 @@ import { upgradeWebSocket } from "hono/bun";
 import type { AppEnv } from "../types.ts";
 import { addConnection, removeConnection } from "../ws/index.ts";
 
-export function createWsRoute() {
+type WsRouteOptions = {
+    useAgentParam?: boolean;
+};
+
+export function createWsRoute(options?: WsRouteOptions) {
     const app = new Hono<AppEnv>();
+    const handleUpgrade = upgradeWebSocket(function (c) {
+        const userId = c.get("user")?.id;
+        const agentId =
+            options?.useAgentParam === true ? c.req.param("agentId") : c.req.query("agentId");
+
+        return {
+            onOpen(_, ws) {
+                if (userId && agentId) {
+                    addConnection(userId, agentId, ws);
+                }
+            },
+
+            onClose(_, ws) {
+                if (userId && agentId) {
+                    removeConnection(userId, agentId, ws);
+                }
+            },
+        };
+    });
+
+    app.get("/ws", function (c, next) {
+        if (!c.get("user")) {
+            return c.json({ error: "Unauthorized" }, 401);
+        }
+
+        const agentId =
+            options?.useAgentParam === true ? c.req.param("agentId") : c.req.query("agentId");
+        if (!agentId) {
+            return c.json({ error: "agentId is required" }, 400);
+        }
 
-    app.get(
-        "/ws",
-        upgradeWebSocket(function (c) {
-            const user = c.get("user");
-            const userId = user?.id;
-
-            return {
-                onOpen(_, ws) {
-                    if (userId) {
-                        addConnection(userId, ws);
-                    }
-                    ws.send("Connected!");
-                },
-
-                onClose(_, ws) {
-                    if (userId) {
-                        removeConnection(userId, ws);
-                    }
-                },
-            };
-        }),
-    );
+        return handleUpgrade(c, next) as unknown as Response;
+    });
 
     return app;
 }