fix gemini image responses, gate screenshot tool appropriately based on api capabilities

Author: sawka

docs/docs/waveai.mdx

@@ -34,7 +34,7 @@ Controls AI's access to your workspace:
 
 ## File Attachments
 
-Drag files onto the AI panel to attach:
+Drag files onto the AI panel to attach (not supported with all models):
 
 | Type | Formats | Size Limit | Notes |
 |------|---------|------------|-------|
@@ -68,7 +68,7 @@ Supports text files, images, PDFs, and directories. Use `-n` for new chat, `-s`
 - **Navigate Web**: Changes URLs in web browser widgets
 
 ### All Widgets
-- **Capture Screenshots**: Takes screenshots of any widget for visual analysis
+- **Capture Screenshots**: Takes screenshots of any widget for visual analysis (not supported on all models)
 
 :::warning Security
 File system operations require explicit approval. You control all file access.

pkg/aiusechat/aiutil/aiutil.go

@@ -209,6 +209,13 @@ func CheckModelSubPrefix(model string, prefix string) bool {
 	return false
 }
 
+// GeminiSupportsImageToolResults returns true if the model supports multimodal function responses (images in tool results)
+// This is only supported by Gemini 3 Pro and later models
+func GeminiSupportsImageToolResults(model string) bool {
+	m := strings.ToLower(model)
+	return strings.Contains(m, "gemini-3") || strings.Contains(m, "gemini-4")
+}
+
 // CreateToolUseData creates a UIMessageDataToolUse from tool call information
 func CreateToolUseData(toolCallID, toolName string, arguments string, chatOpts uctypes.WaveChatOpts) uctypes.UIMessageDataToolUse {
 	toolUseData := uctypes.UIMessageDataToolUse{

pkg/aiusechat/gemini/gemini-convertmessage.go

@@ -13,6 +13,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/wavetermdev/waveterm/pkg/aiusechat/aiutil"
 	"github.com/wavetermdev/waveterm/pkg/aiusechat/uctypes"
+	"github.com/wavetermdev/waveterm/pkg/util/utilfn"
 )
 
 // cleanSchemaForGemini removes fields from JSON Schema that Gemini doesn't accept
@@ -232,9 +233,36 @@ func ConvertToolResultsToGeminiChatMessage(toolResults []uctypes.AIToolResult) (
 		}
 
 		response := make(map[string]any)
+		var nestedParts []GeminiMessagePart
+
 		if result.ErrorText != "" {
 			response["ok"] = false
 			response["error"] = result.ErrorText
+		} else if strings.HasPrefix(result.Text, "data:") {
+			mimeType, base64Data, err := utilfn.DecodeDataURL(result.Text)
+			if err != nil {
+				log.Printf("gemini: failed to decode data URL in tool result: %v\n", err)
+				response["ok"] = false
+				response["error"] = fmt.Sprintf("failed to decode data URL: %v", err)
+			} else if strings.HasPrefix(mimeType, "image/") {
+				// For image data URLs, use multimodal function response (Gemini 3 Pro+)
+				displayName := fmt.Sprintf("result_%s.%s", result.ToolUseID[:8], strings.TrimPrefix(mimeType, "image/"))
+				response["ok"] = true
+				response["image"] = map[string]string{"$ref": displayName}
+
+				// Add the image data as a nested part
+				nestedParts = append(nestedParts, GeminiMessagePart{
+					InlineData: &GeminiInlineData{
+						MimeType:    mimeType,
+						Data:        base64.StdEncoding.EncodeToString(base64Data),
+						DisplayName: displayName,
+					},
+				})
+			} else {
+				log.Printf("gemini: unsupported data URL mimetype in tool result: %s\n", mimeType)
+				response["ok"] = false
+				response["error"] = fmt.Sprintf("unsupported data URL mimetype: %s", mimeType)
+			}
 		} else {
 			response["ok"] = true
 			response["result"] = result.Text
@@ -244,6 +272,7 @@ func ConvertToolResultsToGeminiChatMessage(toolResults []uctypes.AIToolResult) (
 			FunctionResponse: &GeminiFunctionResponse{
 				Name:     result.ToolName,
 				Response: response,
+				Parts:    nestedParts,
 			},
 		})
 	}

pkg/aiusechat/gemini/gemini-types.go

@@ -79,14 +79,16 @@ func (p *GeminiMessagePart) Clean() *GeminiMessagePart {
 
 // GeminiInlineData represents inline binary data
 type GeminiInlineData struct {
-	MimeType string `json:"mimeType"`
-	Data     string `json:"data"` // base64 encoded
+	MimeType    string `json:"mimeType"`
+	Data        string `json:"data"` // base64 encoded
+	DisplayName string `json:"displayName,omitempty"` // for multimodal function responses
 }
 
 // GeminiFileData represents uploaded file reference
 type GeminiFileData struct {
-	MimeType string `json:"mimeType"`
-	FileUri  string `json:"fileUri"` // gs:// URI from file upload
+	MimeType    string `json:"mimeType"`
+	FileUri     string `json:"fileUri"` // gs:// URI from file upload
+	DisplayName string `json:"displayName,omitempty"` // for multimodal function responses
 }
 
 // GeminiFunctionCall represents a function call from the model
@@ -97,8 +99,9 @@ type GeminiFunctionCall struct {
 
 // GeminiFunctionResponse represents a function execution result
 type GeminiFunctionResponse struct {
-	Name     string         `json:"name"`
-	Response map[string]any `json:"response"`
+	Name     string              `json:"name"`
+	Response map[string]any      `json:"response"`
+	Parts    []GeminiMessagePart `json:"parts,omitempty"` // nested parts for multimodal content (Gemini 3 Pro and later)
 }
 
 // GeminiUsageMetadata represents token usage

pkg/aiusechat/tools.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/google/uuid"
+	"github.com/wavetermdev/waveterm/pkg/aiusechat/aiutil"
 	"github.com/wavetermdev/waveterm/pkg/aiusechat/uctypes"
 	"github.com/wavetermdev/waveterm/pkg/blockcontroller"
 	"github.com/wavetermdev/waveterm/pkg/util/utilfn"
@@ -132,7 +133,7 @@ func MakeBlockShortDesc(block *waveobj.Block) string {
 	}
 }
 
-func GenerateTabStateAndTools(ctx context.Context, tabid string, widgetAccess bool) (string, []uctypes.ToolDefinition, error) {
+func GenerateTabStateAndTools(ctx context.Context, tabid string, widgetAccess bool, chatOpts *uctypes.WaveChatOpts) (string, []uctypes.ToolDefinition, error) {
 	if tabid == "" {
 		return "", nil, nil
 	}
@@ -160,7 +161,13 @@ func GenerateTabStateAndTools(ctx context.Context, tabid string, widgetAccess bo
 	// log.Printf("TABPROMPT %s\n", tabState)
 	var tools []uctypes.ToolDefinition
 	if widgetAccess {
-		tools = append(tools, GetCaptureScreenshotToolDefinition(tabid))
+		// Only add screenshot tool for:
+		// - openai-responses API type
+		// - google-gemini API type with Gemini 3+ models
+		if chatOpts.Config.APIType == uctypes.APIType_OpenAIResponses ||
+		   (chatOpts.Config.APIType == uctypes.APIType_GoogleGemini && aiutil.GeminiSupportsImageToolResults(chatOpts.Config.Model)) {
+			tools = append(tools, GetCaptureScreenshotToolDefinition(tabid))
+		}
 		tools = append(tools, GetReadTextFileToolDefinition())
 		tools = append(tools, GetReadDirToolDefinition())
 		tools = append(tools, GetWriteTextFileToolDefinition())

pkg/aiusechat/usechat.go

@@ -645,7 +645,7 @@ func WaveAIPostMessageHandler(w http.ResponseWriter, r *http.Request) {
 
 	if req.TabId != "" {
 		chatOpts.TabStateGenerator = func() (string, []uctypes.ToolDefinition, string, error) {
-			tabState, tabTools, err := GenerateTabStateAndTools(r.Context(), req.TabId, req.WidgetAccess)
+			tabState, tabTools, err := GenerateTabStateAndTools(r.Context(), req.TabId, req.WidgetAccess, &chatOpts)
 			return tabState, tabTools, req.TabId, err
 		}
 	}