Skip to main content

lash_core/tool_dispatch/
context.rs

1use std::sync::{Arc, Mutex};
2
3use tokio::sync::mpsc;
4
5use crate::plugin::{
6    PluginSession, SessionGraphService, SessionLifecycleService, SessionStateService,
7};
8use crate::{
9    PreparedToolCall, SessionEvent, ToolCallRecord, ToolFailure, ToolFailureClass, ToolProvider,
10    ToolResult, ToolSurface,
11};
12
13#[derive(Clone, Default)]
14pub(crate) struct CheckpointMessageBuffer {
15    queue: Arc<Mutex<Vec<crate::PluginMessage>>>,
16}
17
18impl CheckpointMessageBuffer {
19    pub(crate) fn enqueue(&self, messages: Vec<crate::PluginMessage>) -> Result<(), String> {
20        let mut queue = self
21            .queue
22            .lock()
23            .map_err(|_| "checkpoint message buffer poisoned".to_string())?;
24        queue.extend(messages);
25        Ok(())
26    }
27
28    pub(crate) fn drain(&self) -> Result<Vec<crate::PluginMessage>, String> {
29        let mut queue = self
30            .queue
31            .lock()
32            .map_err(|_| "checkpoint message buffer poisoned".to_string())?;
33        Ok(queue.drain(..).collect())
34    }
35}
36
37#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
38pub struct ToolHostEventEffectOutcome {
39    pub resource_type: String,
40    pub alias: String,
41    pub event: String,
42    pub source_type: String,
43    #[serde(default)]
44    pub payload: serde_json::Value,
45    pub started_process_ids: Vec<String>,
46}
47
48#[derive(Clone, Default)]
49pub(crate) struct ToolHostEventOutcomeBuffer {
50    queue: Arc<Mutex<Vec<ToolHostEventEffectOutcome>>>,
51}
52
53impl ToolHostEventOutcomeBuffer {
54    pub(crate) fn enqueue(&self, outcome: ToolHostEventEffectOutcome) -> Result<(), String> {
55        let mut queue = self
56            .queue
57            .lock()
58            .map_err(|_| "tool host event outcome buffer poisoned".to_string())?;
59        queue.push(outcome);
60        Ok(())
61    }
62
63    pub(crate) fn drain(&self) -> Result<Vec<ToolHostEventEffectOutcome>, String> {
64        let mut queue = self
65            .queue
66            .lock()
67            .map_err(|_| "tool host event outcome buffer poisoned".to_string())?;
68        Ok(queue.drain(..).collect())
69    }
70}
71
72#[derive(Clone)]
73pub struct ToolDispatchContext<'run> {
74    pub plugins: Arc<PluginSession>,
75    pub tools: Arc<dyn ToolProvider>,
76    pub surface: Arc<ToolSurface>,
77    pub sessions: Arc<dyn SessionStateService>,
78    pub session_lifecycle: Arc<dyn SessionLifecycleService>,
79    pub session_graph: Arc<dyn SessionGraphService>,
80    pub processes: Arc<dyn crate::ProcessService>,
81    pub process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
82    pub(crate) effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
83    pub(crate) direct_completions: crate::DirectCompletionClient<'run>,
84    pub(crate) parent_invocation: Option<crate::RuntimeInvocation>,
85    pub session_id: String,
86    pub agent_frame_id: crate::AgentFrameId,
87    pub event_tx: mpsc::Sender<SessionEvent>,
88    pub(crate) checkpoint_messages: CheckpointMessageBuffer,
89    pub(crate) host_event_outcomes: ToolHostEventOutcomeBuffer,
90    pub attachment_store: Arc<dyn crate::AttachmentStore>,
91    pub turn_context: crate::TurnContext,
92}
93
94impl<'run> ToolDispatchContext<'run> {
95    pub fn process_scope(&self) -> crate::ProcessOpScope<'_> {
96        crate::ProcessOpScope::new(self.effect_controller.scoped())
97            .with_parent_invocation(self.parent_invocation.clone())
98            .with_agent_frame_id(Some(self.agent_frame_id.clone()))
99    }
100}
101
102#[derive(Clone)]
103pub(crate) struct ToolDispatchOutcome {
104    pub record: ToolCallRecord,
105}
106
107pub(crate) enum ToolPreparationOutcome {
108    Prepared(PreparedToolCall),
109    Completed(Box<ToolDispatchOutcome>),
110}
111
112pub(super) fn completed_preparation(outcome: ToolDispatchOutcome) -> ToolPreparationOutcome {
113    ToolPreparationOutcome::Completed(Box::new(outcome))
114}
115pub(super) fn outcome(
116    tool_name: String,
117    args: serde_json::Value,
118    result: ToolResult,
119    duration_ms: u64,
120) -> ToolDispatchOutcome {
121    let record = ToolCallRecord {
122        call_id: None,
123        tool: tool_name,
124        args,
125        output: result.into_output(),
126        duration_ms,
127    };
128    ToolDispatchOutcome { record }
129}
130
131pub(super) fn runtime_failure(
132    class: ToolFailureClass,
133    code: impl Into<String>,
134    message: impl Into<String>,
135) -> ToolResult {
136    ToolResult::failure(ToolFailure::runtime(class, code, message))
137}