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, FinishReason, SessionId, UserEvent};
7
8#[derive(Clone)]
9pub struct UserMessageCtx {
10    pub session_id: SessionId,
11    pub user_input: String,
12}
13
14pub struct PreLlmCtx {
15    pub session_id: SessionId,
16    pub messages: Vec<ChatMessage>,
17    pub tools: Vec<Value>,
18    /// Unified emit function: sends [`UserEvent`] to both the renderer (real-time)
19    /// and the event bus (persistence).  Set by the react loop; middleware calls
20    /// [`emit()`](Self::emit) which delegates to this closure.
21    ///
22    /// The closure captures the renderer callback and the event bus, so middleware
23    /// never needs to know about broadcast channels or drain mechanics.
24    pub emit_fn: Option<Box<dyn Fn(UserEvent) + Send + Sync>>,
25}
26
27impl PreLlmCtx {
28    /// Emit a [`UserEvent`] to both the renderer and the event bus.
29    ///
30    /// This is the single entry point for middleware to send events.  The event
31    /// is delivered to the renderer callback (real-time display) and to the
32    /// event bus (persistence, event_log, checkpoint) in one call.
33    ///
34    /// Does nothing if the emit function was not set (e.g. in tests).
35    /// Panics inside the emit function are caught and logged as warnings.
36    pub fn emit(&self, event: UserEvent) {
37        if let Some(ref f) = self.emit_fn
38            && let Err(e) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
39                f(event);
40            }))
41        {
42            tracing::warn!("emit failed: {:?}", e);
43        }
44    }
45}
46
47#[derive(Clone)]
48pub struct PostLlmCtx {
49    pub session_id: SessionId,
50    pub full_text: String,
51    pub is_tool_call: bool,
52    pub tool_calls: Vec<(String, String, String)>,
53    pub available_tools: Vec<String>,
54    pub turn_count: u32,
55    pub total_tool_calls: usize,
56    /// Number of tool-enforcement nudges issued in the current turn.
57    /// Read from session; middleware may increment this to track nudge attempts.
58    pub nudge_count: usize,
59    /// Number of tool calls already executed in the current turn.
60    /// Used by `TurnToolLimitMiddleware` to enforce per-turn tool call limits.
61    pub turn_tool_calls: usize,
62    pub skip_push: bool,
63    pub follow_up_message: Option<String>,
64    /// Semantic finish reason from the LLM (Stop / ToolUse / Truncated / Other).
65    /// Middleware can inspect this to implement custom continuation logic
66    /// (e.g. auto-continue on truncation).
67    pub finish_reason: FinishReason,
68}
69
70#[async_trait]
71pub trait Middleware: Send + Sync {
72    async fn on_user_message(&self, _ctx: &mut UserMessageCtx) -> AgentResult<()> {
73        Ok(())
74    }
75
76    async fn on_pre_llm(&self, _ctx: &mut PreLlmCtx) -> AgentResult<()> {
77        Ok(())
78    }
79
80    async fn on_post_llm(&self, _ctx: &mut PostLlmCtx) -> AgentResult<()> {
81        Ok(())
82    }
83}
84
85pub(crate) type MiddlewareRef = Arc<dyn Middleware>;
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::types::SessionId;
91
92    #[test]
93    fn test_post_llm_ctx_new_fields_defaults() {
94        let ctx = PostLlmCtx {
95            session_id: SessionId {
96                id: 1,
97                external_id: None,
98            },
99            full_text: "test".to_string(),
100            is_tool_call: false,
101            tool_calls: vec![],
102            available_tools: vec![],
103            turn_count: 0,
104            total_tool_calls: 0,
105            nudge_count: 0,
106            turn_tool_calls: 0,
107            skip_push: false,
108            follow_up_message: None,
109            finish_reason: FinishReason::Stop,
110        };
111        assert!(ctx.available_tools.is_empty());
112        assert_eq!(ctx.turn_count, 0);
113        assert_eq!(ctx.total_tool_calls, 0);
114        assert!(!ctx.skip_push);
115        assert!(ctx.follow_up_message.is_none());
116        assert_eq!(ctx.finish_reason, FinishReason::Stop);
117    }
118
119    #[test]
120    fn test_post_llm_ctx_skip_push_follow_up_set() {
121        let ctx = PostLlmCtx {
122            session_id: SessionId {
123                id: 2,
124                external_id: None,
125            },
126            full_text: "I will execute...".to_string(),
127            is_tool_call: false,
128            tool_calls: vec![],
129            available_tools: vec!["echo".to_string()],
130            turn_count: 1,
131            total_tool_calls: 0,
132            nudge_count: 0,
133            turn_tool_calls: 0,
134            skip_push: true,
135            follow_up_message: Some("Please call tools now.".to_string()),
136            finish_reason: FinishReason::Stop,
137        };
138        assert!(ctx.skip_push);
139        assert_eq!(
140            ctx.follow_up_message,
141            Some("Please call tools now.".to_string())
142        );
143        assert_eq!(ctx.available_tools, vec!["echo".to_string()]);
144        assert_eq!(ctx.total_tool_calls, 0);
145    }
146
147    #[test]
148    fn test_post_llm_ctx_clone_preserves_new_fields() {
149        let ctx = PostLlmCtx {
150            session_id: SessionId {
151                id: 3,
152                external_id: None,
153            },
154            full_text: "hello".to_string(),
155            is_tool_call: false,
156            tool_calls: vec![],
157            available_tools: vec!["add".to_string(), "subtract".to_string()],
158            turn_count: 5,
159            total_tool_calls: 3,
160            nudge_count: 0,
161            turn_tool_calls: 2,
162            skip_push: true,
163            follow_up_message: Some("nudge".to_string()),
164            finish_reason: FinishReason::Truncated {
165                reason: Some("max_tokens".into()),
166            },
167        };
168        let cloned = ctx.clone();
169        assert_eq!(cloned.available_tools, vec!["add", "subtract"]);
170        assert_eq!(cloned.turn_count, 5);
171        assert_eq!(cloned.total_tool_calls, 3);
172        assert_eq!(cloned.turn_tool_calls, 2);
173        assert!(cloned.skip_push);
174        assert_eq!(cloned.follow_up_message, Some("nudge".to_string()));
175        assert!(cloned.finish_reason.is_truncated());
176    }
177}
agent_base/engine/middleware.rs

agent_base/engine/
middleware.rs
