@m6d/cortex-server 1.3.0 → 1.5.0

package/src/ai/index.ts

@@ -1,18 +1,18 @@
 import {
     type UIMessage,
     type ToolSet,
-    consumeStream,
     convertToModelMessages,
     generateId,
     generateText,
     safeValidateUIMessages,
+    stepCountIs,
     streamText,
 } from "ai";
 import { HTTPException } from "hono/http-exception";
 import type { ResolvedCortexAgentConfig } from "../config.ts";
-import type { Thread } from "../types.ts";
+import type { MessageMetadata, Thread } from "../types.ts";
 import { createModel, createEmbeddingModel } from "./helpers.ts";
-import { buildSystemPrompt } from "./prompt.ts";
+import { buildSystemPrompt, resolveSession } from "./prompt.ts";
 import { createQueryGraphTool } from "./tools/query-graph.tool.ts";
 import { createCallEndpointTool } from "./tools/call-endpoint.tool.ts";
 import { createExecuteCodeTool } from "./tools/execute-code.tool.ts";
@@ -21,6 +21,15 @@ 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";
+import { buildContextMessages } from "./context/builder.ts";
+import { optimizeThreadContext, estimateTokens, trimMessagesToFit } from "./context/index.ts";
+import {
+    registerStream,
+    attachSseStream,
+    removeStream,
+    isStreamRunning,
+} from "./active-streams.ts";
+import { toThreadSummary } from "../types.ts";
 
 export async function stream(
     messages: unknown[],
@@ -28,8 +37,9 @@ export async function stream(
     userId: string,
     token: string,
     config: ResolvedCortexAgentConfig,
-    abortSignal?: AbortSignal,
 ) {
+    const abortController = new AbortController();
+
     const validationResult = await safeValidateUIMessages({ messages });
     if (!validationResult.success) {
         throw new HTTPException(423, { message: "Invalid messages format" });
@@ -37,15 +47,19 @@ export async function stream(
 
     const validatedMessages = validationResult.data;
     await config.db.messages.upsert(thread.id, validatedMessages);
+    const updatedThread = await config.db.threads.touch(thread.id);
 
-    const originalMessages = await config.db.messages
-        .list(userId, thread.id, { limit: 20 })
-        .then((x) => x.map((y) => y.content));
+    const activeStream = registerStream(thread.id, abortController);
 
-    const recentMessages = await convertToModelMessages(originalMessages);
+    notify(userId, thread.agentId, {
+        type: "thread:run-started",
+        payload: { thread: toThreadSummary(updatedThread, true) },
+    });
 
+    // Extract prompt from the just-upserted messages (last user message)
+    // so we can start graph resolution without waiting for history fetch
     const prompt =
-        originalMessages
+        validatedMessages
             .filter((x) => x.role === "user")
             .at(-1)
             ?.parts.find((x) => x.type === "text")?.text ?? "";
@@ -54,15 +68,24 @@ export async function stream(
     const embeddingModel = createEmbeddingModel(config.embedding);
     const neo4j = createNeo4jClient(config.neo4j, embeddingModel);
 
-    // Pre-resolve graph context
-    const resolved = await resolveFromGraph(prompt, {
-        neo4j,
-        embeddingModel,
-        reranker: config.reranker,
-    });
+    // Run independent operations in parallel
+    const [contextResult, resolved, session] = await Promise.all([
+        // Branch A: Load messages + build token-aware context window
+        buildContextMessages(userId, thread, config.db, config.context),
+        // Branch B: Resolve graph context (400-2000ms, the bottleneck)
+        resolveFromGraph(prompt, {
+            neo4j,
+            embeddingModel,
+            reranker: config.reranker,
+        }),
+        // Branch C: Resolve session data
+        resolveSession(config, thread, token),
+    ]);
+
+    const { messages: contextMessages, allMessages: originalMessages } = contextResult;
 
     // Build tools
-    const builtInTools: Record<string, unknown> = {
+    const builtInTools: ToolSet = {
         captureFiles: captureFilesTool,
         queryGraph: createQueryGraphTool(neo4j),
     };
@@ -88,36 +111,105 @@ export async function stream(
         ...config.tools,
     } as ToolSet;
 
-    const systemPrompt = await buildSystemPrompt(config, prompt, thread, token, resolved);
+    const systemPrompt = await buildSystemPrompt(config, resolved, session);
+
+    // The context builder reserved a static token budget for the system prompt + tools.
+    // Now that we have the actual values, verify the reserve was sufficient and trim
+    // the oldest messages if it wasn't.
+    const actualFixedCost = estimateTokens(systemPrompt) + estimateTokens(JSON.stringify(tools));
+    const { reservedTokenBudget, maxContextTokens } = config.context;
+    const trimmedMessages =
+        actualFixedCost > reservedTokenBudget
+            ? trimMessagesToFit(contextMessages, maxContextTokens - actualFixedCost)
+            : contextMessages;
+
+    const recentMessages = await convertToModelMessages(trimmedMessages);
 
     const result = streamText({
         model,
         system: systemPrompt,
         tools,
         messages: recentMessages,
-        abortSignal,
+        abortSignal: abortController.signal,
+        stopWhen: stepCountIs(50),
     });
 
-    return result.toUIMessageStreamResponse({
+    return result.toUIMessageStreamResponse<UIMessage<MessageMetadata>>({
         originalMessages,
         generateMessageId: generateId,
-        consumeSseStream: consumeStream,
+        consumeSseStream: ({ stream: sseStream }) => {
+            attachSseStream(thread.id, sseStream);
+        },
         onFinish: async ({ messages: finishedMessages, isAborted }) => {
             if (isAborted) {
                 finalizeAbortedMessages(finishedMessages);
             }
+
+            // Record token usage (result promises reject on abort, so skip)
+            const lastAssistantMessage = finishedMessages
+                .filter((x) => x.role === "assistant")
+                .at(-1);
+            if (lastAssistantMessage && !isAborted) {
+                const providerMetadata = await result.providerMetadata;
+                const response = await result.response;
+                const usage = await result.totalUsage;
+                let metadata: MessageMetadata = {
+                    isAborted,
+                    providerMetadata,
+                    modelId: response.modelId,
+                    tokenUsage: {
+                        input: {
+                            noCache: usage.inputTokenDetails.noCacheTokens ?? 0,
+                            cacheRead: usage.inputTokenDetails.cacheReadTokens ?? 0,
+                            cacheWrite: usage.inputTokenDetails.cacheWriteTokens ?? 0,
+                            total: usage.inputTokens ?? 0,
+                        },
+                        output: {
+                            reasoning: usage.outputTokenDetails.reasoningTokens ?? 0,
+                            text: usage.outputTokenDetails.textTokens ?? 0,
+                            total: usage.outputTokens ?? 0,
+                        },
+                        total: usage.totalTokens ?? 0,
+                    },
+                };
+
+                lastAssistantMessage.metadata = metadata;
+            } else if (lastAssistantMessage) {
+                lastAssistantMessage.metadata = {
+                    isAborted,
+                    modelId: "",
+                    providerMetadata: undefined,
+                };
+            }
+
+            const persistedThread = await config.db.threads.getById(userId, thread.id);
+            if (!persistedThread) {
+                removeStream(thread.id);
+                return;
+            }
+
             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 },
-                });
+            // immediately because messages may not have been saved yet.
+            notify(userId, persistedThread.agentId, {
+                type: "thread:messages-updated",
+                payload: {
+                    threadId: thread.id,
+                    thread: toThreadSummary(persistedThread, false),
+                },
+            });
+
+            setTimeout(() => removeStream(thread.id, activeStream.id), 10_000);
+
+            // Fire-and-forget: optimize context for next request
+            // Runs after response is delivered — no perceived latency
+            try {
+                optimizeThreadContext(thread, finishedMessages, config);
+            } catch (err) {
+                console.error("[cortex-server] Context optimization failed:", err);
             }
         },
     });
@@ -143,9 +235,12 @@ going to do so or any other speech. Spit out only the title.`,
 
     await config.db.threads.updateTitle(threadId, output ?? "");
 
-    notify(userId, {
+    const thread = await config.db.threads.getById(userId, threadId);
+    if (!thread) return;
+
+    notify(userId, thread.agentId, {
         type: "thread:title-updated",
-        payload: { threadId, title: output ?? "" },
+        payload: { thread: toThreadSummary(thread, isStreamRunning(thread.id)) },
     });
 }
 

package/src/ai/prompt.ts

@@ -2,29 +2,35 @@ import type { ResolvedContext } from "../graph/resolver.ts";
 import type { ResolvedCortexAgentConfig } from "../config.ts";
 import type { Thread } from "../types.ts";
 
-export async function buildSystemPrompt(
+/**
+ * Resolves session data for the thread, loading from the configured
+ * session loader if not already cached on the thread.
+ */
+export async function resolveSession(
     config: ResolvedCortexAgentConfig,
-    prompt: string,
     thread: Thread,
     token: string,
+) {
+    let session = thread.session;
+
+    if (!session && config.loadSessionData) {
+        session = await config.loadSessionData(token);
+        // Persist to DB for future cache hits
+        await config.db.threads.updateSession(thread.id, session);
+        thread.session = session;
+    }
+
+    return session;
+}
+
+export async function buildSystemPrompt(
+    config: ResolvedCortexAgentConfig,
     resolved: ResolvedContext | null,
+    session: Record<string, unknown> | null,
 ) {
     // Resolve the consumer's base system prompt
     let basePrompt: string;
     if (typeof config.systemPrompt === "function") {
-        // Resolve session data with caching
-        let session: Record<string, unknown> | null = thread.session as Record<
-            string,
-            unknown
-        > | null;
-
-        if (!session && config.loadSessionData) {
-            session = await config.loadSessionData(token);
-            // Persist to DB for future cache hits
-            await config.db.threads.updateSession(thread.id, session);
-            thread.session = session;
-        }
-
         basePrompt = await config.systemPrompt(session);
     } else {
         basePrompt = config.systemPrompt;
@@ -59,6 +65,7 @@ The following endpoints were automatically matched to the user's message.`,
                       )
                       .join("\n")
                 : "";
+        const meta = ep.metadata !== "{}" ? `\n- Metadata: ${ep.metadata}` : "";
         parts.push(
             `
 ### ${ep.concept} (read)
@@ -66,7 +73,7 @@ The following endpoints were automatically matched to the user's message.`,
 - ${ep.method} ${ep.path}
 - Params: ${ep.params}
 - Body: ${ep.body}
-- Response: ${ep.response}${rules}${deps}`,
+- Response: ${ep.response}${rules}${deps}${meta}`,
         );
     }
 
@@ -82,6 +89,7 @@ The following endpoints were automatically matched to the user's message.`,
                       )
                       .join("\n")
                 : "";
+        const meta = ep.metadata !== "{}" ? `\n- Metadata: ${ep.metadata}` : "";
         parts.push(
             `
 ### ${ep.concept} (write)
@@ -89,17 +97,18 @@ The following endpoints were automatically matched to the user's message.`,
 - ${ep.method} ${ep.path}
 - Params: ${ep.params}
 - Body: ${ep.body}
-- Response: ${ep.response}${rules}${deps}`,
+- Response: ${ep.response}${rules}${deps}${meta}`,
         );
     }
 
     for (const svc of resolved.services) {
         const rules = svc.rules.length > 0 ? `\n  Rules: ${svc.rules.join("; ")}` : "";
+        const meta = svc.metadata !== "{}" ? `\n- Metadata: ${svc.metadata}` : "";
         parts.push(
             `
 ### ${svc.concept} via ${svc.serviceName} (service)
 - Built-in ID: ${svc.builtInId}
-- Description: ${svc.description || "N/A"}${rules}`,
+- Description: ${svc.description || "N/A"}${rules}${meta}`,
         );
     }
 

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