1use super::*;
2
3fn trace_fields_from_outcome(
4 outcome: &TurnOutcome,
5) -> (
6 &'static str,
7 &'static str,
8 Option<lash_trace::TraceAgentFrameSwitch>,
9) {
10 match outcome {
11 TurnOutcome::Finished(TurnFinish::AssistantMessage { .. }) => {
12 ("completed", "assistant_message", None)
13 }
14 TurnOutcome::Finished(TurnFinish::SubmittedValue { .. }) => {
15 ("completed", "submitted_value", None)
16 }
17 TurnOutcome::Finished(TurnFinish::ToolValue { .. }) => ("completed", "tool_value", None),
18 TurnOutcome::AgentFrameSwitch { frame_id, .. } => (
19 "completed",
20 "agent_frame_switch",
21 Some(lash_trace::TraceAgentFrameSwitch {
22 frame_id: frame_id.clone(),
23 }),
24 ),
25 TurnOutcome::Stopped(stop) => ("failed", trace_stop_reason(stop), None),
26 }
27}
28
29fn trace_stop_reason(stop: &TurnStop) -> &'static str {
30 match stop {
31 TurnStop::Cancelled => "cancelled",
32 TurnStop::Incomplete => "incomplete",
33 TurnStop::InvalidInput => "invalid_input",
34 TurnStop::MaxTurns => "max_turns",
35 TurnStop::ToolFailure => "tool_failure",
36 TurnStop::ProviderError => "provider_error",
37 TurnStop::PluginAbort => "plugin_abort",
38 TurnStop::RuntimeError => "runtime_error",
39 TurnStop::SubmittedError { .. } => "submitted_error",
40 TurnStop::ToolError { .. } => "tool_error",
41 }
42}
43
44fn session_head_refresh_error(err: SessionError) -> RuntimeError {
45 RuntimeError::new(
46 RuntimeErrorCode::Other("session_head_refresh".to_string()),
47 err.to_string(),
48 )
49}
50
51fn queued_work_payload_type(payload: &crate::QueuedWorkPayload) -> &'static str {
52 match payload {
53 crate::QueuedWorkPayload::TurnInput { .. } => "turn_input",
54 crate::QueuedWorkPayload::ProcessWake { .. } => "process_wake",
55 crate::QueuedWorkPayload::SessionCommand { command } => command.kind(),
56 }
57}
58
59fn queued_work_batch_ids(claim: &crate::QueuedWorkClaim) -> Vec<String> {
60 claim
61 .batches
62 .iter()
63 .map(|batch| batch.batch_id.clone())
64 .collect()
65}
66
67fn turn_phase_id(parent_turn_id: &str, phase: &str) -> String {
68 format!("{parent_turn_id}:{phase}")
69}
70
71fn scoped_child_turn_controller<'run>(
72 scoped_effect_controller: &'run ScopedEffectController<'_>,
73 session_id: &str,
74 turn_id: &str,
75) -> Result<ScopedEffectController<'run>, RuntimeError> {
76 ScopedEffectController::borrowed(
77 scoped_effect_controller.controller(),
78 EffectScope::turn(session_id, turn_id),
79 )
80}
81
82pub(in crate::runtime) fn queued_work_trace_payload(
83 boundary: crate::QueuedWorkClaimBoundary,
84 claim: &crate::QueuedWorkClaim,
85 causes: &[crate::TurnCause],
86) -> serde_json::Value {
87 serde_json::json!({
88 "boundary": boundary,
89 "claim_id": claim.claim_id,
90 "owner_id": claim.owner_id,
91 "batch_ids": queued_work_batch_ids(claim),
92 "payload_types": claim.batches.iter()
93 .flat_map(|batch| batch.items.iter())
94 .map(|item| queued_work_payload_type(&item.payload))
95 .collect::<Vec<_>>(),
96 "causes": causes,
97 })
98}
99
100pub(in crate::runtime) fn queued_work_completion_trace_payload(
101 completions: &[crate::QueuedWorkCompletion],
102) -> serde_json::Value {
103 serde_json::json!({
104 "claims": completions.iter().map(|completion| {
105 serde_json::json!({
106 "session_id": completion.session_id,
107 "claim_id": completion.claim_id,
108 "batch_ids": completion.batch_ids,
109 })
110 }).collect::<Vec<_>>(),
111 })
112}
113
114async fn emit_queued_work_started_to_sink(
115 events: &dyn TurnActivitySink,
116 boundary: crate::QueuedWorkClaimBoundary,
117 claim: &crate::QueuedWorkClaim,
118 causes: Vec<crate::TurnCause>,
119) {
120 emit_turn_activity_to_sink(
121 events,
122 TurnActivity::independent(TurnEvent::QueuedWorkStarted {
123 boundary,
124 batch_ids: queued_work_batch_ids(claim),
125 causes,
126 }),
127 )
128 .await;
129}
130
131pub(in crate::runtime) async fn send_queued_work_started_event(
132 event_tx: &mpsc::Sender<RuntimeStreamEvent>,
133 boundary: crate::QueuedWorkClaimBoundary,
134 claim: &crate::QueuedWorkClaim,
135 causes: Vec<crate::TurnCause>,
136) {
137 send_turn_activity(
138 event_tx,
139 TurnActivityId::fresh(),
140 TurnEvent::QueuedWorkStarted {
141 boundary,
142 batch_ids: queued_work_batch_ids(claim),
143 causes,
144 },
145 )
146 .await;
147}
148
149struct TurnFinishInput {
150 turn_pipeline: TurnBoundary,
151 assembler: TurnAssembler,
152 new_messages: crate::MessageSequence,
153 policy: RuntimeSessionPolicy,
154 turn_index: usize,
155 queued_work_completions: Vec<crate::QueuedWorkCompletion>,
156 trace_turn_id: String,
157}
158
159impl LashRuntime {
160 fn max_context_tokens(&self) -> usize {
161 self.state.effective_policy().context_window_tokens()
162 }
163 #[doc(hidden)]
164 pub fn set_turn_phase_probe(&mut self, probe: Arc<dyn RuntimeTurnPhaseProbe>) {
165 self.turn_phase_probe = Some(probe);
166 }
167
168 fn mark_phase_begin(&self, phase: RuntimeTurnPhase) {
169 if let Some(probe) = self.turn_phase_probe.as_ref() {
170 probe.begin(phase);
171 }
172 }
173
174 fn mark_phase_end(&self, phase: RuntimeTurnPhase) {
175 if let Some(probe) = self.turn_phase_probe.as_ref() {
176 probe.end(phase);
177 }
178 }
179
180 async fn finish_turn(
181 &mut self,
182 finish: TurnFinishInput,
183 events: &dyn EventSink,
184 scoped_effect_controller: &ScopedEffectController<'_>,
185 cancel_state: &CancellationToken,
186 ) -> Result<AssembledTurn, RuntimeError> {
187 let TurnFinishInput {
188 mut turn_pipeline,
189 assembler,
190 new_messages,
191 policy,
192 turn_index,
193 queued_work_completions,
194 trace_turn_id,
195 } = finish;
196 self.policy = self.state.effective_policy().clone();
197 turn_pipeline.state_mut().policy = self.policy.clone();
198 turn_pipeline.state_mut().turn_index = turn_index;
199
200 let mut turn_usage_delta = {
201 let mut ledger = self.shared_token_ledger.lock().expect("token ledger lock");
202 std::mem::take(&mut *ledger)
203 };
204 if assembler.token_usage.total() > 0 || assembler.token_usage.cached_input_tokens > 0 {
205 turn_usage_delta.push(TokenLedgerEntry {
206 source: "turn".to_string(),
207 model: policy.model.id.clone(),
208 usage: assembler.token_usage.clone(),
209 });
210 }
211 let turn_usage_delta = merge_usage_delta_entries(turn_usage_delta);
212
213 turn_pipeline.finalize_turn_read_state(new_messages, cancel_state.is_cancelled());
214 if assembler.token_usage.total() > 0 || assembler.token_usage.cached_input_tokens > 0 {
215 turn_pipeline.state_mut().token_usage = assembler.token_usage.clone();
216 }
217
218 let last_prompt_usage = assembler
219 .last_llm_usage()
220 .and_then(|usage| normalize_prompt_usage(policy.provider(), usage));
221 turn_pipeline.state_mut().last_prompt_usage = last_prompt_usage;
222 let assembled_state = turn_pipeline.export_state_for_assembly();
223 let assembled = assembler.finish(
224 assembled_state,
225 cancel_state.is_cancelled(),
226 None,
227 &self.host.core.control.termination,
228 );
229
230 let Some(session) = self.session.as_ref() else {
231 self.state.apply_snapshot(&assembled.state);
232 self.emit_completed_turn_trace(&assembled.state, &assembled.outcome, &trace_turn_id);
233 return Ok(assembled);
234 };
235
236 let plugins = Arc::clone(session.plugins());
237 let manager = match self.runtime_session_services_for_turn(None) {
238 Ok(manager) => manager,
239 Err(err) => {
240 return Err(RuntimeError::new(
241 RuntimeErrorCode::PluginSessionManager,
242 err.to_string(),
243 ));
244 }
245 };
246
247 self.mark_phase_begin(RuntimeTurnPhase::FinalizeTurn);
248 let finalized = match plugins
249 .finalize_turn_with_phase_probe(
250 assembled,
251 manager.state_service(),
252 manager.lifecycle_service(),
253 manager.graph_service(),
254 self.turn_phase_probe.clone(),
255 )
256 .await
257 {
258 Ok(finalized) => finalized,
259 Err(err) => {
260 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
261 return Err(RuntimeError::new(
262 RuntimeErrorCode::PluginFinalizeTurn,
263 err.to_string(),
264 ));
265 }
266 };
267 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
268
269 let mut returned_turn = finalized.turn;
270 self.mark_phase_begin(RuntimeTurnPhase::PersistTurn);
271 self.mark_phase_begin(RuntimeTurnPhase::FinalCommit);
272 let queued_work_completion_trace = queued_work_completions.clone();
273 let pending_attachment_ids = self
274 .host
275 .core
276 .durability
277 .attachment_store
278 .pending_manifest_commit_ids();
279 if let Err(err) = turn_pipeline
280 .final_commit(
281 &mut returned_turn,
282 self.session.as_mut(),
283 &turn_usage_delta,
284 Some(&trace_turn_id),
285 queued_work_completions,
286 pending_attachment_ids.clone(),
287 )
288 .await
289 {
290 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
291 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
292 return Err(err);
293 }
294 self.host
295 .core
296 .durability
297 .attachment_store
298 .mark_manifest_committed(&pending_attachment_ids);
299 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
300
301 emit_session_events_to_sink(events, finalized.events).await;
302 self.state = turn_pipeline.into_final_state();
303 if matches!(returned_turn.outcome, TurnOutcome::AgentFrameSwitch { .. })
304 && let Some(session) = self.session.as_mut()
305 {
306 let protocol_session = Arc::clone(session.plugins().protocol_session());
307 let session_id = self.state.session_id.clone();
308 protocol_session
309 .restore_session(
310 crate::plugin::ProtocolSessionContext::new(session, &session_id),
311 &self.state,
312 )
313 .await
314 .map_err(|err| {
315 RuntimeError::new(
316 RuntimeErrorCode::Other("protocol_restore_session".to_string()),
317 err.to_string(),
318 )
319 })?;
320 }
321 if !queued_work_completion_trace.is_empty() {
322 crate::trace::emit_trace(
323 &self.host.core.tracing.trace_sink,
324 &self.host.core.tracing.trace_context,
325 lash_trace::TraceContext::default()
326 .for_session(returned_turn.state.session_id.clone())
327 .for_turn_index(returned_turn.state.turn_index)
328 .for_turn(trace_turn_id.clone()),
329 lash_trace::TraceEvent::Custom {
330 name: "queued_work.completed".to_string(),
331 payload: queued_work_completion_trace_payload(&queued_work_completion_trace),
332 },
333 );
334 }
335 self.mark_phase_begin(RuntimeTurnPhase::PostPersistHooks);
336 self.emit_turn_persisted_event(&returned_turn, scoped_effect_controller, &trace_turn_id)
337 .await?;
338 self.mark_phase_end(RuntimeTurnPhase::PostPersistHooks);
339 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
340
341 self.emit_completed_turn_trace(
342 &returned_turn.state,
343 &returned_turn.outcome,
344 &trace_turn_id,
345 );
346 Ok(returned_turn)
347 }
348
349 fn emit_completed_turn_trace(
350 &self,
351 state: &SessionSnapshot,
352 outcome: &TurnOutcome,
353 trace_turn_id: &str,
354 ) {
355 if self.host.core.tracing.trace_sink.is_none() {
356 return;
357 }
358
359 let (status, done_reason, agent_frame_switch) = trace_fields_from_outcome(outcome);
360 crate::trace::emit_trace(
361 &self.host.core.tracing.trace_sink,
362 &self.host.core.tracing.trace_context,
363 lash_trace::TraceContext::default()
364 .for_session(state.session_id.clone())
365 .for_turn_index(state.turn_index)
366 .for_turn(trace_turn_id.to_string()),
367 lash_trace::TraceEvent::TurnCompleted {
368 status: status.to_string(),
369 done_reason: done_reason.to_string(),
370 agent_frame_switch,
371 },
372 );
373 }
374
375 async fn emit_turn_persisted_event(
376 &self,
377 returned_turn: &AssembledTurn,
378 scoped_effect_controller: &ScopedEffectController<'_>,
379 trace_turn_id: &str,
380 ) -> Result<(), RuntimeError> {
381 let Some(session) = self.session.as_ref() else {
382 return Ok(());
383 };
384 let Ok(manager) = self.runtime_session_services() else {
385 return Ok(());
386 };
387 let phase_turn_id = turn_phase_id(trace_turn_id, "turn-persisted");
388 let phase_controller = scoped_child_turn_controller(
389 scoped_effect_controller,
390 &self.state.session_id,
391 &phase_turn_id,
392 )?;
393 let direct_completions = manager.direct_completion_client(
394 RuntimeEffectControllerHandle::borrowed(phase_controller),
395 Some(phase_turn_id),
396 );
397
398 session
399 .plugins()
400 .emit_runtime_event_with_phase_probe(
401 crate::PluginLifecycleEvent::TurnPersisted(Box::new(
402 crate::SessionStateChangedContext {
403 session_id: self.state.session_id.clone(),
404 state: crate::SessionReadView::from_snapshot(&returned_turn.state),
405 sessions: manager.state_service(),
406 session_graph: manager.graph_service(),
407 direct_completions,
408 },
409 )),
410 self.turn_phase_probe.clone(),
411 )
412 .await;
413 Ok(())
414 }
415
416 pub async fn stream_turn(
418 &mut self,
419 input: TurnInput,
420 opts: TurnOptions<'_>,
421 ) -> Result<AssembledTurn, RuntimeError> {
422 let cancel = opts.cancel.clone();
423 let scoped_effect_controller = opts.scoped_effect_controller();
424 self.stream_turn_with_scoped_effect_controller_inner(
425 input,
426 opts.events_or_noop(),
427 opts.turn_events_or_noop(),
428 scoped_effect_controller,
429 cancel,
430 None,
431 )
432 .await
433 }
434
435 pub async fn stream_next_queued_work(
436 &mut self,
437 opts: TurnOptions<'_>,
438 ) -> Result<Option<AssembledTurn>, RuntimeError> {
439 self.stream_queued_work(opts, None).await
440 }
441
442 pub async fn stream_selected_queued_work(
443 &mut self,
444 opts: TurnOptions<'_>,
445 batch_ids: &[String],
446 ) -> Result<Option<AssembledTurn>, RuntimeError> {
447 self.stream_queued_work(opts, Some(batch_ids)).await
448 }
449
450 async fn stream_queued_work(
451 &mut self,
452 opts: TurnOptions<'_>,
453 selected_batch_ids: Option<&[String]>,
454 ) -> Result<Option<AssembledTurn>, RuntimeError> {
455 if self.drain_next_session_command().await?.is_some() {
456 return Ok(None);
457 }
458 let Some(store) = self
459 .session
460 .as_ref()
461 .and_then(|session| session.history_store())
462 else {
463 return Ok(None);
464 };
465 let claim = if let Some(batch_ids) = selected_batch_ids {
466 store
467 .claim_ready_queued_work_by_batch_ids(
468 &self.state.session_id,
469 &self.runtime_scope_id,
470 crate::QueuedWorkClaimBoundary::Idle,
471 crate::QUEUED_WORK_CLAIM_TTL_MS,
472 batch_ids,
473 )
474 .await
475 } else {
476 store
477 .claim_ready_queued_work(
478 &self.state.session_id,
479 &self.runtime_scope_id,
480 crate::QueuedWorkClaimBoundary::Idle,
481 crate::QUEUED_WORK_CLAIM_TTL_MS,
482 64,
483 )
484 .await
485 }
486 .map_err(|err| RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string()))?;
487 let Some(claim) = claim else {
488 return Ok(None);
489 };
490 let mut work = claim.materialize_for_turn();
491 let turn_id = work
492 .input
493 .trace_turn_id
494 .clone()
495 .or_else(|| Some(opts.effect_scope_id().to_owned()))
496 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
497 work.input.trace_turn_id = Some(turn_id.clone());
498 let causes = work.turn_causes.clone();
499 emit_queued_work_started_to_sink(
500 opts.turn_events_or_noop(),
501 crate::QueuedWorkClaimBoundary::Idle,
502 &claim,
503 causes.clone(),
504 )
505 .await;
506 crate::trace::emit_trace(
507 &self.host.core.tracing.trace_sink,
508 &self.host.core.tracing.trace_context,
509 lash_trace::TraceContext::default()
510 .for_session(self.state.session_id.clone())
511 .for_turn_index(self.state.turn_index + 1)
512 .for_turn(turn_id.clone()),
513 lash_trace::TraceEvent::Custom {
514 name: "queued_work.claimed".to_string(),
515 payload: queued_work_trace_payload(
516 crate::QueuedWorkClaimBoundary::Idle,
517 &claim,
518 &causes,
519 ),
520 },
521 );
522 let cancel = opts.cancel.clone();
523 let scoped_effect_controller = opts.scoped_effect_controller();
524 self.stream_turn_with_scoped_effect_controller_inner(
525 work.input,
526 opts.events_or_noop(),
527 opts.turn_events_or_noop(),
528 scoped_effect_controller,
529 cancel,
530 Some(claim),
531 )
532 .await
533 .map(Some)
534 }
535
536 fn ensure_durable_store_facets_for_scope(
544 &self,
545 scoped_effect_controller: &ScopedEffectController<'_>,
546 ) -> Result<(), RuntimeError> {
547 if scoped_effect_controller.controller().durability_tier() != crate::DurabilityTier::Durable
548 {
549 return Ok(());
550 }
551 if self
552 .host
553 .core
554 .durability
555 .attachment_store
556 .persistence()
557 .durability_tier()
558 != crate::DurabilityTier::Durable
559 {
560 return Err(RuntimeError::durable_store_required(
561 crate::DurableStoreFacet::AttachmentStore,
562 ));
563 }
564 if self
565 .host
566 .core
567 .durability
568 .lashlang_artifact_store
569 .durability_tier()
570 != crate::DurabilityTier::Durable
571 {
572 return Err(RuntimeError::durable_store_required(
573 crate::DurableStoreFacet::ArtifactStore,
574 ));
575 }
576 if let Some(store) = self
577 .session
578 .as_ref()
579 .and_then(|session| session.history_store())
580 && store.durability_tier() != crate::DurabilityTier::Durable
581 {
582 return Err(RuntimeError::durable_store_required(
583 crate::DurableStoreFacet::SessionStore,
584 ));
585 }
586 if let Some(process_registry) = self.host.process_registry.as_ref()
587 && process_registry.durability_tier() != crate::DurabilityTier::Durable
588 {
589 return Err(RuntimeError::durable_store_required(
590 crate::DurableStoreFacet::ProcessRegistry,
591 ));
592 }
593 Ok(())
594 }
595
596 async fn stream_turn_with_scoped_effect_controller_inner(
597 &mut self,
598 mut input: TurnInput,
599 events: &dyn EventSink,
600 turn_events: &dyn TurnActivitySink,
601 scoped_effect_controller: ScopedEffectController<'_>,
602 cancel: CancellationToken,
603 queued_claim: Option<crate::QueuedWorkClaim>,
604 ) -> Result<AssembledTurn, RuntimeError> {
605 if queued_claim.is_none() {
606 while self.drain_next_session_command().await?.is_some() {}
607 }
608 if let Some(input_turn_id) = input.trace_turn_id.as_deref()
609 && scoped_effect_controller
610 .effect_scope()
611 .validates_turn_trace_id()
612 && input_turn_id != scoped_effect_controller.scope_id()
613 {
614 return Err(RuntimeError::new(
615 RuntimeErrorCode::EffectScopeTurnIdMismatch,
616 format!(
617 "input trace_turn_id `{input_turn_id}` does not match effect scope id `{}`",
618 scoped_effect_controller.scope_id()
619 ),
620 ));
621 }
622 self.ensure_durable_store_facets_for_scope(&scoped_effect_controller)?;
623 input
624 .trace_turn_id
625 .get_or_insert_with(|| scoped_effect_controller.scope_id().to_string());
626 self.stream_turn_inner(
627 input.clone(),
628 events,
629 turn_events,
630 scoped_effect_controller,
631 cancel.clone(),
632 queued_claim,
633 )
634 .await
635 }
636
637 pub async fn stream_turn_with_agent_frames(
645 &mut self,
646 input: TurnInput,
647 opts: TurnOptions<'_>,
648 ) -> Result<AgentFrameRun, RuntimeError> {
649 let cancel = opts.cancel.clone();
650 let scoped_effect_controller = opts.scoped_effect_controller();
651 self.stream_turn_with_agent_frames_inner(
652 input,
653 opts.events_or_noop(),
654 opts.turn_events_or_noop(),
655 scoped_effect_controller,
656 cancel,
657 )
658 .await
659 }
660
661 async fn stream_turn_with_agent_frames_inner(
662 &mut self,
663 mut input: TurnInput,
664 events: &dyn EventSink,
665 turn_events: &dyn TurnActivitySink,
666 scoped_effect_controller: ScopedEffectController<'_>,
667 cancel: CancellationToken,
668 ) -> Result<AgentFrameRun, RuntimeError> {
669 if let Some(input_turn_id) = input.trace_turn_id.as_deref()
670 && scoped_effect_controller
671 .effect_scope()
672 .validates_turn_trace_id()
673 && input_turn_id != scoped_effect_controller.scope_id()
674 {
675 return Err(RuntimeError::new(
676 RuntimeErrorCode::EffectScopeTurnIdMismatch,
677 format!(
678 "input trace_turn_id `{input_turn_id}` does not match effect scope id `{}`",
679 scoped_effect_controller.scope_id()
680 ),
681 ));
682 }
683 let follow_protocol_turn_options = input.protocol_turn_options.clone();
684 let follow_turn_context = input.turn_context.clone();
685 let follow_trace_turn_id = input
686 .trace_turn_id
687 .clone()
688 .unwrap_or_else(|| scoped_effect_controller.scope_id().to_string());
689 input
690 .trace_turn_id
691 .get_or_insert(follow_trace_turn_id.clone());
692 let mut turns = Vec::new();
693 loop {
694 let turn_trace_turn_id = agent_frame_follow_turn_id(&follow_trace_turn_id, turns.len());
695 input.trace_turn_id = Some(turn_trace_turn_id.clone());
696 let turn_effect_controller = if turns.is_empty() {
697 scoped_effect_controller.clone()
698 } else {
699 ScopedEffectController::borrowed(
700 scoped_effect_controller.controller(),
701 EffectScope::turn(&self.state.session_id, &turn_trace_turn_id),
702 )?
703 };
704 let turn = self
705 .stream_turn_with_scoped_effect_controller_inner(
706 input,
707 events,
708 turn_events,
709 turn_effect_controller,
710 cancel.clone(),
711 None,
712 )
713 .await?;
714 let switched_frame = match &turn.outcome {
715 TurnOutcome::AgentFrameSwitch { frame_id, task } => {
716 Some((frame_id.clone(), task.clone()))
717 }
718 _ => None,
719 };
720 turns.push(turn);
721
722 let Some((_frame_id, task)) = switched_frame else {
723 return Ok(AgentFrameRun { turns });
724 };
725 input = turn_input_from_text(task);
726 input.protocol_turn_options = follow_protocol_turn_options.clone();
727 input.turn_context = follow_turn_context.clone();
728 }
729 }
730
731 async fn stream_turn_inner(
732 &mut self,
733 mut input: TurnInput,
734 events: &dyn EventSink,
735 turn_events: &dyn TurnActivitySink,
736 scoped_effect_controller: ScopedEffectController<'_>,
737 cancel: CancellationToken,
738 queued_claim: Option<crate::QueuedWorkClaim>,
739 ) -> Result<AssembledTurn, RuntimeError> {
740 self.refresh_session_graph_from_store()
741 .await
742 .map_err(session_head_refresh_error)?;
743 let input_trace_turn_id = input.trace_turn_id.clone();
744 let queued_turn_work = queued_claim
745 .as_ref()
746 .map(crate::QueuedWorkClaim::materialize_for_turn);
747 if let Some(work) = queued_turn_work.as_ref()
748 && input.items.is_empty()
749 && input.image_blobs.is_empty()
750 {
751 input = work.input.clone();
752 if input.trace_turn_id.is_none() {
753 input.trace_turn_id = input_trace_turn_id;
754 }
755 }
756 if self
757 .session
758 .as_ref()
759 .and_then(|session| session.history_store())
760 .is_some()
761 {
762 ensure_durable_effect_input(&input)?;
763 }
764 if let Some(extension) = &input.protocol_extension
765 && let Some(session) = self.session.as_ref()
766 {
767 let protocol_session = std::sync::Arc::clone(session.plugins().protocol_session());
768 protocol_session
769 .validate_turn_extension(extension)
770 .await
771 .map_err(|err| {
772 RuntimeError::new(RuntimeErrorCode::ProtocolTurnExtension, err.to_string())
773 })?;
774 }
775 let previous_prompt_usage = self.state.last_prompt_usage.clone();
776 let normalized = match self
777 .normalize_input_items(&input.items, &input.image_blobs)
778 .await
779 {
780 Ok(items) => items,
781 Err(e) => {
782 self.state.last_prompt_usage = None;
783 let mut assembler = TurnAssembler::default();
784 let error_event = SessionEvent::Error {
785 message: e.clone(),
786 envelope: Some(crate::session_model::ErrorEnvelope {
787 kind: "input_validation".to_string(),
788 code: Some("invalid_turn_input".to_string()),
789 terminal_reason: None,
790 user_message: e.clone(),
791 raw: None,
792 }),
793 };
794 assembler.push(&error_event);
795 emit_turn_activity_to_sink(
796 turn_events,
797 TurnActivity::independent(TurnEvent::Error { message: e }),
798 )
799 .await;
800 emit_session_event_to_sink(events, error_event).await;
801 let outcome_event = SessionEvent::TurnOutcome {
802 outcome: TurnOutcome::Stopped(TurnStop::InvalidInput),
803 };
804 assembler.push(&outcome_event);
805 emit_session_event_to_sink(events, outcome_event).await;
806 assembler.push(&SessionEvent::Done);
807 emit_session_event_to_sink(events, SessionEvent::Done).await;
808 return Ok(assembler.finish(
809 self.state.to_snapshot(),
810 false,
811 None,
812 &self.host.core.control.termination,
813 ));
814 }
815 };
816 let turn_index = self.state.turn_index + 1;
817 let trace_turn_id = input
818 .trace_turn_id
819 .clone()
820 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
821 if self.host.core.tracing.trace_sink.is_some() {
822 let mut trace_metadata = std::collections::BTreeMap::new();
823 trace_metadata.insert(
824 "input_item_count".to_string(),
825 serde_json::json!(normalized.len()),
826 );
827 crate::trace::emit_trace(
828 &self.host.core.tracing.trace_sink,
829 &self.host.core.tracing.trace_context,
830 lash_trace::TraceContext::default()
831 .for_session(self.state.session_id.clone())
832 .for_turn_index(turn_index)
833 .for_turn(trace_turn_id.clone()),
834 lash_trace::TraceEvent::TurnStarted {
835 metadata: trace_metadata,
836 },
837 );
838 }
839
840 let base_read_model = self.state.read_model();
841 let base_messages = base_read_model.messages;
842 let base_render_cache = base_read_model.prompt_render_cache;
843 let mut turn_delta = Vec::new();
844 let initial_turn_causes = queued_turn_work
845 .as_ref()
846 .map(|work| work.turn_causes.clone())
847 .unwrap_or_default();
848 turn_delta.extend(
849 initial_turn_causes
850 .iter()
851 .map(crate::TurnCause::to_event_message),
852 );
853
854 let user_id = fresh_message_id();
855 let mut user_parts: Vec<Part> = Vec::new();
856 for item in normalized {
857 match item {
858 NormalizedItem::Text(text) => {
859 if text.is_empty() {
860 continue;
861 }
862 user_parts.push(Part {
863 id: format!("{}.p{}", user_id, user_parts.len()),
864 kind: PartKind::Text,
865 content: text,
866 attachment: None,
867 tool_call_id: None,
868 tool_name: None,
869 tool_replay: None,
870 prune_state: PruneState::Intact,
871 reasoning_meta: None,
872 response_meta: None,
873 });
874 }
875 NormalizedItem::Image(reference) => {
876 user_parts.push(Part {
877 id: format!("{}.p{}", user_id, user_parts.len()),
878 kind: PartKind::Image,
879 content: String::new(),
880 attachment: Some(crate::session_model::message::PartAttachment {
881 reference,
882 }),
883 tool_call_id: None,
884 tool_name: None,
885 tool_replay: None,
886 prune_state: PruneState::Intact,
887 reasoning_meta: None,
888 response_meta: None,
889 });
890 }
891 }
892 }
893 if user_parts.is_empty() && initial_turn_causes.is_empty() {
894 user_parts.push(Part {
895 id: format!("{}.p0", user_id),
896 kind: PartKind::Text,
897 content: String::new(),
898 attachment: None,
899 tool_call_id: None,
900 tool_name: None,
901 tool_replay: None,
902 prune_state: PruneState::Intact,
903 reasoning_meta: None,
904 response_meta: None,
905 });
906 }
907 if !user_parts.is_empty() {
908 reassign_part_ids(&user_id, &mut user_parts);
909 turn_delta.push(Message {
910 id: user_id.clone(),
911 role: MessageRole::User,
912 parts: shared_parts(user_parts),
913 origin: None,
914 });
915 }
916
917 let manager = self
918 .runtime_session_services_for_turn(None)
919 .map_err(|err| {
920 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
921 })?;
922 let plugin_session = self
923 .session
924 .as_ref()
925 .map(|s| Arc::clone(s.plugins()))
926 .ok_or_else(|| {
927 RuntimeError::new(
928 RuntimeErrorCode::ContextPrepareTurn,
929 "runtime session not available",
930 )
931 })?;
932 let prepare_phase_turn_id = turn_phase_id(&trace_turn_id, "prepare-turn");
933 let prepare_phase_controller = scoped_child_turn_controller(
934 &scoped_effect_controller,
935 &self.state.session_id,
936 &prepare_phase_turn_id,
937 )?;
938 let turn_ctx = crate::TurnTransformContext {
939 session_id: self.state.session_id.clone(),
940 state: self.read_view(),
941 prompt_usage: previous_prompt_usage.clone(),
942 max_context_tokens: Some(LashRuntime::max_context_tokens(self)),
943 sessions: manager.state_service(),
944 session_lifecycle: manager.lifecycle_service(),
945 session_graph: manager.graph_service(),
946 scoped_effect_controller: scoped_effect_controller.clone(),
947 direct_completions: manager.direct_completion_client(
948 RuntimeEffectControllerHandle::borrowed(prepare_phase_controller),
949 Some(prepare_phase_turn_id),
950 ),
951 };
952 self.mark_phase_begin(RuntimeTurnPhase::ContextTransform);
953 let prepared_context = plugin_session
954 .prepare_turn_context(
955 &turn_ctx,
956 crate::session_model::context::PreparedContext {
957 messages: crate::MessageSequence::from_base_and_delta(
958 base_messages,
959 turn_delta,
960 )
961 .with_base_render_cache(base_render_cache),
962 ..Default::default()
963 },
964 self.turn_phase_probe.clone(),
965 )
966 .await
967 .map_err(|err| {
968 RuntimeError::new(RuntimeErrorCode::ContextPrepareTurn, err.to_string())
969 })?;
970 self.mark_phase_end(RuntimeTurnPhase::ContextTransform);
971 drop(turn_ctx);
976 let messages = prepared_context.messages;
977 if let Some(session) = self.session.as_mut() {
978 session
979 .set_context_surface(
980 prepared_context.tool_providers,
981 prepared_context.prompt_contributions,
982 prepared_context.include_base_tools,
983 )
984 .map_err(|err| {
985 RuntimeError::new(
986 RuntimeErrorCode::Other("session_tool_registry".to_string()),
987 err.to_string(),
988 )
989 })?;
990 }
991
992 self.state.last_prompt_usage = None;
993 self.stream_prepared_turn(
994 messages,
995 previous_prompt_usage,
996 input.protocol_turn_options.clone(),
997 input.protocol_extension.clone(),
998 input.turn_context.clone(),
999 initial_turn_causes,
1000 trace_turn_id,
1001 turn_index,
1002 events,
1003 turn_events,
1004 scoped_effect_controller,
1005 cancel,
1006 queued_claim,
1007 )
1008 .await
1009 }
1010
1011 pub async fn run_turn_assembled(
1013 &mut self,
1014 input: TurnInput,
1015 cancel: CancellationToken,
1016 scoped_effect_controller: ScopedEffectController<'_>,
1017 ) -> Result<AssembledTurn, RuntimeError> {
1018 self.stream_turn(input, TurnOptions::new(cancel, scoped_effect_controller))
1019 .await
1020 }
1021
1022 #[allow(clippy::too_many_arguments)]
1024 pub async fn stream_prepared_turn(
1025 &mut self,
1026 messages: crate::MessageSequence,
1027 _previous_prompt_usage: Option<PromptUsage>,
1028 protocol_turn_options: Option<crate::ProtocolTurnOptions>,
1029 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
1030 turn_context: crate::TurnContext,
1031 initial_turn_causes: Vec<crate::TurnCause>,
1032 trace_turn_id: String,
1033 turn_index: usize,
1034 events: &dyn EventSink,
1035 turn_events: &dyn TurnActivitySink,
1036 scoped_effect_controller: ScopedEffectController<'_>,
1037 cancel: CancellationToken,
1038 initial_queue_claim: Option<crate::QueuedWorkClaim>,
1039 ) -> Result<AssembledTurn, RuntimeError> {
1040 let (event_tx, mut event_rx) = mpsc::channel::<RuntimeStreamEvent>(100);
1041 let child_usage_event_relay = ChildUsageEventRelay::new(event_tx.clone());
1042 let mut turn_policy = self.state.effective_policy().clone();
1043 let turn_provider_override = turn_context.provider().cloned();
1044 if let Some(provider) = turn_provider_override.as_ref() {
1045 turn_policy.provider_id = provider.kind().to_string();
1046 }
1047 if let Some(model) = turn_context.model_spec() {
1048 turn_policy.model = model.clone();
1049 }
1050 let session_protocol_turn_options = self.state.effective_protocol_turn_options().clone();
1051 let effective_protocol_turn_options = protocol_turn_options
1052 .clone()
1053 .map(|options| session_protocol_turn_options.merged_with_override(&options))
1054 .unwrap_or(session_protocol_turn_options);
1055 let manager = self
1056 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
1057 .map_err(|err| {
1058 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1059 })?;
1060 let plugins = {
1061 let session = self
1062 .session
1063 .as_ref()
1064 .expect("lash runtime session must be available");
1065 Arc::clone(session.plugins())
1066 };
1067 let mut assembler = TurnAssembler::new();
1068 self.mark_phase_begin(RuntimeTurnPhase::BeforeTurnHooks);
1069 let prepared = {
1075 let prepare_turn = plugins.prepare_turn_with_phase_probe(
1076 PrepareTurnRequest {
1077 session_id: self.state.session_id.clone(),
1078 state: crate::SessionReadView::from_runtime_state(
1079 &self.state,
1080 turn_policy.clone(),
1081 effective_protocol_turn_options.clone(),
1082 ),
1083 messages,
1084 sessions: manager.state_service(),
1085 session_lifecycle: manager.lifecycle_service(),
1086 session_graph: manager.graph_service(),
1087 turn_context: turn_context.clone(),
1088 },
1089 self.turn_phase_probe.clone(),
1090 );
1091 tokio::pin!(prepare_turn);
1092
1093 loop {
1094 tokio::select! {
1095 prepared = &mut prepare_turn => {
1096 let prepared = prepared.map_err(|err| {
1097 RuntimeError::new(RuntimeErrorCode::PluginPrepareTurn, err.to_string())
1098 })?;
1099 self.mark_phase_end(RuntimeTurnPhase::BeforeTurnHooks);
1100 break prepared;
1101 }
1102 maybe_event = event_rx.recv() => {
1103 if let Some(event) = maybe_event {
1104 emit_runtime_stream_event_to_sinks(
1105 events,
1106 turn_events,
1107 event,
1108 &mut assembler,
1109 )
1110 .await;
1111 }
1112 }
1113 }
1114 }
1115 };
1116 for event in &prepared.events {
1117 assembler.push(event);
1118 }
1119 emit_session_events_to_sink(events, prepared.events).await;
1120 if let Some(abort) = prepared.abort {
1121 drop(event_tx);
1122
1123 let mut turn_pipeline = TurnBoundary::from_state(self.state.clone());
1124 turn_pipeline.apply_prepared_messages(&prepared.messages);
1125 let state = turn_pipeline.into_final_state();
1126 let issue = TurnIssue {
1127 kind: "plugin".to_string(),
1128 code: Some(abort.code),
1129 terminal_reason: None,
1130 message: abort.message.clone(),
1131 raw: None,
1132 };
1133 let error_event = SessionEvent::Error {
1134 message: abort.message,
1135 envelope: Some(crate::session_model::ErrorEnvelope {
1136 kind: "plugin".to_string(),
1137 code: issue.code.clone(),
1138 terminal_reason: None,
1139 user_message: issue.message.clone(),
1140 raw: None,
1141 }),
1142 };
1143 assembler.push(&error_event);
1144 emit_turn_activity_to_sink(
1145 turn_events,
1146 TurnActivity::independent(TurnEvent::Error {
1147 message: issue.message.clone(),
1148 }),
1149 )
1150 .await;
1151 emit_session_event_to_sink(events, error_event).await;
1152 let outcome_event = SessionEvent::TurnOutcome {
1153 outcome: TurnOutcome::Stopped(TurnStop::PluginAbort),
1154 };
1155 assembler.push(&outcome_event);
1156 emit_session_event_to_sink(events, outcome_event).await;
1157 assembler.push(&SessionEvent::Done);
1158 emit_session_event_to_sink(events, SessionEvent::Done).await;
1159 return Ok(assembler.finish(
1160 state.to_snapshot(),
1161 cancel.is_cancelled(),
1162 Some(issue),
1163 &self.host.core.control.termination,
1164 ));
1165 }
1166 let mut turn_pipeline = TurnBoundary::from_state(self.state.clone());
1167 let store = self
1168 .session
1169 .as_ref()
1170 .and_then(|session| session.history_store());
1171 let progress_store = if scoped_effect_controller.controller().durability_tier()
1175 == crate::DurabilityTier::Durable
1176 {
1177 None
1178 } else {
1179 store.as_ref().map(|store| store.as_ref())
1180 };
1181 turn_pipeline
1182 .prepared_checkpoint(
1183 progress_store,
1184 turn_policy.clone(),
1185 turn_index,
1186 &prepared.messages,
1187 self.session.as_mut(),
1188 )
1189 .await
1190 .map_err(|err| {
1191 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
1192 })?;
1193 let resolved_turn_policy = if let Some(provider) = turn_provider_override {
1194 RuntimeSessionPolicy::from_provider(turn_policy.clone(), provider)
1195 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
1196 } else {
1197 self.host
1198 .resolve_session_policy(&self.state.session_id, turn_policy.clone())
1199 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
1200 };
1201 let manager = self
1202 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
1203 .map_err(|err| {
1204 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1205 })?;
1206 let cancel_state = cancel.clone();
1207 let finish_scoped_effect_controller = scoped_effect_controller.clone();
1208 let session = self
1209 .session
1210 .take()
1211 .expect("lash runtime session must be available");
1212 let mut driver = RuntimeTurnDriver {
1213 session,
1214 policy: resolved_turn_policy,
1215 host: self.host.clone(),
1216 turn_id: scoped_effect_controller.scope_id().to_string(),
1217 scoped_effect_controller,
1218 session_id: self.state.session_id.clone(),
1219 turn_index,
1220 turn_pipeline,
1221 llm_stream_summaries: HashMap::new(),
1222 next_llm_ordinal: 0,
1223 session_services: manager,
1224 protocol_turn_options: effective_protocol_turn_options,
1225 protocol_extension,
1226 turn_context,
1227 turn_causes: initial_turn_causes,
1228 pending_queue_claims: initial_queue_claim.into_iter().collect(),
1229 checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer::default(),
1230 turn_phase_probe: self.turn_phase_probe.clone(),
1231 };
1232 let protocol_run_offset = 0;
1233 self.mark_phase_begin(RuntimeTurnPhase::EffectLoop);
1234 let run_result = drive_turn_to_completion(
1235 driver.run(prepared.messages, event_tx, cancel, protocol_run_offset),
1236 &mut event_rx,
1237 &mut assembler,
1238 &child_usage_event_relay,
1239 events,
1240 turn_events,
1241 )
1242 .await;
1243 let (new_messages, _new_protocol_iteration) = match run_result {
1244 Ok(result) => result,
1245 Err(err) => {
1246 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
1247 let RuntimeTurnDriver { session, .. } = driver;
1248 self.session = Some(session);
1249 return Err(err);
1250 }
1251 };
1252 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
1253 tracing::debug!(
1254 new_message_count = new_messages.len(),
1255 tool_call_count = assembler.tool_calls.len(),
1256 "runtime post-run_task"
1257 );
1258
1259 let RuntimeTurnDriver {
1260 session,
1261 policy,
1262 turn_pipeline,
1263 pending_queue_claims,
1264 ..
1265 } = driver;
1266 self.session = Some(session);
1267 self.finish_turn(
1268 TurnFinishInput {
1269 turn_pipeline,
1270 assembler,
1271 new_messages,
1272 policy,
1273 turn_index,
1274 queued_work_completions: pending_queue_claims
1275 .iter()
1276 .map(crate::QueuedWorkClaim::completion)
1277 .collect(),
1278 trace_turn_id,
1279 },
1280 events,
1281 &finish_scoped_effect_controller,
1282 &cancel_state,
1283 )
1284 .await
1285 }
1286 async fn normalize_input_items(
1287 &self,
1288 items: &[InputItem],
1289 image_blobs: &HashMap<String, Vec<u8>>,
1290 ) -> Result<Vec<NormalizedItem>, String> {
1291 normalize_input_items(
1292 items,
1293 image_blobs,
1294 self.host.core.durability.attachment_store.as_ref(),
1295 )
1296 .await
1297 }
1298}
1299
1300fn turn_input_from_text(text: String) -> TurnInput {
1301 TurnInput {
1302 items: vec![InputItem::Text { text }],
1303 image_blobs: HashMap::new(),
1304 protocol_turn_options: None,
1305 trace_turn_id: None,
1306 protocol_extension: None,
1307 turn_context: crate::TurnContext::default(),
1308 }
1309}
1310
1311fn agent_frame_follow_turn_id(root_turn_id: &str, completed_turn_count: usize) -> String {
1312 if completed_turn_count == 0 {
1313 root_turn_id.to_string()
1314 } else {
1315 format!("{root_turn_id}:agent-frame:{completed_turn_count}")
1316 }
1317}
1318
1319pub fn ensure_durable_effect_input(input: &TurnInput) -> Result<(), RuntimeError> {
1320 if input.protocol_extension.is_some() {
1321 return Err(RuntimeError::new(
1322 RuntimeErrorCode::DurableEffectLiveProtocolExtension,
1323 "durable effect hosts do not support live protocol_extension inputs; encode replayable data in protocol_turn_options or persisted plugin state",
1324 ));
1325 }
1326 input
1327 .turn_context
1328 .live_plugin_inputs()
1329 .durable_effect_rejection()?;
1330 Ok(())
1331}
1332
1333async fn emit_turn_activity_to_sink(events: &dyn TurnActivitySink, activity: TurnActivity) {
1334 if !events.is_noop() {
1335 events.emit(activity).await;
1336 }
1337}
1338
1339async fn drive_turn_to_completion<F>(
1349 run_future: F,
1350 event_rx: &mut mpsc::Receiver<RuntimeStreamEvent>,
1351 assembler: &mut TurnAssembler,
1352 child_usage_event_relay: &ChildUsageEventRelay,
1353 events: &dyn EventSink,
1354 turn_events: &dyn TurnActivitySink,
1355) -> Result<(crate::MessageSequence, usize), RuntimeError>
1356where
1357 F: std::future::Future<Output = Result<(crate::MessageSequence, usize), RuntimeError>>,
1358{
1359 let run_result = {
1360 tokio::pin!(run_future);
1361 loop {
1362 tokio::select! {
1363 maybe_event = event_rx.recv() => {
1364 if let Some(event) = maybe_event {
1365 emit_runtime_stream_event_to_sinks(
1366 events,
1367 turn_events,
1368 event,
1369 assembler,
1370 )
1371 .await;
1372 }
1373 }
1374 completed = &mut run_future => {
1375 child_usage_event_relay.clear();
1376 break completed;
1377 }
1378 }
1379 }
1380 };
1381 while let Some(event) = event_rx.recv().await {
1382 emit_runtime_stream_event_to_sinks(events, turn_events, event, assembler).await;
1383 }
1384 run_result
1385}
1386
1387async fn emit_runtime_stream_event_to_sinks(
1388 events: &dyn EventSink,
1389 turn_events: &dyn TurnActivitySink,
1390 event: RuntimeStreamEvent,
1391 assembler: &mut TurnAssembler,
1392) {
1393 match event {
1394 RuntimeStreamEvent::Session(event) => {
1395 assembler.push(&event);
1396 emit_session_event_to_sink(events, event).await;
1397 }
1398 RuntimeStreamEvent::Turn(activity) => {
1399 assembler.push_turn_activity(&activity);
1400 emit_turn_activity_to_sink(turn_events, activity).await;
1401 }
1402 }
1403}
1404
1405#[cfg(test)]
1406mod tests {
1407 use super::agent_frame_follow_turn_id;
1408
1409 #[test]
1410 fn agent_frame_follow_turn_ids_are_distinct_and_deterministic() {
1411 assert_eq!(agent_frame_follow_turn_id("root-turn", 0), "root-turn");
1412 assert_eq!(
1413 agent_frame_follow_turn_id("root-turn", 1),
1414 "root-turn:agent-frame:1"
1415 );
1416 assert_eq!(
1417 agent_frame_follow_turn_id("root-turn", 2),
1418 "root-turn:agent-frame:2"
1419 );
1420 }
1421}