letta-code/src/agent/memory.ts

/**
 * Agent memory block management
 * Loads memory blocks from .mdx files in src/agent/prompts
 */

import type { CreateBlock } from "@letta-ai/letta-client/resources/blocks/blocks";
import { MEMORY_PROMPTS } from "./promptAssets";

/**
 * Block labels that are stored globally (shared across all projects).
 */
export const GLOBAL_BLOCK_LABELS = ["persona", "human"] as const;

/**
 * Block labels that are stored per-project (local to the current directory).
 */
export const PROJECT_BLOCK_LABELS = ["skills", "loaded_skills"] as const;

/**
 * All available memory block labels (derived from global + project blocks)
 */
export const MEMORY_BLOCK_LABELS = [
  ...GLOBAL_BLOCK_LABELS,
  ...PROJECT_BLOCK_LABELS,
] as const;

/**
 * Type for memory block labels
 */
export type MemoryBlockLabel = (typeof MEMORY_BLOCK_LABELS)[number];

/**
 * Block labels that should be read-only (agent cannot modify via memory tools).
 * These blocks are managed by specific tools (e.g., Skill tool for skills/loaded_skills).
 */
export const READ_ONLY_BLOCK_LABELS = [
  "skills",
  "loaded_skills",
  "memory_filesystem",
] as const;

/**
 * Block labels that should be isolated per-conversation.
 * When creating a conversation, these blocks are copied from the agent's blocks
 * to create conversation-specific versions, preventing cross-conversation state pollution.
 */
export const ISOLATED_BLOCK_LABELS = ["skills", "loaded_skills"] as const;

/**
 * Check if a block label is a project-level block
 */
export function isProjectBlock(label: string): boolean {
  return (PROJECT_BLOCK_LABELS as readonly string[]).includes(label);
}

/**
 * Parse frontmatter and content from an .mdx file
 */
export function parseMdxFrontmatter(content: string): {
  frontmatter: Record<string, string>;
  body: string;
} {
  const frontmatterRegex = /^---\n([\s\S]*?)\n---\n([\s\S]*)$/;
  const match = content.match(frontmatterRegex);

  if (!match || !match[1] || !match[2]) {
    return { frontmatter: {}, body: content };
  }

  const frontmatterText = match[1];
  const body = match[2];
  const frontmatter: Record<string, string> = {};

  // Parse YAML-like frontmatter (simple key: value pairs)
  for (const line of frontmatterText.split("\n")) {
    const colonIndex = line.indexOf(":");
    if (colonIndex > 0) {
      const key = line.slice(0, colonIndex).trim();
      const value = line.slice(colonIndex + 1).trim();
      frontmatter[key] = value;
    }
  }

  return { frontmatter, body: body.trim() };
}

/**
 * Load memory blocks from .mdx files in src/agent/prompts
 */
async function loadMemoryBlocksFromMdx(): Promise<CreateBlock[]> {
  const memoryBlocks: CreateBlock[] = [];

  const mdxFiles = MEMORY_BLOCK_LABELS.map((label) => `${label}.mdx`);

  for (const filename of mdxFiles) {
    try {
      const content = MEMORY_PROMPTS[filename];
      if (!content) {
        console.warn(`Missing embedded prompt file: ${filename}`);
        continue;
      }
      const { frontmatter, body } = parseMdxFrontmatter(content);

      const label = frontmatter.label || filename.replace(".mdx", "");
      const block: CreateBlock = {
        label,
        value: body,
      };

      if (frontmatter.description) {
        block.description = frontmatter.description;
      }

      if (frontmatter.limit) {
        const limit = parseInt(frontmatter.limit, 10);
        if (!Number.isNaN(limit) && limit > 0) {
          block.limit = limit;
        }
      }

      // Set read-only for skills blocks (managed by Skill tool, not memory tools)
      if ((READ_ONLY_BLOCK_LABELS as readonly string[]).includes(label)) {
        block.read_only = true;
      }

      memoryBlocks.push(block);
    } catch (error) {
      console.error(`Error loading ${filename}:`, error);
    }
  }

  return memoryBlocks;
}

// Cache for loaded memory blocks
let cachedMemoryBlocks: CreateBlock[] | null = null;

/**
 * Get default starter memory blocks for new agents
 */
export async function getDefaultMemoryBlocks(): Promise<CreateBlock[]> {
  if (!cachedMemoryBlocks) {
    cachedMemoryBlocks = await loadMemoryBlocksFromMdx();
  }
  return cachedMemoryBlocks;
}

/**
 * Ensure an agent has the required skills blocks (skills, loaded_skills).
 * If missing, creates them with default values and attaches to the agent.
 * This is needed for backwards compatibility with agents created before skills were added.
 *
 * @param agentId - The agent ID to check/update
 * @returns Array of block labels that were created (empty if all existed)
 */
export async function ensureSkillsBlocks(agentId: string): Promise<string[]> {
  const { getClient } = await import("./client");
  const client = await getClient();

  // Get current blocks on the agent
  // Response may be paginated or an array depending on SDK version
  const blocksResponse = await client.agents.blocks.list(agentId);
  const blocks = Array.isArray(blocksResponse)
    ? blocksResponse
    : (blocksResponse as { items?: Array<{ label?: string }> }).items || [];
  const existingLabels = new Set(blocks.map((b) => b.label));

  const createdLabels: string[] = [];

  // Check each required skills block
  for (const label of ISOLATED_BLOCK_LABELS) {
    if (existingLabels.has(label)) {
      continue;
    }

    // Load the default block definition from .mdx
    const content = MEMORY_PROMPTS[`${label}.mdx`];
    if (!content) {
      console.warn(`Missing embedded prompt file for ${label}.mdx`);
      continue;
    }

    const { frontmatter, body } = parseMdxFrontmatter(content);

    const blockData: CreateBlock = {
      label,
      value: body,
      read_only: true, // Skills blocks are read-only (managed by Skill tool)
    };

    if (frontmatter.description) {
      blockData.description = frontmatter.description;
    }

    if (frontmatter.limit) {
      const limit = parseInt(frontmatter.limit, 10);
      if (!Number.isNaN(limit) && limit > 0) {
        blockData.limit = limit;
      }
    }

    // Create the block and attach to agent
    const createdBlock = await client.blocks.create(blockData);
    await client.agents.blocks.attach(createdBlock.id, { agent_id: agentId });
    createdLabels.push(label);
  }

  return createdLabels;
}