Merge pull request #30 from okieselbach/claude/add-vector-search-mcp-qLyNT

src/McpServer/autopilot-monitor-mcp/package.json

@@ -1,6 +1,6 @@
 {
   "name": "autopilot-monitor-mcp",
-  "version": "1.0.0",
+  "version": "1.2.0",
   "description": "MCP server for Autopilot-Monitor — enables AI agents to query enrollment sessions, device properties, events and vulnerabilities",
   "type": "module",
   "bin": {
@@ -13,7 +13,9 @@
   },
   "dependencies": {
     "@anthropic-ai/sdk": "^0.80.0",
+    "@huggingface/transformers": "^3.8.1",
     "@modelcontextprotocol/sdk": "^1.0.0",
+    "fuse.js": "^7.1.0",
     "zod": "^3.22.0"
   },
   "devDependencies": {

src/McpServer/autopilot-monitor-mcp/src/fuse-search-provider.ts

@@ -0,0 +1,57 @@
+/**
+ * SearchProvider backed by Fuse.js — lightweight fuzzy text search.
+ *
+ * Pros:  Zero model download, instant startup, tiny memory footprint.
+ * Cons:  Token/substring matching only — no true semantic understanding.
+ *        "timeout" will NOT match "waiting exceeded" (use vector for that).
+ */
+
+import Fuse from 'fuse.js';
+import type { SearchDocument, SearchOptions, SearchProvider, SearchResult } from './search-provider.js';
+
+export class FuseSearchProvider implements SearchProvider {
+  readonly name = 'fuse';
+  private documents: SearchDocument[] = [];
+  private fuse: Fuse<SearchDocument> | null = null;
+
+  get size(): number {
+    return this.documents.length;
+  }
+
+  async index(docs: SearchDocument[]): Promise<void> {
+    this.documents.push(...docs);
+    this.fuse = new Fuse(this.documents, {
+      keys: [
+        { name: 'text', weight: 0.7 },
+        { name: 'id', weight: 0.15 },
+        { name: 'metadata.title', weight: 0.1 },
+        { name: 'metadata.category', weight: 0.05 },
+      ],
+      includeScore: true,
+      threshold: 0.5,
+      ignoreLocation: true,
+      useExtendedSearch: true,
+    });
+  }
+
+  async search(query: string, options: SearchOptions = {}): Promise<SearchResult[]> {
+    const { topK = 5, minScore = 0.3 } = options;
+
+    if (!this.fuse || this.documents.length === 0) {
+      return [];
+    }
+
+    const fuseResults = this.fuse.search(query, { limit: topK * 2 });
+
+    return fuseResults
+      .map((r) => ({
+        id: r.item.id,
+        text: r.item.text,
+        metadata: r.item.metadata,
+        // Fuse score is 0 = perfect, 1 = no match. Invert to 0..1 where 1 = best.
+        score: 1 - (r.score ?? 1),
+      }))
+      .filter((r) => r.score >= minScore)
+      .slice(0, topK);
+  }
+}

src/McpServer/autopilot-monitor-mcp/src/index.ts

@@ -1,14 +1,32 @@
+import { resolve, dirname } from 'node:path';
+import { fileURLToPath } from 'node:url';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
 import { registerTools } from './tools.js';
 import { registerResources } from './resources.js';
+import { loadKnowledgeDocs } from './knowledge-base.js';
+import { createSearchProvider } from './search-factory.js';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+
+const RULES_DIR = process.env.RULES_DIR ?? resolve(__dirname, '..', '..', '..', '..', 'rules');
 
 const server = new McpServer({
   name: 'autopilot-monitor',
-  version: '1.0.0',
+  version: '1.2.0',
 });
 
-registerTools(server);
+// Initialize the knowledge base search provider
+console.error('Loading knowledge base documents…');
+const docs = await loadKnowledgeDocs(RULES_DIR);
+
+console.error(`Initializing search provider (${docs.length} documents)…`);
+const knowledgeBase = await createSearchProvider();
+await knowledgeBase.index(docs);
+
+console.error(`Search provider ready: ${knowledgeBase.name} — ${knowledgeBase.size} documents indexed.`);
+
+registerTools(server, knowledgeBase);
 registerResources(server);
 
 const transport = new StdioServerTransport();
@@ -16,3 +34,4 @@ await server.connect(transport);
 
 console.error('Autopilot-Monitor MCP Server running (stdio transport)');
 console.error(`API URL: ${process.env.AUTOPILOT_API_URL ?? 'https://autopilotmonitor-api.azurewebsites.net'}`);
+console.error(`Search backend: ${knowledgeBase.name} | Documents: ${knowledgeBase.size} | Rules dir: ${RULES_DIR}`);

src/McpServer/autopilot-monitor-mcp/src/knowledge-base.ts

@@ -0,0 +1,157 @@
+/**
+ * Loads analysis rules, gather rules, and IME log patterns from the rules/
+ * directory and indexes them into a SearchProvider for semantic or fuzzy search.
+ *
+ * This module is backend-agnostic — it only deals with document loading and
+ * text preparation. The actual search strategy is determined by the provider.
+ */
+
+import { readdir, readFile } from 'node:fs/promises';
+import { join } from 'node:path';
+import type { SearchDocument } from './search-provider.js';
+
+// ── Rule / pattern types ─────────────────────────────────────
+
+interface AnalyzeRule {
+  ruleId: string;
+  title: string;
+  description: string;
+  severity: string;
+  category: string;
+  explanation?: string;
+  remediation?: Array<{ title: string; steps: string[] }>;
+  tags?: string[];
+}
+
+interface GatherRule {
+  ruleId: string;
+  title: string;
+  description: string;
+  category: string;
+  collectorType?: string;
+  target?: string;
+  trigger?: string;
+  tags?: string[];
+}
+
+interface ImeLogPattern {
+  patternId: string;
+  category: string;
+  description: string;
+  pattern: string;
+  action?: string;
+}
+
+// ── File helpers ─────────────────────────────────────────────
+
+async function loadJsonFiles<T>(dir: string): Promise<T[]> {
+  let entries: string[];
+  try {
+    entries = await readdir(dir);
+  } catch {
+    return [];
+  }
+  const results: T[] = [];
+  for (const f of entries) {
+    if (!f.endsWith('.json')) continue;
+    try {
+      const raw = await readFile(join(dir, f), 'utf-8');
+      results.push(JSON.parse(raw) as T);
+    } catch {
+      // skip malformed files
+    }
+  }
+  return results;
+}
+
+// ── Text builders ────────────────────────────────────────────
+
+function analyzeRuleText(r: AnalyzeRule): string {
+  const parts = [
+    `[${r.ruleId}] ${r.title}`,
+    r.description,
+    r.explanation ?? '',
+  ];
+  if (r.remediation) {
+    for (const rem of r.remediation) {
+      parts.push(`Remediation: ${rem.title} — ${rem.steps.join('; ')}`);
+    }
+  }
+  if (r.tags?.length) parts.push(`Tags: ${r.tags.join(', ')}`);
+  return parts.filter(Boolean).join('\n');
+}
+
+function gatherRuleText(r: GatherRule): string {
+  return [
+    `[${r.ruleId}] ${r.title}`,
+    r.description,
+    r.target ? `Target: ${r.target}` : '',
+    r.tags?.length ? `Tags: ${r.tags.join(', ')}` : '',
+  ].filter(Boolean).join('\n');
+}
+
+function imePatternText(p: ImeLogPattern): string {
+  return [
+    `[${p.patternId}] ${p.description}`,
+    `Category: ${p.category}`,
+    p.action ? `Action: ${p.action}` : '',
+  ].filter(Boolean).join('\n');
+}
+
+// ── Public API ───────────────────────────────────────────────
+
+/**
+ * Load all rule/pattern documents from the rules directory.
+ * Returns SearchDocument[] ready to be fed into any SearchProvider.
+ */
+export async function loadKnowledgeDocs(rulesRoot: string): Promise<SearchDocument[]> {
+  const [analyzeRules, gatherRules, imePatterns] = await Promise.all([
+    loadJsonFiles<AnalyzeRule>(join(rulesRoot, 'analyze')),
+    loadJsonFiles<GatherRule>(join(rulesRoot, 'gather')),
+    loadJsonFiles<ImeLogPattern>(join(rulesRoot, 'ime-log-patterns')),
+  ]);
+
+  const docs: SearchDocument[] = [];
+
+  for (const r of analyzeRules) {
+    docs.push({
+      id: r.ruleId,
+      text: analyzeRuleText(r),
+      metadata: {
+        type: 'analyze-rule',
+        severity: r.severity,
+        category: r.category,
+        title: r.title,
+        tags: r.tags ?? [],
+      },
+    });
+  }
+
+  for (const r of gatherRules) {
+    docs.push({
+      id: r.ruleId,
+      text: gatherRuleText(r),
+      metadata: {
+        type: 'gather-rule',
+        category: r.category,
+        title: r.title,
+        tags: r.tags ?? [],
+      },
+    });
+  }
+
+  for (const p of imePatterns) {
+    docs.push({
+      id: p.patternId,
+      text: imePatternText(p),
+      metadata: {
+        type: 'ime-log-pattern',
+        category: p.category,
+        description: p.description,
+        action: p.action ?? null,
+      },
+    });
+  }
+
+  return docs;
+}

src/McpServer/autopilot-monitor-mcp/src/search-factory.ts

@@ -0,0 +1,46 @@
+/**
+ * Factory for creating SearchProvider instances.
+ *
+ * Selection logic:
+ *   1. Explicit env var SEARCH_BACKEND=vector|fuse  →  use that
+ *   2. Auto-detect: try to load @huggingface/transformers
+ *      - available  →  vector
+ *      - missing    →  fuse (graceful fallback)
+ *
+ * This allows running the MCP server without the ~23 MB model download
+ * by setting SEARCH_BACKEND=fuse or simply not installing @huggingface/transformers.
+ */
+
+import type { SearchBackend, SearchProvider } from './search-provider.js';
+
+async function isTransformersAvailable(): Promise<boolean> {
+  try {
+    await import('@huggingface/transformers');
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export async function resolveBackend(): Promise<SearchBackend> {
+  const explicit = process.env.SEARCH_BACKEND?.toLowerCase();
+  if (explicit === 'vector' || explicit === 'fuse') return explicit;
+
+  return (await isTransformersAvailable()) ? 'vector' : 'fuse';
+}
+
+/**
+ * Create a new SearchProvider instance.
+ * Pass a backend explicitly, or let it auto-detect.
+ */
+export async function createSearchProvider(backend?: SearchBackend): Promise<SearchProvider> {
+  const resolved = backend ?? await resolveBackend();
+
+  if (resolved === 'vector') {
+    const { VectorSearchProvider } = await import('./vector-search-provider.js');
+    return new VectorSearchProvider();
+  }
+
+  const { FuseSearchProvider } = await import('./fuse-search-provider.js');
+  return new FuseSearchProvider();
+}

src/McpServer/autopilot-monitor-mcp/src/search-provider.ts

@@ -0,0 +1,54 @@
+/**
+ * Abstract search provider interface.
+ *
+ * Implementations can use vector embeddings, fuzzy text matching, or any other
+ * ranking strategy. The MCP tools program against this interface so the backend
+ * is swappable without touching tool registration code.
+ */
+
+// ── Shared types ─────────────────────────────────────────────
+
+export interface SearchDocument {
+  id: string;
+  text: string;
+  metadata: Record<string, unknown>;
+}
+
+export interface SearchResult {
+  id: string;
+  text: string;
+  metadata: Record<string, unknown>;
+  /** Relevance score — always normalized to 0..1 regardless of backend. */
+  score: number;
+}
+
+export interface SearchOptions {
+  topK?: number;
+  minScore?: number;
+}
+
+// ── Provider contract ────────────────────────────────────────
+
+export interface SearchProvider {
+  /** Human-readable backend name (e.g. "vector/all-MiniLM-L6-v2", "fuse"). */
+  readonly name: string;
+
+  /** Number of currently indexed documents. */
+  readonly size: number;
+
+  /**
+   * Index a batch of documents. Can be called multiple times to add more docs.
+   * Implementation may precompute embeddings, build a Fuse index, etc.
+   */
+  index(docs: SearchDocument[]): Promise<void>;
+
+  /**
+   * Search indexed documents by a natural-language query string.
+   * Returns results sorted by descending relevance.
+   */
+  search(query: string, options?: SearchOptions): Promise<SearchResult[]>;
+}
+
+// ── Provider identifiers ─────────────────────────────────────
+
+export type SearchBackend = 'vector' | 'fuse';

src/McpServer/autopilot-monitor-mcp/src/tools.ts

@@ -1,8 +1,10 @@
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { z } from 'zod';
 import { apiFetch, buildQuery } from './client.js';
+import type { SearchProvider } from './search-provider.js';
+import { createSearchProvider } from './search-factory.js';
 
-export function registerTools(server: McpServer): void {
+export function registerTools(server: McpServer, knowledgeBase?: SearchProvider): void {
   // Tool 1: search_sessions
   server.tool(
     'search_sessions',
@@ -38,7 +40,6 @@ export function registerTools(server: McpServer): void {
     async (args) => {
       const { deviceProperties, ...rest } = args;
       const queryParams: Record<string, string | number | boolean | undefined | null> = { ...rest };
-      // Convert deviceProperties record into prop.* query parameters
       if (deviceProperties) {
         for (const [key, value] of Object.entries(deviceProperties)) {
           queryParams[`prop.${key}`] = value;
@@ -105,7 +106,6 @@ export function registerTools(server: McpServer): void {
         events?: Array<{ eventType?: string; severity?: string; source?: string }>;
         count?: number;
       };
-      // Client-side filter since backend doesn't support these filters on the events endpoint
       if (data?.events && (filters.eventType || filters.severity || filters.source)) {
         data.events = data.events.filter((e) => {
           if (filters.eventType && e.eventType !== filters.eventType) return false;
@@ -167,4 +167,166 @@ export function registerTools(server: McpServer): void {
       return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] };
     }
   );
+
+  // Tool 8: search_events_semantic
+  server.tool(
+    'search_events_semantic',
+    'Semantic/fuzzy search over enrollment event messages within a session or across recent failed sessions. ' +
+    'Unlike exact filters, this finds events by MEANING — e.g. "network timeout" also matches "connection timed out", "request failed after waiting". ' +
+    'Use this when you need to find events that match a symptom description rather than an exact event type. ' +
+    'Provide a sessionId to search within one session, or omit it to search across recent failed sessions.',
+    {
+      query: z.string().describe('Natural language description of what to find (e.g. "app download stuck", "certificate error", "disk space low")'),
+      sessionId: z.string().optional().describe('Search within a specific session. If omitted, searches across recent failed sessions.'),
+      tenantId: z.string().optional().describe('Tenant ID (global-scoped key only)'),
+      topK: z.number().min(1).max(30).optional().default(10).describe('Number of matching events to return (1-30, default 10)'),
+      minScore: z.number().min(0).max(1).optional().default(0.35)
+        .describe('Minimum similarity score (0-1, default 0.35)'),
+    },
+    async ({ query, sessionId, tenantId, topK, minScore }) => {
+      type EventEntry = { eventType?: string; severity?: string; source?: string; message?: string; timestamp?: string; phase?: string; data?: Record<string, unknown> };
+      let events: EventEntry[] = [];
+      let sessionIds: string[] = [];
+
+      if (sessionId) {
+        const q = buildQuery({ tenantId } as Record<string, string | undefined>);
+        const data = await apiFetch(`/api/sessions/${sessionId}/events${q}`) as { events?: EventEntry[] };
+        events = data?.events ?? [];
+        sessionIds = [sessionId];
+      } else {
+        const searchQ = buildQuery({ status: 'Failed', tenantId, limit: 5 } as Record<string, string | number | undefined>);
+        const sessions = await apiFetch(`/api/sessions/search${searchQ}`) as {
+          sessions?: Array<{ sessionId?: string }>;
+        };
+        const ids = (sessions?.sessions ?? []).map((s) => s.sessionId).filter(Boolean) as string[];
+        sessionIds = ids;
+        const q = buildQuery({ tenantId } as Record<string, string | undefined>);
+        const allEvents = await Promise.all(
+          ids.map(async (sid) => {
+            try {
+              const d = await apiFetch(`/api/sessions/${sid}/events${q}`) as { events?: EventEntry[] };
+              return (d?.events ?? []).map((e) => ({ ...e, _sessionId: sid }));
+            } catch { return []; }
+          })
+        );
+        events = allEvents.flat();
+      }
+
+      const candidates = events.filter((e) => e.message && e.message.length > 5);
+
+      if (candidates.length === 0) {
+        return {
+          content: [{
+            type: 'text' as const,
+            text: JSON.stringify({ query, resultCount: 0, results: [], note: 'No events with messages found.' }),
+          }],
+        };
+      }
+
+      // Build searchable documents from events
+      const docs = candidates.map((e, i) => {
+        const parts = [e.message];
+        if (e.eventType) parts.push(`Event: ${e.eventType}`);
+        if (e.severity) parts.push(`Severity: ${e.severity}`);
+        if (e.source) parts.push(`Source: ${e.source}`);
+        return {
+          id: `event-${i}`,
+          text: parts.join(' | '),
+          metadata: { index: i } as Record<string, unknown>,
+        };
+      });
+
+      // Create a temporary search provider for this batch of events
+      const provider = await createSearchProvider();
+      await provider.index(docs);
+      const searchResults = await provider.search(query, { topK, minScore });
+
+      const results = searchResults.map((r) => {
+        const idx = r.metadata.index as number;
+        const e = candidates[idx];
+        return {
+          score: Math.round(r.score * 1000) / 1000,
+          sessionId: (e as Record<string, unknown>)._sessionId ?? sessionId,
+          eventType: e.eventType,
+          severity: e.severity,
+          source: e.source,
+          phase: e.phase,
+          timestamp: e.timestamp,
+          message: e.message,
+        };
+      });
+
+      return {
+        content: [{
+          type: 'text' as const,
+          text: JSON.stringify({
+            query,
+            searchBackend: provider.name,
+            sessionsSearched: sessionIds,
+            eventsAnalyzed: candidates.length,
+            resultCount: results.length,
+            results,
+          }, null, 2),
+        }],
+      };
+    }
+  );
+
+  // Tool 9: search_knowledge
+  server.tool(
+    'search_knowledge',
+    'Semantic/fuzzy search over the Autopilot Monitor knowledge base: analysis rules, gather rules, and IME log patterns. ' +
+    'Use natural language queries like "app install timeout", "BitLocker issues", "detection script failure". ' +
+    'Returns the most relevant rules and patterns ranked by similarity. ' +
+    'Great for finding remediation steps, understanding error patterns, or discovering relevant diagnostic rules.',
+    {
+      query: z.string().describe('Natural language search query (e.g. "app download timeout", "TPM not ready", "ESP stuck")'),
+      topK: z.number().min(1).max(20).optional().default(5).describe('Number of results to return (1-20, default 5)'),
+      type: z.enum(['all', 'analyze-rule', 'gather-rule', 'ime-log-pattern']).optional().default('all')
+        .describe('Filter by document type. Default: search all types.'),
+      minScore: z.number().min(0).max(1).optional().default(0.3)
+        .describe('Minimum similarity score threshold (0-1, default 0.3). Lower = more results, higher = stricter matching.'),
+    },
+    async ({ query, topK, type, minScore }) => {
+      if (!knowledgeBase || knowledgeBase.size === 0) {
+        return {
+          content: [{
+            type: 'text' as const,
+            text: JSON.stringify({ error: 'Knowledge base not initialized. The server may still be loading.' }),
+          }],
+        };
+      }
+
+      // Over-fetch when filtering by type, then trim
+      const fetchK = type === 'all' ? topK : topK * 3;
+      let results = await knowledgeBase.search(query, { topK: fetchK, minScore });
+
+      if (type !== 'all') {
+        results = results.filter((r) => r.metadata.type === type);
+      }
+
+      results = results.slice(0, topK);
+
+      const formatted = results.map((r) => ({
+        id: r.id,
+        score: Math.round(r.score * 1000) / 1000,
+        type: r.metadata.type,
+        title: r.metadata.title ?? r.metadata.description ?? r.id,
+        content: r.text,
+        metadata: r.metadata,
+      }));
+
+      return {
+        content: [{
+          type: 'text' as const,
+          text: JSON.stringify({
+            query,
+            searchBackend: knowledgeBase.name,
+            resultCount: formatted.length,
+            results: formatted,
+          }, null, 2),
+        }],
+      };
+    }
+  );
 }

src/McpServer/autopilot-monitor-mcp/src/vector-search-provider.ts

@@ -0,0 +1,98 @@
+/**
+ * SearchProvider backed by sentence-transformer embeddings + cosine similarity.
+ * Uses @huggingface/transformers with the all-MiniLM-L6-v2 model (quantized, ~23 MB).
+ *
+ * Pros:  True semantic understanding — "timeout" matches "waiting exceeded".
+ * Cons:  First-run downloads the model; higher memory & CPU at index time.
+ */
+
+import type { FeatureExtractionPipeline } from '@huggingface/transformers';
+import type { SearchDocument, SearchOptions, SearchProvider, SearchResult } from './search-provider.js';
+
+const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2';
+
+// ── Shared singleton embedder ────────────────────────────────
+
+let embedder: FeatureExtractionPipeline | null = null;
+
+async function getEmbedder(): Promise<FeatureExtractionPipeline> {
+  if (!embedder) {
+    // Dynamic import + cast to avoid TS2590 from the overloaded pipeline() signature
+    const { pipeline } = await import('@huggingface/transformers');
+    embedder = (await (pipeline as Function)('feature-extraction', MODEL_NAME, {
+      dtype: 'q8',
+    })) as FeatureExtractionPipeline;
+  }
+  return embedder;
+}
+
+/** Compute a normalized embedding for a single text. */
+export async function embed(text: string): Promise<number[]> {
+  const model = await getEmbedder();
+  const output = await model(text, { pooling: 'mean', normalize: true });
+  return Array.from(output.data as Float32Array);
+}
+
+// ── Helpers ──────────────────────────────────────────────────
+
+function cosineSimilarity(a: number[], b: number[]): number {
+  let dot = 0;
+  let normA = 0;
+  let normB = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    normA += a[i] * a[i];
+    normB += b[i] * b[i];
+  }
+  const denom = Math.sqrt(normA) * Math.sqrt(normB);
+  return denom === 0 ? 0 : dot / denom;
+}
+
+// ── Provider implementation ──────────────────────────────────
+
+interface StoredDocument extends SearchDocument {
+  embedding: number[];
+}
+
+export class VectorSearchProvider implements SearchProvider {
+  readonly name = `vector/${MODEL_NAME}`;
+  private documents: StoredDocument[] = [];
+
+  get size(): number {
+    return this.documents.length;
+  }
+
+  async index(docs: SearchDocument[]): Promise<void> {
+    const model = await getEmbedder();
+    const batchSize = 16;
+    for (let i = 0; i < docs.length; i += batchSize) {
+      const batch = docs.slice(i, i + batchSize);
+      const embeddings = await Promise.all(
+        batch.map(async (d) => {
+          const out = await model(d.text, { pooling: 'mean', normalize: true });
+          return Array.from(out.data as Float32Array);
+        })
+      );
+      for (let j = 0; j < batch.length; j++) {
+        this.documents.push({ ...batch[j], embedding: embeddings[j] });
+      }
+    }
+  }
+
+  async search(query: string, options: SearchOptions = {}): Promise<SearchResult[]> {
+    const { topK = 5, minScore = 0.3 } = options;
+    const queryEmbedding = await embed(query);
+
+    const scored = this.documents.map((doc) => ({
+      id: doc.id,
+      text: doc.text,
+      metadata: doc.metadata,
+      score: cosineSimilarity(queryEmbedding, doc.embedding),
+    }));
+
+    return scored
+      .filter((r) => r.score >= minScore)
+      .sort((a, b) => b.score - a.score)
+      .slice(0, topK);
+  }
+}