feat(plugin): add targeted LLM continuations to native API v2

crates/core/Cargo.toml

@@ -63,6 +63,7 @@ tokio = { version = "1", default-features = false, features = ["rt", "rt-multi-t
 tokio-stream = { version = "0.1", default-features = false, features = ["sync"] }
 typed-builder = "0.23.2"
 futures-util = "0.3"
+async-stream = "0.3"
 opentelemetry = { workspace = true, features = ["trace"] }
 opentelemetry-semantic-conventions.workspace = true
 opentelemetry_sdk = { workspace = true, features = ["trace", "internal-logs"] }

crates/core/src/api/llm.rs

@@ -28,7 +28,8 @@ use crate::api::runtime::subscriber_dispatcher::{
 use crate::api::runtime::{
     EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream,
     LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn,
-    MiddlewareContinuationContext, with_active_event_uuid,
+    MiddlewareContinuationContext, targeted_llm_execution, targeted_llm_stream_execution,
+    with_active_event_uuid,
 };
 use crate::api::runtime::{ScopeStackHandle, current_scope_stack};
 use crate::api::scope::event;
@@ -1494,7 +1495,11 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
                 let state = context
                     .read()
                     .map_err(|error| FlowError::Internal(error.to_string()))?;
-                state.llm_build_execution_chain(&execution_name, func, &scope_locals)
+                state.llm_build_execution_chain(
+                    &execution_name,
+                    targeted_llm_execution(func),
+                    &scope_locals,
+                )
             };
             execution(intercepted_request).await
         }),
@@ -1703,7 +1708,11 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu
                 let state = context
                     .read()
                     .map_err(|error| FlowError::Internal(error.to_string()))?;
-                state.llm_stream_build_execution_chain(&execution_name, func, &scope_locals)
+                state.llm_stream_build_execution_chain(
+                    &execution_name,
+                    targeted_llm_stream_execution(func),
+                    &scope_locals,
+                )
             };
             let execution_context = MiddlewareContinuationContext::capture();
             execution(intercepted_request)

crates/core/src/api/runtime.rs

@@ -6,6 +6,7 @@
 pub mod callbacks;
 mod continuation_context;
 pub mod global;
+mod llm_dispatch_context;
 pub mod scope_stack;
 pub mod state;
 pub mod subscriber_dispatcher;
@@ -23,6 +24,11 @@ pub use continuation_context::MiddlewareContinuationContext;
 #[cfg(test)]
 pub(crate) use continuation_context::MiddlewareContinuationLease;
 pub use global::global_context;
+#[cfg(test)]
+pub(crate) use llm_dispatch_context::current_llm_dispatch_target;
+pub(crate) use llm_dispatch_context::{
+    LlmDispatchTargetContext, targeted_llm_execution, targeted_llm_stream_execution,
+};
 pub use scope_stack::{
     PropagationContext, ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding,
     capture_propagation_context, capture_propagation_context_with_root, capture_thread_scope_stack,

crates/core/src/api/runtime/continuation_context.rs

@@ -8,6 +8,9 @@ use std::future::Future;
 use crate::api::optimization::{
     LlmOptimizationRecorder, current_llm_optimization_recorder, scope_llm_optimization_recorder,
 };
+use crate::api::runtime::llm_dispatch_context::{
+    LlmDispatchTargetContext, current_llm_dispatch_target, scope_llm_dispatch_target,
+};
 use crate::api::runtime::scope_stack::{
     ScopeStackHandle, TASK_SCOPE_STACK, active_event_uuid, current_context_scope_stack,
     current_scope_stack, scope_stack_active, snapshot_scope_stack, with_active_event_uuid,
@@ -31,6 +34,7 @@ pub struct MiddlewareContinuationContext {
     publication_context: Option<PublicationContext>,
     publication_buffer: Option<PublicationBuffer>,
     optimization_recorder: Option<LlmOptimizationRecorder>,
+    llm_dispatch_target: Option<LlmDispatchTargetContext>,
 }
 
 impl MiddlewareContinuationContext {
@@ -44,6 +48,7 @@ impl MiddlewareContinuationContext {
             publication_context: capture_publication_context(),
             publication_buffer: capture_nested_publication_buffer(),
             optimization_recorder: current_llm_optimization_recorder(),
+            llm_dispatch_target: current_llm_dispatch_target(),
         }
     }
 
@@ -69,6 +74,7 @@ impl MiddlewareContinuationContext {
             publication_context: self.publication_context.clone(),
             publication_buffer: self.publication_buffer.clone(),
             optimization_recorder: self.optimization_recorder.clone(),
+            llm_dispatch_target: self.llm_dispatch_target.clone(),
         })
     }
 
@@ -97,12 +103,36 @@ impl MiddlewareContinuationContext {
                 None => published.await,
             }
         };
-        match &self.optimization_recorder {
-            Some(recorder) => scope_llm_optimization_recorder(recorder.clone(), active).await,
-            None => active.await,
+        let optimized = async {
+            match &self.optimization_recorder {
+                Some(recorder) => scope_llm_optimization_recorder(recorder.clone(), active).await,
+                None => active.await,
+            }
+        };
+        match &self.llm_dispatch_target {
+            Some(target) => {
+                scope_llm_dispatch_target(self.active_event_uuid, target.clone(), optimized).await
+            }
+            None => optimized.await,
         }
     }
 
+    /// Invoke a callback and poll its future with the captured Relay context and typed LLM target.
+    #[doc(hidden)]
+    pub(crate) async fn invoke_with_llm_dispatch_target<C, F>(
+        &self,
+        target: LlmDispatchTargetContext,
+        callback: C,
+    ) -> F::Output
+    where
+        C: FnOnce() -> F,
+        F: Future,
+    {
+        let mut context = self.clone();
+        context.llm_dispatch_target = Some(target);
+        context.invoke(callback).await
+    }
+
     /// Invoke a callback and poll its future with the captured Relay context.
     ///
     /// The callback itself can inspect Relay task state before constructing its