@m6d/cortex-server 1.3.0 → 1.5.0

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;
 }

package/src/types.ts

@@ -1,24 +1,38 @@
 import type { ToolUIPart, UIMessage } from "ai";
 import type { ResolvedCortexAgentConfig } from "./config";
+import type { InferSelectModel } from "drizzle-orm";
+import type { messages, threads } from "./db/schema";
 
-export type Thread = {
+export type Thread = InferSelectModel<typeof threads>;
+
+export type StoredMessage = InferSelectModel<typeof messages>;
+
+export type ThreadSummary = {
     id: string;
-    userId: string;
-    agentId: string;
-    title: string | null;
-    session: Record<string, unknown> | null;
+    title?: string;
     createdAt: Date;
     updatedAt: Date;
+    isRunning: boolean;
 };
 
-export type StoredMessage = {
-    id: string;
-    threadId: string;
-    text: string | null;
-    content: UIMessage;
-    role: "system" | "user" | "assistant" | "tool";
-    createdAt: Date;
-    updatedAt: Date;
+export type MessageMetadata = {
+    modelId: string;
+    providerMetadata: unknown;
+    isAborted?: boolean;
+    tokenUsage?: {
+        input: {
+            noCache: number;
+            cacheRead: number;
+            cacheWrite: number;
+            total: number;
+        };
+        output: {
+            reasoning: number;
+            text: number;
+            total: number;
+        };
+        total: number;
+    };
 };
 
 export type CapturedFileInput = {
@@ -48,3 +62,13 @@ export type CortexAppEnv = {
         agentId: string;
     };
 };
+
+export function toThreadSummary(thread: Thread, isRunning: boolean) {
+    return {
+        id: thread.id,
+        title: thread.title ?? undefined,
+        createdAt: thread.createdAt,
+        updatedAt: thread.updatedAt,
+        isRunning,
+    } satisfies ThreadSummary;
+}

package/src/ws/connections.ts

@@ -2,22 +2,28 @@ import type { WSContext } from "hono/ws";
 
 const connections = new Map<string, WSContext[]>();
 
-export function addConnection(userId: string, ws: WSContext) {
-    const sockets = connections.get(userId) ?? [];
+function getConnectionKey(userId: string, agentId: string) {
+    return `${userId}:${agentId}`;
+}
+
+export function addConnection(userId: string, agentId: string, ws: WSContext) {
+    const key = getConnectionKey(userId, agentId);
+    const sockets = connections.get(key) ?? [];
     sockets.push(ws);
-    connections.set(userId, sockets);
+    connections.set(key, sockets);
 }
 
-export function removeConnection(userId: string, ws: WSContext) {
-    let sockets = connections.get(userId) ?? [];
+export function removeConnection(userId: string, agentId: string, ws: WSContext) {
+    const key = getConnectionKey(userId, agentId);
+    let sockets = connections.get(key) ?? [];
     sockets = sockets.filter((x) => x !== ws);
     if (sockets.length) {
-        connections.set(userId, sockets);
+        connections.set(key, sockets);
     } else {
-        connections.delete(userId);
+        connections.delete(key);
     }
 }
 
-export function getConnections(userId: string) {
-    return connections.get(userId) ?? [];
+export function getConnections(userId: string, agentId: string) {
+    return connections.get(getConnectionKey(userId, agentId)) ?? [];
 }

package/src/ws/events.ts

@@ -1,11 +1,39 @@
+import type { ThreadSummary } from "../types.ts";
+
+export type ThreadCreatedEvent = {
+    type: "thread:created";
+    payload: { thread: ThreadSummary };
+};
+
+export type ThreadDeletedEvent = {
+    type: "thread:deleted";
+    payload: { threadId: string };
+};
+
 export type ThreadTitleUpdatedEvent = {
     type: "thread:title-updated";
-    payload: { threadId: string; title: string };
+    payload: { thread: ThreadSummary };
+};
+
+export type ThreadRunStartedEvent = {
+    type: "thread:run-started";
+    payload: { thread: ThreadSummary };
+};
+
+export type ThreadRunFinishedEvent = {
+    type: "thread:run-finished";
+    payload: { thread: ThreadSummary };
 };
 
 export type ThreadMessagesUpdatedEvent = {
     type: "thread:messages-updated";
-    payload: { threadId: string };
+    payload: { threadId: string; thread: ThreadSummary };
 };
 
-export type WsEvent = ThreadTitleUpdatedEvent | ThreadMessagesUpdatedEvent;
+export type WsEvent =
+    | ThreadCreatedEvent
+    | ThreadDeletedEvent
+    | ThreadTitleUpdatedEvent
+    | ThreadRunStartedEvent
+    | ThreadRunFinishedEvent
+    | ThreadMessagesUpdatedEvent;

package/src/ws/index.ts

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

package/src/ws/notify.ts

@@ -1,9 +1,9 @@
 import { getConnections } from "./connections.ts";
 import type { WsEvent } from "./events.ts";
 
-export function notify(userId: string, event: WsEvent) {
+export function notify(userId: string, agentId: string, event: WsEvent) {
     const message = JSON.stringify(event);
-    for (const ws of getConnections(userId)) {
+    for (const ws of getConnections(userId, agentId)) {
         ws.send(message);
     }
 }