fix: fixed more widgets

.gitignore

@@ -12,7 +12,6 @@
 .swiftpm/
 migrate_working_dir/
 AGENTS.md
-flutter_*.png
 
 # IntelliJ related
 *.iml

lib/features/chat/widgets/documentation_message_widget.dart

@@ -7,6 +7,7 @@ import 'dart:io' show Platform;
 import '../../../shared/theme/theme_extensions.dart';
 import '../../../shared/widgets/markdown/streaming_markdown_widget.dart';
 import '../../../core/utils/reasoning_parser.dart';
+import 'enhanced_image_attachment.dart';
 
 class DocumentationMessageWidget extends ConsumerStatefulWidget {
   final dynamic message;
@@ -138,7 +139,7 @@ class _DocumentationMessageWidgetState
 
   void _buildCachedAvatar() {
     _cachedAvatar = Padding(
-      padding: const EdgeInsets.only(bottom: 12),
+      padding: const EdgeInsets.only(bottom: 8),
       child: Row(
         children: [
           Container(
@@ -175,6 +176,7 @@ class _DocumentationMessageWidgetState
   void dispose() {
     _fadeController.dispose();
     _slideController.dispose();
+    _throttleTimer?.cancel();
     super.dispose();
   }
 
@@ -202,22 +204,42 @@ class _DocumentationMessageWidgetState
   }
 
   Widget _buildUserMessage() {
-    return Container(
+    final hasImages = widget.message.attachmentIds != null &&
+        widget.message.attachmentIds!.isNotEmpty;
+    final hasText = widget.message.content.isNotEmpty;
+
+    return GestureDetector(
+      onLongPress: () => _toggleActions(),
+      behavior: HitTestBehavior.translucent,
+      child: Container(
         width: double.infinity,
-          margin: const EdgeInsets.only(bottom: 16, left: 50, right: 12),
+        margin: const EdgeInsets.only(bottom: 12, left: 50, right: 12),
         child: Column(
           crossAxisAlignment: CrossAxisAlignment.end,
           children: [
+            // Display images outside and above the text bubble
+            if (hasImages) ...[
+              Row(
+                mainAxisAlignment: MainAxisAlignment.end,
+                children: [
+                  Flexible(
+                    child: _buildUserAttachmentImages(),
+                  ),
+                ],
+              ),
+              if (hasText) const SizedBox(height: Spacing.xs),
+            ],
+
+            // Display text bubble if there's text content
+            if (hasText)
               Row(
                 mainAxisAlignment: MainAxisAlignment.end,
                 children: [
                   Flexible(
-                    child: GestureDetector(
-                      onLongPress: () => _toggleActions(),
                     child: Container(
                       padding: const EdgeInsets.symmetric(
                         horizontal: 16,
-                          vertical: 12,
+                        vertical: 8,
                       ),
                       decoration: BoxDecoration(
                         color: context.conduitTheme.chatBubbleUser,
@@ -232,13 +254,12 @@ class _DocumentationMessageWidgetState
                         style: TextStyle(
                           color: context.conduitTheme.chatBubbleUserText,
                           fontSize: AppTypography.bodyMedium,
-                            height: 1.5,
+                          height: 1.3,
                           letterSpacing: 0.1,
                         ),
                       ),
                     ),
                   ),
-                  ),
                 ],
               ),
             
@@ -249,6 +270,7 @@ class _DocumentationMessageWidgetState
             ],
           ],
         ),
+      ),
     )
         .animate()
         .fadeIn(duration: const Duration(milliseconds: 400))
@@ -261,9 +283,12 @@ class _DocumentationMessageWidgetState
   }
 
   Widget _buildDocumentationMessage() {
-    return Container(
+    return GestureDetector(
+      onLongPress: () => _toggleActions(),
+      behavior: HitTestBehavior.translucent,
+      child: Container(
         width: double.infinity,
-          margin: const EdgeInsets.only(bottom: 24, left: 12, right: 12),
+        margin: const EdgeInsets.only(bottom: 16, left: 12, right: 12),
         child: Column(
           crossAxisAlignment: CrossAxisAlignment.start,
           children: [
@@ -358,13 +383,18 @@ class _DocumentationMessageWidgetState
             ],
 
             // Documentation-style content without heavy bubble; premium markdown
-              GestureDetector(
+            SizedBox(
-                onLongPress: () => _toggleActions(),
-                child: SizedBox(
               width: double.infinity,
               child: Column(
                 crossAxisAlignment: CrossAxisAlignment.start,
                 children: [
+                  // Display attachment images if any (for user uploaded images)
+                  if (widget.message.attachmentIds != null &&
+                      widget.message.attachmentIds!.isNotEmpty) ...[
+                    _buildAttachmentImages(),
+                    const SizedBox(height: Spacing.md),
+                  ],
+                  
                   if (widget.isStreaming &&
                       (widget.message.content.trim().isEmpty ||
                           widget.message.content == '[TYPING_INDICATOR]'))
@@ -383,7 +413,6 @@ class _DocumentationMessageWidgetState
                 ],
               ),
             ),
-              ),
 
             // Action buttons below the message content
             if (_showActions) ...[
@@ -392,6 +421,7 @@ class _DocumentationMessageWidgetState
             ],
           ],
         ),
+      ),
     )
         .animate()
         .fadeIn(duration: const Duration(milliseconds: 300))
@@ -408,15 +438,150 @@ class _DocumentationMessageWidgetState
       return const SizedBox.shrink();
     }
 
+    // Process content to ensure proper image rendering
+    final processedContent = _processContentForImages(content);
+
     return StreamingMarkdownWidget(
-      staticContent: content,
+      staticContent: processedContent,
       isStreaming: widget.isStreaming,
     );
   }
 
+  String _processContentForImages(String content) {
+    // Check if content contains image markdown or base64 data URLs
+    // This ensures images generated by AI are properly formatted
     
+    // Pattern to detect base64 images that might not be in markdown format
+    final base64Pattern = RegExp(r'data:image/[^;]+;base64,[A-Za-z0-9+/]+=*');
     
-  // Removed lightweight streaming text; we now stream markdown with throttling
+    // If we find base64 images not wrapped in markdown, wrap them
+    if (base64Pattern.hasMatch(content) && !content.contains('![')) {
+      content = content.replaceAllMapped(base64Pattern, (match) {
+        final imageData = match.group(0)!;
+        // Check if this image is already in markdown format
+        final markdownCheck = RegExp(r'!\[.*?\]\(' + RegExp.escape(imageData) + r'\)');
+        if (!markdownCheck.hasMatch(content)) {
+          return '\n![Generated Image]($imageData)\n';
+        }
+        return imageData;
+      });
+    }
+    
+    return content;
+  }
+
+  Widget _buildUserAttachmentImages() {
+    if (widget.message.attachmentIds == null ||
+        widget.message.attachmentIds!.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    final imageCount = widget.message.attachmentIds!.length;
+    
+    // Similar to iMessage style but adapted for documentation widget
+    if (imageCount == 1) {
+      return ClipRRect(
+        borderRadius: BorderRadius.circular(AppBorderRadius.lg),
+        child: EnhancedImageAttachment(
+          attachmentId: widget.message.attachmentIds![0],
+          isUserMessage: true,
+          constraints: const BoxConstraints(
+            maxWidth: 280,
+            maxHeight: 350,
+          ),
+        ),
+      );
+    } else if (imageCount == 2) {
+      return Row(
+        mainAxisSize: MainAxisSize.min,
+        mainAxisAlignment: MainAxisAlignment.end,
+        children: widget.message.attachmentIds!.map((attachmentId) {
+          return Padding(
+            padding: EdgeInsets.only(
+              left: attachmentId == widget.message.attachmentIds!.first 
+                  ? 0 
+                  : Spacing.xs,
+            ),
+            child: ClipRRect(
+              borderRadius: BorderRadius.circular(AppBorderRadius.lg),
+              child: EnhancedImageAttachment(
+                attachmentId: attachmentId,
+                isUserMessage: true,
+                constraints: const BoxConstraints(
+                  maxWidth: 135,
+                  maxHeight: 180,
+                ),
+              ),
+            ),
+          );
+        }).toList(),
+      );
+    } else {
+      return Container(
+        constraints: const BoxConstraints(maxWidth: 280),
+        child: Wrap(
+          alignment: WrapAlignment.end,
+          spacing: Spacing.xs,
+          runSpacing: Spacing.xs,
+          children: widget.message.attachmentIds!.map((attachmentId) {
+            return ClipRRect(
+              borderRadius: BorderRadius.circular(AppBorderRadius.md),
+              child: EnhancedImageAttachment(
+                attachmentId: attachmentId,
+                isUserMessage: true,
+                constraints: BoxConstraints(
+                  maxWidth: imageCount == 3 ? 135 : 90,
+                  maxHeight: imageCount == 3 ? 135 : 90,
+                ),
+              ),
+            );
+          }).toList(),
+        ),
+      );
+    }
+  }
+
+  Widget _buildAttachmentImages() {
+    if (widget.message.attachmentIds == null ||
+        widget.message.attachmentIds!.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    final imageCount = widget.message.attachmentIds!.length;
+
+    // Display images in a clean, modern layout for assistant messages
+    if (imageCount == 1) {
+      return ClipRRect(
+        borderRadius: BorderRadius.circular(AppBorderRadius.md),
+        child: EnhancedImageAttachment(
+          attachmentId: widget.message.attachmentIds![0],
+          isMarkdownFormat: true,
+          constraints: const BoxConstraints(
+            maxWidth: 500,
+            maxHeight: 400,
+          ),
+        ),
+      );
+    } else {
+      return Wrap(
+        spacing: Spacing.sm,
+        runSpacing: Spacing.sm,
+        children: widget.message.attachmentIds!.map<Widget>((attachmentId) {
+          return ClipRRect(
+            borderRadius: BorderRadius.circular(AppBorderRadius.md),
+            child: EnhancedImageAttachment(
+              attachmentId: attachmentId,
+              isMarkdownFormat: true,
+              constraints: BoxConstraints(
+                maxWidth: imageCount == 2 ? 245 : 160,
+                maxHeight: imageCount == 2 ? 245 : 160,
+              ),
+            ),
+          );
+        }).toList(),
+      );
+    }
+  }
 
   Widget _buildTypingIndicator() {
     return Consumer(

lib/features/chat/widgets/enhanced_image_attachment.dart

@@ -0,0 +1,401 @@
+import 'dart:convert';
+import 'package:flutter/material.dart';
+import 'package:flutter_riverpod/flutter_riverpod.dart';
+import 'package:cached_network_image/cached_network_image.dart';
+import '../../../shared/theme/theme_extensions.dart';
+import '../../../core/providers/app_providers.dart';
+
+// Global cache for image data to prevent reloading
+final _globalImageCache = <String, String>{};
+
+class EnhancedImageAttachment extends ConsumerStatefulWidget {
+  final String attachmentId;
+  final bool isMarkdownFormat;
+  final VoidCallback? onTap;
+  final BoxConstraints? constraints;
+  final bool isUserMessage;
+
+  const EnhancedImageAttachment({
+    super.key,
+    required this.attachmentId,
+    this.isMarkdownFormat = false,
+    this.onTap,
+    this.constraints,
+    this.isUserMessage = false,
+  });
+
+  @override
+  ConsumerState<EnhancedImageAttachment> createState() =>
+      _EnhancedImageAttachmentState();
+}
+
+class _EnhancedImageAttachmentState
+    extends ConsumerState<EnhancedImageAttachment>
+    with AutomaticKeepAliveClientMixin {
+  String? _cachedImageData;
+  bool _isLoading = true;
+  String? _errorMessage;
+
+  @override
+  bool get wantKeepAlive => true;
+
+  @override
+  void initState() {
+    super.initState();
+    _loadImage();
+  }
+
+  Future<void> _loadImage() async {
+    // Check global cache first
+    if (_globalImageCache.containsKey(widget.attachmentId)) {
+      if (mounted) {
+        setState(() {
+          _cachedImageData = _globalImageCache[widget.attachmentId];
+          _isLoading = false;
+        });
+      }
+      return;
+    }
+
+    // Check if this is already a data URL or base64 image
+    if (widget.attachmentId.startsWith('data:') ||
+        widget.attachmentId.startsWith('http')) {
+      _globalImageCache[widget.attachmentId] = widget.attachmentId;
+      if (mounted) {
+        setState(() {
+          _cachedImageData = widget.attachmentId;
+          _isLoading = false;
+        });
+      }
+      return;
+    }
+
+    final api = ref.read(apiServiceProvider);
+    if (api == null) {
+      if (mounted) {
+        setState(() {
+          _errorMessage = 'API service not available';
+          _isLoading = false;
+        });
+      }
+      return;
+    }
+
+    try {
+      // Get file info to check if it's an image
+      final fileInfo = await api.getFileInfo(widget.attachmentId);
+      final fileName = _extractFileName(fileInfo);
+      final ext = fileName.toLowerCase().split('.').last;
+
+      if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'svg'].contains(ext)) {
+        if (mounted) {
+          setState(() {
+            _errorMessage = 'Not an image file: $fileName';
+            _isLoading = false;
+          });
+        }
+        return;
+      }
+
+      // Get the image content
+      final fileContent = await api.getFileContent(widget.attachmentId);
+      
+      // Cache globally
+      _globalImageCache[widget.attachmentId] = fileContent;
+      
+      // Limit cache size
+      if (_globalImageCache.length > 50) {
+        _globalImageCache.remove(_globalImageCache.keys.first);
+      }
+      
+      if (mounted) {
+        setState(() {
+          _cachedImageData = fileContent;
+          _isLoading = false;
+        });
+      }
+    } catch (e) {
+      if (mounted) {
+        setState(() {
+          _errorMessage = 'Failed to load image: ${e.toString()}';
+          _isLoading = false;
+        });
+      }
+    }
+  }
+
+  String _extractFileName(Map<String, dynamic> fileInfo) {
+    return fileInfo['filename'] ??
+        fileInfo['meta']?['name'] ??
+        fileInfo['name'] ??
+        fileInfo['file_name'] ??
+        fileInfo['original_name'] ??
+        fileInfo['original_filename'] ??
+        'unknown';
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    super.build(context); // Required for AutomaticKeepAliveClientMixin
+    
+    if (_isLoading) {
+      return _buildLoadingState();
+    }
+
+    if (_errorMessage != null) {
+      return _buildErrorState();
+    }
+
+    if (_cachedImageData == null) {
+      return const SizedBox.shrink();
+    }
+
+    // Handle different image data formats
+    if (_cachedImageData!.startsWith('http')) {
+      return _buildNetworkImage();
+    } else {
+      return _buildBase64Image();
+    }
+  }
+
+  Widget _buildLoadingState() {
+    return Container(
+      constraints: widget.constraints ??
+          const BoxConstraints(
+            maxWidth: 300,
+            maxHeight: 300,
+            minHeight: 150,
+            minWidth: 200,
+          ),
+      margin: const EdgeInsets.only(bottom: Spacing.xs),
+      decoration: BoxDecoration(
+        color: context.conduitTheme.surfaceBackground.withValues(alpha: 0.5),
+        borderRadius: BorderRadius.circular(AppBorderRadius.md),
+        border: Border.all(
+          color: context.conduitTheme.dividerColor.withValues(alpha: 0.3),
+          width: BorderWidth.thin,
+        ),
+      ),
+      child: Center(
+        child: CircularProgressIndicator(
+          color: context.conduitTheme.buttonPrimary,
+          strokeWidth: 2,
+        ),
+      ),
+    );
+  }
+
+  Widget _buildErrorState() {
+    return Container(
+      constraints: widget.constraints ??
+          const BoxConstraints(
+            maxWidth: 300,
+            maxHeight: 150,
+            minHeight: 100,
+            minWidth: 200,
+          ),
+      margin: const EdgeInsets.only(bottom: Spacing.xs),
+      decoration: BoxDecoration(
+        color: context.conduitTheme.surfaceBackground.withValues(alpha: 0.3),
+        borderRadius: BorderRadius.circular(AppBorderRadius.md),
+        border: Border.all(
+          color: context.conduitTheme.error.withValues(alpha: 0.3),
+          width: BorderWidth.thin,
+        ),
+      ),
+      child: Column(
+        mainAxisAlignment: MainAxisAlignment.center,
+        children: [
+          Icon(
+            Icons.broken_image_outlined,
+            color: context.conduitTheme.error,
+            size: 32,
+          ),
+          const SizedBox(height: Spacing.xs),
+          Padding(
+            padding: const EdgeInsets.symmetric(horizontal: Spacing.sm),
+            child: Text(
+              _errorMessage!,
+              style: TextStyle(
+                color: context.conduitTheme.error,
+                fontSize: AppTypography.bodySmall,
+              ),
+              textAlign: TextAlign.center,
+              maxLines: 2,
+              overflow: TextOverflow.ellipsis,
+            ),
+          ),
+        ],
+      ),
+    );
+  }
+
+  Widget _buildNetworkImage() {
+    final imageWidget = CachedNetworkImage(
+      imageUrl: _cachedImageData!,
+      fit: BoxFit.cover,
+      placeholder: (context, url) => _buildLoadingState(),
+      errorWidget: (context, url, error) {
+        _errorMessage = error.toString();
+        return _buildErrorState();
+      },
+    );
+
+    return _wrapImage(imageWidget);
+  }
+
+  Widget _buildBase64Image() {
+    try {
+      // Extract base64 data from data URL if needed
+      String actualBase64;
+      if (_cachedImageData!.startsWith('data:')) {
+        final commaIndex = _cachedImageData!.indexOf(',');
+        if (commaIndex != -1) {
+          actualBase64 = _cachedImageData!.substring(commaIndex + 1);
+        } else {
+          throw Exception('Invalid data URL format');
+        }
+      } else {
+        actualBase64 = _cachedImageData!;
+      }
+
+      final imageBytes = base64.decode(actualBase64);
+      final imageWidget = Image.memory(
+        imageBytes,
+        fit: BoxFit.cover,
+        errorBuilder: (context, error, stackTrace) {
+          _errorMessage = 'Failed to decode image';
+          return _buildErrorState();
+        },
+      );
+
+      return _wrapImage(imageWidget);
+    } catch (e) {
+      _errorMessage = 'Invalid image format';
+      return _buildErrorState();
+    }
+  }
+
+  Widget _wrapImage(Widget imageWidget) {
+    return Container(
+      constraints: widget.constraints ??
+          const BoxConstraints(
+            maxWidth: 400,
+            maxHeight: 400,
+          ),
+      margin: widget.isMarkdownFormat
+          ? const EdgeInsets.symmetric(vertical: Spacing.sm)
+          : EdgeInsets.zero,
+      child: Material(
+        color: Colors.transparent,
+        child: InkWell(
+          onTap: widget.onTap ?? () => _showFullScreenImage(context),
+          child: Hero(
+            tag: 'image_${widget.attachmentId}',
+            child: imageWidget,
+          ),
+        ),
+      ),
+    );
+  }
+
+  void _showFullScreenImage(BuildContext context) {
+    Navigator.of(context).push(
+      MaterialPageRoute(
+        fullscreenDialog: true,
+        builder: (context) => FullScreenImageViewer(
+          imageData: _cachedImageData!,
+          tag: 'image_${widget.attachmentId}',
+        ),
+      ),
+    );
+  }
+}
+
+class FullScreenImageViewer extends StatelessWidget {
+  final String imageData;
+  final String tag;
+
+  const FullScreenImageViewer({
+    super.key,
+    required this.imageData,
+    required this.tag,
+  });
+
+  @override
+  Widget build(BuildContext context) {
+    Widget imageWidget;
+
+    if (imageData.startsWith('http')) {
+      imageWidget = CachedNetworkImage(
+        imageUrl: imageData,
+        fit: BoxFit.contain,
+        placeholder: (context, url) => Center(
+          child: CircularProgressIndicator(
+            color: context.conduitTheme.buttonPrimary,
+          ),
+        ),
+        errorWidget: (context, url, error) => Center(
+          child: Icon(
+            Icons.error_outline,
+            color: context.conduitTheme.error,
+            size: 48,
+          ),
+        ),
+      );
+    } else {
+      try {
+        String actualBase64;
+        if (imageData.startsWith('data:')) {
+          final commaIndex = imageData.indexOf(',');
+          actualBase64 = imageData.substring(commaIndex + 1);
+        } else {
+          actualBase64 = imageData;
+        }
+        final imageBytes = base64.decode(actualBase64);
+        imageWidget = Image.memory(
+          imageBytes,
+          fit: BoxFit.contain,
+        );
+      } catch (e) {
+        imageWidget = Center(
+          child: Icon(
+            Icons.error_outline,
+            color: context.conduitTheme.error,
+            size: 48,
+          ),
+        );
+      }
+    }
+
+    return Scaffold(
+      backgroundColor: Colors.black,
+      body: Stack(
+        children: [
+          Center(
+            child: Hero(
+              tag: tag,
+              child: InteractiveViewer(
+                minScale: 0.5,
+                maxScale: 5.0,
+                child: imageWidget,
+              ),
+            ),
+          ),
+          Positioned(
+            top: MediaQuery.of(context).padding.top + 16,
+            right: 16,
+            child: IconButton(
+              icon: const Icon(
+                Icons.close,
+                color: Colors.white,
+                size: 28,
+              ),
+              onPressed: () => Navigator.of(context).pop(),
+            ),
+          ),
+        ],
+      ),
+    );
+  }
+}

lib/features/chat/widgets/modern_message_bubble.dart

@@ -1,12 +1,12 @@
-import 'dart:convert';
 import 'package:flutter/material.dart';
 import '../../../shared/theme/theme_extensions.dart';
+import '../../../shared/widgets/markdown/streaming_markdown_widget.dart';
+import 'enhanced_image_attachment.dart';
 
 import 'package:flutter/cupertino.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
 import 'package:flutter_animate/flutter_animate.dart';
 import 'dart:io' show Platform;
-import '../../../core/providers/app_providers.dart';
 
 class ModernMessageBubble extends ConsumerStatefulWidget {
   final dynamic message;
@@ -42,14 +42,6 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
   bool _showActions = false;
   late AnimationController _fadeController;
   late AnimationController _slideController;
-  static const int _maxCachedImages = 24;
-
-  // Cache for image base64 data to prevent repeated API calls
-  final Map<String, String?> _imageCache = {};
-  
-  // Cache for rendered image widgets to prevent rebuilding during streaming
-  final Map<String, Widget> _imageWidgetCache = {};
-  String? _lastAttachmentIds;
 
   @override
   void initState() {
@@ -66,202 +58,121 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
 
 
 
-  Widget _buildAttachmentImages() {
+  Widget _buildUserAttachmentImages() {
     if (widget.message.attachmentIds == null ||
         widget.message.attachmentIds!.isEmpty) {
       return const SizedBox.shrink();
     }
 
-    final currentAttachmentIds = widget.message.attachmentIds!.join('_');
+    final imageCount = widget.message.attachmentIds!.length;
     
-    // Clear cache if attachment IDs changed
+    // iMessage-style image layout
-    if (_lastAttachmentIds != currentAttachmentIds) {
+    if (imageCount == 1) {
-      _imageWidgetCache.clear();
+      // Single image - larger display
-      _lastAttachmentIds = currentAttachmentIds;
+      return Row(
-    }
+        mainAxisAlignment: MainAxisAlignment.end,
-
-    return Column(
-      crossAxisAlignment: CrossAxisAlignment.start,
-      children: widget.message.attachmentIds!.map<Widget>((attachmentId) {
-        // Return cached widget if available
-        if (_imageWidgetCache.containsKey(attachmentId)) {
-          return _imageWidgetCache[attachmentId]!;
-        }
-
-        // Build widget and cache it
-        final imageWidget = _buildSingleImageWidget(attachmentId);
-        _imageWidgetCache[attachmentId] = imageWidget;
-        return imageWidget;
-      }).toList(),
-    );
-  }
-
-  Widget _buildSingleImageWidget(String attachmentId) {
-    return Consumer(
-      builder: (context, ref, child) {
-        final api = ref.read(apiServiceProvider);
-        if (api == null) return const SizedBox.shrink();
-
-        return FutureBuilder<String?>(
-          key: ValueKey('img_$attachmentId'),
-          future: _getCachedImageBase64(api, attachmentId),
-          builder: (context, snapshot) {
-            if (snapshot.connectionState == ConnectionState.waiting) {
-              return Container(
-                height: 150,
-                width: 200,
-                margin: const EdgeInsets.only(bottom: Spacing.xs),
-                decoration: BoxDecoration(
-                  color: context.conduitTheme.surfaceBackground.withValues(
-                    alpha: 0.5,
-                  ),
-                  borderRadius: BorderRadius.circular(AppBorderRadius.sm),
-                ),
-                child: Center(
-                  child: CircularProgressIndicator(
-                    color: context.conduitTheme.buttonPrimary,
-                    strokeWidth: 2,
-                  ),
-                ),
-              );
-            }
-
-            if (snapshot.hasError ||
-                !snapshot.hasData ||
-                snapshot.data == null) {
-              return Container(
-                height: 100,
-                width: 150,
-                margin: const EdgeInsets.only(bottom: Spacing.xs),
-                decoration: BoxDecoration(
-                  color: context.conduitTheme.surfaceBackground.withValues(
-                    alpha: 0.3,
-                  ),
-                  borderRadius: BorderRadius.circular(AppBorderRadius.sm),
-                  border: Border.all(
-                    color: context.conduitTheme.textSecondary.withValues(
-                      alpha: 0.3,
-                    ),
-                  ),
-                ),
-                child: Column(
-                  mainAxisAlignment: MainAxisAlignment.center,
         children: [
-                    Icon(
+          ClipRRect(
-                      Icons.broken_image_outlined,
+            borderRadius: BorderRadius.circular(AppBorderRadius.messageBubble),
-                      color: context.conduitTheme.textSecondary,
+            child: EnhancedImageAttachment(
-                      size: 32,
+              attachmentId: widget.message.attachmentIds![0],
+              isUserMessage: true,
+              constraints: const BoxConstraints(
+                maxWidth: 280,
+                maxHeight: 350,
               ),
-                    const SizedBox(height: Spacing.xs),
-                    Text(
-                      'Image unavailable',
-                      style: TextStyle(
-                        color: context.conduitTheme.textSecondary,
-                        fontSize: AppTypography.bodySmall,
             ),
           ),
         ],
-                ),
       );
-            }
+    } else if (imageCount == 2) {
-
+      // Two images side by side
-            final base64Data = snapshot.data!;
+      return Row(
-            try {
+        mainAxisAlignment: MainAxisAlignment.end,
-              // Handle data URLs (data:image/...;base64,...)
+        mainAxisSize: MainAxisSize.min,
-              String actualBase64;
+        children: [
-              if (base64Data.startsWith('data:')) {
+          Flexible(
-                // Extract base64 part from data URL
+            child: Row(
-                final commaIndex = base64Data.indexOf(',');
+              mainAxisSize: MainAxisSize.min,
-                if (commaIndex != -1) {
+              children: widget.message.attachmentIds!.map((attachmentId) {
-                  actualBase64 = base64Data.substring(commaIndex + 1);
+                return Padding(
-                } else {
+                  padding: EdgeInsets.only(
-                  throw Exception('Invalid data URL format');
+                    left: attachmentId == widget.message.attachmentIds!.first 
-                }
+                        ? 0 
-              } else {
+                        : Spacing.xs,
-                // Direct base64 string
-                actualBase64 = base64Data;
-              }
-
-              final imageBytes = base64.decode(actualBase64);
-              return Container(
-                margin: const EdgeInsets.only(bottom: Spacing.xs),
-                constraints: const BoxConstraints(
-                  maxWidth: 300,
-                  maxHeight: 300,
                   ),
                   child: ClipRRect(
-                  borderRadius: BorderRadius.circular(AppBorderRadius.sm),
+                    borderRadius: BorderRadius.circular(AppBorderRadius.messageBubble),
-                  child: Image.memory(
+                    child: EnhancedImageAttachment(
-                    imageBytes,
+                      attachmentId: attachmentId,
-                    fit: BoxFit.cover,
+                      isUserMessage: true,
-                    errorBuilder: (context, error, stackTrace) {
+                      constraints: const BoxConstraints(
-                      return Container(
+                        maxWidth: 135,
-                        height: 100,
+                        maxHeight: 180,
-                        width: 150,
-                        decoration: BoxDecoration(
-                          color: context.conduitTheme.surfaceBackground
-                              .withValues(alpha: 0.3),
-                          borderRadius: BorderRadius.circular(
-                            AppBorderRadius.sm,
                       ),
                     ),
-                        child: Column(
-                          mainAxisAlignment: MainAxisAlignment.center,
-                          children: [
-                            Icon(
-                              Icons.error_outline,
-                              color: context.conduitTheme.error,
-                              size: 32,
-                            ),
-                            const SizedBox(height: Spacing.xs),
-                            Text(
-                              'Failed to load image',
-                              style: TextStyle(
-                                color: context.conduitTheme.error,
-                                fontSize: AppTypography.bodySmall,
-                              ),
-                            ),
-                          ],
-                        ),
-                      );
-                    },
-                  ),
                   ),
                 );
-            } catch (e) {
+              }).toList(),
-              return Container(
-                height: 100,
-                width: 150,
-                margin: const EdgeInsets.only(bottom: Spacing.xs),
-                decoration: BoxDecoration(
-                  color: context.conduitTheme.surfaceBackground.withValues(
-                    alpha: 0.3,
-                  ),
-                  borderRadius: BorderRadius.circular(AppBorderRadius.sm),
-                ),
-                child: Column(
-                  mainAxisAlignment: MainAxisAlignment.center,
-                  children: [
-                    Icon(
-                      Icons.error_outline,
-                      color: context.conduitTheme.error,
-                      size: 32,
-                    ),
-                    const SizedBox(height: Spacing.xs),
-                    Text(
-                      'Invalid image format',
-                      style: TextStyle(
-                        color: context.conduitTheme.error,
-                        fontSize: AppTypography.bodySmall,
             ),
           ),
         ],
+      );
+    } else {
+      // Grid layout for 3+ images
+      return Row(
+        mainAxisAlignment: MainAxisAlignment.end,
+        children: [
+          Flexible(
+            child: Container(
+              constraints: const BoxConstraints(maxWidth: 280),
+              child: Wrap(
+                alignment: WrapAlignment.end,
+                spacing: Spacing.xs,
+                runSpacing: Spacing.xs,
+                children: widget.message.attachmentIds!.map((attachmentId) {
+                  return ClipRRect(
+                    borderRadius: BorderRadius.circular(AppBorderRadius.md),
+                    child: EnhancedImageAttachment(
+                      attachmentId: attachmentId,
+                      isUserMessage: true,
+                      constraints: BoxConstraints(
+                        maxWidth: imageCount == 3 ? 135 : 90,
+                        maxHeight: imageCount == 3 ? 135 : 90,
                       ),
+                    ),
+                  );
+                }).toList(),
+              ),
+            ),
+          ),
+        ],
       );
     }
-          },
+  }
+
+  Widget _buildAssistantAttachmentImages() {
+    if (widget.message.attachmentIds == null ||
+        widget.message.attachmentIds!.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    // Assistant images - similar style but left-aligned
+    return Wrap(
+      spacing: Spacing.sm,
+      runSpacing: Spacing.sm,
+      children: widget.message.attachmentIds!.map<Widget>((attachmentId) {
+        return ClipRRect(
+          borderRadius: BorderRadius.circular(AppBorderRadius.messageBubble),
+          child: EnhancedImageAttachment(
+            attachmentId: attachmentId,
+            constraints: const BoxConstraints(
+              maxWidth: 300,
+              maxHeight: 350,
+            ),
+          ),
         );
-      },
+      }).toList(),
     );
   }
 
@@ -296,26 +207,39 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
   }
 
   Widget _buildUserMessage() {
-    return Container(
+    final hasImages = widget.message.attachmentIds != null &&
+        widget.message.attachmentIds!.isNotEmpty;
+    final hasText = widget.message.content.isNotEmpty;
+
+    return GestureDetector(
+      onLongPress: () => _toggleActions(),
+      behavior: HitTestBehavior.translucent,
+      child: Container(
         width: double.infinity,
         margin: const EdgeInsets.only(
-            bottom: Spacing.messagePadding,
+          bottom: Spacing.sm,
           left: Spacing.xxxl,
           right: Spacing.xs,
         ),
         child: Column(
           crossAxisAlignment: CrossAxisAlignment.end,
           children: [
+            // Display images outside and above the text bubble (iMessage style)
+            if (hasImages) ...[
+              _buildUserAttachmentImages(),
+              if (hasText) const SizedBox(height: Spacing.xs),
+            ],
+
+            // Display text bubble if there's text content
+            if (hasText)
               Row(
                 mainAxisAlignment: MainAxisAlignment.end,
                 children: [
                   Flexible(
-                    child: GestureDetector(
-                      onLongPress: () => _toggleActions(),
                       child: Container(
                         padding: const EdgeInsets.symmetric(
                           horizontal: Spacing.messagePadding,
-                          vertical: Spacing.sm,
+                          vertical: Spacing.xs,
                         ),
                       decoration: BoxDecoration(
                         gradient: LinearGradient(
@@ -335,41 +259,31 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
                           color: context.conduitTheme.chatBubbleUserBorder,
                           width: BorderWidth.regular,
                         ),
-                          boxShadow: ConduitShadows.high,
+                        boxShadow: [
+                          BoxShadow(
+                            color: Colors.black.withValues(alpha: 0.08),
+                            blurRadius: 4,
+                            offset: const Offset(0, 2),
                           ),
-                        child: Column(
+                        ],
-                          crossAxisAlignment: CrossAxisAlignment.start,
+                      ),
-                          children: [
+                      child: _buildCustomText(
-                            // Display images if any
-                            if (widget.message.attachmentIds != null &&
-                                widget.message.attachmentIds!.isNotEmpty)
-                              _buildAttachmentImages(),
-
-                            // Display text content if any
-                            if (widget.message.content.isNotEmpty) ...[
-                              if (widget.message.attachmentIds != null &&
-                                  widget.message.attachmentIds!.isNotEmpty)
-                                const SizedBox(height: Spacing.sm),
-                              _buildCustomText(
                         widget.message.content,
                         context.conduitTheme.chatBubbleUserText,
                       ),
-                            ],
-                          ],
-                        ),
-                      ),
                     ),
                   ),
                 ],
               ),
             
-              // Action buttons below the message bubble
+            // Action buttons below the message
             if (_showActions) ...[
               const SizedBox(height: Spacing.sm),
               _buildUserActionButtons(),
             ],
           ],
         ),
+      ),
     )
         .animate()
         .fadeIn(duration: AnimationDuration.messageAppear)
@@ -382,10 +296,21 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
   }
 
   Widget _buildAssistantMessage() {
-    return Container(
+    final hasImages = widget.message.attachmentIds != null &&
+        widget.message.attachmentIds!.isNotEmpty;
+    final hasContent = widget.message.content.isNotEmpty &&
+        widget.message.content != '[TYPING_INDICATOR]';
+    final showTyping = (widget.message.content.isEmpty ||
+            widget.message.content == '[TYPING_INDICATOR]') &&
+        widget.isStreaming;
+
+    return GestureDetector(
+      onLongPress: () => _toggleActions(),
+      behavior: HitTestBehavior.translucent,
+      child: Container(
         width: double.infinity,
         margin: const EdgeInsets.only(
-            bottom: Spacing.lg,
+          bottom: Spacing.md,
           left: Spacing.xs,
           right: Spacing.xxxl,
         ),
@@ -394,7 +319,7 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
           children: [
             // Simplified AI Name and Avatar
             Padding(
-                padding: const EdgeInsets.only(bottom: 12),
+              padding: const EdgeInsets.only(bottom: 8),
               child: Row(
                 children: [
                   Container(
@@ -433,11 +358,19 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
               ),
             ),
 
-              // Message Content
+            // Display images outside the bubble if any
-              GestureDetector(
+            if (hasImages) ...[
-                onLongPress: () => _toggleActions(),
+              _buildAssistantAttachmentImages(),
-                child: Container(
+              if (hasContent || showTyping) const SizedBox(height: Spacing.xs),
-                  padding: const EdgeInsets.all(Spacing.messagePadding),
+            ],
+
+            // Message Content Bubble
+            if (hasContent || showTyping)
+              Container(
+                padding: EdgeInsets.symmetric(
+                  horizontal: Spacing.messagePadding,
+                  vertical: Spacing.xs,
+                ),
                 decoration: BoxDecoration(
                   color: context.conduitTheme.chatBubbleAssistant,
                   borderRadius: BorderRadius.circular(
@@ -447,27 +380,19 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
                     color: context.conduitTheme.chatBubbleAssistantBorder,
                     width: BorderWidth.regular,
                   ),
-                    boxShadow: ConduitShadows.low,
+                  boxShadow: [
+                    BoxShadow(
+                      color: Colors.black.withValues(alpha: 0.05),
+                      blurRadius: 3,
+                      offset: const Offset(0, 1),
                     ),
-                  child: Column(
-                    crossAxisAlignment: CrossAxisAlignment.start,
-                    children: [
-                      // Check for typing indicator - show for empty content OR explicit typing indicator during streaming
-                      if ((widget.message.content.isEmpty ||
-                              widget.message.content == '[TYPING_INDICATOR]') &&
-                          widget.isStreaming) ...[
-                        _buildTypingIndicator(),
-                      ] else if (widget.message.content.isNotEmpty &&
-                          widget.message.content != '[TYPING_INDICATOR]') ...[
-                        _buildCustomText(
-                          widget.message.content,
-                          context.conduitTheme.chatBubbleAssistantText,
-                        ),
-                      ] else
-                        // Fallback: show empty state for non-streaming empty messages
-                        const SizedBox.shrink(),
                   ],
                 ),
+                child: showTyping
+                    ? _buildTypingIndicator()
+                    : _buildCustomText(
+                        widget.message.content,
+                        context.conduitTheme.chatBubbleAssistantText,
                       ),
               ),
 
@@ -478,6 +403,7 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
             ],
           ],
         ),
+      ),
     )
         .animate()
         .fadeIn(duration: AnimationDuration.messageAppear)
@@ -491,216 +417,19 @@ class _ModernMessageBubbleState extends ConsumerState<ModernMessageBubble>
 
 
 
-  Future<String?> _getCachedImageBase64(dynamic api, String fileId) async {
-    // Check cache first to prevent repeated API calls
-    if (_imageCache.containsKey(fileId)) {
-      return _imageCache[fileId];
-    }
 
-    // If not in cache, get the image and cache it
-    final result = await _getImageBase64(api, fileId);
-    // Simple LRU-like eviction to bound memory
-    if (_imageCache.length >= _maxCachedImages) {
-      _imageCache.remove(_imageCache.keys.first);
-    }
-    _imageCache[fileId] = result;
-    return result;
-  }
-
-  Future<String?> _getImageBase64(dynamic api, String fileId) async {
-    try {
-      // Check if this is already a data URL (for images)
-      if (fileId.startsWith('data:')) {
-        return fileId;
-      }
-
-      // First, get file info to determine if it's an image
-      final fileInfo = await api.getFileInfo(fileId);
-      final fileName =
-          fileInfo['filename'] ??
-          fileInfo['meta']?['name'] ??
-          fileInfo['name'] ??
-          fileInfo['file_name'] ??
-          fileInfo['original_name'] ??
-          fileInfo['original_filename'] ??
-          '';
-      final ext = fileName.toLowerCase().split('.').last;
-
-      // Only process image files
-      if (!['jpg', 'jpeg', 'png', 'gif', 'webp'].contains(ext)) {
-        debugPrint('DEBUG: Skipping non-image file: $fileName');
-        return null;
-      }
-
-      // Get file content as base64 string
-      final fileContent = await api.getFileContent(fileId);
-      return fileContent;
-    } catch (e) {
-      debugPrint('DEBUG: Error getting image content for $fileId: $e');
-      return null;
-    }
-  }
 
   Widget _buildCustomText(String text, [Color? textColor]) {
-    // Simple markdown-like parsing for efficiency
+    // Use the new markdown widget for rich text rendering
-    final lines = text.split('\n');
+    return StreamingMarkdownWidget(
-    final widgets = <Widget>[];
+      staticContent: text,
-
+      isStreaming: widget.isStreaming,
-    for (int i = 0; i < lines.length; i++) {
-      final line = lines[i];
-      if (line.trim().isEmpty) {
-        if (i < lines.length - 1) {
-          widgets.add(const SizedBox(height: Spacing.sm));
-        }
-        continue;
-      }
-
-      // Parse basic markdown
-      Widget textWidget = _parseMarkdownLine(line, textColor);
-
-      if (i < lines.length - 1) {
-        widgets.add(textWidget);
-        widgets.add(const SizedBox(height: Spacing.xs));
-      } else {
-        widgets.add(textWidget);
-      }
-    }
-
-    return Column(
-      crossAxisAlignment: CrossAxisAlignment.start,
-      children: widgets,
     );
   }
 
-  Widget _parseMarkdownLine(String line, [Color? textColor]) {
-    // Handle code blocks
-    if (line.startsWith('```')) {
-      return Container(
-            margin: const EdgeInsets.symmetric(vertical: Spacing.xs),
-            padding: const EdgeInsets.all(Spacing.md),
-            decoration: BoxDecoration(
-              color: context.conduitTheme.surfaceBackground.withValues(
-                alpha: Alpha.badgeBackground,
-              ),
-              borderRadius: BorderRadius.circular(AppBorderRadius.sm),
-              border: Border.all(
-                color: context.conduitTheme.textPrimary.withValues(
-                  alpha: Alpha.subtle,
-                ),
-                width: BorderWidth.regular,
-              ),
-            ),
-            child: Text(
-              line.substring(3),
-              style: AppTypography.chatCodeStyle.copyWith(
-                color: textColor ?? context.conduitTheme.textSecondary,
-              ),
-            ),
-          )
-          .animate()
-          .fadeIn(duration: AnimationDuration.microInteraction)
-          .slideX(
-            begin: 0.1,
-            end: 0,
-            duration: AnimationDuration.microInteraction,
-          );
-    }
 
-    // Handle headers
-    if (line.startsWith('#')) {
-      int level = 0;
-      while (level < line.length && line[level] == '#') {
-        level++;
-      }
-      final fontSize = AppTypography.headlineMedium - (level * 2);
-      return Text(
-            line.substring(level).trim(),
-            style: AppTypography.headlineMediumStyle.copyWith(
-              color: textColor ?? context.conduitTheme.textPrimary,
-              fontSize: fontSize.toDouble(),
-            ),
-          )
-          .animate()
-          .fadeIn(duration: AnimationDuration.microInteraction)
-          .slideX(
-            begin: 0.1,
-            end: 0,
-            duration: AnimationDuration.microInteraction,
-          );
-    }
 
-    // Handle inline code
-    if (line.contains('`')) {
-      final parts = line.split('`');
-      final widgets = <Widget>[];
 
-      for (int i = 0; i < parts.length; i++) {
-        if (parts[i].isNotEmpty) {
-          if (i % 2 == 1) {
-            // Inline code
-            widgets.add(
-              Container(
-                padding: const EdgeInsets.symmetric(
-                  horizontal: Spacing.xs + Spacing.xxs,
-                  vertical: Spacing.xxs,
-                ),
-                decoration: BoxDecoration(
-                  color: context.conduitTheme.textPrimary.withValues(
-                    alpha: Alpha.badgeBackground,
-                  ),
-                  borderRadius: BorderRadius.circular(AppBorderRadius.xs),
-                ),
-                child: Text(
-                  parts[i],
-                  style: AppTypography.chatCodeStyle.copyWith(
-                    color: textColor ?? context.conduitTheme.textSecondary,
-                  ),
-                ),
-              ),
-            );
-          } else {
-            // Regular text
-            widgets.add(
-              Text(
-                parts[i],
-                style: AppTypography.chatMessageStyle.copyWith(
-                  color: textColor ?? context.conduitTheme.textPrimary,
-                ),
-              ),
-            );
-          }
-        }
-      }
-
-      return Wrap(
-            crossAxisAlignment: WrapCrossAlignment.start,
-            children: widgets,
-          )
-          .animate()
-          .fadeIn(duration: AnimationDuration.microInteraction)
-          .slideX(
-            begin: 0.1,
-            end: 0,
-            duration: AnimationDuration.microInteraction,
-          );
-    }
-
-    // Regular text
-    return Text(
-          line,
-          style: AppTypography.chatMessageStyle.copyWith(
-            color: textColor ?? context.conduitTheme.textPrimary,
-            letterSpacing: 0.1,
-          ),
-        )
-        .animate()
-        .fadeIn(duration: AnimationDuration.microInteraction)
-        .slideX(
-          begin: 0.1,
-          end: 0,
-          duration: AnimationDuration.microInteraction,
-        );
-  }
 
   Widget _buildTypingIndicator() {
     return Consumer(

lib/shared/widgets/markdown/markdown_config.dart

@@ -1,3 +1,4 @@
+import 'dart:convert';
 import 'package:flutter/material.dart';
 import 'package:markdown_widget/markdown_widget.dart';
 import 'package:flutter_highlight/themes/atom-one-dark.dart';
@@ -41,7 +42,7 @@ class ConduitMarkdownConfig {
         // Link config
         LinkConfig(
           style: TextStyle(
-            color: AppTheme.brandPrimary,
+            color: theme.buttonPrimary,
             decoration: TextDecoration.underline,
           ),
           onTap: (url) async {
@@ -51,30 +52,60 @@ class ConduitMarkdownConfig {
           },
         ),
 
-        // Image config - optimized for mobile
+        // Image config - optimized for mobile with support for base64 and network images
         ImgConfig(
-          builder: (url, attributes) => CachedNetworkImage(
+          builder: (url, attributes) {
+            // Check if it's a base64 data URL
+            if (url.startsWith('data:')) {
+              return _buildBase64Image(url, theme);
+            }
+            // Network image
+            return CachedNetworkImage(
               imageUrl: url,
               placeholder: (context, url) => Container(
                 height: 200,
-              color: theme.surfaceBackground,
+                decoration: BoxDecoration(
+                  color: theme.surfaceBackground.withValues(alpha: 0.5),
+                  borderRadius: BorderRadius.circular(AppBorderRadius.md),
+                ),
                 child: Center(
                   child: CircularProgressIndicator(
-                  color: AppTheme.brandPrimary,
+                    color: theme.loadingIndicator,
+                    strokeWidth: 2,
                   ),
                 ),
               ),
               errorWidget: (context, url, error) => Container(
                 height: 100,
-              color: theme.surfaceBackground,
+                decoration: BoxDecoration(
-              child: Center(
+                  color: theme.surfaceBackground.withValues(alpha: 0.3),
-                child: Icon(
+                  borderRadius: BorderRadius.circular(AppBorderRadius.md),
-                  Icons.broken_image,
+                  border: Border.all(
-                  color: theme.iconSecondary,
+                    color: theme.error.withValues(alpha: 0.3),
+                    width: BorderWidth.thin,
                   ),
                 ),
+                child: Column(
+                  mainAxisAlignment: MainAxisAlignment.center,
+                  children: [
+                    Icon(
+                      Icons.broken_image_outlined,
+                      color: theme.error,
+                      size: 32,
+                    ),
+                    const SizedBox(height: Spacing.xs),
+                    Text(
+                      'Failed to load image',
+                      style: TextStyle(
+                        color: theme.error,
+                        fontSize: 12,
                       ),
                     ),
+                  ],
+                ),
+              ),
+            );
+          },
         ),
 
         // Table config - mobile responsive
@@ -89,6 +120,7 @@ class ConduitMarkdownConfig {
         PConfig(
           textStyle: AppTypography.chatMessageStyle.copyWith(
             color: theme.textPrimary,
+            height: 1.3,
           ),
         ),
 
@@ -122,6 +154,82 @@ class ConduitMarkdownConfig {
       ],
     );
   }
+
+  static Widget _buildBase64Image(String dataUrl, ConduitThemeExtension theme) {
+    try {
+      // Extract base64 part from data URL
+      final commaIndex = dataUrl.indexOf(',');
+      if (commaIndex == -1) {
+        throw Exception('Invalid data URL format');
+      }
+      
+      final base64String = dataUrl.substring(commaIndex + 1);
+      final imageBytes = base64.decode(base64String);
+      
+      return Container(
+        margin: const EdgeInsets.symmetric(vertical: Spacing.sm),
+        constraints: const BoxConstraints(
+          maxWidth: 500,
+          maxHeight: 500,
+        ),
+        child: ClipRRect(
+          borderRadius: BorderRadius.circular(AppBorderRadius.md),
+          child: Image.memory(
+            imageBytes,
+            fit: BoxFit.contain,
+            errorBuilder: (context, error, stackTrace) {
+              return Container(
+                height: 100,
+                decoration: BoxDecoration(
+                  color: theme.surfaceBackground.withValues(alpha: 0.3),
+                  borderRadius: BorderRadius.circular(AppBorderRadius.md),
+                  border: Border.all(
+                    color: theme.error.withValues(alpha: 0.3),
+                    width: BorderWidth.thin,
+                  ),
+                ),
+                child: Column(
+                  mainAxisAlignment: MainAxisAlignment.center,
+                  children: [
+                    Icon(
+                      Icons.error_outline,
+                      color: theme.error,
+                      size: 32,
+                    ),
+                    const SizedBox(height: Spacing.xs),
+                    Text(
+                      'Invalid image data',
+                      style: TextStyle(
+                        color: theme.error,
+                        fontSize: 12,
+                      ),
+                    ),
+                  ],
+                ),
+              );
+            },
+          ),
+        ),
+      );
+    } catch (e) {
+      return Container(
+        height: 100,
+        decoration: BoxDecoration(
+          color: theme.surfaceBackground.withValues(alpha: 0.3),
+          borderRadius: BorderRadius.circular(AppBorderRadius.md),
+        ),
+        child: Center(
+          child: Text(
+            'Invalid image format',
+            style: TextStyle(
+              color: theme.error,
+              fontSize: 12,
+            ),
+          ),
+        ),
+      );
+    }
+  }
 }
 
 /// Custom wrapper for code blocks with copy functionality
@@ -148,7 +256,7 @@ class CodeBlockWrapper extends StatelessWidget {
           top: 8,
           right: 8,
           child: Material(
-            color: Colors.transparent,
+            color: theme.surfaceBackground.withValues(alpha: 0.0),
             child: InkWell(
               borderRadius: BorderRadius.circular(AppBorderRadius.sm),
               onTap: () {

lib/shared/widgets/markdown/streaming_markdown_widget.dart

@@ -122,13 +122,10 @@ class _StreamingMarkdownWidgetState extends State<StreamingMarkdownWidget> {
 
     if (widget.isStreaming && _renderedContent.isNotEmpty) {
       // Use MarkdownBlock for streaming - it's optimized for live updates
-      return Container(
+      return MarkdownBlock(
-        padding: widget.padding,
-        child: MarkdownBlock(
         data: _renderedContent,
         config: config,
         selectable: true,
-        ),
       );
     } else {
       // Use MarkdownWidget for completed messages
@@ -139,7 +136,7 @@ class _StreamingMarkdownWidgetState extends State<StreamingMarkdownWidget> {
         selectable: true,
         shrinkWrap: true,
         physics: const NeverScrollableScrollPhysics(),
-        padding: widget.padding,
+        padding: EdgeInsets.zero,
       );
     }
   }