refactor(chat): Replace StatusHistoryTimeline with StreamingStatusWidget

lib/shared/widgets/markdown/citation_badge.dart

@@ -5,8 +5,8 @@ import '../../../core/models/chat_message.dart';
 import '../../theme/theme_extensions.dart';
 
 /// Helper utilities for working with source references.
-class _SourceHelper {
-  const _SourceHelper._();
+class SourceHelper {
+  const SourceHelper._();
 
   /// Extracts a URL from a source reference, checking multiple fields.
   static String? getSourceUrl(ChatSourceReference source) {
@@ -27,22 +27,39 @@ class _SourceHelper {
   }
 
   /// Gets a display title for a source.
+  ///
+  /// For web sources (with URLs), shows the domain name like "wikipedia.org".
+  /// This matches OpenWebUI's behavior where web search results show domains.
   static String getSourceTitle(ChatSourceReference source, int index) {
-    if (source.title != null && source.title!.isNotEmpty) {
-      return source.title!;
-    }
+    // For web sources, prefer showing the URL domain
     final url = getSourceUrl(source);
     if (url != null) {
-      return _extractDomain(url);
+      return extractDomain(url);
     }
+
+    // If title is a URL, extract domain
+    if (source.title != null && source.title!.isNotEmpty) {
+      final title = source.title!;
+      if (title.startsWith('http')) {
+        return extractDomain(title);
+      }
+      return title;
+    }
+
+    // Check if ID is a URL
     if (source.id != null && source.id!.isNotEmpty) {
-      return source.id!;
+      final id = source.id!;
+      if (id.startsWith('http')) {
+        return extractDomain(id);
+      }
+      return id;
     }
+
     return 'Source ${index + 1}';
   }
 
   /// Extracts the domain from a URL for display.
-  static String _extractDomain(String url) {
+  static String extractDomain(String url) {
     try {
       final uri = Uri.parse(url);
       String domain = uri.host;
@@ -55,6 +72,16 @@ class _SourceHelper {
     }
   }
 
+  /// Formats a title for display, truncating if needed.
+  /// Matches OpenWebUI's getDisplayTitle behavior.
+  static String formatDisplayTitle(String title) {
+    if (title.isEmpty) return 'N/A';
+    if (title.length > 25) {
+      return '${title.substring(0, 12)}…${title.substring(title.length - 8)}';
+    }
+    return title;
+  }
+
   /// Launches a URL in an external browser.
   static Future<void> launchSourceUrl(String url) async {
     try {
@@ -68,9 +95,9 @@ class _SourceHelper {
   }
 }
 
-/// A compact badge showing a citation reference that links to a source.
+/// A compact inline citation badge showing source domain/title.
 ///
-/// Mirrors OpenWebUI's Source.svelte and SourceToken.svelte behavior.
+/// Uses the app's design system for consistency with other chips and badges.
 class CitationBadge extends StatelessWidget {
   const CitationBadge({
     super.key,
@@ -94,65 +121,63 @@ class CitationBadge extends StatelessWidget {
 
     // Check if index is valid
     if (sourceIndex < 0 || sourceIndex >= sources.length) {
-      // Invalid source index - show placeholder
-      return _buildBadge(
-        theme: theme,
-        displayNumber: sourceIndex + 1,
-        isValid: false,
-      );
+      return const SizedBox.shrink();
     }
 
     final source = sources[sourceIndex];
-    final url = _SourceHelper.getSourceUrl(source);
-    final title = _SourceHelper.getSourceTitle(source, sourceIndex);
+    final url = SourceHelper.getSourceUrl(source);
+    final title = SourceHelper.getSourceTitle(source, sourceIndex);
+    final displayTitle = SourceHelper.formatDisplayTitle(title);
 
     return Tooltip(
       message: title,
       preferBelow: false,
-      child: _buildBadge(
-        theme: theme,
-        displayNumber: sourceIndex + 1,
-        isValid: true,
-        onTap: () {
-          if (onTap != null) {
-            onTap!();
-          } else if (url != null) {
-            _SourceHelper.launchSourceUrl(url);
-          }
-        },
-      ),
-    );
-  }
-
-  Widget _buildBadge({
-    required ConduitThemeExtension theme,
-    required int displayNumber,
-    required bool isValid,
-    VoidCallback? onTap,
-  }) {
-    return GestureDetector(
-      onTap: onTap,
-      child: Container(
-        padding: const EdgeInsets.symmetric(horizontal: 6, vertical: 1),
-        margin: const EdgeInsets.symmetric(horizontal: 1),
-        decoration: BoxDecoration(
-          color: isValid
-              ? theme.surfaceContainer.withValues(alpha: 0.6)
-              : theme.surfaceContainer.withValues(alpha: 0.3),
-          borderRadius: BorderRadius.circular(10),
-          border: Border.all(
-            color: theme.dividerColor.withValues(alpha: 0.3),
-            width: 0.5,
-          ),
-        ),
-        child: Text(
-          displayNumber.toString(),
-          style: TextStyle(
-            fontSize: 10,
-            fontWeight: FontWeight.w600,
-            color: isValid
-                ? theme.textPrimary.withValues(alpha: 0.8)
-                : theme.textSecondary.withValues(alpha: 0.5),
+      child: Material(
+        color: Colors.transparent,
+        child: InkWell(
+          onTap: () {
+            if (onTap != null) {
+              onTap!();
+            } else if (url != null) {
+              SourceHelper.launchSourceUrl(url);
+            }
+          },
+          borderRadius: BorderRadius.circular(AppBorderRadius.chip),
+          child: Container(
+            padding: const EdgeInsets.symmetric(
+              horizontal: Spacing.sm,
+              vertical: Spacing.xxs,
+            ),
+            margin: const EdgeInsets.symmetric(horizontal: 2),
+            decoration: BoxDecoration(
+              color: theme.surfaceContainer.withValues(alpha: 0.6),
+              borderRadius: BorderRadius.circular(AppBorderRadius.chip),
+              border: Border.all(
+                color: theme.cardBorder.withValues(alpha: 0.5),
+                width: BorderWidth.thin,
+              ),
+            ),
+            child: Row(
+              mainAxisSize: MainAxisSize.min,
+              children: [
+                Icon(
+                  Icons.link_rounded,
+                  size: 10,
+                  color: theme.textSecondary.withValues(alpha: 0.7),
+                ),
+                const SizedBox(width: Spacing.xxs),
+                Text(
+                  displayTitle,
+                  style: TextStyle(
+                    fontSize: AppTypography.labelSmall,
+                    fontWeight: FontWeight.w500,
+                    color: theme.textSecondary,
+                  ),
+                  maxLines: 1,
+                  overflow: TextOverflow.ellipsis,
+                ),
+              ],
+            ),
           ),
         ),
       ),
@@ -161,6 +186,8 @@ class CitationBadge extends StatelessWidget {
 }
 
 /// A grouped citation badge for multiple sources like [1,2,3].
+///
+/// Shows first source with +N indicator for additional sources.
 class CitationBadgeGroup extends StatelessWidget {
   const CitationBadgeGroup({
     super.key,
@@ -195,105 +222,131 @@ class CitationBadgeGroup extends StatelessWidget {
       );
     }
 
-    // For multiple citations, show grouped badge
     final theme = context.conduitTheme;
-    final validCount = sourceIndices
-        .where((i) => i >= 0 && i < sources.length)
-        .length;
+
+    // Get first valid source for display
+    final firstIndex = sourceIndices.first;
+    final isFirstValid = firstIndex >= 0 && firstIndex < sources.length;
+
+    if (!isFirstValid) {
+      return const SizedBox.shrink();
+    }
+
+    final firstSource = sources[firstIndex];
+    final firstTitle = SourceHelper.getSourceTitle(firstSource, firstIndex);
+    final displayTitle = SourceHelper.formatDisplayTitle(firstTitle);
+    final additionalCount = sourceIndices.length - 1;
 
     return PopupMenuButton<int>(
       tooltip: 'View sources',
       padding: EdgeInsets.zero,
       constraints: const BoxConstraints(),
       position: PopupMenuPosition.under,
-      shape: RoundedRectangleBorder(borderRadius: BorderRadius.circular(8)),
+      shape: RoundedRectangleBorder(
+        borderRadius: BorderRadius.circular(AppBorderRadius.md),
+      ),
       color: theme.surfaceBackground,
+      surfaceTintColor: Colors.transparent,
+      elevation: Elevation.medium,
       itemBuilder: (context) {
-        return sourceIndices.map((index) {
-          final isValid = index >= 0 && index < sources.length;
-          final title = isValid
-              ? _SourceHelper.getSourceTitle(sources[index], index)
-              : 'Invalid source';
+        return sourceIndices
+            .map((index) {
+              final isValid = index >= 0 && index < sources.length;
+              if (!isValid) return null;
 
-          return PopupMenuItem<int>(
-            value: index,
-            height: 36,
-            child: Row(
-              mainAxisSize: MainAxisSize.min,
-              children: [
-                Container(
-                  width: 20,
-                  height: 20,
-                  decoration: BoxDecoration(
-                    color: theme.surfaceContainer,
-                    borderRadius: BorderRadius.circular(6),
-                  ),
-                  child: Center(
-                    child: Text(
-                      (index + 1).toString(),
-                      style: TextStyle(
-                        fontSize: 11,
-                        fontWeight: FontWeight.w600,
-                        color: theme.textPrimary,
+              final source = sources[index];
+              final title = SourceHelper.getSourceTitle(source, index);
+
+              return PopupMenuItem<int>(
+                value: index,
+                height: 40,
+                child: Row(
+                  mainAxisSize: MainAxisSize.min,
+                  children: [
+                    Icon(
+                      Icons.link_rounded,
+                      size: 14,
+                      color: theme.textSecondary.withValues(alpha: 0.7),
+                    ),
+                    const SizedBox(width: Spacing.sm),
+                    Flexible(
+                      child: Text(
+                        SourceHelper.formatDisplayTitle(title),
+                        style: TextStyle(
+                          fontSize: AppTypography.bodySmall,
+                          fontWeight: FontWeight.w500,
+                          color: theme.textPrimary,
+                        ),
+                        maxLines: 1,
+                        overflow: TextOverflow.ellipsis,
                       ),
                     ),
-                  ),
+                  ],
                 ),
-                const SizedBox(width: 8),
-                Flexible(
-                  child: Text(
-                    title,
-                    style: TextStyle(fontSize: 13, color: theme.textSecondary),
-                    maxLines: 1,
-                    overflow: TextOverflow.ellipsis,
-                  ),
-                ),
-              ],
-            ),
-          );
-        }).toList();
+              );
+            })
+            .whereType<PopupMenuItem<int>>()
+            .toList();
       },
       onSelected: (index) {
         if (onSourceTap != null) {
           onSourceTap!(index);
         } else if (index >= 0 && index < sources.length) {
-          final url = _SourceHelper.getSourceUrl(sources[index]);
+          final url = SourceHelper.getSourceUrl(sources[index]);
           if (url != null) {
-            _SourceHelper.launchSourceUrl(url);
+            SourceHelper.launchSourceUrl(url);
           }
         }
       },
       child: Container(
-        padding: const EdgeInsets.symmetric(horizontal: 6, vertical: 1),
-        margin: const EdgeInsets.symmetric(horizontal: 1),
+        padding: const EdgeInsets.symmetric(
+          horizontal: Spacing.sm,
+          vertical: Spacing.xxs,
+        ),
+        margin: const EdgeInsets.symmetric(horizontal: 2),
         decoration: BoxDecoration(
           color: theme.surfaceContainer.withValues(alpha: 0.6),
-          borderRadius: BorderRadius.circular(10),
+          borderRadius: BorderRadius.circular(AppBorderRadius.chip),
           border: Border.all(
-            color: theme.dividerColor.withValues(alpha: 0.3),
-            width: 0.5,
+            color: theme.cardBorder.withValues(alpha: 0.5),
+            width: BorderWidth.thin,
           ),
         ),
         child: Row(
           mainAxisSize: MainAxisSize.min,
           children: [
-            Text(
-              (sourceIndices.first + 1).toString(),
-              style: TextStyle(
-                fontSize: 10,
-                fontWeight: FontWeight.w600,
-                color: theme.textPrimary.withValues(alpha: 0.8),
-              ),
+            Icon(
+              Icons.link_rounded,
+              size: 10,
+              color: theme.textSecondary.withValues(alpha: 0.7),
             ),
-            if (validCount > 1) ...[
-              Text(
-                '+${validCount - 1}',
+            const SizedBox(width: Spacing.xxs),
+            Text(
+              displayTitle,
+              style: TextStyle(
+                fontSize: AppTypography.labelSmall,
+                fontWeight: FontWeight.w500,
+                color: theme.textSecondary,
+              ),
+              maxLines: 1,
+              overflow: TextOverflow.ellipsis,
+            ),
+            const SizedBox(width: Spacing.xxs),
+            Container(
+              padding: const EdgeInsets.symmetric(horizontal: 4, vertical: 1),
+              decoration: BoxDecoration(
+                color: theme.buttonPrimary.withValues(alpha: 0.15),
+                borderRadius: BorderRadius.circular(AppBorderRadius.small),
+              ),
+              child: Text(
+                '+$additionalCount',
                 style: TextStyle(
                   fontSize: 9,
-                  color: theme.textSecondary.withValues(alpha: 0.7),
+                  fontWeight: FontWeight.w600,
+                  color: theme.buttonPrimary,
                 ),
               ),
-            ],
+            ),
           ],
         ),
       ),

lib/shared/widgets/markdown/inline_citation_text.dart

@@ -0,0 +1,92 @@
+import 'package:flutter/material.dart';
+
+import '../../../core/models/chat_message.dart';
+import '../../../core/utils/citation_parser.dart';
+import '../../theme/theme_extensions.dart';
+import 'citation_badge.dart';
+
+/// Renders text with inline citation badges.
+///
+/// Parses citation patterns like [1], [2,3] and renders them as clickable
+/// badges showing source titles inline with the surrounding text.
+class InlineCitationText extends StatelessWidget {
+  const InlineCitationText({
+    super.key,
+    required this.text,
+    required this.sources,
+    this.style,
+    this.onSourceTap,
+  });
+
+  /// The text content that may contain citation patterns like [1], [2,3].
+  final String text;
+
+  /// Available sources for citation lookup.
+  final List<ChatSourceReference> sources;
+
+  /// Base text style.
+  final TextStyle? style;
+
+  /// Callback when a source badge is tapped.
+  final void Function(int sourceIndex)? onSourceTap;
+
+  @override
+  Widget build(BuildContext context) {
+    final segments = CitationParser.parse(text);
+
+    // If no citations found, render as plain text
+    if (segments == null || segments.isEmpty) {
+      return Text(text, style: style);
+    }
+
+    final theme = context.conduitTheme;
+    final baseStyle =
+        style ??
+        TextStyle(
+          color: theme.textPrimary,
+          fontSize: AppTypography.bodyMedium,
+          height: 1.45,
+        );
+
+    final spans = <InlineSpan>[];
+
+    for (final segment in segments) {
+      if (segment.isText && segment.text != null) {
+        spans.add(TextSpan(text: segment.text, style: baseStyle));
+      } else if (segment.isCitation && segment.citation != null) {
+        final citation = segment.citation!;
+        spans.add(
+          WidgetSpan(
+            alignment: PlaceholderAlignment.middle,
+            child: _buildCitationBadge(context, citation.sourceIds),
+          ),
+        );
+      }
+    }
+
+    return Text.rich(TextSpan(children: spans), style: baseStyle);
+  }
+
+  Widget _buildCitationBadge(BuildContext context, List<int> sourceIds) {
+    if (sourceIds.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    // Convert to 0-based indices
+    final indices = sourceIds.map((id) => id - 1).toList();
+
+    if (indices.length == 1) {
+      return CitationBadge(
+        sourceIndex: indices.first,
+        sources: sources,
+        onTap: onSourceTap != null ? () => onSourceTap!(indices.first) : null,
+      );
+    }
+
+    return CitationBadgeGroup(
+      sourceIndices: indices,
+      sources: sources,
+      onSourceTap: onSourceTap,
+    );
+  }
+}

lib/shared/widgets/markdown/streaming_markdown_widget.dart

@@ -2,6 +2,7 @@ import 'package:flutter/material.dart';
 
 import '../../../core/models/chat_message.dart';
 import '../../../core/utils/citation_parser.dart';
+import '../../theme/theme_extensions.dart';
 import 'citation_badge.dart';
 import 'markdown_config.dart';
 import 'markdown_preprocessor.dart';
@@ -134,11 +135,10 @@ class StreamingMarkdownWidget extends StatelessWidget {
     return SelectionArea(child: result);
   }
 
-  /// Builds markdown content with citation source references.
+  /// Builds markdown content with inline citation badges.
   ///
-  /// Citations like [1], [2] are kept as text in the markdown to preserve
-  /// inline formatting. A source reference footer is added when citations
-  /// are detected, providing clickable access to sources.
+  /// Citations like [1], [2] are rendered as clickable badges inline
+  /// within the text, matching OpenWebUI's behavior.
   Widget _buildMarkdownWithCitations(BuildContext context, String data) {
     // If no sources provided, render plain markdown
     if (sources == null || sources!.isEmpty) {
@@ -160,9 +160,43 @@ class StreamingMarkdownWidget extends StatelessWidget {
       );
     }
 
-    // Extract unique source IDs referenced in the content
-    final referencedIds = CitationParser.extractSourceIds(data);
-    if (referencedIds.isEmpty) {
+    // Render content with inline citation badges
+    return _InlineCitationMarkdown(
+      data: data,
+      sources: sources!,
+      onTapLink: onTapLink,
+      onSourceTap: onSourceTap,
+      imageBuilderOverride: imageBuilderOverride,
+    );
+  }
+}
+
+/// Widget that renders markdown with inline citation badges.
+///
+/// Parses the markdown content, identifies citation patterns, and renders
+/// them as clickable badges inline with the text.
+class _InlineCitationMarkdown extends StatelessWidget {
+  const _InlineCitationMarkdown({
+    required this.data,
+    required this.sources,
+    this.onTapLink,
+    this.onSourceTap,
+    this.imageBuilderOverride,
+  });
+
+  final String data;
+  final List<ChatSourceReference> sources;
+  final MarkdownLinkTapCallback? onTapLink;
+  final void Function(int sourceIndex)? onSourceTap;
+  final Widget Function(Uri uri, String? title, String? alt)?
+  imageBuilderOverride;
+
+  @override
+  Widget build(BuildContext context) {
+    // Split content into lines/paragraphs for processing
+    final segments = _parseContentWithCitations(data);
+
+    if (segments.isEmpty) {
       return ConduitMarkdown.build(
         context: context,
         data: data,
@@ -171,67 +205,219 @@ class StreamingMarkdownWidget extends StatelessWidget {
       );
     }
 
-    // Render markdown content as-is (preserving all formatting)
-    // and add a source references footer
-    return Column(
-      crossAxisAlignment: CrossAxisAlignment.start,
-      mainAxisSize: MainAxisSize.min,
-      children: [
+    // Build widgets for each segment
+    final children = <Widget>[];
+    final buffer = StringBuffer();
+
+    for (final segment in segments) {
+      if (segment.hasCitations) {
+        // Flush any accumulated non-citation content
+        if (buffer.isNotEmpty) {
+          children.add(
+            ConduitMarkdown.build(
+              context: context,
+              data: buffer.toString(),
+              onTapLink: onTapLink,
+              imageBuilderOverride: imageBuilderOverride,
+            ),
+          );
+          buffer.clear();
+        }
+
+        // Render this segment with inline citations
+        children.add(_buildParagraphWithCitations(context, segment.text));
+      } else {
+        // Accumulate non-citation content
+        if (buffer.isNotEmpty) {
+          buffer.writeln();
+        }
+        buffer.write(segment.text);
+      }
+    }
+
+    // Flush remaining content
+    if (buffer.isNotEmpty) {
+      children.add(
         ConduitMarkdown.build(
           context: context,
-          data: data,
+          data: buffer.toString(),
           onTapLink: onTapLink,
           imageBuilderOverride: imageBuilderOverride,
         ),
-        _SourceReferencesFooter(
-          referencedIds: referencedIds,
-          sources: sources!,
-          onSourceTap: onSourceTap,
-        ),
-      ],
+      );
+    }
+
+    if (children.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    if (children.length == 1) {
+      return children.first;
+    }
+
+    return Column(
+      crossAxisAlignment: CrossAxisAlignment.start,
+      mainAxisSize: MainAxisSize.min,
+      children: children,
+    );
+  }
+
+  /// Parses content into segments, identifying which have citations.
+  List<_ContentSegment> _parseContentWithCitations(String content) {
+    final segments = <_ContentSegment>[];
+
+    // Split by double newlines (paragraphs) while preserving structure
+    final paragraphs = content.split(RegExp(r'\n\n+'));
+
+    for (final paragraph in paragraphs) {
+      if (paragraph.trim().isEmpty) continue;
+
+      final hasCitations = CitationParser.hasCitations(paragraph);
+      segments.add(
+        _ContentSegment(text: paragraph, hasCitations: hasCitations),
+      );
+    }
+
+    return segments;
+  }
+
+  /// Builds a paragraph widget with inline citation badges.
+  Widget _buildParagraphWithCitations(BuildContext context, String text) {
+    final theme = context.conduitTheme;
+    final segments = CitationParser.parse(text);
+
+    if (segments == null || segments.isEmpty) {
+      return Text(text);
+    }
+
+    final baseStyle = AppTypography.bodyMediumStyle.copyWith(
+      color: theme.textPrimary,
+      height: 1.45,
+    );
+
+    final spans = <InlineSpan>[];
+
+    for (final segment in segments) {
+      if (segment.isText && segment.text != null) {
+        // Process text for basic markdown formatting
+        spans.add(_buildTextSpan(segment.text!, baseStyle, theme));
+      } else if (segment.isCitation && segment.citation != null) {
+        final citation = segment.citation!;
+        spans.add(
+          WidgetSpan(
+            alignment: PlaceholderAlignment.middle,
+            child: _buildCitationBadge(context, citation.sourceIds),
+          ),
+        );
+      }
+    }
+
+    return Text.rich(TextSpan(children: spans), style: baseStyle);
+  }
+
+  /// Builds a text span with basic markdown formatting support.
+  InlineSpan _buildTextSpan(
+    String text,
+    TextStyle baseStyle,
+    ConduitThemeExtension theme,
+  ) {
+    // Handle basic inline markdown: **bold**, *italic*, `code`
+    final spans = <InlineSpan>[];
+
+    // Pattern for bold, italic, and code
+    final pattern = RegExp(r'(\*\*(.+?)\*\*)|(\*(.+?)\*)|(`(.+?)`)');
+
+    var lastEnd = 0;
+    for (final match in pattern.allMatches(text)) {
+      // Add text before match
+      if (match.start > lastEnd) {
+        spans.add(
+          TextSpan(
+            text: text.substring(lastEnd, match.start),
+            style: baseStyle,
+          ),
+        );
+      }
+
+      if (match.group(1) != null) {
+        // Bold **text**
+        spans.add(
+          TextSpan(
+            text: match.group(2),
+            style: baseStyle.copyWith(fontWeight: FontWeight.bold),
+          ),
+        );
+      } else if (match.group(3) != null) {
+        // Italic *text*
+        spans.add(
+          TextSpan(
+            text: match.group(4),
+            style: baseStyle.copyWith(fontStyle: FontStyle.italic),
+          ),
+        );
+      } else if (match.group(5) != null) {
+        // Code `text`
+        spans.add(
+          TextSpan(
+            text: match.group(6),
+            style: baseStyle.copyWith(
+              fontFamily: AppTypography.monospaceFontFamily,
+              backgroundColor: theme.surfaceContainer.withValues(alpha: 0.3),
+            ),
+          ),
+        );
+      }
+
+      lastEnd = match.end;
+    }
+
+    // Add remaining text
+    if (lastEnd < text.length) {
+      spans.add(TextSpan(text: text.substring(lastEnd), style: baseStyle));
+    }
+
+    if (spans.isEmpty) {
+      return TextSpan(text: text, style: baseStyle);
+    }
+
+    if (spans.length == 1) {
+      return spans.first;
+    }
+
+    return TextSpan(children: spans);
+  }
+
+  /// Builds a citation badge widget.
+  Widget _buildCitationBadge(BuildContext context, List<int> sourceIds) {
+    if (sourceIds.isEmpty) {
+      return const SizedBox.shrink();
+    }
+
+    // Convert to 0-based indices
+    final indices = sourceIds.map((id) => id - 1).toList();
+
+    if (indices.length == 1) {
+      return CitationBadge(
+        sourceIndex: indices.first,
+        sources: sources,
+        onTap: onSourceTap != null ? () => onSourceTap!(indices.first) : null,
+      );
+    }
+
+    return CitationBadgeGroup(
+      sourceIndices: indices,
+      sources: sources,
+      onSourceTap: onSourceTap,
     );
   }
 }
 
-/// Footer widget showing source references with clickable badges.
-class _SourceReferencesFooter extends StatelessWidget {
-  const _SourceReferencesFooter({
-    required this.referencedIds,
-    required this.sources,
-    this.onSourceTap,
-  });
+/// A segment of content that may or may not contain citations.
+class _ContentSegment {
+  final String text;
+  final bool hasCitations;
 
-  /// 1-based source IDs that are referenced in the content.
-  final List<int> referencedIds;
-
-  /// All available sources.
-  final List<ChatSourceReference> sources;
-
-  /// Callback when a source is tapped.
-  final void Function(int sourceIndex)? onSourceTap;
-
-  @override
-  Widget build(BuildContext context) {
-    if (referencedIds.isEmpty) {
-      return const SizedBox.shrink();
-    }
-
-    return Padding(
-      padding: const EdgeInsets.only(top: 8),
-      child: Wrap(
-        spacing: 4,
-        runSpacing: 4,
-        children: [
-          for (final id in referencedIds)
-            CitationBadge(
-              sourceIndex: id - 1, // Convert to 0-based
-              sources: sources,
-              onTap: onSourceTap != null ? () => onSourceTap!(id - 1) : null,
-            ),
-        ],
-      ),
-    );
-  }
+  const _ContentSegment({required this.text, required this.hasCitations});
 }
 
 /// Types of special blocks that need custom rendering