agent_base/engine/
middleware.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::Value;
5use tokio::sync::broadcast;
6
7use crate::types::{AgentResult, AgentEvent, ChatMessage, SessionId};
8
9#[derive(Clone)]
10pub struct UserMessageCtx {
11    pub session_id: SessionId,
12    pub user_input: String,
13    pub event_bus: broadcast::Sender<AgentEvent>,
14}
15
16#[derive(Clone)]
17pub struct PreLlmCtx {
18    pub session_id: SessionId,
19    pub messages: Vec<ChatMessage>,
20    pub tools: Vec<Value>,
21    pub event_bus: broadcast::Sender<AgentEvent>,
22}
23
24#[derive(Clone)]
25pub struct PostLlmCtx {
26    pub session_id: SessionId,
27    pub full_text: String,
28    pub is_tool_call: bool,
29    pub tool_calls: Vec<(String, String, String)>,
30    pub event_bus: broadcast::Sender<AgentEvent>,
31    pub available_tools: Vec<String>,
32    pub turn_count: u32,
33    pub total_tool_calls: usize,
34    pub skip_push: bool,
35    pub follow_up_message: Option<String>,
36}
37
38#[async_trait]
39pub trait Middleware: Send + Sync {
40    async fn on_user_message(&self, _ctx: &mut UserMessageCtx) -> AgentResult<()> {
41        Ok(())
42    }
43
44    async fn on_pre_llm(&self, _ctx: &mut PreLlmCtx) -> AgentResult<()> {
45        Ok(())
46    }
47
48    async fn on_post_llm(&self, _ctx: &mut PostLlmCtx) -> AgentResult<()> {
49        Ok(())
50    }
51}
52
53pub(crate) type MiddlewareRef = Arc<dyn Middleware>;
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use crate::types::SessionId;
59
60    #[test]
61    fn test_post_llm_ctx_new_fields_defaults() {
62        let (tx, _rx) = broadcast::channel(1);
63        let ctx = PostLlmCtx {
64            session_id: SessionId { id: 1, external_id: None },
65            full_text: "test".to_string(),
66            is_tool_call: false,
67            tool_calls: vec![],
68            event_bus: tx,
69            available_tools: vec![],
70            turn_count: 0,
71            total_tool_calls: 0,
72            skip_push: false,
73            follow_up_message: None,
74        };
75        assert!(ctx.available_tools.is_empty());
76        assert_eq!(ctx.turn_count, 0);
77        assert_eq!(ctx.total_tool_calls, 0);
78        assert!(!ctx.skip_push);
79        assert!(ctx.follow_up_message.is_none());
80    }
81
82    #[test]
83    fn test_post_llm_ctx_skip_push_follow_up_set() {
84        let (tx, _rx) = broadcast::channel(1);
85        let ctx = PostLlmCtx {
86            session_id: SessionId { id: 2, external_id: None },
87            full_text: "I will execute...".to_string(),
88            is_tool_call: false,
89            tool_calls: vec![],
90            event_bus: tx,
91            available_tools: vec!["echo".to_string()],
92            turn_count: 1,
93            total_tool_calls: 0,
94            skip_push: true,
95            follow_up_message: Some("Please call tools now.".to_string()),
96        };
97        assert!(ctx.skip_push);
98        assert_eq!(ctx.follow_up_message, Some("Please call tools now.".to_string()));
99        assert_eq!(ctx.available_tools, vec!["echo".to_string()]);
100        assert_eq!(ctx.total_tool_calls, 0);
101    }
102
103    #[test]
104    fn test_post_llm_ctx_clone_preserves_new_fields() {
105        let (tx, _rx) = broadcast::channel(1);
106        let ctx = PostLlmCtx {
107            session_id: SessionId { id: 3, external_id: None },
108            full_text: "hello".to_string(),
109            is_tool_call: false,
110            tool_calls: vec![],
111            event_bus: tx,
112            available_tools: vec!["add".to_string(), "subtract".to_string()],
113            turn_count: 5,
114            total_tool_calls: 3,
115            skip_push: true,
116            follow_up_message: Some("nudge".to_string()),
117        };
118        let cloned = ctx.clone();
119        assert_eq!(cloned.available_tools, vec!["add", "subtract"]);
120        assert_eq!(cloned.turn_count, 5);
121        assert_eq!(cloned.total_tool_calls, 3);
122        assert!(cloned.skip_push);
123        assert_eq!(cloned.follow_up_message, Some("nudge".to_string()));
124    }
125}
agent_base/engine/middleware.rs

agent_base/engine/
middleware.rs
