1use std::sync::Arc;
2
3use tokio::sync::mpsc::Sender;
4use tokio_util::sync::CancellationToken;
5
6use crate::tool_dispatch::ToolDispatchContext;
7use crate::{TurnActivity, TurnActivityId, TurnEvent};
8
9#[derive(Clone)]
10pub struct RuntimeExecutionContext<'run> {
11 pub(super) session_id: String,
12 pub(super) dispatch: Arc<ToolDispatchContext<'run>>,
13 process_env_store: Arc<dyn crate::ProcessExecutionEnvStore>,
14 attachment_store: Arc<dyn crate::AttachmentStore>,
15 chronological_projection: Arc<crate::ChronologicalProjection>,
16 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
17 turn_context: crate::TurnContext,
18 execution_env_spec: crate::ProcessExecutionEnvSpec,
19 process_originator: Option<crate::ProcessOriginator>,
20 pub(super) runtime_process_id: Option<String>,
21 pub(super) process_event_context: Option<RuntimeExecutionProcessEventContext>,
22 process_env_ref: Option<crate::ProcessExecutionEnvRef>,
23 process_wake_target: Option<crate::SessionScope>,
24 pub(super) parent_invocation: Option<crate::RuntimeInvocation>,
25 turn_phase_probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
26 pub(super) turn_event_tx: Option<Sender<TurnActivity>>,
27 pub(super) cancellation_token: Option<CancellationToken>,
28 process_work_driver: Option<crate::ProcessWorkDriver>,
34 started_process_ids: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
39}
40
41#[derive(Clone)]
42pub(super) struct RuntimeExecutionProcessEventContext {
43 pub process_id: String,
44 pub registry: Arc<dyn crate::ProcessRegistry>,
45 pub awaiter: crate::ProcessAwaiter,
46 pub store: Option<Arc<dyn crate::RuntimePersistence>>,
47 pub session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
48 pub queued_work_driver: Option<crate::QueuedWorkDriver>,
49}
50
51impl<'run> RuntimeExecutionContext<'run> {
52 pub(crate) fn drain_tool_trigger_outcomes(
53 &self,
54 ) -> Result<Vec<crate::tool_dispatch::ToolTriggerEffectOutcome>, crate::PluginError> {
55 self.dispatch
56 .trigger_outcomes
57 .drain()
58 .map_err(crate::PluginError::Session)
59 }
60
61 pub(super) fn process_scope(
62 &self,
63 parent_invocation: Option<crate::RuntimeInvocation>,
64 ) -> crate::ProcessOpScope<'_> {
65 crate::ProcessOpScope::new(self.dispatch.effect_controller.scoped())
66 .with_parent_invocation(parent_invocation)
67 .with_agent_frame_id(Some(self.dispatch.agent_frame_id.clone()))
68 }
69
70 #[allow(
71 clippy::too_many_arguments,
72 reason = "code execution bridge carries explicit per-turn runtime dependencies"
73 )]
74 pub(crate) fn new(
75 session_id: String,
76 dispatch: Arc<ToolDispatchContext<'run>>,
77 process_env_store: Arc<dyn crate::ProcessExecutionEnvStore>,
78 attachment_store: Arc<dyn crate::AttachmentStore>,
79 chronological_projection: Arc<crate::ChronologicalProjection>,
80 protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
81 turn_context: crate::TurnContext,
82 ) -> Self {
83 Self {
84 session_id,
85 dispatch,
86 process_env_store,
87 attachment_store,
88 chronological_projection,
89 protocol_extension,
90 turn_context,
91 execution_env_spec: crate::ProcessExecutionEnvSpec::new(
92 crate::PluginOptions::default(),
93 crate::SessionPolicy::default(),
94 ),
95 process_originator: None,
96 runtime_process_id: None,
97 process_event_context: None,
98 started_process_ids: Arc::default(),
99 process_env_ref: None,
100 process_wake_target: None,
101 parent_invocation: None,
102 turn_phase_probe: None,
103 turn_event_tx: None,
104 cancellation_token: None,
105 process_work_driver: None,
106 }
107 }
108
109 pub fn session_id(&self) -> &str {
110 &self.session_id
111 }
112
113 pub fn execution_scope_id(&self) -> String {
114 self.dispatch
115 .effect_controller
116 .scoped()
117 .scope_id()
118 .to_string()
119 }
120
121 pub fn session_scope(&self) -> crate::SessionScope {
122 if self.dispatch.agent_frame_id.is_empty() {
123 crate::SessionScope::new(self.session_id.clone())
124 } else {
125 crate::SessionScope::for_agent_frame(
126 self.session_id.clone(),
127 self.dispatch.agent_frame_id.clone(),
128 )
129 }
130 }
131
132 pub fn trigger_store(&self) -> Option<Arc<dyn crate::TriggerStore>> {
133 self.dispatch
134 .trigger_router
135 .as_ref()
136 .map(crate::TriggerRouter::store)
137 }
138
139 pub fn trigger_registration_originator(&self) -> crate::ProcessOriginator {
140 self.process_originator
141 .clone()
142 .unwrap_or_else(|| crate::ProcessOriginator::session(self.session_scope()))
143 }
144
145 pub fn trigger_registration_wake_target(&self) -> Option<crate::SessionScope> {
146 self.process_wake_target
147 .clone()
148 .or_else(|| Some(self.session_scope()))
149 }
150
151 pub fn attachment_store(&self) -> Arc<dyn crate::AttachmentStore> {
152 Arc::clone(&self.attachment_store)
153 }
154
155 pub fn process_env_store(&self) -> Arc<dyn crate::ProcessExecutionEnvStore> {
156 Arc::clone(&self.process_env_store)
157 }
158
159 pub fn chronological_projection(&self) -> Arc<crate::ChronologicalProjection> {
160 Arc::clone(&self.chronological_projection)
161 }
162
163 pub fn protocol_extension<T: 'static>(&self) -> Option<&T> {
164 self.protocol_extension
165 .as_ref()
166 .and_then(|extension| extension.as_any().downcast_ref::<T>())
167 }
168
169 pub fn turn_context(&self) -> &crate::TurnContext {
170 &self.turn_context
171 }
172
173 pub fn tool_catalog(&self) -> Arc<crate::ToolCatalog> {
174 Arc::clone(&self.dispatch.tool_catalog)
175 }
176
177 pub(crate) fn session_graph_service(&self) -> &dyn crate::plugin::SessionGraphService {
178 self.dispatch.session_graph.as_ref()
179 }
180
181 pub(super) async fn emit_turn_activity(
182 &self,
183 correlation_id: TurnActivityId,
184 event: TurnEvent,
185 ) {
186 if let Some(tx) = &self.turn_event_tx {
187 let _ = tx.send(TurnActivity::new(correlation_id, event)).await;
188 }
189 }
190
191 pub(crate) fn with_turn_event_sender(mut self, turn_event_tx: Sender<TurnActivity>) -> Self {
192 self.turn_event_tx = Some(turn_event_tx);
193 self
194 }
195
196 pub(crate) fn with_parent_invocation(mut self, metadata: crate::RuntimeInvocation) -> Self {
197 self.parent_invocation = Some(metadata);
198 self
199 }
200
201 pub(crate) fn with_execution_env_spec(
202 mut self,
203 execution_env_spec: crate::ProcessExecutionEnvSpec,
204 ) -> Self {
205 self.execution_env_spec = execution_env_spec;
206 self
207 }
208
209 pub(crate) fn with_process_registration_context(
210 mut self,
211 registration: &crate::ProcessRegistration,
212 ) -> Self {
213 self.process_originator = Some(registration.provenance.originator.clone());
214 self.runtime_process_id = Some(registration.id.clone());
215 self.process_env_ref = registration.env_ref.clone();
216 self.process_wake_target = registration.wake_target.clone();
217 self
218 }
219
220 pub(crate) fn with_process_event_context(
221 mut self,
222 process_id: impl Into<String>,
223 registry: Arc<dyn crate::ProcessRegistry>,
224 awaiter: crate::ProcessAwaiter,
225 store: Option<Arc<dyn crate::RuntimePersistence>>,
226 session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
227 queued_work_driver: Option<crate::QueuedWorkDriver>,
228 ) -> Self {
229 self.process_event_context = Some(RuntimeExecutionProcessEventContext {
230 process_id: process_id.into(),
231 registry,
232 awaiter,
233 store,
234 session_store_factory,
235 queued_work_driver,
236 });
237 self
238 }
239
240 pub(super) fn record_started_process(&self, process_id: &str) {
244 self.started_process_ids
245 .lock()
246 .expect("started process ids lock")
247 .insert(process_id.to_string());
248 }
249
250 pub(super) fn is_run_local_process(&self, process_id: &str) -> bool {
251 self.started_process_ids
252 .lock()
253 .expect("started process ids lock")
254 .contains(process_id)
255 }
256
257 pub(crate) fn process_spawn_provenance(&self) -> Option<crate::ProcessSpawnProvenance> {
258 self.process_originator
259 .clone()
260 .map(|originator| crate::ProcessSpawnProvenance {
261 originator,
262 wake_target: self.process_wake_target.clone(),
263 })
264 }
265
266 pub(super) async fn attach_captured_process_execution_env(
267 &self,
268 registration: crate::ProcessRegistration,
269 ) -> Result<crate::ProcessRegistration, crate::PluginError> {
270 if registration.env_ref.is_some() {
271 return Ok(registration);
272 }
273 match registration.input.as_ref() {
274 crate::ProcessInput::ToolCall { .. } | crate::ProcessInput::Engine { .. } => {
275 let env_ref = self.captured_process_execution_env_ref().await?;
276 Ok(registration.with_execution_env_ref(Some(env_ref)))
277 }
278 crate::ProcessInput::External { .. } | crate::ProcessInput::SessionTurn { .. } => {
279 Ok(registration)
280 }
281 }
282 }
283
284 pub async fn captured_process_execution_env_ref(
285 &self,
286 ) -> Result<crate::ProcessExecutionEnvRef, crate::PluginError> {
287 if let Some(env_ref) = self.process_env_ref.clone() {
288 return Ok(env_ref);
289 }
290 crate::persist_process_execution_env(
291 self.process_env_store.as_ref(),
292 &self.execution_env_spec,
293 )
294 .await
295 }
296
297 pub(crate) fn with_turn_phase_probe(
298 mut self,
299 probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
300 ) -> Self {
301 self.turn_phase_probe = probe;
302 self
303 }
304
305 #[doc(hidden)]
306 pub fn named_phase(&self, phase: &'static str) -> crate::runtime::RuntimeNamedPhase {
307 crate::runtime::RuntimeNamedPhase::begin(self.turn_phase_probe.clone(), phase)
308 }
309
310 pub fn parent_invocation(&self) -> Option<&crate::RuntimeInvocation> {
311 self.parent_invocation.as_ref()
312 }
313
314 pub(crate) fn with_cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
315 self.cancellation_token = Some(cancellation_token);
316 self
317 }
318
319 pub(crate) fn with_process_work_driver(
320 mut self,
321 process_work_driver: Option<crate::ProcessWorkDriver>,
322 ) -> Self {
323 self.process_work_driver = process_work_driver;
324 self
325 }
326
327 pub(crate) fn tool_scheduling(&self, name: &str) -> crate::ToolScheduling {
328 crate::tool_dispatch::resolve_tool_scheduling(&self.dispatch, name)
329 }
330
331 pub fn callable_tool_manifest(&self, name: &str) -> Option<crate::ToolManifest> {
332 crate::tool_dispatch::resolve_callable_manifest(&self.dispatch, name)
333 }
334
335 pub fn callable_tool_manifest_by_id(&self, id: &crate::ToolId) -> Option<crate::ToolManifest> {
336 crate::tool_dispatch::resolve_callable_manifest_by_id(&self.dispatch, id)
337 }
338
339 pub fn tool_argument_projection_policy(
340 &self,
341 name: &str,
342 ) -> crate::ToolArgumentProjectionPolicy {
343 crate::tool_dispatch::resolve_tool_argument_projection_policy(&self.dispatch, name)
344 }
345
346 pub async fn start_child_process(
347 &self,
348 registration: crate::ProcessRegistration,
349 kind: impl Into<String>,
350 label: Option<String>,
351 ) -> crate::ToolInvocationReply {
352 let _phase = self.named_phase("process.start_child");
353 let registration = match self
354 .attach_captured_process_execution_env(registration)
355 .await
356 {
357 Ok(registration) => registration,
358 Err(err) => {
359 return crate::ToolInvocationReply::error(serde_json::json!(err.to_string()));
360 }
361 };
362 let process_id = registration.id.clone();
363 let mut options = crate::ProcessStartOptions::new()
364 .with_descriptor(crate::ProcessHandleDescriptor::new(Some(kind), label));
365 if let Some(spawn) = self.process_spawn_provenance() {
366 options = options.with_spawn_provenance(spawn);
367 }
368 match self
369 .dispatch
370 .processes
371 .start(
372 &self.session_id,
373 registration,
374 options,
375 self.process_scope(self.parent_invocation.clone()),
376 )
377 .await
378 {
379 Ok(_) => {
380 self.record_started_process(&process_id);
381 crate::ToolInvocationReply::success(Self::process_handle_json(&process_id))
382 }
383 Err(err) => crate::ToolInvocationReply::error(serde_json::json!(err.to_string())),
384 }
385 }
386
387 pub async fn sleep_process(
388 &self,
389 scope: &str,
390 sequence: u64,
391 duration_ms: u64,
392 ) -> Result<(), crate::RuntimeEffectControllerError> {
393 let cancellation = self.cancellation_token.clone().unwrap_or_default();
394 let invocation = crate::runtime::causal::process_sleep_invocation(
395 &self.session_id,
396 self.parent_invocation.as_ref(),
397 scope,
398 sequence,
399 );
400 let outcome = self
401 .dispatch
402 .effect_controller
403 .controller()
404 .execute_effect(
405 crate::RuntimeEffectEnvelope::new(
406 invocation,
407 crate::RuntimeEffectCommand::Sleep { duration_ms },
408 ),
409 crate::RuntimeEffectLocalExecutor::sleep_with_clock(
410 cancellation,
411 std::sync::Arc::clone(&self.dispatch.clock),
412 ),
413 )
414 .await?;
415 match outcome {
416 crate::RuntimeEffectOutcome::Sleep => Ok(()),
417 other => Err(crate::RuntimeEffectControllerError::new(
418 "runtime_effect_wrong_outcome",
419 format!("expected sleep outcome, got {}", other.kind().as_str()),
420 )),
421 }
422 }
423
424 pub async fn await_process_signal_event(
425 &self,
426 process_id: &str,
427 signal_name: &str,
428 event_ordinal: u64,
429 ) -> Result<serde_json::Value, crate::RuntimeEffectControllerError> {
430 let cancellation = self.cancellation_token.clone().unwrap_or_default();
431 let key = self
432 .dispatch
433 .effect_controller
434 .controller()
435 .await_event_key(
436 &crate::ExecutionScope::process(process_id),
437 crate::AwaitEventWaitIdentity::process_signal(
438 process_id,
439 signal_name,
440 event_ordinal,
441 ),
442 )
443 .await?;
444 let invocation = crate::runtime::causal::process_await_event_invocation(
445 &self.session_id,
446 self.parent_invocation.as_ref(),
447 process_id,
448 signal_name,
449 event_ordinal,
450 );
451 let outcome = self
452 .dispatch
453 .effect_controller
454 .controller()
455 .execute_effect(
456 crate::RuntimeEffectEnvelope::new(
457 invocation,
458 crate::RuntimeEffectCommand::AwaitEvent { key },
459 ),
460 crate::RuntimeEffectLocalExecutor::await_event_with_clock(
461 cancellation,
462 None,
463 std::sync::Arc::clone(&self.dispatch.clock),
464 ),
465 )
466 .await?;
467 match outcome.into_await_event()? {
468 crate::Resolution::Ok(value) => Ok(value),
469 crate::Resolution::Err(err) => Err(crate::RuntimeEffectControllerError::new(
470 err.code,
471 err.message,
472 )),
473 crate::Resolution::Timeout => Err(crate::RuntimeEffectControllerError::new(
474 "process_signal_wait_timeout",
475 "process signal wait timed out",
476 )),
477 crate::Resolution::Cancelled => Err(crate::RuntimeEffectControllerError::new(
478 "process_signal_wait_cancelled",
479 "process signal wait was cancelled",
480 )),
481 }
482 }
483
484 pub async fn signal_process_by_id(
485 &self,
486 registry: Arc<dyn crate::ProcessRegistry>,
487 process_id: &str,
488 signal_name: &str,
489 signal_id: String,
490 payload: serde_json::Value,
491 ) -> Result<crate::ProcessEvent, crate::RuntimeEffectControllerError> {
492 let event_type = crate::process_signal_event_type(signal_name)?;
493 let replay_key = format!("process:{process_id}:signal.{signal_name}:{signal_id}");
494 let signal_payload = payload.clone();
495 let command = crate::ProcessCommand::Signal {
496 process_id: process_id.to_string(),
497 signal_name: signal_name.to_string(),
498 signal_id,
499 request: crate::ProcessEventAppendRequest::new(event_type.clone(), payload)
500 .with_replay_key(replay_key),
501 };
502 let effect_id = command.effect_id();
503 let invocation = crate::runtime::causal::process_effect_invocation(
504 &self.session_id,
505 self.parent_invocation.clone(),
506 &effect_id,
507 );
508 let outcome = self
509 .dispatch
510 .effect_controller
511 .controller()
512 .execute_effect(
513 crate::RuntimeEffectEnvelope::new(
514 invocation,
515 crate::RuntimeEffectCommand::process(command),
516 ),
517 crate::RuntimeEffectLocalExecutor::processes(
518 Arc::clone(®istry),
519 self.process_work_driver.clone(),
520 ),
521 )
522 .await?;
523 match outcome.into_process()? {
524 crate::ProcessEffectOutcome::Signal { event } => {
525 let waiting_ordinal =
526 registry
527 .get_process(process_id)
528 .await
529 .and_then(|record| match record.wait {
530 Some(crate::WaitState {
531 kind:
532 crate::WaitKind::Signal {
533 name,
534 event_type: wait_event_type,
535 ordinal,
536 ..
537 },
538 ..
539 }) if name == signal_name && wait_event_type == event_type => {
540 Some(ordinal)
541 }
542 _ => None,
543 });
544 let ordinal = match waiting_ordinal {
545 Some(ordinal) => ordinal,
546 None => {
547 registry
548 .count_events_through(process_id, &event_type, event.sequence)
549 .await?
550 }
551 };
552 if ordinal > 0 {
553 let key = self
554 .dispatch
555 .effect_controller
556 .controller()
557 .await_event_key(
558 &crate::ExecutionScope::process(process_id),
559 crate::AwaitEventWaitIdentity::process_signal(
560 process_id,
561 signal_name,
562 ordinal,
563 ),
564 )
565 .await?;
566 let _ = self
567 .dispatch
568 .effect_controller
569 .controller()
570 .resolve_await_event(&key, crate::Resolution::Ok(signal_payload))
571 .await?;
572 }
573 Ok(event)
574 }
575 other => Err(crate::RuntimeEffectControllerError::new(
576 "runtime_effect_wrong_outcome",
577 format!("expected signal outcome, got {other:?}"),
578 )),
579 }
580 }
581
582 pub async fn append_process_event(
583 &self,
584 registry: Arc<dyn crate::ProcessRegistry>,
585 process_id: &str,
586 request: crate::ProcessEventAppendRequest,
587 ) -> Result<crate::ProcessEvent, crate::PluginError> {
588 let result = registry.append_event(process_id, request).await?;
589 if let Some(context) = self.process_event_context.as_ref() {
590 crate::tool_provider::process_events::enqueue_wake_delivery(
591 context.store.clone(),
592 context.session_store_factory.as_ref(),
593 result.wake_delivery,
594 Some(self.session_graph_service()),
595 context.queued_work_driver.as_ref(),
596 )
597 .await?;
598 }
599 Ok(result.event)
600 }
601}
602
603#[cfg(test)]
604mod tests {
605 use super::*;
606 use crate::tool_dispatch::ToolDispatchContext;
607 use crate::{ToolCall, ToolProvider, ToolResult};
608
609 struct NoopTools;
610
611 #[async_trait::async_trait]
612 impl ToolProvider for NoopTools {
613 fn tool_manifests(&self) -> Vec<crate::ToolManifest> {
614 Vec::new()
615 }
616
617 fn resolve_contract(&self, _name: &str) -> Option<Arc<crate::ToolContract>> {
618 None
619 }
620
621 async fn execute(&self, _call: ToolCall<'_>) -> ToolResult {
622 ToolResult::err_fmt("not used")
623 }
624 }
625
626 #[test]
627 fn tool_argument_projection_policy_resolves_from_active_catalog_and_defaults_unknown() {
628 let tool = crate::ToolDefinition::raw(
629 "tool:seedy",
630 "seedy",
631 "Seed-aware",
632 crate::ToolDefinition::default_input_schema(),
633 serde_json::json!({ "type": "string" }),
634 )
635 .with_argument_projection(
636 crate::ToolArgumentProjectionPolicy::preserve_projected_refs_in_field("seed"),
637 );
638 let plugins = crate::plugin::PluginHost::empty()
639 .build_session("session", None)
640 .expect("plugin session");
641 let (event_tx, _event_rx) = tokio::sync::mpsc::channel(1);
642 let dispatch = Arc::new(ToolDispatchContext {
643 plugins,
644 tools: Arc::new(NoopTools),
645 tool_catalog: Arc::new(crate::ToolCatalog::from_tools(
646 vec![tool.manifest()],
647 std::collections::BTreeMap::new(),
648 )),
649 sessions: Arc::new(crate::testing::MockSessionManager::default()),
650 session_lifecycle: Arc::new(crate::testing::MockSessionManager::default()),
651 session_graph: Arc::new(crate::testing::MockSessionManager::default()),
652 processes: Arc::new(crate::UnavailableProcessService),
653 process_cancel_ability: Arc::new(crate::DefaultProcessCancelAbility),
654 trigger_router: None,
655 effect_controller: crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
656 crate::InlineRuntimeEffectController,
657 )),
658 direct_completions: crate::DirectCompletionClient::unavailable(
659 "direct completions are unavailable in this test context",
660 ),
661 parent_invocation: None,
662 execution_env_spec: crate::ProcessExecutionEnvSpec::new(
663 crate::PluginOptions::default(),
664 crate::SessionPolicy::default(),
665 ),
666 session_id: "session".to_string(),
667 agent_frame_id: String::new(),
668 event_tx,
669 checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer::default(),
670 trigger_outcomes: crate::tool_dispatch::ToolTriggerOutcomeBuffer::default(),
671 attachment_store: Arc::new(crate::InMemoryAttachmentStore::new()),
672 turn_context: crate::TurnContext::default(),
673 clock: std::sync::Arc::new(crate::SystemClock),
674 });
675 let ctx = RuntimeExecutionContext::new(
676 "session".to_string(),
677 dispatch,
678 Arc::new(crate::InMemoryProcessExecutionEnvStore::new()),
679 Arc::new(crate::InMemoryAttachmentStore::new()),
680 Arc::new(crate::ChronologicalProjection::default()),
681 None,
682 crate::TurnContext::default(),
683 );
684
685 assert_eq!(
686 ctx.tool_argument_projection_policy("seedy"),
687 crate::ToolArgumentProjectionPolicy::preserve_projected_refs_in_field("seed")
688 );
689 assert_eq!(
690 ctx.tool_argument_projection_policy("missing"),
691 crate::ToolArgumentProjectionPolicy::MaterializeProjectedValues
692 );
693 }
694}