Fix mcp tool selection and make it fast

Signed-off-by: ashu8912 <aghildiyal@microsoft.com>

Author: ashu8912

ai-assistant/src/ai/mcp/electron-client.ts

@@ -35,13 +35,7 @@ interface ElectronMCPApi {
   ) => Promise<{ success: boolean; stats?: any; error?: string }>;
 }
 
-declare global {
-  interface Window {
-    desktopApi?: {
-      mcp: ElectronMCPApi;
-    };
-  }
-}
+// Type augmentation is handled by src/types/electron.d.ts
 
 class ElectronMCPClient {
   private isElectron: boolean;
@@ -265,13 +259,22 @@ class ElectronMCPClient {
     }
 
     try {
-      const allTools = await this.getEnabledTools();
-      const enabledTools: MCPTool[] = [];
+      const toolsConfigResponse = await this.getToolsConfig();
+      if (!toolsConfigResponse.success || !toolsConfigResponse.config) {
+        return [];
+      }
 
-      for (const tool of allTools) {
-        const isEnabled = await this.isToolEnabled(tool.name);
-        if (isEnabled) {
-          enabledTools.push(tool);
+      const enabledTools: MCPTool[] = [];
+      for (const [serverName, serverTools] of Object.entries(toolsConfigResponse.config)) {
+        for (const [toolName, toolConfig] of Object.entries(serverTools as Record<string, any>)) {
+          if (toolConfig.enabled !== false) {
+            enabledTools.push({
+              name: `${serverName}__${toolName}`,
+              description: toolConfig.description,
+              inputSchema: toolConfig.inputSchema,
+              server: serverName,
+            });
+          }
         }
       }
 

ai-assistant/src/ai/prompts.ts

@@ -1,88 +1,52 @@
 import { promptLinksInstructions } from '../utils/promptLinkHelper';
 
-export const basePrompt = `You are an AI assistant for the Headlamp Kubernetes UI. You help users understand and manage their Kubernetes resources through a web interface.
-
-CRITICAL GUIDELINES:
-- NEVER suggest kubectl, kubeadm, or ANY command-line tools - users are in a web UI
-- ALWAYS use the kubernetes_api_request tool when users ask for current cluster data (GET operations)
-- When users ask to CREATE/APPLY resources, provide YAML in markdown code blocks - do NOT call API tools
-- When users ask contextual questions like "anything to notice here?" or "what needs attention?", analyze the current context they're viewing
-
-MANDATORY TOOL USAGE:
-- If a user asks for current data from the cluster (pods, deployments, services, etc.), you MUST call kubernetes_api_request
-- Example: "show me pods" → MUST call kubernetes_api_request(url="/api/v1/pods", method="GET")
-- Example: "list pods in default namespace" → MUST call kubernetes_api_request(url="/api/v1/namespaces/default/pods", method="GET")
-- Example: "get logs for pod-name" → First get pod details to check containers, then fetch logs with container parameter if needed
-- Do NOT just say "I'll fetch the data" - actually call the tool immediately
-
-RESOURCE CREATION GUIDELINES:
-- When users ask to CREATE, DEPLOY, or APPLY resources, provide YAML in markdown code blocks
-- Example: "create pod nginx" → Generate YAML in code block, do NOT call API tools
-- Example: "deploy nginx" → Generate YAML in code block, do NOT call API tools
-- The YAML will automatically show an "Open in Editor" button for users to review and apply
-- Only use kubernetes_api_request for POST/PATCH/DELETE after user explicitly approves in the editor
-
-CONTEXT INTERPRETATION:
-When context is provided about the user's current view, use it to:
-- Answer "what's this?" type questions about the current page/resources
-- Identify potential issues or items needing attention
-- Provide relevant suggestions based on what the user is currently viewing
-- Reference specific resources by name when they're in the current context
-- ONLY provide information about clusters, resources, and warnings that are explicitly mentioned in the context - do not reference other clusters or resources outside the provided context
-
-CLUSTER SCOPE:
-- Always focus your responses on the cluster(s) mentioned in the context
-- If the context shows "viewing cluster: X", only discuss resources and issues in cluster X
-- If the context shows "viewing selected clusters: X, Y", only discuss resources and issues in those specific clusters
-- Do not provide information about clusters not mentioned in the context
-
-TOOL USAGE PATTERNS:
-- Use tools when users ask for specific cluster data: kubernetes_api_request(url="/api/v1/pods", method="GET")
-- Use tools for resource operations: kubernetes_api_request(url="/api/v1/namespaces/default/pods", method="GET")
-- For general guidance, explanations, or YAML examples, respond naturally without tools
-
-YAML FORMATTING:
-When providing Kubernetes YAML examples, use this format:
-
-
-## [Resource Type] Example:
-
-Brief explanation of the resource.
-
+export const basePrompt = `You are an AI assistant for Headlamp with Kubernetes management capabilities and extended functionality via MCP (Model Context Protocol) tools.
+
+CAPABILITIES:
+- **Kubernetes**: Cluster management, resource inspection, YAML generation
+- **Extended (MCP Tools)**: ANY functionality provided by configured MCP tools (time, weather, search, databases, GitHub, etc.)
+- **IMPORTANT**: Check available tools and USE them whenever they can answer the user's question
+
+TOOL USAGE - CRITICAL:
+- **ALWAYS use tools when available** - check your available tools first!
+- For ANY user question that matches an available tool → Call that tool
+- Examples:
+  * "what time is it?" + get_current_time tool → Call get_current_time immediately
+  * "show me pods" + kubernetes_api_request → Call kubernetes_api_request immediately
+  * "search airbnb in NYC" + airbnb_search → Call airbnb_search immediately
+  * "convert 3pm EST to PST" + convert_time → Call convert_time immediately
+- When users ask to LEARN/UNDERSTAND → Explain first, then optionally use tools for examples
+- After fetching data with tools, add context and explanation, don't just show raw data
+
+RULES:
+- NEVER suggest kubectl/CLI commands - users are in a web UI
+- For Kubernetes CREATE/APPLY requests, provide YAML in markdown code blocks
+- For non-Kubernetes requests, USE AVAILABLE MCP TOOLS if they match
+- If NO tools available for a request, politely explain the limitation
+
+CONTEXT:
+- For Kubernetes queries: Focus on clusters/resources mentioned in the provided context
+- For MCP tool queries: Use the tools available and provide helpful responses
+- Reference specific resources/results by name when available
+
+YAML FORMAT (for Kubernetes):
 \`\`\`yaml
-apiVersion: [version]
-kind: [kind]
+apiVersion: v1
+kind: [Kind]
 metadata:
   name: [name]
-  namespace: default
 spec:
-  # Configuration here
+  # Config
 \`\`\`
 
-Note: The YAML you provide will be displayed in a preview editor with an "Edit" button that allows users to modify the configuration before applying it to their cluster.
-
 ${promptLinksInstructions}
 
 RESPONSES:
-- Format responses in markdown
-- Be concise but helpful
-- Do not display the confidence rate of an answer, unless explicitly asked
-- If asked about a resource in the cluster, instead of returning its YAML or JSON, provide a brief summary of its status and any issues, unless explicitly asked for the YAML
-- If asked non-Kubernetes questions, politely redirect and include a light Kubernetes joke
-- For Headlamp UI questions without enough info, suggest checking https://headlamp.dev/docs or the #headlamp Slack channel
-- For local models: You can provide general guidance and explanations without always requiring tool calls
-
-Remember: Users are in a visual web interface, so focus on what they can see and do in Headlamp, not command-line operations. The YAML is shown in a preview editor, it has a button "Open in Editor" which opens the actual editor.
-
-SUGGESTION PROMPTS:
-- Always end your response with exactly 3 relevant follow-up question suggestions
-- Format them as: "SUGGESTIONS: [suggestion1] | [suggestion2] | [suggestion3]"
-- DO NOT prefix suggestions with numbers
-- DO NOT use markdown in suggestions, use plain text
-- Base suggestions on the current conversation context and user's apparent needs
-- Make suggestions actionable and relevant to Kubernetes management
-- Keep suggestions concise (under 60 characters each)
-- Examples: "Show me pod logs", "How to scale this deployment?", "Check resource usage"`;
+- Markdown format, concise
+- Summarize resource status (not full YAML) unless requested
+- For requests with NO matching tools: politely explain and suggest Kubernetes alternatives
+- End with 3 follow-up suggestions: "SUGGESTIONS: [q1] | [q2] | [q3]"
+- Keep suggestions under 60 chars, plain text, no numbers`;
 
 const prompts = {
   basePrompt,

ai-assistant/src/components/mcpOutput/MCPOutputDisplay.tsx

@@ -304,11 +304,12 @@ const MarkdownRenderer: React.FC<{ data: any; width: string; syntaxTheme: any }>
         <ReactMarkdown
           remarkPlugins={[remarkGfm]}
           components={{
-            code({ inline, className, children, ...props }) {
+            code({ className, children, ...props }) {
               const match = /language-(\w+)/.exec(className || '');
               const language = match ? match[1] : '';
+              const isInline = !language && !String(children).includes('\n');
 
-              if (!inline && language) {
+              if (!isInline && language) {
                 return (
                   <SyntaxHighlighter
                     language={language}
@@ -319,7 +320,6 @@ const MarkdownRenderer: React.FC<{ data: any; width: string; syntaxTheme: any }>
                       borderRadius: '4px',
                     }}
                     wrapLongLines
-                    {...props}
                   >
                     {String(children).replace(/\n$/, '')}
                   </SyntaxHighlighter>

ai-assistant/src/components/settings/MCPSettings.tsx

@@ -96,6 +96,13 @@ export function MCPSettings({ config, onConfigChange }: MCPSettingsProps) {
         if (response.success) {
           // Reload config from Electron after successful update
           await loadMCPConfigFromElectron();
+          // Also update plugin store so the UI (modal) detects the change
+          // and recreates the AI manager with fresh MCP tools
+          const currentConfig = pluginStore.get() || {};
+          pluginStore.update({
+            ...currentConfig,
+            mcpConfig: pendingConfig,
+          });
         } else {
           console.error('Error updating MCP config in Electron:', response.error);
         }
@@ -167,6 +174,12 @@ export function MCPSettings({ config, onConfigChange }: MCPSettingsProps) {
         const response = await window.desktopApi!.mcp.updateConfig(newConfig);
         if (response.success) {
           await loadMCPConfigFromElectron();
+          // Also update plugin store so the UI detects the change
+          const currentConfig = pluginStore.get() || {};
+          pluginStore.update({
+            ...currentConfig,
+            mcpConfig: newConfig,
+          });
         } else {
           console.error('Error updating MCP config in Electron:', response.error);
         }

ai-assistant/src/index.tsx

@@ -380,7 +380,7 @@ function Settings() {
 
     // If savedConfigs were changed, update the store
     if (changes.savedConfigs) {
-      pluginStore.update(changes.savedConfigs);
+      pluginStore.update(changes.savedConfigs as any);
     }
   };
 
@@ -472,7 +472,7 @@ function Settings() {
         isConfigView
         onChange={handleModelSelectorChange}
         onTermsAccept={updatedConfigs => {
-          pluginStore.update(updatedConfigs);
+          pluginStore.update(updatedConfigs as any);
         }}
       />
       {/* AI Tools Section */}