diff --git a/src/cli/commands/mcp.ts b/src/cli/commands/mcp.ts index 4b96640bf..778d7075e 100644 --- a/src/cli/commands/mcp.ts +++ b/src/cli/commands/mcp.ts @@ -881,119 +881,179 @@ export class MCPCommandFactory { } const sdk = new NeuroLink(); - await this.getMCPStatusWithTimeout(sdk); - - // Check if server exists and is connected - const allServers = await sdk.listMCPServers(); - const server = allServers.find((s) => s.name === serverName); + let exitCode = 0; + try { + await this.getMCPStatusWithTimeout(sdk); - if (!server) { - if (spinner) { - spinner.fail(); - } - logger.error(chalk.red(`❌ Server not found: ${serverName}`)); - process.exit(1); - } + // Check if server exists and is connected + const allServers = await sdk.listMCPServers(); + const server = allServers.find((s) => s.name === serverName); - if (server.status !== "connected") { - if (spinner) { - spinner.fail(); + if (!server) { + if (spinner) { + spinner.fail(); + } + logger.error(chalk.red(`❌ Server not found: ${serverName}`)); + exitCode = 1; + } else if (server.status !== "connected") { + if (spinner) { + spinner.fail(); + } + logger.error(chalk.red(`❌ Server not connected: ${serverName}`)); + logger.always( + chalk.yellow("💡 Try: neurolink mcp test " + serverName), + ); + exitCode = 1; + } else { + // Check if tool exists + const tool = server.tools?.find((t) => t.name === toolName); + if (!tool) { + if (spinner) { + spinner.fail(); + } + logger.error(chalk.red(`❌ Tool not found: ${toolName}`)); + if (server.tools?.length) { + logger.always(chalk.blue("Available tools:")); + server.tools.forEach((t) => { + logger.always(` • ${t.name}: ${t.description}`); + }); + } + exitCode = 1; + } else { + exitCode = await this.runExecTool( + toolName, + serverName, + params, + argv, + spinner, + ); + } } - logger.error(chalk.red(`❌ Server not connected: ${serverName}`)); - logger.always(chalk.yellow("💡 Try: neurolink mcp test " + serverName)); - process.exit(1); - } + } finally { + // Always shut down MCP connections to prevent lingering child processes + await sdk.shutdown().catch(() => undefined); + } + + // Flush stdout/stderr before exit so buffered output (spinner lines, + // result text) is not truncated by process.exit(). process.exit() is + // still required because MCP stdio connections and health-check timers + // keep the Node.js event loop alive even after sdk.shutdown(). + await new Promise((resolve) => { + process.stdout.write("", () => { + process.stderr.write("", () => resolve()); + }); + }); + process.exit(exitCode); + } catch (_error) { + logger.error( + chalk.red(`❌ Exec command failed: ${(_error as Error).message}`), + ); + process.exit(1); + } + } - // Check if tool exists - const tool = server.tools?.find((t) => t.name === toolName); - if (!tool) { - if (spinner) { - spinner.fail(); - } - logger.error(chalk.red(`❌ Tool not found: ${toolName}`)); - if (server.tools?.length) { - logger.always(chalk.blue("Available tools:")); - server.tools.forEach((t) => { - logger.always(` • ${t.name}: ${t.description}`); - }); - } - process.exit(1); - } + /** + * Run the actual MCP tool execution and format/write the output. + * Extracted to keep executeExec nesting within ESLint max-depth limit. + * Returns 0 on success, 1 on failure. + */ + private static async runExecTool( + toolName: string, + serverName: string, + params: UnknownRecord, + argv: MCPCommandArgs, + spinner: ReturnType | null, + ): Promise { + const WARN_BYTES = 50 * 1024; // 50 KB + try { + const { toolRegistry } = await import("../../lib/mcp/toolRegistry.js"); + const executionResult = await withTimeout( + toolRegistry.executeTool(toolName, params, { + sessionId: `cli-${Date.now()}`, + userId: process.env.USER || "cli-user", + config: { domainType: "cli-execution", customData: { serverName } }, + }), + MCP_STATUS_TIMEOUT_MS, + ErrorFactory.toolTimeout(toolName, MCP_STATUS_TIMEOUT_MS), + ); - // Execute the tool using the NeuroLink MCP tool registry - try { - const { toolRegistry } = await import("../../lib/mcp/toolRegistry.js"); - const executionResult = await toolRegistry.executeTool( - toolName, - params, - { - sessionId: `cli-${Date.now()}`, - userId: process.env.USER || "cli-user", - config: { - domainType: "cli-execution", - customData: { serverName }, - }, - }, - ); + const result = { + tool: toolName, + server: serverName, + params, + result: executionResult, + success: true, + timestamp: new Date().toISOString(), + }; - const result = { - tool: toolName, - server: serverName, - params, - result: executionResult, - success: true, - timestamp: new Date().toISOString(), - }; + if (spinner) { + spinner.succeed(chalk.green("✅ Tool executed successfully")); + } - if (spinner) { - spinner.succeed(chalk.green("✅ Tool executed successfully")); - } + const resultJson = JSON.stringify(result, null, 2); + const resultBytes = Buffer.byteLength(resultJson, "utf-8"); + const outputFile = argv.output as string | undefined; - // Display results - if (argv.format === "json") { - logger.always(JSON.stringify(result, null, 2)); - } else { - logger.always(chalk.green("🔧 Tool Execution Results:")); - logger.always(` Tool: ${chalk.cyan(toolName)}`); - logger.always(` Server: ${chalk.cyan(serverName)}`); - logger.always( - ` Result: ${JSON.stringify(executionResult, null, 2)}`, + if (outputFile) { + await fs.promises.writeFile(outputFile, resultJson, "utf-8"); + logger.always( + chalk.green(`✅ Result saved to: ${chalk.cyan(outputFile)}`), + ); + logger.always( + ` Size: ${resultBytes >= 1024 ? `${(resultBytes / 1024).toFixed(1)} KB` : `${resultBytes} B`}`, + ); + } else if (argv.format === "json") { + if (resultBytes > WARN_BYTES) { + logger.warn( + chalk.yellow( + `⚠ Large result (${(resultBytes / 1024).toFixed(1)} KB). ` + + `Use --output to save to disk instead.`, + ), ); - logger.always(` Timestamp: ${result.timestamp}`); } - } catch (toolError) { - const errorMessage = - toolError instanceof Error ? toolError.message : String(toolError); - - if (spinner) { - spinner.fail(chalk.red("❌ Tool execution failed")); + logger.always(resultJson); + } else { + if (resultBytes > WARN_BYTES) { + logger.warn( + chalk.yellow( + `⚠ Large result (${(resultBytes / 1024).toFixed(1)} KB). ` + + `Use --output to save to disk instead.`, + ), + ); } + logger.always(chalk.green("🔧 Tool Execution Results:")); + logger.always(` Tool: ${chalk.cyan(toolName)}`); + logger.always(` Server: ${chalk.cyan(serverName)}`); + logger.always(` Result: ${JSON.stringify(executionResult, null, 2)}`); + logger.always(` Timestamp: ${result.timestamp}`); + } + return 0; + } catch (toolError) { + const errorMessage = + toolError instanceof Error ? toolError.message : String(toolError); - const result = { - tool: toolName, - server: serverName, - params, - _error: errorMessage, - success: false, - timestamp: new Date().toISOString(), - }; + if (spinner) { + spinner.fail(chalk.red("❌ Tool execution failed")); + } - if (argv.format === "json") { - logger.always(JSON.stringify(result, null, 2)); - } else { - logger.error(chalk.red("🔧 Tool Execution Failed:")); - logger.error(` Tool: ${chalk.cyan(toolName)}`); - logger.error(` Server: ${chalk.cyan(serverName)}`); - logger.error(` Error: ${chalk.red(errorMessage)}`); - } + const result = { + tool: toolName, + server: serverName, + params, + _error: errorMessage, + success: false, + timestamp: new Date().toISOString(), + }; - process.exit(1); + if (argv.format === "json") { + logger.always(JSON.stringify(result, null, 2)); + } else { + logger.error(chalk.red("🔧 Tool Execution Failed:")); + logger.error(` Tool: ${chalk.cyan(toolName)}`); + logger.error(` Server: ${chalk.cyan(serverName)}`); + logger.error(` Error: ${chalk.red(errorMessage)}`); } - } catch (_error) { - logger.error( - chalk.red(`❌ Exec command failed: ${(_error as Error).message}`), - ); - process.exit(1); + return 1; } } diff --git a/src/lib/artifacts/artifactStore.ts b/src/lib/artifacts/artifactStore.ts new file mode 100644 index 000000000..cc525b27f --- /dev/null +++ b/src/lib/artifacts/artifactStore.ts @@ -0,0 +1,184 @@ +/** + * Artifact Store + * + * Pluggable storage for externalized MCP tool outputs. + * + * When `mcp.outputLimits.strategy = "externalize"` the full tool payload is + * written here instead of being sent inline to the LLM. The model receives a + * compact surrogate with a preview and an artifact ID. The full payload can be + * retrieved on demand via the `retrieve_context` tool. + * + * Architecture: + * ArtifactStore (interface) — canonical types in src/lib/types/artifactTypes.ts + * LocalTempArtifactStore — single-process, filesystem-backed implementation + * + * Distributed backends (S3, Redis blobs) can be added later by implementing + * ArtifactStore from types/artifactTypes.ts. + * + * @module artifacts/artifactStore + */ + +import { randomUUID } from "node:crypto"; +import { mkdir, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { logger } from "../utils/logger.js"; +import type { + ArtifactMeta, + ArtifactRef, + ArtifactStore, +} from "../types/artifactTypes.js"; + +// Re-export so callers can import everything from one place +export type { + ArtifactMeta, + ArtifactRef, + ArtifactStore, +} from "../types/artifactTypes.js"; + +// --------------------------------------------------------------------------- +// LocalTempArtifactStore +// --------------------------------------------------------------------------- + +/** Characters used for the quick preview embedded in surrogate results. */ +const DEFAULT_PREVIEW_CHARS = 500; + +/** Index entry type (in-memory only). */ +type IndexEntry = ArtifactMeta & { path: string }; + +/** + * Filesystem-backed artifact store using the OS temp directory. + * + * Files are written with mode 0o600 (owner read/write only). + * An in-memory index tracks metadata without a separate index file. + * + * Suitable for: + * - CLI usage + * - Single-process SDK deployments + * - Multi-process deployments where each process manages its own artifacts + * (artifacts created in one process are not visible to others) + * + * @example + * ```typescript + * const store = new LocalTempArtifactStore(); + * const ref = await store.store(largeJson, { + * toolName: "list_files", + * serverId: "filesystem-server", + * sizeBytes: Buffer.byteLength(largeJson), + * contentType: "json", + * }); + * // Later, via retrieve_context: + * const full = await store.retrieve(ref.id); + * ``` + */ +export class LocalTempArtifactStore implements ArtifactStore { + private readonly dir: string; + private readonly index: Map = new Map(); + + constructor(dir?: string) { + this.dir = dir ?? join(tmpdir(), "neurolink-artifacts"); + } + + generatePreview(payload: string): string { + if (payload.length <= DEFAULT_PREVIEW_CHARS) { + return payload; + } + return `${payload.slice(0, DEFAULT_PREVIEW_CHARS)}…`; + } + + async store( + payload: string, + meta: Omit, + ): Promise { + await mkdir(this.dir, { recursive: true, mode: 0o700 }); + + const id = randomUUID(); + const ext = meta.contentType === "json" ? ".json" : ".txt"; + const filePath = join(this.dir, `${id}${ext}`); + + await writeFile(filePath, payload, { encoding: "utf-8", mode: 0o600 }); + + const fullMeta: IndexEntry = { + ...meta, + createdAt: Date.now(), + path: filePath, + }; + this.index.set(id, fullMeta); + + logger.debug( + `[ArtifactStore] Stored artifact ${id} for tool "${meta.toolName}" ` + + `(${formatBytes(meta.sizeBytes)})`, + ); + + return { + id, + preview: this.generatePreview(payload), + sizeBytes: meta.sizeBytes, + meta: { ...meta, createdAt: fullMeta.createdAt }, + }; + } + + async retrieve(id: string): Promise { + const entry = this.index.get(id); + if (!entry) { + logger.debug(`[ArtifactStore] Artifact ${id} not in index`); + return null; + } + try { + const content = await readFile(entry.path, "utf-8"); + logger.debug( + `[ArtifactStore] Retrieved artifact ${id} (${formatBytes(entry.sizeBytes)})`, + ); + return content; + } catch (err) { + logger.warn( + `[ArtifactStore] Failed to read artifact ${id}: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + return null; + } + } + + async delete(id: string): Promise { + const entry = this.index.get(id); + if (!entry) { + return; + } + try { + await rm(entry.path, { force: true }); + } catch { + // Suppress — file may already be gone + } + this.index.delete(id); + } + + async cleanup(olderThanMs: number): Promise { + const cutoff = Date.now() - olderThanMs; + let count = 0; + for (const [id, entry] of this.index.entries()) { + if (entry.createdAt < cutoff) { + await this.delete(id); + count++; + } + } + if (count > 0) { + logger.debug(`[ArtifactStore] Cleaned up ${count} expired artifact(s)`); + } + return count; + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function formatBytes(bytes: number): string { + if (bytes < 1024) { + return `${bytes} B`; + } + if (bytes < 1024 * 1024) { + return `${(bytes / 1024).toFixed(1)} KB`; + } + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +} diff --git a/src/lib/core/redisConversationMemoryManager.ts b/src/lib/core/redisConversationMemoryManager.ts index 6ec504645..fec36e5f2 100644 --- a/src/lib/core/redisConversationMemoryManager.ts +++ b/src/lib/core/redisConversationMemoryManager.ts @@ -8,6 +8,7 @@ import { tracers } from "../telemetry/tracers.js"; import { randomUUID } from "crypto"; import { MESSAGES_PER_TURN } from "../config/conversationMemory.js"; import { generateToolOutputPreview } from "../context/toolOutputLimits.js"; +import { NEUROLINK_ARTIFACT_ID_KEY } from "../mcp/mcpOutputNormalizer.js"; import { SummarizationEngine } from "../context/summarizationEngine.js"; import { NeuroLink } from "../neurolink.js"; import type { @@ -1863,11 +1864,32 @@ User message: "${userMessage}"`; }, ); + // Extract artifact ID if this result was externalized by McpOutputNormalizer. + // The surrogate carries `_meta.neurolinkArtifactId` on the raw result object. + let artifactId: string | undefined; + try { + const rawResult = toolResult.result; + if (rawResult && typeof rawResult === "object") { + const meta = (rawResult as Record)._meta; + if (meta && typeof meta === "object") { + const idValue = (meta as Record)[ + NEUROLINK_ARTIFACT_ID_KEY + ]; + if (typeof idValue === "string") { + artifactId = idValue; + } + } + } + } catch { + // Ignore extraction errors — artifact ID is best-effort metadata + } + // Build metadata — only store preview when truncation occurred (no duplication) const metadata: ChatMessageMetadata = { truncated, ...(truncated && { toolOutputPreview: preview }), ...(truncated && { originalSize }), + ...(artifactId && { artifactId }), }; // Build result — success/error metadata only, NOT the output data diff --git a/src/lib/mcp/externalServerManager.ts b/src/lib/mcp/externalServerManager.ts index 005a3ff03..b7d7802a0 100644 --- a/src/lib/mcp/externalServerManager.ts +++ b/src/lib/mcp/externalServerManager.ts @@ -275,6 +275,20 @@ export class ExternalServerManager extends EventEmitter { process.on("beforeExit", () => this.shutdown()); } + /** + * Attach a McpOutputNormalizer to the underlying ToolDiscoveryService. + * All tool outputs will be measured and (if oversized) replaced with compact + * surrogates before being returned to callers. + */ + setOutputNormalizer( + normalizer: import("./mcpOutputNormalizer.js").McpOutputNormalizer, + ): void { + this.toolDiscovery.setOutputNormalizer(normalizer); + mcpLogger.debug( + "[ExternalServerManager] MCP output normalizer attached to ToolDiscoveryService", + ); + } + /** * Set HITL manager for human-in-the-loop safety mechanisms * @param hitlManager - HITL manager instance (optional, can be undefined to disable) diff --git a/src/lib/mcp/mcpOutputNormalizer.ts b/src/lib/mcp/mcpOutputNormalizer.ts new file mode 100644 index 000000000..e73829427 --- /dev/null +++ b/src/lib/mcp/mcpOutputNormalizer.ts @@ -0,0 +1,236 @@ +/** + * MCP Output Normalizer + * + * Single responsibility: intercept a raw MCP `CallToolResult`, measure it, + * and apply the configured strategy so that oversized payloads never reach + * caches, Redis, or LLM context windows raw. + * + * Two strategies: + * - "inline" Pass through unchanged. The full payload enters the LLM + * context as-is. Emit a warning above warnBytes. + * - "externalize" Write the full payload to the ArtifactStore, return a + * compact surrogate with a head/tail preview and an artifact + * ID. The model uses `retrieve_context` with that ID to read + * the full output on demand, with offset/limit pagination. + * + * The surrogate result is shaped as an MCP `CallToolResult` so it passes + * transparently through any downstream code that expects that format. + * A `_meta` extension carries the artifact ID for structured extraction in + * `redisConversationMemoryManager`. + * + * @module mcp/mcpOutputNormalizer + */ + +import type { ArtifactStore, ArtifactRef } from "../artifacts/artifactStore.js"; +import { generateToolOutputPreview } from "../context/toolOutputLimits.js"; +import { logger } from "../utils/logger.js"; +import { withTimeout } from "../utils/errorHandling.js"; +import type { + McpOutputNormalizerConfig, + McpOutputContext, + NormalizedMcpOutput, +} from "../types/mcpOutputTypes.js"; + +// Re-export so callers can import everything from one place +export type { + McpOutputStrategy, + McpOutputNormalizerConfig, + McpOutputContext, + NormalizedMcpOutput, +} from "../types/mcpOutputTypes.js"; + +// --------------------------------------------------------------------------- +// Public constants +// --------------------------------------------------------------------------- + +/** Default byte ceiling above which externalize fires (100 KB). */ +export const DEFAULT_MAX_MCP_OUTPUT_BYTES = 100 * 1024; + +/** Default byte threshold for emitting a warning while still inline (50 KB). */ +export const DEFAULT_WARN_MCP_OUTPUT_BYTES = 50 * 1024; + +/** Metadata key embedded in surrogate `_meta` and used by memory manager. */ +export const NEUROLINK_ARTIFACT_ID_KEY = "neurolinkArtifactId"; + +// --------------------------------------------------------------------------- +// McpOutputNormalizer +// --------------------------------------------------------------------------- + +/** + * Stateless normalizer (state lives in the injected ArtifactStore). + * + * Construct once per NeuroLink instance and set via + * `ToolDiscoveryService.setOutputNormalizer()`. + */ +export class McpOutputNormalizer { + constructor( + private readonly config: McpOutputNormalizerConfig, + private readonly artifactStore?: ArtifactStore, + ) {} + + /** + * Measure `callResult`, apply strategy if oversized, return normalized output. + * + * Never throws: on any internal failure the raw result is returned unchanged + * with a warning log so tool execution is never broken by the normalizer. + */ + async normalize( + callResult: unknown, + context: McpOutputContext, + ): Promise { + const serialized = serialize(callResult); + const originalBytes = Buffer.byteLength(serialized, "utf-8"); + + // Fast path: below warn threshold — always inline, no logging + if (originalBytes <= this.config.warnBytes) { + return { result: callResult, isExternalized: false, originalBytes }; + } + + // Between warn and max: emit a warning but keep inline regardless of strategy + if (originalBytes <= this.config.maxBytes) { + logger.warn( + `[McpOutputNormalizer] Large MCP output from "${context.toolName}" on ` + + `"${context.serverId}" (${formatBytes(originalBytes)}). ` + + `Approaching limit of ${formatBytes(this.config.maxBytes)}.`, + { + toolName: context.toolName, + serverId: context.serverId, + originalBytes, + }, + ); + return { result: callResult, isExternalized: false, originalBytes }; + } + + // Above max — apply strategy + logger.warn( + `[McpOutputNormalizer] MCP output from "${context.toolName}" on ` + + `"${context.serverId}" exceeds limit ` + + `(${formatBytes(originalBytes)} > ${formatBytes(this.config.maxBytes)}). ` + + `Applying strategy "${this.config.strategy}".`, + { toolName: context.toolName, serverId: context.serverId, originalBytes }, + ); + + if (this.config.strategy === "inline") { + // Caller explicitly opted in to inline regardless of size + return { result: callResult, isExternalized: false, originalBytes }; + } + + // strategy === "externalize" + if (!this.artifactStore) { + // Misconfiguration: externalize was chosen but no store was provided. + // Pass through inline so execution is never broken, but log loudly. + logger.error( + `[McpOutputNormalizer] strategy="externalize" but no ArtifactStore ` + + `configured — passing through raw result for "${context.toolName}". ` + + `Set mcp.outputLimits.strategy="externalize" and ensure the NeuroLink ` + + `constructor creates a LocalTempArtifactStore.`, + ); + return { result: callResult, isExternalized: false, originalBytes }; + } + + let ref: ArtifactRef; + try { + ref = await withTimeout( + this.artifactStore.store(serialized, { + toolName: context.toolName, + serverId: context.serverId, + sessionId: context.sessionId, + sizeBytes: originalBytes, + contentType: isJsonLike(callResult) ? "json" : "text", + }), + 10_000, + new Error(`ArtifactStore.store() timed out for "${context.toolName}"`), + ); + } catch (err) { + // Storage failure or timeout — pass through inline so the call doesn't break. + logger.error( + `[McpOutputNormalizer] ArtifactStore.store() failed for ` + + `"${context.toolName}": ${err instanceof Error ? err.message : String(err)} ` + + `— passing through raw result.`, + ); + return { result: callResult, isExternalized: false, originalBytes }; + } + + // Generate a compact head/tail preview for the surrogate. + // Cap at warnBytes so the surrogate itself is always well within limits. + const { preview } = generateToolOutputPreview(serialized, { + maxBytes: Math.min(this.config.warnBytes, DEFAULT_WARN_MCP_OUTPUT_BYTES), + }); + + return { + result: buildSurrogate(preview, ref.id, context, originalBytes), + isExternalized: true, + artifactId: ref.id, + originalBytes, + }; + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** + * Build a compact MCP-shaped surrogate the LLM receives instead of the + * raw oversized payload. + * + * Shape mirrors `CallToolResult` so downstream code that inspects + * `result.content` keeps working unchanged. + * `_meta` carries the artifact ID for structured extraction in + * `redisConversationMemoryManager`. + */ +function buildSurrogate( + preview: string, + artifactId: string, + context: McpOutputContext, + originalBytes: number, +): unknown { + const text = + `[MCP Tool Output — ${context.toolName} | ${context.serverId}]\n` + + `Original size: ${formatBytes(originalBytes)} | ` + + `Externalized — use retrieve_context with artifactId="${artifactId}" ` + + `to read the full output (supports offset + limit pagination)\n` + + `\n--- Preview (head + tail) ---\n` + + preview + + `\n--- End Preview ---\n` + + `[${NEUROLINK_ARTIFACT_ID_KEY}=${artifactId}]`; + + return { + content: [{ type: "text" as const, text }], + _meta: { + [NEUROLINK_ARTIFACT_ID_KEY]: artifactId, + originalBytes, + toolName: context.toolName, + serverId: context.serverId, + }, + }; +} + +function serialize(value: unknown): string { + if (typeof value === "string") { + return value; + } + try { + // Compact JSON keeps byte measurement accurate and size enforcement honest. + // Pretty-printing with null, 2 inflates every object by ~30–50 % and would + // shift the externalization threshold relative to what the LLM actually + // receives if the payload is ever inlined. + return JSON.stringify(value); + } catch { + return String(value); + } +} + +function isJsonLike(value: unknown): boolean { + return typeof value === "object" && value !== null; +} + +function formatBytes(bytes: number): string { + if (bytes < 1024) { + return `${bytes} B`; + } + if (bytes < 1024 * 1024) { + return `${(bytes / 1024).toFixed(1)} KB`; + } + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +} diff --git a/src/lib/mcp/toolDiscoveryService.ts b/src/lib/mcp/toolDiscoveryService.ts index b280620dc..76c860d89 100644 --- a/src/lib/mcp/toolDiscoveryService.ts +++ b/src/lib/mcp/toolDiscoveryService.ts @@ -32,6 +32,7 @@ import { import { withTimeout } from "../utils/errorHandling.js"; import { SpanKind, SpanStatusCode } from "@opentelemetry/api"; import { tracers } from "../telemetry/tracers.js"; +import type { McpOutputNormalizer } from "./mcpOutputNormalizer.js"; const mcpTracer = tracers.mcp; @@ -54,10 +55,23 @@ export class ToolDiscoveryService extends EventEmitter { private serverTools = new Map>(); private discoveryInProgress = new Set(); + /** Optional normalizer applied to every tool output before it is returned. */ + private outputNormalizer?: McpOutputNormalizer; + constructor() { super(); } + /** + * Attach a McpOutputNormalizer. + * When set, every raw callTool() result is passed through the normalizer + * before being returned. Oversized outputs are replaced with compact + * surrogates according to the configured strategy. + */ + setOutputNormalizer(normalizer: McpOutputNormalizer): void { + this.outputNormalizer = normalizer; + } + /** * Discover tools from an external MCP server */ @@ -517,6 +531,38 @@ export class ToolDiscoveryService extends EventEmitter { new Error(`Tool execution timeout: ${toolName}`), ); callSpan.setStatus({ code: SpanStatusCode.OK }); + + // ── MCP output normalization ────────────────────────────────── + // Intercept here — after receive, before cache, before memory, + // before LLM context injection. Returns a compact surrogate when + // the payload exceeds mcp.outputLimits.maxBytes. + if (this.outputNormalizer) { + try { + const normalized = await this.outputNormalizer.normalize( + callResult, + { toolName, serverId }, + ); + callSpan.setAttribute( + "mcp.output.strategy", + normalized.isExternalized ? "externalize" : "inline", + ); + if (normalized.isExternalized) { + callSpan.setAttribute( + "mcp.output.original_bytes", + normalized.originalBytes, + ); + } + return normalized.result; + } catch (normErr) { + mcpLogger.warn( + `[ToolDiscoveryService] McpOutputNormalizer failed for ` + + `${toolName}: ${normErr instanceof Error ? normErr.message : String(normErr)} ` + + `— returning raw result`, + ); + } + } + // ── end normalization ───────────────────────────────────────── + return callResult; } catch (err) { callSpan.setStatus({ @@ -546,7 +592,7 @@ export class ToolDiscoveryService extends EventEmitter { `[ToolDiscoveryService] Tool execution completed: ${toolName}`, { duration, - hasContent: !!result.content, + hasContent: !!(result as { content?: unknown })?.content, }, ); diff --git a/src/lib/memory/memoryRetrievalTools.ts b/src/lib/memory/memoryRetrievalTools.ts index 13035ec44..5ae10dd96 100644 --- a/src/lib/memory/memoryRetrievalTools.ts +++ b/src/lib/memory/memoryRetrievalTools.ts @@ -10,7 +10,9 @@ import { tool } from "ai"; import { z } from "zod"; import type { RedisConversationMemoryManager } from "../core/redisConversationMemoryManager.js"; +import type { ArtifactStore } from "../artifacts/artifactStore.js"; import { logger } from "../utils/logger.js"; +import { withTimeout } from "../utils/errorHandling.js"; import { SpanSerializer, SpanType, @@ -29,21 +31,44 @@ const MAX_SEARCH_MATCHES = 50; /** * Factory function that creates memory retrieval tools bound to a memory manager. - * @param memoryManager - The Redis conversation memory manager instance + * + * @param memoryManager Redis conversation memory manager instance. + * @param artifactStore Optional artifact store for externalized MCP outputs. + * When provided, retrieve_context gains an `artifactId` + * parameter that fetches the full payload written by + * McpOutputNormalizer under strategy="externalize". * @returns Record of tool name to Vercel AI SDK tool definition */ export function createMemoryRetrievalTools( - memoryManager: RedisConversationMemoryManager, + memoryManager: RedisConversationMemoryManager | undefined, + artifactStore?: ArtifactStore, ) { return { retrieve_context: tool({ description: - "Retrieve messages from conversation memory. Use this to access full tool " + - "outputs when a result was truncated, review previous assistant responses, " + - "or search through conversation history. Supports filtering by role, " + - "pagination for large content, and regex search within messages.", + "Retrieve messages from conversation memory, or fetch the full payload of " + + "an externalized MCP tool output by artifact ID. Use this to:\n" + + "• Access full tool outputs when a result was truncated or externalized\n" + + "• Review previous assistant responses\n" + + "• Search through conversation history\n" + + "Supports filtering by role, pagination for large content, and regex search.\n" + + "To fetch an externalized artifact, provide `artifactId` (omit sessionId).", inputSchema: z.object({ - sessionId: z.string().describe("Session ID for the conversation"), + sessionId: z + .string() + .optional() + .describe( + "Session ID for conversation history retrieval. " + + "Required unless artifactId is provided.", + ), + artifactId: z + .string() + .optional() + .describe( + "Artifact ID from an externalized MCP tool output " + + "(visible in the tool output as neurolinkArtifactId=). " + + "When provided, returns the full stored payload directly.", + ), messageId: z .string() .optional() @@ -81,6 +106,68 @@ export function createMemoryRetrievalTools( ), }), execute: async (args) => { + // ── Artifact resolution path ──────────────────────────────────────── + // When the caller supplies an artifactId we short-circuit to the + // artifact store (bypassing Redis) and return the full payload with + // optional offset/limit pagination. + if (args.artifactId) { + if (!artifactStore) { + logger.warn( + "[MemoryRetrievalTools] retrieve_context called with artifactId " + + "but no ArtifactStore is configured", + ); + return { + error: + "Artifact store not configured — " + + "mcp.outputLimits.strategy must be set to 'externalize' to use artifactId retrieval", + artifactId: args.artifactId, + }; + } + const content = await withTimeout( + artifactStore.retrieve(args.artifactId), + 10_000, + new Error( + `ArtifactStore.retrieve() timed out for artifact "${args.artifactId}"`, + ), + ); + if (content === null) { + return { + error: "Artifact not found or has expired", + artifactId: args.artifactId, + }; + } + const charLimit = Math.min( + args.limit ?? DEFAULT_RETRIEVAL_LIMIT, + MAX_RETRIEVAL_LIMIT, + ); + const start = args.offset ?? 0; + const slice = content.slice(start, start + charLimit); + return { + artifactId: args.artifactId, + content: slice, + totalSize: content.length, + hasMore: start + charLimit < content.length, + offset: start, + limit: charLimit, + }; + } + // ── End artifact resolution ───────────────────────────────────────── + + if (!args.sessionId) { + return { + error: "sessionId is required when artifactId is not provided", + }; + } + + if (!memoryManager) { + return { + error: + "Session history retrieval requires Redis conversation memory — " + + "enable mcp.conversationMemory with a Redis backend, or use " + + "artifactId to retrieve an externalized MCP tool output.", + }; + } + const span = SpanSerializer.createSpan( SpanType.MEMORY, "memory.retrieve", @@ -92,15 +179,20 @@ export function createMemoryRetrievalTools( }, ); const startTime = Date.now(); + // args.sessionId is guaranteed non-null here — we returned early above + // when it was missing. Cast via string coercion to satisfy eslint. + const sessionId = String(args.sessionId); try { - const conversation = await memoryManager.getSessionRaw( - args.sessionId, + const conversation = await withTimeout( + memoryManager.getSessionRaw(sessionId), + 10_000, + new Error(`getSessionRaw() timed out for session "${sessionId}"`), ); if (!conversation) { span.durationMs = Date.now() - startTime; const endedSpan = SpanSerializer.endSpan(span, SpanStatus.OK); getMetricsAggregator().recordSpan(endedSpan); - return { error: "Session not found", sessionId: args.sessionId }; + return { error: "Session not found", sessionId }; } let messages = conversation.messages; diff --git a/src/lib/neurolink.ts b/src/lib/neurolink.ts index d146629d0..6cf6f4b74 100644 --- a/src/lib/neurolink.ts +++ b/src/lib/neurolink.ts @@ -62,6 +62,13 @@ import { ToolCallBatcher } from "./mcp/batching/index.js"; import { ToolResultCache } from "./mcp/caching/index.js"; import { EnhancedToolDiscovery } from "./mcp/enhancedToolDiscovery.js"; import { ExternalServerManager } from "./mcp/externalServerManager.js"; +import { + McpOutputNormalizer, + DEFAULT_MAX_MCP_OUTPUT_BYTES, + DEFAULT_WARN_MCP_OUTPUT_BYTES, +} from "./mcp/mcpOutputNormalizer.js"; +import { LocalTempArtifactStore } from "./artifacts/artifactStore.js"; +import type { ArtifactStore } from "./artifacts/artifactStore.js"; import type { MCPTool, RoutingDecision } from "./mcp/routing/index.js"; import { ToolRouter } from "./mcp/routing/index.js"; // Import direct tools server for automatic registration @@ -493,6 +500,8 @@ export class NeuroLink { private mcpToolBatcher?: ToolCallBatcher; private mcpEnhancedDiscovery?: EnhancedToolDiscovery; private mcpToolMiddlewares: ToolMiddleware[] = []; + /** Artifact store for externalized MCP tool outputs (set when strategy=externalize). */ + private mcpArtifactStore?: ArtifactStore; private _disableToolCacheForCurrentRequest = false; private mcpEnhancementsConfig?: MCPEnhancementsConfig; @@ -1321,6 +1330,33 @@ export class NeuroLink { } // ToolRouter — lazy-initialized when 2+ external servers exist (see addExternalMCPServer) + + // McpOutputNormalizer — active when mcp.outputLimits is configured + if (mcpConfig?.outputLimits) { + const strategy = mcpConfig.outputLimits.strategy ?? "externalize"; + const maxBytes = + mcpConfig.outputLimits.maxBytes ?? DEFAULT_MAX_MCP_OUTPUT_BYTES; + const warnBytes = + mcpConfig.outputLimits.warnBytes ?? DEFAULT_WARN_MCP_OUTPUT_BYTES; + + let artifactStore: ArtifactStore | undefined; + if (strategy === "externalize") { + artifactStore = new LocalTempArtifactStore(); + this.mcpArtifactStore = artifactStore; + logger.debug("[NeuroLink] MCP artifact store initialized (local-temp)"); + } + + const normalizer = new McpOutputNormalizer( + { strategy, maxBytes, warnBytes }, + artifactStore, + ); + this.externalServerManager.setOutputNormalizer(normalizer); + logger.debug("[NeuroLink] MCP output normalizer initialized", { + strategy, + maxBytes, + warnBytes, + }); + } } /** @@ -1464,108 +1500,73 @@ export class NeuroLink { "redis" in memConfig && !!(memConfig as { redis?: unknown }).redis) || process.env.STORAGE_TYPE === "redis"; - if (!memConfig?.enabled || !hasRedisConfig) { + const hasArtifactStore = !!this.mcpArtifactStore; + + // Register when Redis is configured OR when an artifact store exists. + // Artifact store alone is sufficient for the artifactId retrieval path — + // session history retrieval just returns a clear error when Redis is absent. + if ((!memConfig?.enabled || !hasRedisConfig) && !hasArtifactStore) { logger.debug( - "[NeuroLink] Skipping memory retrieval tools — requires Redis conversation memory", + "[NeuroLink] Skipping memory retrieval tools — requires Redis conversation memory or an artifact store", ); return; } - const tools = { - retrieve_context: { - description: - "Retrieve messages from conversation memory. Use this to access full tool " + - "outputs when a result was truncated, review previous assistant responses, " + - "or search through conversation history.", - execute: async (params: unknown) => { - // Lazy access: conversationMemory is initialized on first generate() call - const memoryManager = this.conversationMemory; - if (!memoryManager || !("getSessionRaw" in memoryManager)) { - return { - success: false, - error: - "Memory retrieval not available — Redis memory manager not initialized", - metadata: { - toolName: "retrieve_context", - serverId: "direct", - executionTime: 0, - }, - }; - } - - const actualTools = createMemoryRetrievalTools( - memoryManager as import("./core/redisConversationMemoryManager.js").RedisConversationMemoryManager, - ); - const result = await ( - actualTools.retrieve_context.execute as ( + // Extract the canonical tool definition (schema + description) from the + // memoryRetrievalTools factory. We pass undefined as the memoryManager here + // because we only need the Zod inputSchema and description at registration + // time — the actual manager is resolved lazily at execution time. + const canonicalTools = createMemoryRetrievalTools( + undefined, + this.mcpArtifactStore, + ); + const retrieveContextDef = canonicalTools.retrieve_context; + + // Register via this.registerTool() so the tool ends up in the "user-defined" + // category inside toolRegistry. getCustomTools() returns that category, which + // is what ToolsManager reads to build the tool schema sent to the LLM. + // (Tools registered via toolRegistry.registerTool() directly land in the + // "built-in" category and are never included in the LLM's tool schema.) + this.registerTool("retrieve_context", { + name: "retrieve_context", + description: + retrieveContextDef.description ?? "Retrieve context or artifacts", + // Pass the Zod schema so ToolsManager gives the LLM full parameter types. + // registerTool() detects isZodSchema on inputSchema and preserves it. + inputSchema: (retrieveContextDef as unknown as { inputSchema: object }) + .inputSchema, + execute: async (params: unknown) => { + // Lazy: conversationMemory is initialized on the first generate() call. + // When only an artifact store is present (no Redis), memoryManager is + // undefined — createMemoryRetrievalTools handles that via an explicit guard. + const memoryManager = this.conversationMemory as + | import("./core/redisConversationMemoryManager.js").RedisConversationMemoryManager + | undefined; + const tools = createMemoryRetrievalTools( + memoryManager, + this.mcpArtifactStore, + ); + // Return the result directly so the LLM receives clean output instead + // of a nested { success, data, metadata } wrapper. + // Bounded by TOOL_TIMEOUTS.EXECUTION_DEFAULT_MS so a stalled Redis or + // filesystem backend never hangs the tool call indefinitely. + return await withTimeout( + ( + tools.retrieve_context.execute as ( params: unknown, ctx: unknown, ) => Promise - )(params, { - toolCallId: "memory-retrieval", - messages: [], - }); - // Check if the tool itself reported an error - const hasError = - result && - typeof result === "object" && - "error" in result && - !("messages" in result); - const errorMsg = hasError - ? (result as { error: string }).error - : undefined; - return { - success: !hasError, - data: result, - ...(errorMsg ? { error: errorMsg } : {}), - metadata: { - toolName: "retrieve_context", - serverId: "direct", - executionTime: 0, - }, - }; - }, - }, - }; - - const registrations = Object.entries(tools).map( - async ([toolName, toolDef]) => { - const toolId = `direct.${toolName}`; - const toolInfo: ToolInfo = { - name: toolName, - description: toolDef.description, - inputSchema: {}, - serverId: "direct", - category: "built-in" as MCPServerCategory, - }; - - await this.toolRegistry.registerTool(toolId, toolInfo, { - execute: async (params: unknown) => { - try { - return await toolDef.execute(params); - } catch (error) { - // Known limitation: this non-throwing error path returns - // { success: false } without recording errorCategories in - // toolExecutionMetrics. These are internal memory-tool failures - // (low frequency), so the risk of metric gaps is minimal. - // A full fix would require access to the metrics map here, - // which is not available in the registration closure. - return { - success: false, - error: error instanceof Error ? error.message : String(error), - metadata: { toolName, serverId: "direct", executionTime: 0 }, - }; - } - }, - description: toolDef.description, - inputSchema: {}, - }); + )(params, { toolCallId: "memory-retrieval", messages: [] }), + TOOL_TIMEOUTS.EXECUTION_DEFAULT_MS, + ErrorFactory.toolTimeout( + "retrieve_context", + TOOL_TIMEOUTS.EXECUTION_DEFAULT_MS, + ), + ); }, - ); - - void Promise.all(registrations).then(() => { - logger.info("[NeuroLink] Memory retrieval tools registered"); }); + + logger.info("[NeuroLink] Memory retrieval tools registered"); } /** Format memory context for prompt inclusion */ diff --git a/src/lib/session/globalSessionState.ts b/src/lib/session/globalSessionState.ts index 44c9c19ad..95e8864f3 100644 --- a/src/lib/session/globalSessionState.ts +++ b/src/lib/session/globalSessionState.ts @@ -1,10 +1,44 @@ import { nanoid } from "nanoid"; import { NeuroLink } from "../neurolink.js"; -import type { - ConversationMemoryConfig, - NeurolinkOptions, -} from "../types/conversation.js"; +import type { ConversationMemoryConfig } from "../types/conversation.js"; +import type { NeurolinkConstructorConfig } from "../types/configTypes.js"; import { buildObservabilityConfigFromEnv } from "../utils/observabilityHelpers.js"; +import type { McpOutputStrategy } from "../types/mcpOutputTypes.js"; + +/** + * Build mcp.outputLimits config from environment variables. + * Reads NEUROLINK_MCP_OUTPUT_STRATEGY and NEUROLINK_MCP_MAX_OUTPUT_BYTES. + * Returns undefined when neither variable is set (no overhead). + */ +function buildMcpOutputLimitsFromEnv(): + | { strategy: McpOutputStrategy; maxBytes?: number; warnBytes?: number } + | undefined { + const strategyRaw = process.env.NEUROLINK_MCP_OUTPUT_STRATEGY; + const maxBytesRaw = process.env.NEUROLINK_MCP_MAX_OUTPUT_BYTES; + const warnBytesRaw = process.env.NEUROLINK_MCP_WARN_OUTPUT_BYTES; + + if (!strategyRaw && !maxBytesRaw) { + return undefined; + } + + const strategy: McpOutputStrategy = + strategyRaw === "inline" || strategyRaw === "externalize" + ? strategyRaw + : "externalize"; // safe default when only maxBytes is set + + const maxBytes = maxBytesRaw ? parseInt(maxBytesRaw, 10) : undefined; + const warnBytes = warnBytesRaw ? parseInt(warnBytesRaw, 10) : undefined; + + return { + strategy, + ...(maxBytes !== undefined && Number.isFinite(maxBytes) && maxBytes >= 0 + ? { maxBytes } + : {}), + ...(warnBytes !== undefined && Number.isFinite(warnBytes) && warnBytes >= 0 + ? { warnBytes } + : {}), + }; +} // Define a specific type for session variable values type SessionVariableValue = string | number | boolean; @@ -30,7 +64,7 @@ export class GlobalSessionManager { setLoopSession(config?: ConversationMemoryConfig): string { const sessionId = `NL_${nanoid()}`; - const neurolinkOptions: NeurolinkOptions = {}; + const neurolinkOptions: NeurolinkConstructorConfig = {}; if (config?.enabled) { neurolinkOptions.conversationMemory = { @@ -46,6 +80,15 @@ export class GlobalSessionManager { neurolinkOptions.observability = observabilityConfig; } + // Add MCP output limits from environment variables (CLI usage) + const mcpOutputLimits = buildMcpOutputLimitsFromEnv(); + if (mcpOutputLimits) { + neurolinkOptions.mcp = { + ...neurolinkOptions.mcp, + outputLimits: mcpOutputLimits, + }; + } + this.loopSession = { neurolinkInstance: new NeuroLink(neurolinkOptions), sessionId, @@ -142,9 +185,17 @@ export class GlobalSessionManager { // Create new NeuroLink with observability config from environment (CLI usage) const observabilityConfig = buildObservabilityConfigFromEnv(); - return new NeuroLink( - observabilityConfig ? { observability: observabilityConfig } : undefined, - ); + const mcpOutputLimits = buildMcpOutputLimitsFromEnv(); + + const options: NeurolinkConstructorConfig = {}; + if (observabilityConfig) { + options.observability = observabilityConfig; + } + if (mcpOutputLimits) { + options.mcp = { outputLimits: mcpOutputLimits }; + } + + return new NeuroLink(Object.keys(options).length ? options : undefined); } getCurrentSessionId(): string | undefined { diff --git a/src/lib/types/artifactTypes.ts b/src/lib/types/artifactTypes.ts new file mode 100644 index 000000000..138945400 --- /dev/null +++ b/src/lib/types/artifactTypes.ts @@ -0,0 +1,81 @@ +/** + * Artifact Store Types (canonical location) + * + * Types for the MCP large-output artifact storage system. + * When mcp.outputLimits.strategy = "externalize", oversized MCP tool outputs + * are stored as artifacts and the model receives a compact surrogate instead. + * + * @module types/artifactTypes + */ + +// --------------------------------------------------------------------------- +// Artifact metadata & reference +// --------------------------------------------------------------------------- + +/** Metadata recorded alongside a stored artifact. */ +export type ArtifactMeta = { + /** Tool name that produced the output. */ + toolName: string; + /** MCP server ID. */ + serverId: string; + /** Session that triggered the tool call (optional). */ + sessionId?: string; + /** Serialized byte size of the full payload. */ + sizeBytes: number; + /** Whether the payload is valid JSON or plain text. */ + contentType: "json" | "text"; + /** Unix epoch ms when the artifact was created. */ + createdAt: number; +}; + +/** Lightweight descriptor returned after a successful ArtifactStore.store(). */ +export type ArtifactRef = { + /** UUID v4 — stable identifier used in surrogate results and metadata. */ + id: string; + /** First N characters of the payload (for surrogate headers). */ + preview: string; + /** Full serialized byte size. */ + sizeBytes: number; + /** Stored metadata. */ + meta: ArtifactMeta; +}; + +// --------------------------------------------------------------------------- +// ArtifactStore interface +// --------------------------------------------------------------------------- + +/** + * Pluggable storage contract for externalized MCP tool outputs. + * + * Default backend: LocalTempArtifactStore (filesystem, single-process). + * Future backends can implement this interface for S3, Redis blobs, etc. + */ +export interface ArtifactStore { + /** + * Persist a payload and return a lightweight reference. + * @param payload Serialized tool output (JSON string or plain text). + * @param meta Descriptor without `createdAt` (assigned internally). + */ + store( + payload: string, + meta: Omit, + ): Promise; + + /** + * Retrieve the full payload by artifact ID. + * Returns `null` if the artifact is not found or has been cleaned up. + */ + retrieve(id: string): Promise; + + /** Delete a single artifact. No-op if the ID does not exist. */ + delete(id: string): Promise; + + /** + * Delete all artifacts older than `olderThanMs` milliseconds. + * Returns the number of artifacts deleted. + */ + cleanup(olderThanMs: number): Promise; + + /** Generate a short preview string from a serialized payload. */ + generatePreview(payload: string): string; +} diff --git a/src/lib/types/configTypes.ts b/src/lib/types/configTypes.ts index 323e970a6..265183d4b 100644 --- a/src/lib/types/configTypes.ts +++ b/src/lib/types/configTypes.ts @@ -109,6 +109,39 @@ export type MCPEnhancementsConfig = { }; /** Global tool middleware applied to every tool execution. Default: empty. */ middleware?: ToolMiddleware[]; + + /** + * Large MCP tool output handling. + * + * MCP servers can return arbitrarily large payloads. Without limits these + * are loaded entirely into memory, cached in full, stored whole in Redis, and + * injected into the LLM context window — all of which silently fail at scale. + * + * When configured, NeuroLink intercepts oversized outputs at the tool boundary + * (before caching and before memory persistence) and applies the chosen + * strategy so the model receives a compact surrogate instead of a firehose. + * + * Two strategies: + * - "inline" Keep sending the full payload to the model regardless of + * size. A warning is emitted above warnBytes. + * - "externalize" Store the full payload on disk as an artifact and return a + * compact surrogate with a head/tail preview and an artifact + * ID. The model uses `retrieve_context` with that ID to read + * the full output on demand, with offset/limit pagination. + * + * Defaults (when `outputLimits` is set): + * strategy = "externalize" + * maxBytes = 100 KB (100 * 1024) + * warnBytes = 50 KB (50 * 1024) + */ + outputLimits?: { + /** What to do when output exceeds maxBytes. Default: "externalize". */ + strategy?: "inline" | "externalize"; + /** Byte ceiling above which the strategy fires. Default: 102400 (100 KB). */ + maxBytes?: number; + /** Bytes at which a warning is emitted even when still inline. Default: 51200 (50 KB). */ + warnBytes?: number; + }; }; /** diff --git a/src/lib/types/conversation.ts b/src/lib/types/conversation.ts index e11c3f008..ca7c27c56 100644 --- a/src/lib/types/conversation.ts +++ b/src/lib/types/conversation.ts @@ -273,6 +273,13 @@ export type ChatMessageMetadata = { toolOutputPreview?: string; /** Original byte size of the full tool output before any truncation */ originalSize?: number; + /** + * Artifact store ID for an externalized MCP tool output. + * Set when `mcp.outputLimits.strategy = "externalize"` and the tool output + * exceeded `maxBytes`. Use retrieve_context with this ID to fetch the full + * payload from the local artifact store. + */ + artifactId?: string; }; /** diff --git a/src/lib/types/index.ts b/src/lib/types/index.ts index ddb1a1215..a315b87f2 100644 --- a/src/lib/types/index.ts +++ b/src/lib/types/index.ts @@ -36,6 +36,19 @@ export type { ExternalMCPToolInfo, ExternalMCPToolResult, } from "./externalMcp.js"; +// Artifact store types +export type { + ArtifactMeta, + ArtifactRef, + ArtifactStore, +} from "./artifactTypes.js"; +// MCP output normalizer types +export type { + McpOutputContext, + McpOutputNormalizerConfig, + McpOutputStrategy, + NormalizedMcpOutput, +} from "./mcpOutputTypes.js"; // MCP domain types export type { AuthorizationUrlResult, diff --git a/src/lib/types/mcpOutputTypes.ts b/src/lib/types/mcpOutputTypes.ts new file mode 100644 index 000000000..ea2994626 --- /dev/null +++ b/src/lib/types/mcpOutputTypes.ts @@ -0,0 +1,44 @@ +/** + * MCP Output Normalizer Types (canonical location) + * + * Types for the large MCP response handling system. + * + * @module types/mcpOutputTypes + */ + +/** + * Two honest strategies for oversized MCP tool outputs: + * - "inline" Full payload always sent to the model (warning logged above warnBytes). + * - "externalize" Full payload stored as an artifact; model receives a compact + * surrogate with head/tail preview and an artifact ID it can + * resolve via retrieve_context with offset/limit pagination. + */ +export type McpOutputStrategy = "inline" | "externalize"; + +/** Configuration for McpOutputNormalizer. */ +export type McpOutputNormalizerConfig = { + strategy: McpOutputStrategy; + /** Byte ceiling above which the strategy fires. */ + maxBytes: number; + /** Bytes at which a warning is emitted while still inline. */ + warnBytes: number; +}; + +/** Contextual info passed alongside the raw MCP callResult. */ +export type McpOutputContext = { + toolName: string; + serverId: string; + sessionId?: string; +}; + +/** Value returned by McpOutputNormalizer.normalize(). */ +export type NormalizedMcpOutput = { + /** The result to substitute for the raw callResult. May be a surrogate. */ + result: unknown; + /** Whether the full payload was written to the artifact store. */ + isExternalized: boolean; + /** Artifact ID when isExternalized === true. */ + artifactId?: string; + /** Serialized byte size of the original payload. */ + originalBytes: number; +}; diff --git a/test/continuous-test-suite-mcp-output-limits.ts b/test/continuous-test-suite-mcp-output-limits.ts new file mode 100644 index 000000000..cb45233f3 --- /dev/null +++ b/test/continuous-test-suite-mcp-output-limits.ts @@ -0,0 +1,511 @@ +#!/usr/bin/env tsx + +/** + * Continuous Test Suite — Large MCP Response Handling + * + * Tests the MCP output normalizer, artifact store, and retrieve_context + * integration without requiring a live AI provider. + * + * 5 sections: + * 1. ArtifactStore — store, retrieve, cleanup, permissions + * 2. McpOutputNormalizer — inline and externalize strategies + * 3. Surrogate result shape — MCP-compatible output, _meta fields + * 4. Error safety — normalizer never throws + * 5. retrieve_context artifact resolution — artifactId param + * + * Run with: npx tsx test/continuous-test-suite-mcp-output-limits.ts + */ + +import assert from "node:assert/strict"; +import { mkdtemp } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { rm } from "node:fs/promises"; + +import { LocalTempArtifactStore } from "../src/lib/artifacts/artifactStore.js"; +import { + McpOutputNormalizer, + DEFAULT_MAX_MCP_OUTPUT_BYTES, + DEFAULT_WARN_MCP_OUTPUT_BYTES, + NEUROLINK_ARTIFACT_ID_KEY, +} from "../src/lib/mcp/mcpOutputNormalizer.js"; + +// ─── helpers ──────────────────────────────────────────────────────────────── + +let passed = 0; +let failed = 0; +const errors: string[] = []; + +async function test(name: string, fn: () => Promise | void) { + try { + await fn(); + console.log(` ✓ ${name}`); + passed++; + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + console.log(` ✗ ${name}\n ${msg}`); + errors.push(`${name}: ${msg}`); + failed++; + } +} + +function makePayload(sizeBytes: number): string { + return "x".repeat(sizeBytes); +} + +function makeMcpResult(text: string): unknown { + return { content: [{ type: "text", text }] }; +} + +// ─── Test root ─────────────────────────────────────────────────────────────── +// All per-test artifact stores are nested under a single mkdtemp root so a +// single rm() at teardown cleans everything, with no leaks into tmpdir(). + +const testRoot = await mkdtemp(join(tmpdir(), "neurolink-test-")); + +/** Return a sub-directory path under the test root. */ +function underTestDir(name: string): string { + return join(testRoot, name); +} + +// ─── Section 1: ArtifactStore ─────────────────────────────────────────────── + +console.log("\nSection 1: ArtifactStore"); + +const testDir = underTestDir("s1-main"); +const store = new LocalTempArtifactStore(testDir); + +await test("store() returns ArtifactRef with correct metadata", async () => { + const payload = JSON.stringify({ key: "value", items: [1, 2, 3] }); + const ref = await store.store(payload, { + toolName: "test_tool", + serverId: "test-server", + sizeBytes: Buffer.byteLength(payload), + contentType: "json", + }); + + assert.ok(typeof ref.id === "string" && ref.id.length > 0, "id is set"); + assert.equal(ref.meta.toolName, "test_tool"); + assert.equal(ref.meta.serverId, "test-server"); + assert.equal(ref.meta.contentType, "json"); + assert.ok(ref.meta.createdAt > 0, "createdAt is set"); +}); + +await test("retrieve() returns the exact stored payload", async () => { + const payload = "hello world\nline 2\nline 3"; + const ref = await store.store(payload, { + toolName: "echo_tool", + serverId: "srv", + sizeBytes: Buffer.byteLength(payload), + contentType: "text", + }); + + const retrieved = await store.retrieve(ref.id); + assert.equal(retrieved, payload); +}); + +await test("retrieve() returns null for unknown id", async () => { + const result = await store.retrieve("00000000-0000-0000-0000-000000000000"); + assert.equal(result, null); +}); + +await test("delete() removes the artifact", async () => { + const payload = "to be deleted"; + const ref = await store.store(payload, { + toolName: "t", + serverId: "s", + sizeBytes: Buffer.byteLength(payload), + contentType: "text", + }); + + await store.delete(ref.id); + const result = await store.retrieve(ref.id); + assert.equal(result, null); +}); + +await test("cleanup() removes only expired artifacts", async () => { + const payload = "cleanup test"; + const ref = await store.store(payload, { + toolName: "t", + serverId: "s", + sizeBytes: Buffer.byteLength(payload), + contentType: "text", + }); + + // Cleanup with 1-hour TTL — the artifact was just created, should survive + const deleted = await store.cleanup(60 * 60 * 1000); + assert.equal(deleted, 0); + + // Cleanup with -1000ms TTL — cutoff is 1 second in the future, + // so every artifact (including ones just created) is treated as expired. + const deleted2 = await store.cleanup(-1000); + assert.ok(deleted2 >= 1, "at least one artifact cleaned up"); + + const result = await store.retrieve(ref.id); + assert.equal(result, null, "artifact was cleaned up"); +}); + +await test("generatePreview() truncates long payloads", () => { + const long = "a".repeat(2000); + const preview = store.generatePreview(long); + assert.ok(preview.length < long.length, "preview is shorter"); + assert.ok(preview.endsWith("…"), "preview ends with ellipsis"); +}); + +// ─── Section 2: McpOutputNormalizer strategies ────────────────────────────── + +console.log("\nSection 2: McpOutputNormalizer strategies"); + +const ctx = { toolName: "big_tool", serverId: "test-server" }; + +await test("small payload (below warnBytes) passes through inline with no logging", async () => { + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 1000, warnBytes: 500 }, + new LocalTempArtifactStore(underTestDir("s2a")), + ); + const payload = makeMcpResult("hello"); + const out = await normalizer.normalize(payload, ctx); + + assert.equal(out.isExternalized, false); + assert.deepEqual(out.result, payload); +}); + +await test("payload between warnBytes and maxBytes passes through inline", async () => { + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 1000, warnBytes: 100 }, + new LocalTempArtifactStore(underTestDir("s2b")), + ); + // ~200 bytes — above warnBytes but below maxBytes + const payload = makeMcpResult(makePayload(200)); + const out = await normalizer.normalize(payload, ctx); + + assert.equal(out.isExternalized, false); + assert.deepEqual(out.result, payload); +}); + +await test("inline strategy: oversized payload always passes through", async () => { + const normalizer = new McpOutputNormalizer({ + strategy: "inline", + maxBytes: 100, + warnBytes: 50, + }); + const payload = makeMcpResult(makePayload(5000)); + const out = await normalizer.normalize(payload, ctx); + + assert.equal(out.isExternalized, false); + assert.deepEqual(out.result, payload); +}); + +await test("externalize strategy: oversized payload is stored and surrogate returned", async () => { + const extStore = new LocalTempArtifactStore(underTestDir("s2c")); + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + extStore, + ); + const bigPayload = makePayload(500); + const out = await normalizer.normalize(makeMcpResult(bigPayload), ctx); + + assert.equal(out.isExternalized, true); + assert.ok(typeof out.artifactId === "string", "artifactId is set"); + + // Full payload is retrievable from the store + const stored = await extStore.retrieve(out.artifactId!); + assert.ok(stored !== null, "artifact is retrievable"); + assert.ok( + stored!.includes(bigPayload.slice(0, 100)), + "stored payload matches", + ); +}); + +await test("externalize without store: falls through inline + logs error", async () => { + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + // deliberately no artifact store + ); + const payload = makeMcpResult(makePayload(500)); + const out = await normalizer.normalize(payload, ctx); + + // Must not throw and must not lose the data + assert.equal(out.isExternalized, false); + assert.deepEqual(out.result, payload, "raw result returned on store failure"); +}); + +// ─── Section 3: Surrogate result shape ────────────────────────────────────── + +console.log("\nSection 3: Surrogate result shape"); + +await test("externalize surrogate is valid MCP CallToolResult shape", async () => { + const extStore = new LocalTempArtifactStore(underTestDir("s3a")); + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + extStore, + ); + const out = await normalizer.normalize(makeMcpResult(makePayload(500)), ctx); + + const surrogate = out.result as { + content: Array<{ type: string; text: string }>; + _meta: Record; + }; + assert.ok(Array.isArray(surrogate.content), "content is array"); + assert.equal(surrogate.content[0].type, "text"); + assert.ok(surrogate.content[0].text.includes("big_tool"), "toolName in text"); + assert.ok( + surrogate.content[0].text.includes("retrieve_context"), + "retrieve_context hint present", + ); +}); + +await test("externalize surrogate carries _meta with artifactId", async () => { + const extStore = new LocalTempArtifactStore(underTestDir("s3b")); + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + extStore, + ); + const out = await normalizer.normalize(makeMcpResult(makePayload(500)), ctx); + + const surrogate = out.result as Record; + const meta = surrogate._meta as Record; + assert.ok(meta, "_meta is present"); + assert.equal(typeof meta[NEUROLINK_ARTIFACT_ID_KEY], "string"); + assert.equal(meta[NEUROLINK_ARTIFACT_ID_KEY], out.artifactId); + assert.ok(typeof meta.originalBytes === "number"); + assert.equal(meta.toolName, "big_tool"); +}); + +await test("surrogate text embeds the artifact ID marker for memory manager extraction", async () => { + const extStore = new LocalTempArtifactStore(underTestDir("s3c")); + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + extStore, + ); + const out = await normalizer.normalize(makeMcpResult(makePayload(500)), ctx); + + const surrogate = out.result as { + content: Array<{ type: string; text: string }>; + }; + const text = surrogate.content[0].text; + assert.ok( + text.includes(`${NEUROLINK_ARTIFACT_ID_KEY}=${out.artifactId}`), + "artifact ID marker embedded in surrogate text", + ); +}); + +// ─── Section 4: Error safety ───────────────────────────────────────────────── + +console.log("\nSection 4: Error safety"); + +await test("normalize() does not throw on circular reference input", async () => { + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 1, warnBytes: 0 }, + new LocalTempArtifactStore(underTestDir("s4a")), + ); + + const circ: Record = { a: 1 }; + circ.self = circ; + + let threw = false; + try { + await normalizer.normalize(circ, ctx); + } catch { + threw = true; + } + assert.equal(threw, false, "normalize() did not throw"); +}); + +await test("normalize() returns raw result when artifact store throws", async () => { + // Store that always throws on store() + const brokenStore = { + store: async () => { + throw new Error("disk full"); + }, + retrieve: async () => null, + delete: async () => {}, + cleanup: async () => 0, + generatePreview: (p: string) => p.slice(0, 100), + }; + + const normalizer = new McpOutputNormalizer( + { strategy: "externalize", maxBytes: 100, warnBytes: 50 }, + brokenStore, + ); + const payload = makeMcpResult(makePayload(500)); + const out = await normalizer.normalize(payload, ctx); + + assert.equal(out.isExternalized, false); + assert.deepEqual(out.result, payload, "raw result returned when store fails"); +}); + +// ─── Section 5: retrieve_context artifact resolution ──────────────────────── + +console.log("\nSection 5: retrieve_context artifact resolution"); + +await test("retrieve() returns exact stored payload", async () => { + const s = new LocalTempArtifactStore(underTestDir("s5a")); + const payload = JSON.stringify({ result: "full data", rows: 100 }); + const ref = await s.store(payload, { + toolName: "query_db", + serverId: "db-server", + sizeBytes: Buffer.byteLength(payload), + contentType: "json", + }); + + const content = await s.retrieve(ref.id); + assert.equal(content, payload); +}); + +await test("retrieve() + slice simulates offset+limit pagination", async () => { + const s = new LocalTempArtifactStore(underTestDir("s5b")); + const payload = "abcdefghijklmnopqrstuvwxyz"; + const ref = await s.store(payload, { + toolName: "t", + serverId: "s", + sizeBytes: Buffer.byteLength(payload), + contentType: "text", + }); + + const content = await s.retrieve(ref.id); + assert.ok(content !== null); + // Simulate offset=5, limit=5 + assert.equal(content!.slice(5, 10), "fghij"); +}); + +await test("retrieve() returns null for missing artifact", async () => { + const s = new LocalTempArtifactStore(underTestDir("s5c")); + const result = await s.retrieve("non-existent-id"); + assert.equal(result, null); +}); + +// ─── Config constants ───────────────────────────────────────────────────────── + +console.log("\nConfig constants"); + +await test("default constants are correct", () => { + assert.equal(DEFAULT_MAX_MCP_OUTPUT_BYTES, 100 * 1024); + assert.equal(DEFAULT_WARN_MCP_OUTPUT_BYTES, 50 * 1024); +}); + +// ─── Section 6: retrieve_context without Redis ──────────────────────────────── +// Verifies Fix 1: artifact retrieval works with no memory manager (no Redis). + +console.log("\nSection 6: retrieve_context without Redis"); + +import { createMemoryRetrievalTools } from "../src/lib/memory/memoryRetrievalTools.js"; + +await test("retrieve_context with artifactId works without memory manager", async () => { + const s = new LocalTempArtifactStore(underTestDir("s6a")); + const payload = JSON.stringify({ key: "value", rows: [1, 2, 3] }); + const ref = await s.store(payload, { + toolName: "my_tool", + serverId: "srv", + sizeBytes: Buffer.byteLength(payload), + contentType: "json", + }); + + // createMemoryRetrievalTools with undefined manager — artifact path only + const tools = createMemoryRetrievalTools(undefined, s); + const execute = tools.retrieve_context.execute as ( + params: unknown, + ctx: unknown, + ) => Promise; + + const result = await execute( + { artifactId: ref.id }, + { toolCallId: "t1", messages: [] }, + ); + + const r = result as Record; + assert.equal(r.artifactId, ref.id); + assert.equal(r.content, payload); + assert.equal(r.totalSize, payload.length); + assert.equal(r.hasMore, false); +}); + +await test("retrieve_context with offset+limit pagination works without memory manager", async () => { + const s = new LocalTempArtifactStore(underTestDir("s6b")); + const payload = "abcdefghijklmnopqrstuvwxyz"; + const ref = await s.store(payload, { + toolName: "t", + serverId: "s", + sizeBytes: Buffer.byteLength(payload), + contentType: "text", + }); + + const tools = createMemoryRetrievalTools(undefined, s); + const execute = tools.retrieve_context.execute as ( + params: unknown, + ctx: unknown, + ) => Promise; + + const result = await execute( + { artifactId: ref.id, offset: 5, limit: 5 }, + { toolCallId: "t2", messages: [] }, + ); + + const r = result as Record; + assert.equal(r.content, "fghij"); + assert.equal(r.offset, 5); + assert.equal(r.limit, 5); + assert.equal(r.totalSize, 26); + assert.equal(r.hasMore, true); +}); + +await test("retrieve_context returns error for missing artifact without memory manager", async () => { + const s = new LocalTempArtifactStore(underTestDir("s6c")); + const tools = createMemoryRetrievalTools(undefined, s); + const execute = tools.retrieve_context.execute as ( + params: unknown, + ctx: unknown, + ) => Promise; + + const result = await execute( + { artifactId: "non-existent-id" }, + { toolCallId: "t3", messages: [] }, + ); + + const r = result as Record; + assert.ok(typeof r.error === "string", "error is reported"); +}); + +await test("retrieve_context with sessionId returns error when no memory manager", async () => { + const s = new LocalTempArtifactStore(underTestDir("s6d")); + const tools = createMemoryRetrievalTools(undefined, s); + const execute = tools.retrieve_context.execute as ( + params: unknown, + ctx: unknown, + ) => Promise; + + // sessionId without Redis — should get a clear error, not a crash + const result = await execute( + { sessionId: "some-session" }, + { toolCallId: "t4", messages: [] }, + ); + + const r = result as Record; + assert.ok(typeof r.error === "string", "error reported for missing manager"); + assert.ok( + (r.error as string).includes("Redis"), + "error mentions Redis requirement", + ); +}); + +// ─── Cleanup ───────────────────────────────────────────────────────────────── + +try { + await rm(testRoot, { recursive: true, force: true }); +} catch { + // ignore +} + +// ─── Summary ───────────────────────────────────────────────────────────────── + +console.log(`\n${"─".repeat(60)}`); +console.log(`Results: ${passed} passed, ${failed} failed`); +if (errors.length > 0) { + console.log("\nFailed tests:"); + for (const e of errors) { + console.log(` • ${e}`); + } + process.exit(1); +} else { + console.log("All tests passed."); +}