Skip to main content

lash_core/session/
tool_execution.rs

1use super::execution_context::RuntimeExecutionContext;
2use crate::tool_dispatch::{
3    ToolAttemptEffectIdentity, ToolCallLaunch, ToolDispatchOutcome, ToolPreparationOutcome,
4    coordinate_tool_invocation, finalize_tool_result_with_execution_context,
5    prepare_granted_tool_call_with_context, prepare_tool_call_with_context, schedule_tool_batch,
6};
7use crate::{
8    ModelToolReturn, SessionStreamEvent, ToolCallOutput, ToolCallRecord, ToolCancellation,
9    ToolFailure, ToolFailureClass, ToolResult, TurnActivityId, TurnEvent,
10};
11use std::collections::HashMap;
12
13#[derive(Clone)]
14pub struct ToolInvocation {
15    pub id: String,
16    pub tool_id: crate::ToolId,
17    pub args: serde_json::Value,
18    pub execution_grant: Option<Box<crate::ToolExecutionGrant>>,
19    pub child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
20}
21
22impl ToolInvocation {
23    pub fn new(id: impl Into<String>, tool_id: crate::ToolId, args: serde_json::Value) -> Self {
24        Self {
25            id: id.into(),
26            tool_id,
27            args,
28            execution_grant: None,
29            child_execution_trace_hook: None,
30        }
31    }
32
33    pub fn with_execution_grant(mut self, grant: crate::ToolExecutionGrant) -> Self {
34        self.execution_grant = Some(Box::new(grant));
35        self
36    }
37
38    pub fn label(&self) -> String {
39        self.tool_id.to_string()
40    }
41
42    pub fn with_child_execution_trace_hook(
43        mut self,
44        hook: crate::ToolChildExecutionTraceHook,
45    ) -> Self {
46        self.child_execution_trace_hook = Some(hook);
47        self
48    }
49}
50
51impl std::fmt::Debug for ToolInvocation {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("ToolInvocation")
54            .field("id", &self.id)
55            .field("tool_id", &self.tool_id)
56            .field("args", &self.args)
57            .field(
58                "execution_grant",
59                &self.execution_grant.as_ref().map(|_| "<grant>"),
60            )
61            .field(
62                "child_execution_trace_hook",
63                &self.child_execution_trace_hook.as_ref().map(|_| "<hook>"),
64            )
65            .finish()
66    }
67}
68
69#[derive(Clone, Debug)]
70pub struct ToolInvocationReply {
71    pub output: ToolCallOutput,
72    pub record: Option<ToolCallRecord>,
73}
74
75impl ToolInvocationReply {
76    pub fn success(value: serde_json::Value) -> Self {
77        Self {
78            output: ToolCallOutput::success(value),
79            record: None,
80        }
81    }
82
83    pub fn error(value: serde_json::Value) -> Self {
84        let message = value
85            .as_str()
86            .map(ToOwned::to_owned)
87            .unwrap_or_else(|| value.to_string());
88        let mut failure = ToolFailure::tool(ToolFailureClass::Execution, "tool_error", message);
89        failure.raw =
90            Some(serde_json::from_value(value).unwrap_or_else(|_| {
91                crate::ToolValue::String("unserializable tool error".to_string())
92            }));
93        Self {
94            output: ToolCallOutput::failure(failure),
95            record: None,
96        }
97    }
98
99    pub fn from_output(output: ToolCallOutput) -> Self {
100        Self {
101            output,
102            record: None,
103        }
104    }
105
106    pub fn cancelled(message: impl Into<String>) -> Self {
107        Self::from_output(ToolCallOutput::cancelled(ToolCancellation::runtime(
108            message,
109        )))
110    }
111
112    pub(crate) fn with_record(mut self, record: ToolCallRecord) -> Self {
113        self.record = Some(record);
114        self
115    }
116}
117
118#[derive(Clone, Debug)]
119pub(crate) struct CompletedProtocolToolCall {
120    pub completed: crate::sansio::CompletedToolCall,
121    pub record: ToolCallRecord,
122}
123
124fn cancelled_runtime_tool_call_launch(
125    call_id: String,
126    tool_name: String,
127    args: serde_json::Value,
128    replay: Option<crate::llm::types::ProviderReplayMeta>,
129) -> crate::runtime::ToolCallLaunch {
130    crate::runtime::ToolCallLaunch::Done {
131        result: cancelled_completed_tool_call(call_id, tool_name, args, replay),
132    }
133}
134
135fn cancelled_completed_tool_call(
136    call_id: String,
137    tool_name: String,
138    args: serde_json::Value,
139    replay: Option<crate::llm::types::ProviderReplayMeta>,
140) -> crate::sansio::CompletedToolCall {
141    let output = ToolCallOutput::cancelled(ToolCancellation::runtime("tool call cancelled"));
142    crate::sansio::CompletedToolCall {
143        call_id: call_id.clone(),
144        tool_name: tool_name.clone(),
145        args,
146        model_return: ModelToolReturn {
147            call_id,
148            tool_name,
149            parts: vec![crate::ModelToolReturnPart::text(
150                "[Tool execution cancelled]\ntool call cancelled".to_string(),
151            )],
152        },
153        output,
154        duration_ms: 0,
155        replay,
156    }
157}
158
159fn deterministic_tool_invocation_batch_id(calls: &[ToolInvocation]) -> String {
160    let identity = calls
161        .iter()
162        .map(|call| {
163            serde_json::json!({
164                "id": call.id.clone(),
165                "tool_id": call.tool_id.to_string(),
166                "args": call.args.clone(),
167                "execution_grant": call.execution_grant.as_ref().map(|grant| serde_json::json!({
168                    "tool_id": grant.manifest.id.to_string(),
169                    "source_id": grant.source_id.clone(),
170                    "execution_binding": grant.execution_binding.clone(),
171                })),
172            })
173        })
174        .collect::<Vec<_>>();
175    let digest = crate::stable_hash::stable_json_sha256_hex(&identity)
176        .unwrap_or_else(|_| format!("len-{}", calls.len()));
177    format!("tool-batch:{digest}")
178}
179
180struct CoordinatedToolLaunch {
181    launch: crate::runtime::ToolCallLaunch,
182    triggers: Vec<crate::tool_dispatch::ToolTriggerEffectOutcome>,
183}
184
185impl RuntimeExecutionContext<'_> {
186    fn tool_batch_invocation(&self, batch_id: &str) -> crate::RuntimeInvocation {
187        let suffix = format!("tool-batch:{batch_id}");
188        if let Some(parent) = self.parent_invocation.as_ref() {
189            let parent_effect_id = parent.effect_id().unwrap_or("effect");
190            return crate::runtime::causal::child_effect_invocation(
191                parent,
192                format!("{parent_effect_id}:{suffix}"),
193                crate::RuntimeEffectKind::ToolBatch,
194                suffix,
195            );
196        }
197        let replay_key = format!("{}:{suffix}", self.execution_scope_id());
198        crate::RuntimeInvocation::effect(
199            crate::RuntimeScope::new(self.session_id.clone()),
200            suffix,
201            crate::RuntimeEffectKind::ToolBatch,
202            replay_key,
203        )
204    }
205
206    pub(crate) async fn execute_prepared_tool_batch_launches(
207        &self,
208        batch: crate::PreparedToolBatch,
209        parent_invocation: crate::RuntimeInvocation,
210        child_trace_hooks: HashMap<String, crate::ToolChildExecutionTraceHook>,
211    ) -> Result<crate::ToolBatchEffectOutcome, crate::RuntimeEffectControllerError> {
212        let indexed_tools = batch.calls.into_iter().enumerate().collect::<Vec<_>>();
213        let cancellation = self.cancellation_token.clone().unwrap_or_default();
214        let tool_cancel = cancellation.child_token();
215        let child_trace_hooks = std::sync::Arc::new(child_trace_hooks);
216        if !self
217            .dispatch
218            .effect_controller
219            .controller()
220            .supports_concurrent_effects()
221        {
222            let mut launches = Vec::with_capacity(indexed_tools.len());
223            let mut triggers = Vec::new();
224            let mut context = self.clone().with_cancellation_token(tool_cancel.clone());
225            for (index, child) in indexed_tools {
226                if cancellation.is_cancelled() {
227                    tool_cancel.cancel();
228                    launches.push(cancelled_runtime_tool_call_launch(
229                        child.call.call_id,
230                        child.call.tool_name,
231                        child.call.args,
232                        child.call.replay,
233                    ));
234                    continue;
235                }
236                let child_execution_trace_hook =
237                    child_trace_hooks.get(&child.call.call_id).cloned();
238                let outcome = context
239                    .execute_prepared_tool_batch_child(
240                        child,
241                        index,
242                        parent_invocation.clone(),
243                        child_execution_trace_hook,
244                    )
245                    .await;
246                launches.push(outcome.launch);
247                triggers.extend(outcome.triggers);
248                context = context.with_cancellation_token(tool_cancel.clone());
249            }
250            return Ok(crate::ToolBatchEffectOutcome { launches, triggers });
251        }
252        let child_outcomes = schedule_tool_batch(
253            indexed_tools,
254            |(index, _)| *index,
255            |(_, child)| self.tool_scheduling(&child.call.tool_name),
256            {
257                let context = self.clone();
258                let cancellation = cancellation.clone();
259                let tool_cancel = tool_cancel.clone();
260                let child_trace_hooks = std::sync::Arc::clone(&child_trace_hooks);
261                move |(index, child)| {
262                    let context = context.clone().with_cancellation_token(tool_cancel.clone());
263                    let cancellation = cancellation.clone();
264                    let tool_cancel = tool_cancel.clone();
265                    let parent_invocation = parent_invocation.clone();
266                    let cancelled_tool = child.call.clone();
267                    let child_execution_trace_hook =
268                        child_trace_hooks.get(&child.call.call_id).cloned();
269                    async move {
270                        let tool_call = context.execute_prepared_tool_batch_child(
271                            child,
272                            index,
273                            parent_invocation,
274                            child_execution_trace_hook,
275                        );
276                        tokio::pin!(tool_call);
277                        tokio::select! {
278                            biased;
279                            _ = cancellation.cancelled() => {
280                                tool_cancel.cancel();
281                                let grace = context
282                                    .dispatch
283                                    .clock
284                                    .sleep(std::time::Duration::from_millis(50));
285                                tokio::pin!(grace);
286                                tokio::select! {
287                                    biased;
288                                    outcome = &mut tool_call => outcome,
289                                    _ = &mut grace => CoordinatedToolLaunch {
290                                        launch: cancelled_runtime_tool_call_launch(
291                                            cancelled_tool.call_id,
292                                            cancelled_tool.tool_name,
293                                            cancelled_tool.args,
294                                            cancelled_tool.replay,
295                                        ),
296                                        triggers: Vec::new(),
297                                    },
298                                }
299                            }
300                            outcome = &mut tool_call => outcome,
301                        }
302                    }
303                }
304            },
305        )
306        .await;
307
308        let mut launches = Vec::with_capacity(child_outcomes.len());
309        let mut triggers = Vec::new();
310        for outcome in child_outcomes {
311            launches.push(outcome.launch);
312            triggers.extend(outcome.triggers);
313        }
314        Ok(crate::ToolBatchEffectOutcome { launches, triggers })
315    }
316
317    async fn execute_prepared_tool_batch_child(
318        &self,
319        child: crate::PreparedToolBatchCall,
320        index: usize,
321        parent_invocation: crate::RuntimeInvocation,
322        child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
323    ) -> CoordinatedToolLaunch {
324        let call_id = child.call.call_id.clone();
325        let tool_name = child.call.tool_name.clone();
326        let args = child.call.args.clone();
327        let replay = child.call.replay.clone();
328        let activity_id = TurnActivityId::new(format!("tool:{call_id}"));
329        self.emit_tool_call_started(&call_id, &tool_name, args.clone(), activity_id.clone())
330            .await;
331
332        let retry_policy = crate::tool_dispatch::resolve_callable_manifest_by_id(
333            self.dispatch.as_ref(),
334            &child.call.tool_id,
335        )
336        .map(|manifest| manifest.retry_policy)
337        .or_else(|| {
338            child
339                .execution_grant
340                .as_ref()
341                .map(|grant| grant.manifest.retry_policy)
342        })
343        .unwrap_or(crate::ToolRetryPolicy::Never);
344        let trace_hooks: HashMap<String, crate::ToolChildExecutionTraceHook> =
345            child_execution_trace_hook
346                .map(|hook| std::iter::once((call_id.clone(), hook)).collect())
347                .unwrap_or_default();
348        let coordinated = coordinate_tool_invocation(
349            self.dispatch.as_ref(),
350            child.call.clone(),
351            child.execution_grant,
352            retry_policy,
353            ToolAttemptEffectIdentity::Batch {
354                parent: parent_invocation.clone(),
355                replay_suffix: child.replay_suffix.clone(),
356            },
357            self.cancellation_token.clone(),
358            || crate::RuntimeEffectLocalExecutor::tool_batch(self.clone(), trace_hooks.clone()),
359        )
360        .await;
361        let outcome = match coordinated.launch {
362            ToolCallLaunch::Done(outcome) => outcome,
363            ToolCallLaunch::Pending(pending) => {
364                self.await_pending_tool_dispatch_outcome_with_suffix(
365                    &call_id,
366                    Some(parent_invocation),
367                    format!("{}:await", child.replay_suffix),
368                    pending,
369                    self.cancellation_token.clone(),
370                )
371                .await
372            }
373        };
374        let completed = self
375            .complete_tool_call(index, call_id, replay, outcome, activity_id)
376            .await;
377        CoordinatedToolLaunch {
378            launch: crate::runtime::ToolCallLaunch::Done {
379                result: completed.completed,
380            },
381            triggers: coordinated.triggers,
382        }
383    }
384
385    async fn emit_tool_call_started(
386        &self,
387        call_id: &str,
388        name: &str,
389        args: serde_json::Value,
390        activity_id: TurnActivityId,
391    ) {
392        let _ = self
393            .dispatch
394            .event_tx
395            .send(SessionStreamEvent::ToolCallStart {
396                call_id: Some(call_id.to_string()),
397                name: name.to_string(),
398                args: args.clone(),
399            })
400            .await;
401        self.emit_tool_call_started_trace(call_id, name, &args);
402        self.emit_turn_activity(
403            activity_id,
404            TurnEvent::ToolCallStarted {
405                call_id: Some(call_id.to_string()),
406                name: name.to_string(),
407                args,
408                graph_key: self.code_block_graph_key(),
409                parent_call_id: self.batch_parent_call_id(),
410            },
411        )
412        .await;
413    }
414
415    #[expect(
416        clippy::too_many_arguments,
417        reason = "tool execution carries explicit runtime call metadata"
418    )]
419    pub(crate) async fn execute_tool_call(
420        &self,
421        call_id: String,
422        name: String,
423        args: serde_json::Value,
424        index: usize,
425        replay: Option<crate::llm::types::ProviderReplayMeta>,
426        parent_invocation: Option<crate::RuntimeInvocation>,
427        child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
428    ) -> CompletedProtocolToolCall {
429        let _ = self
430            .dispatch
431            .event_tx
432            .send(SessionStreamEvent::ToolCallStart {
433                call_id: Some(call_id.clone()),
434                name: name.clone(),
435                args: args.clone(),
436            })
437            .await;
438        let tool_correlation_id = TurnActivityId::new(format!("tool:{call_id}"));
439        self.emit_tool_call_started_trace(&call_id, &name, &args);
440        self.emit_turn_activity(
441            tool_correlation_id.clone(),
442            TurnEvent::ToolCallStarted {
443                call_id: Some(call_id.clone()),
444                name: name.clone(),
445                args: args.clone(),
446                graph_key: self.code_block_graph_key(),
447                parent_call_id: self.batch_parent_call_id(),
448            },
449        )
450        .await;
451
452        let parent_invocation = parent_invocation.or_else(|| self.parent_invocation.clone());
453        let mut dispatch = (*self.dispatch).clone();
454        dispatch.parent_invocation = parent_invocation.clone();
455        let pending = crate::sansio::PendingToolCall {
456            call_id: call_id.clone(),
457            tool_name: name,
458            args,
459            replay: replay.clone(),
460        };
461        let launch =
462            match prepare_tool_call_with_context(&dispatch, pending, Some(call_id.clone())).await {
463                ToolPreparationOutcome::Prepared(prepared) => {
464                    let retry_policy = crate::tool_dispatch::resolve_callable_manifest_by_id(
465                        &dispatch,
466                        &prepared.tool_id,
467                    )
468                    .map(|manifest| manifest.retry_policy)
469                    .unwrap_or(crate::ToolRetryPolicy::Never);
470                    let trace_hooks: HashMap<String, crate::ToolChildExecutionTraceHook> =
471                        child_execution_trace_hook
472                            .map(|hook| std::iter::once((call_id.clone(), hook)).collect())
473                            .unwrap_or_default();
474                    let coordinated = coordinate_tool_invocation(
475                        &dispatch,
476                        prepared,
477                        None,
478                        retry_policy,
479                        ToolAttemptEffectIdentity::Scalar {
480                            parent: parent_invocation.clone(),
481                        },
482                        self.cancellation_token.clone(),
483                        || {
484                            crate::RuntimeEffectLocalExecutor::tool_batch(
485                                self.clone(),
486                                trace_hooks.clone(),
487                            )
488                        },
489                    )
490                    .await;
491                    self.restore_tool_trigger_outcomes(coordinated.triggers);
492                    coordinated.launch
493                }
494                ToolPreparationOutcome::Completed(outcome) => ToolCallLaunch::Done(*outcome),
495            };
496        let mut outcome = match launch {
497            ToolCallLaunch::Done(outcome) => outcome,
498            ToolCallLaunch::Pending(pending) => {
499                self.await_pending_tool_dispatch_outcome(
500                    &call_id,
501                    parent_invocation.clone(),
502                    pending,
503                    self.cancellation_token.clone(),
504                )
505                .await
506            }
507        };
508        outcome.record.call_id = Some(call_id.clone());
509
510        self.complete_tool_call(index, call_id, replay, outcome, tool_correlation_id)
511            .await
512    }
513
514    #[expect(
515        clippy::too_many_arguments,
516        reason = "tool execution carries explicit runtime call metadata"
517    )]
518    pub(crate) async fn execute_tool_call_by_id(
519        &self,
520        call_id: String,
521        tool_id: crate::ToolId,
522        args: serde_json::Value,
523        index: usize,
524        replay: Option<crate::llm::types::ProviderReplayMeta>,
525        parent_invocation: Option<crate::RuntimeInvocation>,
526        child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
527    ) -> CompletedProtocolToolCall {
528        let Some(manifest) =
529            crate::tool_dispatch::resolve_callable_manifest_by_id(self.dispatch.as_ref(), &tool_id)
530        else {
531            let outcome = ToolDispatchOutcome {
532                record: ToolCallRecord {
533                    call_id: Some(call_id.clone()),
534                    tool: tool_id.to_string(),
535                    args,
536                    output: ToolCallOutput::failure(ToolFailure::runtime(
537                        ToolFailureClass::Unavailable,
538                        "tool_unavailable",
539                        format!("Tool id `{tool_id}` is unavailable in this session"),
540                    )),
541                    duration_ms: 0,
542                },
543            };
544            let activity_id = TurnActivityId::new(format!("tool:{call_id}"));
545            return self
546                .complete_tool_call(index, call_id, replay, outcome, activity_id)
547                .await;
548        };
549        self.execute_tool_call(
550            call_id,
551            manifest.name,
552            args,
553            index,
554            replay,
555            parent_invocation,
556            child_execution_trace_hook,
557        )
558        .await
559    }
560
561    #[expect(
562        clippy::too_many_arguments,
563        reason = "tool execution carries explicit runtime call metadata"
564    )]
565    pub(crate) async fn execute_tool_call_by_grant(
566        &self,
567        call_id: String,
568        grant: crate::ToolExecutionGrant,
569        args: serde_json::Value,
570        index: usize,
571        replay: Option<crate::llm::types::ProviderReplayMeta>,
572        parent_invocation: Option<crate::RuntimeInvocation>,
573        child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
574    ) -> CompletedProtocolToolCall {
575        let name = grant.manifest.name.clone();
576        let _ = self
577            .dispatch
578            .event_tx
579            .send(SessionStreamEvent::ToolCallStart {
580                call_id: Some(call_id.clone()),
581                name: name.clone(),
582                args: args.clone(),
583            })
584            .await;
585        let tool_correlation_id = TurnActivityId::new(format!("tool:{call_id}"));
586        self.emit_tool_call_started_trace(&call_id, &name, &args);
587        self.emit_turn_activity(
588            tool_correlation_id.clone(),
589            TurnEvent::ToolCallStarted {
590                call_id: Some(call_id.clone()),
591                name: name.clone(),
592                args: args.clone(),
593                graph_key: self.code_block_graph_key(),
594                parent_call_id: self.batch_parent_call_id(),
595            },
596        )
597        .await;
598
599        let parent_invocation = parent_invocation.or_else(|| self.parent_invocation.clone());
600        let mut dispatch = (*self.dispatch).clone();
601        dispatch.parent_invocation = parent_invocation.clone();
602        let pending = crate::sansio::PendingToolCall {
603            call_id: call_id.clone(),
604            tool_name: name,
605            args,
606            replay: replay.clone(),
607        };
608        let launch = match prepare_granted_tool_call_with_context(
609            &dispatch,
610            &grant,
611            pending,
612            Some(call_id.clone()),
613        )
614        .await
615        {
616            ToolPreparationOutcome::Prepared(prepared) => {
617                let trace_hooks: HashMap<String, crate::ToolChildExecutionTraceHook> =
618                    child_execution_trace_hook
619                        .map(|hook| std::iter::once((call_id.clone(), hook)).collect())
620                        .unwrap_or_default();
621                let coordinated = coordinate_tool_invocation(
622                    &dispatch,
623                    prepared,
624                    Some(Box::new(grant.clone())),
625                    grant.manifest.retry_policy,
626                    ToolAttemptEffectIdentity::Scalar {
627                        parent: parent_invocation.clone(),
628                    },
629                    self.cancellation_token.clone(),
630                    || {
631                        crate::RuntimeEffectLocalExecutor::tool_batch(
632                            self.clone(),
633                            trace_hooks.clone(),
634                        )
635                    },
636                )
637                .await;
638                self.restore_tool_trigger_outcomes(coordinated.triggers);
639                coordinated.launch
640            }
641            ToolPreparationOutcome::Completed(outcome) => ToolCallLaunch::Done(*outcome),
642        };
643        let mut outcome = match launch {
644            ToolCallLaunch::Done(outcome) => outcome,
645            ToolCallLaunch::Pending(pending) => {
646                self.await_pending_tool_dispatch_outcome(
647                    &call_id,
648                    parent_invocation.clone(),
649                    pending,
650                    self.cancellation_token.clone(),
651                )
652                .await
653            }
654        };
655        outcome.record.call_id = Some(call_id.clone());
656
657        self.complete_tool_call(index, call_id, replay, outcome, tool_correlation_id)
658            .await
659    }
660
661    pub(crate) async fn prepare_tool_call(
662        &self,
663        pending: crate::sansio::PendingToolCall,
664    ) -> ToolPreparationOutcome {
665        let call_id = Some(pending.call_id.clone());
666        prepare_tool_call_with_context(self.dispatch.as_ref(), pending, call_id).await
667    }
668
669    pub(crate) async fn execute_prepared_tool_attempt_effect(
670        &self,
671        prepared: crate::PreparedToolCall,
672        execution_grant: Option<Box<crate::ToolExecutionGrant>>,
673        attempt: u32,
674        max_attempts: u32,
675        attempt_invocation: crate::RuntimeInvocation,
676        child_execution_trace_hook: Option<crate::ToolChildExecutionTraceHook>,
677    ) -> Result<crate::ToolAttemptEffectOutcome, crate::RuntimeEffectControllerError> {
678        let mut attempt_dispatch = (*self.dispatch).clone();
679        attempt_dispatch.parent_invocation = Some(attempt_invocation.clone());
680        attempt_dispatch.trigger_outcomes =
681            crate::tool_dispatch::ToolTriggerOutcomeBuffer::default();
682        let attempt_dispatch = std::sync::Arc::new(attempt_dispatch);
683        let mut attempt_context = self.clone();
684        attempt_context.dispatch = std::sync::Arc::clone(&attempt_dispatch);
685        attempt_context.parent_invocation = Some(attempt_invocation.clone());
686
687        let mut tool_context =
688            crate::ToolContext::from_dispatch(std::sync::Arc::clone(&attempt_dispatch))
689                .runtime_execution_context(attempt_context.clone())
690                .prepared_call(&prepared)
691                .cancellation_token(self.cancellation_token.clone())
692                .runtime_process_id(self.runtime_process_id.clone())
693                .parent_invocation(Some(attempt_invocation))
694                .child_execution_trace_hook(child_execution_trace_hook);
695        if let Some(process_events) = self.process_event_context.as_ref() {
696            tool_context = tool_context.process_events(
697                process_events.process_id.clone(),
698                std::sync::Arc::clone(&process_events.registry),
699                process_events.awaiter.clone(),
700                process_events.store.clone(),
701                process_events.session_store_factory.clone(),
702                process_events.queued_work_driver.clone(),
703            );
704        }
705        let tool_context = tool_context.build();
706        crate::tool_dispatch::execute_prepared_tool_attempt_effect(
707            attempt_dispatch.as_ref(),
708            prepared,
709            execution_grant,
710            attempt,
711            max_attempts,
712            tool_context,
713        )
714        .await
715    }
716
717    fn restore_tool_trigger_outcomes(
718        &self,
719        outcomes: Vec<crate::tool_dispatch::ToolTriggerEffectOutcome>,
720    ) {
721        for outcome in outcomes {
722            let _ = self.dispatch.trigger_outcomes.enqueue(outcome);
723        }
724    }
725
726    pub(super) async fn await_process_with_cancellation(
727        &self,
728        process_id: &str,
729        parent_invocation: Option<crate::RuntimeInvocation>,
730        cancellation: Option<tokio_util::sync::CancellationToken>,
731    ) -> Result<crate::ProcessAwaitOutput, crate::PluginError> {
732        let _phase = self.named_phase("process.await_handle");
733        if let Some(cancellation) = cancellation {
734            tokio::select! {
735                result = self.dispatch.processes.await_process(
736                    process_id,
737                    self.process_scope(parent_invocation.clone()),
738                ) => result,
739                _ = cancellation.cancelled() => {
740                    let _ = self.dispatch.processes.cancel(
741                        &self.dispatch.session_id,
742                        process_id,
743                        self.process_scope(parent_invocation.clone()),
744                    ).await;
745                    self.dispatch.processes.await_process(
746                        process_id,
747                        self.process_scope(parent_invocation),
748                    ).await
749                }
750            }
751        } else {
752            self.dispatch
753                .processes
754                .await_process(process_id, self.process_scope(parent_invocation))
755                .await
756        }
757    }
758
759    pub(crate) async fn complete_tool_call(
760        &self,
761        _index: usize,
762        call_id: String,
763        replay: Option<crate::llm::types::ProviderReplayMeta>,
764        outcome: ToolDispatchOutcome,
765        tool_correlation_id: TurnActivityId,
766    ) -> CompletedProtocolToolCall {
767        let output = outcome.record.output.clone();
768        let projection_output = output.clone();
769        let projection_tool_name = outcome.record.tool.clone();
770        let projection_args = outcome.record.args.clone();
771        let projection_duration_ms = outcome.record.duration_ms;
772        let projection_call_id = call_id.clone();
773        tokio::task::yield_now().await;
774        let plugins = std::sync::Arc::clone(&self.dispatch.plugins);
775        let projection_context = crate::plugin::ToolResultProjectionContext {
776            session_id: self.dispatch.session_id.clone(),
777            tool_name: projection_tool_name,
778            args: projection_args,
779            output: projection_output,
780            duration_ms: projection_duration_ms,
781            call_id: projection_call_id,
782        };
783        let model_return = match plugins.project_tool_result(projection_context).await {
784            Ok(projected) => projected,
785            Err(err) => ModelToolReturn::text(
786                call_id.clone(),
787                outcome.record.tool.clone(),
788                err.to_string(),
789            ),
790        };
791
792        let record = ToolCallRecord {
793            call_id: Some(call_id.clone()),
794            tool: outcome.record.tool.clone(),
795            args: outcome.record.args.clone(),
796            output: output.clone(),
797            duration_ms: outcome.record.duration_ms,
798        };
799        self.emit_tool_call_completed_trace(&record);
800        self.emit_turn_activity(
801            tool_correlation_id,
802            TurnEvent::ToolCallCompleted {
803                call_id: Some(call_id.clone()),
804                name: outcome.record.tool.clone(),
805                args: outcome.record.args.clone(),
806                output: output.clone(),
807                duration_ms: outcome.record.duration_ms,
808                graph_key: self.code_block_graph_key(),
809                parent_call_id: self.batch_parent_call_id(),
810            },
811        )
812        .await;
813        CompletedProtocolToolCall {
814            completed: crate::sansio::CompletedToolCall {
815                call_id,
816                tool_name: outcome.record.tool,
817                args: outcome.record.args,
818                output,
819                model_return,
820                duration_ms: outcome.record.duration_ms,
821                replay,
822            },
823            record,
824        }
825    }
826
827    pub(crate) async fn pending_completion_dispatch_outcome(
828        &self,
829        tool_name: String,
830        args: serde_json::Value,
831        resolution: crate::Resolution,
832        duration_ms: u64,
833    ) -> ToolDispatchOutcome {
834        let output = crate::tool_result::tool_output_from_completion_resolution(resolution);
835        let result = finalize_tool_result_with_execution_context(
836            self.dispatch.as_ref(),
837            &tool_name,
838            &args,
839            ToolResult::from_output(output),
840            duration_ms,
841        )
842        .await;
843        let output = result.into_done_output().unwrap_or_else(|_| {
844            ToolCallOutput::failure(ToolFailure::runtime(
845                ToolFailureClass::Internal,
846                "pending_tool_not_finalized",
847                "pending tool result reached a completed-output projection path",
848            ))
849        });
850        ToolDispatchOutcome {
851            record: ToolCallRecord {
852                call_id: None,
853                tool: tool_name,
854                args,
855                output,
856                duration_ms,
857            },
858        }
859    }
860
861    async fn await_pending_tool_dispatch_outcome(
862        &self,
863        call_id: &str,
864        parent_invocation: Option<crate::RuntimeInvocation>,
865        pending: crate::tool_dispatch::PendingToolDispatchOutcome,
866        cancellation: Option<tokio_util::sync::CancellationToken>,
867    ) -> ToolDispatchOutcome {
868        self.await_pending_tool_dispatch_outcome_with_suffix(
869            call_id,
870            parent_invocation,
871            format!("{call_id}:await"),
872            pending,
873            cancellation,
874        )
875        .await
876    }
877
878    async fn await_pending_tool_dispatch_outcome_with_suffix(
879        &self,
880        call_id: &str,
881        parent_invocation: Option<crate::RuntimeInvocation>,
882        replay_suffix: String,
883        pending: crate::tool_dispatch::PendingToolDispatchOutcome,
884        cancellation: Option<tokio_util::sync::CancellationToken>,
885    ) -> ToolDispatchOutcome {
886        let fallback;
887        let parent = if let Some(parent) = parent_invocation.as_ref() {
888            parent
889        } else {
890            fallback = crate::RuntimeInvocation::effect(
891                crate::RuntimeScope::new(&self.dispatch.session_id),
892                format!("tool:{call_id}:await"),
893                crate::RuntimeEffectKind::AwaitEvent,
894                format!("tool:{call_id}:await"),
895            );
896            &fallback
897        };
898        let parent_effect_id = parent.effect_id().unwrap_or("tool");
899        let invocation = crate::runtime::causal::child_effect_invocation(
900            parent,
901            format!("{parent_effect_id}:{replay_suffix}"),
902            crate::RuntimeEffectKind::AwaitEvent,
903            replay_suffix,
904        );
905        let cancellation = cancellation.unwrap_or_default();
906        let deadline = pending
907            .pending
908            .deadline
909            .map(|duration| self.dispatch.clock.now() + duration);
910        let outcome = self
911            .dispatch
912            .effect_controller
913            .controller()
914            .execute_effect(
915                crate::RuntimeEffectEnvelope::new(
916                    invocation,
917                    crate::RuntimeEffectCommand::AwaitEvent { key: pending.key },
918                ),
919                crate::RuntimeEffectLocalExecutor::await_event_with_clock(
920                    cancellation,
921                    deadline,
922                    std::sync::Arc::clone(&self.dispatch.clock),
923                ),
924            )
925            .await;
926        let resolution = match outcome.and_then(crate::RuntimeEffectOutcome::into_await_event) {
927            Ok(resolution) => resolution,
928            Err(err) => {
929                return ToolDispatchOutcome {
930                    record: ToolCallRecord {
931                        call_id: None,
932                        tool: pending.tool_name,
933                        args: pending.args,
934                        output: ToolCallOutput::failure(ToolFailure::runtime(
935                            ToolFailureClass::Internal,
936                            "pending_tool_completion_failed",
937                            err.to_string(),
938                        )),
939                        duration_ms: pending.duration_ms,
940                    },
941                };
942            }
943        };
944        self.pending_completion_dispatch_outcome(
945            pending.tool_name,
946            pending.args,
947            resolution,
948            pending.duration_ms,
949        )
950        .await
951    }
952
953    pub async fn call_tool_by_id(
954        &self,
955        call_id: String,
956        tool_id: crate::ToolId,
957        args: serde_json::Value,
958        index: usize,
959    ) -> ToolInvocationReply {
960        let executed = self
961            .execute_tool_call_by_id(call_id, tool_id, args, index, None, None, None)
962            .await;
963        let reply = ToolInvocationReply::from_output(executed.completed.output);
964        reply.with_record(executed.record)
965    }
966
967    pub async fn call_tool_by_id_with_child_execution_trace_hook(
968        &self,
969        call_id: String,
970        tool_id: crate::ToolId,
971        args: serde_json::Value,
972        index: usize,
973        trace_hook: crate::ToolChildExecutionTraceHook,
974    ) -> ToolInvocationReply {
975        let executed = self
976            .execute_tool_call_by_id(call_id, tool_id, args, index, None, None, Some(trace_hook))
977            .await;
978        let reply = ToolInvocationReply::from_output(executed.completed.output);
979        reply.with_record(executed.record)
980    }
981
982    pub async fn call_tool_with_execution_grant(
983        &self,
984        call_id: String,
985        grant: crate::ToolExecutionGrant,
986        args: serde_json::Value,
987        index: usize,
988    ) -> ToolInvocationReply {
989        let executed = self
990            .execute_tool_call_by_grant(call_id, grant, args, index, None, None, None)
991            .await;
992        let reply = ToolInvocationReply::from_output(executed.completed.output);
993        reply.with_record(executed.record)
994    }
995
996    pub async fn call_tool_with_execution_grant_and_child_execution_trace_hook(
997        &self,
998        call_id: String,
999        grant: crate::ToolExecutionGrant,
1000        args: serde_json::Value,
1001        index: usize,
1002        trace_hook: crate::ToolChildExecutionTraceHook,
1003    ) -> ToolInvocationReply {
1004        let executed = self
1005            .execute_tool_call_by_grant(call_id, grant, args, index, None, None, Some(trace_hook))
1006            .await;
1007        let reply = ToolInvocationReply::from_output(executed.completed.output);
1008        reply.with_record(executed.record)
1009    }
1010
1011    pub async fn call_tool_batch(&self, calls: Vec<ToolInvocation>) -> Vec<ToolInvocationReply> {
1012        if calls.is_empty() {
1013            return Vec::new();
1014        }
1015
1016        let batch_id = deterministic_tool_invocation_batch_id(&calls);
1017        let mut replies = vec![None; calls.len()];
1018        let mut prepared_entries = Vec::new();
1019
1020        for (index, call) in calls.into_iter().enumerate() {
1021            let preparation = if let Some(grant) = call.execution_grant.as_deref().cloned() {
1022                let pending = crate::sansio::PendingToolCall {
1023                    call_id: call.id.clone(),
1024                    tool_name: grant.manifest.name.clone(),
1025                    args: call.args,
1026                    replay: None,
1027                };
1028                (
1029                    Some(grant.clone()),
1030                    prepare_granted_tool_call_with_context(
1031                        self.dispatch.as_ref(),
1032                        &grant,
1033                        pending,
1034                        Some(call.id.clone()),
1035                    )
1036                    .await,
1037                )
1038            } else {
1039                let Some(manifest) = crate::tool_dispatch::resolve_callable_manifest_by_id(
1040                    self.dispatch.as_ref(),
1041                    &call.tool_id,
1042                ) else {
1043                    let outcome = ToolDispatchOutcome {
1044                        record: ToolCallRecord {
1045                            call_id: Some(call.id.clone()),
1046                            tool: call.tool_id.to_string(),
1047                            args: call.args,
1048                            output: ToolCallOutput::failure(ToolFailure::runtime(
1049                                ToolFailureClass::Unavailable,
1050                                "tool_unavailable",
1051                                format!(
1052                                    "Tool id `{}` is unavailable in this session",
1053                                    call.tool_id
1054                                ),
1055                            )),
1056                            duration_ms: 0,
1057                        },
1058                    };
1059                    let completed = self
1060                        .complete_tool_call(
1061                            index,
1062                            call.id,
1063                            None,
1064                            outcome,
1065                            TurnActivityId::new(format!("tool:{}", batch_id)),
1066                        )
1067                        .await;
1068                    replies[index] = Some(
1069                        ToolInvocationReply::from_output(completed.completed.output)
1070                            .with_record(completed.record),
1071                    );
1072                    continue;
1073                };
1074
1075                let pending = crate::sansio::PendingToolCall {
1076                    call_id: call.id.clone(),
1077                    tool_name: manifest.name,
1078                    args: call.args,
1079                    replay: None,
1080                };
1081                (None, self.prepare_tool_call(pending).await)
1082            };
1083            let (execution_grant, preparation) = preparation;
1084            match preparation {
1085                ToolPreparationOutcome::Prepared(prepared) => {
1086                    prepared_entries.push((
1087                        index,
1088                        prepared,
1089                        execution_grant,
1090                        call.child_execution_trace_hook,
1091                    ));
1092                }
1093                ToolPreparationOutcome::Completed(outcome) => {
1094                    let completed = self
1095                        .complete_tool_call(
1096                            index,
1097                            call.id,
1098                            None,
1099                            *outcome,
1100                            TurnActivityId::new(format!("tool:{}", batch_id)),
1101                        )
1102                        .await;
1103                    replies[index] = Some(
1104                        ToolInvocationReply::from_output(completed.completed.output)
1105                            .with_record(completed.record),
1106                    );
1107                }
1108            }
1109        }
1110
1111        if !prepared_entries.is_empty() {
1112            let invocation = self.tool_batch_invocation(&batch_id);
1113            let batch = crate::PreparedToolBatch::new_with_grants(
1114                batch_id.clone(),
1115                prepared_entries
1116                    .iter()
1117                    .map(|(_, prepared, grant, _)| (prepared.clone(), grant.clone()))
1118                    .collect(),
1119            );
1120            let child_trace_hooks = prepared_entries
1121                .iter()
1122                .filter_map(|(_, prepared, _, hook)| {
1123                    hook.clone().map(|hook| (prepared.call_id.clone(), hook))
1124                })
1125                .collect();
1126            let envelope = crate::RuntimeEffectEnvelope::new(
1127                invocation.clone(),
1128                crate::RuntimeEffectCommand::ToolBatch { batch },
1129            );
1130            let local_executor =
1131                crate::RuntimeEffectLocalExecutor::tool_batch(self.clone(), child_trace_hooks);
1132            let raw_outcome = self
1133                .dispatch
1134                .effect_controller
1135                .controller()
1136                .execute_effect(envelope, local_executor)
1137                .await;
1138            let outcome =
1139                match raw_outcome.and_then(crate::RuntimeEffectOutcome::into_tool_batch_effect) {
1140                    Ok(outcome) => outcome,
1141                    Err(err) => {
1142                        for (index, prepared, _, _) in prepared_entries {
1143                            replies[index] = Some(ToolInvocationReply::error(serde_json::json!(
1144                                format!("tool batch failed: {err}")
1145                            )));
1146                            let _ = prepared;
1147                        }
1148                        return replies
1149                            .into_iter()
1150                            .map(|reply| reply.expect("every batch reply slot should be filled"))
1151                            .collect();
1152                    }
1153                };
1154            if outcome.launches.len() != prepared_entries.len() {
1155                let message = format!(
1156                    "tool batch returned {} launches for {} prepared calls",
1157                    outcome.launches.len(),
1158                    prepared_entries.len()
1159                );
1160                for (index, _, _, _) in prepared_entries {
1161                    replies[index] = Some(ToolInvocationReply::error(serde_json::json!(message)));
1162                }
1163            } else {
1164                for ((index, prepared, _, _), launch) in
1165                    prepared_entries.into_iter().zip(outcome.launches)
1166                {
1167                    let call_id = prepared.call_id.clone();
1168                    let reply = match launch {
1169                        crate::runtime::ToolCallLaunch::Done { result } => {
1170                            let record = ToolCallRecord {
1171                                call_id: Some(result.call_id.clone()),
1172                                tool: result.tool_name.clone(),
1173                                args: result.args.clone(),
1174                                output: result.output.clone(),
1175                                duration_ms: result.duration_ms,
1176                            };
1177                            ToolInvocationReply::from_output(result.output).with_record(record)
1178                        }
1179                        crate::runtime::ToolCallLaunch::Pending {
1180                            key,
1181                            pending,
1182                            duration_ms,
1183                        } => {
1184                            let dispatch_outcome = self
1185                                .await_pending_tool_dispatch_outcome(
1186                                    &call_id,
1187                                    Some(invocation.clone()),
1188                                    crate::tool_dispatch::PendingToolDispatchOutcome {
1189                                        tool_name: prepared.tool_name.clone(),
1190                                        args: prepared.args.clone(),
1191                                        key,
1192                                        pending,
1193                                        duration_ms,
1194                                    },
1195                                    self.cancellation_token.clone(),
1196                                )
1197                                .await;
1198                            let completed = self
1199                                .complete_tool_call(
1200                                    index,
1201                                    call_id.clone(),
1202                                    prepared.replay.clone(),
1203                                    dispatch_outcome,
1204                                    TurnActivityId::new(format!("tool:{call_id}")),
1205                                )
1206                                .await;
1207                            ToolInvocationReply::from_output(completed.completed.output)
1208                                .with_record(completed.record)
1209                        }
1210                    };
1211                    replies[index] = Some(reply);
1212                }
1213            }
1214        }
1215
1216        replies
1217            .into_iter()
1218            .map(|reply| reply.expect("every batch reply slot should be filled"))
1219            .collect()
1220    }
1221
1222    pub async fn start_tool_call(
1223        &self,
1224        call_id: String,
1225        name: String,
1226        args: serde_json::Value,
1227    ) -> ToolInvocationReply {
1228        self.start_tool_process(call_id, name, args).await
1229    }
1230
1231    pub async fn await_tool_handle(
1232        &self,
1233        call_id: String,
1234        handle: serde_json::Value,
1235    ) -> ToolInvocationReply {
1236        self.await_process_handle(call_id, handle).await
1237    }
1238
1239    pub async fn cancel_tool_handle(
1240        &self,
1241        call_id: String,
1242        handle: serde_json::Value,
1243    ) -> ToolInvocationReply {
1244        self.cancel_process_handle(call_id, handle).await
1245    }
1246
1247    pub async fn signal_tool_handle(
1248        &self,
1249        call_id: String,
1250        handle: serde_json::Value,
1251        signal_name: String,
1252        payload: serde_json::Value,
1253    ) -> ToolInvocationReply {
1254        self.signal_process_handle(call_id, handle, signal_name, payload)
1255            .await
1256    }
1257}