feat!: pass codec context to LLM execution intercepts

crates/adaptive/src/acg_component.rs

@@ -666,7 +666,7 @@ pub(crate) fn create_acg_llm_execution_intercept(
     plugin: Arc<dyn ProviderPlugin>,
 ) -> LlmExecutionFn {
     Arc::new(
-        move |_name: &str, request: LlmRequest, next: LlmExecutionNextFn| {
+        move |_name: &str, request: LlmRequest, _context, next: LlmExecutionNextFn| {
             let cache = hot_cache.clone();
             let agent_id = agent_id.clone();
             let provider = provider.clone();
@@ -689,7 +689,7 @@ pub(crate) fn create_acg_llm_stream_execution_intercept(
     plugin: Arc<dyn ProviderPlugin>,
 ) -> LlmStreamExecutionFn {
     Arc::new(
-        move |_name: &str, request: LlmRequest, next: LlmStreamExecutionNextFn| {
+        move |_name: &str, request: LlmRequest, _context, next: LlmStreamExecutionNextFn| {
             let cache = hot_cache.clone();
             let agent_id = agent_id.clone();
             let provider = provider.clone();

crates/adaptive/src/response_cache/intercept.rs

@@ -128,7 +128,7 @@ pub(crate) fn make_intercept(
     config: Arc<ResponseCacheConfig>,
 ) -> LlmExecutionFn {
     Arc::new(
-        move |provider: &str, request: LlmRequest, next: LlmExecutionNextFn| {
+        move |provider: &str, request: LlmRequest, _context, next: LlmExecutionNextFn| {
             let store = Arc::clone(&store);
             let config = Arc::clone(&config);
             let provider = provider.to_string();
@@ -150,7 +150,7 @@ pub(crate) fn make_stream_intercept(
     config: Arc<ResponseCacheConfig>,
 ) -> LlmStreamExecutionFn {
     Arc::new(
-        move |provider: &str, request: LlmRequest, next: LlmStreamExecutionNextFn| {
+        move |provider: &str, request: LlmRequest, _context, next: LlmStreamExecutionNextFn| {
             let store = Arc::clone(&store);
             let config = Arc::clone(&config);
             let provider = provider.to_string();

crates/adaptive/src/runtime/features.rs

@@ -806,7 +806,7 @@ impl AdaptiveFeature for AcgFeature {
             ctx.register_llm_execution_intercept(
                 &self.execution_name,
                 self.priority,
-                Arc::new(move |name, request, next| {
+                Arc::new(move |name, request, context, next| {
                     let execution_intercept = execution_intercept.clone();
                     let bound_scopes = bound_scopes.clone();
                     let name = name.to_string();
@@ -818,7 +818,7 @@ impl AdaptiveFeature for AcgFeature {
                         if has_bound_scopes {
                             return next(request).await;
                         }
-                        execution_intercept(&name, request, next).await
+                        execution_intercept(&name, request, context, next).await
                     })
                 }),
             )?;
@@ -832,7 +832,7 @@ impl AdaptiveFeature for AcgFeature {
             ctx.register_llm_stream_execution_intercept(
                 &self.stream_name,
                 self.priority,
-                Arc::new(move |name, request, next| {
+                Arc::new(move |name, request, context, next| {
                     let stream_intercept = stream_intercept.clone();
                     let bound_scopes = bound_scopes.clone();
                     let name = name.to_string();
@@ -844,7 +844,7 @@ impl AdaptiveFeature for AcgFeature {
                         if has_bound_scopes {
                             return next(request).await;
                         }
-                        stream_intercept(&name, request, next).await
+                        stream_intercept(&name, request, context, next).await
                     })
                 }),
             )

crates/adaptive/tests/integration/runtime_integration_tests.rs

@@ -753,7 +753,7 @@ impl Plugin for HeaderPlugin {
             ctx.register_llm_execution_intercept(
                 "llm_exec_plugin",
                 priority,
-                Arc::new(|_name, request, next| {
+                Arc::new(|_name, request, _context, next| {
                     Box::pin(async move {
                         let mut response = next(request).await?;
                         if let Json::Object(ref mut map) = response {
@@ -766,7 +766,7 @@ impl Plugin for HeaderPlugin {
             ctx.register_llm_stream_execution_intercept(
                 "llm_stream_exec_plugin",
                 priority,
-                Arc::new(|_name, request, next| {
+                Arc::new(|_name, request, _context, next| {
                     Box::pin(async move {
                         let mut stream = next(request).await?;
                         let mut chunks = Vec::new();

crates/adaptive/tests/unit/acg_component_tests.rs

@@ -17,7 +17,9 @@ use crate::config::AcgComponentConfig;
 use crate::storage::memory::InMemoryBackend;
 use crate::storage::traits::StorageBackendDyn;
 use nemo_relay::api::llm::LlmRequest;
-use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn};
+use nemo_relay::api::runtime::{
+    LlmExecutionContext, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn,
+};
 use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent};
 use serde_json::{Value, json};
 use tokio_stream::StreamExt;
@@ -698,7 +700,7 @@ async fn acg_component_stream_execution_intercept_rewrites_streaming_requests()
         })
     });
 
-    let mut stream = intercept("anthropic", request, next)
+    let mut stream = intercept("anthropic", request, LlmExecutionContext::default(), next)
         .await
         .expect("stream intercept should succeed");
     let first = stream
@@ -1302,7 +1304,7 @@ async fn acg_component_execution_intercept_rewrites_non_streaming_requests() {
     );
     let next: LlmExecutionNextFn = Arc::new(|req| Box::pin(async move { Ok(req.content) }));
 
-    let result = intercept("anthropic", request, next)
+    let result = intercept("anthropic", request, LlmExecutionContext::default(), next)
         .await
         .expect("execution intercept should succeed");
 
@@ -1515,9 +1517,14 @@ async fn acg_component_execution_intercept_passes_original_request_when_translat
     );
     let next: LlmExecutionNextFn = Arc::new(|req| Box::pin(async move { Ok(req.content) }));
 
-    let result = intercept("anthropic", invalid_request.clone(), next)
-        .await
-        .expect("execution intercept should succeed");
+    let result = intercept(
+        "anthropic",
+        invalid_request.clone(),
+        LlmExecutionContext::default(),
+        next,
+    )
+    .await
+    .expect("execution intercept should succeed");
 
     assert_eq!(result, invalid_request.content);
 }
@@ -1544,9 +1551,14 @@ async fn acg_component_stream_execution_intercept_passes_original_request_when_t
         })
     });
 
-    let mut stream = intercept("anthropic", invalid_request.clone(), next)
-        .await
-        .expect("stream intercept should pass through");
+    let mut stream = intercept(
+        "anthropic",
+        invalid_request.clone(),
+        LlmExecutionContext::default(),
+        next,
+    )
+    .await
+    .expect("stream intercept should pass through");
     let first = stream
         .next()
         .await

crates/adaptive/tests/unit/runtime_features_tests.rs

@@ -211,7 +211,7 @@ fn assert_llm_execution_intercept_registered(name: &str) {
         register_llm_execution_intercept(
             name,
             i32::MAX,
-            Arc::new(|_name, request, next| next(request)),
+            Arc::new(|_name, request, _context, next| next(request)),
         ),
         name,
     );
@@ -221,7 +221,7 @@ fn assert_llm_execution_intercept_absent(name: &str) {
     register_llm_execution_intercept(
         name,
         i32::MAX,
-        Arc::new(|_name, request, next| next(request)),
+        Arc::new(|_name, request, _context, next| next(request)),
     )
     .unwrap();
     deregister_llm_execution_intercept(name).unwrap();
@@ -232,7 +232,7 @@ fn assert_llm_stream_execution_intercept_registered(name: &str) {
         register_llm_stream_execution_intercept(
             name,
             i32::MAX,
-            Arc::new(|_name, request, next| next(request)),
+            Arc::new(|_name, request, _context, next| next(request)),
         ),
         name,
     );
@@ -242,7 +242,7 @@ fn assert_llm_stream_execution_intercept_absent(name: &str) {
     register_llm_stream_execution_intercept(
         name,
         i32::MAX,
-        Arc::new(|_name, request, next| next(request)),
+        Arc::new(|_name, request, _context, next| next(request)),
     )
     .unwrap();
     deregister_llm_stream_execution_intercept(name).unwrap();
@@ -807,13 +807,13 @@ async fn registration_context_registers_all_supported_callback_types() {
     ctx.register_llm_execution_intercept(
         "adaptive_test_execution",
         6,
-        Arc::new(|_name, request, _next| Box::pin(async move { Ok(request.content) })),
+        Arc::new(|_name, request, _context, _next| Box::pin(async move { Ok(request.content) })),
     )
     .unwrap();
     ctx.register_llm_stream_execution_intercept(
         "adaptive_test_stream",
         7,
-        Arc::new(|_name, request, _next| {
+        Arc::new(|_name, request, _context, _next| {
             Box::pin(async move {
                 Ok(LlmJsonStream::new(tokio_stream::iter(vec![Ok(
                     request.content
@@ -1017,7 +1017,7 @@ async fn acg_feature_reports_execution_registration_conflicts() {
     register_llm_execution_intercept(
         &execution_name,
         1,
-        Arc::new(|_name, request, next| next(request)),
+        Arc::new(|_name, request, _context, next| next(request)),
     )
     .unwrap();
 
@@ -1146,8 +1146,12 @@ async fn response_cache_feature_cleans_up_when_llm_registration_conflicts() {
         Uuid::now_v7(),
     );
     let name = feature.name.clone();
-    register_llm_execution_intercept(&name, 1, Arc::new(|_name, request, next| next(request)))
-        .unwrap();
+    register_llm_execution_intercept(
+        &name,
+        1,
+        Arc::new(|_name, request, _context, next| next(request)),
+    )
+    .unwrap();
 
     let error = {
         let mut ctx = RegistrationContext::new(&mut runtime);
@@ -1180,7 +1184,7 @@ async fn response_cache_feature_cleans_up_when_stream_registration_conflicts() {
     register_llm_stream_execution_intercept(
         &stream_name,
         1,
-        Arc::new(|_name, request, next| next(request)),
+        Arc::new(|_name, request, _context, next| next(request)),
     )
     .unwrap();
 

crates/cli/tests/coverage/daemon/worker_managed_tests.rs

@@ -1479,7 +1479,7 @@ async fn managed_runtime_rejects_response_mutating_execution_middleware() {
     register_llm_execution_intercept(
         INTERCEPT,
         1,
-        Arc::new(|_name, _request, _next| Box::pin(async { Ok(json!({})) })),
+        Arc::new(|_name, _request, _context, _next| Box::pin(async { Ok(json!({})) })),
     )
     .expect("register execution middleware");
 

crates/cli/tests/coverage/shared/server_tests.rs

@@ -5084,7 +5084,7 @@ async fn gateway_concurrent_next_uses_canonical_selected_buffered_response() {
     register_llm_execution_intercept(
         INTERCEPT_NAME,
         1,
-        Arc::new(move |_name, request, next| {
+        Arc::new(move |_name, request, _context, next| {
             let release_losing = release_from_intercept.clone();
             Box::pin(async move {
                 if request.content["messages"][0]["content"].as_str() != Some(MARKER) {
@@ -5172,7 +5172,7 @@ async fn gateway_concurrent_next_uses_canonical_selected_streaming_response() {
     register_llm_stream_execution_intercept(
         INTERCEPT_NAME,
         1,
-        Arc::new(move |_name, request, next| {
+        Arc::new(move |_name, request, _context, next| {
             let release_losing = release_from_intercept.clone();
             Box::pin(async move {
                 if request.content["messages"][0]["content"].as_str() != Some(MARKER) {
@@ -5255,7 +5255,7 @@ async fn gateway_concurrent_next_relays_selected_buffered_failure() {
     register_llm_execution_intercept(
         INTERCEPT_NAME,
         1,
-        Arc::new(move |_name, request, next| {
+        Arc::new(move |_name, request, _context, next| {
             let release_losing = release_from_intercept.clone();
             Box::pin(async move {
                 if request.content["messages"][0]["content"].as_str() != Some(MARKER) {
@@ -5334,7 +5334,7 @@ async fn gateway_concurrent_next_relays_selected_streaming_failure() {
     register_llm_stream_execution_intercept(
         INTERCEPT_NAME,
         1,
-        Arc::new(move |_name, request, next| {
+        Arc::new(move |_name, request, _context, next| {
             let release_losing = release_from_intercept.clone();
             Box::pin(async move {
                 if request.content["messages"][0]["content"].as_str() != Some(MARKER) {
@@ -5634,7 +5634,7 @@ async fn gateway_surfaces_post_upstream_intercept_rejection_instead_of_relaying_
     register_llm_execution_intercept(
         INTERCEPT_NAME,
         1,
-        Arc::new(|_name, request, next| {
+        Arc::new(|_name, request, _context, next| {
             Box::pin(async move {
                 let marked = request.content["messages"][0]["content"].as_str() == Some(MARKER);
                 let response = next(request).await?;
@@ -6285,7 +6285,7 @@ async fn model_call_policy_still_applies_on_a_named_upstream() {
     register_llm_execution_intercept(
         INTERCEPT,
         1,
-        Arc::new(move |_name, request, next| {
+        Arc::new(move |_name, request, _context, next| {
             Box::pin(async move {
                 if request.content["messages"][0]["content"].as_str() == Some("blocked by policy") {
                     return Err(nemo_relay::error::FlowError::GuardrailRejected(

crates/core/src/api/llm.rs

@@ -29,8 +29,8 @@ use crate::api::runtime::subscriber_dispatcher::{
     dispatch_sanitized_event, dispatch_transformed_event, register_pending_publication,
 };
 use crate::api::runtime::{
-    EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream,
-    LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn,
+    EventSubscriberFn, LlmCollectorFn, LlmExecutionContext, LlmExecutionNextFn, LlmFinalizerFn,
+    LlmJsonStream, LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn,
     MiddlewareContinuationContext,
 };
 use crate::api::runtime::{ScopeStackHandle, capture_trace_context, current_scope_stack};
@@ -1739,6 +1739,8 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
     );
     let execution_name = name.clone();
     let event_uuid = handle.uuid;
+    let execution_context =
+        LlmExecutionContext::for_non_streaming(request_codec, response_codec.clone());
     let execution = with_active_event_trace_context(
         event_uuid,
         Some(active_trace_context),
@@ -1757,7 +1759,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
                     .read()
                     .map_err(|error| FlowError::Internal(error.to_string()))?
                     .registry_snapshot(&[RuntimeRegistrationKind::LlmExecutionIntercept]);
-                state.llm_build_execution_chain(&execution_name, func, &scope_local_refs)
+                state.llm_build_execution_chain(
+                    &execution_name,
+                    func,
+                    &scope_local_refs,
+                    execution_context,
+                )
             };
             execution(intercepted_request).await
         }),
@@ -1966,6 +1973,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu
     let execution_name = name.clone();
     let event_uuid = handle.uuid;
     let stream_started_at = Instant::now();
+    let execution_context = LlmExecutionContext::for_streaming(request_codec);
     let execution = with_active_event_trace_context(
         event_uuid,
         Some(active_trace_context),
@@ -1984,12 +1992,17 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu
                     .read()
                     .map_err(|error| FlowError::Internal(error.to_string()))?
                     .registry_snapshot(&[RuntimeRegistrationKind::LlmStreamExecutionIntercept]);
-                state.llm_stream_build_execution_chain(&execution_name, func, &scope_local_refs)
+                state.llm_stream_build_execution_chain(
+                    &execution_name,
+                    func,
+                    &scope_local_refs,
+                    execution_context,
+                )
             };
-            let execution_context = MiddlewareContinuationContext::capture();
+            let continuation_context = MiddlewareContinuationContext::capture();
             execution(intercepted_request)
                 .await
-                .map(|stream| contextualize_stream(stream, execution_context))
+                .map(|stream| contextualize_stream(stream, continuation_context))
         }),
     )
     .await;

crates/core/src/api/runtime.rs

@@ -6,24 +6,26 @@
 pub mod callbacks;
 mod continuation_context;
 pub mod global;
+mod llm_execution_context;
 pub mod scope_stack;
 pub mod state;
 pub mod subscriber_dispatcher;
 
 pub use callbacks::{
     BuiltinLlmCodec, ConditionalMiddlewareGuardrailFn, EventMetadataInjectorFn, EventSanitizeFn,
     EventSubscriberFn, LlmCodecIdentity, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn,
-    LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn,
-    LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext,
-    LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamInner,
-    ToolConditionalFn, ToolExecutionContext, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn,
-    ToolSanitizeFn,
+    LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestContext, LlmRequestInterceptFn,
+    LlmResponseContext, LlmSanitizeRequestContext, LlmSanitizeRequestFn,
+    LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn,
+    LlmStreamExecutionNextFn, LlmStreamInner, ToolConditionalFn, ToolExecutionContext,
+    ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn,
 };
 #[doc(hidden)]
 pub use continuation_context::MiddlewareContinuationContext;
 #[cfg(test)]
 pub(crate) use continuation_context::MiddlewareContinuationLease;
 pub use global::global_context;
+pub use llm_execution_context::LlmExecutionContext;
 pub(crate) use scope_stack::capture_trace_context;
 pub use scope_stack::{
     PropagationContext, ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding,