starweaver_runtime/capability/
hooks.rs

1//! Agent capability hook interface.
2
3use async_trait::async_trait;
4use starweaver_context::AgentContext;
5use starweaver_model::{
6    ModelMessage, ModelRequest, ModelResponse, ModelSettings, ToolCallPart, ToolDefinition,
7    ToolReturnPart,
8};
9use starweaver_tools::ToolContext;
10
11use crate::{
12    agent::AgentInput, executor::AgentCheckpoint, run::AgentRunState, stream::AgentStreamRecord,
13};
14
15use super::{CapabilityResult, CapabilitySpec, RetryEventKind};
16
17/// Hook interface for runtime extension points.
18#[async_trait]
19pub trait AgentCapability: Send + Sync {
20    /// Stable capability spec used for ordering and reconstruction evidence.
21    fn spec(&self) -> CapabilitySpec {
22        CapabilitySpec::new(std::any::type_name::<Self>())
23    }
24
25    /// Called after a run state is created and before the first request is prepared.
26    async fn on_run_start(&self, _state: &mut AgentRunState) -> CapabilityResult<()> {
27        Ok(())
28    }
29
30    /// Context-aware run-start hook.
31    async fn on_run_start_with_context(
32        &self,
33        state: &mut AgentRunState,
34        _context: &mut AgentContext,
35    ) -> CapabilityResult<()> {
36        self.on_run_start(state).await
37    }
38
39    /// Called after run state is initialized and before the first model request is built.
40    async fn prepare_run_input(
41        &self,
42        _state: &mut AgentRunState,
43        input: AgentInput,
44    ) -> CapabilityResult<AgentInput> {
45        Ok(input)
46    }
47
48    /// Context-aware run-input preparation hook.
49    async fn prepare_run_input_with_context(
50        &self,
51        state: &mut AgentRunState,
52        _context: &mut AgentContext,
53        input: AgentInput,
54    ) -> CapabilityResult<AgentInput> {
55        self.prepare_run_input(state, input).await
56    }
57
58    /// Called after message history is assembled and before provider-bound preparation/model call.
59    ///
60    /// Mutations from this hook are captured in canonical session history. Use this hook
61    /// for model-visible context that must remain part of future request prefixes; use
62    /// `prepare_provider_messages` for provider-only transient rewrites. The runtime
63    /// context injector is implemented as a built-in capability in this canonical pipeline.
64    async fn prepare_model_messages(
65        &self,
66        _state: &mut AgentRunState,
67        messages: Vec<ModelMessage>,
68    ) -> CapabilityResult<Vec<ModelMessage>> {
69        Ok(messages)
70    }
71
72    /// Context-aware model-message preparation hook.
73    ///
74    /// Mutations from this hook are captured in canonical session history. Use this hook
75    /// for model-visible context that must remain part of future request prefixes; use
76    /// `prepare_provider_messages_with_context` for provider-only transient rewrites.
77    async fn prepare_model_messages_with_context(
78        &self,
79        state: &mut AgentRunState,
80        _context: &mut AgentContext,
81        messages: Vec<ModelMessage>,
82    ) -> CapabilityResult<Vec<ModelMessage>> {
83        self.prepare_model_messages(state, messages).await
84    }
85
86    /// Called after canonical model-message capabilities and durable history capture, before the model call.
87    ///
88    /// Mutations from this hook are provider-bound only and are not copied back into the
89    /// session message history.
90    async fn prepare_provider_messages(
91        &self,
92        _state: &mut AgentRunState,
93        messages: Vec<ModelMessage>,
94    ) -> CapabilityResult<Vec<ModelMessage>> {
95        Ok(messages)
96    }
97
98    /// Context-aware provider-bound message preparation hook.
99    async fn prepare_provider_messages_with_context(
100        &self,
101        state: &mut AgentRunState,
102        _context: &mut AgentContext,
103        messages: Vec<ModelMessage>,
104    ) -> CapabilityResult<Vec<ModelMessage>> {
105        self.prepare_provider_messages(state, messages).await
106    }
107
108    /// Called after tool definitions are collected and before request parameters are finalized.
109    async fn prepare_tools(
110        &self,
111        _state: &AgentRunState,
112        tools: Vec<ToolDefinition>,
113    ) -> CapabilityResult<Vec<ToolDefinition>> {
114        Ok(tools)
115    }
116
117    /// Context-aware prepare-tools hook.
118    async fn prepare_tools_with_context(
119        &self,
120        state: &AgentRunState,
121        _context: &AgentContext,
122        tools: Vec<ToolDefinition>,
123    ) -> CapabilityResult<Vec<ToolDefinition>> {
124        self.prepare_tools(state, tools).await
125    }
126
127    /// Called after a request is prepared and before the model call.
128    async fn before_model_request(
129        &self,
130        _state: &mut AgentRunState,
131        _request: &mut ModelRequest,
132        _settings: &mut Option<ModelSettings>,
133    ) -> CapabilityResult<()> {
134        Ok(())
135    }
136
137    /// Context-aware before-model-request hook.
138    async fn before_model_request_with_context(
139        &self,
140        state: &mut AgentRunState,
141        context: &mut AgentContext,
142        request: &mut ModelRequest,
143        settings: &mut Option<ModelSettings>,
144    ) -> CapabilityResult<()> {
145        let _ = context;
146        self.before_model_request(state, request, settings).await
147    }
148
149    /// Called after a model response is received.
150    async fn after_model_response(
151        &self,
152        _state: &mut AgentRunState,
153        _response: &mut ModelResponse,
154    ) -> CapabilityResult<()> {
155        Ok(())
156    }
157
158    /// Called before a tool call is executed.
159    async fn before_tool_execution(
160        &self,
161        _state: &mut AgentRunState,
162        _tool_context: &mut ToolContext,
163        _call: &ToolCallPart,
164    ) -> CapabilityResult<()> {
165        Ok(())
166    }
167
168    /// Context-aware before-tool-execution hook.
169    async fn before_tool_execution_with_context(
170        &self,
171        state: &mut AgentRunState,
172        _context: &mut AgentContext,
173        tool_context: &mut ToolContext,
174        call: &ToolCallPart,
175    ) -> CapabilityResult<()> {
176        self.before_tool_execution(state, tool_context, call).await
177    }
178
179    /// Called after a tool result is produced and before it is applied to run state.
180    async fn after_tool_result(
181        &self,
182        _state: &mut AgentRunState,
183        _call: &ToolCallPart,
184        _tool_return: &mut ToolReturnPart,
185    ) -> CapabilityResult<()> {
186        Ok(())
187    }
188
189    /// Context-aware after-tool-result hook.
190    async fn after_tool_result_with_context(
191        &self,
192        state: &mut AgentRunState,
193        _context: &mut AgentContext,
194        call: &ToolCallPart,
195        tool_return: &mut ToolReturnPart,
196    ) -> CapabilityResult<()> {
197        self.after_tool_result(state, call, tool_return).await
198    }
199
200    /// Context-aware after-model-response hook.
201    async fn after_model_response_with_context(
202        &self,
203        state: &mut AgentRunState,
204        context: &mut AgentContext,
205        response: &mut ModelResponse,
206    ) -> CapabilityResult<()> {
207        let _ = context;
208        self.after_model_response(state, response).await
209    }
210
211    /// Called before final output validation begins.
212    async fn before_output_validation(
213        &self,
214        _state: &mut AgentRunState,
215        _output: &str,
216    ) -> CapabilityResult<()> {
217        Ok(())
218    }
219
220    /// Context-aware before-output-validation hook.
221    async fn before_output_validation_with_context(
222        &self,
223        state: &mut AgentRunState,
224        _context: &mut AgentContext,
225        output: &str,
226    ) -> CapabilityResult<()> {
227        self.before_output_validation(state, output).await
228    }
229
230    /// Called after output text is selected and before finalization.
231    async fn validate_output(
232        &self,
233        _state: &mut AgentRunState,
234        _output: &str,
235    ) -> CapabilityResult<()> {
236        Ok(())
237    }
238
239    /// Context-aware output validation hook.
240    async fn validate_output_with_context(
241        &self,
242        state: &mut AgentRunState,
243        context: &mut AgentContext,
244        output: &str,
245    ) -> CapabilityResult<()> {
246        let _ = context;
247        self.validate_output(state, output).await
248    }
249
250    /// Called after output validation accepts the output.
251    async fn after_output_validation(
252        &self,
253        _state: &mut AgentRunState,
254        _output: &str,
255    ) -> CapabilityResult<()> {
256        Ok(())
257    }
258
259    /// Context-aware after-output-validation hook.
260    async fn after_output_validation_with_context(
261        &self,
262        state: &mut AgentRunState,
263        _context: &mut AgentContext,
264        output: &str,
265    ) -> CapabilityResult<()> {
266        self.after_output_validation(state, output).await
267    }
268
269    /// Called after an executor checkpoint is emitted.
270    async fn on_checkpoint(
271        &self,
272        _state: &AgentRunState,
273        _checkpoint: &AgentCheckpoint,
274    ) -> CapabilityResult<()> {
275        Ok(())
276    }
277
278    /// Context-aware checkpoint hook.
279    async fn on_checkpoint_with_context(
280        &self,
281        state: &AgentRunState,
282        _context: &AgentContext,
283        checkpoint: &AgentCheckpoint,
284    ) -> CapabilityResult<()> {
285        self.on_checkpoint(state, checkpoint).await
286    }
287
288    /// Called when semantic retry is scheduled.
289    async fn on_retry(
290        &self,
291        _state: &mut AgentRunState,
292        _kind: RetryEventKind,
293        _retries: usize,
294        _message: &str,
295    ) -> CapabilityResult<()> {
296        Ok(())
297    }
298
299    /// Called after a stream event is recorded.
300    async fn on_stream_event(
301        &self,
302        _state: &AgentRunState,
303        _event: &AgentStreamRecord,
304    ) -> CapabilityResult<()> {
305        Ok(())
306    }
307
308    /// Context-aware stream observer hook.
309    async fn on_stream_event_with_context(
310        &self,
311        state: &AgentRunState,
312        _context: &AgentContext,
313        event: &AgentStreamRecord,
314    ) -> CapabilityResult<()> {
315        self.on_stream_event(state, event).await
316    }
317
318    /// Context-aware retry hook.
319    async fn on_retry_with_context(
320        &self,
321        state: &mut AgentRunState,
322        _context: &mut AgentContext,
323        kind: RetryEventKind,
324        retries: usize,
325        message: &str,
326    ) -> CapabilityResult<()> {
327        self.on_retry(state, kind, retries, message).await
328    }
329
330    /// Called when a run completes.
331    async fn on_run_complete(&self, _state: &mut AgentRunState) -> CapabilityResult<()> {
332        Ok(())
333    }
334
335    /// Context-aware run-complete hook.
336    async fn on_run_complete_with_context(
337        &self,
338        state: &mut AgentRunState,
339        context: &mut AgentContext,
340    ) -> CapabilityResult<()> {
341        let _ = context;
342        self.on_run_complete(state).await
343    }
344}
starweaver_runtime/capability/hooks.rs

starweaver_runtime/capability/
hooks.rs
