@m6d/cortex-server 1.1.2 → 1.3.0

package/dist/src/ai/index.d.ts

@@ -1,4 +1,4 @@
 import type { ResolvedCortexAgentConfig } from "../config.ts";
 import type { Thread } from "../types.ts";
-export declare function stream(messages: unknown[], thread: Thread, userId: string, token: string, config: ResolvedCortexAgentConfig): Promise<Response>;
+export declare function stream(messages: unknown[], thread: Thread, userId: string, token: string, config: ResolvedCortexAgentConfig, abortSignal?: AbortSignal): Promise<Response>;
 export declare function generateTitle(threadId: string, prompt: string, userId: string, config: ResolvedCortexAgentConfig): Promise<void>;

package/dist/src/ai/interceptors/{resolve-captured-files.d.ts → request-interceptor.d.ts}

@@ -4,8 +4,9 @@ export type ResolvedFile = {
     name: string;
     bytes: string;
 };
-export declare function createCapturedFileInterceptor(db: DatabaseAdapter, storage: StorageAdapter, options?: {
+export type RequestInterceptorOptions = {
     transformFile?: (file: ResolvedFile) => unknown;
-}): (body: Record<string, unknown>, context: {
+};
+export declare function createRequestInterceptor(db: DatabaseAdapter, storage: StorageAdapter, options?: RequestInterceptorOptions): (body: Record<string, unknown>, context: {
     token: string;
 }) => Promise<Record<string, unknown>>;

package/dist/src/config.d.ts

@@ -1,7 +1,8 @@
-import type { UIMessage } from "ai";
+import type { ToolSet, UIMessage } from "ai";
 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";
 export type KnowledgeConfig = {
     swagger?: {
         url: string;
@@ -22,7 +23,7 @@ export type StorageConfig = {
 };
 export type CortexAgentDefinition = {
     systemPrompt: string | ((session: Record<string, unknown> | null) => string | Promise<string>);
-    tools?: Record<string, unknown>;
+    tools?: ToolSet;
     backendFetch?: {
         baseUrl: string;
         apiKey: string;
@@ -30,6 +31,7 @@ export type CortexAgentDefinition = {
         transformRequestBody?: (body: Record<string, unknown>, context: {
             token: string;
         }) => Promise<Record<string, unknown>>;
+        interceptor?: RequestInterceptorOptions;
     };
     loadSessionData?: (token: string) => Promise<Record<string, unknown>>;
     onToolCall?: (toolCall: {
@@ -89,7 +91,7 @@ export type ResolvedCortexAgentConfig = {
         apiKey: string;
     };
     systemPrompt: string | ((session: Record<string, unknown> | null) => string | Promise<string>);
-    tools?: Record<string, unknown>;
+    tools?: ToolSet;
     backendFetch?: {
         baseUrl: string;
         apiKey: string;
@@ -97,6 +99,7 @@ export type ResolvedCortexAgentConfig = {
         transformRequestBody?: (body: Record<string, unknown>, context: {
             token: string;
         }) => Promise<Record<string, unknown>>;
+        interceptor?: RequestInterceptorOptions;
     };
     loadSessionData?: (token: string) => Promise<Record<string, unknown>>;
     onToolCall?: (toolCall: {

package/dist/src/index.d.ts

@@ -2,6 +2,8 @@ export type { CortexConfig, CortexAgentDefinition, KnowledgeConfig, DatabaseConf
 export type { Thread, AppEnv } from "./types";
 export type { CortexInstance } from "./factory";
 export { createCortex } from "./factory";
+export type { ResolvedFile, RequestInterceptorOptions, } from "./ai/interceptors/request-interceptor";
+export { createRequestInterceptor } from "./ai/interceptors/request-interceptor";
 export { captureFilesTool } from "./ai/tools/capture-files.tool";
 export { createQueryGraphTool } from "./ai/tools/query-graph.tool";
 export { createCallEndpointTool } from "./ai/tools/call-endpoint.tool";

package/dist/src/ws/events.d.ts

@@ -5,4 +5,10 @@ export type ThreadTitleUpdatedEvent = {
         title: string;
     };
 };
-export type WsEvent = ThreadTitleUpdatedEvent;
+export type ThreadMessagesUpdatedEvent = {
+    type: "thread:messages-updated";
+    payload: {
+        threadId: string;
+    };
+};
+export type WsEvent = ThreadTitleUpdatedEvent | ThreadMessagesUpdatedEvent;

package/dist/src/ws/index.d.ts

@@ -1,3 +1,3 @@
 export { addConnection, removeConnection, getConnections } from "./connections.ts";
 export { notify } from "./notify.ts";
-export type { WsEvent, ThreadTitleUpdatedEvent } from "./events.ts";
+export type { WsEvent, ThreadTitleUpdatedEvent, ThreadMessagesUpdatedEvent } from "./events.ts";

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@m6d/cortex-server",
-    "version": "1.1.2",
+    "version": "1.3.0",
     "description": "Reusable AI agent chat server library for Hono + Bun",
     "license": "MIT",
     "type": "module",
@@ -29,7 +29,8 @@
         "@hono/zod-validator": "^0.5.0",
         "drizzle-orm": "^1.0.0-beta.15-859cf75",
         "minio": "^8.0.7",
-        "mssql": "^12.2.0"
+        "mssql": "^12.2.0",
+        "quickjs-emscripten": "^0.32.0"
     },
     "devDependencies": {
         "@ai-sdk/openai-compatible": "^2.0.0",

package/src/adapters/mssql.ts

@@ -109,19 +109,15 @@ export function createMssqlAdapter(connectionString: string, storage: StorageAda
                 }
 
                 const existingMessages = messagesToInsert.filter((x) => existingIds.includes(x.id));
-                if (existingMessages.length) {
-                    await db.transaction(async (tx) => {
-                        for (const message of existingMessages) {
-                            await tx
-                                .update(messages)
-                                .set({
-                                    content: message,
-                                    text: message.parts.find((x) => x.type === "text")?.text,
-                                })
-                                .where(eq(messages.id, message.id))
-                                .execute();
-                        }
-                    });
+                for (const message of existingMessages) {
+                    await db
+                        .update(messages)
+                        .set({
+                            content: message,
+                            text: message.parts.find((x) => x.type === "text")?.text,
+                        })
+                        .where(eq(messages.id, message.id))
+                        .execute();
                 }
             },
         },

package/src/ai/index.ts

@@ -1,5 +1,7 @@
 import {
+    type UIMessage,
     type ToolSet,
+    consumeStream,
     convertToModelMessages,
     generateId,
     generateText,
@@ -15,7 +17,7 @@ import { createQueryGraphTool } from "./tools/query-graph.tool.ts";
 import { createCallEndpointTool } from "./tools/call-endpoint.tool.ts";
 import { createExecuteCodeTool } from "./tools/execute-code.tool.ts";
 import { captureFilesTool } from "./tools/capture-files.tool.ts";
-import { createCapturedFileInterceptor } from "./interceptors/resolve-captured-files.ts";
+import { createRequestInterceptor } from "./interceptors/request-interceptor.ts";
 import { createNeo4jClient } from "../graph/neo4j.ts";
 import { resolveFromGraph } from "../graph/resolver.ts";
 import { notify } from "../ws/index.ts";
@@ -26,6 +28,7 @@ export async function stream(
     userId: string,
     token: string,
     config: ResolvedCortexAgentConfig,
+    abortSignal?: AbortSignal,
 ) {
     const validationResult = await safeValidateUIMessages({ messages });
     if (!validationResult.success) {
@@ -69,7 +72,11 @@ export async function stream(
             ...config.backendFetch,
             transformRequestBody:
                 config.backendFetch.transformRequestBody ??
-                createCapturedFileInterceptor(config.db, config.storage),
+                createRequestInterceptor(
+                    config.db,
+                    config.storage,
+                    config.backendFetch.interceptor,
+                ),
         };
 
         builtInTools["callEndpoint"] = createCallEndpointTool(backendFetchWithInterceptor, token);
@@ -88,20 +95,30 @@ export async function stream(
         system: systemPrompt,
         tools,
         messages: recentMessages,
-        onAbort: () => {
-            console.log("Stream aborted");
-        },
+        abortSignal,
     });
 
     return result.toUIMessageStreamResponse({
         originalMessages,
         generateMessageId: generateId,
+        consumeSseStream: consumeStream,
         onFinish: async ({ messages: finishedMessages, isAborted }) => {
             if (isAborted) {
-                console.log("Stream was aborted");
+                finalizeAbortedMessages(finishedMessages);
             }
             await config.db.messages.upsert(thread.id, finishedMessages);
             config.onStreamFinish?.({ messages: finishedMessages, isAborted });
+
+            // XXX: we need to notify the user so that the client can
+            // fetch new messages. The client can't fetch messages
+            // immediately after abort because messages may not have been
+            // saved yet.
+            if (isAborted) {
+                notify(userId, {
+                    type: "thread:messages-updated",
+                    payload: { threadId: thread.id },
+                });
+            }
         },
     });
 }
@@ -131,3 +148,31 @@ going to do so or any other speech. Spit out only the title.`,
         payload: { threadId, title: output ?? "" },
     });
 }
+
+const TERMINAL_TOOL_STATES = new Set(["output-available", "output-error", "output-denied"]);
+
+function finalizeAbortedMessages(messages: UIMessage[]) {
+    const lastMessage = messages.at(-1);
+    if (!lastMessage || lastMessage.role !== "assistant") return;
+
+    lastMessage.parts = lastMessage.parts.map((part) => {
+        if ((part.type === "text" || part.type === "reasoning") && part.state === "streaming") {
+            return { ...part, state: "done" as const };
+        }
+
+        if ("toolCallId" in part && "state" in part) {
+            const toolState = part.state;
+            if (!TERMINAL_TOOL_STATES.has(toolState)) {
+                const { approval: _, ...rest } = part;
+                return {
+                    ...rest,
+                    state: "output-error" as const,
+                    errorText: "Generation was aborted",
+                    output: undefined,
+                };
+            }
+        }
+
+        return part;
+    });
+}

package/src/ai/interceptors/{resolve-captured-files.ts → request-interceptor.ts}

@@ -7,14 +7,18 @@ export type ResolvedFile = {
     bytes: string;
 };
 
-export function createCapturedFileInterceptor(
+export type RequestInterceptorOptions = {
+    transformFile?: (file: ResolvedFile) => unknown;
+};
+
+export function createRequestInterceptor(
     db: DatabaseAdapter,
     storage: StorageAdapter,
-    options?: { transformFile?: (file: ResolvedFile) => unknown },
+    options?: RequestInterceptorOptions,
 ) {
     const transformFile = options?.transformFile ?? ((file: ResolvedFile) => file);
 
-    return async function resolveCapturedFiles(
+    return async function resolveRequestBody(
         body: Record<string, unknown>,
         context: { token: string },
     ) {

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/config.ts

@@ -1,7 +1,8 @@
-import type { UIMessage } from "ai";
+import type { Tool, ToolSet, UIMessage } from "ai";
 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";
 
 export type KnowledgeConfig = {
     swagger?: { url: string };
@@ -24,7 +25,7 @@ export type StorageConfig = {
 
 export type CortexAgentDefinition = {
     systemPrompt: string | ((session: Record<string, unknown> | null) => string | Promise<string>);
-    tools?: Record<string, unknown>;
+    tools?: ToolSet;
     backendFetch?: {
         baseUrl: string;
         apiKey: string;
@@ -33,6 +34,7 @@ export type CortexAgentDefinition = {
             body: Record<string, unknown>,
             context: { token: string },
         ) => Promise<Record<string, unknown>>;
+        interceptor?: RequestInterceptorOptions;
     };
     loadSessionData?: (token: string) => Promise<Record<string, unknown>>;
     onToolCall?: (toolCall: {
@@ -90,7 +92,7 @@ export type ResolvedCortexAgentConfig = {
         apiKey: string;
     };
     systemPrompt: string | ((session: Record<string, unknown> | null) => string | Promise<string>);
-    tools?: Record<string, unknown>;
+    tools?: ToolSet;
     backendFetch?: {
         baseUrl: string;
         apiKey: string;
@@ -99,6 +101,7 @@ export type ResolvedCortexAgentConfig = {
             body: Record<string, unknown>,
             context: { token: string },
         ) => Promise<Record<string, unknown>>;
+        interceptor?: RequestInterceptorOptions;
     };
     loadSessionData?: (token: string) => Promise<Record<string, unknown>>;
     onToolCall?: (toolCall: {

package/src/index.ts

@@ -14,6 +14,13 @@ export type { Thread, AppEnv } from "./types";
 export type { CortexInstance } from "./factory";
 export { createCortex } from "./factory";
 
+// Request interceptor (consumers may customize or replace the default)
+export type {
+    ResolvedFile,
+    RequestInterceptorOptions,
+} from "./ai/interceptors/request-interceptor";
+export { createRequestInterceptor } from "./ai/interceptors/request-interceptor";
+
 // Tools (consumers may register custom tools or use built-in ones)
 export { captureFilesTool } from "./ai/tools/capture-files.tool";
 export { createQueryGraphTool } from "./ai/tools/query-graph.tool";

package/src/routes/chat.ts

@@ -30,7 +30,7 @@ export function createChatRoutes() {
                 throw new HTTPException(404, { message: "Not found" });
             }
 
-            return stream(messages, thread, userId, token, config);
+            return stream(messages, thread, userId, token, config, c.req.raw.signal);
         },
     );
 

package/src/ws/events.ts

@@ -3,4 +3,9 @@ export type ThreadTitleUpdatedEvent = {
     payload: { threadId: string; title: string };
 };
 
-export type WsEvent = ThreadTitleUpdatedEvent;
+export type ThreadMessagesUpdatedEvent = {
+    type: "thread:messages-updated";
+    payload: { threadId: string };
+};
+
+export type WsEvent = ThreadTitleUpdatedEvent | ThreadMessagesUpdatedEvent;

package/src/ws/index.ts

@@ -1,3 +1,3 @@
 export { addConnection, removeConnection, getConnections } from "./connections.ts";
 export { notify } from "./notify.ts";
-export type { WsEvent, ThreadTitleUpdatedEvent } from "./events.ts";
+export type { WsEvent, ThreadTitleUpdatedEvent, ThreadMessagesUpdatedEvent } from "./events.ts";