@m6d/cortex-server 1.8.0 → 2.0.1

package/src/lib/ai/context/intra-turn-compressor.ts

@@ -1,132 +1,56 @@
-import type {
-    LanguageModelV3Prompt,
-    LanguageModelV3Message,
-    LanguageModelV3ToolResultPart,
-} from "@ai-sdk/provider";
-import { generateText } from "ai";
-import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
-import { estimateTokens } from "./token-estimator.ts";
-import { summarizeValue } from "./compressor.ts";
-import type { ContextConfig } from "./types.ts";
-
-type ToolRoleMessage = LanguageModelV3Message & { role: "tool" };
+import { chat } from "@tanstack/ai";
+import type { ModelMessage } from "@tanstack/ai";
+import { createModel } from "@/ai/helpers";
+import { estimateTokens } from "./token-estimator";
+import { summarizeValue } from "./compressor";
+import type { ContextConfig } from "./types";
 
 type SummarizationModelConfig = NonNullable<ContextConfig["summarizationModel"]>;
 
+type ToolCallOrigin = { name: string; step: number };
+
 /**
- * Compresses tool results from older steps in the prompt.
- *
- * On step N, tool results from steps 1..N-2 are aggressively compressed
- * to short one-liners. Only the most recent tool result stays in full.
- * Does not mutate the input — returns a new prompt array.
+ * A tool message names only the call it answers, so both the tool's own name and the
+ * step it belongs to have to be recovered from the assistant message that asked for
+ * it. One tool message carries one result, so a step that issued three calls in
+ * parallel arrives as three messages; regrouping them into their step is what stops
+ * a step's own results from ageing each other out.
  */
-export function compressIntraTurnToolResults(
-    prompt: LanguageModelV3Prompt,
-    toolResultMaxTokens: number,
-) {
-    // Collect indices of all tool-role messages
-    const toolMessageIndices: number[] = [];
-    for (let i = 0; i < prompt.length; i++) {
-        if (prompt[i]!.role === "tool") {
-            toolMessageIndices.push(i);
-        }
-    }
-
-    // Need at least 3 tool messages to have "old" ones to compress
-    // (keep last 1 fully intact, compress the rest)
-    if (toolMessageIndices.length < 3) return prompt;
-
-    // Clone the prompt so we don't mutate the original
-    const result = prompt.map((msg) => ({ ...msg }));
+function toolStepsOf(messages: ModelMessage[]) {
+    const calls = new Map<string, ToolCallOrigin>();
+    const steps: number[][] = [];
 
-    // Compress all tool messages except the last one
-    const indicesToCompress = toolMessageIndices.slice(0, -1);
-
-    for (const idx of indicesToCompress) {
-        const msg = result[idx] as ToolRoleMessage;
-        const compressedContent = msg.content.map((part) => {
-            if (part.type !== "tool-result") return part;
-            return compressToolResultPart(part, toolResultMaxTokens);
-        });
+    for (const [index, message] of messages.entries()) {
+        for (const call of message.toolCalls ?? []) {
+            calls.set(call.id, { name: call.function.name, step: steps.length });
+        }
+        if (message.role !== "tool") continue;
 
-        result[idx] = { ...msg, content: compressedContent } as LanguageModelV3Message;
+        // A result whose call was never announced belongs to the step in progress;
+        // filing it anywhere else would invent a step that issued nothing.
+        const step = calls.get(message.toolCallId ?? "")?.step ?? steps.length;
+        (steps[step] ??= []).push(index);
     }
 
-    return result;
-}
-
-function compressToolResultPart(
-    part: LanguageModelV3ToolResultPart,
-    maxTokens: number,
-): LanguageModelV3ToolResultPart {
-    const output = part.output;
-
-    // Already small enough
-    const outputStr = stringifyOutput(output);
-    if (estimateTokens(outputStr) <= maxTokens) return part;
-
-    const compressed = compressOutputToOneLiner(output, part.toolName);
-
-    return {
-        ...part,
-        output: { type: "text", value: compressed },
-    };
+    return { calls, steps };
 }
 
-function stringifyOutput(output: LanguageModelV3ToolResultPart["output"]): string {
-    switch (output.type) {
-        case "text":
-        case "error-text":
-            return output.value;
-        case "json":
-        case "error-json":
-            return JSON.stringify(output.value);
-        case "execution-denied":
-            return output.reason ?? "denied";
-        case "content":
-            return JSON.stringify(output.value);
-        default:
-            return "";
-    }
+function toolNameOf(message: ModelMessage, calls: Map<string, ToolCallOrigin>) {
+    return calls.get(message.toolCallId ?? "")?.name ?? "tool";
 }
 
-function compressOutputToOneLiner(
-    output: LanguageModelV3ToolResultPart["output"],
-    toolName: string,
-): string {
-    const tag = `[COMPRESSED prior ${toolName} result]`;
+/**
+ * A message's content is either a plain string or content parts, of which only the
+ * text ones can be digested — which is all a one-line stand-in was going to keep.
+ */
+function contentText(content: ModelMessage["content"]) {
+    if (typeof content === "string") return content;
+    if (content === null) return "";
 
-    switch (output.type) {
-        case "text": {
-            const preview = extractPreview(output.value);
-            return `${tag} ${preview}`;
-        }
-        case "json":
-            return `${tag} ${summarizeValue(output.value)}`;
-        case "error-text": {
-            const firstLine = output.value.split("\n")[0] ?? "error";
-            return `${tag} ERROR - ${firstLine}`;
-        }
-        case "error-json": {
-            const errStr = JSON.stringify(output.value).slice(0, 200);
-            return `${tag} ERROR - ${errStr}`;
-        }
-        case "execution-denied":
-            return `${tag} execution denied${output.reason ? ` - ${output.reason}` : ""}`;
-        case "content": {
-            const textParts = output.value
-                .filter((p): p is Extract<typeof p, { type: "text" }> => p.type === "text")
-                .map((p) => p.text);
-            const joined = textParts.join(" ");
-            const preview = extractPreview(joined);
-            return `${tag} ${preview}`;
-        }
-        default:
-            return `${tag} completed`;
-    }
+    return content.flatMap((part) => (part.type === "text" ? [part.content] : [])).join(" ");
 }
 
-function extractPreview(text: string): string {
+function extractPreview(text: string) {
     // Try to parse as JSON for structured summary
     try {
         const parsed: unknown = JSON.parse(text);
@@ -141,124 +65,117 @@ function extractPreview(text: string): string {
     return text.slice(0, 147) + "...";
 }
 
+/**
+ * Compresses tool results from older steps in the prompt.
+ *
+ * On step N, tool results from steps 1..N-2 are aggressively compressed
+ * to short one-liners. Only the most recent step's results stay in full.
+ * Does not mutate the input — returns a new array.
+ */
+export function compressIntraTurnToolResults(
+    messages: ModelMessage[],
+    toolResultMaxTokens: number,
+) {
+    const { calls, steps } = toolStepsOf(messages);
+
+    // Below three there is nothing genuinely "old" to compress once the most
+    // recent step is kept intact.
+    if (steps.length < 3) return messages;
+
+    const result = [...messages];
+
+    for (const index of steps.slice(0, -1).flat()) {
+        const message = messages[index]!;
+        const text = contentText(message.content);
+        if (estimateTokens(text) <= toolResultMaxTokens) continue;
+
+        const tag = `[COMPRESSED prior ${toolNameOf(message, calls)} result]`;
+        result[index] = { ...message, content: `${tag} ${extractPreview(text)}` };
+    }
+
+    return result;
+}
+
 /**
  * Estimates the total token count of tool results in the prompt.
  */
-export function estimateToolResultTokens(prompt: LanguageModelV3Prompt): number {
+export function estimateToolResultTokens(messages: ModelMessage[]) {
     let total = 0;
-
-    for (const msg of prompt) {
-        if (msg.role !== "tool") continue;
-        for (const part of msg.content) {
-            if (part.type === "tool-result") {
-                total += estimateTokens(stringifyOutput(part.output));
-            }
-        }
+    for (const message of messages) {
+        if (message.role === "tool") total += estimateTokens(contentText(message.content));
     }
 
     return total;
 }
 
+const SUMMARIZER_SYSTEM_PROMPT = `You are a tool result summarizer. Given previous tool call results from an AI agent's workflow, produce a concise summary preserving:
+- Key data values, IDs, and counts
+- Success/failure status of each operation
+- Any error messages
+- Data that subsequent tool calls may need
+
+Maximum 300 tokens. Use bullet points. No preamble.`;
+
+function requestSummary(resultTexts: string[], modelConfig: SummarizationModelConfig) {
+    return chat({
+        adapter: createModel(modelConfig, { name: "summarization" }),
+        systemPrompts: [SUMMARIZER_SYSTEM_PROMPT],
+        messages: [
+            {
+                role: "user",
+                content: `Summarize these tool results:\n\n${resultTexts.join("\n\n")}`,
+            },
+        ],
+        stream: false,
+    });
+}
+
 /**
  * When accumulated tool result tokens exceed the threshold, summarizes
  * older step results using a lightweight LLM call.
  *
- * Replaces all tool results except the last two with a single synthetic
- * summary in the oldest tool message.
+ * Replaces every tool result before the last two steps with a single
+ * synthetic summary in the oldest tool message.
  */
 export async function summarizeOldStepResults(
-    prompt: LanguageModelV3Prompt,
+    messages: ModelMessage[],
     thresholdTokens: number,
     modelConfig: SummarizationModelConfig,
 ) {
-    const totalToolTokens = estimateToolResultTokens(prompt);
-    if (totalToolTokens <= thresholdTokens) return prompt;
+    if (estimateToolResultTokens(messages) <= thresholdTokens) return messages;
 
-    // Collect tool-role message indices
-    const toolMessageIndices: number[] = [];
-    for (let i = 0; i < prompt.length; i++) {
-        if (prompt[i]!.role === "tool") {
-            toolMessageIndices.push(i);
-        }
-    }
+    const { calls, steps } = toolStepsOf(messages);
 
-    // Keep the last 2 tool messages intact, summarize the rest
-    if (toolMessageIndices.length < 3) return prompt;
+    // Summarization keeps the last two steps intact — one more than compression
+    // does, because a summary loses more detail than a truncated one-liner.
+    if (steps.length < 3) return messages;
 
-    const indicesToSummarize = toolMessageIndices.slice(0, -2);
+    // An emptied tool message has nothing to contribute to the digest, and nothing
+    // to gain from being pointed at one.
+    const indices = steps
+        .slice(0, -2)
+        .flat()
+        .filter((index) => contentText(messages[index]!.content) !== "");
+    if (indices.length === 0) return messages;
 
-    // Collect tool result text from older steps
-    const resultTexts: string[] = [];
-    for (const idx of indicesToSummarize) {
-        const msg = prompt[idx] as ToolRoleMessage;
-        for (const part of msg.content) {
-            if (part.type === "tool-result") {
-                resultTexts.push(`[${part.toolName}]: ${stringifyOutput(part.output)}`);
-            }
-        }
-    }
-
-    if (resultTexts.length === 0) return prompt;
-
-    // Summarize with LLM
-    const provider = createOpenAICompatible({
-        name: modelConfig.providerName ?? "summarization-provider",
-        baseURL: modelConfig.baseURL,
-        apiKey: modelConfig.apiKey,
+    const resultTexts = indices.map((index) => {
+        const message = messages[index]!;
+        return `[${toolNameOf(message, calls)}]: ${contentText(message.content)}`;
     });
 
-    const model = provider.chatModel(modelConfig.modelName);
+    const summary = await requestSummary(resultTexts, modelConfig);
 
-    const { text: summary } = await generateText({
-        model,
-        system: `You are a tool result summarizer. Given previous tool call results from an AI agent's workflow, produce a concise summary preserving:
-- Key data values, IDs, and counts
-- Success/failure status of each operation
-- Any error messages
-- Data that subsequent tool calls may need
+    // The first old tool message carries the summary and the rest point at it, so
+    // the model reads the digest once rather than once per collapsed message.
+    const result = [...messages];
+    indices.forEach((index, position) => {
+        const content =
+            position === 0
+                ? `[Summary of ${indices.length} earlier tool results]:\n${summary}`
+                : "[see summary above]";
 
-Maximum 300 tokens. Use bullet points. No preamble.`,
-        prompt: `Summarize these tool results:\n\n${resultTexts.join("\n\n")}`,
+        result[index] = { ...messages[index]!, content };
     });
 
-    // Replace old tool messages with collapsed versions containing just the summary
-    const result = [...prompt];
-
-    for (let i = 0; i < indicesToSummarize.length; i++) {
-        const idx = indicesToSummarize[i]!;
-        const originalMsg = prompt[idx] as ToolRoleMessage;
-
-        if (i === 0) {
-            // First old tool message gets the summary
-            const summarizedContent = originalMsg.content.map((part) => {
-                if (part.type !== "tool-result") return part;
-                return {
-                    ...part,
-                    output: {
-                        type: "text" as const,
-                        value: `[Summary of ${resultTexts.length} earlier tool results]:\n${summary}`,
-                    },
-                };
-            });
-
-            // Keep only the first tool-result part with the summary
-            const firstResultIdx = summarizedContent.findIndex((p) => p.type === "tool-result");
-            const collapsedContent =
-                firstResultIdx >= 0 ? [summarizedContent[firstResultIdx]!] : summarizedContent;
-
-            result[idx] = { ...originalMsg, content: collapsedContent } as LanguageModelV3Message;
-        } else {
-            // Subsequent old tool messages get one-liners
-            const collapsedContent = originalMsg.content.map((part) => {
-                if (part.type !== "tool-result") return part;
-                return {
-                    ...part,
-                    output: { type: "text" as const, value: `[see summary above]` },
-                };
-            });
-            result[idx] = { ...originalMsg, content: collapsedContent } as LanguageModelV3Message;
-        }
-    }
-
     return result;
 }

package/src/lib/ai/context/optimize-thread-context.ts

@@ -0,0 +1,50 @@
+import type { ResolvedCortexAgentConfig } from "@/config";
+import type { ChatMessage, Thread } from "@/types";
+import type { ThreadContextMeta } from "./types";
+import { estimateMessagesTokens } from "./token-estimator";
+import { summarizeMessages } from "./summarizer";
+
+/**
+ * Post-response context optimization, called fire-and-forget from onFinish.
+ *
+ * Below the summarization threshold this only refreshes the token estimate —
+ * summarizing early would burn a model call to compress context that still fits.
+ */
+export async function optimizeThreadContext(
+    thread: Thread,
+    messages: ChatMessage[],
+    config: ResolvedCortexAgentConfig,
+) {
+    const contextConfig = config.context;
+    const estimates = estimateMessagesTokens(messages);
+    const totalEstimatedTokens = estimates.reduce((sum, e) => sum + e.tokens, 0);
+    const threshold = contextConfig.maxContextTokens * contextConfig.summarizationThreshold;
+
+    if (totalEstimatedTokens <= threshold) {
+        await config.db.threads.updateContextMeta(thread.id, {
+            summary: thread.contextMeta?.summary ?? null,
+            summaryUpToMessageId: thread.contextMeta?.summaryUpToMessageId ?? null,
+            totalEstimatedTokens,
+            lastOptimizedAt: new Date().toISOString(),
+        } satisfies ThreadContextMeta);
+        return;
+    }
+
+    // The most recent messages stay 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;
+
+    const summary = await summarizeMessages(
+        messagesToSummarize,
+        thread.contextMeta?.summary ?? null,
+        contextConfig.summarizationModel ?? config.model,
+    );
+
+    await config.db.threads.updateContextMeta(thread.id, {
+        summary,
+        summaryUpToMessageId: messagesToSummarize.at(-1)?.id ?? null,
+        totalEstimatedTokens,
+        lastOptimizedAt: new Date().toISOString(),
+    } satisfies ThreadContextMeta);
+}

package/src/lib/ai/context/summarizer.ts

@@ -1,29 +1,29 @@
-import { generateText } from "ai";
-import type { UIMessage } from "ai";
-import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
-import type { ContextConfig } from "./types.ts";
+import { chat } from "@tanstack/ai";
+import { createModel } from "@/ai/helpers";
+import type { ChatMessage } from "@/types";
+import type { ContextConfig } from "./types";
 
 type SummarizationModelConfig = NonNullable<ContextConfig["summarizationModel"]>;
 
-export async function summarizeMessages(
-    messages: UIMessage[],
+const SYSTEM_PROMPT = `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.`;
+
+export function summarizeMessages(
+    messages: ChatMessage[],
     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(" ")}`;
+            const text = msg.parts
+                .flatMap((part) => (part.type === "text" ? [part.content] : []))
+                .join(" ");
+            return `[${msg.role}]: ${text}`;
         })
         .join("\n");
 
@@ -31,20 +31,18 @@ export async function summarizeMessages(
         ? `\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}
+    return chat({
+        adapter: createModel(modelConfig, { name: "summarization" }),
+        systemPrompts: [SYSTEM_PROMPT],
+        messages: [
+            {
+                role: "user",
+                content: `${summaryContext}
 Summarize the following conversation:
 
 ${conversationText}`,
+            },
+        ],
+        stream: false,
     });
-
-    return text;
 }

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

@@ -1,4 +1,5 @@
-import type { UIMessage } from "ai";
+import type { MessagePart } from "@tanstack/ai";
+import type { ChatMessage } from "@/types";
 
 /** Average characters per token for English text. Used by the heuristic estimator. */
 export const CHARS_PER_TOKEN = 4;
@@ -12,45 +13,49 @@ export function estimateTokens(text: string) {
 }
 
 /**
- * Estimates token count for a single UIMessage by walking its parts.
+ * Media and UI parts carry a source, not text. What the provider actually bills
+ * for them is set by its own tokenizer — tiles for an image, pages for a document —
+ * and the length of a base64 payload says nothing useful about it, so they are
+ * charged a flat rate instead.
  */
-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;
+const OPAQUE_PART_TOKENS = 5;
+
+function estimatePartTokens(part: MessagePart) {
+    switch (part.type) {
+        case "text":
+        case "thinking":
+            return estimateTokens(part.content);
+        case "tool-call": {
+            // `arguments` is already the JSON the model emitted. `output` is only
+            // set for client-executed tools and for calls resolved after approval.
+            const output: unknown = part.output;
+            const outputTokens = output === undefined ? 0 : estimateTokens(JSON.stringify(output));
+
+            return estimateTokens(part.arguments) + outputTokens;
         }
+        case "tool-result":
+            return estimateTokens(
+                typeof part.content === "string" ? part.content : JSON.stringify(part.content),
+            );
+        case "structured-output":
+            return estimateTokens(part.raw);
+        default:
+            return OPAQUE_PART_TOKENS;
     }
+}
 
-    // Per-message overhead (role, metadata framing)
-    tokens += 4;
+export function estimateMessageTokens(message: ChatMessage) {
+    const parts = message.parts.reduce((sum, part) => sum + estimatePartTokens(part), 0);
 
-    return tokens;
+    // Per-message overhead (role, metadata framing)
+    return parts + 4;
 }
 
 /**
- * Estimates token counts for an array of UIMessages.
+ * Estimates token counts for an array of messages.
  * Returns per-message estimates in the same order.
  */
-export function estimateMessagesTokens(messages: UIMessage[]) {
+export function estimateMessagesTokens(messages: ChatMessage[]) {
     return messages.map(function (message) {
         return {
             message,

package/src/lib/ai/fetch.ts

@@ -1,4 +1,5 @@
-import type { ResolvedCortexAgentConfig } from "../config.ts";
+import type { ResolvedCortexAgentConfig } from "@/config";
+import { resolveBackendUrl } from "./backend-url";
 
 export async function fetchBackend(
     path: string,
@@ -14,18 +15,22 @@ export async function fetchBackend(
         finalOptions?.body &&
         typeof finalOptions.body === "string"
     ) {
+        let parsed: Record<string, unknown> | undefined;
         try {
-            const parsed = JSON.parse(finalOptions.body) as Record<string, unknown>;
+            parsed = JSON.parse(finalOptions.body) as Record<string, unknown>;
+        } catch {
+            // Non-JSON bodies pass through unchanged.
+        }
+
+        if (parsed !== undefined) {
             const transformed = await backendFetch.transformRequestBody(parsed, {
                 token,
             });
             finalOptions = { ...finalOptions, body: JSON.stringify(transformed) };
-        } catch {
-            // If body isn't valid JSON, pass through as-is
         }
     }
 
-    return fetch(`${backendFetch.baseUrl}${path}`, {
+    return fetch(resolveBackendUrl(backendFetch.baseUrl, path), {
         ...finalOptions,
         headers: {
             "Content-Type": "application/json",