1#[cfg(test)]
2use super::logical_turn::agent_frame_follow_turn_id;
3use super::logical_turn::{
4 LogicalTurnClaims, LogicalTurnStart, PhysicalTurnExecution, PreparedLogicalTurn,
5};
6use super::*;
7
8fn trace_fields_from_outcome(
9 outcome: &TurnOutcome,
10) -> (
11 &'static str,
12 &'static str,
13 Option<lash_trace::TraceAgentFrameSwitch>,
14) {
15 match outcome {
16 TurnOutcome::Finished(TurnFinish::AssistantMessage { .. }) => {
17 ("completed", "assistant_message", None)
18 }
19 TurnOutcome::Finished(TurnFinish::FinalValue { .. }) => ("completed", "final_value", None),
20 TurnOutcome::Finished(TurnFinish::ToolValue { .. }) => ("completed", "tool_value", None),
21 TurnOutcome::AgentFrameSwitch { frame_id, .. } => (
22 "completed",
23 "agent_frame_switch",
24 Some(lash_trace::TraceAgentFrameSwitch {
25 frame_id: frame_id.clone(),
26 }),
27 ),
28 TurnOutcome::Stopped(stop) => ("failed", trace_stop_reason(stop), None),
29 }
30}
31
32fn trace_stop_reason(stop: &TurnStop) -> &'static str {
33 match stop {
34 TurnStop::Cancelled => "cancelled",
35 TurnStop::Incomplete => "incomplete",
36 TurnStop::InvalidInput => "invalid_input",
37 TurnStop::MaxTurns => "max_turns",
38 TurnStop::ToolFailure => "tool_failure",
39 TurnStop::ProviderError => "provider_error",
40 TurnStop::PluginAbort => "plugin_abort",
41 TurnStop::RuntimeError => "runtime_error",
42 TurnStop::SubmittedError { .. } => "submitted_error",
43 TurnStop::ToolError { .. } => "tool_error",
44 }
45}
46
47fn session_head_refresh_error(err: SessionError) -> RuntimeError {
48 RuntimeError::new(
49 RuntimeErrorCode::Other("session_head_refresh".to_string()),
50 err.to_string(),
51 )
52}
53
54#[derive(Clone, Copy)]
55pub(super) enum SessionExecutionLeaseReleasePolicy {
56 KeepOnAgentFrameSwitch,
57}
58
59impl SessionExecutionLeaseReleasePolicy {
60 fn should_release(self, outcome: &TurnOutcome) -> bool {
61 match self {
62 Self::KeepOnAgentFrameSwitch => {
63 !matches!(outcome, TurnOutcome::AgentFrameSwitch { .. })
64 }
65 }
66 }
67}
68
69fn queued_work_payload_type(payload: &crate::QueuedWorkPayload) -> &'static str {
70 match payload {
71 crate::QueuedWorkPayload::ProcessWake { .. } => "process_wake",
72 crate::QueuedWorkPayload::AgentFrameTask { .. } => "agent_frame_task",
73 crate::QueuedWorkPayload::SessionCommand { command } => command.kind(),
74 }
75}
76
77fn queued_work_batch_ids(claim: &crate::QueuedWorkClaim) -> Vec<String> {
78 claim
79 .batches
80 .iter()
81 .map(|batch| batch.batch_id.clone())
82 .collect()
83}
84
85#[derive(Clone, Copy)]
95pub(super) struct TurnStopwatch {
96 started: std::time::Instant,
97 started_at_ms: u64,
98}
99
100impl TurnStopwatch {
101 pub(super) fn start(clock: &dyn crate::Clock) -> Self {
102 Self {
103 started: clock.now(),
104 started_at_ms: clock.timestamp_ms(),
105 }
106 }
107
108 pub(super) fn stamp(&self, turn: &mut AssembledTurn, clock: &dyn crate::Clock) {
109 turn.execution.started_at_ms = self.started_at_ms;
110 turn.execution.duration_ms = clock
111 .now()
112 .saturating_duration_since(self.started)
113 .as_millis() as u64;
114 }
115}
116
117fn turn_phase_id(parent_turn_id: &str, phase: &str) -> String {
118 format!("{parent_turn_id}:{phase}")
119}
120
121fn scoped_child_turn_controller<'run>(
122 scoped_effect_controller: &'run ScopedEffectController<'_>,
123 session_id: &str,
124 turn_id: &str,
125) -> Result<ScopedEffectController<'run>, RuntimeError> {
126 ScopedEffectController::borrowed(
127 scoped_effect_controller.controller(),
128 ExecutionScope::turn(session_id, turn_id),
129 )
130}
131
132pub(in crate::runtime) fn queued_work_trace_payload(
133 boundary: crate::QueuedWorkClaimBoundary,
134 claim: &crate::QueuedWorkClaim,
135 causes: &[crate::TurnCause],
136) -> serde_json::Value {
137 serde_json::json!({
138 "boundary": boundary,
139 "claim_id": claim.claim_id,
140 "owner_id": claim.owner.owner_id,
141 "incarnation_id": claim.owner.incarnation_id,
142 "batch_ids": queued_work_batch_ids(claim),
143 "payload_types": claim.batches.iter()
144 .flat_map(|batch| batch.items.iter())
145 .map(|item| queued_work_payload_type(&item.payload))
146 .collect::<Vec<_>>(),
147 "causes": causes,
148 })
149}
150
151pub(in crate::runtime) fn queued_work_completion_trace_payload(
152 completions: &[crate::QueuedWorkCompletion],
153) -> serde_json::Value {
154 serde_json::json!({
155 "claims": completions.iter().map(|completion| {
156 serde_json::json!({
157 "session_id": completion.session_id,
158 "claim_id": completion.claim_id,
159 "batch_ids": completion.batch_ids,
160 })
161 }).collect::<Vec<_>>(),
162 })
163}
164
165pub(in crate::runtime) fn turn_input_completion_trace_payload(
166 completions: &[crate::TurnInputCompletion],
167) -> serde_json::Value {
168 serde_json::json!({
169 "claims": completions.iter().map(|completion| {
170 serde_json::json!({
171 "session_id": completion.session_id,
172 "claim_id": completion.claim_id,
173 "input_ids": completion.input_ids,
174 })
175 }).collect::<Vec<_>>(),
176 })
177}
178
179async fn emit_queued_work_started_to_sink(
180 events: &dyn TurnActivitySink,
181 boundary: crate::QueuedWorkClaimBoundary,
182 claim: &crate::QueuedWorkClaim,
183 causes: Vec<crate::TurnCause>,
184) {
185 emit_turn_activity_to_sink(
186 events,
187 TurnActivity::independent(TurnEvent::QueuedWorkStarted {
188 boundary,
189 batch_ids: queued_work_batch_ids(claim),
190 causes,
191 }),
192 )
193 .await;
194}
195
196pub(in crate::runtime) async fn send_queued_work_started_event(
197 event_tx: &mpsc::Sender<RuntimeStreamEvent>,
198 boundary: crate::QueuedWorkClaimBoundary,
199 claim: &crate::QueuedWorkClaim,
200 causes: Vec<crate::TurnCause>,
201) {
202 send_turn_activity(
203 event_tx,
204 TurnActivityId::fresh(),
205 TurnEvent::QueuedWorkStarted {
206 boundary,
207 batch_ids: queued_work_batch_ids(claim),
208 causes,
209 },
210 )
211 .await;
212}
213
214struct TurnFinishInput {
215 turn_pipeline: TurnBoundary,
216 assembler: TurnAssembler,
217 new_messages: crate::MessageSequence,
218 policy: RuntimeSessionPolicy,
219 turn_index: usize,
220 queued_work_claims: Vec<crate::QueuedWorkClaim>,
221 turn_input_claims: Vec<crate::TurnInputClaim>,
222 trace_turn_id: String,
223}
224
225impl LashRuntime {
226 fn max_context_tokens(&self) -> usize {
227 self.state.effective_policy().context_window_tokens()
228 }
229
230 async fn claim_session_execution_lease(
231 &self,
232 cancel: CancellationToken,
233 busy_is_error: bool,
234 ) -> Result<Option<SessionExecutionLeaseGuard>, RuntimeError> {
235 let Some(store) = self
236 .session
237 .as_ref()
238 .and_then(|session| session.history_store())
239 else {
240 return Ok(None);
241 };
242 match SessionExecutionLeaseGuard::try_acquire(
243 store,
244 &self.state.session_id,
245 &self.runtime_lease_owner,
246 self.host.core.control.lease_timings,
247 Arc::clone(&self.host.core.clock),
248 cancel,
249 )
250 .await
251 .map_err(|err| RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string()))?
252 {
253 Some(lease) => Ok(Some(lease)),
254 None if busy_is_error => Err(RuntimeError::new(
255 RuntimeErrorCode::SessionExecutionBusy,
256 format!(
257 "session `{}` is already executing on another runtime owner",
258 self.state.session_id
259 ),
260 )),
261 None => Ok(None),
262 }
263 }
264
265 async fn settle_session_execution_lease<T>(
266 &self,
267 guard: Option<&SessionExecutionLeaseGuard>,
268 result: Result<T, RuntimeError>,
269 ) -> Result<T, RuntimeError> {
270 match result {
271 Ok(value) => {
272 if let Some(guard) = guard {
273 guard.release_if_live().await.map_err(|err| {
274 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
275 })?;
276 }
277 Ok(value)
278 }
279 Err(err) => {
280 if err.code != RuntimeErrorCode::StoreCommitFailed
281 && let Some(guard) = guard
282 && let Err(release_err) = guard.release_if_live().await
283 {
284 tracing::warn!(
285 error = %release_err,
286 "failed to release session execution lease after runtime error"
287 );
288 }
289 Err(err)
290 }
291 }
292 }
293
294 async fn ensure_session_execution_lease_live(
295 &self,
296 guard: Option<&SessionExecutionLeaseGuard>,
297 ) -> Result<(), RuntimeError> {
298 let Some(guard) = guard else {
299 return Ok(());
300 };
301 guard.refresh_or_mark_lost().await.map_err(|err| {
302 RuntimeError::new(
303 RuntimeErrorCode::SessionExecutionLeaseLost,
304 format!(
305 "session execution lease for session `{}` was lost before commit: {err}",
306 self.state.session_id
307 ),
308 )
309 })
310 }
311
312 async fn abandon_queued_work_claims_after_lease_loss(
319 &self,
320 err: &RuntimeError,
321 claims: &[crate::QueuedWorkClaim],
322 ) {
323 if err.code != RuntimeErrorCode::SessionExecutionLeaseLost || claims.is_empty() {
324 return;
325 }
326 let Some(store) = self
327 .session
328 .as_ref()
329 .and_then(|session| session.history_store())
330 else {
331 return;
332 };
333 for claim in claims {
334 if let Err(abandon_err) = store.abandon_queued_work_claim(claim).await {
335 tracing::warn!(
336 error = %abandon_err,
337 session_id = %claim.session_id,
338 claim_id = %claim.claim_id,
339 "failed to abandon queued work claim after session execution lease loss"
340 );
341 }
342 }
343 }
344
345 async fn abandon_turn_input_claims_after_lease_loss(
346 &self,
347 err: &RuntimeError,
348 claims: &[crate::TurnInputClaim],
349 ) {
350 if err.code != RuntimeErrorCode::SessionExecutionLeaseLost || claims.is_empty() {
351 return;
352 }
353 let Some(store) = self
354 .session
355 .as_ref()
356 .and_then(|session| session.history_store())
357 else {
358 return;
359 };
360 for claim in claims {
361 if let Err(abandon_err) = store.abandon_turn_input_claim(claim).await {
362 tracing::warn!(
363 error = %abandon_err,
364 session_id = %claim.session_id,
365 claim_id = %claim.claim_id,
366 "failed to abandon turn input claim after session execution lease loss"
367 );
368 }
369 }
370 }
371
372 #[doc(hidden)]
373 pub fn set_turn_phase_probe(&mut self, probe: Arc<dyn RuntimeTurnPhaseProbe>) {
374 self.turn_phase_probe = Some(probe);
375 }
376
377 fn mark_phase_begin(&self, phase: RuntimeTurnPhase) {
378 if let Some(probe) = self.turn_phase_probe.as_ref() {
379 probe.begin(phase);
380 }
381 }
382
383 fn mark_phase_end(&self, phase: RuntimeTurnPhase) {
384 if let Some(probe) = self.turn_phase_probe.as_ref() {
385 probe.end(phase);
386 }
387 }
388
389 async fn finish_turn(
390 &mut self,
391 finish: TurnFinishInput,
392 events: &dyn EventSink,
393 scoped_effect_controller: &ScopedEffectController<'_>,
394 cancel_state: &CancellationToken,
395 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
396 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
397 ) -> Result<PhysicalTurnExecution, RuntimeError> {
398 let TurnFinishInput {
399 mut turn_pipeline,
400 assembler,
401 new_messages,
402 policy,
403 turn_index,
404 queued_work_claims,
405 turn_input_claims,
406 trace_turn_id,
407 } = finish;
408 self.policy = self.state.effective_policy().clone();
409 turn_pipeline.state_mut().policy = self.policy.clone();
410 turn_pipeline.state_mut().turn_index = turn_index;
411
412 let mut turn_usage_delta = {
413 let mut ledger = self.shared_token_ledger.lock().expect("token ledger lock");
414 std::mem::take(&mut *ledger)
415 };
416 if assembler.token_usage.total() > 0 {
417 turn_usage_delta.push(TokenLedgerEntry {
418 source: "turn".to_string(),
419 model: policy.model.id.clone(),
420 usage: assembler.token_usage.clone(),
421 });
422 }
423 let turn_usage_delta = merge_usage_delta_entries(turn_usage_delta);
424
425 turn_pipeline.finalize_turn_read_state(new_messages, cancel_state.is_cancelled());
426 if assembler.token_usage.total() > 0 {
427 turn_pipeline.state_mut().token_usage = assembler.token_usage.clone();
428 }
429
430 let last_prompt_usage = assembler.last_llm_usage().and_then(normalize_prompt_usage);
431 turn_pipeline.state_mut().last_prompt_usage = last_prompt_usage;
432 let assembled_state = turn_pipeline.export_state_for_assembly();
433 let assembled = assembler.finish(
434 assembled_state,
435 cancel_state.is_cancelled(),
436 None,
437 &self.host.core.control.termination,
438 );
439
440 let Some(session) = self.session.as_ref() else {
441 self.state.apply_snapshot(&assembled.state);
442 self.emit_completed_turn_trace(&assembled.state, &assembled.outcome, &trace_turn_id);
443 return Ok(PhysicalTurnExecution {
444 turn: assembled,
445 enqueued_queue_batches: Vec::new(),
446 });
447 };
448 self.ensure_session_execution_lease_live(session_execution_lease)
449 .await?;
450
451 let plugins = Arc::clone(session.plugins());
452 let manager = match self.runtime_session_services_for_turn(None) {
453 Ok(manager) => manager,
454 Err(err) => {
455 return Err(RuntimeError::new(
456 RuntimeErrorCode::PluginSessionManager,
457 err.to_string(),
458 ));
459 }
460 };
461
462 self.mark_phase_begin(RuntimeTurnPhase::FinalizeTurn);
463 let finalized = match plugins
464 .finalize_turn_with_phase_probe(
465 assembled,
466 manager.state_service(),
467 manager.lifecycle_service(),
468 manager.graph_service(),
469 self.turn_phase_probe.clone(),
470 )
471 .await
472 {
473 Ok(finalized) => finalized,
474 Err(err) => {
475 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
476 return Err(RuntimeError::new(
477 RuntimeErrorCode::PluginFinalizeTurn,
478 err.to_string(),
479 ));
480 }
481 };
482 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
483 self.ensure_session_execution_lease_live(session_execution_lease)
484 .await?;
485
486 let mut returned_turn = finalized.turn;
487 let release_session_execution_lease =
488 session_execution_lease_release_policy.should_release(&returned_turn.outcome);
489 let commit_effects = LogicalTurnClaims::new(queued_work_claims, turn_input_claims)
490 .into_commit_effects(
491 &returned_turn.outcome,
492 &self.state.session_id,
493 &trace_turn_id,
494 Some(self.state.effective_protocol_turn_options().clone()),
495 );
496 self.mark_phase_begin(RuntimeTurnPhase::PersistTurn);
497 self.mark_phase_begin(RuntimeTurnPhase::FinalCommit);
498 let queued_work_completion_trace = commit_effects.completed_queue_claims.clone();
499 let turn_input_completion_trace = commit_effects.completed_turn_input_claims.clone();
500 let pending_attachment_ids = self
501 .host
502 .core
503 .durability
504 .attachment_store
505 .pending_manifest_commit_ids();
506 let enqueued_queue_batches = match turn_pipeline
507 .final_commit(
508 &mut returned_turn,
509 self.session.as_mut(),
510 &turn_usage_delta,
511 Some(&trace_turn_id),
512 commit_effects.completed_queue_claims,
513 commit_effects.completed_turn_input_claims,
514 commit_effects.enqueued_queue_batches,
515 cancel_state.is_cancelled().then(|| trace_turn_id.clone()),
516 pending_attachment_ids.clone(),
517 release_session_execution_lease
518 .then(|| session_execution_lease.map(SessionExecutionLeaseGuard::completion))
519 .flatten(),
520 )
521 .await
522 {
523 Ok(batches) => batches,
524 Err(err) => {
525 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
526 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
527 return Err(err);
528 }
529 };
530 if release_session_execution_lease && let Some(lease) = session_execution_lease {
531 lease.mark_released();
532 }
533 self.host
534 .core
535 .durability
536 .attachment_store
537 .mark_manifest_committed(&pending_attachment_ids);
538 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
539
540 emit_session_events_to_sink(events, finalized.events).await;
541 self.state = turn_pipeline.into_final_state();
542 if matches!(returned_turn.outcome, TurnOutcome::AgentFrameSwitch { .. })
543 && let Some(session) = self.session.as_mut()
544 {
545 let protocol_session = Arc::clone(session.plugins().protocol_session());
546 let session_id = self.state.session_id.clone();
547 protocol_session
548 .restore_session(
549 crate::plugin::ProtocolSessionContext::new(session, &session_id),
550 &self.state,
551 )
552 .await
553 .map_err(|err| {
554 RuntimeError::new(
555 RuntimeErrorCode::Other("protocol_restore_session".to_string()),
556 err.to_string(),
557 )
558 })?;
559 }
560 if !queued_work_completion_trace.is_empty() {
561 crate::trace::emit_trace(
562 &self.host.core.tracing.trace_sink,
563 &self.host.core.tracing.trace_context,
564 lash_trace::TraceContext::default()
565 .for_session(returned_turn.state.session_id.clone())
566 .for_turn_index(returned_turn.state.turn_index)
567 .for_turn(trace_turn_id.clone()),
568 lash_trace::TraceEvent::Custom {
569 name: "queued_work.completed".to_string(),
570 payload: queued_work_completion_trace_payload(&queued_work_completion_trace),
571 },
572 self.host.core.clock.as_ref(),
573 );
574 }
575 if !turn_input_completion_trace.is_empty() {
576 crate::trace::emit_trace(
577 &self.host.core.tracing.trace_sink,
578 &self.host.core.tracing.trace_context,
579 lash_trace::TraceContext::default()
580 .for_session(returned_turn.state.session_id.clone())
581 .for_turn_index(returned_turn.state.turn_index)
582 .for_turn(trace_turn_id.clone()),
583 lash_trace::TraceEvent::Custom {
584 name: "turn_input.completed".to_string(),
585 payload: turn_input_completion_trace_payload(&turn_input_completion_trace),
586 },
587 self.host.core.clock.as_ref(),
588 );
589 }
590 self.mark_phase_begin(RuntimeTurnPhase::PostPersistHooks);
591 self.emit_turn_persisted_event(&returned_turn, scoped_effect_controller, &trace_turn_id)
592 .await?;
593 self.mark_phase_end(RuntimeTurnPhase::PostPersistHooks);
594 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
595
596 self.emit_completed_turn_trace(
597 &returned_turn.state,
598 &returned_turn.outcome,
599 &trace_turn_id,
600 );
601 Ok(PhysicalTurnExecution {
602 turn: returned_turn,
603 enqueued_queue_batches,
604 })
605 }
606
607 fn emit_completed_turn_trace(
608 &self,
609 state: &SessionSnapshot,
610 outcome: &TurnOutcome,
611 trace_turn_id: &str,
612 ) {
613 if self.host.core.tracing.trace_sink.is_none() {
614 return;
615 }
616
617 let (status, done_reason, agent_frame_switch) = trace_fields_from_outcome(outcome);
618 crate::trace::emit_trace(
619 &self.host.core.tracing.trace_sink,
620 &self.host.core.tracing.trace_context,
621 lash_trace::TraceContext::default()
622 .for_session(state.session_id.clone())
623 .for_turn_index(state.turn_index)
624 .for_turn(trace_turn_id.to_string()),
625 lash_trace::TraceEvent::TurnCompleted {
626 status: status.to_string(),
627 done_reason: done_reason.to_string(),
628 agent_frame_switch,
629 },
630 self.host.core.clock.as_ref(),
631 );
632 }
633
634 #[allow(clippy::too_many_arguments)]
635 pub(super) async fn finish_logical_turn_error(
636 &mut self,
637 message: String,
638 trace_turn_id: String,
639 events: &dyn EventSink,
640 turn_events: &dyn TurnActivitySink,
641 scoped_effect_controller: ScopedEffectController<'_>,
642 cancel: CancellationToken,
643 claims: LogicalTurnClaims,
644 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
645 ) -> Result<PhysicalTurnExecution, RuntimeError> {
646 let mut assembler = TurnAssembler::default();
647 let error_event = SessionEvent::Error {
648 message: message.clone(),
649 envelope: Some(crate::session_model::ErrorEnvelope {
650 kind: "runtime".to_string(),
651 code: Some("agent_frame_switch_limit".to_string()),
652 terminal_reason: None,
653 user_message: message.clone(),
654 raw: None,
655 retryable: Some(false),
656 provider_failure_kind: None,
657 }),
658 };
659 assembler.push(&error_event);
660 emit_turn_activity_to_sink(
661 turn_events,
662 TurnActivity::independent(TurnEvent::Error {
663 message: message.clone(),
664 }),
665 )
666 .await;
667 emit_session_event_to_sink(events, error_event).await;
668 let outcome_event = SessionEvent::TurnOutcome {
669 outcome: TurnOutcome::Stopped(TurnStop::RuntimeError),
670 };
671 assembler.push(&outcome_event);
672 emit_session_event_to_sink(events, outcome_event).await;
673 assembler.push(&SessionEvent::Done);
674 emit_session_event_to_sink(events, SessionEvent::Done).await;
675
676 let messages = crate::MessageSequence::from_base(self.state.read_model().messages);
677 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
678 self.state.clone(),
679 Arc::clone(&self.host.core.clock),
680 )
681 .with_session_execution_lease(
682 session_execution_lease.map(SessionExecutionLeaseGuard::fence),
683 );
684 turn_pipeline.apply_prepared_messages(&messages);
685 self.finish_turn(
686 TurnFinishInput {
687 turn_pipeline,
688 assembler,
689 new_messages: messages,
690 policy: RuntimeSessionPolicy::new(
691 self.state.effective_policy().clone(),
692 Default::default(),
693 ),
694 turn_index: self.state.turn_index + 1,
695 queued_work_claims: claims.queued,
696 turn_input_claims: claims.turn_inputs,
697 trace_turn_id,
698 },
699 events,
700 &scoped_effect_controller,
701 &cancel,
702 session_execution_lease,
703 SessionExecutionLeaseReleasePolicy::KeepOnAgentFrameSwitch,
704 )
705 .await
706 }
707
708 async fn emit_turn_persisted_event(
709 &self,
710 returned_turn: &AssembledTurn,
711 scoped_effect_controller: &ScopedEffectController<'_>,
712 trace_turn_id: &str,
713 ) -> Result<(), RuntimeError> {
714 let Some(session) = self.session.as_ref() else {
715 return Ok(());
716 };
717 let Ok(manager) = self.runtime_session_services() else {
718 return Ok(());
719 };
720 let phase_turn_id = turn_phase_id(trace_turn_id, "turn-persisted");
721 let phase_controller = scoped_child_turn_controller(
722 scoped_effect_controller,
723 &self.state.session_id,
724 &phase_turn_id,
725 )?;
726 let direct_completions = manager.direct_completion_client(
727 RuntimeEffectControllerHandle::borrowed(phase_controller),
728 Some(phase_turn_id),
729 );
730
731 session
732 .plugins()
733 .emit_runtime_event_with_phase_probe(
734 crate::PluginLifecycleEvent::TurnPersisted(Box::new(
735 crate::SessionStateChangedContext {
736 session_id: self.state.session_id.clone(),
737 state: crate::SessionReadView::from_snapshot(&returned_turn.state),
738 sessions: manager.state_service(),
739 session_graph: manager.graph_service(),
740 direct_completions,
741 },
742 )),
743 self.turn_phase_probe.clone(),
744 )
745 .await;
746 Ok(())
747 }
748
749 pub async fn stream_turn(
751 &mut self,
752 input: TurnInput,
753 opts: TurnOptions<'_>,
754 ) -> Result<AssembledTurn, RuntimeError> {
755 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
756 let cancel = opts.cancel.clone();
757 let session_execution_lease = self
758 .claim_session_execution_lease(cancel.clone(), true)
759 .await?;
760 let scoped_effect_controller = opts.scoped_effect_controller();
761 let result = Box::pin(self.drive_logical_turn(
762 LogicalTurnStart::Input(input),
763 opts.events_or_noop(),
764 opts.turn_events_or_noop(),
765 scoped_effect_controller,
766 cancel,
767 LogicalTurnClaims::new(Vec::new(), Vec::new()),
768 session_execution_lease.as_ref(),
769 stopwatch,
770 ))
771 .await
772 .map(|run| {
773 run.into_final_turn()
774 .expect("logical turn always contains a terminal physical turn")
775 });
776 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
777 .await
778 }
779
780 pub async fn stream_next_queued_work(
781 &mut self,
782 opts: TurnOptions<'_>,
783 ) -> Result<Option<AssembledTurn>, RuntimeError> {
784 self.stream_queued_work(opts, None).await
785 }
786
787 pub async fn stream_selected_queued_work(
788 &mut self,
789 opts: TurnOptions<'_>,
790 batch_ids: &[String],
791 ) -> Result<Option<AssembledTurn>, RuntimeError> {
792 self.stream_queued_work(opts, Some(batch_ids)).await
793 }
794
795 async fn stream_queued_work(
796 &mut self,
797 opts: TurnOptions<'_>,
798 selected_batch_ids: Option<&[String]>,
799 ) -> Result<Option<AssembledTurn>, RuntimeError> {
800 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
801 let cancel = opts.cancel.clone();
802 let Some(session_execution_lease) = self
803 .claim_session_execution_lease(cancel.clone(), false)
804 .await?
805 else {
806 return Ok(None);
807 };
808 let session_execution_fence = session_execution_lease.fence();
809 let Some(store) = self
810 .session
811 .as_ref()
812 .and_then(|session| session.history_store())
813 else {
814 session_execution_lease
815 .release_if_live()
816 .await
817 .map_err(|err| {
818 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
819 })?;
820 return Ok(None);
821 };
822 let drain_commands_before_turn_input = if selected_batch_ids.is_some() {
823 true
824 } else {
825 self.session_commands_precede_pending_turn_input(store.as_ref())
826 .await?
827 };
828 if drain_commands_before_turn_input {
829 loop {
830 match self
831 .drain_next_session_command(&session_execution_fence)
832 .await
833 {
834 Ok(Some(_)) => {}
835 Ok(None) => break,
836 Err(err) => {
837 let _ = session_execution_lease.release_if_live().await;
838 return Err(err);
839 }
840 }
841 }
842 }
843 if selected_batch_ids.is_none() {
844 let input_claim = store
845 .claim_next_turn_inputs(
846 &self.state.session_id,
847 &session_execution_fence,
848 &self.runtime_lease_owner,
849 64,
850 )
851 .await
852 .map_err(super::runtime_error_from_store_commit)?;
853 if let Some(input_claim) = input_claim {
854 let mut input = input_claim.materialize_for_turn();
855 let turn_id = input
856 .trace_turn_id
857 .clone()
858 .or_else(|| Some(opts.execution_scope_id().to_owned()))
859 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
860 input.trace_turn_id = Some(turn_id.clone());
861 crate::trace::emit_trace(
862 &self.host.core.tracing.trace_sink,
863 &self.host.core.tracing.trace_context,
864 lash_trace::TraceContext::default()
865 .for_session(self.state.session_id.clone())
866 .for_turn_index(self.state.turn_index + 1)
867 .for_turn(turn_id.clone()),
868 lash_trace::TraceEvent::Custom {
869 name: "turn_input.claimed".to_string(),
870 payload: serde_json::json!({
871 "claim_id": &input_claim.claim_id,
872 "input_ids": input_claim.inputs.iter().map(|input| input.input_id.clone()).collect::<Vec<_>>(),
873 }),
874 },
875 self.host.core.clock.as_ref(),
876 );
877 let claim_for_abandon = input_claim.clone();
878 let scoped_effect_controller = opts.scoped_effect_controller();
879 let result = Box::pin(self.drive_logical_turn(
880 LogicalTurnStart::Input(input),
881 opts.events_or_noop(),
882 opts.turn_events_or_noop(),
883 scoped_effect_controller,
884 cancel,
885 LogicalTurnClaims::new(Vec::new(), vec![input_claim]),
886 Some(&session_execution_lease),
887 stopwatch,
888 ))
889 .await
890 .map(AgentFrameRun::into_final_turn);
891 if let Err(err) = &result {
892 self.abandon_turn_input_claims_after_lease_loss(
893 err,
894 std::slice::from_ref(&claim_for_abandon),
895 )
896 .await;
897 }
898 return self
899 .settle_session_execution_lease(Some(&session_execution_lease), result)
900 .await;
901 }
902 }
903 let claim = if let Some(batch_ids) = selected_batch_ids {
904 store
905 .claim_ready_queued_work_by_batch_ids(
906 &self.state.session_id,
907 &session_execution_fence,
908 &self.runtime_lease_owner,
909 crate::QueuedWorkClaimBoundary::Idle,
910 batch_ids,
911 )
912 .await
913 } else {
914 store
915 .claim_ready_queued_work(
916 &self.state.session_id,
917 &session_execution_fence,
918 &self.runtime_lease_owner,
919 crate::QueuedWorkClaimBoundary::Idle,
920 64,
921 )
922 .await
923 }
924 .map_err(super::runtime_error_from_store_commit)?;
925 let Some(claim) = claim else {
926 session_execution_lease
927 .release_if_live()
928 .await
929 .map_err(|err| {
930 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
931 })?;
932 return Ok(None);
933 };
934 let mut work = claim.materialize_for_turn();
935 let turn_id = work
936 .input
937 .trace_turn_id
938 .clone()
939 .or_else(|| Some(opts.execution_scope_id().to_owned()))
940 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
941 work.input.trace_turn_id = Some(turn_id.clone());
942 let causes = work.turn_causes.clone();
943 emit_queued_work_started_to_sink(
944 opts.turn_events_or_noop(),
945 crate::QueuedWorkClaimBoundary::Idle,
946 &claim,
947 causes.clone(),
948 )
949 .await;
950 crate::trace::emit_trace(
951 &self.host.core.tracing.trace_sink,
952 &self.host.core.tracing.trace_context,
953 lash_trace::TraceContext::default()
954 .for_session(self.state.session_id.clone())
955 .for_turn_index(self.state.turn_index + 1)
956 .for_turn(turn_id.clone()),
957 lash_trace::TraceEvent::Custom {
958 name: "queued_work.claimed".to_string(),
959 payload: queued_work_trace_payload(
960 crate::QueuedWorkClaimBoundary::Idle,
961 &claim,
962 &causes,
963 ),
964 },
965 self.host.core.clock.as_ref(),
966 );
967 let claim_for_abandon = claim.clone();
968 let scoped_effect_controller = opts.scoped_effect_controller();
969 let result = Box::pin(self.drive_logical_turn(
970 LogicalTurnStart::Input(work.input),
971 opts.events_or_noop(),
972 opts.turn_events_or_noop(),
973 scoped_effect_controller,
974 cancel,
975 LogicalTurnClaims::new(vec![claim], Vec::new()),
976 Some(&session_execution_lease),
977 stopwatch,
978 ))
979 .await
980 .map(AgentFrameRun::into_final_turn);
981 if let Err(err) = &result {
982 self.abandon_queued_work_claims_after_lease_loss(
983 err,
984 std::slice::from_ref(&claim_for_abandon),
985 )
986 .await;
987 }
988 self.settle_session_execution_lease(Some(&session_execution_lease), result)
989 .await
990 }
991
992 async fn session_commands_precede_pending_turn_input(
993 &self,
994 store: &dyn crate::RuntimePersistence,
995 ) -> Result<bool, RuntimeError> {
996 let pending_inputs = store
997 .list_pending_turn_inputs(&self.state.session_id)
998 .await
999 .map_err(super::runtime_error_from_store_commit)?;
1000 let earliest_input = pending_inputs
1001 .iter()
1002 .filter(|input| input.state.is_next_turn_pending())
1003 .min_by_key(|input| (input.enqueued_at_ms, input.enqueue_seq));
1004 let queued_work = store
1005 .list_pending_queued_work(&self.state.session_id)
1006 .await
1007 .map_err(super::runtime_error_from_store_commit)?;
1008 let earliest_command = queued_work
1009 .iter()
1010 .filter(|batch| batch.is_session_command_work())
1011 .min_by_key(|batch| (batch.enqueued_at_ms, batch.enqueue_seq));
1012 Ok(match (earliest_command, earliest_input) {
1013 (Some(command), Some(input)) => command.enqueued_at_ms < input.enqueued_at_ms,
1014 (Some(_), None) => true,
1015 _ => false,
1016 })
1017 }
1018
1019 fn ensure_durable_store_facets_for_scope(
1027 &self,
1028 scoped_effect_controller: &ScopedEffectController<'_>,
1029 ) -> Result<(), RuntimeError> {
1030 if scoped_effect_controller.controller().durability_tier() != crate::DurabilityTier::Durable
1031 {
1032 return Ok(());
1033 }
1034 if self
1035 .host
1036 .core
1037 .durability
1038 .attachment_store
1039 .persistence()
1040 .durability_tier()
1041 != crate::DurabilityTier::Durable
1042 {
1043 return Err(RuntimeError::durable_store_required(
1044 crate::DurableStoreFacet::AttachmentStore,
1045 ));
1046 }
1047 if self
1048 .host
1049 .core
1050 .durability
1051 .process_env_store
1052 .durability_tier()
1053 != crate::DurabilityTier::Durable
1054 {
1055 return Err(RuntimeError::durable_store_required(
1056 crate::DurableStoreFacet::ProcessEnvStore,
1057 ));
1058 }
1059 if let Some(store) = self
1060 .session
1061 .as_ref()
1062 .and_then(|session| session.history_store())
1063 && store.durability_tier() != crate::DurabilityTier::Durable
1064 {
1065 return Err(RuntimeError::durable_store_required(
1066 crate::DurableStoreFacet::SessionStore,
1067 ));
1068 }
1069 if let Some(process_registry) = self.host.process_registry.as_ref()
1070 && process_registry.durability_tier() != crate::DurabilityTier::Durable
1071 {
1072 return Err(RuntimeError::durable_store_required(
1073 crate::DurableStoreFacet::ProcessRegistry,
1074 ));
1075 }
1076 Ok(())
1077 }
1078
1079 #[allow(clippy::too_many_arguments)]
1080 pub(super) async fn stream_turn_with_scoped_effect_controller_inner(
1081 &mut self,
1082 mut input: TurnInput,
1083 events: &dyn EventSink,
1084 turn_events: &dyn TurnActivitySink,
1085 scoped_effect_controller: ScopedEffectController<'_>,
1086 cancel: CancellationToken,
1087 queued_claims: Vec<crate::QueuedWorkClaim>,
1088 turn_input_claims: Vec<crate::TurnInputClaim>,
1089 materialize_initial_claims: bool,
1090 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1091 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1092 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1093 if queued_claims.is_empty() && turn_input_claims.is_empty() {
1094 if let Some(lease) = session_execution_lease {
1095 while self
1096 .drain_next_session_command(&lease.fence())
1097 .await?
1098 .is_some()
1099 {}
1100 } else if self
1101 .session
1102 .as_ref()
1103 .and_then(|session| session.history_store())
1104 .is_some()
1105 {
1106 return Err(RuntimeError::new(
1107 RuntimeErrorCode::StoreCommitFailed,
1108 "session command drain requires a session execution lease",
1109 ));
1110 }
1111 }
1112 if let Some(input_turn_id) = input.trace_turn_id.as_deref()
1113 && scoped_effect_controller
1114 .execution_scope()
1115 .validates_turn_trace_id()
1116 && input_turn_id != scoped_effect_controller.scope_id()
1117 {
1118 return Err(RuntimeError::new(
1119 RuntimeErrorCode::ExecutionScopeTurnIdMismatch,
1120 format!(
1121 "input trace_turn_id `{input_turn_id}` does not match execution scope id `{}`",
1122 scoped_effect_controller.scope_id()
1123 ),
1124 ));
1125 }
1126 self.ensure_durable_store_facets_for_scope(&scoped_effect_controller)?;
1127 input
1128 .trace_turn_id
1129 .get_or_insert_with(|| scoped_effect_controller.scope_id().to_string());
1130 self.stream_turn_inner(
1131 input.clone(),
1132 events,
1133 turn_events,
1134 scoped_effect_controller,
1135 cancel.clone(),
1136 queued_claims,
1137 turn_input_claims,
1138 materialize_initial_claims,
1139 session_execution_lease,
1140 session_execution_lease_release_policy,
1141 )
1142 .await
1143 }
1144
1145 pub async fn stream_turn_with_agent_frames(
1154 &mut self,
1155 input: TurnInput,
1156 opts: TurnOptions<'_>,
1157 ) -> Result<AgentFrameRun, RuntimeError> {
1158 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
1159 let cancel = opts.cancel.clone();
1160 let session_execution_lease = self
1161 .claim_session_execution_lease(cancel.clone(), true)
1162 .await?;
1163 let scoped_effect_controller = opts.scoped_effect_controller();
1164 let result = Box::pin(self.drive_logical_turn(
1165 LogicalTurnStart::Input(input),
1166 opts.events_or_noop(),
1167 opts.turn_events_or_noop(),
1168 scoped_effect_controller,
1169 cancel,
1170 LogicalTurnClaims::new(Vec::new(), Vec::new()),
1171 session_execution_lease.as_ref(),
1172 stopwatch,
1173 ))
1174 .await;
1175 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
1176 .await
1177 }
1178
1179 #[allow(clippy::too_many_arguments)]
1180 async fn stream_turn_inner(
1181 &mut self,
1182 mut input: TurnInput,
1183 events: &dyn EventSink,
1184 turn_events: &dyn TurnActivitySink,
1185 scoped_effect_controller: ScopedEffectController<'_>,
1186 cancel: CancellationToken,
1187 queued_claims: Vec<crate::QueuedWorkClaim>,
1188 turn_input_claims: Vec<crate::TurnInputClaim>,
1189 materialize_initial_claims: bool,
1190 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1191 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1192 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1193 self.refresh_session_graph_from_store()
1194 .await
1195 .map_err(session_head_refresh_error)?;
1196 let input_trace_turn_id = input.trace_turn_id.clone();
1197 let queued_turn_work = materialize_initial_claims
1198 .then(|| queued_claims.first())
1199 .flatten()
1200 .map(crate::QueuedWorkClaim::materialize_for_turn);
1201 let pending_turn_input = materialize_initial_claims
1202 .then(|| turn_input_claims.first())
1203 .flatten()
1204 .map(crate::TurnInputClaim::materialize_for_turn);
1205 if let Some(work) = pending_turn_input.as_ref()
1206 && input.items.is_empty()
1207 && input.image_blobs.is_empty()
1208 {
1209 input = work.clone();
1210 if input.trace_turn_id.is_none() {
1211 input.trace_turn_id = input_trace_turn_id.clone();
1212 }
1213 }
1214 if let Some(work) = queued_turn_work.as_ref()
1215 && input.items.is_empty()
1216 && input.image_blobs.is_empty()
1217 {
1218 input = work.input.clone();
1219 if input.trace_turn_id.is_none() {
1220 input.trace_turn_id = input_trace_turn_id;
1221 }
1222 }
1223 if self
1224 .session
1225 .as_ref()
1226 .and_then(|session| session.history_store())
1227 .is_some()
1228 {
1229 ensure_durable_effect_input(&input)?;
1230 }
1231 if let Some(extension) = &input.protocol_extension
1232 && let Some(session) = self.session.as_ref()
1233 {
1234 let protocol_session = std::sync::Arc::clone(session.plugins().protocol_session());
1235 protocol_session
1236 .validate_turn_extension(extension)
1237 .await
1238 .map_err(|err| {
1239 RuntimeError::new(RuntimeErrorCode::ProtocolTurnExtension, err.to_string())
1240 })?;
1241 }
1242 let previous_prompt_usage = self.state.last_prompt_usage.clone();
1243 let normalized = match self
1244 .normalize_input_items(&input.items, &input.image_blobs)
1245 .await
1246 {
1247 Ok(items) => items,
1248 Err(e) => {
1249 self.state.last_prompt_usage = None;
1250 let mut assembler = TurnAssembler::default();
1251 let error_event = SessionEvent::Error {
1252 message: e.clone(),
1253 envelope: Some(crate::session_model::ErrorEnvelope {
1254 kind: "input_validation".to_string(),
1255 code: Some("invalid_turn_input".to_string()),
1256 terminal_reason: None,
1257 user_message: e.clone(),
1258 raw: None,
1259 retryable: Some(false),
1260 provider_failure_kind: None,
1261 }),
1262 };
1263 assembler.push(&error_event);
1264 emit_turn_activity_to_sink(
1265 turn_events,
1266 TurnActivity::independent(TurnEvent::Error { message: e }),
1267 )
1268 .await;
1269 emit_session_event_to_sink(events, error_event).await;
1270 let outcome_event = SessionEvent::TurnOutcome {
1271 outcome: TurnOutcome::Stopped(TurnStop::InvalidInput),
1272 };
1273 assembler.push(&outcome_event);
1274 emit_session_event_to_sink(events, outcome_event).await;
1275 assembler.push(&SessionEvent::Done);
1276 emit_session_event_to_sink(events, SessionEvent::Done).await;
1277 let turn_index = self.state.turn_index + 1;
1278 let trace_turn_id = input
1279 .trace_turn_id
1280 .clone()
1281 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1282 let messages = crate::MessageSequence::from_base(self.state.read_model().messages);
1283 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
1284 self.state.clone(),
1285 Arc::clone(&self.host.core.clock),
1286 )
1287 .with_session_execution_lease(
1288 session_execution_lease.map(SessionExecutionLeaseGuard::fence),
1289 );
1290 turn_pipeline.apply_prepared_messages(&messages);
1291 return self
1292 .finish_turn(
1293 TurnFinishInput {
1294 turn_pipeline,
1295 assembler,
1296 new_messages: messages,
1297 policy: RuntimeSessionPolicy::new(
1298 self.state.effective_policy().clone(),
1299 Default::default(),
1300 ),
1301 turn_index,
1302 queued_work_claims: queued_claims,
1303 turn_input_claims,
1304 trace_turn_id,
1305 },
1306 events,
1307 &scoped_effect_controller,
1308 &cancel,
1309 session_execution_lease,
1310 session_execution_lease_release_policy,
1311 )
1312 .await;
1313 }
1314 };
1315 let turn_index = self.state.turn_index + 1;
1316 let trace_turn_id = input
1317 .trace_turn_id
1318 .clone()
1319 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1320 if self.host.core.tracing.trace_sink.is_some() {
1321 let mut trace_metadata = std::collections::BTreeMap::new();
1322 trace_metadata.insert(
1323 "input_item_count".to_string(),
1324 serde_json::json!(normalized.len()),
1325 );
1326 crate::trace::emit_trace(
1327 &self.host.core.tracing.trace_sink,
1328 &self.host.core.tracing.trace_context,
1329 lash_trace::TraceContext::default()
1330 .for_session(self.state.session_id.clone())
1331 .for_turn_index(turn_index)
1332 .for_turn(trace_turn_id.clone()),
1333 lash_trace::TraceEvent::TurnStarted {
1334 metadata: trace_metadata,
1335 },
1336 self.host.core.clock.as_ref(),
1337 );
1338 }
1339
1340 let base_read_model = self.state.read_model();
1341 let base_messages = base_read_model.messages;
1342 let base_render_cache = base_read_model.prompt_render_cache;
1343 let mut turn_delta = Vec::new();
1344 let initial_turn_causes = queued_turn_work
1345 .as_ref()
1346 .map(|work| work.turn_causes.clone())
1347 .unwrap_or_default();
1348 turn_delta.extend(
1349 initial_turn_causes
1350 .iter()
1351 .map(crate::TurnCause::to_event_message),
1352 );
1353
1354 let user_id = fresh_message_id();
1355 let mut user_parts: Vec<Part> = Vec::new();
1356 for item in normalized {
1357 match item {
1358 NormalizedItem::Text(text) => {
1359 if text.is_empty() {
1360 continue;
1361 }
1362 user_parts.push(Part {
1363 id: format!("{}.p{}", user_id, user_parts.len()),
1364 kind: PartKind::Text,
1365 content: text,
1366 attachment: None,
1367 tool_call_id: None,
1368 tool_name: None,
1369 tool_replay: None,
1370 prune_state: PruneState::Intact,
1371 reasoning_meta: None,
1372 response_meta: None,
1373 });
1374 }
1375 NormalizedItem::Image(reference) => {
1376 user_parts.push(Part {
1377 id: format!("{}.p{}", user_id, user_parts.len()),
1378 kind: PartKind::Image,
1379 content: String::new(),
1380 attachment: Some(crate::session_model::message::PartAttachment {
1381 reference,
1382 }),
1383 tool_call_id: None,
1384 tool_name: None,
1385 tool_replay: None,
1386 prune_state: PruneState::Intact,
1387 reasoning_meta: None,
1388 response_meta: None,
1389 });
1390 }
1391 }
1392 }
1393 if user_parts.is_empty() && initial_turn_causes.is_empty() {
1394 user_parts.push(Part {
1395 id: format!("{}.p0", user_id),
1396 kind: PartKind::Text,
1397 content: String::new(),
1398 attachment: None,
1399 tool_call_id: None,
1400 tool_name: None,
1401 tool_replay: None,
1402 prune_state: PruneState::Intact,
1403 reasoning_meta: None,
1404 response_meta: None,
1405 });
1406 }
1407 if !user_parts.is_empty() {
1408 reassign_part_ids(&user_id, &mut user_parts);
1409 turn_delta.push(Message {
1410 id: user_id.clone(),
1411 role: MessageRole::User,
1412 parts: shared_parts(user_parts),
1413 origin: None,
1414 });
1415 }
1416
1417 let manager = self
1418 .runtime_session_services_for_turn(None)
1419 .map_err(|err| {
1420 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1421 })?;
1422 let plugin_session = self
1423 .session
1424 .as_ref()
1425 .map(|s| Arc::clone(s.plugins()))
1426 .ok_or_else(|| {
1427 RuntimeError::new(
1428 RuntimeErrorCode::ContextPrepareTurn,
1429 "runtime session not available",
1430 )
1431 })?;
1432 let prepare_phase_turn_id = turn_phase_id(&trace_turn_id, "prepare-turn");
1433 let prepare_phase_controller = scoped_child_turn_controller(
1434 &scoped_effect_controller,
1435 &self.state.session_id,
1436 &prepare_phase_turn_id,
1437 )?;
1438 let turn_ctx = crate::TurnTransformContext {
1439 session_id: self.state.session_id.clone(),
1440 state: self.read_view(),
1441 prompt_usage: previous_prompt_usage.clone(),
1442 max_context_tokens: Some(LashRuntime::max_context_tokens(self)),
1443 sessions: manager.state_service(),
1444 session_lifecycle: manager.lifecycle_service(),
1445 session_graph: manager.graph_service(),
1446 scoped_effect_controller: scoped_effect_controller.clone(),
1447 direct_completions: manager.direct_completion_client(
1448 RuntimeEffectControllerHandle::borrowed(prepare_phase_controller),
1449 Some(prepare_phase_turn_id),
1450 ),
1451 };
1452 self.mark_phase_begin(RuntimeTurnPhase::ContextTransform);
1453 let prepared_context = plugin_session
1454 .prepare_turn_context(
1455 &turn_ctx,
1456 crate::session_model::context::PreparedContext {
1457 messages: crate::MessageSequence::from_base_and_delta(
1458 base_messages,
1459 turn_delta,
1460 )
1461 .with_base_render_cache(base_render_cache),
1462 ..Default::default()
1463 },
1464 self.turn_phase_probe.clone(),
1465 )
1466 .await
1467 .map_err(|err| {
1468 RuntimeError::new(RuntimeErrorCode::ContextPrepareTurn, err.to_string())
1469 })?;
1470 self.mark_phase_end(RuntimeTurnPhase::ContextTransform);
1471 drop(turn_ctx);
1476 let messages = prepared_context.messages;
1477 if let Some(session) = self.session.as_mut() {
1478 session
1479 .set_context_overlay(
1480 prepared_context.tool_providers,
1481 prepared_context.prompt_contributions,
1482 prepared_context.include_base_tools,
1483 )
1484 .map_err(|err| {
1485 RuntimeError::new(
1486 RuntimeErrorCode::Other("session_tool_registry".to_string()),
1487 err.to_string(),
1488 )
1489 })?;
1490 }
1491
1492 self.state.last_prompt_usage = None;
1493 Box::pin(self.stream_prepared_turn_inner(
1494 messages,
1495 previous_prompt_usage,
1496 input.protocol_turn_options.clone(),
1497 input.protocol_extension.clone(),
1498 input.turn_context.clone(),
1499 initial_turn_causes,
1500 trace_turn_id,
1501 turn_index,
1502 events,
1503 turn_events,
1504 scoped_effect_controller,
1505 cancel,
1506 queued_claims,
1507 turn_input_claims,
1508 session_execution_lease,
1509 session_execution_lease_release_policy,
1510 ))
1511 .await
1512 }
1513
1514 pub async fn run_turn_assembled(
1516 &mut self,
1517 input: TurnInput,
1518 cancel: CancellationToken,
1519 scoped_effect_controller: ScopedEffectController<'_>,
1520 ) -> Result<AssembledTurn, RuntimeError> {
1521 self.stream_turn(input, TurnOptions::new(cancel, scoped_effect_controller))
1522 .await
1523 }
1524
1525 #[allow(clippy::too_many_arguments)]
1527 pub async fn stream_prepared_turn(
1528 &mut self,
1529 messages: crate::MessageSequence,
1530 previous_prompt_usage: Option<PromptUsage>,
1531 protocol_turn_options: Option<crate::ProtocolTurnOptions>,
1532 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
1533 turn_context: crate::TurnContext,
1534 initial_turn_causes: Vec<crate::TurnCause>,
1535 trace_turn_id: String,
1536 turn_index: usize,
1537 events: &dyn EventSink,
1538 turn_events: &dyn TurnActivitySink,
1539 scoped_effect_controller: ScopedEffectController<'_>,
1540 cancel: CancellationToken,
1541 initial_queue_claim: Option<crate::QueuedWorkClaim>,
1542 initial_turn_input_claim: Option<crate::TurnInputClaim>,
1543 ) -> Result<AssembledTurn, RuntimeError> {
1544 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
1545 let session_execution_lease = self
1546 .claim_session_execution_lease(cancel.clone(), true)
1547 .await?;
1548 let result = Box::pin(self.drive_logical_turn(
1549 LogicalTurnStart::Prepared(PreparedLogicalTurn {
1550 messages,
1551 previous_prompt_usage,
1552 protocol_turn_options,
1553 protocol_extension,
1554 turn_context,
1555 initial_turn_causes,
1556 trace_turn_id,
1557 turn_index,
1558 }),
1559 events,
1560 turn_events,
1561 scoped_effect_controller,
1562 cancel,
1563 LogicalTurnClaims::new(
1564 initial_queue_claim.into_iter().collect(),
1565 initial_turn_input_claim.into_iter().collect(),
1566 ),
1567 session_execution_lease.as_ref(),
1568 stopwatch,
1569 ))
1570 .await
1571 .map(|run| {
1572 run.into_final_turn()
1573 .expect("logical turn always contains a terminal physical turn")
1574 });
1575 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
1576 .await
1577 }
1578
1579 #[allow(clippy::too_many_arguments)]
1580 pub(super) async fn stream_prepared_turn_inner(
1581 &mut self,
1582 messages: crate::MessageSequence,
1583 _previous_prompt_usage: Option<PromptUsage>,
1584 protocol_turn_options: Option<crate::ProtocolTurnOptions>,
1585 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
1586 turn_context: crate::TurnContext,
1587 initial_turn_causes: Vec<crate::TurnCause>,
1588 trace_turn_id: String,
1589 turn_index: usize,
1590 events: &dyn EventSink,
1591 turn_events: &dyn TurnActivitySink,
1592 scoped_effect_controller: ScopedEffectController<'_>,
1593 cancel: CancellationToken,
1594 initial_queue_claims: Vec<crate::QueuedWorkClaim>,
1595 initial_turn_input_claims: Vec<crate::TurnInputClaim>,
1596 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1597 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1598 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1599 if session_execution_lease.is_none()
1600 && self
1601 .session
1602 .as_ref()
1603 .and_then(|session| session.history_store())
1604 .is_some()
1605 {
1606 return Err(RuntimeError::new(
1607 RuntimeErrorCode::StoreCommitFailed,
1608 "prepared turn requires a session execution lease",
1609 ));
1610 }
1611 let session_execution_fence =
1612 session_execution_lease.map(SessionExecutionLeaseGuard::fence);
1613 let (event_tx, mut event_rx) = mpsc::channel::<RuntimeStreamEvent>(100);
1614 let child_usage_event_relay = ChildUsageEventRelay::new(event_tx.clone());
1615 let mut turn_policy = self.state.effective_policy().clone();
1616 let turn_provider_override = turn_context.provider().cloned();
1617 if let Some(provider) = turn_provider_override.as_ref() {
1618 turn_policy.provider_id = provider.kind().to_string();
1619 }
1620 let session_protocol_turn_options = self.state.effective_protocol_turn_options().clone();
1621 let effective_protocol_turn_options = protocol_turn_options
1622 .clone()
1623 .map(|options| session_protocol_turn_options.merged_with_override(&options))
1624 .unwrap_or(session_protocol_turn_options);
1625 let manager = self
1626 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
1627 .map_err(|err| {
1628 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1629 })?;
1630 let plugins = {
1631 let session = self
1632 .session
1633 .as_ref()
1634 .expect("lash runtime session must be available");
1635 Arc::clone(session.plugins())
1636 };
1637 let mut assembler = TurnAssembler::new();
1638 self.mark_phase_begin(RuntimeTurnPhase::BeforeTurnHooks);
1639 let prepared = {
1645 let prepare_turn = plugins.prepare_turn_with_phase_probe(
1646 PrepareTurnRequest {
1647 session_id: self.state.session_id.clone(),
1648 state: crate::SessionReadView::from_runtime_state(
1649 &self.state,
1650 turn_policy.clone(),
1651 effective_protocol_turn_options.clone(),
1652 ),
1653 messages,
1654 sessions: manager.state_service(),
1655 session_lifecycle: manager.lifecycle_service(),
1656 session_graph: manager.graph_service(),
1657 turn_context: turn_context.clone(),
1658 },
1659 self.turn_phase_probe.clone(),
1660 );
1661 let mut prepare_turn = Box::pin(prepare_turn);
1662
1663 loop {
1664 tokio::select! {
1665 prepared = prepare_turn.as_mut() => {
1666 let prepared = prepared.map_err(|err| {
1667 RuntimeError::new(RuntimeErrorCode::PluginPrepareTurn, err.to_string())
1668 })?;
1669 self.mark_phase_end(RuntimeTurnPhase::BeforeTurnHooks);
1670 break prepared;
1671 }
1672 maybe_event = event_rx.recv() => {
1673 if let Some(event) = maybe_event {
1674 emit_runtime_stream_event_to_sinks(
1675 events,
1676 turn_events,
1677 event,
1678 &mut assembler,
1679 )
1680 .await;
1681 }
1682 }
1683 }
1684 }
1685 };
1686 for event in &prepared.events {
1687 assembler.push(event);
1688 }
1689 emit_session_events_to_sink(events, prepared.events).await;
1690 if let Some(abort) = prepared.abort {
1691 drop(event_tx);
1692
1693 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
1694 self.state.clone(),
1695 Arc::clone(&self.host.core.clock),
1696 )
1697 .with_session_execution_lease(session_execution_fence.clone());
1698 turn_pipeline.apply_prepared_messages(&prepared.messages);
1699 let issue = TurnIssue {
1700 kind: "plugin".to_string(),
1701 code: Some(abort.code),
1702 terminal_reason: None,
1703 message: abort.message.clone(),
1704 raw: None,
1705 retryable: None,
1706 provider_failure_kind: None,
1707 };
1708 let error_event = SessionEvent::Error {
1709 message: abort.message,
1710 envelope: Some(crate::session_model::ErrorEnvelope {
1711 kind: "plugin".to_string(),
1712 code: issue.code.clone(),
1713 terminal_reason: None,
1714 user_message: issue.message.clone(),
1715 raw: None,
1716 retryable: None,
1717 provider_failure_kind: None,
1718 }),
1719 };
1720 assembler.push(&error_event);
1721 emit_turn_activity_to_sink(
1722 turn_events,
1723 TurnActivity::independent(TurnEvent::Error {
1724 message: issue.message.clone(),
1725 }),
1726 )
1727 .await;
1728 emit_session_event_to_sink(events, error_event).await;
1729 let outcome_event = SessionEvent::TurnOutcome {
1730 outcome: TurnOutcome::Stopped(TurnStop::PluginAbort),
1731 };
1732 assembler.push(&outcome_event);
1733 emit_session_event_to_sink(events, outcome_event).await;
1734 assembler.push(&SessionEvent::Done);
1735 emit_session_event_to_sink(events, SessionEvent::Done).await;
1736 return self
1737 .finish_turn(
1738 TurnFinishInput {
1739 turn_pipeline,
1740 assembler,
1741 new_messages: prepared.messages,
1742 policy: RuntimeSessionPolicy::new(
1743 self.state.effective_policy().clone(),
1744 Default::default(),
1745 ),
1746 turn_index,
1747 queued_work_claims: initial_queue_claims,
1748 turn_input_claims: initial_turn_input_claims,
1749 trace_turn_id,
1750 },
1751 events,
1752 &scoped_effect_controller,
1753 &cancel,
1754 session_execution_lease,
1755 session_execution_lease_release_policy,
1756 )
1757 .await;
1758 }
1759 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
1760 self.state.clone(),
1761 Arc::clone(&self.host.core.clock),
1762 )
1763 .with_session_execution_lease(session_execution_fence.clone());
1764 let store = self
1765 .session
1766 .as_ref()
1767 .and_then(|session| session.history_store());
1768 let progress_store = if scoped_effect_controller.controller().durability_tier()
1772 == crate::DurabilityTier::Durable
1773 {
1774 None
1775 } else {
1776 store.as_ref().map(|store| store.as_ref())
1777 };
1778 turn_pipeline
1779 .prepared_checkpoint(
1780 progress_store,
1781 turn_policy.clone(),
1782 turn_index,
1783 &prepared.messages,
1784 self.session.as_mut(),
1785 )
1786 .await
1787 .map_err(super::runtime_error_from_store_commit)?;
1788 let resolved_turn_policy = if let Some(provider) = turn_provider_override {
1789 RuntimeSessionPolicy::from_provider(
1790 turn_policy.clone(),
1791 provider.with_clock(Arc::clone(&self.host.core.clock)),
1792 )
1793 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
1794 } else {
1795 self.host
1796 .resolve_session_policy(&self.state.session_id, turn_policy.clone())
1797 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
1798 };
1799 let manager = self
1800 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
1801 .map_err(|err| {
1802 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1803 })?;
1804 let cancel_state = cancel.clone();
1805 let finish_scoped_effect_controller = scoped_effect_controller.clone();
1806 let session = self
1807 .session
1808 .take()
1809 .expect("lash runtime session must be available");
1810 let mut driver = RuntimeTurnDriver {
1811 session,
1812 policy: resolved_turn_policy,
1813 host: self.host.clone(),
1814 turn_id: scoped_effect_controller.scope_id().to_string(),
1815 scoped_effect_controller,
1816 session_id: self.state.session_id.clone(),
1817 turn_index,
1818 turn_pipeline,
1819 llm_stream_summaries: HashMap::new(),
1820 llm_calls: Vec::new(),
1821 next_llm_ordinal: 0,
1822 session_services: manager,
1823 protocol_turn_options: effective_protocol_turn_options,
1824 protocol_extension,
1825 turn_context,
1826 turn_causes: initial_turn_causes,
1827 pending_queue_claims: initial_queue_claims,
1828 pending_turn_input_claims: initial_turn_input_claims,
1829 checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer::default(),
1830 session_execution_lease: session_execution_fence,
1831 runtime_lease_owner: self.runtime_lease_owner.clone(),
1832 turn_phase_probe: self.turn_phase_probe.clone(),
1833 };
1834 let protocol_run_offset = 0;
1835 self.mark_phase_begin(RuntimeTurnPhase::EffectLoop);
1836 let run_result = drive_turn_to_completion(
1837 driver.run(prepared.messages, event_tx, cancel, protocol_run_offset),
1838 &mut event_rx,
1839 &mut assembler,
1840 &child_usage_event_relay,
1841 events,
1842 turn_events,
1843 )
1844 .await;
1845 let (new_messages, _new_protocol_iteration) = match run_result {
1846 Ok(result) => result,
1847 Err(err) => {
1848 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
1849 let RuntimeTurnDriver {
1850 session,
1851 pending_queue_claims,
1852 pending_turn_input_claims,
1853 ..
1854 } = driver;
1855 self.session = Some(session);
1856 self.abandon_queued_work_claims_after_lease_loss(&err, &pending_queue_claims)
1857 .await;
1858 self.abandon_turn_input_claims_after_lease_loss(&err, &pending_turn_input_claims)
1859 .await;
1860 return Err(err);
1861 }
1862 };
1863 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
1864 tracing::debug!(
1865 new_message_count = new_messages.len(),
1866 tool_call_count = assembler.tool_calls.len(),
1867 "runtime post-run_task"
1868 );
1869
1870 let RuntimeTurnDriver {
1871 session,
1872 policy,
1873 turn_pipeline,
1874 llm_calls,
1875 pending_queue_claims,
1876 pending_turn_input_claims,
1877 ..
1878 } = driver;
1879 self.session = Some(session);
1880 let pending_queue_claims_for_abandon = pending_queue_claims.clone();
1881 let pending_turn_input_claims_for_abandon = pending_turn_input_claims.clone();
1882 let finish_result = self
1883 .finish_turn(
1884 TurnFinishInput {
1885 turn_pipeline,
1886 assembler: assembler.with_llm_calls(llm_calls),
1887 new_messages,
1888 policy,
1889 turn_index,
1890 queued_work_claims: pending_queue_claims,
1891 turn_input_claims: pending_turn_input_claims,
1892 trace_turn_id,
1893 },
1894 events,
1895 &finish_scoped_effect_controller,
1896 &cancel_state,
1897 session_execution_lease,
1898 session_execution_lease_release_policy,
1899 )
1900 .await;
1901 if let Err(err) = &finish_result {
1902 self.abandon_queued_work_claims_after_lease_loss(
1903 err,
1904 &pending_queue_claims_for_abandon,
1905 )
1906 .await;
1907 self.abandon_turn_input_claims_after_lease_loss(
1908 err,
1909 &pending_turn_input_claims_for_abandon,
1910 )
1911 .await;
1912 }
1913 finish_result
1914 }
1915 async fn normalize_input_items(
1916 &self,
1917 items: &[InputItem],
1918 image_blobs: &HashMap<String, Vec<u8>>,
1919 ) -> Result<Vec<NormalizedItem>, String> {
1920 normalize_input_items(
1921 items,
1922 image_blobs,
1923 self.host.core.durability.attachment_store.as_ref(),
1924 )
1925 .await
1926 }
1927}
1928
1929pub fn ensure_durable_effect_input(input: &TurnInput) -> Result<(), RuntimeError> {
1930 if input.protocol_extension.is_some() {
1931 return Err(RuntimeError::new(
1932 RuntimeErrorCode::DurableEffectLiveProtocolExtension,
1933 "durable effect hosts do not support live protocol_extension inputs; encode replayable data in protocol_turn_options or persisted plugin state",
1934 ));
1935 }
1936 input
1937 .turn_context
1938 .live_plugin_inputs()
1939 .durable_effect_rejection()?;
1940 Ok(())
1941}
1942
1943async fn emit_turn_activity_to_sink(events: &dyn TurnActivitySink, activity: TurnActivity) {
1944 if !events.is_noop() {
1945 events.emit(activity).await;
1946 }
1947}
1948
1949async fn drive_turn_to_completion<F>(
1959 run_future: F,
1960 event_rx: &mut mpsc::Receiver<RuntimeStreamEvent>,
1961 assembler: &mut TurnAssembler,
1962 child_usage_event_relay: &ChildUsageEventRelay,
1963 events: &dyn EventSink,
1964 turn_events: &dyn TurnActivitySink,
1965) -> Result<(crate::MessageSequence, usize), RuntimeError>
1966where
1967 F: std::future::Future<Output = Result<(crate::MessageSequence, usize), RuntimeError>>,
1968{
1969 let run_result = {
1970 let mut run_future = Box::pin(run_future);
1971 loop {
1972 tokio::select! {
1973 maybe_event = event_rx.recv() => {
1974 if let Some(event) = maybe_event {
1975 emit_runtime_stream_event_to_sinks(
1976 events,
1977 turn_events,
1978 event,
1979 assembler,
1980 )
1981 .await;
1982 }
1983 }
1984 completed = run_future.as_mut() => {
1985 child_usage_event_relay.clear();
1986 break completed;
1987 }
1988 }
1989 }
1990 };
1991 while let Some(event) = event_rx.recv().await {
1992 emit_runtime_stream_event_to_sinks(events, turn_events, event, assembler).await;
1993 }
1994 run_result
1995}
1996
1997async fn emit_runtime_stream_event_to_sinks(
1998 events: &dyn EventSink,
1999 turn_events: &dyn TurnActivitySink,
2000 event: RuntimeStreamEvent,
2001 assembler: &mut TurnAssembler,
2002) {
2003 match event {
2004 RuntimeStreamEvent::Session(event) => {
2005 assembler.push(&event);
2006 emit_session_event_to_sink(events, event).await;
2007 }
2008 RuntimeStreamEvent::Turn(activity) => {
2009 emit_turn_activity_to_sink(turn_events, activity).await;
2010 }
2011 }
2012}
2013
2014#[cfg(test)]
2015mod tests {
2016 use super::agent_frame_follow_turn_id;
2017
2018 #[test]
2019 fn agent_frame_follow_turn_ids_are_distinct_and_deterministic() {
2020 assert_eq!(agent_frame_follow_turn_id("root-turn", 0), "root-turn");
2021 assert_eq!(
2022 agent_frame_follow_turn_id("root-turn", 1),
2023 "root-turn:agent-frame:1"
2024 );
2025 assert_eq!(
2026 agent_frame_follow_turn_id("root-turn", 2),
2027 "root-turn:agent-frame:2"
2028 );
2029 }
2030}