NE-2451: Implement probe_dns_local diagnostic tool

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/go-logr/logr v1.4.3
 	github.com/google/gnostic-models v0.7.1
 	github.com/google/jsonschema-go v0.4.2
+	github.com/miekg/dns v1.1.57
 	github.com/modelcontextprotocol/go-sdk v1.4.0
 	github.com/prometheus/client_golang v1.23.2
 	github.com/spf13/afero v1.15.0
@@ -139,11 +140,13 @@ require (
 	go.yaml.in/yaml/v2 v2.4.3 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect
 	golang.org/x/crypto v0.48.0 // indirect
+	golang.org/x/mod v0.32.0 // indirect
 	golang.org/x/net v0.51.0 // indirect
 	golang.org/x/sys v0.41.0 // indirect
 	golang.org/x/term v0.40.0 // indirect
 	golang.org/x/text v0.34.0 // indirect
 	golang.org/x/time v0.12.0 // indirect
+	golang.org/x/tools v0.41.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20260209200024-4cfbd4190f57 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20260209200024-4cfbd4190f57 // indirect
 	google.golang.org/grpc v1.79.2 // indirect

pkg/toolsets/netedge/probe_dns_local.go

@@ -0,0 +1,137 @@
+package netedge
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+	"strings"
+	"time"
+
+	"github.com/containers/kubernetes-mcp-server/pkg/api"
+	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/miekg/dns"
+	"k8s.io/utils/ptr"
+)
+
+func initProbeDNSLocal() []api.ServerTool {
+	return []api.ServerTool{
+		{
+			Tool: api.Tool{
+				Name:        "probe_dns_local",
+				Description: "Run a DNS query using local libraries on the MCP server host to verify connectivity and resolution.",
+				InputSchema: &jsonschema.Schema{
+					Type: "object",
+					Properties: map[string]*jsonschema.Schema{
+						"server": {
+							Type:        "string",
+							Description: "DNS server IP (e.g. 8.8.8.8, 10.0.0.10)",
+						},
+						"name": {
+							Type:        "string",
+							Description: "FQDN to query",
+						},
+						"type": {
+							Type:        "string",
+							Description: "Record type (A, AAAA, CNAME, TXT, SRV, etc.). Defaults to A.",
+							Default:     json.RawMessage(`"A"`),
+						},
+					},
+					Required: []string{"server", "name"},
+				},
+				Annotations: api.ToolAnnotations{
+					Title:           "Probe DNS (Local)",
+					ReadOnlyHint:    ptr.To(true),
+					DestructiveHint: ptr.To(false),
+					OpenWorldHint:   ptr.To(true),
+				},
+			},
+			ClusterAware: ptr.To(false),
+			Handler:      probeDNSLocalHandler,
+		},
+	}
+}
+
+// dnsExchange interface allows mocking the dns client for tests.
+type dnsExchange interface {
+	Exchange(m *dns.Msg, a string) (r *dns.Msg, rtt time.Duration, err error)
+}
+
+// defaultDNSClient wraps the miekg/dns client.
+type defaultDNSClient struct {
+	client *dns.Client
+}
+
+func (d *defaultDNSClient) Exchange(m *dns.Msg, a string) (*dns.Msg, time.Duration, error) {
+	return d.client.Exchange(m, a)
+}
+
+// activeDNSClient is the client used by the handler, allows injection during tests
+var activeDNSClient dnsExchange = &defaultDNSClient{client: new(dns.Client)}
+
+// DNSResult represents the required JSON response format for probe_dns_local
+type DNSResult struct {
+	Answers   []string `json:"answers"`
+	Rcode     string   `json:"rcode"`
+	LatencyMS int64    `json:"latency_ms"`
+}
+
+func probeDNSLocalHandler(params api.ToolHandlerParams) (*api.ToolCallResult, error) {
+	serverParam, ok := params.GetArguments()["server"].(string)
+	if !ok || serverParam == "" {
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("server parameter is required")), nil
+	}
+
+	nameParam, ok := params.GetArguments()["name"].(string)
+	if !ok || nameParam == "" {
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("name parameter is required")), nil
+	}
+
+	typeParam, ok := params.GetArguments()["type"].(string)
+	if !ok || typeParam == "" {
+		typeParam = "A"
+	}
+
+	// Ensure name falls back to a FQDN format
+	fqdn := dns.Fqdn(nameParam)
+
+	// Ensure server parameter has a port
+	if _, _, err := net.SplitHostPort(serverParam); err != nil {
+		// If port is missing, try adding default DNS port 53
+		appended := net.JoinHostPort(serverParam, "53")
+		if _, _, err2 := net.SplitHostPort(appended); err2 == nil {
+			serverParam = appended
+		} else {
+			return api.NewToolCallResultStructured(nil, fmt.Errorf("invalid server address format: %w", err)), nil
+		}
+	}
+
+	recordType, ok := dns.StringToType[strings.ToUpper(typeParam)]
+	if !ok {
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("invalid or unsupported DNS record type: %s", typeParam)), nil
+	}
+
+	msg := new(dns.Msg)
+	msg.SetQuestion(fqdn, recordType)
+	msg.RecursionDesired = true
+
+	resp, rtt, err := activeDNSClient.Exchange(msg, serverParam)
+
+	if err != nil {
+		// Log network level errors directly to the tool output so agent can interpret it
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("DNS query failed: %w", err)), nil
+	}
+
+	result := DNSResult{
+		Answers:   make([]string, 0, len(resp.Answer)),
+		Rcode:     dns.RcodeToString[resp.Rcode],
+		LatencyMS: rtt.Milliseconds(),
+	}
+
+	for _, answer := range resp.Answer {
+		// Replace tabs with spaces for prettier JSON presentation if needed
+		ans := strings.ReplaceAll(answer.String(), "\t", " ")
+		result.Answers = append(result.Answers, ans)
+	}
+
+	return api.NewToolCallResultStructured(result, nil), nil
+}

pkg/toolsets/netedge/probe_dns_local_test.go

@@ -0,0 +1,197 @@
+package netedge
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/containers/kubernetes-mcp-server/pkg/api"
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type mockDNSClient struct {
+	msg        *dns.Msg
+	rtt        time.Duration
+	err        error
+	lastServer string
+}
+
+func (m *mockDNSClient) Exchange(msg *dns.Msg, server string) (*dns.Msg, time.Duration, error) {
+	m.lastServer = server
+	return m.msg, m.rtt, m.err
+}
+
+func TestProbeDNSLocalHandler(t *testing.T) {
+	// Setup static success response
+	successMsg := new(dns.Msg)
+	successMsg.Rcode = dns.RcodeSuccess
+
+	aRecord, _ := dns.NewRR("example.com. 3600 IN A 93.184.216.34")
+	successMsg.Answer = append(successMsg.Answer, aRecord)
+
+	tests := []struct {
+		name          string
+		args          map[string]interface{}
+		mockClient    *mockDNSClient
+		expectedError string
+		validate      func(t *testing.T, result *api.ToolCallResult)
+	}{
+		{
+			name: "success query A record",
+			args: map[string]interface{}{
+				"server": "8.8.8.8",
+				"name":   "example.com",
+				"type":   "A",
+			},
+			mockClient: &mockDNSClient{
+				msg: successMsg,
+				rtt: 10 * time.Millisecond,
+				err: nil,
+			},
+			validate: func(t *testing.T, result *api.ToolCallResult) {
+				var res DNSResult
+				err := json.Unmarshal([]byte(result.Content), &res)
+				require.NoError(t, err)
+				assert.Equal(t, "NOERROR", res.Rcode)
+				assert.Equal(t, int64(10), res.LatencyMS)
+				assert.Len(t, res.Answers, 1)
+				assert.Contains(t, res.Answers[0], "93.184.216.34")
+
+				// Also check structured content
+				structured, ok := result.StructuredContent.(DNSResult)
+				require.True(t, ok)
+				assert.Equal(t, "NOERROR", structured.Rcode)
+			},
+		},
+		{
+			name: "missing name parameter",
+			args: map[string]interface{}{
+				"server": "8.8.8.8",
+			},
+			expectedError: "name parameter is required",
+		},
+		{
+			name: "missing server parameter",
+			args: map[string]interface{}{
+				"name": "example.com",
+			},
+			expectedError: "server parameter is required",
+		},
+		{
+			name: "invalid record type",
+			args: map[string]interface{}{
+				"server": "8.8.8.8",
+				"name":   "example.com",
+				"type":   "INVALID",
+			},
+			expectedError: "invalid or unsupported DNS record type: INVALID",
+		},
+		{
+			name: "network failure from library",
+			args: map[string]interface{}{
+				"server": "8.8.8.8",
+				"name":   "example.com",
+				"type":   "A",
+			},
+			mockClient: &mockDNSClient{
+				msg: nil,
+				rtt: 0,
+				err: &net.OpError{Op: "dial", Net: "udp", Err: net.UnknownNetworkError("timeout")},
+			},
+			expectedError: "DNS query failed",
+		},
+		{
+			name: "default type is A if omitted",
+			args: map[string]interface{}{
+				"server": "8.8.8.8",
+				"name":   "example.com",
+			},
+			mockClient: &mockDNSClient{
+				msg: successMsg,
+				rtt: 5 * time.Millisecond,
+				err: nil,
+			},
+			validate: func(t *testing.T, result *api.ToolCallResult) {
+				var res DNSResult
+				err := json.Unmarshal([]byte(result.Content), &res)
+				require.NoError(t, err)
+				assert.Equal(t, "NOERROR", res.Rcode)
+
+				// Also check structured content
+				structured, ok := result.StructuredContent.(DNSResult)
+				require.True(t, ok)
+				assert.Equal(t, "NOERROR", structured.Rcode)
+			},
+		},
+		{
+			name: "ipv4 address appends default port",
+			args: map[string]interface{}{
+				"server": "1.1.1.1",
+				"name":   "example.com",
+			},
+			mockClient: &mockDNSClient{
+				msg: successMsg,
+				rtt: 5 * time.Millisecond,
+			},
+			validate: func(t *testing.T, result *api.ToolCallResult) {
+				assert.Equal(t, "1.1.1.1:53", activeDNSClient.(*mockDNSClient).lastServer)
+			},
+		},
+		{
+			name: "ipv6 address appends default port",
+			args: map[string]interface{}{
+				"server": "2001:4860:4860::8888",
+				"name":   "example.com",
+			},
+			mockClient: &mockDNSClient{
+				msg: successMsg,
+				rtt: 5 * time.Millisecond,
+			},
+			validate: func(t *testing.T, result *api.ToolCallResult) {
+				assert.Equal(t, "[2001:4860:4860::8888]:53", activeDNSClient.(*mockDNSClient).lastServer)
+			},
+		},
+	}
+
+	// Stash original
+	origClient := activeDNSClient
+	defer func() {
+		activeDNSClient = origClient
+	}()
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.mockClient != nil {
+				activeDNSClient = tt.mockClient
+			}
+
+			// Mock toolcall params
+			toolReq := &mockToolCallRequest{args: tt.args}
+			params := api.ToolHandlerParams{
+				Context:         context.Background(),
+				ToolCallRequest: toolReq,
+			}
+
+			result, err := probeDNSLocalHandler(params)
+
+			// The handler wraps the validation error into NewToolCallResult("", err) or returns err directly?
+			// Actually our handler returns api.NewToolCallResult("", err), nil for errors so err is always nil.
+			require.NoError(t, err)
+
+			if tt.expectedError != "" {
+				require.NotNil(t, result.Error)
+				assert.Contains(t, result.Error.Error(), tt.expectedError)
+			} else {
+				require.NotNil(t, result)
+				require.NoError(t, result.Error)
+				if tt.validate != nil {
+					tt.validate(t, result)
+				}
+			}
+		})
+	}
+}

pkg/toolsets/netedge/toolset.go

@@ -26,6 +26,7 @@ func (t *Toolset) GetTools(_ api.Openshift) []api.ServerTool {
 		InitQueryPrometheus(),
 		initCoreDNS(),
 		initEndpoints(),
+		initProbeDNSLocal(),
 	)
 }