lash_core/
tool_provider.rs

1use std::sync::{Arc, Mutex};
2
3use lash_sansio::llm::types::ProviderReplayMeta;
4use serde::{Deserialize, Serialize};
5
6use crate::plugin::{
7    PluginError, SessionGraphService, SessionLifecycleService, SessionSnapshot, SessionStateService,
8};
9use crate::{ToolContract, ToolDefinition, ToolId, ToolManifest, ToolResult};
10
11mod attachments;
12mod direct_completion;
13mod dispatch;
14mod process;
15pub(crate) mod process_events;
16mod session;
17mod triggers;
18
19pub use attachments::ToolAttachmentClient;
20pub use direct_completion::ToolDirectCompletionClient;
21pub use dispatch::ToolDispatchClient;
22pub use process::ToolSessionProcessAdmin;
23pub use process_events::ToolProcessEventClient;
24pub use session::{ToolSessionAdmin, ToolSessionModel};
25pub use triggers::ToolTriggerClient;
26
27/// A message sent from the sandbox to the host during execution.
28#[derive(Clone, Debug)]
29pub struct SandboxMessage {
30    pub text: String,
31    /// "tool_output" or another host-rendered progress event kind.
32    pub kind: String,
33}
34
35/// Sender for streaming progress messages from tools (e.g. live bash output).
36pub type ProgressSender = tokio::sync::mpsc::UnboundedSender<SandboxMessage>;
37
38#[derive(Clone, Default)]
39pub(crate) struct ToolCompletionState {
40    key: Arc<Mutex<Option<crate::AwaitEventKey>>>,
41}
42
43impl ToolCompletionState {
44    fn store(
45        &self,
46        key: crate::AwaitEventKey,
47    ) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
48        let mut guard = self.key.lock().map_err(|_| {
49            crate::RuntimeError::new(
50                "tool_completion_state_poisoned",
51                "tool completion key state lock poisoned",
52            )
53        })?;
54        if let Some(existing) = guard.as_ref() {
55            return Ok(existing.clone());
56        }
57        *guard = Some(key.clone());
58        Ok(key)
59    }
60
61    pub(crate) fn take(&self) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
62        self.key.lock().map(|mut guard| guard.take()).map_err(|_| {
63            crate::RuntimeError::new(
64                "tool_completion_state_poisoned",
65                "tool completion key state lock poisoned",
66            )
67        })
68    }
69}
70
71/// Per-call environment for [`ToolProvider::execute`]. Fields are sealed so
72/// the runtime can add capabilities without breaking tool authors.
73#[derive(Clone)]
74pub struct ToolContext<'run> {
75    pub(crate) session_id: String,
76    pub(crate) agent_frame_id: crate::AgentFrameId,
77    pub(crate) sessions: Arc<dyn SessionStateService>,
78    pub(crate) session_lifecycle: Arc<dyn SessionLifecycleService>,
79    pub(crate) processes: Arc<dyn crate::ProcessService>,
80    pub(crate) process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
81    pub(crate) effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
82    pub(crate) runtime_dispatch: Option<Arc<crate::tool_dispatch::ToolDispatchContext<'run>>>,
83    pub(crate) runtime_execution_context: Option<crate::RuntimeExecutionContext<'run>>,
84    pub(crate) cancellation_token: Option<tokio_util::sync::CancellationToken>,
85    pub(crate) async_process_id: Option<String>,
86    pub(crate) runtime_process_id: Option<String>,
87    pub(crate) process_events: Option<ToolProcessEventContext>,
88    pub(crate) attachment_store: Arc<crate::SessionAttachmentStore>,
89    pub(crate) direct_completions: crate::DirectCompletionClient<'run>,
90    pub(crate) prepared_payload: serde_json::Value,
91    pub(crate) tool_execution_binding: serde_json::Value,
92    /// The id of the in-flight tool call that is invoking this tool.
93    pub(crate) tool_call_id: Option<String>,
94    pub(crate) attempt_number: u32,
95    pub(crate) max_attempts: u32,
96    pub(crate) replay_key: Option<String>,
97    pub(crate) completion: ToolCompletionState,
98    pub(crate) parent_invocation: Option<crate::RuntimeInvocation>,
99    pub(crate) execution_env_spec: crate::ProcessExecutionEnvSpec,
100    pub(crate) child_execution_trace_hook: Option<ToolChildExecutionTraceHook>,
101}
102
103#[derive(Clone)]
104pub struct ToolChildProcessStarted {
105    pub process_id: String,
106    pub child_entry_name: Option<String>,
107}
108
109#[derive(Clone)]
110pub struct ToolChildExecutionTraceHook {
111    on_child_process_started: Arc<dyn Fn(ToolChildProcessStarted) + Send + Sync>,
112}
113
114impl ToolChildExecutionTraceHook {
115    pub fn new(
116        on_child_process_started: impl Fn(ToolChildProcessStarted) + Send + Sync + 'static,
117    ) -> Self {
118        Self {
119            on_child_process_started: Arc::new(on_child_process_started),
120        }
121    }
122
123    pub fn child_process_started(&self, event: ToolChildProcessStarted) {
124        (self.on_child_process_started)(event);
125    }
126}
127
128#[derive(Clone)]
129pub(crate) struct ToolProcessEventContext {
130    process_id: String,
131    registry: Arc<dyn crate::ProcessRegistry>,
132    awaiter: crate::ProcessAwaiter,
133    store: Option<Arc<dyn crate::RuntimePersistence>>,
134    session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
135    session_graph: Arc<dyn SessionGraphService>,
136    queued_work_driver: Option<crate::QueuedWorkDriver>,
137}
138
139pub(crate) struct ToolContextBuilder<'run> {
140    session_id: String,
141    agent_frame_id: crate::AgentFrameId,
142    sessions: Arc<dyn SessionStateService>,
143    session_lifecycle: Arc<dyn SessionLifecycleService>,
144    session_graph: Arc<dyn SessionGraphService>,
145    processes: Arc<dyn crate::ProcessService>,
146    process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
147    effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
148    runtime_dispatch: Option<Arc<crate::tool_dispatch::ToolDispatchContext<'run>>>,
149    runtime_execution_context: Option<crate::RuntimeExecutionContext<'run>>,
150    cancellation_token: Option<tokio_util::sync::CancellationToken>,
151    async_process_id: Option<String>,
152    runtime_process_id: Option<String>,
153    process_events: Option<ToolProcessEventContext>,
154    attachment_store: Arc<crate::SessionAttachmentStore>,
155    direct_completions: crate::DirectCompletionClient<'run>,
156    prepared_payload: serde_json::Value,
157    tool_execution_binding: serde_json::Value,
158    tool_call_id: Option<String>,
159    completion: ToolCompletionState,
160    parent_invocation: Option<crate::RuntimeInvocation>,
161    execution_env_spec: crate::ProcessExecutionEnvSpec,
162    child_execution_trace_hook: Option<ToolChildExecutionTraceHook>,
163}
164
165impl<'run> ToolContextBuilder<'run> {
166    pub(crate) fn from_dispatch(
167        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
168    ) -> Self {
169        Self {
170            session_id: dispatch.session_id.clone(),
171            agent_frame_id: dispatch.agent_frame_id.clone(),
172            sessions: Arc::clone(&dispatch.sessions),
173            session_lifecycle: Arc::clone(&dispatch.session_lifecycle),
174            session_graph: Arc::clone(&dispatch.session_graph),
175            processes: Arc::clone(&dispatch.processes),
176            process_cancel_ability: Arc::clone(&dispatch.process_cancel_ability),
177            effect_controller: dispatch.effect_controller.clone(),
178            runtime_dispatch: Some(Arc::clone(&dispatch)),
179            runtime_execution_context: None,
180            cancellation_token: None,
181            async_process_id: None,
182            runtime_process_id: None,
183            process_events: None,
184            attachment_store: Arc::clone(&dispatch.attachment_store),
185            direct_completions: dispatch.direct_completions.clone(),
186            prepared_payload: serde_json::Value::Null,
187            tool_execution_binding: serde_json::Value::Null,
188            tool_call_id: None,
189            completion: ToolCompletionState::default(),
190            parent_invocation: dispatch.parent_invocation.clone(),
191            execution_env_spec: dispatch.execution_env_spec.clone(),
192            child_execution_trace_hook: None,
193        }
194    }
195
196    #[cfg(any(test, feature = "testing"))]
197    pub(crate) fn tool_call_id(mut self, tool_call_id: impl Into<Option<String>>) -> Self {
198        self.tool_call_id = tool_call_id.into();
199        self
200    }
201
202    pub(crate) fn prepared_call(mut self, call: &PreparedToolCall) -> Self {
203        self.tool_call_id = Some(call.call_id.clone());
204        self.prepared_payload = call.prepared_payload.clone();
205        self
206    }
207
208    pub(crate) fn tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
209        self.tool_execution_binding = binding;
210        self
211    }
212
213    pub(crate) fn cancellation_token(
214        mut self,
215        cancellation_token: Option<tokio_util::sync::CancellationToken>,
216    ) -> Self {
217        self.cancellation_token = cancellation_token;
218        self
219    }
220
221    pub(crate) fn runtime_execution_context(
222        mut self,
223        context: crate::RuntimeExecutionContext<'run>,
224    ) -> Self {
225        self.runtime_execution_context = Some(context);
226        self
227    }
228
229    pub(crate) fn runtime_process_id(mut self, process_id: Option<String>) -> Self {
230        self.runtime_process_id = process_id;
231        self
232    }
233
234    pub(crate) fn async_process(
235        mut self,
236        process_id: impl Into<String>,
237        cancellation_token: tokio_util::sync::CancellationToken,
238    ) -> Self {
239        self.async_process_id = Some(process_id.into());
240        self.cancellation_token = Some(cancellation_token);
241        self
242    }
243
244    pub(crate) fn process_events(
245        mut self,
246        process_id: impl Into<String>,
247        registry: Arc<dyn crate::ProcessRegistry>,
248        awaiter: crate::ProcessAwaiter,
249        store: Option<Arc<dyn crate::RuntimePersistence>>,
250        session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
251        queued_work_driver: Option<crate::QueuedWorkDriver>,
252    ) -> Self {
253        self.process_events = Some(ToolProcessEventContext {
254            process_id: process_id.into(),
255            registry,
256            awaiter,
257            store,
258            session_store_factory,
259            session_graph: Arc::clone(&self.session_graph),
260            queued_work_driver,
261        });
262        self
263    }
264
265    pub(crate) fn parent_invocation(mut self, metadata: Option<crate::RuntimeInvocation>) -> Self {
266        self.parent_invocation = metadata;
267        self
268    }
269
270    pub(crate) fn child_execution_trace_hook(
271        mut self,
272        hook: Option<ToolChildExecutionTraceHook>,
273    ) -> Self {
274        self.child_execution_trace_hook = hook;
275        self
276    }
277
278    pub(crate) fn build(self) -> ToolContext<'run> {
279        ToolContext {
280            session_id: self.session_id,
281            agent_frame_id: self.agent_frame_id,
282            sessions: self.sessions,
283            session_lifecycle: self.session_lifecycle,
284            processes: self.processes,
285            process_cancel_ability: self.process_cancel_ability,
286            effect_controller: self.effect_controller,
287            runtime_dispatch: self.runtime_dispatch,
288            runtime_execution_context: self.runtime_execution_context,
289            cancellation_token: self.cancellation_token,
290            async_process_id: self.async_process_id,
291            runtime_process_id: self.runtime_process_id,
292            process_events: self.process_events,
293            attachment_store: self.attachment_store,
294            direct_completions: self.direct_completions,
295            prepared_payload: self.prepared_payload,
296            tool_execution_binding: self.tool_execution_binding,
297            tool_call_id: self.tool_call_id,
298            attempt_number: 1,
299            max_attempts: 1,
300            replay_key: None,
301            completion: self.completion,
302            parent_invocation: self.parent_invocation,
303            execution_env_spec: self.execution_env_spec,
304            child_execution_trace_hook: self.child_execution_trace_hook,
305        }
306    }
307}
308
309impl<'run> ToolContext<'run> {
310    #[cfg(any(test, feature = "testing"))]
311    #[expect(
312        clippy::too_many_arguments,
313        reason = "testing constructor mirrors the sealed runtime tool context dependencies"
314    )]
315    pub(crate) fn builder(
316        session_id: String,
317        sessions: Arc<dyn SessionStateService>,
318        session_lifecycle: Arc<dyn SessionLifecycleService>,
319        session_graph: Arc<dyn SessionGraphService>,
320        processes: Arc<dyn crate::ProcessService>,
321        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
322        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
323        attachment_store: Arc<crate::SessionAttachmentStore>,
324        direct_completions: crate::DirectCompletionClient<'run>,
325    ) -> ToolContextBuilder<'run> {
326        ToolContextBuilder {
327            session_id,
328            agent_frame_id: String::new(),
329            sessions,
330            session_lifecycle,
331            session_graph,
332            processes,
333            process_cancel_ability,
334            effect_controller,
335            runtime_dispatch: None,
336            runtime_execution_context: None,
337            cancellation_token: None,
338            async_process_id: None,
339            runtime_process_id: None,
340            process_events: None,
341            attachment_store,
342            direct_completions,
343            prepared_payload: serde_json::Value::Null,
344            tool_execution_binding: serde_json::Value::Null,
345            tool_call_id: None,
346            completion: ToolCompletionState::default(),
347            parent_invocation: None,
348            execution_env_spec: crate::ProcessExecutionEnvSpec::new(
349                crate::PluginOptions::default(),
350                crate::SessionPolicy::default(),
351            ),
352            child_execution_trace_hook: None,
353        }
354    }
355
356    pub(crate) fn from_dispatch(
357        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
358    ) -> ToolContextBuilder<'run> {
359        ToolContextBuilder::from_dispatch(dispatch)
360    }
361
362    pub fn session_id(&self) -> &str {
363        &self.session_id
364    }
365
366    pub fn agent_frame_id(&self) -> &str {
367        &self.agent_frame_id
368    }
369
370    pub fn sessions(&self) -> ToolSessionAdmin<'run> {
371        ToolSessionAdmin {
372            session_id: self.session_id.clone(),
373            sessions: Arc::clone(&self.sessions),
374            session_lifecycle: Arc::clone(&self.session_lifecycle),
375            effect_controller: self.effect_controller.clone(),
376        }
377    }
378
379    pub fn dispatch(&self) -> ToolDispatchClient<'run> {
380        ToolDispatchClient {
381            context: self.clone(),
382        }
383    }
384
385    pub fn triggers(&self) -> ToolTriggerClient<'run> {
386        ToolTriggerClient {
387            context: self.clone(),
388        }
389    }
390
391    pub fn processes(&self) -> ToolSessionProcessAdmin<'run> {
392        ToolSessionProcessAdmin {
393            session_id: self.session_id.clone(),
394            agent_frame_id: self.agent_frame_id.clone(),
395            processes: Arc::clone(&self.processes),
396            process_cancel_ability: Arc::clone(&self.process_cancel_ability),
397            effect_controller: self.effect_controller.clone(),
398            parent_invocation: self.parent_invocation.clone(),
399            tool_call_id: self.tool_call_id.clone(),
400            execution_env_spec: self.execution_env_spec.clone(),
401        }
402    }
403
404    pub fn emit_child_process_started(
405        &self,
406        process_id: impl Into<String>,
407        child_entry_name: Option<String>,
408    ) {
409        let Some(hook) = &self.child_execution_trace_hook else {
410            return;
411        };
412        hook.child_process_started(ToolChildProcessStarted {
413            process_id: process_id.into(),
414            child_entry_name,
415        });
416    }
417
418    pub fn direct_completions(&self) -> ToolDirectCompletionClient<'run> {
419        ToolDirectCompletionClient {
420            session_id: self.session_id.clone(),
421            tool_call_id: self.tool_call_id.clone(),
422            direct_completions: self.direct_completions.clone(),
423            parent_invocation: self.parent_invocation.clone(),
424        }
425    }
426
427    pub fn attachments(&self) -> ToolAttachmentClient {
428        ToolAttachmentClient {
429            store: Arc::clone(&self.attachment_store),
430        }
431    }
432
433    pub fn process_events(&self) -> ToolProcessEventClient {
434        ToolProcessEventClient {
435            context: self.process_events.clone(),
436        }
437    }
438
439    pub fn cancellation_token(&self) -> Option<&tokio_util::sync::CancellationToken> {
440        self.cancellation_token.as_ref()
441    }
442
443    #[doc(hidden)]
444    pub fn named_phase(&self, phase: &'static str) -> crate::runtime::RuntimeNamedPhase {
445        match self.runtime_execution_context.as_ref() {
446            Some(context) => context.named_phase(phase),
447            None => crate::runtime::RuntimeNamedPhase::begin(None, phase),
448        }
449    }
450
451    pub fn async_process_id(&self) -> Option<&str> {
452        self.async_process_id.as_deref()
453    }
454
455    pub fn runtime_process_id(&self) -> Option<&str> {
456        self.async_process_id
457            .as_deref()
458            .or(self.runtime_process_id.as_deref())
459            .or_else(|| {
460                self.process_events
461                    .as_ref()
462                    .map(|context| context.process_id.as_str())
463            })
464    }
465
466    pub fn tool_call_id(&self) -> Option<&str> {
467        self.tool_call_id.as_deref()
468    }
469
470    pub fn prepared_payload(&self) -> &serde_json::Value {
471        &self.prepared_payload
472    }
473
474    pub fn tool_execution_binding(&self) -> &serde_json::Value {
475        &self.tool_execution_binding
476    }
477
478    pub fn decode_prepared_payload<T>(&self) -> Result<T, serde_json::Error>
479    where
480        T: serde::de::DeserializeOwned,
481    {
482        serde_json::from_value(self.prepared_payload.clone())
483    }
484
485    pub fn attempt_number(&self) -> u32 {
486        self.attempt_number
487    }
488
489    pub fn max_attempts(&self) -> u32 {
490        self.max_attempts
491    }
492
493    pub fn replay_key(&self) -> Option<&str> {
494        self.replay_key.as_deref()
495    }
496
497    /// Obtain the durable completion key for this call, required before returning
498    /// [`ToolResult::Pending`](crate::ToolResult::Pending).
499    ///
500    /// A tool that defers its outcome (waiting on a webhook, human approval, or another
501    /// service) calls this, hands the returned [`AwaitEventKey`](crate::AwaitEventKey)
502    /// to whatever will complete the work out-of-band, and then returns
503    /// `ToolResult::Pending(..)`. The key names the durable wait the runtime parks the
504    /// call on; the external resolver delivers the result against it later.
505    ///
506    /// The key is stored on the context and consumed by the dispatcher when the tool
507    /// returns `Pending`. Returning `Pending` without first calling this fails the call
508    /// with `pending_tool_missing_completion_key`. Calls made outside a prepared tool
509    /// invocation (no tool call id) fail with `tool_completion_key_missing_call_id`.
510    pub async fn completion_key(&self) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
511        let tool_call_id = self.tool_call_id.clone().ok_or_else(|| {
512            crate::RuntimeError::new(
513                "tool_completion_key_missing_call_id",
514                "completion keys require a prepared tool call id",
515            )
516        })?;
517        let scoped = self.effect_controller.scoped();
518        let key = scoped
519            .controller()
520            .await_event_key(
521                scoped.execution_scope(),
522                crate::AwaitEventWaitIdentity::tool_completion(tool_call_id),
523            )
524            .await?;
525        self.completion.store(key)
526    }
527
528    pub(crate) fn take_completion_key(
529        &self,
530    ) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
531        self.completion.take()
532    }
533
534    pub fn with_async_process(
535        mut self,
536        process_id: impl Into<String>,
537        cancellation_token: tokio_util::sync::CancellationToken,
538    ) -> Self {
539        self.async_process_id = Some(process_id.into());
540        self.runtime_process_id = self.async_process_id.clone();
541        self.cancellation_token = Some(cancellation_token);
542        self
543    }
544
545    #[cfg(any(test, feature = "testing"))]
546    #[doc(hidden)]
547    pub fn with_process_events_for_testing(
548        mut self,
549        process_id: impl Into<String>,
550        registry: Arc<dyn crate::ProcessRegistry>,
551    ) -> Self {
552        let awaiter = crate::ProcessAwaiter::polling(Arc::clone(&registry));
553        self.process_events = Some(ToolProcessEventContext {
554            process_id: process_id.into(),
555            registry,
556            awaiter,
557            store: None,
558            session_store_factory: None,
559            session_graph: Arc::new(crate::plugin::NoopSessionManager),
560            queued_work_driver: None,
561        });
562        self
563    }
564
565    pub(crate) fn with_retry_context(
566        mut self,
567        tool_name: &str,
568        attempt_number: u32,
569        max_attempts: u32,
570    ) -> Self {
571        self.attempt_number = attempt_number.max(1);
572        self.max_attempts = max_attempts.max(1);
573        self.replay_key = self
574            .tool_call_id
575            .as_ref()
576            .map(|call_id| format!("lash-tool:{}:{call_id}:{tool_name}", self.session_id));
577        self
578    }
579
580    pub(crate) fn with_prepared_payload(mut self, payload: serde_json::Value) -> Self {
581        self.prepared_payload = payload;
582        self
583    }
584
585    pub(crate) fn with_tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
586        self.tool_execution_binding = binding;
587        self
588    }
589
590    /// Constructor reserved for `lash_core::testing` helpers. Do not call directly;
591    /// use [`lash_core::testing::mock_tool_context`] instead.
592    #[cfg(any(test, feature = "testing"))]
593    #[doc(hidden)]
594    #[expect(
595        clippy::too_many_arguments,
596        reason = "test-only constructor mirrors the sealed runtime tool context"
597    )]
598    pub fn __for_testing(
599        session_id: String,
600        sessions: Arc<dyn SessionStateService>,
601        session_lifecycle: Arc<dyn SessionLifecycleService>,
602        session_graph: Arc<dyn SessionGraphService>,
603        processes: Arc<dyn crate::ProcessService>,
604        attachment_store: Arc<crate::SessionAttachmentStore>,
605        direct_completions: crate::DirectCompletionClient<'static>,
606        tool_call_id: Option<String>,
607    ) -> ToolContext<'static> {
608        ToolContext::builder(
609            session_id,
610            sessions,
611            session_lifecycle,
612            session_graph,
613            processes,
614            Arc::new(crate::DefaultProcessCancelAbility),
615            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
616                crate::InlineRuntimeEffectController,
617            )),
618            attachment_store,
619            direct_completions,
620        )
621        .tool_call_id(tool_call_id)
622        .build()
623    }
624
625    /// Constructor reserved for tests that need a custom process-cancel host
626    /// ability. Do not call directly; prefer public testing helpers when they
627    /// cover the case.
628    #[cfg(any(test, feature = "testing"))]
629    #[doc(hidden)]
630    #[expect(
631        clippy::too_many_arguments,
632        reason = "test-only constructor mirrors the sealed runtime context"
633    )]
634    pub fn __for_testing_with_process_cancel_ability(
635        session_id: String,
636        sessions: Arc<dyn SessionStateService>,
637        session_lifecycle: Arc<dyn SessionLifecycleService>,
638        session_graph: Arc<dyn SessionGraphService>,
639        processes: Arc<dyn crate::ProcessService>,
640        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
641        attachment_store: Arc<crate::SessionAttachmentStore>,
642        direct_completions: crate::DirectCompletionClient<'static>,
643        tool_call_id: Option<String>,
644    ) -> ToolContext<'static> {
645        ToolContext::builder(
646            session_id,
647            sessions,
648            session_lifecycle,
649            session_graph,
650            processes,
651            process_cancel_ability,
652            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
653                crate::InlineRuntimeEffectController,
654            )),
655            attachment_store,
656            direct_completions,
657        )
658        .tool_call_id(tool_call_id)
659        .build()
660    }
661}
662
663/// Runtime-prepared executable tool call.
664///
665/// The raw model/provider identity remains visible, but any argument rewrites
666/// and provider-owned context projections are frozen before the call crosses a
667/// runtime effect or process boundary.
668#[derive(Clone, Debug, Serialize, Deserialize)]
669pub struct PreparedToolCall {
670    pub call_id: String,
671    pub tool_id: ToolId,
672    pub tool_name: String,
673    pub args: serde_json::Value,
674    #[serde(default, skip_serializing_if = "Option::is_none")]
675    pub replay: Option<ProviderReplayMeta>,
676    #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
677    pub prepared_payload: serde_json::Value,
678}
679
680impl PreparedToolCall {
681    pub fn identity(tool_id: ToolId, call: crate::sansio::PendingToolCall) -> Self {
682        Self {
683            call_id: call.call_id,
684            tool_id,
685            tool_name: call.tool_name,
686            args: call.args,
687            replay: call.replay,
688            prepared_payload: serde_json::Value::Null,
689        }
690    }
691
692    pub fn from_parts(
693        call_id: impl Into<String>,
694        tool_id: impl Into<ToolId>,
695        tool_name: impl Into<String>,
696        args: serde_json::Value,
697        replay: Option<ProviderReplayMeta>,
698        prepared_payload: serde_json::Value,
699    ) -> Self {
700        Self {
701            call_id: call_id.into(),
702            tool_id: tool_id.into(),
703            tool_name: tool_name.into(),
704            args,
705            replay,
706            prepared_payload,
707        }
708    }
709}
710
711/// One ordered child inside a runtime-prepared tool batch.
712///
713/// The call itself carries the executable provider payload. `replay_suffix`
714/// is the deterministic suffix used for child effects such as retry sleeps or
715/// pending completion awaits when the batch is the durable parent.
716#[derive(Clone, Debug, Serialize, Deserialize)]
717pub struct PreparedToolBatchCall {
718    pub call: PreparedToolCall,
719    pub replay_suffix: String,
720    #[serde(default, skip_serializing_if = "Option::is_none")]
721    pub execution_grant: Option<Box<ToolExecutionGrant>>,
722}
723
724/// Runtime-prepared executable tool batch.
725///
726/// The vector order is source order. Scheduling may run parallel-safe tools
727/// concurrently, but launches and pending completion consumption are projected
728/// back through this order.
729#[derive(Clone, Debug, Serialize, Deserialize)]
730pub struct PreparedToolBatch {
731    pub batch_id: String,
732    pub calls: Vec<PreparedToolBatchCall>,
733}
734
735impl PreparedToolBatch {
736    pub fn new(batch_id: impl Into<String>, calls: Vec<PreparedToolCall>) -> Self {
737        let batch_id = batch_id.into();
738        let calls = calls
739            .into_iter()
740            .enumerate()
741            .map(|(index, call)| PreparedToolBatchCall {
742                replay_suffix: format!("child:{index}:{}", call.call_id),
743                call,
744                execution_grant: None,
745            })
746            .collect();
747        Self { batch_id, calls }
748    }
749
750    pub fn new_with_grants(
751        batch_id: impl Into<String>,
752        calls: Vec<(PreparedToolCall, Option<ToolExecutionGrant>)>,
753    ) -> Self {
754        let batch_id = batch_id.into();
755        let calls = calls
756            .into_iter()
757            .enumerate()
758            .map(|(index, (call, execution_grant))| PreparedToolBatchCall {
759                replay_suffix: format!("child:{index}:{}", call.call_id),
760                call,
761                execution_grant: execution_grant.map(Box::new),
762            })
763            .collect();
764        Self { batch_id, calls }
765    }
766
767    pub fn is_empty(&self) -> bool {
768        self.calls.is_empty()
769    }
770
771    pub fn len(&self) -> usize {
772        self.calls.len()
773    }
774}
775
776/// Explicit authority to execute a tool outside Tool Catalog membership.
777///
778/// Normal tool calls are authorized by catalog membership. A grant is a
779/// separate, caller-provided capability used by deferred resolution flows: it
780/// carries the manifest/contract to validate the call plus an opaque host
781/// execution binding that providers can inspect from the prepare and execute
782/// contexts.
783#[derive(Clone, Debug, Serialize, Deserialize)]
784pub struct ToolExecutionGrant {
785    /// Tool identity and model-facing metadata authorized by the grant.
786    pub manifest: ToolManifest,
787    /// Contract used to validate granted call arguments without consulting the
788    /// current Tool Catalog.
789    pub contract: Box<ToolContract>,
790    /// Explicit registry source route for registry-backed execution. Direct
791    /// non-registry providers may ignore this; [`ToolRegistry`] requires it.
792    pub source_id: Option<String>,
793    /// Opaque host routing payload passed to prepare and execute contexts.
794    pub execution_binding: serde_json::Value,
795}
796
797impl ToolExecutionGrant {
798    pub fn new(manifest: ToolManifest, contract: ToolContract) -> Self {
799        Self {
800            manifest,
801            contract: Box::new(contract),
802            source_id: None,
803            execution_binding: serde_json::Value::Null,
804        }
805    }
806
807    pub fn from_definition(definition: ToolDefinition) -> Self {
808        Self::new(definition.manifest(), definition.contract())
809    }
810
811    pub fn with_source_id(mut self, source_id: impl Into<String>) -> Self {
812        self.source_id = Some(source_id.into());
813        self
814    }
815
816    pub fn with_execution_binding(mut self, execution_binding: serde_json::Value) -> Self {
817        self.execution_binding = execution_binding;
818        self
819    }
820}
821
822#[derive(Clone)]
823pub struct ToolPrepareContext {
824    session_id: String,
825    sessions: Arc<dyn SessionStateService>,
826    turn_context: crate::TurnContext,
827    tool_call_id: Option<String>,
828    tool_execution_binding: serde_json::Value,
829}
830
831impl ToolPrepareContext {
832    pub(crate) fn with_execution_binding(
833        session_id: String,
834        sessions: Arc<dyn SessionStateService>,
835        turn_context: crate::TurnContext,
836        tool_call_id: Option<String>,
837        tool_execution_binding: serde_json::Value,
838    ) -> Self {
839        Self {
840            session_id,
841            sessions,
842            turn_context,
843            tool_call_id,
844            tool_execution_binding,
845        }
846    }
847
848    pub fn session_id(&self) -> &str {
849        &self.session_id
850    }
851
852    pub fn tool_call_id(&self) -> Option<&str> {
853        self.tool_call_id.as_deref()
854    }
855
856    pub fn tool_execution_binding(&self) -> &serde_json::Value {
857        &self.tool_execution_binding
858    }
859
860    pub fn turn_context(&self) -> &crate::TurnContext {
861        &self.turn_context
862    }
863
864    pub fn plugin_input<T>(&self, plugin_id: &'static str) -> Option<&T>
865    where
866        T: 'static,
867    {
868        self.turn_context.plugin_input::<T>(plugin_id)
869    }
870
871    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
872        self.sessions.snapshot_session(&self.session_id).await
873    }
874
875    pub async fn tool_catalog(&self) -> Result<Vec<serde_json::Value>, PluginError> {
876        self.sessions.tool_catalog(&self.session_id).await
877    }
878
879    pub async fn shared_tool_catalog(
880        &self,
881    ) -> Result<std::sync::Arc<Vec<serde_json::Value>>, PluginError> {
882        self.sessions.shared_tool_catalog(&self.session_id).await
883    }
884}
885
886/// Inputs handed to [`ToolProvider::prepare_tool_call`].
887pub struct ToolPrepareCall<'a> {
888    pub tool_id: ToolId,
889    pub pending: crate::sansio::PendingToolCall,
890    pub context: &'a ToolPrepareContext,
891}
892
893/// Per-call inputs handed to [`ToolProvider::execute`].
894///
895/// Fields are `pub` because `ToolCall` is a transient borrow; consumers
896/// typically destructure (`let ToolCall { name, args, .. } = call`). The
897/// stable surface lives on [`ToolContext`] (sealed) and the runtime's
898/// dispatcher, which constructs `ToolCall` values.
899pub struct ToolCall<'a> {
900    pub name: &'a str,
901    pub args: &'a serde_json::Value,
902    pub context: &'a ToolContext<'a>,
903    pub progress: Option<&'a ProgressSender>,
904}
905
906/// Trait for providing tools to the sandbox. Implement this per-project.
907///
908/// Implementations supply cheap [`ToolManifest`]s, lazily resolved
909/// [`ToolContract`]s, and a single
910/// [`execute`](Self::execute) method that handles every call. Tools that
911/// need session state read it from `call.context`; tools that stream
912/// progress send through `call.progress`.
913#[async_trait::async_trait]
914pub trait ToolProvider: Send + Sync + 'static {
915    fn tool_manifests(&self) -> Vec<ToolManifest>;
916    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
917        self.tool_manifests()
918            .into_iter()
919            .find(|manifest| manifest.name == name)
920    }
921    fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
922        self.tool_manifests()
923            .into_iter()
924            .find(|manifest| manifest.id == *id)
925    }
926    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
927    fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
928        let manifest = self.resolve_manifest_by_id(id)?;
929        self.resolve_contract(&manifest.name)
930    }
931    async fn prepare_tool_call(
932        &self,
933        call: ToolPrepareCall<'_>,
934    ) -> Result<PreparedToolCall, ToolResult> {
935        Ok(PreparedToolCall::identity(call.tool_id, call.pending))
936    }
937    async fn prepare_granted_tool_call(
938        &self,
939        grant: &ToolExecutionGrant,
940        call: ToolPrepareCall<'_>,
941    ) -> Result<PreparedToolCall, ToolResult> {
942        let _ = call;
943        Err(ToolResult::err_fmt(format_args!(
944            "Granted execution is unsupported for tool id `{}`",
945            grant.manifest.id
946        )))
947    }
948    async fn execute(&self, call: ToolCall<'_>) -> ToolResult;
949    async fn execute_granted(
950        &self,
951        grant: &ToolExecutionGrant,
952        args: &serde_json::Value,
953        context: &ToolContext<'_>,
954        progress: Option<&ProgressSender>,
955    ) -> ToolResult {
956        let _ = (args, context, progress);
957        ToolResult::err_fmt(format_args!(
958            "Granted execution is unsupported for tool id `{}`",
959            grant.manifest.id
960        ))
961    }
962    async fn execute_by_id(
963        &self,
964        tool_id: &ToolId,
965        args: &serde_json::Value,
966        context: &ToolContext<'_>,
967        progress: Option<&ProgressSender>,
968    ) -> ToolResult {
969        let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
970            return ToolResult::err_fmt(format!("Unknown tool id: {tool_id}"));
971        };
972        self.execute(ToolCall {
973            name: &manifest.name,
974            args,
975            context,
976            progress,
977        })
978        .await
979    }
980}
981
982#[cfg(test)]
983mod tests {
984    use super::*;
985
986    #[test]
987    fn tool_context_builder_carries_call_payload_and_cancellation_state() {
988        let cancellation = tokio_util::sync::CancellationToken::new();
989        let prepared = PreparedToolCall::from_parts(
990            "call-1",
991            "tool:demo_tool",
992            "demo_tool",
993            serde_json::json!({ "input": true }),
994            None,
995            serde_json::json!({ "prepared": true }),
996        );
997
998        let context = ToolContext::builder(
999            "session-1".to_string(),
1000            Arc::new(crate::testing::MockSessionManager::default()),
1001            Arc::new(crate::testing::MockSessionManager::default()),
1002            Arc::new(crate::testing::MockSessionManager::default()),
1003            Arc::new(crate::UnavailableProcessService),
1004            Arc::new(crate::DefaultProcessCancelAbility),
1005            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
1006                crate::InlineRuntimeEffectController,
1007            )),
1008            Arc::new(crate::SessionAttachmentStore::in_memory()),
1009            crate::DirectCompletionClient::unavailable(
1010                "direct completions are unavailable in this test context",
1011            ),
1012        )
1013        .prepared_call(&prepared)
1014        .cancellation_token(Some(cancellation.clone()))
1015        .async_process("process-1", cancellation.clone())
1016        .build();
1017
1018        assert_eq!(context.session_id(), "session-1");
1019        assert_eq!(context.tool_call_id(), Some("call-1"));
1020        assert_eq!(
1021            context.prepared_payload(),
1022            &serde_json::json!({ "prepared": true })
1023        );
1024        assert_eq!(context.async_process_id(), Some("process-1"));
1025        assert!(context.cancellation_token().is_some());
1026    }
1027}
lash_core/tool_provider.rs

lash_core/
tool_provider.rs
