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 {
66                id: 1,
67                external_id: None,
68            },
69            full_text: "test".to_string(),
70            is_tool_call: false,
71            tool_calls: vec![],
72            available_tools: vec![],
73            turn_count: 0,
74            total_tool_calls: 0,
75            nudge_count: 0,
76            turn_tool_calls: 0,
77            skip_push: false,
78            follow_up_message: None,
79        };
80        assert!(ctx.available_tools.is_empty());
81        assert_eq!(ctx.turn_count, 0);
82        assert_eq!(ctx.total_tool_calls, 0);
83        assert!(!ctx.skip_push);
84        assert!(ctx.follow_up_message.is_none());
85    }
86
87    #[test]
88    fn test_post_llm_ctx_skip_push_follow_up_set() {
89        let ctx = PostLlmCtx {
90            session_id: SessionId {
91                id: 2,
92                external_id: None,
93            },
94            full_text: "I will execute...".to_string(),
95            is_tool_call: false,
96            tool_calls: vec![],
97            available_tools: vec!["echo".to_string()],
98            turn_count: 1,
99            total_tool_calls: 0,
100            nudge_count: 0,
101            turn_tool_calls: 0,
102            skip_push: true,
103            follow_up_message: Some("Please call tools now.".to_string()),
104        };
105        assert!(ctx.skip_push);
106        assert_eq!(
107            ctx.follow_up_message,
108            Some("Please call tools now.".to_string())
109        );
110        assert_eq!(ctx.available_tools, vec!["echo".to_string()]);
111        assert_eq!(ctx.total_tool_calls, 0);
112    }
113
114    #[test]
115    fn test_post_llm_ctx_clone_preserves_new_fields() {
116        let ctx = PostLlmCtx {
117            session_id: SessionId {
118                id: 3,
119                external_id: None,
120            },
121            full_text: "hello".to_string(),
122            is_tool_call: false,
123            tool_calls: vec![],
124            available_tools: vec!["add".to_string(), "subtract".to_string()],
125            turn_count: 5,
126            total_tool_calls: 3,
127            nudge_count: 0,
128            turn_tool_calls: 2,
129            skip_push: true,
130            follow_up_message: Some("nudge".to_string()),
131        };
132        let cloned = ctx.clone();
133        assert_eq!(cloned.available_tools, vec!["add", "subtract"]);
134        assert_eq!(cloned.turn_count, 5);
135        assert_eq!(cloned.total_tool_calls, 3);
136        assert_eq!(cloned.turn_tool_calls, 2);
137        assert!(cloned.skip_push);
138        assert_eq!(cloned.follow_up_message, Some("nudge".to_string()));
139    }
140}
agent_base/engine/middleware.rs

agent_base/engine/
middleware.rs
