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