@m6d/cortex-server 1.2.0 → 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/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.2.0",
+    "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/ai/index.ts

@@ -1,5 +1,7 @@
 import {
+    type UIMessage,
     type ToolSet,
+    consumeStream,
     convertToModelMessages,
     generateId,
     generateText,
@@ -26,6 +28,7 @@ export async function stream(
     userId: string,
     token: string,
     config: ResolvedCortexAgentConfig,
+    abortSignal?: AbortSignal,
 ) {
     const validationResult = await safeValidateUIMessages({ messages });
     if (!validationResult.success) {
@@ -92,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 },
+                });
+            }
         },
     });
 }
@@ -135,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/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/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";