refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表

crates/aether-scheduler-core/src/candidate.rs

@@ -45,6 +45,20 @@ pub struct BuildMinimalCandidateSelectionInput<'a> {
     pub priority_mode: SchedulerPriorityMode,
 }
 
+#[derive(Debug, Clone, Copy)]
+struct RequiredCapabilityDescriptor<'a> {
+    name: &'a str,
+    compatible: bool,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct CandidateOrderingState {
+    capability_priority: (u32, u32),
+    affinity_hash: Option<u64>,
+    health_bucket: Option<crate::ProviderKeyHealthBucket>,
+    health_score: f64,
+}
+
 pub fn candidate_supports_required_capability(
     candidate: &SchedulerMinimalCandidateSelectionCandidate,
     required_capability: &str,
@@ -91,23 +105,17 @@ pub fn requested_capability_priority_for_candidate(
         return (0, 0);
     };
 
-    let mut exclusive_misses = 0u32;
-    let mut compatible_misses = 0u32;
-    for (capability, value) in required_capabilities {
-        if !requested_capability_is_enabled(value) {
-            continue;
-        }
-        if candidate_supports_required_capability(candidate, capability) {
-            continue;
-        }
-        if requested_capability_is_compatible(capability) {
-            compatible_misses += 1;
-        } else {
-            exclusive_misses += 1;
-        }
-    }
-
-    (exclusive_misses, compatible_misses)
+    requested_capability_priority_for_candidate_descriptors(
+        required_capabilities
+            .iter()
+            .filter_map(|(capability, value)| {
+                requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
+                    name: capability.as_str(),
+                    compatible: requested_capability_is_compatible(capability),
+                })
+            }),
+        candidate,
+    )
 }
 
 pub fn auth_api_key_concurrency_limit_reached(
@@ -153,7 +161,8 @@ pub fn build_minimal_candidate_selection(
         return Ok(Vec::new());
     }
 
-    let mut candidates = Vec::new();
+    let required_capabilities = enabled_required_capabilities(required_capabilities);
+    let mut candidates = Vec::with_capacity(rows.len());
     for row in rows {
         if !crate::auth_constraints_allow_provider(
             auth_constraints,
@@ -196,16 +205,9 @@ pub fn build_minimal_candidate_selection(
         });
     }
 
-    candidates.sort_by(|left, right| {
-        requested_capability_priority_for_candidate(required_capabilities, left)
-            .cmp(&requested_capability_priority_for_candidate(
-                required_capabilities,
-                right,
-            ))
-            .then_with(|| {
-                compare_candidates_by_priority_mode(left, right, priority_mode, affinity_key)
-            })
-    });
+    let ordering_states =
+        build_candidate_ordering_states(&candidates, &required_capabilities, affinity_key, None);
+    sort_candidates_by_ordering_state(&mut candidates, &ordering_states, priority_mode, false);
 
     Ok(candidates)
 }
@@ -298,28 +300,27 @@ pub fn collect_selectable_candidates_from_keys(
     selectable_keys: &BTreeSet<(String, String, String)>,
     cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
 ) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
-    let mut selected = Vec::new();
+    let mut promoted = None;
+    let mut selected = Vec::with_capacity(candidates.len());
     let mut emitted_keys = BTreeSet::new();
 
-    if let Some(target) = cached_affinity_target {
-        if let Some(candidate) = candidates
-            .iter()
-            .find(|candidate| crate::matches_affinity_target(candidate, target))
-            .cloned()
-        {
-            let key = crate::candidate_key(&candidate);
-            if selectable_keys.contains(&key) && emitted_keys.insert(key) {
-                selected.push(candidate);
-            }
-        }
-    }
-
     for candidate in candidates {
         let key = crate::candidate_key(&candidate);
         if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
             continue;
         }
-        selected.push(candidate);
+        if promoted.is_none()
+            && cached_affinity_target
+                .is_some_and(|target| crate::matches_affinity_target(&candidate, target))
+        {
+            promoted = Some(candidate);
+        } else {
+            selected.push(candidate);
+        }
+    }
+
+    if let Some(candidate) = promoted {
+        selected.insert(0, candidate);
     }
 
     selected
@@ -332,35 +333,14 @@ pub fn reorder_candidates_by_scheduler_health(
     affinity_key: Option<&str>,
     priority_mode: SchedulerPriorityMode,
 ) {
-    candidates.sort_by(|left, right| {
-        requested_capability_priority_for_candidate(required_capabilities, left)
-            .cmp(&requested_capability_priority_for_candidate(
-                required_capabilities,
-                right,
-            ))
-            .then_with(|| match priority_mode {
-                SchedulerPriorityMode::Provider => left
-                    .provider_priority
-                    .cmp(&right.provider_priority)
-                    .then(left.key_internal_priority.cmp(&right.key_internal_priority))
-                    .then_with(|| {
-                        compare_provider_key_health_order(left, right, provider_key_rpm_states)
-                    })
-                    .then_with(|| crate::compare_affinity_order(left, right, affinity_key))
-                    .then_with(|| compare_candidate_identity(left, right)),
-                SchedulerPriorityMode::GlobalKey => left
-                    .key_global_priority_for_format
-                    .unwrap_or(i32::MAX)
-                    .cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
-                    .then_with(|| {
-                        compare_provider_key_health_order(left, right, provider_key_rpm_states)
-                    })
-                    .then_with(|| crate::compare_affinity_order(left, right, affinity_key))
-                    .then(left.provider_priority.cmp(&right.provider_priority))
-                    .then(left.key_internal_priority.cmp(&right.key_internal_priority))
-                    .then_with(|| compare_candidate_identity(left, right)),
-            })
-    });
+    let required_capabilities = enabled_required_capabilities(required_capabilities);
+    let ordering_states = build_candidate_ordering_states(
+        candidates,
+        &required_capabilities,
+        affinity_key,
+        Some(provider_key_rpm_states),
+    );
+    sort_candidates_by_ordering_state(candidates, &ordering_states, priority_mode, true);
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -456,33 +436,6 @@ pub fn candidate_runtime_skip_reason_with_state(
     None
 }
 
-fn compare_provider_key_health_order(
-    left: &SchedulerMinimalCandidateSelectionCandidate,
-    right: &SchedulerMinimalCandidateSelectionCandidate,
-    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
-) -> std::cmp::Ordering {
-    let left_bucket = candidate_provider_key_health_bucket(left, provider_key_rpm_states);
-    let right_bucket = candidate_provider_key_health_bucket(right, provider_key_rpm_states);
-    right_bucket.cmp(&left_bucket).then_with(|| {
-        let left_score = candidate_provider_key_health_score(left, provider_key_rpm_states);
-        let right_score = candidate_provider_key_health_score(right, provider_key_rpm_states);
-        right_score
-            .partial_cmp(&left_score)
-            .unwrap_or(std::cmp::Ordering::Equal)
-    })
-}
-
-fn candidate_provider_key_health_bucket(
-    candidate: &SchedulerMinimalCandidateSelectionCandidate,
-    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
-) -> Option<crate::ProviderKeyHealthBucket> {
-    provider_key_rpm_states
-        .get(&candidate.key_id)
-        .and_then(|key| {
-            crate::provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
-        })
-}
-
 fn compare_candidate_identity(
     left: &SchedulerMinimalCandidateSelectionCandidate,
     right: &SchedulerMinimalCandidateSelectionCandidate,
@@ -497,12 +450,190 @@ fn compare_candidate_identity(
         )
 }
 
+fn enabled_required_capabilities(
+    required_capabilities: Option<&serde_json::Value>,
+) -> Vec<RequiredCapabilityDescriptor<'_>> {
+    let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
+    else {
+        return Vec::new();
+    };
+
+    required_capabilities
+        .iter()
+        .filter_map(|(capability, value)| {
+            requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
+                name: capability.as_str(),
+                compatible: requested_capability_is_compatible(capability),
+            })
+        })
+        .collect()
+}
+
+fn requested_capability_priority_for_candidate_descriptors<'a, I>(
+    required_capabilities: I,
+    candidate: &SchedulerMinimalCandidateSelectionCandidate,
+) -> (u32, u32)
+where
+    I: IntoIterator<Item = RequiredCapabilityDescriptor<'a>>,
+{
+    let mut exclusive_misses = 0u32;
+    let mut compatible_misses = 0u32;
+    for capability in required_capabilities {
+        if candidate_supports_required_capability(candidate, capability.name) {
+            continue;
+        }
+        if capability.compatible {
+            compatible_misses += 1;
+        } else {
+            exclusive_misses += 1;
+        }
+    }
+
+    (exclusive_misses, compatible_misses)
+}
+
+fn build_candidate_ordering_states(
+    candidates: &[SchedulerMinimalCandidateSelectionCandidate],
+    required_capabilities: &[RequiredCapabilityDescriptor<'_>],
+    affinity_key: Option<&str>,
+    provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
+) -> Vec<CandidateOrderingState> {
+    candidates
+        .iter()
+        .map(|candidate| CandidateOrderingState {
+            capability_priority: requested_capability_priority_for_candidate_descriptors(
+                required_capabilities.iter().copied(),
+                candidate,
+            ),
+            affinity_hash: affinity_key.map(|key| crate::candidate_affinity_hash(key, candidate)),
+            health_bucket: provider_key_rpm_states.and_then(|states| {
+                states.get(&candidate.key_id).and_then(|key| {
+                    crate::provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
+                })
+            }),
+            health_score: candidate_provider_key_health_score(candidate, provider_key_rpm_states),
+        })
+        .collect()
+}
+
+fn sort_candidates_by_ordering_state(
+    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
+    ordering_states: &[CandidateOrderingState],
+    priority_mode: SchedulerPriorityMode,
+    include_health: bool,
+) {
+    if candidates.len() < 2 {
+        return;
+    }
+
+    let mut order = (0..candidates.len()).collect::<Vec<_>>();
+    order.sort_by(|left, right| {
+        compare_candidates_with_ordering_state(
+            &ordering_states[*left],
+            &candidates[*left],
+            &ordering_states[*right],
+            &candidates[*right],
+            priority_mode,
+            include_health,
+        )
+    });
+    apply_candidate_order(candidates, order);
+}
+
+fn compare_candidates_with_ordering_state(
+    left_state: &CandidateOrderingState,
+    left_candidate: &SchedulerMinimalCandidateSelectionCandidate,
+    right_state: &CandidateOrderingState,
+    right_candidate: &SchedulerMinimalCandidateSelectionCandidate,
+    priority_mode: SchedulerPriorityMode,
+    include_health: bool,
+) -> std::cmp::Ordering {
+    left_state
+        .capability_priority
+        .cmp(&right_state.capability_priority)
+        .then_with(|| {
+            compare_priority_before_health(left_candidate, right_candidate, priority_mode)
+        })
+        .then_with(|| {
+            if include_health {
+                compare_provider_key_health_state(left_state, right_state)
+            } else {
+                std::cmp::Ordering::Equal
+            }
+        })
+        .then_with(|| left_state.affinity_hash.cmp(&right_state.affinity_hash))
+        .then_with(|| {
+            compare_priority_after_affinity(left_candidate, right_candidate, priority_mode)
+        })
+        .then_with(|| compare_candidate_identity(left_candidate, right_candidate))
+}
+
+fn compare_priority_before_health(
+    left: &SchedulerMinimalCandidateSelectionCandidate,
+    right: &SchedulerMinimalCandidateSelectionCandidate,
+    priority_mode: SchedulerPriorityMode,
+) -> std::cmp::Ordering {
+    match priority_mode {
+        SchedulerPriorityMode::Provider => left
+            .provider_priority
+            .cmp(&right.provider_priority)
+            .then(left.key_internal_priority.cmp(&right.key_internal_priority)),
+        SchedulerPriorityMode::GlobalKey => left
+            .key_global_priority_for_format
+            .unwrap_or(i32::MAX)
+            .cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX)),
+    }
+}
+
+fn compare_priority_after_affinity(
+    left: &SchedulerMinimalCandidateSelectionCandidate,
+    right: &SchedulerMinimalCandidateSelectionCandidate,
+    priority_mode: SchedulerPriorityMode,
+) -> std::cmp::Ordering {
+    match priority_mode {
+        SchedulerPriorityMode::Provider => std::cmp::Ordering::Equal,
+        SchedulerPriorityMode::GlobalKey => left
+            .provider_priority
+            .cmp(&right.provider_priority)
+            .then(left.key_internal_priority.cmp(&right.key_internal_priority)),
+    }
+}
+
+fn compare_provider_key_health_state(
+    left: &CandidateOrderingState,
+    right: &CandidateOrderingState,
+) -> std::cmp::Ordering {
+    right
+        .health_bucket
+        .cmp(&left.health_bucket)
+        .then_with(|| right.health_score.total_cmp(&left.health_score))
+}
+
+fn apply_candidate_order(
+    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
+    sorted_old_indices: Vec<usize>,
+) {
+    let mut target_positions = vec![0usize; sorted_old_indices.len()];
+    for (new_position, old_position) in sorted_old_indices.into_iter().enumerate() {
+        target_positions[old_position] = new_position;
+    }
+
+    for index in 0..candidates.len() {
+        let current = index;
+        while target_positions[current] != current {
+            let target = target_positions[current];
+            candidates.swap(current, target);
+            target_positions.swap(current, target);
+        }
+    }
+}
+
 fn candidate_provider_key_health_score(
     candidate: &SchedulerMinimalCandidateSelectionCandidate,
-    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
+    provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
 ) -> f64 {
     provider_key_rpm_states
-        .get(&candidate.key_id)
+        .and_then(|states| states.get(&candidate.key_id))
         .and_then(|key| {
             crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
         })

crates/aether-scheduler-core/src/model.rs

@@ -44,15 +44,17 @@ fn resolve_global_model_name_by<F>(
 where
     F: Fn(&StoredMinimalCandidateSelectionRow) -> bool,
 {
-    let mut matches = rows
-        .iter()
-        .filter(|row| matches(row))
-        .map(|row| row.global_model_name.trim())
-        .filter(|value| !value.is_empty())
-        .map(ToOwned::to_owned)
-        .collect::<BTreeSet<_>>()
-        .into_iter();
-    matches.next()
+    let mut best_match = None::<&str>;
+    for row in rows.iter().filter(|row| matches(row)) {
+        let candidate = row.global_model_name.trim();
+        if candidate.is_empty() {
+            continue;
+        }
+        if best_match.is_none_or(|current| candidate < current) {
+            best_match = Some(candidate);
+        }
+    }
+    best_match.map(ToOwned::to_owned)
 }
 
 pub fn resolve_provider_model_name(
@@ -75,25 +77,26 @@ pub fn resolve_provider_model_name(
         return Some((selected_provider_model_name, None));
     }
 
-    let candidate_models = candidate_model_names(row, api_format);
     let mut sorted_allowed_models = key_allowed_models
         .iter()
-        .map(|value| value.trim())
+        .map(String::as_str)
+        .map(str::trim)
         .filter(|value| !value.is_empty())
-        .map(ToOwned::to_owned)
         .collect::<Vec<_>>();
-    sorted_allowed_models.sort();
+    sorted_allowed_models.sort_unstable();
 
-    for allowed_model in &sorted_allowed_models {
-        if candidate_models.contains(allowed_model.as_str()) {
-            return Some((allowed_model.clone(), Some(allowed_model.clone())));
+    for &allowed_model in &sorted_allowed_models {
+        if row_has_candidate_model_name(row, api_format, allowed_model) {
+            let allowed_model = allowed_model.to_owned();
+            return Some((allowed_model.clone(), Some(allowed_model)));
         }
     }
 
     let global_model_mappings = row.global_model_mappings.as_ref()?;
-    for allowed_model in sorted_allowed_models {
+    for &allowed_model in &sorted_allowed_models {
         for pattern in global_model_mappings {
-            if matches_model_mapping(pattern, &allowed_model) {
+            if matches_model_mapping(pattern, allowed_model) {
+                let allowed_model = allowed_model.to_owned();
                 return Some((allowed_model.clone(), Some(allowed_model)));
             }
         }
@@ -110,23 +113,14 @@ pub fn select_provider_model_name(
         return row.model_provider_model_name.clone();
     };
 
-    let mut scoped = mappings
+    mappings
         .iter()
         .filter(|mapping| mapping_scope_matches(mapping, api_format))
-        .collect::<Vec<_>>();
-    if scoped.is_empty() {
-        return row.model_provider_model_name.clone();
-    }
-
-    scoped.sort_by(|left, right| {
-        left.priority
-            .cmp(&right.priority)
-            .then(left.name.cmp(&right.name))
-    });
-    let top_priority = scoped[0].priority;
-    scoped
-        .into_iter()
-        .find(|mapping| mapping.priority == top_priority)
+        .min_by(|left, right| {
+            left.priority
+                .cmp(&right.priority)
+                .then(left.name.cmp(&right.name))
+        })
         .map(|mapping| mapping.name.clone())
         .unwrap_or_else(|| row.model_provider_model_name.clone())
 }
@@ -153,7 +147,7 @@ fn mapping_scope_matches(mapping: &StoredProviderModelMapping, api_format: &str)
 
     api_formats
         .iter()
-        .any(|value| normalize_api_format(value) == api_format)
+        .any(|value| api_format_matches(value, api_format))
 }
 
 pub fn row_supports_required_capability(
@@ -230,7 +224,7 @@ pub fn extract_global_priority_for_format(
 
     let Some(value) = object
         .iter()
-        .find(|(key, _)| normalize_api_format(key) == api_format)
+        .find(|(key, _)| api_format_matches(key, api_format))
         .map(|(_, value)| value)
     else {
         return Ok(None);
@@ -262,6 +256,26 @@ pub fn normalize_api_format(value: &str) -> String {
     value.trim().to_ascii_lowercase()
 }
 
+fn row_has_candidate_model_name(
+    row: &StoredMinimalCandidateSelectionRow,
+    api_format: &str,
+    model_name: &str,
+) -> bool {
+    row.model_provider_model_name == model_name
+        || row
+            .model_provider_model_mappings
+            .as_ref()
+            .is_some_and(|mappings| {
+                mappings.iter().any(|mapping| {
+                    mapping_scope_matches(mapping, api_format) && mapping.name == model_name
+                })
+            })
+}
+
+fn api_format_matches(left: &str, right: &str) -> bool {
+    left.trim().eq_ignore_ascii_case(right.trim())
+}
+
 #[cfg(test)]
 mod tests {
     use super::matches_model_mapping;

crates/aether-scheduler-core/src/request_candidate.rs

@@ -2,7 +2,7 @@ use aether_contracts::{ExecutionError, ExecutionPlan};
 use aether_data_contracts::repository::candidates::{
     RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
 };
-use serde_json::Value;
+use serde_json::{Map, Value};
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct SchedulerRequestCandidateReportContext {
@@ -91,20 +91,13 @@ pub fn parse_request_candidate_report_context(
     report_context: Option<&Value>,
 ) -> Option<SchedulerRequestCandidateReportContext> {
     let report_context = report_context?;
-    let retry_index = report_context
-        .get("retry_index")
-        .and_then(Value::as_u64)
-        .unwrap_or_default();
     Some(SchedulerRequestCandidateReportContext {
         request_id: string_field(report_context, "request_id"),
         candidate_id: string_field(report_context, "candidate_id"),
         user_id: string_field(report_context, "user_id"),
         api_key_id: string_field(report_context, "api_key_id"),
-        candidate_index: report_context
-            .get("candidate_index")
-            .and_then(Value::as_u64)
-            .and_then(|value| u32::try_from(value).ok()),
-        retry_index: u32::try_from(retry_index).unwrap_or(u32::MAX),
+        candidate_index: u32_field(report_context, "candidate_index"),
+        retry_index: u32_field(report_context, "retry_index").unwrap_or_default(),
         provider_id: string_field(report_context, "provider_id"),
         endpoint_id: string_field(report_context, "endpoint_id"),
         key_id: string_field(report_context, "key_id"),
@@ -123,9 +116,24 @@ pub fn resolve_report_request_candidate_slot(
     now_unix_ms: u64,
     generated_candidate_id: String,
 ) -> Option<SchedulerResolvedReportRequestCandidateSlot> {
-    let request_id = metadata.request_id.clone()?;
     let matched_candidate = match_existing_report_candidate(existing_candidates, &metadata);
-    let synthesized_extra_data = build_report_candidate_extra_data(&metadata);
+    let SchedulerRequestCandidateReportContext {
+        request_id,
+        candidate_id,
+        user_id,
+        api_key_id,
+        candidate_index: metadata_candidate_index,
+        retry_index,
+        provider_id,
+        endpoint_id,
+        key_id,
+        client_api_format,
+        provider_api_format,
+        proxy,
+    } = metadata;
+    let request_id = request_id?;
+    let synthesized_extra_data =
+        build_report_candidate_extra_data(client_api_format, provider_api_format, proxy);
     let created_at_unix_ms = matched_candidate
         .as_ref()
         .map(|candidate| candidate.created_at_unix_ms)
@@ -133,42 +141,42 @@ pub fn resolve_report_request_candidate_slot(
     let candidate_index = matched_candidate
         .as_ref()
         .map(|candidate| candidate.candidate_index)
-        .or(metadata.candidate_index)
+        .or(metadata_candidate_index)
         .unwrap_or_else(|| next_candidate_index(existing_candidates));
     let retry_index = matched_candidate
         .as_ref()
         .map(|candidate| candidate.retry_index)
-        .unwrap_or(metadata.retry_index);
+        .unwrap_or(retry_index);
 
     Some(SchedulerResolvedReportRequestCandidateSlot {
         id: matched_candidate
             .as_ref()
             .map(|candidate| candidate.id.clone())
-            .or(metadata.candidate_id)
+            .or(candidate_id)
             .unwrap_or(generated_candidate_id),
         request_id,
         user_id: matched_candidate
             .as_ref()
             .and_then(|candidate| candidate.user_id.clone())
-            .or(metadata.user_id),
+            .or(user_id),
         api_key_id: matched_candidate
             .as_ref()
             .and_then(|candidate| candidate.api_key_id.clone())
-            .or(metadata.api_key_id),
+            .or(api_key_id),
         candidate_index,
         retry_index,
         provider_id: matched_candidate
             .as_ref()
             .and_then(|candidate| candidate.provider_id.clone())
-            .or(metadata.provider_id),
+            .or(provider_id),
         endpoint_id: matched_candidate
             .as_ref()
             .and_then(|candidate| candidate.endpoint_id.clone())
-            .or(metadata.endpoint_id),
+            .or(endpoint_id),
         key_id: matched_candidate
             .as_ref()
             .and_then(|candidate| candidate.key_id.clone())
-            .or(metadata.key_id),
+            .or(key_id),
         extra_data: merge_request_candidate_extra_data(
             matched_candidate
                 .as_ref()
@@ -195,22 +203,14 @@ pub fn build_execution_request_candidate_seed(
         .and_then(Value::as_object)
         .cloned()
         .unwrap_or_default();
-    let request_id = string_field(&Value::Object(context.clone()), "request_id")
-        .unwrap_or_else(|| plan.request_id.clone());
-    let candidate_index = context
-        .get("candidate_index")
-        .and_then(Value::as_u64)
-        .and_then(|value| u32::try_from(value).ok())
-        .unwrap_or(0);
-    let retry_index = context
-        .get("retry_index")
-        .and_then(Value::as_u64)
-        .and_then(|value| u32::try_from(value).ok())
-        .unwrap_or(0);
-    let candidate_id = string_field(&Value::Object(context.clone()), "candidate_id")
-        .unwrap_or(generated_candidate_id);
-    let user_id = string_field(&Value::Object(context.clone()), "user_id");
-    let api_key_id = string_field(&Value::Object(context.clone()), "api_key_id");
+    let request_id =
+        string_field_from_object(&context, "request_id").unwrap_or_else(|| plan.request_id.clone());
+    let candidate_index = u32_field_from_object(&context, "candidate_index").unwrap_or(0);
+    let retry_index = u32_field_from_object(&context, "retry_index").unwrap_or(0);
+    let candidate_id =
+        string_field_from_object(&context, "candidate_id").unwrap_or(generated_candidate_id);
+    let user_id = string_field_from_object(&context, "user_id");
+    let api_key_id = string_field_from_object(&context, "api_key_id");
 
     context.insert("request_id".to_string(), Value::String(request_id.clone()));
     context.insert(
@@ -374,7 +374,10 @@ pub fn finalize_execution_request_candidate_report_context(
     report_context: Value,
     candidate_id: &str,
 ) -> Value {
-    let mut context = report_context.as_object().cloned().unwrap_or_default();
+    let mut context = match report_context {
+        Value::Object(context) => context,
+        _ => Map::new(),
+    };
     let candidate_id = candidate_id.trim();
     if !candidate_id.is_empty() {
         context.insert(
@@ -410,6 +413,12 @@ fn extract_error_message(body_json: &Value) -> Option<&str> {
 
 fn string_field(value: &Value, key: &str) -> Option<String> {
     value
+        .as_object()
+        .and_then(|object| string_field_from_object(object, key))
+}
+
+fn string_field_from_object(object: &Map<String, Value>, key: &str) -> Option<String> {
+    object
         .get(key)
         .and_then(Value::as_str)
         .map(str::trim)
@@ -417,6 +426,19 @@ fn string_field(value: &Value, key: &str) -> Option<String> {
         .map(ToOwned::to_owned)
 }
 
+fn u32_field(value: &Value, key: &str) -> Option<u32> {
+    value
+        .as_object()
+        .and_then(|object| u32_field_from_object(object, key))
+}
+
+fn u32_field_from_object(object: &Map<String, Value>, key: &str) -> Option<u32> {
+    object
+        .get(key)
+        .and_then(Value::as_u64)
+        .and_then(|value| u32::try_from(value).ok())
+}
+
 fn match_existing_report_candidate<'a>(
     candidates: &'a [StoredRequestCandidate],
     metadata: &SchedulerRequestCandidateReportContext,
@@ -465,24 +487,26 @@ fn next_candidate_index(candidates: &[StoredRequestCandidate]) -> u32 {
 }
 
 fn build_report_candidate_extra_data(
-    metadata: &SchedulerRequestCandidateReportContext,
+    client_api_format: Option<String>,
+    provider_api_format: Option<String>,
+    proxy: Option<Value>,
 ) -> Option<Value> {
-    let mut extra_data = serde_json::Map::new();
+    let mut extra_data = Map::with_capacity(5);
     extra_data.insert("gateway_execution_runtime".to_string(), Value::Bool(true));
     extra_data.insert("phase".to_string(), Value::String("3c_trial".to_string()));
-    if let Some(client_api_format) = metadata.client_api_format.clone() {
+    if let Some(client_api_format) = client_api_format {
         extra_data.insert(
             "client_api_format".to_string(),
             Value::String(client_api_format),
         );
     }
-    if let Some(provider_api_format) = metadata.provider_api_format.clone() {
+    if let Some(provider_api_format) = provider_api_format {
         extra_data.insert(
             "provider_api_format".to_string(),
             Value::String(provider_api_format),
         );
     }
-    if let Some(proxy) = metadata.proxy.clone() {
+    if let Some(proxy) = proxy {
         extra_data.insert("proxy".to_string(), proxy);
     }
     (!extra_data.is_empty()).then_some(Value::Object(extra_data))