fix: harden MCP statement output surface

internal/mycli/cli_mcp.go

@@ -30,11 +30,48 @@ import (
 	"github.com/samber/lo"
 )
 
+const mcpMaxOutputBytes = 1 << 20
+
 // ExecuteStatementArgs represents the arguments for the execute_statement tool
 type ExecuteStatementArgs struct {
 	Statement string `json:"statement" jsonschema:"Valid spanner-mycli statement to execute"`
 }
 
+type mcpOutputCapture struct {
+	limit     int
+	builder   strings.Builder
+	truncated bool
+}
+
+func newMCPOutputCapture(limit int) *mcpOutputCapture {
+	if limit < 0 {
+		limit = 0
+	}
+	return &mcpOutputCapture{limit: limit}
+}
+
+func (c *mcpOutputCapture) Write(p []byte) (int, error) {
+	if c.truncated {
+		return len(p), nil
+	}
+
+	if remaining := c.limit - c.builder.Len(); remaining > 0 {
+		if len(p) <= remaining {
+			_, _ = c.builder.Write(p)
+			return len(p), nil
+		}
+		_, _ = c.builder.Write(p[:remaining])
+	}
+
+	c.truncated = true
+	fmt.Fprintf(&c.builder, "\n\n[spanner-mycli MCP output truncated after %d bytes]\n", c.limit)
+	return len(p), nil
+}
+
+func (c *mcpOutputCapture) String() string {
+	return c.builder.String()
+}
+
 // executeStatementHandler handles the execute_statement tool
 func executeStatementHandler(cli *Cli) func(context.Context, *mcp.CallToolRequest, ExecuteStatementArgs) (*mcp.CallToolResult, any, error) {
 	// Mutex to protect concurrent access to cli.executeStatement
@@ -72,11 +109,22 @@ func executeStatementHandler(cli *Cli) func(context.Context, *mcp.CallToolReques
 			}, nil, nil
 		}
 
-		// Create a string builder to capture the output
-		var sb strings.Builder
+		if _, ok := stmt.(MetaCommandStatement); ok {
+			slog.Debug("MCP request rejected meta command",
+				"duration", time.Since(start))
+			// Per MCP spec, return execution errors as tool output, not protocol errors.
+			return &mcp.CallToolResult{
+				Content: []mcp.Content{
+					&mcp.TextContent{Text: "ERROR: meta commands are not supported by MCP execute_statement"},
+				},
+			}, nil, nil
+		}
+
+		// Capture output without allowing one MCP call to grow memory unbounded.
+		output := newMCPOutputCapture(mcpMaxOutputBytes)
 
-		// Execute the statement with the string builder as the output
-		_, err = cli.executeStatement(ctx, stmt, false, statement, &sb)
+		// Execute the statement with the capped capture as the output
+		_, err = cli.executeStatement(ctx, stmt, false, statement, output)
 		if err != nil {
 			slog.Debug("MCP execution failed",
 				"error", err.Error(),
@@ -89,17 +137,18 @@ func executeStatementHandler(cli *Cli) func(context.Context, *mcp.CallToolReques
 			}, nil, nil
 		}
 
+		text := output.String()
 		result := &mcp.CallToolResult{
 			Content: []mcp.Content{
-				&mcp.TextContent{Text: sb.String()},
+				&mcp.TextContent{Text: text},
 			},
 		}
 
 		// Log response if debug logging is enabled
 		slog.Debug("MCP response sent",
-			"output_length", len(sb.String()),
+			"output_length", len(text),
 			"duration", time.Since(start),
-			"output_preview", tabwrap.Truncate(sb.String(), 100, "..."))
+			"output_preview", tabwrap.Truncate(text, 100, "..."))
 
 		return result, nil, nil
 	}
@@ -121,7 +170,9 @@ Supports:
 - GQL queries for graph-based data
 - Client-side commands: HELP, SHOW TABLES/COLUMNS/INDEX, SET variables, USE database
 
-The result is returned as an ASCII-formatted table. When displaying in Markdown, wrap the output in a code block.
+Backslash meta commands are not supported by this MCP tool.
+
+The result is returned as captured text using the active CLI output format (table by default). When displaying in Markdown, wrap the output in a code block.
 
 Use "HELP" command to see all available statements and their syntax.`)
 

internal/mycli/execute_ddl.go

@@ -41,6 +41,17 @@ var replacerForProgress = strings.NewReplacer(
 	"\t", " ",
 )
 
+func newProgressWithTTY(ctx context.Context, session *Session) *mpb.Progress {
+	if session == nil || session.systemVariables == nil || session.systemVariables.StreamManager == nil {
+		return nil
+	}
+	ttyStream := session.systemVariables.StreamManager.GetTtyStream()
+	if ttyStream == nil {
+		return nil
+	}
+	return mpb.NewWithContext(ctx, mpb.WithOutput(ttyStream))
+}
+
 func executeDdlStatements(ctx context.Context, session *Session, ddls []string) (*Result, error) {
 	if len(ddls) == 0 {
 		return &Result{
@@ -64,19 +75,21 @@ func executeDdlStatements(ctx context.Context, session *Session, ddls []string)
 		}
 	}
 	if session.systemVariables.Display.EnableProgressBar {
-		p = mpb.NewWithContext(ctx)
-
-		for _, ddl := range ddls {
-			bar := p.AddBar(int64(100),
-				mpb.PrependDecorators(
-					decor.Spinner(nil, decor.WCSyncSpaceR),
-					decor.Name(tabwrap.Truncate(replacerForProgress.Replace(ddl), 40, "..."), decor.WCSyncSpaceR),
-					decor.Percentage(decor.WCSyncSpace),
-					decor.Elapsed(decor.ET_STYLE_MMSS, decor.WCSyncSpace)),
-				mpb.BarRemoveOnComplete(),
-			)
-			bar.EnableTriggerComplete()
-			bars = append(bars, bar)
+		p = newProgressWithTTY(ctx, session)
+
+		if p != nil {
+			for _, ddl := range ddls {
+				bar := p.AddBar(int64(100),
+					mpb.PrependDecorators(
+						decor.Spinner(nil, decor.WCSyncSpaceR),
+						decor.Name(tabwrap.Truncate(replacerForProgress.Replace(ddl), 40, "..."), decor.WCSyncSpaceR),
+						decor.Percentage(decor.WCSyncSpace),
+						decor.Elapsed(decor.ET_STYLE_MMSS, decor.WCSyncSpace)),
+					mpb.BarRemoveOnComplete(),
+				)
+				bar.EnableTriggerComplete()
+				bars = append(bars, bar)
+			}
 		}
 	}
 

internal/mycli/output_context_test.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"context"
 	"io"
-	"runtime"
 	"strings"
 	"testing"
 
@@ -170,38 +169,133 @@ func TestMCPHandler_streamedFormatDoesNotLeakToGlobalStream(t *testing.T) {
 	}
 }
 
-// TestMCPHandler_shellMetaCommandDoesNotLeakToGlobalStream pins that \!
-// shell command stdout goes to the per-statement writer, not the global
-// stream: the MCP handler can parse and execute meta commands, so shell
-// output written to the StreamManager writer would corrupt the JSON-RPC
-// stream under --mcp.
-func TestMCPHandler_shellMetaCommandDoesNotLeakToGlobalStream(t *testing.T) {
+func TestMCPHandler_rejectsMetaCommands(t *testing.T) {
 	t.Parallel()
-	if runtime.GOOS == "windows" {
-		t.Skip("shell command syntax differs on Windows")
+
+	tests := []struct {
+		name              string
+		statement         string
+		skipSystemCommand bool
+	}{
+		{
+			name:      "shell command",
+			statement: `\! echo mcp-shell-probe`,
+		},
+		{
+			name:              "shell command with skip system command",
+			statement:         `\! echo mcp-shell-probe`,
+			skipSystemCommand: true,
+		},
+		{
+			name:      "output redirect",
+			statement: `\o out.log`,
+		},
+		{
+			name:      "tee output",
+			statement: `\T out.log`,
+		},
 	}
 
-	var global bytes.Buffer
-	session := newDetachedTestSession(&global)
-	// ShellMetaCommand is not detached-compatible; the shell command itself
-	// never touches the (nil) client.
-	session.mode = DatabaseConnected
-	cli := &Cli{
-		SessionHandler:  NewSessionHandler(session),
-		SystemVariables: session.systemVariables,
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var global bytes.Buffer
+			session := newDetachedTestSession(&global)
+			session.systemVariables.Config.SkipSystemCommand = tt.skipSystemCommand
+			cli := &Cli{
+				SessionHandler:  NewSessionHandler(session),
+				SystemVariables: session.systemVariables,
+			}
+
+			handler := executeStatementHandler(cli)
+			result, _, err := handler(context.Background(), nil, ExecuteStatementArgs{Statement: tt.statement})
+			if err != nil {
+				t.Fatalf("handler: %v", err)
+			}
+
+			text := mcpResultText(t, result)
+			if !strings.Contains(text, "ERROR: meta commands are not supported by MCP execute_statement") {
+				t.Errorf("tool result = %q, want meta-command rejection", text)
+			}
+			if strings.Contains(text, "mcp-shell-probe") {
+				t.Errorf("tool result = %q, shell command appears to have executed", text)
+			}
+			if global.Len() != 0 {
+				t.Errorf("StreamManager writer got %q, want empty (would corrupt the MCP protocol stream)", global.String())
+			}
+		})
 	}
+}
 
-	handler := executeStatementHandler(cli)
-	result, _, err := handler(context.Background(), nil, ExecuteStatementArgs{Statement: `\! echo mcp-shell-probe`})
-	if err != nil {
-		t.Fatalf("handler: %v", err)
+func TestMCPOutputCaptureTruncates(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		limit      int
+		chunks     []string
+		wantPrefix string
+		notWant    string
+	}{
+		{
+			name:       "partial chunk at limit",
+			limit:      8,
+			chunks:     []string{"12345", "67890"},
+			wantPrefix: "12345678",
+			notWant:    "90",
+		},
+		{
+			name:       "discard writes after truncation",
+			limit:      10,
+			chunks:     []string{"12345", "67890", "overflow", "more"},
+			wantPrefix: "1234567890",
+			notWant:    "overflow",
+		},
 	}
 
-	text := mcpResultText(t, result)
-	if !strings.Contains(text, "mcp-shell-probe") {
-		t.Errorf("tool result %q does not contain shell output", text)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			capture := newMCPOutputCapture(tt.limit)
+			for _, chunk := range tt.chunks {
+				n, err := capture.Write([]byte(chunk))
+				if err != nil {
+					t.Fatalf("Write(%q): %v", chunk, err)
+				}
+				if n != len(chunk) {
+					t.Fatalf("Write(%q) = %d, want %d", chunk, n, len(chunk))
+				}
+			}
+
+			got := capture.String()
+			if !strings.HasPrefix(got, tt.wantPrefix) {
+				t.Errorf("capture prefix = %q, want %q", got, tt.wantPrefix)
+			}
+			if !strings.Contains(got, "spanner-mycli MCP output truncated after") {
+				t.Errorf("capture = %q, want truncation marker", got)
+			}
+			if strings.Contains(got, tt.notWant) {
+				t.Errorf("capture = %q, should not contain discarded content %q", got, tt.notWant)
+			}
+		})
+	}
+}
+
+func TestProgressWithTTYDisabledWithoutTTY(t *testing.T) {
+	t.Parallel()
+
+	var global bytes.Buffer
+	session := newDetachedTestSession(&global)
+	session.systemVariables.Display.EnableProgressBar = true
+
+	progress := newProgressWithTTY(context.Background(), session)
+	if progress != nil {
+		progress.Wait()
+		t.Fatal("newProgressWithTTY returned a progress container without a TTY")
 	}
 	if global.Len() != 0 {
-		t.Errorf("StreamManager writer got %q, want empty (would corrupt the MCP protocol stream)", global.String())
+		t.Errorf("StreamManager writer got %q, want empty", global.String())
 	}
 }

internal/mycli/statements.go

@@ -345,17 +345,19 @@ func (s *ShowOperationStatement) executeSyncMode(ctx context.Context, session *S
 	operationDesc := s.getOperationDescription(op)
 
 	if session.systemVariables.Display.EnableProgressBar {
-		p = mpb.NewWithContext(ctx)
-		bar = p.AddBar(int64(100),
-			mpb.PrependDecorators(
-				decor.Spinner(nil, decor.WCSyncSpaceR),
-				decor.Name(tabwrap.Truncate(replacerForProgress.Replace(operationDesc), 40, "..."), decor.WCSyncSpaceR),
-				decor.Percentage(decor.WCSyncSpace),
-				decor.Elapsed(decor.ET_STYLE_MMSS, decor.WCSyncSpace)),
-			mpb.BarRemoveOnComplete(),
-		)
-		bar.EnableTriggerComplete()
-		defer teardown()
+		p = newProgressWithTTY(ctx, session)
+		if p != nil {
+			bar = p.AddBar(int64(100),
+				mpb.PrependDecorators(
+					decor.Spinner(nil, decor.WCSyncSpaceR),
+					decor.Name(tabwrap.Truncate(replacerForProgress.Replace(operationDesc), 40, "..."), decor.WCSyncSpaceR),
+					decor.Percentage(decor.WCSyncSpace),
+					decor.Elapsed(decor.ET_STYLE_MMSS, decor.WCSyncSpace)),
+				mpb.BarRemoveOnComplete(),
+			)
+			bar.EnableTriggerComplete()
+			defer teardown()
+		}
 	}
 
 	// Update progress bar with initial status