@m6d/cortex-server 1.6.0 → 1.8.0

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

@@ -0,0 +1,259 @@
+import type { UIMessage } from "ai";
+import { estimateTokens, CHARS_PER_TOKEN } from "./token-estimator.ts";
+import type { MessageMetadata } from "../../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;
+
+            const charBudget = maxTokensPerResult * CHARS_PER_TOKEN;
+            const truncatedOutput = smartStructuralCompress(part.output, charBudget);
+
+            return { ...part, output: truncatedOutput } as typeof part;
+        });
+
+        return { ...message, parts: compressedParts };
+    });
+}
+
+/**
+ * Structurally compresses a value to fit within a character budget.
+ * Understands JSON arrays, error objects, and nested objects — preserving
+ * semantic meaning better than naive string truncation.
+ */
+export function smartStructuralCompress(value: unknown, charBudget: number) {
+    if (typeof value === "string") {
+        return compressString(value, charBudget);
+    }
+
+    if (Array.isArray(value)) {
+        return compressArray(value, charBudget);
+    }
+
+    if (typeof value === "object" && value !== null) {
+        return compressObject(value as Record<string, unknown>, charBudget);
+    }
+
+    // Primitives
+    const str = String(value);
+    if (str.length <= charBudget) return str;
+    return str.slice(0, charBudget - 20) + "\n[...truncated]";
+}
+
+function compressString(value: string, charBudget: number) {
+    if (value.length <= charBudget) return value;
+
+    // Try to parse as JSON first — if it's a stringified structure, compress structurally
+    try {
+        const parsed: unknown = JSON.parse(value);
+        if (typeof parsed === "object" && parsed !== null) {
+            const result = smartStructuralCompress(parsed, charBudget) as string;
+            return result;
+        }
+    } catch {
+        // Not JSON, just truncate
+    }
+
+    return value.slice(0, charBudget - 20) + "\n[...truncated]";
+}
+
+function compressArray(value: unknown[], charBudget: number) {
+    const totalCount = value.length;
+
+    if (totalCount === 0) return "[]";
+
+    // Infer schema from first object element (if items are objects)
+    const firstObj = value.find(
+        (item): item is Record<string, unknown> =>
+            typeof item === "object" && item !== null && !Array.isArray(item),
+    );
+    const schemaKeys = firstObj ? Object.keys(firstObj) : null;
+
+    // Keep as many items as fit within ~60% of budget, leaving room for the summary footer
+    const itemBudget = Math.floor(charBudget * 0.6);
+    const items: unknown[] = [];
+    let accumulated = 0;
+
+    for (const item of value) {
+        const itemStr = JSON.stringify(item);
+        if (accumulated + itemStr.length > itemBudget && items.length > 0) break;
+        accumulated += itemStr.length;
+        items.push(item);
+    }
+
+    const shownStr = JSON.stringify(items, null, 2);
+    const remaining = totalCount - items.length;
+
+    if (remaining === 0 && shownStr.length <= charBudget) return shownStr;
+
+    const schemaHint = schemaKeys ? ` Schema: {${schemaKeys.join(", ")}}` : "";
+    const footer = `\n[COMPRESSED: ${remaining} more items of ${totalCount} total not shown.${schemaHint} Re-execute with filters or pagination to access specific items.]`;
+
+    const result = shownStr + footer;
+    if (result.length <= charBudget) return result;
+
+    // If still too large, fall back to a tighter representation
+    const tightStr = JSON.stringify(items);
+    return tightStr.slice(0, charBudget - footer.length - 20) + footer;
+}
+
+function compressObject(value: Record<string, unknown>, charBudget: number) {
+    // Check if it looks like an error — keep message, drop stack
+    if (isErrorLike(value)) {
+        return compressErrorObject(value, charBudget);
+    }
+
+    const fullStr = JSON.stringify(value, null, 2);
+    if (fullStr.length <= charBudget) return fullStr;
+
+    const allKeys = Object.keys(value);
+
+    // Phase 1: Partition keys into scalar (cheap) and complex (expensive)
+    const scalarEntries: string[] = [];
+    const complexKeys: string[] = [];
+    let scalarCharsUsed = 0;
+
+    for (const key of allKeys) {
+        const val = value[key];
+        const valStr = JSON.stringify(val);
+
+        if (valStr === undefined) {
+            const entry = `  "${key}": null`;
+            scalarEntries.push(entry);
+            scalarCharsUsed += entry.length + 2; // +2 for comma and newline
+        } else if (isScalar(val)) {
+            const entry = `  "${key}": ${valStr}`;
+            scalarEntries.push(entry);
+            scalarCharsUsed += entry.length + 2;
+        } else {
+            complexKeys.push(key);
+        }
+    }
+
+    // Phase 2: Distribute remaining budget to complex values
+    const compressionHeader = `  "__compressed": "Some values were compressed. If you need the full data for a specific key, re-execute the tool returning only that key."`;
+    const headerCost = compressionHeader.length + 2;
+    const structuralOverhead = 4; // { } and newlines
+    const remainingBudget = charBudget - scalarCharsUsed - headerCost - structuralOverhead;
+    const complexBudget = Math.floor(
+        Math.max(remainingBudget, 100) / Math.max(complexKeys.length, 1),
+    );
+
+    const complexEntries: string[] = [];
+    const truncatedKeys: string[] = [];
+
+    for (const key of complexKeys) {
+        const val = value[key];
+        const valStr = JSON.stringify(val);
+
+        if (valStr !== undefined && valStr.length <= complexBudget) {
+            complexEntries.push(`  "${key}": ${valStr}`);
+        } else {
+            truncatedKeys.push(key);
+            complexEntries.push(`  "${key}": ${summarizeValue(val)}`);
+        }
+    }
+
+    // Build result with compression notice
+    const allEntries = [...scalarEntries, ...complexEntries];
+
+    if (truncatedKeys.length > 0) {
+        allEntries.push(compressionHeader);
+        allEntries.push(`  "__truncatedKeys": ${JSON.stringify(truncatedKeys)}`);
+    }
+
+    const result = `{\n${allEntries.join(",\n")}\n}`;
+    if (result.length <= charBudget) return result;
+
+    // Final fallback: keep scalars + summaries, trim from the complex entries end
+    return result.slice(0, charBudget - 20) + "\n[...truncated]";
+}
+
+function isScalar(val: unknown): boolean {
+    return (
+        val === null ||
+        typeof val === "string" ||
+        typeof val === "number" ||
+        typeof val === "boolean"
+    );
+}
+
+function isErrorLike(value: Record<string, unknown>) {
+    return "error" in value || "message" in value || "stack" in value || "stackTrace" in value;
+}
+
+function compressErrorObject(value: Record<string, unknown>, charBudget: number) {
+    // Keep error/message fields, drop stack traces
+    const { stack: _stack, stackTrace: _stackTrace, ...rest } = value;
+    const compressed = JSON.stringify(rest, null, 2);
+
+    if (compressed.length <= charBudget) return compressed;
+    return compressed.slice(0, charBudget - 20) + "\n[...truncated]";
+}
+
+/**
+ * Produces a short human-readable summary of a value's shape and key data.
+ * Used both for inline object-key placeholders and for tool result one-liners.
+ */
+export function summarizeValue(val: unknown) {
+    if (Array.isArray(val)) {
+        const count = val.length;
+        const firstItem = val[0];
+        const schemaHint =
+            typeof firstItem === "object" && firstItem !== null
+                ? ` {${Object.keys(firstItem).slice(0, 5).join(", ")}}`
+                : "";
+        return `[Array(${count})${schemaHint}]`;
+    }
+
+    if (typeof val === "object" && val !== null) {
+        const obj = val as Record<string, unknown>;
+        const keys = Object.keys(obj);
+
+        // Detect common status/success response patterns
+        if ("status" in obj || "success" in obj) {
+            const status = obj["status"] ?? obj["success"];
+            const dataKey = keys.find((k) => k === "data" || k === "result" || k === "results");
+            const dataVal = dataKey ? obj[dataKey] : null;
+            const dataSummary = Array.isArray(dataVal) ? `, ${dataVal.length} items` : "";
+            return `{${String(status)}${dataSummary}}`;
+        }
+
+        return `{Object(${keys.length} keys: ${keys.slice(0, 5).join(", ")}${keys.length > 5 ? ", ..." : ""})}`;
+    }
+
+    if (typeof val === "string") {
+        return val.length > 80 ? `"${val.slice(0, 77)}..."` : JSON.stringify(val);
+    }
+
+    return String(val);
+}

package/src/{ai → lib/ai}/context/index.ts

@@ -6,9 +6,14 @@ export {
     estimateMessageTokens,
     estimateMessagesTokens,
 } from "./token-estimator.ts";
-export { compressToolResults } from "./compressor.ts";
+export { compressToolResults, smartStructuralCompress, summarizeValue } from "./compressor.ts";
 export { summarizeMessages } from "./summarizer.ts";
 export { buildContextMessages, trimMessagesToFit } from "./builder.ts";
+export {
+    compressIntraTurnToolResults,
+    estimateToolResultTokens,
+    summarizeOldStepResults,
+} from "./intra-turn-compressor.ts";
 
 import type { UIMessage } from "ai";
 import type { ResolvedCortexAgentConfig } from "../../config.ts";

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

@@ -0,0 +1,264 @@
+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" };
+
+type SummarizationModelConfig = NonNullable<ContextConfig["summarizationModel"]>;
+
+/**
+ * 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.
+ */
+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 }));
+
+    // 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);
+        });
+
+        result[idx] = { ...msg, content: compressedContent } as LanguageModelV3Message;
+    }
+
+    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 },
+    };
+}
+
+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 compressOutputToOneLiner(
+    output: LanguageModelV3ToolResultPart["output"],
+    toolName: string,
+): string {
+    const tag = `[COMPRESSED prior ${toolName} result]`;
+
+    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`;
+    }
+}
+
+function extractPreview(text: string): string {
+    // Try to parse as JSON for structured summary
+    try {
+        const parsed: unknown = JSON.parse(text);
+        if (typeof parsed === "object" && parsed !== null) {
+            return summarizeValue(parsed);
+        }
+    } catch {
+        // Not JSON
+    }
+
+    if (text.length <= 150) return text;
+    return text.slice(0, 147) + "...";
+}
+
+/**
+ * Estimates the total token count of tool results in the prompt.
+ */
+export function estimateToolResultTokens(prompt: LanguageModelV3Prompt): number {
+    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));
+            }
+        }
+    }
+
+    return total;
+}
+
+/**
+ * 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.
+ */
+export async function summarizeOldStepResults(
+    prompt: LanguageModelV3Prompt,
+    thresholdTokens: number,
+    modelConfig: SummarizationModelConfig,
+) {
+    const totalToolTokens = estimateToolResultTokens(prompt);
+    if (totalToolTokens <= thresholdTokens) return prompt;
+
+    // Collect tool-role message indices
+    const toolMessageIndices: number[] = [];
+    for (let i = 0; i < prompt.length; i++) {
+        if (prompt[i]!.role === "tool") {
+            toolMessageIndices.push(i);
+        }
+    }
+
+    // Keep the last 2 tool messages intact, summarize the rest
+    if (toolMessageIndices.length < 3) return prompt;
+
+    const indicesToSummarize = toolMessageIndices.slice(0, -2);
+
+    // 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 model = provider.chatModel(modelConfig.modelName);
+
+    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
+
+Maximum 300 tokens. Use bullet points. No preamble.`,
+        prompt: `Summarize these tool results:\n\n${resultTexts.join("\n\n")}`,
+    });
+
+    // 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/{ai → lib/ai}/context/types.ts

@@ -10,6 +10,8 @@ export type ContextConfig = {
     };
     toolResultMaxTokens: number;
     recentMessagesToKeep: number;
+    /** Token threshold for triggering intra-turn summarization of older tool results */
+    intraTurnSummarizationThresholdTokens: number;
 };
 
 export type ThreadContextMeta = {
@@ -23,6 +25,7 @@ export const DEFAULT_CONTEXT_CONFIG: ContextConfig = {
     maxContextTokens: 120_000,
     reservedTokenBudget: 8_000,
     summarizationThreshold: 0.75,
-    toolResultMaxTokens: 2_000,
+    toolResultMaxTokens: 1_000,
     recentMessagesToKeep: 6,
+    intraTurnSummarizationThresholdTokens: 20_000,
 };