🚧 Fix race condition in streaming output detection

Signed-off-by: m-misiura <mmisiura@redhat.com>

Author: m-misiura

src/orchestrator/handlers/chat_completions_detection/streaming.rs

@@ -563,14 +563,29 @@ async fn handle_whole_doc_detection(
 }
 
 /// Builds a response with output detections.
-fn output_detection_response(
+async fn output_detection_response(
     completion_state: &Arc<CompletionState<ChatCompletionChunk>>,
     choice_index: u32,
     chunk: Chunk,
     detections: Vec<Detection>,
 ) -> Result<ChatCompletionChunk, Error> {
-    // Get chat completions for this choice index
-    let chat_completions = completion_state.completions.get(&choice_index).unwrap();
+    // Wait for entry to exist (yields to other tasks until ready)
+    let chat_completions = {
+        let mut attempts = 0;
+        loop {
+            if let Some(entry) = completion_state.completions.get(&choice_index) {
+                break entry;
+            }
+            if attempts > 1000 {
+                return Err(Error::Other(format!(
+                    "completion entry for choice_index {} not ready after 1000 yields",
+                    choice_index
+                )));
+            }
+            attempts += 1;
+            tokio::task::yield_now().await;
+        }
+    };
     // Get range of chat completions for this chunk
     let chat_completions = chat_completions
         .range(chunk.input_start_index..=chunk.input_end_index)
@@ -709,6 +724,7 @@ async fn process_detection_batch_stream(
             Ok((choice_index, chunk, detections)) => {
                 let input_end_index = chunk.input_end_index;
                 match output_detection_response(&completion_state, choice_index, chunk, detections)
+                    .await
                 {
                     Ok(chat_completion) => {
                         // Send chat completion to response channel
@@ -718,8 +734,22 @@ async fn process_detection_batch_stream(
                             return;
                         }
                         // If this is the final chat completion chunk with content, send chat completion chunk with finish reason
-                        let chat_completions =
-                            completion_state.completions.get(&choice_index).unwrap();
+                        // Wait for entry to exist (yields to other tasks until ready)
+                        let chat_completions = {
+                            let mut attempts = 0;
+                            loop {
+                                if let Some(entry) = completion_state.completions.get(&choice_index)
+                                {
+                                    break entry;
+                                }
+                                if attempts > 1000 {
+                                    error!(%trace_id, %choice_index, "completion entry not ready after 1000 yields");
+                                    return;
+                                }
+                                attempts += 1;
+                                tokio::task::yield_now().await;
+                            }
+                        };
                         if chat_completions.keys().rev().nth(1) == Some(&input_end_index)
                             && let Some((_, chat_completion)) = chat_completions.last_key_value()
                             && chat_completion