letta-code/src/tests/agent/memoryPrompt.test.ts

import { describe, expect, test } from "bun:test";

import {
  buildSystemPrompt,
  isKnownPreset,
  SYSTEM_PROMPT_MEMFS_ADDON,
  SYSTEM_PROMPT_MEMORY_ADDON,
  swapMemoryAddon,
} from "../../agent/promptAssets";

function countOccurrences(haystack: string, needle: string): number {
  if (!needle) return 0;
  return haystack.split(needle).length - 1;
}

describe("isKnownPreset", () => {
  test("returns true for known preset IDs", () => {
    expect(isKnownPreset("default")).toBe(true);
    expect(isKnownPreset("letta-claude")).toBe(true);
    expect(isKnownPreset("letta-codex")).toBe(true);
  });

  test("returns false for unknown IDs", () => {
    expect(isKnownPreset("explore")).toBe(false);
    expect(isKnownPreset("nonexistent")).toBe(false);
  });
});

describe("buildSystemPrompt", () => {
  test("builds standard prompt with memory addon", () => {
    const result = buildSystemPrompt("letta-claude", "standard");
    expect(result).toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
    expect(result).not.toContain("## Memory Filesystem");
  });

  test("builds memfs prompt with memfs addon", () => {
    const result = buildSystemPrompt("letta-claude", "memfs");
    expect(result).toContain("## Memory Filesystem");
    expect(result).not.toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
  });

  test("throws on unknown preset", () => {
    expect(() => buildSystemPrompt("unknown-id", "standard")).toThrow(
      'Unknown preset "unknown-id"',
    );
  });

  test("is idempotent — same inputs always produce same output", () => {
    const first = buildSystemPrompt("default", "memfs");
    const second = buildSystemPrompt("default", "memfs");
    expect(first).toBe(second);
  });

  test("default preset uses SYSTEM_PROMPT content", () => {
    const result = buildSystemPrompt("default", "standard");
    expect(result).toContain("You are a self-improving AI agent");
    // default is NOT letta-claude — it uses the Letta-tuned system prompt
    const lettaClaudeResult = buildSystemPrompt("letta-claude", "standard");
    expect(result).not.toBe(lettaClaudeResult);
  });
});

describe("swapMemoryAddon", () => {
  test("swaps standard to memfs", () => {
    const base = "You are a test agent.";
    const standard = `${base}\n\n${SYSTEM_PROMPT_MEMORY_ADDON.trimStart()}`;

    const result = swapMemoryAddon(standard, "memfs");

    expect(result).toContain("## Memory Filesystem");
    expect(result).not.toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
    expect(countOccurrences(result, "## Memory Filesystem")).toBe(1);
  });

  test("swaps memfs to standard without orphan fragments", () => {
    const base = "You are a test agent.";
    const memfs = `${base}\n\n${SYSTEM_PROMPT_MEMFS_ADDON.trimStart()}`;

    const result = swapMemoryAddon(memfs, "standard");

    expect(result).toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
    expect(result).not.toContain("## Memory Filesystem");
    expect(result).not.toContain("# See what changed");
    expect(result).not.toContain('git commit -m "<type>: <what changed>"');
  });

  test("handles duplicate addons", () => {
    const base = "You are a test agent.";
    const doubled = `${base}\n\n${SYSTEM_PROMPT_MEMORY_ADDON}\n\n${SYSTEM_PROMPT_MEMORY_ADDON}`;

    const result = swapMemoryAddon(doubled, "memfs");

    expect(countOccurrences(result, "## Memory Filesystem")).toBe(1);
    expect(result).not.toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
  });

  test("strips orphan memfs tail fragment", () => {
    const tailStart = SYSTEM_PROMPT_MEMFS_ADDON.indexOf("# See what changed");
    expect(tailStart).toBeGreaterThanOrEqual(0);
    const orphanTail = SYSTEM_PROMPT_MEMFS_ADDON.slice(tailStart).trim();

    const drifted = `Header text\n\n${orphanTail}`;
    const result = swapMemoryAddon(drifted, "standard");

    expect(result).toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
    expect(result).not.toContain("# See what changed");
  });

  test("strips legacy heading-based ## Memory section", () => {
    const legacy =
      "You are a test agent.\n\n## Memory\nLegacy memory instructions here.\n\nSome other details.";

    const result = swapMemoryAddon(legacy, "memfs");

    expect(result).toContain("## Memory Filesystem");
    expect(result).not.toContain("Legacy memory instructions");
    expect(countOccurrences(result, "## Memory Filesystem")).toBe(1);
  });

  test("strips legacy heading-based ## Memory Filesystem section", () => {
    const legacy =
      "You are a test agent.\n\n## Memory Filesystem\nOld memfs instructions.";

    const result = swapMemoryAddon(legacy, "standard");

    expect(result).toContain(
      "Your memory consists of core memory (composed of memory blocks)",
    );
    expect(result).not.toContain("Old memfs instructions");
  });

  test("is idempotent", () => {
    const base = "You are a test agent.";
    const once = swapMemoryAddon(base, "memfs");
    const twice = swapMemoryAddon(once, "memfs");

    expect(twice).toBe(once);
    expect(countOccurrences(twice, "## Syncing")).toBe(1);
    expect(countOccurrences(twice, "# See what changed")).toBe(1);
  });
});