@m6d/cortex-server 1.2.0 → 1.4.0

package/src/ai/context/builder.ts

@@ -0,0 +1,94 @@
+import type { UIMessage } from "ai";
+import { generateId } from "ai";
+import type { DatabaseAdapter } from "../../adapters/database.ts";
+import type { MessageMetadata, Thread } from "../../types.ts";
+import type { ContextConfig, ThreadContextMeta } from "./types.ts";
+import { compressToolResults } from "./compressor.ts";
+import { estimateMessageTokens } from "./token-estimator.ts";
+
+type ContextBuildResult = {
+    messages: UIMessage<MessageMetadata>[];
+    allMessages: UIMessage<MessageMetadata>[];
+    summary: string | null;
+};
+
+/**
+ * Builds a token-aware context window from stored messages.
+ *
+ * 1. Loads messages from DB with generous limit
+ * 2. Reads existing summary from thread.contextMeta
+ * 3. Compresses large tool results
+ * 4. Walks messages newest-to-oldest, accumulating token estimates
+ * 5. Stops when adding the next message would exceed the budget
+ * 6. Prepends summary as synthetic message if older messages were trimmed
+ */
+export async function buildContextMessages(
+    userId: string,
+    thread: Thread,
+    db: DatabaseAdapter,
+    contextConfig: ContextConfig,
+) {
+    // 1. Load messages with generous limit
+    const storedMessages = await db.messages.list(userId, thread.id, { limit: 50 });
+    const allMessages = storedMessages.map((m) => m.content);
+
+    // 2. Read existing summary
+    const contextMeta = thread.contextMeta;
+    const summary = contextMeta?.summary ?? null;
+
+    // 3. Compress large tool results
+    const compressed = compressToolResults(allMessages, contextConfig.toolResultMaxTokens);
+
+    // 4. Walk newest-to-oldest, accumulating token estimates.
+    // Reserve space for the system prompt + tool definitions (they share the context window).
+    const budget = contextConfig.maxContextTokens - contextConfig.reservedTokenBudget;
+    const selected: UIMessage<MessageMetadata>[] = [];
+    let accumulated = 0;
+
+    for (let i = compressed.length - 1; i >= 0; i--) {
+        const msgTokens = estimateMessageTokens(compressed[i]!);
+
+        if (accumulated + msgTokens > budget && selected.length > 0) {
+            break;
+        }
+
+        accumulated += msgTokens;
+        selected.unshift(compressed[i]!);
+    }
+
+    // 5. If we trimmed messages and a summary exists, prepend it.
+    // Make room for the summary by evicting the oldest messages if needed.
+    if (summary && selected.length < allMessages.length) {
+        const summaryMessage = {
+            id: generateId(),
+            role: "user",
+            parts: [{ type: "text", text: `[Previous conversation summary]: ${summary}` }],
+        } satisfies UIMessage;
+
+        const summaryTokens = estimateMessageTokens(summaryMessage);
+        const trimmed = trimMessagesToFit(selected, budget - summaryTokens);
+        trimmed.unshift(summaryMessage);
+        return { messages: trimmed, allMessages, summary };
+    }
+
+    return { messages: selected, allMessages, summary } satisfies ContextBuildResult;
+}
+
+/**
+ * Drops the oldest messages until the total estimated tokens fit within `budget`.
+ * Always keeps at least the most recent message.
+ */
+export function trimMessagesToFit(messages: UIMessage<MessageMetadata>[], budget: number) {
+    let total = 0;
+    for (const msg of messages) {
+        total += estimateMessageTokens(msg);
+    }
+
+    const trimmed = [...messages];
+    while (total > budget && trimmed.length > 1) {
+        const evicted = trimmed.shift()!;
+        total -= estimateMessageTokens(evicted);
+    }
+
+    return trimmed;
+}

package/src/ai/context/compressor.ts

@@ -0,0 +1,47 @@
+import type { UIMessage } from "ai";
+import { estimateTokens, CHARS_PER_TOKEN } from "./token-estimator.ts";
+import type { MessageMetadata } from "src/types.ts";
+
+/**
+ * Returns a new array of messages with large tool outputs truncated
+ * to `maxTokensPerResult`. Does not mutate the input messages.
+ */
+export function compressToolResults(
+    messages: UIMessage<MessageMetadata>[],
+    maxTokensPerResult: number,
+) {
+    return messages.map((message) => {
+        let hasLargeToolOutput = false;
+
+        for (const part of message.parts) {
+            if ("toolCallId" in part && "output" in part && part.output != null) {
+                const outputTokens = estimateTokens(JSON.stringify(part.output));
+                if (outputTokens > maxTokensPerResult) {
+                    hasLargeToolOutput = true;
+                    break;
+                }
+            }
+        }
+
+        if (!hasLargeToolOutput) return message;
+
+        const compressedParts = message.parts.map((part) => {
+            if (!("toolCallId" in part) || !("output" in part) || part.output == null) {
+                return part;
+            }
+
+            const outputStr = JSON.stringify(part.output);
+            const outputTokens = estimateTokens(outputStr);
+
+            if (outputTokens <= maxTokensPerResult) return part;
+
+            // Convert token budget back to character budget
+            const charBudget = maxTokensPerResult * CHARS_PER_TOKEN;
+            const truncatedOutput = outputStr.slice(0, charBudget) + "\n[...truncated]";
+
+            return { ...part, output: truncatedOutput } as typeof part;
+        });
+
+        return { ...message, parts: compressedParts };
+    });
+}

package/src/ai/context/index.ts

@@ -0,0 +1,75 @@
+export type { ContextConfig, ThreadContextMeta } from "./types.ts";
+export { DEFAULT_CONTEXT_CONFIG } from "./types.ts";
+export {
+    CHARS_PER_TOKEN,
+    estimateTokens,
+    estimateMessageTokens,
+    estimateMessagesTokens,
+} from "./token-estimator.ts";
+export { compressToolResults } from "./compressor.ts";
+export { summarizeMessages } from "./summarizer.ts";
+export { buildContextMessages, trimMessagesToFit } from "./builder.ts";
+
+import type { UIMessage } from "ai";
+import type { ResolvedCortexAgentConfig } from "../../config.ts";
+import type { Thread } from "../../types.ts";
+import type { ThreadContextMeta } from "./types.ts";
+import { estimateMessagesTokens } from "./token-estimator.ts";
+import { summarizeMessages } from "./summarizer.ts";
+
+/**
+ * Post-response context optimization.
+ * Called fire-and-forget from onFinish — summarizes older messages
+ * when token usage exceeds the configured threshold.
+ */
+export async function optimizeThreadContext(
+    thread: Thread,
+    messages: UIMessage[],
+    config: ResolvedCortexAgentConfig,
+) {
+    const contextConfig = config.context;
+
+    // 1. Estimate tokens for all messages
+    const estimates = estimateMessagesTokens(messages);
+    const totalEstimatedTokens = estimates.reduce((sum, e) => sum + e.tokens, 0);
+
+    // 2. Check if over summarization threshold
+    const threshold = contextConfig.maxContextTokens * contextConfig.summarizationThreshold;
+
+    if (totalEstimatedTokens <= threshold) {
+        // Update token estimate but skip summarization
+        const meta: ThreadContextMeta = {
+            summary: thread.contextMeta?.summary ?? null,
+            summaryUpToMessageId: thread.contextMeta?.summaryUpToMessageId ?? null,
+            totalEstimatedTokens,
+            lastOptimizedAt: new Date().toISOString(),
+        };
+        await config.db.threads.updateContextMeta(thread.id, meta);
+        return;
+    }
+
+    // 3. Keep the most recent messages unsummarized (they're the hot context)
+    const recentCount = Math.min(contextConfig.recentMessagesToKeep, messages.length);
+    const messagesToSummarize = messages.slice(0, messages.length - recentCount);
+
+    if (messagesToSummarize.length === 0) return;
+
+    // 4. Determine model config for summarization
+    const modelConfig = contextConfig.summarizationModel ?? config.model;
+
+    // 5. Generate summary incorporating any existing summary
+    const existingSummary = thread.contextMeta?.summary ?? null;
+
+    const summary = await summarizeMessages(messagesToSummarize, existingSummary, modelConfig);
+
+    // 6. Update thread context meta
+    const lastSummarizedMessage = messagesToSummarize.at(-1);
+    const meta = {
+        summary,
+        summaryUpToMessageId: lastSummarizedMessage?.id ?? null,
+        totalEstimatedTokens,
+        lastOptimizedAt: new Date().toISOString(),
+    } satisfies ThreadContextMeta;
+
+    await config.db.threads.updateContextMeta(thread.id, meta);
+}

package/src/ai/context/summarizer.ts

@@ -0,0 +1,50 @@
+import { generateText } from "ai";
+import type { UIMessage } from "ai";
+import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
+import type { ContextConfig } from "./types.ts";
+
+type SummarizationModelConfig = NonNullable<ContextConfig["summarizationModel"]>;
+
+export async function summarizeMessages(
+    messages: UIMessage[],
+    existingSummary: string | null,
+    modelConfig: SummarizationModelConfig,
+) {
+    const provider = createOpenAICompatible({
+        name: modelConfig.providerName ?? "summarization-provider",
+        baseURL: modelConfig.baseURL,
+        apiKey: modelConfig.apiKey,
+    });
+
+    const model = provider.chatModel(modelConfig.modelName);
+
+    const conversationText = messages
+        .map(function (msg) {
+            const textParts = msg.parts
+                .filter((p): p is Extract<typeof p, { type: "text" }> => p.type === "text")
+                .map((p) => p.text);
+            return `[${msg.role}]: ${textParts.join(" ")}`;
+        })
+        .join("\n");
+
+    const summaryContext = existingSummary
+        ? `\nPrior summary of earlier messages:\n${existingSummary}\n`
+        : "";
+
+    const { text } = await generateText({
+        model,
+        system: `You are a precise conversation summarizer. Produce a concise summary that preserves:
+- Key decisions and conclusions
+- Important entities (names, IDs, URLs, values)
+- User intent and goals
+- Any unresolved questions or next steps
+
+Maximum 500 tokens. Use bullet points. Do not include preamble.`,
+        prompt: `${summaryContext}
+Summarize the following conversation:
+
+${conversationText}`,
+    });
+
+    return text;
+}

package/src/ai/context/token-estimator.ts

@@ -0,0 +1,60 @@
+import type { UIMessage } from "ai";
+
+/** Average characters per token for English text. Used by the heuristic estimator. */
+export const CHARS_PER_TOKEN = 4;
+
+/**
+ * Estimates token count for a string using the chars/4 heuristic.
+ * ~10% accuracy for English text — good enough for budget decisions.
+ */
+export function estimateTokens(text: string) {
+    return Math.ceil(text.length / CHARS_PER_TOKEN);
+}
+
+/**
+ * Estimates token count for a single UIMessage by walking its parts.
+ */
+export function estimateMessageTokens(message: UIMessage) {
+    let tokens = 0;
+
+    for (const part of message.parts) {
+        if (part.type === "text" || part.type === "reasoning") {
+            tokens += estimateTokens(part.text);
+        } else if ("toolCallId" in part) {
+            // Tool invocation parts (tool-${name})
+            if ("input" in part && part.input != null) {
+                tokens += estimateTokens(JSON.stringify(part.input));
+            }
+            if ("output" in part && part.output != null) {
+                tokens += estimateTokens(JSON.stringify(part.output));
+            }
+        } else if (part.type === "source-url") {
+            tokens += estimateTokens(part.url);
+        } else if (part.type === "source-document") {
+            tokens += estimateTokens(part.title);
+        } else if (part.type === "file") {
+            tokens += estimateTokens(part.url);
+        } else {
+            // step-start, data parts, etc.
+            tokens += 5;
+        }
+    }
+
+    // Per-message overhead (role, metadata framing)
+    tokens += 4;
+
+    return tokens;
+}
+
+/**
+ * Estimates token counts for an array of UIMessages.
+ * Returns per-message estimates in the same order.
+ */
+export function estimateMessagesTokens(messages: UIMessage[]) {
+    return messages.map(function (message) {
+        return {
+            message,
+            tokens: estimateMessageTokens(message),
+        };
+    });
+}

package/src/ai/context/types.ts

@@ -0,0 +1,28 @@
+export type ContextConfig = {
+    maxContextTokens: number;
+    reservedTokenBudget: number;
+    summarizationThreshold: number;
+    summarizationModel?: {
+        baseURL: string;
+        apiKey: string;
+        modelName: string;
+        providerName?: string;
+    };
+    toolResultMaxTokens: number;
+    recentMessagesToKeep: number;
+};
+
+export type ThreadContextMeta = {
+    summary: string | null;
+    summaryUpToMessageId: string | null;
+    totalEstimatedTokens: number;
+    lastOptimizedAt: string | null;
+};
+
+export const DEFAULT_CONTEXT_CONFIG: ContextConfig = {
+    maxContextTokens: 120_000,
+    reservedTokenBudget: 8_000,
+    summarizationThreshold: 0.75,
+    toolResultMaxTokens: 2_000,
+    recentMessagesToKeep: 6,
+};

package/src/ai/index.ts

@@ -1,16 +1,18 @@
 import {
+    type UIMessage,
     type ToolSet,
     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";
@@ -19,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[],
@@ -27,6 +38,8 @@ export async function stream(
     token: string,
     config: ResolvedCortexAgentConfig,
 ) {
+    const abortController = new AbortController();
+
     const validationResult = await safeValidateUIMessages({ messages });
     if (!validationResult.success) {
         throw new HTTPException(423, { message: "Invalid messages format" });
@@ -34,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 ?? "";
@@ -51,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),
     };
@@ -85,27 +111,106 @@ 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,
-        onAbort: () => {
-            console.log("Stream aborted");
-        },
+        abortSignal: abortController.signal,
+        stopWhen: stepCountIs(50),
     });
 
-    return result.toUIMessageStreamResponse({
+    return result.toUIMessageStreamResponse<UIMessage<MessageMetadata>>({
         originalMessages,
         generateMessageId: generateId,
+        consumeSseStream: ({ stream: sseStream }) => {
+            attachSseStream(thread.id, sseStream);
+        },
         onFinish: async ({ messages: finishedMessages, isAborted }) => {
             if (isAborted) {
-                console.log("Stream was aborted");
+                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 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);
+            }
         },
     });
 }
@@ -130,8 +235,39 @@ 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)) },
+    });
+}
+
+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/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;