From 29fe3afe8b697a0d12fb11f86522ae2c08cbc29b Mon Sep 17 00:00:00 2001 From: Will Killian Date: Wed, 15 Jul 2026 10:17:49 -0400 Subject: [PATCH 1/3] refactor: resolve core and FFI Sonar findings Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 50 +- crates/core/src/api/optimization.rs | 495 ++++++++++++------ crates/core/src/observability/mod.rs | 148 ++++++ .../core/src/observability/openinference.rs | 137 +---- crates/core/src/observability/otel.rs | 138 +---- .../src/observability/plugin_component.rs | 99 ++-- crates/ffi/src/api/mod.rs | 74 ++- 7 files changed, 604 insertions(+), 537 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 7c587bf5e..971c8f183 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -688,25 +688,13 @@ fn llm_call_end_with_behavior( } else { Some(sanitized_response) }; - let mut decode_error = None; - let mut annotated_response = match annotated_response { - Some(annotated_response) => Some((*annotated_response).clone()), - None => match (response_codec.as_ref(), data.as_ref()) { - (Some(codec), Some(response)) => match codec.decode_response(response) { - Ok(mut decoded) => { - if behavior.attach_estimated_cost { - attach_estimated_cost_for_provider(&mut decoded, Some(&handle.name)); - } - Some(decoded) - } - Err(error) => { - decode_error = Some(error); - None - } - }, - _ => None, - }, - }; + let (mut annotated_response, decode_error) = resolve_llm_end_annotation( + annotated_response, + response_codec, + data.as_ref(), + &behavior, + &handle.name, + ); handle.optimization_recorder.close_for_finalization(None); emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); @@ -753,6 +741,30 @@ fn llm_call_end_with_behavior( } } +fn resolve_llm_end_annotation( + annotated_response: Option>, + response_codec: Option>, + data: Option<&Json>, + behavior: &LlmCallEndBehavior, + provider_name: &str, +) -> (Option, Option) { + if let Some(annotated_response) = annotated_response { + return (Some((*annotated_response).clone()), None); + } + let (Some(codec), Some(response)) = (response_codec, data) else { + return (None, None); + }; + match codec.decode_response(response) { + Ok(mut decoded) => { + if behavior.attach_estimated_cost { + attach_estimated_cost_for_provider(&mut decoded, Some(provider_name)); + } + (Some(decoded), None) + } + Err(error) => (None, Some(error)), + } +} + fn emit_llm_end_without_output( handle: &LlmHandle, metadata: Option, diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index 2f7bba7ac..4a1d940aa 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -14,7 +14,9 @@ use crate::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationModel, LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokens, }; -use crate::codec::response::{AnnotatedLlmResponse, CostSource, PricingResolver}; +use crate::codec::response::{ + AnnotatedLlmResponse, CostEstimate, CostSource, PricingResolver, Usage, +}; /// Maximum contributions retained for one LLM call. pub const MAX_LLM_OPTIMIZATION_CONTRIBUTIONS: usize = 64; @@ -59,6 +61,30 @@ impl LlmOptimizationRecorder { /// Rejection never affects LLM execution and does not consume a sequence. #[must_use] pub fn record(&self, mut contribution: LlmOptimizationContribution) -> bool { + if !self.reserve_record_attempt() || !self.payload_schema_is_valid(&contribution) { + return false; + } + // Relay always replaces producer-supplied identity. Serialization is + // deliberately outside the accumulator lock; if another writer wins + // the next sequence while we measure, retry with the new sequence. + contribution.id = Some(Uuid::now_v7()); + loop { + let Some(sequence) = self.next_contribution_sequence() else { + return false; + }; + contribution.sequence = Some(sequence); + let Some(contribution_bytes) = self.serialized_contribution_size(&contribution) else { + return false; + }; + match self.commit_contribution(contribution.clone(), contribution_bytes, sequence) { + ContributionCommit::Committed => return true, + ContributionCommit::Retry => continue, + ContributionCommit::Rejected => return false, + } + } + } + + fn reserve_record_attempt(&self) -> bool { let Ok(mut state) = self.state.lock() else { return false; }; @@ -70,83 +96,86 @@ impl LlmOptimizationRecorder { return false; } state.attempted_contributions += 1; - drop(state); - - match contribution.payload.as_ref() { - Some(_payload) if contribution.payload_schema.is_none() => { - if let Ok(mut state) = self.state.lock() - && !state.closed - { - state.invalid_payload_schema = true; - } - return false; - } - _ => {} + true + } + + fn payload_schema_is_valid(&self, contribution: &LlmOptimizationContribution) -> bool { + if contribution.payload.is_some() && contribution.payload_schema.is_none() { + self.note_invalid_payload_schema(); + return false; } + true + } - // Relay always replaces producer-supplied identity. Serialization is - // deliberately outside the accumulator lock; if another writer wins - // the next sequence while we measure, retry with the new sequence. - contribution.id = Some(Uuid::now_v7()); - loop { - let sequence = { - let Ok(state) = self.state.lock() else { - return false; - }; - if state.closed { - return false; - } - if state.contributions.len() >= MAX_LLM_OPTIMIZATION_CONTRIBUTIONS { - drop(state); - self.note_contribution_limit_exceeded(); - return false; - } - state.contributions.len() as u64 - }; - contribution.sequence = Some(sequence); + fn next_contribution_sequence(&self) -> Option { + let Ok(state) = self.state.lock() else { + return None; + }; + if state.closed { + return None; + } + if state.contributions.len() >= MAX_LLM_OPTIMIZATION_CONTRIBUTIONS { + drop(state); + self.note_contribution_limit_exceeded(); + return None; + } + Some(state.contributions.len() as u64) + } - let contribution_bytes = - match bounded_json_size(&contribution, MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES) { - Ok(size) => size, - Err(SerializedSizeError::LimitExceeded) => { - self.note_contribution_limit_exceeded(); - return false; - } - Err(SerializedSizeError::Serialization) => { - if let Ok(mut state) = self.state.lock() - && !state.closed - { - state.invalid_payload_schema = true; - } - return false; - } - }; - - let Ok(mut state) = self.state.lock() else { - return false; - }; - if state.closed { - return false; - } - if state.contributions.len() as u64 != sequence { - continue; + fn serialized_contribution_size( + &self, + contribution: &LlmOptimizationContribution, + ) -> Option { + match bounded_json_size(contribution, MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES) { + Ok(size) => Some(size), + Err(SerializedSizeError::LimitExceeded) => { + self.note_contribution_limit_exceeded(); + None } - let Some(total_contribution_bytes) = state - .total_contribution_bytes - .checked_add(contribution_bytes) - else { - seal_for_contribution_limit(&mut state); - return false; - }; - if total_contribution_bytes > MAX_LLM_OPTIMIZATION_TOTAL_CONTRIBUTION_BYTES { - seal_for_contribution_limit(&mut state); - return false; + Err(SerializedSizeError::Serialization) => { + self.note_invalid_payload_schema(); + None } + } + } - state.total_contribution_bytes = total_contribution_bytes; - state.contributions.push(contribution); - state.recorded_at.push(Utc::now()); - return true; + fn commit_contribution( + &self, + contribution: LlmOptimizationContribution, + contribution_bytes: usize, + sequence: u64, + ) -> ContributionCommit { + let Ok(mut state) = self.state.lock() else { + return ContributionCommit::Rejected; + }; + if state.closed { + return ContributionCommit::Rejected; + } + if state.contributions.len() as u64 != sequence { + return ContributionCommit::Retry; + } + let Some(total_contribution_bytes) = state + .total_contribution_bytes + .checked_add(contribution_bytes) + else { + seal_for_contribution_limit(&mut state); + return ContributionCommit::Rejected; + }; + if total_contribution_bytes > MAX_LLM_OPTIMIZATION_TOTAL_CONTRIBUTION_BYTES { + seal_for_contribution_limit(&mut state); + return ContributionCommit::Rejected; + } + state.total_contribution_bytes = total_contribution_bytes; + state.contributions.push(contribution); + state.recorded_at.push(Utc::now()); + ContributionCommit::Committed + } + + fn note_invalid_payload_schema(&self) { + if let Ok(mut state) = self.state.lock() + && !state.closed + { + state.invalid_payload_schema = true; } } @@ -280,6 +309,12 @@ impl LlmOptimizationRecorder { } } +enum ContributionCommit { + Committed, + Retry, + Rejected, +} + impl AccumulatorState { fn has_evidence(&self) -> bool { !self.contributions.is_empty() @@ -371,17 +406,15 @@ pub(crate) async fn scope_llm_optimization_recorder( .await } -pub(crate) fn finalize_optimization_summary( - recorder: &LlmOptimizationRecorder, - mut response: Option<&mut AnnotatedLlmResponse>, - requested_model: Option<&str>, - pricing: &PricingResolver, -) -> Option { - let finished = recorder.finish(); - if finished.contributions.is_empty() && finished.limitations.is_empty() { - return None; - } +struct ContributionAnalysis { + limitations: Vec, + token_totals: CheckedTokenTotals, + baseline_model: Option, + contributed_effective_model: Option, + use_effective_as_baseline: bool, +} +fn analyze_contributions(finished: &FinishedContributions) -> ContributionAnalysis { let applied_routing = finished .contributions .iter() @@ -391,12 +424,11 @@ pub(crate) fn finalize_optimization_summary( == crate::codec::optimization::LlmOptimizationKind::MODEL_ROUTING }) .collect::>(); - let mut limitations = finished.limitations; let routing_ambiguous = applied_routing.len() > 1; + let mut limitations = finished.limitations.clone(); if routing_ambiguous { limitations.push("multiple_routing_contributions".to_string()); } - let mut token_totals = CheckedTokenTotals::default(); for contribution in finished .contributions @@ -416,63 +448,63 @@ pub(crate) fn finalize_optimization_summary( token_totals.add_contribution(saved); } } - let mut token_count_overflow = token_totals.overflow.any(); if token_totals.missing_total { limitations.push("missing_token_savings_total".to_string()); } if token_totals.inconsistent_total { limitations.push("inconsistent_token_savings_total".to_string()); } - let tokens_saved = token_totals.values.clone(); - let authoritative_transition = (applied_routing.len() == 1) .then(|| applied_routing[0].model_transition.as_ref()) .flatten(); - let mut baseline_model = authoritative_transition.and_then(|route| route.baseline.clone()); - let contributed_effective_model = - authoritative_transition.and_then(|route| route.effective.clone()); - - // An applied routing contribution names the model Relay actually - // dispatched. Prefer it over provider response aliases or deployment - // names; fall back to response/request attribution when no router applies. - let effective_model = contributed_effective_model + ContributionAnalysis { + limitations, + token_totals, + baseline_model: authoritative_transition.and_then(|route| route.baseline.clone()), + contributed_effective_model: authoritative_transition + .and_then(|route| route.effective.clone()), + use_effective_as_baseline: applied_routing.is_empty() || routing_ambiguous, + } +} + +fn resolve_effective_model( + contributed: Option, + response: Option<&AnnotatedLlmResponse>, + requested_model: Option<&str>, +) -> Option { + contributed .or_else(|| { response - .as_ref() .and_then(|response| response.model.as_ref()) .map(|model| LlmOptimizationModel::new(model.clone())) }) - .or_else(|| requested_model.map(LlmOptimizationModel::new)); - if (applied_routing.is_empty() || routing_ambiguous) && baseline_model.is_none() { - baseline_model = effective_model.clone(); - } + .or_else(|| requested_model.map(LlmOptimizationModel::new)) +} - let mut effective_usage = response +struct UsageAnalysis { + effective: Option, + baseline: Option, + token_count_overflow: bool, + baseline_derivation_incomplete: bool, +} + +fn derive_optimization_usage( + mut response: Option<&mut AnnotatedLlmResponse>, + tokens_saved: &LlmOptimizationTokens, + token_totals: &CheckedTokenTotals, + limitations: &mut Vec, +) -> UsageAnalysis { + let mut effective = response .as_ref() .and_then(|response| response.usage.clone()); + let mut token_count_overflow = token_totals.overflow.any(); let mut baseline_derivation_incomplete = token_totals.missing_total || token_totals.inconsistent_total; - if let Some(usage) = effective_usage.as_mut() { - if usage.prompt_tokens.is_none() { - limitations.push("missing_effective_prompt_tokens".to_string()); - } - if usage.completion_tokens.is_none() { - limitations.push("missing_effective_completion_tokens".to_string()); - } - if usage.total_tokens.is_none() { - match (usage.prompt_tokens, usage.completion_tokens) { - (Some(prompt), Some(completion)) => match prompt.checked_add(completion) { - Some(total) => usage.total_tokens = Some(total), - None => token_count_overflow = true, - }, - _ => limitations.push("missing_effective_total_tokens".to_string()), - } - } + if let Some(usage) = effective.as_mut() { + note_missing_core_usage(usage, limitations, &mut token_count_overflow); } if let (Some(inferred), Some(response_usage)) = ( - effective_usage - .as_ref() - .and_then(|usage| usage.total_tokens), + effective.as_ref().and_then(|usage| usage.total_tokens), response .as_mut() .and_then(|response| response.usage.as_mut()), @@ -480,58 +512,121 @@ pub(crate) fn finalize_optimization_summary( { response_usage.total_tokens = Some(inferred); } - let baseline_usage = effective_usage.as_ref().map(|usage| { - let mut baseline = usage.clone(); - baseline.cost = None; - token_count_overflow |= checked_add_observed_tokens( + let baseline = effective.as_ref().map(|usage| { + derive_baseline_usage( + usage, + tokens_saved, + token_totals, + limitations, + &mut token_count_overflow, + &mut baseline_derivation_incomplete, + ) + }); + if token_count_overflow { + limitations.push("token_count_overflow".to_string()); + } + UsageAnalysis { + effective, + baseline, + token_count_overflow, + baseline_derivation_incomplete, + } +} + +fn note_missing_core_usage( + usage: &mut Usage, + limitations: &mut Vec, + token_count_overflow: &mut bool, +) { + if usage.prompt_tokens.is_none() { + limitations.push("missing_effective_prompt_tokens".to_string()); + } + if usage.completion_tokens.is_none() { + limitations.push("missing_effective_completion_tokens".to_string()); + } + if usage.total_tokens.is_none() { + match (usage.prompt_tokens, usage.completion_tokens) { + (Some(prompt), Some(completion)) => match prompt.checked_add(completion) { + Some(total) => usage.total_tokens = Some(total), + None => *token_count_overflow = true, + }, + _ => limitations.push("missing_effective_total_tokens".to_string()), + } + } +} + +fn derive_baseline_usage( + usage: &Usage, + tokens_saved: &LlmOptimizationTokens, + token_totals: &CheckedTokenTotals, + limitations: &mut Vec, + token_count_overflow: &mut bool, + baseline_derivation_incomplete: &mut bool, +) -> Usage { + let mut baseline = usage.clone(); + baseline.cost = None; + let fields = [ + ( &mut baseline.prompt_tokens, tokens_saved.prompt_tokens, token_totals.overflow.prompt, "missing_effective_prompt_tokens", - &mut limitations, - &mut baseline_derivation_incomplete, - ); - token_count_overflow |= checked_add_observed_tokens( + ), + ( &mut baseline.completion_tokens, tokens_saved.completion_tokens, token_totals.overflow.completion, "missing_effective_completion_tokens", - &mut limitations, - &mut baseline_derivation_incomplete, - ); - token_count_overflow |= checked_add_observed_tokens( + ), + ( &mut baseline.cache_read_tokens, tokens_saved.cache_read_tokens, token_totals.overflow.cache_read, "missing_effective_cache_read_tokens", - &mut limitations, - &mut baseline_derivation_incomplete, - ); - token_count_overflow |= checked_add_observed_tokens( + ), + ( &mut baseline.cache_write_tokens, tokens_saved.cache_write_tokens, token_totals.overflow.cache_write, "missing_effective_cache_write_tokens", - &mut limitations, - &mut baseline_derivation_incomplete, - ); - token_count_overflow |= checked_add_observed_tokens( + ), + ( &mut baseline.total_tokens, tokens_saved.total_tokens, token_totals.overflow.total, "missing_effective_total_tokens", - &mut limitations, - &mut baseline_derivation_incomplete, + ), + ]; + for (observed, saved, overflowed, missing_limitation) in fields { + *token_count_overflow |= checked_add_observed_tokens( + observed, + saved, + overflowed, + missing_limitation, + limitations, + baseline_derivation_incomplete, ); - baseline - }); - if token_count_overflow { - limitations.push("token_count_overflow".to_string()); } + baseline +} + +struct PricingAnalysis { + baseline_cost: Option, + actual_cost: Option, + complete_core_usage: bool, +} - // A provider-reported amount remains authoritative. A model-pricing - // estimate may have been calculated from a provider alias, so recompute - // it against the route Relay actually dispatched. +#[allow(clippy::too_many_arguments)] +fn price_optimization_usage( + effective_usage: &mut Option, + response: Option<&mut AnnotatedLlmResponse>, + effective_model: Option<&LlmOptimizationModel>, + baseline_model: Option<&LlmOptimizationModel>, + baseline_usage: Option<&Usage>, + token_count_overflow: bool, + baseline_derivation_incomplete: bool, + pricing: &PricingResolver, +) -> PricingAnalysis { let provider_reported_cost = effective_usage .as_ref() .and_then(|usage| usage.cost.as_ref()) @@ -541,35 +636,49 @@ pub(crate) fn finalize_optimization_summary( .as_ref() .is_some_and(|usage| usage.prompt_tokens.is_some() && usage.completion_tokens.is_some()); let actual_cost = provider_reported_cost.or_else(|| { - if !complete_core_usage { - return None; - } - let model = effective_model.as_ref()?; - let usage = effective_usage.as_ref()?; - pricing.estimate_cost_for_provider(model.provider.as_deref(), &model.model, usage) + let model = complete_core_usage.then_some(effective_model).flatten()?; + pricing.estimate_cost_for_provider( + model.provider.as_deref(), + &model.model, + effective_usage.as_ref()?, + ) }); if let Some(usage) = effective_usage.as_mut() { usage.cost.clone_from(&actual_cost); } - if let Some(usage) = response - .as_mut() - .and_then(|response| response.usage.as_mut()) - { + if let Some(usage) = response.and_then(|response| response.usage.as_mut()) { usage.cost.clone_from(&actual_cost); } - let baseline_cost = (!token_count_overflow && !baseline_derivation_incomplete && complete_core_usage) - .then_some(baseline_model.as_ref()) + .then_some(baseline_model) .flatten() .and_then(|model| { pricing.estimate_cost_for_provider( model.provider.as_deref(), &model.model, - baseline_usage.as_ref()?, + baseline_usage?, ) }); + PricingAnalysis { + baseline_cost, + actual_cost, + complete_core_usage, + } +} +#[allow(clippy::too_many_arguments)] +fn add_summary_limitations( + limitations: &mut Vec, + effective_usage: Option<&Usage>, + effective_model: Option<&LlmOptimizationModel>, + baseline_model: Option<&LlmOptimizationModel>, + baseline_cost: Option<&CostEstimate>, + actual_cost: Option<&CostEstimate>, + token_count_overflow: bool, + baseline_derivation_incomplete: bool, + complete_core_usage: bool, +) { if effective_usage.is_none() { limitations.push("missing_effective_usage".to_string()); } @@ -590,6 +699,68 @@ pub(crate) fn finalize_optimization_summary( if actual_cost.is_none() { limitations.push("missing_actual_cost".to_string()); } +} + +pub(crate) fn finalize_optimization_summary( + recorder: &LlmOptimizationRecorder, + mut response: Option<&mut AnnotatedLlmResponse>, + requested_model: Option<&str>, + pricing: &PricingResolver, +) -> Option { + let finished = recorder.finish(); + if finished.contributions.is_empty() && finished.limitations.is_empty() { + return None; + } + + let mut analysis = analyze_contributions(&finished); + let effective_model = resolve_effective_model( + analysis.contributed_effective_model.take(), + response.as_deref(), + requested_model, + ); + if analysis.use_effective_as_baseline && analysis.baseline_model.is_none() { + analysis.baseline_model = effective_model.clone(); + } + let baseline_model = analysis.baseline_model; + let tokens_saved = analysis.token_totals.values.clone(); + let mut limitations = analysis.limitations; + let token_totals = analysis.token_totals; + let mut token_count_overflow = token_totals.overflow.any(); + + let usage = derive_optimization_usage( + response.as_deref_mut(), + &tokens_saved, + &token_totals, + &mut limitations, + ); + let mut effective_usage = usage.effective; + let baseline_usage = usage.baseline; + token_count_overflow |= usage.token_count_overflow; + let baseline_derivation_incomplete = usage.baseline_derivation_incomplete; + + let pricing_analysis = price_optimization_usage( + &mut effective_usage, + response.as_deref_mut(), + effective_model.as_ref(), + baseline_model.as_ref(), + baseline_usage.as_ref(), + token_count_overflow, + baseline_derivation_incomplete, + pricing, + ); + let baseline_cost = pricing_analysis.baseline_cost; + let actual_cost = pricing_analysis.actual_cost; + add_summary_limitations( + &mut limitations, + effective_usage.as_ref(), + effective_model.as_ref(), + baseline_model.as_ref(), + baseline_cost.as_ref(), + actual_cost.as_ref(), + token_count_overflow, + baseline_derivation_incomplete, + pricing_analysis.complete_core_usage, + ); let (estimated_cost_saved, currency) = calculate_estimated_cost_saved( baseline_cost.as_ref(), diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index c3c6bf74c..7f62be147 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -72,6 +72,154 @@ pub(crate) fn default_mark_exclude_names() -> Vec { vec!["llm.chunk".to_string()] } +#[cfg(any(feature = "otel", feature = "openinference"))] +pub(crate) fn push_common_optimization_attributes( + attributes: &mut Vec, + summary: &crate::codec::optimization::LlmOptimizationSummary, +) { + push_optimization_models_and_tokens(attributes, summary); + push_optimization_cost(attributes, "baseline", summary.baseline_cost.as_ref()); + push_optimization_cost(attributes, "actual", summary.actual_cost.as_ref()); + push_optimization_savings_and_status(attributes, summary); + push_optimization_pricing_provenance(attributes, summary); +} + +#[cfg(any(feature = "otel", feature = "openinference"))] +fn push_optimization_models_and_tokens( + attributes: &mut Vec, + summary: &crate::codec::optimization::LlmOptimizationSummary, +) { + if let Some(model) = summary.baseline_model.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.baseline_model", + model.model.clone(), + )); + } + if let Some(model) = summary.effective_model.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.effective_model", + model.model.clone(), + )); + } + if let Some(tokens) = summary.tokens_saved.prompt_tokens { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.prompt_tokens_saved", + i64::try_from(tokens).unwrap_or(i64::MAX), + )); + } + if let Some(tokens) = summary.tokens_saved.total_tokens { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.total_tokens_saved", + i64::try_from(tokens).unwrap_or(i64::MAX), + )); + } +} + +#[cfg(any(feature = "otel", feature = "openinference"))] +fn push_optimization_cost( + attributes: &mut Vec, + label: &str, + cost: Option<&crate::codec::response::CostEstimate>, +) { + let Some(cost) = cost else { + return; + }; + if let Some(total) = cost.total_or_component_sum() { + attributes.push(opentelemetry::KeyValue::new( + format!("nemo_relay.llm.optimization.{label}_cost"), + total, + )); + } + attributes.push(opentelemetry::KeyValue::new( + format!("nemo_relay.llm.optimization.{label}_cost_currency"), + cost.currency.clone(), + )); + if let Some(source) = cost.pricing_source.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + format!("nemo_relay.llm.optimization.{label}_pricing_source"), + source.clone(), + )); + } + if let Some(as_of) = cost.pricing_as_of.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + format!("nemo_relay.llm.optimization.{label}_pricing_as_of"), + as_of.clone(), + )); + } +} + +#[cfg(any(feature = "otel", feature = "openinference"))] +fn push_optimization_savings_and_status( + attributes: &mut Vec, + summary: &crate::codec::optimization::LlmOptimizationSummary, +) { + if let Some(saved) = summary.estimated_cost_saved { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.estimated_cost_saved", + saved, + )); + if let Some(currency) = summary.currency.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.estimated_cost_saved_currency", + currency.clone(), + )); + } + } + if let Some(currency) = summary.currency.as_ref() { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.currency", + currency.clone(), + )); + } + let status = match summary.status { + crate::codec::optimization::LlmOptimizationSummaryStatus::Complete => "complete", + crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", + }; + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.status", + status, + )); +} + +#[cfg(any(feature = "otel", feature = "openinference"))] +fn push_optimization_pricing_provenance( + attributes: &mut Vec, + summary: &crate::codec::optimization::LlmOptimizationSummary, +) { + let source = summary + .baseline_cost + .as_ref() + .and_then(|cost| cost.pricing_source.as_ref()) + .or_else(|| { + summary + .actual_cost + .as_ref() + .and_then(|cost| cost.pricing_source.as_ref()) + }); + if let Some(source) = source { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.pricing_source", + source.clone(), + )); + } + let as_of = summary + .baseline_cost + .as_ref() + .and_then(|cost| cost.pricing_as_of.as_ref()) + .or_else(|| { + summary + .actual_cost + .as_ref() + .and_then(|cost| cost.pricing_as_of.as_ref()) + }); + if let Some(as_of) = as_of { + attributes.push(opentelemetry::KeyValue::new( + "nemo_relay.llm.optimization.pricing_as_of", + as_of.clone(), + )); + } +} + /// Validates OTLP attribute mappings shared by exporter configuration surfaces. pub fn validate_attribute_mappings( mappings: &[OtlpAttributeMapping], diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 4e3d9c1a0..b1690456c 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -1045,142 +1045,7 @@ fn push_optimization_attributes( attributes: &mut Vec, summary: &crate::codec::optimization::LlmOptimizationSummary, ) { - let string_fields = [ - ( - "nemo_relay.llm.optimization.baseline_model", - summary - .baseline_model - .as_ref() - .map(|model| model.model.clone()), - ), - ( - "nemo_relay.llm.optimization.effective_model", - summary - .effective_model - .as_ref() - .map(|model| model.model.clone()), - ), - ( - "nemo_relay.llm.optimization.currency", - summary.currency.clone(), - ), - ]; - for (key, value) in string_fields { - if let Some(value) = value { - attributes.push(KeyValue::new(key, value)); - } - } - if let Some(tokens) = summary.tokens_saved.prompt_tokens { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.prompt_tokens_saved", - i64::try_from(tokens).unwrap_or(i64::MAX), - )); - } - if let Some(tokens) = summary.tokens_saved.total_tokens { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.total_tokens_saved", - i64::try_from(tokens).unwrap_or(i64::MAX), - )); - } - if let Some(cost) = summary.baseline_cost.as_ref() { - if let Some(total) = cost.total_or_component_sum() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_cost", - total, - )); - } - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_cost_currency", - cost.currency.clone(), - )); - if let Some(source) = cost.pricing_source.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_pricing_source", - source.clone(), - )); - } - if let Some(as_of) = cost.pricing_as_of.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_pricing_as_of", - as_of.clone(), - )); - } - } - if let Some(cost) = summary.actual_cost.as_ref() { - if let Some(total) = cost.total_or_component_sum() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_cost", - total, - )); - } - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_cost_currency", - cost.currency.clone(), - )); - if let Some(source) = cost.pricing_source.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_pricing_source", - source.clone(), - )); - } - if let Some(as_of) = cost.pricing_as_of.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_pricing_as_of", - as_of.clone(), - )); - } - } - if let Some(saved) = summary.estimated_cost_saved { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.estimated_cost_saved", - saved, - )); - if let Some(currency) = summary.currency.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.estimated_cost_saved_currency", - currency.clone(), - )); - } - } - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.status", - match summary.status { - crate::codec::optimization::LlmOptimizationSummaryStatus::Complete => "complete", - crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", - }, - )); - if let Some(source) = summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.pricing_source.as_ref()) - .or_else(|| { - summary - .actual_cost - .as_ref() - .and_then(|cost| cost.pricing_source.as_ref()) - }) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.pricing_source", - source.clone(), - )); - } - if let Some(as_of) = summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.pricing_as_of.as_ref()) - .or_else(|| { - summary - .actual_cost - .as_ref() - .and_then(|cost| cost.pricing_as_of.as_ref()) - }) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.pricing_as_of", - as_of.clone(), - )); - } + crate::observability::push_common_optimization_attributes(attributes, summary); } fn push_annotated_input_messages(attributes: &mut Vec, messages: &[Message]) { diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index 671367036..aaff8b4f2 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -842,143 +842,7 @@ fn push_optimization_attributes( attributes: &mut Vec, summary: &crate::codec::optimization::LlmOptimizationSummary, ) { - if let Some(model) = summary.baseline_model.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_model", - model.model.clone(), - )); - } - if let Some(model) = summary.effective_model.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.effective_model", - model.model.clone(), - )); - } - if let Some(tokens) = summary.tokens_saved.prompt_tokens { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.prompt_tokens_saved", - i64::try_from(tokens).unwrap_or(i64::MAX), - )); - } - if let Some(tokens) = summary.tokens_saved.total_tokens { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.total_tokens_saved", - i64::try_from(tokens).unwrap_or(i64::MAX), - )); - } - if let Some(cost) = summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.total_or_component_sum()) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_cost", - cost, - )); - } - if let Some(cost) = summary.baseline_cost.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_cost_currency", - cost.currency.clone(), - )); - if let Some(source) = cost.pricing_source.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_pricing_source", - source.clone(), - )); - } - if let Some(as_of) = cost.pricing_as_of.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.baseline_pricing_as_of", - as_of.clone(), - )); - } - } - if let Some(cost) = summary - .actual_cost - .as_ref() - .and_then(|cost| cost.total_or_component_sum()) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_cost", - cost, - )); - } - if let Some(cost) = summary.actual_cost.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_cost_currency", - cost.currency.clone(), - )); - if let Some(source) = cost.pricing_source.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_pricing_source", - source.clone(), - )); - } - if let Some(as_of) = cost.pricing_as_of.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.actual_pricing_as_of", - as_of.clone(), - )); - } - } - if let Some(cost) = summary.estimated_cost_saved { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.estimated_cost_saved", - cost, - )); - if let Some(currency) = summary.currency.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.estimated_cost_saved_currency", - currency.clone(), - )); - } - } - if let Some(currency) = summary.currency.as_ref() { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.currency", - currency.clone(), - )); - } - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.status", - match summary.status { - crate::codec::optimization::LlmOptimizationSummaryStatus::Complete => "complete", - crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", - }, - )); - if let Some(source) = summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.pricing_source.as_ref()) - .or_else(|| { - summary - .actual_cost - .as_ref() - .and_then(|cost| cost.pricing_source.as_ref()) - }) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.pricing_source", - source.clone(), - )); - } - if let Some(as_of) = summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.pricing_as_of.as_ref()) - .or_else(|| { - summary - .actual_cost - .as_ref() - .and_then(|cost| cost.pricing_as_of.as_ref()) - }) - { - attributes.push(KeyValue::new( - "nemo_relay.llm.optimization.pricing_as_of", - as_of.clone(), - )); - } + crate::observability::push_common_optimization_attributes(attributes, summary); } fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 94f025390..041c57cb1 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -1944,47 +1944,68 @@ fn validate_atof_values( } let mut stream_sink_names = HashSet::new(); for (index, sink) in section.sinks.iter().enumerate() { - match sink { - AtofSinkSectionConfig::File(file) => { - if AtofExporterMode::parse(&file.mode).is_none() { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "observability.unsupported_value", - Some("atof".to_string()), - Some(format!("sinks[{index}].mode")), - format!("ATOF sinks[{index}].mode must be 'append' or 'overwrite'"), - ); - } - } - AtofSinkSectionConfig::Stream(stream) => { - if let Some(name) = stream.name.as_deref() { - let trimmed = name.trim(); - let message = if trimmed.is_empty() { - Some(format!("ATOF sinks[{index}].name must be non-empty")) - } else if name != trimmed { - Some(format!( - "ATOF sinks[{index}].name must not have leading or trailing whitespace" - )) - } else if !stream_sink_names.insert(name) { - Some(format!("ATOF stream sink name {name:?} must be unique")) - } else { - None - }; - if let Some(message) = message { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "observability.unsupported_value", - Some("atof".to_string()), - Some(format!("sinks[{index}].name")), - message, - ); - } - } - validate_atof_stream_sink_values(diagnostics, policy, index, stream); + validate_atof_sink(diagnostics, policy, index, sink, &mut stream_sink_names); + } +} + +fn validate_atof_sink<'a>( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + index: usize, + sink: &'a AtofSinkSectionConfig, + stream_sink_names: &mut HashSet<&'a str>, +) { + match sink { + AtofSinkSectionConfig::File(file) => { + if AtofExporterMode::parse(&file.mode).is_none() { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "observability.unsupported_value", + Some("atof".to_string()), + Some(format!("sinks[{index}].mode")), + format!("ATOF sinks[{index}].mode must be 'append' or 'overwrite'"), + ); } } + AtofSinkSectionConfig::Stream(stream) => { + validate_atof_stream_sink_name(diagnostics, policy, index, stream, stream_sink_names); + validate_atof_stream_sink_values(diagnostics, policy, index, stream); + } + } +} + +fn validate_atof_stream_sink_name<'a>( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + index: usize, + stream: &'a AtofStreamSinkSectionConfig, + stream_sink_names: &mut HashSet<&'a str>, +) { + let Some(name) = stream.name.as_deref() else { + return; + }; + let trimmed = name.trim(); + let message = if trimmed.is_empty() { + Some(format!("ATOF sinks[{index}].name must be non-empty")) + } else if name != trimmed { + Some(format!( + "ATOF sinks[{index}].name must not have leading or trailing whitespace" + )) + } else if !stream_sink_names.insert(name) { + Some(format!("ATOF stream sink name {name:?} must be unique")) + } else { + None + }; + if let Some(message) = message { + push_policy_diag( + diagnostics, + policy.unsupported_value, + "observability.unsupported_value", + Some("atof".to_string()), + Some(format!("sinks[{index}].name")), + message, + ); } } diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index 27817a82b..7759b6e3b 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -41,7 +41,6 @@ use crate::types::{ }; pub use crate::types::{nemo_relay_openinference_subscriber_free, nemo_relay_otel_subscriber_free}; use libc::c_char; -use nemo_relay::api::event::PendingMarkSpec; use nemo_relay::api::llm as core_llm_api; use nemo_relay::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; use nemo_relay::api::registry as core_registry_api; @@ -55,7 +54,6 @@ use nemo_relay::api::scope::ScopeAttributes; use nemo_relay::api::subscriber as core_subscriber_api; use nemo_relay::api::tool as core_tool_api; use nemo_relay::api::tool::ToolAttributes; -use nemo_relay::codec::optimization::LlmOptimizationContribution; use nemo_relay::error::Result as FlowResult; use nemo_relay::plugin::dynamic::{DynamicPluginActivationSpec, PluginHostActivation}; use nemo_relay::plugin::{ @@ -315,50 +313,24 @@ pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new_v2( set_last_error("request must be non-null"); return NemoRelayStatus::NullPointer; } - let annotated_request = if annotated_json.is_null() { - None - } else { - let value = match c_str_to_json(annotated_json) { - Some(value) => value, - None => return NemoRelayStatus::InvalidJson, - }; - match serde_json::from_value(value) { - Ok(value) => Some(value), - Err(error) => { - set_last_error(&format!("invalid annotated request JSON: {error}")); - return NemoRelayStatus::InvalidJson; - } - } - }; - let pending_marks = if pending_marks_json.is_null() { - Vec::new() - } else { - let value = match c_str_to_json(pending_marks_json) { - Some(value) => value, - None => return NemoRelayStatus::InvalidJson, - }; - match serde_json::from_value::>(value) { + let annotated_request = + match unsafe { parse_optional_intercept_json(annotated_json, "annotated request") } { Ok(value) => value, - Err(error) => { - set_last_error(&format!("invalid pending marks JSON: {error}")); - return NemoRelayStatus::InvalidJson; - } - } - }; - let optimization_contributions = if optimization_contributions_json.is_null() { - Vec::new() - } else { - let value = match c_str_to_json(optimization_contributions_json) { - Some(value) => value, - None => return NemoRelayStatus::InvalidJson, + Err(status) => return status, }; - match serde_json::from_value::>(value) { - Ok(value) => value, - Err(error) => { - set_last_error(&format!("invalid optimization contributions JSON: {error}")); - return NemoRelayStatus::InvalidJson; - } - } + let pending_marks = + match unsafe { parse_optional_intercept_json(pending_marks_json, "pending marks") } { + Ok(value) => value.unwrap_or_default(), + Err(status) => return status, + }; + let optimization_contributions = match unsafe { + parse_optional_intercept_json( + optimization_contributions_json, + "optimization contributions", + ) + } { + Ok(value) => value.unwrap_or_default(), + Err(status) => return status, }; let outcome = LlmRequestInterceptOutcome { request: unsafe { &*request }.0.clone(), @@ -378,6 +350,20 @@ pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new_v2( } } +unsafe fn parse_optional_intercept_json( + input: *const c_char, + description: &str, +) -> std::result::Result, NemoRelayStatus> { + if input.is_null() { + return Ok(None); + } + let value = c_str_to_json(input).ok_or(NemoRelayStatus::InvalidJson)?; + serde_json::from_value(value).map(Some).map_err(|error| { + set_last_error(&format!("invalid {description} JSON: {error}")); + NemoRelayStatus::InvalidJson + }) +} + /// Run the registered LLM conditional execution guardrail chain. /// /// Returns `NemoRelayStatus::Ok` if all guardrails pass, or From c12b4866a238c22ca4da8e5b33fadf6d936fbe17 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Wed, 15 Jul 2026 11:23:35 -0400 Subject: [PATCH 2/3] test: cover ATOF stream sink name validation Signed-off-by: Will Killian --- .../observability/plugin_component_tests.rs | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index 419f32121..41ea9e032 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -774,6 +774,65 @@ fn atof_endpoint_validation_rejects_bad_values() { ); } +#[test] +fn atof_stream_sink_name_validation_reports_each_invalid_name() { + let _guard = crate::observability::test_mutex().lock().unwrap(); + reset_runtime(); + + let report = validate_plugin_config(&plugin_config(json!({ + "atof": { + "enabled": true, + "sinks": [ + {"type": "stream", "url": "http://localhost/missing"}, + {"type": "stream", "name": "", "url": "http://localhost/empty"}, + {"type": "stream", "name": " leading", "url": "http://localhost/leading"}, + {"type": "stream", "name": "trailing ", "url": "http://localhost/trailing"}, + {"type": "stream", "name": "duplicate", "url": "http://localhost/first"}, + {"type": "stream", "name": "duplicate", "url": "http://localhost/second"}, + {"type": "stream", "name": "valid", "url": "http://localhost/valid"} + ] + } + }))); + + let name_diagnostics = report + .diagnostics + .iter() + .filter(|diagnostic| { + diagnostic + .field + .as_deref() + .is_some_and(|field| field.ends_with(".name")) + }) + .collect::>(); + assert_eq!(name_diagnostics.len(), 4); + assert!(name_diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("sinks[1].name") + && diagnostic.message.contains("must be non-empty") + })); + assert!(name_diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("sinks[2].name") + && diagnostic + .message + .contains("leading or trailing whitespace") + })); + assert!(name_diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("sinks[3].name") + && diagnostic + .message + .contains("leading or trailing whitespace") + })); + assert!(name_diagnostics.iter().any(|diagnostic| { + diagnostic.field.as_deref() == Some("sinks[5].name") + && diagnostic.message.contains("must be unique") + })); + assert!(!name_diagnostics.iter().any(|diagnostic| { + matches!( + diagnostic.field.as_deref(), + Some("sinks[0].name" | "sinks[4].name" | "sinks[6].name") + ) + })); +} + #[test] fn build_atof_sink_config_maps_headers_timeout_and_rejects_transport() { let mut headers = std::collections::HashMap::new(); From bd32f8f1d76bfba2923338bedcd3b85479cb97a5 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Wed, 15 Jul 2026 16:57:14 -0400 Subject: [PATCH 3/3] fix: finalize in-flight optimization records Signed-off-by: Will Killian --- crates/core/src/api/optimization.rs | 120 ++++++++++++++----- crates/core/tests/unit/optimization_tests.rs | 28 +++++ 2 files changed, 121 insertions(+), 27 deletions(-) diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index 4a1d940aa..183d8c542 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -4,7 +4,7 @@ //! Managed, bounded LLM optimization accounting. use std::collections::BTreeSet; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Condvar, Mutex}; use chrono::{DateTime, Utc}; use serde::Serialize; @@ -36,6 +36,7 @@ struct AccumulatorState { recorded_at: Vec>, total_contribution_bytes: usize, attempted_contributions: usize, + in_flight_records: usize, emitted: usize, closed: bool, finished: bool, @@ -44,13 +45,19 @@ struct AccumulatorState { invalid_payload_schema: bool, } +#[derive(Debug, Default)] +struct Accumulator { + state: Mutex, + records_settled: Condvar, +} + /// Cloneable capability for adding evidence to the current managed LLM call. /// /// A streaming execution intercept may capture this value before returning its /// stream and use it when the route is committed by the first upstream item. #[derive(Debug, Clone, Default)] pub struct LlmOptimizationRecorder { - state: Arc>, + state: Arc, } impl LlmOptimizationRecorder { @@ -60,23 +67,33 @@ impl LlmOptimizationRecorder { /// invariant, a per-call bound, or because accounting has already closed. /// Rejection never affects LLM execution and does not consume a sequence. #[must_use] - pub fn record(&self, mut contribution: LlmOptimizationContribution) -> bool { - if !self.reserve_record_attempt() || !self.payload_schema_is_valid(&contribution) { + pub fn record(&self, contribution: LlmOptimizationContribution) -> bool { + let Some(_attempt) = self.reserve_record_attempt() else { + return false; + }; + if !self.payload_schema_is_valid(&contribution) { return false; } // Relay always replaces producer-supplied identity. Serialization is // deliberately outside the accumulator lock; if another writer wins // the next sequence while we measure, retry with the new sequence. - contribution.id = Some(Uuid::now_v7()); + let mut contribution = Some(contribution); + let Some(record) = contribution.as_mut() else { + return false; + }; + record.id = Some(Uuid::now_v7()); loop { let Some(sequence) = self.next_contribution_sequence() else { return false; }; - contribution.sequence = Some(sequence); - let Some(contribution_bytes) = self.serialized_contribution_size(&contribution) else { + let Some(record) = contribution.as_mut() else { return false; }; - match self.commit_contribution(contribution.clone(), contribution_bytes, sequence) { + record.sequence = Some(sequence); + let Some(contribution_bytes) = self.serialized_contribution_size(record) else { + return false; + }; + match self.commit_contribution(&mut contribution, contribution_bytes, sequence) { ContributionCommit::Committed => return true, ContributionCommit::Retry => continue, ContributionCommit::Rejected => return false, @@ -84,19 +101,27 @@ impl LlmOptimizationRecorder { } } - fn reserve_record_attempt(&self) -> bool { - let Ok(mut state) = self.state.lock() else { - return false; + fn reserve_record_attempt(&self) -> Option { + let Ok(mut state) = self.state.state.lock() else { + return None; }; if state.closed { - return false; + return None; } if state.attempted_contributions >= MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS { seal_for_contribution_limit(&mut state); - return false; + return None; } state.attempted_contributions += 1; - true + state.in_flight_records += 1; + Some(RecordAttempt { + state: Arc::clone(&self.state), + }) + } + + #[cfg(test)] + pub(crate) fn reserve_record_attempt_for_test(&self) -> Option { + self.reserve_record_attempt() } fn payload_schema_is_valid(&self, contribution: &LlmOptimizationContribution) -> bool { @@ -108,7 +133,7 @@ impl LlmOptimizationRecorder { } fn next_contribution_sequence(&self) -> Option { - let Ok(state) = self.state.lock() else { + let Ok(state) = self.state.state.lock() else { return None; }; if state.closed { @@ -141,11 +166,11 @@ impl LlmOptimizationRecorder { fn commit_contribution( &self, - contribution: LlmOptimizationContribution, + contribution: &mut Option, contribution_bytes: usize, sequence: u64, ) -> ContributionCommit { - let Ok(mut state) = self.state.lock() else { + let Ok(mut state) = self.state.state.lock() else { return ContributionCommit::Rejected; }; if state.closed { @@ -165,6 +190,9 @@ impl LlmOptimizationRecorder { seal_for_contribution_limit(&mut state); return ContributionCommit::Rejected; } + let Some(contribution) = contribution.take() else { + return ContributionCommit::Rejected; + }; state.total_contribution_bytes = total_contribution_bytes; state.contributions.push(contribution); state.recorded_at.push(Utc::now()); @@ -172,16 +200,16 @@ impl LlmOptimizationRecorder { } fn note_invalid_payload_schema(&self) { - if let Ok(mut state) = self.state.lock() - && !state.closed + if let Ok(mut state) = self.state.state.lock() + && !state.finished { state.invalid_payload_schema = true; } } fn note_contribution_limit_exceeded(&self) { - if let Ok(mut state) = self.state.lock() - && !state.closed + if let Ok(mut state) = self.state.state.lock() + && !state.finished { seal_for_contribution_limit(&mut state); } @@ -199,7 +227,16 @@ impl LlmOptimizationRecorder { } fn is_closed(&self) -> bool { - self.state.lock().map(|state| state.closed).unwrap_or(true) + self.state + .state + .lock() + .map(|state| state.closed) + .unwrap_or(true) + } + + #[cfg(test)] + pub(crate) fn is_closed_for_test(&self) -> bool { + self.is_closed() } /// Snapshot contributions not yet accepted by mark delivery. @@ -222,7 +259,7 @@ impl LlmOptimizationRecorder { pub(crate) fn unemitted_with_timestamps( &self, ) -> Vec<(LlmOptimizationContribution, DateTime)> { - let Ok(state) = self.state.lock() else { + let Ok(state) = self.state.state.lock() else { return Vec::new(); }; let start = state.emitted.min(state.contributions.len()); @@ -235,7 +272,7 @@ impl LlmOptimizationRecorder { /// Advance the delivery cursor for a bounded number of accepted marks. pub(crate) fn mark_emitted(&self, count: usize) { - let Ok(mut state) = self.state.lock() else { + let Ok(mut state) = self.state.state.lock() else { return; }; state.emitted = state @@ -247,7 +284,7 @@ impl LlmOptimizationRecorder { /// Add a best-effort lifecycle limitation to the eventual summary. #[cfg(test)] pub(crate) fn note_limitation(&self, limitation: impl Into) { - if let Ok(mut state) = self.state.lock() + if let Ok(mut state) = self.state.state.lock() && !state.closed { state.limitations.insert(limitation.into()); @@ -261,7 +298,7 @@ impl LlmOptimizationRecorder { /// an interrupted but otherwise unoptimized stream from manufacturing an /// optimization summary. pub(crate) fn close_for_finalization(&self, conditional_limitation: Option<&str>) -> bool { - let Ok(mut state) = self.state.lock() else { + let Ok(mut state) = self.state.state.lock() else { return false; }; if state.finished { @@ -276,7 +313,7 @@ impl LlmOptimizationRecorder { } fn finish(&self) -> FinishedContributions { - let Ok(mut state) = self.state.lock() else { + let Ok(mut state) = self.state.state.lock() else { return FinishedContributions { contributions: Vec::new(), limitations: vec!["optimization_accumulator_unavailable".to_string()], @@ -289,6 +326,21 @@ impl LlmOptimizationRecorder { }; } state.closed = true; + while state.in_flight_records > 0 { + let Ok(waiting) = self.state.records_settled.wait(state) else { + return FinishedContributions { + contributions: Vec::new(), + limitations: vec!["optimization_accumulator_unavailable".to_string()], + }; + }; + state = waiting; + if state.finished { + return FinishedContributions { + contributions: Vec::new(), + limitations: Vec::new(), + }; + } + } state.finished = true; let mut limitations = std::mem::take(&mut state.limitations) .into_iter() @@ -315,6 +367,20 @@ enum ContributionCommit { Rejected, } +pub(crate) struct RecordAttempt { + state: Arc, +} + +impl Drop for RecordAttempt { + fn drop(&mut self) { + let Ok(mut state) = self.state.state.lock() else { + return; + }; + state.in_flight_records = state.in_flight_records.saturating_sub(1); + self.state.records_settled.notify_all(); + } +} + impl AccumulatorState { fn has_evidence(&self) -> bool { !self.contributions.is_empty() diff --git a/crates/core/tests/unit/optimization_tests.rs b/crates/core/tests/unit/optimization_tests.rs index ff5291719..8b993c124 100644 --- a/crates/core/tests/unit/optimization_tests.rs +++ b/crates/core/tests/unit/optimization_tests.rs @@ -5,6 +5,8 @@ use serde_json::json; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc; +use std::time::{Duration, Instant}; use super::*; use crate::api::event::DataSchema; @@ -730,6 +732,32 @@ fn concurrent_finish_is_an_atomic_acceptance_boundary() { assert!(!recorder.record(LlmOptimizationContribution::new("late", "custom"))); } +#[test] +fn finalization_waits_for_reserved_record_attempts() { + let recorder = LlmOptimizationRecorder::default(); + let attempt = recorder.reserve_record_attempt_for_test().unwrap(); + let finalizer = recorder.clone(); + let (sender, receiver) = mpsc::channel(); + std::thread::spawn(move || { + sender.send(finalizer.finish()).unwrap(); + }); + + let deadline = Instant::now() + Duration::from_secs(1); + while !recorder.is_closed_for_test() && Instant::now() < deadline { + std::thread::yield_now(); + } + assert!(recorder.is_closed_for_test()); + assert!(matches!( + receiver.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + + drop(attempt); + let finished = receiver.recv_timeout(Duration::from_secs(1)).unwrap(); + assert!(finished.contributions.is_empty()); + assert!(finished.limitations.is_empty()); +} + #[tokio::test] async fn recorder_task_local_is_not_implicitly_inherited_by_spawned_tasks() { let recorder = LlmOptimizationRecorder::default();