kv-store

/**
 * @file toolsProvider.ts
 * Registers tools with LM Studio.
 *
 * Tools:
 *   1. Remember        - store a new memory
 */
import { text, tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";
import { KvDatabase, MAX_VALUE_LENGTH_CHARS } from "./db";

function readConfig(ctl: ToolsProviderController) {
  const c = ctl.getPluginConfig(configSchematics);
  return {
    dbStorageFilename: (c.get("dbStorageFilename") as string | undefined) ?? "",
    enableTimestampUtilityTool: (c.get("enableTimestampUtilityTool") as "on" | "off" | undefined) ??
      "off",
  };
}

// singletons, reused across calls
let db: KvDatabase | null = null;
let currentFilename: string = "";
let initPromise: Promise<{
  db: KvDatabase;
}> | null = null;

async function ensureInitialized(storageFilename: string): Promise<{ db: KvDatabase }> {
  const resolvedFilename = storageFilename.trim();
  if (db && currentFilename === resolvedFilename) return { db };

  if (initPromise && currentFilename === resolvedFilename) return initPromise;

  initPromise = (async () => {
    if (db) {
      try {
        db.close();
      } catch {}
    }
    db = new KvDatabase(resolvedFilename || undefined);
    await db.init();

    currentFilename = resolvedFilename;
    return { db: db! };
  })();

  return initPromise;
}

async function runKvOperation<T>(
  operation: () => Promise<T> | T,
): Promise<({ ok: true } & T) | { ok: false; error: string }> {
  try {
    const result = await operation();
    return {
      ok: true,
      ...result,
    };
  } catch (error) {
    return {
      ok: false,
      error: error instanceof Error ? error.message : "Unknown error",
    };
  }
}

function getValueSize(value: string | null): number | null {
  return value === null ? null : value.length;
}

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

  const getCurrentTimestampTool = tool({
    name: "get_current_timestamp",
    description: text`
      Get current UTC time.
      Returns { timestamp } as an ISO 8601 string.
    `,
    parameters: {},
    implementation: async () => {
      return { timestamp: new Date().toISOString() };
    },
  });

  const setTool = tool({
    name: "kv_set",
    description: text`
      Create or overwrite one key-value pair.
      Input: key (1-255 chars, trimmed), value (string, max ${MAX_VALUE_LENGTH_CHARS} chars).
      Returns { ok: true } or { ok: false, error }.
    `,
    parameters: {
      key: z.string().min(1).max(255).describe("Key name, 1-255 chars."),
      value: z
        .string()
        .max(MAX_VALUE_LENGTH_CHARS)
        .describe(`String value to store (max ${MAX_VALUE_LENGTH_CHARS} chars).`),
    },
    implementation: async ({ key, value }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        db.set(key.trim(), value);
        return {
          size: value.length,
        };
      });
    },
  });

  const getTool = tool({
    name: "kv_get",
    description: text`
      Get value for one key.
      Returns { ok: true, found, value } or { ok: false, error }.
      Missing key: found = false, value = null.
    `,
    parameters: {
      key: z.string().describe("Key name to read."),
    },
    implementation: async ({ key }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const value = db.get(key);
        return {
          found: value !== null,
          size: getValueSize(value),
          value,
        };
      });
    },
  });

  const appendTool = tool({
    name: "kv_append",
    description: text`
      Append text to a key's current value.
      Creates the key if it does not already exist.
      Input: key (1-255 chars, trimmed), value (string, max combined length ${MAX_VALUE_LENGTH_CHARS} chars).
      Returns { ok: true, found, value } or { ok: false, error }.
    `,
    parameters: {
      key: z.string().min(1).max(255).describe("Key name, 1-255 chars."),
      value: z
        .string()
        .max(MAX_VALUE_LENGTH_CHARS)
        .describe(`String value to append (max ${MAX_VALUE_LENGTH_CHARS} chars before combining).`),
    },
    implementation: async ({ key, value }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const appendedValue = db.append(key.trim(), value);
        return {
          found: true,
          size: appendedValue.length,
          value: appendedValue,
        };
      });
    },
  });

  const getManyTool = tool({
    name: "kv_get_many",
    description: text`
      Get values for multiple keys.
      Input: 1-10 keys.
      Returns { ok: true, pairs } or { ok: false, error }.
      Each pair is { key, value }. Only existing keys are returned.
    `,
    parameters: {
      keys: z.array(z.string()).min(1).max(10).describe("Keys to read (1-10)."),
    },
    implementation: async ({ keys }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const pairs = db.getMany(keys).map((pair) => ({
          ...pair,
          size: getValueSize(pair.value),
        }));
        return { pairs };
      });
    },
  });

  const deleteTool = tool({
    name: "kv_delete",
    description: text`
      Delete one key.
      No-op if missing.
      Returns { ok: true } or { ok: false, error }.
    `,
    parameters: {
      key: z.string().describe("Key name to delete."),
    },
    implementation: async ({ key }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        db.delete(key);
        return {};
      });
    },
  });

  const deleteManyTool = tool({
    name: "kv_delete_many",
    description: text`
      Delete multiple keys.
      Input: 1-10 keys.
      Missing keys are ignored.
      Returns { ok: true } or { ok: false, error }.
    `,
    parameters: {
      keys: z.array(z.string()).min(1).max(10).describe("Keys to delete (1-10)."),
    },
    implementation: async ({ keys }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        db.deleteMany(keys);
        return {};
      });
    },
  });

  const renameTool = tool({
    name: "kv_rename",
    description: text`
      Rename one key.
      Preserves the stored value.
      Fails if the source key does not exist or the destination key already exists.
      Returns { ok: true } or { ok: false, error }.
    `,
    parameters: {
      oldKey: z.string().min(1).max(255).describe("Existing key name to rename."),
      newKey: z.string().min(1).max(255).describe("New key name to assign."),
    },
    implementation: async ({ oldKey, newKey }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        db.rename(oldKey.trim(), newKey.trim());
        return {};
      });
    },
  });

  const listKeysTool = tool({
    name: "kv_list_keys",
    description: text`
      List keys in ascending sort order.
      Optional pagination: limit and offset.
      Max 200 keys per call.
      Returns { ok: true, keys } or { ok: false, error }.
    `,
    parameters: {
      limit: z.number().int().min(0).max(200).optional().describe("Max keys to return (0-200)."),
      offset: z.number().int().min(0).optional().describe("Number of sorted keys to skip."),
    },
    implementation: async ({ limit, offset }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const keys = db.listKeys(undefined, limit, offset);
        return { keys };
      });
    },
  });

  const keysFilterTool = tool({
    name: "kv_filter_keys",
    description: text`
      List keys that contain a substring (case-insensitive), sorted ascending.
      Optional pagination: limit and offset.
      Max 200 keys per call.
      Returns { ok: true, keys } or { ok: false, error }.
    `,
    parameters: {
      match: z.string().describe("Substring to match in key names (case-insensitive)."),
      limit: z.number().int().min(0).max(200).optional().describe("Max keys to return (0-200)."),
      offset: z.number().int().min(0).optional().describe("Number of matching sorted keys to skip."),
    },
    implementation: async ({ match, limit, offset }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const keys = db.listKeys(match, limit, offset);
        return { keys };
      });
    },
  });

  const countKeysTool = tool({
    name: "kv_count_keys",
    description: text`
      Count keys in the store.
      Optional filter: case-insensitive substring match.
      Returns { ok: true, count } or { ok: false, error }.
    `,
    parameters: {
      filter: z.string().optional().describe("Optional case-insensitive key substring filter."),
    },
    implementation: async ({ filter }) => {
      return runKvOperation(async () => {
        const { db } = await ensureInitialized(cfg.dbStorageFilename);
        const count = db.countKeys(filter);
        return { count };
      });
    },
  });

  const tools = [
    setTool,
    getTool,
    appendTool,
    getManyTool,
    deleteTool,
    deleteManyTool,
    renameTool,
    listKeysTool,
    keysFilterTool,
    countKeysTool,
  ];

  if (cfg.enableTimestampUtilityTool === "on") {
    tools.push(getCurrentTimestampTool);
  }

  return tools;
}