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

lash/
turn.rs
