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