Skip to main content

lash_core/plugin/
hooks.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use super::*;
6
7pub type PluginFuture<T> = Pin<Box<dyn Future<Output = Result<T, PluginError>> + Send>>;
8pub type PluginLifecycleFuture<'run> =
9    Pin<Box<dyn Future<Output = Result<(), PluginError>> + Send + 'run>>;
10pub type PluginLifecycleEventHook =
11    Arc<dyn for<'run> Fn(PluginLifecycleEvent<'run>) -> PluginLifecycleFuture<'run> + Send + Sync>;
12pub type PluginSessionTask = PluginFuture<()>;
13pub type SessionConfigMutator = Arc<
14    dyn Fn(SessionConfigChangedContext, SessionPolicy) -> PluginFuture<SessionPolicy> + Send + Sync,
15>;
16pub type BeforeTurnHook =
17    Arc<dyn Fn(TurnHookContext) -> PluginFuture<Vec<PluginDirective>> + Send + Sync>;
18pub type BeforeToolCallHook =
19    Arc<dyn Fn(ToolCallHookContext) -> PluginFuture<Vec<PluginDirective>> + Send + Sync>;
20pub type AfterToolCallHook =
21    Arc<dyn Fn(ToolResultHookContext) -> PluginFuture<Vec<PluginDirective>> + Send + Sync>;
22pub type ToolResultProjector =
23    Arc<dyn Fn(ToolResultProjectionContext) -> PluginFuture<crate::ModelToolReturn> + Send + Sync>;
24pub type AfterTurnHook =
25    Arc<dyn Fn(TurnResultHookContext) -> PluginFuture<Vec<PluginDirective>> + Send + Sync>;
26pub type CheckpointHook =
27    Arc<dyn Fn(CheckpointHookContext) -> PluginFuture<Vec<PluginDirective>> + Send + Sync>;
28pub type PromptContributor =
29    Arc<dyn Fn(PromptHookContext) -> PluginFuture<Vec<PromptContribution>> + Send + Sync>;
30pub type ToolSurfaceContributor =
31    Arc<dyn Fn(ToolSurfaceContext) -> Result<ToolSurfaceContribution, PluginError> + Send + Sync>;
32pub type ToolDiscoveryContributor = Arc<
33    dyn Fn(ToolDiscoveryContext) -> Result<ToolDiscoveryContribution, PluginError> + Send + Sync,
34>;
35pub type AssistantStreamHook =
36    Arc<dyn Fn(AssistantStreamHookContext) -> PluginFuture<AssistantStreamTransform> + Send + Sync>;
37pub type AssistantResponseHook = Arc<
38    dyn Fn(AssistantResponseHookContext) -> PluginFuture<AssistantResponseTransform> + Send + Sync,
39>;
40
41#[derive(Clone)]
42pub struct PromptHookContext {
43    pub session_id: String,
44    pub sessions: Arc<dyn SessionStateService>,
45    pub state: SessionReadView,
46    pub protocol_turn_options: ProtocolTurnOptions,
47    pub turn_context: crate::TurnContext,
48}
49
50#[derive(Clone)]
51pub struct TurnHookContext {
52    pub session_id: String,
53    pub state: SessionReadView,
54    pub sessions: Arc<dyn SessionStateService>,
55    pub turn_context: crate::TurnContext,
56}
57
58#[derive(Clone)]
59pub struct SessionConfigChangedContext {
60    pub session_id: String,
61    pub previous: SessionPolicy,
62    pub current: SessionPolicy,
63    pub sessions: Arc<dyn SessionStateService>,
64}
65
66#[derive(Clone)]
67pub struct SessionStateChangedContext<'run> {
68    pub session_id: String,
69    pub state: SessionReadView,
70    pub sessions: Arc<dyn SessionStateService>,
71    pub session_graph: Arc<dyn SessionGraphService>,
72    pub direct_completions: crate::DirectCompletionClient<'run>,
73}
74
75#[derive(Clone)]
76pub enum PluginLifecycleEvent<'run> {
77    TurnFinalized(Arc<AssembledTurn>),
78    /// Best-effort observer hook emitted after durable session state advances.
79    ///
80    /// Hook failures are isolated from the foreground turn: the turn has already
81    /// committed, and observers must not affect that commit.
82    TurnPersisted(Box<SessionStateChangedContext<'run>>),
83    SessionRestored(SessionReadView),
84    SessionConfigChanged(Box<SessionConfigChangedContext>),
85}
86
87#[derive(Clone, Debug)]
88pub struct TurnResultSummary {
89    pub outcome: crate::TurnOutcome,
90    pub assistant_output: crate::runtime::AssistantOutput,
91    pub execution: crate::runtime::ExecutionSummary,
92    pub token_usage: crate::TokenUsage,
93    pub tool_calls: Arc<Vec<crate::ToolCallRecord>>,
94    pub errors: Arc<Vec<crate::runtime::TurnIssue>>,
95}
96
97impl TurnResultSummary {
98    pub fn from_assembled(turn: &AssembledTurn) -> Self {
99        Self {
100            outcome: turn.outcome.clone(),
101            assistant_output: turn.assistant_output.clone(),
102            execution: turn.execution.clone(),
103            token_usage: turn.token_usage.clone(),
104            tool_calls: Arc::new(turn.tool_calls.clone()),
105            errors: Arc::new(turn.errors.clone()),
106        }
107    }
108}
109
110#[derive(Clone)]
111pub struct ToolCallHookContext {
112    pub session_id: String,
113    pub tool_name: String,
114    pub args: serde_json::Value,
115    pub argument_projection: crate::ToolArgumentProjectionPolicy,
116    pub turn_context: crate::TurnContext,
117    pub(crate) sessions: Arc<dyn SessionStateService>,
118}
119
120impl ToolCallHookContext {
121    pub fn new(
122        session_id: String,
123        tool_name: String,
124        args: serde_json::Value,
125        argument_projection: crate::ToolArgumentProjectionPolicy,
126        turn_context: crate::TurnContext,
127        sessions: Arc<dyn SessionStateService>,
128    ) -> Self {
129        Self {
130            session_id,
131            tool_name,
132            args,
133            argument_projection,
134            turn_context,
135            sessions,
136        }
137    }
138
139    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
140        self.sessions.snapshot_session(&self.session_id).await
141    }
142
143    pub async fn set_tools_availability(
144        &self,
145        names: &[String],
146        availability: Option<crate::ToolAvailability>,
147    ) -> Result<u64, PluginError> {
148        self.sessions
149            .set_tools_availability(&self.session_id, names, availability)
150            .await
151    }
152}
153
154#[derive(Clone)]
155pub struct ToolResultHookContext {
156    pub session_id: String,
157    pub tool_name: String,
158    pub args: serde_json::Value,
159    pub result: ToolResult,
160    pub duration_ms: u64,
161    pub turn_context: crate::TurnContext,
162    pub(crate) sessions: Arc<dyn SessionStateService>,
163}
164
165impl ToolResultHookContext {
166    pub fn new(
167        session_id: String,
168        tool_name: String,
169        args: serde_json::Value,
170        result: ToolResult,
171        duration_ms: u64,
172        turn_context: crate::TurnContext,
173        sessions: Arc<dyn SessionStateService>,
174    ) -> Self {
175        Self {
176            session_id,
177            tool_name,
178            args,
179            result,
180            duration_ms,
181            turn_context,
182            sessions,
183        }
184    }
185
186    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
187        self.sessions.snapshot_session(&self.session_id).await
188    }
189
190    pub async fn set_tools_availability(
191        &self,
192        names: &[String],
193        availability: Option<crate::ToolAvailability>,
194    ) -> Result<u64, PluginError> {
195        self.sessions
196            .set_tools_availability(&self.session_id, names, availability)
197            .await
198    }
199}
200
201#[derive(Clone)]
202pub struct ToolResultProjectionContext {
203    pub session_id: String,
204    pub call_id: String,
205    pub tool_name: String,
206    pub args: serde_json::Value,
207    pub output: crate::ToolCallOutput,
208    pub duration_ms: u64,
209}
210
211#[derive(Clone)]
212pub struct TurnResultHookContext {
213    pub session_id: String,
214    pub turn: Arc<TurnResultSummary>,
215    pub sessions: Arc<dyn SessionStateService>,
216}
217
218#[derive(Clone)]
219pub struct CheckpointHookContext {
220    pub session_id: String,
221    pub checkpoint: CheckpointKind,
222    pub state: SessionReadView,
223    pub sessions: Arc<dyn SessionStateService>,
224    pub session_lifecycle: Arc<dyn SessionLifecycleService>,
225    pub session_graph: Arc<dyn SessionGraphService>,
226}
227
228#[derive(Clone)]
229pub struct AssistantStreamHookContext {
230    pub session_id: String,
231    pub chunk: String,
232}
233
234#[derive(Clone, Debug, Default)]
235pub struct AssistantStreamTransform {
236    pub chunk: String,
237    pub reasoning_deltas: Vec<String>,
238    pub events: Vec<PluginRuntimeEvent>,
239    /// When `true`, the runtime cancels the in-flight LLM call the
240    /// moment this hook returns and finalizes the turn using whatever
241    /// text has been streamed so far. Any plugin may set this — the
242    /// first to raise it wins. Used by protocol plugins to enforce
243    /// one-block-per-turn contracts (e.g. the RLM stream mask aborts
244    /// as soon as the first lashlang fence closes).
245    pub abort_stream: bool,
246}
247
248#[derive(Clone)]
249pub struct AssistantResponseHookContext {
250    pub session_id: String,
251    pub response: crate::LlmResponse,
252}
253
254#[derive(Clone, Debug)]
255pub struct AssistantResponseTransform {
256    pub response: crate::LlmResponse,
257    pub events: Vec<PluginRuntimeEvent>,
258}