fix: resolve sanitized LLM request from annotations

crates/core/src/api/llm.rs

@@ -276,6 +276,7 @@ fn emit_llm_start(
     handle: &LlmHandle,
     request: &LlmRequest,
     annotated_request: Option<Arc<AnnotatedLlmRequest>>,
+    request_codec: Option<&dyn LlmCodec>,
 ) -> Result<()> {
     ensure_runtime_owner()?;
     let (event, subscribers) = {
@@ -292,6 +293,15 @@ fn emit_llm_start(
             .map_err(|error| FlowError::Internal(error.to_string()))?;
 
         let sanitized_request = state.llm_sanitize_request_chain(request.clone(), &scope_locals);
+        let annotated_request = match request_codec {
+            Some(codec)
+                if sanitized_request.headers != request.headers
+                    || sanitized_request.content != request.content =>
+            {
+                codec.decode(&sanitized_request).ok().map(Arc::new)
+            }
+            _ => annotated_request,
+        };
         let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null);
         let event = state.build_llm_start_event(handle, Some(input), annotated_request);
         (event, subscribers)
@@ -341,10 +351,16 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result<LlmHandle> {
         .timestamp_opt(params.timestamp)
         .build();
     let handle = create_llm_handle(handle_params)?;
-    emit_llm_start(&handle, params.request, params.annotated_request)?;
+    emit_llm_start(&handle, params.request, params.annotated_request, None)?;
     Ok(handle)
 }
 
+#[derive(Clone, Copy)]
+struct LlmCallEndBehavior {
+    response_codec_errors_fatal: bool,
+    attach_estimated_cost: bool,
+}
+
 /// Finish a manual LLM lifecycle span.
 ///
 /// This emits an LLM-end event for a handle previously returned by
@@ -377,6 +393,19 @@ pub fn llm_call(params: LlmCallParams<'_>) -> Result<LlmHandle> {
 /// Sanitize-response guardrails affect only the emitted end-event payload, not
 /// the caller-owned `response` value.
 pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> {
+    llm_call_end_with_behavior(
+        params,
+        LlmCallEndBehavior {
+            response_codec_errors_fatal: true,
+            attach_estimated_cost: false,
+        },
+    )
+}
+
+fn llm_call_end_with_behavior(
+    params: LlmCallEndParams<'_>,
+    behavior: LlmCallEndBehavior,
+) -> Result<()> {
     let LlmCallEndParams {
         handle,
         response,
@@ -411,7 +440,12 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> {
             Some(annotated_response) => Some(annotated_response),
             None => match (response_codec.as_ref(), data.as_ref()) {
                 (Some(codec), Some(response)) => match codec.decode_response(response) {
-                    Ok(decoded) => Some(Arc::new(decoded)),
+                    Ok(mut decoded) => {
+                        if behavior.attach_estimated_cost {
+                            attach_estimated_cost_for_provider(&mut decoded, Some(&handle.name));
+                        }
+                        Some(Arc::new(decoded))
+                    }
                     Err(error) => {
                         decode_error = Some(error);
                         None
@@ -432,7 +466,9 @@ pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> {
         (event, subscribers)
     };
     NemoRelayContextState::emit_event(&event, &subscribers);
-    if let Some(error) = decode_error {
+    if let Some(error) = decode_error
+        && behavior.response_codec_errors_fatal
+    {
         Err(error)
     } else {
         Ok(())
@@ -554,6 +590,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
         }
     }
 
+    let request_codec = codec.clone();
     let (intercepted_request, annotated_request) =
         run_request_intercepts_with_codec(&name, request, codec)?;
 
@@ -567,7 +604,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
             .model_name_opt(model_name)
             .build(),
     )?;
-    emit_llm_start(&handle, &intercepted_request, annotated_request.clone())?;
+    emit_llm_start(
+        &handle,
+        &intercepted_request,
+        annotated_request.clone(),
+        request_codec.as_deref(),
+    )?;
 
     let execution = {
         let scope_stack = current_scope_stack();
@@ -583,22 +625,18 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
 
     match execution(intercepted_request).await {
         Ok(response) => {
-            let annotated_response = response_codec
-                .as_ref()
-                .and_then(|codec| {
-                    let mut decoded = codec.decode_response(&response).ok()?;
-                    attach_estimated_cost_for_provider(&mut decoded, Some(&name));
-                    Some(decoded)
-                })
-                .map(Arc::new);
-            llm_call_end(
+            llm_call_end_with_behavior(
                 LlmCallEndParams::builder()
                     .handle(&handle)
                     .response(response.clone())
                     .data_opt(data)
                     .metadata_opt(metadata)
-                    .annotated_response_opt(annotated_response)
+                    .response_codec_opt(response_codec)
                     .build(),
+                LlmCallEndBehavior {
+                    response_codec_errors_fatal: false,
+                    attach_estimated_cost: true,
+                },
             )?;
             Ok(response)
         }
@@ -706,6 +744,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu
         }
     }
 
+    let request_codec = codec.clone();
     let (intercepted_request, annotated_request) =
         run_request_intercepts_with_codec(&name, request, codec)?;
 
@@ -719,7 +758,12 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu
             .model_name_opt(model_name)
             .build(),
     )?;
-    emit_llm_start(&handle, &intercepted_request, annotated_request)?;
+    emit_llm_start(
+        &handle,
+        &intercepted_request,
+        annotated_request,
+        request_codec.as_deref(),
+    )?;
 
     let execution = {
         let scope_stack = current_scope_stack();

crates/core/src/stream.rs

@@ -140,17 +140,6 @@ impl LlmStreamWrapper {
             None => Json::Null,
         };
 
-        // Decode aggregated response if response codec is present (non-fatal)
-        let annotated_response: Option<Arc<AnnotatedLlmResponse>> = self
-            .response_codec
-            .as_ref()
-            .and_then(|c| {
-                let mut decoded = c.decode_response(&aggregated).ok()?;
-                attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name));
-                Some(decoded)
-            })
-            .map(Arc::new);
-
         let event_snapshot = {
             let ss_guard = self.scope_stack.read().expect("scope stack lock poisoned");
             let sl =
@@ -167,6 +156,18 @@ impl LlmStreamWrapper {
                     } else {
                         Some(sanitized)
                     };
+                    let annotated_response: Option<Arc<AnnotatedLlmResponse>> = 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 event = state.end_llm_handle(
                         &self.handle,
                         data,