Merge pull request #124 from AxmeAI/feat/semantic-search-mcp-tool-20260429

feat(search): semantic search MCP tools + tiered context mode (B-005 Phase 2)

Author: George-iam

src/cli.ts

@@ -796,6 +796,89 @@ Do NOT skip — without context you will miss critical project rules.
       process.exit(1);
     }
 
+    case "config": {
+      // axme-code config get <key>
+      // axme-code config set <key> <value>
+      // Currently supported keys: context.mode (full|search). The set path
+      // for context.mode = search atomically installs the transformers
+      // runtime + builds the initial embeddings index, and rolls the
+      // config back to "full" if either step fails (D-136 / B-005 design).
+      const sub = args[1];
+      const key = args[2];
+      const value = args[3];
+      const projectPath = resolve(".");
+      const { readConfig: rc, writeConfig: wc } = await import("./storage/config.js");
+
+      if (sub === "get") {
+        if (!key) {
+          console.error("usage: axme-code config get <key>  (e.g. context.mode)");
+          process.exit(1);
+        }
+        const cfg = rc(projectPath);
+        if (key === "context.mode") console.log(cfg.contextMode);
+        else if (key === "model") console.log(cfg.model);
+        else if (key === "auditor_model") console.log(cfg.auditorModel);
+        else if (key === "review_enabled") console.log(String(cfg.reviewEnabled));
+        else { console.error(`Unknown config key: ${key}`); process.exit(1); }
+        break;
+      }
+
+      if (sub === "set") {
+        if (!key || value === undefined) {
+          console.error("usage: axme-code config set <key> <value>");
+          process.exit(1);
+        }
+        if (key !== "context.mode") {
+          console.error(`Set is currently supported only for context.mode. Got: ${key}`);
+          process.exit(1);
+        }
+        if (value !== "full" && value !== "search") {
+          console.error(`context.mode must be 'full' or 'search'. Got: ${value}`);
+          process.exit(1);
+        }
+        const cfg = rc(projectPath);
+        const prevMode = cfg.contextMode;
+
+        if (value === "full") {
+          wc(projectPath, { ...cfg, contextMode: "full" });
+          console.log("Saved: context.mode = full");
+          break;
+        }
+
+        // value === "search" — install runtime if missing, then reindex.
+        const { runConfigSetSearch } = await import("./tools/search-install.js");
+        const result = await runConfigSetSearch(projectPath);
+        if (result.ok) {
+          wc(projectPath, { ...cfg, contextMode: "search" });
+          console.log(`Saved: context.mode = search (indexed ${result.indexed} entries)`);
+        } else {
+          // Rollback config to whatever it was before — we never changed it
+          // in the failure path, but be explicit so future edits don't drop
+          // this guarantee.
+          wc(projectPath, { ...cfg, contextMode: prevMode });
+          console.error(`\nFailed to enable search mode: ${result.error}`);
+          console.error(`Config left at context.mode = ${prevMode}.`);
+          process.exit(1);
+        }
+        break;
+      }
+
+      console.error("Unknown 'config' subcommand. Available: get <key>, set <key> <value>");
+      process.exit(1);
+    }
+
+    case "reindex": {
+      // Rebuild the embeddings index from every memory + decision on disk.
+      // No-op if context.mode = full and runtime is missing — caller should
+      // run `axme-code config set context.mode search` first.
+      const projectPath = resolve(args[1] || ".");
+      const { reindexAll } = await import("./tools/search-install.js");
+      const result = await reindexAll(projectPath);
+      if (result.ok) console.log(`Reindexed ${result.indexed} entries.`);
+      else { console.error(`Reindex failed: ${result.error}`); process.exit(1); }
+      break;
+    }
+
     case "help":
     case "--help":
     case "-h":

src/server.ts

@@ -22,6 +22,9 @@ import { saveMemoryTool } from "./tools/memory-tools.js";
 import { saveDecisionTool } from "./tools/decision-tools.js";
 import { updateSafetyTool, showSafetyTool } from "./tools/safety-tools.js";
 import { statusTool, worklogTool } from "./tools/status.js";
+import { getMemoryTool, getDecisionTool, searchKbTool } from "./tools/kb-search.js";
+import { embedKbEntry } from "./storage/embeddings.js";
+import { readConfig } from "./storage/config.js";
 import { detectWorkspace } from "./utils/workspace-detector.js";
 import {
   findOrphanSessions,
@@ -432,6 +435,10 @@ server.tool(
     const sid = getOwnedSessionIdForLogging();
     const resolved = ppWithScope(project_path, scope);
     const result = saveMemoryTool(resolved, { type, title, description, body, keywords, scope }, sid);
+    // Update the embeddings index when search mode is on. Awaited so the
+    // index is consistent on return; ~50-200ms once the embedder is warm.
+    // Skips silently in full mode and on missing runtime.
+    await embedKbEntry(resolved, result.slug, "memory", title, description, readConfig(resolved).contextMode);
     return { content: [{ type: "text" as const, text: `Memory saved: ${result.slug} (${type}) -> ${resolved}` }] };
   },
 );
@@ -452,6 +459,9 @@ server.tool(
 
     const resolved = ppWithScope(project_path, scope);
     const result = saveDecisionTool(resolved, { title, decision, reasoning, enforce, scope });
+    // Use decision text as description so the search index returns hits
+    // ranked by the actual rule, not just the title.
+    await embedKbEntry(resolved, result.id, "decision", title, decision, readConfig(resolved).contextMode);
     return { content: [{ type: "text" as const, text: `Decision saved: ${result.id} - ${title} -> ${resolved}` }] };
   },
 );
@@ -485,6 +495,48 @@ server.tool(
   },
 );
 
+// --- axme_get_memory ---
+server.tool(
+  "axme_get_memory",
+  "Fetch the full body of one memory by slug. Use after seeing the slug in axme_context (search mode catalog) or axme_search_kb results.",
+  {
+    project_path: z.string().optional().describe("Absolute path to the project root (defaults to server cwd)"),
+    slug: z.string().describe("Memory slug, e.g. 'always-call-axme-context-first'"),
+  },
+  async ({ project_path, slug }) => {
+    return { content: [{ type: "text" as const, text: getMemoryTool(pp(project_path), slug) }] };
+  },
+);
+
+// --- axme_get_decision ---
+server.tool(
+  "axme_get_decision",
+  "Fetch the full body of one decision by ID (e.g. 'D-110') or slug. Use after seeing it in axme_context (search mode catalog) or axme_search_kb results.",
+  {
+    project_path: z.string().optional().describe("Absolute path to the project root (defaults to server cwd)"),
+    id_or_slug: z.string().describe("Decision ID like 'D-110' or its slug"),
+  },
+  async ({ project_path, id_or_slug }) => {
+    return { content: [{ type: "text" as const, text: getDecisionTool(pp(project_path), id_or_slug) }] };
+  },
+);
+
+// --- axme_search_kb ---
+server.tool(
+  "axme_search_kb",
+  "Semantic search across memories and decisions. Useful for fuzzy lookups mid-session ('how did we handle X?'). Requires the embeddings runtime — install with `axme-code config set context.mode search`.",
+  {
+    project_path: z.string().optional().describe("Absolute path to the project root (defaults to server cwd)"),
+    query: z.string().describe("Search query in natural language"),
+    k: z.number().int().min(1).max(50).optional().describe("Top K results to return (default 5, max 50)"),
+    type: z.enum(["memory", "decision"]).optional().describe("Filter results to one type. Omit to search both."),
+  },
+  async ({ project_path, query, k, type }) => {
+    const text = await searchKbTool(pp(project_path), { query, k, type });
+    return { content: [{ type: "text" as const, text }] };
+  },
+);
+
 // --- axme_backlog ---
 server.tool(
   "axme_backlog",

src/storage/config.ts

@@ -49,13 +49,27 @@ function formatConfig(config: ProjectConfig): string {
     `presets:`,
     ...config.presets.map(p => `  - ${p}`),
     "",
+    "# Context-loading mode at session start.",
+    "#   full   — every memory and decision body loaded (default; best for KBs <=100 entries)",
+    "#   search — only catalog loaded, bodies fetched via axme_get_memory / axme_get_decision /",
+    "#            axme_search_kb. Recommended for KBs >100 entries. Requires embeddings runtime,",
+    "#            installed by: axme-code config set context.mode search",
+    "context:",
+    `  mode: ${config.contextMode}`,
+    "",
   ].join("\n");
 }
 
 function parseConfig(content: string): ProjectConfig {
   const doc = yaml.load(content) as Record<string, any> | null;
   if (!doc || typeof doc !== "object") return { ...DEFAULT_PROJECT_CONFIG };
 
+  let contextMode: "full" | "search" = DEFAULT_PROJECT_CONFIG.contextMode;
+  const ctxRaw = doc.context;
+  if (ctxRaw && typeof ctxRaw === "object" && (ctxRaw.mode === "full" || ctxRaw.mode === "search")) {
+    contextMode = ctxRaw.mode;
+  }
+
   return {
     model: String(doc.model ?? DEFAULT_PROJECT_CONFIG.model),
     auditorModel: String(doc.auditor_model ?? DEFAULT_PROJECT_CONFIG.auditorModel),
@@ -71,5 +85,6 @@ function parseConfig(content: string): ProjectConfig {
           return true; // keep all, just warn
         })
       : DEFAULT_PROJECT_CONFIG.presets,
+    contextMode,
   };
 }