@m6d/cortex-server 1.6.0 → 1.8.0

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

@@ -7,14 +7,14 @@ import {
     safeValidateUIMessages,
     stepCountIs,
     streamText,
+    wrapLanguageModel,
 } from "ai";
 import { HTTPException } from "hono/http-exception";
 import type { ResolvedCortexAgentConfig } from "../config.ts";
-import type { MessageMetadata, Thread } from "../types.ts";
+import type { MessageMetadata, Thread, TokenUsage } 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";
@@ -23,6 +23,11 @@ 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 {
+    compressIntraTurnToolResults,
+    estimateToolResultTokens,
+    summarizeOldStepResults,
+} from "./context/intra-turn-compressor.ts";
 import {
     registerStream,
     attachSseStream,
@@ -67,18 +72,20 @@ export async function stream(
 
     const model = createModel(config.model);
     const embeddingModel = createEmbeddingModel(config.embedding);
-    const neo4j = createNeo4jClient(config.neo4j, embeddingModel);
+    const neo4j = config.neo4j ? createNeo4jClient(config.neo4j, embeddingModel) : undefined;
 
     // 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,
-        }),
+        neo4j
+            ? resolveFromGraph(prompt, {
+                  neo4j,
+                  embeddingModel,
+                  reranker: config.reranker,
+              })
+            : Promise.resolve(null),
         // Branch C: Resolve session data
         resolveSession(config, thread, token),
     ]);
@@ -88,9 +95,12 @@ export async function stream(
     // Build tools
     const builtInTools: ToolSet = {
         captureFiles: captureFilesTool,
-        queryGraph: createQueryGraphTool(neo4j),
     };
 
+    if (neo4j) {
+        builtInTools["queryGraph"] = createQueryGraphTool(neo4j);
+    }
+
     if (config.backendFetch) {
         const backendFetchWithInterceptor = {
             ...config.backendFetch,
@@ -103,13 +113,17 @@ export async function stream(
                 ),
         };
 
-        builtInTools["callEndpoint"] = createCallEndpointTool(backendFetchWithInterceptor, token);
         builtInTools["executeCode"] = createExecuteCodeTool(backendFetchWithInterceptor, token);
     }
 
+    const resolvedUserTools =
+        typeof config.tools === "function"
+            ? config.tools({ thread, userId, token, session, requestContext })
+            : config.tools;
+
     const tools = {
         ...builtInTools,
-        ...config.tools,
+        ...resolvedUserTools,
     } as ToolSet;
 
     const systemPrompt = await buildSystemPrompt(config, resolved, {
@@ -129,13 +143,82 @@ export async function stream(
 
     const recentMessages = await convertToModelMessages(trimmedMessages);
 
+    const { context: contextConfig } = config;
+
+    const capturedSteps: {
+        prompt: string;
+        output: string | null;
+        tokenUsage: TokenUsage | null;
+    }[] = [];
+
     const result = streamText({
-        model,
+        model: wrapLanguageModel({
+            model,
+            middleware: {
+                specificationVersion: "v3",
+                transformParams: async ({ params }) => {
+                    // Strategy 1: Compress older step tool results to one-liners
+                    let optimizedPrompt = compressIntraTurnToolResults(
+                        params.prompt,
+                        contextConfig.toolResultMaxTokens,
+                    );
+
+                    // Strategy 5: If still over threshold, summarize with LLM
+                    const summarizationModel = contextConfig.summarizationModel ?? config.model;
+                    const toolTokens = estimateToolResultTokens(optimizedPrompt);
+
+                    if (toolTokens > contextConfig.intraTurnSummarizationThresholdTokens) {
+                        try {
+                            optimizedPrompt = await summarizeOldStepResults(
+                                optimizedPrompt,
+                                contextConfig.intraTurnSummarizationThresholdTokens,
+                                summarizationModel,
+                            );
+                        } catch (err) {
+                            console.error("[cortex-server] Intra-turn summarization failed:", err);
+                        }
+                    }
+
+                    return { ...params, prompt: optimizedPrompt };
+                },
+            },
+        }),
         system: systemPrompt,
         tools,
         messages: recentMessages,
         abortSignal: abortController.signal,
-        stopWhen: stepCountIs(50),
+        stopWhen: stepCountIs(25),
+        onStepFinish: (step) => {
+            const usage = step.usage;
+            capturedSteps.push({
+                prompt: JSON.stringify(step.request.body),
+                output: JSON.stringify({
+                    text: step.text,
+                    content: step.content,
+                    finishReason: step.finishReason,
+                    response: {
+                        id: step.response.id,
+                        modelId: step.response.modelId,
+                        messages: step.response.messages,
+                    },
+                    providerMetadata: step.providerMetadata,
+                }),
+                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,
+                },
+            });
+        },
     });
 
     return result.toUIMessageStreamResponse<UIMessage<MessageMetadata>>({
@@ -193,6 +276,24 @@ export async function stream(
             }
 
             await config.db.messages.upsert(thread.id, finishedMessages);
+
+            // Persist captured LLM request/response data per step
+            if (capturedSteps.length > 0 && lastAssistantMessage) {
+                const llmRequestInserts = capturedSteps.map((step, index) => ({
+                    messageId: lastAssistantMessage.id,
+                    stepNumber: index,
+                    prompt: step.prompt,
+                    output: step.output,
+                    tokenUsage: step.tokenUsage,
+                }));
+
+                try {
+                    await config.db.llmRequests.insert(llmRequestInserts);
+                } catch (err) {
+                    console.error("[cortex-server] Failed to persist AI requests:", err);
+                }
+            }
+
             config.onStreamFinish?.({ messages: finishedMessages, isAborted });
 
             // XXX: we need to notify the user so that the client can

package/src/lib/ai/prompt.ts

@@ -0,0 +1,354 @@
+import type { ResolvedContext, ResolvedEndpoint, ResolvedService } from "../graph/resolver.ts";
+import type { PromptContext, ResolvedCortexAgentConfig } from "../config.ts";
+import type { Thread } from "../types.ts";
+
+/**
+ * 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,
+    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,
+    promptContext: PromptContext,
+) {
+    // Resolve the consumer's base system prompt
+    let basePrompt: string;
+    if (typeof config.systemPrompt === "function") {
+        basePrompt = await config.systemPrompt(promptContext);
+    } else {
+        basePrompt = config.systemPrompt;
+    }
+
+    const sections: string[] = [basePrompt];
+
+    // Pre-resolved endpoints from knowledge graph
+    if (resolved) {
+        sections.push(buildResolvedSection(resolved));
+    }
+
+    return sections.join("\n");
+}
+
+// ---------------------------------------------------------------------------
+// Schema field type (parsed from endpoint JSON strings)
+// ---------------------------------------------------------------------------
+
+type SchemaField = {
+    name: string;
+    required?: boolean;
+    type?: string;
+    isArray?: boolean;
+    properties?: SchemaField[];
+};
+
+// ---------------------------------------------------------------------------
+// Layer 2: Schema compression
+// ---------------------------------------------------------------------------
+
+/** Compact param/body: `name?: type, nested?: {a: string, b: number}[]` */
+function compactParamFields(fields: SchemaField[]): string {
+    return fields
+        .map((f) => {
+            const opt = f.required ? "" : "?";
+            const arr = f.isArray ? "[]" : "";
+            if (f.properties && f.properties.length > 0) {
+                return `${f.name}${opt}: {${compactParamFields(f.properties)}}${arr}`;
+            }
+            return `${f.name}${opt}: ${f.type ?? "unknown"}${arr}`;
+        })
+        .join(", ");
+}
+
+function compactParams(jsonStr: string) {
+    try {
+        const fields = JSON.parse(jsonStr) as SchemaField[];
+        if (!Array.isArray(fields) || fields.length === 0) return "(none)";
+        return compactParamFields(fields);
+    } catch {
+        return jsonStr;
+    }
+}
+
+/** Compact response: `items[].{id, employee.{id, name.{ar, en}}}, count` */
+function compactResponseFields(fields: SchemaField[]): string {
+    return fields
+        .map((f) => {
+            const arr = f.isArray ? "[]" : "";
+            if (f.properties && f.properties.length > 0) {
+                return `${f.name}${arr}.{${compactResponseFields(f.properties)}}`;
+            }
+            return f.name;
+        })
+        .join(", ");
+}
+
+function compactResponse(jsonStr: string) {
+    try {
+        const fields = JSON.parse(jsonStr) as SchemaField[];
+        if (!Array.isArray(fields) || fields.length === 0) return "(none)";
+        return compactResponseFields(fields);
+    } catch {
+        return jsonStr;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Layer 1: Group endpoints by method + path
+// ---------------------------------------------------------------------------
+
+type GroupedEndpoint = Omit<ResolvedEndpoint, "concept"> & {
+    concepts: string[];
+};
+
+function groupEndpointsBySignature(endpoints: ResolvedEndpoint[]) {
+    const groups = new Map<string, GroupedEndpoint>();
+
+    for (const ep of endpoints) {
+        const key = `${ep.method}|${ep.path}`;
+        const existing = groups.get(key);
+
+        if (existing) {
+            if (!existing.concepts.includes(ep.concept)) {
+                existing.concepts.push(ep.concept);
+            }
+            for (const rule of ep.rules) {
+                if (!existing.rules.includes(rule)) {
+                    existing.rules.push(rule);
+                }
+            }
+            for (const dep of ep.dependencies) {
+                const isDuplicate = existing.dependencies.some(
+                    (d) => d.depPath === dep.depPath && d.paramName === dep.paramName,
+                );
+                if (!isDuplicate) {
+                    existing.dependencies.push(dep);
+                }
+            }
+        } else {
+            groups.set(key, {
+                concepts: [ep.concept],
+                relation: ep.relation,
+                name: ep.name,
+                path: ep.path,
+                method: ep.method,
+                params: ep.params,
+                body: ep.body,
+                response: ep.response,
+                dependencies: [...ep.dependencies],
+                rules: [...ep.rules],
+                metadata: ep.metadata,
+            });
+        }
+    }
+
+    return Array.from(groups.values());
+}
+
+// ---------------------------------------------------------------------------
+// Layer 3: Extract shared rules into a global section
+// ---------------------------------------------------------------------------
+
+function extractSharedRules(groups: GroupedEndpoint[], services: ResolvedService[]) {
+    const ruleCounts = new Map<string, number>();
+
+    for (const ep of groups) {
+        for (const rule of ep.rules) {
+            ruleCounts.set(rule, (ruleCounts.get(rule) ?? 0) + 1);
+        }
+    }
+    for (const svc of services) {
+        for (const rule of svc.rules) {
+            ruleCounts.set(rule, (ruleCounts.get(rule) ?? 0) + 1);
+        }
+    }
+
+    const sharedRules = [...ruleCounts.entries()]
+        .filter(([, count]) => count >= 2)
+        .map(([rule]) => rule);
+
+    const sharedSet = new Set(sharedRules);
+
+    return {
+        sharedRules,
+        groups: groups.map((ep) => ({
+            ...ep,
+            rules: ep.rules.filter((r) => !sharedSet.has(r)),
+        })),
+        services: services.map((svc) => ({
+            ...svc,
+            rules: svc.rules.filter((r) => !sharedSet.has(r)),
+        })),
+    };
+}
+
+// ---------------------------------------------------------------------------
+// Layer 4: Deduplicate identical response schemas across endpoints
+// ---------------------------------------------------------------------------
+
+type ResponseShapeRefs = {
+    shapeDefs: { name: string; content: string }[];
+    refMap: Map<string, string>;
+};
+
+function buildResponseShapeRefs(groups: GroupedEndpoint[]): ResponseShapeRefs {
+    // Key on the raw JSON string (byte-identical for endpoints sharing the same
+    // Endpoint node in Neo4j). Fall back to the compact string for cases where
+    // different raw JSON produces the same compact shape.
+    const rawToEndpoints = new Map<string, string[]>();
+
+    for (const ep of groups) {
+        if (ep.response === "[]") continue;
+        const existing = rawToEndpoints.get(ep.response);
+        if (existing) {
+            existing.push(`${ep.method} ${ep.path}`);
+        } else {
+            rawToEndpoints.set(ep.response, [`${ep.method} ${ep.path}`]);
+        }
+    }
+
+    // Second pass: also group by compact output to catch different raw JSON
+    // that normalizes to the same shape
+    const compactToRaws = new Map<string, Set<string>>();
+    for (const raw of rawToEndpoints.keys()) {
+        const compact = compactResponse(raw);
+        const existing = compactToRaws.get(compact);
+        if (existing) {
+            existing.add(raw);
+        } else {
+            compactToRaws.set(compact, new Set([raw]));
+        }
+    }
+
+    // Merge: a shape is shared if its raw string appears on 2+ endpoints,
+    // OR if multiple raw strings compact to the same output
+    const shapeDefs: { name: string; content: string }[] = [];
+    const refMap = new Map<string, string>();
+    let idx = 1;
+
+    for (const [compact, raws] of compactToRaws) {
+        let totalEndpoints = 0;
+        for (const raw of raws) {
+            totalEndpoints += rawToEndpoints.get(raw)!.length;
+        }
+
+        if (totalEndpoints >= 2) {
+            const name = `$R${idx++}`;
+            shapeDefs.push({ name, content: compact });
+            refMap.set(compact, name);
+        }
+    }
+
+    return { shapeDefs, refMap };
+}
+
+// ---------------------------------------------------------------------------
+// Pre-resolved section builder
+// ---------------------------------------------------------------------------
+
+function buildResolvedSection(resolved: ResolvedContext) {
+    const allEndpoints = [...resolved.readEndpoints, ...resolved.writeEndpoints];
+
+    if (allEndpoints.length === 0 && resolved.services.length === 0) {
+        return "\n## Pre-resolved API Endpoints\nNo matching endpoints found. Use queryGraph to search the knowledge graph manually.";
+    }
+
+    // Layer 1: deduplicate endpoints by method+path
+    const grouped = groupEndpointsBySignature(allEndpoints);
+
+    // Layer 3: hoist repeated rules to a shared section
+    const {
+        sharedRules,
+        groups: cleanedGroups,
+        services: cleanedServices,
+    } = extractSharedRules(grouped, resolved.services);
+
+    // Layer 4: deduplicate identical response schemas
+    const { shapeDefs, refMap } = buildResponseShapeRefs(cleanedGroups);
+
+    const parts: string[] = [
+        `
+## Pre-resolved API Endpoints
+The following endpoints were automatically matched to the user's message.`,
+    ];
+
+    // Shared rules
+    if (sharedRules.length > 0) {
+        parts.push("\n### General Rules");
+        for (const rule of sharedRules) {
+            parts.push(`\n- ${rule}`);
+        }
+    }
+
+    // Response shape definitions
+    if (shapeDefs.length > 0) {
+        parts.push("\n### Response Shapes");
+        for (const def of shapeDefs) {
+            parts.push(`\n${def.name}: ${def.content}`);
+        }
+    }
+
+    // Endpoints (Layer 2: compact schemas)
+    for (const ep of cleanedGroups) {
+        const concepts = ep.concepts.join(", ");
+        const params = compactParams(ep.params);
+        const body = compactParams(ep.body);
+        const compactResp = compactResponse(ep.response);
+        const response = refMap.get(compactResp) ?? compactResp;
+
+        const rules = ep.rules.length > 0 ? `\n  Rules: ${ep.rules.join("; ")}` : "";
+        const deps =
+            ep.dependencies.length > 0
+                ? "\n  Dependencies:\n" +
+                  ep.dependencies
+                      .map(
+                          (d) =>
+                              `    - Call ${d.depMethod} ${d.depPath} first → use its "${d.fromField}" as "${d.paramName}"`,
+                      )
+                      .join("\n")
+                : "";
+        const meta = ep.metadata !== "{}" ? `\n  Metadata: ${ep.metadata}` : "";
+
+        let block = `\n### ${ep.name} — ${ep.method} ${ep.path} (${ep.relation})`;
+        block += `\n- Concepts: ${concepts}`;
+        block += `\n- Params: ${params}`;
+        if (body !== "(none)") {
+            block += `\n- Body: ${body}`;
+        }
+        block += `\n- Response: ${response}`;
+        block += rules;
+        block += deps;
+        block += meta;
+
+        parts.push(block);
+    }
+
+    // Services
+    for (const svc of cleanedServices) {
+        const rules = svc.rules.length > 0 ? `\n  Rules: ${svc.rules.join("; ")}` : "";
+        const meta = svc.metadata !== "{}" ? `\n- Metadata: ${svc.metadata}` : "";
+        parts.push(
+            `\n### ${svc.concept} via ${svc.serviceName} (service)` +
+                `\n- Built-in ID: ${svc.builtInId}` +
+                `\n- Description: ${svc.description || "N/A"}${rules}${meta}`,
+        );
+    }
+
+    return parts.join("");
+}

package/src/{ai → lib/ai}/tools/execute-code.tool.ts

@@ -131,8 +131,11 @@ export function createExecuteCodeTool(
 ) {
     return tool({
         title: "Executes JavaScript code",
-        description:
-            "Run a JavaScript script that calls APIs and returns only relevant data. The script has an `api` helper: api.get(path, params?), api.post(path, body?), api.put(path, body?), api.del(path). Each returns parsed JSON and throws on error. For parameters of type file, use `capturedFile#[uploadId]` as the value.",
+        description: `Run a JavaScript script that calls APIs and returns only relevant data. 
+The script has an \`api\` helper: api.get(path, params?), api.post(path, body?), api.put(path, body?), api.del(path). 
+Each returns parsed JSON and throws on error. 
+For parameters of type file, use \`capturedFile#[uploadId]\` as the value. If you try to upload files in any other way
+you are going to get an error. All file uploads follow this rule without exception.`,
         inputSchema: z.object({
             code: z
                 .string()

package/src/{auth → lib/auth}/middleware.ts

@@ -6,6 +6,13 @@ import type { AppEnv, AuthedAppEnv } from "../types";
 import type { CortexConfig } from "../config";
 
 export function createUserLoaderMiddleware(authConfig: CortexConfig["auth"]) {
+    if (!authConfig) {
+        return createMiddleware<AppEnv>(async (c, next) => {
+            c.set("user", { id: "00000000-0000-0000-0000-000000000000", token: "" });
+            await next();
+        });
+    }
+
     const jwks = createRemoteJWKSet(new URL(authConfig.jwksUri));
 
     return createMiddleware<AppEnv>(async (c, next) => {