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