memind/mcp-result-compactor.mjs

import crypto from 'node:crypto';

export const MCP_COMPACT_MODES = Object.freeze(['off', 'shadow', 'active']);
export const DEFAULT_MCP_COMPACT_BUDGET_CHARS = 4000;
export const DEFAULT_MCP_COMPACT_TTL_SECONDS = 3600;
export const MCP_COMPACT_HANDLE_PREFIX = 'mcpctx';

const memoryEntries = new Map();

function normalizeMode(value, fallback = 'off') {
  const raw = String(value ?? fallback).trim().toLowerCase();
  return MCP_COMPACT_MODES.includes(raw) ? raw : fallback;
}

export function resolveMcpCompactMode(env = process.env) {
  return normalizeMode(env.MEMIND_MCP_COMPACT_MODE, 'off');
}

export function isMcpCompactActive(env = process.env) {
  return resolveMcpCompactMode(env) === 'active';
}

export function resolveMcpCompactBudgetChars(env = process.env) {
  const raw = Number(env.MEMIND_MCP_COMPACT_BUDGET_CHARS ?? DEFAULT_MCP_COMPACT_BUDGET_CHARS);
  return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : DEFAULT_MCP_COMPACT_BUDGET_CHARS;
}

export function resolveMcpCompactTtlSeconds(env = process.env) {
  const raw = Number(env.MEMIND_MCP_COMPACT_TTL_SECONDS ?? DEFAULT_MCP_COMPACT_TTL_SECONDS);
  return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : DEFAULT_MCP_COMPACT_TTL_SECONDS;
}

function byteLength(text) {
  return Buffer.byteLength(String(text ?? ''), 'utf8');
}

function normalizePayload(payload) {
  if (payload == null) return '';
  if (typeof payload === 'string') return payload;
  try {
    return JSON.stringify(payload);
  } catch {
    return String(payload);
  }
}

export function buildCompactSummary(text, budgetChars) {
  const raw = String(text ?? '');
  const budget = Math.max(200, Number(budgetChars) || DEFAULT_MCP_COMPACT_BUDGET_CHARS);
  if (raw.length <= budget) return raw;

  const markerReserve = 96;
  const bodyBudget = Math.max(160, budget - markerReserve);
  const headLen = Math.floor(bodyBudget * 0.62);
  const tailLen = Math.max(48, bodyBudget - headLen);
  const omitted = Math.max(0, raw.length - headLen - tailLen);
  return [
    raw.slice(0, headLen),
    '',
    `...[memind compact omitted ${omitted} chars]...`,
    '',
    raw.slice(-tailLen),
  ].join('\n');
}

function formatCompactEnvelope({ summary, handle, stats }) {
  const savedPct = Math.round((stats.savedRatio ?? 0) * 100);
  return [
    summary,
    '',
    `[memind_ctx handle=${handle} saved=${savedPct}% raw_bytes=${stats.rawBytes}; use ctx_fetch to retrieve full payload]`,
  ].join('\n');
}

export function createMemoryCompactStore({ now = Date.now } = {}) {
  return {
    kind: 'memory',
    async put(text, { kind = 'generic', ttlSeconds = DEFAULT_MCP_COMPACT_TTL_SECONDS } = {}) {
      const handle = `${MCP_COMPACT_HANDLE_PREFIX}_${crypto.randomBytes(8).toString('hex')}`;
      const expiresAt = now() + ttlSeconds * 1000;
      memoryEntries.set(handle, {
        text: String(text ?? ''),
        kind,
        expiresAt,
      });
      return handle;
    },
    async get(handle) {
      const entry = memoryEntries.get(String(handle ?? '').trim());
      if (!entry) return null;
      if (entry.expiresAt <= now()) {
        memoryEntries.delete(String(handle ?? '').trim());
        return null;
      }
      return entry.text;
    },
    async clear() {
      memoryEntries.clear();
    },
  };
}

let defaultStore = createMemoryCompactStore();
let redisStorePromise = null;

export function resetMcpCompactStoresForTests() {
  memoryEntries.clear();
  defaultStore = createMemoryCompactStore();
  redisStorePromise = null;
}

export function setDefaultMcpCompactStore(store) {
  defaultStore = store ?? createMemoryCompactStore();
}

async function resolveDefaultStore(env = process.env) {
  const redisUrl = String(env.MEMIND_RUNTIME_REDIS_URL ?? env.REDIS_URL ?? '').trim();
  if (!redisUrl) return defaultStore;
  if (!redisStorePromise) {
    redisStorePromise = createRedisCompactStore({ url: redisUrl, env }).catch(() => defaultStore);
  }
  return redisStorePromise;
}

export async function createRedisCompactStore({
  url,
  env = process.env,
  keyPrefix = 'memind:mcp-compact:',
} = {}) {
  const { createClient } = await import('redis');
  const client = createClient({ url });
  client.on('error', () => {});
  await client.connect();
  const ttlSeconds = resolveMcpCompactTtlSeconds(env);
  return {
    kind: 'redis',
    async put(text, { kind = 'generic', ttlSeconds: itemTtl = ttlSeconds } = {}) {
      const handle = `${MCP_COMPACT_HANDLE_PREFIX}_${crypto.randomBytes(8).toString('hex')}`;
      const key = `${keyPrefix}${handle}`;
      await client.set(key, JSON.stringify({ text: String(text ?? ''), kind }), {
        EX: itemTtl,
      });
      return handle;
    },
    async get(handle) {
      const key = `${keyPrefix}${String(handle ?? '').trim()}`;
      const raw = await client.get(key);
      if (!raw) return null;
      try {
        const parsed = JSON.parse(raw);
        return parsed?.text ?? null;
      } catch {
        return null;
      }
    },
  };
}

export async function compact(payload, {
  kind = 'generic',
  budget,
  mode,
  env = process.env,
  store,
} = {}) {
  const resolvedMode = normalizeMode(mode ?? resolveMcpCompactMode(env), 'off');
  const resolvedBudget = budget ?? resolveMcpCompactBudgetChars(env);
  const text = normalizePayload(payload);
  const rawBytes = byteLength(text);

  if (resolvedMode === 'off' || text.length <= resolvedBudget) {
    return {
      text,
      compacted: false,
      mode: resolvedMode,
      stats: {
        rawBytes,
        sentBytes: rawBytes,
        savedRatio: 0,
        kind,
      },
    };
  }

  const summary = buildCompactSummary(text, resolvedBudget);
  const summaryBytes = byteLength(summary);
  const stats = {
    rawBytes,
    sentBytes: summaryBytes,
    savedRatio: rawBytes > 0 ? Number(((rawBytes - summaryBytes) / rawBytes).toFixed(4)) : 0,
    kind,
    mode: resolvedMode,
  };

  if (resolvedMode === 'shadow') {
    return {
      text,
      compacted: false,
      mode: 'shadow',
      stats,
      shadowSummary: summary,
    };
  }

  const resolvedStore = store ?? await resolveDefaultStore(env);
  let handle = null;
  try {
    handle = await resolvedStore.put(text, {
      kind,
      ttlSeconds: resolveMcpCompactTtlSeconds(env),
    });
  } catch {
    handle = null;
  }
  if (!handle) {
    return {
      text,
      compacted: false,
      mode: 'active',
      stats,
      fallback: 'store_unavailable',
    };
  }

  const compactText = formatCompactEnvelope({ summary, handle, stats });
  return {
    text: compactText,
    compacted: true,
    mode: 'active',
    handle,
    stats: {
      ...stats,
      sentBytes: byteLength(compactText),
      savedRatio: rawBytes > 0
        ? Number(((rawBytes - byteLength(compactText)) / rawBytes).toFixed(4))
        : 0,
    },
  };
}

export async function applyMcpCompaction(payload, { kind = 'generic', env = process.env, store } = {}) {
  return compact(payload, { kind, env, store });
}

export async function fetchCompactPayload(handle, { env = process.env, store } = {}) {
  const normalized = String(handle ?? '').trim();
  if (!normalized.startsWith(`${MCP_COMPACT_HANDLE_PREFIX}_`)) {
    return { ok: false, message: 'invalid handle' };
  }
  const resolvedStore = store ?? await resolveDefaultStore(env);
  try {
    const text = await resolvedStore.get(normalized);
    if (!text) return { ok: false, message: 'handle expired or not found' };
    return { ok: true, text, handle: normalized };
  } catch (err) {
    return {
      ok: false,
      message: err instanceof Error ? err.message : String(err),
    };
  }
}

export const ctxFetchToolDefinition = {
  name: 'ctx_fetch',
  description: 'Retrieve the full MCP tool payload previously compacted by Memind context runtime.',
  inputSchema: {
    type: 'object',
    properties: {
      handle: { type: 'string', description: 'Handle emitted by memind_ctx compact envelope.' },
    },
    required: ['handle'],
  },
};