@m6d/cortex-server 1.3.0 → 1.5.0

package/src/ai/active-streams.ts

@@ -0,0 +1,123 @@
+type ActiveStream = {
+    id: string;
+    abortController: AbortController;
+    buffer: string[];
+    subscribers: Set<ReadableStreamDefaultController<string>>;
+    isComplete: boolean;
+};
+
+const streams = new Map<string, ActiveStream>();
+
+export function registerStream(threadId: string, abortController: AbortController) {
+    const existing = streams.get(threadId);
+    if (existing) {
+        existing.abortController.abort();
+        for (const controller of existing.subscribers) {
+            try {
+                controller.close();
+            } catch {
+                /* already closed */
+            }
+        }
+        streams.delete(threadId);
+    }
+
+    const entry: ActiveStream = {
+        id: crypto.randomUUID(),
+        abortController,
+        buffer: [],
+        subscribers: new Set(),
+        isComplete: false,
+    };
+
+    streams.set(threadId, entry);
+    return entry;
+}
+
+export function attachSseStream(threadId: string, sseStream: ReadableStream<string>) {
+    const entry = streams.get(threadId);
+    if (!entry) return;
+
+    const reader = sseStream.getReader();
+    (async function consume() {
+        try {
+            while (true) {
+                const { done, value } = await reader.read();
+                if (done) break;
+                entry.buffer.push(value);
+                for (const controller of entry.subscribers) {
+                    try {
+                        controller.enqueue(value);
+                    } catch {
+                        /* subscriber cancelled */
+                    }
+                }
+            }
+        } catch {
+            /* stream aborted or errored */
+        } finally {
+            entry.isComplete = true;
+            for (const controller of entry.subscribers) {
+                try {
+                    controller.close();
+                } catch {
+                    /* already closed */
+                }
+            }
+        }
+    })();
+}
+
+export function subscribe(threadId: string) {
+    const entry = streams.get(threadId);
+    if (!entry) return null;
+
+    let savedController: ReadableStreamDefaultController<string>;
+
+    return new ReadableStream<string>({
+        start(controller) {
+            savedController = controller;
+            controller.enqueue("[START]"); // XXX: we need to add this flag for an immediate response for the subscriber, otherwise we'd have to wait for the first token from the llm.
+            for (const chunk of entry.buffer) {
+                controller.enqueue(chunk);
+            }
+            if (entry.isComplete) {
+                controller.close();
+                return;
+            }
+            entry.subscribers.add(controller);
+        },
+        cancel() {
+            entry.subscribers.delete(savedController);
+        },
+    });
+}
+
+export function abortStream(threadId: string) {
+    const entry = streams.get(threadId);
+    if (!entry) return false;
+    queueMicrotask(() => entry.abortController.abort());
+    return true;
+}
+
+export function isStreamRunning(threadId: string) {
+    const entry = streams.get(threadId);
+    return entry ? !entry.isComplete : false;
+}
+
+export function removeStream(threadId: string, streamId?: string) {
+    const entry = streams.get(threadId);
+    if (!entry) return;
+    if (streamId && entry.id !== streamId) return;
+
+    if (entry) {
+        for (const controller of entry.subscribers) {
+            try {
+                controller.close();
+            } catch {
+                /* already closed */
+            }
+        }
+    }
+    streams.delete(threadId);
+}

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