agent_base/engine/
middleware.rs

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

agent_base/engine/
middleware.rs
