rebase on top of main

Author: biswapanda

lib/llm/src/kv_router.rs

@@ -337,14 +337,16 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
         match self.inner.client.instance_source.as_ref() {
             InstanceSource::Static => self.inner.r#static(request).await,
             InstanceSource::Dynamic(_) => {
+                let context_id = request.context().id().to_string();
                 let (instance_id, overlap_amount) = if let Some(id) = request.backend_instance_id {
                     // If instance_id is set, use it
                     (id, 0)
                 } else {
                     // Otherwise, find the best match
-                    self.chooser.find_best_match(&request.token_ids).await?
+                    self.chooser
+                        .find_best_match(&context_id, &request.token_ids)
+                        .await?
                 };
-
                 let query_instance_id = request.has_annotation("query_instance_id");
                 // Extract context information before moving the request
                 let stream_context = request.context().clone();

lib/llm/src/migration.rs

@@ -188,6 +188,7 @@ mod tests {
             mdc_sum: None,
             annotations: vec![],
             estimated_prefix_hit_num_blocks: None,
+            backend_instance_id: None,
         }
     }
 
@@ -202,6 +203,7 @@ mod tests {
             top_logprobs: None,
             finish_reason: None,
             index: None,
+            backend_instance_id: None,
         })
     }