Reimplement wsh ai, fix text file attaching format

cmd/wsh/cmd/wshcmd-ai.go

@@ -4,166 +4,197 @@
 package cmd
 
 import (
+	"encoding/base64"
 	"fmt"
 	"io"
+	"mime"
 	"os"
-	"strings"
+	"path/filepath"
 
 	"github.com/spf13/cobra"
-	"github.com/wavetermdev/waveterm/pkg/waveobj"
 	"github.com/wavetermdev/waveterm/pkg/wshrpc"
 	"github.com/wavetermdev/waveterm/pkg/wshrpc/wshclient"
 	"github.com/wavetermdev/waveterm/pkg/wshutil"
 )
 
 var aiCmd = &cobra.Command{
-	Use:                   "ai [-] [message...]",
-	Short:                 "Send a message to an AI block",
+	Use:   "ai [options] [files...]",
+	Short: "Append content to Wave AI sidebar prompt",
+	Long: `Append content to Wave AI sidebar prompt (does not auto-submit by default)
+
+Arguments:
+  files...               Files to attach (use '-' for stdin)
+
+Examples:
+  git diff | wsh ai -                    # Pipe diff to AI, ask question in UI
+  wsh ai main.go                         # Attach file, ask question in UI
+  wsh ai *.go -m "find bugs"             # Attach files with message
+  wsh ai -s - -m "review" < log.txt      # Stdin + message, auto-submit
+  wsh ai -n config.json                  # New chat with file attached`,
 	RunE:                  aiRun,
 	PreRunE:               preRunSetupRpcClient,
 	DisableFlagsInUseLine: true,
 }
 
-var aiFileFlags []string
+var aiMessageFlag string
+var aiSubmitFlag bool
 var aiNewBlockFlag bool
 
 func init() {
 	rootCmd.AddCommand(aiCmd)
-	aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI block")
-	aiCmd.Flags().StringArrayVarP(&aiFileFlags, "file", "f", nil, "attach file content (use '-' for stdin)")
+	aiCmd.Flags().StringVarP(&aiMessageFlag, "message", "m", "", "optional message/question to append after files")
+	aiCmd.Flags().BoolVarP(&aiSubmitFlag, "submit", "s", false, "submit the prompt immediately after appending")
+	aiCmd.Flags().BoolVarP(&aiNewBlockFlag, "new", "n", false, "create a new AI chat instead of using existing")
 }
 
-func encodeFile(builder *strings.Builder, file io.Reader, fileName string) error {
-	data, err := io.ReadAll(file)
-	if err != nil {
-		return fmt.Errorf("error reading file: %w", err)
+func getMimeType(filename string) string {
+	ext := filepath.Ext(filename)
+	if ext == "" {
+		return "text/plain"
 	}
-	// Start delimiter with the file name
-	builder.WriteString(fmt.Sprintf("\n@@@start file %q\n", fileName))
-	// Read the file content and write it to the builder
-	builder.Write(data)
-	// End delimiter with the file name
-	builder.WriteString(fmt.Sprintf("\n@@@end file %q\n\n", fileName))
-	return nil
+	mimeType := mime.TypeByExtension(ext)
+	if mimeType == "" {
+		return "text/plain"
+	}
+	return mimeType
+}
+
+func getMaxFileSize(mimeType string) (int, string) {
+	if mimeType == "application/pdf" {
+		return 5 * 1024 * 1024, "5MB"
+	}
+	if mimeType[:6] == "image/" {
+		return 7 * 1024 * 1024, "7MB"
+	}
+	return 200 * 1024, "200KB"
 }
 
 func aiRun(cmd *cobra.Command, args []string) (rtnErr error) {
 	defer func() {
 		sendActivity("ai", rtnErr == nil)
 	}()
 
-	if len(args) == 0 {
+	if len(args) == 0 && aiMessageFlag == "" {
 		OutputHelpMessage(cmd)
-		return fmt.Errorf("no message provided")
+		return fmt.Errorf("no files or message provided")
 	}
 
+	const maxBatchSize = 7 * 1024 * 1024
+	const largeFileThreshold = 1 * 1024 * 1024
+	const maxFileCount = 15
+	const rpcTimeout = 30000
+
+	var allFiles []wshrpc.AIAttachedFile
 	var stdinUsed bool
-	var message strings.Builder
 
-	// Handle file attachments first
-	for _, file := range aiFileFlags {
-		if file == "-" {
+	if len(args) > maxFileCount {
+		return fmt.Errorf("too many files (maximum %d files allowed)", maxFileCount)
+	}
+
+	for _, filePath := range args {
+		var data []byte
+		var fileName string
+		var mimeType string
+		var err error
+
+		if filePath == "-" {
 			if stdinUsed {
 				return fmt.Errorf("stdin (-) can only be used once")
 			}
 			stdinUsed = true
-			if err := encodeFile(&message, os.Stdin, "<stdin>"); err != nil {
+
+			data, err = io.ReadAll(os.Stdin)
+			if err != nil {
 				return fmt.Errorf("reading from stdin: %w", err)
 			}
+			fileName = "stdin"
+			mimeType = "text/plain"
 		} else {
-			fd, err := os.Open(file)
+			fileInfo, err := os.Stat(filePath)
 			if err != nil {
-				return fmt.Errorf("opening file %s: %w", file, err)
+				return fmt.Errorf("accessing file %s: %w", filePath, err)
 			}
-			defer fd.Close()
-			if err := encodeFile(&message, fd, file); err != nil {
-				return fmt.Errorf("reading file %s: %w", file, err)
+			if fileInfo.IsDir() {
+				return fmt.Errorf("%s is a directory, not a file", filePath)
 			}
+
+			data, err = os.ReadFile(filePath)
+			if err != nil {
+				return fmt.Errorf("reading file %s: %w", filePath, err)
+			}
+			fileName = filepath.Base(filePath)
+			mimeType = getMimeType(filePath)
 		}
-	}
 
-	// Default to "waveai" block
-	isDefaultBlock := blockArg == ""
-	if isDefaultBlock {
-		blockArg = "view@waveai"
+		maxSize, sizeStr := getMaxFileSize(mimeType)
+		if len(data) > maxSize {
+			return fmt.Errorf("file %s exceeds maximum size of %s for %s files", fileName, sizeStr, mimeType)
+		}
+
+		allFiles = append(allFiles, wshrpc.AIAttachedFile{
+			Name:   fileName,
+			Type:   mimeType,
+			Size:   len(data),
+			Data64: base64.StdEncoding.EncodeToString(data),
+		})
 	}
-	var fullORef *waveobj.ORef
-	var err error
-	if !aiNewBlockFlag {
-		fullORef, err = resolveSimpleId(blockArg)
+
+	tabId := os.Getenv("WAVETERM_TABID")
+	if tabId == "" {
+		return fmt.Errorf("WAVETERM_TABID environment variable not set")
 	}
-	if (err != nil && isDefaultBlock) || aiNewBlockFlag {
-		// Create new AI block if default block doesn't exist
-		data := &wshrpc.CommandCreateBlockData{
-			BlockDef: &waveobj.BlockDef{
-				Meta: map[string]interface{}{
-					waveobj.MetaKey_View: "waveai",
-				},
-			},
-			Focused: true,
-		}
 
-		newORef, err := wshclient.CreateBlockCommand(RpcClient, *data, &wshrpc.RpcOpts{Timeout: 2000})
-		if err != nil {
-			return fmt.Errorf("creating AI block: %w", err)
+	route := wshutil.MakeTabRouteId(tabId)
+
+	if aiNewBlockFlag {
+		newChatData := wshrpc.CommandWaveAIAddContextData{
+			NewChat: true,
 		}
-		fullORef = &newORef
-		// Wait for the block's route to be available
-		gotRoute, err := wshclient.WaitForRouteCommand(RpcClient, wshrpc.CommandWaitForRouteData{
-			RouteId: wshutil.MakeFeBlockRouteId(fullORef.OID),
-			WaitMs:  4000,
-		}, &wshrpc.RpcOpts{Timeout: 5000})
+		err := wshclient.WaveAIAddContextCommand(RpcClient, newChatData, &wshrpc.RpcOpts{
+			Route:   route,
+			Timeout: rpcTimeout,
+		})
 		if err != nil {
-			return fmt.Errorf("waiting for AI block: %w", err)
+			return fmt.Errorf("creating new chat: %w", err)
 		}
-		if !gotRoute {
-			return fmt.Errorf("AI block route could not be established")
-		}
-	} else if err != nil {
-		return fmt.Errorf("resolving block: %w", err)
 	}
 
-	// Create the route for this block
-	route := wshutil.MakeFeBlockRouteId(fullORef.OID)
+	var smallFiles []wshrpc.AIAttachedFile
+	var smallFilesSize int
 
-	// Then handle main message
-	if args[0] == "-" {
-		if stdinUsed {
-			return fmt.Errorf("stdin (-) can only be used once")
-		}
-		data, err := io.ReadAll(os.Stdin)
-		if err != nil {
-			return fmt.Errorf("reading from stdin: %w", err)
-		}
-		message.Write(data)
-
-		// Also include any remaining arguments (excluding the "-" itself)
-		if len(args) > 1 {
-			if message.Len() > 0 {
-				message.WriteString(" ")
+	for _, file := range allFiles {
+		if file.Size > largeFileThreshold {
+			contextData := wshrpc.CommandWaveAIAddContextData{
+				Files: []wshrpc.AIAttachedFile{file},
+			}
+			err := wshclient.WaveAIAddContextCommand(RpcClient, contextData, &wshrpc.RpcOpts{
+				Route:   route,
+				Timeout: rpcTimeout,
+			})
+			if err != nil {
+				return fmt.Errorf("adding file %s: %w", file.Name, err)
 			}
-			message.WriteString(strings.Join(args[1:], " "))
+		} else {
+			smallFilesSize += file.Size
+			if smallFilesSize > maxBatchSize {
+				return fmt.Errorf("small files total size exceeds maximum batch size of 7MB")
+			}
+			smallFiles = append(smallFiles, file)
 		}
-	} else {
-		message.WriteString(strings.Join(args, " "))
 	}
 
-	if message.Len() == 0 {
-		return fmt.Errorf("message is empty")
-	}
-	if message.Len() > 50*1024 {
-		return fmt.Errorf("current max message size is 50k")
+	finalContextData := wshrpc.CommandWaveAIAddContextData{
+		Files:  smallFiles,
+		Text:   aiMessageFlag,
+		Submit: aiSubmitFlag,
 	}
 
-	messageData := wshrpc.AiMessageData{
-		Message: message.String(),
-	}
-	err = wshclient.AiSendMessageCommand(RpcClient, messageData, &wshrpc.RpcOpts{
+	err := wshclient.WaveAIAddContextCommand(RpcClient, finalContextData, &wshrpc.RpcOpts{
 		Route:   route,
-		Timeout: 2000,
+		Timeout: rpcTimeout,
 	})
 	if err != nil {
-		return fmt.Errorf("sending message: %w", err)
+		return fmt.Errorf("adding context: %w", err)
 	}
 
 	return nil

docs/docs/wsh-reference.mdx

@@ -110,25 +110,45 @@ wsh getmeta -b [other-tab-id] "bg:*" --clear-prefix | wsh setmeta -b tab --json
 
 ## ai
 
-Send messages to new or existing AI blocks directly from the CLI. `-f` passes a file. note that there is a maximum size of 10k for messages and files, so use a tail/grep to cut down file sizes before passing. The `-f` option works great for small files though like shell scripts or `.zshrc` etc. You can use "-" to read input from stdin.
+Append content to the Wave AI sidebar. Files are attached as proper file attachments (supporting images, PDFs, and text), not encoded as text. By default, content is added to the sidebar without auto-submitting, allowing you to review and add more context before sending to the AI.
 
-By default the messages get sent to the first AI block (by blocknum). If no AI block exists, then a new one will be created. Use `-n` to force creation of a new AI block. Use `-b` to target a specific AI block.
+You can attach multiple files at once (up to 15 files). Use `-m` to add a message along with files, `-s` to auto-submit immediately, and `-n` to start a new chat conversation. Use "-" to read from stdin.
 
 ```sh
-wsh ai "how do i write an ls command that sorts files in reverse size order"
-wsh ai -f <(tail -n 20 "my.log") -- "any idea what these error messages mean"
-wsh ai -f README.md "help me update this readme file"
+# Pipe command output to AI (ask question in UI)
+git diff | wsh ai -
+docker logs mycontainer | wsh ai -
 
-# creates a new AI block
-wsh ai -n "tell me a story"
+# Attach files without auto-submit (review in UI first)
+wsh ai main.go utils.go
+wsh ai screenshot.png logs.txt
 
-# targets block number 5
-wsh ai -b 5 "tell me more"
+# Attach files with message
+wsh ai app.py -m "find potential bugs"
+wsh ai *.log -m "analyze these error logs"
 
-# read from stdin and also supply a message
-tail -n 50 mylog.log | wsh ai - "can you tell me what this error means?"
+# Auto-submit immediately
+wsh ai config.json -s -m "explain this configuration"
+tail -n 50 app.log | wsh ai -s - -m "what's causing these errors?"
+
+# Start new chat and attach files
+wsh ai -n report.pdf data.csv -m "summarize these reports"
+
+# Attach different file types (images, PDFs, code)
+wsh ai architecture.png api-spec.pdf server.go -m "review the system design"
 ```
 
+**File Size Limits:**
+- Text files: 200KB maximum
+- PDF files: 5MB maximum
+- Image files: 7MB maximum (accounts for base64 encoding overhead)
+- Maximum 15 files per command
+
+**Flags:**
+- `-m, --message <text>` - Add message text along with files
+- `-s, --submit` - Auto-submit immediately (default waits for user)
+- `-n, --new` - Clear current chat and start fresh conversation
+
 ---
 
 ## editconfig

docs/docs/wsh.mdx

@@ -116,17 +116,40 @@ wsh setvar -b tab SHARED_ENV=staging
 
 ### AI-Assisted Development
 
+The `wsh ai` command appends content to the Wave AI sidebar. By default, files are attached without auto-submitting, allowing you to review and add more context before sending.
+
 ```bash
-# Get AI help with code (uses "-" to read from stdin)
-git diff | wsh ai - "review these changes"
+# Pipe output to AI sidebar (ask question in UI)
+git diff | wsh ai -
+
+# Attach files with a message
+wsh ai main.go utils.go -m "find bugs in these files"
+
+# Auto-submit with message
+wsh ai config.json -s -m "explain this config"
 
-# Get help with a file
-wsh ai -f .zshrc "help me add ~/bin to my path"
+# Start new chat with attached files
+wsh ai -n *.log -m "analyze these logs"
 
-# Debug issues (uses "-" to read from stdin)
-dmesg | wsh ai - "help me understand these errors"
+# Attach multiple file types (images, PDFs, code)
+wsh ai screenshot.png report.pdf app.py -m "review these"
+
+# Debug with stdin and auto-submit
+dmesg | wsh ai -s - -m "help me understand these errors"
 ```
 
+**Flags:**
+- `-` - Read from stdin instead of a file
+- `-m, --message` - Add message text along with files
+- `-s, --submit` - Auto-submit immediately (default is to wait for user)
+- `-n, --new` - Clear chat and start fresh conversation
+
+**File Limits:**
+- Text files: 200KB max
+- PDFs: 5MB max
+- Images: 7MB max
+- Maximum 15 files per command
+
 ## Tips & Features
 
 1. **Working with Blocks**