@m6d/cortex-server 1.7.0 → 2.0.0

package/src/ai/active-streams.ts

@@ -1,123 +0,0 @@
-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/compressor.ts

@@ -1,47 +0,0 @@
-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

@@ -1,75 +0,0 @@
-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

@@ -1,50 +0,0 @@
-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

@@ -1,60 +0,0 @@
-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/helpers.ts

@@ -1,32 +0,0 @@
-import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
-import type { ResolvedCortexAgentConfig } from "../config.ts";
-
-export function createModel(config: ResolvedCortexAgentConfig["model"]) {
-    const provider = createOpenAICompatible({
-        name: config.providerName ?? "default",
-        baseURL: config.baseURL,
-        apiKey: config.apiKey,
-    });
-
-    return provider(config.modelName);
-}
-
-export function createEmbeddingModel(config: ResolvedCortexAgentConfig["embedding"]) {
-    const provider = createOpenAICompatible({
-        name: "embedding",
-        baseURL: config.baseURL,
-        apiKey: config.apiKey,
-    });
-
-    return provider.textEmbeddingModel(config.modelName);
-}
-
-export async function streamToBase64(stream: ReadableStream<Uint8Array>) {
-    return new Response(stream)
-        .arrayBuffer()
-        .then((buffer) => Buffer.from(buffer).toString("base64"));
-}
-
-export function tokenToUserId(token: string) {
-    return JSON.parse(Buffer.from(token.split(".")[1]!, "base64").toString())["sub"] as string;
-}

package/src/ai/index.ts

@@ -1,277 +0,0 @@
-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 { MessageMetadata, Thread } from "../types.ts";
-import { createModel, createEmbeddingModel } from "./helpers.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";
-import { captureFilesTool } from "./tools/capture-files.tool.ts";
-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[],
-    thread: Thread,
-    userId: string,
-    token: string,
-    requestContext: Record<string, unknown>,
-    config: ResolvedCortexAgentConfig,
-) {
-    const abortController = new AbortController();
-
-    const validationResult = await safeValidateUIMessages({ messages });
-    if (!validationResult.success) {
-        throw new HTTPException(423, { message: "Invalid messages format" });
-    }
-
-    const validatedMessages = validationResult.data;
-    await config.db.messages.upsert(thread.id, validatedMessages);
-    const updatedThread = await config.db.threads.touch(thread.id);
-
-    const activeStream = registerStream(thread.id, abortController);
-
-    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 =
-        validatedMessages
-            .filter((x) => x.role === "user")
-            .at(-1)
-            ?.parts.find((x) => x.type === "text")?.text ?? "";
-
-    const model = createModel(config.model);
-    const embeddingModel = createEmbeddingModel(config.embedding);
-    const neo4j = createNeo4jClient(config.neo4j, embeddingModel);
-
-    // 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: ToolSet = {
-        captureFiles: captureFilesTool,
-        queryGraph: createQueryGraphTool(neo4j),
-    };
-
-    if (config.backendFetch) {
-        const backendFetchWithInterceptor = {
-            ...config.backendFetch,
-            transformRequestBody:
-                config.backendFetch.transformRequestBody ??
-                createRequestInterceptor(
-                    config.db,
-                    config.storage,
-                    config.backendFetch.interceptor,
-                ),
-        };
-
-        builtInTools["callEndpoint"] = createCallEndpointTool(backendFetchWithInterceptor, token);
-        builtInTools["executeCode"] = createExecuteCodeTool(backendFetchWithInterceptor, token);
-    }
-
-    const tools = {
-        ...builtInTools,
-        ...config.tools,
-    } as ToolSet;
-
-    const systemPrompt = await buildSystemPrompt(config, resolved, {
-        session,
-        requestContext,
-    });
-
-    // 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: abortController.signal,
-        stopWhen: stepCountIs(50),
-    });
-
-    return result.toUIMessageStreamResponse<UIMessage<MessageMetadata>>({
-        originalMessages,
-        generateMessageId: generateId,
-        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 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);
-            }
-        },
-    });
-}
-
-export async function generateTitle(
-    threadId: string,
-    prompt: string,
-    userId: string,
-    config: ResolvedCortexAgentConfig,
-) {
-    const model = createModel(config.model);
-
-    const { output } = await generateText({
-        model,
-        system: `You are an expert in generating titles for threads of chats
-given the first message in that thread.
-
-When asked, only respond with the title without saying you're
-going to do so or any other speech. Spit out only the title.`,
-        prompt: `Generate a title for this prompt: ${prompt}`,
-    });
-
-    await config.db.threads.updateTitle(threadId, output ?? "");
-
-    const thread = await config.db.threads.getById(userId, threadId);
-    if (!thread) return;
-
-    notify(userId, thread.agentId, {
-        type: "thread:title-updated",
-        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;
-    });
-}