lash/
turn.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use crate::support::*;
5use futures_util::Stream;
6
7pub use lash_core::{AssistantOutput, TurnIssue};
8
9/// The two internal event sinks threaded through the turn-execution helpers.
10///
11/// `events` is the raw lower-level runtime event stream, reachable from app
12/// code only via [`TurnBuilder::advanced`]; `turn_events` is the semantic
13/// [`TurnActivity`] stream used by the primary builder API.
14/// Bundling them keeps the internal turn fns to a single sink parameter.
15#[derive(Clone, Copy, Default)]
16pub(crate) struct TurnSinks<'a> {
17    events: Option<&'a dyn EventSink>,
18    turn_events: Option<&'a dyn TurnActivitySink>,
19}
20
21impl<'a> TurnSinks<'a> {
22    pub(crate) fn turn(events: &'a dyn TurnActivitySink) -> Self {
23        Self {
24            events: None,
25            turn_events: Some(events),
26        }
27    }
28
29    pub(crate) fn session(events: &'a dyn EventSink) -> Self {
30        Self {
31            events: Some(events),
32            turn_events: None,
33        }
34    }
35
36    fn events(&self) -> Option<&'a dyn EventSink> {
37        self.events
38    }
39
40    fn turn_events(&self) -> Option<&'a dyn TurnActivitySink> {
41        self.turn_events
42    }
43}
44
45/// Cancellation tokens of the turns currently executing through one opened
46/// [`LashSession`](crate::LashSession) (shared by its clones).
47/// [`LashSession::cancel_running_turns`](crate::LashSession::cancel_running_turns)
48/// cancels them without the caller having to thread a token around.
49#[derive(Clone, Default)]
50pub(crate) struct TurnCancelRegistry {
51    inner: Arc<StdMutex<TurnCancelRegistryInner>>,
52}
53
54#[derive(Default)]
55struct TurnCancelRegistryInner {
56    next_id: u64,
57    active: BTreeMap<u64, CancellationToken>,
58}
59
60impl TurnCancelRegistry {
61    /// Register the token of a turn that is about to execute. The guard
62    /// removes the entry when the turn finishes, however it finishes.
63    fn register(&self, token: CancellationToken) -> TurnCancelGuard {
64        let mut inner = self.inner.lock().expect("turn cancel registry");
65        let id = inner.next_id;
66        inner.next_id += 1;
67        inner.active.insert(id, token);
68        TurnCancelGuard {
69            registry: Arc::clone(&self.inner),
70            id,
71        }
72    }
73
74    pub(crate) fn cancel_all(&self) -> usize {
75        let inner = self.inner.lock().expect("turn cancel registry");
76        for token in inner.active.values() {
77            token.cancel();
78        }
79        inner.active.len()
80    }
81}
82
83pub(crate) struct TurnCancelGuard {
84    registry: Arc<StdMutex<TurnCancelRegistryInner>>,
85    id: u64,
86}
87
88impl Drop for TurnCancelGuard {
89    fn drop(&mut self) {
90        self.registry
91            .lock()
92            .expect("turn cancel registry")
93            .active
94            .remove(&self.id);
95    }
96}
97
98pub struct TurnBuilder {
99    pub(crate) runtime: RuntimeHandle,
100    pub(crate) effect_host: Arc<dyn EffectHost>,
101    pub(crate) active_plugins: Vec<ActivePluginBinding>,
102    pub(crate) input: TurnInput,
103    pub(crate) cancel: CancellationToken,
104    pub(crate) cancels: TurnCancelRegistry,
105    pub(crate) protocol_turn_options: Option<ProtocolTurnOptions>,
106    pub(crate) provider: Option<ProviderHandle>,
107    pub(crate) turn_id: Option<String>,
108}
109
110impl TurnBuilder {
111    pub fn cancel(mut self, cancel: CancellationToken) -> Self {
112        self.cancel = cancel;
113        self
114    }
115
116    pub fn protocol_turn_options(mut self, options: ProtocolTurnOptions) -> Self {
117        self.protocol_turn_options = Some(options);
118        self
119    }
120
121    pub fn provider(mut self, provider: ProviderHandle) -> Self {
122        self.provider = Some(provider);
123        self
124    }
125
126    pub fn turn_id(mut self, id: impl Into<String>) -> Self {
127        self.turn_id = Some(id.into());
128        self
129    }
130
131    pub fn prompt_template(mut self, template: PromptTemplate) -> Self {
132        self.input.turn_context.set_prompt_template(template);
133        self
134    }
135
136    pub fn prompt_contribution(mut self, contribution: PromptContribution) -> Self {
137        self.input
138            .turn_context
139            .add_prompt_contribution(contribution);
140        self
141    }
142
143    pub fn replace_prompt_slot(
144        mut self,
145        slot: PromptSlot,
146        contributions: impl IntoIterator<Item = PromptContribution>,
147    ) -> Self {
148        self.input
149            .turn_context
150            .replace_prompt_slot(slot, contributions);
151        self
152    }
153
154    pub fn clear_prompt_slot(mut self, slot: PromptSlot) -> Self {
155        self.input.turn_context.clear_prompt_slot(slot);
156        self
157    }
158
159    pub fn prompt_layer(mut self, layer: PromptLayer) -> Self {
160        self.input.turn_context.set_prompt_layer(layer);
161        self
162    }
163
164    /// Attach typed per-turn input for an activated plugin binding.
165    ///
166    /// This is the generic primitive. Plugin crates should usually wrap it in a
167    /// domain extension trait such as `.with_tone(tone)` or `.with_board(board)`
168    /// so application code stays typed in its own vocabulary.
169    pub fn with_plugin_input<P: PluginBinding>(mut self, input: P::Input) -> Self {
170        self.input.turn_context.insert_plugin_input(P::ID, input);
171        self
172    }
173
174    pub fn effects(self, controller: &dyn RuntimeEffectController) -> ScopedTurnBuilder<'_> {
175        ScopedTurnBuilder {
176            builder: self,
177            controller,
178        }
179    }
180
181    pub async fn run(self) -> Result<TurnOutput> {
182        let collector = RunActivityCollector::default();
183        let result = self.stream_to(&collector).await?;
184        Ok(TurnOutput {
185            result,
186            activities: collector.into_activities(),
187        })
188    }
189
190    pub async fn stream_to(self, events: &dyn TurnActivitySink) -> Result<TurnResult> {
191        let effect_host = Arc::clone(&self.effect_host);
192        reject_configured_durable_effect_host(effect_host.as_ref(), "turn")?;
193        self.stream_to_with_effect_host(events, effect_host.as_ref())
194            .await
195    }
196
197    pub fn stream(self) -> Result<TurnStream> {
198        let effect_host = Arc::clone(&self.effect_host);
199        reject_configured_durable_effect_host(effect_host.as_ref(), "turn stream")?;
200        self.stream_with_effect_host(effect_host.as_ref())
201    }
202
203    /// Access lower-level turn execution that bypasses the semantic
204    /// [`TurnActivity`] tier.
205    pub fn advanced(self) -> AdvancedTurn {
206        AdvancedTurn { builder: self }
207    }
208
209    fn resolved_turn_id(&self) -> String {
210        self.turn_id
211            .clone()
212            .or_else(|| self.input.trace_turn_id.clone())
213            .unwrap_or_else(fresh_turn_id)
214    }
215
216    fn turn_scope(&self, turn_id: &str) -> lash_core::ExecutionScope {
217        lash_core::ExecutionScope::turn(self.runtime.observe().session_id(), turn_id)
218    }
219
220    pub(crate) fn prepare(
221        mut self,
222        trace_turn_id: Option<String>,
223    ) -> Result<(RuntimeHandle, TurnInput, CancellationToken, TurnCancelGuard)> {
224        if let Some(options) = self.protocol_turn_options {
225            self.input.protocol_turn_options = Some(options);
226        }
227        if let Some(provider) = self.provider {
228            self.input.turn_context.set_provider(provider);
229        }
230        if let Some(trace_turn_id) = trace_turn_id {
231            self.input.trace_turn_id = Some(trace_turn_id);
232        }
233        validate_required_plugin_inputs(&self.active_plugins, &self.input)?;
234        let cancel_guard = self.cancels.register(self.cancel.clone());
235        Ok((self.runtime, self.input, self.cancel, cancel_guard))
236    }
237
238    async fn stream_to_with_effect_host(
239        self,
240        events: &dyn TurnActivitySink,
241        effect_host: &dyn EffectHost,
242    ) -> Result<TurnResult> {
243        let turn_id = self.resolved_turn_id();
244        let scoped_effect_controller = effect_host.scoped(self.turn_scope(&turn_id))?;
245        self.stream_to_with_scope(events, scoped_effect_controller, Some(turn_id))
246            .await
247    }
248
249    async fn stream_to_with_effect_controller(
250        self,
251        events: &dyn TurnActivitySink,
252        controller: &dyn RuntimeEffectController,
253    ) -> Result<TurnResult> {
254        let turn_id = self.resolved_turn_id();
255        let scoped_effect_controller =
256            ScopedEffectController::borrowed(controller, self.turn_scope(&turn_id))?;
257        self.stream_to_with_scope(events, scoped_effect_controller, Some(turn_id))
258            .await
259    }
260
261    async fn stream_to_with_scope(
262        self,
263        events: &dyn TurnActivitySink,
264        scoped_effect_controller: ScopedEffectController<'_>,
265        trace_turn_id: Option<String>,
266    ) -> Result<TurnResult> {
267        let (runtime, input, cancel, _cancel_guard) = self.prepare(trace_turn_id)?;
268        stream_prepared_turn(
269            &runtime,
270            input,
271            TurnSinks::turn(events),
272            scoped_effect_controller,
273            cancel,
274        )
275        .await
276    }
277
278    fn stream_with_effect_host(self, effect_host: &dyn EffectHost) -> Result<TurnStream> {
279        let turn_id = self.resolved_turn_id();
280        let scoped_effect_controller = effect_host
281            .scoped_static(self.turn_scope(&turn_id))?
282            .ok_or(EmbedError::StaticTurnStreamRequiresStaticEffectHost)?;
283        self.stream_with_scope(scoped_effect_controller, Some(turn_id))
284    }
285
286    fn stream_with_scope(
287        self,
288        scoped_effect_controller: ScopedEffectController<'static>,
289        trace_turn_id: Option<String>,
290    ) -> Result<TurnStream> {
291        let (runtime, input, cancel, cancel_guard) = self.prepare(trace_turn_id)?;
292        let (tx, rx) = mpsc::channel(64);
293        let sink = ChannelTurnActivitySink { tx };
294        let completion = tokio::spawn(async move {
295            let _cancel_guard = cancel_guard;
296            stream_prepared_turn(
297                &runtime,
298                input,
299                TurnSinks::turn(&sink),
300                scoped_effect_controller,
301                cancel,
302            )
303            .await
304        });
305        Ok(TurnStream {
306            activities: rx,
307            completion,
308        })
309    }
310}
311
312pub struct ScopedTurnBuilder<'run> {
313    builder: TurnBuilder,
314    controller: &'run dyn RuntimeEffectController,
315}
316
317impl<'run> ScopedTurnBuilder<'run> {
318    pub fn cancel(mut self, cancel: CancellationToken) -> Self {
319        self.builder = self.builder.cancel(cancel);
320        self
321    }
322
323    pub fn protocol_turn_options(mut self, options: ProtocolTurnOptions) -> Self {
324        self.builder = self.builder.protocol_turn_options(options);
325        self
326    }
327
328    pub fn provider(mut self, provider: ProviderHandle) -> Self {
329        self.builder = self.builder.provider(provider);
330        self
331    }
332
333    pub fn turn_id(mut self, id: impl Into<String>) -> Self {
334        self.builder = self.builder.turn_id(id);
335        self
336    }
337
338    pub fn prompt_template(mut self, template: PromptTemplate) -> Self {
339        self.builder = self.builder.prompt_template(template);
340        self
341    }
342
343    pub fn prompt_contribution(mut self, contribution: PromptContribution) -> Self {
344        self.builder = self.builder.prompt_contribution(contribution);
345        self
346    }
347
348    pub fn replace_prompt_slot(
349        mut self,
350        slot: PromptSlot,
351        contributions: impl IntoIterator<Item = PromptContribution>,
352    ) -> Self {
353        self.builder = self.builder.replace_prompt_slot(slot, contributions);
354        self
355    }
356
357    pub fn clear_prompt_slot(mut self, slot: PromptSlot) -> Self {
358        self.builder = self.builder.clear_prompt_slot(slot);
359        self
360    }
361
362    pub fn prompt_layer(mut self, layer: PromptLayer) -> Self {
363        self.builder = self.builder.prompt_layer(layer);
364        self
365    }
366
367    pub fn with_plugin_input<P: PluginBinding>(mut self, input: P::Input) -> Self {
368        self.builder = self.builder.with_plugin_input::<P>(input);
369        self
370    }
371
372    pub async fn run(self) -> Result<TurnOutput> {
373        let collector = RunActivityCollector::default();
374        let result = self.stream_to(&collector).await?;
375        Ok(TurnOutput {
376            result,
377            activities: collector.into_activities(),
378        })
379    }
380
381    pub async fn stream_to(self, events: &dyn TurnActivitySink) -> Result<TurnResult> {
382        self.builder
383            .stream_to_with_effect_controller(events, self.controller)
384            .await
385    }
386}
387
388/// Lower-level turn execution that exposes the raw runtime event stream.
389///
390/// Reachable via [`TurnBuilder::advanced`]. Most applications should use
391/// [`TurnBuilder::stream_to`] for semantic turn activity; benchmarks and
392/// diagnostics use this when they need the same low-level event stream as the
393/// runtime trace.
394pub struct AdvancedTurn {
395    builder: TurnBuilder,
396}
397
398impl AdvancedTurn {
399    pub async fn run_with_scope(
400        self,
401        scoped_effect_controller: ScopedEffectController<'_>,
402    ) -> Result<TurnOutput> {
403        let collector = RunActivityCollector::default();
404        let result = self
405            .stream_to_with_scope(&collector, scoped_effect_controller)
406            .await?;
407        Ok(TurnOutput {
408            result,
409            activities: collector.into_activities(),
410        })
411    }
412
413    pub async fn collect_with_scope(
414        self,
415        events: &dyn TurnActivitySink,
416        scoped_effect_controller: ScopedEffectController<'_>,
417    ) -> Result<TurnOutput> {
418        let collector = RunActivityCollector::default();
419        let fanout = BorrowedTurnActivityFanout {
420            live: events,
421            collector: &collector,
422        };
423        let result = self
424            .stream_to_with_scope(&fanout, scoped_effect_controller)
425            .await?;
426        Ok(TurnOutput {
427            result,
428            activities: collector.into_activities(),
429        })
430    }
431
432    pub async fn stream_to_with_scope(
433        self,
434        events: &dyn TurnActivitySink,
435        scoped_effect_controller: ScopedEffectController<'_>,
436    ) -> Result<TurnResult> {
437        let trace_turn_id = trace_turn_id_for_scope(&self.builder, &scoped_effect_controller);
438        self.builder
439            .stream_to_with_scope(events, scoped_effect_controller, trace_turn_id)
440            .await
441    }
442
443    pub fn stream_with_scope(
444        self,
445        scoped_effect_controller: ScopedEffectController<'static>,
446    ) -> Result<TurnStream> {
447        let trace_turn_id = trace_turn_id_for_scope(&self.builder, &scoped_effect_controller);
448        self.builder
449            .stream_with_scope(scoped_effect_controller, trace_turn_id)
450    }
451
452    /// Run the turn while sending raw lower-level runtime events to `events`.
453    pub async fn collect_session_events_with_scope(
454        self,
455        events: &dyn EventSink,
456        scoped_effect_controller: ScopedEffectController<'_>,
457    ) -> Result<TurnResult> {
458        let trace_turn_id = trace_turn_id_for_scope(&self.builder, &scoped_effect_controller);
459        let (runtime, input, cancel, _cancel_guard) = self.builder.prepare(trace_turn_id)?;
460        stream_prepared_turn(
461            &runtime,
462            input,
463            TurnSinks::session(events),
464            scoped_effect_controller,
465            cancel,
466        )
467        .await
468    }
469}
470
471pub struct TurnStream {
472    activities: mpsc::Receiver<Result<TurnActivity>>,
473    completion: JoinHandle<Result<TurnResult>>,
474}
475
476impl TurnStream {
477    pub async fn next_activity(&mut self) -> Option<Result<TurnActivity>> {
478        self.activities.recv().await
479    }
480
481    pub async fn finish(self) -> Result<TurnResult> {
482        self.completion.await.map_err(|err| {
483            EmbedError::Runtime(lash_core::RuntimeError::new(
484                RuntimeErrorCode::TurnStreamJoin,
485                format!("turn stream task failed: {err}"),
486            ))
487        })?
488    }
489}
490
491impl Stream for TurnStream {
492    type Item = Result<TurnActivity>;
493
494    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
495        self.activities.poll_recv(cx)
496    }
497}
498
499pub struct QueuedTurnBuilder {
500    pub(crate) runtime: RuntimeHandle,
501    pub(crate) effect_host: Arc<dyn EffectHost>,
502    pub(crate) cancel: CancellationToken,
503    pub(crate) cancels: TurnCancelRegistry,
504    pub(crate) batch_ids: Vec<String>,
505    pub(crate) drain_id: Option<String>,
506}
507
508impl QueuedTurnBuilder {
509    pub fn cancel(mut self, cancel: CancellationToken) -> Self {
510        self.cancel = cancel;
511        self
512    }
513
514    pub fn batch_ids(mut self, batch_ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
515        self.batch_ids = batch_ids.into_iter().map(Into::into).collect();
516        self
517    }
518
519    pub fn drain_id(mut self, drain_id: impl Into<String>) -> Self {
520        self.drain_id = Some(drain_id.into());
521        self
522    }
523
524    pub fn effects(self, controller: &dyn RuntimeEffectController) -> ScopedQueuedTurnBuilder<'_> {
525        ScopedQueuedTurnBuilder {
526            builder: self,
527            controller,
528        }
529    }
530
531    pub async fn run(self) -> Result<Option<TurnOutput>> {
532        let collector = RunActivityCollector::default();
533        let Some(result) = self.stream_to(&collector).await? else {
534            return Ok(None);
535        };
536        Ok(Some(TurnOutput {
537            result,
538            activities: collector.into_activities(),
539        }))
540    }
541
542    pub async fn stream_to(self, events: &dyn TurnActivitySink) -> Result<Option<TurnResult>> {
543        let effect_host = Arc::clone(&self.effect_host);
544        reject_configured_durable_effect_host(effect_host.as_ref(), "queued turn")?;
545        self.stream_to_with_effect_host(events, effect_host.as_ref())
546            .await
547    }
548
549    pub fn advanced(self) -> AdvancedQueuedTurn {
550        AdvancedQueuedTurn { builder: self }
551    }
552
553    fn resolved_drain_id(&self) -> String {
554        self.drain_id
555            .clone()
556            .or_else(|| self.batch_ids.first().cloned())
557            .unwrap_or_else(fresh_queue_drain_id)
558    }
559
560    async fn stream_to_with_effect_host(
561        self,
562        events: &dyn TurnActivitySink,
563        effect_host: &dyn EffectHost,
564    ) -> Result<Option<TurnResult>> {
565        let drain_id = self.resolved_drain_id();
566        let scope =
567            lash_core::ExecutionScope::queue_drain(self.runtime.observe().session_id(), drain_id);
568        let scoped_effect_controller = effect_host.scoped(scope)?;
569        self.stream_to_with_scope(events, scoped_effect_controller)
570            .await
571    }
572
573    async fn stream_to_with_effect_controller(
574        self,
575        events: &dyn TurnActivitySink,
576        controller: &dyn RuntimeEffectController,
577    ) -> Result<Option<TurnResult>> {
578        let drain_id = self.resolved_drain_id();
579        let scope =
580            lash_core::ExecutionScope::queue_drain(self.runtime.observe().session_id(), drain_id);
581        let scoped_effect_controller = ScopedEffectController::borrowed(controller, scope)?;
582        self.stream_to_with_scope(events, scoped_effect_controller)
583            .await
584    }
585
586    async fn stream_to_with_scope(
587        self,
588        events: &dyn TurnActivitySink,
589        scoped_effect_controller: ScopedEffectController<'_>,
590    ) -> Result<Option<TurnResult>> {
591        let Self {
592            runtime,
593            effect_host: _,
594            cancel,
595            cancels,
596            batch_ids,
597            drain_id: _,
598        } = self;
599        let _cancel_guard = cancels.register(cancel.clone());
600        stream_next_queued_prepared_turn(
601            &runtime,
602            TurnSinks::turn(events),
603            scoped_effect_controller,
604            cancel,
605            &batch_ids,
606        )
607        .await
608    }
609}
610
611pub struct ScopedQueuedTurnBuilder<'run> {
612    builder: QueuedTurnBuilder,
613    controller: &'run dyn RuntimeEffectController,
614}
615
616impl<'run> ScopedQueuedTurnBuilder<'run> {
617    pub fn cancel(mut self, cancel: CancellationToken) -> Self {
618        self.builder = self.builder.cancel(cancel);
619        self
620    }
621
622    pub fn batch_ids(mut self, batch_ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
623        self.builder = self.builder.batch_ids(batch_ids);
624        self
625    }
626
627    pub fn drain_id(mut self, drain_id: impl Into<String>) -> Self {
628        self.builder = self.builder.drain_id(drain_id);
629        self
630    }
631
632    pub async fn run(self) -> Result<Option<TurnOutput>> {
633        let collector = RunActivityCollector::default();
634        let Some(result) = self.stream_to(&collector).await? else {
635            return Ok(None);
636        };
637        Ok(Some(TurnOutput {
638            result,
639            activities: collector.into_activities(),
640        }))
641    }
642
643    pub async fn stream_to(self, events: &dyn TurnActivitySink) -> Result<Option<TurnResult>> {
644        self.builder
645            .stream_to_with_effect_controller(events, self.controller)
646            .await
647    }
648}
649
650pub struct AdvancedQueuedTurn {
651    builder: QueuedTurnBuilder,
652}
653
654impl AdvancedQueuedTurn {
655    pub async fn stream_to_with_scope(
656        self,
657        events: &dyn TurnActivitySink,
658        scoped_effect_controller: ScopedEffectController<'_>,
659    ) -> Result<Option<TurnResult>> {
660        self.builder
661            .stream_to_with_scope(events, scoped_effect_controller)
662            .await
663    }
664}
665
666fn fresh_turn_id() -> String {
667    lash_core::TurnActivityId::fresh().0
668}
669
670fn fresh_queue_drain_id() -> String {
671    format!("queue-drain:{}", fresh_turn_id())
672}
673
674fn trace_turn_id_for_scope(
675    builder: &TurnBuilder,
676    scoped_effect_controller: &ScopedEffectController<'_>,
677) -> Option<String> {
678    if scoped_effect_controller
679        .execution_scope()
680        .validates_turn_trace_id()
681    {
682        Some(
683            builder
684                .turn_id
685                .clone()
686                .unwrap_or_else(|| scoped_effect_controller.scope_id().to_string()),
687        )
688    } else {
689        builder
690            .turn_id
691            .clone()
692            .or_else(|| builder.input.trace_turn_id.clone())
693    }
694}
695
696fn reject_configured_durable_effect_host(
697    effect_host: &dyn EffectHost,
698    operation: &'static str,
699) -> Result<()> {
700    if effect_host.durability_tier() == DurabilityTier::Durable {
701        return Err(EmbedError::DurableEffectHostRequiresHandlerContext { operation });
702    }
703    Ok(())
704}
705
706pub(crate) async fn stream_next_queued_prepared_turn(
707    runtime: &RuntimeHandle,
708    sinks: TurnSinks<'_>,
709    scoped_effect_controller: ScopedEffectController<'_>,
710    cancel: CancellationToken,
711    batch_ids: &[String],
712) -> Result<Option<TurnResult>> {
713    let turn = Box::pin(stream_next_queued_prepared_assembled(
714        runtime,
715        sinks,
716        scoped_effect_controller,
717        cancel,
718        batch_ids,
719    ))
720    .await?;
721    Ok(turn.map(TurnResult::from_assembled))
722}
723
724pub(crate) async fn stream_next_queued_prepared_assembled(
725    runtime: &RuntimeHandle,
726    sinks: TurnSinks<'_>,
727    scoped_effect_controller: ScopedEffectController<'_>,
728    cancel: CancellationToken,
729    batch_ids: &[String],
730) -> Result<Option<AssembledTurn>> {
731    let writer_handle = runtime.writer();
732    let mut writer = writer_handle.lock().await;
733    let observation_sink = SessionObservationTurnActivitySink {
734        runtime: runtime.clone(),
735        live: sinks.turn_events(),
736    };
737    let opts = turn_options(
738        sinks.events(),
739        &observation_sink,
740        scoped_effect_controller,
741        cancel,
742    );
743    let turn = if batch_ids.is_empty() {
744        writer.stream_next_queued_work(opts).await?
745    } else {
746        writer.stream_selected_queued_work(opts, batch_ids).await?
747    };
748    runtime.publish_from(&writer);
749    Ok(turn)
750}
751
752fn turn_options<'a>(
753    events: Option<&'a dyn EventSink>,
754    turn_events: &'a dyn TurnActivitySink,
755    scoped_effect_controller: ScopedEffectController<'a>,
756    cancel: CancellationToken,
757) -> lash_core::TurnOptions<'a> {
758    let mut opts = lash_core::TurnOptions::new(cancel, scoped_effect_controller);
759    if let Some(events) = events {
760        opts = opts.with_events(events);
761    }
762    opts.with_turn_events(turn_events)
763}
764
765struct SessionObservationTurnActivitySink<'a> {
766    runtime: RuntimeHandle,
767    live: Option<&'a dyn TurnActivitySink>,
768}
769
770#[async_trait]
771impl TurnActivitySink for SessionObservationTurnActivitySink<'_> {
772    fn is_noop(&self) -> bool {
773        false
774    }
775
776    async fn emit(&self, activity: TurnActivity) {
777        self.runtime.record_turn_activity(activity.clone());
778        if let Some(live) = self.live {
779            live.emit(activity).await;
780        }
781    }
782}
783
784struct ChannelTurnActivitySink {
785    tx: mpsc::Sender<Result<TurnActivity>>,
786}
787
788#[async_trait]
789impl TurnActivitySink for ChannelTurnActivitySink {
790    async fn emit(&self, activity: TurnActivity) {
791        let _ = self.tx.send(Ok(activity)).await;
792    }
793}
794fn validate_required_plugin_inputs(
795    active_plugins: &[ActivePluginBinding],
796    input: &TurnInput,
797) -> Result<()> {
798    for plugin in active_plugins {
799        if plugin.requires_turn_input && !input.turn_context.has_plugin_input(plugin.id) {
800            return Err(EmbedError::MissingPluginTurnInput {
801                plugin_id: plugin.id,
802            });
803        }
804    }
805    Ok(())
806}
807
808pub(crate) async fn stream_prepared_turn(
809    runtime: &RuntimeHandle,
810    input: TurnInput,
811    sinks: TurnSinks<'_>,
812    scoped_effect_controller: ScopedEffectController<'_>,
813    cancel: CancellationToken,
814) -> Result<TurnResult> {
815    let turn = Box::pin(stream_prepared_assembled(
816        runtime,
817        input,
818        sinks,
819        scoped_effect_controller,
820        cancel,
821    ))
822    .await?;
823    Ok(TurnResult::from_assembled(turn))
824}
825
826pub(crate) async fn stream_prepared_assembled(
827    runtime: &RuntimeHandle,
828    input: TurnInput,
829    sinks: TurnSinks<'_>,
830    scoped_effect_controller: ScopedEffectController<'_>,
831    cancel: CancellationToken,
832) -> Result<AssembledTurn> {
833    let turn = Box::pin(stream_prepared_agent_frame_run(
834        runtime,
835        input,
836        sinks,
837        scoped_effect_controller,
838        cancel,
839    ))
840    .await?;
841    turn.into_final_turn().ok_or_else(|| {
842        EmbedError::Runtime(lash_core::RuntimeError::new(
843            RuntimeErrorCode::EmptyAgentFrameRun,
844            "runtime completed without an assembled turn",
845        ))
846    })
847}
848
849pub(crate) async fn stream_prepared_agent_frame_run(
850    runtime: &RuntimeHandle,
851    input: TurnInput,
852    sinks: TurnSinks<'_>,
853    scoped_effect_controller: ScopedEffectController<'_>,
854    cancel: CancellationToken,
855) -> Result<lash_core::AgentFrameRun> {
856    let writer_handle = runtime.writer();
857    let mut writer = writer_handle.lock().await;
858    if let Some(extension) = input.protocol_extension.as_ref() {
859        writer
860            .validate_protocol_turn_extension(extension)
861            .await
862            .map_err(EmbedError::Session)?;
863    }
864    let observation_sink = SessionObservationTurnActivitySink {
865        runtime: runtime.clone(),
866        live: sinks.turn_events(),
867    };
868    let turn = Box::pin(writer.stream_turn_with_agent_frames(
869        input,
870        turn_options(
871            sinks.events(),
872            &observation_sink,
873            scoped_effect_controller,
874            cancel,
875        ),
876    ))
877    .await?;
878    runtime.publish_from(&writer);
879    Ok(turn)
880}
881
882#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
883pub struct TurnResult {
884    pub state: SessionSnapshot,
885    pub outcome: TurnOutcome,
886    pub assistant_output: AssistantOutput,
887    /// Parent's own LLM tokens for this turn. Does **not** include child
888    /// sessions; see [`children_usage`](Self::children_usage) and
889    /// [`total_usage`](Self::total_usage).
890    pub usage: TokenUsage,
891    /// Per-`(source, model)` ledger entries for child sessions whose LLM
892    /// calls completed during this turn (subagents, compaction, observers,
893    /// etc.). Empty unless the turn spawned children.
894    #[serde(default)]
895    pub children_usage: Vec<TokenLedgerEntry>,
896    /// Provider calls made by the parent session during this turn, in protocol
897    /// order. Child-session calls remain on each child's result. This is the
898    /// complete lash-side model attribution surface: a turn has no single
899    /// producing model, so lash exposes the per-call ledger and the host
900    /// composes any higher-level view from it (ADR 0033).
901    #[serde(default)]
902    pub llm_calls: Vec<LlmCallRecord>,
903    pub tool_calls: Vec<ToolCallRecord>,
904    pub execution: ExecutionSummary,
905    pub errors: Vec<TurnIssue>,
906}
907
908impl TurnResult {
909    fn from_assembled(turn: lash_core::AssembledTurn) -> Self {
910        Self {
911            state: turn.state,
912            outcome: turn.outcome,
913            assistant_output: turn.assistant_output,
914            usage: turn.token_usage,
915            children_usage: turn.children_usage,
916            llm_calls: turn.llm_calls,
917            tool_calls: turn.tool_calls,
918            execution: turn.execution,
919            errors: turn.errors,
920        }
921    }
922
923    /// Sum of parent's own LLM tokens and every child session's LLM tokens
924    /// for this turn.
925    pub fn total_usage(&self) -> TokenUsage {
926        let mut total = self.usage.clone();
927        for entry in &self.children_usage {
928            total.add(&entry.usage);
929        }
930        total
931    }
932
933    /// Wall-clock instant the runtime started this turn (claim of the
934    /// session-execution lease / queued-work claim), read from the runtime
935    /// clock. Backed by [`ExecutionSummary::started_at_ms`] on
936    /// [`execution`](Self::execution).
937    pub fn started_at(&self) -> std::time::SystemTime {
938        std::time::UNIX_EPOCH + std::time::Duration::from_millis(self.execution.started_at_ms)
939    }
940
941    /// Whole-turn duration — claim through final commit and post-persist
942    /// hooks — measured on the runtime clock's monotonic source. Backed by
943    /// [`ExecutionSummary::duration_ms`] on [`execution`](Self::execution).
944    pub fn duration(&self) -> std::time::Duration {
945        std::time::Duration::from_millis(self.execution.duration_ms)
946    }
947
948    pub fn assistant_message(&self) -> Option<&str> {
949        match &self.outcome {
950            TurnOutcome::Finished(lash_core::TurnFinish::AssistantMessage { text }) => Some(text),
951            _ => None,
952        }
953    }
954
955    pub fn final_value(&self) -> Option<&serde_json::Value> {
956        match &self.outcome {
957            TurnOutcome::Finished(lash_core::TurnFinish::FinalValue { value }) => Some(value),
958            _ => None,
959        }
960    }
961
962    pub fn tool_value(&self) -> Option<(&str, &serde_json::Value)> {
963        match &self.outcome {
964            TurnOutcome::Finished(lash_core::TurnFinish::ToolValue { tool_name, value }) => {
965                Some((tool_name.as_str(), value))
966            }
967            _ => None,
968        }
969    }
970
971    pub fn is_success(&self) -> bool {
972        matches!(
973            self.outcome,
974            TurnOutcome::Finished(_) | TurnOutcome::AgentFrameSwitch { .. }
975        )
976    }
977}
978
979#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
980pub struct TurnOutput {
981    pub result: TurnResult,
982    pub activities: Vec<TurnActivity>,
983}
984
985impl TurnOutput {
986    pub fn assistant_message(&self) -> Option<&str> {
987        self.result.assistant_message()
988    }
989
990    pub fn final_value(&self) -> Option<&serde_json::Value> {
991        self.result.final_value()
992    }
993
994    pub fn tool_value(&self) -> Option<(&str, &serde_json::Value)> {
995        self.result.tool_value()
996    }
997
998    pub fn is_success(&self) -> bool {
999        self.result.is_success()
1000    }
1001}
1002
1003struct BorrowedTurnActivityFanout<'a> {
1004    live: &'a dyn TurnActivitySink,
1005    collector: &'a RunActivityCollector,
1006}
1007
1008#[async_trait]
1009impl TurnActivitySink for BorrowedTurnActivityFanout<'_> {
1010    async fn emit(&self, activity: TurnActivity) {
1011        self.live.emit(activity.clone()).await;
1012        self.collector.emit(activity).await;
1013    }
1014}
1015
1016#[derive(Default)]
1017pub(crate) struct RunActivityCollector {
1018    activities: Arc<StdMutex<Vec<TurnActivity>>>,
1019}
1020
1021impl RunActivityCollector {
1022    fn into_activities(self) -> Vec<TurnActivity> {
1023        self.activities
1024            .lock()
1025            .expect("run activity collector lock")
1026            .clone()
1027    }
1028
1029    #[cfg(all(test, feature = "rlm"))]
1030    pub(crate) fn snapshot(&self) -> Vec<TurnActivity> {
1031        self.activities
1032            .lock()
1033            .expect("run activity collector lock")
1034            .clone()
1035    }
1036}
1037
1038#[async_trait]
1039impl TurnActivitySink for RunActivityCollector {
1040    async fn emit(&self, activity: TurnActivity) {
1041        self.activities
1042            .lock()
1043            .expect("run activity collector lock")
1044            .push(activity);
1045    }
1046}
1047
1048pub struct TurnActivityFanout {
1049    sinks: Vec<Arc<dyn TurnActivitySink>>,
1050}
1051
1052impl TurnActivityFanout {
1053    pub fn new(sinks: impl IntoIterator<Item = Arc<dyn TurnActivitySink>>) -> Self {
1054        Self {
1055            sinks: sinks.into_iter().collect(),
1056        }
1057    }
1058}
1059
1060#[async_trait]
1061impl TurnActivitySink for TurnActivityFanout {
1062    async fn emit(&self, activity: TurnActivity) {
1063        for sink in &self.sinks {
1064            sink.emit(activity.clone()).await;
1065        }
1066    }
1067}
1068
1069pub fn message_text(message: &Message) -> String {
1070    message
1071        .parts
1072        .iter()
1073        .map(|part| part.content.as_str())
1074        .collect::<Vec<_>>()
1075        .join("\n")
1076}
1077
1078pub fn message_role(message: &Message) -> &'static str {
1079    match message.role {
1080        MessageRole::User => "user",
1081        MessageRole::Assistant => "assistant",
1082        MessageRole::System => "system",
1083        MessageRole::Event => "event",
1084    }
1085}
lash/turn.rs

lash/
turn.rs
