1#[cfg(test)]
2use super::logical_turn::agent_frame_follow_turn_id;
3use super::logical_turn::{
4 LogicalTurnClaims, LogicalTurnStart, PhysicalTurnExecution, PreparedLogicalTurn,
5};
6use super::turn_control::ActiveTurnControl;
7use super::*;
8use std::pin::Pin;
9
10fn trace_fields_from_outcome(
11 outcome: &TurnOutcome,
12) -> (
13 &'static str,
14 &'static str,
15 Option<lash_trace::TraceAgentFrameSwitch>,
16) {
17 match outcome {
18 TurnOutcome::Finished(TurnFinish::AssistantMessage { .. }) => {
19 ("completed", "assistant_message", None)
20 }
21 TurnOutcome::Finished(TurnFinish::FinalValue { .. }) => ("completed", "final_value", None),
22 TurnOutcome::Finished(TurnFinish::ToolValue { .. }) => ("completed", "tool_value", None),
23 TurnOutcome::AgentFrameSwitch { frame_id, .. } => (
24 "completed",
25 "agent_frame_switch",
26 Some(lash_trace::TraceAgentFrameSwitch {
27 frame_id: frame_id.clone(),
28 }),
29 ),
30 TurnOutcome::Stopped(stop) => ("failed", trace_stop_reason(stop), None),
31 }
32}
33
34fn trace_stop_reason(stop: &TurnStop) -> &'static str {
35 match stop {
36 TurnStop::Cancelled => "cancelled",
37 TurnStop::Incomplete => "incomplete",
38 TurnStop::InvalidInput => "invalid_input",
39 TurnStop::MaxTurns => "max_turns",
40 TurnStop::ToolFailure => "tool_failure",
41 TurnStop::ProviderError => "provider_error",
42 TurnStop::PluginAbort => "plugin_abort",
43 TurnStop::RuntimeError => "runtime_error",
44 TurnStop::SubmittedError { .. } => "submitted_error",
45 TurnStop::ToolError { .. } => "tool_error",
46 }
47}
48
49fn session_head_refresh_error(err: SessionError) -> RuntimeError {
50 RuntimeError::new(
51 RuntimeErrorCode::Other("session_head_refresh".to_string()),
52 err.to_string(),
53 )
54}
55
56#[derive(Clone, Copy)]
57pub(super) enum SessionExecutionLeaseReleasePolicy {
58 KeepOnAgentFrameSwitch,
59}
60
61impl SessionExecutionLeaseReleasePolicy {
62 fn should_release(self, outcome: &TurnOutcome) -> bool {
63 match self {
64 Self::KeepOnAgentFrameSwitch => {
65 !matches!(outcome, TurnOutcome::AgentFrameSwitch { .. })
66 }
67 }
68 }
69}
70
71fn queued_work_payload_type(payload: &crate::QueuedWorkPayload) -> &'static str {
72 match payload {
73 crate::QueuedWorkPayload::ProcessWake { .. } => "process_wake",
74 crate::QueuedWorkPayload::AgentFrameTask { .. } => "agent_frame_task",
75 crate::QueuedWorkPayload::SessionCommand { command } => command.kind(),
76 }
77}
78
79fn queued_work_batch_ids(claim: &crate::QueuedWorkClaim) -> Vec<String> {
80 claim
81 .batches
82 .iter()
83 .map(|batch| batch.batch_id.clone())
84 .collect()
85}
86
87#[derive(Clone, Copy)]
97pub(super) struct TurnStopwatch {
98 started: std::time::Instant,
99 started_at_ms: u64,
100}
101
102impl TurnStopwatch {
103 pub(super) fn start(clock: &dyn crate::Clock) -> Self {
104 Self {
105 started: clock.now(),
106 started_at_ms: clock.timestamp_ms(),
107 }
108 }
109
110 pub(super) fn stamp(&self, turn: &mut AssembledTurn, clock: &dyn crate::Clock) {
111 turn.execution.started_at_ms = self.started_at_ms;
112 turn.execution.duration_ms = clock
113 .now()
114 .saturating_duration_since(self.started)
115 .as_millis() as u64;
116 }
117}
118
119fn turn_phase_id(parent_turn_id: &str, phase: &str) -> String {
120 format!("{parent_turn_id}:{phase}")
121}
122
123fn scoped_child_turn_controller<'run>(
124 scoped_effect_controller: &ScopedEffectController<'run>,
125 session_id: &str,
126 turn_id: &str,
127) -> Result<ScopedEffectController<'run>, RuntimeError> {
128 scoped_effect_controller.rescope(ExecutionScope::turn(session_id, turn_id))
129}
130
131fn turn_control_resolver<'a>(
138 effect_host: &'a dyn EffectHost,
139 scoped_effect_controller: &'a ScopedEffectController<'_>,
140) -> &'a dyn AwaitEventResolver {
141 if effect_host.durability_tier() == crate::DurabilityTier::Inline {
142 effect_host
143 } else {
144 scoped_effect_controller.controller()
145 }
146}
147
148pub(in crate::runtime) fn queued_work_trace_payload(
149 boundary: crate::QueuedWorkClaimBoundary,
150 claim: &crate::QueuedWorkClaim,
151 causes: &[crate::TurnCause],
152) -> serde_json::Value {
153 serde_json::json!({
154 "boundary": boundary,
155 "claim_id": claim.claim_id,
156 "owner_id": claim.owner.owner_id,
157 "incarnation_id": claim.owner.incarnation_id,
158 "batch_ids": queued_work_batch_ids(claim),
159 "payload_types": claim.batches.iter()
160 .flat_map(|batch| batch.items.iter())
161 .map(|item| queued_work_payload_type(&item.payload))
162 .collect::<Vec<_>>(),
163 "causes": causes,
164 })
165}
166
167pub(in crate::runtime) fn queued_work_completion_trace_payload(
168 completions: &[crate::QueuedWorkCompletion],
169) -> serde_json::Value {
170 serde_json::json!({
171 "claims": completions.iter().map(|completion| {
172 serde_json::json!({
173 "session_id": completion.session_id,
174 "claim_id": completion.claim_id,
175 "batch_ids": completion.batch_ids,
176 })
177 }).collect::<Vec<_>>(),
178 })
179}
180
181pub(in crate::runtime) fn turn_input_completion_trace_payload(
182 completions: &[crate::TurnInputCompletion],
183) -> serde_json::Value {
184 serde_json::json!({
185 "claims": completions.iter().map(|completion| {
186 serde_json::json!({
187 "session_id": completion.session_id,
188 "claim_id": completion.claim_id,
189 "input_ids": completion.input_ids,
190 })
191 }).collect::<Vec<_>>(),
192 })
193}
194
195async fn emit_queued_work_started_to_sink(
196 events: &dyn TurnActivitySink,
197 boundary: crate::QueuedWorkClaimBoundary,
198 claim: &crate::QueuedWorkClaim,
199 causes: Vec<crate::TurnCause>,
200) {
201 emit_turn_activity_to_sink(
202 events,
203 TurnActivity::independent(TurnEvent::QueuedWorkStarted {
204 boundary,
205 batch_ids: queued_work_batch_ids(claim),
206 causes,
207 }),
208 )
209 .await;
210}
211
212pub(in crate::runtime) async fn send_queued_work_started_event(
213 event_tx: &mpsc::Sender<RuntimeStreamEvent>,
214 boundary: crate::QueuedWorkClaimBoundary,
215 claim: &crate::QueuedWorkClaim,
216 causes: Vec<crate::TurnCause>,
217) {
218 send_turn_activity(
219 event_tx,
220 TurnActivityId::fresh(),
221 TurnEvent::QueuedWorkStarted {
222 boundary,
223 batch_ids: queued_work_batch_ids(claim),
224 causes,
225 },
226 )
227 .await;
228}
229
230struct TurnFinishInput {
231 turn_pipeline: TurnBoundary,
232 assembler: TurnAssembler,
233 new_messages: crate::MessageSequence,
234 policy: RuntimeSessionPolicy,
235 turn_index: usize,
236 trace_turn_id: String,
237}
238
239impl LashRuntime {
240 fn max_context_tokens(&self) -> usize {
241 self.state.effective_policy().context_window_tokens()
242 }
243
244 async fn claim_session_execution_lease(
245 &self,
246 cancel: CancellationToken,
247 busy_is_error: bool,
248 ) -> Result<Option<SessionExecutionLeaseGuard>, RuntimeError> {
249 let Some(store) = self
250 .session
251 .as_ref()
252 .and_then(|session| session.history_store())
253 else {
254 return Ok(None);
255 };
256 match SessionExecutionLeaseGuard::try_acquire(
257 store,
258 &self.state.session_id,
259 &self.runtime_lease_owner,
260 self.host.core.control.lease_timings,
261 Arc::clone(&self.host.core.clock),
262 cancel,
263 )
264 .await
265 .map_err(|err| RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string()))?
266 {
267 Some(lease) => Ok(Some(lease)),
268 None if busy_is_error => Err(RuntimeError::new(
269 RuntimeErrorCode::SessionExecutionBusy,
270 format!(
271 "session `{}` is already executing on another runtime owner",
272 self.state.session_id
273 ),
274 )),
275 None => Ok(None),
276 }
277 }
278
279 async fn settle_session_execution_lease<T>(
280 &self,
281 guard: Option<&SessionExecutionLeaseGuard>,
282 result: Result<T, RuntimeError>,
283 ) -> Result<T, RuntimeError> {
284 match result {
285 Ok(value) => {
286 if let Some(guard) = guard {
287 guard.release_if_live().await.map_err(|err| {
288 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
289 })?;
290 }
291 Ok(value)
292 }
293 Err(err) => {
294 if err.code != RuntimeErrorCode::StoreCommitFailed
295 && let Some(guard) = guard
296 && let Err(release_err) = guard.release_if_live().await
297 {
298 tracing::warn!(
299 error = %release_err,
300 "failed to release session execution lease after runtime error"
301 );
302 }
303 Err(err)
304 }
305 }
306 }
307
308 async fn abandon_queued_work_claims_after_lease_loss(
315 &self,
316 err: &RuntimeError,
317 claims: &[crate::QueuedWorkClaim],
318 ) {
319 if err.code != RuntimeErrorCode::SessionExecutionLeaseLost || claims.is_empty() {
320 return;
321 }
322 let Some(store) = self
323 .session
324 .as_ref()
325 .and_then(|session| session.history_store())
326 else {
327 return;
328 };
329 if let Err(abandon_err) = store.abandon_queued_work_claims(claims).await {
330 tracing::warn!(
331 error = %abandon_err,
332 claim_count = claims.len(),
333 "failed to abandon queued work claims after session execution lease loss"
334 );
335 }
336 }
337
338 async fn abandon_turn_input_claims_after_lease_loss(
339 &self,
340 err: &RuntimeError,
341 claims: &[crate::TurnInputClaim],
342 ) {
343 if err.code != RuntimeErrorCode::SessionExecutionLeaseLost || claims.is_empty() {
344 return;
345 }
346 let Some(store) = self
347 .session
348 .as_ref()
349 .and_then(|session| session.history_store())
350 else {
351 return;
352 };
353 if let Err(abandon_err) = store.abandon_turn_input_claims(claims).await {
354 tracing::warn!(
355 error = %abandon_err,
356 claim_count = claims.len(),
357 "failed to abandon turn input claims after session execution lease loss"
358 );
359 }
360 }
361
362 #[doc(hidden)]
363 pub fn set_turn_phase_probe(&mut self, probe: Arc<dyn RuntimeTurnPhaseProbe>) {
364 self.turn_phase_probe = Some(probe);
365 }
366
367 fn mark_phase_begin(&self, phase: RuntimeTurnPhase) {
368 if let Some(probe) = self.turn_phase_probe.as_ref() {
369 probe.begin(phase);
370 }
371 }
372
373 fn mark_phase_end(&self, phase: RuntimeTurnPhase) {
374 if let Some(probe) = self.turn_phase_probe.as_ref() {
375 probe.end(phase);
376 }
377 }
378
379 #[allow(clippy::too_many_arguments)]
380 async fn finish_turn(
381 &mut self,
382 finish: TurnFinishInput,
383 claims: &LogicalTurnClaims,
384 events: &dyn EventSink,
385 scoped_effect_controller: &ScopedEffectController<'_>,
386 cancel_state: &CancellationToken,
387 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
388 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
389 turn_control: &ActiveTurnControl,
390 ) -> Result<PhysicalTurnExecution, RuntimeError> {
391 let turn_control_host = Arc::clone(&self.host.core.control.effect_host);
392 let turn_control_resolver =
393 turn_control_resolver(turn_control_host.as_ref(), scoped_effect_controller);
394 let TurnFinishInput {
395 mut turn_pipeline,
396 assembler,
397 new_messages,
398 policy,
399 turn_index,
400 trace_turn_id,
401 } = finish;
402 self.policy = self.state.effective_policy().clone();
403 turn_pipeline.state_mut().policy = self.policy.clone();
404 turn_pipeline.state_mut().turn_index = turn_index;
405
406 let mut turn_usage_delta = {
407 let mut ledger = self.shared_token_ledger.lock().expect("token ledger lock");
408 std::mem::take(&mut *ledger)
409 };
410 if assembler.token_usage.total() > 0 {
411 turn_usage_delta.push(TokenLedgerEntry {
412 source: "turn".to_string(),
413 model: policy.model.id.clone(),
414 usage: assembler.token_usage.clone(),
415 });
416 }
417 let turn_usage_delta = merge_usage_delta_entries(turn_usage_delta);
418
419 let assembled_cancelled = matches!(
420 assembler.outcome,
421 Some(TurnOutcome::Stopped(TurnStop::Cancelled))
422 );
423 let lease_was_lost = session_execution_lease.is_some_and(|lease| lease.is_lost());
424 let cancellation = turn_control
425 .settle_before_commit(
426 turn_control_resolver,
427 assembled_cancelled || (cancel_state.is_cancelled() && !lease_was_lost),
428 )
429 .await?;
430 if cancellation.is_some() {
431 cancel_state.cancel();
432 }
433 let interrupted = cancellation.is_some();
439
440 turn_pipeline.finalize_turn_read_state(new_messages, interrupted);
441 if assembler.token_usage.total() > 0 {
442 turn_pipeline.state_mut().token_usage = assembler.token_usage.clone();
443 }
444
445 let last_prompt_usage = assembler.last_llm_usage().and_then(normalize_prompt_usage);
446 turn_pipeline.state_mut().last_prompt_usage = last_prompt_usage;
447 let assembled_state = turn_pipeline.export_state_for_assembly();
448 let mut assembled = assembler.finish(
449 assembled_state,
450 interrupted,
451 None,
452 &self.host.core.control.termination,
453 );
454 assembled.cancellation = cancellation;
455
456 let Some(session) = self.session.as_ref() else {
457 self.state.apply_snapshot(&assembled.state);
458 self.emit_completed_turn_trace(&assembled.state, &assembled.outcome, &trace_turn_id);
459 publish_terminal_after_commit(
460 turn_control,
461 turn_control_resolver,
462 &TurnTerminal::Committed {
463 outcome: assembled.outcome.clone(),
464 cancellation: assembled.cancellation.clone(),
465 session_revision: None,
466 },
467 &self.state.session_id,
468 &trace_turn_id,
469 )
470 .await;
471 return Ok(PhysicalTurnExecution {
472 turn: assembled,
473 enqueued_queue_batches: Vec::new(),
474 });
475 };
476
477 let plugins = Arc::clone(session.plugins());
478 let manager = match self.runtime_session_services_for_turn(None) {
479 Ok(manager) => manager,
480 Err(err) => {
481 return Err(RuntimeError::new(
482 RuntimeErrorCode::PluginSessionManager,
483 err.to_string(),
484 ));
485 }
486 };
487
488 self.mark_phase_begin(RuntimeTurnPhase::FinalizeTurn);
489 let finalized = match plugins
490 .finalize_turn_with_phase_probe(
491 assembled,
492 manager.state_service(),
493 manager.lifecycle_service(),
494 manager.graph_service(),
495 self.turn_phase_probe.clone(),
496 )
497 .await
498 {
499 Ok(finalized) => finalized,
500 Err(err) => {
501 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
502 return Err(RuntimeError::new(
503 RuntimeErrorCode::PluginFinalizeTurn,
504 err.to_string(),
505 ));
506 }
507 };
508 self.mark_phase_end(RuntimeTurnPhase::FinalizeTurn);
509 let mut returned_turn = finalized.turn;
510 if returned_turn.cancellation.is_some()
511 && !matches!(
512 returned_turn.outcome,
513 TurnOutcome::Stopped(TurnStop::Cancelled)
514 )
515 {
516 returned_turn.outcome = TurnOutcome::Stopped(TurnStop::Cancelled);
517 }
518 if matches!(
519 returned_turn.outcome,
520 TurnOutcome::Stopped(TurnStop::Cancelled)
521 ) && returned_turn.cancellation.is_none()
522 {
523 return Err(RuntimeError::new(
524 "turn_cancellation_evidence_missing",
525 "cancelled turns must carry cancellation evidence",
526 ));
527 }
528 let release_session_execution_lease =
529 session_execution_lease_release_policy.should_release(&returned_turn.outcome);
530 let commit_effects = claims.commit_effects(
531 &returned_turn.outcome,
532 &self.state.session_id,
533 &trace_turn_id,
534 Some(self.state.effective_protocol_turn_options().clone()),
535 );
536 self.mark_phase_begin(RuntimeTurnPhase::PersistTurn);
537 self.mark_phase_begin(RuntimeTurnPhase::FinalCommit);
538 let queued_work_completion_trace = commit_effects.completed_queue_claims.clone();
539 let turn_input_completion_trace = commit_effects.completed_turn_input_claims.clone();
540 let enqueued_queue_batches = match turn_pipeline
541 .final_commit(
542 &mut returned_turn,
543 self.session.as_mut(),
544 &turn_usage_delta,
545 Some(&trace_turn_id),
546 commit_effects.originating_queue_claims,
547 commit_effects.originating_turn_input_claims,
548 commit_effects.completed_queue_claims,
549 commit_effects.completed_turn_input_claims,
550 commit_effects.enqueued_queue_batches,
551 Some(trace_turn_id.clone()),
557 release_session_execution_lease
558 .then(|| session_execution_lease.map(SessionExecutionLeaseGuard::completion))
559 .flatten(),
560 )
561 .await
562 {
563 Ok(batches) => batches,
564 Err(err) => {
565 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
566 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
567 return Err(err);
568 }
569 };
570 if release_session_execution_lease && let Some(lease) = session_execution_lease {
571 lease.mark_released();
572 }
573 self.last_committed_lease_continuity = if release_session_execution_lease {
574 None
575 } else {
576 session_execution_lease.and_then(SessionExecutionLeaseGuard::continuity)
577 };
578 self.mark_phase_end(RuntimeTurnPhase::FinalCommit);
579
580 emit_session_events_to_sink(events, finalized.events).await;
581 self.state = turn_pipeline.into_final_state();
582 publish_terminal_after_commit(
583 turn_control,
584 turn_control_resolver,
585 &TurnTerminal::Committed {
586 outcome: returned_turn.outcome.clone(),
587 cancellation: returned_turn.cancellation.clone(),
588 session_revision: None,
589 },
590 &self.state.session_id,
591 &trace_turn_id,
592 )
593 .await;
594 if matches!(returned_turn.outcome, TurnOutcome::AgentFrameSwitch { .. })
595 && let Some(session) = self.session.as_mut()
596 {
597 let protocol_session = Arc::clone(session.plugins().protocol_session());
598 let session_id = self.state.session_id.clone();
599 protocol_session
600 .restore_session(
601 crate::plugin::ProtocolSessionContext::new(session, &session_id),
602 &self.state,
603 )
604 .await
605 .map_err(|err| {
606 RuntimeError::new(
607 RuntimeErrorCode::Other("protocol_restore_session".to_string()),
608 err.to_string(),
609 )
610 })?;
611 }
612 if !queued_work_completion_trace.is_empty() {
613 crate::trace::emit_trace(
614 &self.host.core.tracing.trace_sink,
615 &self.host.core.tracing.trace_context,
616 lash_trace::TraceContext::default()
617 .for_session(returned_turn.state.session_id.clone())
618 .for_turn_index(returned_turn.state.turn_index)
619 .for_turn(trace_turn_id.clone()),
620 lash_trace::TraceEvent::Custom {
621 name: "queued_work.completed".to_string(),
622 payload: queued_work_completion_trace_payload(&queued_work_completion_trace),
623 },
624 self.host.core.clock.as_ref(),
625 );
626 }
627 if !turn_input_completion_trace.is_empty() {
628 crate::trace::emit_trace(
629 &self.host.core.tracing.trace_sink,
630 &self.host.core.tracing.trace_context,
631 lash_trace::TraceContext::default()
632 .for_session(returned_turn.state.session_id.clone())
633 .for_turn_index(returned_turn.state.turn_index)
634 .for_turn(trace_turn_id.clone()),
635 lash_trace::TraceEvent::Custom {
636 name: "turn_input.completed".to_string(),
637 payload: turn_input_completion_trace_payload(&turn_input_completion_trace),
638 },
639 self.host.core.clock.as_ref(),
640 );
641 }
642 self.mark_phase_begin(RuntimeTurnPhase::PostPersistHooks);
643 self.emit_turn_persisted_event(&returned_turn, scoped_effect_controller, &trace_turn_id)
644 .await?;
645 self.mark_phase_end(RuntimeTurnPhase::PostPersistHooks);
646 self.mark_phase_end(RuntimeTurnPhase::PersistTurn);
647
648 self.emit_completed_turn_trace(
649 &returned_turn.state,
650 &returned_turn.outcome,
651 &trace_turn_id,
652 );
653 Ok(PhysicalTurnExecution {
654 turn: returned_turn,
655 enqueued_queue_batches,
656 })
657 }
658
659 #[allow(clippy::too_many_arguments)]
660 async fn finish_cancelled_turn_after_effect_abort(
661 &mut self,
662 driver: RuntimeTurnDriver<'_>,
663 mut assembler: TurnAssembler,
664 cancellation_messages: crate::MessageSequence,
665 events: &dyn EventSink,
666 finish_scoped_effect_controller: &ScopedEffectController<'_>,
667 cancel: &CancellationToken,
668 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
669 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
670 turn_control: &ActiveTurnControl,
671 turn_index: usize,
672 trace_turn_id: String,
673 ) -> Result<PhysicalTurnExecution, RuntimeError> {
674 let RuntimeTurnDriver {
675 session,
676 policy,
677 turn_pipeline,
678 pending_queue_claims,
679 pending_turn_input_claims,
680 ..
681 } = driver;
682 self.session = Some(session);
683 let outcome_event = SessionStreamEvent::TurnOutcome {
684 outcome: TurnOutcome::Stopped(TurnStop::Cancelled),
685 };
686 assembler.push(&outcome_event);
687 emit_session_event_to_sink(events, outcome_event).await;
688 assembler.push(&SessionStreamEvent::Done);
689 emit_session_event_to_sink(events, SessionStreamEvent::Done).await;
690 let claims = LogicalTurnClaims::new(pending_queue_claims, pending_turn_input_claims);
691 self.finish_turn(
692 TurnFinishInput {
693 turn_pipeline,
694 assembler,
695 new_messages: cancellation_messages,
696 policy,
697 turn_index,
698 trace_turn_id,
699 },
700 &claims,
701 events,
702 finish_scoped_effect_controller,
703 cancel,
704 session_execution_lease,
705 session_execution_lease_release_policy,
706 turn_control,
707 )
708 .await
709 }
710
711 fn emit_completed_turn_trace(
712 &self,
713 state: &SessionSnapshot,
714 outcome: &TurnOutcome,
715 trace_turn_id: &str,
716 ) {
717 if self.host.core.tracing.trace_sink.is_none() {
718 return;
719 }
720
721 let (status, done_reason, agent_frame_switch) = trace_fields_from_outcome(outcome);
722 crate::trace::emit_trace(
723 &self.host.core.tracing.trace_sink,
724 &self.host.core.tracing.trace_context,
725 lash_trace::TraceContext::default()
726 .for_session(state.session_id.clone())
727 .for_turn_index(state.turn_index)
728 .for_turn(trace_turn_id.to_string()),
729 lash_trace::TraceEvent::TurnCompleted {
730 status: status.to_string(),
731 done_reason: done_reason.to_string(),
732 agent_frame_switch,
733 },
734 self.host.core.clock.as_ref(),
735 );
736 }
737
738 #[allow(clippy::too_many_arguments)]
739 pub(super) async fn finish_logical_turn_error(
740 &mut self,
741 message: String,
742 trace_turn_id: String,
743 events: &dyn EventSink,
744 turn_events: &dyn TurnActivitySink,
745 scoped_effect_controller: ScopedEffectController<'_>,
746 cancel: CancellationToken,
747 claims: LogicalTurnClaims,
748 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
749 ) -> Result<PhysicalTurnExecution, RuntimeError> {
750 let turn_control_host = Arc::clone(&self.host.core.control.effect_host);
751 let turn_control_resolver =
752 turn_control_resolver(turn_control_host.as_ref(), &scoped_effect_controller);
753 let turn_control = Arc::new(
754 ActiveTurnControl::new(
755 turn_control_resolver,
756 TurnAddress::new(&self.state.session_id, &trace_turn_id),
757 )
758 .await?,
759 );
760 let mut assembler = TurnAssembler::default();
761 let error_event = SessionStreamEvent::Error {
762 message: message.clone(),
763 envelope: Some(crate::session_model::ErrorEnvelope {
764 kind: "runtime".to_string(),
765 code: Some("agent_frame_switch_limit".to_string()),
766 terminal_reason: None,
767 user_message: message.clone(),
768 raw: None,
769 retryable: Some(false),
770 provider_failure_kind: None,
771 }),
772 };
773 assembler.push(&error_event);
774 emit_turn_activity_to_sink(
775 turn_events,
776 TurnActivity::independent(TurnEvent::Error {
777 message: message.clone(),
778 }),
779 )
780 .await;
781 emit_session_event_to_sink(events, error_event).await;
782 let outcome_event = SessionStreamEvent::TurnOutcome {
783 outcome: TurnOutcome::Stopped(TurnStop::RuntimeError),
784 };
785 assembler.push(&outcome_event);
786 emit_session_event_to_sink(events, outcome_event).await;
787 assembler.push(&SessionStreamEvent::Done);
788 emit_session_event_to_sink(events, SessionStreamEvent::Done).await;
789
790 let messages = crate::MessageSequence::from_base(self.state.read_model().messages);
791 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
792 self.state.clone(),
793 Arc::clone(&self.host.core.clock),
794 )
795 .with_session_execution_lease(
796 session_execution_lease.map(SessionExecutionLeaseGuard::fence),
797 );
798 turn_pipeline.apply_prepared_messages(&messages);
799 self.finish_turn(
800 TurnFinishInput {
801 turn_pipeline,
802 assembler,
803 new_messages: messages,
804 policy: RuntimeSessionPolicy::new(
805 self.state.effective_policy().clone(),
806 Default::default(),
807 ),
808 turn_index: self.state.turn_index + 1,
809 trace_turn_id,
810 },
811 &claims,
812 events,
813 &scoped_effect_controller,
814 &cancel,
815 session_execution_lease,
816 SessionExecutionLeaseReleasePolicy::KeepOnAgentFrameSwitch,
817 &turn_control,
818 )
819 .await
820 }
821
822 async fn emit_turn_persisted_event(
823 &self,
824 returned_turn: &AssembledTurn,
825 scoped_effect_controller: &ScopedEffectController<'_>,
826 trace_turn_id: &str,
827 ) -> Result<(), RuntimeError> {
828 let Some(session) = self.session.as_ref() else {
829 return Ok(());
830 };
831 let Ok(manager) = self.runtime_session_services() else {
832 return Ok(());
833 };
834 let phase_turn_id = turn_phase_id(trace_turn_id, "turn-persisted");
835 let phase_controller = scoped_child_turn_controller(
836 scoped_effect_controller,
837 &self.state.session_id,
838 &phase_turn_id,
839 )?;
840 let direct_completions = manager.direct_completion_client(
841 RuntimeEffectControllerHandle::borrowed(phase_controller),
842 Some(phase_turn_id),
843 );
844
845 session
846 .plugins()
847 .emit_runtime_event_with_phase_probe(
848 crate::PluginLifecycleEvent::TurnPersisted(Box::new(
849 crate::SessionStateChangedContext {
850 session_id: self.state.session_id.clone(),
851 state: crate::SessionReadView::from_snapshot(&returned_turn.state),
852 sessions: manager.state_service(),
853 session_graph: manager.graph_service(),
854 direct_completions,
855 },
856 )),
857 self.turn_phase_probe.clone(),
858 )
859 .await;
860 Ok(())
861 }
862
863 pub async fn stream_turn(
865 &mut self,
866 mut input: TurnInput,
867 opts: TurnOptions<'_>,
868 ) -> Result<AssembledTurn, RuntimeError> {
869 if let Some(hint) = opts.local_cancel_origin_hint() {
870 input.turn_context.set_local_cancel_origin_hint(hint);
871 }
872 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
873 let cancel = opts.cancel.clone();
874 let session_execution_lease = self
875 .claim_session_execution_lease(cancel.clone(), true)
876 .await?;
877 let scoped_effect_controller = opts.scoped_effect_controller();
878 let result = Box::pin(self.drive_logical_turn(
879 LogicalTurnStart::Input(input),
880 opts.events_or_noop(),
881 opts.turn_events_or_noop(),
882 scoped_effect_controller,
883 cancel,
884 LogicalTurnClaims::new(Vec::new(), Vec::new()),
885 session_execution_lease.as_ref(),
886 stopwatch,
887 ))
888 .await
889 .map(|run| {
890 run.into_final_turn()
891 .expect("logical turn always contains a terminal physical turn")
892 });
893 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
894 .await
895 }
896
897 pub async fn stream_next_queued_work(
898 &mut self,
899 opts: TurnOptions<'_>,
900 ) -> Result<Option<AssembledTurn>, RuntimeError> {
901 self.stream_queued_work(opts, None).await
902 }
903
904 pub async fn stream_selected_queued_work(
905 &mut self,
906 opts: TurnOptions<'_>,
907 batch_ids: &[String],
908 ) -> Result<Option<AssembledTurn>, RuntimeError> {
909 self.stream_queued_work(opts, Some(batch_ids)).await
910 }
911
912 async fn stream_queued_work(
913 &mut self,
914 opts: TurnOptions<'_>,
915 selected_batch_ids: Option<&[String]>,
916 ) -> Result<Option<AssembledTurn>, RuntimeError> {
917 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
918 let cancel = opts.cancel.clone();
919 let Some(session_execution_lease) = self
920 .claim_session_execution_lease(cancel.clone(), false)
921 .await?
922 else {
923 return Ok(None);
924 };
925 let session_execution_fence = session_execution_lease.fence();
926 let Some(store) = self
927 .session
928 .as_ref()
929 .and_then(|session| session.history_store())
930 else {
931 session_execution_lease
932 .release_if_live()
933 .await
934 .map_err(|err| {
935 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
936 })?;
937 return Ok(None);
938 };
939 let drain_commands_before_turn_input = if selected_batch_ids.is_some() {
940 true
941 } else {
942 self.session_commands_precede_pending_turn_input(store.as_ref())
943 .await?
944 };
945 if drain_commands_before_turn_input {
946 loop {
947 match self
948 .drain_next_session_command(&session_execution_fence)
949 .await
950 {
951 Ok(Some(_)) => {}
952 Ok(None) => break,
953 Err(err) => {
954 let _ = session_execution_lease.release_if_live().await;
955 return Err(err);
956 }
957 }
958 }
959 }
960 if selected_batch_ids.is_none() {
961 let input_claim = store
962 .claim_next_turn_inputs(
963 &self.state.session_id,
964 &session_execution_fence,
965 &self.runtime_lease_owner,
966 64,
967 )
968 .await
969 .map_err(super::runtime_error_from_store_commit)?;
970 if let Some(input_claim) = input_claim {
971 let mut input = input_claim.materialize_for_turn();
972 if let Some(hint) = opts.local_cancel_origin_hint() {
973 input.turn_context.set_local_cancel_origin_hint(hint);
974 }
975 let turn_id = input
976 .trace_turn_id
977 .clone()
978 .or_else(|| Some(opts.execution_scope_id().to_owned()))
979 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
980 input.trace_turn_id = Some(turn_id.clone());
981 crate::trace::emit_trace(
982 &self.host.core.tracing.trace_sink,
983 &self.host.core.tracing.trace_context,
984 lash_trace::TraceContext::default()
985 .for_session(self.state.session_id.clone())
986 .for_turn_index(self.state.turn_index + 1)
987 .for_turn(turn_id.clone()),
988 lash_trace::TraceEvent::Custom {
989 name: "turn_input.claimed".to_string(),
990 payload: serde_json::json!({
991 "claim_id": &input_claim.claim_id,
992 "input_ids": input_claim.inputs.iter().map(|input| input.input_id.clone()).collect::<Vec<_>>(),
993 }),
994 },
995 self.host.core.clock.as_ref(),
996 );
997 let claim_for_abandon = input_claim.clone();
998 let scoped_effect_controller = opts.scoped_effect_controller();
999 let result = Box::pin(self.drive_logical_turn(
1000 LogicalTurnStart::Input(input),
1001 opts.events_or_noop(),
1002 opts.turn_events_or_noop(),
1003 scoped_effect_controller,
1004 cancel,
1005 LogicalTurnClaims::new(Vec::new(), vec![input_claim]),
1006 Some(&session_execution_lease),
1007 stopwatch,
1008 ))
1009 .await
1010 .map(AgentFrameRun::into_final_turn);
1011 if let Err(err) = &result {
1012 self.abandon_turn_input_claims_after_lease_loss(
1013 err,
1014 std::slice::from_ref(&claim_for_abandon),
1015 )
1016 .await;
1017 }
1018 return self
1019 .settle_session_execution_lease(Some(&session_execution_lease), result)
1020 .await;
1021 }
1022 }
1023 let claim = if let Some(batch_ids) = selected_batch_ids {
1024 store
1025 .claim_ready_queued_work_by_batch_ids(
1026 &self.state.session_id,
1027 &session_execution_fence,
1028 &self.runtime_lease_owner,
1029 crate::QueuedWorkClaimBoundary::Idle,
1030 batch_ids,
1031 )
1032 .await
1033 } else {
1034 store
1035 .claim_ready_queued_work(
1036 &self.state.session_id,
1037 &session_execution_fence,
1038 &self.runtime_lease_owner,
1039 crate::QueuedWorkClaimBoundary::Idle,
1040 64,
1041 )
1042 .await
1043 }
1044 .map_err(super::runtime_error_from_store_commit)?;
1045 let Some(claim) = claim else {
1046 session_execution_lease
1047 .release_if_live()
1048 .await
1049 .map_err(|err| {
1050 RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
1051 })?;
1052 return Ok(None);
1053 };
1054 let mut work = claim.materialize_for_turn();
1055 if let Some(hint) = opts.local_cancel_origin_hint() {
1056 work.input.turn_context.set_local_cancel_origin_hint(hint);
1057 }
1058 let turn_id = work
1059 .input
1060 .trace_turn_id
1061 .clone()
1062 .or_else(|| Some(opts.execution_scope_id().to_owned()))
1063 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1064 work.input.trace_turn_id = Some(turn_id.clone());
1065 let causes = work.turn_causes.clone();
1066 emit_queued_work_started_to_sink(
1067 opts.turn_events_or_noop(),
1068 crate::QueuedWorkClaimBoundary::Idle,
1069 &claim,
1070 causes.clone(),
1071 )
1072 .await;
1073 crate::trace::emit_trace(
1074 &self.host.core.tracing.trace_sink,
1075 &self.host.core.tracing.trace_context,
1076 lash_trace::TraceContext::default()
1077 .for_session(self.state.session_id.clone())
1078 .for_turn_index(self.state.turn_index + 1)
1079 .for_turn(turn_id.clone()),
1080 lash_trace::TraceEvent::Custom {
1081 name: "queued_work.claimed".to_string(),
1082 payload: queued_work_trace_payload(
1083 crate::QueuedWorkClaimBoundary::Idle,
1084 &claim,
1085 &causes,
1086 ),
1087 },
1088 self.host.core.clock.as_ref(),
1089 );
1090 let claim_for_abandon = claim.clone();
1091 let scoped_effect_controller = opts.scoped_effect_controller();
1092 let result = Box::pin(self.drive_logical_turn(
1093 LogicalTurnStart::Input(work.input),
1094 opts.events_or_noop(),
1095 opts.turn_events_or_noop(),
1096 scoped_effect_controller,
1097 cancel,
1098 LogicalTurnClaims::new(vec![claim], Vec::new()),
1099 Some(&session_execution_lease),
1100 stopwatch,
1101 ))
1102 .await
1103 .map(AgentFrameRun::into_final_turn);
1104 if let Err(err) = &result {
1105 self.abandon_queued_work_claims_after_lease_loss(
1106 err,
1107 std::slice::from_ref(&claim_for_abandon),
1108 )
1109 .await;
1110 }
1111 self.settle_session_execution_lease(Some(&session_execution_lease), result)
1112 .await
1113 }
1114
1115 async fn session_commands_precede_pending_turn_input(
1116 &self,
1117 store: &dyn crate::RuntimePersistence,
1118 ) -> Result<bool, RuntimeError> {
1119 let pending_inputs = store
1120 .list_pending_turn_inputs(&self.state.session_id)
1121 .await
1122 .map_err(super::runtime_error_from_store_commit)?;
1123 let earliest_input = pending_inputs
1124 .iter()
1125 .filter(|input| input.state.is_next_turn_pending())
1126 .min_by_key(|input| (input.enqueued_at_ms, input.enqueue_seq));
1127 let queued_work = store
1128 .list_pending_queued_work(&self.state.session_id)
1129 .await
1130 .map_err(super::runtime_error_from_store_commit)?;
1131 let earliest_command = queued_work
1132 .iter()
1133 .filter(|batch| batch.is_session_command_work())
1134 .min_by_key(|batch| (batch.enqueued_at_ms, batch.enqueue_seq));
1135 Ok(match (earliest_command, earliest_input) {
1136 (Some(command), Some(input)) => command.enqueued_at_ms < input.enqueued_at_ms,
1137 (Some(_), None) => true,
1138 _ => false,
1139 })
1140 }
1141
1142 fn ensure_durable_store_facets_for_scope(
1150 &self,
1151 scoped_effect_controller: &ScopedEffectController<'_>,
1152 ) -> Result<(), RuntimeError> {
1153 if scoped_effect_controller.controller().durability_tier() != crate::DurabilityTier::Durable
1154 {
1155 return Ok(());
1156 }
1157 if self
1158 .host
1159 .core
1160 .durability
1161 .attachment_store
1162 .persistence()
1163 .durability_tier()
1164 != crate::DurabilityTier::Durable
1165 {
1166 return Err(RuntimeError::durable_store_required(
1167 crate::DurableStoreFacet::AttachmentStore,
1168 ));
1169 }
1170 if self
1171 .host
1172 .core
1173 .durability
1174 .process_env_store
1175 .durability_tier()
1176 != crate::DurabilityTier::Durable
1177 {
1178 return Err(RuntimeError::durable_store_required(
1179 crate::DurableStoreFacet::ProcessEnvStore,
1180 ));
1181 }
1182 if let Some(store) = self
1183 .session
1184 .as_ref()
1185 .and_then(|session| session.history_store())
1186 && store.durability_tier() != crate::DurabilityTier::Durable
1187 {
1188 return Err(RuntimeError::durable_store_required(
1189 crate::DurableStoreFacet::SessionStore,
1190 ));
1191 }
1192 if let Some(process_registry) = self.host.process_registry.as_ref()
1193 && process_registry.durability_tier() != crate::DurabilityTier::Durable
1194 {
1195 return Err(RuntimeError::durable_store_required(
1196 crate::DurableStoreFacet::ProcessRegistry,
1197 ));
1198 }
1199 Ok(())
1200 }
1201
1202 #[allow(clippy::too_many_arguments)]
1203 pub(super) async fn stream_turn_with_scoped_effect_controller_inner(
1204 &mut self,
1205 mut input: TurnInput,
1206 events: &dyn EventSink,
1207 turn_events: &dyn TurnActivitySink,
1208 scoped_effect_controller: ScopedEffectController<'_>,
1209 cancel: CancellationToken,
1210 queued_claims: Vec<crate::QueuedWorkClaim>,
1211 turn_input_claims: Vec<crate::TurnInputClaim>,
1212 materialize_initial_claims: bool,
1213 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1214 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1215 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1216 if queued_claims.is_empty() && turn_input_claims.is_empty() {
1217 if let Some(lease) = session_execution_lease {
1218 while self
1219 .drain_next_session_command(&lease.fence())
1220 .await?
1221 .is_some()
1222 {}
1223 } else if self
1224 .session
1225 .as_ref()
1226 .and_then(|session| session.history_store())
1227 .is_some()
1228 {
1229 return Err(RuntimeError::new(
1230 RuntimeErrorCode::StoreCommitFailed,
1231 "session command drain requires a session execution lease",
1232 ));
1233 }
1234 }
1235 if let Some(input_turn_id) = input.trace_turn_id.as_deref()
1236 && scoped_effect_controller
1237 .execution_scope()
1238 .validates_turn_trace_id()
1239 && input_turn_id != scoped_effect_controller.scope_id()
1240 {
1241 return Err(RuntimeError::new(
1242 RuntimeErrorCode::ExecutionScopeTurnIdMismatch,
1243 format!(
1244 "input trace_turn_id `{input_turn_id}` does not match execution scope id `{}`",
1245 scoped_effect_controller.scope_id()
1246 ),
1247 ));
1248 }
1249 self.ensure_durable_store_facets_for_scope(&scoped_effect_controller)?;
1250 let turn_id = input
1251 .trace_turn_id
1252 .get_or_insert_with(|| scoped_effect_controller.scope_id().to_string())
1253 .clone();
1254 let _attachment_owner_binding = self
1258 .host
1259 .core
1260 .durability
1261 .attachment_store
1262 .bind_turn_scoped(turn_id);
1263 Box::pin(self.stream_turn_inner(
1264 input.clone(),
1265 events,
1266 turn_events,
1267 scoped_effect_controller,
1268 cancel.clone(),
1269 queued_claims,
1270 turn_input_claims,
1271 materialize_initial_claims,
1272 session_execution_lease,
1273 session_execution_lease_release_policy,
1274 ))
1275 .await
1276 }
1277
1278 pub async fn stream_turn_with_agent_frames(
1287 &mut self,
1288 mut input: TurnInput,
1289 opts: TurnOptions<'_>,
1290 ) -> Result<AgentFrameRun, RuntimeError> {
1291 if let Some(hint) = opts.local_cancel_origin_hint() {
1292 input.turn_context.set_local_cancel_origin_hint(hint);
1293 }
1294 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
1295 let cancel = opts.cancel.clone();
1296 let session_execution_lease = self
1297 .claim_session_execution_lease(cancel.clone(), true)
1298 .await?;
1299 let scoped_effect_controller = opts.scoped_effect_controller();
1300 let result = Box::pin(self.drive_logical_turn(
1301 LogicalTurnStart::Input(input),
1302 opts.events_or_noop(),
1303 opts.turn_events_or_noop(),
1304 scoped_effect_controller,
1305 cancel,
1306 LogicalTurnClaims::new(Vec::new(), Vec::new()),
1307 session_execution_lease.as_ref(),
1308 stopwatch,
1309 ))
1310 .await;
1311 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
1312 .await
1313 }
1314
1315 #[allow(clippy::too_many_arguments)]
1316 async fn stream_turn_inner(
1317 &mut self,
1318 mut input: TurnInput,
1319 events: &dyn EventSink,
1320 turn_events: &dyn TurnActivitySink,
1321 scoped_effect_controller: ScopedEffectController<'_>,
1322 cancel: CancellationToken,
1323 queued_claims: Vec<crate::QueuedWorkClaim>,
1324 turn_input_claims: Vec<crate::TurnInputClaim>,
1325 materialize_initial_claims: bool,
1326 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1327 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1328 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1329 let lease_continuity =
1330 session_execution_lease.and_then(SessionExecutionLeaseGuard::continuity);
1331 let may_reuse_resident_graph = self.graph_loaded_from_store
1332 && lease_continuity.is_some()
1333 && lease_continuity == self.last_committed_lease_continuity;
1334 if !may_reuse_resident_graph {
1335 self.refresh_session_graph_from_store()
1336 .await
1337 .map_err(session_head_refresh_error)?;
1338 }
1339 let input_trace_turn_id = input.trace_turn_id.clone();
1344 let queued_turn_work = materialize_initial_claims
1345 .then(|| queued_claims.first())
1346 .flatten()
1347 .map(crate::QueuedWorkClaim::materialize_for_turn);
1348 let pending_turn_input = materialize_initial_claims
1349 .then(|| turn_input_claims.first())
1350 .flatten()
1351 .map(crate::TurnInputClaim::materialize_for_turn);
1352 if let Some(work) = pending_turn_input.as_ref()
1353 && input.items.is_empty()
1354 {
1355 input = work.clone();
1356 if input.trace_turn_id.is_none() {
1357 input.trace_turn_id = input_trace_turn_id.clone();
1358 }
1359 }
1360 if let Some(work) = queued_turn_work.as_ref()
1361 && input.items.is_empty()
1362 {
1363 input = work.input.clone();
1364 if input.trace_turn_id.is_none() {
1365 input.trace_turn_id = input_trace_turn_id;
1366 }
1367 }
1368 if self
1369 .session
1370 .as_ref()
1371 .and_then(|session| session.history_store())
1372 .is_some()
1373 {
1374 ensure_durable_effect_input(&input)?;
1375 }
1376 if let Some(extension) = &input.protocol_extension
1377 && let Some(session) = self.session.as_ref()
1378 {
1379 let protocol_session = std::sync::Arc::clone(session.plugins().protocol_session());
1380 protocol_session
1381 .validate_turn_extension(extension)
1382 .await
1383 .map_err(|err| {
1384 RuntimeError::new(RuntimeErrorCode::ProtocolTurnExtension, err.to_string())
1385 })?;
1386 }
1387 let previous_prompt_usage = self.state.last_prompt_usage.clone();
1388 let normalized = match self.normalize_input_items(&input.items).await {
1389 Ok(items) => items,
1390 Err(e) => {
1391 self.state.last_prompt_usage = None;
1392 let mut assembler = TurnAssembler::default();
1393 let error_event = SessionStreamEvent::Error {
1394 message: e.clone(),
1395 envelope: Some(crate::session_model::ErrorEnvelope {
1396 kind: "input_validation".to_string(),
1397 code: Some("invalid_turn_input".to_string()),
1398 terminal_reason: None,
1399 user_message: e.clone(),
1400 raw: None,
1401 retryable: Some(false),
1402 provider_failure_kind: None,
1403 }),
1404 };
1405 assembler.push(&error_event);
1406 emit_turn_activity_to_sink(
1407 turn_events,
1408 TurnActivity::independent(TurnEvent::Error { message: e }),
1409 )
1410 .await;
1411 emit_session_event_to_sink(events, error_event).await;
1412 let outcome_event = SessionStreamEvent::TurnOutcome {
1413 outcome: TurnOutcome::Stopped(TurnStop::InvalidInput),
1414 };
1415 assembler.push(&outcome_event);
1416 emit_session_event_to_sink(events, outcome_event).await;
1417 assembler.push(&SessionStreamEvent::Done);
1418 emit_session_event_to_sink(events, SessionStreamEvent::Done).await;
1419 let turn_index = self.state.turn_index + 1;
1420 let trace_turn_id = input
1421 .trace_turn_id
1422 .clone()
1423 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1424 let turn_control_host = Arc::clone(&self.host.core.control.effect_host);
1425 let turn_control_resolver =
1426 turn_control_resolver(turn_control_host.as_ref(), &scoped_effect_controller);
1427 let turn_control = ActiveTurnControl::new(
1428 turn_control_resolver,
1429 TurnAddress::new(&self.state.session_id, &trace_turn_id),
1430 )
1431 .await?
1432 .with_local_cancel_origin(input.turn_context.local_cancel_origin_hint());
1433 let messages = crate::MessageSequence::from_base(self.state.read_model().messages);
1434 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
1435 self.state.clone(),
1436 Arc::clone(&self.host.core.clock),
1437 )
1438 .with_session_execution_lease(
1439 session_execution_lease.map(SessionExecutionLeaseGuard::fence),
1440 );
1441 turn_pipeline.apply_prepared_messages(&messages);
1442 let claims = LogicalTurnClaims::new(queued_claims, turn_input_claims);
1443 return self
1444 .finish_turn(
1445 TurnFinishInput {
1446 turn_pipeline,
1447 assembler,
1448 new_messages: messages,
1449 policy: RuntimeSessionPolicy::new(
1450 self.state.effective_policy().clone(),
1451 Default::default(),
1452 ),
1453 turn_index,
1454 trace_turn_id,
1455 },
1456 &claims,
1457 events,
1458 &scoped_effect_controller,
1459 &cancel,
1460 session_execution_lease,
1461 session_execution_lease_release_policy,
1462 &turn_control,
1463 )
1464 .await;
1465 }
1466 };
1467 let turn_index = self.state.turn_index + 1;
1468 let trace_turn_id = input
1469 .trace_turn_id
1470 .clone()
1471 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
1472 if self.host.core.tracing.trace_sink.is_some() {
1473 let mut trace_metadata = std::collections::BTreeMap::new();
1474 trace_metadata.insert(
1475 "input_item_count".to_string(),
1476 serde_json::json!(normalized.len()),
1477 );
1478 crate::trace::emit_trace(
1479 &self.host.core.tracing.trace_sink,
1480 &self.host.core.tracing.trace_context,
1481 lash_trace::TraceContext::default()
1482 .for_session(self.state.session_id.clone())
1483 .for_turn_index(turn_index)
1484 .for_turn(trace_turn_id.clone()),
1485 lash_trace::TraceEvent::TurnStarted {
1486 metadata: trace_metadata,
1487 },
1488 self.host.core.clock.as_ref(),
1489 );
1490 }
1491
1492 let base_read_model = self.state.read_model();
1493 let base_messages = base_read_model.messages;
1494 let base_render_cache = base_read_model.prompt_render_cache;
1495 let mut turn_delta = Vec::new();
1496 let initial_turn_causes = queued_turn_work
1497 .as_ref()
1498 .map(|work| work.turn_causes.clone())
1499 .unwrap_or_default();
1500 turn_delta.extend(
1501 initial_turn_causes
1502 .iter()
1503 .map(crate::TurnCause::to_event_message),
1504 );
1505
1506 let user_id = fresh_message_id();
1507 let mut user_parts: Vec<Part> = Vec::new();
1508 for item in normalized {
1509 match item {
1510 NormalizedItem::Text(text) => {
1511 if text.is_empty() {
1512 continue;
1513 }
1514 user_parts.push(Part {
1515 id: format!("{}.p{}", user_id, user_parts.len()),
1516 kind: PartKind::Text,
1517 content: text,
1518 attachment: None,
1519 tool_call_id: None,
1520 tool_name: None,
1521 tool_replay: None,
1522 prune_state: PruneState::Intact,
1523 reasoning_meta: None,
1524 response_meta: None,
1525 });
1526 }
1527 NormalizedItem::Attachment(source) => {
1528 user_parts.push(Part {
1529 id: format!("{}.p{}", user_id, user_parts.len()),
1530 kind: PartKind::Attachment,
1531 content: String::new(),
1532 attachment: Some(crate::session_model::message::PartAttachment { source }),
1533 tool_call_id: None,
1534 tool_name: None,
1535 tool_replay: None,
1536 prune_state: PruneState::Intact,
1537 reasoning_meta: None,
1538 response_meta: None,
1539 });
1540 }
1541 }
1542 }
1543 if user_parts.is_empty() && initial_turn_causes.is_empty() {
1544 user_parts.push(Part {
1545 id: format!("{}.p0", user_id),
1546 kind: PartKind::Text,
1547 content: String::new(),
1548 attachment: None,
1549 tool_call_id: None,
1550 tool_name: None,
1551 tool_replay: None,
1552 prune_state: PruneState::Intact,
1553 reasoning_meta: None,
1554 response_meta: None,
1555 });
1556 }
1557 if !user_parts.is_empty() {
1558 reassign_part_ids(&user_id, &mut user_parts);
1559 turn_delta.push(Message {
1560 id: user_id.clone(),
1561 role: MessageRole::User,
1562 parts: shared_parts(user_parts),
1563 origin: None,
1564 });
1565 }
1566
1567 let manager = self
1568 .runtime_session_services_for_turn(None)
1569 .map_err(|err| {
1570 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1571 })?;
1572 let plugin_session = self
1573 .session
1574 .as_ref()
1575 .map(|s| Arc::clone(s.plugins()))
1576 .ok_or_else(|| {
1577 RuntimeError::new(
1578 RuntimeErrorCode::ContextPrepareTurn,
1579 "runtime session not available",
1580 )
1581 })?;
1582 let prepare_phase_turn_id = turn_phase_id(&trace_turn_id, "prepare-turn");
1583 let prepare_phase_controller = scoped_child_turn_controller(
1584 &scoped_effect_controller,
1585 &self.state.session_id,
1586 &prepare_phase_turn_id,
1587 )?;
1588 let turn_ctx = crate::TurnTransformContext {
1589 session_id: self.state.session_id.clone(),
1590 state: self.read_view(),
1591 prompt_usage: previous_prompt_usage.clone(),
1592 max_context_tokens: Some(LashRuntime::max_context_tokens(self)),
1593 sessions: manager.state_service(),
1594 session_lifecycle: manager.lifecycle_service(),
1595 session_graph: manager.graph_service(),
1596 scoped_effect_controller: scoped_effect_controller.clone(),
1597 direct_completions: manager.direct_completion_client(
1598 RuntimeEffectControllerHandle::borrowed(prepare_phase_controller),
1599 Some(prepare_phase_turn_id),
1600 ),
1601 };
1602 self.mark_phase_begin(RuntimeTurnPhase::ContextTransform);
1603 let prepared_context = plugin_session
1604 .prepare_turn_context(
1605 &turn_ctx,
1606 crate::session_model::context::PreparedContext {
1607 messages: crate::MessageSequence::from_base_and_delta(
1608 base_messages,
1609 turn_delta,
1610 )
1611 .with_base_render_cache(base_render_cache),
1612 ..Default::default()
1613 },
1614 self.turn_phase_probe.clone(),
1615 )
1616 .await
1617 .map_err(|err| {
1618 RuntimeError::new(RuntimeErrorCode::ContextPrepareTurn, err.to_string())
1619 })?;
1620 self.mark_phase_end(RuntimeTurnPhase::ContextTransform);
1621 drop(turn_ctx);
1626 let messages = prepared_context.messages;
1627 if let Some(session) = self.session.as_mut() {
1628 session
1629 .set_context_overlay(
1630 prepared_context.tool_providers,
1631 prepared_context.prompt_contributions,
1632 prepared_context.include_base_tools,
1633 )
1634 .map_err(|err| {
1635 RuntimeError::new(
1636 RuntimeErrorCode::Other("session_tool_registry".to_string()),
1637 err.to_string(),
1638 )
1639 })?;
1640 }
1641
1642 self.state.last_prompt_usage = None;
1643 Box::pin(self.stream_prepared_turn_inner(
1644 messages,
1645 previous_prompt_usage,
1646 input.protocol_turn_options.clone(),
1647 input.protocol_extension.clone(),
1648 input.turn_context.clone(),
1649 initial_turn_causes,
1650 trace_turn_id,
1651 turn_index,
1652 events,
1653 turn_events,
1654 scoped_effect_controller,
1655 cancel,
1656 queued_claims,
1657 turn_input_claims,
1658 session_execution_lease,
1659 session_execution_lease_release_policy,
1660 ))
1661 .await
1662 }
1663
1664 pub async fn run_turn_assembled(
1666 &mut self,
1667 input: TurnInput,
1668 cancel: CancellationToken,
1669 scoped_effect_controller: ScopedEffectController<'_>,
1670 ) -> Result<AssembledTurn, RuntimeError> {
1671 self.stream_turn(input, TurnOptions::new(cancel, scoped_effect_controller))
1672 .await
1673 }
1674
1675 #[allow(clippy::too_many_arguments)]
1677 pub async fn stream_prepared_turn(
1678 &mut self,
1679 messages: crate::MessageSequence,
1680 previous_prompt_usage: Option<PromptUsage>,
1681 protocol_turn_options: Option<crate::ProtocolTurnOptions>,
1682 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
1683 turn_context: crate::TurnContext,
1684 initial_turn_causes: Vec<crate::TurnCause>,
1685 trace_turn_id: String,
1686 turn_index: usize,
1687 events: &dyn EventSink,
1688 turn_events: &dyn TurnActivitySink,
1689 scoped_effect_controller: ScopedEffectController<'_>,
1690 cancel: CancellationToken,
1691 initial_queue_claim: Option<crate::QueuedWorkClaim>,
1692 initial_turn_input_claim: Option<crate::TurnInputClaim>,
1693 ) -> Result<AssembledTurn, RuntimeError> {
1694 let stopwatch = TurnStopwatch::start(self.host.core.clock.as_ref());
1695 let session_execution_lease = self
1696 .claim_session_execution_lease(cancel.clone(), true)
1697 .await?;
1698 let result = Box::pin(self.drive_logical_turn(
1699 LogicalTurnStart::Prepared(PreparedLogicalTurn {
1700 messages,
1701 previous_prompt_usage,
1702 protocol_turn_options,
1703 protocol_extension,
1704 turn_context,
1705 initial_turn_causes,
1706 trace_turn_id,
1707 turn_index,
1708 }),
1709 events,
1710 turn_events,
1711 scoped_effect_controller,
1712 cancel,
1713 LogicalTurnClaims::new(
1714 initial_queue_claim.into_iter().collect(),
1715 initial_turn_input_claim.into_iter().collect(),
1716 ),
1717 session_execution_lease.as_ref(),
1718 stopwatch,
1719 ))
1720 .await
1721 .map(|run| {
1722 run.into_final_turn()
1723 .expect("logical turn always contains a terminal physical turn")
1724 });
1725 self.settle_session_execution_lease(session_execution_lease.as_ref(), result)
1726 .await
1727 }
1728
1729 #[allow(clippy::too_many_arguments)]
1730 async fn prepare_turn_preamble(
1731 &mut self,
1732 plugins: &crate::PluginSession,
1733 manager: &Arc<RuntimeSessionServices>,
1734 messages: crate::MessageSequence,
1735 turn_policy: &crate::SessionPolicy,
1736 effective_protocol_turn_options: &crate::ProtocolTurnOptions,
1737 turn_context: &crate::TurnContext,
1738 event_rx: &mut mpsc::Receiver<RuntimeStreamEvent>,
1739 assembler: &mut TurnAssembler,
1740 events: &dyn EventSink,
1741 turn_events: &dyn TurnActivitySink,
1742 ) -> Result<crate::plugin::TurnPreparation, RuntimeError> {
1743 self.mark_phase_begin(RuntimeTurnPhase::BeforeTurnHooks);
1744 let prepare_turn = plugins.prepare_turn_with_phase_probe(
1745 PrepareTurnRequest {
1746 session_id: self.state.session_id.clone(),
1747 state: crate::SessionReadView::from_runtime_state(
1748 &self.state,
1749 turn_policy.clone(),
1750 effective_protocol_turn_options.clone(),
1751 ),
1752 messages,
1753 sessions: manager.state_service(),
1754 session_lifecycle: manager.lifecycle_service(),
1755 session_graph: manager.graph_service(),
1756 turn_context: turn_context.clone(),
1757 },
1758 self.turn_phase_probe.clone(),
1759 );
1760 let mut prepare_turn = Box::pin(prepare_turn);
1761
1762 loop {
1763 tokio::select! {
1764 prepared = prepare_turn.as_mut() => {
1765 let prepared = prepared.map_err(|err| {
1766 RuntimeError::new(RuntimeErrorCode::PluginPrepareTurn, err.to_string())
1767 })?;
1768 self.mark_phase_end(RuntimeTurnPhase::BeforeTurnHooks);
1769 return Ok(prepared);
1770 }
1771 maybe_event = event_rx.recv() => {
1772 if let Some(event) = maybe_event {
1773 emit_runtime_stream_event_to_sinks(
1774 events,
1775 turn_events,
1776 event,
1777 assembler,
1778 )
1779 .await;
1780 }
1781 }
1782 }
1783 }
1784 }
1785
1786 #[allow(clippy::too_many_arguments)]
1787 async fn finish_prepared_turn_abort(
1788 &mut self,
1789 prepared: crate::plugin::TurnPreparation,
1790 event_tx: mpsc::Sender<RuntimeStreamEvent>,
1791 mut assembler: TurnAssembler,
1792 turn_index: usize,
1793 trace_turn_id: String,
1794 claims: &LogicalTurnClaims,
1795 events: &dyn EventSink,
1796 turn_events: &dyn TurnActivitySink,
1797 scoped_effect_controller: &ScopedEffectController<'_>,
1798 cancel: &CancellationToken,
1799 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1800 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1801 session_execution_fence: Option<crate::SessionExecutionLeaseFence>,
1802 turn_control: &ActiveTurnControl,
1803 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1804 let Some(abort) = prepared.abort else {
1805 unreachable!("abort finisher requires a prepared plugin abort");
1806 };
1807 drop(event_tx);
1808
1809 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
1813 self.state.clone(),
1814 Arc::clone(&self.host.core.clock),
1815 )
1816 .with_session_execution_lease(session_execution_fence);
1817 turn_pipeline.apply_prepared_messages(&prepared.messages);
1818 let issue = TurnIssue {
1819 kind: "plugin".to_string(),
1820 code: Some(abort.code),
1821 terminal_reason: None,
1822 message: abort.message.clone(),
1823 raw: None,
1824 retryable: None,
1825 provider_failure_kind: None,
1826 };
1827 let error_event = SessionStreamEvent::Error {
1828 message: abort.message,
1829 envelope: Some(crate::session_model::ErrorEnvelope {
1830 kind: "plugin".to_string(),
1831 code: issue.code.clone(),
1832 terminal_reason: None,
1833 user_message: issue.message.clone(),
1834 raw: None,
1835 retryable: None,
1836 provider_failure_kind: None,
1837 }),
1838 };
1839 assembler.push(&error_event);
1840 emit_turn_activity_to_sink(
1841 turn_events,
1842 TurnActivity::independent(TurnEvent::Error {
1843 message: issue.message.clone(),
1844 }),
1845 )
1846 .await;
1847 emit_session_event_to_sink(events, error_event).await;
1848 let outcome_event = SessionStreamEvent::TurnOutcome {
1849 outcome: TurnOutcome::Stopped(TurnStop::PluginAbort),
1850 };
1851 assembler.push(&outcome_event);
1852 emit_session_event_to_sink(events, outcome_event).await;
1853 assembler.push(&SessionStreamEvent::Done);
1854 emit_session_event_to_sink(events, SessionStreamEvent::Done).await;
1855 self.finish_turn(
1856 TurnFinishInput {
1857 turn_pipeline,
1858 assembler,
1859 new_messages: prepared.messages,
1860 policy: RuntimeSessionPolicy::new(
1861 self.state.effective_policy().clone(),
1862 Default::default(),
1863 ),
1864 turn_index,
1865 trace_turn_id,
1866 },
1867 claims,
1868 events,
1869 scoped_effect_controller,
1870 cancel,
1871 session_execution_lease,
1872 session_execution_lease_release_policy,
1873 turn_control,
1874 )
1875 .await
1876 }
1877
1878 #[allow(clippy::too_many_arguments)]
1879 pub(super) async fn stream_prepared_turn_inner(
1880 &mut self,
1881 messages: crate::MessageSequence,
1882 _previous_prompt_usage: Option<PromptUsage>,
1883 protocol_turn_options: Option<crate::ProtocolTurnOptions>,
1884 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
1885 turn_context: crate::TurnContext,
1886 initial_turn_causes: Vec<crate::TurnCause>,
1887 trace_turn_id: String,
1888 turn_index: usize,
1889 events: &dyn EventSink,
1890 turn_events: &dyn TurnActivitySink,
1891 scoped_effect_controller: ScopedEffectController<'_>,
1892 cancel: CancellationToken,
1893 initial_queue_claims: Vec<crate::QueuedWorkClaim>,
1894 initial_turn_input_claims: Vec<crate::TurnInputClaim>,
1895 session_execution_lease: Option<&SessionExecutionLeaseGuard>,
1896 session_execution_lease_release_policy: SessionExecutionLeaseReleasePolicy,
1897 ) -> Result<PhysicalTurnExecution, RuntimeError> {
1898 let turn_control_host = Arc::clone(&self.host.core.control.effect_host);
1899 let turn_control_resolver =
1900 turn_control_resolver(turn_control_host.as_ref(), &scoped_effect_controller);
1901 let inline_turn_control_controller =
1902 if turn_control_host.durability_tier() == crate::DurabilityTier::Inline {
1903 Some(turn_control_host.scoped(scoped_effect_controller.execution_scope().clone())?)
1904 } else {
1905 None
1906 };
1907 let turn_control = Arc::new(
1908 ActiveTurnControl::new(
1909 turn_control_resolver,
1910 TurnAddress::new(&self.state.session_id, &trace_turn_id),
1911 )
1912 .await?
1913 .with_local_cancel_origin(turn_context.local_cancel_origin_hint()),
1914 );
1915 if session_execution_lease.is_none()
1916 && self
1917 .session
1918 .as_ref()
1919 .and_then(|session| session.history_store())
1920 .is_some()
1921 {
1922 return Err(RuntimeError::new(
1923 RuntimeErrorCode::StoreCommitFailed,
1924 "prepared turn requires a session execution lease",
1925 ));
1926 }
1927 let session_execution_fence =
1928 session_execution_lease.map(SessionExecutionLeaseGuard::fence);
1929 let (event_tx, mut event_rx) = mpsc::channel::<RuntimeStreamEvent>(100);
1930 let child_usage_event_relay = ChildUsageEventRelay::new(event_tx.clone());
1931 let mut turn_policy = self.state.effective_policy().clone();
1932 let turn_provider_override = turn_context.provider().cloned();
1933 if let Some(provider) = turn_provider_override.as_ref() {
1934 turn_policy.provider_id = provider.kind().to_string();
1935 }
1936 let session_protocol_turn_options = self.state.effective_protocol_turn_options().clone();
1937 let effective_protocol_turn_options = protocol_turn_options
1938 .clone()
1939 .map(|options| session_protocol_turn_options.merged_with_override(&options))
1940 .unwrap_or(session_protocol_turn_options);
1941 let manager = self
1942 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
1943 .map_err(|err| {
1944 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
1945 })?;
1946 let plugins = {
1947 let session = self
1948 .session
1949 .as_ref()
1950 .expect("lash runtime session must be available");
1951 Arc::clone(session.plugins())
1952 };
1953 let mut assembler = TurnAssembler::new();
1954 let initial_claims =
1955 LogicalTurnClaims::new(initial_queue_claims, initial_turn_input_claims);
1956 let mut prepared = self
1960 .prepare_turn_preamble(
1961 plugins.as_ref(),
1962 &manager,
1963 messages,
1964 &turn_policy,
1965 &effective_protocol_turn_options,
1966 &turn_context,
1967 &mut event_rx,
1968 &mut assembler,
1969 events,
1970 turn_events,
1971 )
1972 .await?;
1973 for event in &prepared.events {
1974 assembler.push(event);
1975 }
1976 emit_session_events_to_sink(events, std::mem::take(&mut prepared.events)).await;
1977 if prepared.abort.is_some() {
1978 return self
1979 .finish_prepared_turn_abort(
1980 prepared,
1981 event_tx,
1982 assembler,
1983 turn_index,
1984 trace_turn_id,
1985 &initial_claims,
1986 events,
1987 turn_events,
1988 &scoped_effect_controller,
1989 &cancel,
1990 session_execution_lease,
1991 session_execution_lease_release_policy,
1992 session_execution_fence,
1993 turn_control.as_ref(),
1994 )
1995 .await;
1996 }
1997 let mut turn_pipeline = TurnBoundary::from_state_with_clock(
2000 self.state.clone(),
2001 Arc::clone(&self.host.core.clock),
2002 )
2003 .with_session_execution_lease(session_execution_fence.clone());
2004 let store = self
2005 .session
2006 .as_ref()
2007 .and_then(|session| session.history_store());
2008 let progress_store = if scoped_effect_controller.controller().durability_tier()
2012 == crate::DurabilityTier::Durable
2013 {
2014 None
2015 } else {
2016 store.as_ref().map(|store| store.as_ref())
2017 };
2018 turn_pipeline
2019 .prepared_checkpoint(
2020 progress_store,
2021 turn_policy.clone(),
2022 turn_index,
2023 &prepared.messages,
2024 self.session.as_mut(),
2025 )
2026 .await
2027 .map_err(super::runtime_error_from_store_commit)?;
2028 let resolved_turn_policy = if let Some(provider) = turn_provider_override {
2029 RuntimeSessionPolicy::from_provider(
2030 turn_policy.clone(),
2031 provider.with_clock(Arc::clone(&self.host.core.clock)),
2032 )
2033 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
2034 } else {
2035 self.host
2036 .resolve_session_policy(&self.state.session_id, turn_policy.clone())
2037 .map_err(|err| RuntimeError::new("llm_provider", err.to_string()))?
2038 };
2039 let manager = self
2040 .runtime_session_services_for_turn(Some(child_usage_event_relay.clone()))
2041 .map_err(|err| {
2042 RuntimeError::new(RuntimeErrorCode::PluginSessionManager, err.to_string())
2043 })?;
2044 let cancel_state = cancel.clone();
2045 let finish_scoped_effect_controller = scoped_effect_controller.clone();
2046 let turn_cancel_peek_controller = inline_turn_control_controller
2047 .as_ref()
2048 .map(ScopedEffectController::controller)
2049 .unwrap_or_else(|| finish_scoped_effect_controller.controller());
2050 let session = self
2051 .session
2052 .take()
2053 .expect("lash runtime session must be available");
2054 let mut driver = Box::new(RuntimeTurnDriver {
2055 session,
2056 policy: resolved_turn_policy,
2057 host: self.host.clone(),
2058 turn_id: scoped_effect_controller.scope_id().to_string(),
2059 scoped_effect_controller,
2060 session_id: self.state.session_id.clone(),
2061 turn_index,
2062 turn_pipeline,
2063 llm_stream_summaries: HashMap::new(),
2064 llm_calls: Vec::new(),
2065 next_llm_ordinal: 0,
2066 session_services: manager,
2067 protocol_turn_options: effective_protocol_turn_options,
2068 protocol_extension,
2069 turn_context,
2070 turn_causes: initial_turn_causes,
2071 pending_queue_claims: initial_claims.queued,
2072 pending_turn_input_claims: initial_claims.turn_inputs,
2073 checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer::default(),
2074 session_execution_lease: session_execution_fence,
2075 runtime_lease_owner: self.runtime_lease_owner.clone(),
2076 turn_phase_probe: self.turn_phase_probe.clone(),
2077 });
2078 let protocol_run_offset = 0;
2079 self.mark_phase_begin(RuntimeTurnPhase::EffectLoop);
2080 let run_result = Box::pin(run_turn_effect_loop(
2081 &mut driver,
2082 prepared.messages,
2083 event_tx,
2084 cancel.clone(),
2085 protocol_run_offset,
2086 Arc::clone(&turn_control),
2087 Arc::clone(&turn_control_host),
2088 turn_cancel_peek_controller,
2089 &mut event_rx,
2090 &mut assembler,
2091 &child_usage_event_relay,
2092 events,
2093 turn_events,
2094 ))
2095 .await;
2096 let (new_messages, _new_protocol_iteration) = match run_result {
2097 Ok(result) => result,
2098 Err(err) if cancel.is_cancelled() => {
2099 if turn_control.evidence().is_none() {
2100 turn_control
2101 .observe_pending_cancel(
2102 turn_cancel_peek_controller,
2103 crate::runtime::turn_control::TurnCancelPeekIdentity::PostAbortGate,
2104 )
2105 .await?;
2106 }
2107 if turn_control.evidence().is_some() {
2108 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
2109 let cancellation_messages = driver.turn_pipeline.message_sequence();
2110 return Box::pin(self.finish_cancelled_turn_after_effect_abort(
2111 *driver,
2112 assembler,
2113 cancellation_messages,
2114 events,
2115 &finish_scoped_effect_controller,
2116 &cancel,
2117 session_execution_lease,
2118 session_execution_lease_release_policy,
2119 turn_control.as_ref(),
2120 turn_index,
2121 trace_turn_id,
2122 ))
2123 .await;
2124 }
2125 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
2126 let RuntimeTurnDriver {
2127 session,
2128 pending_queue_claims,
2129 pending_turn_input_claims,
2130 ..
2131 } = *driver;
2132 self.session = Some(session);
2133 self.abandon_queued_work_claims_after_lease_loss(&err, &pending_queue_claims)
2134 .await;
2135 self.abandon_turn_input_claims_after_lease_loss(&err, &pending_turn_input_claims)
2136 .await;
2137 return Err(err);
2138 }
2139 Err(err) => {
2140 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
2141 let RuntimeTurnDriver {
2142 session,
2143 pending_queue_claims,
2144 pending_turn_input_claims,
2145 ..
2146 } = *driver;
2147 self.session = Some(session);
2148 self.abandon_queued_work_claims_after_lease_loss(&err, &pending_queue_claims)
2149 .await;
2150 self.abandon_turn_input_claims_after_lease_loss(&err, &pending_turn_input_claims)
2151 .await;
2152 return Err(err);
2153 }
2154 };
2155 self.mark_phase_end(RuntimeTurnPhase::EffectLoop);
2156 tracing::debug!(
2157 new_message_count = new_messages.len(),
2158 tool_call_count = assembler.tool_calls.len(),
2159 "runtime post-run_task"
2160 );
2161
2162 let RuntimeTurnDriver {
2163 session,
2164 policy,
2165 turn_pipeline,
2166 llm_calls,
2167 pending_queue_claims,
2168 pending_turn_input_claims,
2169 ..
2170 } = *driver;
2171 self.session = Some(session);
2172 let pending_claims =
2173 LogicalTurnClaims::new(pending_queue_claims, pending_turn_input_claims);
2174 let finish_result = Box::pin(self.finish_turn(
2175 TurnFinishInput {
2176 turn_pipeline,
2177 assembler: assembler.with_llm_calls(llm_calls),
2178 new_messages,
2179 policy,
2180 turn_index,
2181 trace_turn_id,
2182 },
2183 &pending_claims,
2184 events,
2185 &finish_scoped_effect_controller,
2186 &cancel_state,
2187 session_execution_lease,
2188 session_execution_lease_release_policy,
2189 turn_control.as_ref(),
2190 ))
2191 .await;
2192 if let Err(err) = &finish_result {
2193 self.abandon_queued_work_claims_after_lease_loss(err, &pending_claims.queued)
2194 .await;
2195 self.abandon_turn_input_claims_after_lease_loss(err, &pending_claims.turn_inputs)
2196 .await;
2197 }
2198 finish_result
2199 }
2200 async fn normalize_input_items(
2201 &self,
2202 items: &[InputItem],
2203 ) -> Result<Vec<NormalizedItem>, String> {
2204 normalize_input_items(
2205 items,
2206 self.host.core.durability.attachment_store.as_ref(),
2207 self.host.core.attachment_source_policy.as_ref(),
2208 )
2209 .await
2210 }
2211}
2212
2213pub fn ensure_durable_effect_input(input: &TurnInput) -> Result<(), RuntimeError> {
2214 if input.protocol_extension.is_some() {
2215 return Err(RuntimeError::new(
2216 RuntimeErrorCode::DurableEffectLiveProtocolExtension,
2217 "durable effect hosts do not support live protocol_extension inputs; encode replayable data in protocol_turn_options or persisted plugin state",
2218 ));
2219 }
2220 input
2221 .turn_context
2222 .live_plugin_inputs()
2223 .durable_effect_rejection()?;
2224 Ok(())
2225}
2226
2227async fn emit_turn_activity_to_sink(events: &dyn TurnActivitySink, activity: TurnActivity) {
2228 if !events.is_noop() {
2229 events.emit(activity).await;
2230 }
2231}
2232
2233async fn publish_terminal_after_commit(
2234 turn_control: &ActiveTurnControl,
2235 resolver: &dyn AwaitEventResolver,
2236 terminal: &TurnTerminal,
2237 session_id: &str,
2238 turn_id: &str,
2239) {
2240 if let Err(err) = turn_control.publish_terminal(resolver, terminal).await {
2241 tracing::warn!(
2242 error = %err,
2243 session_id,
2244 turn_id,
2245 "turn committed but terminal publication failed"
2246 );
2247 }
2248}
2249
2250#[allow(clippy::too_many_arguments)]
2251async fn run_turn_effect_loop(
2252 driver: &mut RuntimeTurnDriver<'_>,
2253 messages: crate::MessageSequence,
2254 event_tx: mpsc::Sender<RuntimeStreamEvent>,
2255 cancellation: CancellationToken,
2256 protocol_run_offset: usize,
2257 turn_control: Arc<ActiveTurnControl>,
2258 turn_control_host: Arc<dyn EffectHost>,
2259 cancel_controller: &dyn RuntimeEffectController,
2260 event_rx: &mut mpsc::Receiver<RuntimeStreamEvent>,
2261 assembler: &mut TurnAssembler,
2262 child_usage_event_relay: &ChildUsageEventRelay,
2263 events: &dyn EventSink,
2264 turn_events: &dyn TurnActivitySink,
2265) -> Result<(crate::MessageSequence, usize), RuntimeError> {
2266 let start_gate = crate::runtime::RuntimeNamedPhase::begin(
2274 driver.turn_phase_probe.clone(),
2275 "turn_cancel.start_gate",
2276 );
2277 let pending_cancel = await_turn_cancellation_start_gate(|| {
2278 turn_control.observe_pending_cancel(
2279 cancel_controller,
2280 crate::runtime::turn_control::TurnCancelPeekIdentity::StartGate,
2281 )
2282 })
2283 .await?;
2284 drop(start_gate);
2285 if pending_cancel.is_some() {
2286 cancellation.cancel();
2287 }
2288 let cancel_watcher = crate::task::spawn({
2289 let turn_control = Arc::clone(&turn_control);
2290 let cancellation = cancellation.clone();
2291 async move {
2292 if await_turn_cancellation_with_retry(|| {
2293 turn_control.await_cancel(turn_control_host.as_ref(), CancellationToken::new())
2294 })
2295 .await
2296 .is_some()
2297 {
2298 cancellation.cancel();
2299 }
2300 }
2301 });
2302 let run_future = Box::pin(driver.run(
2306 messages,
2307 event_tx,
2308 cancellation.clone(),
2309 protocol_run_offset,
2310 ));
2311 let result = drive_turn_to_completion(
2312 run_future,
2313 event_rx,
2314 assembler,
2315 child_usage_event_relay,
2316 events,
2317 turn_events,
2318 )
2319 .await;
2320 cancel_watcher.abort();
2321 result
2322}
2323
2324const TURN_CANCEL_WATCH_RETRY_INITIAL: std::time::Duration = std::time::Duration::from_millis(25);
2325const TURN_CANCEL_WATCH_RETRY_MAX: std::time::Duration = std::time::Duration::from_secs(1);
2326const TURN_CANCEL_START_GATE_ATTEMPTS: usize = 3;
2334
2335async fn await_turn_cancellation_start_gate<F, C>(
2336 mut watch: F,
2337) -> Result<Option<TurnCancellationEvidence>, RuntimeError>
2338where
2339 F: FnMut() -> C,
2340 C: std::future::Future<Output = Result<Option<TurnCancellationEvidence>, RuntimeError>>,
2341{
2342 let mut backoff = TURN_CANCEL_WATCH_RETRY_INITIAL;
2343 for attempt in 1..=TURN_CANCEL_START_GATE_ATTEMPTS {
2344 match watch().await {
2345 Ok(observation) => return Ok(observation),
2346 Err(err) if attempt == TURN_CANCEL_START_GATE_ATTEMPTS => return Err(err),
2347 Err(err) => {
2348 tracing::warn!(
2349 error = %err,
2350 attempt,
2351 max_attempts = TURN_CANCEL_START_GATE_ATTEMPTS,
2352 retry_after_ms = backoff.as_millis(),
2353 "turn cancellation start gate failed; retrying before failing the invocation"
2354 );
2355 tokio::time::sleep(backoff).await;
2356 backoff = backoff.saturating_mul(2).min(TURN_CANCEL_WATCH_RETRY_MAX);
2357 }
2358 }
2359 }
2360 unreachable!("positive start-gate attempt limit")
2361}
2362
2363async fn await_turn_cancellation_with_retry<F, C>(mut watch: F) -> Option<TurnCancellationEvidence>
2364where
2365 F: FnMut() -> C,
2366 C: std::future::Future<Output = Result<Option<TurnCancellationEvidence>, RuntimeError>>,
2367{
2368 let mut backoff = TURN_CANCEL_WATCH_RETRY_INITIAL;
2369 loop {
2370 match watch().await {
2371 Ok(observation) => return observation,
2372 Err(err) => {
2373 tracing::warn!(
2374 error = %err,
2375 retry_after_ms = backoff.as_millis(),
2376 "turn cancellation watcher failed; retrying while the turn remains active"
2377 );
2378 tokio::time::sleep(backoff).await;
2379 backoff = backoff.saturating_mul(2).min(TURN_CANCEL_WATCH_RETRY_MAX);
2380 }
2381 }
2382 }
2383}
2384
2385async fn drive_turn_to_completion<F>(
2395 mut run_future: Pin<Box<F>>,
2396 event_rx: &mut mpsc::Receiver<RuntimeStreamEvent>,
2397 assembler: &mut TurnAssembler,
2398 child_usage_event_relay: &ChildUsageEventRelay,
2399 events: &dyn EventSink,
2400 turn_events: &dyn TurnActivitySink,
2401) -> Result<(crate::MessageSequence, usize), RuntimeError>
2402where
2403 F: std::future::Future<Output = Result<(crate::MessageSequence, usize), RuntimeError>> + ?Sized,
2404{
2405 let run_result = {
2406 loop {
2407 tokio::select! {
2408 biased;
2412
2413 completed = run_future.as_mut() => {
2414 child_usage_event_relay.clear();
2415 break completed;
2416 }
2417 maybe_event = event_rx.recv() => {
2418 if let Some(event) = maybe_event {
2419 emit_runtime_stream_event_to_sinks(
2420 events,
2421 turn_events,
2422 event,
2423 assembler,
2424 )
2425 .await;
2426 }
2427 }
2428 }
2429 }
2430 };
2431 while let Some(event) = event_rx.recv().await {
2432 emit_runtime_stream_event_to_sinks(events, turn_events, event, assembler).await;
2433 }
2434 run_result
2435}
2436
2437#[allow(clippy::too_many_arguments)]
2438async fn emit_runtime_stream_event_to_sinks(
2439 events: &dyn EventSink,
2440 turn_events: &dyn TurnActivitySink,
2441 event: RuntimeStreamEvent,
2442 assembler: &mut TurnAssembler,
2443) {
2444 match event {
2445 RuntimeStreamEvent::Session(event) => {
2446 assembler.push(&event);
2447 emit_session_event_to_sink(events, event).await;
2448 }
2449 RuntimeStreamEvent::Turn(activity) => {
2450 emit_turn_activity_to_sink(turn_events, activity).await;
2451 }
2452 }
2453}
2454
2455#[cfg(test)]
2456mod tests {
2457 use std::sync::Arc;
2458 use std::sync::atomic::{AtomicUsize, Ordering};
2459
2460 use super::{
2461 ActiveTurnControl, TURN_CANCEL_START_GATE_ATTEMPTS, agent_frame_follow_turn_id,
2462 await_turn_cancellation_start_gate, await_turn_cancellation_with_retry,
2463 publish_terminal_after_commit,
2464 };
2465 use crate::{
2466 AwaitEventKey, AwaitEventResolver, AwaitEventWaitIdentity, ExecutionScope,
2467 InlineRuntimeEffectController, Resolution, ResolveOutcome, RuntimeError, TurnAddress,
2468 TurnCancellationEvidence, TurnFinish, TurnOutcome, TurnTerminal,
2469 };
2470
2471 #[derive(Default)]
2472 struct RejectTerminalPublication {
2473 attempts: AtomicUsize,
2474 inline: InlineRuntimeEffectController,
2475 }
2476
2477 #[async_trait::async_trait]
2478 impl AwaitEventResolver for RejectTerminalPublication {
2479 async fn await_event_key(
2480 &self,
2481 scope: &ExecutionScope,
2482 wait: AwaitEventWaitIdentity,
2483 ) -> Result<AwaitEventKey, RuntimeError> {
2484 self.inline.await_event_key(scope, wait).await
2485 }
2486
2487 async fn resolve_await_event(
2488 &self,
2489 _key: &AwaitEventKey,
2490 _resolution: Resolution,
2491 ) -> Result<ResolveOutcome, RuntimeError> {
2492 self.attempts.fetch_add(1, Ordering::SeqCst);
2493 Err(RuntimeError::new(
2494 "transient_terminal_publication",
2495 "terminal backend unavailable",
2496 ))
2497 }
2498
2499 async fn peek_await_event(
2500 &self,
2501 key: &AwaitEventKey,
2502 ) -> Result<Option<Resolution>, RuntimeError> {
2503 self.inline.peek_await_event(key).await
2504 }
2505
2506 async fn await_await_event(
2507 &self,
2508 key: &AwaitEventKey,
2509 cancel: tokio_util::sync::CancellationToken,
2510 deadline: Option<std::time::Instant>,
2511 ) -> Result<Resolution, RuntimeError> {
2512 self.inline.await_await_event(key, cancel, deadline).await
2513 }
2514
2515 async fn revoke_await_events_for_session(
2516 &self,
2517 session_id: &str,
2518 ) -> Result<(), RuntimeError> {
2519 self.inline
2520 .revoke_await_events_for_session(session_id)
2521 .await
2522 }
2523
2524 async fn cancel_await_events_for_session(
2525 &self,
2526 session_id: &str,
2527 ) -> Result<(), RuntimeError> {
2528 self.inline
2529 .cancel_await_events_for_session(session_id)
2530 .await
2531 }
2532 }
2533
2534 #[test]
2535 fn agent_frame_follow_turn_ids_are_distinct_and_deterministic() {
2536 assert_eq!(agent_frame_follow_turn_id("root-turn", 0), "root-turn");
2537 assert_eq!(
2538 agent_frame_follow_turn_id("root-turn", 1),
2539 "root-turn:agent-frame:1"
2540 );
2541 assert_eq!(
2542 agent_frame_follow_turn_id("root-turn", 2),
2543 "root-turn:agent-frame:2"
2544 );
2545 }
2546
2547 #[tokio::test]
2548 async fn cancellation_watch_retries_transient_errors_until_evidence_arrives() {
2549 let attempts = Arc::new(AtomicUsize::new(0));
2550 let observed_attempts = Arc::clone(&attempts);
2551 let evidence = await_turn_cancellation_with_retry(move || {
2552 let attempt = observed_attempts.fetch_add(1, Ordering::SeqCst);
2553 async move {
2554 if attempt < 2 {
2555 Err(RuntimeError::new(
2556 "transient_cancel_watch",
2557 "temporary ingress failure",
2558 ))
2559 } else {
2560 Ok(Some(TurnCancellationEvidence {
2561 request_id: "retry-request".to_string(),
2562 origin: Some("test-user".to_string()),
2563 reason: None,
2564 }))
2565 }
2566 }
2567 })
2568 .await
2569 .expect("cancellation evidence after retries");
2570
2571 assert_eq!(attempts.load(Ordering::SeqCst), 3);
2572 assert_eq!(evidence.request_id, "retry-request");
2573 }
2574
2575 #[tokio::test]
2576 async fn cancellation_start_gate_fails_after_bounded_retries() {
2577 let attempts = Arc::new(AtomicUsize::new(0));
2578 let observed_attempts = Arc::clone(&attempts);
2579 let err = await_turn_cancellation_start_gate(move || {
2580 observed_attempts.fetch_add(1, Ordering::SeqCst);
2581 async {
2582 Err(RuntimeError::new(
2583 "cancel_start_gate_unavailable",
2584 "temporary ingress failure",
2585 ))
2586 }
2587 })
2588 .await
2589 .expect_err("start gate must fail closed after its retry budget");
2590
2591 assert_eq!(
2592 attempts.load(Ordering::SeqCst),
2593 TURN_CANCEL_START_GATE_ATTEMPTS
2594 );
2595 assert_eq!(err.code.to_string(), "cancel_start_gate_unavailable");
2596 }
2597
2598 #[tokio::test]
2599 async fn terminal_publication_failure_is_non_fatal_after_commit() {
2600 let resolver = RejectTerminalPublication::default();
2601 let control = ActiveTurnControl::new(
2602 &resolver,
2603 TurnAddress::new("committed-session", "committed-turn"),
2604 )
2605 .await
2606 .expect("active turn control");
2607 publish_terminal_after_commit(
2608 &control,
2609 &resolver,
2610 &TurnTerminal::Committed {
2611 outcome: TurnOutcome::Finished(TurnFinish::AssistantMessage {
2612 text: "committed".to_string(),
2613 }),
2614 cancellation: None,
2615 session_revision: Some(1),
2616 },
2617 "committed-session",
2618 "committed-turn",
2619 )
2620 .await;
2621 assert_eq!(resolver.attempts.load(Ordering::SeqCst), 1);
2622 }
2623}