From d4e3b75f39f0fa421eba72009063d51315704cfe Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Tue, 3 Mar 2026 17:47:20 +0100
Subject: [PATCH 1/6] Add output format routing, interactive table browser, and
 multi-chunk results

Replace the single JSON-to-stderr output with format-aware routing:
- Interactive terminal: table browser (viewport-based, with search)
- Non-interactive / --output json: JSON array of objects
- Small results (<=30 rows): static formatted table

Add libs/tableview, a reusable interactive table browser component
using bubbles/viewport. Features: horizontal + vertical scrolling,
cursor row highlighting, / search with match navigation (n/N),
g/G for top/bottom.

Add multi-chunk result fetching via GetStatementResultChunkN for
queries returning more data than fits in a single response chunk.
---
 experimental/aitools/cmd/query.go       |  96 ++++---
 experimental/aitools/cmd/query_test.go  |  54 ++--
 experimental/aitools/cmd/render.go      |  96 +++++++
 experimental/aitools/cmd/render_test.go |  95 +++++++
 libs/tableview/tableview.go             | 342 ++++++++++++++++++++++++
 libs/tableview/tableview_test.go        |  72 +++++
 6 files changed, 693 insertions(+), 62 deletions(-)
 create mode 100644 experimental/aitools/cmd/render.go
 create mode 100644 experimental/aitools/cmd/render_test.go
 create mode 100644 libs/tableview/tableview.go
 create mode 100644 libs/tableview/tableview_test.go

diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go
index ba391a4b24..b77b01da4f 100644
--- a/experimental/aitools/cmd/query.go
+++ b/experimental/aitools/cmd/query.go
@@ -2,7 +2,6 @@ package mcp
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -17,6 +16,7 @@ import (
 	"github.com/databricks/cli/experimental/aitools/lib/session"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -34,6 +34,10 @@ const (
 
 	// cancelTimeout is how long to wait for server-side cancellation.
 	cancelTimeout = 10 * time.Second
+
+	// staticTableThreshold is the maximum number of rows rendered as a static table.
+	// Beyond this, an interactive scrollable table is used.
+	staticTableThreshold = 30
 )
 
 func newQueryCmd() *cobra.Command {
@@ -79,13 +83,32 @@ Output includes the query results as JSON and row count.`,
 				return err
 			}
 
-			output, err := formatQueryResult(resp)
+			columns := extractColumns(resp.Manifest)
+			rows, err := fetchAllRows(ctx, w.StatementExecution, resp)
 			if err != nil {
 				return err
 			}
 
-			cmdio.LogString(ctx, output)
-			return nil
+			if len(columns) == 0 && len(rows) == 0 {
+				fmt.Fprintln(cmd.OutOrStdout(), "Query executed successfully (no results)")
+				return nil
+			}
+
+			// Use table output only when stdout is a TTY with color support.
+			// cmdio.IsInteractive checks stderr (for spinners), but output
+			// format depends on whether stdout itself is a terminal.
+			stdoutInteractive := cmdio.SupportsColor(ctx, cmd.OutOrStdout())
+
+			switch {
+			case root.OutputType(cmd) == flags.OutputJSON:
+				return renderJSON(cmd.OutOrStdout(), columns, rows)
+			case !stdoutInteractive:
+				return renderJSON(cmd.OutOrStdout(), columns, rows)
+			case len(rows) <= staticTableThreshold:
+				return renderStaticTable(cmd.OutOrStdout(), columns, rows)
+			default:
+				return renderInteractiveTable(cmd.OutOrStdout(), columns, rows)
+			}
 		},
 	}
 
@@ -273,6 +296,31 @@ func executeAndPoll(ctx context.Context, api sql.StatementExecutionInterface, wa
 	}
 }
 
+// fetchAllRows collects all result rows, fetching additional chunks if needed.
+func fetchAllRows(ctx context.Context, api sql.StatementExecutionInterface, resp *sql.StatementResponse) ([][]string, error) {
+	if resp.Result == nil {
+		return nil, nil
+	}
+
+	rows := append([][]string{}, resp.Result.DataArray...)
+
+	totalChunks := 0
+	if resp.Manifest != nil {
+		totalChunks = resp.Manifest.TotalChunkCount
+	}
+
+	for chunk := 1; chunk < totalChunks; chunk++ {
+		log.Debugf(ctx, "Fetching result chunk %d/%d for statement %s", chunk+1, totalChunks, resp.StatementId)
+		chunkResp, err := api.GetStatementResultChunkNByStatementIdAndChunkIndex(ctx, resp.StatementId, chunk)
+		if err != nil {
+			return nil, fmt.Errorf("fetch result chunk %d: %w", chunk, err)
+		}
+		rows = append(rows, chunkResp.DataArray...)
+	}
+
+	return rows, nil
+}
+
 // isTerminalState returns true if the statement has reached a final state.
 func isTerminalState(status *sql.StatementStatus) bool {
 	if status == nil {
@@ -333,43 +381,3 @@ func cleanSQL(s string) string {
 
 	return strings.Join(lines, "\n")
 }
-
-func formatQueryResult(resp *sql.StatementResponse) (string, error) {
-	var sb strings.Builder
-
-	if resp.Manifest == nil || resp.Result == nil {
-		sb.WriteString("Query executed successfully (no results)\n")
-		return sb.String(), nil
-	}
-
-	var columns []string
-	if resp.Manifest.Schema != nil {
-		for _, col := range resp.Manifest.Schema.Columns {
-			columns = append(columns, col.Name)
-		}
-	}
-
-	var rows []map[string]any
-	if resp.Result.DataArray != nil {
-		for _, row := range resp.Result.DataArray {
-			rowMap := make(map[string]any)
-			for i, val := range row {
-				if i < len(columns) {
-					rowMap[columns[i]] = val
-				}
-			}
-			rows = append(rows, rowMap)
-		}
-	}
-
-	output, err := json.MarshalIndent(rows, "", "  ")
-	if err != nil {
-		return "", fmt.Errorf("marshal results: %w", err)
-	}
-
-	sb.Write(output)
-	sb.WriteString("\n\n")
-	sb.WriteString(fmt.Sprintf("Row count: %d\n", len(rows)))
-
-	return sb.String(), nil
-}
diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go
index 8b8686aee7..3928c7634a 100644
--- a/experimental/aitools/cmd/query_test.go
+++ b/experimental/aitools/cmd/query_test.go
@@ -159,32 +159,50 @@ func TestResolveWarehouseIDWithFlag(t *testing.T) {
 	assert.Equal(t, "explicit-id", id)
 }
 
-func TestFormatQueryResultNoResults(t *testing.T) {
+func TestFetchAllRowsSingleChunk(t *testing.T) {
+	ctx := cmdio.MockDiscard(context.Background())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
 	resp := &sql.StatementResponse{
-		Status: &sql.StatementStatus{State: sql.StatementStateSucceeded},
+		StatementId: "stmt-1",
+		Manifest:    &sql.ResultManifest{TotalChunkCount: 1},
+		Result:      &sql.ResultData{DataArray: [][]string{{"1", "alice"}, {"2", "bob"}}},
 	}
-	output, err := formatQueryResult(resp)
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
 	require.NoError(t, err)
-	assert.Contains(t, output, "no results")
+	assert.Equal(t, [][]string{{"1", "alice"}, {"2", "bob"}}, rows)
 }
 
-func TestFormatQueryResultWithData(t *testing.T) {
+func TestFetchAllRowsMultiChunk(t *testing.T) {
+	ctx := cmdio.MockDiscard(context.Background())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
 	resp := &sql.StatementResponse{
-		Status: &sql.StatementStatus{State: sql.StatementStateSucceeded},
-		Manifest: &sql.ResultManifest{
-			Schema: &sql.ResultSchema{
-				Columns: []sql.ColumnInfo{{Name: "id"}, {Name: "name"}},
-			},
-		},
-		Result: &sql.ResultData{
-			DataArray: [][]string{{"1", "alice"}, {"2", "bob"}},
-		},
+		StatementId: "stmt-1",
+		Manifest:    &sql.ResultManifest{TotalChunkCount: 3},
+		Result:      &sql.ResultData{DataArray: [][]string{{"1", "a"}}},
 	}
-	output, err := formatQueryResult(resp)
+
+	mockAPI.EXPECT().GetStatementResultChunkNByStatementIdAndChunkIndex(mock.Anything, "stmt-1", 1).
+		Return(&sql.ResultData{DataArray: [][]string{{"2", "b"}}}, nil).Once()
+	mockAPI.EXPECT().GetStatementResultChunkNByStatementIdAndChunkIndex(mock.Anything, "stmt-1", 2).
+		Return(&sql.ResultData{DataArray: [][]string{{"3", "c"}}}, nil).Once()
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
+	require.NoError(t, err)
+	assert.Equal(t, [][]string{{"1", "a"}, {"2", "b"}, {"3", "c"}}, rows)
+}
+
+func TestFetchAllRowsNilResult(t *testing.T) {
+	ctx := cmdio.MockDiscard(context.Background())
+	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
+
+	resp := &sql.StatementResponse{StatementId: "stmt-1"}
+
+	rows, err := fetchAllRows(ctx, mockAPI, resp)
 	require.NoError(t, err)
-	assert.Contains(t, output, "alice")
-	assert.Contains(t, output, "bob")
-	assert.Contains(t, output, "Row count: 2")
+	assert.Nil(t, rows)
 }
 
 func TestIsTerminalState(t *testing.T) {
diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go
new file mode 100644
index 0000000000..47fac571db
--- /dev/null
+++ b/experimental/aitools/cmd/render.go
@@ -0,0 +1,96 @@
+package mcp
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"strings"
+	"text/tabwriter"
+
+	"github.com/databricks/cli/libs/tableview"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+)
+
+const (
+	// maxColumnWidth is the maximum display width for any single column in static table output.
+	maxColumnWidth = 40
+)
+
+// extractColumns returns column names from the query result manifest.
+func extractColumns(manifest *sql.ResultManifest) []string {
+	if manifest == nil || manifest.Schema == nil {
+		return nil
+	}
+	columns := make([]string, len(manifest.Schema.Columns))
+	for i, col := range manifest.Schema.Columns {
+		columns[i] = col.Name
+	}
+	return columns
+}
+
+// renderJSON writes query results as a JSON array of objects with a row count footer.
+func renderJSON(w io.Writer, columns []string, rows [][]string) error {
+	objects := make([]map[string]any, len(rows))
+	for i, row := range rows {
+		obj := make(map[string]any, len(columns))
+		for j, val := range row {
+			if j < len(columns) {
+				obj[columns[j]] = val
+			}
+		}
+		objects[i] = obj
+	}
+
+	output, err := json.MarshalIndent(objects, "", "  ")
+	if err != nil {
+		return fmt.Errorf("marshal results: %w", err)
+	}
+
+	fmt.Fprintf(w, "%s\n\nRow count: %d\n", output, len(rows))
+	return nil
+}
+
+// renderStaticTable writes query results as a formatted text table.
+func renderStaticTable(w io.Writer, columns []string, rows [][]string) error {
+	tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0)
+
+	// Header row.
+	fmt.Fprintln(tw, strings.Join(columns, "\t"))
+
+	// Separator.
+	seps := make([]string, len(columns))
+	for i, col := range columns {
+		width := len(col)
+		for _, row := range rows {
+			if i < len(row) {
+				width = max(width, len(row[i]))
+			}
+		}
+		width = min(width, maxColumnWidth)
+		seps[i] = strings.Repeat("-", width)
+	}
+	fmt.Fprintln(tw, strings.Join(seps, "\t"))
+
+	// Data rows.
+	for _, row := range rows {
+		vals := make([]string, len(columns))
+		for i := range columns {
+			if i < len(row) {
+				vals[i] = row[i]
+			}
+		}
+		fmt.Fprintln(tw, strings.Join(vals, "\t"))
+	}
+
+	if err := tw.Flush(); err != nil {
+		return err
+	}
+
+	fmt.Fprintf(w, "\n%d rows\n", len(rows))
+	return nil
+}
+
+// renderInteractiveTable displays query results in the interactive table browser.
+func renderInteractiveTable(w io.Writer, columns []string, rows [][]string) error {
+	return tableview.Run(w, columns, rows)
+}
diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go
new file mode 100644
index 0000000000..05c251bf92
--- /dev/null
+++ b/experimental/aitools/cmd/render_test.go
@@ -0,0 +1,95 @@
+package mcp
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExtractColumns(t *testing.T) {
+	tests := []struct {
+		name     string
+		manifest *sql.ResultManifest
+		want     []string
+	}{
+		{
+			"with columns",
+			&sql.ResultManifest{Schema: &sql.ResultSchema{
+				Columns: []sql.ColumnInfo{{Name: "id"}, {Name: "name"}},
+			}},
+			[]string{"id", "name"},
+		},
+		{"nil manifest", nil, nil},
+		{"nil schema", &sql.ResultManifest{}, nil},
+		{
+			"empty columns",
+			&sql.ResultManifest{Schema: &sql.ResultSchema{}},
+			[]string{},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := extractColumns(tc.manifest)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestRenderJSON(t *testing.T) {
+	var buf bytes.Buffer
+	columns := []string{"id", "name"}
+	rows := [][]string{{"1", "alice"}, {"2", "bob"}}
+
+	err := renderJSON(&buf, columns, rows)
+	require.NoError(t, err)
+
+	output := buf.String()
+	assert.Contains(t, output, `"alice"`)
+	assert.Contains(t, output, `"bob"`)
+	assert.Contains(t, output, "Row count: 2")
+}
+
+func TestRenderJSONNoRows(t *testing.T) {
+	var buf bytes.Buffer
+	columns := []string{"id"}
+	var rows [][]string
+
+	err := renderJSON(&buf, columns, rows)
+	require.NoError(t, err)
+
+	output := buf.String()
+	assert.Contains(t, output, "Row count: 0")
+}
+
+func TestRenderStaticTable(t *testing.T) {
+	var buf bytes.Buffer
+	columns := []string{"id", "name"}
+	rows := [][]string{{"1", "alice"}, {"2", "bob"}}
+
+	err := renderStaticTable(&buf, columns, rows)
+	require.NoError(t, err)
+
+	output := buf.String()
+	assert.Contains(t, output, "id")
+	assert.Contains(t, output, "name")
+	assert.Contains(t, output, "alice")
+	assert.Contains(t, output, "bob")
+	assert.Contains(t, output, "---")
+	assert.Contains(t, output, "2 rows")
+}
+
+func TestRenderStaticTableEmpty(t *testing.T) {
+	var buf bytes.Buffer
+	columns := []string{"id", "name"}
+	var rows [][]string
+
+	err := renderStaticTable(&buf, columns, rows)
+	require.NoError(t, err)
+
+	output := buf.String()
+	assert.Contains(t, output, "id")
+	assert.Contains(t, output, "0 rows")
+}
diff --git a/libs/tableview/tableview.go b/libs/tableview/tableview.go
new file mode 100644
index 0000000000..ff337fd1b6
--- /dev/null
+++ b/libs/tableview/tableview.go
@@ -0,0 +1,342 @@
+// Package tableview provides an interactive table browser with scrolling and search.
+package tableview
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"text/tabwriter"
+
+	"github.com/charmbracelet/bubbles/viewport"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+)
+
+const (
+	horizontalScrollStep = 4
+	footerHeight         = 1
+	searchFooterHeight   = 2
+)
+
+const (
+	// headerLines is the number of non-data lines at the top (header + separator).
+	headerLines = 2
+)
+
+var (
+	searchHighlightStyle = lipgloss.NewStyle().Background(lipgloss.Color("228")).Foreground(lipgloss.Color("0"))
+	cursorStyle          = lipgloss.NewStyle().Background(lipgloss.Color("57")).Foreground(lipgloss.Color("229"))
+	footerStyle          = lipgloss.NewStyle().Foreground(lipgloss.Color("241"))
+	searchStyle          = lipgloss.NewStyle().Foreground(lipgloss.Color("229"))
+)
+
+// Run displays tabular data in an interactive browser.
+// Writes to w (typically stdout). Blocks until user quits.
+func Run(w io.Writer, columns []string, rows [][]string) error {
+	lines := renderTableLines(columns, rows)
+
+	m := model{
+		lines:     lines,
+		totalRows: len(rows),
+		cursor:    headerLines, // Start on first data row.
+	}
+
+	p := tea.NewProgram(m, tea.WithOutput(w))
+	_, err := p.Run()
+	return err
+}
+
+// renderTableLines produces aligned table text as individual lines.
+func renderTableLines(columns []string, rows [][]string) []string {
+	var buf strings.Builder
+	tw := tabwriter.NewWriter(&buf, 0, 4, 2, ' ', 0)
+
+	// Header.
+	fmt.Fprintln(tw, strings.Join(columns, "\t"))
+
+	// Separator: compute widths from header + data for dash line.
+	widths := make([]int, len(columns))
+	for i, col := range columns {
+		widths[i] = len(col)
+	}
+	for _, row := range rows {
+		for i := range columns {
+			if i < len(row) {
+				widths[i] = max(widths[i], len(row[i]))
+			}
+		}
+	}
+	seps := make([]string, len(columns))
+	for i, w := range widths {
+		seps[i] = strings.Repeat("─", w)
+	}
+	fmt.Fprintln(tw, strings.Join(seps, "\t"))
+
+	// Data rows.
+	for _, row := range rows {
+		vals := make([]string, len(columns))
+		for i := range columns {
+			if i < len(row) {
+				vals[i] = row[i]
+			}
+		}
+		fmt.Fprintln(tw, strings.Join(vals, "\t"))
+	}
+
+	tw.Flush()
+
+	// Split into lines, drop trailing empty.
+	lines := strings.Split(buf.String(), "\n")
+	if len(lines) > 0 && lines[len(lines)-1] == "" {
+		lines = lines[:len(lines)-1]
+	}
+	return lines
+}
+
+// findMatches returns line indices containing the query (case-insensitive).
+func findMatches(lines []string, query string) []int {
+	if query == "" {
+		return nil
+	}
+	lower := strings.ToLower(query)
+	var matches []int
+	for i, line := range lines {
+		if strings.Contains(strings.ToLower(line), lower) {
+			matches = append(matches, i)
+		}
+	}
+	return matches
+}
+
+// highlightSearch applies search match highlighting to a single line.
+func highlightSearch(line, query string) string {
+	if query == "" {
+		return line
+	}
+	lower := strings.ToLower(query)
+	qLen := len(query)
+	lineLower := strings.ToLower(line)
+
+	var b strings.Builder
+	pos := 0
+	for {
+		idx := strings.Index(lineLower[pos:], lower)
+		if idx < 0 {
+			b.WriteString(line[pos:])
+			break
+		}
+		b.WriteString(line[pos : pos+idx])
+		b.WriteString(searchHighlightStyle.Render(line[pos+idx : pos+idx+qLen]))
+		pos += idx + qLen
+	}
+	return b.String()
+}
+
+// renderContent builds the viewport content with cursor and search highlighting.
+func (m model) renderContent() string {
+	result := make([]string, len(m.lines))
+	for i, line := range m.lines {
+		rendered := highlightSearch(line, m.searchQuery)
+		if i == m.cursor {
+			rendered = cursorStyle.Render(line)
+			if m.searchQuery != "" {
+				// Apply search highlight on top of cursor for the current line.
+				rendered = highlightSearch(cursorStyle.Render(line), m.searchQuery)
+			}
+		}
+		result[i] = rendered
+	}
+	return strings.Join(result, "\n")
+}
+
+type model struct {
+	viewport  viewport.Model
+	lines     []string
+	totalRows int
+	ready     bool
+	cursor    int // line index of the highlighted row
+
+	// Search state.
+	searching   bool
+	searchInput string
+	searchQuery string
+	matchLines  []int
+	matchIdx    int
+}
+
+func (m model) Init() tea.Cmd {
+	return nil
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		fh := footerHeight
+		if m.searching {
+			fh = searchFooterHeight
+		}
+		if !m.ready {
+			m.viewport = viewport.New(msg.Width, msg.Height-fh)
+			m.viewport.SetHorizontalStep(horizontalScrollStep)
+			m.viewport.SetContent(m.renderContent())
+			m.ready = true
+		} else {
+			m.viewport.Width = msg.Width
+			m.viewport.Height = msg.Height - fh
+		}
+		return m, nil
+
+	case tea.KeyMsg:
+		if m.searching {
+			return m.updateSearch(msg)
+		}
+		return m.updateNormal(msg)
+	}
+
+	var cmd tea.Cmd
+	m.viewport, cmd = m.viewport.Update(msg)
+	return m, cmd
+}
+
+func (m model) updateNormal(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+	switch msg.String() {
+	case "q", "esc", "ctrl+c":
+		return m, tea.Quit
+	case "/":
+		m.searching = true
+		m.searchInput = ""
+		m.viewport.Height--
+		return m, nil
+	case "n":
+		if len(m.matchLines) > 0 {
+			m.matchIdx = (m.matchIdx + 1) % len(m.matchLines)
+			m.cursor = m.matchLines[m.matchIdx]
+			m.viewport.SetContent(m.renderContent())
+			m.scrollToCursor()
+		}
+		return m, nil
+	case "N":
+		if len(m.matchLines) > 0 {
+			m.matchIdx = (m.matchIdx - 1 + len(m.matchLines)) % len(m.matchLines)
+			m.cursor = m.matchLines[m.matchIdx]
+			m.viewport.SetContent(m.renderContent())
+			m.scrollToCursor()
+		}
+		return m, nil
+	case "up", "k":
+		m.moveCursor(-1)
+		return m, nil
+	case "down", "j":
+		m.moveCursor(1)
+		return m, nil
+	case "pgup", "b":
+		m.moveCursor(-m.viewport.Height)
+		return m, nil
+	case "pgdown", "f", " ":
+		m.moveCursor(m.viewport.Height)
+		return m, nil
+	case "g":
+		m.cursor = headerLines
+		m.viewport.SetContent(m.renderContent())
+		m.viewport.GotoTop()
+		return m, nil
+	case "G":
+		m.cursor = len(m.lines) - 1
+		m.viewport.SetContent(m.renderContent())
+		m.viewport.GotoBottom()
+		return m, nil
+	}
+
+	// Let viewport handle horizontal scroll and other keys.
+	var cmd tea.Cmd
+	m.viewport, cmd = m.viewport.Update(msg)
+	return m, cmd
+}
+
+// moveCursor moves the cursor by delta lines, clamped to data rows.
+func (m *model) moveCursor(delta int) {
+	m.cursor += delta
+	m.cursor = max(m.cursor, headerLines)
+	m.cursor = min(m.cursor, len(m.lines)-1)
+	m.viewport.SetContent(m.renderContent())
+	m.scrollToCursor()
+}
+
+// scrollToCursor ensures the cursor line is visible in the viewport.
+func (m *model) scrollToCursor() {
+	top := m.viewport.YOffset
+	bottom := top + m.viewport.Height - 1
+	if m.cursor < top {
+		m.viewport.SetYOffset(m.cursor)
+	} else if m.cursor > bottom {
+		m.viewport.SetYOffset(m.cursor - m.viewport.Height + 1)
+	}
+}
+
+func (m model) updateSearch(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+	switch msg.String() {
+	case "enter":
+		m.searching = false
+		m.searchQuery = m.searchInput
+		m.matchLines = findMatches(m.lines, m.searchQuery)
+		m.matchIdx = 0
+		m.viewport.Height++
+		// Move cursor to first match and re-render.
+		if len(m.matchLines) > 0 {
+			m.cursor = m.matchLines[0]
+		}
+		m.viewport.SetContent(m.renderContent())
+		if len(m.matchLines) > 0 {
+			m.scrollToCursor()
+		}
+		return m, nil
+	case "esc", "ctrl+c":
+		m.searching = false
+		m.searchInput = ""
+		m.viewport.Height++
+		return m, nil
+	case "backspace":
+		if len(m.searchInput) > 0 {
+			m.searchInput = m.searchInput[:len(m.searchInput)-1]
+		}
+		return m, nil
+	default:
+		// Only accept printable characters.
+		if len(msg.String()) == 1 || msg.Type == tea.KeyRunes {
+			m.searchInput += msg.String()
+		}
+		return m, nil
+	}
+}
+
+func (m model) View() string {
+	if !m.ready {
+		return "Loading..."
+	}
+
+	footer := m.renderFooter()
+	return m.viewport.View() + "\n" + footer
+}
+
+func (m model) renderFooter() string {
+	if m.searching {
+		prompt := searchStyle.Render("/ " + m.searchInput + "█")
+		return footerStyle.Render(fmt.Sprintf("%d rows", m.totalRows)) + "\n" + prompt
+	}
+
+	parts := []string{fmt.Sprintf("%d rows", m.totalRows)}
+
+	if m.searchQuery != "" && len(m.matchLines) > 0 {
+		parts = append(parts, fmt.Sprintf("match %d/%d", m.matchIdx+1, len(m.matchLines)))
+		parts = append(parts, "n/N next/prev")
+	} else if m.searchQuery != "" {
+		parts = append(parts, "no matches")
+	}
+
+	parts = append(parts, "←→↑↓ scroll", "g/G top/bottom", "/ search", "q quit")
+
+	pct := int(m.viewport.ScrollPercent() * 100)
+	parts = append(parts, fmt.Sprintf("%d%%", pct))
+
+	return footerStyle.Render(strings.Join(parts, " | "))
+}
diff --git a/libs/tableview/tableview_test.go b/libs/tableview/tableview_test.go
new file mode 100644
index 0000000000..c761a9cf00
--- /dev/null
+++ b/libs/tableview/tableview_test.go
@@ -0,0 +1,72 @@
+package tableview
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRenderTableLines(t *testing.T) {
+	columns := []string{"id", "name"}
+	rows := [][]string{{"1", "alice"}, {"2", "bob"}}
+
+	lines := renderTableLines(columns, rows)
+	require.GreaterOrEqual(t, len(lines), 4)
+
+	assert.Contains(t, lines[0], "id")
+	assert.Contains(t, lines[0], "name")
+	assert.Contains(t, lines[1], "──")
+	assert.Contains(t, lines[2], "alice")
+	assert.Contains(t, lines[3], "bob")
+}
+
+func TestRenderTableLinesEmpty(t *testing.T) {
+	columns := []string{"id", "name"}
+	var rows [][]string
+
+	lines := renderTableLines(columns, rows)
+	require.GreaterOrEqual(t, len(lines), 2)
+	assert.Contains(t, lines[0], "id")
+	assert.Contains(t, lines[1], "──")
+}
+
+func TestFindMatches(t *testing.T) {
+	lines := []string{"header", "---", "alice", "bob", "alice again"}
+	matches := findMatches(lines, "alice")
+	assert.Equal(t, []int{2, 4}, matches)
+}
+
+func TestFindMatchesCaseInsensitive(t *testing.T) {
+	lines := []string{"Alice", "BOB", "alice"}
+	matches := findMatches(lines, "ALICE")
+	assert.Equal(t, []int{0, 2}, matches)
+}
+
+func TestFindMatchesNoResults(t *testing.T) {
+	lines := []string{"alice", "bob"}
+	matches := findMatches(lines, "charlie")
+	assert.Nil(t, matches)
+}
+
+func TestFindMatchesEmptyQuery(t *testing.T) {
+	lines := []string{"alice", "bob"}
+	matches := findMatches(lines, "")
+	assert.Nil(t, matches)
+}
+
+func TestHighlightSearchEmptyQuery(t *testing.T) {
+	result := highlightSearch("hello alice", "")
+	assert.Equal(t, "hello alice", result)
+}
+
+func TestHighlightSearchWithMatch(t *testing.T) {
+	result := highlightSearch("hello alice", "alice")
+	assert.Contains(t, result, "alice")
+	assert.Contains(t, result, "hello")
+}
+
+func TestHighlightSearchNoMatch(t *testing.T) {
+	result := highlightSearch("hello bob", "alice")
+	assert.Equal(t, "hello bob", result)
+}

From a2750ad1d021433005a096d91cad76ec6bd28300 Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Tue, 3 Mar 2026 17:56:47 +0100
Subject: [PATCH 2/6] Fix query output routing for non-interactive stdin

Avoid opening the interactive table browser when stdin cannot accept input by falling back to static tables, and add output routing tests with updated command help text.
---
 experimental/aitools/cmd/query.go      | 43 ++++++++++++++----
 experimental/aitools/cmd/query_test.go | 60 ++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 9 deletions(-)

diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go
index b77b01da4f..2b255b28e7 100644
--- a/experimental/aitools/cmd/query.go
+++ b/experimental/aitools/cmd/query.go
@@ -40,6 +40,32 @@ const (
 	staticTableThreshold = 30
 )
 
+type queryOutputMode int
+
+const (
+	queryOutputModeJSON queryOutputMode = iota
+	queryOutputModeStaticTable
+	queryOutputModeInteractiveTable
+)
+
+func selectQueryOutputMode(outputType flags.Output, stdoutInteractive bool, promptSupported bool, rowCount int) queryOutputMode {
+	if outputType == flags.OutputJSON {
+		return queryOutputModeJSON
+	}
+	if !stdoutInteractive {
+		return queryOutputModeJSON
+	}
+	// Interactive table browsing requires keyboard input from stdin.
+	// If prompts are not supported, prefer static table output instead.
+	if !promptSupported {
+		return queryOutputModeStaticTable
+	}
+	if rowCount <= staticTableThreshold {
+		return queryOutputModeStaticTable
+	}
+	return queryOutputModeInteractiveTable
+}
+
 func newQueryCmd() *cobra.Command {
 	var warehouseID string
 	var filePath string
@@ -56,7 +82,8 @@ exists, it is read as a SQL file automatically.
 The command auto-detects an available warehouse unless --warehouse is set
 or the DATABRICKS_WAREHOUSE_ID environment variable is configured.
 
-Output includes the query results as JSON and row count.`,
+Output is JSON in non-interactive contexts. In interactive terminals it renders
+tables, and large results open an interactive table browser.`,
 		Example: `  databricks experimental aitools tools query "SELECT * FROM samples.nyctaxi.trips LIMIT 5"
   databricks experimental aitools tools query --warehouse abc123 "SELECT 1"
   databricks experimental aitools tools query --file report.sql
@@ -94,17 +121,15 @@ Output includes the query results as JSON and row count.`,
 				return nil
 			}
 
-			// Use table output only when stdout is a TTY with color support.
-			// cmdio.IsInteractive checks stderr (for spinners), but output
-			// format depends on whether stdout itself is a terminal.
+			// Output format depends on stdout capabilities.
+			// Interactive table browsing also requires prompt-capable stdin.
 			stdoutInteractive := cmdio.SupportsColor(ctx, cmd.OutOrStdout())
+			promptSupported := cmdio.IsPromptSupported(ctx)
 
-			switch {
-			case root.OutputType(cmd) == flags.OutputJSON:
-				return renderJSON(cmd.OutOrStdout(), columns, rows)
-			case !stdoutInteractive:
+			switch selectQueryOutputMode(root.OutputType(cmd), stdoutInteractive, promptSupported, len(rows)) {
+			case queryOutputModeJSON:
 				return renderJSON(cmd.OutOrStdout(), columns, rows)
-			case len(rows) <= staticTableThreshold:
+			case queryOutputModeStaticTable:
 				return renderStaticTable(cmd.OutOrStdout(), columns, rows)
 			default:
 				return renderInteractiveTable(cmd.OutOrStdout(), columns, rows)
diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go
index 3928c7634a..741e23d3a7 100644
--- a/experimental/aitools/cmd/query_test.go
+++ b/experimental/aitools/cmd/query_test.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/flags"
 	mocksql "github.com/databricks/databricks-sdk-go/experimental/mocks/service/sql"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -159,6 +160,65 @@ func TestResolveWarehouseIDWithFlag(t *testing.T) {
 	assert.Equal(t, "explicit-id", id)
 }
 
+func TestSelectQueryOutputMode(t *testing.T) {
+	tests := []struct {
+		name              string
+		outputType        flags.Output
+		stdoutInteractive bool
+		promptSupported   bool
+		rowCount          int
+		want              queryOutputMode
+	}{
+		{
+			name:              "json flag always returns json",
+			outputType:        flags.OutputJSON,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          999,
+			want:              queryOutputModeJSON,
+		},
+		{
+			name:              "non interactive stdout returns json",
+			outputType:        flags.OutputText,
+			stdoutInteractive: false,
+			promptSupported:   true,
+			rowCount:          5,
+			want:              queryOutputModeJSON,
+		},
+		{
+			name:              "missing stdin interactivity falls back to static table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   false,
+			rowCount:          staticTableThreshold + 10,
+			want:              queryOutputModeStaticTable,
+		},
+		{
+			name:              "small results use static table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          staticTableThreshold,
+			want:              queryOutputModeStaticTable,
+		},
+		{
+			name:              "large results use interactive table",
+			outputType:        flags.OutputText,
+			stdoutInteractive: true,
+			promptSupported:   true,
+			rowCount:          staticTableThreshold + 1,
+			want:              queryOutputModeInteractiveTable,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := selectQueryOutputMode(tc.outputType, tc.stdoutInteractive, tc.promptSupported, tc.rowCount)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
 func TestFetchAllRowsSingleChunk(t *testing.T) {
 	ctx := cmdio.MockDiscard(context.Background())
 	mockAPI := mocksql.NewMockStatementExecutionInterface(t)

From 7172334df1985c5b1f886afabf0c253fdeaada8a Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Tue, 3 Mar 2026 17:59:52 +0100
Subject: [PATCH 3/6] Fix search highlight on cursor line and clean up
 tableview

Fix a bug where search highlighting on the cursor line would fail
because cursorStyle.Render() wrapped the line in ANSI codes before
highlightSearch tried to match against it. Now search highlighting
is applied first on clean text, then cursor style wraps the result.

Also merge two const blocks and derive totalRows from lines instead
of storing it as redundant state.
---
 experimental/aitools/cmd/query.go |  2 +-
 libs/tableview/tableview.go       | 32 ++++++++++++++-----------------
 2 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go
index 2b255b28e7..5b1f3477ff 100644
--- a/experimental/aitools/cmd/query.go
+++ b/experimental/aitools/cmd/query.go
@@ -48,7 +48,7 @@ const (
 	queryOutputModeInteractiveTable
 )
 
-func selectQueryOutputMode(outputType flags.Output, stdoutInteractive bool, promptSupported bool, rowCount int) queryOutputMode {
+func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSupported bool, rowCount int) queryOutputMode {
 	if outputType == flags.OutputJSON {
 		return queryOutputModeJSON
 	}
diff --git a/libs/tableview/tableview.go b/libs/tableview/tableview.go
index ff337fd1b6..18eca554ce 100644
--- a/libs/tableview/tableview.go
+++ b/libs/tableview/tableview.go
@@ -16,9 +16,6 @@ const (
 	horizontalScrollStep = 4
 	footerHeight         = 1
 	searchFooterHeight   = 2
-)
-
-const (
 	// headerLines is the number of non-data lines at the top (header + separator).
 	headerLines = 2
 )
@@ -36,9 +33,8 @@ func Run(w io.Writer, columns []string, rows [][]string) error {
 	lines := renderTableLines(columns, rows)
 
 	m := model{
-		lines:     lines,
-		totalRows: len(rows),
-		cursor:    headerLines, // Start on first data row.
+		lines:  lines,
+		cursor: headerLines, // Start on first data row.
 	}
 
 	p := tea.NewProgram(m, tea.WithOutput(w))
@@ -133,16 +129,13 @@ func highlightSearch(line, query string) string {
 }
 
 // renderContent builds the viewport content with cursor and search highlighting.
+// Search highlighting is applied first on clean text, then cursor style wraps the result.
 func (m model) renderContent() string {
 	result := make([]string, len(m.lines))
 	for i, line := range m.lines {
 		rendered := highlightSearch(line, m.searchQuery)
 		if i == m.cursor {
-			rendered = cursorStyle.Render(line)
-			if m.searchQuery != "" {
-				// Apply search highlight on top of cursor for the current line.
-				rendered = highlightSearch(cursorStyle.Render(line), m.searchQuery)
-			}
+			rendered = cursorStyle.Render(rendered)
 		}
 		result[i] = rendered
 	}
@@ -150,11 +143,10 @@ func (m model) renderContent() string {
 }
 
 type model struct {
-	viewport  viewport.Model
-	lines     []string
-	totalRows int
-	ready     bool
-	cursor    int // line index of the highlighted row
+	viewport viewport.Model
+	lines    []string
+	ready    bool
+	cursor   int // line index of the highlighted row
 
 	// Search state.
 	searching   bool
@@ -164,6 +156,10 @@ type model struct {
 	matchIdx    int
 }
 
+func (m model) dataRowCount() int {
+	return max(len(m.lines)-headerLines, 0)
+}
+
 func (m model) Init() tea.Cmd {
 	return nil
 }
@@ -321,10 +317,10 @@ func (m model) View() string {
 func (m model) renderFooter() string {
 	if m.searching {
 		prompt := searchStyle.Render("/ " + m.searchInput + "█")
-		return footerStyle.Render(fmt.Sprintf("%d rows", m.totalRows)) + "\n" + prompt
+		return footerStyle.Render(fmt.Sprintf("%d rows", m.dataRowCount())) + "\n" + prompt
 	}
 
-	parts := []string{fmt.Sprintf("%d rows", m.totalRows)}
+	parts := []string{fmt.Sprintf("%d rows", m.dataRowCount())}
 
 	if m.searchQuery != "" && len(m.matchLines) > 0 {
 		parts = append(parts, fmt.Sprintf("match %d/%d", m.matchIdx+1, len(m.matchLines)))

From 7c9a4bc355bc0c68fda91269693919d443669a73 Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Fri, 6 Mar 2026 10:21:32 +0100
Subject: [PATCH 4/6] Use t.Context() instead of context.Background() in tests

Fix gocritic lint: ruleguard forbids context.Background() in tests,
use t.Context() instead.
---
 experimental/aitools/cmd/query_test.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go
index 4a904b4949..22b4fbaf13 100644
--- a/experimental/aitools/cmd/query_test.go
+++ b/experimental/aitools/cmd/query_test.go
@@ -220,7 +220,7 @@ func TestSelectQueryOutputMode(t *testing.T) {
 }
 
 func TestFetchAllRowsSingleChunk(t *testing.T) {
-	ctx := cmdio.MockDiscard(context.Background())
+	ctx := cmdio.MockDiscard(t.Context())
 	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
 
 	resp := &sql.StatementResponse{
@@ -235,7 +235,7 @@ func TestFetchAllRowsSingleChunk(t *testing.T) {
 }
 
 func TestFetchAllRowsMultiChunk(t *testing.T) {
-	ctx := cmdio.MockDiscard(context.Background())
+	ctx := cmdio.MockDiscard(t.Context())
 	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
 
 	resp := &sql.StatementResponse{
@@ -255,7 +255,7 @@ func TestFetchAllRowsMultiChunk(t *testing.T) {
 }
 
 func TestFetchAllRowsNilResult(t *testing.T) {
-	ctx := cmdio.MockDiscard(context.Background())
+	ctx := cmdio.MockDiscard(t.Context())
 	mockAPI := mocksql.NewMockStatementExecutionInterface(t)
 
 	resp := &sql.StatementResponse{StatementId: "stmt-1"}

From 1f199bf37397bd87b3eb5ef20b2456e2e8698319 Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Fri, 6 Mar 2026 13:11:59 +0100
Subject: [PATCH 5/6] Move JSON row count footer to stderr for parseable output

The Row count footer after JSON output made stdout non-parseable
by tools like jq. Move it to stderr via cmdio.LogString so stdout
is clean JSON that can be piped directly.

Addresses review feedback from pietern.
---
 experimental/aitools/cmd/query.go       |  2 +-
 experimental/aitools/cmd/render.go      | 10 +++++++---
 experimental/aitools/cmd/render_test.go | 12 ++++++++----
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go
index dd5b97c761..fb44d89c15 100644
--- a/experimental/aitools/cmd/query.go
+++ b/experimental/aitools/cmd/query.go
@@ -128,7 +128,7 @@ tables, and large results open an interactive table browser.`,
 
 			switch selectQueryOutputMode(root.OutputType(cmd), stdoutInteractive, promptSupported, len(rows)) {
 			case queryOutputModeJSON:
-				return renderJSON(cmd.OutOrStdout(), columns, rows)
+				return renderJSON(ctx, cmd.OutOrStdout(), columns, rows)
 			case queryOutputModeStaticTable:
 				return renderStaticTable(cmd.OutOrStdout(), columns, rows)
 			default:
diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go
index 47fac571db..7c206e6810 100644
--- a/experimental/aitools/cmd/render.go
+++ b/experimental/aitools/cmd/render.go
@@ -1,12 +1,14 @@
 package mcp
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"strings"
 	"text/tabwriter"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/tableview"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 )
@@ -28,8 +30,9 @@ func extractColumns(manifest *sql.ResultManifest) []string {
 	return columns
 }
 
-// renderJSON writes query results as a JSON array of objects with a row count footer.
-func renderJSON(w io.Writer, columns []string, rows [][]string) error {
+// renderJSON writes query results as a parseable JSON array to stdout.
+// Row count is written to stderr so stdout remains valid JSON for piping.
+func renderJSON(ctx context.Context, w io.Writer, columns []string, rows [][]string) error {
 	objects := make([]map[string]any, len(rows))
 	for i, row := range rows {
 		obj := make(map[string]any, len(columns))
@@ -46,7 +49,8 @@ func renderJSON(w io.Writer, columns []string, rows [][]string) error {
 		return fmt.Errorf("marshal results: %w", err)
 	}
 
-	fmt.Fprintf(w, "%s\n\nRow count: %d\n", output, len(rows))
+	fmt.Fprintf(w, "%s\n", output)
+	cmdio.LogString(ctx, fmt.Sprintf("\nRow count: %d", len(rows)))
 	return nil
 }
 
diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go
index 05c251bf92..dc557c2c78 100644
--- a/experimental/aitools/cmd/render_test.go
+++ b/experimental/aitools/cmd/render_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -39,29 +40,32 @@ func TestExtractColumns(t *testing.T) {
 }
 
 func TestRenderJSON(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
 	var buf bytes.Buffer
 	columns := []string{"id", "name"}
 	rows := [][]string{{"1", "alice"}, {"2", "bob"}}
 
-	err := renderJSON(&buf, columns, rows)
+	err := renderJSON(ctx, &buf, columns, rows)
 	require.NoError(t, err)
 
 	output := buf.String()
 	assert.Contains(t, output, `"alice"`)
 	assert.Contains(t, output, `"bob"`)
-	assert.Contains(t, output, "Row count: 2")
+	// Row count goes to stderr, not stdout. Stdout should be valid JSON.
+	assert.NotContains(t, output, "Row count")
 }
 
 func TestRenderJSONNoRows(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
 	var buf bytes.Buffer
 	columns := []string{"id"}
 	var rows [][]string
 
-	err := renderJSON(&buf, columns, rows)
+	err := renderJSON(ctx, &buf, columns, rows)
 	require.NoError(t, err)
 
 	output := buf.String()
-	assert.Contains(t, output, "Row count: 0")
+	assert.NotContains(t, output, "Row count")
 }
 
 func TestRenderStaticTable(t *testing.T) {

From 24ae985c182ad8f5c5b4d5764a4ee2ee33e8c028 Mon Sep 17 00:00:00 2001
From: simon <simon.faltum@databricks.com>
Date: Fri, 6 Mar 2026 16:29:44 +0100
Subject: [PATCH 6/6] Drop row count footer from JSON output entirely

Clean JSON output with no footer for scripting and piping to jq.
---
 experimental/aitools/cmd/query.go       | 2 +-
 experimental/aitools/cmd/render.go      | 5 +----
 experimental/aitools/cmd/render_test.go | 8 ++------
 3 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go
index fb44d89c15..dd5b97c761 100644
--- a/experimental/aitools/cmd/query.go
+++ b/experimental/aitools/cmd/query.go
@@ -128,7 +128,7 @@ tables, and large results open an interactive table browser.`,
 
 			switch selectQueryOutputMode(root.OutputType(cmd), stdoutInteractive, promptSupported, len(rows)) {
 			case queryOutputModeJSON:
-				return renderJSON(ctx, cmd.OutOrStdout(), columns, rows)
+				return renderJSON(cmd.OutOrStdout(), columns, rows)
 			case queryOutputModeStaticTable:
 				return renderStaticTable(cmd.OutOrStdout(), columns, rows)
 			default:
diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go
index 7c206e6810..b7eadb401c 100644
--- a/experimental/aitools/cmd/render.go
+++ b/experimental/aitools/cmd/render.go
@@ -1,14 +1,12 @@
 package mcp
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"strings"
 	"text/tabwriter"
 
-	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/tableview"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 )
@@ -32,7 +30,7 @@ func extractColumns(manifest *sql.ResultManifest) []string {
 
 // renderJSON writes query results as a parseable JSON array to stdout.
 // Row count is written to stderr so stdout remains valid JSON for piping.
-func renderJSON(ctx context.Context, w io.Writer, columns []string, rows [][]string) error {
+func renderJSON(w io.Writer, columns []string, rows [][]string) error {
 	objects := make([]map[string]any, len(rows))
 	for i, row := range rows {
 		obj := make(map[string]any, len(columns))
@@ -50,7 +48,6 @@ func renderJSON(ctx context.Context, w io.Writer, columns []string, rows [][]str
 	}
 
 	fmt.Fprintf(w, "%s\n", output)
-	cmdio.LogString(ctx, fmt.Sprintf("\nRow count: %d", len(rows)))
 	return nil
 }
 
diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go
index dc557c2c78..f559609749 100644
--- a/experimental/aitools/cmd/render_test.go
+++ b/experimental/aitools/cmd/render_test.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"testing"
 
-	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -40,28 +39,25 @@ func TestExtractColumns(t *testing.T) {
 }
 
 func TestRenderJSON(t *testing.T) {
-	ctx := cmdio.MockDiscard(t.Context())
 	var buf bytes.Buffer
 	columns := []string{"id", "name"}
 	rows := [][]string{{"1", "alice"}, {"2", "bob"}}
 
-	err := renderJSON(ctx, &buf, columns, rows)
+	err := renderJSON(&buf, columns, rows)
 	require.NoError(t, err)
 
 	output := buf.String()
 	assert.Contains(t, output, `"alice"`)
 	assert.Contains(t, output, `"bob"`)
-	// Row count goes to stderr, not stdout. Stdout should be valid JSON.
 	assert.NotContains(t, output, "Row count")
 }
 
 func TestRenderJSONNoRows(t *testing.T) {
-	ctx := cmdio.MockDiscard(t.Context())
 	var buf bytes.Buffer
 	columns := []string{"id"}
 	var rows [][]string
 
-	err := renderJSON(ctx, &buf, columns, rows)
+	err := renderJSON(&buf, columns, rows)
 	require.NoError(t, err)
 
 	output := buf.String()