Extend JMESPath support to Message methods (#6735)

* Extend JMESPath support to Message methods

* Update cli/azd/pkg/input/console.go

Co-authored-by: Copilot <[email protected]>

* Update cli/azd/pkg/input/console.go

Co-authored-by: Copilot <[email protected]>

* Update TestAskerConsole_Message_JsonQueryFilter

---------

Co-authored-by: Copilot <[email protected]>

Author: scottaddie

cli/azd/pkg/input/console.go

@@ -217,7 +217,15 @@ func (c *AskerConsole) Message(ctx context.Context, message string) {
 	if c.formatter != nil && c.formatter.Kind() == output.JsonFormat {
 		// we call json.Marshal directly, because the formatter marshalls using indentation, and we would prefer
 		// these objects be written on a single line.
-		jsonMessage, err := json.Marshal(output.EventForMessage(message))
+		var obj interface{} = output.EventForMessage(message)
+		if q, ok := c.formatter.(output.Queryable); ok {
+			if filtered, err := q.QueryFilter(obj); err == nil {
+				obj = filtered
+			} else {
+				log.Printf("failed to apply query filter in Message: %v", err)
+			}
+		}
+		jsonMessage, err := json.Marshal(obj)
 		if err != nil {
 			panic(fmt.Sprintf("Message: unexpected error during marshaling for a valid object: %v", err))
 		}
@@ -270,8 +278,24 @@ func (c *AskerConsole) MessageUxItem(ctx context.Context, item ux.UxItem) {
 	if c.formatter != nil && c.formatter.Kind() == output.JsonFormat {
 		// no need to check the spinner for json format, as the spinner won't start when using json format
 		// instead, there would be a message about starting spinner
-		json, _ := json.Marshal(item)
-		fmt.Fprintln(c.writer, string(json))
+		var obj interface{} = item
+		if q, ok := c.formatter.(output.Queryable); ok {
+			// UxItem.MarshalJSON() produces an EventEnvelope. To apply JMESPath we
+			// need the unmarshaled structure, so round-trip through JSON first.
+			if raw, err := json.Marshal(item); err == nil {
+				var envelope interface{}
+				if err := json.Unmarshal(raw, &envelope); err == nil {
+					if filtered, err := q.QueryFilter(envelope); err == nil {
+						obj = filtered
+					}
+				}
+			}
+		}
+		jsonBytes, err := json.Marshal(obj)
+		if err != nil {
+			panic(err)
+		}
+		fmt.Fprintln(c.writer, string(jsonBytes))
 		return
 	}
 

cli/azd/pkg/input/console_test.go

@@ -14,6 +14,7 @@ import (
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/azure/azure-dev/cli/azd/pkg/contracts"
 	"github.com/azure/azure-dev/cli/azd/pkg/output"
 	"github.com/stretchr/testify/require"
 )
@@ -266,3 +267,105 @@ func newTestExternalPromptServer(handler func(promptOptions) json.RawMessage) *h
 		_, _ = w.Write(respBody)
 	}))
 }
+
+func TestAskerConsole_Message_JsonQueryFilter(t *testing.T) {
+	tests := []struct {
+		name   string
+		query  string
+		assert func(t *testing.T, got string)
+	}{
+		{
+			name:  "NoQuery",
+			query: "",
+			assert: func(t *testing.T, got string) {
+				// Unmarshal into the full envelope and verify structure
+				var env contracts.EventEnvelope
+				err := json.Unmarshal([]byte(strings.TrimSpace(got)), &env)
+				require.NoError(t, err, "output should be valid JSON envelope")
+				require.Equal(t, contracts.ConsoleMessageEventDataType, env.Type)
+
+				data, ok := env.Data.(map[string]any)
+				require.True(t, ok, "Data should be a map, got %T", env.Data)
+				// EventForMessage appends a trailing newline to the message
+				require.Equal(t, "hello world\n", data["message"])
+			},
+		},
+		{
+			name:  "QueryDataMessage",
+			query: "data.message",
+			assert: func(t *testing.T, got string) {
+				// Query should extract a bare JSON string, not an object
+				var s string
+				err := json.Unmarshal([]byte(strings.TrimSpace(got)), &s)
+				require.NoError(t, err, "output should unmarshal as a JSON string")
+				// EventForMessage appends a trailing newline to the message
+				require.Equal(t, "hello world\n", s)
+			},
+		},
+		{
+			name:  "QueryType",
+			query: "type",
+			assert: func(t *testing.T, got string) {
+				// Query should extract a bare JSON string, not an object
+				var s string
+				err := json.Unmarshal([]byte(strings.TrimSpace(got)), &s)
+				require.NoError(t, err, "output should unmarshal as a JSON string")
+				require.Equal(t, "consoleMessage", s)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			buf := &strings.Builder{}
+			formatter := &output.JsonFormatter{Query: tc.query}
+
+			c := NewConsole(
+				true,
+				false,
+				Writers{Output: writerAdapter{buf}},
+				ConsoleHandles{
+					Stderr: os.Stderr,
+					Stdin:  os.Stdin,
+					Stdout: writerAdapter{buf},
+				},
+				formatter,
+				nil,
+			)
+
+			c.Message(context.Background(), "hello world")
+
+			tc.assert(t, buf.String())
+		})
+	}
+}
+
+func TestAskerConsole_Message_InvalidQuery_FallsBack(t *testing.T) {
+	buf := &strings.Builder{}
+	formatter := &output.JsonFormatter{Query: "[invalid"}
+
+	c := NewConsole(
+		true,
+		false,
+		Writers{Output: writerAdapter{buf}},
+		ConsoleHandles{
+			Stderr: os.Stderr,
+			Stdin:  os.Stdin,
+			Stdout: writerAdapter{buf},
+		},
+		formatter,
+		nil,
+	)
+
+	// Should not panic; falls back to unfiltered output
+	c.Message(context.Background(), "hello world")
+
+	got := buf.String()
+	require.Contains(t, got, `"consoleMessage"`,
+		"invalid query should fall back to full envelope")
+}
+
+// writerAdapter wraps *strings.Builder to satisfy io.Writer for test purposes.
+type writerAdapter struct {
+	*strings.Builder
+}

cli/azd/pkg/output/formatter.go

@@ -22,6 +22,13 @@ type Formatter interface {
 	Format(obj interface{}, writer io.Writer, opts interface{}) error
 }
 
+// Queryable is an optional interface that a Formatter may implement to support
+// JMESPath filtering on arbitrary objects. Console messages (e.g. error output)
+// use this to apply --query without going through the full Format() pipeline.
+type Queryable interface {
+	QueryFilter(obj interface{}) (interface{}, error)
+}
+
 func NewFormatter(format string) (Formatter, error) {
 	switch format {
 	case string(JsonFormat):

cli/azd/pkg/output/json.go

@@ -53,6 +53,16 @@ func (f *JsonFormatter) Format(obj interface{}, writer io.Writer, _ interface{})
 }
 
 var _ Formatter = (*JsonFormatter)(nil)
+var _ Queryable = (*JsonFormatter)(nil)
+
+// QueryFilter applies the JMESPath query (if any) to the given object.
+// When no query is configured, the object is returned unchanged.
+func (f *JsonFormatter) QueryFilter(obj interface{}) (interface{}, error) {
+	if f.Query == "" {
+		return obj, nil
+	}
+	return ApplyQuery(obj, f.Query)
+}
 
 // jsonObjectForMessage creates a json object representing a message. Any ANSI control sequences from the message are
 // removed. A trailing newline is added to the message.

cli/azd/pkg/output/query_test.go

@@ -200,3 +200,39 @@ func TestApplyQuery_DocumentationHintInError(t *testing.T) {
 		})
 	}
 }
+
+func TestJsonFormatter_QueryFilter(t *testing.T) {
+	t.Run("NoQuery", func(t *testing.T) {
+		f := &JsonFormatter{}
+		data := map[string]interface{}{"name": "test"}
+		result, err := f.QueryFilter(data)
+		require.NoError(t, err)
+		require.Equal(t, data, result)
+	})
+
+	t.Run("WithQuery", func(t *testing.T) {
+		f := &JsonFormatter{Query: "data.message"}
+		data := map[string]interface{}{
+			"type": "consoleMessage",
+			"data": map[string]interface{}{
+				"message": "hello world",
+			},
+		}
+		result, err := f.QueryFilter(data)
+		require.NoError(t, err)
+		require.Equal(t, "hello world", result)
+	})
+
+	t.Run("InvalidQuery", func(t *testing.T) {
+		f := &JsonFormatter{Query: "[invalid"}
+		data := map[string]interface{}{"name": "test"}
+		_, err := f.QueryFilter(data)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "jmespath.org")
+	})
+}
+
+func TestQueryable_Interface(t *testing.T) {
+	// Verify JsonFormatter implements Queryable
+	var _ Queryable = (*JsonFormatter)(nil)
+}