lash_core/
tool_provider.rs

1use std::collections::HashSet;
2use std::future::Future;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex};
5
6use lash_sansio::llm::types::ProviderReplayMeta;
7use serde::{Deserialize, Serialize};
8
9use crate::plugin::{
10    PluginError, SessionGraphService, SessionLifecycleService, SessionSnapshot, SessionStateService,
11};
12use crate::{ToolContract, ToolDefinition, ToolId, ToolManifest, ToolResult};
13
14mod attachments;
15mod direct_completion;
16mod dispatch;
17mod process;
18pub(crate) mod process_events;
19mod session;
20mod triggers;
21
22pub use attachments::ToolAttachmentClient;
23pub use direct_completion::ToolDirectCompletionClient;
24pub use dispatch::ToolDispatchClient;
25pub use process::ToolSessionProcessAdmin;
26pub use process_events::ToolProcessEventClient;
27pub use session::{ToolSessionAdmin, ToolSessionModel};
28pub use triggers::ToolTriggerClient;
29
30/// A message sent from the sandbox to the host during execution.
31#[derive(Clone, Debug)]
32pub struct SandboxMessage {
33    pub text: String,
34    /// "tool_output" or another host-rendered progress event kind.
35    pub kind: String,
36}
37
38/// Sender for streaming progress messages from tools (e.g. live bash output).
39pub type ProgressSender = tokio::sync::mpsc::UnboundedSender<SandboxMessage>;
40
41#[derive(Clone, Default)]
42pub(crate) struct ToolCompletionState {
43    key: Arc<Mutex<Option<crate::AwaitEventKey>>>,
44}
45
46impl ToolCompletionState {
47    fn store(
48        &self,
49        key: crate::AwaitEventKey,
50    ) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
51        let mut guard = self.key.lock().map_err(|_| {
52            crate::RuntimeError::new(
53                "tool_completion_state_poisoned",
54                "tool completion key state lock poisoned",
55            )
56        })?;
57        if let Some(existing) = guard.as_ref() {
58            return Ok(existing.clone());
59        }
60        *guard = Some(key.clone());
61        Ok(key)
62    }
63
64    pub(crate) fn take(&self) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
65        self.key.lock().map(|mut guard| guard.take()).map_err(|_| {
66            crate::RuntimeError::new(
67                "tool_completion_state_poisoned",
68                "tool completion key state lock poisoned",
69            )
70        })
71    }
72}
73
74#[derive(Clone, Default)]
75pub(crate) struct ToolDurableEffectState {
76    step_ids: Arc<Mutex<HashSet<String>>>,
77    process_event_sequence: Arc<AtomicU64>,
78}
79
80impl ToolDurableEffectState {
81    fn reserve_step(&self, step_id: &str) -> Result<(), crate::RuntimeError> {
82        let mut guard = self.step_ids.lock().map_err(|_| {
83            crate::RuntimeError::new(
84                "durable_effect_state_poisoned",
85                "durable effect step state lock poisoned",
86            )
87        })?;
88        if !guard.insert(step_id.to_string()) {
89            return Err(crate::RuntimeError::new(
90                "durable_effect_duplicate_step_id",
91                format!("durable effect step id `{step_id}` was already used by this tool call"),
92            ));
93        }
94        Ok(())
95    }
96
97    fn next_process_event_sequence(&self) -> u64 {
98        self.process_event_sequence.fetch_add(1, Ordering::Relaxed)
99    }
100}
101
102/// Per-call environment for [`ToolProvider::execute`]. Fields are sealed so
103/// the runtime can add capabilities without breaking tool authors.
104#[derive(Clone)]
105pub struct ToolContext<'run> {
106    pub(crate) session_id: String,
107    pub(crate) agent_frame_id: crate::AgentFrameId,
108    pub(crate) sessions: Arc<dyn SessionStateService>,
109    pub(crate) session_lifecycle: Arc<dyn SessionLifecycleService>,
110    pub(crate) processes: Arc<dyn crate::ProcessService>,
111    pub(crate) process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
112    pub(crate) effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
113    pub(crate) runtime_dispatch: Option<Arc<crate::tool_dispatch::ToolDispatchContext<'run>>>,
114    pub(crate) runtime_execution_context: Option<crate::RuntimeExecutionContext<'run>>,
115    pub(crate) cancellation_token: Option<tokio_util::sync::CancellationToken>,
116    pub(crate) async_process_id: Option<String>,
117    pub(crate) runtime_process_id: Option<String>,
118    pub(crate) process_events: Option<ToolProcessEventContext>,
119    pub(crate) attachment_store: Arc<crate::SessionAttachmentStore>,
120    pub(crate) direct_completions: crate::DirectCompletionClient<'run>,
121    pub(crate) prepared_payload: serde_json::Value,
122    pub(crate) tool_execution_binding: serde_json::Value,
123    /// The id of the in-flight tool call that is invoking this tool.
124    pub(crate) tool_call_id: Option<String>,
125    pub(crate) attempt_number: u32,
126    pub(crate) max_attempts: u32,
127    pub(crate) replay_key: Option<String>,
128    pub(crate) completion: ToolCompletionState,
129    pub(crate) durable_effects: ToolDurableEffectState,
130    pub(crate) parent_invocation: Option<crate::RuntimeInvocation>,
131    pub(crate) execution_env_spec: crate::ProcessExecutionEnvSpec,
132    pub(crate) child_execution_trace_hook: Option<ToolChildExecutionTraceHook>,
133}
134
135#[derive(Clone)]
136pub struct ToolChildProcessStarted {
137    pub process_id: String,
138    pub child_entry_name: Option<String>,
139}
140
141#[derive(Clone)]
142pub struct ToolChildExecutionTraceHook {
143    on_child_process_started: Arc<dyn Fn(ToolChildProcessStarted) + Send + Sync>,
144}
145
146impl ToolChildExecutionTraceHook {
147    pub fn new(
148        on_child_process_started: impl Fn(ToolChildProcessStarted) + Send + Sync + 'static,
149    ) -> Self {
150        Self {
151            on_child_process_started: Arc::new(on_child_process_started),
152        }
153    }
154
155    pub fn child_process_started(&self, event: ToolChildProcessStarted) {
156        (self.on_child_process_started)(event);
157    }
158}
159
160#[derive(Clone)]
161pub(crate) struct ToolProcessEventContext {
162    process_id: String,
163    registry: Arc<dyn crate::ProcessRegistry>,
164    awaiter: crate::ProcessAwaiter,
165    store: Option<Arc<dyn crate::RuntimePersistence>>,
166    session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
167    session_graph: Arc<dyn SessionGraphService>,
168    queued_work_driver: Option<crate::QueuedWorkDriver>,
169}
170
171pub(crate) struct ToolContextBuilder<'run> {
172    session_id: String,
173    agent_frame_id: crate::AgentFrameId,
174    sessions: Arc<dyn SessionStateService>,
175    session_lifecycle: Arc<dyn SessionLifecycleService>,
176    session_graph: Arc<dyn SessionGraphService>,
177    processes: Arc<dyn crate::ProcessService>,
178    process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
179    effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
180    runtime_dispatch: Option<Arc<crate::tool_dispatch::ToolDispatchContext<'run>>>,
181    runtime_execution_context: Option<crate::RuntimeExecutionContext<'run>>,
182    cancellation_token: Option<tokio_util::sync::CancellationToken>,
183    async_process_id: Option<String>,
184    runtime_process_id: Option<String>,
185    process_events: Option<ToolProcessEventContext>,
186    attachment_store: Arc<crate::SessionAttachmentStore>,
187    direct_completions: crate::DirectCompletionClient<'run>,
188    prepared_payload: serde_json::Value,
189    tool_execution_binding: serde_json::Value,
190    tool_call_id: Option<String>,
191    completion: ToolCompletionState,
192    durable_effects: ToolDurableEffectState,
193    parent_invocation: Option<crate::RuntimeInvocation>,
194    execution_env_spec: crate::ProcessExecutionEnvSpec,
195    child_execution_trace_hook: Option<ToolChildExecutionTraceHook>,
196}
197
198impl<'run> ToolContextBuilder<'run> {
199    pub(crate) fn from_dispatch(
200        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
201    ) -> Self {
202        Self {
203            session_id: dispatch.session_id.clone(),
204            agent_frame_id: dispatch.agent_frame_id.clone(),
205            sessions: Arc::clone(&dispatch.sessions),
206            session_lifecycle: Arc::clone(&dispatch.session_lifecycle),
207            session_graph: Arc::clone(&dispatch.session_graph),
208            processes: Arc::clone(&dispatch.processes),
209            process_cancel_ability: Arc::clone(&dispatch.process_cancel_ability),
210            effect_controller: dispatch.effect_controller.clone(),
211            runtime_dispatch: Some(Arc::clone(&dispatch)),
212            runtime_execution_context: None,
213            cancellation_token: None,
214            async_process_id: None,
215            runtime_process_id: None,
216            process_events: None,
217            attachment_store: Arc::clone(&dispatch.attachment_store),
218            direct_completions: dispatch.direct_completions.clone(),
219            prepared_payload: serde_json::Value::Null,
220            tool_execution_binding: serde_json::Value::Null,
221            tool_call_id: None,
222            completion: ToolCompletionState::default(),
223            durable_effects: ToolDurableEffectState::default(),
224            parent_invocation: dispatch.parent_invocation.clone(),
225            execution_env_spec: dispatch.execution_env_spec.clone(),
226            child_execution_trace_hook: None,
227        }
228    }
229
230    #[cfg(any(test, feature = "testing"))]
231    pub(crate) fn tool_call_id(mut self, tool_call_id: impl Into<Option<String>>) -> Self {
232        self.tool_call_id = tool_call_id.into();
233        self
234    }
235
236    pub(crate) fn prepared_call(mut self, call: &PreparedToolCall) -> Self {
237        self.tool_call_id = Some(call.call_id.clone());
238        self.prepared_payload = call.prepared_payload.clone();
239        self
240    }
241
242    pub(crate) fn tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
243        self.tool_execution_binding = binding;
244        self
245    }
246
247    pub(crate) fn cancellation_token(
248        mut self,
249        cancellation_token: Option<tokio_util::sync::CancellationToken>,
250    ) -> Self {
251        self.cancellation_token = cancellation_token;
252        self
253    }
254
255    pub(crate) fn runtime_execution_context(
256        mut self,
257        context: crate::RuntimeExecutionContext<'run>,
258    ) -> Self {
259        self.runtime_execution_context = Some(context);
260        self
261    }
262
263    pub(crate) fn runtime_process_id(mut self, process_id: Option<String>) -> Self {
264        self.runtime_process_id = process_id;
265        self
266    }
267
268    pub(crate) fn async_process(
269        mut self,
270        process_id: impl Into<String>,
271        cancellation_token: tokio_util::sync::CancellationToken,
272    ) -> Self {
273        self.async_process_id = Some(process_id.into());
274        self.cancellation_token = Some(cancellation_token);
275        self
276    }
277
278    pub(crate) fn process_events(
279        mut self,
280        process_id: impl Into<String>,
281        registry: Arc<dyn crate::ProcessRegistry>,
282        awaiter: crate::ProcessAwaiter,
283        store: Option<Arc<dyn crate::RuntimePersistence>>,
284        session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
285        queued_work_driver: Option<crate::QueuedWorkDriver>,
286    ) -> Self {
287        self.process_events = Some(ToolProcessEventContext {
288            process_id: process_id.into(),
289            registry,
290            awaiter,
291            store,
292            session_store_factory,
293            session_graph: Arc::clone(&self.session_graph),
294            queued_work_driver,
295        });
296        self
297    }
298
299    pub(crate) fn parent_invocation(mut self, metadata: Option<crate::RuntimeInvocation>) -> Self {
300        self.parent_invocation = metadata;
301        self
302    }
303
304    pub(crate) fn child_execution_trace_hook(
305        mut self,
306        hook: Option<ToolChildExecutionTraceHook>,
307    ) -> Self {
308        self.child_execution_trace_hook = hook;
309        self
310    }
311
312    pub(crate) fn build(self) -> ToolContext<'run> {
313        ToolContext {
314            session_id: self.session_id,
315            agent_frame_id: self.agent_frame_id,
316            sessions: self.sessions,
317            session_lifecycle: self.session_lifecycle,
318            processes: self.processes,
319            process_cancel_ability: self.process_cancel_ability,
320            effect_controller: self.effect_controller,
321            runtime_dispatch: self.runtime_dispatch,
322            runtime_execution_context: self.runtime_execution_context,
323            cancellation_token: self.cancellation_token,
324            async_process_id: self.async_process_id,
325            runtime_process_id: self.runtime_process_id,
326            process_events: self.process_events,
327            attachment_store: self.attachment_store,
328            direct_completions: self.direct_completions,
329            prepared_payload: self.prepared_payload,
330            tool_execution_binding: self.tool_execution_binding,
331            tool_call_id: self.tool_call_id,
332            attempt_number: 1,
333            max_attempts: 1,
334            replay_key: None,
335            completion: self.completion,
336            durable_effects: self.durable_effects,
337            parent_invocation: self.parent_invocation,
338            execution_env_spec: self.execution_env_spec,
339            child_execution_trace_hook: self.child_execution_trace_hook,
340        }
341    }
342}
343
344impl<'run> ToolContext<'run> {
345    #[cfg(any(test, feature = "testing"))]
346    #[expect(
347        clippy::too_many_arguments,
348        reason = "testing constructor mirrors the sealed runtime tool context dependencies"
349    )]
350    pub(crate) fn builder(
351        session_id: String,
352        sessions: Arc<dyn SessionStateService>,
353        session_lifecycle: Arc<dyn SessionLifecycleService>,
354        session_graph: Arc<dyn SessionGraphService>,
355        processes: Arc<dyn crate::ProcessService>,
356        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
357        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
358        attachment_store: Arc<crate::SessionAttachmentStore>,
359        direct_completions: crate::DirectCompletionClient<'run>,
360    ) -> ToolContextBuilder<'run> {
361        ToolContextBuilder {
362            session_id,
363            agent_frame_id: String::new(),
364            sessions,
365            session_lifecycle,
366            session_graph,
367            processes,
368            process_cancel_ability,
369            effect_controller,
370            runtime_dispatch: None,
371            runtime_execution_context: None,
372            cancellation_token: None,
373            async_process_id: None,
374            runtime_process_id: None,
375            process_events: None,
376            attachment_store,
377            direct_completions,
378            prepared_payload: serde_json::Value::Null,
379            tool_execution_binding: serde_json::Value::Null,
380            tool_call_id: None,
381            completion: ToolCompletionState::default(),
382            durable_effects: ToolDurableEffectState::default(),
383            parent_invocation: None,
384            execution_env_spec: crate::ProcessExecutionEnvSpec::new(
385                crate::PluginOptions::default(),
386                crate::SessionPolicy::default(),
387            ),
388            child_execution_trace_hook: None,
389        }
390    }
391
392    pub(crate) fn from_dispatch(
393        dispatch: Arc<crate::tool_dispatch::ToolDispatchContext<'run>>,
394    ) -> ToolContextBuilder<'run> {
395        ToolContextBuilder::from_dispatch(dispatch)
396    }
397
398    pub fn session_id(&self) -> &str {
399        &self.session_id
400    }
401
402    pub fn agent_frame_id(&self) -> &str {
403        &self.agent_frame_id
404    }
405
406    pub fn sessions(&self) -> ToolSessionAdmin<'run> {
407        ToolSessionAdmin {
408            session_id: self.session_id.clone(),
409            sessions: Arc::clone(&self.sessions),
410            session_lifecycle: Arc::clone(&self.session_lifecycle),
411            effect_controller: self.effect_controller.clone(),
412        }
413    }
414
415    pub fn dispatch(&self) -> ToolDispatchClient<'run> {
416        ToolDispatchClient {
417            context: self.clone(),
418        }
419    }
420
421    pub fn triggers(&self) -> ToolTriggerClient<'run> {
422        ToolTriggerClient {
423            context: self.clone(),
424        }
425    }
426
427    pub fn processes(&self) -> ToolSessionProcessAdmin<'run> {
428        ToolSessionProcessAdmin {
429            session_id: self.session_id.clone(),
430            agent_frame_id: self.agent_frame_id.clone(),
431            processes: Arc::clone(&self.processes),
432            process_cancel_ability: Arc::clone(&self.process_cancel_ability),
433            effect_controller: self.effect_controller.clone(),
434            parent_invocation: self.parent_invocation.clone(),
435            tool_call_id: self.tool_call_id.clone(),
436            execution_env_spec: self.execution_env_spec.clone(),
437        }
438    }
439
440    pub fn emit_child_process_started(
441        &self,
442        process_id: impl Into<String>,
443        child_entry_name: Option<String>,
444    ) {
445        let Some(hook) = &self.child_execution_trace_hook else {
446            return;
447        };
448        hook.child_process_started(ToolChildProcessStarted {
449            process_id: process_id.into(),
450            child_entry_name,
451        });
452    }
453
454    pub fn direct_completions(&self) -> ToolDirectCompletionClient<'run> {
455        ToolDirectCompletionClient {
456            session_id: self.session_id.clone(),
457            tool_call_id: self.tool_call_id.clone(),
458            direct_completions: self.direct_completions.clone(),
459            parent_invocation: self.parent_invocation.clone(),
460            parent_tool_attempt_is_durable: self.parent_invocation.as_ref().is_some_and(
461                |invocation| {
462                    invocation.effect_kind() == Some(crate::RuntimeEffectKind::ToolAttempt)
463                },
464            ) && self
465                .effect_controller
466                .controller()
467                .durability_tier()
468                == crate::DurabilityTier::Durable,
469        }
470    }
471
472    pub fn attachments(&self) -> ToolAttachmentClient {
473        ToolAttachmentClient {
474            store: Arc::clone(&self.attachment_store),
475        }
476    }
477
478    pub fn process_events(&self) -> ToolProcessEventClient {
479        ToolProcessEventClient {
480            context: self.process_events.clone(),
481        }
482    }
483
484    /// Borrow this tool call's durable effect boundary.
485    ///
486    /// This is available only while executing a prepared tool call under a
487    /// controller that explicitly supports durable tool effects. The returned
488    /// facade records JSON steps and await-event waits in the caller's existing
489    /// effect log; it does not expose the underlying workflow engine context.
490    pub fn durable_effects(&self) -> Result<ToolDurableEffects<'_, 'run>, crate::RuntimeError> {
491        let Some(tool_call_id) = self.tool_call_id.as_deref() else {
492            return Err(crate::RuntimeError::new(
493                "durable_effects_missing_call_id",
494                "durable effects require a prepared tool call id",
495            ));
496        };
497        if tool_call_id.trim().is_empty() {
498            return Err(crate::RuntimeError::new(
499                "durable_effects_missing_call_id",
500                "durable effects require a non-empty prepared tool call id",
501            ));
502        }
503        let scoped = self.effect_controller.scoped();
504        if !scoped.controller().supports_durable_effects() {
505            return Err(crate::RuntimeError::new(
506                "durable_effects_unavailable",
507                "this effect controller does not support durable tool effects",
508            ));
509        }
510        if self.parent_invocation.as_ref().is_some_and(|invocation| {
511            invocation.effect_kind() == Some(crate::RuntimeEffectKind::ToolAttempt)
512        }) && scoped.controller().durability_tier() == crate::DurabilityTier::Durable
513        {
514            return Err(crate::RuntimeError::new(
515                "durable_effects_unavailable_in_tool_attempt",
516                "durable tool sub-effects are not available inside a journaled tool attempt",
517            ));
518        }
519        Ok(ToolDurableEffects { context: self })
520    }
521
522    pub fn cancellation_token(&self) -> Option<&tokio_util::sync::CancellationToken> {
523        self.cancellation_token.as_ref()
524    }
525
526    #[doc(hidden)]
527    pub fn named_phase(&self, phase: &'static str) -> crate::runtime::RuntimeNamedPhase {
528        match self.runtime_execution_context.as_ref() {
529            Some(context) => context.named_phase(phase),
530            None => crate::runtime::RuntimeNamedPhase::begin(None, phase),
531        }
532    }
533
534    pub fn async_process_id(&self) -> Option<&str> {
535        self.async_process_id.as_deref()
536    }
537
538    pub fn runtime_process_id(&self) -> Option<&str> {
539        self.async_process_id
540            .as_deref()
541            .or(self.runtime_process_id.as_deref())
542            .or_else(|| {
543                self.process_events
544                    .as_ref()
545                    .map(|context| context.process_id.as_str())
546            })
547    }
548
549    pub fn tool_call_id(&self) -> Option<&str> {
550        self.tool_call_id.as_deref()
551    }
552
553    pub fn prepared_payload(&self) -> &serde_json::Value {
554        &self.prepared_payload
555    }
556
557    pub fn tool_execution_binding(&self) -> &serde_json::Value {
558        &self.tool_execution_binding
559    }
560
561    pub fn decode_prepared_payload<T>(&self) -> Result<T, serde_json::Error>
562    where
563        T: serde::de::DeserializeOwned,
564    {
565        serde_json::from_value(self.prepared_payload.clone())
566    }
567
568    pub fn attempt_number(&self) -> u32 {
569        self.attempt_number
570    }
571
572    pub fn max_attempts(&self) -> u32 {
573        self.max_attempts
574    }
575
576    pub fn replay_key(&self) -> Option<&str> {
577        self.replay_key.as_deref()
578    }
579
580    /// Obtain the durable completion key for this call, required before returning
581    /// [`ToolResult::Pending`](crate::ToolResult::Pending).
582    ///
583    /// A tool that defers its outcome (waiting on a webhook, human approval, or another
584    /// service) calls this, hands the returned [`AwaitEventKey`](crate::AwaitEventKey)
585    /// to whatever will complete the work out-of-band, and then returns
586    /// `ToolResult::Pending(..)`. The key names the durable wait the runtime parks the
587    /// call on; the external resolver delivers the result against it later.
588    ///
589    /// The key is stored on the context and consumed by the dispatcher when the tool
590    /// returns `Pending`. Returning `Pending` without first calling this fails the call
591    /// with `pending_tool_missing_completion_key`. Calls made outside a prepared tool
592    /// invocation (no tool call id) fail with `tool_completion_key_missing_call_id`.
593    pub async fn completion_key(&self) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
594        let tool_call_id = self.tool_call_id.clone().ok_or_else(|| {
595            crate::RuntimeError::new(
596                "tool_completion_key_missing_call_id",
597                "completion keys require a prepared tool call id",
598            )
599        })?;
600        let scoped = self.effect_controller.scoped();
601        let key = scoped
602            .controller()
603            .await_event_key(
604                scoped.execution_scope(),
605                crate::AwaitEventWaitIdentity::tool_completion(tool_call_id),
606            )
607            .await?;
608        self.completion.store(key)
609    }
610
611    pub(crate) fn take_completion_key(
612        &self,
613    ) -> Result<Option<crate::AwaitEventKey>, crate::RuntimeError> {
614        self.completion.take()
615    }
616
617    pub fn with_async_process(
618        mut self,
619        process_id: impl Into<String>,
620        cancellation_token: tokio_util::sync::CancellationToken,
621    ) -> Self {
622        self.async_process_id = Some(process_id.into());
623        self.runtime_process_id = self.async_process_id.clone();
624        self.cancellation_token = Some(cancellation_token);
625        self
626    }
627
628    #[cfg(any(test, feature = "testing"))]
629    #[doc(hidden)]
630    pub fn with_process_events_for_testing(
631        mut self,
632        process_id: impl Into<String>,
633        registry: Arc<dyn crate::ProcessRegistry>,
634    ) -> Self {
635        let awaiter = crate::ProcessAwaiter::polling(Arc::clone(&registry));
636        self.process_events = Some(ToolProcessEventContext {
637            process_id: process_id.into(),
638            registry,
639            awaiter,
640            store: None,
641            session_store_factory: None,
642            session_graph: Arc::new(crate::plugin::NoopSessionManager),
643            queued_work_driver: None,
644        });
645        self
646    }
647
648    pub(crate) fn with_retry_context(
649        mut self,
650        tool_name: &str,
651        attempt_number: u32,
652        max_attempts: u32,
653    ) -> Self {
654        self.attempt_number = attempt_number.max(1);
655        self.max_attempts = max_attempts.max(1);
656        self.replay_key = self
657            .tool_call_id
658            .as_ref()
659            .map(|call_id| format!("lash-tool:{}:{call_id}:{tool_name}", self.session_id));
660        self
661    }
662
663    pub(crate) fn with_prepared_payload(mut self, payload: serde_json::Value) -> Self {
664        self.prepared_payload = payload;
665        self
666    }
667
668    pub(crate) fn with_tool_execution_binding(mut self, binding: serde_json::Value) -> Self {
669        self.tool_execution_binding = binding;
670        self
671    }
672
673    /// Constructor reserved for `lash_core::testing` helpers. Do not call directly;
674    /// use [`lash_core::testing::mock_tool_context`] instead.
675    #[cfg(any(test, feature = "testing"))]
676    #[doc(hidden)]
677    #[expect(
678        clippy::too_many_arguments,
679        reason = "test-only constructor mirrors the sealed runtime tool context"
680    )]
681    pub fn __for_testing(
682        session_id: String,
683        sessions: Arc<dyn SessionStateService>,
684        session_lifecycle: Arc<dyn SessionLifecycleService>,
685        session_graph: Arc<dyn SessionGraphService>,
686        processes: Arc<dyn crate::ProcessService>,
687        attachment_store: Arc<crate::SessionAttachmentStore>,
688        direct_completions: crate::DirectCompletionClient<'static>,
689        tool_call_id: Option<String>,
690    ) -> ToolContext<'static> {
691        ToolContext::builder(
692            session_id,
693            sessions,
694            session_lifecycle,
695            session_graph,
696            processes,
697            Arc::new(crate::DefaultProcessCancelAbility),
698            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
699                crate::InlineRuntimeEffectController,
700            )),
701            attachment_store,
702            direct_completions,
703        )
704        .tool_call_id(tool_call_id)
705        .build()
706    }
707
708    /// Constructor reserved for tests that need a custom process-cancel host
709    /// ability. Do not call directly; prefer public testing helpers when they
710    /// cover the case.
711    #[cfg(any(test, feature = "testing"))]
712    #[doc(hidden)]
713    #[expect(
714        clippy::too_many_arguments,
715        reason = "test-only constructor mirrors the sealed runtime context"
716    )]
717    pub fn __for_testing_with_process_cancel_ability(
718        session_id: String,
719        sessions: Arc<dyn SessionStateService>,
720        session_lifecycle: Arc<dyn SessionLifecycleService>,
721        session_graph: Arc<dyn SessionGraphService>,
722        processes: Arc<dyn crate::ProcessService>,
723        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
724        attachment_store: Arc<crate::SessionAttachmentStore>,
725        direct_completions: crate::DirectCompletionClient<'static>,
726        tool_call_id: Option<String>,
727    ) -> ToolContext<'static> {
728        ToolContext::builder(
729            session_id,
730            sessions,
731            session_lifecycle,
732            session_graph,
733            processes,
734            process_cancel_ability,
735            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
736                crate::InlineRuntimeEffectController,
737            )),
738            attachment_store,
739            direct_completions,
740        )
741        .tool_call_id(tool_call_id)
742        .build()
743    }
744}
745
746/// Durable effect operations available to advanced in-process tools.
747///
748/// The facade borrows the caller's existing runtime effect boundary. It records
749/// JSON-only local steps and awaits host-signed event keys without exposing
750/// Restate, Temporal, or any other workflow-native context.
751pub struct ToolDurableEffects<'ctx, 'run> {
752    context: &'ctx ToolContext<'run>,
753}
754
755impl<'ctx, 'run> ToolDurableEffects<'ctx, 'run> {
756    pub async fn run_json<F, Fut>(
757        &self,
758        step_id: impl Into<String>,
759        input: serde_json::Value,
760        run: F,
761    ) -> Result<serde_json::Value, crate::RuntimeError>
762    where
763        F: FnOnce(serde_json::Value) -> Fut + Send + 'run,
764        Fut: Future<Output = Result<serde_json::Value, crate::RuntimeError>> + Send + 'run,
765    {
766        let step_id = step_id.into();
767        if step_id.trim().is_empty() {
768            return Err(crate::RuntimeError::new(
769                "durable_effect_empty_step_id",
770                "durable effect step id must be non-empty",
771            ));
772        }
773        self.context.durable_effects.reserve_step(&step_id)?;
774        let invocation = self.step_invocation(
775            format!("durable-step:{step_id}"),
776            crate::RuntimeEffectKind::DurableStep,
777            format!("durable-step:{step_id}"),
778        )?;
779        let outcome = self
780            .context
781            .effect_controller
782            .controller()
783            .execute_effect(
784                crate::RuntimeEffectEnvelope::new(
785                    invocation,
786                    crate::RuntimeEffectCommand::DurableStep {
787                        step_id,
788                        input: input.clone(),
789                    },
790                ),
791                crate::RuntimeEffectLocalExecutor::durable_step(run),
792            )
793            .await
794            .map_err(crate::RuntimeEffectControllerError::into_runtime_error)?;
795        outcome
796            .into_durable_step()
797            .map_err(crate::RuntimeEffectControllerError::into_runtime_error)
798    }
799
800    pub async fn external_event_key(
801        &self,
802        key: impl Into<String>,
803    ) -> Result<crate::AwaitEventKey, crate::RuntimeError> {
804        let key = key.into();
805        if key.trim().is_empty() {
806            return Err(crate::RuntimeError::new(
807                "durable_effect_empty_event_key",
808                "durable effect external event key must be non-empty",
809            ));
810        }
811        let scoped = self.context.effect_controller.scoped();
812        scoped
813            .controller()
814            .await_event_key(
815                scoped.execution_scope(),
816                crate::AwaitEventWaitIdentity::Custom { key },
817            )
818            .await
819    }
820
821    pub async fn await_event_json(
822        &self,
823        key: crate::AwaitEventKey,
824    ) -> Result<serde_json::Value, crate::RuntimeError> {
825        let invocation = self.step_invocation(
826            format!("await-event:{}", key.key_id),
827            crate::RuntimeEffectKind::AwaitEvent,
828            format!("await-event:{}", key.key_id),
829        )?;
830        let cancellation = self.context.cancellation_token.clone().unwrap_or_default();
831        let clock = self
832            .context
833            .runtime_dispatch
834            .as_ref()
835            .map(|dispatch| Arc::clone(&dispatch.clock))
836            .unwrap_or_else(|| Arc::new(crate::SystemClock));
837        let outcome = self
838            .context
839            .effect_controller
840            .controller()
841            .execute_effect(
842                crate::RuntimeEffectEnvelope::new(
843                    invocation,
844                    crate::RuntimeEffectCommand::AwaitEvent { key },
845                ),
846                crate::RuntimeEffectLocalExecutor::await_event_with_clock(
847                    cancellation,
848                    None,
849                    clock,
850                ),
851            )
852            .await
853            .map_err(crate::RuntimeEffectControllerError::into_runtime_error)?;
854        match outcome
855            .into_await_event()
856            .map_err(crate::RuntimeEffectControllerError::into_runtime_error)?
857        {
858            crate::Resolution::Ok(value) => Ok(value),
859            crate::Resolution::Err(err) => Err(crate::RuntimeError::new(err.code, err.message)),
860            crate::Resolution::Timeout => Err(crate::RuntimeError::new(
861                "durable_effect_event_timeout",
862                "durable effect external event wait timed out",
863            )),
864            crate::Resolution::Cancelled => Err(crate::RuntimeError::new(
865                "durable_effect_event_cancelled",
866                "durable effect external event wait was cancelled",
867            )),
868        }
869    }
870
871    pub async fn emit_process_event(
872        &self,
873        event_type: impl Into<String>,
874        payload: serde_json::Value,
875    ) -> Result<crate::ProcessEvent, crate::RuntimeError> {
876        let Some(process) = self.context.process_events.as_ref() else {
877            return Err(crate::RuntimeError::new(
878                "durable_effect_process_event_unavailable",
879                "durable effect process events are unavailable outside a durable process",
880            ));
881        };
882        let event_type = event_type.into();
883        if event_type.trim().is_empty() {
884            return Err(crate::RuntimeError::new(
885                "durable_effect_empty_process_event_type",
886                "durable effect process event type must be non-empty",
887            ));
888        }
889        let tool_call_id = self.context.tool_call_id.as_deref().ok_or_else(|| {
890            crate::RuntimeError::new(
891                "durable_effects_missing_call_id",
892                "durable effects require a prepared tool call id",
893            )
894        })?;
895        let sequence = self.context.durable_effects.next_process_event_sequence();
896        let request = crate::ProcessEventAppendRequest::new(event_type, payload).with_replay_key(
897            format!("tool:{tool_call_id}:durable-process-event:{sequence}"),
898        );
899        self.context
900            .process_events()
901            .emit_request(request)
902            .await
903            .map_err(|err| {
904                crate::RuntimeError::new(
905                    "durable_effect_process_event_append_failed",
906                    err.to_string(),
907                )
908            })
909            .and_then(|event| {
910                if event.process_id == process.process_id {
911                    Ok(event)
912                } else {
913                    Err(crate::RuntimeError::new(
914                        "durable_effect_process_event_process_mismatch",
915                        "process event append returned an event for a different process",
916                    ))
917                }
918            })
919    }
920
921    fn step_invocation(
922        &self,
923        effect_id_suffix: impl Into<String>,
924        kind: crate::RuntimeEffectKind,
925        replay_suffix: impl AsRef<str>,
926    ) -> Result<crate::RuntimeInvocation, crate::RuntimeError> {
927        let tool_call_id = self.context.tool_call_id.as_deref().ok_or_else(|| {
928            crate::RuntimeError::new(
929                "durable_effects_missing_call_id",
930                "durable effects require a prepared tool call id",
931            )
932        })?;
933        let effect_id_suffix = effect_id_suffix.into();
934        if let Some(parent) = self.context.parent_invocation.as_ref() {
935            return Ok(crate::runtime::causal::child_effect_invocation(
936                parent,
937                format!("{tool_call_id}:{effect_id_suffix}"),
938                kind,
939                replay_suffix,
940            ));
941        }
942        let scoped = self.context.effect_controller.scoped();
943        let replay_key = format!(
944            "{}:tool:{tool_call_id}:{}",
945            scoped.scope_id(),
946            replay_suffix.as_ref()
947        );
948        Ok(crate::RuntimeInvocation::effect(
949            crate::RuntimeScope::new(self.context.session_id.clone()),
950            format!("{tool_call_id}:{effect_id_suffix}"),
951            kind,
952            replay_key,
953        ))
954    }
955}
956
957/// Runtime-prepared executable tool call.
958///
959/// The raw model/provider identity remains visible, but any argument rewrites
960/// and provider-owned context projections are frozen before the call crosses a
961/// runtime effect or process boundary.
962#[derive(Clone, Debug, Serialize, Deserialize)]
963pub struct PreparedToolCall {
964    pub call_id: String,
965    pub tool_id: ToolId,
966    pub tool_name: String,
967    pub args: serde_json::Value,
968    #[serde(default, skip_serializing_if = "Option::is_none")]
969    pub replay: Option<ProviderReplayMeta>,
970    #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
971    pub prepared_payload: serde_json::Value,
972}
973
974impl PreparedToolCall {
975    pub fn identity(tool_id: ToolId, call: crate::sansio::PendingToolCall) -> Self {
976        Self {
977            call_id: call.call_id,
978            tool_id,
979            tool_name: call.tool_name,
980            args: call.args,
981            replay: call.replay,
982            prepared_payload: serde_json::Value::Null,
983        }
984    }
985
986    pub fn from_parts(
987        call_id: impl Into<String>,
988        tool_id: impl Into<ToolId>,
989        tool_name: impl Into<String>,
990        args: serde_json::Value,
991        replay: Option<ProviderReplayMeta>,
992        prepared_payload: serde_json::Value,
993    ) -> Self {
994        Self {
995            call_id: call_id.into(),
996            tool_id: tool_id.into(),
997            tool_name: tool_name.into(),
998            args,
999            replay,
1000            prepared_payload,
1001        }
1002    }
1003}
1004
1005/// One ordered child inside a runtime-prepared tool batch.
1006///
1007/// The call itself carries the executable provider payload. `replay_suffix`
1008/// is the deterministic suffix used for child effects such as retry sleeps or
1009/// pending completion awaits when the batch is the durable parent.
1010#[derive(Clone, Debug, Serialize, Deserialize)]
1011pub struct PreparedToolBatchCall {
1012    pub call: PreparedToolCall,
1013    pub replay_suffix: String,
1014    #[serde(default, skip_serializing_if = "Option::is_none")]
1015    pub execution_grant: Option<Box<ToolExecutionGrant>>,
1016}
1017
1018/// Runtime-prepared executable tool batch.
1019///
1020/// The vector order is source order. Scheduling may run parallel-safe tools
1021/// concurrently, but launches and pending completion consumption are projected
1022/// back through this order.
1023#[derive(Clone, Debug, Serialize, Deserialize)]
1024pub struct PreparedToolBatch {
1025    pub batch_id: String,
1026    pub calls: Vec<PreparedToolBatchCall>,
1027}
1028
1029impl PreparedToolBatch {
1030    pub fn new(batch_id: impl Into<String>, calls: Vec<PreparedToolCall>) -> Self {
1031        let batch_id = batch_id.into();
1032        let calls = calls
1033            .into_iter()
1034            .enumerate()
1035            .map(|(index, call)| PreparedToolBatchCall {
1036                replay_suffix: format!("child:{index}:{}", call.call_id),
1037                call,
1038                execution_grant: None,
1039            })
1040            .collect();
1041        Self { batch_id, calls }
1042    }
1043
1044    pub fn new_with_grants(
1045        batch_id: impl Into<String>,
1046        calls: Vec<(PreparedToolCall, Option<ToolExecutionGrant>)>,
1047    ) -> Self {
1048        let batch_id = batch_id.into();
1049        let calls = calls
1050            .into_iter()
1051            .enumerate()
1052            .map(|(index, (call, execution_grant))| PreparedToolBatchCall {
1053                replay_suffix: format!("child:{index}:{}", call.call_id),
1054                call,
1055                execution_grant: execution_grant.map(Box::new),
1056            })
1057            .collect();
1058        Self { batch_id, calls }
1059    }
1060
1061    pub fn is_empty(&self) -> bool {
1062        self.calls.is_empty()
1063    }
1064
1065    pub fn len(&self) -> usize {
1066        self.calls.len()
1067    }
1068}
1069
1070/// Explicit authority to execute a tool outside Tool Catalog membership.
1071///
1072/// Normal tool calls are authorized by catalog membership. A grant is a
1073/// separate, caller-provided capability used by deferred resolution flows: it
1074/// carries the manifest/contract to validate the call plus an opaque host
1075/// execution binding that providers can inspect from the prepare and execute
1076/// contexts.
1077#[derive(Clone, Debug, Serialize, Deserialize)]
1078pub struct ToolExecutionGrant {
1079    /// Tool identity and model-facing metadata authorized by the grant.
1080    pub manifest: ToolManifest,
1081    /// Contract used to validate granted call arguments without consulting the
1082    /// current Tool Catalog.
1083    pub contract: Box<ToolContract>,
1084    /// Explicit registry source route for registry-backed execution. Direct
1085    /// non-registry providers may ignore this; [`ToolRegistry`] requires it.
1086    pub source_id: Option<String>,
1087    /// Opaque host routing payload passed to prepare and execute contexts.
1088    pub execution_binding: serde_json::Value,
1089}
1090
1091impl ToolExecutionGrant {
1092    pub fn new(manifest: ToolManifest, contract: ToolContract) -> Self {
1093        Self {
1094            manifest,
1095            contract: Box::new(contract),
1096            source_id: None,
1097            execution_binding: serde_json::Value::Null,
1098        }
1099    }
1100
1101    pub fn from_definition(definition: ToolDefinition) -> Self {
1102        Self::new(definition.manifest(), definition.contract())
1103    }
1104
1105    pub fn with_source_id(mut self, source_id: impl Into<String>) -> Self {
1106        self.source_id = Some(source_id.into());
1107        self
1108    }
1109
1110    pub fn with_execution_binding(mut self, execution_binding: serde_json::Value) -> Self {
1111        self.execution_binding = execution_binding;
1112        self
1113    }
1114}
1115
1116#[derive(Clone)]
1117pub struct ToolPrepareContext {
1118    session_id: String,
1119    sessions: Arc<dyn SessionStateService>,
1120    turn_context: crate::TurnContext,
1121    tool_call_id: Option<String>,
1122    tool_execution_binding: serde_json::Value,
1123}
1124
1125impl ToolPrepareContext {
1126    pub(crate) fn with_execution_binding(
1127        session_id: String,
1128        sessions: Arc<dyn SessionStateService>,
1129        turn_context: crate::TurnContext,
1130        tool_call_id: Option<String>,
1131        tool_execution_binding: serde_json::Value,
1132    ) -> Self {
1133        Self {
1134            session_id,
1135            sessions,
1136            turn_context,
1137            tool_call_id,
1138            tool_execution_binding,
1139        }
1140    }
1141
1142    pub fn session_id(&self) -> &str {
1143        &self.session_id
1144    }
1145
1146    pub fn tool_call_id(&self) -> Option<&str> {
1147        self.tool_call_id.as_deref()
1148    }
1149
1150    pub fn tool_execution_binding(&self) -> &serde_json::Value {
1151        &self.tool_execution_binding
1152    }
1153
1154    pub fn turn_context(&self) -> &crate::TurnContext {
1155        &self.turn_context
1156    }
1157
1158    pub fn plugin_input<T>(&self, plugin_id: &'static str) -> Option<&T>
1159    where
1160        T: 'static,
1161    {
1162        self.turn_context.plugin_input::<T>(plugin_id)
1163    }
1164
1165    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
1166        self.sessions.snapshot_session(&self.session_id).await
1167    }
1168
1169    pub async fn tool_catalog(&self) -> Result<Vec<serde_json::Value>, PluginError> {
1170        self.sessions.tool_catalog(&self.session_id).await
1171    }
1172
1173    pub async fn shared_tool_catalog(
1174        &self,
1175    ) -> Result<std::sync::Arc<Vec<serde_json::Value>>, PluginError> {
1176        self.sessions.shared_tool_catalog(&self.session_id).await
1177    }
1178}
1179
1180/// Inputs handed to [`ToolProvider::prepare_tool_call`].
1181pub struct ToolPrepareCall<'a> {
1182    pub tool_id: ToolId,
1183    pub pending: crate::sansio::PendingToolCall,
1184    pub context: &'a ToolPrepareContext,
1185}
1186
1187/// Per-call inputs handed to [`ToolProvider::execute`].
1188///
1189/// Fields are `pub` because `ToolCall` is a transient borrow; consumers
1190/// typically destructure (`let ToolCall { name, args, .. } = call`). The
1191/// stable surface lives on [`ToolContext`] (sealed) and the runtime's
1192/// dispatcher, which constructs `ToolCall` values.
1193pub struct ToolCall<'a> {
1194    pub name: &'a str,
1195    pub args: &'a serde_json::Value,
1196    pub context: &'a ToolContext<'a>,
1197    pub progress: Option<&'a ProgressSender>,
1198}
1199
1200/// Trait for providing tools to the sandbox. Implement this per-project.
1201///
1202/// Implementations supply cheap [`ToolManifest`]s, lazily resolved
1203/// [`ToolContract`]s, and a single
1204/// [`execute`](Self::execute) method that handles every call. Tools that
1205/// need session state read it from `call.context`; tools that stream
1206/// progress send through `call.progress`.
1207#[async_trait::async_trait]
1208pub trait ToolProvider: Send + Sync + 'static {
1209    fn tool_manifests(&self) -> Vec<ToolManifest>;
1210    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
1211        self.tool_manifests()
1212            .into_iter()
1213            .find(|manifest| manifest.name == name)
1214    }
1215    fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
1216        self.tool_manifests()
1217            .into_iter()
1218            .find(|manifest| manifest.id == *id)
1219    }
1220    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
1221    fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
1222        let manifest = self.resolve_manifest_by_id(id)?;
1223        self.resolve_contract(&manifest.name)
1224    }
1225    async fn prepare_tool_call(
1226        &self,
1227        call: ToolPrepareCall<'_>,
1228    ) -> Result<PreparedToolCall, ToolResult> {
1229        Ok(PreparedToolCall::identity(call.tool_id, call.pending))
1230    }
1231    async fn prepare_granted_tool_call(
1232        &self,
1233        grant: &ToolExecutionGrant,
1234        call: ToolPrepareCall<'_>,
1235    ) -> Result<PreparedToolCall, ToolResult> {
1236        let _ = call;
1237        Err(ToolResult::err_fmt(format_args!(
1238            "Granted execution is unsupported for tool id `{}`",
1239            grant.manifest.id
1240        )))
1241    }
1242    async fn execute(&self, call: ToolCall<'_>) -> ToolResult;
1243    async fn execute_granted(
1244        &self,
1245        grant: &ToolExecutionGrant,
1246        args: &serde_json::Value,
1247        context: &ToolContext<'_>,
1248        progress: Option<&ProgressSender>,
1249    ) -> ToolResult {
1250        let _ = (args, context, progress);
1251        ToolResult::err_fmt(format_args!(
1252            "Granted execution is unsupported for tool id `{}`",
1253            grant.manifest.id
1254        ))
1255    }
1256    async fn execute_by_id(
1257        &self,
1258        tool_id: &ToolId,
1259        args: &serde_json::Value,
1260        context: &ToolContext<'_>,
1261        progress: Option<&ProgressSender>,
1262    ) -> ToolResult {
1263        let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
1264            return ToolResult::err_fmt(format!("Unknown tool id: {tool_id}"));
1265        };
1266        self.execute(ToolCall {
1267            name: &manifest.name,
1268            args,
1269            context,
1270            progress,
1271        })
1272        .await
1273    }
1274}
1275
1276#[cfg(test)]
1277mod tests {
1278    use super::*;
1279    use crate::AwaitEventResolver;
1280    use crate::ProcessRegistry;
1281    use std::sync::atomic::{AtomicU64, Ordering};
1282
1283    struct NoDurableEffectController;
1284
1285    impl crate::AwaitEventResolver for NoDurableEffectController {}
1286
1287    #[async_trait::async_trait]
1288    impl crate::RuntimeEffectController for NoDurableEffectController {
1289        async fn execute_effect(
1290            &self,
1291            _envelope: crate::RuntimeEffectEnvelope,
1292            _local_executor: crate::RuntimeEffectLocalExecutor<'_>,
1293        ) -> Result<crate::RuntimeEffectOutcome, crate::RuntimeEffectControllerError> {
1294            Err(crate::RuntimeEffectControllerError::new(
1295                "unexpected_effect",
1296                "test controller should not execute effects",
1297            ))
1298        }
1299    }
1300
1301    fn test_context_with_controller(
1302        tool_call_id: Option<String>,
1303        controller: Arc<dyn crate::RuntimeEffectController>,
1304    ) -> ToolContext<'static> {
1305        ToolContext::builder(
1306            "session-1".to_string(),
1307            Arc::new(crate::testing::MockSessionManager::default()),
1308            Arc::new(crate::testing::MockSessionManager::default()),
1309            Arc::new(crate::testing::MockSessionManager::default()),
1310            Arc::new(crate::UnavailableProcessService),
1311            Arc::new(crate::DefaultProcessCancelAbility),
1312            crate::runtime::RuntimeEffectControllerHandle::shared(controller),
1313            Arc::new(crate::SessionAttachmentStore::in_memory()),
1314            crate::DirectCompletionClient::unavailable(
1315                "direct completions are unavailable in this test context",
1316            ),
1317        )
1318        .tool_call_id(tool_call_id)
1319        .build()
1320    }
1321
1322    #[test]
1323    fn tool_context_builder_carries_call_payload_and_cancellation_state() {
1324        let cancellation = tokio_util::sync::CancellationToken::new();
1325        let prepared = PreparedToolCall::from_parts(
1326            "call-1",
1327            "tool:demo_tool",
1328            "demo_tool",
1329            serde_json::json!({ "input": true }),
1330            None,
1331            serde_json::json!({ "prepared": true }),
1332        );
1333
1334        let context = ToolContext::builder(
1335            "session-1".to_string(),
1336            Arc::new(crate::testing::MockSessionManager::default()),
1337            Arc::new(crate::testing::MockSessionManager::default()),
1338            Arc::new(crate::testing::MockSessionManager::default()),
1339            Arc::new(crate::UnavailableProcessService),
1340            Arc::new(crate::DefaultProcessCancelAbility),
1341            crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
1342                crate::InlineRuntimeEffectController,
1343            )),
1344            Arc::new(crate::SessionAttachmentStore::in_memory()),
1345            crate::DirectCompletionClient::unavailable(
1346                "direct completions are unavailable in this test context",
1347            ),
1348        )
1349        .prepared_call(&prepared)
1350        .cancellation_token(Some(cancellation.clone()))
1351        .async_process("process-1", cancellation.clone())
1352        .build();
1353
1354        assert_eq!(context.session_id(), "session-1");
1355        assert_eq!(context.tool_call_id(), Some("call-1"));
1356        assert_eq!(
1357            context.prepared_payload(),
1358            &serde_json::json!({ "prepared": true })
1359        );
1360        assert_eq!(context.async_process_id(), Some("process-1"));
1361        assert!(context.cancellation_token().is_some());
1362    }
1363
1364    #[test]
1365    fn durable_effects_requires_prepared_call_id_and_supporting_controller() {
1366        let missing_call =
1367            test_context_with_controller(None, Arc::new(crate::InlineRuntimeEffectController));
1368        let err = match missing_call.durable_effects() {
1369            Ok(_) => panic!("missing prepared tool call id should fail"),
1370            Err(err) => err,
1371        };
1372        assert_eq!(err.code.as_str(), "durable_effects_missing_call_id");
1373
1374        let unsupported = test_context_with_controller(
1375            Some("call-1".to_string()),
1376            Arc::new(NoDurableEffectController),
1377        );
1378        let err = match unsupported.durable_effects() {
1379            Ok(_) => panic!("unsupported controller should fail"),
1380            Err(err) => err,
1381        };
1382        assert_eq!(err.code.as_str(), "durable_effects_unavailable");
1383    }
1384
1385    #[tokio::test]
1386    async fn durable_run_json_executes_and_maps_closure_errors() {
1387        let context = test_context_with_controller(
1388            Some("call-run-json".to_string()),
1389            Arc::new(crate::InlineRuntimeEffectController),
1390        );
1391        let durable = context.durable_effects().expect("durable effects");
1392        let value = durable
1393            .run_json(
1394                "create",
1395                serde_json::json!({ "x": 1 }),
1396                |input| async move { Ok(serde_json::json!({ "seen": input["x"] })) },
1397            )
1398            .await
1399            .expect("durable step");
1400        assert_eq!(value, serde_json::json!({ "seen": 1 }));
1401
1402        let err = durable
1403            .run_json("fail", serde_json::json!({}), |_| async {
1404                Err(crate::RuntimeError::new(
1405                    "durable_step_failed",
1406                    "step failed",
1407                ))
1408            })
1409            .await
1410            .expect_err("closure error");
1411        assert_eq!(err.code.as_str(), "durable_step_failed");
1412        assert_eq!(err.message, "step failed");
1413    }
1414
1415    #[tokio::test]
1416    async fn durable_run_json_rejects_empty_or_duplicate_step_ids_before_running() {
1417        let context = test_context_with_controller(
1418            Some("call-step-ids".to_string()),
1419            Arc::new(crate::InlineRuntimeEffectController),
1420        );
1421        let durable = context.durable_effects().expect("durable effects");
1422        let runs = Arc::new(AtomicU64::new(0));
1423
1424        let err = durable
1425            .run_json("", serde_json::Value::Null, {
1426                let runs = Arc::clone(&runs);
1427                move |_| async move {
1428                    runs.fetch_add(1, Ordering::Relaxed);
1429                    Ok(serde_json::Value::Null)
1430                }
1431            })
1432            .await
1433            .expect_err("empty step id");
1434        assert_eq!(err.code.as_str(), "durable_effect_empty_step_id");
1435        assert_eq!(runs.load(Ordering::Relaxed), 0);
1436
1437        durable
1438            .run_json("same", serde_json::Value::Null, {
1439                let runs = Arc::clone(&runs);
1440                move |_| async move {
1441                    runs.fetch_add(1, Ordering::Relaxed);
1442                    Ok(serde_json::Value::Null)
1443                }
1444            })
1445            .await
1446            .expect("first step");
1447        let err = durable
1448            .run_json("same", serde_json::Value::Null, {
1449                let runs = Arc::clone(&runs);
1450                move |_| async move {
1451                    runs.fetch_add(1, Ordering::Relaxed);
1452                    Ok(serde_json::Value::Null)
1453                }
1454            })
1455            .await
1456            .expect_err("duplicate step id");
1457        assert_eq!(err.code.as_str(), "durable_effect_duplicate_step_id");
1458        assert_eq!(runs.load(Ordering::Relaxed), 1);
1459    }
1460
1461    #[tokio::test]
1462    async fn durable_external_event_key_is_custom_and_stable() {
1463        let context = test_context_with_controller(
1464            Some("call-event-key".to_string()),
1465            Arc::new(crate::InlineRuntimeEffectController),
1466        );
1467        let durable = context.durable_effects().expect("durable effects");
1468        let first = durable
1469            .external_event_key("tool-event-stable")
1470            .await
1471            .expect("first key");
1472        let second = durable
1473            .external_event_key("tool-event-stable")
1474            .await
1475            .expect("second key");
1476
1477        assert_eq!(first, second);
1478        assert_eq!(
1479            first.wait,
1480            crate::AwaitEventWaitIdentity::Custom {
1481                key: "tool-event-stable".to_string()
1482            }
1483        );
1484    }
1485
1486    #[tokio::test]
1487    async fn durable_await_event_json_maps_terminal_resolutions() {
1488        let controller = Arc::new(crate::InlineRuntimeEffectController);
1489        let context =
1490            test_context_with_controller(Some("call-await-event".to_string()), controller.clone());
1491        let durable = context.durable_effects().expect("durable effects");
1492
1493        let ok_key = durable
1494            .external_event_key("tool-event-ok")
1495            .await
1496            .expect("ok key");
1497        controller
1498            .resolve_await_event(
1499                &ok_key,
1500                crate::Resolution::Ok(serde_json::json!({ "answer": 42 })),
1501            )
1502            .await
1503            .expect("resolve ok");
1504        let value = durable
1505            .await_event_json(ok_key)
1506            .await
1507            .expect("await ok value");
1508        assert_eq!(value, serde_json::json!({ "answer": 42 }));
1509
1510        let err_key = durable
1511            .external_event_key("tool-event-err")
1512            .await
1513            .expect("err key");
1514        controller
1515            .resolve_await_event(
1516                &err_key,
1517                crate::Resolution::Err(crate::ExternalCompletionError::new(
1518                    "external_bad",
1519                    "external failed",
1520                )),
1521            )
1522            .await
1523            .expect("resolve err");
1524        let err = durable
1525            .await_event_json(err_key)
1526            .await
1527            .expect_err("await err value");
1528        assert_eq!(err.code.as_str(), "external_bad");
1529
1530        let cancelled_key = durable
1531            .external_event_key("tool-event-cancelled")
1532            .await
1533            .expect("cancelled key");
1534        controller
1535            .resolve_await_event(&cancelled_key, crate::Resolution::Cancelled)
1536            .await
1537            .expect("resolve cancelled");
1538        let err = durable
1539            .await_event_json(cancelled_key)
1540            .await
1541            .expect_err("await cancelled value");
1542        assert_eq!(err.code.as_str(), "durable_effect_event_cancelled");
1543
1544        let timeout_key = durable
1545            .external_event_key("tool-event-timeout")
1546            .await
1547            .expect("timeout key");
1548        controller
1549            .resolve_await_event(&timeout_key, crate::Resolution::Timeout)
1550            .await
1551            .expect("resolve timeout");
1552        let err = durable
1553            .await_event_json(timeout_key)
1554            .await
1555            .expect_err("await timeout value");
1556        assert_eq!(err.code.as_str(), "durable_effect_event_timeout");
1557    }
1558
1559    #[tokio::test]
1560    async fn durable_emit_process_event_requires_process_and_appends_inside_process() {
1561        let context = test_context_with_controller(
1562            Some("call-no-process".to_string()),
1563            Arc::new(crate::InlineRuntimeEffectController),
1564        );
1565        let err = context
1566            .durable_effects()
1567            .expect("durable effects")
1568            .emit_process_event("tool.event", serde_json::json!({}))
1569            .await
1570            .expect_err("outside process");
1571        assert_eq!(
1572            err.code.as_str(),
1573            "durable_effect_process_event_unavailable"
1574        );
1575
1576        let registry = Arc::new(crate::TestLocalProcessRegistry::default());
1577        let process_id = "process:durable-tool-event";
1578        registry
1579            .register_process(
1580                crate::ProcessRegistration::new(
1581                    process_id,
1582                    crate::ProcessInput::External {
1583                        metadata: serde_json::json!({}),
1584                    },
1585                    crate::RecoveryDisposition::ExternallyOwned,
1586                    crate::ProcessProvenance::host(),
1587                )
1588                .with_extra_event_types([crate::ProcessEventType {
1589                    name: "tool.event".to_string(),
1590                    payload_schema: crate::LashSchema::any(),
1591                    semantics: crate::ProcessEventSemanticsSpec::default(),
1592                }]),
1593            )
1594            .await
1595            .expect("register process");
1596        let registry_dyn: Arc<dyn crate::ProcessRegistry> = registry;
1597        let context = test_context_with_controller(
1598            Some("call-process-event".to_string()),
1599            Arc::new(crate::InlineRuntimeEffectController),
1600        )
1601        .with_process_events_for_testing(process_id, registry_dyn);
1602
1603        let event = context
1604            .durable_effects()
1605            .expect("durable effects")
1606            .emit_process_event("tool.event", serde_json::json!({ "ok": true }))
1607            .await
1608            .expect("process event");
1609        assert_eq!(event.process_id, process_id);
1610        assert_eq!(event.event_type, "tool.event");
1611        assert_eq!(event.payload, serde_json::json!({ "ok": true }));
1612        assert_eq!(
1613            event.invocation.replay_key(),
1614            Some("tool:call-process-event:durable-process-event:0")
1615        );
1616
1617        let append_err = context
1618            .durable_effects()
1619            .expect("durable effects")
1620            .emit_process_event("undeclared.event", serde_json::json!({}))
1621            .await
1622            .expect_err("undeclared event type must fail the append");
1623        assert_eq!(
1624            append_err.code.as_str(),
1625            "durable_effect_process_event_append_failed"
1626        );
1627    }
1628}
lash_core/tool_provider.rs

lash_core/
tool_provider.rs
