From 8ed542a6efd684085226374991c8ace17d339e0c Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Wed, 8 Jul 2026 12:50:04 -0600 Subject: [PATCH 01/10] feat: add extensible LLM optimization accounting Signed-off-by: Bryan Bednarski --- crates/core/src/api/llm.rs | 121 +++- crates/core/src/api/mod.rs | 2 + crates/core/src/api/optimization.rs | 648 ++++++++++++++++++ crates/core/src/api/runtime/state.rs | 3 + crates/core/src/api/shared.rs | 28 +- crates/core/src/codec/anthropic.rs | 1 + crates/core/src/codec/mod.rs | 1 + crates/core/src/codec/openai_chat.rs | 1 + crates/core/src/codec/openai_responses.rs | 1 + crates/core/src/codec/optimization.rs | 6 + crates/core/src/observability/atif.rs | 67 +- .../core/src/observability/openinference.rs | 94 +++ crates/core/src/observability/otel.rs | 90 +++ crates/core/src/stream.rs | 28 +- .../tests/integration/middleware_tests.rs | 118 +++- .../tests/integration/native_plugin_tests.rs | 1 + .../core/tests/integration/pipeline_tests.rs | 1 + crates/core/tests/unit/atif_tests.rs | 63 ++ .../core/tests/unit/codec/response_tests.rs | 3 + .../unit/observability/openinference_tests.rs | 36 + .../tests/unit/observability/otel_tests.rs | 34 + crates/core/tests/unit/shared_tests.rs | 12 +- crates/core/tests/unit/types_tests.rs | 1 + crates/ffi/src/api/mod.rs | 1 + crates/node/src/api/mod.rs | 3 +- crates/node/src/callable.rs | 4 + crates/plugin/src/lib.rs | 5 + crates/python/src/py_types/codecs.rs | 1 + crates/python/src/py_types/core.rs | 30 +- crates/types/src/api/llm.rs | 15 + crates/types/src/codec/mod.rs | 2 + crates/types/src/codec/optimization.rs | 353 ++++++++++ crates/types/src/codec/response.rs | 5 + crates/types/tests/serialization_tests.rs | 1 + 34 files changed, 1739 insertions(+), 41 deletions(-) create mode 100644 crates/core/src/api/optimization.rs create mode 100644 crates/core/src/codec/optimization.rs create mode 100644 crates/types/src/codec/optimization.rs diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index a3d823fc4..f493ce636 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -9,7 +9,12 @@ use serde_json::json; use typed_builder::TypedBuilder; use uuid::Uuid; -use crate::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; +use crate::api::event::{ + BaseEvent, CategoryProfile, DataSchema, Event, EventCategory, MarkEvent, PendingMarkSpec, +}; +use crate::api::optimization::{ + LlmOptimizationRecorder, finalize_optimization_summary, scope_llm_optimization_recorder, +}; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::current_scope_stack; use crate::api::runtime::global_context; @@ -65,6 +70,10 @@ pub struct LlmHandle { /// Optional normalized model name for observability. #[builder(default, setter(into))] pub model_name: Option, + /// Bounded, in-memory optimization evidence recorder for this call. + #[serde(skip, default)] + #[builder(default)] + pub optimization_recorder: LlmOptimizationRecorder, } /// Builder parameters for [`NemoRelayContextState::create_llm_handle`]. @@ -356,6 +365,42 @@ fn emit_pending_request_marks( Ok(()) } +pub(crate) fn emit_optimization_marks( + handle: &LlmHandle, + subscribers: &[EventSubscriberFn], +) -> Result<()> { + let contributions = handle.optimization_recorder.take_unemitted(); + if contributions.is_empty() { + return Ok(()); + } + ensure_runtime_owner()?; + for contribution in contributions { + let offset = contribution.sequence.unwrap_or(0).saturating_add(2); + let offset = i64::try_from(offset).unwrap_or(i64::MAX); + let data = serde_json::to_value(&contribution).unwrap_or(Json::Null); + let event = Event::Mark(MarkEvent::new( + BaseEvent::builder() + .name("nemo_relay.llm.optimization") + .parent_uuid(handle.uuid) + .timestamp(handle.started_at + TimeDelta::microseconds(offset)) + .data(data) + .data_schema(DataSchema { + name: "nemo.relay.llm_optimization_contribution".to_string(), + version: "1".to_string(), + }) + .build(), + Some(EventCategory::custom()), + Some( + CategoryProfile::builder() + .subtype("nemo_relay.llm.optimization") + .build(), + ), + )); + NemoRelayContextState::emit_event(&event, subscribers); + } + Ok(()) +} + /// Start a manual LLM lifecycle span. /// /// This emits an LLM-start event after applying sanitize-request guardrails to @@ -490,15 +535,15 @@ fn llm_call_end_with_behavior( Some(sanitized_response) }; let mut decode_error = None; - let annotated_response = match annotated_response { - Some(annotated_response) => Some(annotated_response), + 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(Arc::new(decoded)) + Some(decoded) } Err(error) => { decode_error = Some(error); @@ -508,6 +553,23 @@ fn llm_call_end_with_behavior( _ => None, }, }; + emit_optimization_marks(handle, &subscribers)?; + let pricing = crate::codec::response::active_pricing_resolver(); + let summary = finalize_optimization_summary( + &handle.optimization_recorder, + annotated_response.as_mut(), + handle.model_name.as_deref(), + &pricing, + ); + if annotated_response.is_none() + && let Some(summary) = summary + { + annotated_response = Some(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }); + } + let annotated_response = annotated_response.map(Arc::new); let event = { let context = global_context(); let state = context @@ -542,20 +604,35 @@ fn emit_llm_end_without_output( lifecycle_subscribers: Option<&[EventSubscriberFn]>, ) -> Result<()> { ensure_runtime_owner()?; - let (event, subscribers) = { + let subscribers = { let scope_stack = current_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); let scope_subscribers = scope_guard.collect_scope_local_subscribers(); - let subscribers = match lifecycle_subscribers { + match lifecycle_subscribers { Some(subscribers) => subscribers.to_vec(), None => snapshot_event_subscribers(scope_subscribers)?, - }; + } + }; + emit_optimization_marks(handle, &subscribers)?; + let pricing = crate::codec::response::active_pricing_resolver(); + let annotated_response = finalize_optimization_summary( + &handle.optimization_recorder, + None, + handle.model_name.as_deref(), + &pricing, + ) + .map(|summary| { + Arc::new(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }) + }); + let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - let event = state.end_llm_handle(handle, handle.data.clone(), metadata, None); - (event, subscribers) + state.end_llm_handle(handle, handle.data.clone(), metadata, annotated_response) }; if let Some(event) = sanitize_event(event) { NemoRelayContextState::emit_event(&event, &subscribers); @@ -661,7 +738,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { } let request_codec = codec.clone(); - let (intercepted_request, annotated_request, pending_marks) = + let (intercepted_request, annotated_request, pending_marks, optimization_contributions) = run_request_intercepts_with_codec(&name, request, codec)?; let handle = create_llm_handle( @@ -687,6 +764,10 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; + handle + .optimization_recorder + .record_all(optimization_contributions); + emit_optimization_marks(&handle, &lifecycle_subscribers)?; let execution = { let scope_stack = current_scope_stack(); @@ -700,7 +781,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { state.llm_build_execution_chain(&name, func, &scope_locals) }; - match execution(intercepted_request).await { + match scope_llm_optimization_recorder( + handle.optimization_recorder.clone(), + execution(intercepted_request), + ) + .await + { Ok(response) => { llm_call_end_with_behavior( LlmCallEndParams::builder() @@ -826,7 +912,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu } let request_codec = codec.clone(); - let (intercepted_request, annotated_request, pending_marks) = + let (intercepted_request, annotated_request, pending_marks, optimization_contributions) = run_request_intercepts_with_codec(&name, request, codec)?; let handle = create_llm_handle( @@ -852,6 +938,10 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; + handle + .optimization_recorder + .record_all(optimization_contributions); + emit_optimization_marks(&handle, &lifecycle_subscribers)?; let execution = { let scope_stack = current_scope_stack(); @@ -866,7 +956,12 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu state.llm_stream_build_execution_chain(&name, func, &scope_locals) }; - match execution(intercepted_request).await { + match scope_llm_optimization_recorder( + handle.optimization_recorder.clone(), + execution(intercepted_request), + ) + .await + { Ok(raw_stream) => { let wrapper = LlmStreamWrapper::new_managed( raw_stream, diff --git a/crates/core/src/api/mod.rs b/crates/core/src/api/mod.rs index 1932063c0..f8a914e45 100644 --- a/crates/core/src/api/mod.rs +++ b/crates/core/src/api/mod.rs @@ -7,6 +7,8 @@ pub mod event; /// LLM lifecycle helpers and managed execution entry points. pub mod llm; +/// Plugin-neutral evidence recording for managed LLM calls. +pub mod optimization; /// Global and scope-local middleware registration helpers. pub mod registry; /// Advanced runtime state, callbacks, and scope-stack helpers. diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs new file mode 100644 index 000000000..92f0b6123 --- /dev/null +++ b/crates/core/src/api/optimization.rs @@ -0,0 +1,648 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Managed, bounded LLM optimization accounting. + +use std::sync::{Arc, Mutex}; + +use uuid::Uuid; + +use crate::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationModel, LlmOptimizationSummary, + LlmOptimizationSummaryStatus, LlmOptimizationTokens, +}; +use crate::codec::response::{AnnotatedLlmResponse, PricingResolver}; + +/// Maximum contributions retained for one LLM call. +pub const MAX_LLM_OPTIMIZATION_CONTRIBUTIONS: usize = 64; +/// Maximum serialized custom payload size for one contribution. +pub const MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES: usize = 16 * 1024; +/// Maximum aggregate serialized custom payload size for one LLM call. +pub const MAX_LLM_OPTIMIZATION_TOTAL_PAYLOAD_BYTES: usize = 256 * 1024; + +#[derive(Debug, Default)] +struct AccumulatorState { + contributions: Vec, + total_payload_bytes: usize, + emitted: usize, + contribution_limit_exceeded: bool, + invalid_payload_schema: bool, +} + +/// 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>, +} + +impl LlmOptimizationRecorder { + /// Record one contribution without blocking on I/O or exporter delivery. + /// + /// Returns `false` when the contribution is rejected by a payload/schema + /// invariant or a per-call bound. Rejection never affects LLM execution. + #[must_use] + pub fn record(&self, mut contribution: LlmOptimizationContribution) -> bool { + let payload_bytes = match contribution.payload.as_ref() { + Some(_payload) if contribution.payload_schema.is_none() => { + if let Ok(mut state) = self.state.lock() { + state.invalid_payload_schema = true; + } + return false; + } + Some(payload) => match serde_json::to_vec(payload) { + Ok(serialized) => serialized.len(), + Err(_) => { + if let Ok(mut state) = self.state.lock() { + state.invalid_payload_schema = true; + } + return false; + } + }, + None => 0, + }; + + let Ok(mut state) = self.state.lock() else { + return false; + }; + if state.contributions.len() >= MAX_LLM_OPTIMIZATION_CONTRIBUTIONS + || payload_bytes > MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES + || state.total_payload_bytes.saturating_add(payload_bytes) + > MAX_LLM_OPTIMIZATION_TOTAL_PAYLOAD_BYTES + { + state.contribution_limit_exceeded = true; + return false; + } + + contribution.id = Some(Uuid::now_v7()); + contribution.sequence = Some(state.contributions.len() as u64); + state.total_payload_bytes += payload_bytes; + state.contributions.push(contribution); + true + } + + pub(crate) fn record_all( + &self, + contributions: impl IntoIterator, + ) { + for contribution in contributions { + let _ = self.record(contribution); + } + } + + pub(crate) fn take_unemitted(&self) -> Vec { + let Ok(mut state) = self.state.lock() else { + return Vec::new(); + }; + let start = state.emitted.min(state.contributions.len()); + let contributions = state.contributions[start..].to_vec(); + state.emitted = state.contributions.len(); + contributions + } + + fn finish(&self) -> FinishedContributions { + let Ok(mut state) = self.state.lock() else { + return FinishedContributions { + contributions: Vec::new(), + limitations: vec!["optimization_accumulator_unavailable".to_string()], + }; + }; + let mut limitations = Vec::new(); + if state.contribution_limit_exceeded { + limitations.push("contribution_limit_exceeded".to_string()); + } + if state.invalid_payload_schema { + limitations.push("invalid_contribution_payload_schema".to_string()); + } + FinishedContributions { + contributions: std::mem::take(&mut state.contributions), + limitations, + } + } +} + +struct FinishedContributions { + contributions: Vec, + limitations: Vec, +} + +tokio::task_local! { + static CURRENT_LLM_OPTIMIZATION_RECORDER: LlmOptimizationRecorder; +} + +/// Return a recorder for the current execution intercept, if it is managed by Relay. +#[must_use] +pub fn current_llm_optimization_recorder() -> Option { + CURRENT_LLM_OPTIMIZATION_RECORDER + .try_with(Clone::clone) + .ok() +} + +/// Best-effort shorthand for recording evidence on the current managed call. +#[must_use] +pub fn record_llm_optimization_contribution(contribution: LlmOptimizationContribution) -> bool { + current_llm_optimization_recorder().is_some_and(|recorder| recorder.record(contribution)) +} + +pub(crate) async fn scope_llm_optimization_recorder( + recorder: LlmOptimizationRecorder, + future: F, +) -> F::Output { + CURRENT_LLM_OPTIMIZATION_RECORDER + .scope(recorder, future) + .await +} + +pub(crate) fn finalize_optimization_summary( + recorder: &LlmOptimizationRecorder, + 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 tokens_saved = LlmOptimizationTokens::default(); + let mut baseline_model = None; + let mut contributed_effective_model = None; + for contribution in finished + .contributions + .iter() + .filter(|contribution| contribution.applied) + { + if let Some(saved) = contribution + .token_impact + .as_ref() + .and_then(|impact| impact.saved.as_ref()) + { + tokens_saved.add_assign(saved); + } + if contribution.kind.as_str() + == crate::codec::optimization::LlmOptimizationKind::MODEL_ROUTING + && let Some(transition) = contribution.model_transition.as_ref() + { + if baseline_model.is_none() { + baseline_model = transition.baseline.clone(); + } + if transition.effective.is_some() { + contributed_effective_model = transition.effective.clone(); + } + } + } + + let effective_model = response + .as_ref() + .and_then(|response| response.model.as_ref()) + .map(|model| LlmOptimizationModel::new(model.clone())) + .or(contributed_effective_model) + .or_else(|| requested_model.map(LlmOptimizationModel::new)); + if baseline_model.is_none() { + baseline_model = effective_model.clone(); + } + + let effective_usage = response + .as_ref() + .and_then(|response| response.usage.clone()); + let baseline_usage = effective_usage.as_ref().map(|usage| { + let mut baseline = usage.clone(); + baseline.cost = None; + add_tokens(&mut baseline.prompt_tokens, tokens_saved.prompt_tokens); + add_tokens( + &mut baseline.completion_tokens, + tokens_saved.completion_tokens, + ); + add_tokens( + &mut baseline.cache_read_tokens, + tokens_saved.cache_read_tokens, + ); + add_tokens( + &mut baseline.cache_write_tokens, + tokens_saved.cache_write_tokens, + ); + let total_saved = tokens_saved + .total_tokens + .or_else(|| option_sum([tokens_saved.prompt_tokens, tokens_saved.completion_tokens])); + add_tokens(&mut baseline.total_tokens, total_saved); + baseline + }); + + let actual_cost = effective_usage + .as_ref() + .and_then(|usage| usage.cost.clone()) + .or_else(|| { + 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 baseline_cost = baseline_model.as_ref().and_then(|model| { + pricing.estimate_cost_for_provider( + model.provider.as_deref(), + &model.model, + baseline_usage.as_ref()?, + ) + }); + + let mut limitations = finished.limitations; + if effective_usage.is_none() { + limitations.push("missing_effective_usage".to_string()); + } + if baseline_model.is_none() { + limitations.push("missing_baseline_model".to_string()); + } + if baseline_cost.is_none() { + limitations.push("missing_baseline_pricing".to_string()); + } + if actual_cost.is_none() { + limitations.push("missing_actual_cost".to_string()); + } + + let (estimated_cost_saved, currency) = match (&baseline_cost, &actual_cost) { + (Some(baseline), Some(actual)) + if baseline.currency.eq_ignore_ascii_case(&actual.currency) => + { + ( + baseline + .total_or_component_sum() + .zip(actual.total_or_component_sum()) + .map(|(baseline, actual)| baseline - actual), + Some(baseline.currency.clone()), + ) + } + (Some(_), Some(_)) => { + limitations.push("cost_currency_mismatch".to_string()); + (None, None) + } + _ => (None, None), + }; + + limitations.sort(); + limitations.dedup(); + let summary = LlmOptimizationSummary { + schema_version: "1".to_string(), + calculation_version: "1".to_string(), + status: if limitations.is_empty() { + LlmOptimizationSummaryStatus::Complete + } else { + LlmOptimizationSummaryStatus::Partial + }, + limitations, + baseline_model, + effective_model, + effective_usage, + baseline_usage, + tokens_saved, + baseline_cost, + actual_cost, + estimated_cost_saved, + currency, + contributions: finished.contributions, + }; + if let Some(response) = response { + response.optimization_summary = Some(summary.clone()); + } + Some(summary) +} + +fn add_tokens(target: &mut Option, value: Option) { + if let Some(value) = value { + *target = Some(target.unwrap_or(0).saturating_add(value)); + } +} + +fn option_sum(values: impl IntoIterator>) -> Option { + let mut present = false; + let total = values.into_iter().flatten().fold(0_u64, |total, value| { + present = true; + total.saturating_add(value) + }); + present.then_some(total) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::api::event::DataSchema; + use crate::codec::optimization::{ + LlmOptimizationEvidenceQuality, LlmOptimizationModelTransition, LlmOptimizationTokenImpact, + }; + use crate::codec::response::{PricingCatalog, Usage}; + use crate::json::Json; + + fn resolver() -> PricingResolver { + resolver_with_rates(2.0, 1.0) + } + + fn resolver_with_rates(baseline_input: f64, effective_input: f64) -> PricingResolver { + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + {"provider":"test","model_id":"baseline","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":baseline_input,"output_per_million":4.0,"cache_read_per_million":0.5,"cache_write_per_million":3.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}}, + {"provider":"test","model_id":"effective","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":effective_input,"output_per_million":2.0,"cache_read_per_million":0.25,"cache_write_per_million":2.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}} + ] + }) + .to_string(), + ) + .unwrap(); + PricingResolver::from_catalogs(vec![catalog]) + } + + fn contribution() -> LlmOptimizationContribution { + let mut contribution = LlmOptimizationContribution::new( + "test.optimizer", + crate::codec::optimization::LlmOptimizationKind::model_routing(), + ); + contribution.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new("baseline").with_provider("test")), + effective: Some(LlmOptimizationModel::new("effective").with_provider("test")), + }); + contribution.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(200)), + quality: Some(LlmOptimizationEvidenceQuality::Estimated), + estimation_method: Some("test-tokenizer".to_string()), + ..LlmOptimizationTokenImpact::default() + }); + contribution + } + + #[test] + fn combined_summary_retains_token_evidence_and_snapshot_pricing() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + Some("baseline"), + &resolver(), + ) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Complete); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().prompt_tokens, + Some(1000) + ); + assert_eq!(summary.baseline_cost.as_ref().unwrap().total, Some(0.0024)); + assert_eq!(summary.actual_cost.as_ref().unwrap().total, Some(0.001)); + assert!((summary.estimated_cost_saved.unwrap() - 0.0014).abs() < 1e-12); + } + + #[test] + fn unpriced_summary_is_partial_without_losing_tokens() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(8), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + None, + &PricingResolver::default(), + ) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); + assert!(summary.estimated_cost_saved.is_none()); + } + + #[test] + fn zero_and_negative_savings_are_preserved() { + for (baseline_rate, effective_rate, expected_sign) in [(0.0, 0.0, 0_i8), (0.5, 2.0, -1_i8)] + { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + total_tokens: Some(800), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + None, + &resolver_with_rates(baseline_rate, effective_rate), + ) + .unwrap(); + let saved = summary.estimated_cost_saved.unwrap(); + match expected_sign { + 0 => assert_eq!(saved, 0.0), + -1 => assert!(saved < 0.0), + _ => unreachable!(), + } + } + } + + #[test] + fn multiple_contributions_and_cache_savings_aggregate_explicitly() { + let recorder = LlmOptimizationRecorder::default(); + for (producer, prompt, cache_read, cache_write) in + [("test.one", 5, 7, 0), ("test.two", 11, 13, 17)] + { + let mut item = LlmOptimizationContribution::new( + producer, + crate::codec::optimization::LlmOptimizationKind::input_compression(), + ); + item.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens { + prompt_tokens: Some(prompt), + cache_read_tokens: Some(cache_read), + cache_write_tokens: Some(cache_write), + total_tokens: Some(prompt), + ..LlmOptimizationTokens::default() + }), + ..LlmOptimizationTokenImpact::default() + }); + assert!(recorder.record(item)); + } + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(100), + completion_tokens: Some(10), + total_tokens: Some(110), + cache_read_tokens: Some(20), + cache_write_tokens: Some(3), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()) + .unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(16)); + assert_eq!(summary.tokens_saved.cache_read_tokens, Some(20)); + assert_eq!(summary.tokens_saved.cache_write_tokens, Some(17)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().cache_read_tokens, + Some(40) + ); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().cache_write_tokens, + Some(20) + ); + assert_eq!(summary.contributions[0].sequence, Some(0)); + assert_eq!(summary.contributions[1].sequence, Some(1)); + } + + #[test] + fn serialized_summary_can_be_repriced_with_a_new_catalog() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let original = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()) + .unwrap(); + let restored: LlmOptimizationSummary = + serde_json::from_value(serde_json::to_value(&original).unwrap()).unwrap(); + let newer = resolver_with_rates(10.0, 5.0); + let baseline = newer + .estimate_cost_for_provider( + Some("test"), + "baseline", + restored.baseline_usage.as_ref().unwrap(), + ) + .unwrap() + .total_or_component_sum() + .unwrap(); + let actual = newer + .estimate_cost_for_provider( + Some("test"), + "effective", + restored.effective_usage.as_ref().unwrap(), + ) + .unwrap() + .total_or_component_sum() + .unwrap(); + assert_ne!(baseline - actual, original.estimated_cost_saved.unwrap()); + assert_eq!(restored.tokens_saved.prompt_tokens, Some(200)); + } + + #[test] + fn no_usage_is_an_explicit_partial_summary() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let summary = + finalize_optimization_summary(&recorder, None, Some("effective"), &resolver()).unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!( + summary + .limitations + .contains(&"missing_effective_usage".to_string()) + ); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); + } + + #[test] + fn payload_byte_limits_are_enforced_without_unbounded_work() { + let oversized = LlmOptimizationRecorder::default(); + let mut item = LlmOptimizationContribution::new("test", "custom"); + item.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + item.payload = Some(Json::String("x".repeat(MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES))); + assert!(!oversized.record(item)); + + let aggregate = LlmOptimizationRecorder::default(); + for index in 0..17 { + let mut item = LlmOptimizationContribution::new(format!("test.{index}"), "custom"); + item.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + item.payload = Some(Json::String("x".repeat(15_000))); + assert!(aggregate.record(item)); + } + let mut overflow = LlmOptimizationContribution::new("test.overflow", "custom"); + overflow.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + overflow.payload = Some(Json::String("x".repeat(15_000))); + assert!(!aggregate.record(overflow)); + assert!( + aggregate + .finish() + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + } + + #[test] + fn bounds_and_invalid_payloads_are_best_effort_and_visible() { + let recorder = LlmOptimizationRecorder::default(); + let mut invalid = LlmOptimizationContribution::new("test", "custom"); + invalid.payload = Some(json!({"evidence": true})); + assert!(!recorder.record(invalid)); + for index in 0..MAX_LLM_OPTIMIZATION_CONTRIBUTIONS { + assert!(recorder.record(LlmOptimizationContribution::new( + format!("test.{index}"), + "custom" + ))); + } + assert!(!recorder.record(LlmOptimizationContribution::new("overflow", "custom"))); + let summary = + finalize_optimization_summary(&recorder, None, None, &PricingResolver::default()) + .unwrap(); + assert_eq!( + summary.contributions.len(), + MAX_LLM_OPTIMIZATION_CONTRIBUTIONS + ); + assert!( + summary + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + assert!( + summary + .limitations + .contains(&"invalid_contribution_payload_schema".to_string()) + ); + } + + #[tokio::test] + async fn recorder_can_be_captured_for_stream_commit() { + let recorder = LlmOptimizationRecorder::default(); + let captured = scope_llm_optimization_recorder(recorder.clone(), async { + current_llm_optimization_recorder().unwrap() + }) + .await; + assert!(captured.record(LlmOptimizationContribution::new("test.stream", "commit"))); + assert_eq!(recorder.finish().contributions.len(), 1); + } +} diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 43f63dfd7..68afc9882 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -1137,6 +1137,7 @@ impl NemoRelayContextState { let mut request_value = request; let mut annotated_value = annotated; let mut pending_marks = Vec::new(); + let mut optimization_contributions = Vec::new(); for entry in entries { let input_content = request_value.content.clone(); let outcome = (entry.payload.callable)(name, request_value, annotated_value)?; @@ -1155,6 +1156,7 @@ impl NemoRelayContextState { request_value = outcome.request; annotated_value = outcome.annotated_request; pending_marks.extend(outcome.pending_marks); + optimization_contributions.extend(outcome.optimization_contributions); if entry.payload.break_chain { break; } @@ -1163,6 +1165,7 @@ impl NemoRelayContextState { request: request_value, annotated_request: annotated_value, pending_marks, + optimization_contributions, }) } diff --git a/crates/core/src/api/shared.rs b/crates/core/src/api/shared.rs index d9145d434..a1f3fc5d0 100644 --- a/crates/core/src/api/shared.rs +++ b/crates/core/src/api/shared.rs @@ -188,15 +188,18 @@ pub(crate) fn metadata_with_otel_status( metadata } +pub(crate) type InterceptedLlmRequest = ( + LlmRequest, + Option>, + Vec, + Vec, +); + pub(crate) fn run_request_intercepts_with_codec( name: &str, request: LlmRequest, codec: Option>, -) -> Result<( - LlmRequest, - Option>, - Vec, -)> { +) -> Result { let annotated = match &codec { Some(codec) => Some(codec.decode(&request)?), None => None, @@ -226,14 +229,25 @@ pub(crate) fn run_request_intercepts_with_codec( let mut request = outcome.request; inject_dynamo_session_ids(&mut request); let pending_marks = outcome.pending_marks; + let optimization_contributions = outcome.optimization_contributions; match (codec, outcome.annotated_request) { (Some(codec), Some(annotated)) => { let mut encoded = codec.encode(&annotated, &request)?; encoded.headers.extend(request.headers); - Ok((encoded, Some(Arc::new(annotated)), pending_marks)) + Ok(( + encoded, + Some(Arc::new(annotated)), + pending_marks, + optimization_contributions, + )) } - (_, annotated) => Ok((request, annotated.map(Arc::new), pending_marks)), + (_, annotated) => Ok(( + request, + annotated.map(Arc::new), + pending_marks, + optimization_contributions, + )), } } diff --git a/crates/core/src/codec/anthropic.rs b/crates/core/src/codec/anthropic.rs index 3243888c8..c82d22925 100644 --- a/crates/core/src/codec/anthropic.rs +++ b/crates/core/src/codec/anthropic.rs @@ -410,6 +410,7 @@ impl LlmResponseCodec for AnthropicMessagesCodec { tool_calls, finish_reason, usage, + optimization_summary: None, api_specific, extra: raw.extra, }) diff --git a/crates/core/src/codec/mod.rs b/crates/core/src/codec/mod.rs index 743521bb5..8c103f97d 100644 --- a/crates/core/src/codec/mod.rs +++ b/crates/core/src/codec/mod.rs @@ -18,6 +18,7 @@ pub mod anthropic; pub mod model_pricing; pub mod openai_chat; pub mod openai_responses; +pub mod optimization; pub mod request; pub mod resolve; pub mod response; diff --git a/crates/core/src/codec/openai_chat.rs b/crates/core/src/codec/openai_chat.rs index ccdc7d60e..2c7c66c7e 100644 --- a/crates/core/src/codec/openai_chat.rs +++ b/crates/core/src/codec/openai_chat.rs @@ -216,6 +216,7 @@ impl LlmResponseCodec for OpenAIChatCodec { tool_calls, finish_reason, usage, + optimization_summary: None, api_specific, extra: raw.extra, }) diff --git a/crates/core/src/codec/openai_responses.rs b/crates/core/src/codec/openai_responses.rs index 7f64d28c3..57b05661c 100644 --- a/crates/core/src/codec/openai_responses.rs +++ b/crates/core/src/codec/openai_responses.rs @@ -544,6 +544,7 @@ impl LlmResponseCodec for OpenAIResponsesCodec { tool_calls, finish_reason, usage, + optimization_summary: None, api_specific, extra: raw.extra, }) diff --git a/crates/core/src/codec/optimization.rs b/crates/core/src/codec/optimization.rs new file mode 100644 index 000000000..75dcbe23d --- /dev/null +++ b/crates/core/src/codec/optimization.rs @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Re-exported plugin-neutral LLM optimization data types. + +pub use nemo_relay_types::codec::optimization::*; diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 96ffe8665..f8de8fa14 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -770,7 +770,7 @@ fn extract_metrics( .and_then(Json::as_array) .map(|a| a.iter().filter_map(Json::as_f64).collect()); let known: std::collections::HashSet<&str> = TOKEN_USAGE_KNOWN_KEYS.iter().copied().collect(); - let extra_map: serde_json::Map = output + let mut extra_map: serde_json::Map = output .as_object() .into_iter() .flat_map(|output| { @@ -782,12 +782,31 @@ fn extract_metrics( .filter(|(k, _)| !known.contains(k.as_str())) .map(|(k, v)| (k.clone(), v.clone())) .collect(); + if let Some(summary) = + normalized_response.and_then(|response| response.optimization_summary.as_ref()) + && let Ok(summary) = serde_json::to_value(summary) + { + let relay = extra_map + .entry("nemo_relay".to_string()) + .or_insert_with(|| Json::Object(serde_json::Map::new())); + if !relay.is_object() { + *relay = Json::Object(serde_json::Map::new()); + } + if let Some(relay) = relay.as_object_mut() { + relay.insert("optimization".to_string(), summary); + } + } let extra = if extra_map.is_empty() { None } else { Some(Json::Object(extra_map)) }; - if prompt.is_none() && completion.is_none() && cached.is_none() && cost.is_none() { + if prompt.is_none() + && completion.is_none() + && cached.is_none() + && cost.is_none() + && extra.is_none() + { return None; } Some(AtifMetrics { @@ -1213,6 +1232,10 @@ struct FinalMetricsTotals { completion_tokens: Option, cached_tokens: Option, cost_usd: Option, + optimization_prompt_tokens_saved: Option, + optimization_total_tokens_saved: Option, + optimization_estimated_cost_saved_usd: Option, + optimization_calls: u64, } impl FinalMetricsTotals { @@ -1221,16 +1244,54 @@ impl FinalMetricsTotals { add_u64_total(&mut self.completion_tokens, metrics.completion_tokens); add_u64_total(&mut self.cached_tokens, metrics.cached_tokens); add_f64_total(&mut self.cost_usd, metrics.cost_usd); + let optimization = metrics + .extra + .as_ref() + .and_then(|extra| extra.pointer("/nemo_relay/optimization")); + if let Some(optimization) = optimization { + self.optimization_calls = self.optimization_calls.saturating_add(1); + add_u64_total( + &mut self.optimization_prompt_tokens_saved, + optimization + .pointer("/tokens_saved/prompt_tokens") + .and_then(Json::as_u64), + ); + add_u64_total( + &mut self.optimization_total_tokens_saved, + optimization + .pointer("/tokens_saved/total_tokens") + .and_then(Json::as_u64), + ); + let saved_usd = optimization + .get("currency") + .and_then(Json::as_str) + .filter(|currency| currency.eq_ignore_ascii_case("USD")) + .and_then(|_| optimization.get("estimated_cost_saved")) + .and_then(Json::as_f64); + add_f64_total(&mut self.optimization_estimated_cost_saved_usd, saved_usd); + } } fn into_final_metrics(self, step_count: usize) -> AtifFinalMetrics { + let optimization_extra = (self.optimization_calls > 0).then(|| { + serde_json::json!({ + "nemo_relay": { + "optimization": { + "llm_call_count": self.optimization_calls, + "prompt_tokens_saved": self.optimization_prompt_tokens_saved, + "total_tokens_saved": self.optimization_total_tokens_saved, + "estimated_cost_saved_usd": self.optimization_estimated_cost_saved_usd, + } + } + }) + }); AtifFinalMetrics { total_prompt_tokens: self.prompt_tokens, total_completion_tokens: self.completion_tokens, total_cached_tokens: self.cached_tokens, total_cost_usd: self.cost_usd, total_steps: Some(step_count as u64), - extra: None, + extra: optimization_extra, } } } diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 15676277d..5dd4b98a0 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -908,6 +908,100 @@ fn push_annotated_response_attributes( if let Some(tool_calls) = response.tool_calls.as_deref() { push_response_tool_calls(attributes, 0, tool_calls); } + if let Some(summary) = response.optimization_summary.as_ref() { + push_optimization_attributes(attributes, summary); + } +} + +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), + )); + } + let costs = [ + ( + "nemo_relay.llm.optimization.baseline_cost", + summary + .baseline_cost + .as_ref() + .and_then(|cost| cost.total_or_component_sum()), + ), + ( + "nemo_relay.llm.optimization.actual_cost", + summary + .actual_cost + .as_ref() + .and_then(|cost| cost.total_or_component_sum()), + ), + ( + "nemo_relay.llm.optimization.estimated_cost_saved", + summary.estimated_cost_saved, + ), + ]; + for (key, value) in costs { + if let Some(value) = value { + attributes.push(KeyValue::new(key, value)); + } + } + attributes.push(KeyValue::new( + "nemo_relay.llm.optimization.status", + match summary.status { + crate::codec::optimization::LlmOptimizationSummaryStatus::Complete => "complete", + crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", + }, + )); + let provenance = summary + .baseline_cost + .as_ref() + .or(summary.actual_cost.as_ref()); + if let Some(source) = provenance.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) = provenance.and_then(|cost| cost.pricing_as_of.as_ref()) { + attributes.push(KeyValue::new( + "nemo_relay.llm.optimization.pricing_as_of", + as_of.clone(), + )); + } } 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 727b60a79..f3b5d85bd 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -707,9 +707,99 @@ fn end_attributes(event: &Event) -> Vec { attributes.push(KeyValue::new("nemo_relay.llm.cost.total", cost)); attributes.push(KeyValue::new("nemo_relay.llm.cost.currency", currency)); } + if let Some(response) = event.annotated_response() + && let Some(summary) = response.optimization_summary.as_ref() + { + push_optimization_attributes(&mut attributes, summary); + } attributes } +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 + .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.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.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", + }, + )); + let provenance = summary + .baseline_cost + .as_ref() + .or(summary.actual_cost.as_ref()); + if let Some(source) = provenance.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) = provenance.and_then(|cost| cost.pricing_as_of.as_ref()) { + attributes.push(KeyValue::new( + "nemo_relay.llm.optimization.pricing_as_of", + as_of.clone(), + )); + } +} + fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { if let Some(response) = event.normalized_llm_response() { let response = response.as_ref(); diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index f06e25f0d..2b380fd86 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -33,6 +33,8 @@ use tokio_stream::Stream; use crate::api::event::{BaseEvent, MarkEvent}; use crate::api::llm::LlmHandle; +use crate::api::llm::emit_optimization_marks; +use crate::api::optimization::finalize_optimization_summary; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{EventSubscriberFn, ScopeStackHandle, current_scope_stack}; @@ -208,15 +210,29 @@ impl LlmStreamWrapper { } else { Some(sanitized) }; - let annotated_response: Option> = self - .response_codec - .as_ref() - .and_then(|codec| { + let mut annotated_response: Option = + self.response_codec.as_ref().and_then(|codec| { let mut decoded = codec.decode_response(data.as_ref()?).ok()?; attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); Some(decoded) - }) - .map(Arc::new); + }); + let _ = emit_optimization_marks(&self.handle, &self.subscribers); + let pricing = crate::codec::response::active_pricing_resolver(); + let summary = finalize_optimization_summary( + &self.handle.optimization_recorder, + annotated_response.as_mut(), + self.handle.model_name.as_deref(), + &pricing, + ); + if annotated_response.is_none() + && let Some(summary) = summary + { + annotated_response = Some(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }); + } + let annotated_response = annotated_response.map(Arc::new); let event_snapshot = { let ctx = global_context(); let state = ctx.read(); diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index a59b4956d..451eb3515 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -22,6 +22,7 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; +use nemo_relay::api::optimization::record_llm_optimization_contribution; use nemo_relay::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, @@ -55,6 +56,10 @@ use nemo_relay::api::tool::{ ToolExecutionInterceptOutcome, tool_call, tool_call_end, tool_call_execute, tool_conditional_execution, tool_request_intercepts, }; +use nemo_relay::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationTokenImpact, + LlmOptimizationTokens, +}; use nemo_relay::error::FlowError; #[cfg(all(feature = "otel", feature = "openinference"))] use nemo_relay::observability::MarkProjection; @@ -2370,7 +2375,6 @@ fn test_llm_request_intercept_registry_mutations_apply_to_later_calls() { }), ) .unwrap(); - let request = llm_request_intercepts( "llm", LlmRequest { @@ -3424,6 +3428,118 @@ async fn test_managed_llm_emits_pending_marks_under_started_scope() { deregister_subscriber("pending_mark_observer").unwrap(); } +#[tokio::test] +async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "optimization_observer", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + register_llm_request_intercept( + "optimization_contributor", + 1, + false, + Arc::new(|_name, request, annotated| { + let mut contribution = + LlmOptimizationContribution::new("test.optimizer", "test_custom_kind"); + contribution.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(12)), + quality: Some(LlmOptimizationEvidenceQuality::Estimated), + estimation_method: Some("test-counter".to_string()), + ..LlmOptimizationTokenImpact::default() + }); + Ok(LlmRequestInterceptOutcome::new(request, annotated) + .with_optimization_contribution(contribution)) + }), + ) + .unwrap(); + register_llm_execution_intercept( + "optimization_execution_contributor", + 1, + Arc::new(|_name, request, next| { + Box::pin(async move { + let contribution = + LlmOptimizationContribution::new("test.execution", "test_execution_kind"); + assert!(record_llm_optimization_contribution(contribution)); + next(request).await + }) + }), + ) + .unwrap(); + + llm_call_execute( + LlmCallExecuteParams::builder() + .name("optimized-managed-llm") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { Ok(json!({"response": "done"})) }) + })) + .build(), + ) + .await + .unwrap(); + + let captured = captured_events_snapshot(&events); + let start = captured + .iter() + .find(|event| { + event.name() == "optimized-managed-llm" + && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap(); + let marks = captured + .iter() + .filter(|event| event.name() == "nemo_relay.llm.optimization") + .collect::>(); + assert_eq!(marks.len(), 2); + assert!( + marks + .iter() + .all(|mark| mark.parent_uuid() == Some(start.uuid())) + ); + assert!(marks.iter().all(|mark| { + mark.data_schema().unwrap().name == "nemo.relay.llm_optimization_contribution" + })); + assert_eq!( + marks[0].data().unwrap()["token_impact"]["saved"]["prompt_tokens"], + 12 + ); + assert_eq!(marks[0].data().unwrap()["sequence"], 0); + assert_eq!(marks[1].data().unwrap()["sequence"], 1); + + let end = captured + .iter() + .find(|event| { + event.name() == "optimized-managed-llm" + && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap(); + let summary = end + .annotated_response() + .unwrap() + .optimization_summary + .as_ref() + .unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(12)); + assert_eq!(summary.contributions.len(), 2); + assert_eq!(summary.contributions[0].producer, "test.optimizer"); + assert_eq!(summary.contributions[1].producer, "test.execution"); + + deregister_llm_request_intercept("optimization_contributor").unwrap(); + deregister_llm_execution_intercept("optimization_execution_contributor").unwrap(); + deregister_subscriber("optimization_observer").unwrap(); +} + #[tokio::test] async fn test_failed_request_intercept_does_not_emit_pending_marks_or_start_scope() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 73496a910..978716d51 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -522,6 +522,7 @@ async fn sdk_cdylib_registers_tool_request_intercept() { tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra, })) diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index cf411e509..6f8f09df9 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -1181,6 +1181,7 @@ impl LlmResponseCodec for MockResponseCodec { cache_write_tokens: None, cost: None, }), + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), }) diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 6a994b27e..c683fb2fb 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -108,6 +108,7 @@ fn annotated_response_with_usage(model: &str, usage: Usage) -> AnnotatedLlmRespo tool_calls: None, finish_reason: None, usage: Some(usage), + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), } @@ -1495,6 +1496,68 @@ fn test_final_metrics_preserve_explicit_zero_cost_without_fabricating_tokens() { assert_eq!(final_metrics.total_cost_usd, Some(0.0)); } +#[test] +fn test_optimization_summary_projects_to_step_and_final_metrics() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version": "1", + "calculation_version": "1", + "status": "complete", + "baseline_model": {"model": "baseline"}, + "effective_model": {"model": "effective"}, + "tokens_saved": {"prompt_tokens": 12, "total_tokens": 12}, + "baseline_cost": {"total": 0.02, "currency": "USD", "source": "model_pricing", "pricing_as_of": "2026-07-08", "pricing_source": "test"}, + "actual_cost": {"total": 0.01, "currency": "USD", "source": "model_pricing", "pricing_as_of": "2026-07-08", "pricing_source": "test"}, + "estimated_cost_saved": 0.01, + "currency": "USD", + "contributions": [] + })) + .unwrap(); + let mut response = annotated_response_with_usage( + "effective", + Usage { + prompt_tokens: Some(8), + ..Usage::default() + }, + ); + response.optimization_summary = Some(summary); + + let metrics = extract_metrics(&json!({}), None, None, Some(&response)).unwrap(); + assert_eq!( + metrics + .extra + .as_ref() + .unwrap() + .pointer("/nemo_relay/optimization/tokens_saved/prompt_tokens"), + Some(&json!(12)) + ); + let final_metrics = compute_final_metrics(&[AtifStep { + step_id: 1, + source: "agent".to_string(), + message: json!("done"), + timestamp: None, + model_name: Some("effective".to_string()), + reasoning_effort: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: Some(metrics), + llm_call_count: Some(1), + is_copied_context: None, + extra: None, + }]) + .unwrap(); + let optimization = final_metrics + .extra + .as_ref() + .unwrap() + .pointer("/nemo_relay/optimization") + .unwrap(); + assert_eq!(optimization["prompt_tokens_saved"], 12); + assert_eq!(optimization["total_tokens_saved"], 12); + assert_eq!(optimization["estimated_cost_saved_usd"], 0.01); +} + #[test] fn test_exporter_llm_lifecycle_plain_input() { // Input without LlmRequest envelope — passed through unchanged. diff --git a/crates/core/tests/unit/codec/response_tests.rs b/crates/core/tests/unit/codec/response_tests.rs index 1fa6eaaa7..93d2935b1 100644 --- a/crates/core/tests/unit/codec/response_tests.rs +++ b/crates/core/tests/unit/codec/response_tests.rs @@ -55,6 +55,7 @@ fn full_response() -> AnnotatedLlmResponse { cache_write_tokens: Some(3), cost: None, }), + optimization_summary: None, api_specific: Some(ApiSpecificResponse::OpenAIChat { logprobs: None, system_fingerprint: Some("fp_abc123".into()), @@ -73,6 +74,7 @@ fn minimal_response() -> AnnotatedLlmResponse { tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), } @@ -1698,6 +1700,7 @@ impl LlmResponseCodec for MockResponseCodec { tool_calls: None, finish_reason: Some(FinishReason::Complete), usage: None, + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), }) diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index f9581cf6b..28d4f5bdc 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -126,11 +126,45 @@ fn empty_annotated_response() -> AnnotatedLlmResponse { tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), } } +#[test] +fn optimization_summary_emits_namespaced_openinference_attributes() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"complete", + "baseline_model":{"model":"baseline"}, "effective_model":{"model":"effective"}, + "tokens_saved":{"prompt_tokens":12,"total_tokens":12}, + "baseline_cost":{"total":0.02,"currency":"USD","source":"model_pricing","pricing_as_of":"2026-07-08","pricing_source":"test"}, + "actual_cost":{"total":0.01,"currency":"USD","source":"model_pricing"}, + "estimated_cost_saved":0.01, "currency":"USD", "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + assert_eq!( + attributes["nemo_relay.llm.optimization.effective_model"], + "effective" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.total_tokens_saved"], + "12" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost"], + "0.01" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_source"], + "test" + ); +} + fn install_test_pricing(model_id: &str) { let catalog = PricingCatalog::from_json_str( &json!({ @@ -2939,6 +2973,7 @@ fn llm_end_with_usage_emits_token_count_attributes() { cache_write_tokens: Some(10), cost: None, }), + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), })) @@ -3909,6 +3944,7 @@ fn llm_end_with_partial_usage_emits_only_present_fields() { cache_write_tokens: None, cost: None, }), + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), })) diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index b03e3e7f5..c82dbbefd 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -45,11 +45,45 @@ fn empty_annotated_response() -> AnnotatedLlmResponse { tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra: serde_json::Map::new(), } } +#[test] +fn optimization_summary_emits_namespaced_otel_attributes() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"complete", + "baseline_model":{"model":"baseline"}, "effective_model":{"model":"effective"}, + "tokens_saved":{"prompt_tokens":12,"total_tokens":12}, + "baseline_cost":{"total":0.02,"currency":"USD","source":"model_pricing","pricing_as_of":"2026-07-08","pricing_source":"test"}, + "actual_cost":{"total":0.01,"currency":"USD","source":"model_pricing"}, + "estimated_cost_saved":0.01, "currency":"USD", "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_model"], + "baseline" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.prompt_tokens_saved"], + "12" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.estimated_cost_saved"], + "0.01" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_as_of"], + "2026-07-08" + ); +} + fn install_test_pricing(model_id: &str) { let catalog = PricingCatalog::from_json_str( &json!({ diff --git a/crates/core/tests/unit/shared_tests.rs b/crates/core/tests/unit/shared_tests.rs index 196efd03e..ced67814b 100644 --- a/crates/core/tests/unit/shared_tests.rs +++ b/crates/core/tests/unit/shared_tests.rs @@ -162,7 +162,7 @@ fn test_run_request_intercepts_with_codec_none_and_codec_paths() { ) .unwrap(); - let (request_without_codec, annotated_without_codec, pending_marks_without_codec) = + let (request_without_codec, annotated_without_codec, pending_marks_without_codec, _) = run_request_intercepts_with_codec( "shared", LlmRequest { @@ -199,7 +199,7 @@ fn test_run_request_intercepts_with_codec_none_and_codec_paths() { .unwrap(); let codec: Arc = Arc::new(SharedTestCodec); - let (request_with_codec, annotated_with_codec, pending_marks_with_codec) = + let (request_with_codec, annotated_with_codec, pending_marks_with_codec, _) = run_request_intercepts_with_codec( "shared", LlmRequest { @@ -269,7 +269,7 @@ fn test_run_request_intercepts_injects_dynamo_agent_lineage() { ) .unwrap(); - let (request, _, _) = run_request_intercepts_with_codec( + let (request, _, _, _) = run_request_intercepts_with_codec( "openai.responses", LlmRequest { headers: Map::new(), @@ -296,7 +296,7 @@ fn test_run_request_intercepts_injects_dynamo_agent_lineage() { .build(), ) .unwrap(); - let (request_with_codec, _, _) = run_request_intercepts_with_codec( + let (request_with_codec, _, _, _) = run_request_intercepts_with_codec( "openai.responses", LlmRequest { headers: Map::new(), @@ -333,7 +333,7 @@ fn test_run_request_intercepts_injects_dynamo_agent_lineage() { ) .unwrap(); - let (request, _, _) = run_request_intercepts_with_codec( + let (request, _, _, _) = run_request_intercepts_with_codec( "openai.responses", LlmRequest { headers: Map::new(), @@ -375,7 +375,7 @@ fn test_run_request_intercepts_injects_dynamo_agent_lineage() { .build(), ) .unwrap(); - let (request, _, _) = run_request_intercepts_with_codec( + let (request, _, _, _) = run_request_intercepts_with_codec( "openai.responses", LlmRequest { headers: Map::new(), diff --git a/crates/core/tests/unit/types_tests.rs b/crates/core/tests/unit/types_tests.rs index dc899e46c..d9375b497 100644 --- a/crates/core/tests/unit/types_tests.rs +++ b/crates/core/tests/unit/types_tests.rs @@ -123,6 +123,7 @@ fn annotated_response(id: &str, model: &str, text: &str) -> AnnotatedLlmResponse tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra: Map::new(), } diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index a0d9ee5e2..11f63871d 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -308,6 +308,7 @@ pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new( request: unsafe { &*request }.0.clone(), annotated_request, pending_marks, + optimization_contributions: Vec::new(), }; match serde_json::to_value(outcome) { Ok(value) => { diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index efd757fd7..1e3f43a60 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -2955,7 +2955,7 @@ pub fn tool_conditional_execution(env: Env, name: String, args: Json) -> Result< /// The `request` should be a JSON object with `headers` and `content` fields matching /// the `LlmRequest` schema. Returns the transformed request as JSON. #[napi( - ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> }>" + ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions: Json[] }>" )] pub fn llm_request_intercepts(env: Env, name: String, request: Json) -> Result { let llm_request: LlmRequest = serde_json::from_value(request) @@ -2971,6 +2971,7 @@ pub fn llm_request_intercepts(env: Env, name: String, request: Json) -> Result, #[serde(default)] pending_marks: Vec, + #[serde(default)] + optimization_contributions: Vec, } let outcome: JsOutcome = serde_json::from_value(result).map_err(|e| { FlowError::Internal(format!("invalid JS LLM request intercept outcome: {e}")) @@ -310,6 +313,7 @@ pub fn wrap_js_llm_request_intercept_fn( request: outcome.request, annotated_request: outcome.annotated, pending_marks: outcome.pending_marks.into_iter().map(Into::into).collect(), + optimization_contributions: outcome.optimization_contributions, }) }, ) diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 9f1dfd918..477566a39 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -22,6 +22,11 @@ pub use nemo_relay_types::api::event::{ pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; pub use nemo_relay_types::api::tool::{ToolAttributes, ToolExecutionInterceptOutcome}; +pub use nemo_relay_types::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, + LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload, + LlmOptimizationTokenImpact, LlmOptimizationTokens, +}; pub use nemo_relay_types::codec::request::AnnotatedLlmRequest; pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index d8b9d1a84..1a27c6742 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -539,6 +539,7 @@ impl PyAnnotatedLLMResponse { tool_calls: optional_py_json(tool_calls, "tool_calls")?, finish_reason: optional_finish_reason(finish_reason)?, usage: optional_py_json(usage, "usage")?, + optimization_summary: None, api_specific: optional_py_json(api_specific, "api_specific")?, extra: optional_py_json(extra, "extra")?.unwrap_or_default(), }, diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index aa6fc16b5..442e7f3b4 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -689,19 +689,32 @@ pub struct PyLLMRequestInterceptOutcome { #[pymethods] impl PyLLMRequestInterceptOutcome { #[new] - #[pyo3(signature = (request, annotated_request=None, pending_marks=Vec::new()))] + #[pyo3(signature = (request, annotated_request=None, pending_marks=Vec::new(), optimization_contributions=None))] fn new( request: PyLLMRequest, annotated_request: Option, pending_marks: Vec, - ) -> Self { - Self { + optimization_contributions: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let optimization_contributions = optimization_contributions + .map(py_to_json) + .transpose()? + .map(serde_json::from_value) + .transpose() + .map_err(|error| { + pyo3::exceptions::PyValueError::new_err(format!( + "invalid optimization_contributions: {error}" + )) + })? + .unwrap_or_default(); + Ok(Self { inner: LlmRequestInterceptOutcome { request: request.inner, annotated_request: annotated_request.map(|value| value.inner), pending_marks: pending_marks.into_iter().map(|value| value.inner).collect(), + optimization_contributions, }, - } + }) } #[getter] @@ -728,6 +741,15 @@ impl PyLLMRequestInterceptOutcome { .map(|inner| PyPendingMarkSpec { inner }) .collect() } + + #[getter] + fn optimization_contributions(&self, py: Python<'_>) -> PyResult> { + json_to_py( + py, + &serde_json::to_value(&self.inner.optimization_contributions) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?, + ) + } } /// Canonical result returned by Python tool execution intercepts. diff --git a/crates/types/src/api/llm.rs b/crates/types/src/api/llm.rs index 121754818..359244bee 100644 --- a/crates/types/src/api/llm.rs +++ b/crates/types/src/api/llm.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use crate::Json; use crate::api::event::PendingMarkSpec; +use crate::codec::optimization::LlmOptimizationContribution; use crate::codec::request::AnnotatedLlmRequest; bitflags! { @@ -48,6 +49,9 @@ pub struct LlmRequestInterceptOutcome { /// Ordered marks to emit after Relay creates and starts the LLM scope. #[serde(default)] pub pending_marks: Vec, + /// Ordered plugin-neutral optimization evidence for this LLM call. + #[serde(default)] + pub optimization_contributions: Vec, } impl LlmRequestInterceptOutcome { @@ -57,6 +61,7 @@ impl LlmRequestInterceptOutcome { request, annotated_request, pending_marks: Vec::new(), + optimization_contributions: Vec::new(), } } @@ -66,6 +71,16 @@ impl LlmRequestInterceptOutcome { self.pending_marks.push(mark); self } + + /// Append one optimization contribution while preserving interceptor order. + #[must_use] + pub fn with_optimization_contribution( + mut self, + contribution: LlmOptimizationContribution, + ) -> Self { + self.optimization_contributions.push(contribution); + self + } } impl From for LlmRequestInterceptOutcome { diff --git a/crates/types/src/codec/mod.rs b/crates/types/src/codec/mod.rs index 5ff75f032..79d61e9ba 100644 --- a/crates/types/src/codec/mod.rs +++ b/crates/types/src/codec/mod.rs @@ -3,6 +3,8 @@ //! Shared normalized LLM request and response data types. +/// Plugin-neutral LLM optimization evidence and summaries. +pub mod optimization; /// Normalized LLM request data types. pub mod request; /// Normalized LLM response data types. diff --git a/crates/types/src/codec/optimization.rs b/crates/types/src/codec/optimization.rs new file mode 100644 index 000000000..cf5cc9a6b --- /dev/null +++ b/crates/types/src/codec/optimization.rs @@ -0,0 +1,353 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Plugin-neutral evidence and summaries for LLM optimizations. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::Json; +use crate::api::event::DataSchema; + +use super::response::{CostEstimate, Usage}; + +/// Open, forward-compatible optimization classification. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct LlmOptimizationKind(String); + +impl LlmOptimizationKind { + /// A request transformation that reduces input tokens. + pub const INPUT_COMPRESSION: &'static str = "input_compression"; + /// A routing decision that changes the model serving a request. + pub const MODEL_ROUTING: &'static str = "model_routing"; + + /// Preserve an arbitrary producer-defined kind on the wire. + #[must_use] + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + /// Return the exact wire value. + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } + + /// Construct the standard input-compression kind. + #[must_use] + pub fn input_compression() -> Self { + Self::new(Self::INPUT_COMPRESSION) + } + + /// Construct the standard model-routing kind. + #[must_use] + pub fn model_routing() -> Self { + Self::new(Self::MODEL_ROUTING) + } +} + +impl From for LlmOptimizationKind { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl From<&str> for LlmOptimizationKind { + fn from(value: &str) -> Self { + Self::new(value) + } +} + +/// Model identity used for counterfactual and effective pricing. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LlmOptimizationModel { + /// Model identifier understood by Relay's pricing resolver. + pub model: String, + /// Optional pricing-provider namespace. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider: Option, +} + +impl LlmOptimizationModel { + /// Construct a model identity without a provider namespace. + #[must_use] + pub fn new(model: impl Into) -> Self { + Self { + model: model.into(), + provider: None, + } + } + + /// Add the pricing-provider namespace. + #[must_use] + pub fn with_provider(mut self, provider: impl Into) -> Self { + self.provider = Some(provider.into()); + self + } +} + +/// A model change proposed or applied by an optimizer. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct LlmOptimizationModelTransition { + /// Counterfactual model that would otherwise have served the request. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline: Option, + /// Model selected by the optimization. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effective: Option, +} + +/// Token quantities retained independently from pricing arithmetic. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct LlmOptimizationTokens { + /// Input/prompt tokens. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_tokens: Option, + /// Output/completion tokens. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub completion_tokens: Option, + /// Tokens read from a provider prompt cache. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_read_tokens: Option, + /// Tokens written to a provider prompt cache. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_write_tokens: Option, + /// Total tokens when supplied by the producer. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total_tokens: Option, +} + +impl LlmOptimizationTokens { + /// Construct explicit prompt and total token savings. + #[must_use] + pub fn saved_prompt(prompt_tokens: u64) -> Self { + Self { + prompt_tokens: Some(prompt_tokens), + total_tokens: Some(prompt_tokens), + ..Self::default() + } + } + + /// Saturating field-wise addition, preserving absent fields when neither side supplies them. + pub fn add_assign(&mut self, other: &Self) { + fn add(target: &mut Option, value: Option) { + if let Some(value) = value { + *target = Some(target.unwrap_or(0).saturating_add(value)); + } + } + add(&mut self.prompt_tokens, other.prompt_tokens); + add(&mut self.completion_tokens, other.completion_tokens); + add(&mut self.cache_read_tokens, other.cache_read_tokens); + add(&mut self.cache_write_tokens, other.cache_write_tokens); + add(&mut self.total_tokens, other.total_tokens); + } +} + +/// Whether token evidence was observed directly or estimated. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LlmOptimizationEvidenceQuality { + /// Directly observed token counts. + Observed, + /// Counts produced by a tokenizer or estimator. + Estimated, +} + +/// Shared token evidence that Relay can aggregate without understanding a plugin payload. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct LlmOptimizationTokenImpact { + /// Token counts before the optimization. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline: Option, + /// Token counts after the optimization. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effective: Option, + /// Explicit reduction retained for downstream repricing. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub saved: Option, + /// Evidence quality for these counts. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub quality: Option, + /// Tokenizer, counter, or estimation method. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub estimation_method: Option, +} + +/// Trait implemented by typed plugin payloads embedded in an optimization contribution. +pub trait LlmOptimizationPayload: Serialize { + /// Stable schema name for the payload. + const SCHEMA_NAME: &'static str; + /// Schema version for the payload. + const SCHEMA_VERSION: &'static str; +} + +/// One optimizer's evidence for a change to an LLM call. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct LlmOptimizationContribution { + /// Relay-assigned contribution identifier. Producer values are replaced on ingestion. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub id: Option, + /// Relay-assigned order within the LLM call. Producer values are replaced on ingestion. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sequence: Option, + /// Stable producer or plugin identity. + pub producer: String, + /// Open optimization classification. + pub kind: LlmOptimizationKind, + /// Whether the optimization affected the executed call. + #[serde(default)] + pub applied: bool, + /// Optional shared model transition. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_transition: Option, + /// Optional shared token evidence. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_impact: Option, + /// Schema of `payload`; required whenever `payload` is present. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload_schema: Option, + /// Opaque producer payload retained for audit and future consumers. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload: Option, + /// Unknown top-level fields from future producers. + #[serde(flatten)] + pub extra: BTreeMap, +} + +impl LlmOptimizationContribution { + /// Construct an applied contribution with no model, token, or custom evidence. + #[must_use] + pub fn new(producer: impl Into, kind: impl Into) -> Self { + Self { + id: None, + sequence: None, + producer: producer.into(), + kind: kind.into(), + applied: true, + model_transition: None, + token_impact: None, + payload_schema: None, + payload: None, + extra: BTreeMap::new(), + } + } + + /// Attach a schema-tagged custom payload. + pub fn with_payload( + mut self, + payload: &T, + ) -> Result { + self.payload_schema = Some(DataSchema { + name: T::SCHEMA_NAME.to_string(), + version: T::SCHEMA_VERSION.to_string(), + }); + self.payload = Some(serde_json::to_value(payload)?); + Ok(self) + } +} + +/// Completeness of close-time optimization accounting. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LlmOptimizationSummaryStatus { + /// All requested token and monetary calculations were available. + Complete, + /// Evidence remains useful, but one or more calculations were unavailable. + Partial, +} + +/// Close-time accounting attached to the normalized LLM response. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct LlmOptimizationSummary { + /// Wire schema version. + pub schema_version: String, + /// Arithmetic implementation version. + pub calculation_version: String, + /// Whether all accounting inputs were available. + pub status: LlmOptimizationSummaryStatus, + /// Machine-readable reasons for partial accounting. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub limitations: Vec, + /// Counterfactual model used for baseline pricing. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline_model: Option, + /// Model that actually served the terminal response. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effective_model: Option, + /// Usage observed on the effective response. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effective_usage: Option, + /// Counterfactual usage derived from observed usage plus explicit savings. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline_usage: Option, + /// Explicit aggregate token reductions, retained independently from pricing. + pub tokens_saved: LlmOptimizationTokens, + /// Estimated cost for baseline model/usage. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline_cost: Option, + /// Provider-reported or Relay-estimated actual cost. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub actual_cost: Option, + /// Baseline cost minus actual cost when both share a currency. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub estimated_cost_saved: Option, + /// Currency for `estimated_cost_saved`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub currency: Option, + /// Ordered, bounded source evidence used by the calculation. + pub contributions: Vec, +} + +#[cfg(test)] +mod tests { + use serde::Serialize; + use serde_json::json; + + use super::*; + + #[derive(Serialize)] + struct CustomPayload { + evidence: String, + } + + impl LlmOptimizationPayload for CustomPayload { + const SCHEMA_NAME: &'static str = "example.custom_optimization"; + const SCHEMA_VERSION: &'static str = "3"; + } + + #[test] + fn custom_kinds_payloads_and_future_fields_round_trip() { + let mut contribution = LlmOptimizationContribution::new("example", "energy_reduction") + .with_payload(&CustomPayload { + evidence: "measured".to_string(), + }) + .unwrap(); + contribution + .extra + .insert("future_field".to_string(), json!({"v": 2})); + let decoded: LlmOptimizationContribution = + serde_json::from_value(serde_json::to_value(&contribution).unwrap()).unwrap(); + assert_eq!(decoded.kind.as_str(), "energy_reduction"); + assert_eq!( + decoded.payload_schema.as_ref().unwrap().name, + "example.custom_optimization" + ); + assert_eq!(decoded.extra["future_field"], json!({"v": 2})); + } + + #[test] + fn saved_prompt_tokens_remain_explicit_on_the_wire() { + let impact = LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(42)), + quality: Some(LlmOptimizationEvidenceQuality::Estimated), + estimation_method: Some("tokenizer-v1".to_string()), + ..LlmOptimizationTokenImpact::default() + }; + let wire = serde_json::to_value(impact).unwrap(); + assert_eq!(wire["saved"]["prompt_tokens"], 42); + assert_eq!(wire["saved"]["total_tokens"], 42); + } +} diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs index ac1bc2482..87000ee0e 100644 --- a/crates/types/src/codec/response.rs +++ b/crates/types/src/codec/response.rs @@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize}; use crate::Json; +use super::optimization::LlmOptimizationSummary; use super::request::MessageContent; // --------------------------------------------------------------------------- @@ -50,6 +51,10 @@ pub struct AnnotatedLlmResponse { #[serde(skip_serializing_if = "Option::is_none")] pub usage: Option, + /// Plugin-neutral optimization accounting computed by Relay at LLM close. + #[serde(skip_serializing_if = "Option::is_none")] + pub optimization_summary: Option, + /// API-specific response data that cannot be normalized across providers. #[serde(skip_serializing_if = "Option::is_none")] pub api_specific: Option, diff --git a/crates/types/tests/serialization_tests.rs b/crates/types/tests/serialization_tests.rs index bb0507444..ddf6feedf 100644 --- a/crates/types/tests/serialization_tests.rs +++ b/crates/types/tests/serialization_tests.rs @@ -48,6 +48,7 @@ fn event_round_trips_with_annotated_llm_profiles() { tool_calls: None, finish_reason: None, usage: None, + optimization_summary: None, api_specific: None, extra: Map::new(), }; From 13bef5e783111f9969e0a0e99a166d47d76e869f Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Wed, 8 Jul 2026 12:51:10 -0600 Subject: [PATCH 02/10] fix: bound optimization payload measurement Signed-off-by: Bryan Bednarski --- crates/core/src/api/optimization.rs | 54 +++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index 92f0b6123..7f986f178 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -52,9 +52,15 @@ impl LlmOptimizationRecorder { } return false; } - Some(payload) => match serde_json::to_vec(payload) { - Ok(serialized) => serialized.len(), - Err(_) => { + Some(payload) => match bounded_json_size(payload, MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES) { + Ok(size) => size, + Err(PayloadSizeError::LimitExceeded) => { + if let Ok(mut state) = self.state.lock() { + state.contribution_limit_exceeded = true; + } + return false; + } + Err(PayloadSizeError::Serialization) => { if let Ok(mut state) = self.state.lock() { state.invalid_payload_schema = true; } @@ -123,6 +129,48 @@ impl LlmOptimizationRecorder { } } +enum PayloadSizeError { + LimitExceeded, + Serialization, +} + +fn bounded_json_size(value: &serde_json::Value, limit: usize) -> Result { + struct CountingWriter { + size: usize, + limit: usize, + exceeded: bool, + } + + impl std::io::Write for CountingWriter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if self.size.saturating_add(bytes.len()) > self.limit { + self.exceeded = true; + return Err(std::io::Error::other("optimization payload limit exceeded")); + } + self.size += bytes.len(); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let mut writer = CountingWriter { + size: 0, + limit, + exceeded: false, + }; + if serde_json::to_writer(&mut writer, value).is_err() { + return Err(if writer.exceeded { + PayloadSizeError::LimitExceeded + } else { + PayloadSizeError::Serialization + }); + } + Ok(writer.size) +} + struct FinishedContributions { contributions: Vec, limitations: Vec, From 11e7abe04d92dc5d64bee7e65c52cddbb57df278 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Wed, 8 Jul 2026 15:33:32 -0600 Subject: [PATCH 03/10] fix: price the model actually dispatched Signed-off-by: Bryan Bednarski --- crates/core/src/api/optimization.rs | 42 +++++++++++++++++++++++++---- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index 7f986f178..3fbccb26e 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -242,11 +242,16 @@ pub(crate) fn finalize_optimization_summary( } } - let effective_model = response - .as_ref() - .and_then(|response| response.model.as_ref()) - .map(|model| LlmOptimizationModel::new(model.clone())) - .or(contributed_effective_model) + // 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 + .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 baseline_model.is_none() { baseline_model = effective_model.clone(); @@ -451,6 +456,33 @@ mod tests { assert!((summary.estimated_cost_saved.unwrap() - 0.0014).abs() < 1e-12); } + #[test] + fn applied_route_is_the_authoritative_effective_model() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + // Providers may return an alias or deployment name rather than + // the exact model Relay selected and sent upstream. + model: Some("provider-response-alias".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + Some("original-request-model"), + &resolver(), + ) + .unwrap(); + assert_eq!(summary.baseline_model.as_ref().unwrap().model, "baseline"); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + } + #[test] fn unpriced_summary_is_partial_without_losing_tokens() { let recorder = LlmOptimizationRecorder::default(); From cc6573896af71f56660d368e5fbd5b35a9823d11 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Wed, 8 Jul 2026 16:39:45 -0600 Subject: [PATCH 04/10] test: initialize optimization summaries in fixtures Signed-off-by: Bryan Bednarski --- crates/adaptive/tests/coverage/subscriber_tests.rs | 1 + crates/adaptive/tests/unit/drain_tests.rs | 2 ++ crates/ffi/tests/unit/types_tests.rs | 1 + crates/python/tests/coverage/py_types_coverage_tests.rs | 4 ++++ 4 files changed, 8 insertions(+) diff --git a/crates/adaptive/tests/coverage/subscriber_tests.rs b/crates/adaptive/tests/coverage/subscriber_tests.rs index 045625686..6aaf6651b 100644 --- a/crates/adaptive/tests/coverage/subscriber_tests.rs +++ b/crates/adaptive/tests/coverage/subscriber_tests.rs @@ -133,6 +133,7 @@ fn test_event_to_call_record_llm_end_with_annotated_response_stays_observability finish_reason: Some(FinishReason::Complete), usage: None, api_specific: None, + optimization_summary: None, extra: serde_json::Map::new(), })) .build(), diff --git a/crates/adaptive/tests/unit/drain_tests.rs b/crates/adaptive/tests/unit/drain_tests.rs index 25c8b9b0d..c6e10658d 100644 --- a/crates/adaptive/tests/unit/drain_tests.rs +++ b/crates/adaptive/tests/unit/drain_tests.rs @@ -873,6 +873,7 @@ fn test_accumulator_extracts_annotated_response() { cost: None, }), api_specific: None, + optimization_summary: None, extra: serde_json::Map::new(), }; @@ -990,6 +991,7 @@ fn test_accumulator_annotated_response_partial_data() { finish_reason: None, usage: None, api_specific: None, + optimization_summary: None, extra: serde_json::Map::new(), }; diff --git a/crates/ffi/tests/unit/types_tests.rs b/crates/ffi/tests/unit/types_tests.rs index 7eb87f057..6ba6a0cfe 100644 --- a/crates/ffi/tests/unit/types_tests.rs +++ b/crates/ffi/tests/unit/types_tests.rs @@ -636,6 +636,7 @@ fn test_annotated_event_accessors_and_codec_handles() { }), }), api_specific: None, + optimization_summary: None, extra: serde_json::Map::from_iter([("trace".into(), json!(true))]), }; let llm_end = make_scope_event(ScopeEventFixture { diff --git a/crates/python/tests/coverage/py_types_coverage_tests.rs b/crates/python/tests/coverage/py_types_coverage_tests.rs index 3198eac80..fb59353fc 100644 --- a/crates/python/tests/coverage/py_types_coverage_tests.rs +++ b/crates/python/tests/coverage/py_types_coverage_tests.rs @@ -623,6 +623,7 @@ fn test_stream_request_event_and_handle_wrappers_cover_remaining_methods() { api_name: "custom".into(), data: json!({"ok": true}), }), + optimization_summary: None, extra: serde_json::Map::from_iter([("extra".into(), json!(true))]), }; @@ -1217,6 +1218,7 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { api_name: "custom".into(), data: json!({"debug": true}), }), + optimization_summary: None, extra: serde_json::Map::from_iter([("trace".into(), json!("abc"))]), }, }; @@ -1260,6 +1262,7 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { finish_reason: None, usage: None, api_specific: None, + optimization_summary: None, extra: serde_json::Map::new(), }, }; @@ -1478,6 +1481,7 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { api_name: "custom".into(), data: json!({"debug": true}), }), + optimization_summary: None, extra: serde_json::Map::new(), }, }; From 968fdda16461b81b8649c6def8f3d10442441bbb Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 00:37:04 -0600 Subject: [PATCH 05/10] fix: harden optimization accounting lifecycle Signed-off-by: Bryan Bednarski --- crates/core/src/api/llm.rs | 205 ++- crates/core/src/api/optimization.rs | 951 +++++------ crates/core/src/api/runtime/state.rs | 37 +- .../src/api/runtime/subscriber_dispatcher.rs | 12 +- crates/core/src/api/shared.rs | 25 +- crates/core/src/stream.rs | 45 +- .../tests/integration/middleware_tests.rs | 282 +++- .../core/tests/integration/pipeline_tests.rs | 231 ++- crates/core/tests/integration/stream_tests.rs | 163 ++ crates/core/tests/unit/llm_api_tests.rs | 97 +- crates/core/tests/unit/optimization_tests.rs | 1390 +++++++++++++++++ crates/core/tests/unit/shared_tests.rs | 71 + crates/core/tests/unit/stream_tests.rs | 21 + crates/types/src/codec/optimization.rs | 70 +- 14 files changed, 2980 insertions(+), 620 deletions(-) create mode 100644 crates/core/tests/unit/optimization_tests.rs diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index f493ce636..f7aee3f72 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -16,18 +16,18 @@ use crate::api::optimization::{ LlmOptimizationRecorder, finalize_optimization_summary, scope_llm_optimization_recorder, }; use crate::api::runtime::NemoRelayContextState; -use crate::api::runtime::current_scope_stack; use crate::api::runtime::global_context; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmStreamExecutionNextFn, }; +use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; use crate::api::scope::{EmitMarkEventParams, ScopeHandle}; use crate::api::shared::{ ensure_runtime_owner, inject_dynamo_session_ids, metadata_with_otel_status, - resolve_parent_uuid, run_request_intercepts_with_codec, sanitize_event, - snapshot_event_subscribers, + resolve_parent_uuid, run_request_intercepts_with_codec_and_recorder, + sanitize_event_with_scope_stack, snapshot_event_subscribers, }; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider}; @@ -38,6 +38,21 @@ use crate::stream::LlmStreamWrapper; pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; +#[derive(Clone)] +struct CapturedLlmScopeStack(ScopeStackHandle); + +impl Default for CapturedLlmScopeStack { + fn default() -> Self { + Self(current_scope_stack()) + } +} + +impl std::fmt::Debug for CapturedLlmScopeStack { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("CapturedLlmScopeStack(..)") + } +} + /// Runtime-owned handle identifying an active or completed LLM call. #[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)] #[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] @@ -74,6 +89,19 @@ pub struct LlmHandle { #[serde(skip, default)] #[builder(default)] pub optimization_recorder: LlmOptimizationRecorder, + /// Scope stack captured when the LLM lifecycle starts. + /// + /// Close-time work can run from a different task, especially for streams, + /// so optimization marks must not consult the poller's ambient scope. + #[serde(skip, default)] + #[builder(setter(skip), default)] + captured_scope_stack: CapturedLlmScopeStack, +} + +impl LlmHandle { + pub(crate) fn captured_scope_stack(&self) -> &ScopeStackHandle { + &self.captured_scope_stack.0 + } } /// Builder parameters for [`NemoRelayContextState::create_llm_handle`]. @@ -278,7 +306,7 @@ fn emit_llm_start( ) -> Result<()> { ensure_runtime_owner()?; let subscribers = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; @@ -300,7 +328,7 @@ fn emit_llm_start_with_subscribers( ) -> Result<()> { ensure_runtime_owner()?; let entries = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); let scope_locals = scope_guard.collect_scope_local_registries(|registries| { ®istries.llm_sanitize_request_guardrails @@ -330,7 +358,7 @@ fn emit_llm_start_with_subscribers( .map_err(|error| FlowError::Internal(error.to_string()))?; state.build_llm_start_event(handle, Some(input), annotated_request) }; - if let Some(event) = sanitize_event(event) { + if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, subscribers); } Ok(()) @@ -358,22 +386,36 @@ fn emit_pending_request_marks( mark.category, mark.category_profile, )); - if let Some(event) = sanitize_event(event) { + if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, subscribers); } } Ok(()) } -pub(crate) fn emit_optimization_marks( +pub(crate) fn emit_optimization_marks(handle: &LlmHandle, subscribers: &[EventSubscriberFn]) { + emit_optimization_marks_with( + handle, + subscribers, + |event| sanitize_event_with_scope_stack(event, handle.captured_scope_stack()), + |event, subscribers| NemoRelayContextState::try_emit_event(event, subscribers), + ); +} + +fn emit_optimization_marks_with( handle: &LlmHandle, subscribers: &[EventSubscriberFn], -) -> Result<()> { - let contributions = handle.optimization_recorder.take_unemitted(); + mut sanitize: impl FnMut(Event) -> Option, + mut enqueue: impl FnMut(&Event, &[EventSubscriberFn]) -> bool, +) { + let contributions = handle.optimization_recorder.unemitted(); if contributions.is_empty() { - return Ok(()); + return; + } + if let Err(error) = ensure_runtime_owner() { + eprintln!("nemo_relay: unable to emit LLM optimization marks: {error}"); + return; } - ensure_runtime_owner()?; for contribution in contributions { let offset = contribution.sequence.unwrap_or(0).saturating_add(2); let offset = i64::try_from(offset).unwrap_or(i64::MAX); @@ -396,9 +438,21 @@ pub(crate) fn emit_optimization_marks( .build(), ), )); - NemoRelayContextState::emit_event(&event, subscribers); + let Some(event) = sanitize(event) else { + // Sanitizers currently rewrite fields rather than intentionally + // dropping events. `None` means the sanitizer context was + // unavailable, so preserve this ordered suffix for a later retry. + break; + }; + if enqueue(&event, subscribers) { + handle.optimization_recorder.mark_emitted(1); + } else { + // Preserve this item and the remaining ordered suffix for a later + // lifecycle boundary. Accounting remains best effort and must not + // alter the provider result. + break; + } } - Ok(()) } /// Start a manual LLM lifecycle span. @@ -510,7 +564,7 @@ fn llm_call_end_with_behavior( } = params; ensure_runtime_owner()?; let (entries, subscribers) = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); let scope_locals = scope_guard.collect_scope_local_registries(|registries| { ®istries.llm_sanitize_response_guardrails @@ -553,7 +607,8 @@ fn llm_call_end_with_behavior( _ => None, }, }; - emit_optimization_marks(handle, &subscribers)?; + handle.optimization_recorder.close_for_finalization(None); + emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); let summary = finalize_optimization_summary( &handle.optimization_recorder, @@ -586,7 +641,7 @@ fn llm_call_end_with_behavior( .build(), ) }; - if let Some(event) = sanitize_event(event) { + if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, &subscribers); } if let Some(error) = decode_error @@ -605,7 +660,7 @@ fn emit_llm_end_without_output( ) -> Result<()> { ensure_runtime_owner()?; let subscribers = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); let scope_subscribers = scope_guard.collect_scope_local_subscribers(); match lifecycle_subscribers { @@ -613,7 +668,8 @@ fn emit_llm_end_without_output( None => snapshot_event_subscribers(scope_subscribers)?, } }; - emit_optimization_marks(handle, &subscribers)?; + handle.optimization_recorder.close_for_finalization(None); + emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); let annotated_response = finalize_optimization_summary( &handle.optimization_recorder, @@ -634,7 +690,7 @@ fn emit_llm_end_without_output( .map_err(|error| FlowError::Internal(error.to_string()))?; state.end_llm_handle(handle, handle.data.clone(), metadata, annotated_response) }; - if let Some(event) = sanitize_event(event) { + if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, &subscribers); } Ok(()) @@ -738,10 +794,19 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { } let request_codec = codec.clone(); + let optimization_recorder = LlmOptimizationRecorder::default(); let (intercepted_request, annotated_request, pending_marks, optimization_contributions) = - run_request_intercepts_with_codec(&name, request, codec)?; + scope_llm_optimization_recorder(optimization_recorder.clone(), async { + run_request_intercepts_with_codec_and_recorder( + &name, + request, + codec, + &optimization_recorder, + ) + }) + .await?; - let handle = create_llm_handle( + let mut handle = create_llm_handle( CreateLlmHandleParams::builder() .name(name.as_str()) .parent_uuid_opt(resolve_parent_uuid(parent.as_ref())) @@ -751,8 +816,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { .model_name_opt(model_name) .build(), )?; + handle.optimization_recorder = optimization_recorder; let lifecycle_subscribers = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; @@ -767,26 +833,28 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { handle .optimization_recorder .record_all(optimization_contributions); - emit_optimization_marks(&handle, &lifecycle_subscribers)?; + emit_optimization_marks(&handle, &lifecycle_subscribers); - let execution = { - let scope_stack = current_scope_stack(); - let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); - let scope_locals = scope_guard - .collect_scope_local_registries(|registries| ®istries.llm_execution_intercepts); - let context = global_context(); - let state = context - .read() - .map_err(|error| FlowError::Internal(error.to_string()))?; - state.llm_build_execution_chain(&name, func, &scope_locals) - }; + let execution_name = name.clone(); + let execution = + scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { + let execution = { + let scope_stack = current_scope_stack(); + let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + let scope_locals = scope_guard.collect_scope_local_registries(|registries| { + ®istries.llm_execution_intercepts + }); + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + state.llm_build_execution_chain(&execution_name, func, &scope_locals) + }; + execution(intercepted_request).await + }) + .await; - match scope_llm_optimization_recorder( - handle.optimization_recorder.clone(), - execution(intercepted_request), - ) - .await - { + match execution { Ok(response) => { llm_call_end_with_behavior( LlmCallEndParams::builder() @@ -912,10 +980,19 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu } let request_codec = codec.clone(); + let optimization_recorder = LlmOptimizationRecorder::default(); let (intercepted_request, annotated_request, pending_marks, optimization_contributions) = - run_request_intercepts_with_codec(&name, request, codec)?; + scope_llm_optimization_recorder(optimization_recorder.clone(), async { + run_request_intercepts_with_codec_and_recorder( + &name, + request, + codec, + &optimization_recorder, + ) + }) + .await?; - let handle = create_llm_handle( + let mut handle = create_llm_handle( CreateLlmHandleParams::builder() .name(name.as_str()) .parent_uuid_opt(resolve_parent_uuid(parent.as_ref())) @@ -925,8 +1002,9 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu .model_name_opt(model_name) .build(), )?; + handle.optimization_recorder = optimization_recorder; let lifecycle_subscribers = { - let scope_stack = current_scope_stack(); + let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? }; @@ -941,27 +1019,28 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu handle .optimization_recorder .record_all(optimization_contributions); - emit_optimization_marks(&handle, &lifecycle_subscribers)?; + emit_optimization_marks(&handle, &lifecycle_subscribers); - let execution = { - let scope_stack = current_scope_stack(); - let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); - let scope_locals = scope_guard.collect_scope_local_registries(|registries| { - ®istries.llm_stream_execution_intercepts - }); - let context = global_context(); - let state = context - .read() - .map_err(|error| FlowError::Internal(error.to_string()))?; - state.llm_stream_build_execution_chain(&name, func, &scope_locals) - }; + let execution_name = name.clone(); + let execution = + scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { + let execution = { + let scope_stack = current_scope_stack(); + let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + let scope_locals = scope_guard.collect_scope_local_registries(|registries| { + ®istries.llm_stream_execution_intercepts + }); + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + state.llm_stream_build_execution_chain(&execution_name, func, &scope_locals) + }; + execution(intercepted_request).await + }) + .await; - match scope_llm_optimization_recorder( - handle.optimization_recorder.clone(), - execution(intercepted_request), - ) - .await - { + match execution { Ok(raw_stream) => { let wrapper = LlmStreamWrapper::new_managed( raw_stream, diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index 3fbccb26e..cf8880457 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -3,28 +3,39 @@ //! Managed, bounded LLM optimization accounting. +use std::collections::BTreeSet; use std::sync::{Arc, Mutex}; +use serde::Serialize; use uuid::Uuid; use crate::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationModel, LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokens, }; -use crate::codec::response::{AnnotatedLlmResponse, PricingResolver}; +use crate::codec::response::{AnnotatedLlmResponse, CostSource, PricingResolver}; /// Maximum contributions retained for one LLM call. pub const MAX_LLM_OPTIMIZATION_CONTRIBUTIONS: usize = 64; -/// Maximum serialized custom payload size for one contribution. -pub const MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES: usize = 16 * 1024; -/// Maximum aggregate serialized custom payload size for one LLM call. -pub const MAX_LLM_OPTIMIZATION_TOTAL_PAYLOAD_BYTES: usize = 256 * 1024; +/// Maximum serialized size of one complete contribution envelope. +pub const MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES: usize = 16 * 1024; +/// Maximum aggregate serialized size of all contribution envelopes for one call. +pub const MAX_LLM_OPTIMIZATION_TOTAL_CONTRIBUTION_BYTES: usize = 256 * 1024; +/// Maximum contribution records inspected before the recorder seals itself. +/// +/// Invalid records count toward this bound even though they do not consume an +/// accepted sequence number. +pub const MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS: usize = 64; #[derive(Debug, Default)] struct AccumulatorState { contributions: Vec, - total_payload_bytes: usize, + total_contribution_bytes: usize, + attempted_contributions: usize, emitted: usize, + closed: bool, + finished: bool, + limitations: BTreeSet, contribution_limit_exceeded: bool, invalid_payload_schema: bool, } @@ -42,51 +53,106 @@ impl LlmOptimizationRecorder { /// Record one contribution without blocking on I/O or exporter delivery. /// /// Returns `false` when the contribution is rejected by a payload/schema - /// invariant or a per-call bound. Rejection never affects LLM execution. + /// 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 { - let payload_bytes = match contribution.payload.as_ref() { + let Ok(mut state) = self.state.lock() else { + return false; + }; + if state.closed { + return false; + } + if state.attempted_contributions >= MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS { + seal_for_contribution_limit(&mut state); + 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() { + if let Ok(mut state) = self.state.lock() + && !state.closed + { state.invalid_payload_schema = true; } return false; } - Some(payload) => match bounded_json_size(payload, MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES) { - Ok(size) => size, - Err(PayloadSizeError::LimitExceeded) => { - if let Ok(mut state) = self.state.lock() { - state.contribution_limit_exceeded = 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; } - Err(PayloadSizeError::Serialization) => { - if let Ok(mut state) = self.state.lock() { - state.invalid_payload_schema = true; - } + if state.contributions.len() >= MAX_LLM_OPTIMIZATION_CONTRIBUTIONS { + drop(state); + self.note_contribution_limit_exceeded(); return false; } - }, - None => 0, - }; + state.contributions.len() as u64 + }; + contribution.sequence = Some(sequence); + + 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.contributions.len() >= MAX_LLM_OPTIMIZATION_CONTRIBUTIONS - || payload_bytes > MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES - || state.total_payload_bytes.saturating_add(payload_bytes) - > MAX_LLM_OPTIMIZATION_TOTAL_PAYLOAD_BYTES - { - state.contribution_limit_exceeded = 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; + } + 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; + } + + state.total_contribution_bytes = total_contribution_bytes; + state.contributions.push(contribution); + return true; } + } - contribution.id = Some(Uuid::now_v7()); - contribution.sequence = Some(state.contributions.len() as u64); - state.total_payload_bytes += payload_bytes; - state.contributions.push(contribution); - true + fn note_contribution_limit_exceeded(&self) { + if let Ok(mut state) = self.state.lock() + && !state.closed + { + seal_for_contribution_limit(&mut state); + } } pub(crate) fn record_all( @@ -94,18 +160,68 @@ impl LlmOptimizationRecorder { contributions: impl IntoIterator, ) { for contribution in contributions { - let _ = self.record(contribution); + if !self.record(contribution) && self.is_closed() { + break; + } } } - pub(crate) fn take_unemitted(&self) -> Vec { - let Ok(mut state) = self.state.lock() else { + fn is_closed(&self) -> bool { + self.state.lock().map(|state| state.closed).unwrap_or(true) + } + + /// Snapshot contributions not yet accepted by mark delivery. + /// + /// This does not move the cursor. Call [`Self::mark_emitted`] only after + /// the asynchronous dispatcher accepts an item. + pub(crate) fn unemitted(&self) -> Vec { + let Ok(state) = self.state.lock() else { return Vec::new(); }; let start = state.emitted.min(state.contributions.len()); - let contributions = state.contributions[start..].to_vec(); - state.emitted = state.contributions.len(); - contributions + state.contributions[start..].to_vec() + } + + /// 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 { + return; + }; + state.emitted = state + .emitted + .saturating_add(count) + .min(state.contributions.len()); + } + + /// 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() + && !state.closed + { + state.limitations.insert(limitation.into()); + } + } + + /// Atomically seal contribution acceptance at an LLM close boundary. + /// + /// When `conditional_limitation` is supplied, it is added only if the call + /// already has optimization evidence or accounting limitations. This keeps + /// 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 { + return false; + }; + if state.finished { + return false; + } + let has_evidence = state.has_evidence(); + if has_evidence && let Some(limitation) = conditional_limitation { + state.limitations.insert(limitation.to_string()); + } + state.closed = true; + has_evidence } fn finish(&self) -> FinishedContributions { @@ -115,12 +231,24 @@ impl LlmOptimizationRecorder { limitations: vec!["optimization_accumulator_unavailable".to_string()], }; }; - let mut limitations = Vec::new(); + if state.finished { + return FinishedContributions { + contributions: Vec::new(), + limitations: Vec::new(), + }; + } + state.closed = true; + state.finished = true; + let mut limitations = std::mem::take(&mut state.limitations) + .into_iter() + .collect::>(); if state.contribution_limit_exceeded { limitations.push("contribution_limit_exceeded".to_string()); + state.contribution_limit_exceeded = false; } if state.invalid_payload_schema { limitations.push("invalid_contribution_payload_schema".to_string()); + state.invalid_payload_schema = false; } FinishedContributions { contributions: std::mem::take(&mut state.contributions), @@ -129,12 +257,27 @@ impl LlmOptimizationRecorder { } } -enum PayloadSizeError { +impl AccumulatorState { + fn has_evidence(&self) -> bool { + !self.contributions.is_empty() + || !self.limitations.is_empty() + || self.contribution_limit_exceeded + || self.invalid_payload_schema + } +} + +fn seal_for_contribution_limit(state: &mut AccumulatorState) { + state.contribution_limit_exceeded = true; + state.closed = true; +} + +#[derive(Debug)] +enum SerializedSizeError { LimitExceeded, Serialization, } -fn bounded_json_size(value: &serde_json::Value, limit: usize) -> Result { +fn bounded_json_size(value: &T, limit: usize) -> Result { struct CountingWriter { size: usize, limit: usize, @@ -145,7 +288,9 @@ fn bounded_json_size(value: &serde_json::Value, limit: usize) -> Result std::io::Result { if self.size.saturating_add(bytes.len()) > self.limit { self.exceeded = true; - return Err(std::io::Error::other("optimization payload limit exceeded")); + return Err(std::io::Error::other( + "optimization contribution limit exceeded", + )); } self.size += bytes.len(); Ok(bytes.len()) @@ -163,9 +308,9 @@ fn bounded_json_size(value: &serde_json::Value, limit: usize) -> Result( pub(crate) fn finalize_optimization_summary( recorder: &LlmOptimizationRecorder, - response: Option<&mut AnnotatedLlmResponse>, + mut response: Option<&mut AnnotatedLlmResponse>, requested_model: Option<&str>, pricing: &PricingResolver, ) -> Option { @@ -214,33 +359,55 @@ pub(crate) fn finalize_optimization_summary( return None; } - let mut tokens_saved = LlmOptimizationTokens::default(); - let mut baseline_model = None; - let mut contributed_effective_model = None; + let applied_routing = finished + .contributions + .iter() + .filter(|contribution| contribution.applied) + .filter(|contribution| { + contribution.kind.as_str() + == crate::codec::optimization::LlmOptimizationKind::MODEL_ROUTING + }) + .collect::>(); + let mut limitations = finished.limitations; + let routing_ambiguous = applied_routing.len() > 1; + if routing_ambiguous { + limitations.push("multiple_routing_contributions".to_string()); + } + + let mut token_totals = CheckedTokenTotals::default(); for contribution in finished .contributions .iter() .filter(|contribution| contribution.applied) { + let is_routing = contribution.kind.as_str() + == crate::codec::optimization::LlmOptimizationKind::MODEL_ROUTING; + if routing_ambiguous && is_routing { + continue; + } if let Some(saved) = contribution .token_impact .as_ref() .and_then(|impact| impact.saved.as_ref()) { - tokens_saved.add_assign(saved); - } - if contribution.kind.as_str() - == crate::codec::optimization::LlmOptimizationKind::MODEL_ROUTING - && let Some(transition) = contribution.model_transition.as_ref() - { - if baseline_model.is_none() { - baseline_model = transition.baseline.clone(); - } - if transition.effective.is_some() { - contributed_effective_model = transition.effective.clone(); - } + 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 @@ -253,84 +420,159 @@ pub(crate) fn finalize_optimization_summary( .map(|model| LlmOptimizationModel::new(model.clone())) }) .or_else(|| requested_model.map(LlmOptimizationModel::new)); - if baseline_model.is_none() { + if (applied_routing.is_empty() || routing_ambiguous) && baseline_model.is_none() { baseline_model = effective_model.clone(); } - let effective_usage = response + let mut effective_usage = response .as_ref() .and_then(|response| response.usage.clone()); + 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(inferred), Some(response_usage)) = ( + effective_usage + .as_ref() + .and_then(|usage| usage.total_tokens), + response + .as_mut() + .and_then(|response| response.usage.as_mut()), + ) && response_usage.total_tokens.is_none() + { + response_usage.total_tokens = Some(inferred); + } let baseline_usage = effective_usage.as_ref().map(|usage| { let mut baseline = usage.clone(); baseline.cost = None; - add_tokens(&mut baseline.prompt_tokens, tokens_saved.prompt_tokens); - add_tokens( + token_count_overflow |= checked_add_observed_tokens( + &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, ); - add_tokens( + 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, ); - add_tokens( + 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, ); - let total_saved = tokens_saved - .total_tokens - .or_else(|| option_sum([tokens_saved.prompt_tokens, tokens_saved.completion_tokens])); - add_tokens(&mut baseline.total_tokens, total_saved); baseline }); + if token_count_overflow { + limitations.push("token_count_overflow".to_string()); + } - let actual_cost = effective_usage + // 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. + let provider_reported_cost = effective_usage .as_ref() - .and_then(|usage| usage.cost.clone()) - .or_else(|| { - 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 baseline_cost = baseline_model.as_ref().and_then(|model| { - pricing.estimate_cost_for_provider( - model.provider.as_deref(), - &model.model, - baseline_usage.as_ref()?, - ) + .and_then(|usage| usage.cost.as_ref()) + .filter(|cost| cost.source == CostSource::ProviderReported) + .cloned(); + let complete_core_usage = effective_usage + .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) }); + 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()) + { + usage.cost.clone_from(&actual_cost); + } + + let baseline_cost = + (!token_count_overflow && !baseline_derivation_incomplete && complete_core_usage) + .then_some(baseline_model.as_ref()) + .flatten() + .and_then(|model| { + pricing.estimate_cost_for_provider( + model.provider.as_deref(), + &model.model, + baseline_usage.as_ref()?, + ) + }); - let mut limitations = finished.limitations; if effective_usage.is_none() { limitations.push("missing_effective_usage".to_string()); } + if effective_model.is_none() { + limitations.push("missing_effective_model".to_string()); + } if baseline_model.is_none() { limitations.push("missing_baseline_model".to_string()); } - if baseline_cost.is_none() { + if baseline_cost.is_none() + && baseline_model.is_some() + && !token_count_overflow + && !baseline_derivation_incomplete + && complete_core_usage + { limitations.push("missing_baseline_pricing".to_string()); } if actual_cost.is_none() { limitations.push("missing_actual_cost".to_string()); } - let (estimated_cost_saved, currency) = match (&baseline_cost, &actual_cost) { - (Some(baseline), Some(actual)) - if baseline.currency.eq_ignore_ascii_case(&actual.currency) => - { - ( - baseline - .total_or_component_sum() - .zip(actual.total_or_component_sum()) - .map(|(baseline, actual)| baseline - actual), - Some(baseline.currency.clone()), - ) - } - (Some(_), Some(_)) => { - limitations.push("cost_currency_mismatch".to_string()); - (None, None) - } - _ => (None, None), - }; + let (estimated_cost_saved, currency) = calculate_estimated_cost_saved( + baseline_cost.as_ref(), + actual_cost.as_ref(), + &mut limitations, + ); limitations.sort(); limitations.dedup(); @@ -360,369 +602,174 @@ pub(crate) fn finalize_optimization_summary( Some(summary) } -fn add_tokens(target: &mut Option, value: Option) { - if let Some(value) = value { - *target = Some(target.unwrap_or(0).saturating_add(value)); +fn calculate_estimated_cost_saved( + baseline_cost: Option<&crate::codec::response::CostEstimate>, + actual_cost: Option<&crate::codec::response::CostEstimate>, + limitations: &mut Vec, +) -> (Option, Option) { + let baseline_total = baseline_cost.and_then(|cost| cost.total_or_component_sum()); + let actual_total = actual_cost.and_then(|cost| cost.total_or_component_sum()); + if baseline_cost.is_some() && baseline_total.is_none() { + limitations.push("missing_baseline_cost_total".to_string()); + } + if actual_cost.is_some() && actual_total.is_none() { + limitations.push("missing_actual_cost_total".to_string()); + } + + match (baseline_cost, actual_cost) { + (Some(baseline), Some(actual)) + if baseline.currency.eq_ignore_ascii_case(&actual.currency) => + { + let saved = baseline_total + .zip(actual_total) + .map(|(baseline, actual)| baseline - actual); + let currency = saved.is_some().then(|| baseline.currency.clone()); + (saved, currency) + } + (Some(_), Some(_)) => { + limitations.push("cost_currency_mismatch".to_string()); + (None, None) + } + _ => (None, None), } } -fn option_sum(values: impl IntoIterator>) -> Option { - let mut present = false; - let total = values.into_iter().flatten().fold(0_u64, |total, value| { - present = true; - total.saturating_add(value) - }); - present.then_some(total) +#[derive(Debug, Clone, Copy, Default)] +struct TokenOverflow { + prompt: bool, + completion: bool, + cache_read: bool, + cache_write: bool, + total: bool, } -#[cfg(test)] -mod tests { - use serde_json::json; +impl TokenOverflow { + fn any(self) -> bool { + self.prompt || self.completion || self.cache_read || self.cache_write || self.total + } +} - use super::*; - use crate::api::event::DataSchema; - use crate::codec::optimization::{ - LlmOptimizationEvidenceQuality, LlmOptimizationModelTransition, LlmOptimizationTokenImpact, - }; - use crate::codec::response::{PricingCatalog, Usage}; - use crate::json::Json; - - fn resolver() -> PricingResolver { - resolver_with_rates(2.0, 1.0) - } - - fn resolver_with_rates(baseline_input: f64, effective_input: f64) -> PricingResolver { - let catalog = PricingCatalog::from_json_str( - &json!({ - "version": 1, - "entries": [ - {"provider":"test","model_id":"baseline","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":baseline_input,"output_per_million":4.0,"cache_read_per_million":0.5,"cache_write_per_million":3.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}}, - {"provider":"test","model_id":"effective","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":effective_input,"output_per_million":2.0,"cache_read_per_million":0.25,"cache_write_per_million":2.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}} - ] - }) - .to_string(), - ) - .unwrap(); - PricingResolver::from_catalogs(vec![catalog]) - } - - fn contribution() -> LlmOptimizationContribution { - let mut contribution = LlmOptimizationContribution::new( - "test.optimizer", - crate::codec::optimization::LlmOptimizationKind::model_routing(), +#[derive(Debug, Default)] +struct CheckedTokenTotals { + values: LlmOptimizationTokens, + overflow: TokenOverflow, + missing_total: bool, + inconsistent_total: bool, +} + +impl CheckedTokenTotals { + fn add_contribution(&mut self, other: &LlmOptimizationTokens) { + checked_accumulate( + &mut self.values.prompt_tokens, + other.prompt_tokens, + &mut self.overflow.prompt, ); - contribution.model_transition = Some(LlmOptimizationModelTransition { - baseline: Some(LlmOptimizationModel::new("baseline").with_provider("test")), - effective: Some(LlmOptimizationModel::new("effective").with_provider("test")), - }); - contribution.token_impact = Some(LlmOptimizationTokenImpact { - saved: Some(LlmOptimizationTokens::saved_prompt(200)), - quality: Some(LlmOptimizationEvidenceQuality::Estimated), - estimation_method: Some("test-tokenizer".to_string()), - ..LlmOptimizationTokenImpact::default() - }); - contribution - } - - #[test] - fn combined_summary_retains_token_evidence_and_snapshot_pricing() { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let mut response = AnnotatedLlmResponse { - model: Some("effective".to_string()), - usage: Some(Usage { - prompt_tokens: Some(800), - completion_tokens: Some(100), - total_tokens: Some(900), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let summary = finalize_optimization_summary( - &recorder, - Some(&mut response), - Some("baseline"), - &resolver(), - ) - .unwrap(); - assert_eq!(summary.status, LlmOptimizationSummaryStatus::Complete); - assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); - assert_eq!( - summary.baseline_usage.as_ref().unwrap().prompt_tokens, - Some(1000) + checked_accumulate( + &mut self.values.completion_tokens, + other.completion_tokens, + &mut self.overflow.completion, + ); + checked_accumulate( + &mut self.values.cache_read_tokens, + other.cache_read_tokens, + &mut self.overflow.cache_read, + ); + checked_accumulate( + &mut self.values.cache_write_tokens, + other.cache_write_tokens, + &mut self.overflow.cache_write, ); - assert_eq!(summary.baseline_cost.as_ref().unwrap().total, Some(0.0024)); - assert_eq!(summary.actual_cost.as_ref().unwrap().total, Some(0.001)); - assert!((summary.estimated_cost_saved.unwrap() - 0.0014).abs() < 1e-12); - } - - #[test] - fn applied_route_is_the_authoritative_effective_model() { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let mut response = AnnotatedLlmResponse { - // Providers may return an alias or deployment name rather than - // the exact model Relay selected and sent upstream. - model: Some("provider-response-alias".to_string()), - usage: Some(Usage { - prompt_tokens: Some(800), - completion_tokens: Some(100), - total_tokens: Some(900), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let summary = finalize_optimization_summary( - &recorder, - Some(&mut response), - Some("original-request-model"), - &resolver(), - ) - .unwrap(); - assert_eq!(summary.baseline_model.as_ref().unwrap().model, "baseline"); - assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); - } - - #[test] - fn unpriced_summary_is_partial_without_losing_tokens() { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let mut response = AnnotatedLlmResponse { - model: Some("effective".to_string()), - usage: Some(Usage { - prompt_tokens: Some(8), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let summary = finalize_optimization_summary( - &recorder, - Some(&mut response), - None, - &PricingResolver::default(), - ) - .unwrap(); - assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); - assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); - assert!(summary.estimated_cost_saved.is_none()); - } - - #[test] - fn zero_and_negative_savings_are_preserved() { - for (baseline_rate, effective_rate, expected_sign) in [(0.0, 0.0, 0_i8), (0.5, 2.0, -1_i8)] - { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let mut response = AnnotatedLlmResponse { - model: Some("effective".to_string()), - usage: Some(Usage { - prompt_tokens: Some(800), - total_tokens: Some(800), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let summary = finalize_optimization_summary( - &recorder, - Some(&mut response), - None, - &resolver_with_rates(baseline_rate, effective_rate), - ) - .unwrap(); - let saved = summary.estimated_cost_saved.unwrap(); - match expected_sign { - 0 => assert_eq!(saved, 0.0), - -1 => assert!(saved < 0.0), - _ => unreachable!(), - } - } - } - #[test] - fn multiple_contributions_and_cache_savings_aggregate_explicitly() { - let recorder = LlmOptimizationRecorder::default(); - for (producer, prompt, cache_read, cache_write) in - [("test.one", 5, 7, 0), ("test.two", 11, 13, 17)] + let (derived_total, derived_overflow) = + checked_option_sum([other.prompt_tokens, other.completion_tokens]); + self.overflow.total |= derived_overflow; + if derived_overflow { + self.values.total_tokens = None; + } + if let (Some(explicit), Some(prompt), Some(completion)) = ( + other.total_tokens, + other.prompt_tokens, + other.completion_tokens, + ) && prompt + .checked_add(completion) + .is_some_and(|derived| derived != explicit) { - let mut item = LlmOptimizationContribution::new( - producer, - crate::codec::optimization::LlmOptimizationKind::input_compression(), - ); - item.token_impact = Some(LlmOptimizationTokenImpact { - saved: Some(LlmOptimizationTokens { - prompt_tokens: Some(prompt), - cache_read_tokens: Some(cache_read), - cache_write_tokens: Some(cache_write), - total_tokens: Some(prompt), - ..LlmOptimizationTokens::default() - }), - ..LlmOptimizationTokenImpact::default() - }); - assert!(recorder.record(item)); + self.inconsistent_total = true; } - let mut response = AnnotatedLlmResponse { - model: Some("effective".to_string()), - usage: Some(Usage { - prompt_tokens: Some(100), - completion_tokens: Some(10), - total_tokens: Some(110), - cache_read_tokens: Some(20), - cache_write_tokens: Some(3), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let summary = - finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()) - .unwrap(); - assert_eq!(summary.tokens_saved.prompt_tokens, Some(16)); - assert_eq!(summary.tokens_saved.cache_read_tokens, Some(20)); - assert_eq!(summary.tokens_saved.cache_write_tokens, Some(17)); - assert_eq!( - summary.baseline_usage.as_ref().unwrap().cache_read_tokens, - Some(40) - ); - assert_eq!( - summary.baseline_usage.as_ref().unwrap().cache_write_tokens, - Some(20) - ); - assert_eq!(summary.contributions[0].sequence, Some(0)); - assert_eq!(summary.contributions[1].sequence, Some(1)); - } - - #[test] - fn serialized_summary_can_be_repriced_with_a_new_catalog() { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let mut response = AnnotatedLlmResponse { - model: Some("effective".to_string()), - usage: Some(Usage { - prompt_tokens: Some(800), - completion_tokens: Some(100), - total_tokens: Some(900), - ..Usage::default() - }), - ..AnnotatedLlmResponse::default() - }; - let original = - finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()) - .unwrap(); - let restored: LlmOptimizationSummary = - serde_json::from_value(serde_json::to_value(&original).unwrap()).unwrap(); - let newer = resolver_with_rates(10.0, 5.0); - let baseline = newer - .estimate_cost_for_provider( - Some("test"), - "baseline", - restored.baseline_usage.as_ref().unwrap(), - ) - .unwrap() - .total_or_component_sum() - .unwrap(); - let actual = newer - .estimate_cost_for_provider( - Some("test"), - "effective", - restored.effective_usage.as_ref().unwrap(), - ) - .unwrap() - .total_or_component_sum() - .unwrap(); - assert_ne!(baseline - actual, original.estimated_cost_saved.unwrap()); - assert_eq!(restored.tokens_saved.prompt_tokens, Some(200)); - } - - #[test] - fn no_usage_is_an_explicit_partial_summary() { - let recorder = LlmOptimizationRecorder::default(); - assert!(recorder.record(contribution())); - let summary = - finalize_optimization_summary(&recorder, None, Some("effective"), &resolver()).unwrap(); - assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); - assert!( - summary - .limitations - .contains(&"missing_effective_usage".to_string()) + let contribution_total = other.total_tokens.or(derived_total); + if contribution_total.is_none() { + self.missing_total = true; + } + checked_accumulate( + &mut self.values.total_tokens, + contribution_total, + &mut self.overflow.total, ); - assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); } +} - #[test] - fn payload_byte_limits_are_enforced_without_unbounded_work() { - let oversized = LlmOptimizationRecorder::default(); - let mut item = LlmOptimizationContribution::new("test", "custom"); - item.payload_schema = Some(DataSchema { - name: "test.payload".to_string(), - version: "1".to_string(), - }); - item.payload = Some(Json::String("x".repeat(MAX_LLM_OPTIMIZATION_PAYLOAD_BYTES))); - assert!(!oversized.record(item)); - - let aggregate = LlmOptimizationRecorder::default(); - for index in 0..17 { - let mut item = LlmOptimizationContribution::new(format!("test.{index}"), "custom"); - item.payload_schema = Some(DataSchema { - name: "test.payload".to_string(), - version: "1".to_string(), - }); - item.payload = Some(Json::String("x".repeat(15_000))); - assert!(aggregate.record(item)); +fn checked_accumulate(target: &mut Option, value: Option, overflowed: &mut bool) { + if *overflowed { + return; + } + let Some(value) = value else { + return; + }; + match target.unwrap_or(0).checked_add(value) { + Some(total) => *target = Some(total), + None => { + *target = None; + *overflowed = true; } - let mut overflow = LlmOptimizationContribution::new("test.overflow", "custom"); - overflow.payload_schema = Some(DataSchema { - name: "test.payload".to_string(), - version: "1".to_string(), - }); - overflow.payload = Some(Json::String("x".repeat(15_000))); - assert!(!aggregate.record(overflow)); - assert!( - aggregate - .finish() - .limitations - .contains(&"contribution_limit_exceeded".to_string()) - ); } +} - #[test] - fn bounds_and_invalid_payloads_are_best_effort_and_visible() { - let recorder = LlmOptimizationRecorder::default(); - let mut invalid = LlmOptimizationContribution::new("test", "custom"); - invalid.payload = Some(json!({"evidence": true})); - assert!(!recorder.record(invalid)); - for index in 0..MAX_LLM_OPTIMIZATION_CONTRIBUTIONS { - assert!(recorder.record(LlmOptimizationContribution::new( - format!("test.{index}"), - "custom" - ))); +fn checked_add_observed_tokens( + target: &mut Option, + value: Option, + value_overflowed: bool, + missing_limitation: &'static str, + limitations: &mut Vec, + baseline_derivation_incomplete: &mut bool, +) -> bool { + if value_overflowed { + *target = None; + *baseline_derivation_incomplete = true; + return true; + } + let Some(value) = value else { + return false; + }; + let Some(observed) = *target else { + limitations.push(missing_limitation.to_string()); + *baseline_derivation_incomplete = true; + return false; + }; + match observed.checked_add(value) { + Some(total) => { + *target = Some(total); + false + } + None => { + *target = None; + true } - assert!(!recorder.record(LlmOptimizationContribution::new("overflow", "custom"))); - let summary = - finalize_optimization_summary(&recorder, None, None, &PricingResolver::default()) - .unwrap(); - assert_eq!( - summary.contributions.len(), - MAX_LLM_OPTIMIZATION_CONTRIBUTIONS - ); - assert!( - summary - .limitations - .contains(&"contribution_limit_exceeded".to_string()) - ); - assert!( - summary - .limitations - .contains(&"invalid_contribution_payload_schema".to_string()) - ); } +} - #[tokio::test] - async fn recorder_can_be_captured_for_stream_commit() { - let recorder = LlmOptimizationRecorder::default(); - let captured = scope_llm_optimization_recorder(recorder.clone(), async { - current_llm_optimization_recorder().unwrap() - }) - .await; - assert!(captured.record(LlmOptimizationContribution::new("test.stream", "commit"))); - assert_eq!(recorder.finish().contributions.len(), 1); +fn checked_option_sum(values: impl IntoIterator>) -> (Option, bool) { + let mut present = false; + let mut total = 0_u64; + for value in values.into_iter().flatten() { + present = true; + let Some(next) = total.checked_add(value) else { + return (None, true); + }; + total = next; } + (present.then_some(total), false) } + +#[cfg(test)] +#[path = "../../tests/unit/optimization_tests.rs"] +mod tests; diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 68afc9882..60a01aa8e 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -190,7 +190,16 @@ impl NemoRelayContextState { /// - `event`: Fully constructed lifecycle event to deliver. /// - `subscribers`: Subscribers that should observe the event. pub(crate) fn emit_event(event: &Event, subscribers: &[EventSubscriberFn]) { - subscriber_dispatcher::dispatch_event(event, subscribers); + let _ = subscriber_dispatcher::dispatch_event(event, subscribers); + } + + /// Queue an event and report whether the asynchronous dispatcher accepted it. + /// + /// Subscriber callbacks still run asynchronously. This acknowledgement only + /// covers queue acceptance and is used by bounded observability cursors so a + /// transient dispatcher failure does not permanently discard evidence. + pub(crate) fn try_emit_event(event: &Event, subscribers: &[EventSubscriberFn]) -> bool { + subscriber_dispatcher::dispatch_event(event, subscribers) } /// Build a standalone mark event. @@ -1133,6 +1142,26 @@ impl NemoRelayContextState { annotated: Option, entries: &[Intercept], codec_active: bool, + ) -> crate::error::Result { + Self::llm_request_intercepts_snapshot_chain_with_recorder( + name, + request, + annotated, + entries, + codec_active, + None, + ) + } + + /// Run a request-intercept snapshot while ingesting optimization evidence + /// directly into the managed call's bounded accumulator. + pub(crate) fn llm_request_intercepts_snapshot_chain_with_recorder( + name: &str, + request: LlmRequest, + annotated: Option, + entries: &[Intercept], + codec_active: bool, + optimization_recorder: Option<&crate::api::optimization::LlmOptimizationRecorder>, ) -> crate::error::Result { let mut request_value = request; let mut annotated_value = annotated; @@ -1156,7 +1185,11 @@ impl NemoRelayContextState { request_value = outcome.request; annotated_value = outcome.annotated_request; pending_marks.extend(outcome.pending_marks); - optimization_contributions.extend(outcome.optimization_contributions); + if let Some(recorder) = optimization_recorder { + recorder.record_all(outcome.optimization_contributions); + } else { + optimization_contributions.extend(outcome.optimization_contributions); + } if entry.payload.break_chain { break; } diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs index a5be4f70e..2d741072c 100644 --- a/crates/core/src/api/runtime/subscriber_dispatcher.rs +++ b/crates/core/src/api/runtime/subscriber_dispatcher.rs @@ -38,9 +38,9 @@ mod native { static IN_DISPATCHER: Cell = const { Cell::new(false) }; } - pub(super) fn dispatch_event(event: &Event, subscribers: &[EventSubscriberFn]) { + pub(super) fn dispatch_event(event: &Event, subscribers: &[EventSubscriberFn]) -> bool { if subscribers.is_empty() { - return; + return true; } let message = DispatcherMessage::Deliver { event: Box::new(event.clone()), @@ -51,10 +51,14 @@ mod native { Ok(sender) => { if let Err(error) = sender.send(message) { eprintln!("nemo_relay: failed to queue subscriber event: {error}"); + false + } else { + true } } Err(error) => { eprintln!("nemo_relay: failed to start subscriber dispatcher: {error}"); + false } } } @@ -152,8 +156,8 @@ mod native { } /// Queue an event for subscriber delivery. -pub(crate) fn dispatch_event(event: &Event, subscribers: &[EventSubscriberFn]) { - native::dispatch_event(event, subscribers); +pub(crate) fn dispatch_event(event: &Event, subscribers: &[EventSubscriberFn]) -> bool { + native::dispatch_event(event, subscribers) } /// Wait for all queued subscriber callbacks submitted before this call. diff --git a/crates/core/src/api/shared.rs b/crates/core/src/api/shared.rs index a1f3fc5d0..a3baf051d 100644 --- a/crates/core/src/api/shared.rs +++ b/crates/core/src/api/shared.rs @@ -195,10 +195,31 @@ pub(crate) type InterceptedLlmRequest = ( Vec, ); +#[cfg(test)] pub(crate) fn run_request_intercepts_with_codec( name: &str, request: LlmRequest, codec: Option>, +) -> Result { + run_request_intercepts_with_codec_inner(name, request, codec, None) +} + +/// Run request intercepts and record optimization contributions directly into +/// the managed call's bounded accumulator as each intercept completes. +pub(crate) fn run_request_intercepts_with_codec_and_recorder( + name: &str, + request: LlmRequest, + codec: Option>, + recorder: &crate::api::optimization::LlmOptimizationRecorder, +) -> Result { + run_request_intercepts_with_codec_inner(name, request, codec, Some(recorder)) +} + +fn run_request_intercepts_with_codec_inner( + name: &str, + request: LlmRequest, + codec: Option>, + recorder: Option<&crate::api::optimization::LlmOptimizationRecorder>, ) -> Result { let annotated = match &codec { Some(codec) => Some(codec.decode(&request)?), @@ -218,13 +239,13 @@ pub(crate) fn run_request_intercepts_with_codec( state.llm_request_intercept_entries(&scope_locals) }; - let outcome = - crate::api::runtime::NemoRelayContextState::llm_request_intercepts_snapshot_chain( + let outcome = crate::api::runtime::NemoRelayContextState::llm_request_intercepts_snapshot_chain_with_recorder( name, request, annotated, &entries, codec.is_some(), + recorder, )?; let mut request = outcome.request; inject_dynamo_session_ids(&mut request); diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 2b380fd86..6d1d13dfb 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -132,10 +132,11 @@ impl LlmStreamWrapper { response_codec: Option>, subscribers: Vec, ) -> Self { + let scope_stack = handle.captured_scope_stack().clone(); Self { inner, handle, - scope_stack: current_scope_stack(), + scope_stack, collector, finalizer: Some(finalizer), response_codec, @@ -163,24 +164,34 @@ impl LlmStreamWrapper { return; } self.ended = true; - self.emit_end_event(self.metadata.clone()); + let metadata = metadata_with_otel_status( + self.metadata.clone(), + "ERROR", + Some("stream dropped before clean completion".to_string()), + ); + self.emit_end_event(metadata, true); } - fn finish_with_status(&mut self, status_code: &'static str, status_message: Option) { + fn finish_with_status( + &mut self, + status_code: &'static str, + status_message: Option, + interrupted: bool, + ) { if self.ended { return; } self.ended = true; let metadata = metadata_with_otel_status(self.metadata.clone(), status_code, status_message); - self.emit_end_event(metadata); + self.emit_end_event(metadata, interrupted); } /// Emit the LLM END event with aggregated response data. /// /// Calls the finalizer to produce the aggregated response, runs sanitize /// response guardrails, and emits the END event. - fn emit_end_event(&mut self, metadata: Option) { + fn emit_end_event(&mut self, metadata: Option, interrupted: bool) { let aggregated = match self.finalizer.take() { Some(finalizer) => finalizer(), None => Json::Null, @@ -216,7 +227,13 @@ impl LlmStreamWrapper { attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); Some(decoded) }); - let _ = emit_optimization_marks(&self.handle, &self.subscribers); + let interruption = (interrupted + && !has_authoritative_final_usage(annotated_response.as_ref())) + .then_some("stream_interrupted"); + self.handle + .optimization_recorder + .close_for_finalization(interruption); + emit_optimization_marks(&self.handle, &self.subscribers); let pricing = crate::codec::response::active_pricing_resolver(); let summary = finalize_optimization_summary( &self.handle.optimization_recorder, @@ -301,18 +318,18 @@ impl Stream for LlmStreamWrapper { Ok(()) => Poll::Ready(Some(Ok(raw_chunk))), Err(e) => { let message = e.to_string(); - this.finish_with_status("ERROR", Some(message)); + this.finish_with_status("ERROR", Some(message), true); Poll::Ready(Some(Err(e))) } } } Poll::Ready(Some(Err(e))) => { let message = e.to_string(); - this.finish_with_status("ERROR", Some(message)); + this.finish_with_status("ERROR", Some(message), true); Poll::Ready(Some(Err(e))) } Poll::Ready(None) => { - this.finish_with_status("OK", None); + this.finish_with_status("OK", None, false); Poll::Ready(None) } Poll::Pending => Poll::Pending, @@ -320,6 +337,16 @@ impl Stream for LlmStreamWrapper { } } +fn has_authoritative_final_usage(response: Option<&AnnotatedLlmResponse>) -> bool { + response.is_some_and(|response| { + response.finish_reason.is_some() + && response.usage.as_ref().is_some_and(|usage| { + usage.total_tokens.is_some() + || (usage.prompt_tokens.is_some() && usage.completion_tokens.is_some()) + }) + }) +} + fn llm_chunk_mark_data(chunk_index: u64, raw_chunk: &Json) -> Json { if let Some(data) = summarize_openai_chat_chunk(chunk_index, raw_chunk) { return data; diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 451eb3515..d8882f04a 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -15,7 +15,7 @@ use std::sync::{Arc, Mutex}; use futures::StreamExt; use nemo_relay::api::event::{ - CategoryProfile, Event, EventCategory, PendingMarkSpec, ScopeCategory, + CategoryProfile, DataSchema, Event, EventCategory, PendingMarkSpec, ScopeCategory, }; use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_request_intercepts, @@ -27,20 +27,22 @@ use nemo_relay::api::registry::{ deregister_llm_conditional_execution_guardrail, deregister_llm_execution_intercept, deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, deregister_llm_stream_execution_intercept, - deregister_mark_sanitize_guardrail, deregister_tool_conditional_execution_guardrail, - deregister_tool_execution_intercept, deregister_tool_request_intercept, - deregister_tool_sanitize_request_guardrail, deregister_tool_sanitize_response_guardrail, - register_llm_conditional_execution_guardrail, register_llm_execution_intercept, - register_llm_request_intercept, register_llm_sanitize_request_guardrail, - register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, - register_mark_sanitize_guardrail, register_tool_conditional_execution_guardrail, + deregister_mark_sanitize_guardrail, deregister_scope_sanitize_end_guardrail, + deregister_tool_conditional_execution_guardrail, deregister_tool_execution_intercept, + deregister_tool_request_intercept, deregister_tool_sanitize_request_guardrail, + deregister_tool_sanitize_response_guardrail, register_llm_conditional_execution_guardrail, + register_llm_execution_intercept, register_llm_request_intercept, + register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, + register_llm_stream_execution_intercept, register_mark_sanitize_guardrail, + register_scope_sanitize_end_guardrail, register_tool_conditional_execution_guardrail, register_tool_execution_intercept, register_tool_request_intercept, register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, scope_register_llm_conditional_execution_guardrail, scope_register_llm_execution_intercept, scope_register_llm_request_intercept, scope_register_llm_sanitize_request_guardrail, scope_register_llm_sanitize_response_guardrail, scope_register_llm_stream_execution_intercept, - scope_register_tool_conditional_execution_guardrail, scope_register_tool_execution_intercept, - scope_register_tool_request_intercept, scope_register_tool_sanitize_request_guardrail, + scope_register_mark_sanitize_guardrail, scope_register_tool_conditional_execution_guardrail, + scope_register_tool_execution_intercept, scope_register_tool_request_intercept, + scope_register_tool_sanitize_request_guardrail, scope_register_tool_sanitize_response_guardrail, }; use nemo_relay::api::runtime::NemoRelayContextState; @@ -61,6 +63,7 @@ use nemo_relay::codec::optimization::{ LlmOptimizationTokens, }; use nemo_relay::error::FlowError; +use nemo_relay::json::Json; #[cfg(all(feature = "otel", feature = "openinference"))] use nemo_relay::observability::MarkProjection; #[cfg(all(feature = "otel", feature = "openinference"))] @@ -697,7 +700,6 @@ async fn test_tool_execution_outcome_marks_follow_end_with_tool_parentage() { }), ) .unwrap(); - let mut plugin_ctx = PluginRegistrationContext::new(); plugin_ctx .register_tool_execution_intercept( @@ -3327,7 +3329,6 @@ async fn test_managed_llm_emits_pending_marks_under_started_scope() { }), ) .unwrap(); - register_llm_request_intercept( "pending_managed", 1, @@ -3441,6 +3442,37 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), ) .unwrap(); + register_mark_sanitize_guardrail( + "optimization_sanitizer", + 1, + Arc::new(|event, mut fields| { + if event.name() == "nemo_relay.llm.optimization" + && let Some(data) = fields.data.as_mut().and_then(Json::as_object_mut) + { + data.insert("payload".to_string(), json!({"secret": "[redacted]"})); + data.remove("future_secret"); + } + fields + }), + ) + .unwrap(); + register_scope_sanitize_end_guardrail( + "optimization_end_sanitizer", + 1, + Arc::new(|event, mut fields| { + if event.name() == "optimized-managed-llm" + && let Some(profile) = fields.category_profile.as_mut() + && let Some(response) = profile.annotated_response.as_mut() + && let Some(summary) = Arc::make_mut(response).optimization_summary.as_mut() + && let Some(contribution) = summary.contributions.first_mut() + { + contribution.payload = Some(json!({"secret": "[scope-end-redacted]"})); + contribution.extra.remove("future_secret"); + } + fields + }), + ) + .unwrap(); register_llm_request_intercept( "optimization_contributor", @@ -3455,6 +3487,14 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { estimation_method: Some("test-counter".to_string()), ..LlmOptimizationTokenImpact::default() }); + contribution.payload_schema = Some(DataSchema { + name: "test.optimizer_evidence".to_string(), + version: "1".to_string(), + }); + contribution.payload = Some(json!({"secret": "classified"})); + contribution + .extra + .insert("future_secret".to_string(), json!("classified")); Ok(LlmRequestInterceptOutcome::new(request, annotated) .with_optimization_contribution(contribution)) }), @@ -3464,12 +3504,10 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { "optimization_execution_contributor", 1, Arc::new(|_name, request, next| { - Box::pin(async move { - let contribution = - LlmOptimizationContribution::new("test.execution", "test_execution_kind"); - assert!(record_llm_optimization_contribution(contribution)); - next(request).await - }) + let contribution = + LlmOptimizationContribution::new("test.execution", "test_execution_kind"); + assert!(record_llm_optimization_contribution(contribution)); + Box::pin(async move { next(request).await }) }), ) .unwrap(); @@ -3516,6 +3554,8 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { ); assert_eq!(marks[0].data().unwrap()["sequence"], 0); assert_eq!(marks[1].data().unwrap()["sequence"], 1); + assert_eq!(marks[0].data().unwrap()["payload"]["secret"], "[redacted]"); + assert!(marks[0].data().unwrap().get("future_secret").is_none()); let end = captured .iter() @@ -3533,13 +3573,219 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { assert_eq!(summary.tokens_saved.prompt_tokens, Some(12)); assert_eq!(summary.contributions.len(), 2); assert_eq!(summary.contributions[0].producer, "test.optimizer"); + assert_eq!( + summary.contributions[0].payload.as_ref().unwrap()["secret"], + "[scope-end-redacted]" + ); + assert!(!summary.contributions[0].extra.contains_key("future_secret")); assert_eq!(summary.contributions[1].producer, "test.execution"); deregister_llm_request_intercept("optimization_contributor").unwrap(); deregister_llm_execution_intercept("optimization_execution_contributor").unwrap(); + deregister_mark_sanitize_guardrail("optimization_sanitizer").unwrap(); + deregister_scope_sanitize_end_guardrail("optimization_end_sanitizer").unwrap(); deregister_subscriber("optimization_observer").unwrap(); } +#[tokio::test] +async fn test_stream_optimization_mark_uses_the_llm_captured_sanitizer_scope() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let original_stack = current_scope_stack(); + let owner = push_scope( + nemo_relay::api::scope::PushScopeParams::builder() + .name("optimization-sanitizer-owner") + .scope_type(ScopeType::Agent) + .build(), + ) + .unwrap(); + scope_register_mark_sanitize_guardrail( + &owner.uuid, + "stream-optimization-sanitizer", + 1, + Arc::new(|event, mut fields| { + if event.name() == "nemo_relay.llm.optimization" + && let Some(data) = fields.data.as_mut().and_then(Json::as_object_mut) + { + data.insert("payload".to_string(), json!({"secret": "[redacted]"})); + } + fields + }), + ) + .unwrap(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "stream_optimization_sanitizer_observer", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_stream_execution_intercept( + "stream_optimization_sanitizer_contributor", + 1, + Arc::new(|_name, request, next| { + let mut contribution = LlmOptimizationContribution::new("test.stream", "stream_test"); + contribution.payload_schema = Some(DataSchema { + name: "test.stream_evidence".to_string(), + version: "1".to_string(), + }); + contribution.payload = Some(json!({"secret": "classified"})); + assert!(record_llm_optimization_contribution(contribution)); + Box::pin(async move { next(request).await }) + }), + ) + .unwrap(); + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("stream-optimization-sanitized") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { + Ok( + Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": "done"}))])) + as LlmJsonStream, + ) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"response": "done"}))) + .build(), + ) + .await + .unwrap(); + + // Poll under a different ambient stack. The wrapper must continue using + // the scope captured with the LLM handle, where the sanitizer is registered. + set_thread_scope_stack(create_scope_stack()); + while let Some(item) = stream.next().await { + item.unwrap(); + } + set_thread_scope_stack(original_stack); + + let captured = captured_events_snapshot(&events); + let mark = captured + .iter() + .find(|event| event.name() == "nemo_relay.llm.optimization") + .expect("expected a canonical optimization mark"); + assert_eq!(mark.data().unwrap()["payload"]["secret"], "[redacted]"); + + deregister_llm_stream_execution_intercept("stream_optimization_sanitizer_contributor").unwrap(); + deregister_subscriber("stream_optimization_sanitizer_observer").unwrap(); + pop_scope( + nemo_relay::api::scope::PopScopeParams::builder() + .handle_uuid(&owner.uuid) + .build(), + ) + .unwrap(); +} + +#[tokio::test] +async fn test_concurrent_managed_llm_calls_keep_optimization_evidence_isolated() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "concurrent_optimization_observer", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_request_intercept( + "concurrent_optimization_contributor", + 1, + false, + Arc::new(|_name, request, annotated| { + let call = request.content["call"].as_str().unwrap().to_string(); + let saved_tokens = if call == "a" { 11 } else { 22 }; + let mut contribution = + LlmOptimizationContribution::new(format!("test.{call}"), "concurrency_test"); + contribution.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(saved_tokens)), + ..LlmOptimizationTokenImpact::default() + }); + Ok(LlmRequestInterceptOutcome::new(request, annotated) + .with_optimization_contribution(contribution)) + }), + ) + .unwrap(); + + let call = |name: &'static str, label: &'static str| { + llm_call_execute( + LlmCallExecuteParams::builder() + .name(name) + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"call": label}), + }) + .func(Arc::new(move |_| { + Box::pin(async move { + tokio::task::yield_now().await; + Ok(json!({"call": label})) + }) + })) + .build(), + ) + }; + let (a, b) = tokio::join!(call("concurrent-a", "a"), call("concurrent-b", "b")); + a.unwrap(); + b.unwrap(); + + let captured = captured_events_snapshot(&events); + for (name, expected_producer, expected_tokens) in [ + ("concurrent-a", "test.a", 11), + ("concurrent-b", "test.b", 22), + ] { + let end = captured + .iter() + .find(|event| { + event.name() == name && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap_or_else(|| panic!("missing end event for {name}")); + let contributions = &end + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + .unwrap_or_else(|| panic!("missing optimization summary for {name}")) + .contributions; + assert_eq!(contributions.len(), 1); + assert_eq!(contributions[0].producer, expected_producer); + assert_eq!( + end.annotated_response() + .unwrap() + .optimization_summary + .as_ref() + .unwrap() + .tokens_saved + .prompt_tokens, + Some(expected_tokens) + ); + + let mark = captured + .iter() + .find(|event| { + event.name() == "nemo_relay.llm.optimization" + && event.parent_uuid() == Some(end.uuid()) + }) + .unwrap_or_else(|| panic!("missing optimization mark for {name}")); + assert_eq!(mark.data().unwrap()["producer"], expected_producer); + assert_eq!( + mark.data().unwrap()["token_impact"]["saved"]["prompt_tokens"], + expected_tokens + ); + } + + deregister_llm_request_intercept("concurrent_optimization_contributor").unwrap(); + deregister_subscriber("concurrent_optimization_observer").unwrap(); +} + #[tokio::test] async fn test_failed_request_intercept_does_not_emit_pending_marks_or_start_scope() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 6f8f09df9..083171054 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -19,6 +19,7 @@ use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, }; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; +use nemo_relay::api::optimization::record_llm_optimization_contribution; use nemo_relay::api::registry::{ deregister_llm_request_intercept, deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, register_llm_request_intercept, @@ -26,16 +27,20 @@ use nemo_relay::api::registry::{ }; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn}; +use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn}; use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::ScopeType; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::codec::openai_chat::OpenAIChatCodec; +use nemo_relay::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationKind, LlmOptimizationModel, + LlmOptimizationModelTransition, +}; use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; use nemo_relay::codec::response::FinishReason; use nemo_relay::codec::response::{ - AnnotatedLlmResponse, PricingCatalog, PricingResolver, Usage, reset_active_pricing_resolver, - set_active_pricing_resolver, + AnnotatedLlmResponse, CostEstimate, CostSource, PricingCatalog, PricingResolver, Usage, + reset_active_pricing_resolver, set_active_pricing_resolver, }; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, Result}; @@ -94,6 +99,41 @@ fn install_mock_response_pricing() { set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); } +fn install_routed_response_pricing() { + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + { + "provider": "test", "model_id": "baseline", + "pricing_as_of": "2026-07-08", "pricing_source": "test", + "rates": {"input_per_million": 4.0, "output_per_million": 8.0}, + "prompt_cache": {"read_accounting": "included_in_prompt_tokens"} + }, + { + "provider": "test", "model_id": "effective", + "pricing_as_of": "2026-07-08", "pricing_source": "test", + "rates": {"input_per_million": 1.0, "output_per_million": 2.0}, + "prompt_cache": {"read_accounting": "included_in_prompt_tokens"} + } + ] + }) + .to_string(), + ) + .unwrap(); + set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); +} + +fn routed_model_contribution() -> LlmOptimizationContribution { + let mut contribution = + LlmOptimizationContribution::new("test.router", LlmOptimizationKind::model_routing()); + contribution.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new("baseline").with_provider("test")), + effective: Some(LlmOptimizationModel::new("effective").with_provider("test")), + }); + contribution +} + // --------------------------------------------------------------------------- // TrackingCodec — records decode/encode calls and performs real transformations // --------------------------------------------------------------------------- @@ -1188,6 +1228,48 @@ impl LlmResponseCodec for MockResponseCodec { } } +struct RoutedAliasResponseCodec { + cost_source: CostSource, +} + +impl LlmResponseCodec for RoutedAliasResponseCodec { + fn decode_response(&self, _response: &Json) -> Result { + let provider_reported = self.cost_source == CostSource::ProviderReported; + Ok(AnnotatedLlmResponse { + model: Some("provider-response-alias".to_string()), + finish_reason: Some(FinishReason::Complete), + usage: Some(Usage { + prompt_tokens: Some(1_000), + completion_tokens: Some(500), + total_tokens: Some(1_500), + cost: Some(CostEstimate { + total: Some(if provider_reported { 0.42 } else { 99.0 }), + currency: "USD".to_string(), + input: None, + output: None, + cache_read: None, + cache_write: None, + source: self.cost_source.clone(), + pricing_provider: Some("test".to_string()), + pricing_model: Some("provider-response-alias".to_string()), + pricing_as_of: Some(if provider_reported { + "2026-07-01".to_string() + } else { + "2020-01-01".to_string() + }), + pricing_source: Some(if provider_reported { + "provider".to_string() + } else { + "stale-alias-pricing".to_string() + }), + }), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }) + } +} + /// Mock response codec that always fails. struct FailingResponseCodec; @@ -1575,6 +1657,149 @@ async fn test_stream_response_codec_populates_annotated_response() { reset_active_pricing_resolver().unwrap(); } +#[tokio::test] +async fn managed_buffered_and_streaming_close_price_the_committed_route_not_response_alias() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + install_routed_response_pricing(); + + let events = Arc::new(Mutex::new(Vec::new())); + let captured_events = events.clone(); + register_subscriber( + "routed_alias_pricing_sub", + Arc::new(move |event: &Event| captured_events.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + llm_call_execute( + LlmCallExecuteParams::builder() + .name("routed-buffered") + .request(make_llm_request(json!({"model": "requested"}))) + .func(Arc::new(|_| { + Box::pin(async { + assert!(record_llm_optimization_contribution( + routed_model_contribution() + )); + Ok(json!({"response": "done"})) + }) + })) + .response_codec(Arc::new(RoutedAliasResponseCodec { + cost_source: CostSource::ModelPricing, + })) + .build(), + ) + .await + .unwrap(); + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("routed-streaming") + .request(make_llm_request(json!({"model": "requested"}))) + .func(Arc::new(|_| { + Box::pin(async { + assert!(record_llm_optimization_contribution( + routed_model_contribution() + )); + Ok( + Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": "done"}))])) + as LlmJsonStream, + ) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"response": "done"}))) + .response_codec(Arc::new(RoutedAliasResponseCodec { + cost_source: CostSource::ModelPricing, + })) + .build(), + ) + .await + .unwrap(); + while let Some(item) = stream.next().await { + item.unwrap(); + } + + llm_call_execute( + LlmCallExecuteParams::builder() + .name("routed-provider-reported") + .request(make_llm_request(json!({"model": "requested"}))) + .func(Arc::new(|_| { + Box::pin(async { + assert!(record_llm_optimization_contribution( + routed_model_contribution() + )); + Ok(json!({"response": "done"})) + }) + })) + .response_codec(Arc::new(RoutedAliasResponseCodec { + cost_source: CostSource::ProviderReported, + })) + .build(), + ) + .await + .unwrap(); + + let captured = captured_events_snapshot(&events); + for name in ["routed-buffered", "routed-streaming"] { + let response = captured + .iter() + .find(|event| { + event.name() == name && is_scope_event(event, ScopeType::Llm, ScopeCategory::End) + }) + .and_then(Event::annotated_response) + .unwrap_or_else(|| panic!("missing annotated response for {name}")); + let summary = response.optimization_summary.as_ref().unwrap(); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.actual_cost.as_ref().unwrap().total, Some(0.002)); + assert_eq!( + summary + .actual_cost + .as_ref() + .unwrap() + .pricing_model + .as_deref(), + Some("effective") + ); + assert_eq!( + response + .usage + .as_ref() + .and_then(|usage| usage.cost.as_ref()) + .and_then(|cost| cost.total), + Some(0.002) + ); + } + + let provider_response = captured + .iter() + .find(|event| { + event.name() == "routed-provider-reported" + && is_scope_event(event, ScopeType::Llm, ScopeCategory::End) + }) + .and_then(Event::annotated_response) + .expect("missing provider-reported annotated response"); + let provider_cost = provider_response + .optimization_summary + .as_ref() + .and_then(|summary| summary.actual_cost.as_ref()) + .unwrap(); + assert_eq!(provider_cost.total, Some(0.42)); + assert_eq!(provider_cost.source, CostSource::ProviderReported); + assert_eq!(provider_cost.pricing_source.as_deref(), Some("provider")); + assert_eq!( + provider_response + .usage + .as_ref() + .and_then(|usage| usage.cost.as_ref()) + .map(|cost| cost.source.clone()), + Some(CostSource::ProviderReported) + ); + + deregister_subscriber("routed_alias_pricing_sub").unwrap(); + reset_active_pricing_resolver().unwrap(); +} + #[tokio::test] async fn test_stream_response_codec_annotation_uses_sanitized_aggregated_response() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/stream_tests.rs b/crates/core/tests/integration/stream_tests.rs index c4c9c35c7..3dbbee616 100644 --- a/crates/core/tests/integration/stream_tests.rs +++ b/crates/core/tests/integration/stream_tests.rs @@ -11,9 +11,11 @@ use std::sync::{Arc, Mutex}; use nemo_relay::api::event::{Event, ScopeCategory}; use nemo_relay::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; use nemo_relay::api::llm::{LlmCallParams, llm_call}; +use nemo_relay::api::optimization::LlmOptimizationRecorder; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::codec::optimization::LlmOptimizationContribution; use nemo_relay::error::FlowError; use nemo_relay::error::Result; use nemo_relay::json::Json; @@ -42,6 +44,13 @@ fn make_llm_handle(name: &str) -> LlmHandle { .build() } +fn make_optimized_llm_handle(name: &str, producer: &str) -> (LlmHandle, LlmOptimizationRecorder) { + let handle = make_llm_handle(name); + let recorder = handle.optimization_recorder.clone(); + assert!(recorder.record(LlmOptimizationContribution::new(producer, "stream_test"))); + (handle, recorder) +} + fn make_stream(items: Vec>) -> Pin> + Send>> { Box::pin(tokio_stream::iter(items)) } @@ -255,10 +264,164 @@ async fn test_stream_wrapper_drop_emits_end_event_for_partial_stream() { .find(|event| is_llm_end(event)) .expect("expected END event when a partial stream is dropped"); assert_eq!(end_event.output(), Some(&json!([{"token": "partial"}]))); + assert_eq!( + end_event.metadata().unwrap()["otel.status_code"], + json!("ERROR") + ); + assert!( + end_event.annotated_response().is_none(), + "an interrupted stream without optimization evidence must not manufacture a summary" + ); deregister_subscriber("stream_drop_end_test").unwrap(); } +#[tokio::test] +async fn stream_termination_modes_close_accounting_without_losing_evidence() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "stream_termination_accounting", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let (clean_handle, clean_recorder) = make_optimized_llm_handle("stream-clean", "test.clean"); + let (collector, finalizer, _) = make_collector_finalizer(); + let mut clean = LlmStreamWrapper::new( + make_stream(vec![Ok(json!({"chunk": "clean"}))]), + clean_handle, + collector, + finalizer, + None, + None, + None, + ); + while let Some(item) = clean.next().await { + item.unwrap(); + } + assert!(!clean_recorder.record(LlmOptimizationContribution::new("late", "test"))); + + let (before_error_handle, before_error_recorder) = + make_optimized_llm_handle("stream-error-before", "test.error_before"); + let (collector, finalizer, _) = make_collector_finalizer(); + let mut error_before = LlmStreamWrapper::new( + make_stream(vec![Err(FlowError::Internal("before first".into()))]), + before_error_handle, + collector, + finalizer, + None, + None, + None, + ); + assert!(error_before.next().await.unwrap().is_err()); + assert!(!before_error_recorder.record(LlmOptimizationContribution::new("late", "test"))); + + let (after_error_handle, after_error_recorder) = + make_optimized_llm_handle("stream-error-after", "test.error_after"); + let (collector, finalizer, _) = make_collector_finalizer(); + let mut error_after = LlmStreamWrapper::new( + make_stream(vec![ + Ok(json!({"chunk": "committed"})), + Err(FlowError::Internal("after first".into())), + ]), + after_error_handle, + collector, + finalizer, + None, + None, + None, + ); + assert!(error_after.next().await.unwrap().is_ok()); + assert!(error_after.next().await.unwrap().is_err()); + assert!(!after_error_recorder.record(LlmOptimizationContribution::new("late", "test"))); + + let (drop_before_handle, drop_before_recorder) = + make_optimized_llm_handle("stream-drop-before", "test.drop_before"); + let (collector, finalizer, _) = make_collector_finalizer(); + let drop_before = LlmStreamWrapper::new( + make_stream(vec![Ok(json!({"chunk": "unread"}))]), + drop_before_handle, + collector, + finalizer, + None, + None, + None, + ); + drop(drop_before); + assert!(!drop_before_recorder.record(LlmOptimizationContribution::new("late", "test"))); + + let (drop_after_handle, drop_after_recorder) = + make_optimized_llm_handle("stream-drop-after", "test.route_commit"); + let drop_after_uuid = drop_after_handle.uuid; + let (collector, finalizer, _) = make_collector_finalizer(); + let mut drop_after = LlmStreamWrapper::new( + make_stream(vec![ + Ok(json!({"chunk": "route committed"})), + Ok(json!({"chunk": "unread"})), + ]), + drop_after_handle, + collector, + finalizer, + None, + None, + None, + ); + assert!(drop_after.next().await.unwrap().is_ok()); + drop(drop_after); + assert!(!drop_after_recorder.record(LlmOptimizationContribution::new("late", "test"))); + + let events = captured_snapshot(&events); + let summary_for = |name: &str| { + events + .iter() + .find(|event| event.name() == name && is_llm_end(event)) + .and_then(Event::annotated_response) + .and_then(|response| response.optimization_summary.as_ref()) + .unwrap_or_else(|| panic!("missing optimization summary for {name}")) + }; + assert!( + !summary_for("stream-clean") + .limitations + .iter() + .any(|item| item == "stream_interrupted") + ); + for name in [ + "stream-error-before", + "stream-error-after", + "stream-drop-before", + "stream-drop-after", + ] { + assert!( + summary_for(name) + .limitations + .iter() + .any(|item| item == "stream_interrupted"), + "{name} should be marked interrupted" + ); + } + assert_eq!( + summary_for("stream-drop-after").contributions[0].producer, + "test.route_commit" + ); + let committed_mark = events + .iter() + .find(|event| { + event.name() == "nemo_relay.llm.optimization" + && event.parent_uuid() == Some(drop_after_uuid) + }) + .expect("route-commit contribution mark should survive an interrupted stream"); + assert_eq!( + committed_mark.data().unwrap()["producer"], + "test.route_commit" + ); + + deregister_subscriber("stream_termination_accounting").unwrap(); +} + #[tokio::test] async fn test_stream_wrapper_error_propagation() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index c14e0cee4..1e5c94350 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -11,15 +11,17 @@ use serde_json::json; use tokio_stream::StreamExt; use super::{ - LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, - llm_stream_call_execute, + LlmCallExecuteParams, LlmHandle, LlmRequest, LlmStreamCallExecuteParams, + emit_optimization_marks_with, llm_call_execute, llm_stream_call_execute, }; use crate::api::event::ScopeCategory; +use crate::api::optimization::finalize_optimization_summary; use crate::api::runtime::LlmJsonStream; use crate::api::runtime::{NemoRelayContextState, global_context}; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::error::FlowError; use crate::json::Json; +use crate::{codec::optimization::LlmOptimizationContribution, codec::response::PricingResolver}; fn reset_global() { crate::shared_runtime::reset_runtime_owner_for_tests(); @@ -40,6 +42,97 @@ fn request() -> LlmRequest { } } +#[test] +fn rejected_optimization_mark_queue_keeps_cursor_and_summary_evidence() { + let _guard = lock_global_runtime(); + reset_global(); + + let handle = LlmHandle::builder().name("queue-rejection-test").build(); + assert!( + handle + .optimization_recorder + .record(LlmOptimizationContribution::new("test", "queue_rejection")) + ); + emit_optimization_marks_with(&handle, &[], Some, |_event, _subscribers| false); + assert_eq!(handle.optimization_recorder.unemitted().len(), 1); + + let summary = finalize_optimization_summary( + &handle.optimization_recorder, + None, + None, + &PricingResolver::default(), + ) + .expect("queue rejection must not discard close-time evidence"); + assert_eq!(summary.contributions.len(), 1); + assert_eq!(summary.contributions[0].producer, "test"); +} + +#[test] +fn unavailable_mark_sanitizer_does_not_acknowledge_the_delivery_cursor() { + let _guard = lock_global_runtime(); + reset_global(); + + let handle = LlmHandle::builder().name("sanitizer-unavailable").build(); + assert!( + handle + .optimization_recorder + .record(LlmOptimizationContribution::new("test", "sanitize_retry")) + ); + handle.optimization_recorder.close_for_finalization(None); + emit_optimization_marks_with( + &handle, + &[], + |_event| None, + |_event, _subscribers| panic!("unavailable sanitization must not enqueue"), + ); + assert_eq!(handle.optimization_recorder.unemitted().len(), 1); + + emit_optimization_marks_with(&handle, &[], Some, |_event, _subscribers| true); + assert!(handle.optimization_recorder.unemitted().is_empty()); +} + +#[test] +fn close_boundary_freezes_identical_mark_and_summary_contributions() { + let _guard = lock_global_runtime(); + reset_global(); + + let handle = LlmHandle::builder().name("close-boundary").build(); + assert!( + handle + .optimization_recorder + .record(LlmOptimizationContribution::new( + "accepted", + "close_boundary" + )) + ); + assert!(handle.optimization_recorder.close_for_finalization(None)); + assert!( + !handle + .optimization_recorder + .record(LlmOptimizationContribution::new("late", "close_boundary")) + ); + + let mut marks = Vec::new(); + emit_optimization_marks_with(&handle, &[], Some, |event, _subscribers| { + marks.push(event.clone()); + true + }); + let summary = finalize_optimization_summary( + &handle.optimization_recorder, + None, + None, + &PricingResolver::default(), + ) + .unwrap(); + assert_eq!(marks.len(), 1); + assert_eq!(summary.contributions.len(), 1); + assert_eq!(marks[0].data().unwrap()["producer"], "accepted"); + assert_eq!( + marks[0].data().unwrap()["id"], + json!(summary.contributions[0].id.unwrap()) + ); +} + #[test] fn llm_call_execute_adds_otel_status_metadata_to_end_events() { let _guard = lock_global_runtime(); diff --git a/crates/core/tests/unit/optimization_tests.rs b/crates/core/tests/unit/optimization_tests.rs new file mode 100644 index 000000000..ff5291719 --- /dev/null +++ b/crates/core/tests/unit/optimization_tests.rs @@ -0,0 +1,1390 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unit tests for managed, bounded LLM optimization accounting. + +use serde_json::json; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use super::*; +use crate::api::event::DataSchema; +use crate::codec::optimization::{ + LlmOptimizationEvidenceQuality, LlmOptimizationModelTransition, LlmOptimizationTokenImpact, +}; +use crate::codec::response::{ + CostEstimate, CostSource, PricingCatalog, PricingCatalogError, PricingSource, Usage, +}; +use crate::json::Json; + +fn resolver() -> PricingResolver { + resolver_with_rates(2.0, 1.0) +} + +fn resolver_with_rates(baseline_input: f64, effective_input: f64) -> PricingResolver { + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + {"provider":"test","model_id":"baseline","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":baseline_input,"output_per_million":4.0,"cache_read_per_million":0.5,"cache_write_per_million":3.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}}, + {"provider":"test","model_id":"effective","pricing_as_of":"2026-07-08","pricing_source":"test-snapshot","rates":{"input_per_million":effective_input,"output_per_million":2.0,"cache_read_per_million":0.25,"cache_write_per_million":2.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}} + ] + }) + .to_string(), + ) + .unwrap(); + PricingResolver::from_catalogs(vec![catalog]) +} + +fn contribution() -> LlmOptimizationContribution { + let mut contribution = LlmOptimizationContribution::new( + "test.optimizer", + crate::codec::optimization::LlmOptimizationKind::model_routing(), + ); + contribution.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new("baseline").with_provider("test")), + effective: Some(LlmOptimizationModel::new("effective").with_provider("test")), + }); + contribution.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(200)), + quality: Some(LlmOptimizationEvidenceQuality::Estimated), + estimation_method: Some("test-tokenizer".to_string()), + ..LlmOptimizationTokenImpact::default() + }); + contribution +} + +#[test] +fn combined_summary_retains_token_evidence_and_snapshot_pricing() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + Some("baseline"), + &resolver(), + ) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Complete); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().prompt_tokens, + Some(1000) + ); + assert_eq!(summary.baseline_cost.as_ref().unwrap().total, Some(0.0024)); + assert_eq!(summary.actual_cost.as_ref().unwrap().total, Some(0.001)); + assert!((summary.estimated_cost_saved.unwrap() - 0.0014).abs() < 1e-12); +} + +#[test] +fn applied_route_is_the_authoritative_effective_model() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + // Providers may return an alias or deployment name rather than + // the exact model Relay selected and sent upstream. + model: Some("provider-response-alias".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + cost: Some(CostEstimate { + total: Some(99.0), + currency: "USD".to_string(), + input: None, + output: None, + cache_read: None, + cache_write: None, + source: CostSource::ModelPricing, + pricing_provider: Some("test".to_string()), + pricing_model: Some("provider-response-alias".to_string()), + pricing_as_of: Some("2020-01-01".to_string()), + pricing_source: Some("stale-alias-pricing".to_string()), + }), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + Some("original-request-model"), + &resolver(), + ) + .unwrap(); + assert_eq!(summary.baseline_model.as_ref().unwrap().model, "baseline"); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.actual_cost.as_ref().unwrap().total, Some(0.001)); + assert_eq!( + summary + .actual_cost + .as_ref() + .unwrap() + .pricing_model + .as_deref(), + Some("effective") + ); + assert_eq!( + response + .usage + .as_ref() + .and_then(|usage| usage.cost.as_ref()) + .and_then(|cost| cost.pricing_model.as_deref()), + Some("effective") + ); +} + +#[test] +fn routed_model_preserves_provider_reported_cost_and_provenance() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let provider_cost = CostEstimate { + total: Some(0.75), + currency: "EUR".to_string(), + input: Some(0.5), + output: Some(0.25), + cache_read: None, + cache_write: None, + source: CostSource::ProviderReported, + pricing_provider: Some("provider-billing".to_string()), + pricing_model: Some("provider-response-alias".to_string()), + pricing_as_of: Some("2026-07-08".to_string()), + pricing_source: Some("provider-invoice".to_string()), + }; + let mut response = AnnotatedLlmResponse { + model: Some("provider-response-alias".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + cost: Some(provider_cost.clone()), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.actual_cost.as_ref(), Some(&provider_cost)); + assert_eq!( + summary + .effective_usage + .as_ref() + .and_then(|usage| usage.cost.as_ref()), + Some(&provider_cost) + ); + assert_eq!( + response + .usage + .as_ref() + .and_then(|usage| usage.cost.as_ref()), + Some(&provider_cost) + ); + assert!( + summary + .limitations + .contains(&"cost_currency_mismatch".to_string()) + ); +} + +#[test] +fn unpriced_summary_is_partial_without_losing_tokens() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(8), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + None, + &PricingResolver::default(), + ) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); + assert!(summary.estimated_cost_saved.is_none()); +} + +#[test] +fn zero_and_negative_savings_are_preserved() { + for (baseline_rate, effective_rate, expected_sign) in [(0.0, 0.0, 0_i8), (0.5, 2.0, -1_i8)] { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(0), + total_tokens: Some(800), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + None, + &resolver_with_rates(baseline_rate, effective_rate), + ) + .unwrap(); + let saved = summary.estimated_cost_saved.unwrap(); + match expected_sign { + 0 => assert_eq!(saved, 0.0), + -1 => assert!(saved < 0.0), + _ => unreachable!(), + } + } +} + +#[test] +fn multiple_contributions_and_cache_savings_aggregate_explicitly() { + let recorder = LlmOptimizationRecorder::default(); + for (producer, prompt, cache_read, cache_write) in + [("test.one", 5, 7, 0), ("test.two", 11, 13, 17)] + { + let mut item = LlmOptimizationContribution::new( + producer, + crate::codec::optimization::LlmOptimizationKind::input_compression(), + ); + item.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens { + prompt_tokens: Some(prompt), + cache_read_tokens: Some(cache_read), + cache_write_tokens: Some(cache_write), + total_tokens: Some(prompt), + ..LlmOptimizationTokens::default() + }), + ..LlmOptimizationTokenImpact::default() + }); + assert!(recorder.record(item)); + } + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(100), + completion_tokens: Some(10), + total_tokens: Some(110), + cache_read_tokens: Some(20), + cache_write_tokens: Some(3), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(16)); + assert_eq!(summary.tokens_saved.cache_read_tokens, Some(20)); + assert_eq!(summary.tokens_saved.cache_write_tokens, Some(17)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().cache_read_tokens, + Some(40) + ); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().cache_write_tokens, + Some(20) + ); + assert_eq!(summary.contributions[0].sequence, Some(0)); + assert_eq!(summary.contributions[1].sequence, Some(1)); +} + +#[test] +fn serialized_summary_can_be_repriced_with_a_new_catalog() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + completion_tokens: Some(100), + total_tokens: Some(900), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let original = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + let restored: LlmOptimizationSummary = + serde_json::from_value(serde_json::to_value(&original).unwrap()).unwrap(); + let newer = resolver_with_rates(10.0, 5.0); + let baseline = newer + .estimate_cost_for_provider( + Some("test"), + "baseline", + restored.baseline_usage.as_ref().unwrap(), + ) + .unwrap() + .total_or_component_sum() + .unwrap(); + let actual = newer + .estimate_cost_for_provider( + Some("test"), + "effective", + restored.effective_usage.as_ref().unwrap(), + ) + .unwrap() + .total_or_component_sum() + .unwrap(); + assert_ne!(baseline - actual, original.estimated_cost_saved.unwrap()); + assert_eq!(restored.tokens_saved.prompt_tokens, Some(200)); +} + +#[test] +fn close_time_pricing_uses_the_loaded_resolver_without_reloading_sources() { + struct CountingSource { + loads: std::sync::Arc, + catalog: PricingCatalog, + } + + impl PricingSource for CountingSource { + fn source_name(&self) -> &str { + "counting-test-source" + } + + fn load_catalog(&self) -> Result, PricingCatalogError> { + self.loads.fetch_add(1, Ordering::SeqCst); + Ok(Some(self.catalog.clone())) + } + } + + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + {"provider":"test","model_id":"baseline","pricing_as_of":"2026-07-08","pricing_source":"test","rates":{"input_per_million":2.0,"output_per_million":2.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}}, + {"provider":"test","model_id":"effective","pricing_as_of":"2026-07-08","pricing_source":"test","rates":{"input_per_million":1.0,"output_per_million":1.0},"prompt_cache":{"read_accounting":"included_in_prompt_tokens"}} + ] + }) + .to_string(), + ) + .unwrap(); + let loads = std::sync::Arc::new(AtomicUsize::new(0)); + let pricing = PricingResolver::from_sources(vec![Box::new(CountingSource { + loads: loads.clone(), + catalog, + })]) + .unwrap(); + assert_eq!(loads.load(Ordering::SeqCst), 1); + + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(10), + completion_tokens: Some(0), + total_tokens: Some(10), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &pricing).unwrap(); + assert!(summary.actual_cost.is_some()); + assert_eq!(loads.load(Ordering::SeqCst), 1); +} + +#[test] +fn no_usage_is_an_explicit_partial_summary() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let summary = + finalize_optimization_summary(&recorder, None, Some("effective"), &resolver()).unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!( + summary + .limitations + .contains(&"missing_effective_usage".to_string()) + ); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(200)); +} + +#[test] +fn full_contribution_byte_limits_are_enforced_without_unbounded_work() { + let oversized = LlmOptimizationRecorder::default(); + let mut item = LlmOptimizationContribution::new("test", "custom"); + item.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + item.payload = Some(Json::String( + "x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES), + )); + assert!(!oversized.record(item)); + + let aggregate = LlmOptimizationRecorder::default(); + for index in 0..17 { + let mut item = LlmOptimizationContribution::new(format!("test.{index}"), "custom"); + item.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + item.payload = Some(Json::String("x".repeat(15_000))); + assert!(aggregate.record(item)); + } + let mut overflow = LlmOptimizationContribution::new("test.overflow", "custom"); + overflow.payload_schema = Some(DataSchema { + name: "test.payload".to_string(), + version: "1".to_string(), + }); + overflow.payload = Some(Json::String("x".repeat(15_000))); + assert!(!aggregate.record(overflow)); + assert!( + aggregate + .finish() + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); +} + +#[test] +fn bounds_and_invalid_payloads_are_best_effort_and_visible() { + let recorder = LlmOptimizationRecorder::default(); + let mut invalid = LlmOptimizationContribution::new("test", "custom"); + invalid.payload = Some(json!({"evidence": true})); + assert!(!recorder.record(invalid)); + for index in 0..(MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS - 1) { + assert!(recorder.record(LlmOptimizationContribution::new( + format!("test.{index}"), + "custom" + ))); + } + assert!(!recorder.record(LlmOptimizationContribution::new("overflow", "custom"))); + let summary = + finalize_optimization_summary(&recorder, None, None, &PricingResolver::default()).unwrap(); + assert_eq!( + summary.contributions.len(), + MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS - 1 + ); + assert!( + summary + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + assert!( + summary + .limitations + .contains(&"invalid_contribution_payload_schema".to_string()) + ); +} + +#[tokio::test] +async fn recorder_can_be_captured_for_stream_commit() { + let recorder = LlmOptimizationRecorder::default(); + let captured = scope_llm_optimization_recorder(recorder.clone(), async { + current_llm_optimization_recorder().unwrap() + }) + .await; + assert!(captured.record(LlmOptimizationContribution::new("test.stream", "commit"))); + assert_eq!(recorder.finish().contributions.len(), 1); +} + +#[test] +fn full_envelope_fields_are_bounded_and_rejections_do_not_consume_sequences() { + let oversized_producer_recorder = LlmOptimizationRecorder::default(); + let mut oversized_producer = LlmOptimizationContribution::new( + "x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES), + "custom", + ); + oversized_producer.id = Some(uuid::Uuid::nil()); + oversized_producer.sequence = Some(99); + assert!(!oversized_producer_recorder.record(oversized_producer)); + assert!( + !oversized_producer_recorder.record(LlmOptimizationContribution::new("sealed", "custom")) + ); + assert!( + oversized_producer_recorder + .finish() + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + + let oversized_extra_recorder = LlmOptimizationRecorder::default(); + let mut oversized_extra = LlmOptimizationContribution::new("test", "custom"); + oversized_extra.extra.insert( + "future_evidence".to_string(), + Json::String("x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES)), + ); + assert!(!oversized_extra_recorder.record(oversized_extra)); + assert!(!oversized_extra_recorder.record(LlmOptimizationContribution::new("sealed", "custom"))); + + let oversized_kind = LlmOptimizationContribution::new( + "test", + "x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES), + ); + let mut oversized_model = LlmOptimizationContribution::new("test", "custom"); + oversized_model.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new( + "x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES), + )), + effective: None, + }); + let mut oversized_method = LlmOptimizationContribution::new("test", "custom"); + oversized_method.token_impact = Some(LlmOptimizationTokenImpact { + estimation_method: Some("x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES)), + ..LlmOptimizationTokenImpact::default() + }); + for oversized in [oversized_kind, oversized_model, oversized_method] { + let recorder = LlmOptimizationRecorder::default(); + assert!(!recorder.record(oversized)); + assert!(!recorder.record(LlmOptimizationContribution::new("sealed", "custom"))); + assert!( + recorder + .finish() + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + } + + let recorder = LlmOptimizationRecorder::default(); + let mut malformed = LlmOptimizationContribution::new("malformed", "custom"); + malformed.payload = Some(json!({"missing": "schema"})); + assert!(!recorder.record(malformed)); + let mut accepted = LlmOptimizationContribution::new("test", "custom"); + accepted.id = Some(uuid::Uuid::nil()); + accepted.sequence = Some(99); + assert!(recorder.record(accepted)); + + let finished = recorder.finish(); + assert_eq!(finished.contributions.len(), 1); + assert_eq!(finished.contributions[0].sequence, Some(0)); + assert_ne!(finished.contributions[0].id, Some(uuid::Uuid::nil())); + assert!( + finished + .limitations + .contains(&"invalid_contribution_payload_schema".to_string()) + ); +} + +#[test] +fn record_all_bounds_malformed_attempts_and_keeps_accepted_sequences_dense() { + struct CountingMalformed { + next_calls: std::sync::Arc, + } + + impl Iterator for CountingMalformed { + type Item = LlmOptimizationContribution; + + fn next(&mut self) -> Option { + self.next_calls.fetch_add(1, Ordering::SeqCst); + let mut contribution = LlmOptimizationContribution::new("malformed", "custom"); + contribution.payload = Some(json!({"schema": "missing"})); + Some(contribution) + } + } + + let next_calls = std::sync::Arc::new(AtomicUsize::new(0)); + let recorder = LlmOptimizationRecorder::default(); + recorder.record_all(CountingMalformed { + next_calls: next_calls.clone(), + }); + assert_eq!( + next_calls.load(Ordering::SeqCst), + MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS + 1 + ); + assert!(!recorder.record(LlmOptimizationContribution::new("late", "custom"))); + let finished = recorder.finish(); + assert!(finished.contributions.is_empty()); + assert!( + finished + .limitations + .contains(&"contribution_limit_exceeded".to_string()) + ); + assert!( + finished + .limitations + .contains(&"invalid_contribution_payload_schema".to_string()) + ); + + let dense = LlmOptimizationRecorder::default(); + let mut invalid = LlmOptimizationContribution::new("invalid", "custom"); + invalid.payload = Some(json!({"schema": "missing"})); + dense.record_all([ + LlmOptimizationContribution::new("first", "custom"), + invalid, + LlmOptimizationContribution::new("second", "custom"), + ]); + let finished = dense.finish(); + assert_eq!(finished.contributions.len(), 2); + assert_eq!(finished.contributions[0].sequence, Some(0)); + assert_eq!(finished.contributions[1].sequence, Some(1)); +} + +#[test] +fn delivery_cursor_advances_only_when_explicitly_acknowledged() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(LlmOptimizationContribution::new("one", "custom"))); + assert!(recorder.record(LlmOptimizationContribution::new("two", "custom"))); + + assert_eq!(recorder.unemitted().len(), 2); + assert_eq!(recorder.unemitted().len(), 2); + recorder.mark_emitted(1); + let remaining = recorder.unemitted(); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].producer, "two"); + recorder.mark_emitted(usize::MAX); + assert!(recorder.unemitted().is_empty()); + + let finished = recorder.finish(); + assert_eq!(finished.contributions.len(), 2); + recorder.mark_emitted(1); + assert!(recorder.unemitted().is_empty()); +} + +#[test] +fn finish_closes_the_recorder_and_late_writes_are_rejected() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(LlmOptimizationContribution::new("before", "custom"))); + recorder.note_limitation("before_close"); + + let finished = recorder.finish(); + assert_eq!(finished.contributions.len(), 1); + assert_eq!(finished.limitations, vec!["before_close"]); + assert!(!recorder.record(LlmOptimizationContribution::new("late", "custom"))); + recorder.note_limitation("after_close"); + + let second_finish = recorder.finish(); + assert!(second_finish.contributions.is_empty()); + assert!(second_finish.limitations.is_empty()); +} + +#[test] +fn concurrent_recording_assigns_unique_dense_order() { + let recorder = LlmOptimizationRecorder::default(); + let workers = (0..MAX_LLM_OPTIMIZATION_CONTRIBUTIONS) + .map(|index| { + let recorder = recorder.clone(); + std::thread::spawn(move || { + recorder.record(LlmOptimizationContribution::new( + format!("worker.{index}"), + "custom", + )) + }) + }) + .collect::>(); + for worker in workers { + assert!(worker.join().unwrap()); + } + + let finished = recorder.finish(); + let sequences = finished + .contributions + .iter() + .map(|contribution| contribution.sequence.unwrap()) + .collect::>(); + assert_eq!( + sequences, + (0..MAX_LLM_OPTIMIZATION_CONTRIBUTIONS as u64).collect::>() + ); + let ids = finished + .contributions + .iter() + .map(|contribution| contribution.id.unwrap()) + .collect::>(); + assert_eq!(ids.len(), MAX_LLM_OPTIMIZATION_CONTRIBUTIONS); +} + +#[test] +fn concurrent_finish_is_an_atomic_acceptance_boundary() { + let recorder = LlmOptimizationRecorder::default(); + let barrier = std::sync::Arc::new(std::sync::Barrier::new( + MAX_LLM_OPTIMIZATION_CONTRIBUTIONS + 1, + )); + let workers = (0..MAX_LLM_OPTIMIZATION_CONTRIBUTIONS) + .map(|index| { + let recorder = recorder.clone(); + let barrier = std::sync::Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + recorder.record(LlmOptimizationContribution::new( + format!("worker.{index}"), + "custom", + )) + }) + }) + .collect::>(); + + barrier.wait(); + let finished = recorder.finish(); + let accepted = workers + .into_iter() + .map(|worker| worker.join().unwrap()) + .filter(|accepted| *accepted) + .count(); + assert_eq!(accepted, finished.contributions.len()); + assert!(!recorder.record(LlmOptimizationContribution::new("late", "custom"))); +} + +#[tokio::test] +async fn recorder_task_local_is_not_implicitly_inherited_by_spawned_tasks() { + let recorder = LlmOptimizationRecorder::default(); + scope_llm_optimization_recorder(recorder, async { + assert!(current_llm_optimization_recorder().is_some()); + assert!( + tokio::spawn(async { current_llm_optimization_recorder().is_none() }) + .await + .unwrap() + ); + }) + .await; +} + +fn compression_contribution( + producer: &str, + saved: LlmOptimizationTokens, +) -> LlmOptimizationContribution { + let mut contribution = LlmOptimizationContribution::new( + producer, + crate::codec::optimization::LlmOptimizationKind::input_compression(), + ); + contribution.token_impact = Some(LlmOptimizationTokenImpact { + saved: Some(saved), + ..LlmOptimizationTokenImpact::default() + }); + contribution +} + +#[test] +fn missing_observed_fields_are_not_fabricated_for_baseline_usage() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "compressor", + LlmOptimizationTokens::saved_prompt(5), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: None, + completion_tokens: Some(2), + total_tokens: Some(2), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + + assert_eq!(summary.baseline_usage.as_ref().unwrap().prompt_tokens, None); + assert!(summary.baseline_cost.is_none()); + assert!(summary.actual_cost.is_none()); + assert!( + summary + .limitations + .contains(&"missing_effective_prompt_tokens".to_string()) + ); +} + +#[test] +fn empty_and_partial_usage_never_produce_complete_zero_cost_accounting() { + for usage in [ + Usage::default(), + Usage { + prompt_tokens: Some(10), + ..Usage::default() + }, + ] { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "compressor", + LlmOptimizationTokens::saved_prompt(2), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(usage), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!(summary.actual_cost.is_none()); + assert!(summary.baseline_cost.is_none()); + assert!(summary.estimated_cost_saved.is_none()); + assert!( + summary + .limitations + .contains(&"missing_effective_completion_tokens".to_string()) + ); + assert!( + summary + .limitations + .contains(&"missing_effective_total_tokens".to_string()) + ); + } +} + +#[test] +fn effective_total_is_inferred_from_complete_core_usage() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "compressor", + LlmOptimizationTokens::saved_prompt(3), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(10), + completion_tokens: Some(2), + total_tokens: None, + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!( + summary.effective_usage.as_ref().unwrap().total_tokens, + Some(12) + ); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().total_tokens, + Some(15) + ); + assert!( + !summary + .limitations + .contains(&"missing_effective_total_tokens".to_string()) + ); +} + +#[test] +fn mixed_and_inconsistent_saved_totals_have_explicit_semantics() { + let mixed = LlmOptimizationRecorder::default(); + assert!(mixed.record(compression_contribution( + "explicit", + LlmOptimizationTokens { + prompt_tokens: Some(5), + total_tokens: Some(5), + ..LlmOptimizationTokens::default() + }, + ))); + assert!(mixed.record(compression_contribution( + "derived", + LlmOptimizationTokens { + prompt_tokens: Some(7), + total_tokens: None, + ..LlmOptimizationTokens::default() + }, + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(20), + completion_tokens: Some(0), + total_tokens: Some(20), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&mixed, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.tokens_saved.total_tokens, Some(12)); + assert!( + !summary + .limitations + .contains(&"missing_token_savings_total".to_string()) + ); + + let inconsistent = LlmOptimizationRecorder::default(); + assert!(inconsistent.record(compression_contribution( + "inconsistent", + LlmOptimizationTokens { + prompt_tokens: Some(3), + completion_tokens: Some(2), + total_tokens: Some(9), + ..LlmOptimizationTokens::default() + }, + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(20), + completion_tokens: Some(5), + total_tokens: Some(25), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&inconsistent, Some(&mut response), None, &resolver()) + .unwrap(); + assert_eq!(summary.tokens_saved.total_tokens, Some(9)); + assert!(summary.baseline_cost.is_none()); + assert!( + summary + .limitations + .contains(&"inconsistent_token_savings_total".to_string()) + ); + + let cache_only = LlmOptimizationRecorder::default(); + assert!(cache_only.record(compression_contribution( + "cache-only", + LlmOptimizationTokens { + cache_read_tokens: Some(4), + ..LlmOptimizationTokens::default() + }, + ))); + let summary = finalize_optimization_summary( + &cache_only, + None, + Some("effective"), + &PricingResolver::default(), + ) + .unwrap(); + assert!(summary.tokens_saved.total_tokens.is_none()); + assert!( + summary + .limitations + .contains(&"missing_token_savings_total".to_string()) + ); +} + +#[test] +fn checked_token_aggregation_reports_overflow_without_clamping_evidence() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "large", + LlmOptimizationTokens { + prompt_tokens: Some(u64::MAX), + total_tokens: Some(u64::MAX), + ..LlmOptimizationTokens::default() + }, + ))); + assert!(recorder.record(compression_contribution( + "overflow", + LlmOptimizationTokens::saved_prompt(1), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(1), + total_tokens: Some(1), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!( + summary + .limitations + .contains(&"token_count_overflow".to_string()) + ); + assert_eq!(summary.tokens_saved.prompt_tokens, None); + assert_eq!(summary.tokens_saved.total_tokens, None); + assert_eq!(summary.baseline_usage.as_ref().unwrap().prompt_tokens, None); + assert!(summary.baseline_cost.is_none()); + assert!(summary.estimated_cost_saved.is_none()); + assert_eq!(summary.contributions.len(), 2); +} + +#[test] +fn checked_token_aggregation_accepts_the_exact_u64_boundary() { + let recorder = LlmOptimizationRecorder::default(); + for (producer, saved) in [("almost-max", u64::MAX - 1), ("last-token", 1)] { + assert!(recorder.record(compression_contribution( + producer, + LlmOptimizationTokens { + prompt_tokens: Some(saved), + ..LlmOptimizationTokens::default() + }, + ))); + } + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(0), + total_tokens: Some(0), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(u64::MAX)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().prompt_tokens, + Some(u64::MAX) + ); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().total_tokens, + Some(u64::MAX) + ); + assert!( + !summary + .limitations + .contains(&"token_count_overflow".to_string()) + ); +} + +#[test] +fn checked_baseline_derivation_reports_effective_plus_saved_overflow() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "one-token", + LlmOptimizationTokens::saved_prompt(1), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(u64::MAX), + completion_tokens: Some(0), + total_tokens: Some(u64::MAX), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert!( + summary + .limitations + .contains(&"token_count_overflow".to_string()) + ); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(1)); + assert_eq!(summary.baseline_usage.as_ref().unwrap().prompt_tokens, None); + assert!(summary.baseline_cost.is_none()); + assert!(summary.actual_cost.is_some()); + assert!(summary.estimated_cost_saved.is_none()); +} + +#[test] +fn multiple_routing_authorities_are_ignored_but_compression_evidence_survives() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut second_route = contribution(); + second_route.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new("other-baseline")), + effective: Some(LlmOptimizationModel::new("other-effective")), + }); + assert!(recorder.record(second_route)); + assert!(recorder.record(compression_contribution( + "test.compressor", + LlmOptimizationTokens::saved_prompt(7), + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(100), + completion_tokens: Some(0), + total_tokens: Some(100), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert!( + summary + .limitations + .contains(&"multiple_routing_contributions".to_string()) + ); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(7)); + assert_eq!(summary.baseline_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.contributions.len(), 3); +} + +#[test] +fn nonapplied_routing_evidence_does_not_create_multiple_authorities() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + for index in 0..5 { + let mut hypothetical = contribution(); + hypothetical.producer = format!("hypothetical.{index}"); + hypothetical.applied = false; + assert!(recorder.record(hypothetical)); + } + let mut response = AnnotatedLlmResponse { + model: Some("provider-alias".to_string()), + usage: Some(Usage { + prompt_tokens: Some(10), + completion_tokens: Some(0), + total_tokens: Some(10), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert!( + !summary + .limitations + .contains(&"multiple_routing_contributions".to_string()) + ); + assert_eq!(summary.baseline_model.as_ref().unwrap().model, "baseline"); + assert_eq!(summary.effective_model.as_ref().unwrap().model, "effective"); + assert_eq!(summary.contributions.len(), 6); +} + +#[test] +fn incomplete_cost_object_makes_the_summary_partial() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(contribution())); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(800), + total_tokens: Some(800), + cost: Some(CostEstimate { + total: None, + currency: "USD".to_string(), + input: None, + output: None, + cache_read: None, + cache_write: None, + source: CostSource::ProviderReported, + pricing_provider: Some("test".to_string()), + pricing_model: Some("effective".to_string()), + pricing_as_of: None, + pricing_source: Some("provider".to_string()), + }), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!( + summary + .limitations + .contains(&"missing_actual_cost_total".to_string()) + ); + assert!(summary.estimated_cost_saved.is_none()); + assert!(summary.currency.is_none()); +} + +#[test] +fn exact_contribution_envelope_boundary_is_accepted() { + let mut template = LlmOptimizationContribution::new("", "custom"); + template.id = Some(uuid::Uuid::nil()); + template.sequence = Some(0); + let fixed_size = bounded_json_size(&template, usize::MAX).unwrap(); + let boundary_producer_len = MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES - fixed_size; + + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(LlmOptimizationContribution::new( + "x".repeat(boundary_producer_len), + "custom", + ))); + let accepted = recorder.finish(); + assert_eq!(accepted.contributions.len(), 1); + assert_eq!( + bounded_json_size( + &accepted.contributions[0], + MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES, + ) + .unwrap(), + MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES + ); + + let overflow = LlmOptimizationRecorder::default(); + assert!(!overflow.record(LlmOptimizationContribution::new( + "x".repeat(boundary_producer_len + 1), + "custom", + ))); +} + +#[test] +fn prompt_completion_fallback_total_uses_checked_arithmetic() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "fallback-overflow", + LlmOptimizationTokens { + prompt_tokens: Some(u64::MAX), + completion_tokens: Some(1), + total_tokens: None, + ..LlmOptimizationTokens::default() + }, + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(1), + completion_tokens: Some(1), + total_tokens: Some(2), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, Some(u64::MAX)); + assert_eq!(summary.tokens_saved.completion_tokens, Some(1)); + assert_eq!(summary.tokens_saved.total_tokens, None); + assert_eq!(summary.baseline_usage.as_ref().unwrap().total_tokens, None); + assert!( + summary + .limitations + .contains(&"token_count_overflow".to_string()) + ); +} + +#[test] +fn one_overflowed_token_category_does_not_erase_unaffected_categories() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "first", + LlmOptimizationTokens { + prompt_tokens: Some(u64::MAX), + completion_tokens: Some(3), + ..LlmOptimizationTokens::default() + }, + ))); + assert!(recorder.record(compression_contribution( + "second", + LlmOptimizationTokens { + prompt_tokens: Some(1), + completion_tokens: Some(4), + ..LlmOptimizationTokens::default() + }, + ))); + let mut response = AnnotatedLlmResponse { + model: Some("effective".to_string()), + usage: Some(Usage { + prompt_tokens: Some(10), + completion_tokens: Some(5), + total_tokens: Some(15), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = + finalize_optimization_summary(&recorder, Some(&mut response), None, &resolver()).unwrap(); + assert_eq!(summary.tokens_saved.prompt_tokens, None); + assert_eq!(summary.tokens_saved.completion_tokens, Some(7)); + assert_eq!( + summary.baseline_usage.as_ref().unwrap().completion_tokens, + Some(12) + ); + assert!( + summary + .limitations + .contains(&"token_count_overflow".to_string()) + ); +} + +fn synthetic_cost(total: Option, input: Option, currency: &str) -> CostEstimate { + CostEstimate { + total, + currency: currency.to_string(), + input, + output: None, + cache_read: None, + cache_write: None, + source: CostSource::ProviderReported, + pricing_provider: None, + pricing_model: None, + pricing_as_of: None, + pricing_source: None, + } +} + +#[test] +fn cost_savings_requires_complete_totals_and_matching_currency() { + let empty_usd = synthetic_cost(None, None, "USD"); + let complete_usd = synthetic_cost(Some(2.0), None, "USD"); + + let mut baseline_empty = Vec::new(); + assert_eq!( + calculate_estimated_cost_saved(Some(&empty_usd), Some(&complete_usd), &mut baseline_empty,), + (None, None) + ); + assert_eq!(baseline_empty, vec!["missing_baseline_cost_total"]); + + let mut actual_empty = Vec::new(); + assert_eq!( + calculate_estimated_cost_saved(Some(&complete_usd), Some(&empty_usd), &mut actual_empty,), + (None, None) + ); + assert_eq!(actual_empty, vec!["missing_actual_cost_total"]); + + let mut both_empty = Vec::new(); + assert_eq!( + calculate_estimated_cost_saved(Some(&empty_usd), Some(&empty_usd), &mut both_empty), + (None, None) + ); + assert_eq!( + both_empty, + vec!["missing_baseline_cost_total", "missing_actual_cost_total"] + ); + + let baseline_components = synthetic_cost(None, Some(3.5), "USD"); + let actual_components = synthetic_cost(None, Some(1.25), "USD"); + let mut component_only = Vec::new(); + assert_eq!( + calculate_estimated_cost_saved( + Some(&baseline_components), + Some(&actual_components), + &mut component_only, + ), + (Some(2.25), Some("USD".to_string())) + ); + assert!(component_only.is_empty()); + + let eur = synthetic_cost(Some(1.0), None, "EUR"); + let mut mismatch = Vec::new(); + assert_eq!( + calculate_estimated_cost_saved(Some(&complete_usd), Some(&eur), &mut mismatch), + (None, None) + ); + assert_eq!(mismatch, vec!["cost_currency_mismatch"]); +} + +#[test] +fn missing_effective_model_is_an_explicit_limitation() { + let recorder = LlmOptimizationRecorder::default(); + assert!(recorder.record(compression_contribution( + "compressor", + LlmOptimizationTokens::saved_prompt(3), + ))); + let mut response = AnnotatedLlmResponse { + model: None, + usage: Some(Usage { + prompt_tokens: Some(7), + total_tokens: Some(7), + cost: Some(synthetic_cost(Some(0.1), None, "USD")), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + + let summary = finalize_optimization_summary( + &recorder, + Some(&mut response), + None, + &PricingResolver::default(), + ) + .unwrap(); + assert!(summary.effective_model.is_none()); + assert!(summary.baseline_model.is_none()); + assert!( + summary + .limitations + .contains(&"missing_effective_model".to_string()) + ); +} diff --git a/crates/core/tests/unit/shared_tests.rs b/crates/core/tests/unit/shared_tests.rs index ced67814b..d10e4f2e8 100644 --- a/crates/core/tests/unit/shared_tests.rs +++ b/crates/core/tests/unit/shared_tests.rs @@ -9,12 +9,14 @@ use std::sync::Arc; use serde_json::{Map, json}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; +use crate::api::optimization::{LlmOptimizationRecorder, MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES}; use crate::api::registry::{deregister_llm_request_intercept, register_llm_request_intercept}; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{create_scope_stack, set_thread_scope_stack}; use crate::api::scope::ScopeType; use crate::api::scope::{pop_scope, push_scope}; +use crate::codec::optimization::LlmOptimizationContribution; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; use crate::codec::traits::LlmCodec; use crate::error::Result; @@ -79,6 +81,8 @@ fn reset_global() { set_thread_scope_stack(create_scope_stack()); let _ = deregister_llm_request_intercept("shared-none"); let _ = deregister_llm_request_intercept("shared-codec"); + let _ = deregister_llm_request_intercept("shared-contribution-oversized"); + let _ = deregister_llm_request_intercept("shared-contribution-accepted"); } #[test] @@ -234,6 +238,73 @@ fn test_run_request_intercepts_with_codec_none_and_codec_paths() { reset_global(); } +#[test] +fn managed_request_chain_records_contributions_incrementally_while_standalone_retains_them() { + let _guard = lock_runtime_owner(); + reset_global(); + + register_llm_request_intercept( + "shared-contribution-accepted", + 1, + false, + Arc::new(|_name, request, annotated| { + Ok( + LlmRequestInterceptOutcome::new(request, annotated).with_optimization_contribution( + LlmOptimizationContribution::new("accepted", "custom"), + ), + ) + }), + ) + .unwrap(); + register_llm_request_intercept( + "shared-contribution-oversized", + 2, + false, + Arc::new(|_name, request, annotated| { + Ok( + LlmRequestInterceptOutcome::new(request, annotated).with_optimization_contribution( + LlmOptimizationContribution::new( + "x".repeat(MAX_LLM_OPTIMIZATION_CONTRIBUTION_BYTES), + "custom", + ), + ), + ) + }), + ) + .unwrap(); + + let standalone = run_request_intercepts_with_codec( + "shared", + LlmRequest { + headers: Map::new(), + content: json!({"prompt": "hello"}), + }, + None, + ) + .unwrap(); + assert_eq!(standalone.3.len(), 2); + assert!(standalone.3.iter().all(|item| item.sequence.is_none())); + + let recorder = LlmOptimizationRecorder::default(); + let managed = run_request_intercepts_with_codec_and_recorder( + "shared", + LlmRequest { + headers: Map::new(), + content: json!({"prompt": "hello"}), + }, + None, + &recorder, + ) + .unwrap(); + assert!(managed.3.is_empty()); + let recorded = recorder.unemitted(); + assert_eq!(recorded.len(), 1); + assert_eq!(recorded[0].producer, "accepted"); + assert_eq!(recorded[0].sequence, Some(0)); + + reset_global(); +} + #[test] fn test_run_request_intercepts_injects_dynamo_agent_lineage() { let _guard = lock_runtime_owner(); diff --git a/crates/core/tests/unit/stream_tests.rs b/crates/core/tests/unit/stream_tests.rs index 52615a1c7..0eb68537d 100644 --- a/crates/core/tests/unit/stream_tests.rs +++ b/crates/core/tests/unit/stream_tests.rs @@ -4,6 +4,7 @@ //! Unit tests for stream in the NeMo Relay core crate. use super::*; +use crate::codec::response::{FinishReason, Usage}; use serde_json::json; fn assert_no_top_level_fields(data: &Json, fields: &[&str]) { @@ -16,6 +17,26 @@ fn assert_no_top_level_fields(data: &Json, fields: &[&str]) { } } +#[test] +fn partial_stream_usage_is_not_treated_as_authoritative_without_terminal_evidence() { + let partial = AnnotatedLlmResponse { + usage: Some(Usage { + prompt_tokens: Some(10), + completion_tokens: Some(2), + total_tokens: Some(12), + ..Usage::default() + }), + ..AnnotatedLlmResponse::default() + }; + assert!(!has_authoritative_final_usage(Some(&partial))); + + let terminal = AnnotatedLlmResponse { + finish_reason: Some(FinishReason::Complete), + ..partial + }; + assert!(has_authoritative_final_usage(Some(&terminal))); +} + #[test] fn openai_chat_chunk_summary_keeps_only_compact_metadata() { let data = llm_chunk_mark_data( diff --git a/crates/types/src/codec/optimization.rs b/crates/types/src/codec/optimization.rs index cf5cc9a6b..e16c883b6 100644 --- a/crates/types/src/codec/optimization.rs +++ b/crates/types/src/codec/optimization.rs @@ -14,6 +14,11 @@ use crate::api::event::DataSchema; use super::response::{CostEstimate, Usage}; /// Open, forward-compatible optimization classification. +/// +/// This is intentionally a string-backed newtype rather than a closed enum: +/// third-party optimization kinds must deserialize and round-trip losslessly +/// before Relay knows about them. Standard constants and constructors provide +/// enum-like ergonomics without making new producers wait for a core release. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] #[serde(transparent)] pub struct LlmOptimizationKind(String); @@ -130,20 +135,6 @@ impl LlmOptimizationTokens { ..Self::default() } } - - /// Saturating field-wise addition, preserving absent fields when neither side supplies them. - pub fn add_assign(&mut self, other: &Self) { - fn add(target: &mut Option, value: Option) { - if let Some(value) = value { - *target = Some(target.unwrap_or(0).saturating_add(value)); - } - } - add(&mut self.prompt_tokens, other.prompt_tokens); - add(&mut self.completion_tokens, other.completion_tokens); - add(&mut self.cache_read_tokens, other.cache_read_tokens); - add(&mut self.cache_write_tokens, other.cache_write_tokens); - add(&mut self.total_tokens, other.total_tokens); - } } /// Whether token evidence was observed directly or estimated. @@ -300,54 +291,3 @@ pub struct LlmOptimizationSummary { /// Ordered, bounded source evidence used by the calculation. pub contributions: Vec, } - -#[cfg(test)] -mod tests { - use serde::Serialize; - use serde_json::json; - - use super::*; - - #[derive(Serialize)] - struct CustomPayload { - evidence: String, - } - - impl LlmOptimizationPayload for CustomPayload { - const SCHEMA_NAME: &'static str = "example.custom_optimization"; - const SCHEMA_VERSION: &'static str = "3"; - } - - #[test] - fn custom_kinds_payloads_and_future_fields_round_trip() { - let mut contribution = LlmOptimizationContribution::new("example", "energy_reduction") - .with_payload(&CustomPayload { - evidence: "measured".to_string(), - }) - .unwrap(); - contribution - .extra - .insert("future_field".to_string(), json!({"v": 2})); - let decoded: LlmOptimizationContribution = - serde_json::from_value(serde_json::to_value(&contribution).unwrap()).unwrap(); - assert_eq!(decoded.kind.as_str(), "energy_reduction"); - assert_eq!( - decoded.payload_schema.as_ref().unwrap().name, - "example.custom_optimization" - ); - assert_eq!(decoded.extra["future_field"], json!({"v": 2})); - } - - #[test] - fn saved_prompt_tokens_remain_explicit_on_the_wire() { - let impact = LlmOptimizationTokenImpact { - saved: Some(LlmOptimizationTokens::saved_prompt(42)), - quality: Some(LlmOptimizationEvidenceQuality::Estimated), - estimation_method: Some("tokenizer-v1".to_string()), - ..LlmOptimizationTokenImpact::default() - }; - let wire = serde_json::to_value(impact).unwrap(); - assert_eq!(wire["saved"]["prompt_tokens"], 42); - assert_eq!(wire["saved"]["total_tokens"], 42); - } -} From e095fa648dde81941b1a93724c2c15d48ad23d09 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 00:37:22 -0600 Subject: [PATCH 06/10] feat: complete optimization SDK contract surfaces Signed-off-by: Bryan Bednarski --- crates/ffi/nemo_relay.h | 26 ++ crates/ffi/src/api/mod.rs | 55 +++- crates/ffi/tests/unit/api/core_tests.rs | 39 +++ crates/node/plugin.d.ts | 76 ++++- crates/node/src/api/mod.rs | 6 +- crates/node/src/callable.rs | 2 +- crates/node/tests/llm_tests.mjs | 9 + crates/plugin/src/lib.rs | 6 +- crates/plugin/tests/optimization_contract.rs | 51 +++ crates/python/src/py_types/codecs.rs | 16 + .../tests/coverage/py_types_coverage_tests.rs | 51 ++- .../llm_optimization_contribution_v1.json | 53 +++ crates/types/tests/optimization_tests.rs | 93 ++++++ crates/worker/src/lib.rs | 11 +- crates/worker/tests/optimization_contract.rs | 51 +++ go/nemo_relay/callbacks.go | 10 +- go/nemo_relay/optimization.go | 172 ++++++++++ go/nemo_relay/optimization_test.go | 88 +++++ python/nemo_relay/_native.pyi | 9 + .../plugin/src/nemo_relay_plugin/__init__.py | 21 ++ python/plugin/src/nemo_relay_plugin/_api.py | 306 ++++++++++++++++++ python/tests/plugin/test_worker_sdk.py | 80 +++++ 22 files changed, 1213 insertions(+), 18 deletions(-) create mode 100644 crates/plugin/tests/optimization_contract.rs create mode 100644 crates/types/tests/fixtures/llm_optimization_contribution_v1.json create mode 100644 crates/types/tests/optimization_tests.rs create mode 100644 crates/worker/tests/optimization_contract.rs create mode 100644 go/nemo_relay/optimization.go create mode 100644 go/nemo_relay/optimization_test.go diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 0f5402c39..24dbffa15 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -437,6 +437,32 @@ NemoRelayStatus nemo_relay_llm_request_intercept_outcome_json_new(const struct F const char *pending_marks_json, char **out_outcome_json); +/** + * Allocate canonical JSON for a C LLM request-intercept callback result, + * including optional plugin-neutral optimization contributions. + * + * `annotated_json`, `pending_marks_json`, and + * `optimization_contributions_json` may be null. Null list pointers serialize + * as empty lists. Contributions use the canonical + * `LlmOptimizationContribution` JSON shape; custom `kind` strings and unknown + * top-level fields are preserved. The existing unversioned helper remains + * ABI-compatible and behaves as though this function received a null + * `optimization_contributions_json` pointer. + * + * # Safety + * + * `request` must point to a live `FfiLLMRequest`, optional JSON inputs must + * be valid null-terminated strings when non-null, and `out_outcome_json` must + * be writable. A successful output must either be transferred through a + * callback's `out_outcome_json` or freed by its caller with + * `nemo_relay_string_free`. + */ +NemoRelayStatus nemo_relay_llm_request_intercept_outcome_json_new_v2(const struct FfiLLMRequest *request, + const char *annotated_json, + const char *pending_marks_json, + const char *optimization_contributions_json, + char **out_outcome_json); + /** * Run the registered LLM conditional execution guardrail chain. * diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index 11f63871d..2efe19a67 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -55,6 +55,7 @@ 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::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginConfig, PluginError, @@ -263,6 +264,43 @@ pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new( annotated_json: *const c_char, pending_marks_json: *const c_char, out_outcome_json: *mut *mut c_char, +) -> NemoRelayStatus { + unsafe { + nemo_relay_llm_request_intercept_outcome_json_new_v2( + request, + annotated_json, + pending_marks_json, + std::ptr::null(), + out_outcome_json, + ) + } +} + +/// Allocate canonical JSON for a C LLM request-intercept callback result, +/// including optional plugin-neutral optimization contributions. +/// +/// `annotated_json`, `pending_marks_json`, and +/// `optimization_contributions_json` may be null. Null list pointers serialize +/// as empty lists. Contributions use the canonical +/// `LlmOptimizationContribution` JSON shape; custom `kind` strings and unknown +/// top-level fields are preserved. The existing unversioned helper remains +/// ABI-compatible and behaves as though this function received a null +/// `optimization_contributions_json` pointer. +/// +/// # Safety +/// +/// `request` must point to a live `FfiLLMRequest`, optional JSON inputs must +/// be valid null-terminated strings when non-null, and `out_outcome_json` must +/// be writable. A successful output must either be transferred through a +/// callback's `out_outcome_json` or freed by its caller with +/// `nemo_relay_string_free`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new_v2( + request: *const FfiLLMRequest, + annotated_json: *const c_char, + pending_marks_json: *const c_char, + optimization_contributions_json: *const c_char, + out_outcome_json: *mut *mut c_char, ) -> NemoRelayStatus { clear_last_error(); if out_outcome_json.is_null() { @@ -304,11 +342,26 @@ pub unsafe extern "C" fn nemo_relay_llm_request_intercept_outcome_json_new( } } }; + 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, + }; + match serde_json::from_value::>(value) { + Ok(value) => value, + Err(error) => { + set_last_error(&format!("invalid optimization contributions JSON: {error}")); + return NemoRelayStatus::InvalidJson; + } + } + }; let outcome = LlmRequestInterceptOutcome { request: unsafe { &*request }.0.clone(), annotated_request, pending_marks, - optimization_contributions: Vec::new(), + optimization_contributions, }; match serde_json::to_value(outcome) { Ok(value) => { diff --git a/crates/ffi/tests/unit/api/core_tests.rs b/crates/ffi/tests/unit/api/core_tests.rs index 261ef8a0e..4977e0540 100644 --- a/crates/ffi/tests/unit/api/core_tests.rs +++ b/crates/ffi/tests/unit/api/core_tests.rs @@ -30,6 +30,45 @@ fn test_ffi_llm_request_intercept_outcome_json_allocation_and_validation() { assert_eq!(outcome["annotated_request"], Json::Null); assert_eq!(outcome["pending_marks"][0]["name"], "first"); assert_eq!(outcome["pending_marks"][1]["data"]["order"], 2); + assert_eq!(outcome["optimization_contributions"], json!([])); + + let malformed_contributions = cstring(r#"[{"producer":1}]"#); + assert_eq!( + unsafe { + api::nemo_relay_llm_request_intercept_outcome_json_new_v2( + request, + ptr::null(), + ptr::null(), + malformed_contributions.as_ptr(), + &mut outcome_json, + ) + }, + NemoRelayStatus::InvalidJson + ); + assert!(outcome_json.is_null()); + + let contributions = cstring(&format!( + "[{}]", + include_str!("../../../../types/tests/fixtures/llm_optimization_contribution_v1.json") + )); + assert_eq!( + unsafe { + api::nemo_relay_llm_request_intercept_outcome_json_new_v2( + request, + ptr::null(), + ptr::null(), + contributions.as_ptr(), + &mut outcome_json, + ) + }, + NemoRelayStatus::Ok + ); + let outcome = unsafe { returned_json(outcome_json) }; + let expected: Json = serde_json::from_str(include_str!( + "../../../../types/tests/fixtures/llm_optimization_contribution_v1.json" + )) + .unwrap(); + assert_eq!(outcome["optimization_contributions"], json!([expected])); let malformed_marks = cstring(r#"{"name":"not-an-array"}"#); assert_eq!( diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index 7d8ac88d6..b8f199bf5 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -54,6 +54,76 @@ export interface PendingMarkSpec { metadata?: Json; } +/** Schema tag attached to an opaque optimization contribution payload. */ +export interface LlmOptimizationDataSchema { + name: string; + version: string; +} + +/** Model identity retained for counterfactual pricing and downstream repricing. */ +export interface LlmOptimizationModel { + model: string; + provider?: string; +} + +/** Baseline and effective model identities for a routing optimization. */ +export interface LlmOptimizationModelTransition { + baseline?: LlmOptimizationModel; + effective?: LlmOptimizationModel; +} + +/** Explicit token evidence, independent from a pricing catalog. */ +export interface LlmOptimizationTokens { + /** Token counts must be non-negative JavaScript safe integers. */ + prompt_tokens?: number; + /** Token counts must be non-negative JavaScript safe integers. */ + completion_tokens?: number; + /** Token counts must be non-negative JavaScript safe integers. */ + cache_read_tokens?: number; + /** Token counts must be non-negative JavaScript safe integers. */ + cache_write_tokens?: number; + /** Token counts must be non-negative JavaScript safe integers. */ + total_tokens?: number; +} + +/** Baseline, effective, and saved token evidence for one optimization. */ +export interface LlmOptimizationTokenImpact { + baseline?: LlmOptimizationTokens; + effective?: LlmOptimizationTokens; + saved?: LlmOptimizationTokens; + quality?: 'observed' | 'estimated'; + estimation_method?: string; +} + +/** + * One plugin's optimization evidence. + * + * `kind` is deliberately an open string so new optimizer categories round-trip + * without a Relay release. Unknown top-level fields are retained by the wire + * contract and represented by this interface's JSON extension surface. + */ +export interface LlmOptimizationContribution { + id?: string; + /** Relay ordering must remain within JavaScript's safe-integer range. */ + sequence?: number; + producer: string; + kind: 'input_compression' | 'model_routing' | (string & {}); + applied: boolean; + model_transition?: LlmOptimizationModelTransition; + token_impact?: LlmOptimizationTokenImpact; + payload_schema?: LlmOptimizationDataSchema; + payload?: Json; + [key: string]: Json | undefined; +} + +/** Canonical result returned by an LLM request intercept. */ +export interface LlmRequestInterceptOutcome { + request: Json; + annotated?: Json | null; + pendingMarks?: PendingMarkSpec[]; + optimizationContributions?: LlmOptimizationContribution[]; +} + /** * Canonical result returned by a tool execution intercept. * @@ -103,11 +173,7 @@ export interface PluginContext { name: string, priority: number, breakChain: boolean, - callback: (args: { name: string; request: Json; annotated: Json | null }) => { - request: Json; - annotated?: Json | null; - pendingMarks?: PendingMarkSpec[]; - }, + callback: (args: { name: string; request: Json; annotated: Json | null }) => LlmRequestInterceptOutcome, ): void; /** Register an LLM execution intercept for this component. */ registerLlmExecutionIntercept( diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 1e3f43a60..e6c6ce848 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -2263,7 +2263,7 @@ pub fn register_llm_request_intercept( priority: i32, break_chain: bool, #[napi( - ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> }" + ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" )] callable: ThreadsafeFunction, ) -> Result<()> { @@ -2739,7 +2739,7 @@ pub fn scope_register_llm_request_intercept( priority: i32, break_chain: bool, #[napi( - ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> }" + ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" )] callable: ThreadsafeFunction, ) -> Result<()> { @@ -2955,7 +2955,7 @@ pub fn tool_conditional_execution(env: Env, name: String, args: Json) -> Result< /// The `request` should be a JSON object with `headers` and `content` fields matching /// the `LlmRequest` schema. Returns the transformed request as JSON. #[napi( - ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions: Json[] }>" + ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }>" )] pub fn llm_request_intercepts(env: Env, name: String, request: Json) -> Result { let llm_request: LlmRequest = serde_json::from_value(request) diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 6e20b4486..29fb61644 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -256,7 +256,7 @@ pub fn wrap_js_tool_exec_fn( /// /// The JS callback receives a single JSON object /// `{ name: string, request: LlmRequest, annotated: AnnotatedLlmRequest | null }` -/// and must return `{ request, annotated?, pendingMarks? }`. +/// and must return `{ request, annotated?, pendingMarks?, optimizationContributions? }`. /// When `annotated` is non-null, request content is read-only and provider-body /// edits must be made through the returned annotation; headers remain writable. pub fn wrap_js_llm_request_intercept_fn( diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index 022afb8ad..dbae5c2ae 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -4,6 +4,7 @@ import { describe, it } from 'node:test'; import assert from 'node:assert/strict'; import { createRequire } from 'node:module'; +import { readFileSync } from 'node:fs'; const require = createRequire(import.meta.url); const lib = require('../index.js'); @@ -821,6 +822,12 @@ describe('LLM intercepts', () => { }); it('standalone request intercepts helper applies intercept chain', async () => { + const contributionFixture = JSON.parse( + readFileSync( + new URL('../../types/tests/fixtures/llm_optimization_contribution_v1.json', import.meta.url), + 'utf8', + ), + ); registerLlmRequestIntercept('node_llm_req_helper', 10, false, ({ request, annotated }) => { request.content.helper = true; return { @@ -834,6 +841,7 @@ describe('LLM intercepts', () => { }, { name: 'request.second', metadata: { source: 'node' } }, ], + optimizationContributions: [contributionFixture], }; }); @@ -856,6 +864,7 @@ describe('LLM intercepts', () => { metadata: { source: 'node' }, }, ]); + assert.deepEqual(result.optimizationContributions, [contributionFixture]); deregisterLlmRequestIntercept('node_llm_req_helper'); }); diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 477566a39..056661db0 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -17,7 +17,8 @@ use std::sync::Mutex; pub use nemo_relay_types::Json; pub use nemo_relay_types::api::event::{ - CategoryProfile, Event, EventCategory, EventSanitizeFields, PendingMarkSpec, ScopeCategory, + CategoryProfile, DataSchema, Event, EventCategory, EventSanitizeFields, PendingMarkSpec, + ScopeCategory, }; pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; @@ -25,7 +26,8 @@ pub use nemo_relay_types::api::tool::{ToolAttributes, ToolExecutionInterceptOutc pub use nemo_relay_types::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload, - LlmOptimizationTokenImpact, LlmOptimizationTokens, + LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, + LlmOptimizationTokens, }; pub use nemo_relay_types::codec::request::AnnotatedLlmRequest; pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; diff --git a/crates/plugin/tests/optimization_contract.rs b/crates/plugin/tests/optimization_contract.rs new file mode 100644 index 000000000..7aa84d703 --- /dev/null +++ b/crates/plugin/tests/optimization_contract.rs @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Conformance tests for the native plugin optimization contribution surface. + +use nemo_relay_plugin::{ + DataSchema, LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, + LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationSummary, + LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, LlmOptimizationTokens, LlmRequest, + LlmRequestInterceptOutcome, +}; +use serde_json::{Value as Json, json}; + +#[test] +fn native_plugin_sdk_round_trips_the_canonical_contribution_contract() { + let fixture: Json = serde_json::from_str(include_str!( + "../../types/tests/fixtures/llm_optimization_contribution_v1.json" + )) + .unwrap(); + let contribution: LlmOptimizationContribution = + serde_json::from_value(fixture.clone()).unwrap(); + + assert_ne!(contribution.kind, LlmOptimizationKind::input_compression()); + assert_ne!(contribution.kind, LlmOptimizationKind::model_routing()); + assert_eq!( + contribution.extra["future_top_level_field"], + fixture["future_top_level_field"] + ); + + let outcome = LlmRequestInterceptOutcome::new( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "test"}), + }, + None, + ) + .with_optimization_contribution(contribution); + let wire = serde_json::to_value(outcome).unwrap(); + assert_eq!(wire["optimization_contributions"][0], fixture); + + // Keep every nested public contract type covered by the SDK's compile-time surface. + fn exported() {} + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); +} diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 1a27c6742..c30b83196 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -618,6 +618,22 @@ impl PyAnnotatedLLMResponse { } } + /// Return Relay's plugin-neutral optimization accounting, if present. + #[getter] + pub(crate) fn optimization_summary(&self, py: Python<'_>) -> PyResult> { + match &self.inner.optimization_summary { + Some(summary) => { + let value = serde_json::to_value(summary).map_err(|error| { + pyo3::exceptions::PyValueError::new_err(format!( + "optimization summary serialization error: {error}" + )) + })?; + json_to_py(py, &value) + } + None => Ok(py.None()), + } + } + #[getter] pub(crate) fn api_specific(&self, py: Python<'_>) -> PyResult> { match &self.inner.api_specific { diff --git a/crates/python/tests/coverage/py_types_coverage_tests.rs b/crates/python/tests/coverage/py_types_coverage_tests.rs index fb59353fc..92cd15fc5 100644 --- a/crates/python/tests/coverage/py_types_coverage_tests.rs +++ b/crates/python/tests/coverage/py_types_coverage_tests.rs @@ -978,6 +978,33 @@ fn test_python_side_core_type_constructors_cover_exposed_entrypoints() { json!({"model": "demo", "messages": []}) ); + let contribution_fixture: serde_json::Value = serde_json::from_str(include_str!( + "../../../types/tests/fixtures/llm_optimization_contribution_v1.json" + )) + .unwrap(); + let contributions = json_to_py(py, &json!([contribution_fixture.clone()])).unwrap(); + let kwargs = PyDict::new(py); + kwargs + .set_item("optimization_contributions", contributions) + .unwrap(); + let outcome = module + .getattr("LLMRequestInterceptOutcome") + .unwrap() + .call((request.clone(),), Some(&kwargs)) + .unwrap(); + let round_trip = py_to_json( + outcome + .getattr("optimization_contributions") + .unwrap() + .as_any(), + ) + .unwrap(); + assert_eq!(round_trip, json!([contribution_fixture])); + assert_eq!( + round_trip[0]["future_top_level_field"], + json!({"preserved": true}) + ); + let bad_headers = PyList::empty(py); let err = module .getattr("LLMRequest") @@ -1218,7 +1245,17 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { api_name: "custom".into(), data: json!({"debug": true}), }), - optimization_summary: None, + optimization_summary: Some( + serde_json::from_value(json!({ + "schema_version": "1", + "calculation_version": "1", + "status": "partial", + "limitations": ["missing_pricing"], + "tokens_saved": {"prompt_tokens": 2, "total_tokens": 2}, + "contributions": [] + })) + .unwrap(), + ), extra: serde_json::Map::from_iter([("trace".into(), json!("abc"))]), }, }; @@ -1241,6 +1278,11 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { py_to_json(response.usage(py).unwrap().bind(py)).unwrap()["cost"]["pricing_provider"], json!("test-provider") ); + assert_eq!( + py_to_json(response.optimization_summary(py).unwrap().bind(py)).unwrap()["tokens_saved"] + ["prompt_tokens"], + json!(2) + ); assert_eq!( py_to_json(response.api_specific(py).unwrap().bind(py)).unwrap()["api_name"], json!("custom") @@ -1273,6 +1315,13 @@ fn test_annotated_llm_types_and_builtin_codecs_cover_mutators_and_codecs() { .bind(py) .is_none() ); + assert!( + response_without_api_specific + .optimization_summary(py) + .unwrap() + .bind(py) + .is_none() + ); let chat_request = PyLLMRequest { inner: nemo_relay::api::llm::LlmRequest { diff --git a/crates/types/tests/fixtures/llm_optimization_contribution_v1.json b/crates/types/tests/fixtures/llm_optimization_contribution_v1.json new file mode 100644 index 000000000..817195a57 --- /dev/null +++ b/crates/types/tests/fixtures/llm_optimization_contribution_v1.json @@ -0,0 +1,53 @@ +{ + "id": "018f0e8d-8c2f-7a74-bf43-8c3e9eaa0001", + "sequence": 7, + "producer": "example.optimizer", + "kind": "custom_energy_optimization", + "applied": true, + "model_transition": { + "baseline": { + "model": "example-baseline", + "provider": "example-provider" + }, + "effective": { + "model": "example-effective", + "provider": "example-provider" + } + }, + "token_impact": { + "baseline": { + "prompt_tokens": 120, + "completion_tokens": 24, + "cache_read_tokens": 12, + "cache_write_tokens": 8, + "total_tokens": 144 + }, + "effective": { + "prompt_tokens": 80, + "completion_tokens": 20, + "cache_read_tokens": 6, + "cache_write_tokens": 4, + "total_tokens": 100 + }, + "saved": { + "prompt_tokens": 40, + "completion_tokens": 4, + "cache_read_tokens": 6, + "cache_write_tokens": 4, + "total_tokens": 44 + }, + "quality": "observed", + "estimation_method": "provider_usage_delta" + }, + "payload_schema": { + "name": "example.optimizer.evidence", + "version": "1" + }, + "payload": { + "experiment": "fixture-v1", + "confidence": 0.98 + }, + "future_top_level_field": { + "preserved": true + } +} diff --git a/crates/types/tests/optimization_tests.rs b/crates/types/tests/optimization_tests.rs new file mode 100644 index 000000000..fd68e6478 --- /dev/null +++ b/crates/types/tests/optimization_tests.rs @@ -0,0 +1,93 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compatibility tests for the public LLM optimization wire contract. + +use nemo_relay_types::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, + LlmOptimizationPayload, LlmOptimizationTokenImpact, LlmOptimizationTokens, +}; +use serde::Serialize; +use serde_json::{Value, json}; + +#[derive(Serialize)] +struct CustomPayload { + evidence: String, +} + +impl LlmOptimizationPayload for CustomPayload { + const SCHEMA_NAME: &'static str = "example.custom_optimization"; + const SCHEMA_VERSION: &'static str = "3"; +} + +#[test] +fn standard_kind_constants_keep_their_exact_wire_values() { + assert_eq!(LlmOptimizationKind::INPUT_COMPRESSION, "input_compression"); + assert_eq!(LlmOptimizationKind::MODEL_ROUTING, "model_routing"); + assert_eq!( + serde_json::to_value(LlmOptimizationKind::input_compression()).unwrap(), + json!("input_compression") + ); + assert_eq!( + serde_json::to_value(LlmOptimizationKind::model_routing()).unwrap(), + json!("model_routing") + ); +} + +#[test] +fn custom_kinds_payloads_and_future_fields_round_trip() { + let mut contribution = LlmOptimizationContribution::new("example", "energy_reduction") + .with_payload(&CustomPayload { + evidence: "measured".to_string(), + }) + .unwrap(); + contribution + .extra + .insert("future_field".to_string(), json!({"v": 2})); + let decoded: LlmOptimizationContribution = + serde_json::from_value(serde_json::to_value(&contribution).unwrap()).unwrap(); + assert_eq!(decoded.kind.as_str(), "energy_reduction"); + assert_eq!( + decoded.payload_schema.as_ref().unwrap().name, + "example.custom_optimization" + ); + assert_eq!(decoded.extra["future_field"], json!({"v": 2})); +} + +#[test] +fn saved_prompt_tokens_remain_explicit_on_the_wire() { + let impact = LlmOptimizationTokenImpact { + saved: Some(LlmOptimizationTokens::saved_prompt(42)), + quality: Some(LlmOptimizationEvidenceQuality::Estimated), + estimation_method: Some("tokenizer-v1".to_string()), + ..LlmOptimizationTokenImpact::default() + }; + let wire = serde_json::to_value(impact).unwrap(); + assert_eq!(wire["saved"]["prompt_tokens"], 42); + assert_eq!(wire["saved"]["total_tokens"], 42); +} + +#[test] +fn canonical_all_fields_fixture_is_lossless_and_open() { + let fixture: Value = serde_json::from_str(include_str!( + "fixtures/llm_optimization_contribution_v1.json" + )) + .unwrap(); + let contribution: LlmOptimizationContribution = + serde_json::from_value(fixture.clone()).unwrap(); + + assert_eq!(contribution.kind.as_str(), "custom_energy_optimization"); + assert_eq!( + contribution + .token_impact + .as_ref() + .and_then(|impact| impact.saved.as_ref()) + .and_then(|saved| saved.cache_write_tokens), + Some(4) + ); + assert_eq!( + contribution.extra["future_top_level_field"], + json!({"preserved": true}) + ); + assert_eq!(serde_json::to_value(contribution).unwrap(), fixture); +} diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 4a456fc1c..86a91a49b 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -34,11 +34,18 @@ use futures_util::{Stream, StreamExt}; #[cfg(unix)] use hyper_util::rt::TokioIo; pub use nemo_relay_types::Json; -pub use nemo_relay_types::api::event::{Event, EventSanitizeFields, PendingMarkSpec}; +pub use nemo_relay_types::api::event::{DataSchema, Event, EventSanitizeFields, PendingMarkSpec}; pub use nemo_relay_types::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::ScopeType; pub use nemo_relay_types::api::tool::ToolExecutionInterceptOutcome; -use nemo_relay_types::codec::request::AnnotatedLlmRequest; +pub use nemo_relay_types::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, + LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload, + LlmOptimizationSummary, LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, + LlmOptimizationTokens, +}; +pub use nemo_relay_types::codec::request::AnnotatedLlmRequest; +pub use nemo_relay_types::codec::response::AnnotatedLlmResponse; pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; use nemo_relay_worker_proto::v1::relay_host_runtime_client::RelayHostRuntimeClient; diff --git a/crates/worker/tests/optimization_contract.rs b/crates/worker/tests/optimization_contract.rs new file mode 100644 index 000000000..0a39da40c --- /dev/null +++ b/crates/worker/tests/optimization_contract.rs @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Conformance tests for the worker SDK optimization contribution surface. + +use nemo_relay_worker::{ + DataSchema, LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, + LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationSummary, + LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, LlmOptimizationTokens, LlmRequest, + LlmRequestInterceptOutcome, +}; +use serde_json::{Value as Json, json}; + +#[test] +fn worker_sdk_round_trips_the_canonical_contribution_contract() { + let fixture: Json = serde_json::from_str(include_str!( + "../../types/tests/fixtures/llm_optimization_contribution_v1.json" + )) + .unwrap(); + let contribution: LlmOptimizationContribution = + serde_json::from_value(fixture.clone()).unwrap(); + + assert_ne!(contribution.kind, LlmOptimizationKind::input_compression()); + assert_ne!(contribution.kind, LlmOptimizationKind::model_routing()); + assert_eq!( + contribution.extra["future_top_level_field"], + fixture["future_top_level_field"] + ); + + let outcome = LlmRequestInterceptOutcome::new( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "test"}), + }, + None, + ) + .with_optimization_contribution(contribution); + let wire = serde_json::to_value(outcome).unwrap(); + assert_eq!(wire["optimization_contributions"][0], fixture); + + // Keep every nested public contract type covered by the SDK's compile-time surface. + fn exported() {} + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); + exported::(); +} diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index 8d771043e..1ebab4914 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -230,9 +230,10 @@ type PendingMarkSpec struct { // LLMRequestInterceptOutcome is the canonical result of an LLM request intercept. type LLMRequestInterceptOutcome struct { - Request LLMRequestDTO `json:"request"` - AnnotatedRequest json.RawMessage `json:"annotated_request"` - PendingMarks []PendingMarkSpec `json:"pending_marks"` + Request LLMRequestDTO `json:"request"` + AnnotatedRequest json.RawMessage `json:"annotated_request"` + PendingMarks []PendingMarkSpec `json:"pending_marks"` + OptimizationContributions []LLMOptimizationContribution `json:"optimization_contributions"` } // ToolExecutionInterceptOutcome is the canonical result of a tool execution @@ -573,6 +574,9 @@ func goLlmRequestInterceptTrampoline( if outcome.PendingMarks == nil { outcome.PendingMarks = []PendingMarkSpec{} } + if outcome.OptimizationContributions == nil { + outcome.OptimizationContributions = []LLMOptimizationContribution{} + } outcomeJSON, err := jsonMarshal(outcome) if err != nil { setLastErrorMessage(err.Error()) diff --git a/go/nemo_relay/optimization.go b/go/nemo_relay/optimization.go new file mode 100644 index 000000000..c36349248 --- /dev/null +++ b/go/nemo_relay/optimization.go @@ -0,0 +1,172 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nemo_relay + +import ( + "encoding/json" + "errors" +) + +// LLMOptimizationKind is an open optimization classification. The constants +// cover Relay's standard kinds; custom string values remain wire-compatible. +type LLMOptimizationKind string + +const ( + // LLMOptimizationKindInputCompression identifies request token reduction. + LLMOptimizationKindInputCompression LLMOptimizationKind = "input_compression" + // LLMOptimizationKindModelRouting identifies a model-routing decision. + LLMOptimizationKindModelRouting LLMOptimizationKind = "model_routing" +) + +// LLMOptimizationDataSchema identifies an opaque contribution payload schema. +type LLMOptimizationDataSchema struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// LLMOptimizationModel identifies one model for accounting and repricing. +type LLMOptimizationModel struct { + Model string `json:"model"` + Provider *string `json:"provider,omitempty"` +} + +// LLMOptimizationModelTransition describes baseline and effective models. +type LLMOptimizationModelTransition struct { + Baseline *LLMOptimizationModel `json:"baseline,omitempty"` + Effective *LLMOptimizationModel `json:"effective,omitempty"` +} + +// LLMOptimizationTokens retains token evidence independently from pricing. +type LLMOptimizationTokens struct { + PromptTokens *uint64 `json:"prompt_tokens,omitempty"` + CompletionTokens *uint64 `json:"completion_tokens,omitempty"` + CacheReadTokens *uint64 `json:"cache_read_tokens,omitempty"` + CacheWriteTokens *uint64 `json:"cache_write_tokens,omitempty"` + TotalTokens *uint64 `json:"total_tokens,omitempty"` +} + +// LLMOptimizationEvidenceQuality classifies observed versus estimated counts. +type LLMOptimizationEvidenceQuality string + +const ( + // LLMOptimizationEvidenceObserved means a provider or runtime observed the count. + LLMOptimizationEvidenceObserved LLMOptimizationEvidenceQuality = "observed" + // LLMOptimizationEvidenceEstimated means a tokenizer or estimator produced the count. + LLMOptimizationEvidenceEstimated LLMOptimizationEvidenceQuality = "estimated" +) + +// LLMOptimizationTokenImpact describes baseline, effective, and saved tokens. +type LLMOptimizationTokenImpact struct { + Baseline *LLMOptimizationTokens `json:"baseline,omitempty"` + Effective *LLMOptimizationTokens `json:"effective,omitempty"` + Saved *LLMOptimizationTokens `json:"saved,omitempty"` + Quality *LLMOptimizationEvidenceQuality `json:"quality,omitempty"` + EstimationMethod *string `json:"estimation_method,omitempty"` +} + +// LLMOptimizationContribution is one plugin's optimization evidence. Extra +// captures unknown top-level fields and MarshalJSON flattens them again, so a +// newer producer can round-trip through this SDK without losing information. +type LLMOptimizationContribution struct { + ID *string `json:"id,omitempty"` + Sequence *uint64 `json:"sequence,omitempty"` + Producer string `json:"producer"` + Kind LLMOptimizationKind `json:"kind"` + Applied bool `json:"applied"` + ModelTransition *LLMOptimizationModelTransition `json:"model_transition,omitempty"` + TokenImpact *LLMOptimizationTokenImpact `json:"token_impact,omitempty"` + PayloadSchema *LLMOptimizationDataSchema `json:"payload_schema,omitempty"` + Payload json.RawMessage `json:"payload,omitempty"` + Extra map[string]json.RawMessage `json:"-"` +} + +type llmOptimizationContributionWire struct { + ID *string `json:"id,omitempty"` + Sequence *uint64 `json:"sequence,omitempty"` + Producer string `json:"producer"` + Kind LLMOptimizationKind `json:"kind"` + Applied bool `json:"applied"` + ModelTransition *LLMOptimizationModelTransition `json:"model_transition,omitempty"` + TokenImpact *LLMOptimizationTokenImpact `json:"token_impact,omitempty"` + PayloadSchema *LLMOptimizationDataSchema `json:"payload_schema,omitempty"` + Payload json.RawMessage `json:"payload,omitempty"` +} + +var llmOptimizationContributionKnownFields = [...]string{ + "id", + "sequence", + "producer", + "kind", + "applied", + "model_transition", + "token_impact", + "payload_schema", + "payload", +} + +// MarshalJSON preserves and flattens forward-compatible top-level fields. +func (c LLMOptimizationContribution) MarshalJSON() ([]byte, error) { + known, err := json.Marshal(llmOptimizationContributionWire{ + ID: c.ID, + Sequence: c.Sequence, + Producer: c.Producer, + Kind: c.Kind, + Applied: c.Applied, + ModelTransition: c.ModelTransition, + TokenImpact: c.TokenImpact, + PayloadSchema: c.PayloadSchema, + Payload: c.Payload, + }) + if err != nil { + return nil, err + } + fields := make(map[string]json.RawMessage, len(c.Extra)+len(llmOptimizationContributionKnownFields)) + for name, value := range c.Extra { + fields[name] = value + } + for _, name := range llmOptimizationContributionKnownFields { + delete(fields, name) + } + if err := json.Unmarshal(known, &fields); err != nil { + return nil, err + } + return json.Marshal(fields) +} + +// UnmarshalJSON decodes known fields and retains unknown top-level fields. +func (c *LLMOptimizationContribution) UnmarshalJSON(data []byte) error { + var known llmOptimizationContributionWire + if err := json.Unmarshal(data, &known); err != nil { + return err + } + var extra map[string]json.RawMessage + if err := json.Unmarshal(data, &extra); err != nil { + return err + } + if extra == nil { + return errors.New("LLM optimization contribution must be a JSON object") + } + if _, ok := extra["producer"]; !ok { + return errors.New("LLM optimization contribution requires producer") + } + if _, ok := extra["kind"]; !ok { + return errors.New("LLM optimization contribution requires kind") + } + for _, name := range llmOptimizationContributionKnownFields { + delete(extra, name) + } + *c = LLMOptimizationContribution{ + ID: known.ID, + Sequence: known.Sequence, + Producer: known.Producer, + Kind: known.Kind, + Applied: known.Applied, + ModelTransition: known.ModelTransition, + TokenImpact: known.TokenImpact, + PayloadSchema: known.PayloadSchema, + Payload: known.Payload, + Extra: extra, + } + return nil +} diff --git a/go/nemo_relay/optimization_test.go b/go/nemo_relay/optimization_test.go new file mode 100644 index 000000000..b89bcbdbb --- /dev/null +++ b/go/nemo_relay/optimization_test.go @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nemo_relay + +import ( + "encoding/json" + "os" + "testing" +) + +func optimizationContributionFixture(t *testing.T) ([]byte, LLMOptimizationContribution) { + t.Helper() + fixture, err := os.ReadFile("../../crates/types/tests/fixtures/llm_optimization_contribution_v1.json") + if err != nil { + t.Fatalf("read optimization fixture: %v", err) + } + var contribution LLMOptimizationContribution + if err := json.Unmarshal(fixture, &contribution); err != nil { + t.Fatalf("decode optimization fixture: %v", err) + } + return fixture, contribution +} + +func assertSemanticJSONEqual(t *testing.T, expected, actual []byte) { + t.Helper() + var expectedValue any + var actualValue any + if err := json.Unmarshal(expected, &expectedValue); err != nil { + t.Fatalf("decode expected JSON: %v", err) + } + if err := json.Unmarshal(actual, &actualValue); err != nil { + t.Fatalf("decode actual JSON: %v", err) + } + if !jsonValuesEqual(expectedValue, actualValue) { + t.Fatalf("JSON mismatch\nexpected: %s\nactual: %s", expected, actual) + } +} + +func jsonValuesEqual(left, right any) bool { + leftJSON, leftErr := json.Marshal(left) + rightJSON, rightErr := json.Marshal(right) + return leftErr == nil && rightErr == nil && string(leftJSON) == string(rightJSON) +} + +func TestLLMOptimizationContributionFixtureRoundTrip(t *testing.T) { + fixture, contribution := optimizationContributionFixture(t) + if contribution.Kind == LLMOptimizationKindInputCompression || contribution.Kind == LLMOptimizationKindModelRouting { + t.Fatalf("fixture kind must exercise the open-string contract: %q", contribution.Kind) + } + if _, ok := contribution.Extra["future_top_level_field"]; !ok { + t.Fatal("expected future_top_level_field to be retained in Extra") + } + wire, err := json.Marshal(contribution) + if err != nil { + t.Fatalf("encode optimization fixture: %v", err) + } + assertSemanticJSONEqual(t, fixture, wire) +} + +func TestLLMRequestInterceptOptimizationContributionsRoundTrip(t *testing.T) { + fixture, contribution := optimizationContributionFixture(t) + const interceptName = "go_optimization_fixture" + if err := RegisterLlmRequestIntercept(interceptName, 1, false, + func(_ string, request LLMRequestDTO, annotated json.RawMessage) (LLMRequestInterceptOutcome, error) { + return LLMRequestInterceptOutcome{ + Request: request, + AnnotatedRequest: annotated, + OptimizationContributions: []LLMOptimizationContribution{contribution}, + }, nil + }); err != nil { + t.Fatalf("register request intercept: %v", err) + } + t.Cleanup(func() { _ = DeregisterLlmRequestIntercept(interceptName) }) + + outcome, err := LlmRequestIntercepts(interceptName, json.RawMessage(`{"headers":{},"content":{}}`)) + if err != nil { + t.Fatalf("run request intercepts: %v", err) + } + if len(outcome.OptimizationContributions) != 1 { + t.Fatalf("expected one optimization contribution, got %d", len(outcome.OptimizationContributions)) + } + wire, err := json.Marshal(outcome.OptimizationContributions[0]) + if err != nil { + t.Fatalf("encode returned contribution: %v", err) + } + assertSemanticJSONEqual(t, fixture, wire) +} diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index 81a211cff..90a725b90 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -410,6 +410,7 @@ class LLMRequestInterceptOutcome: request: LLMRequest, annotated_request: Optional[AnnotatedLLMRequest] = ..., pending_marks: list[PendingMarkSpec] = ..., + optimization_contributions: Optional[Sequence[Mapping[str, _JsonValue]]] = ..., ) -> None: ... @property def request(self) -> LLMRequest: ... @@ -417,6 +418,10 @@ class LLMRequestInterceptOutcome: def annotated_request(self) -> Optional[AnnotatedLLMRequest]: ... @property def pending_marks(self) -> list[PendingMarkSpec]: ... + @property + def optimization_contributions(self) -> list[_JsonObject]: + """Return ordered plugin-neutral optimization contribution objects.""" + ... class ToolExecutionInterceptOutcome: """Canonical result returned by a tool execution intercept. @@ -601,6 +606,10 @@ class AnnotatedLLMResponse: """Return normalized usage accounting, if present.""" ... @property + def optimization_summary(self) -> Optional[_JsonObject]: + """Return Relay's close-time optimization accounting, if present.""" + ... + @property def api_specific(self) -> Optional[_JsonObject]: """Return provider-specific response fields, if present.""" ... diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index 7d11e0876..cf44e8e6e 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -22,6 +22,13 @@ AnnotatedLlmRequest: An annotated Relay LLM request represented as a JSON object. PendingMarkSpec: A mark Relay emits under its managed lifecycle scope. + LlmOptimizationContribution: Plugin-neutral LLM optimization evidence. + LlmOptimizationDataSchema: Schema tag for custom optimization evidence. + LlmOptimizationEvidenceQuality: Whether token evidence was observed or estimated. + LlmOptimizationModel: Model identity used by optimization accounting. + LlmOptimizationModelTransition: Baseline and effective model identities. + LlmOptimizationTokens: Explicit token evidence by category. + LlmOptimizationTokenImpact: Baseline, effective, and saved token evidence. LlmRequestInterceptOutcome: Canonical LLM request-intercept result. ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. @@ -63,6 +70,13 @@ LlmConditionalCallback, LlmExecutionCallback, LlmNext, + LlmOptimizationContribution, + LlmOptimizationDataSchema, + LlmOptimizationEvidenceQuality, + LlmOptimizationModel, + LlmOptimizationModelTransition, + LlmOptimizationTokenImpact, + LlmOptimizationTokens, LlmRequest, LlmRequestCallback, LlmRequestInterceptOutcome, @@ -94,6 +108,13 @@ "Json", "LlmConditionalCallback", "LlmExecutionCallback", + "LlmOptimizationContribution", + "LlmOptimizationDataSchema", + "LlmOptimizationEvidenceQuality", + "LlmOptimizationModel", + "LlmOptimizationModelTransition", + "LlmOptimizationTokenImpact", + "LlmOptimizationTokens", "LlmNext", "LlmRequest", "LlmRequestCallback", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index 615e8d868..6f8d9b387 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -17,6 +17,13 @@ AnnotatedLlmRequest: An annotated Relay LLM request represented as a JSON object. PendingMarkSpec: A mark Relay emits under its managed lifecycle scope. + LlmOptimizationContribution: Plugin-neutral LLM optimization evidence. + LlmOptimizationDataSchema: Schema tag for custom optimization evidence. + LlmOptimizationEvidenceQuality: Whether token evidence was observed or estimated. + LlmOptimizationModel: Model identity used by optimization accounting. + LlmOptimizationModelTransition: Baseline and effective model identities. + LlmOptimizationTokens: Explicit token evidence by category. + LlmOptimizationTokenImpact: Baseline, effective, and saved token evidence. LlmRequestInterceptOutcome: Canonical LLM request-intercept result. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. @@ -118,6 +125,32 @@ class WorkerSdkError(Exception): """ +def _optimization_object(value: Json, field_name: str) -> Mapping[str, Json]: + if not isinstance(value, Mapping): + raise WorkerSdkError(f"optimization contribution {field_name} must be a JSON object") + return value + + +def _optimization_string(value: Json, field_name: str) -> str: + if not isinstance(value, str): + raise WorkerSdkError(f"optimization contribution {field_name} must be a string") + return value + + +def _optimization_optional_string(value: Json, field_name: str) -> str | None: + if value is None: + return None + return _optimization_string(value, field_name) + + +def _optimization_optional_u64(value: Json, field_name: str) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 0 or value > 2**64 - 1: + raise WorkerSdkError(f"optimization contribution {field_name} must be an unsigned 64-bit integer") + return value + + def _sdk_version() -> str: try: return metadata.version("nemo-relay-plugin") @@ -200,6 +233,269 @@ def to_json(self) -> dict[str, Json]: return asdict(self) +@dataclass(slots=True) +class LlmOptimizationDataSchema: + """Identify the schema of an opaque optimization contribution payload.""" + + name: str + version: str + + def to_json(self) -> dict[str, Json]: + """Convert the schema tag to its canonical wire representation.""" + return {"name": self.name, "version": self.version} + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationDataSchema: + """Decode a canonical schema tag.""" + return cls( + name=_optimization_string(value.get("name"), "payload_schema.name"), + version=_optimization_string(value.get("version"), "payload_schema.version"), + ) + + +@dataclass(slots=True) +class LlmOptimizationModel: + """Identify one model for optimization accounting and repricing.""" + + model: str + provider: str | None = None + + def to_json(self) -> dict[str, Json]: + """Convert the model identity to canonical JSON.""" + value: dict[str, Json] = {"model": self.model} + if self.provider is not None: + value["provider"] = self.provider + return value + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationModel: + """Decode a canonical model identity.""" + return cls( + model=_optimization_string(value.get("model"), "model"), + provider=_optimization_optional_string(value.get("provider"), "provider"), + ) + + +@dataclass(slots=True) +class LlmOptimizationModelTransition: + """Describe the counterfactual and effective models for one optimization.""" + + baseline: LlmOptimizationModel | None = None + effective: LlmOptimizationModel | None = None + + def to_json(self) -> dict[str, Json]: + """Convert the model transition to canonical JSON.""" + value: dict[str, Json] = {} + if self.baseline is not None: + value["baseline"] = self.baseline.to_json() + if self.effective is not None: + value["effective"] = self.effective.to_json() + return value + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationModelTransition: + """Decode a canonical model transition.""" + baseline = value.get("baseline") + effective = value.get("effective") + return cls( + baseline=( + LlmOptimizationModel.from_json(_optimization_object(baseline, "model_transition.baseline")) + if baseline is not None + else None + ), + effective=( + LlmOptimizationModel.from_json(_optimization_object(effective, "model_transition.effective")) + if effective is not None + else None + ), + ) + + +@dataclass(slots=True) +class LlmOptimizationTokens: + """Retain token evidence independently from monetary pricing.""" + + prompt_tokens: int | None = None + completion_tokens: int | None = None + cache_read_tokens: int | None = None + cache_write_tokens: int | None = None + total_tokens: int | None = None + + def to_json(self) -> dict[str, Json]: + """Convert populated token fields to canonical JSON.""" + return { + name: value + for name, value in ( + ("prompt_tokens", self.prompt_tokens), + ("completion_tokens", self.completion_tokens), + ("cache_read_tokens", self.cache_read_tokens), + ("cache_write_tokens", self.cache_write_tokens), + ("total_tokens", self.total_tokens), + ) + if value is not None + } + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationTokens: + """Decode canonical token evidence.""" + return cls( + prompt_tokens=_optimization_optional_u64(value.get("prompt_tokens"), "prompt_tokens"), + completion_tokens=_optimization_optional_u64(value.get("completion_tokens"), "completion_tokens"), + cache_read_tokens=_optimization_optional_u64(value.get("cache_read_tokens"), "cache_read_tokens"), + cache_write_tokens=_optimization_optional_u64(value.get("cache_write_tokens"), "cache_write_tokens"), + total_tokens=_optimization_optional_u64(value.get("total_tokens"), "total_tokens"), + ) + + +class LlmOptimizationEvidenceQuality(str, Enum): + """Classify token evidence as directly observed or estimated.""" + + OBSERVED = "observed" + ESTIMATED = "estimated" + + +@dataclass(slots=True) +class LlmOptimizationTokenImpact: + """Describe baseline, effective, and explicitly saved token evidence.""" + + baseline: LlmOptimizationTokens | None = None + effective: LlmOptimizationTokens | None = None + saved: LlmOptimizationTokens | None = None + quality: LlmOptimizationEvidenceQuality | str | None = None + estimation_method: str | None = None + + def to_json(self) -> dict[str, Json]: + """Convert the token impact to canonical JSON.""" + value: dict[str, Json] = {} + if self.baseline is not None: + value["baseline"] = self.baseline.to_json() + if self.effective is not None: + value["effective"] = self.effective.to_json() + if self.saved is not None: + value["saved"] = self.saved.to_json() + if self.quality is not None: + value["quality"] = self.quality.value if isinstance(self.quality, Enum) else self.quality + if self.estimation_method is not None: + value["estimation_method"] = self.estimation_method + return value + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationTokenImpact: + """Decode a canonical token impact.""" + + def tokens(name: str) -> LlmOptimizationTokens | None: + item = value.get(name) + if item is None: + return None + return LlmOptimizationTokens.from_json(_optimization_object(item, f"token_impact.{name}")) + + quality = value.get("quality") + if quality is not None: + try: + quality = LlmOptimizationEvidenceQuality(_optimization_string(quality, "token_impact.quality")) + except ValueError as exc: + raise WorkerSdkError( + "optimization contribution token_impact.quality must be 'observed' or 'estimated'" + ) from exc + return cls( + baseline=tokens("baseline"), + effective=tokens("effective"), + saved=tokens("saved"), + quality=quality, + estimation_method=_optimization_optional_string( + value.get("estimation_method"), "token_impact.estimation_method" + ), + ) + + +@dataclass(slots=True) +class LlmOptimizationContribution: + """One plugin's lossless, forward-compatible optimization evidence.""" + + producer: str + kind: str + applied: bool = True + id: str | None = None + sequence: int | None = None + model_transition: LlmOptimizationModelTransition | None = None + token_impact: LlmOptimizationTokenImpact | None = None + payload_schema: LlmOptimizationDataSchema | None = None + payload: Json | None = None + extra: dict[str, Json] = field(default_factory=dict) + + INPUT_COMPRESSION = "input_compression" + MODEL_ROUTING = "model_routing" + + def to_json(self) -> dict[str, Json]: + """Convert this contribution while flattening unknown top-level fields.""" + if self.payload is not None and self.payload_schema is None: + raise WorkerSdkError("optimization contribution payload requires payload_schema") + value = dict(self.extra) + value.update({"producer": self.producer, "kind": self.kind, "applied": self.applied}) + if self.id is not None: + value["id"] = self.id + if self.sequence is not None: + value["sequence"] = self.sequence + if self.model_transition is not None: + value["model_transition"] = self.model_transition.to_json() + if self.token_impact is not None: + value["token_impact"] = self.token_impact.to_json() + if self.payload_schema is not None: + value["payload_schema"] = self.payload_schema.to_json() + if self.payload is not None: + value["payload"] = self.payload + return value + + @classmethod + def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationContribution: + """Decode all known fields and retain unknown top-level fields in ``extra``.""" + value = _optimization_object(value, "value") + known = { + "id", + "sequence", + "producer", + "kind", + "applied", + "model_transition", + "token_impact", + "payload_schema", + "payload", + } + transition = value.get("model_transition") + impact = value.get("token_impact") + schema = value.get("payload_schema") + applied = value.get("applied", False) + if not isinstance(applied, bool): + raise WorkerSdkError("optimization contribution applied must be a boolean") + contribution = cls( + id=_optimization_optional_string(value.get("id"), "id"), + sequence=_optimization_optional_u64(value.get("sequence"), "sequence"), + producer=_optimization_string(value.get("producer"), "producer"), + kind=_optimization_string(value.get("kind"), "kind"), + applied=applied, + model_transition=( + LlmOptimizationModelTransition.from_json(_optimization_object(transition, "model_transition")) + if transition is not None + else None + ), + token_impact=( + LlmOptimizationTokenImpact.from_json(_optimization_object(impact, "token_impact")) + if impact is not None + else None + ), + payload_schema=( + LlmOptimizationDataSchema.from_json(_optimization_object(schema, "payload_schema")) + if schema is not None + else None + ), + payload=value.get("payload"), + extra={name: item for name, item in value.items() if name not in known}, + ) + if contribution.payload is not None and contribution.payload_schema is None: + raise WorkerSdkError("optimization contribution payload requires payload_schema") + return contribution + + @dataclass(slots=True) class LlmRequestInterceptOutcome: """Canonical result returned by a Python worker LLM request intercept.""" @@ -207,6 +503,7 @@ class LlmRequestInterceptOutcome: request: LlmRequest annotated_request: AnnotatedLlmRequest | None = None pending_marks: list[PendingMarkSpec] = field(default_factory=list) + optimization_contributions: list[LlmOptimizationContribution] = field(default_factory=list) def to_json(self) -> dict[str, Json]: """Convert this outcome to the canonical worker-envelope payload.""" @@ -219,10 +516,19 @@ def to_json(self) -> dict[str, Json]: if not isinstance(mark, PendingMarkSpec): raise WorkerSdkError("LLM request intercept outcome pending_marks must contain PendingMarkSpec values") marks.append(mark.to_json()) + contributions = [] + for contribution in self.optimization_contributions: + if not isinstance(contribution, LlmOptimizationContribution): + raise WorkerSdkError( + "LLM request intercept outcome optimization_contributions must contain " + "LlmOptimizationContribution values" + ) + contributions.append(contribution.to_json()) return { "request": self.request, "annotated_request": self.annotated_request, "pending_marks": marks, + "optimization_contributions": contributions, } diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 3f33fbb27..3acfad427 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -27,6 +27,7 @@ ConfigDiagnostic, DiagnosticLevel, Json, + LlmOptimizationContribution, LlmRequestInterceptOutcome, PendingMarkSpec, PluginContext, @@ -64,6 +65,48 @@ AUTH_TOKEN = "token" +def _optimization_contribution_fixture() -> Json: + fixture_path = ( + Path(__file__).resolve().parents[3] + / "crates" + / "types" + / "tests" + / "fixtures" + / "llm_optimization_contribution_v1.json" + ) + return json.loads(fixture_path.read_text(encoding="utf-8")) + + +def test_optimization_contribution_fixture_round_trips_losslessly(): + fixture = _optimization_contribution_fixture() + contribution = LlmOptimizationContribution.from_json(fixture) + + outcome = LlmRequestInterceptOutcome( + request={"headers": {}, "content": {"model": "test"}}, + optimization_contributions=[contribution], + ).to_json() + + assert outcome["optimization_contributions"] == [fixture] + assert contribution.kind == fixture["kind"] + assert contribution.kind not in { + LlmOptimizationContribution.INPUT_COMPRESSION, + LlmOptimizationContribution.MODEL_ROUTING, + } + assert contribution.extra["future_top_level_field"] == fixture["future_top_level_field"] + assert ( + LlmRequestInterceptOutcome(request={"headers": {}, "content": {}}).to_json()["optimization_contributions"] == [] + ) + + +def test_optimization_contribution_requires_schema_for_payload(): + with pytest.raises(WorkerSdkError, match="payload_schema"): + LlmOptimizationContribution.from_json( + {"producer": "test", "kind": "custom", "applied": True, "payload": {"value": 1}} + ) + with pytest.raises(WorkerSdkError, match="producer must be a string"): + LlmOptimizationContribution.from_json({"producer": 1, "kind": "custom"}) + + class GrpcAbort(Exception): def __init__(self, code: object, details: str) -> None: super().__init__(f"{code}: {details}") @@ -1211,6 +1254,43 @@ def llm_request(name: str, request: Json, annotated: Json | None) -> LlmRequestI assert outcome["pending_marks"] == [] +async def test_llm_request_intercept_preserves_optimization_contribution_worker_envelope(): + fixture = _optimization_contribution_fixture() + + class OptimizationPlugin(WorkerPlugin): + plugin_id = "tests.optimization_contribution" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def llm_request(name: str, request: Json, annotated: Json | None) -> LlmRequestInterceptOutcome: + del name + return LlmRequestInterceptOutcome( + request=request, + annotated_request=annotated, + optimization_contributions=[LlmOptimizationContribution.from_json(fixture)], + ) + + ctx.register_llm_request_intercept("optimization", llm_request) + + service = _service(OptimizationPlugin(), RecordingHostStub()) + await _register(service) + response = await service.Invoke( + _invoke_request( + "optimization", + pb.LLM_REQUEST_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + + assert response.WhichOneof("result") == "llm_request" + assert response.llm_request.outcome.schema == LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA + outcome = _envelope_value(response.llm_request.outcome) + assert outcome["optimization_contributions"] == [fixture] + assert outcome["optimization_contributions"][0]["future_top_level_field"] == {"preserved": True} + + async def test_stream_invoke_success_and_failures(service: _WorkerService, host_stub: RecordingHostStub): await _register(service) From 066663ce13c2b30616b47216eaadbc81ccb48488 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 00:37:32 -0600 Subject: [PATCH 07/10] fix: disambiguate optimization exporter costs Signed-off-by: Bryan Bednarski --- .../core/src/observability/openinference.rs | 103 ++++++++++++---- crates/core/src/observability/otel.rs | 66 ++++++++++- crates/core/tests/unit/atif_tests.rs | 92 ++++++++++++++ .../unit/observability/openinference_tests.rs | 112 ++++++++++++++++++ .../tests/unit/observability/otel_tests.rs | 105 ++++++++++++++++ 5 files changed, 448 insertions(+), 30 deletions(-) diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 5dd4b98a0..37aa8ea16 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -954,29 +954,64 @@ fn push_optimization_attributes( i64::try_from(tokens).unwrap_or(i64::MAX), )); } - let costs = [ - ( - "nemo_relay.llm.optimization.baseline_cost", - summary - .baseline_cost - .as_ref() - .and_then(|cost| cost.total_or_component_sum()), - ), - ( - "nemo_relay.llm.optimization.actual_cost", - summary - .actual_cost - .as_ref() - .and_then(|cost| cost.total_or_component_sum()), - ), - ( + 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", - summary.estimated_cost_saved, - ), - ]; - for (key, value) in costs { - if let Some(value) = value { - attributes.push(KeyValue::new(key, value)); + 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( @@ -986,17 +1021,33 @@ fn push_optimization_attributes( crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", }, )); - let provenance = summary + if let Some(source) = summary .baseline_cost .as_ref() - .or(summary.actual_cost.as_ref()); - if let Some(source) = provenance.and_then(|cost| cost.pricing_source.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) = provenance.and_then(|cost| cost.pricing_as_of.as_ref()) { + 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(), diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index f3b5d85bd..9624b0527 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -753,6 +753,24 @@ fn push_optimization_attributes( 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() @@ -763,11 +781,35 @@ fn push_optimization_attributes( 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( @@ -782,17 +824,33 @@ fn push_optimization_attributes( crate::codec::optimization::LlmOptimizationSummaryStatus::Partial => "partial", }, )); - let provenance = summary + if let Some(source) = summary .baseline_cost .as_ref() - .or(summary.actual_cost.as_ref()); - if let Some(source) = provenance.and_then(|cost| cost.pricing_source.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) = provenance.and_then(|cost| cost.pricing_as_of.as_ref()) { + 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(), diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index c683fb2fb..701e6fd21 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -1558,6 +1558,98 @@ fn test_optimization_summary_projects_to_step_and_final_metrics() { assert_eq!(optimization["estimated_cost_saved_usd"], 0.01); } +#[test] +fn test_atif_preserves_non_usd_summary_without_labeling_savings_as_usd() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version": "1", + "calculation_version": "1", + "status": "complete", + "baseline_model": {"model": "baseline"}, + "effective_model": {"model": "effective"}, + "tokens_saved": {"prompt_tokens": 12, "total_tokens": 12}, + "baseline_cost": {"total": 0.02, "currency": "EUR", "source": "model_pricing"}, + "actual_cost": {"total": 0.01, "currency": "EUR", "source": "model_pricing"}, + "estimated_cost_saved": 0.01, + "currency": "EUR", + "contributions": [] + })) + .unwrap(); + let mut response = annotated_response_with_usage( + "effective", + Usage { + prompt_tokens: Some(8), + ..Usage::default() + }, + ); + response.optimization_summary = Some(summary); + + let metrics = extract_metrics(&json!({}), None, None, Some(&response)).unwrap(); + assert_eq!( + metrics + .extra + .as_ref() + .unwrap() + .pointer("/nemo_relay/optimization/currency"), + Some(&json!("EUR")) + ); + let final_metrics = compute_final_metrics(&[AtifStep { + step_id: 1, + source: "agent".to_string(), + message: json!("done"), + timestamp: None, + model_name: Some("effective".to_string()), + reasoning_effort: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: Some(metrics), + llm_call_count: Some(1), + is_copied_context: None, + extra: None, + }]) + .unwrap(); + assert_eq!( + final_metrics + .extra + .as_ref() + .unwrap() + .pointer("/nemo_relay/optimization/estimated_cost_saved_usd"), + Some(&Json::Null) + ); +} + +#[test] +fn test_atif_without_optimization_summary_does_not_add_optimization_extra() { + let response = annotated_response_with_usage( + "effective", + Usage { + prompt_tokens: Some(8), + ..Usage::default() + }, + ); + let metrics = extract_metrics(&json!({}), None, None, Some(&response)).unwrap(); + assert!(metrics.extra.is_none()); + + let final_metrics = compute_final_metrics(&[AtifStep { + step_id: 1, + source: "agent".to_string(), + message: json!("done"), + timestamp: None, + model_name: Some("effective".to_string()), + reasoning_effort: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: Some(metrics), + llm_call_count: Some(1), + is_copied_context: None, + extra: None, + }]) + .unwrap(); + assert!(final_metrics.extra.is_none()); +} + #[test] fn test_exporter_llm_lifecycle_plain_input() { // Input without LlmRequest envelope — passed through unchanged. diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 28d4f5bdc..0d2f5c107 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -163,6 +163,118 @@ fn optimization_summary_emits_namespaced_openinference_attributes() { attributes["nemo_relay.llm.optimization.pricing_source"], "test" ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_cost_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.estimated_cost_saved_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_pricing_as_of"], + "2026-07-08" + ); +} + +#[test] +fn optimization_cost_attributes_preserve_independent_currency_and_provenance() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"partial", + "limitations":["cost_currency_mismatch"], + "tokens_saved":{}, + "baseline_cost":{"total":2.0,"currency":"EUR","source":"model_pricing","pricing_as_of":"2026-01-01"}, + "actual_cost":{"total":1.0,"currency":"GBP","source":"provider_reported","pricing_as_of":"2026-02-02","pricing_source":"provider"}, + "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_cost_currency"], + "EUR" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost_currency"], + "GBP" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_source"], + "provider" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_as_of"], + "2026-02-02" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_source"], + "provider" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_as_of"], + "2026-01-01" + ); + assert!(!attributes.contains_key("nemo_relay.llm.optimization.estimated_cost_saved")); + assert!(!attributes.contains_key("nemo_relay.llm.optimization.estimated_cost_saved_currency")); +} + +#[test] +fn complete_non_usd_optimization_costs_keep_currency_and_independent_provenance() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"complete", + "tokens_saved":{}, + "baseline_cost":{"total":2.0,"currency":"EUR","source":"model_pricing","pricing_as_of":"2026-01-01","pricing_source":"baseline-catalog"}, + "actual_cost":{"total":1.0,"currency":"EUR","source":"provider_reported","pricing_as_of":"2026-02-02","pricing_source":"provider-meter"}, + "estimated_cost_saved":1.0, "currency":"EUR", "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + + assert_attr( + &attributes, + "nemo_relay.llm.optimization.baseline_cost_currency", + "EUR", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.actual_cost_currency", + "EUR", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.estimated_cost_saved_currency", + "EUR", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.baseline_pricing_source", + "baseline-catalog", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.actual_pricing_source", + "provider-meter", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.baseline_pricing_as_of", + "2026-01-01", + ); + assert_attr( + &attributes, + "nemo_relay.llm.optimization.actual_pricing_as_of", + "2026-02-02", + ); } fn install_test_pricing(model_id: &str) { diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index c82dbbefd..b764c29a6 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -82,6 +82,111 @@ fn optimization_summary_emits_namespaced_otel_attributes() { attributes["nemo_relay.llm.optimization.pricing_as_of"], "2026-07-08" ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_cost_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.estimated_cost_saved_currency"], + "USD" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_pricing_source"], + "test" + ); +} + +#[test] +fn optimization_cost_attributes_preserve_independent_currency_and_provenance() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"partial", + "limitations":["cost_currency_mismatch"], + "tokens_saved":{}, + "baseline_cost":{"total":2.0,"currency":"EUR","source":"model_pricing","pricing_as_of":"2026-01-01"}, + "actual_cost":{"total":1.0,"currency":"GBP","source":"provider_reported","pricing_as_of":"2026-02-02","pricing_source":"provider"}, + "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_cost_currency"], + "EUR" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost_currency"], + "GBP" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_source"], + "provider" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_as_of"], + "2026-02-02" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_source"], + "provider" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.pricing_as_of"], + "2026-01-01" + ); + assert!(!attributes.contains_key("nemo_relay.llm.optimization.estimated_cost_saved")); + assert!(!attributes.contains_key("nemo_relay.llm.optimization.estimated_cost_saved_currency")); +} + +#[test] +fn complete_non_usd_optimization_costs_keep_currency_and_independent_provenance() { + let summary: crate::codec::optimization::LlmOptimizationSummary = + serde_json::from_value(json!({ + "schema_version":"1", "calculation_version":"1", "status":"complete", + "tokens_saved":{}, + "baseline_cost":{"total":2.0,"currency":"EUR","source":"model_pricing","pricing_as_of":"2026-01-01","pricing_source":"baseline-catalog"}, + "actual_cost":{"total":1.0,"currency":"EUR","source":"provider_reported","pricing_as_of":"2026-02-02","pricing_source":"provider-meter"}, + "estimated_cost_saved":1.0, "currency":"EUR", "contributions":[] + })) + .unwrap(); + let mut attributes = Vec::new(); + push_optimization_attributes(&mut attributes, &summary); + let attributes = attr_map(&attributes); + + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_cost_currency"], + "EUR" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_cost_currency"], + "EUR" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.estimated_cost_saved_currency"], + "EUR" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_pricing_source"], + "baseline-catalog" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_source"], + "provider-meter" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.baseline_pricing_as_of"], + "2026-01-01" + ); + assert_eq!( + attributes["nemo_relay.llm.optimization.actual_pricing_as_of"], + "2026-02-02" + ); } fn install_test_pricing(model_id: &str) { From aca78de493fd61d7cf36aa3323a249442469486f Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 01:32:19 -0600 Subject: [PATCH 08/10] fix(optimization): align bindings and mark chronology Signed-off-by: Bryan Bednarski --- crates/core/src/api/llm.rs | 8 +- crates/core/src/api/optimization.rs | 25 ++++- .../tests/integration/middleware_tests.rs | 92 ++++++++++++++++++- .../core/tests/integration/pipeline_tests.rs | 14 ++- crates/node/src/api/mod.rs | 6 +- crates/node/tests/llm_tests.mjs | 7 ++ crates/python/src/py_types/codecs.rs | 12 ++- crates/python/src/py_types/mod.rs | 2 + .../tests/coverage/py_types_coverage_tests.rs | 17 +++- crates/worker/tests/optimization_contract.rs | 10 +- go/nemo_relay/optimization.go | 12 +++ go/nemo_relay/optimization_test.go | 27 ++++++ python/plugin/src/nemo_relay_plugin/_api.py | 2 +- python/tests/plugin/test_worker_sdk.py | 10 ++ 14 files changed, 220 insertions(+), 24 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index f7aee3f72..0dd76ef85 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -408,7 +408,7 @@ fn emit_optimization_marks_with( mut sanitize: impl FnMut(Event) -> Option, mut enqueue: impl FnMut(&Event, &[EventSubscriberFn]) -> bool, ) { - let contributions = handle.optimization_recorder.unemitted(); + let contributions = handle.optimization_recorder.unemitted_with_timestamps(); if contributions.is_empty() { return; } @@ -416,15 +416,17 @@ fn emit_optimization_marks_with( eprintln!("nemo_relay: unable to emit LLM optimization marks: {error}"); return; } - for contribution in contributions { + for (contribution, recorded_at) in contributions { let offset = contribution.sequence.unwrap_or(0).saturating_add(2); let offset = i64::try_from(offset).unwrap_or(i64::MAX); + let request_ordered_timestamp = handle.started_at + TimeDelta::microseconds(offset); + let timestamp = recorded_at.max(request_ordered_timestamp); let data = serde_json::to_value(&contribution).unwrap_or(Json::Null); let event = Event::Mark(MarkEvent::new( BaseEvent::builder() .name("nemo_relay.llm.optimization") .parent_uuid(handle.uuid) - .timestamp(handle.started_at + TimeDelta::microseconds(offset)) + .timestamp(timestamp) .data(data) .data_schema(DataSchema { name: "nemo.relay.llm_optimization_contribution".to_string(), diff --git a/crates/core/src/api/optimization.rs b/crates/core/src/api/optimization.rs index cf8880457..2f7bba7ac 100644 --- a/crates/core/src/api/optimization.rs +++ b/crates/core/src/api/optimization.rs @@ -6,6 +6,7 @@ use std::collections::BTreeSet; use std::sync::{Arc, Mutex}; +use chrono::{DateTime, Utc}; use serde::Serialize; use uuid::Uuid; @@ -30,6 +31,7 @@ pub const MAX_LLM_OPTIMIZATION_CONTRIBUTION_ATTEMPTS: usize = 64; #[derive(Debug, Default)] struct AccumulatorState { contributions: Vec, + recorded_at: Vec>, total_contribution_bytes: usize, attempted_contributions: usize, emitted: usize, @@ -143,6 +145,7 @@ impl LlmOptimizationRecorder { state.total_contribution_bytes = total_contribution_bytes; state.contributions.push(contribution); + state.recorded_at.push(Utc::now()); return true; } } @@ -174,12 +177,31 @@ impl LlmOptimizationRecorder { /// /// This does not move the cursor. Call [`Self::mark_emitted`] only after /// the asynchronous dispatcher accepts an item. + #[cfg(test)] pub(crate) fn unemitted(&self) -> Vec { + self.unemitted_with_timestamps() + .into_iter() + .map(|(contribution, _)| contribution) + .collect() + } + + /// Snapshot unacknowledged contributions with their acceptance time. + /// + /// The timestamp is captured only after the contribution wins its final + /// sequence and size checks, so execution-time marks retain commit-time + /// ordering even when they are emitted at the LLM close boundary. + pub(crate) fn unemitted_with_timestamps( + &self, + ) -> Vec<(LlmOptimizationContribution, DateTime)> { let Ok(state) = self.state.lock() else { return Vec::new(); }; let start = state.emitted.min(state.contributions.len()); - state.contributions[start..].to_vec() + state.contributions[start..] + .iter() + .cloned() + .zip(state.recorded_at[start..].iter().copied()) + .collect() } /// Advance the delivery cursor for a bounded number of accepted marks. @@ -250,6 +272,7 @@ impl LlmOptimizationRecorder { limitations.push("invalid_contribution_payload_schema".to_string()); state.invalid_payload_schema = false; } + state.recorded_at.clear(); FinishedContributions { contributions: std::mem::take(&mut state.contributions), limitations, diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index d8882f04a..e504a6481 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -51,7 +51,7 @@ use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, }; use nemo_relay::api::runtime::{create_scope_stack, current_scope_stack, set_thread_scope_stack}; -use nemo_relay::api::scope::{ScopeHandle, ScopeType}; +use nemo_relay::api::scope::{EmitMarkEventParams, ScopeHandle, ScopeType, event}; use nemo_relay::api::scope::{pop_scope, push_scope}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::api::tool::{ @@ -3564,6 +3564,9 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { && event.scope_category() == Some(ScopeCategory::End) }) .unwrap(); + assert!(marks[0].timestamp() > start.timestamp()); + assert!(marks[0].timestamp() <= marks[1].timestamp()); + assert!(marks[1].timestamp() <= end.timestamp()); let summary = end .annotated_response() .unwrap() @@ -3587,6 +3590,93 @@ async fn test_managed_llm_materializes_optimization_mark_and_end_summary() { deregister_subscriber("optimization_observer").unwrap(); } +#[tokio::test] +async fn execution_optimization_mark_keeps_decision_commit_timestamp_order() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "optimization_timestamp_observer", + Arc::new(move |event: &Event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_execution_intercept( + "optimization_timestamp_contributor", + 1, + Arc::new(|_name, request, next| { + Box::pin(async move { + event( + EmitMarkEventParams::builder() + .name("test.router.requested") + .build(), + )?; + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + event( + EmitMarkEventParams::builder() + .name("test.router.decision") + .build(), + )?; + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + assert!(record_llm_optimization_contribution( + LlmOptimizationContribution::new("test.execution.timestamp", "model_routing") + )); + next(request).await + }) + }), + ) + .unwrap(); + + llm_call_execute( + LlmCallExecuteParams::builder() + .name("optimized-timestamp-llm") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"prompt": "hello"}), + }) + .func(Arc::new(|_| { + Box::pin(async { Ok(json!({"response": "done"})) }) + })) + .build(), + ) + .await + .unwrap(); + + let captured = captured_events_snapshot(&events); + let requested = captured + .iter() + .find(|event| event.name() == "test.router.requested") + .unwrap(); + let decision = captured + .iter() + .find(|event| event.name() == "test.router.decision") + .unwrap(); + let contribution = captured + .iter() + .find(|event| { + event.name() == "nemo_relay.llm.optimization" + && event.data().and_then(|data| data["producer"].as_str()) + == Some("test.execution.timestamp") + }) + .unwrap(); + let end = captured + .iter() + .find(|event| { + event.name() == "optimized-timestamp-llm" + && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap(); + + assert!(requested.timestamp() < decision.timestamp()); + assert!(decision.timestamp() < contribution.timestamp()); + assert!(contribution.timestamp() <= end.timestamp()); + + deregister_llm_execution_intercept("optimization_timestamp_contributor").unwrap(); + deregister_subscriber("optimization_timestamp_observer").unwrap(); +} + #[tokio::test] async fn test_stream_optimization_mark_uses_the_llm_captured_sanitizer_scope() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 083171054..53f0c7b14 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -52,6 +52,14 @@ use nemo_relay::json::Json; static TEST_MUTEX: Mutex<()> = Mutex::new(()); +struct ResetPricingResolverGuard; + +impl Drop for ResetPricingResolverGuard { + fn drop(&mut self) { + let _ = reset_active_pricing_resolver(); + } +} + fn is_scope_event(event: &Event, scope_type: ScopeType, scope_category: ScopeCategory) -> bool { event.scope_type() == Some(scope_type) && event.scope_category() == Some(scope_category) } @@ -1282,6 +1290,7 @@ impl LlmResponseCodec for FailingResponseCodec { #[tokio::test] async fn test_response_codec_populates_annotated_response() { let _lock = TEST_MUTEX.lock().unwrap(); + let _pricing_guard = ResetPricingResolverGuard; reset_global(); setup_isolated_thread(); install_mock_response_pricing(); @@ -1330,7 +1339,6 @@ async fn test_response_codec_populates_annotated_response() { ); deregister_subscriber("resp_codec_sub").unwrap(); - reset_active_pricing_resolver().unwrap(); } #[tokio::test] @@ -1597,6 +1605,7 @@ async fn test_request_codec_annotation_uses_sanitized_start_payload() { #[tokio::test] async fn test_stream_response_codec_populates_annotated_response() { let _lock = TEST_MUTEX.lock().unwrap(); + let _pricing_guard = ResetPricingResolverGuard; reset_global(); setup_isolated_thread(); install_mock_response_pricing(); @@ -1654,12 +1663,12 @@ async fn test_stream_response_codec_populates_annotated_response() { ); deregister_subscriber("stream_resp_codec_sub").unwrap(); - reset_active_pricing_resolver().unwrap(); } #[tokio::test] async fn managed_buffered_and_streaming_close_price_the_committed_route_not_response_alias() { let _lock = TEST_MUTEX.lock().unwrap(); + let _pricing_guard = ResetPricingResolverGuard; reset_global(); setup_isolated_thread(); install_routed_response_pricing(); @@ -1797,7 +1806,6 @@ async fn managed_buffered_and_streaming_close_price_the_committed_route_not_resp ); deregister_subscriber("routed_alias_pricing_sub").unwrap(); - reset_active_pricing_resolver().unwrap(); } #[tokio::test] diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index e6c6ce848..fa30cd609 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -2263,7 +2263,7 @@ pub fn register_llm_request_intercept( priority: i32, break_chain: bool, #[napi( - ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" + ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: 'input_compression' | 'model_routing' | (string & {}); applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" )] callable: ThreadsafeFunction, ) -> Result<()> { @@ -2739,7 +2739,7 @@ pub fn scope_register_llm_request_intercept( priority: i32, break_chain: bool, #[napi( - ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" + ts_arg_type = "(args: { name: string; request: Json; annotated: Json | null }) => { request: Json; annotated?: Json | null; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions?: Array<{ id?: string; sequence?: number; producer: string; kind: 'input_compression' | 'model_routing' | (string & {}); applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }" )] callable: ThreadsafeFunction, ) -> Result<()> { @@ -2955,7 +2955,7 @@ pub fn tool_conditional_execution(env: Env, name: String, args: Json) -> Result< /// The `request` should be a JSON object with `headers` and `content` fields matching /// the `LlmRequest` schema. Returns the transformed request as JSON. #[napi( - ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions: Array<{ id?: string; sequence?: number; producer: string; kind: string; applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }>" + ts_return_type = "Promise<{ request: Json; annotated: Json | null; pendingMarks: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }>; optimizationContributions: Array<{ id?: string; sequence?: number; producer: string; kind: 'input_compression' | 'model_routing' | (string & {}); applied: boolean; model_transition?: { baseline?: { model: string; provider?: string }; effective?: { model: string; provider?: string } }; token_impact?: { baseline?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; effective?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; saved?: { prompt_tokens?: number; completion_tokens?: number; cache_read_tokens?: number; cache_write_tokens?: number; total_tokens?: number }; quality?: 'observed' | 'estimated'; estimation_method?: string }; payload_schema?: { name: string; version: string }; payload?: Json; [key: string]: Json | undefined }> }>" )] pub fn llm_request_intercepts(env: Env, name: String, request: Json) -> Result { let llm_request: LlmRequest = serde_json::from_value(request) diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index dbae5c2ae..f48144941 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -868,6 +868,13 @@ describe('LLM intercepts', () => { deregisterLlmRequestIntercept('node_llm_req_helper'); }); + it('generated request-intercept declarations preserve the open optimization kind', () => { + const declarations = readFileSync(new URL('../index.d.ts', import.meta.url), 'utf8'); + const openKind = "kind: 'input_compression' | 'model_routing' | (string & {})"; + + assert.equal(declarations.split(openKind).length - 1, 3); + }); + it('standalone conditional execution helper throws on rejection', async () => { registerLlmConditionalExecutionGuardrail('node_llm_cond_helper', 10, () => 'llm blocked by helper'); try { diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index c30b83196..636f82fb1 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -18,6 +18,7 @@ use super::{ FORCE_ANNOTATED_REQUEST_TOOLS_SERIALIZATION_ERROR, FORCE_ANNOTATED_RESPONSE_API_SPECIFIC_SERIALIZATION_ERROR, FORCE_ANNOTATED_RESPONSE_MESSAGE_SERIALIZATION_ERROR, + FORCE_ANNOTATED_RESPONSE_OPTIMIZATION_SUMMARY_SERIALIZATION_ERROR, FORCE_ANNOTATED_RESPONSE_TOOL_CALLS_SERIALIZATION_ERROR, FORCE_ANNOTATED_RESPONSE_USAGE_SERIALIZATION_ERROR, }; @@ -623,11 +624,12 @@ impl PyAnnotatedLLMResponse { pub(crate) fn optimization_summary(&self, py: Python<'_>) -> PyResult> { match &self.inner.optimization_summary { Some(summary) => { - let value = serde_json::to_value(summary).map_err(|error| { - pyo3::exceptions::PyValueError::new_err(format!( - "optimization summary serialization error: {error}" - )) - })?; + let value = to_python_json_value( + summary, + "serialization error", + #[cfg(test)] + FORCE_ANNOTATED_RESPONSE_OPTIMIZATION_SUMMARY_SERIALIZATION_ERROR, + )?; json_to_py(py, &value) } None => Ok(py.None()), diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index 6fcd1dece..b713d630a 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -53,6 +53,8 @@ pub(crate) const FORCE_ANNOTATED_RESPONSE_TOOL_CALLS_SERIALIZATION_ERROR: u64 = pub(crate) const FORCE_ANNOTATED_RESPONSE_USAGE_SERIALIZATION_ERROR: u64 = 1 << 8; #[cfg(test)] pub(crate) const FORCE_ANNOTATED_RESPONSE_API_SPECIFIC_SERIALIZATION_ERROR: u64 = 1 << 9; +#[cfg(test)] +pub(crate) const FORCE_ANNOTATED_RESPONSE_OPTIMIZATION_SUMMARY_SERIALIZATION_ERROR: u64 = 1 << 10; #[cfg(test)] pub(crate) fn set_forced_serialization_mask_for_tests(mask: u64) { diff --git a/crates/python/tests/coverage/py_types_coverage_tests.rs b/crates/python/tests/coverage/py_types_coverage_tests.rs index 92cd15fc5..e8783fcb9 100644 --- a/crates/python/tests/coverage/py_types_coverage_tests.rs +++ b/crates/python/tests/coverage/py_types_coverage_tests.rs @@ -1530,7 +1530,17 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { api_name: "custom".into(), data: json!({"debug": true}), }), - optimization_summary: None, + optimization_summary: Some( + serde_json::from_value(json!({ + "schema_version": "1", + "calculation_version": "1", + "status": "partial", + "limitations": ["test"], + "tokens_saved": {}, + "contributions": [] + })) + .unwrap(), + ), extra: serde_json::Map::new(), }, }; @@ -1589,6 +1599,11 @@ fn test_forced_serialization_error_hooks_cover_unreachable_wrappers() { "forced serialization failure", |py, _, _, response| response.usage(py).map(|_| ()), ), + ( + FORCE_ANNOTATED_RESPONSE_OPTIMIZATION_SUMMARY_SERIALIZATION_ERROR, + "forced serialization failure", + |py, _, _, response| response.optimization_summary(py).map(|_| ()), + ), ( FORCE_ANNOTATED_RESPONSE_API_SPECIFIC_SERIALIZATION_ERROR, "forced serialization failure", diff --git a/crates/worker/tests/optimization_contract.rs b/crates/worker/tests/optimization_contract.rs index 0a39da40c..d7e808b80 100644 --- a/crates/worker/tests/optimization_contract.rs +++ b/crates/worker/tests/optimization_contract.rs @@ -4,10 +4,9 @@ //! Conformance tests for the worker SDK optimization contribution surface. use nemo_relay_worker::{ - DataSchema, LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, - LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationSummary, - LlmOptimizationSummaryStatus, LlmOptimizationTokenImpact, LlmOptimizationTokens, LlmRequest, - LlmRequestInterceptOutcome, + DataSchema, LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationModel, + LlmOptimizationModelTransition, LlmOptimizationSummary, LlmOptimizationSummaryStatus, + LlmOptimizationTokenImpact, LlmOptimizationTokens, LlmRequest, LlmRequestInterceptOutcome, }; use serde_json::{Value as Json, json}; @@ -20,8 +19,7 @@ fn worker_sdk_round_trips_the_canonical_contribution_contract() { let contribution: LlmOptimizationContribution = serde_json::from_value(fixture.clone()).unwrap(); - assert_ne!(contribution.kind, LlmOptimizationKind::input_compression()); - assert_ne!(contribution.kind, LlmOptimizationKind::model_routing()); + assert_eq!(contribution.kind.as_str(), "custom_energy_optimization"); assert_eq!( contribution.extra["future_top_level_field"], fixture["future_top_level_field"] diff --git a/go/nemo_relay/optimization.go b/go/nemo_relay/optimization.go index c36349248..bc162cc5d 100644 --- a/go/nemo_relay/optimization.go +++ b/go/nemo_relay/optimization.go @@ -4,6 +4,7 @@ package nemo_relay import ( + "bytes" "encoding/json" "errors" ) @@ -105,8 +106,16 @@ var llmOptimizationContributionKnownFields = [...]string{ "payload", } +func hasLLMOptimizationPayload(payload json.RawMessage) bool { + trimmed := bytes.TrimSpace(payload) + return len(trimmed) > 0 && !bytes.Equal(trimmed, []byte("null")) +} + // MarshalJSON preserves and flattens forward-compatible top-level fields. func (c LLMOptimizationContribution) MarshalJSON() ([]byte, error) { + if hasLLMOptimizationPayload(c.Payload) && c.PayloadSchema == nil { + return nil, errors.New("LLM optimization contribution payload requires payload_schema") + } known, err := json.Marshal(llmOptimizationContributionWire{ ID: c.ID, Sequence: c.Sequence, @@ -153,6 +162,9 @@ func (c *LLMOptimizationContribution) UnmarshalJSON(data []byte) error { if _, ok := extra["kind"]; !ok { return errors.New("LLM optimization contribution requires kind") } + if hasLLMOptimizationPayload(known.Payload) && known.PayloadSchema == nil { + return errors.New("LLM optimization contribution payload requires payload_schema") + } for _, name := range llmOptimizationContributionKnownFields { delete(extra, name) } diff --git a/go/nemo_relay/optimization_test.go b/go/nemo_relay/optimization_test.go index b89bcbdbb..d1052c273 100644 --- a/go/nemo_relay/optimization_test.go +++ b/go/nemo_relay/optimization_test.go @@ -6,6 +6,7 @@ package nemo_relay import ( "encoding/json" "os" + "strings" "testing" ) @@ -58,6 +59,32 @@ func TestLLMOptimizationContributionFixtureRoundTrip(t *testing.T) { assertSemanticJSONEqual(t, fixture, wire) } +func TestLLMOptimizationContributionRequiresPayloadSchema(t *testing.T) { + contribution := LLMOptimizationContribution{ + Producer: "test", + Kind: "custom", + Payload: json.RawMessage(`{"value":1}`), + } + if _, err := json.Marshal(contribution); err == nil || !strings.Contains(err.Error(), "payload_schema") { + t.Fatalf("expected marshal payload_schema error, got %v", err) + } + + var decoded LLMOptimizationContribution + if err := json.Unmarshal( + []byte(`{"producer":"test","kind":"custom","payload":{"value":1}}`), + &decoded, + ); err == nil || !strings.Contains(err.Error(), "payload_schema") { + t.Fatalf("expected unmarshal payload_schema error, got %v", err) + } + + if err := json.Unmarshal( + []byte(`{"producer":"test","kind":"custom","payload":null}`), + &decoded, + ); err != nil { + t.Fatalf("null payload should behave like an absent payload: %v", err) + } +} + func TestLLMRequestInterceptOptimizationContributionsRoundTrip(t *testing.T) { fixture, contribution := optimizationContributionFixture(t) const interceptName = "go_optimization_fixture" diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index 6f8d9b387..3880991fb 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -464,7 +464,7 @@ def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationContribution: transition = value.get("model_transition") impact = value.get("token_impact") schema = value.get("payload_schema") - applied = value.get("applied", False) + applied = value.get("applied", True) if not isinstance(applied, bool): raise WorkerSdkError("optimization contribution applied must be a boolean") contribution = cls( diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 3acfad427..94d45b2c9 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -107,6 +107,16 @@ def test_optimization_contribution_requires_schema_for_payload(): LlmOptimizationContribution.from_json({"producer": 1, "kind": "custom"}) +def test_optimization_contribution_omitted_applied_defaults_consistently(): + direct = LlmOptimizationContribution(producer="test", kind="custom") + decoded = LlmOptimizationContribution.from_json({"producer": "test", "kind": "custom"}) + + assert direct.applied is True + assert decoded.applied is True + assert direct.to_json()["applied"] is True + assert decoded.to_json()["applied"] is True + + class GrpcAbort(Exception): def __init__(self, code: object, details: str) -> None: super().__init__(f"{code}: {details}") From ded7a343f5f36dd597dd3e66b8711e13eb23f32c Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 02:00:23 -0600 Subject: [PATCH 09/10] fix(plugin-sdk): preserve optimization contribution compatibility Signed-off-by: Bryan Bednarski --- python/plugin/src/nemo_relay_plugin/_api.py | 22 ++++++--- python/tests/plugin/test_worker_sdk.py | 50 ++++++++++++++++++--- 2 files changed, 61 insertions(+), 11 deletions(-) diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index 3880991fb..cc1ebd781 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -391,12 +391,11 @@ def tokens(name: str) -> LlmOptimizationTokens | None: quality = value.get("quality") if quality is not None: + quality_string = _optimization_string(quality, "token_impact.quality") try: - quality = LlmOptimizationEvidenceQuality(_optimization_string(quality, "token_impact.quality")) - except ValueError as exc: - raise WorkerSdkError( - "optimization contribution token_impact.quality must be 'observed' or 'estimated'" - ) from exc + quality = LlmOptimizationEvidenceQuality(quality_string) + except ValueError: + quality = quality_string return cls( baseline=tokens("baseline"), effective=tokens("effective"), @@ -430,7 +429,18 @@ def to_json(self) -> dict[str, Json]: """Convert this contribution while flattening unknown top-level fields.""" if self.payload is not None and self.payload_schema is None: raise WorkerSdkError("optimization contribution payload requires payload_schema") - value = dict(self.extra) + known_fields = { + "id", + "sequence", + "producer", + "kind", + "applied", + "model_transition", + "token_impact", + "payload_schema", + "payload", + } + value = {name: item for name, item in self.extra.items() if name not in known_fields} value.update({"producer": self.producer, "kind": self.kind, "applied": self.applied}) if self.id is not None: value["id"] = self.id diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 94d45b2c9..80e914540 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -65,7 +65,8 @@ AUTH_TOKEN = "token" -def _optimization_contribution_fixture() -> Json: +@pytest.fixture(name="optimization_contribution_fixture") +def optimization_contribution_fixture_fixture() -> Json: fixture_path = ( Path(__file__).resolve().parents[3] / "crates" @@ -77,8 +78,8 @@ def _optimization_contribution_fixture() -> Json: return json.loads(fixture_path.read_text(encoding="utf-8")) -def test_optimization_contribution_fixture_round_trips_losslessly(): - fixture = _optimization_contribution_fixture() +def test_optimization_contribution_fixture_round_trips_losslessly(optimization_contribution_fixture: Json): + fixture = optimization_contribution_fixture contribution = LlmOptimizationContribution.from_json(fixture) outcome = LlmRequestInterceptOutcome( @@ -117,6 +118,43 @@ def test_optimization_contribution_omitted_applied_defaults_consistently(): assert decoded.to_json()["applied"] is True +def test_optimization_contribution_preserves_future_quality_strings(): + fixture = { + "producer": "test", + "kind": "custom", + "token_impact": {"quality": "provider_observed_v2"}, + } + + contribution = LlmOptimizationContribution.from_json(fixture) + + assert contribution.token_impact is not None + assert contribution.token_impact.quality == "provider_observed_v2" + assert contribution.to_json() == {**fixture, "applied": True} + + +def test_optimization_contribution_drops_known_fields_from_extra(): + contribution = LlmOptimizationContribution( + producer="test", + kind="custom", + extra={ + "id": "stale-id", + "sequence": 99, + "model_transition": {"baseline": {"model_id": "stale"}}, + "token_impact": {"quality": "stale"}, + "payload_schema": {"name": "stale", "version": "1"}, + "payload": {"stale": True}, + "future_field": "preserved", + }, + ) + + assert contribution.to_json() == { + "producer": "test", + "kind": "custom", + "applied": True, + "future_field": "preserved", + } + + class GrpcAbort(Exception): def __init__(self, code: object, details: str) -> None: super().__init__(f"{code}: {details}") @@ -1264,8 +1302,10 @@ def llm_request(name: str, request: Json, annotated: Json | None) -> LlmRequestI assert outcome["pending_marks"] == [] -async def test_llm_request_intercept_preserves_optimization_contribution_worker_envelope(): - fixture = _optimization_contribution_fixture() +async def test_llm_request_intercept_preserves_optimization_contribution_worker_envelope( + optimization_contribution_fixture: Json, +): + fixture = optimization_contribution_fixture class OptimizationPlugin(WorkerPlugin): plugin_id = "tests.optimization_contribution" From 232d66cfdac7fe6ee2f22262acd6b71db00f123f Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 9 Jul 2026 02:21:28 -0600 Subject: [PATCH 10/10] fix(plugin-sdk): align omitted contribution defaults Signed-off-by: Bryan Bednarski --- crates/types/tests/optimization_tests.rs | 15 +++++++++++++++ go/nemo_relay/optimization_test.go | 13 +++++++++++++ python/plugin/src/nemo_relay_plugin/_api.py | 4 ++-- python/tests/plugin/test_worker_sdk.py | 12 ++++++------ 4 files changed, 36 insertions(+), 8 deletions(-) diff --git a/crates/types/tests/optimization_tests.rs b/crates/types/tests/optimization_tests.rs index fd68e6478..3f4a9f174 100644 --- a/crates/types/tests/optimization_tests.rs +++ b/crates/types/tests/optimization_tests.rs @@ -67,6 +67,21 @@ fn saved_prompt_tokens_remain_explicit_on_the_wire() { assert_eq!(wire["saved"]["total_tokens"], 42); } +#[test] +fn omitted_applied_is_non_applied() { + let contribution: LlmOptimizationContribution = serde_json::from_value(json!({ + "producer": "example", + "kind": "custom" + })) + .unwrap(); + + assert!(!contribution.applied); + assert_eq!( + serde_json::to_value(contribution).unwrap()["applied"], + false + ); +} + #[test] fn canonical_all_fields_fixture_is_lossless_and_open() { let fixture: Value = serde_json::from_str(include_str!( diff --git a/go/nemo_relay/optimization_test.go b/go/nemo_relay/optimization_test.go index d1052c273..60712a8da 100644 --- a/go/nemo_relay/optimization_test.go +++ b/go/nemo_relay/optimization_test.go @@ -85,6 +85,19 @@ func TestLLMOptimizationContributionRequiresPayloadSchema(t *testing.T) { } } +func TestLLMOptimizationContributionOmittedAppliedIsNonApplied(t *testing.T) { + var contribution LLMOptimizationContribution + if err := json.Unmarshal( + []byte(`{"producer":"test","kind":"custom"}`), + &contribution, + ); err != nil { + t.Fatalf("decode optimization contribution: %v", err) + } + if contribution.Applied { + t.Fatal("omitted applied must decode as false") + } +} + func TestLLMRequestInterceptOptimizationContributionsRoundTrip(t *testing.T) { fixture, contribution := optimizationContributionFixture(t) const interceptName = "go_optimization_fixture" diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index cc1ebd781..99d89b390 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -413,7 +413,7 @@ class LlmOptimizationContribution: producer: str kind: str - applied: bool = True + applied: bool = False id: str | None = None sequence: int | None = None model_transition: LlmOptimizationModelTransition | None = None @@ -474,7 +474,7 @@ def from_json(cls, value: Mapping[str, Json]) -> LlmOptimizationContribution: transition = value.get("model_transition") impact = value.get("token_impact") schema = value.get("payload_schema") - applied = value.get("applied", True) + applied = value.get("applied", False) if not isinstance(applied, bool): raise WorkerSdkError("optimization contribution applied must be a boolean") contribution = cls( diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 80e914540..bdb70028a 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -112,10 +112,10 @@ def test_optimization_contribution_omitted_applied_defaults_consistently(): direct = LlmOptimizationContribution(producer="test", kind="custom") decoded = LlmOptimizationContribution.from_json({"producer": "test", "kind": "custom"}) - assert direct.applied is True - assert decoded.applied is True - assert direct.to_json()["applied"] is True - assert decoded.to_json()["applied"] is True + assert direct.applied is False + assert decoded.applied is False + assert direct.to_json()["applied"] is False + assert decoded.to_json()["applied"] is False def test_optimization_contribution_preserves_future_quality_strings(): @@ -129,7 +129,7 @@ def test_optimization_contribution_preserves_future_quality_strings(): assert contribution.token_impact is not None assert contribution.token_impact.quality == "provider_observed_v2" - assert contribution.to_json() == {**fixture, "applied": True} + assert contribution.to_json() == {**fixture, "applied": False} def test_optimization_contribution_drops_known_fields_from_extra(): @@ -150,7 +150,7 @@ def test_optimization_contribution_drops_known_fields_from_extra(): assert contribution.to_json() == { "producer": "test", "kind": "custom", - "applied": True, + "applied": False, "future_field": "preserved", }