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::{AttachmentStore, ToolContract, 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) cancellation_token: Option<tokio_util::sync::CancellationToken>,
84    pub(crate) async_process_id: Option<String>,
85    pub(crate) runtime_process_id: Option<String>,
86    pub(crate) process_events: Option<ToolProcessEventContext>,
87    pub(crate) attachment_store: Arc<dyn AttachmentStore>,
88    pub(crate) direct_completions: crate::DirectCompletionClient<'run>,
89    pub(crate) prepared_payload: serde_json::Value,
90    /// The id of the in-flight tool call that is invoking this tool.
91    pub(crate) tool_call_id: Option<String>,
92    pub(crate) attempt_number: u32,
93    pub(crate) max_attempts: u32,
94    pub(crate) replay_key: Option<String>,
95    pub(crate) completion: ToolCompletionState,
96    pub(crate) parent_invocation: Option<crate::RuntimeInvocation>,
97    pub(crate) execution_env_spec: crate::ProcessExecutionEnvSpec,
98    pub(crate) lashlang_execution_call_site: Option<ToolLashlangExecutionCallSite>,
99}
100
101#[derive(Clone)]
102pub struct ToolLashlangExecutionCallSite {
103    sink: Arc<dyn lash_trace::TraceSink>,
104    base_context: lash_trace::TraceContext,
105    identity: lash_trace::TraceLashlangExecutionIdentity,
106    parent_node_id: String,
107    occurrence: u64,
108}
109
110impl ToolLashlangExecutionCallSite {
111    pub fn new(
112        sink: Arc<dyn lash_trace::TraceSink>,
113        base_context: lash_trace::TraceContext,
114        identity: lash_trace::TraceLashlangExecutionIdentity,
115        parent_node_id: impl Into<String>,
116        occurrence: u64,
117    ) -> Self {
118        Self {
119            sink,
120            base_context,
121            identity,
122            parent_node_id: parent_node_id.into(),
123            occurrence,
124        }
125    }
126}
127
128#[derive(Clone)]
129pub(crate) struct ToolProcessEventContext {
130    process_id: String,
131    registry: Arc<dyn crate::ProcessRegistry>,
132    store: Option<Arc<dyn crate::RuntimePersistence>>,
133    session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
134    session_graph: Arc<dyn SessionGraphService>,
135    queued_work_poke: Option<crate::QueuedWorkPoke>,
136}
137
138pub(crate) struct ToolContextBuilder<'run> {
139    session_id: String,
140    agent_frame_id: crate::AgentFrameId,
141    sessions: Arc<dyn SessionStateService>,
142    session_lifecycle: Arc<dyn SessionLifecycleService>,
143    session_graph: Arc<dyn SessionGraphService>,
144    processes: Arc<dyn crate::ProcessService>,
145    process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
146    effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
147    runtime_dispatch: Option<Arc<crate::tool_dispatch::ToolDispatchContext<'run>>>,
148    cancellation_token: Option<tokio_util::sync::CancellationToken>,
149    async_process_id: Option<String>,
150    runtime_process_id: Option<String>,
151    process_events: Option<ToolProcessEventContext>,
152    attachment_store: Arc<dyn AttachmentStore>,
153    direct_completions: crate::DirectCompletionClient<'run>,
154    prepared_payload: serde_json::Value,
155    tool_call_id: Option<String>,
156    completion: ToolCompletionState,
157    parent_invocation: Option<crate::RuntimeInvocation>,
158    execution_env_spec: crate::ProcessExecutionEnvSpec,
159    lashlang_execution_call_site: Option<ToolLashlangExecutionCallSite>,
160}
161
162impl<'run> ToolContextBuilder<'run> {
163    pub(crate) fn from_dispatch(
164        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
165    ) -> Self {
166        Self {
167            session_id: dispatch.session_id.clone(),
168            agent_frame_id: dispatch.agent_frame_id.clone(),
169            sessions: Arc::clone(&dispatch.sessions),
170            session_lifecycle: Arc::clone(&dispatch.session_lifecycle),
171            session_graph: Arc::clone(&dispatch.session_graph),
172            processes: Arc::clone(&dispatch.processes),
173            process_cancel_ability: Arc::clone(&dispatch.process_cancel_ability),
174            effect_controller: dispatch.effect_controller.clone(),
175            runtime_dispatch: Some(Arc::clone(&dispatch)),
176            cancellation_token: None,
177            async_process_id: None,
178            runtime_process_id: None,
179            process_events: None,
180            attachment_store: Arc::clone(&dispatch.attachment_store),
181            direct_completions: dispatch.direct_completions.clone(),
182            prepared_payload: serde_json::Value::Null,
183            tool_call_id: None,
184            completion: ToolCompletionState::default(),
185            parent_invocation: dispatch.parent_invocation.clone(),
186            execution_env_spec: dispatch.execution_env_spec.clone(),
187            lashlang_execution_call_site: None,
188        }
189    }
190
191    #[cfg(any(test, feature = "testing"))]
192    pub(crate) fn tool_call_id(mut self, tool_call_id: impl Into<Option<String>>) -> Self {
193        self.tool_call_id = tool_call_id.into();
194        self
195    }
196
197    pub(crate) fn prepared_call(mut self, call: &PreparedToolCall) -> Self {
198        self.tool_call_id = Some(call.call_id.clone());
199        self.prepared_payload = call.prepared_payload.clone();
200        self
201    }
202
203    pub(crate) fn cancellation_token(
204        mut self,
205        cancellation_token: Option<tokio_util::sync::CancellationToken>,
206    ) -> Self {
207        self.cancellation_token = cancellation_token;
208        self
209    }
210
211    pub(crate) fn runtime_process_id(mut self, process_id: Option<String>) -> Self {
212        self.runtime_process_id = process_id;
213        self
214    }
215
216    pub(crate) fn async_process(
217        mut self,
218        process_id: impl Into<String>,
219        cancellation_token: tokio_util::sync::CancellationToken,
220    ) -> Self {
221        self.async_process_id = Some(process_id.into());
222        self.cancellation_token = Some(cancellation_token);
223        self
224    }
225
226    pub(crate) fn process_events(
227        mut self,
228        process_id: impl Into<String>,
229        registry: Arc<dyn crate::ProcessRegistry>,
230        store: Option<Arc<dyn crate::RuntimePersistence>>,
231        session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
232        queued_work_poke: Option<crate::QueuedWorkPoke>,
233    ) -> Self {
234        self.process_events = Some(ToolProcessEventContext {
235            process_id: process_id.into(),
236            registry,
237            store,
238            session_store_factory,
239            session_graph: Arc::clone(&self.session_graph),
240            queued_work_poke,
241        });
242        self
243    }
244
245    pub(crate) fn parent_invocation(mut self, metadata: Option<crate::RuntimeInvocation>) -> Self {
246        self.parent_invocation = metadata;
247        self
248    }
249
250    pub(crate) fn lashlang_execution_call_site(
251        mut self,
252        call_site: Option<ToolLashlangExecutionCallSite>,
253    ) -> Self {
254        self.lashlang_execution_call_site = call_site;
255        self
256    }
257
258    pub(crate) fn build(self) -> ToolContext<'run> {
259        ToolContext {
260            session_id: self.session_id,
261            agent_frame_id: self.agent_frame_id,
262            sessions: self.sessions,
263            session_lifecycle: self.session_lifecycle,
264            processes: self.processes,
265            process_cancel_ability: self.process_cancel_ability,
266            effect_controller: self.effect_controller,
267            runtime_dispatch: self.runtime_dispatch,
268            cancellation_token: self.cancellation_token,
269            async_process_id: self.async_process_id,
270            runtime_process_id: self.runtime_process_id,
271            process_events: self.process_events,
272            attachment_store: self.attachment_store,
273            direct_completions: self.direct_completions,
274            prepared_payload: self.prepared_payload,
275            tool_call_id: self.tool_call_id,
276            attempt_number: 1,
277            max_attempts: 1,
278            replay_key: None,
279            completion: self.completion,
280            parent_invocation: self.parent_invocation,
281            execution_env_spec: self.execution_env_spec,
282            lashlang_execution_call_site: self.lashlang_execution_call_site,
283        }
284    }
285}
286
287impl<'run> ToolContext<'run> {
288    #[cfg(any(test, feature = "testing"))]
289    #[expect(
290        clippy::too_many_arguments,
291        reason = "testing constructor mirrors the sealed runtime tool context dependencies"
292    )]
293    pub(crate) fn builder(
294        session_id: String,
295        sessions: Arc<dyn SessionStateService>,
296        session_lifecycle: Arc<dyn SessionLifecycleService>,
297        session_graph: Arc<dyn SessionGraphService>,
298        processes: Arc<dyn crate::ProcessService>,
299        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
300        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
301        attachment_store: Arc<dyn AttachmentStore>,
302        direct_completions: crate::DirectCompletionClient<'run>,
303    ) -> ToolContextBuilder<'run> {
304        ToolContextBuilder {
305            session_id,
306            agent_frame_id: String::new(),
307            sessions,
308            session_lifecycle,
309            session_graph,
310            processes,
311            process_cancel_ability,
312            effect_controller,
313            runtime_dispatch: None,
314            cancellation_token: None,
315            async_process_id: None,
316            runtime_process_id: None,
317            process_events: None,
318            attachment_store,
319            direct_completions,
320            prepared_payload: serde_json::Value::Null,
321            tool_call_id: None,
322            completion: ToolCompletionState::default(),
323            parent_invocation: None,
324            execution_env_spec: crate::ProcessExecutionEnvSpec::new(
325                crate::PluginOptions::default(),
326                crate::SessionPolicy::default(),
327            ),
328            lashlang_execution_call_site: None,
329        }
330    }
331
332    pub(crate) fn from_dispatch(
333        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
334    ) -> ToolContextBuilder<'run> {
335        ToolContextBuilder::from_dispatch(dispatch)
336    }
337
338    pub fn session_id(&self) -> &str {
339        &self.session_id
340    }
341
342    pub fn agent_frame_id(&self) -> &str {
343        &self.agent_frame_id
344    }
345
346    pub fn sessions(&self) -> ToolSessionAdmin<'run> {
347        ToolSessionAdmin {
348            session_id: self.session_id.clone(),
349            sessions: Arc::clone(&self.sessions),
350            session_lifecycle: Arc::clone(&self.session_lifecycle),
351            effect_controller: self.effect_controller.clone(),
352        }
353    }
354
355    pub fn dispatch(&self) -> ToolDispatchClient<'run> {
356        ToolDispatchClient {
357            context: self.clone(),
358        }
359    }
360
361    pub fn triggers(&self) -> ToolTriggerClient<'run> {
362        ToolTriggerClient {
363            context: self.clone(),
364        }
365    }
366
367    pub fn processes(&self) -> ToolSessionProcessAdmin<'run> {
368        ToolSessionProcessAdmin {
369            session_id: self.session_id.clone(),
370            agent_frame_id: self.agent_frame_id.clone(),
371            processes: Arc::clone(&self.processes),
372            process_cancel_ability: Arc::clone(&self.process_cancel_ability),
373            effect_controller: self.effect_controller.clone(),
374            parent_invocation: self.parent_invocation.clone(),
375            tool_call_id: self.tool_call_id.clone(),
376            execution_env_spec: self.execution_env_spec.clone(),
377        }
378    }
379
380    pub fn emit_lashlang_child_process_started(
381        &self,
382        process_id: impl Into<String>,
383        child_entry_name: Option<String>,
384    ) {
385        let Some(call_site) = &self.lashlang_execution_call_site else {
386            return;
387        };
388        let child = lash_trace::TraceLashlangChildExecution {
389            scope: call_site.identity.scope.clone(),
390            subject: lash_trace::TraceRuntimeSubject::Process {
391                process_id: process_id.into(),
392            },
393            module_ref: None,
394            entry_ref: None,
395            entry_name: child_entry_name,
396        };
397        let child_graph_key = child.graph_key();
398        let event = lash_trace::TraceLashlangExecutionEvent::ChildStarted {
399            event_key: format!(
400                "lashlang_execution:{}:child:{}:{}:{}",
401                call_site.identity.graph_key(),
402                call_site.parent_node_id,
403                call_site.occurrence,
404                child_graph_key
405            ),
406            identity: call_site.identity.clone(),
407            parent_node_id: call_site.parent_node_id.clone(),
408            occurrence: call_site.occurrence,
409            child,
410        };
411        let mut context = lash_trace::TraceContext::default()
412            .for_session(call_site.identity.scope.session_id.clone());
413        if let Some(turn_id) = &call_site.identity.scope.turn_id {
414            context = context.for_turn(turn_id.clone());
415        }
416        if let Some(turn_index) = call_site.identity.scope.turn_index {
417            context = context.for_turn_index(turn_index);
418        }
419        if let Some(protocol_iteration) = call_site.identity.scope.protocol_iteration {
420            context = context.for_protocol_iteration(protocol_iteration);
421        }
422        if let lash_trace::TraceRuntimeSubject::Effect { effect_id, .. } =
423            &call_site.identity.subject
424        {
425            context.effect_id = Some(effect_id.clone());
426        }
427        crate::trace::emit_trace(
428            &Some(Arc::clone(&call_site.sink)),
429            &call_site.base_context,
430            context,
431            lash_trace::TraceEvent::LashlangExecution { event },
432        );
433    }
434
435    pub fn direct_completions(&self) -> ToolDirectCompletionClient<'run> {
436        ToolDirectCompletionClient {
437            session_id: self.session_id.clone(),
438            tool_call_id: self.tool_call_id.clone(),
439            direct_completions: self.direct_completions.clone(),
440        }
441    }
442
443    pub fn attachments(&self) -> ToolAttachmentClient {
444        ToolAttachmentClient {
445            store: Arc::clone(&self.attachment_store),
446        }
447    }
448
449    pub fn process_events(&self) -> ToolProcessEventClient {
450        ToolProcessEventClient {
451            context: self.process_events.clone(),
452        }
453    }
454
455    pub fn cancellation_token(&self) -> Option<&tokio_util::sync::CancellationToken> {
456        self.cancellation_token.as_ref()
457    }
458
459    pub fn async_process_id(&self) -> Option<&str> {
460        self.async_process_id.as_deref()
461    }
462
463    pub fn runtime_process_id(&self) -> Option<&str> {
464        self.async_process_id
465            .as_deref()
466            .or(self.runtime_process_id.as_deref())
467            .or_else(|| {
468                self.process_events
469                    .as_ref()
470                    .map(|context| context.process_id.as_str())
471            })
472    }
473
474    pub fn tool_call_id(&self) -> Option<&str> {
475        self.tool_call_id.as_deref()
476    }
477
478    pub fn prepared_payload(&self) -> &serde_json::Value {
479        &self.prepared_payload
480    }
481
482    pub fn decode_prepared_payload<T>(&self) -> Result<T, serde_json::Error>
483    where
484        T: serde::de::DeserializeOwned,
485    {
486        serde_json::from_value(self.prepared_payload.clone())
487    }
488
489    pub fn attempt_number(&self) -> u32 {
490        self.attempt_number
491    }
492
493    pub fn max_attempts(&self) -> u32 {
494        self.max_attempts
495    }
496
497    pub fn replay_key(&self) -> Option<&str> {
498        self.replay_key.as_deref()
499    }
500
501    /// Obtain the durable completion key for this call, required before returning
502    /// [`ToolResult::Pending`](crate::ToolResult::Pending).
503    ///
504    /// A tool that defers its outcome (waiting on a webhook, human approval, or another
505    /// service) calls this, hands the returned [`AwaitEventKey`](crate::AwaitEventKey)
506    /// to whatever will complete the work out-of-band, and then returns
507    /// `ToolResult::Pending(..)`. The key names the durable wait the runtime parks the
508    /// call on; the external resolver delivers the result against it later.
509    ///
510    /// The key is stored on the context and consumed by the dispatcher when the tool
511    /// returns `Pending`. Returning `Pending` without first calling this fails the call
512    /// with `pending_tool_missing_completion_key`. Calls made outside a prepared tool
513    /// invocation (no tool call id) fail with `tool_completion_key_missing_call_id`.
514    pub async fn completion_key(&self) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
515        let tool_call_id = self.tool_call_id.clone().ok_or_else(|| {
516            crate::RuntimeError::new(
517                "tool_completion_key_missing_call_id",
518                "completion keys require a prepared tool call id",
519            )
520        })?;
521        let scoped = self.effect_controller.scoped();
522        let key = scoped
523            .controller()
524            .await_event_key(
525                scoped.execution_scope(),
526                crate::AwaitEventWaitIdentity::tool_completion(tool_call_id),
527            )
528            .await?;
529        self.completion.store(key)
530    }
531
532    pub(crate) fn take_completion_key(
533        &self,
534    ) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
535        self.completion.take()
536    }
537
538    pub fn with_async_process(
539        mut self,
540        process_id: impl Into<String>,
541        cancellation_token: tokio_util::sync::CancellationToken,
542    ) -> Self {
543        self.async_process_id = Some(process_id.into());
544        self.runtime_process_id = self.async_process_id.clone();
545        self.cancellation_token = Some(cancellation_token);
546        self
547    }
548
549    #[cfg(any(test, feature = "testing"))]
550    #[doc(hidden)]
551    pub fn with_process_events_for_testing(
552        mut self,
553        process_id: impl Into<String>,
554        registry: Arc<dyn crate::ProcessRegistry>,
555    ) -> Self {
556        self.process_events = Some(ToolProcessEventContext {
557            process_id: process_id.into(),
558            registry,
559            store: None,
560            session_store_factory: None,
561            session_graph: Arc::new(crate::plugin::NoopSessionManager),
562            queued_work_poke: None,
563        });
564        self
565    }
566
567    pub(crate) fn with_retry_context(
568        mut self,
569        tool_name: &str,
570        attempt_number: u32,
571        max_attempts: u32,
572    ) -> Self {
573        self.attempt_number = attempt_number.max(1);
574        self.max_attempts = max_attempts.max(1);
575        self.replay_key = self
576            .tool_call_id
577            .as_ref()
578            .map(|call_id| format!("lash-tool:{}:{call_id}:{tool_name}", self.session_id));
579        self
580    }
581
582    pub(crate) fn with_prepared_payload(mut self, payload: serde_json::Value) -> Self {
583        self.prepared_payload = payload;
584        self
585    }
586
587    /// Constructor reserved for `lash_core::testing` helpers. Do not call directly;
588    /// use [`lash_core::testing::mock_tool_context`] instead.
589    #[cfg(any(test, feature = "testing"))]
590    #[doc(hidden)]
591    #[expect(
592        clippy::too_many_arguments,
593        reason = "test-only constructor mirrors the sealed runtime tool context"
594    )]
595    pub fn __for_testing(
596        session_id: String,
597        sessions: Arc<dyn SessionStateService>,
598        session_lifecycle: Arc<dyn SessionLifecycleService>,
599        session_graph: Arc<dyn SessionGraphService>,
600        processes: Arc<dyn crate::ProcessService>,
601        attachment_store: Arc<dyn AttachmentStore>,
602        direct_completions: crate::DirectCompletionClient<'static>,
603        tool_call_id: Option<String>,
604    ) -> ToolContext<'static> {
605        ToolContext::builder(
606            session_id,
607            sessions,
608            session_lifecycle,
609            session_graph,
610            processes,
611            Arc::new(crate::DefaultProcessCancelAbility),
612            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
613                crate::InlineRuntimeEffectController,
614            )),
615            attachment_store,
616            direct_completions,
617        )
618        .tool_call_id(tool_call_id)
619        .build()
620    }
621
622    /// Constructor reserved for tests that need a custom process-cancel host
623    /// ability. Do not call directly; prefer public testing helpers when they
624    /// cover the case.
625    #[cfg(any(test, feature = "testing"))]
626    #[doc(hidden)]
627    #[expect(
628        clippy::too_many_arguments,
629        reason = "test-only constructor mirrors the sealed runtime context"
630    )]
631    pub fn __for_testing_with_process_cancel_ability(
632        session_id: String,
633        sessions: Arc<dyn SessionStateService>,
634        session_lifecycle: Arc<dyn SessionLifecycleService>,
635        session_graph: Arc<dyn SessionGraphService>,
636        processes: Arc<dyn crate::ProcessService>,
637        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
638        attachment_store: Arc<dyn AttachmentStore>,
639        direct_completions: crate::DirectCompletionClient<'static>,
640        tool_call_id: Option<String>,
641    ) -> ToolContext<'static> {
642        ToolContext::builder(
643            session_id,
644            sessions,
645            session_lifecycle,
646            session_graph,
647            processes,
648            process_cancel_ability,
649            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
650                crate::InlineRuntimeEffectController,
651            )),
652            attachment_store,
653            direct_completions,
654        )
655        .tool_call_id(tool_call_id)
656        .build()
657    }
658}
659
660/// Runtime-prepared executable tool call.
661///
662/// The raw model/provider identity remains visible, but any argument rewrites
663/// and provider-owned context projections are frozen before the call crosses a
664/// runtime effect or process boundary.
665#[derive(Clone, Debug, Serialize, Deserialize)]
666pub struct PreparedToolCall {
667    pub call_id: String,
668    pub tool_name: String,
669    pub args: serde_json::Value,
670    #[serde(default, skip_serializing_if = "Option::is_none")]
671    pub replay: Option<ProviderReplayMeta>,
672    #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
673    pub prepared_payload: serde_json::Value,
674}
675
676impl PreparedToolCall {
677    pub fn identity(call: crate::sansio::PendingToolCall) -> Self {
678        Self {
679            call_id: call.call_id,
680            tool_name: call.tool_name,
681            args: call.args,
682            replay: call.replay,
683            prepared_payload: serde_json::Value::Null,
684        }
685    }
686
687    pub fn from_parts(
688        call_id: impl Into<String>,
689        tool_name: impl Into<String>,
690        args: serde_json::Value,
691        replay: Option<ProviderReplayMeta>,
692        prepared_payload: serde_json::Value,
693    ) -> Self {
694        Self {
695            call_id: call_id.into(),
696            tool_name: tool_name.into(),
697            args,
698            replay,
699            prepared_payload,
700        }
701    }
702}
703
704#[derive(Clone)]
705pub struct ToolPrepareContext {
706    session_id: String,
707    sessions: Arc<dyn SessionStateService>,
708    turn_context: crate::TurnContext,
709    tool_call_id: Option<String>,
710}
711
712impl ToolPrepareContext {
713    pub(crate) fn new(
714        session_id: String,
715        sessions: Arc<dyn SessionStateService>,
716        turn_context: crate::TurnContext,
717        tool_call_id: Option<String>,
718    ) -> Self {
719        Self {
720            session_id,
721            sessions,
722            turn_context,
723            tool_call_id,
724        }
725    }
726
727    pub fn session_id(&self) -> &str {
728        &self.session_id
729    }
730
731    pub fn tool_call_id(&self) -> Option<&str> {
732        self.tool_call_id.as_deref()
733    }
734
735    pub fn turn_context(&self) -> &crate::TurnContext {
736        &self.turn_context
737    }
738
739    pub fn plugin_input<T>(&self, plugin_id: &'static str) -> Option<&T>
740    where
741        T: 'static,
742    {
743        self.turn_context.plugin_input::<T>(plugin_id)
744    }
745
746    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
747        self.sessions.snapshot_session(&self.session_id).await
748    }
749
750    pub async fn tool_catalog(&self) -> Result<Vec<serde_json::Value>, PluginError> {
751        self.sessions.tool_catalog(&self.session_id).await
752    }
753
754    pub async fn shared_tool_catalog(
755        &self,
756    ) -> Result<std::sync::Arc<Vec<serde_json::Value>>, PluginError> {
757        self.sessions.shared_tool_catalog(&self.session_id).await
758    }
759}
760
761/// Inputs handed to [`ToolProvider::prepare_tool_call`].
762pub struct ToolPrepareCall<'a> {
763    pub pending: crate::sansio::PendingToolCall,
764    pub context: &'a ToolPrepareContext,
765}
766
767/// Per-call inputs handed to [`ToolProvider::execute`].
768///
769/// Fields are `pub` because `ToolCall` is a transient borrow; consumers
770/// typically destructure (`let ToolCall { name, args, .. } = call`). The
771/// stable surface lives on [`ToolContext`] (sealed) and the runtime's
772/// dispatcher, which constructs `ToolCall` values.
773pub struct ToolCall<'a> {
774    pub name: &'a str,
775    pub args: &'a serde_json::Value,
776    pub context: &'a ToolContext<'a>,
777    pub progress: Option<&'a ProgressSender>,
778}
779
780/// Trait for providing tools to the sandbox. Implement this per-project.
781///
782/// Implementations supply cheap [`ToolManifest`]s, lazily resolved
783/// [`ToolContract`]s, and a single
784/// [`execute`](Self::execute) method that handles every call. Tools that
785/// need session state read it from `call.context`; tools that stream
786/// progress send through `call.progress`.
787#[async_trait::async_trait]
788pub trait ToolProvider: Send + Sync + 'static {
789    fn tool_manifests(&self) -> Vec<ToolManifest>;
790    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
791        self.tool_manifests()
792            .into_iter()
793            .find(|manifest| manifest.name == name)
794    }
795    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
796    async fn prepare_tool_call(
797        &self,
798        call: ToolPrepareCall<'_>,
799    ) -> Result<PreparedToolCall, ToolResult> {
800        Ok(PreparedToolCall::identity(call.pending))
801    }
802    async fn execute(&self, call: ToolCall<'_>) -> ToolResult;
803}
804
805#[cfg(test)]
806mod tests {
807    use super::*;
808
809    #[test]
810    fn tool_context_builder_carries_call_payload_and_cancellation_state() {
811        let cancellation = tokio_util::sync::CancellationToken::new();
812        let prepared = PreparedToolCall::from_parts(
813            "call-1",
814            "demo_tool",
815            serde_json::json!({ "input": true }),
816            None,
817            serde_json::json!({ "prepared": true }),
818        );
819
820        let context = ToolContext::builder(
821            "session-1".to_string(),
822            Arc::new(crate::testing::MockSessionManager::default()),
823            Arc::new(crate::testing::MockSessionManager::default()),
824            Arc::new(crate::testing::MockSessionManager::default()),
825            Arc::new(crate::UnavailableProcessService),
826            Arc::new(crate::DefaultProcessCancelAbility),
827            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
828                crate::InlineRuntimeEffectController,
829            )),
830            Arc::new(crate::InMemoryAttachmentStore::new()),
831            crate::DirectCompletionClient::unavailable(
832                "direct completions are unavailable in this test context",
833            ),
834        )
835        .prepared_call(&prepared)
836        .cancellation_token(Some(cancellation.clone()))
837        .async_process("process-1", cancellation.clone())
838        .build();
839
840        assert_eq!(context.session_id(), "session-1");
841        assert_eq!(context.tool_call_id(), Some("call-1"));
842        assert_eq!(
843            context.prepared_payload(),
844            &serde_json::json!({ "prepared": true })
845        );
846        assert_eq!(context.async_process_id(), Some("process-1"));
847        assert!(context.cancellation_token().is_some());
848    }
849}
lash_core/tool_provider.rs

lash_core/
tool_provider.rs
