@pi-unipi/memory 0.1.5 → 0.1.8

package/commands.ts

@@ -198,7 +198,7 @@ For each item, use the memory_store tool to save it with an appropriate title an
         return;
       }
 
-      await showMemorySettings(pi);
+      await showMemorySettings(ctx);
     },
   });
 }

package/embedding.ts

@@ -8,7 +8,7 @@
  * openai/text-embedding-3 supports custom dimensions via API param.
  */
 
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
 import {
   loadEmbeddingConfig,
   getApiKey,
@@ -174,7 +174,7 @@ export async function generateEmbeddingsBatch(
  * Re-embed all memories across all projects.
  * Returns count of successfully re-embedded memories.
  */
-export async function reembedAllMemories(pi: ExtensionAPI): Promise<number> {
+export async function reembedAllMemories(ctx: ExtensionCommandContext): Promise<number> {
   const { getAllProjectDirs, MemoryStorage } = await import("./storage.js");
   const projectDirs = getAllProjectDirs();
   let count = 0;
@@ -197,7 +197,7 @@ export async function reembedAllMemories(pi: ExtensionAPI): Promise<number> {
 
       // Generate embeddings in batch
       const texts = fullRecords.map((r) => `${r.title} ${r.content}`);
-      const embeddings = await generateEmbeddingsBatch(texts, pi);
+      const embeddings = await generateEmbeddingsBatch(texts);
 
       // Update records
       for (let i = 0; i < fullRecords.length; i++) {

package/index.ts

@@ -73,7 +73,18 @@ export default function (pi: ExtensionAPI) {
     // Initialize project storage
     const projectName = getProjectName(ctx.cwd);
     projectStorage = new MemoryStorage(projectName);
-    projectStorage.init();
+    try {
+      projectStorage.init();
+      
+      // Sync any orphaned markdown files into the database
+      const synced = projectStorage.syncOrphanedFiles();
+      if (synced > 0) {
+        console.warn(`[unipi/memory] Synced ${synced} orphaned memory files into database`);
+      }
+    } catch (err) {
+      console.warn("[unipi/memory] Failed to initialize storage, running without memory:", (err as any)?.message ?? err);
+      projectStorage = null;
+    }
 
 
     // Announce module
@@ -127,8 +138,14 @@ export default function (pi: ExtensionAPI) {
             };
           }
 
-          const projectMemories = projectStorage.listAll();
-          const allMemories = listAllProjects();
+          let projectMemories: Array<{ id: string; title: string; type: string }> = [];
+          let allMemories: Array<{ project: string; id: string; title: string; type: string }> = [];
+          try {
+            projectMemories = projectStorage.listAll();
+            allMemories = listAllProjects();
+          } catch (err) {
+            console.warn("[unipi/memory] Failed to list memories for info panel:", err);
+          }
           const uniqueProjects = [...new Set(allMemories.map((m) => m.project))];
 
           // Get 3 most recent memories (sorted by updated DESC in listAll)
@@ -153,9 +170,14 @@ export default function (pi: ExtensionAPI) {
 
     // Show memory status in UI
     if (ctx.hasUI) {
-      const projectCount = projectStorage.listAll().length;
-      const allMemories = listAllProjects();
-      const projectCountAll = allMemories.length;
+      let projectCount = 0;
+      let projectCountAll = 0;
+      try {
+        projectCount = projectStorage?.listAll()?.length ?? 0;
+        projectCountAll = listAllProjects().length;
+      } catch (err) {
+        console.warn("[unipi/memory] Failed to count memories for status:", err);
+      }
       const vecReady = isEmbeddingReady();
       const vecIcon = vecReady ? "⚡" : "📝";
       ctx.ui.setStatus(
@@ -171,7 +193,14 @@ export default function (pi: ExtensionAPI) {
     if (!projectStorage) return;
 
     const projectName = getProjectName(ctx.cwd);
-    const projectMemories = projectStorage.listAll();
+    let projectMemories: Array<{ id: string; title: string; type: string }> = [];
+    try {
+      projectMemories = projectStorage.listAll();
+    } catch (err) {
+      console.warn("[unipi/memory] Failed to list memories for recall:", err);
+      recallDone = true; // Skip recall on error
+      return;
+    }
 
     if (projectMemories.length === 0) {
       recallDone = true; // Nothing to recall, skip

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi-unipi/memory",
-  "version": "0.1.5",
+  "version": "0.1.8",
   "description": "Persistent cross-session memory with vector search for Pi coding agent",
   "type": "module",
   "license": "MIT",

package/storage.ts

@@ -44,7 +44,22 @@ interface MemoryFrontmatter {
 }
 
 const MEMORY_DB_NAME = "memory.db";
-const MEMORY_EMBEDDING_DIM = 384;
+/**
+ * Get the configured embedding dimensions.
+ * Reads from config, falls back to 384.
+ */
+function getEmbeddingDims(): number {
+  try {
+    const configPath = path.join(os.homedir(), ".unipi", "memory", "config.json");
+    if (fs.existsSync(configPath)) {
+      const raw = JSON.parse(fs.readFileSync(configPath, "utf-8"));
+      if (typeof raw.dimensions === "number" && raw.dimensions >= 64) {
+        return raw.dimensions;
+      }
+    }
+  } catch { /* ignore */ }
+  return 384;
+}
 
 /**
  * Get the base memory directory (~/.unipi/memory/)
@@ -173,6 +188,12 @@ export class MemoryStorage {
 
   /**
    * Initialize the storage (create DB, tables, load extension).
+   *
+   * Uses retry logic to handle concurrent access from multiple Pi sessions,
+   * especially on WSL/Windows filesystem where SQLite locking can be flaky.
+   *
+   * IMPORTANT: We never delete the DB here — another session may have it open.
+   * If all retries fail, we throw and let this session run without memory.
    */
   init(): void {
     // Ensure directory exists
@@ -181,10 +202,56 @@ export class MemoryStorage {
     }
 
     const dbPath = path.join(this.scopeDir, MEMORY_DB_NAME);
-    this.db = new Database(dbPath);
+    const maxRetries = 5;
+
+    for (let attempt = 1; attempt <= maxRetries; attempt++) {
+      try {
+        this.initDb(dbPath);
+        return; // Success
+      } catch (err: any) {
+        const isTransient =
+          err?.message?.includes("disk I/O error") ||
+          err?.code === "SQLITE_IOERR" ||
+          err?.code === "SQLITE_BUSY" ||
+          err?.message?.includes("database is locked");
+
+        this.close();
+
+        if (isTransient && attempt < maxRetries) {
+          // Likely concurrent access — back off and retry.
+          // Do NOT delete the DB: another session may have it open
+          // and deleting open files on WSL/Windows is unsafe.
+          const delayMs = 50 * Math.pow(2, attempt - 1); // 50, 100, 200, 400
+          console.warn(
+            `[unipi/memory] Transient error on attempt ${attempt}/${maxRetries}, retrying in ${delayMs}ms...`
+          );
+          const end = Date.now() + delayMs;
+          while (Date.now() < end) { /* busy wait */ }
+          continue;
+        }
+
+        // Either non-transient error, or retries exhausted.
+        // Log and throw — this session will run without memory.
+        if (isTransient) {
+          console.warn(
+            "[unipi/memory] Could not open database after retries. " +
+            "Another session may have the DB locked. Memory unavailable this session."
+          );
+        }
+        throw err;
+      }
+    }
+  }
+
+  /**
+   * Open database and set up schema. Called by init() with retry logic.
+   */
+  private initDb(dbPath: string): void {
+    this.db = new Database(dbPath, { timeout: 5000 });
 
     // Enable WAL mode for concurrent reads
     this.db.pragma("journal_mode = WAL");
+    this.db.pragma("busy_timeout = 5000");
 
     // Load sqlite-vec extension
     try {
@@ -211,11 +278,14 @@ export class MemoryStorage {
     // Create vector table if sqlite-vec loaded
     try {
       this.db.exec(`
-        CREATE VIRTUAL TABLE IF NOT EXISTS memories_vec USING vec0(embedding float[${MEMORY_EMBEDDING_DIM}])
+        CREATE VIRTUAL TABLE IF NOT EXISTS memories_vec USING vec0(embedding float[${getEmbeddingDims()}])
       `);
     } catch {
       // vec0 table may already exist or sqlite-vec not loaded
     }
+
+    // Verify database is usable
+    this.db.prepare("SELECT 1 FROM memories LIMIT 0").get();
   }
 
   /**
@@ -228,8 +298,40 @@ export class MemoryStorage {
     }
   }
 
+  /**
+   * Remove corrupted database files (db, wal, shm).
+   */
+  private removeCorruptedDb(): void {
+    const dbPath = path.join(this.scopeDir, MEMORY_DB_NAME);
+    const files = [dbPath, `${dbPath}-wal`, `${dbPath}-shm`];
+    for (const file of files) {
+      try {
+        if (fs.existsSync(file)) {
+          fs.unlinkSync(file);
+          console.warn(`[unipi/memory] Removed corrupted file: ${file}`);
+        }
+      } catch {
+        // Ignore removal errors
+      }
+    }
+  }
+
+  /**
+   * Check if database is healthy.
+   */
+  isHealthy(): boolean {
+    if (!this.db) return false;
+    try {
+      this.db.prepare("SELECT 1").get();
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
   /**
    * Store or update a memory record.
+   * Uses transaction to ensure atomicity — either all writes succeed or none do.
    */
   store(record: MemoryRecord): void {
     if (!this.db) throw new Error("Storage not initialized");
@@ -247,52 +349,181 @@ export class MemoryStorage {
     // Set project if not provided
     if (!record.project) record.project = this.projectName;
 
-    // Upsert into memories table
-    const stmt = this.db.prepare(`
-      INSERT OR REPLACE INTO memories (id, title, content, tags, project, type, created, updated, embedding)
-      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
-    `);
+    // Prepare markdown content BEFORE transaction (fail fast)
+    const mdPath = path.join(this.scopeDir, `${record.id}.md`);
+    const frontmatter: MemoryFrontmatter = {
+      title: record.title,
+      tags: record.tags,
+      project: record.project,
+      created: record.created,
+      updated: record.updated,
+      type: record.type,
+    };
+    const mdContent = `---\n${yaml.dump(frontmatter, { lineWidth: -1 })}---\n\n${record.content}\n`;
+
+    // Use transaction for atomicity
+    const storeInTx = this.db.transaction(() => {
+      // Upsert into memories table
+      const stmt = this.db!.prepare(`
+        INSERT OR REPLACE INTO memories (id, title, content, tags, project, type, created, updated, embedding)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+      `);
 
-    const tagsJson = JSON.stringify(record.tags);
-    const embeddingBuf = record.embedding ? Buffer.from(record.embedding.buffer) : null;
-
-    stmt.run(
-      record.id,
-      record.title,
-      record.content,
-      tagsJson,
-      record.project,
-      record.type,
-      record.created,
-      record.updated,
-      embeddingBuf
-    );
+      const tagsJson = JSON.stringify(record.tags);
+      const embeddingBuf = record.embedding ? Buffer.from(record.embedding.buffer) : null;
+
+      stmt.run(
+        record.id,
+        record.title,
+        record.content,
+        tagsJson,
+        record.project,
+        record.type,
+        record.created,
+        record.updated,
+        embeddingBuf
+      );
 
-    // Update vector table
-    if (record.embedding) {
-      try {
-        // Delete old vector if exists
-        this.db.prepare("DELETE FROM memories_vec WHERE rowid = ?").run(BigInt(this.idToRowid(record.id)));
-      } catch {
-        // Ignore if not found
+      // Update vector table
+      if (record.embedding) {
+        try {
+          // Delete old vector if exists
+          this.db!.prepare("DELETE FROM memories_vec WHERE rowid = ?").run(BigInt(this.idToRowid(record.id)));
+        } catch {
+          // Ignore if not found
+        }
+
+        try {
+          const vecStmt = this.db!.prepare(
+            "INSERT INTO memories_vec(rowid, embedding) VALUES (?, ?)"
+          );
+          vecStmt.run(
+            BigInt(this.idToRowid(record.id)),
+            Buffer.from(record.embedding.buffer)
+          );
+        } catch (err) {
+          console.warn("[unipi/memory] Failed to insert vector:", err);
+        }
       }
+    });
 
+    // Execute transaction
+    storeInTx();
+
+    // Write markdown file AFTER successful DB write
+    try {
+      // Ensure directory exists
+      const dir = path.dirname(mdPath);
+      if (!fs.existsSync(dir)) {
+        fs.mkdirSync(dir, { recursive: true });
+      }
+      fs.writeFileSync(mdPath, mdContent, "utf-8");
+    } catch (err) {
+      // DB write succeeded but file write failed — log but don't throw
+      // Memory is still in DB and searchable
+      console.warn("[unipi/memory] Failed to write markdown file:", err);
+    }
+  }
+
+  /**
+   * Sync orphaned markdown files into the database.
+   * Reads all .md files in the project dir, parses frontmatter,
+   * and inserts any that are missing from the DB.
+   * Returns count of synced files.
+   */
+  syncOrphanedFiles(): number {
+    if (!this.db) throw new Error("Storage not initialized");
+
+    const files = fs.readdirSync(this.scopeDir)
+      .filter(f => f.endsWith(".md") && !f.startsWith("."));
+
+    // Get existing IDs from DB
+    const existingIds = new Set(
+      (this.db.prepare("SELECT id FROM memories").all() as any[])
+        .map(r => r.id)
+    );
+
+    let synced = 0;
+    for (const file of files) {
+      const filePath = path.join(this.scopeDir, file);
+      const record = parseMemoryFile(filePath);
+      if (!record) continue;
+
+      // Generate ID from title (same logic as store())
+      const id = record.title.toLowerCase().replace(/[^a-z0-9]+/g, "_");
+
+      if (existingIds.has(id)) continue; // Already in DB
+
+      // Insert into DB
       try {
-        const vecStmt = this.db.prepare(
-          "INSERT INTO memories_vec(rowid, embedding) VALUES (?, ?)"
-        );
-        vecStmt.run(
-          BigInt(this.idToRowid(record.id)),
-          Buffer.from(record.embedding.buffer)
+        record.id = id;
+        const tagsJson = JSON.stringify(record.tags);
+        
+        this.db.prepare(`
+          INSERT OR IGNORE INTO memories (id, title, content, tags, project, type, created, updated, embedding)
+          VALUES (?, ?, ?, ?, ?, ?, ?, ?, NULL)
+        `).run(
+          id,
+          record.title,
+          record.content,
+          tagsJson,
+          record.project || this.projectName,
+          record.type,
+          record.created,
+          record.updated
         );
+
+        synced++;
+        console.warn(`[unipi/memory] Synced orphaned file: ${file}`);
       } catch (err) {
-        console.warn("[unipi/memory] Failed to insert vector:", err);
+        console.warn(`[unipi/memory] Failed to sync ${file}:`, err);
       }
     }
 
-    // Write markdown file
-    const mdPath = path.join(this.scopeDir, `${record.id}.md`);
-    writeMemoryFile(mdPath, record);
+    return synced;
+  }
+
+  /**
+   * Check if a memory with the given title already exists.
+   */
+  hasByTitle(title: string): boolean {
+    if (!this.db) throw new Error("Storage not initialized");
+    const id = title.toLowerCase().replace(/[^a-z0-9]+/g, "_");
+    const row = this.db.prepare("SELECT 1 FROM memories WHERE id = ?").get(id);
+    return !!row;
+  }
+
+  /**
+   * Find memories with similar titles (fuzzy match).
+   * Returns array of { record, similarity } sorted by similarity desc.
+   */
+  findSimilarByTitle(title: string, threshold = 0.6): Array<{ record: MemoryRecord; similarity: number }> {
+    if (!this.db) throw new Error("Storage not initialized");
+
+    const allRows = this.db.prepare("SELECT id, title FROM memories").all() as any[];
+    const results: Array<{ record: MemoryRecord; similarity: number }> = [];
+
+    const normalizedTitle = title.toLowerCase().replace(/[^a-z0-9]+/g, " ");
+    const titleWords = new Set(normalizedTitle.split(/\s+/).filter((w: string) => w.length > 2));
+
+    for (const row of allRows) {
+      const normalizedRowTitle = row.title.toLowerCase().replace(/[^a-z0-9]+/g, " ");
+      const rowWords = new Set(normalizedRowTitle.split(/\s+/).filter((w: string) => w.length > 2));
+
+      // Calculate Jaccard similarity
+      const intersection = new Set([...titleWords].filter((w: string) => rowWords.has(w)));
+      const union = new Set([...titleWords, ...rowWords]);
+      const similarity = union.size > 0 ? intersection.size / union.size : 0;
+
+      if (similarity >= threshold) {
+        const record = this.getById(row.id);
+        if (record) {
+          results.push({ record, similarity });
+        }
+      }
+    }
+
+    return results.sort((a, b) => b.similarity - a.similarity);
   }
 
   /**

package/tools.ts

@@ -72,10 +72,28 @@ export function registerMemoryTools(
       const storage = getStorage();
       onActivity?.(); // Mark store as done for lifecycle
 
-      // Check if similar memory exists
+      // Step 1: Check for exact duplicate
       const existing = storage.getByTitle(params.title);
       if (existing) {
-        // Update existing
+        // Exact match found — check if content is also the same
+        const isSameContent = existing.content.trim() === params.content.trim();
+        
+        if (isSameContent) {
+          // Duplicate with same content — gentle error asking to read first
+          return {
+            content: [
+              {
+                type: "text",
+                text: `⚠️ Memory already exists with this title and content: "${params.title}"\n\n` +
+                      `Please read the existing memory first using memory_search before saving.\n` +
+                      `If you want to update it, provide new or modified content.`,
+              },
+            ],
+            details: { action: "duplicate_detected", id: existing.id },
+          };
+        }
+
+        // Same title but different content — update existing
         const updated: MemoryRecord = {
           ...existing,
           content: params.content,
@@ -102,7 +120,50 @@ export function registerMemoryTools(
         };
       }
 
-      // Create new memory
+      // Step 2: Check for similar memories
+      const similarMemories = storage.findSimilarByTitle(params.title, 0.6);
+      
+      if (similarMemories.length > 0) {
+        // Found similar memories — save but notify
+        const similarList = similarMemories
+          .slice(0, 3)
+          .map(s => `  - "${s.record.title}" (${Math.round(s.similarity * 100)}% similar)`)
+          .join("\n");
+
+        // Create new memory
+        const record: MemoryRecord = {
+          id: "",
+          title: params.title,
+          content: params.content,
+          tags: params.tags || [],
+          project: getProjectName(ctx.cwd),
+          type: (params.type as MemoryRecord["type"]) || "summary",
+          created: "",
+          updated: "",
+        };
+
+        const embedding = await generateEmbedding(
+          params.title + " " + params.content,
+          pi
+        );
+        record.embedding = embedding;
+
+        storage.store(record);
+
+        return {
+          content: [
+            {
+              type: "text",
+              text: `Stored memory: ${params.title}\n\n` +
+                    `⚠️ Similar memories found:\n${similarList}\n\n` +
+                    `Consider reviewing these to avoid redundancy.`,
+            },
+          ],
+          details: { action: "created_with_similar", id: record.id, similar: similarMemories.map(s => s.record.id) },
+        };
+      }
+
+      // Step 3: No duplicates or similar — create new memory
       const record: MemoryRecord = {
         id: "",
         title: params.title,

package/tui/settings-tui.ts

@@ -2,10 +2,10 @@
  * @unipi/memory — Settings TUI
  *
  * Interactive settings dialog for embedding configuration.
- * Uses pi's UI primitives (select, input, notify).
+ * Uses ctx.ui primitives (select, input, notify).
  */
 
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
 import {
   loadEmbeddingConfig,
   saveEmbeddingConfig,
@@ -19,20 +19,12 @@ import {
   type EmbeddingConfig,
 } from "../settings.js";
 
-/** pi.ui type that's available when TUI is present */
-type PiUI = {
-  select: (opts: { title: string; message: string; options: Array<{ label: string; value: string; description?: string }> }) => Promise<string | null | undefined>;
-  input: (opts: { title: string; message: string; placeholder?: string; validate?: (value: string) => Promise<string | null> }) => Promise<string | null | undefined>;
-  notify: (opts: { message: string; level: string }) => Promise<void>;
-};
-
 /**
  * Show memory settings dialog.
  * Main entry point for /unipi:memory-settings command.
  */
-export async function showMemorySettings(pi: ExtensionAPI): Promise<void> {
-  // Cast to access pi.ui which exists at runtime but isn't typed
-  const ui = (pi as any).ui as PiUI;
+export async function showMemorySettings(ctx: ExtensionCommandContext): Promise<void> {
+  const ui = ctx.ui;
   let running = true;
 
   while (running) {
@@ -115,28 +107,28 @@ export async function showMemorySettings(pi: ExtensionAPI): Promise<void> {
       description: "Exit settings",
     });
 
-    const selected = await ui.select({
-      title: "🧠 Memory Settings",
-      message: statusLines.join("\n"),
-      options,
-    });
+    const labels = options.map(o => `${o.label} — ${o.description}`);
+    const selected = await ui.select(
+      "🧠 Memory Settings",
+      labels,
+    );
+    // Map selected label back to value
+    const selectedOpt = options.find(o => `${o.label} — ${o.description}` === selected);
+    const selectedValue = selectedOpt?.value;
 
-    if (!selected || selected === "__exit__") {
+    if (!selectedValue || selectedValue === "__exit__") {
       running = false;
       continue;
     }
 
-    switch (selected) {
+    switch (selectedValue) {
       case "__add_key__":
       case "__update_key__":
         await handleApiKeyInput(ui);
         break;
       case "__remove_key__":
         clearApiKey();
-        await ui.notify({
-          message: "API key removed. Vector search disabled.",
-          level: "info",
-        });
+        ui.notify("API key removed. Vector search disabled.", "info");
         break;
       case "__select_model__":
         await handleModelSelection(ui);
@@ -145,16 +137,13 @@ export async function showMemorySettings(pi: ExtensionAPI): Promise<void> {
         await handleDimensionsInput(ui);
         break;
       case "__reembed__":
-        await handleReembed(ui, pi);
+        await handleReembed(ui, ctx);
         break;
       case "__suppress__":
         const cfg = loadEmbeddingConfig();
         cfg.suppressMigrationWarning = true;
         saveEmbeddingConfig(cfg);
-        await ui.notify({
-          message: "Migration warning suppressed.",
-          level: "info",
-        });
+        ui.notify("Migration warning suppressed.", "info");
         break;
     }
   }
@@ -163,66 +152,65 @@ export async function showMemorySettings(pi: ExtensionAPI): Promise<void> {
 /**
  * Handle API key input.
  */
-async function handleApiKeyInput(ui: PiUI): Promise<void> {
-  const key = await ui.input({
-    title: "OpenRouter API Key",
-    message: "Enter your OpenRouter API key (sk-or-v1-...):",
-    placeholder: "sk-or-v1-...",
-    validate: async (value: string) => {
-      if (!value || value.trim().length === 0) {
-        return "API key cannot be empty";
-      }
-      if (!value.startsWith("sk-or-") && !value.startsWith("sk-")) {
-        return "Key should start with sk-or- or sk-";
-      }
-      return null;
-    },
-  });
+async function handleApiKeyInput(ui: ExtensionCommandContext["ui"]): Promise<void> {
+  const key = await ui.input(
+    "Enter your OpenRouter API key (sk-or-v1-...):",
+    "sk-or-v1-...",
+  );
 
   if (key) {
-    setApiKey(key.trim());
-    await ui.notify({
-      message: "API key saved. Vector search enabled.",
-      level: "success",
-    });
+    const trimmed = key.trim();
+    if (trimmed.length === 0) {
+      ui.notify("API key cannot be empty.", "warning");
+      return;
+    }
+    if (!trimmed.startsWith("sk-or-") && !trimmed.startsWith("sk-")) {
+      ui.notify("Key should start with sk-or- or sk-.", "warning");
+      return;
+    }
+    setApiKey(trimmed);
+    ui.notify("API key saved. Vector search enabled.", "info");
   }
 }
 
 /**
  * Handle model selection.
  */
-async function handleModelSelection(ui: PiUI): Promise<void> {
+async function handleModelSelection(ui: ExtensionCommandContext["ui"]): Promise<void> {
   const config = loadEmbeddingConfig();
 
-  const options = OPENROUTER_EMBEDDING_MODELS.map((m) => ({
+  const modelOptions = OPENROUTER_EMBEDDING_MODELS.map((m) => ({
     label: `${m.name}${m.id === config.model ? " ✓" : ""}`,
     value: m.id,
     description: `${m.description} (${m.dimensions}d, ~${m.costPer1k}/1k tokens)`,
   }));
 
   // Add custom option
-  options.push({
+  modelOptions.push({
     label: "✏️ Custom Model ID",
     value: "__custom__",
     description: "Enter a custom OpenRouter model ID",
   });
 
-  const selected = await ui.select({
-    title: "Select Embedding Model",
-    message: "Choose an embedding model. ⚠ Changing model invalidates existing embeddings.",
-    options,
-  });
+  const labels = modelOptions.map(o => `${o.label} — ${o.description}`);
+  const selected = await ui.select(
+    "Select Embedding Model",
+    labels,
+  );
 
   if (!selected) return;
 
-  let modelId = selected;
+  // Map label back to value
+  const selectedOpt = modelOptions.find(o => `${o.label} — ${o.description}` === selected);
+  let modelId = selectedOpt?.value;
 
-  if (selected === "__custom__") {
-    const custom = await ui.input({
-      title: "Custom Model ID",
-      message: "Enter the OpenRouter model ID:",
-      placeholder: "openai/text-embedding-3-small",
-    });
+  if (!modelId) return;
+
+  if (modelId === "__custom__") {
+    const custom = await ui.input(
+      "Enter the OpenRouter model ID:",
+      "openai/text-embedding-3-small",
+    );
     if (!custom) return;
     modelId = custom.trim();
   }
@@ -235,40 +223,33 @@ async function handleModelSelection(ui: PiUI): Promise<void> {
   config.dimensions = dimensions;
   saveEmbeddingConfig(config);
 
-  await ui.notify({
-    message: `Model set to ${modelId} (${dimensions}d).${hasModelChanged() ? " Re-embed existing memories to use new model." : ""}`,
-    level: "success",
-  });
+  ui.notify(
+    `Model set to ${modelId} (${dimensions}d).${hasModelChanged() ? " Re-embed existing memories to use new model." : ""}`,
+    "info",
+  );
 }
 
 /**
  * Handle dimensions input.
  */
-async function handleDimensionsInput(ui: PiUI): Promise<void> {
+async function handleDimensionsInput(ui: ExtensionCommandContext["ui"]): Promise<void> {
   const config = loadEmbeddingConfig();
 
-  const dimStr = await ui.input({
-    title: "Embedding Dimensions",
-    message: `Enter dimensions (default: 384). Lower = faster, less storage.\nNote: openai/text-embedding-3 supports 256-3072.\nada-002 only supports 1536.`,
-    placeholder: "384",
-    validate: async (value: string) => {
-      const num = parseInt(value, 10);
-      if (isNaN(num) || num < 64 || num > 3072) {
-        return "Must be a number between 64 and 3072";
-      }
-      return null;
-    },
-  });
+  const dimStr = await ui.input(
+    `Enter dimensions (default: 384). Lower = faster, less storage.\nNote: openai/text-embedding-3 supports 256-3072.\nada-002 only supports 1536.`,
+    "384",
+  );
 
   if (dimStr) {
     const dims = parseInt(dimStr, 10);
+    if (isNaN(dims) || dims < 64 || dims > 3072) {
+      ui.notify("Must be a number between 64 and 3072.", "warning");
+      return;
+    }
     config.dimensions = dims;
     saveEmbeddingConfig(config);
 
-    await ui.notify({
-      message: `Dimensions set to ${dims}. Re-embed existing memories to apply.`,
-      level: "success",
-    });
+    ui.notify(`Dimensions set to ${dims}. Re-embed existing memories to apply.`, "info");
   }
 }
 
@@ -276,26 +257,25 @@ async function handleDimensionsInput(ui: PiUI): Promise<void> {
  * Handle re-embedding all memories.
  * This is a destructive operation — warns user first.
  */
-async function handleReembed(ui: PiUI, pi: ExtensionAPI): Promise<void> {
-  const confirm = await ui.select({
-    title: "Re-embed All Memories",
-    message: "⚠ This will re-generate ALL embeddings using the current model.\nOld embeddings will be overwritten.\nThis may take a while and costs API calls.",
-    options: [
-      { label: "Yes, re-embed all", value: "yes", description: "Proceed with re-embedding" },
-      { label: "Cancel", value: "no", description: "Abort" },
-    ],
-  });
-
-  if (confirm !== "yes") return;
+async function handleReembed(ui: ExtensionCommandContext["ui"], ctx: ExtensionCommandContext): Promise<void> {
+  const confirmOptions = [
+    { label: "Yes, re-embed all — Proceed with re-embedding", value: "yes" },
+    { label: "Cancel — Abort", value: "no" },
+  ];
+  const confirmLabels = confirmOptions.map(o => o.label);
+  const confirm = await ui.select(
+    "Re-embed All Memories",
+    confirmLabels,
+  );
+
+  const confirmOpt = confirmOptions.find(o => o.label === confirm);
+  if (confirmOpt?.value !== "yes") return;
 
   // Import here to avoid circular deps
   const { reembedAllMemories } = await import("../embedding.js");
-  const count = await reembedAllMemories(pi);
+  const count = await reembedAllMemories(ctx);
 
   markModelUsed();
 
-  await ui.notify({
-    message: `Re-embedded ${count} memories with current model.`,
-    level: "success",
-  });
+  ui.notify(`Re-embedded ${count} memories with current model.`, "info");
 }