vv-agent 0.6.2

VectorVein agent runtime, SDK, CLI, tools, and workspace backends
Documentation
use std::collections::BTreeMap;
use std::sync::Arc;

use serde_json::Value;

use crate::tools::ToolContext;
use crate::types::{AgentTask, LLMResponse, Message, ToolCall, ToolExecutionResult};

use super::events::{
    AfterLlmEvent, AfterToolCallEvent, BeforeLlmEvent, BeforeMemoryCompactEvent,
    BeforeToolCallEvent,
};
use super::traits::RuntimeHook;

#[derive(Default)]
pub struct RuntimeHookManager {
    hooks: Vec<Arc<dyn RuntimeHook>>,
}

impl RuntimeHookManager {
    pub fn new(hooks: Vec<Arc<dyn RuntimeHook>>) -> Self {
        Self { hooks }
    }

    pub fn is_empty(&self) -> bool {
        self.hooks.is_empty()
    }

    pub fn has_hooks(&self) -> bool {
        !self.hooks.is_empty()
    }

    pub fn apply_before_memory_compact(
        &self,
        task: &AgentTask,
        cycle_index: u32,
        messages: Vec<Message>,
        shared_state: &BTreeMap<String, Value>,
    ) -> Vec<Message> {
        let mut current_messages = messages;
        for hook in &self.hooks {
            if let Some(patched) = hook.before_memory_compact(BeforeMemoryCompactEvent {
                task,
                cycle_index,
                messages: &current_messages,
                shared_state,
            }) {
                current_messages = patched;
            }
        }
        current_messages
    }

    pub fn apply_before_llm(
        &self,
        task: &AgentTask,
        cycle_index: u32,
        messages: Vec<Message>,
        tool_schemas: Vec<Value>,
        shared_state: &BTreeMap<String, Value>,
    ) -> (Vec<Message>, Vec<Value>) {
        let mut current_messages = messages;
        let mut current_tool_schemas = tool_schemas;
        for hook in &self.hooks {
            let patch = hook.before_llm(BeforeLlmEvent {
                task,
                cycle_index,
                messages: &current_messages,
                tool_schemas: &current_tool_schemas,
                shared_state,
            });
            let Some(patch) = patch else {
                continue;
            };
            if let Some(messages) = patch.messages {
                current_messages = messages;
            }
            if let Some(tool_schemas) = patch.tool_schemas {
                current_tool_schemas = tool_schemas;
            }
        }
        (current_messages, current_tool_schemas)
    }

    pub fn apply_after_llm(
        &self,
        task: &AgentTask,
        cycle_index: u32,
        messages: &[Message],
        tool_schemas: &[Value],
        response: LLMResponse,
        shared_state: &BTreeMap<String, Value>,
    ) -> LLMResponse {
        let mut current = response;
        for hook in &self.hooks {
            if let Some(patched) = hook.after_llm(AfterLlmEvent {
                task,
                cycle_index,
                messages,
                tool_schemas,
                response: &current,
                shared_state,
            }) {
                current = patched;
            }
        }
        current
    }

    pub fn apply_before_tool_call(
        &self,
        task: &AgentTask,
        cycle_index: u32,
        call: ToolCall,
        context: &ToolContext,
    ) -> (ToolCall, Option<ToolExecutionResult>) {
        let mut current_call = call;
        let mut short_circuit = None;
        for hook in &self.hooks {
            let patch = hook.before_tool_call(BeforeToolCallEvent {
                task,
                cycle_index,
                call: &current_call,
                context,
            });
            let Some(patch) = patch else {
                continue;
            };
            if let Some(call) = patch.call {
                current_call = call;
            }
            if let Some(result) = patch.result {
                short_circuit = Some(result);
                break;
            }
        }
        (current_call, short_circuit)
    }

    pub fn apply_after_tool_call(
        &self,
        task: &AgentTask,
        cycle_index: u32,
        call: &ToolCall,
        context: &ToolContext,
        result: ToolExecutionResult,
    ) -> ToolExecutionResult {
        let mut current = result;
        for hook in &self.hooks {
            if let Some(patched) = hook.after_tool_call(AfterToolCallEvent {
                task,
                cycle_index,
                call,
                context,
                result: &current,
            }) {
                current = patched;
            }
        }
        current
    }
}