Add screenshot paste support (#2679)

2025-05-28 11:11:22 -07:00
parent d8fa740067
commit feb7b15c76
10 changed files with 828 additions and 58 deletions
@@ -2,12 +2,24 @@ import React, { useRef, useState, useEffect, useCallback } from 'react';
 import { Button } from './ui/button';
 import type { View } from '../App';
 import Stop from './ui/Stop';
-import { Attach, Send } from './icons';
+import { Attach, Send, Close } from './icons';
 import { debounce } from 'lodash';
 import BottomMenu from './bottom_menu/BottomMenu';
 import { LocalMessageStorage } from '../utils/localMessageStorage';
 import { Message } from '../types/message';

+interface PastedImage {
+  id: string;
+  dataUrl: string; // For immediate preview
+  filePath?: string; // Path on filesystem after saving
+  isLoading: boolean;
+  error?: string;
+}
+
+// Constants for image handling
+const MAX_IMAGES_PER_MESSAGE = 5;
+const MAX_IMAGE_SIZE_MB = 5;
+
 interface ChatInputProps {
  handleSubmit: (e: React.FormEvent) => void;
  isLoading?: boolean;
@@ -37,15 +49,28 @@ export default function ChatInput({
  const [_value, setValue] = useState(initialValue);
  const [displayValue, setDisplayValue] = useState(initialValue); // For immediate visual feedback
  const [isFocused, setIsFocused] = useState(false);
+  const [pastedImages, setPastedImages] = useState<PastedImage[]>([]);

  // Update internal value when initialValue changes
  useEffect(() => {
    setValue(initialValue);
    setDisplayValue(initialValue);
+
+    // Use a functional update to get the current pastedImages
+    // and perform cleanup. This avoids needing pastedImages in the deps.
+    setPastedImages((currentPastedImages) => {
+      currentPastedImages.forEach((img) => {
+        if (img.filePath) {
+          window.electron.deleteTempFile(img.filePath);
+        }
+      });
+      return []; // Return a new empty array
+    });
+
    // Reset history index when input is cleared
    setHistoryIndex(-1);
    setIsInGlobalHistory(false);
-  }, [initialValue]);
+  }, [initialValue]); // Keep only initialValue as a dependency

  // State to track if the IME is composing (i.e., in the middle of Japanese IME input)
  const [isComposing, setIsComposing] = useState(false);
@@ -55,6 +80,48 @@ export default function ChatInput({
  const textAreaRef = useRef<HTMLTextAreaElement>(null);
  const [processedFilePaths, setProcessedFilePaths] = useState<string[]>([]);

+  const handleRemovePastedImage = (idToRemove: string) => {
+    const imageToRemove = pastedImages.find((img) => img.id === idToRemove);
+    if (imageToRemove?.filePath) {
+      window.electron.deleteTempFile(imageToRemove.filePath);
+    }
+    setPastedImages((currentImages) => currentImages.filter((img) => img.id !== idToRemove));
+  };
+
+  const handleRetryImageSave = async (imageId: string) => {
+    const imageToRetry = pastedImages.find((img) => img.id === imageId);
+    if (!imageToRetry || !imageToRetry.dataUrl) return;
+
+    // Set the image to loading state
+    setPastedImages((prev) =>
+      prev.map((img) =>
+        img.id === imageId
+          ? { ...img, isLoading: true, error: undefined }
+          : img
+      )
+    );
+
+    try {
+      const result = await window.electron.saveDataUrlToTemp(imageToRetry.dataUrl, imageId);
+      setPastedImages((prev) =>
+        prev.map((img) =>
+          img.id === result.id
+            ? { ...img, filePath: result.filePath, error: result.error, isLoading: false }
+            : img
+        )
+      );
+    } catch (err) {
+      console.error('Error retrying image save:', err);
+      setPastedImages((prev) =>
+        prev.map((img) =>
+          img.id === imageId
+            ? { ...img, error: 'Failed to save image via Electron.', isLoading: false }
+            : img
+        )
+      );
+    }
+  };
+
  useEffect(() => {
    if (textAreaRef.current) {
      textAreaRef.current.focus();
@@ -111,6 +178,89 @@ export default function ChatInput({
    debouncedSetValue(val); // Debounce the actual state update
  };

+  const handlePaste = async (evt: React.ClipboardEvent<HTMLTextAreaElement>) => {
+    const files = Array.from(evt.clipboardData.files || []);
+    const imageFiles = files.filter(file => file.type.startsWith('image/'));
+    
+    if (imageFiles.length === 0) return;
+    
+    // Check if adding these images would exceed the limit
+    if (pastedImages.length + imageFiles.length > MAX_IMAGES_PER_MESSAGE) {
+      // Show error message to user
+      setPastedImages((prev) => [
+        ...prev,
+        {
+          id: `error-${Date.now()}`,
+          dataUrl: '',
+          isLoading: false,
+          error: `Cannot paste ${imageFiles.length} image(s). Maximum ${MAX_IMAGES_PER_MESSAGE} images per message allowed.`
+        }
+      ]);
+      
+      // Remove the error message after 3 seconds
+      setTimeout(() => {
+        setPastedImages((prev) => prev.filter(img => !img.id.startsWith('error-')));
+      }, 3000);
+      
+      return;
+    }
+    
+    evt.preventDefault();
+    
+    for (const file of imageFiles) {
+      // Check individual file size before processing
+      if (file.size > MAX_IMAGE_SIZE_MB * 1024 * 1024) {
+        const errorId = `error-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
+        setPastedImages((prev) => [
+          ...prev,
+          {
+            id: errorId,
+            dataUrl: '',
+            isLoading: false,
+            error: `Image too large (${Math.round(file.size / (1024 * 1024))}MB). Maximum ${MAX_IMAGE_SIZE_MB}MB allowed.`
+          }
+        ]);
+        
+        // Remove the error message after 3 seconds
+        setTimeout(() => {
+          setPastedImages((prev) => prev.filter(img => img.id !== errorId));
+        }, 3000);
+        
+        continue;
+      }
+      
+      const reader = new FileReader();
+      reader.onload = async (e) => {
+        const dataUrl = e.target?.result as string;
+        if (dataUrl) {
+          const imageId = `img-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
+          setPastedImages((prev) => [...prev, { id: imageId, dataUrl, isLoading: true }]);
+
+          try {
+            const result = await window.electron.saveDataUrlToTemp(dataUrl, imageId);
+            setPastedImages((prev) =>
+              prev.map((img) =>
+                img.id === result.id
+                  ? { ...img, filePath: result.filePath, error: result.error, isLoading: false }
+                  : img
+              )
+            );
+          } catch (err) {
+            console.error('Error saving pasted image:', err);
+            setPastedImages((prev) =>
+              prev.map((img) =>
+                img.id === imageId
+                  ? { ...img, error: 'Failed to save image via Electron.', isLoading: false }
+                  : img
+              )
+            );
+          }
+        }
+      };
+      reader.readAsDataURL(file);
+    }
+  };
+
  // Cleanup debounced functions on unmount
  useEffect(() => {
    return () => {
@@ -197,6 +347,36 @@ export default function ChatInput({
    }
  };

+  const performSubmit = () => {
+    const validPastedImageFilesPaths = pastedImages
+      .filter((img) => img.filePath && !img.error && !img.isLoading)
+      .map((img) => img.filePath as string);
+
+    let textToSend = displayValue.trim();
+
+    if (validPastedImageFilesPaths.length > 0) {
+      const pathsString = validPastedImageFilesPaths.join(' ');
+      textToSend = textToSend ? `${textToSend} ${pathsString}` : pathsString;
+    }
+
+    if (textToSend) {
+      if (displayValue.trim()) {
+        LocalMessageStorage.addMessage(displayValue);
+      } else if (validPastedImageFilesPaths.length > 0) {
+        LocalMessageStorage.addMessage(validPastedImageFilesPaths.join(' '));
+      }
+
+      handleSubmit(new CustomEvent('submit', { detail: { value: textToSend } }));
+
+      setDisplayValue('');
+      setValue('');
+      setPastedImages([]);
+      setHistoryIndex(-1);
+      setSavedInput('');
+      setIsInGlobalHistory(false);
+    }
+  };
+
  const handleKeyDown = (evt: React.KeyboardEvent<HTMLTextAreaElement>) => {
    // Handle history navigation first
    handleHistoryNavigation(evt);
@@ -207,6 +387,7 @@ export default function ChatInput({
        // Allow line break for Shift+Enter, or during IME composition
        return;
      }
+
      if (evt.altKey) {
        const newValue = displayValue + '\n';
        setDisplayValue(newValue);
@@ -214,44 +395,31 @@ export default function ChatInput({
        return;
      }

-      // Prevent default Enter behavior when loading or when not loading but has content
-      // So it won't trigger a new line
      evt.preventDefault();
-
-      // Only submit if not loading and has content
-      if (!isLoading && displayValue.trim()) {
-        // Always add to global chat storage before submitting
-        LocalMessageStorage.addMessage(displayValue);
-
-        handleSubmit(new CustomEvent('submit', { detail: { value: displayValue } }));
-        setDisplayValue('');
-        setValue('');
-        setHistoryIndex(-1);
-        setSavedInput('');
-        setIsInGlobalHistory(false);
+      const canSubmit =
+        !isLoading &&
+        (displayValue.trim() ||
+          pastedImages.some((img) => img.filePath && !img.error && !img.isLoading));
+      if (canSubmit) {
+        performSubmit();
      }
    }
  };

  const onFormSubmit = (e: React.FormEvent) => {
    e.preventDefault();
-    if (displayValue.trim() && !isLoading) {
-      // Always add to global chat storage before submitting
-      LocalMessageStorage.addMessage(displayValue);
-
-      handleSubmit(new CustomEvent('submit', { detail: { value: displayValue } }));
-      setDisplayValue('');
-      setValue('');
-      setHistoryIndex(-1);
-      setSavedInput('');
-      setIsInGlobalHistory(false);
+    const canSubmit =
+      !isLoading &&
+      (displayValue.trim() ||
+        pastedImages.some((img) => img.filePath && !img.error && !img.isLoading));
+    if (canSubmit) {
+      performSubmit();
    }
  };

  const handleFileSelect = async () => {
    const path = await window.electron.selectFileOrDirectory();
    if (path) {
-      // Append the path to existing text, with a space if there's existing text
      const newValue = displayValue.trim() ? `${displayValue.trim()} ${path}` : path;
      setDisplayValue(newValue);
      setValue(newValue);
@@ -259,6 +427,10 @@ export default function ChatInput({
    }
  };

+  const hasSubmittableContent =
+    displayValue.trim() || pastedImages.some((img) => img.filePath && !img.error && !img.isLoading);
+  const isAnyImageLoading = pastedImages.some((img) => img.isLoading);
+
  return (
    <div
      className={`flex flex-col relative h-auto border rounded-lg transition-colors ${
@@ -278,6 +450,7 @@ export default function ChatInput({
          onCompositionStart={handleCompositionStart}
          onCompositionEnd={handleCompositionEnd}
          onKeyDown={handleKeyDown}
+          onPaste={handlePaste}
          onFocus={() => setIsFocused(true)}
          onBlur={() => setIsFocused(false)}
          ref={textAreaRef}
@@ -290,6 +463,54 @@ export default function ChatInput({
          className="w-full pl-4 pr-[68px] outline-none border-none focus:ring-0 bg-transparent pt-3 pb-1.5 text-sm resize-none text-textStandard placeholder:text-textPlaceholder"
        />

+        {pastedImages.length > 0 && (
+          <div className="flex flex-wrap gap-2 p-2 border-t border-borderSubtle">
+            {pastedImages.map((img) => (
+              <div key={img.id} className="relative group w-20 h-20">
+                {img.dataUrl && (
+                  <img
+                    src={img.dataUrl} // Use dataUrl for instant preview
+                    alt={`Pasted image ${img.id}`}
+                    className={`w-full h-full object-cover rounded border ${img.error ? 'border-red-500' : 'border-borderStandard'}`}
+                  />
+                )}
+                {img.isLoading && (
+                  <div className="absolute inset-0 flex items-center justify-center bg-black bg-opacity-50 rounded">
+                    <div className="animate-spin rounded-full h-6 w-6 border-t-2 border-b-2 border-white"></div>
+                  </div>
+                )}
+                {img.error && !img.isLoading && (
+                  <div className="absolute inset-0 flex flex-col items-center justify-center bg-black bg-opacity-75 rounded p-1 text-center">
+                    <p className="text-red-400 text-[10px] leading-tight break-all mb-1">
+                      {img.error.substring(0, 50)}
+                    </p>
+                    {img.dataUrl && (
+                      <button
+                        type="button"
+                        onClick={() => handleRetryImageSave(img.id)}
+                        className="bg-blue-600 hover:bg-blue-700 text-white rounded px-1 py-0.5 text-[8px] leading-none"
+                        title="Retry saving image"
+                      >
+                        Retry
+                      </button>
+                    )}
+                  </div>
+                )}
+                {!img.isLoading && (
+                  <button
+                    type="button"
+                    onClick={() => handleRemovePastedImage(img.id)}
+                    className="absolute -top-1 -right-1 bg-gray-700 hover:bg-red-600 text-white rounded-full w-5 h-5 flex items-center justify-center text-xs leading-none opacity-0 group-hover:opacity-100 focus:opacity-100 transition-opacity z-10"
+                    aria-label="Remove image"
+                  >
+                    <Close size={14} />
+                  </button>
+                )}
+              </div>
+            ))}
+          </div>
+        )}
+
        {isLoading ? (
          <Button
            type="button"
@@ -309,12 +530,13 @@ export default function ChatInput({
            type="submit"
            size="icon"
            variant="ghost"
-            disabled={!displayValue.trim()}
+            disabled={!hasSubmittableContent || isAnyImageLoading} // Disable if no content or if images are still loading/saving
            className={`absolute right-3 top-2 transition-colors rounded-full w-7 h-7 [&_svg]:size-4 ${
-              !displayValue.trim()
+              !hasSubmittableContent || isAnyImageLoading
                ? 'text-textSubtle cursor-not-allowed'
                : 'bg-bgAppInverse text-textProminentInverse hover:cursor-pointer'
            }`}
+            title={isAnyImageLoading ? 'Waiting for images to save...' : 'Send'}
          >
            <Send />
          </Button>
@@ -296,13 +296,17 @@ function ChatContent({
  const handleSubmit = (e: React.FormEvent) => {
    window.electron.startPowerSaveBlocker();
    const customEvent = e as unknown as CustomEvent;
-    const content = customEvent.detail?.value || '';
+    // ChatInput now sends a single 'value' field with text and appended image paths
+    const combinedTextFromInput = customEvent.detail?.value || '';

-    if (content.trim()) {
+    if (combinedTextFromInput.trim()) {
      setLastInteractionTime(Date.now());

+      // createUserMessage was reverted to only accept text.
+      // It will create a Message with a single TextContent part containing text + paths.
+      const userMessage = createUserMessage(combinedTextFromInput.trim());
+
      if (summarizedThread.length > 0) {
-        // move current `messages` to `ancestorMessages` and `messages` to `summarizedThread`
        resetMessagesWithSummary(
          messages,
          setMessages,
@@ -310,23 +314,21 @@ function ChatContent({
          setAncestorMessages,
          summaryContent
        );
-
-        // update the chat with new sessionId
-
-        // now call the llm
        setTimeout(() => {
-          append(createUserMessage(content));
+          append(userMessage);
          if (scrollRef.current?.scrollToBottom) {
            scrollRef.current.scrollToBottom();
          }
        }, 150);
      } else {
-        // Normal flow (existing code)
-        append(createUserMessage(content));
+        append(userMessage);
        if (scrollRef.current?.scrollToBottom) {
          scrollRef.current.scrollToBottom();
        }
      }
+    } else {
+      // If nothing was actually submitted (e.g. empty input and no images pasted)
+      window.electron.stopPowerSaveBlocker();
    }
  };

@@ -1,7 +1,9 @@
 import React, { useEffect, useMemo, useRef } from 'react';
 import LinkPreview from './LinkPreview';
+import ImagePreview from './ImagePreview';
 import GooseResponseForm from './GooseResponseForm';
 import { extractUrls } from '../utils/urlUtils';
+import { extractImagePaths, removeImagePathsFromText } from '../utils/imageUtils';
 import { formatMessageTimestamp } from '../utils/timeUtils';
 import MarkdownContent from './MarkdownContent';
 import ToolCallWithResponse from './ToolCallWithResponse';
@@ -40,6 +42,13 @@ export default function GooseMessage({
  // Extract text content from the message
  let textContent = getTextContent(message);

+  // Extract image paths from the message
+  const imagePaths = extractImagePaths(textContent);
+
+  // Remove image paths from text for display
+  const displayText =
+    imagePaths.length > 0 ? removeImagePathsFromText(textContent, imagePaths) : textContent;
+
  // Memoize the timestamp
  const timestamp = useMemo(() => formatMessageTimestamp(message.created), [message.created]);

@@ -53,7 +62,7 @@ export default function GooseMessage({
  const messageIndex = messages?.findIndex((msg) => msg.id === message.id);
  const previousMessage = messageIndex > 0 ? messages[messageIndex - 1] : null;
  const previousUrls = previousMessage ? extractUrls(getTextContent(previousMessage)) : [];
-  const urls = toolRequests.length === 0 ? extractUrls(textContent, previousUrls) : [];
+  const urls = toolRequests.length === 0 ? extractUrls(displayText, previousUrls) : [];

  const toolConfirmationContent = getToolConfirmationContent(message);
  const hasToolConfirmation = toolConfirmationContent !== undefined;
@@ -106,8 +115,18 @@ export default function GooseMessage({
        {textContent && (
          <div className="flex flex-col group">
            <div className={`goose-message-content pt-2`}>
-              <div ref={contentRef}>{<MarkdownContent content={textContent} />}</div>
+              <div ref={contentRef}>{<MarkdownContent content={displayText} />}</div>
            </div>
+
+            {/* Render images if any */}
+            {imagePaths.length > 0 && (
+              <div className="flex flex-wrap gap-2 mt-2 mb-2">
+                {imagePaths.map((imagePath, index) => (
+                  <ImagePreview key={index} src={imagePath} alt={`Image ${index + 1}`} />
+                ))}
+              </div>
+            )}
+
            {/* Only show MessageCopyLink if there's text content and no tool requests/responses */}
            <div className="relative flex justify-start">
              {toolRequests.length === 0 && (
@@ -117,7 +136,7 @@ export default function GooseMessage({
              )}
              {textContent && message.content.every((content) => content.type === 'text') && (
                <div className="absolute left-0 pt-1">
-                  <MessageCopyLink text={textContent} contentRef={contentRef} />
+                  <MessageCopyLink text={displayText} contentRef={contentRef} />
                </div>
              )}
            </div>
@@ -0,0 +1,88 @@
+import React, { useState, useEffect } from 'react';
+
+interface ImagePreviewProps {
+  src: string;
+  alt?: string;
+  className?: string;
+}
+
+export default function ImagePreview({
+  src,
+  alt = 'Pasted image',
+  className = '',
+}: ImagePreviewProps) {
+  const [isExpanded, setIsExpanded] = useState(false);
+  const [error, setError] = useState(false);
+  const [isLoading, setIsLoading] = useState(true);
+  const [imageData, setImageData] = useState<string | null>(null);
+
+  useEffect(() => {
+    const loadImage = async () => {
+      try {
+        // Use the IPC handler to get the image data
+        const data = await window.electron.getTempImage(src);
+        if (data) {
+          setImageData(data);
+          setIsLoading(false);
+        } else {
+          setError(true);
+          setIsLoading(false);
+        }
+      } catch (err) {
+        console.error('Error loading image:', err);
+        setError(true);
+        setIsLoading(false);
+      }
+    };
+
+    loadImage();
+  }, [src]);
+
+  const handleError = () => {
+    setError(true);
+    setIsLoading(false);
+  };
+
+  const toggleExpand = () => {
+    if (!error) {
+      setIsExpanded(!isExpanded);
+    }
+  };
+
+  // Validate that this is a safe file path (should contain goose-pasted-images)
+  if (!src.includes('goose-pasted-images')) {
+    return <div className="text-red-500 text-xs italic mt-1 mb-1">Invalid image path: {src}</div>;
+  }
+
+  if (error) {
+    return <div className="text-red-500 text-xs italic mt-1 mb-1">Unable to load image: {src}</div>;
+  }
+
+  return (
+    <div className={`image-preview mt-2 mb-2 ${className}`}>
+      {isLoading && (
+        <div className="animate-pulse bg-gray-200 rounded w-40 h-40 flex items-center justify-center">
+          <span className="text-gray-500 text-xs">Loading...</span>
+        </div>
+      )}
+      {imageData && (
+        <img
+          src={imageData}
+          alt={alt}
+          onError={handleError}
+          onClick={toggleExpand}
+          className={`rounded border border-borderSubtle cursor-pointer hover:border-borderStandard transition-all ${
+            isExpanded ? 'max-w-full max-h-96' : 'max-h-40 max-w-40'
+          } ${isLoading ? 'hidden' : ''}`}
+          style={{ objectFit: 'contain' }}
+        />
+      )}
+      {isExpanded && !error && !isLoading && imageData && (
+        <div className="text-xs text-textSubtle mt-1">Click to collapse</div>
+      )}
+      {!isExpanded && !error && !isLoading && imageData && (
+        <div className="text-xs text-textSubtle mt-1">Click to expand</div>
+      )}
+    </div>
+  );
+}
@@ -1,6 +1,8 @@
 import React, { useRef, useMemo } from 'react';
 import LinkPreview from './LinkPreview';
+import ImagePreview from './ImagePreview';
 import { extractUrls } from '../utils/urlUtils';
+import { extractImagePaths, removeImagePathsFromText } from '../utils/imageUtils';
 import MarkdownContent from './MarkdownContent';
 import { Message, getTextContent } from '../types/message';
 import MessageCopyLink from './MessageCopyLink';
@@ -16,11 +18,17 @@ export default function UserMessage({ message }: UserMessageProps) {
  // Extract text content from the message
  const textContent = getTextContent(message);

+  // Extract image paths from the message
+  const imagePaths = extractImagePaths(textContent);
+
+  // Remove image paths from text for display
+  const displayText = removeImagePathsFromText(textContent, imagePaths);
+
  // Memoize the timestamp
  const timestamp = useMemo(() => formatMessageTimestamp(message.created), [message.created]);

  // Extract URLs which explicitly contain the http:// or https:// protocol
-  const urls = extractUrls(textContent, []);
+  const urls = extractUrls(displayText, []);

  return (
    <div className="flex justify-end mt-[16px] w-full opacity-0 animate-[appear_150ms_ease-in_forwards]">
@@ -29,17 +37,27 @@ export default function UserMessage({ message }: UserMessageProps) {
          <div className="flex bg-slate text-white rounded-xl rounded-br-none py-2 px-3">
            <div ref={contentRef}>
              <MarkdownContent
-                content={textContent}
+                content={displayText}
                className="text-white prose-a:text-white user-message"
              />
            </div>
          </div>
+
+          {/* Render images if any */}
+          {imagePaths.length > 0 && (
+            <div className="flex flex-wrap gap-2 mt-2">
+              {imagePaths.map((imagePath, index) => (
+                <ImagePreview key={index} src={imagePath} alt={`Pasted image ${index + 1}`} />
+              ))}
+            </div>
+          )}
+
          <div className="relative h-[22px] flex justify-end">
            <div className="absolute right-0 text-xs text-textSubtle pt-1 transition-all duration-200 group-hover:-translate-y-4 group-hover:opacity-0">
              {timestamp}
            </div>
            <div className="absolute right-0 pt-1">
-              <MessageCopyLink text={textContent} contentRef={contentRef} />
+              <MessageCopyLink text={displayText} contentRef={contentRef} />
            </div>
          </div>
        </div>
@@ -6,9 +6,11 @@ import BackButton from '../ui/BackButton';
 import { ScrollArea } from '../ui/scroll-area';
 import MarkdownContent from '../MarkdownContent';
 import ToolCallWithResponse from '../ToolCallWithResponse';
+import ImagePreview from '../ImagePreview';
 import { ToolRequestMessageContent, ToolResponseMessageContent } from '../../types/message';
 import { type Message } from '../../types/message';
 import { formatMessageTimestamp } from '../../utils/timeUtils';
+import { extractImagePaths, removeImagePathsFromText } from '../../utils/imageUtils';

 /**
 * Get tool responses map from messages
@@ -106,11 +108,20 @@ export const SessionMessages: React.FC<SessionMessagesProps> = ({
              messages
                .map((message, index) => {
                  // Extract text content from the message
-                  const textContent = message.content
+                  let textContent = message.content
                    .filter((c) => c.type === 'text')
                    .map((c) => c.text)
                    .join('\n');

+                  // Extract image paths from the message
+                  const imagePaths = extractImagePaths(textContent);
+
+                  // Remove image paths from text for display
+                  const displayText =
+                    imagePaths.length > 0
+                      ? removeImagePathsFromText(textContent, imagePaths)
+                      : textContent;
+
                  // Get tool requests from the message
                  const toolRequests = message.content
                    .filter((c) => c.type === 'toolRequest')
@@ -148,9 +159,24 @@ export const SessionMessages: React.FC<SessionMessagesProps> = ({

                      <div className="flex flex-col w-full">
                        {/* Text content */}
-                        {textContent && (
-                          <div className={`${toolRequests.length > 0 ? 'mb-4' : ''}`}>
-                            <MarkdownContent content={textContent} />
+                        {displayText && (
+                          <div
+                            className={`${toolRequests.length > 0 || imagePaths.length > 0 ? 'mb-4' : ''}`}
+                          >
+                            <MarkdownContent content={displayText} />
+                          </div>
+                        )}
+
+                        {/* Render images if any */}
+                        {imagePaths.length > 0 && (
+                          <div className="flex flex-wrap gap-2 mt-2 mb-2">
+                            {imagePaths.map((imagePath, imageIndex) => (
+                              <ImagePreview
+                                key={imageIndex}
+                                src={imagePath}
+                                alt={`Image ${imageIndex + 1}`}
+                              />
+                            ))}
                          </div>
                        )}