persistent-memory

/**
 * @file toolsProvider.ts
 * Registers all five memory tools with LM Studio.
 */

import { tool } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";
import { MemoryDatabase } from "./storage/db";
import { RetrievalEngine } from "./retrieval/engine";
import { extractFacts, detectConflicts } from "./processing/ai";
import {
  VALID_CATEGORIES,
  VALID_SCOPES,
  MAX_MEMORY_CONTENT_LENGTH,
  MAX_TAGS_PER_MEMORY,
  MAX_TAG_LENGTH,
  MAX_SEARCH_RESULTS,
  DEFAULT_SEARCH_RESULTS,
  MAX_SESSION_MEMORIES,
  MAX_PROJECT_NAME_LENGTH,
} from "./constants";
import type { PluginController } from "./pluginTypes";
import type { MemoryRecord, ScoredMemory } from "../types";

// ── Config reading (shared) ───────────────────────────────────────────

interface ResolvedConfig {
  autoInject: boolean;
  contextCount: number;
  enableAI: boolean;
  enableConflict: boolean;
  decayDays: number;
  storagePath: string;
}

function readConfig(ctl: PluginController): ResolvedConfig {
  const c = ctl.getPluginConfig(configSchematics);
  return {
    autoInject: c.get("autoInjectMemories") === "on",
    contextCount: Number(c.get("contextMemoryCount")) || 5,
    enableAI: c.get("enableAIExtraction") === "on",
    enableConflict: c.get("enableConflictDetection") === "on",
    decayDays: Number(c.get("decayHalfLifeDays")) || 30,
    storagePath: String(c.get("memoryStoragePath") || ""),
  };
}

// ── Shared singleton instances ──────────────────────────────────────────

let db: MemoryDatabase | null = null;
let engine: RetrievalEngine | null = null;
let currentPath = "";
let initPromise: Promise<{ db: MemoryDatabase; engine: RetrievalEngine }> | null = null;

async function ensureInitialized(storagePath: string): Promise<{ db: MemoryDatabase; engine: RetrievalEngine }> {
  const resolved = storagePath || "";
  if (db && engine && currentPath === resolved) return { db, engine };
  if (initPromise && currentPath === resolved) return initPromise;

  initPromise = (async () => {
    if (db) {
      try { db.close(); } catch { /* ignore */ }
    }
    db = new MemoryDatabase(resolved || undefined);
    await db.init();
    engine = new RetrievalEngine(db);
    engine.rebuildIndex();
    currentPath = resolved;
    return { db: db!, engine: engine! };
  })();

  return initPromise;
}

async function getSharedInstances(storagePath: string) {
  return ensureInitialized(storagePath);
}

// ── Session memory store (LRU eviction) ─────────────────────────────────

class SessionStore {
  private store = new Map<string, MemoryRecord>();

  get size(): number { return this.store.size; }

  add(content: string, category: MemoryCategory, tags: string[]): string {
    if (this.store.size >= MAX_SESSION_MEMORIES) {
      const oldestKey = this.store.keys().next().value;
      if (oldestKey) this.store.delete(oldestKey);
    }
    const now = Date.now();
    const id = `sess_${now}_${Math.random().toString(36).slice(2, 8)}`;
    this.store.set(id, {
      id, content, category, tags,
      confidence: 1.0, source: "tool-call", scope: "session",
      project: null, createdAt: now, updatedAt: now,
      lastAccessedAt: now, accessCount: 0, supersedes: null,
    });
    return id;
  }

  search(query: string): MemoryRecord[] {
    const lower = query.toLowerCase();
    return [...this.store.values()].filter(
      (m) => m.content.toLowerCase().includes(lower) || m.tags.some((t) => t.includes(lower)),
    );
  }

  delete(id: string): boolean { return this.store.delete(id); }

  clear(): void { this.store.clear(); }

  deleteByPattern(pattern: string): number {
    let count = 0;
    const lower = pattern.toLowerCase();
    for (const [id, mem] of this.store) {
      if (mem.content.toLowerCase().includes(lower)) {
        this.store.delete(id);
        count++;
      }
    }
    return count;
  }
}

const sessionStore = new SessionStore();

// ── Conflict resolution helper ──────────────────────────────────────────

async function handleConflicts(
  client: any, // Added client
  db: MemoryDatabase,
  engine: RetrievalEngine,
  content: string,
  confidence: number,
  tags: string[],
  cfg: ResolvedConfig,
): Promise<{ conflictInfo: { type: string; existingContent: string; resolution: string } | null }> {
  if (!cfg.enableConflict) return { conflictInfo: null };

  try {
    const existing = await engine.retrieve(content, 5, cfg.decayDays); // Added await
    const others = existing.memories.filter((m) => m.id !== content);
    if (others.length === 0) return { conflictInfo: null };

    const conflicts = await detectConflicts(client, content, others); // Passed client
    let conflictInfo: { type: string; existingContent: string; resolution: string } | null = null;

    for (const conflict of conflicts) {
      if (conflict.resolution === "skip" || conflict.conflictType === "duplicate") {
        return { conflictInfo: null };
      }
      if (conflict.conflictType === "contradiction") {
        conflictInfo = {
          type: "contradiction",
          existingContent: conflict.existingContent,
          resolution: "Both memories kept — you may want to resolve this",
        };
      }
    }
    return { conflictInfo };
  } catch {
    return { conflictInfo: null };
  }
}

// ── Tool implementations ────────────────────────────────────────────────

async function rememberImpl(
  params: {
    content: string; category: MemoryCategory; tags?: string[];
    confidence?: number; scope?: MemoryScope; project?: string;
  },
  ctl: PluginController,
): Promise<unknown> {
  const cfg = readConfig(ctl);
  const { db, engine } = await ensureInitialized(cfg.storagePath);
  const { content, category, tags, confidence, scope, project } = params;
  const memScope: MemoryScope = scope ?? "global";

  if (memScope === "project" && !project) {
    return { stored: false, error: "scope='project' requires a project name" };
  }

  if (memScope === "session") {
    const id = sessionStore.add(content, category, tags ?? []);
    return { stored: true, id, scope: "session", content, category, tags: tags ?? [] };
  }

  try {
    const id = db.store(
      content, category, tags ?? [], confidence ?? 1.0,
      "tool-call", null, memScope, project ?? null,
    );
    engine.indexMemory(id, content, tags ?? [], category);

    const conflictResult = await handleConflicts(ctl, db, engine, content, confidence ?? 1.0, tags ?? [], cfg); // Passed ctl

    return {
      stored: true, id, scope: memScope, content, category,
      tags: tags ?? [],
      ...(project ? { project } : {}),
      ...(conflictResult.conflictInfo ? { conflict: conflictResult.conflictInfo } : {}),
    };
  } catch (err) {
    const msg = err instanceof Error ? err.message : String(err);
    return { stored: false, error: msg };
  }
}

async function recallImpl(
  params: {
    query: string; limit?: number; category?: MemoryCategory;
    scope?: MemoryScope; project?: string;
  },
  ctl: PluginController,
): Promise<unknown> {
  const cfg = readConfig(ctl);
  const { db, engine } = await ensureInitialized(cfg.storagePath);
  const { query, limit, category, scope, project } = params;
  const maxResults = limit ?? DEFAULT_SEARCH_RESULTS;

  const result = cfg.enableAI
    ? await engine.retrieveWithSRLM(query, maxResults, cfg.decayDays, 3)
    : engine.retrieve(query, maxResults, cfg.decayDays);

  let memories: Array<ScoredMemory | MemoryRecord> = [...result.memories];
  const sessionHits = sessionStore.search(query);
  if (sessionHits.length > 0) memories.push(...sessionHits);

  if (category) memories = memories.filter((m) => m.category === category);
  if (scope) memories = memories.filter((m) => m.scope === scope);
  if (project) memories = memories.filter((m) => m.project === project);

  memories.sort((a, b) => {
    const scoreA = "compositeScore" in a ? a.compositeScore : 0.5;
    const scoreB = "compositeScore" in b ? b.compositeScore : 0.5;
    return scoreB - scoreA;
  });
  memories = memories.slice(0, maxResults);

  if (memories.length === 0) {
    return { found: 0, memories: [], suggestion: "No memories match this query." };
  }

  return {
    found: memories.length,
    totalMatched: result.totalMatched + sessionHits.length,
    memories: memories.map((m) => ({
      id: m.id, content: m.content, category: m.category, tags: m.tags,
      confidence: m.confidence, scope: m.scope,
      ...(m.project ? { project: m.project } : {}),
      relevance: "compositeScore" in m ? Math.round(m.compositeScore * 100) : 50,
      lastAccessed: new Date(m.lastAccessedAt).toISOString(),
      accessCount: m.accessCount,
    })),
  };
}

async function searchImpl(
  params: {
    query?: string; category?: MemoryCategory; tag?: string;
    recent?: number; limit?: number;
  },
  ctl: PluginController,
): Promise<unknown> {
  const cfg = readConfig(ctl);
  const { db, engine } = await ensureInitialized(cfg.storagePath);
  const { query, category, tag, recent, limit } = params;
  const maxResults = limit ?? DEFAULT_SEARCH_RESULTS;

  if (recent) {
    const memories = db.getRecent(recent);
    return { mode: "recent", found: memories.length, memories: memories.map(normalizeSearchResult) };
  }

  if (tag && !query) {
    const memories = db.getByTag(tag, maxResults);
    return { mode: "tag-filter", tag, found: memories.length, memories: memories.map(normalizeSearchResult) };
  }

  if (category && !query) {
    const memories = db.getByCategory(category, maxResults);
    return { mode: "category-filter", category, found: memories.length, memories: memories.map(normalizeSearchResult) };
  }

  if (query) {
    const result = engine.retrieve(query, maxResults, cfg.decayDays);
    let memories = result.memories;
    if (category) memories = memories.filter((m) => m.category === category);
    if (tag) memories = memories.filter((m) => m.tags.includes(tag.toLowerCase()));
    return {
      mode: "semantic", found: memories.length,
      memories: memories.map((m) => ({
        id: m.id, content: m.content, category: m.category, tags: m.tags,
        relevance: Math.round(m.compositeScore * 100),
      })),
    };
  }

  const memories = db.getRecent(maxResults);
  return { mode: "recent-fallback", found: memories.length, memories: memories.map(normalizeSearchResult) };
}

function normalizeSearchResult(m: MemoryRecord) {
  return { id: m.id, content: m.content, category: m.category, tags: m.tags };
}

async function forgetImpl(
  params: { id?: string; pattern?: string; deleteAll?: boolean },
  ctl: PluginController,
): Promise<unknown> {
  const { id, pattern, deleteAll } = params;

  if (deleteAll === true) {
    const dbCount = db!.deleteAll();
    const sessCount = sessionStore.size;
    sessionStore.clear();
    engine!.rebuildIndex();
    return { deleted: dbCount + sessCount, mode: "delete-all" };
  }

  if (id) {
    if (sessionStore.delete(id)) return { deleted: 1, id, scope: "session" as const };
    const existed = db!.delete(id);
    if (existed) {
      engine!.removeFromIndex(id);
      return { deleted: 1, id };
    }
    return { deleted: 0, error: "Memory ID not found" };
  }

  if (pattern) {
    const sessDeleted = sessionStore.deleteByPattern(pattern);
    const dbDeleted = db!.deleteByPattern(pattern);
    if (dbDeleted > 0) engine!.rebuildIndex();
    return { deleted: dbDeleted + sessDeleted, pattern };
  }

  return { deleted: 0, error: "Specify an id, pattern, or set deleteAll: true" };
}

async function statusImpl(_params: object): Promise<unknown> {
  const stats = db!.getStats();
  const idxStats = engine!.indexStats;
  return {
    ...stats,
    sessionMemories: sessionStore.size,
    dbSizeBytes: Math.round(stats.dbSizeBytes / 1024),
    indexVocabularySize: idxStats.vocabSize,
    indexedDocuments: idxStats.docCount,
    runtime: {
      nodeVersion: process.version,
      platform: process.platform,
      arch: process.arch,
      moduleVersion: process.versions.modules,
    },
    scopes: {
      global: "Persistent across all conversations",
      project: "Persistent, filtered by project name",
      session: `Temporary, in-memory only (${sessionStore.size} active)`,
    },
  };
}

// ── Exported provider ────────────────────────────────────────────────────

export async function toolsProvider(ctl: PluginController) {
  const cfg = readConfig(ctl);

  return [
    tool({
      name: "Remember",
      description:
        `Store a fact, preference, project detail, or note in memory.\\n\\n` +
        `SCOPES:\\n` +
        `• "global" (default) — persists forever across ALL conversations\\n` +
        `• "project" — persists but only surfaces when that project is active (requires project name)\\n` +
        `• "session" — temporary, lost when LM Studio closes. Use for short-lived context.\\n\\n` +
        `USE THIS when:\\n` +
        `• The user tells you something about themselves (name, job, preferences)\\n` +
        `• The user mentions a project they're working on\\n` +
        `• The user asks you to remember something\\n` +
        `• You learn an important fact that would be useful to recall later\\n\\n` +
        `Categories: ${VALID_CATEGORIES.join(", ")}`,
      parameters: {
        content: z.string().min(3).max(MAX_MEMORY_CONTENT_LENGTH)
          .describe("The fact/preference/note to remember. Be concise but complete."),
        category: z.enum(VALID_CATEGORIES).describe("Category of this memory."),
        tags: z.array(z.string().max(MAX_TAG_LENGTH)).max(MAX_TAGS_PER_MEMORY)
          .optional().describe("Optional tags for easier retrieval."),
        confidence: z.number().min(0).max(1).optional()
          .describe("How confident you are in this fact (0.0–1.0). Default 1.0."),
        scope: z.enum(VALID_SCOPES).optional()
          .describe("Memory scope: 'global' (default), 'project', 'session'."),
        project: z.string().max(MAX_PROJECT_NAME_LENGTH).optional()
          .describe("Project name (required when scope='project')."),
      },
      implementation: async (params, ctx) => {
        const result = await rememberImpl(params, ctl);
        if (result && typeof result === "object" && "stored" in result && result.stored) {
          ctx.status("Memory stored successfully");
        }
        return result;
      },
    }),

    tool({
      name: "Recall",
      description:
        `Search memory for relevant facts, preferences, or notes. ` +
        `Searches across all scopes (global + project + session) by default.\\n\\n` +
        `USE THIS when:\\n` +
        `• You need to check what you know about the user\\n` +
        `• The user references something from a past conversation\\n` +
        `• You want context before answering a question\\n` +
        `• The user asks \"do you remember…\" or \"what do you know about…\"`,
      parameters: {
        query: z.string().min(2).describe("What to search for — natural language topic, keyword, or question."),
        limit: z.number().int().min(1).max(MAX_SEARCH_RESULTS).optional()
          .describe(`Max results to return (default: ${DEFAULT_SEARCH_RESULTS}).`),
        category: z.enum(VALID_CATEGORIES).optional().describe("Filter by category."),
        scope: z.enum(VALID_SCOPES).optional().describe("Filter by scope. Omit to search all."),
        project: z.string().max(MAX_PROJECT_NAME_LENGTH).optional()
          .describe("Filter by project name."),
      },
      implementation: async (params, ctx) => {
        const result = await recallImpl(params, ctl);
        if (result && typeof result === "object" && "found" in result) {
          ctx.status(`Found ${(result as { found: number }).found} relevant memories`);
        }
        return result;
      },
    }),

    tool({
      name: "Search Memory",
      description:
        `Advanced memory search with filters by category, tag, or recency. ` +
        `Use 'Recall' for simple topic-based retrieval. Use this for precise filtering.`,
      parameters: {
        query: z.string().optional().describe("Optional text query for semantic search."),
        category: z.enum(VALID_CATEGORIES).optional().describe("Filter by memory category."),
        tag: z.string().optional().describe("Filter by tag (exact match)."),
        recent: z.number().int().min(1).max(50).optional()
          .describe("Get the N most recently created memories."),
        limit: z.number().int().min(1).max(MAX_SEARCH_RESULTS).optional()
          .describe(`Max results (default: ${DEFAULT_SEARCH_RESULTS}).`),
      },
      implementation: async (params, ctx) => searchImpl(params, ctl),
    }),

    tool({
      name: "Forget",
      description:
        `Delete memories by ID, pattern, or clear all. ` +
        `Use when the user asks you to forget something or when information is outdated.\\n\\n` +
        `USE THIS when:\\n` +
        `• User says \"forget that\" or \"delete that memory\"\\n` +
        `• User corrects a fact (delete old, store new)\\n` +
        `• You want to clear their data`,
      parameters: {
        id: z.string().optional().describe("Exact memory ID to delete."),
        pattern: z.string().optional().describe("Delete all memories whose content contains this text."),
        deleteAll: z.boolean().optional()
          .describe("Set to true to delete ALL memories. Use with extreme caution."),
      },
      implementation: async (params, ctx) => {
        const result = await forgetImpl(params, ctl);
        return result;
      },
    }),

    tool({
      name: "Memory Status",
      description:
        `Get statistics about the memory system: total count by scope, categories, ` +
        `session memory count, most accessed memory, database size, and index health.`,
      parameters: {},
      implementation: async (_, ctx) => {
        ctx.status("Gathering memory statistics…");
        return statusImpl({});
      },
    }),
  ];
}