fix: preserve stream candidate order in done events

Author: Wangnov

rust-genai/src/models.rs

@@ -122,7 +122,8 @@ impl GenerateContentEventStream {
                 None => {
                     self.finished = true;
                     if self.saw_done.load(Ordering::Relaxed) {
-                        if let Some(response) = self.aggregate_response.take() {
+                        if let Some(mut response) = self.aggregate_response.take() {
+                            normalize_stream_candidate_order(&mut response);
                             return Ok(Some(GenerateContentStreamEvent::Done(response)));
                         }
                     }
@@ -243,6 +244,39 @@ fn late_index_stream_position(
         .filter(|&candidate_position| aggregate_candidates[candidate_position].index.is_none())
 }
 
+fn normalize_stream_candidate_order(response: &mut GenerateContentResponse) {
+    if response.candidates.len() < 2 {
+        return;
+    }
+
+    let mut ordered = vec![None; response.candidates.len()];
+    let mut unindexed = VecDeque::new();
+    let mut overflow = VecDeque::new();
+
+    for candidate in std::mem::take(&mut response.candidates) {
+        match candidate
+            .index
+            .and_then(|index| usize::try_from(index).ok())
+            .filter(|&index| index < ordered.len())
+        {
+            Some(index) if ordered[index].is_none() => ordered[index] = Some(candidate),
+            Some(_) => overflow.push_back(candidate),
+            None if candidate.index.is_none() => unindexed.push_back(candidate),
+            None => overflow.push_back(candidate),
+        }
+    }
+
+    for slot in &mut ordered {
+        if slot.is_none() {
+            *slot = unindexed.pop_front().or_else(|| overflow.pop_front());
+        }
+    }
+
+    response.candidates = ordered.into_iter().flatten().collect();
+    response.candidates.extend(unindexed);
+    response.candidates.extend(overflow);
+}
+
 fn merge_candidate(
     existing: &mut rust_genai_types::response::Candidate,
     next: &rust_genai_types::response::Candidate,

rust-genai/src/models/tests.rs

@@ -1565,6 +1565,100 @@ fn test_merge_stream_response_merges_late_index_into_sparse_multi_candidate() {
     );
 }
 
+#[test]
+fn test_normalize_stream_candidate_order_reorders_indexed_candidates() {
+    let mut response = GenerateContentResponse {
+        sdk_http_response: None,
+        candidates: vec![
+            Candidate {
+                content: Some(Content::from_parts(vec![Part::text("second")], Role::Model)),
+                citation_metadata: None,
+                finish_message: None,
+                token_count: None,
+                finish_reason: None,
+                avg_logprobs: None,
+                grounding_metadata: None,
+                index: Some(1),
+                logprobs_result: None,
+                safety_ratings: Vec::new(),
+                url_context_metadata: None,
+            },
+            Candidate {
+                content: Some(Content::from_parts(vec![Part::text("first")], Role::Model)),
+                citation_metadata: None,
+                finish_message: None,
+                token_count: None,
+                finish_reason: None,
+                avg_logprobs: None,
+                grounding_metadata: None,
+                index: Some(0),
+                logprobs_result: None,
+                safety_ratings: Vec::new(),
+                url_context_metadata: None,
+            },
+        ],
+        create_time: None,
+        automatic_function_calling_history: None,
+        prompt_feedback: None,
+        usage_metadata: None,
+        model_version: None,
+        response_id: None,
+    };
+
+    normalize_stream_candidate_order(&mut response);
+
+    assert_eq!(response.text().as_deref(), Some("first"));
+    assert_eq!(response.candidates[0].index, Some(0));
+    assert_eq!(response.candidates[1].index, Some(1));
+}
+
+#[test]
+fn test_normalize_stream_candidate_order_preserves_unindexed_gap_positions() {
+    let mut response = GenerateContentResponse {
+        sdk_http_response: None,
+        candidates: vec![
+            Candidate {
+                content: Some(Content::from_parts(vec![Part::text("second")], Role::Model)),
+                citation_metadata: None,
+                finish_message: None,
+                token_count: None,
+                finish_reason: None,
+                avg_logprobs: None,
+                grounding_metadata: None,
+                index: Some(1),
+                logprobs_result: None,
+                safety_ratings: Vec::new(),
+                url_context_metadata: None,
+            },
+            Candidate {
+                content: Some(Content::from_parts(vec![Part::text("first")], Role::Model)),
+                citation_metadata: None,
+                finish_message: None,
+                token_count: None,
+                finish_reason: None,
+                avg_logprobs: None,
+                grounding_metadata: None,
+                index: None,
+                logprobs_result: None,
+                safety_ratings: Vec::new(),
+                url_context_metadata: None,
+            },
+        ],
+        create_time: None,
+        automatic_function_calling_history: None,
+        prompt_feedback: None,
+        usage_metadata: None,
+        model_version: None,
+        response_id: None,
+    };
+
+    normalize_stream_candidate_order(&mut response);
+
+    assert_eq!(response.text().as_deref(), Some("first"));
+    assert_eq!(response.candidates[0].index, None);
+    assert_eq!(response.candidates[1].index, Some(1));
+}
+
 #[test]
 fn test_stream_merge_helpers_respect_context_and_targets() {
     let resolution_low = PartMediaResolution {