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    #[cfg(test)]
209    pub(crate) fn tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
210        self.tool_execution_binding = binding;
211        self
212    }
213
214    pub(crate) fn cancellation_token(
215        mut self,
216        cancellation_token: Option<tokio_util::sync::CancellationToken>,
217    ) -> Self {
218        self.cancellation_token = cancellation_token;
219        self
220    }
221
222    pub(crate) fn runtime_execution_context(
223        mut self,
224        context: crate::RuntimeExecutionContext<'run>,
225    ) -> Self {
226        self.runtime_execution_context = Some(context);
227        self
228    }
229
230    pub(crate) fn runtime_process_id(mut self, process_id: Option<String>) -> Self {
231        self.runtime_process_id = process_id;
232        self
233    }
234
235    pub(crate) fn async_process(
236        mut self,
237        process_id: impl Into<String>,
238        cancellation_token: tokio_util::sync::CancellationToken,
239    ) -> Self {
240        self.async_process_id = Some(process_id.into());
241        self.cancellation_token = Some(cancellation_token);
242        self
243    }
244
245    pub(crate) fn process_events(
246        mut self,
247        process_id: impl Into<String>,
248        registry: Arc<dyn crate::ProcessRegistry>,
249        awaiter: crate::ProcessAwaiter,
250        store: Option<Arc<dyn crate::RuntimePersistence>>,
251        session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
252        queued_work_driver: Option<crate::QueuedWorkDriver>,
253    ) -> Self {
254        self.process_events = Some(ToolProcessEventContext {
255            process_id: process_id.into(),
256            registry,
257            awaiter,
258            store,
259            session_store_factory,
260            session_graph: Arc::clone(&self.session_graph),
261            queued_work_driver,
262        });
263        self
264    }
265
266    pub(crate) fn parent_invocation(mut self, metadata: Option<crate::RuntimeInvocation>) -> Self {
267        self.parent_invocation = metadata;
268        self
269    }
270
271    pub(crate) fn child_execution_trace_hook(
272        mut self,
273        hook: Option<ToolChildExecutionTraceHook>,
274    ) -> Self {
275        self.child_execution_trace_hook = hook;
276        self
277    }
278
279    pub(crate) fn build(self) -> ToolContext<'run> {
280        ToolContext {
281            session_id: self.session_id,
282            agent_frame_id: self.agent_frame_id,
283            sessions: self.sessions,
284            session_lifecycle: self.session_lifecycle,
285            processes: self.processes,
286            process_cancel_ability: self.process_cancel_ability,
287            effect_controller: self.effect_controller,
288            runtime_dispatch: self.runtime_dispatch,
289            runtime_execution_context: self.runtime_execution_context,
290            cancellation_token: self.cancellation_token,
291            async_process_id: self.async_process_id,
292            runtime_process_id: self.runtime_process_id,
293            process_events: self.process_events,
294            attachment_store: self.attachment_store,
295            direct_completions: self.direct_completions,
296            prepared_payload: self.prepared_payload,
297            tool_execution_binding: self.tool_execution_binding,
298            tool_call_id: self.tool_call_id,
299            attempt_number: 1,
300            max_attempts: 1,
301            replay_key: None,
302            completion: self.completion,
303            parent_invocation: self.parent_invocation,
304            execution_env_spec: self.execution_env_spec,
305            child_execution_trace_hook: self.child_execution_trace_hook,
306        }
307    }
308}
309
310impl<'run> ToolContext<'run> {
311    pub(crate) fn replay_validation_trace(&self) -> Option<crate::RuntimeEffectReplayTrace> {
312        self.runtime_execution_context
313            .as_ref()
314            .and_then(crate::RuntimeExecutionContext::replay_validation_trace)
315    }
316
317    pub(crate) fn to_static(&self) -> Option<ToolContext<'static>> {
318        Some(ToolContext {
319            session_id: self.session_id.clone(),
320            agent_frame_id: self.agent_frame_id.clone(),
321            sessions: Arc::clone(&self.sessions),
322            session_lifecycle: Arc::clone(&self.session_lifecycle),
323            processes: Arc::clone(&self.processes),
324            process_cancel_ability: Arc::clone(&self.process_cancel_ability),
325            effect_controller: self.effect_controller.to_static()?,
326            runtime_dispatch: match self.runtime_dispatch.as_ref() {
327                Some(dispatch) => Some(Arc::new(dispatch.to_static()?)),
328                None => None,
329            },
330            runtime_execution_context: match self.runtime_execution_context.as_ref() {
331                Some(context) => Some(context.to_static()?),
332                None => None,
333            },
334            cancellation_token: self.cancellation_token.clone(),
335            async_process_id: self.async_process_id.clone(),
336            runtime_process_id: self.runtime_process_id.clone(),
337            process_events: self.process_events.clone(),
338            attachment_store: Arc::clone(&self.attachment_store),
339            direct_completions: self.direct_completions.to_static()?,
340            prepared_payload: self.prepared_payload.clone(),
341            tool_execution_binding: self.tool_execution_binding.clone(),
342            tool_call_id: self.tool_call_id.clone(),
343            attempt_number: self.attempt_number,
344            max_attempts: self.max_attempts,
345            replay_key: self.replay_key.clone(),
346            completion: self.completion.clone(),
347            parent_invocation: self.parent_invocation.clone(),
348            execution_env_spec: self.execution_env_spec.clone(),
349            child_execution_trace_hook: self.child_execution_trace_hook.clone(),
350        })
351    }
352
353    #[cfg(any(test, feature = "testing"))]
354    #[expect(
355        clippy::too_many_arguments,
356        reason = "testing constructor mirrors the sealed runtime tool context dependencies"
357    )]
358    pub(crate) fn builder(
359        session_id: String,
360        sessions: Arc<dyn SessionStateService>,
361        session_lifecycle: Arc<dyn SessionLifecycleService>,
362        session_graph: Arc<dyn SessionGraphService>,
363        processes: Arc<dyn crate::ProcessService>,
364        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
365        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
366        attachment_store: Arc<crate::SessionAttachmentStore>,
367        direct_completions: crate::DirectCompletionClient<'run>,
368    ) -> ToolContextBuilder<'run> {
369        ToolContextBuilder {
370            session_id,
371            agent_frame_id: String::new(),
372            sessions,
373            session_lifecycle,
374            session_graph,
375            processes,
376            process_cancel_ability,
377            effect_controller,
378            runtime_dispatch: None,
379            runtime_execution_context: None,
380            cancellation_token: None,
381            async_process_id: None,
382            runtime_process_id: None,
383            process_events: None,
384            attachment_store,
385            direct_completions,
386            prepared_payload: serde_json::Value::Null,
387            tool_execution_binding: serde_json::Value::Null,
388            tool_call_id: None,
389            completion: ToolCompletionState::default(),
390            parent_invocation: None,
391            execution_env_spec: crate::ProcessExecutionEnvSpec::new(
392                crate::PluginOptions::default(),
393                crate::SessionPolicy::default(),
394            ),
395            child_execution_trace_hook: None,
396        }
397    }
398
399    pub(crate) fn from_dispatch(
400        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
401    ) -> ToolContextBuilder<'run> {
402        ToolContextBuilder::from_dispatch(dispatch)
403    }
404
405    pub fn session_id(&self) -> &str {
406        &self.session_id
407    }
408
409    pub fn agent_frame_id(&self) -> &str {
410        &self.agent_frame_id
411    }
412
413    pub fn sessions(&self) -> ToolSessionAdmin<'run> {
414        ToolSessionAdmin {
415            session_id: self.session_id.clone(),
416            sessions: Arc::clone(&self.sessions),
417            session_lifecycle: Arc::clone(&self.session_lifecycle),
418            effect_controller: self.effect_controller.clone(),
419        }
420    }
421
422    pub fn dispatch(&self) -> ToolDispatchClient<'run> {
423        ToolDispatchClient {
424            context: self.clone(),
425        }
426    }
427
428    pub fn triggers(&self) -> ToolTriggerClient<'run> {
429        ToolTriggerClient {
430            context: self.clone(),
431        }
432    }
433
434    pub fn processes(&self) -> ToolSessionProcessAdmin<'run> {
435        ToolSessionProcessAdmin {
436            session_id: self.session_id.clone(),
437            agent_frame_id: self.agent_frame_id.clone(),
438            processes: Arc::clone(&self.processes),
439            process_cancel_ability: Arc::clone(&self.process_cancel_ability),
440            effect_controller: self.effect_controller.clone(),
441            parent_invocation: self.parent_invocation.clone(),
442            tool_call_id: self.tool_call_id.clone(),
443            execution_env_spec: self.execution_env_spec.clone(),
444        }
445    }
446
447    pub fn emit_child_process_started(
448        &self,
449        process_id: impl Into<String>,
450        child_entry_name: Option<String>,
451    ) {
452        let Some(hook) = &self.child_execution_trace_hook else {
453            return;
454        };
455        hook.child_process_started(ToolChildProcessStarted {
456            process_id: process_id.into(),
457            child_entry_name,
458        });
459    }
460
461    pub fn direct_completions(&self) -> ToolDirectCompletionClient<'run> {
462        ToolDirectCompletionClient {
463            session_id: self.session_id.clone(),
464            tool_call_id: self.tool_call_id.clone(),
465            direct_completions: self.direct_completions.clone(),
466            parent_invocation: self.parent_invocation.clone(),
467        }
468    }
469
470    pub fn attachments(&self) -> ToolAttachmentClient {
471        ToolAttachmentClient {
472            store: Arc::clone(&self.attachment_store),
473        }
474    }
475
476    pub fn process_events(&self) -> ToolProcessEventClient {
477        ToolProcessEventClient {
478            context: self.process_events.clone(),
479        }
480    }
481
482    pub fn cancellation_token(&self) -> Option<&tokio_util::sync::CancellationToken> {
483        self.cancellation_token.as_ref()
484    }
485
486    #[doc(hidden)]
487    pub fn named_phase(&self, phase: &'static str) -> crate::runtime::RuntimeNamedPhase {
488        match self.runtime_execution_context.as_ref() {
489            Some(context) => context.named_phase(phase),
490            None => crate::runtime::RuntimeNamedPhase::begin(None, phase),
491        }
492    }
493
494    pub fn async_process_id(&self) -> Option<&str> {
495        self.async_process_id.as_deref()
496    }
497
498    pub fn runtime_process_id(&self) -> Option<&str> {
499        self.async_process_id
500            .as_deref()
501            .or(self.runtime_process_id.as_deref())
502            .or_else(|| {
503                self.process_events
504                    .as_ref()
505                    .map(|context| context.process_id.as_str())
506            })
507    }
508
509    pub fn tool_call_id(&self) -> Option<&str> {
510        self.tool_call_id.as_deref()
511    }
512
513    pub fn prepared_payload(&self) -> &serde_json::Value {
514        &self.prepared_payload
515    }
516
517    pub fn tool_execution_binding(&self) -> &serde_json::Value {
518        &self.tool_execution_binding
519    }
520
521    pub fn decode_prepared_payload<T>(&self) -> Result<T, serde_json::Error>
522    where
523        T: serde::de::DeserializeOwned,
524    {
525        serde_json::from_value(self.prepared_payload.clone())
526    }
527
528    pub fn attempt_number(&self) -> u32 {
529        self.attempt_number
530    }
531
532    pub fn max_attempts(&self) -> u32 {
533        self.max_attempts
534    }
535
536    pub fn replay_key(&self) -> Option<&str> {
537        self.replay_key.as_deref()
538    }
539
540    /// Obtain the durable completion key for this call, required before returning
541    /// [`ToolResult::Pending`](crate::ToolResult::Pending).
542    ///
543    /// A tool that defers its outcome (waiting on a webhook, human approval, or another
544    /// service) calls this, hands the returned [`AwaitEventKey`](crate::AwaitEventKey)
545    /// to whatever will complete the work out-of-band, and then returns
546    /// `ToolResult::Pending(..)`. The key names the durable wait the runtime parks the
547    /// call on; the external resolver delivers the result against it later.
548    ///
549    /// The key is stored on the context and consumed by the dispatcher when the tool
550    /// returns `Pending`. Returning `Pending` without first calling this fails the call
551    /// with `pending_tool_missing_completion_key`. Calls made outside a prepared tool
552    /// invocation (no tool call id) fail with `tool_completion_key_missing_call_id`.
553    pub async fn completion_key(&self) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
554        let tool_call_id = self.tool_call_id.clone().ok_or_else(|| {
555            crate::RuntimeError::new(
556                "tool_completion_key_missing_call_id",
557                "completion keys require a prepared tool call id",
558            )
559        })?;
560        let scoped = self.effect_controller.scoped();
561        if scoped.controller().durability_tier() == crate::DurabilityTier::Inline
562            && !scoped
563                .controller()
564                .allows_process_lifetime_completion_keys()
565        {
566            return Err(crate::RuntimeError::new(
567                "tool_completion_key_process_lifetime",
568                "completion keys on an Inline-tier host die with the current process; construct the InlineEffectHost with allow_process_lifetime_completion_keys() only for an explicitly single-process deployment",
569            ));
570        }
571        let key = scoped
572            .controller()
573            .await_event_key(
574                scoped.execution_scope(),
575                crate::AwaitEventWaitIdentity::tool_completion(tool_call_id),
576            )
577            .await?;
578        self.completion.store(key)
579    }
580
581    pub(crate) fn take_completion_key(
582        &self,
583    ) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
584        self.completion.take()
585    }
586
587    pub fn with_async_process(
588        mut self,
589        process_id: impl Into<String>,
590        cancellation_token: tokio_util::sync::CancellationToken,
591    ) -> Self {
592        self.async_process_id = Some(process_id.into());
593        self.runtime_process_id = self.async_process_id.clone();
594        self.cancellation_token = Some(cancellation_token);
595        self
596    }
597
598    #[cfg(any(test, feature = "testing"))]
599    #[doc(hidden)]
600    pub fn with_process_events_for_testing(
601        mut self,
602        process_id: impl Into<String>,
603        registry: Arc<dyn crate::ProcessRegistry>,
604    ) -> Self {
605        let awaiter = crate::ProcessAwaiter::polling(Arc::clone(&registry));
606        self.process_events = Some(ToolProcessEventContext {
607            process_id: process_id.into(),
608            registry,
609            awaiter,
610            store: None,
611            session_store_factory: None,
612            session_graph: Arc::new(crate::plugin::NoopSessionManager),
613            queued_work_driver: None,
614        });
615        self
616    }
617
618    pub(crate) fn with_retry_context(
619        mut self,
620        tool_name: &str,
621        attempt_number: u32,
622        max_attempts: u32,
623    ) -> Self {
624        self.attempt_number = attempt_number.max(1);
625        self.max_attempts = max_attempts.max(1);
626        self.replay_key = self
627            .tool_call_id
628            .as_ref()
629            .map(|call_id| format!("lash-tool:{}:{call_id}:{tool_name}", self.session_id));
630        self
631    }
632
633    pub(crate) fn with_prepared_payload(mut self, payload: serde_json::Value) -> Self {
634        self.prepared_payload = payload;
635        self
636    }
637
638    pub(crate) fn with_tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
639        self.tool_execution_binding = binding;
640        self
641    }
642
643    pub(crate) fn with_attempt_dispatch(
644        mut self,
645        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
646        parent_invocation: crate::RuntimeInvocation,
647    ) -> Self {
648        self.effect_controller = dispatch.effect_controller.clone();
649        self.runtime_dispatch = Some(dispatch);
650        self.parent_invocation = Some(parent_invocation);
651        self
652    }
653
654    /// Constructor reserved for `lash_core::testing` helpers. Do not call directly;
655    /// use [`lash_core::testing::mock_tool_context`] instead.
656    #[cfg(any(test, feature = "testing"))]
657    #[doc(hidden)]
658    #[expect(
659        clippy::too_many_arguments,
660        reason = "test-only constructor mirrors the sealed runtime tool context"
661    )]
662    pub fn __for_testing(
663        session_id: String,
664        sessions: Arc<dyn SessionStateService>,
665        session_lifecycle: Arc<dyn SessionLifecycleService>,
666        session_graph: Arc<dyn SessionGraphService>,
667        processes: Arc<dyn crate::ProcessService>,
668        attachment_store: Arc<crate::SessionAttachmentStore>,
669        direct_completions: crate::DirectCompletionClient<'static>,
670        tool_call_id: Option<String>,
671    ) -> ToolContext<'static> {
672        ToolContext::builder(
673            session_id,
674            sessions,
675            session_lifecycle,
676            session_graph,
677            processes,
678            Arc::new(crate::DefaultProcessCancelAbility),
679            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
680                crate::InlineRuntimeEffectController::default()
681                    .allow_process_lifetime_completion_keys(),
682            )),
683            attachment_store,
684            direct_completions,
685        )
686        .tool_call_id(tool_call_id)
687        .build()
688    }
689
690    /// Constructor reserved for tests that need a custom process-cancel host
691    /// ability. Do not call directly; prefer public testing helpers when they
692    /// cover the case.
693    #[cfg(any(test, feature = "testing"))]
694    #[doc(hidden)]
695    #[expect(
696        clippy::too_many_arguments,
697        reason = "test-only constructor mirrors the sealed runtime context"
698    )]
699    pub fn __for_testing_with_process_cancel_ability(
700        session_id: String,
701        sessions: Arc<dyn SessionStateService>,
702        session_lifecycle: Arc<dyn SessionLifecycleService>,
703        session_graph: Arc<dyn SessionGraphService>,
704        processes: Arc<dyn crate::ProcessService>,
705        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
706        attachment_store: Arc<crate::SessionAttachmentStore>,
707        direct_completions: crate::DirectCompletionClient<'static>,
708        tool_call_id: Option<String>,
709    ) -> ToolContext<'static> {
710        ToolContext::builder(
711            session_id,
712            sessions,
713            session_lifecycle,
714            session_graph,
715            processes,
716            process_cancel_ability,
717            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
718                crate::InlineRuntimeEffectController::default()
719                    .allow_process_lifetime_completion_keys(),
720            )),
721            attachment_store,
722            direct_completions,
723        )
724        .tool_call_id(tool_call_id)
725        .build()
726    }
727}
728
729/// Runtime-prepared executable tool call.
730///
731/// The raw model/provider identity remains visible, but any argument rewrites
732/// and provider-owned context projections are frozen before the call crosses a
733/// runtime effect or process boundary.
734#[derive(Clone, Debug, Serialize, Deserialize)]
735pub struct PreparedToolCall {
736    pub call_id: String,
737    pub tool_id: ToolId,
738    pub tool_name: String,
739    pub args: serde_json::Value,
740    #[serde(default, skip_serializing_if = "Option::is_none")]
741    pub replay: Option<ProviderReplayMeta>,
742    #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
743    pub prepared_payload: serde_json::Value,
744}
745
746impl PreparedToolCall {
747    pub fn identity(tool_id: ToolId, call: crate::sansio::PendingToolCall) -> Self {
748        Self {
749            call_id: call.call_id,
750            tool_id,
751            tool_name: call.tool_name,
752            args: call.args,
753            replay: call.replay,
754            prepared_payload: serde_json::Value::Null,
755        }
756    }
757
758    pub fn from_parts(
759        call_id: impl Into<String>,
760        tool_id: impl Into<ToolId>,
761        tool_name: impl Into<String>,
762        args: serde_json::Value,
763        replay: Option<ProviderReplayMeta>,
764        prepared_payload: serde_json::Value,
765    ) -> Self {
766        Self {
767            call_id: call_id.into(),
768            tool_id: tool_id.into(),
769            tool_name: tool_name.into(),
770            args,
771            replay,
772            prepared_payload,
773        }
774    }
775}
776
777/// One ordered child inside a runtime-prepared tool batch.
778///
779/// The call itself carries the executable provider payload. `replay_suffix`
780/// is the deterministic suffix used for child effects such as retry sleeps or
781/// pending completion awaits when the batch is the durable parent.
782#[derive(Clone, Debug, Serialize, Deserialize)]
783pub struct PreparedToolBatchCall {
784    pub call: PreparedToolCall,
785    pub replay_suffix: String,
786    #[serde(default, skip_serializing_if = "Option::is_none")]
787    pub execution_grant: Option<Box<ToolExecutionGrant>>,
788}
789
790/// Runtime-prepared executable tool batch.
791///
792/// The vector order is source order. Calls run concurrently, but launches and
793/// pending completion consumption are projected back through this order.
794#[derive(Clone, Debug, Serialize, Deserialize)]
795pub struct PreparedToolBatch {
796    pub batch_id: String,
797    pub calls: Vec<PreparedToolBatchCall>,
798}
799
800impl PreparedToolBatch {
801    pub fn new(batch_id: impl Into<String>, calls: Vec<PreparedToolCall>) -> Self {
802        let batch_id = batch_id.into();
803        let calls = calls
804            .into_iter()
805            .enumerate()
806            .map(|(index, call)| PreparedToolBatchCall {
807                replay_suffix: format!("child:{index}:{}", call.call_id),
808                call,
809                execution_grant: None,
810            })
811            .collect();
812        Self { batch_id, calls }
813    }
814
815    pub fn new_with_grants(
816        batch_id: impl Into<String>,
817        calls: Vec<(PreparedToolCall, Option<ToolExecutionGrant>)>,
818    ) -> Self {
819        let batch_id = batch_id.into();
820        let calls = calls
821            .into_iter()
822            .enumerate()
823            .map(|(index, (call, execution_grant))| PreparedToolBatchCall {
824                replay_suffix: format!("child:{index}:{}", call.call_id),
825                call,
826                execution_grant: execution_grant.map(Box::new),
827            })
828            .collect();
829        Self { batch_id, calls }
830    }
831
832    pub fn is_empty(&self) -> bool {
833        self.calls.is_empty()
834    }
835
836    pub fn len(&self) -> usize {
837        self.calls.len()
838    }
839}
840
841/// Explicit authority to execute a tool outside Tool Catalog membership.
842///
843/// Normal tool calls are authorized by catalog membership. A grant is a
844/// separate, caller-provided capability used by deferred resolution flows: it
845/// carries the manifest/contract to validate the call plus an opaque host
846/// execution binding that providers can inspect from the prepare and execute
847/// contexts.
848#[derive(Clone, Debug, Serialize, Deserialize)]
849pub struct ToolExecutionGrant {
850    /// Tool identity and model-facing metadata authorized by the grant.
851    pub manifest: ToolManifest,
852    /// Contract used to validate granted call arguments without consulting the
853    /// current Tool Catalog.
854    pub contract: Box<ToolContract>,
855    /// Explicit registry source route for registry-backed execution. Direct
856    /// non-registry providers may ignore this; [`ToolRegistry`] requires it.
857    pub source_id: Option<String>,
858    /// Opaque host routing payload passed to prepare and execute contexts.
859    pub execution_binding: serde_json::Value,
860}
861
862impl ToolExecutionGrant {
863    pub fn new(manifest: ToolManifest, contract: ToolContract) -> Self {
864        Self {
865            manifest,
866            contract: Box::new(contract),
867            source_id: None,
868            execution_binding: serde_json::Value::Null,
869        }
870    }
871
872    pub fn from_definition(definition: ToolDefinition) -> Self {
873        Self::new(definition.manifest(), definition.contract())
874    }
875
876    pub fn with_source_id(mut self, source_id: impl Into<String>) -> Self {
877        self.source_id = Some(source_id.into());
878        self
879    }
880
881    pub fn with_execution_binding(mut self, execution_binding: serde_json::Value) -> Self {
882        self.execution_binding = execution_binding;
883        self
884    }
885}
886
887#[derive(Clone)]
888pub struct ToolPrepareContext {
889    session_id: String,
890    sessions: Arc<dyn SessionStateService>,
891    turn_context: crate::TurnContext,
892    tool_call_id: Option<String>,
893    tool_execution_binding: serde_json::Value,
894}
895
896impl ToolPrepareContext {
897    pub(crate) fn with_execution_binding(
898        session_id: String,
899        sessions: Arc<dyn SessionStateService>,
900        turn_context: crate::TurnContext,
901        tool_call_id: Option<String>,
902        tool_execution_binding: serde_json::Value,
903    ) -> Self {
904        Self {
905            session_id,
906            sessions,
907            turn_context,
908            tool_call_id,
909            tool_execution_binding,
910        }
911    }
912
913    pub fn session_id(&self) -> &str {
914        &self.session_id
915    }
916
917    pub fn tool_call_id(&self) -> Option<&str> {
918        self.tool_call_id.as_deref()
919    }
920
921    pub fn tool_execution_binding(&self) -> &serde_json::Value {
922        &self.tool_execution_binding
923    }
924
925    pub fn turn_context(&self) -> &crate::TurnContext {
926        &self.turn_context
927    }
928
929    pub fn plugin_input<T>(&self, plugin_id: &'static str) -> Option<&T>
930    where
931        T: 'static,
932    {
933        self.turn_context.plugin_input::<T>(plugin_id)
934    }
935
936    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
937        self.sessions.snapshot_session(&self.session_id).await
938    }
939
940    pub async fn tool_catalog(&self) -> Result<Vec<serde_json::Value>, PluginError> {
941        self.sessions.tool_catalog(&self.session_id).await
942    }
943
944    pub async fn shared_tool_catalog(
945        &self,
946    ) -> Result<std::sync::Arc<Vec<serde_json::Value>>, PluginError> {
947        self.sessions.shared_tool_catalog(&self.session_id).await
948    }
949}
950
951/// Inputs handed to [`ToolProvider::prepare_tool_call`].
952pub struct ToolPrepareCall<'a> {
953    pub tool_id: ToolId,
954    pub pending: crate::sansio::PendingToolCall,
955    pub context: &'a ToolPrepareContext,
956}
957
958/// Per-call inputs handed to [`ToolProvider::execute`].
959///
960/// Fields are `pub` because `ToolCall` is a transient borrow; consumers
961/// typically destructure (`let ToolCall { name, args, .. } = call`). The
962/// stable surface lives on [`ToolContext`] (sealed) and the runtime's
963/// dispatcher, which constructs `ToolCall` values.
964pub struct ToolCall<'a> {
965    pub name: &'a str,
966    pub args: &'a serde_json::Value,
967    pub context: &'a ToolContext<'a>,
968    pub progress: Option<&'a ProgressSender>,
969}
970
971/// Trait for providing tools to the sandbox. Implement this per-project.
972///
973/// Implementations supply cheap [`ToolManifest`]s, lazily resolved
974/// [`ToolContract`]s, and a single
975/// [`execute`](Self::execute) method that handles every call. Tools that
976/// need session state read it from `call.context`; tools that stream
977/// progress send through `call.progress`.
978#[async_trait::async_trait]
979pub trait ToolProvider: Send + Sync + 'static {
980    fn tool_manifests(&self) -> Vec<ToolManifest>;
981    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
982        self.tool_manifests()
983            .into_iter()
984            .find(|manifest| manifest.name == name)
985    }
986    fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
987        self.tool_manifests()
988            .into_iter()
989            .find(|manifest| manifest.id == *id)
990    }
991    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
992    fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
993        let manifest = self.resolve_manifest_by_id(id)?;
994        self.resolve_contract(&manifest.name)
995    }
996    async fn prepare_tool_call(
997        &self,
998        call: ToolPrepareCall<'_>,
999    ) -> Result<PreparedToolCall, ToolResult> {
1000        Ok(PreparedToolCall::identity(call.tool_id, call.pending))
1001    }
1002    async fn prepare_granted_tool_call(
1003        &self,
1004        grant: &ToolExecutionGrant,
1005        call: ToolPrepareCall<'_>,
1006    ) -> Result<PreparedToolCall, ToolResult> {
1007        let _ = call;
1008        Err(ToolResult::err_fmt(format_args!(
1009            "Granted execution is unsupported for tool id `{}`",
1010            grant.manifest.id
1011        )))
1012    }
1013    async fn execute(&self, call: ToolCall<'_>) -> ToolResult;
1014    async fn execute_granted(
1015        &self,
1016        grant: &ToolExecutionGrant,
1017        args: &serde_json::Value,
1018        context: &ToolContext<'_>,
1019        progress: Option<&ProgressSender>,
1020    ) -> ToolResult {
1021        let _ = (args, context, progress);
1022        ToolResult::err_fmt(format_args!(
1023            "Granted execution is unsupported for tool id `{}`",
1024            grant.manifest.id
1025        ))
1026    }
1027    async fn execute_by_id(
1028        &self,
1029        tool_id: &ToolId,
1030        args: &serde_json::Value,
1031        context: &ToolContext<'_>,
1032        progress: Option<&ProgressSender>,
1033    ) -> ToolResult {
1034        let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
1035            return ToolResult::err_fmt(format!("Unknown tool id: {tool_id}"));
1036        };
1037        self.execute(ToolCall {
1038            name: &manifest.name,
1039            args,
1040            context,
1041            progress,
1042        })
1043        .await
1044    }
1045}
1046
1047#[cfg(test)]
1048mod tests {
1049    use super::*;
1050
1051    #[test]
1052    fn tool_context_builder_carries_call_payload_and_cancellation_state() {
1053        let cancellation = tokio_util::sync::CancellationToken::new();
1054        let prepared = PreparedToolCall::from_parts(
1055            "call-1",
1056            "tool:demo_tool",
1057            "demo_tool",
1058            serde_json::json!({ "input": true }),
1059            None,
1060            serde_json::json!({ "prepared": true }),
1061        );
1062
1063        let context = ToolContext::builder(
1064            "session-1".to_string(),
1065            Arc::new(crate::testing::MockSessionManager::default()),
1066            Arc::new(crate::testing::MockSessionManager::default()),
1067            Arc::new(crate::testing::MockSessionManager::default()),
1068            Arc::new(crate::UnavailableProcessService),
1069            Arc::new(crate::DefaultProcessCancelAbility),
1070            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
1071                crate::InlineRuntimeEffectController::default(),
1072            )),
1073            Arc::new(crate::SessionAttachmentStore::in_memory()),
1074            crate::DirectCompletionClient::unavailable(
1075                "direct completions are unavailable in this test context",
1076            ),
1077        )
1078        .prepared_call(&prepared)
1079        .cancellation_token(Some(cancellation.clone()))
1080        .async_process("process-1", cancellation.clone())
1081        .build();
1082
1083        assert_eq!(context.session_id(), "session-1");
1084        assert_eq!(context.tool_call_id(), Some("call-1"));
1085        assert_eq!(
1086            context.prepared_payload(),
1087            &serde_json::json!({ "prepared": true })
1088        );
1089        assert_eq!(context.async_process_id(), Some("process-1"));
1090        assert!(context.cancellation_token().is_some());
1091    }
1092
1093    #[tokio::test]
1094    async fn inline_completion_key_requires_process_lifetime_opt_in() {
1095        let prepared = PreparedToolCall::from_parts(
1096            "call-inline-risk",
1097            "tool:demo_tool",
1098            "demo_tool",
1099            serde_json::json!({}),
1100            None,
1101            serde_json::json!({}),
1102        );
1103        let context = ToolContext::builder(
1104            "session-inline-risk".to_string(),
1105            Arc::new(crate::testing::MockSessionManager::default()),
1106            Arc::new(crate::testing::MockSessionManager::default()),
1107            Arc::new(crate::testing::MockSessionManager::default()),
1108            Arc::new(crate::UnavailableProcessService),
1109            Arc::new(crate::DefaultProcessCancelAbility),
1110            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
1111                crate::InlineRuntimeEffectController::default(),
1112            )),
1113            Arc::new(crate::SessionAttachmentStore::in_memory()),
1114            crate::DirectCompletionClient::unavailable(
1115                "direct completions are unavailable in this test context",
1116            ),
1117        )
1118        .prepared_call(&prepared)
1119        .build();
1120
1121        let error = context
1122            .completion_key()
1123            .await
1124            .expect_err("Inline completion keys must refuse by default");
1125        assert_eq!(error.code.as_str(), "tool_completion_key_process_lifetime");
1126        assert!(error.message.contains("die with the current process"));
1127        assert!(
1128            error
1129                .message
1130                .contains("allow_process_lifetime_completion_keys")
1131        );
1132    }
1133}
lash_core/tool_provider.rs

lash_core/
tool_provider.rs
