perf(router): look up host step answers by id instead of scanning every candidate (#50704)

answered_by_id and the select answer check compared every returned id against every candidate, about 1M string compares per pick at 1,000 deployments

litellm-rust/crates/python-bridge/src/router/host.rs

@@ -357,7 +357,11 @@ impl HostCallResponse {
             Self::Select(ids) | Self::Filtered(Ok((ids, _))),
             HostCall::Select { candidates, .. },
         ) = (self, call)
-            && let Some(id) = ids.iter().find(|id| !candidates.contains(id))
+            && let Some(id) = {
+                let known: std::collections::HashSet<&str> =
+                    candidates.iter().map(String::as_str).collect();
+                ids.iter().find(|id| !known.contains(id.as_str()))
+            }
         {
             return Err(DriveError::UnknownCandidate { id: id.clone() });
         }
@@ -688,14 +692,22 @@ fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
 /// The candidates a host answer names, in its order. An id names every candidate that carries
 /// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
 fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
-    let mut taken = vec![false; candidates.len()];
+    let mut positions: std::collections::HashMap<&str, Vec<usize>> =
+        std::collections::HashMap::new();
+    for (position, deployment) in candidates.iter().enumerate() {
+        positions
+            .entry(deployment.id.as_str())
+            .or_default()
+            .push(position);
+    }
     let mut kept = Vec::new();
     for id in ids {
-        for (position, deployment) in candidates.iter().enumerate() {
-            if !taken[position] && deployment.id == *id {
-                taken[position] = true;
-                kept.push(deployment.clone());
-            }
+        if let Some(found) = positions.remove(id.as_str()) {
+            kept.extend(
+                found
+                    .into_iter()
+                    .map(|position| candidates[position].clone()),
+            );
         }
     }
     kept
@@ -1540,6 +1552,26 @@ mod tests {
         );
     }
 
+    #[rstest]
+    fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() {
+        let candidates: Vec<Arc<Deployment>> = ["a", "b", "a", "c"]
+            .into_iter()
+            .map(|id| {
+                let spec: litellm_router_core::DeploymentSpec =
+                    serde_json::from_value(deployment("g", id)).unwrap();
+                Arc::new(Deployment::from_spec(spec).unwrap())
+            })
+            .collect();
+        let ids = ["c", "a", "x", "a"].map(String::from);
+
+        let kept = answered_by_id(&candidates, &ids);
+
+        let positions: Vec<usize> = kept
+            .iter()
+            .map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap())
+            .collect();
+        assert_eq!(positions, [3, 0, 2]);
+    }
     #[rstest]
     fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
         let mut second = deployment("primary", "a1");