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