Azure · feiskyer · Jul 29, 2025 · Jul 25, 2025 · Jul 25, 2025 · Jul 28, 2025
diff --git a/go.mod b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.10.1
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v4 v4.2.1
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v2 v2.4.0
+	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/monitor/armmonitor v0.11.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v2 v2.2.1
 	github.com/Azure/mcp-kubernetes v0.0.5-0.20250724094522-0e7f5ad3fde1
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510

diff --git a/go.sum b/go.sum
@@ -12,12 +12,16 @@ github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontai
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v2 v2.4.0/go.mod h1:U5gpsREQZE6SLk1t/cFfc1eMhYAlYpEzvaYXuDfefy8=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal v1.1.2 h1:mLY+pNLjCUeKhgnAJWAKhEUQM+RJQo2H1fuGSw1Ky1E=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal v1.1.2/go.mod h1:FbdwsQ2EzwvXxOPcMFYO8ogEc9uMMIj3YkmCdXdAFmk=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v2 v2.0.0 h1:PTFGRSlMKCQelWwxUyYVEUqseBJVemLyqWJjvMyt0do=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/internal/v2 v2.0.0/go.mod h1:LRr2FzBTQlONPPa5HREE5+RjSCTXl7BwOvYOaWTqCaI=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/monitor/armmonitor v0.11.0 h1:Ds0KRF8ggpEGg4Vo42oX1cIt/IfOhHWJBikksZbVxeg=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/monitor/armmonitor v0.11.0/go.mod h1:jj6P8ybImR+5topJ+eH6fgcemSFBmU6/6bFF8KkwuDI=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork v1.0.0 h1:nBy98uKOIfun5z6wx6jwWLrULcM0+cjBalBFZlEZ7CA=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork v1.0.0/go.mod h1:243D9iHbcQXoFUtgHJwL7gl2zx1aDuDMjvBZVGr2uW0=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v2 v2.2.1 h1:bWh0Z2rOEDfB/ywv/l0iHN1JgyazE6kW/aIA89+CEK0=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v2 v2.2.1/go.mod h1:Bzf34hhAE9NSxailk8xVeLEZbUjOXcC+GnU1mMKdhLw=
-github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.0.0 h1:ECsQtyERDVz3NP3kvDOTLvbQhqWp/x9EsGKtb4ogUr8=
-github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.0.0/go.mod h1:s1tW/At+xHqjNFvWU4G0c0Qv33KOhvbGNj0RCTQDV8s=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.1.1 h1:7CBQ+Ei8SP2c6ydQTGCCrS35bDxgTMfoP2miAwK++OU=
+github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.1.1/go.mod h1:c/wcGeGx5FUPbM/JltUYHZcKmigwyVLJlDq+4HdtXaw=
 github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEKWjV8V+WSxDXJ4NFATAsZjh8iIbsQIg=
 github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
 github.com/Azure/mcp-kubernetes v0.0.5-0.20250724094522-0e7f5ad3fde1 h1:JLBw146trv/quPn4oEnaiUkG3sef+q602OtXxjeGIYQ=

diff --git a/internal/azureclient/client.go b/internal/azureclient/client.go
@@ -11,21 +11,23 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v4"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v2"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/monitor/armmonitor"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v2"
 )
 
 // SubscriptionClients contains Azure clients for a specific subscription.
 type SubscriptionClients struct {
-	SubscriptionID         string
-	ContainerServiceClient *armcontainerservice.ManagedClustersClient
-	VNetClient             *armnetwork.VirtualNetworksClient
-	SubnetsClient          *armnetwork.SubnetsClient
-	RouteTableClient       *armnetwork.RouteTablesClient
-	NSGClient              *armnetwork.SecurityGroupsClient
-	LoadBalancerClient     *armnetwork.LoadBalancersClient
-	PrivateEndpointsClient *armnetwork.PrivateEndpointsClient
-	VMSSClient             *armcompute.VirtualMachineScaleSetsClient
-	VMSSVMsClient          *armcompute.VirtualMachineScaleSetVMsClient
+	SubscriptionID           string
+	ContainerServiceClient   *armcontainerservice.ManagedClustersClient
+	VNetClient               *armnetwork.VirtualNetworksClient
+	SubnetsClient            *armnetwork.SubnetsClient
+	RouteTableClient         *armnetwork.RouteTablesClient
+	NSGClient                *armnetwork.SecurityGroupsClient
+	LoadBalancerClient       *armnetwork.LoadBalancersClient
+	PrivateEndpointsClient   *armnetwork.PrivateEndpointsClient
+	VMSSClient               *armcompute.VirtualMachineScaleSetsClient
+	VMSSVMsClient            *armcompute.VirtualMachineScaleSetVMsClient
+	DiagnosticSettingsClient *armmonitor.DiagnosticSettingsClient
 }
 
 // AzureClient represents an Azure API client that can handle multiple subscriptions.
@@ -121,18 +123,24 @@ func (c *AzureClient) GetOrCreateClientsForSubscription(subscriptionID string) (
 		return nil, fmt.Errorf("failed to create VMSS VMs client for subscription %s: %v", subscriptionID, err)
 	}
 
+	diagnosticSettingsClient, err := armmonitor.NewDiagnosticSettingsClient(c.credential, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create diagnostic settings client for subscription %s: %v", subscriptionID, err)
+	}
+
 	// Create and store the clients
 	clients = &SubscriptionClients{
-		SubscriptionID:         subscriptionID,
-		ContainerServiceClient: containerServiceClient,
-		VNetClient:             vnetClient,
-		SubnetsClient:          subnetsClient,
-		RouteTableClient:       routeTableClient,
-		NSGClient:              nsgClient,
-		LoadBalancerClient:     loadBalancerClient,
-		PrivateEndpointsClient: privateEndpointsClient,
-		VMSSClient:             vmssClient,
-		VMSSVMsClient:          vmssVMsClient,
+		SubscriptionID:           subscriptionID,
+		ContainerServiceClient:   containerServiceClient,
+		VNetClient:               vnetClient,
+		SubnetsClient:            subnetsClient,
+		RouteTableClient:         routeTableClient,
+		NSGClient:                nsgClient,
+		LoadBalancerClient:       loadBalancerClient,
+		PrivateEndpointsClient:   privateEndpointsClient,
+		VMSSClient:               vmssClient,
+		VMSSVMsClient:            vmssVMsClient,
+		DiagnosticSettingsClient: diagnosticSettingsClient,
 	}
 
 	c.clientsMap[subscriptionID] = clients
@@ -426,3 +434,37 @@ func (c *AzureClient) GetResourceByID(ctx context.Context, resourceID string) (i
 func (c *AzureClient) GetCache() *AzureCache {
 	return c.cache
 }
+
+// GetDiagnosticSettings retrieves diagnostic settings for the specified resource.
+func (c *AzureClient) GetDiagnosticSettings(ctx context.Context, subscriptionID, resourceURI string) ([]*armmonitor.DiagnosticSettingsResource, error) {
+	// Create cache key
+	cacheKey := fmt.Sprintf("resource:diagnosticsettings:%s:%s", subscriptionID, resourceURI)
+
+	// Check cache first
+	if cached, found := c.cache.Get(cacheKey); found {
+		if settings, ok := cached.([]*armmonitor.DiagnosticSettingsResource); ok {
+			return settings, nil
+		}
+	}
+
+	clients, err := c.GetOrCreateClientsForSubscription(subscriptionID)
+	if err != nil {
+		return nil, err
+	}
+
+	pager := clients.DiagnosticSettingsClient.NewListPager(resourceURI, nil)
+	var diagnosticSettings []*armmonitor.DiagnosticSettingsResource
+
+	for pager.More() {
+		page, err := pager.NextPage(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get diagnostic settings: %v", err)
+		}
+		diagnosticSettings = append(diagnosticSettings, page.Value...)
+	}
+
+	// Store in cache
+	c.cache.Set(cacheKey, diagnosticSettings)
+
+	return diagnosticSettings, nil
+}
diff --git a/internal/components/monitor/diagnostics/handlers.go b/internal/components/monitor/diagnostics/handlers.go
@@ -1,11 +1,13 @@
 package diagnostics
 
 import (
+	"context"
+	"encoding/json"
 	"fmt"
 	"log"
-	"strings"
 
 	"github.com/Azure/aks-mcp/internal/azcli"
+	"github.com/Azure/aks-mcp/internal/azureclient"
 	"github.com/Azure/aks-mcp/internal/components/common"
 	"github.com/Azure/aks-mcp/internal/config"
 	"github.com/Azure/aks-mcp/internal/tools"
@@ -18,7 +20,7 @@ func buildClusterResourceID(subscriptionID, resourceGroup, clusterName string) s
 }
 
 // HandleControlPlaneDiagnosticSettings checks diagnostic settings for AKS cluster
-func HandleControlPlaneDiagnosticSettings(params map[string]interface{}, cfg *config.ConfigData) (string, error) {
+func HandleControlPlaneDiagnosticSettings(params map[string]interface{}, azClient *azureclient.AzureClient, cfg *config.ConfigData) (string, error) {
 	// Extract and validate parameters using common helper
 	subscriptionID, resourceGroup, clusterName, err := common.ExtractAKSParameters(params)
 	if err != nil {
@@ -28,29 +30,29 @@ func HandleControlPlaneDiagnosticSettings(params map[string]interface{}, cfg *co
 	// Build cluster resource ID using utility function
 	clusterResourceID := buildClusterResourceID(subscriptionID, resourceGroup, clusterName)
 
-	// Execute Azure CLI command to get diagnostic settings
-	executor := azcli.NewExecutor()
-	args := []string{
-		"monitor", "diagnostic-settings", "list",
-		"--resource", clusterResourceID,
-		"--output", "json",
+	// Azure client is required
+	if azClient == nil {
+		return "", fmt.Errorf("azure client is required but not provided")
 	}
 
-	cmdParams := map[string]interface{}{
-		"command": "az " + strings.Join(args, " "),
+	// Get diagnostic settings using Azure SDK
+	ctx := context.Background()
+	diagnosticSettings, err := azClient.GetDiagnosticSettings(ctx, subscriptionID, clusterResourceID)
+	if err != nil {
+		return "", fmt.Errorf("failed to get diagnostic settings for cluster %s in resource group %s: %w", clusterName, resourceGroup, err)
 	}
 
-	result, err := executor.Execute(cmdParams, cfg)
+	// Convert to JSON for backward compatibility
+	result, err := json.Marshal(diagnosticSettings)
 	if err != nil {
-		return "", fmt.Errorf("failed to get diagnostic settings for cluster %s in resource group %s: %w", clusterName, resourceGroup, err)
+		return "", fmt.Errorf("failed to marshal diagnostic settings to JSON: %w", err)
 	}
 
-	// Return raw JSON result from Azure CLI
-	return result, nil
+	return string(result), nil
 }
 
 // HandleControlPlaneLogs queries specific control plane logs
-func HandleControlPlaneLogs(params map[string]interface{}, cfg *config.ConfigData) (string, error) {
+func HandleControlPlaneLogs(params map[string]interface{}, azClient *azureclient.AzureClient, cfg *config.ConfigData) (string, error) {
 	// Extract and validate AKS parameters using common helper
 	subscriptionID, resourceGroup, clusterName, err := common.ExtractAKSParameters(params)
 	if err != nil {
@@ -71,7 +73,7 @@ func HandleControlPlaneLogs(params map[string]interface{}, cfg *config.ConfigDat
 
 	// Find the diagnostic setting that has the requested log category enabled
 	// This handles cases where multiple diagnostic settings exist for the same cluster
-	workspaceResourceID, isResourceSpecific, err := FindDiagnosticSettingForCategory(subscriptionID, resourceGroup, clusterName, logCategory, cfg)
+	workspaceResourceID, isResourceSpecific, err := FindDiagnosticSettingForCategory(subscriptionID, resourceGroup, clusterName, logCategory, azClient, cfg)
 	if err != nil {
 		return "", fmt.Errorf("failed to find diagnostic setting for log category %s in cluster %s: %w", logCategory, clusterName, err)
 	}
@@ -123,15 +125,15 @@ func HandleControlPlaneLogs(params map[string]interface{}, cfg *config.ConfigDat
 // Resource handler functions for control plane diagnostics tools
 
 // GetControlPlaneDiagnosticSettingsHandler returns handler for diagnostic settings tool
-func GetControlPlaneDiagnosticSettingsHandler(cfg *config.ConfigData) tools.ResourceHandler {
+func GetControlPlaneDiagnosticSettingsHandler(azClient *azureclient.AzureClient, cfg *config.ConfigData) tools.ResourceHandler {
 	return tools.ResourceHandlerFunc(func(params map[string]interface{}, _ *config.ConfigData) (string, error) {
-		return HandleControlPlaneDiagnosticSettings(params, cfg)
+		return HandleControlPlaneDiagnosticSettings(params, azClient, cfg)
 	})
 }
 
 // GetControlPlaneLogsHandler returns handler for logs querying tool
-func GetControlPlaneLogsHandler(cfg *config.ConfigData) tools.ResourceHandler {
+func GetControlPlaneLogsHandler(azClient *azureclient.AzureClient, cfg *config.ConfigData) tools.ResourceHandler {
 	return tools.ResourceHandlerFunc(func(params map[string]interface{}, _ *config.ConfigData) (string, error) {
-		return HandleControlPlaneLogs(params, cfg)
+		return HandleControlPlaneLogs(params, azClient, cfg)
 	})
 }
diff --git a/internal/components/monitor/diagnostics/handlers_test.go b/internal/components/monitor/diagnostics/handlers_test.go
@@ -14,7 +14,7 @@ func TestGetControlPlaneDiagnosticSettingsHandler(t *testing.T) {
 			AccessLevel: "readonly",
 		},
 	}
-	handler := GetControlPlaneDiagnosticSettingsHandler(cfg)
+	handler := GetControlPlaneDiagnosticSettingsHandler(nil, cfg) // Pass nil Azure client for testing
 
 	if handler == nil {
 		t.Error("Expected handler to be created, got nil")
@@ -38,7 +38,7 @@ func TestGetControlPlaneLogsHandler(t *testing.T) {
 			AccessLevel: "readonly",
 		},
 	}
-	handler := GetControlPlaneLogsHandler(cfg)
+	handler := GetControlPlaneLogsHandler(nil, cfg) // Pass nil Azure client for testing
 
 	if handler == nil {
 		t.Error("Expected handler to be created, got nil")
@@ -119,7 +119,7 @@ func TestHandleControlPlaneDiagnosticSettings_ParameterValidation(t *testing.T)
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			_, err := HandleControlPlaneDiagnosticSettings(tt.params, cfg)
+			_, err := HandleControlPlaneDiagnosticSettings(tt.params, nil, cfg) // Pass nil Azure client for testing
 
 			if tt.wantError {
 				if err == nil {
@@ -130,9 +130,9 @@ func TestHandleControlPlaneDiagnosticSettings_ParameterValidation(t *testing.T)
 					t.Errorf("Expected error to contain '%s', got '%s'", tt.errorMsg, err.Error())
 				}
 			} else {
-				// Note: This test will fail in unit testing because it tries to execute Azure CLI
-				// In a real unit test environment, we would mock the Azure CLI executor
-				if err != nil && !strings.Contains(err.Error(), "failed to get diagnostic settings") {
+				// Note: This test will now fail because Azure client is required
+				// In a real unit test environment, we would provide a valid Azure client
+				if err != nil && !strings.Contains(err.Error(), "azure client is required but not provided") {
 					t.Errorf("Unexpected error type: %v", err)
 				}
 			}
@@ -149,7 +149,7 @@ func TestHandleControlPlaneLogs_ParameterValidation(t *testing.T) {
 
 	// Test with invalid params to ensure validation works
 	params := map[string]interface{}{}
-	_, err := HandleControlPlaneLogs(params, cfg)
+	_, err := HandleControlPlaneLogs(params, nil, cfg) // Pass nil Azure client for testing
 	if err == nil {
 		t.Error("Expected error for missing parameters, got nil")
 	}

diff --git a/internal/components/monitor/diagnostics/integration_test.go b/internal/components/monitor/diagnostics/integration_test.go
@@ -19,7 +19,7 @@ func TestDiagnosticSettings_Integration(t *testing.T) {
 
 	// Test with invalid params to ensure delegation works
 	params := map[string]interface{}{}
-	_, err := HandleControlPlaneDiagnosticSettings(params, cfg)
+	_, err := HandleControlPlaneDiagnosticSettings(params, nil, cfg) // Pass nil Azure client for testing
 	if err == nil {
 		t.Error("Expected error for missing parameters, got nil")
 	}
@@ -38,7 +38,7 @@ func TestControlPlaneLogs_Integration(t *testing.T) {
 
 	// Test with invalid params to ensure delegation works
 	params := map[string]interface{}{}
-	_, err := HandleControlPlaneLogs(params, cfg)
+	_, err := HandleControlPlaneLogs(params, nil, cfg) // Pass nil Azure client for testing
 	if err == nil {
 		t.Error("Expected error for missing parameters, got nil")
 	}
@@ -54,7 +54,7 @@ func TestDiagnosticSettingsHandler_Integration(t *testing.T) {
 			AccessLevel: "readonly",
 		},
 	}
-	handler := GetControlPlaneDiagnosticSettingsHandler(cfg)
+	handler := GetControlPlaneDiagnosticSettingsHandler(nil, cfg) // Pass nil Azure client for testing
 
 	if handler == nil {
 		t.Error("Expected handler to be created, got nil")
@@ -78,7 +78,7 @@ func TestControlPlaneLogsHandler_Integration(t *testing.T) {
 			AccessLevel: "readonly",
 		},
 	}
-	handler := GetControlPlaneLogsHandler(cfg)
+	handler := GetControlPlaneLogsHandler(nil, cfg) // Pass nil Azure client for testing
 
 	if handler == nil {
 		t.Error("Expected handler to be created, got nil")