vv_agent/runtime/hooks/
manager.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use serde_json::Value;
5
6use crate::tools::ToolContext;
7use crate::types::{AgentTask, LLMResponse, Message, ToolCall, ToolExecutionResult};
8
9use super::events::{
10    AfterLlmEvent, AfterToolCallEvent, BeforeLlmEvent, BeforeMemoryCompactEvent,
11    BeforeToolCallEvent,
12};
13use super::traits::RuntimeHook;
14
15#[derive(Default)]
16pub struct RuntimeHookManager {
17    hooks: Vec<Arc<dyn RuntimeHook>>,
18}
19
20impl RuntimeHookManager {
21    pub fn new(hooks: Vec<Arc<dyn RuntimeHook>>) -> Self {
22        Self { hooks }
23    }
24
25    pub fn is_empty(&self) -> bool {
26        self.hooks.is_empty()
27    }
28
29    pub fn has_hooks(&self) -> bool {
30        !self.hooks.is_empty()
31    }
32
33    pub fn apply_before_memory_compact(
34        &self,
35        task: &AgentTask,
36        cycle_index: u32,
37        messages: Vec<Message>,
38        shared_state: &BTreeMap<String, Value>,
39    ) -> Vec<Message> {
40        let mut current_messages = messages;
41        for hook in &self.hooks {
42            if let Some(patched) = hook.before_memory_compact(BeforeMemoryCompactEvent {
43                task,
44                cycle_index,
45                messages: &current_messages,
46                shared_state,
47            }) {
48                current_messages = patched;
49            }
50        }
51        current_messages
52    }
53
54    pub fn apply_before_llm(
55        &self,
56        task: &AgentTask,
57        cycle_index: u32,
58        messages: Vec<Message>,
59        tool_schemas: Vec<Value>,
60        shared_state: &BTreeMap<String, Value>,
61    ) -> (Vec<Message>, Vec<Value>) {
62        let mut current_messages = messages;
63        let mut current_tool_schemas = tool_schemas;
64        for hook in &self.hooks {
65            let patch = hook.before_llm(BeforeLlmEvent {
66                task,
67                cycle_index,
68                messages: &current_messages,
69                tool_schemas: &current_tool_schemas,
70                shared_state,
71            });
72            let Some(patch) = patch else {
73                continue;
74            };
75            if let Some(messages) = patch.messages {
76                current_messages = messages;
77            }
78            if let Some(tool_schemas) = patch.tool_schemas {
79                current_tool_schemas = tool_schemas;
80            }
81        }
82        (current_messages, current_tool_schemas)
83    }
84
85    pub fn apply_after_llm(
86        &self,
87        task: &AgentTask,
88        cycle_index: u32,
89        messages: &[Message],
90        tool_schemas: &[Value],
91        response: LLMResponse,
92        shared_state: &BTreeMap<String, Value>,
93    ) -> LLMResponse {
94        let mut current = response;
95        for hook in &self.hooks {
96            if let Some(patched) = hook.after_llm(AfterLlmEvent {
97                task,
98                cycle_index,
99                messages,
100                tool_schemas,
101                response: &current,
102                shared_state,
103            }) {
104                current = patched;
105            }
106        }
107        current
108    }
109
110    pub fn apply_before_tool_call(
111        &self,
112        task: &AgentTask,
113        cycle_index: u32,
114        call: ToolCall,
115        context: &ToolContext,
116    ) -> (ToolCall, Option<ToolExecutionResult>) {
117        let mut current_call = call;
118        let mut short_circuit = None;
119        for hook in &self.hooks {
120            let patch = hook.before_tool_call(BeforeToolCallEvent {
121                task,
122                cycle_index,
123                call: &current_call,
124                context,
125            });
126            let Some(patch) = patch else {
127                continue;
128            };
129            if let Some(call) = patch.call {
130                current_call = call;
131            }
132            if let Some(result) = patch.result {
133                short_circuit = Some(result);
134                break;
135            }
136        }
137        (current_call, short_circuit)
138    }
139
140    pub fn apply_after_tool_call(
141        &self,
142        task: &AgentTask,
143        cycle_index: u32,
144        call: &ToolCall,
145        context: &ToolContext,
146        result: ToolExecutionResult,
147    ) -> ToolExecutionResult {
148        let mut current = result;
149        for hook in &self.hooks {
150            if let Some(patched) = hook.after_tool_call(AfterToolCallEvent {
151                task,
152                cycle_index,
153                call,
154                context,
155                result: &current,
156            }) {
157                current = patched;
158            }
159        }
160        current
161    }
162}
vv_agent/runtime/hooks/manager.rs

vv_agent/runtime/hooks/
manager.rs
