refactor(hook): add lifecycle hook system (#2284)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: laststylebender14 <ran.mahadik14@gmail.com>
Co-authored-by: laststylebender <43403528+laststylebender14@users.noreply.github.com>

Author: 3 people

crates/forge_app/src/orch.rs

@@ -27,6 +27,7 @@ pub struct Orchestrator<S> {
     agent: Agent,
     event: Event,
     error_tracker: ToolErrorTracker,
+    hook: Arc<Hook>,
 }
 
 impl<S: AgentService> Orchestrator<S> {
@@ -47,6 +48,7 @@ impl<S: AgentService> Orchestrator<S> {
             tool_definitions: Default::default(),
             models: Default::default(),
             error_tracker: Default::default(),
+            hook: Arc::new(Hook::default()),
         }
     }
 
@@ -58,11 +60,10 @@ impl<S: AgentService> Orchestrator<S> {
     // Helper function to get all tool results from a vector of tool calls
     #[async_recursion]
     async fn execute_tool_calls<'a>(
-        &self,
+        &mut self,
         tool_calls: &[ToolCallFull],
         tool_context: &ToolCallContext,
     ) -> anyhow::Result<Vec<(ToolCallFull, ToolResult)>> {
-        let agent = &self.agent;
         // Always process tool calls sequentially
         let mut tool_call_records = Vec::with_capacity(tool_calls.len());
 
@@ -80,22 +81,42 @@ impl<S: AgentService> Orchestrator<S> {
                     .await?;
             }
 
+            // Fire the ToolcallStart lifecycle event
+            let toolcall_start_event = LifecycleEvent::ToolcallStart(EventData::new(
+                self.agent.clone(),
+                self.agent.model.clone(),
+                ToolcallStartPayload::new(tool_call.clone()),
+            ));
+            self.hook
+                .handle(&toolcall_start_event, &mut self.conversation)
+                .await?;
+
             // Execute the tool
             let tool_result = self
                 .services
-                .call(agent, tool_context, tool_call.clone())
+                .call(&self.agent, tool_context, tool_call.clone())
                 .await;
 
             if tool_result.is_error() {
                 warn!(
-                    agent_id = %agent.id,
+                    agent_id = %self.agent.id,
                     name = %tool_call.name,
                     arguments = %tool_call.arguments.to_owned().into_string(),
                     output = ?tool_result.output,
                     "Tool call failed",
                 );
             }
 
+            // Fire the ToolcallEnd lifecycle event
+            let toolcall_end_event = LifecycleEvent::ToolcallEnd(EventData::new(
+                self.agent.clone(),
+                self.agent.model.clone(),
+                ToolcallEndPayload::new(tool_result.clone()),
+            ));
+            self.hook
+                .handle(&toolcall_end_event, &mut self.conversation)
+                .await?;
+
             // Send the end notification for system tools and not agent as a tool
             if is_system_tool {
                 self.send(ChatResponse::ToolCallEnd(tool_result.clone()))
@@ -118,11 +139,10 @@ impl<S: AgentService> Orchestrator<S> {
 
     // Returns if agent supports tool or not.
     fn is_tool_supported(&self) -> anyhow::Result<bool> {
-        let agent = &self.agent;
-        let model_id = &agent.model;
+        let model_id = &self.agent.model;
 
         // Check if at agent level tool support is defined
-        let tool_supported = match agent.tool_supported {
+        let tool_supported = match self.agent.tool_supported {
             Some(tool_supported) => tool_supported,
             None => {
                 // If not defined at agent level, check model level
@@ -135,7 +155,7 @@ impl<S: AgentService> Orchestrator<S> {
         };
 
         debug!(
-            agent_id = %agent.id,
+            agent_id = %self.agent.id,
             model_id = %model_id,
             tool_supported,
             "Tool support check"
@@ -172,16 +192,15 @@ impl<S: AgentService> Orchestrator<S> {
     }
     /// Checks if compaction is needed and performs it if necessary
     fn check_and_compact(&self, context: &Context) -> anyhow::Result<Option<Context>> {
-        let agent = &self.agent;
         // Estimate token count for compaction decision
         let token_count = context.token_count();
-        if agent.compact.should_compact(context, *token_count) {
-            info!(agent_id = %agent.id, "Compaction needed");
-            Compactor::new(agent.compact.clone(), self.environment.clone())
+        if self.agent.compact.should_compact(context, *token_count) {
+            info!(agent_id = %self.agent.id, "Compaction needed");
+            Compactor::new(self.agent.compact.clone(), self.environment.clone())
                 .compact(context.clone(), false)
                 .map(Some)
         } else {
-            debug!(agent_id = %agent.id, "Compaction not needed");
+            debug!(agent_id = %self.agent.id, "Compaction not needed");
             Ok(None)
         }
     }
@@ -207,8 +226,15 @@ impl<S: AgentService> Orchestrator<S> {
 
         let mut context = self.conversation.context.clone().unwrap_or_default();
 
-        // Create agent reference for the rest of the method
-        let agent = &self.agent;
+        // Fire the Start lifecycle event
+        let start_event = LifecycleEvent::Start(EventData::new(
+            self.agent.clone(),
+            model_id.clone(),
+            StartPayload,
+        ));
+        self.hook
+            .handle(&start_event, &mut self.conversation)
+            .await?;
 
         // Signals that the loop should suspend (task may or may not be completed)
         let mut should_yield = false;
@@ -219,8 +245,7 @@ impl<S: AgentService> Orchestrator<S> {
         let mut request_count = 0;
 
         // Retrieve the number of requests allowed per tick.
-        let max_requests_per_turn = agent.max_requests_per_turn;
-
+        let max_requests_per_turn = self.agent.max_requests_per_turn;
         let tool_context =
             ToolCallContext::new(self.conversation.metrics.clone()).sender(self.sender.clone());
 
@@ -233,12 +258,22 @@ impl<S: AgentService> Orchestrator<S> {
             self.conversation.context = Some(context.clone());
             self.services.update(self.conversation.clone()).await?;
 
+            // Fire the Request lifecycle event
+            let request_event = LifecycleEvent::Request(EventData::new(
+                self.agent.clone(),
+                model_id.clone(),
+                RequestPayload::new(request_count),
+            ));
+            self.hook
+                .handle(&request_event, &mut self.conversation)
+                .await?;
+
             let message = crate::retry::retry_with_config(
                 &self.environment.retry_config,
                 || self.execute_chat_turn(&model_id, context.clone(), context.is_reasoning_supported()),
                 self.sender.as_ref().map(|sender| {
                     let sender = sender.clone();
-                    let agent_id = agent.id.clone();
+                    let agent_id = self.agent.id.clone();
                     let model_id = model_id.clone();
                     move |error: &anyhow::Error, duration: Duration| {
                         let root_cause = error.root_cause();
@@ -252,15 +287,25 @@ impl<S: AgentService> Orchestrator<S> {
                 }),
             ).await?;
 
+            // Fire the Response lifecycle event
+            let response_event = LifecycleEvent::Response(EventData::new(
+                self.agent.clone(),
+                model_id.clone(),
+                ResponsePayload::new(message.clone()),
+            ));
+            self.hook
+                .handle(&response_event, &mut self.conversation)
+                .await?;
+
             // TODO: Add a unit test in orch spec, to guarantee that compaction is
-            // triggered after receiving the response Trigger compaction after
+            // triggered after receiving the response
             // making a request NOTE: Ideally compaction should be implemented
             // as a transformer
             if let Some(c_context) = self.check_and_compact(&context)? {
-                info!(agent_id = %agent.id, "Using compacted context from execution");
+                info!(agent_id = %self.agent.id, "Using compacted context from execution");
                 context = c_context;
             } else {
-                debug!(agent_id = %agent.id, "No compaction was needed");
+                debug!(agent_id = %self.agent.id, "No compaction was needed");
             }
 
             info!(
@@ -274,8 +319,7 @@ impl<S: AgentService> Orchestrator<S> {
                 "Processing usage information"
             );
 
-            debug!(agent_id = %agent.id, tool_call_count = message.tool_calls.len(), "Tool call count");
-
+            debug!(agent_id = %self.agent.id, tool_call_count = message.tool_calls.len(), "Tool call count");
             // Turn is completed, if finish_reason is 'stop'. Gemini models return stop as
             // finish reason with tool calls.
             is_complete =
@@ -340,7 +384,7 @@ impl<S: AgentService> Orchestrator<S> {
                 // Check if agent has reached the maximum request per turn limit
                 if request_count >= max_request_allowed {
                     warn!(
-                        agent_id = %agent.id,
+                        agent_id = %self.agent.id,
                         model_id = %model_id,
                         request_count,
                         max_request_allowed,
@@ -372,6 +416,18 @@ impl<S: AgentService> Orchestrator<S> {
 
         self.services.update(self.conversation.clone()).await?;
 
+        // Fire the End lifecycle event
+        self.hook
+            .handle(
+                &LifecycleEvent::End(EventData::new(
+                    self.agent.clone(),
+                    model_id.clone(),
+                    EndPayload,
+                )),
+                &mut self.conversation,
+            )
+            .await?;
+
         // Signal Task Completion
         if is_complete {
             self.send(ChatResponse::TaskComplete).await?;

crates/forge_domain/src/agent.rs

@@ -9,7 +9,7 @@ use crate::{
 
 /// Runtime agent representation with required model and provider
 /// Created by converting AgentDefinition with resolved defaults
-#[derive(Debug, Clone, Setters)]
+#[derive(Debug, Clone, PartialEq, Setters)]
 #[setters(strip_option, into)]
 pub struct Agent {
     /// Flag to enable/disable tool support for this agent.

crates/forge_domain/src/chat_response.rs

@@ -79,7 +79,7 @@ impl ChatResponse {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum InterruptionReason {
     MaxToolFailurePerTurnLimitReached {
         limit: u64,

crates/forge_domain/src/env.rs

@@ -13,7 +13,7 @@ const VERSION: &str = match option_env!("APP_VERSION") {
     None => env!("CARGO_PKG_VERSION"),
 };
 
-#[derive(Debug, Setters, Clone, Serialize, Deserialize, fake::Dummy)]
+#[derive(Debug, Setters, Clone, PartialEq, Serialize, Deserialize, fake::Dummy)]
 #[serde(rename_all = "camelCase")]
 #[setters(strip_option)]
 /// Represents the environment in which the application is running.

crates/forge_domain/src/event.rs

@@ -84,15 +84,15 @@ impl EventValue {
 #[serde(transparent)]
 pub struct UserPrompt(String);
 
-#[derive(Clone, Serialize, Deserialize, Debug, Setters)]
+#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Setters)]
 pub struct EventContext {
     event: EventContextValue,
     suggestions: Vec<String>,
     variables: HashMap<String, Value>,
     current_date: String,
 }
 
-#[derive(Clone, Serialize, Deserialize, Debug, Setters)]
+#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Setters)]
 pub struct EventContextValue {
     pub name: String,
     pub value: String,