Address PR #176 review comments

- Set ClusterAware to false for probe_dns_local

- Improve robust IPv6-aware address port handling

- Use api.NewToolCallResultStructured for structured tool output

Author: bentito

pkg/toolsets/netedge/probe_dns_local.go

@@ -45,7 +45,8 @@ func initProbeDNSLocal() []api.ServerTool {
 					OpenWorldHint:   ptr.To(true),
 				},
 			},
-			Handler: probeDNSLocalHandler,
+			ClusterAware: ptr.To(false),
+			Handler:      probeDNSLocalHandler,
 		},
 	}
 }
@@ -77,12 +78,12 @@ type DNSResult struct {
 func probeDNSLocalHandler(params api.ToolHandlerParams) (*api.ToolCallResult, error) {
 	serverParam, ok := params.GetArguments()["server"].(string)
 	if !ok || serverParam == "" {
-		return api.NewToolCallResult("", fmt.Errorf("server parameter is required")), nil
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("server parameter is required")), nil
 	}
 
 	nameParam, ok := params.GetArguments()["name"].(string)
 	if !ok || nameParam == "" {
-		return api.NewToolCallResult("", fmt.Errorf("name parameter is required")), nil
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("name parameter is required")), nil
 	}
 
 	typeParam, ok := params.GetArguments()["type"].(string)
@@ -95,16 +96,18 @@ func probeDNSLocalHandler(params api.ToolHandlerParams) (*api.ToolCallResult, er
 
 	// Ensure server parameter has a port
 	if _, _, err := net.SplitHostPort(serverParam); err != nil {
-		if strings.Contains(err.Error(), "missing port in address") {
-			serverParam = net.JoinHostPort(serverParam, "53")
+		// 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.NewToolCallResult("", fmt.Errorf("invalid server address format: %w", err)), nil
+			return api.NewToolCallResultStructured(nil, fmt.Errorf("invalid server address format: %w", err)), nil
 		}
 	}
 
 	recordType, ok := dns.StringToType[strings.ToUpper(typeParam)]
 	if !ok {
-		return api.NewToolCallResult("", fmt.Errorf("invalid or unsupported DNS record type: %s", typeParam)), nil
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("invalid or unsupported DNS record type: %s", typeParam)), nil
 	}
 
 	msg := new(dns.Msg)
@@ -115,7 +118,7 @@ func probeDNSLocalHandler(params api.ToolHandlerParams) (*api.ToolCallResult, er
 
 	if err != nil {
 		// Log network level errors directly to the tool output so agent can interpret it
-		return api.NewToolCallResult("", fmt.Errorf("DNS query failed: %w", err)), nil
+		return api.NewToolCallResultStructured(nil, fmt.Errorf("DNS query failed: %w", err)), nil
 	}
 
 	result := DNSResult{
@@ -130,10 +133,5 @@ func probeDNSLocalHandler(params api.ToolHandlerParams) (*api.ToolCallResult, er
 		result.Answers = append(result.Answers, ans)
 	}
 
-	jsonData, err := json.MarshalIndent(result, "", "  ")
-	if err != nil {
-		return api.NewToolCallResult("", fmt.Errorf("failed to marshal DNS result: %w", err)), nil
-	}
-
-	return api.NewToolCallResult(string(jsonData), nil), nil
+	return api.NewToolCallResultStructured(result, nil), nil
 }

pkg/toolsets/netedge/probe_dns_local_test.go

@@ -14,12 +14,14 @@ import (
 )
 
 type mockDNSClient struct {
-	msg *dns.Msg
-	rtt time.Duration
-	err error
+	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
 }
 
@@ -36,7 +38,7 @@ func TestProbeDNSLocalHandler(t *testing.T) {
 		args          map[string]interface{}
 		mockClient    *mockDNSClient
 		expectedError string
-		validate      func(t *testing.T, result string)
+		validate      func(t *testing.T, result *api.ToolCallResult)
 	}{
 		{
 			name: "success query A record",
@@ -50,14 +52,19 @@ func TestProbeDNSLocalHandler(t *testing.T) {
 				rtt: 10 * time.Millisecond,
 				err: nil,
 			},
-			validate: func(t *testing.T, content string) {
+			validate: func(t *testing.T, result *api.ToolCallResult) {
 				var res DNSResult
-				err := json.Unmarshal([]byte(content), &res)
+				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)
 			},
 		},
 		{
@@ -108,11 +115,44 @@ func TestProbeDNSLocalHandler(t *testing.T) {
 				rtt: 5 * time.Millisecond,
 				err: nil,
 			},
-			validate: func(t *testing.T, content string) {
+			validate: func(t *testing.T, result *api.ToolCallResult) {
 				var res DNSResult
-				err := json.Unmarshal([]byte(content), &res)
+				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)
 			},
 		},
 	}
@@ -149,7 +189,7 @@ func TestProbeDNSLocalHandler(t *testing.T) {
 				require.NotNil(t, result)
 				require.NoError(t, result.Error)
 				if tt.validate != nil {
-					tt.validate(t, result.Content)
+					tt.validate(t, result)
 				}
 			}
 		})