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 ToolCatalogContributor =
31    Arc<dyn Fn(ToolCatalogContext) -> Result<ToolCatalogContribution, PluginError> + Send + Sync>;
32pub type AssistantStreamHook =
33    Arc<dyn Fn(AssistantStreamHookContext) -> PluginFuture<AssistantStreamTransform> + Send + Sync>;
34pub type AssistantResponseHook = Arc<
35    dyn Fn(AssistantResponseHookContext) -> PluginFuture<AssistantResponseTransform> + Send + Sync,
36>;
37pub type AssistantStreamFinishedHook =
38    Arc<dyn Fn(AssistantStreamFinishedContext) -> PluginFuture<()> + Send + Sync>;
39
40#[derive(Clone)]
41pub struct PromptHookContext {
42    pub session_id: String,
43    pub sessions: Arc<dyn SessionStateService>,
44    pub state: SessionReadView,
45    pub protocol_turn_options: ProtocolTurnOptions,
46    pub turn_context: crate::TurnContext,
47}
48
49#[derive(Clone)]
50pub struct TurnHookContext {
51    pub session_id: String,
52    pub state: SessionReadView,
53    pub sessions: Arc<dyn SessionStateService>,
54    pub turn_context: crate::TurnContext,
55}
56
57#[derive(Clone)]
58pub struct SessionConfigChangedContext {
59    pub session_id: String,
60    pub previous: SessionPolicy,
61    pub current: SessionPolicy,
62    pub sessions: Arc<dyn SessionStateService>,
63}
64
65#[derive(Clone)]
66pub struct SessionStateChangedContext<'run> {
67    pub session_id: String,
68    pub state: SessionReadView,
69    pub sessions: Arc<dyn SessionStateService>,
70    pub session_graph: Arc<dyn SessionGraphService>,
71    pub direct_completions: crate::DirectCompletionClient<'run>,
72}
73
74#[derive(Clone)]
75pub enum PluginLifecycleEvent<'run> {
76    TurnFinalized(Arc<AssembledTurn>),
77    /// Best-effort observer hook emitted after durable session state advances.
78    ///
79    /// Hook failures are isolated from the foreground turn: the turn has already
80    /// committed, and observers must not affect that commit.
81    TurnPersisted(Box<SessionStateChangedContext<'run>>),
82    SessionRestored(SessionReadView),
83    SessionConfigChanged(Box<SessionConfigChangedContext>),
84}
85
86#[derive(Clone, Debug)]
87pub struct TurnResultSummary {
88    pub outcome: crate::TurnOutcome,
89    pub assistant_output: crate::runtime::AssistantOutput,
90    pub execution: crate::runtime::ExecutionSummary,
91    pub token_usage: crate::TokenUsage,
92    pub tool_calls: Arc<Vec<crate::ToolCallRecord>>,
93    pub errors: Arc<Vec<crate::runtime::TurnIssue>>,
94}
95
96impl TurnResultSummary {
97    pub fn from_assembled(turn: &AssembledTurn) -> Self {
98        Self {
99            outcome: turn.outcome.clone(),
100            assistant_output: turn.assistant_output.clone(),
101            execution: turn.execution.clone(),
102            token_usage: turn.token_usage.clone(),
103            tool_calls: Arc::new(turn.tool_calls.clone()),
104            errors: Arc::new(turn.errors.clone()),
105        }
106    }
107}
108
109#[derive(Clone)]
110pub struct ToolCallHookContext {
111    pub session_id: String,
112    pub tool_name: String,
113    pub args: serde_json::Value,
114    pub argument_projection: crate::ToolArgumentProjectionPolicy,
115    pub turn_context: crate::TurnContext,
116    pub(crate) sessions: Arc<dyn SessionStateService>,
117}
118
119impl ToolCallHookContext {
120    pub fn new(
121        session_id: String,
122        tool_name: String,
123        args: serde_json::Value,
124        argument_projection: crate::ToolArgumentProjectionPolicy,
125        turn_context: crate::TurnContext,
126        sessions: Arc<dyn SessionStateService>,
127    ) -> Self {
128        Self {
129            session_id,
130            tool_name,
131            args,
132            argument_projection,
133            turn_context,
134            sessions,
135        }
136    }
137
138    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
139        self.sessions.snapshot_session(&self.session_id).await
140    }
141
142    pub async fn set_tool_membership(
143        &self,
144        names: &[String],
145        present: bool,
146    ) -> Result<u64, PluginError> {
147        self.sessions
148            .set_tool_membership(&self.session_id, names, present)
149            .await
150    }
151}
152
153#[derive(Clone)]
154pub struct ToolResultHookContext {
155    pub session_id: String,
156    pub tool_name: String,
157    pub args: serde_json::Value,
158    pub result: ToolResult,
159    pub duration_ms: u64,
160    pub turn_context: crate::TurnContext,
161    pub(crate) sessions: Arc<dyn SessionStateService>,
162}
163
164impl ToolResultHookContext {
165    pub fn new(
166        session_id: String,
167        tool_name: String,
168        args: serde_json::Value,
169        result: ToolResult,
170        duration_ms: u64,
171        turn_context: crate::TurnContext,
172        sessions: Arc<dyn SessionStateService>,
173    ) -> Self {
174        Self {
175            session_id,
176            tool_name,
177            args,
178            result,
179            duration_ms,
180            turn_context,
181            sessions,
182        }
183    }
184
185    pub async fn session_snapshot(&self) -> Result<SessionSnapshot, PluginError> {
186        self.sessions.snapshot_session(&self.session_id).await
187    }
188
189    pub async fn set_tool_membership(
190        &self,
191        names: &[String],
192        present: bool,
193    ) -> Result<u64, PluginError> {
194        self.sessions
195            .set_tool_membership(&self.session_id, names, present)
196            .await
197    }
198}
199
200#[derive(Clone)]
201pub struct ToolResultProjectionContext {
202    pub session_id: String,
203    pub call_id: String,
204    pub tool_name: String,
205    pub args: serde_json::Value,
206    pub output: crate::ToolCallOutput,
207    pub duration_ms: u64,
208}
209
210#[derive(Clone)]
211pub struct TurnResultHookContext {
212    pub session_id: String,
213    pub turn: Arc<TurnResultSummary>,
214    pub sessions: Arc<dyn SessionStateService>,
215}
216
217#[derive(Clone)]
218pub struct CheckpointHookContext {
219    pub session_id: String,
220    pub checkpoint: CheckpointKind,
221    pub state: SessionReadView,
222    pub sessions: Arc<dyn SessionStateService>,
223    pub session_lifecycle: Arc<dyn SessionLifecycleService>,
224    pub session_graph: Arc<dyn SessionGraphService>,
225}
226
227#[derive(Clone)]
228pub struct AssistantStreamHookContext {
229    pub session_id: String,
230    pub chunk: String,
231}
232
233#[derive(Clone, Debug, Default)]
234pub struct AssistantStreamTransform {
235    pub chunk: String,
236    pub reasoning_deltas: Vec<String>,
237    pub events: Vec<PluginRuntimeEvent>,
238    /// When `true`, the runtime cancels the in-flight LLM call the
239    /// moment this hook returns and finalizes the turn using whatever
240    /// text has been streamed so far. Any plugin may set this — the
241    /// first to raise it wins. Used by protocol plugins to enforce
242    /// one-block-per-turn contracts (for example, aborting as soon as
243    /// the first protocol-owned code fence closes).
244    pub abort_stream: bool,
245}
246
247#[derive(Clone)]
248pub struct AssistantResponseHookContext {
249    pub session_id: String,
250    pub response: crate::LlmResponse,
251}
252
253#[derive(Clone, Debug)]
254pub struct AssistantResponseTransform {
255    pub response: crate::LlmResponse,
256    pub events: Vec<PluginRuntimeEvent>,
257}
258
259#[derive(Clone, Copy, Debug, PartialEq, Eq)]
260pub enum AssistantStreamFinishReason {
261    Complete,
262    Aborted,
263    Cancelled,
264    ProviderError,
265}
266
267#[derive(Clone)]
268pub struct AssistantStreamFinishedContext {
269    pub session_id: String,
270    pub reason: AssistantStreamFinishReason,
271}