agent_base/engine/
max_turns_nudge.rs

1use async_trait::async_trait;
2
3use crate::types::AgentResult;
4
5use super::middleware::{Middleware, PreLlmCtx};
6
7/// Configuration for the max-turns nudge middleware.
8#[derive(Clone, Debug)]
9pub struct MaxTurnsNudgeConfig {
10    /// When remaining turns <= this threshold, inject a nudge message.
11    /// Set to 0 to disable the nudge.
12    pub threshold: u32,
13    /// The nudge message to inject as a user turn.
14    pub message: String,
15}
16
17impl Default for MaxTurnsNudgeConfig {
18    fn default() -> Self {
19        Self {
20            threshold: 3,
21            message: "You are approaching the maximum number of turns. \
22                Please wrap up your current work and provide a final answer. \
23                Summarize what you've accomplished and any remaining tasks."
24                .to_string(),
25        }
26    }
27}
28
29/// Middleware that injects a nudge message when approaching the max turns limit.
30///
31/// This is a **soft intervention**: the LLM can choose to continue working,
32/// but the nudge encourages it to wrap up gracefully before hitting the hard
33/// limit (which causes an error).
34///
35/// # How it works
36/// - On each LLM call, checks `remaining_turns = max_turns - turn_count`
37/// - If `remaining_turns <= threshold`, injects the nudge as a user message
38/// - The LLM sees the nudge **before** generating its response (this turn)
39///
40/// # Usage
41/// ```ignore
42/// use agent_base::engine::{AgentBuilder, max_turns_nudge::{MaxTurnsNudgeMiddleware, MaxTurnsNudgeConfig}};
43///
44/// let runtime = AgentBuilder::new(client)
45///     .middleware(MaxTurnsNudgeMiddleware::new(MaxTurnsNudgeConfig {
46///         threshold: 3,
47///         message: "Please wrap up soon.".to_string(),
48///     }))
49///     .build()?;
50/// ```
51pub struct MaxTurnsNudgeMiddleware {
52    config: MaxTurnsNudgeConfig,
53}
54
55impl MaxTurnsNudgeMiddleware {
56    pub fn new(config: MaxTurnsNudgeConfig) -> Self {
57        Self { config }
58    }
59}
60
61#[async_trait]
62impl Middleware for MaxTurnsNudgeMiddleware {
63    async fn on_pre_llm(&self, ctx: &mut PreLlmCtx) -> AgentResult<()> {
64        if self.config.threshold == 0 {
65            return Ok(());
66        }
67
68        let remaining_turns = ctx.max_turns.saturating_sub(ctx.turn_count);
69        if remaining_turns <= self.config.threshold {
70            tracing::info!(
71                session_id = ctx.session_id.id,
72                turn = ctx.turn_count,
73                remaining_turns,
74                max_turns = ctx.max_turns,
75                "max turns nudge: injecting nudge message"
76            );
77            ctx.messages.push(crate::types::ChatMessage::User {
78                content: self.config.message.clone(),
79                images: Vec::new(),
80                ephemeral: false,
81            });
82        }
83
84        Ok(())
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::types::SessionId;
92
93    fn make_ctx(turn_count: u32, max_turns: u32) -> PreLlmCtx {
94        PreLlmCtx {
95            session_id: SessionId::new(1),
96            messages: vec![],
97            tools: vec![],
98            emit_fn: None,
99            turn_count,
100            max_turns,
101        }
102    }
103
104    #[tokio::test]
105    async fn test_nudge_not_injected_when_far_from_limit() {
106        let mw = MaxTurnsNudgeMiddleware::new(MaxTurnsNudgeConfig {
107            threshold: 3,
108            message: "nudge".to_string(),
109        });
110        let mut ctx = make_ctx(1, 10);
111        mw.on_pre_llm(&mut ctx).await.unwrap();
112        assert!(ctx.messages.is_empty());
113    }
114
115    #[tokio::test]
116    async fn test_nudge_injected_at_threshold() {
117        let mw = MaxTurnsNudgeMiddleware::new(MaxTurnsNudgeConfig {
118            threshold: 3,
119            message: "nudge".to_string(),
120        });
121        let mut ctx = make_ctx(8, 10); // remaining = 2 <= 3
122        mw.on_pre_llm(&mut ctx).await.unwrap();
123        assert_eq!(ctx.messages.len(), 1);
124        match &ctx.messages[0] {
125            crate::types::ChatMessage::User { content, .. } => assert_eq!(content, "nudge"),
126            _ => panic!("Expected User message"),
127        }
128    }
129
130    #[tokio::test]
131    async fn test_nudge_injected_at_exact_threshold() {
132        let mw = MaxTurnsNudgeMiddleware::new(MaxTurnsNudgeConfig {
133            threshold: 3,
134            message: "nudge".to_string(),
135        });
136        let mut ctx = make_ctx(7, 10); // remaining = 3 <= 3
137        mw.on_pre_llm(&mut ctx).await.unwrap();
138        assert_eq!(ctx.messages.len(), 1);
139    }
140
141    #[tokio::test]
142    async fn test_nudge_disabled_when_threshold_zero() {
143        let mw = MaxTurnsNudgeMiddleware::new(MaxTurnsNudgeConfig {
144            threshold: 0,
145            message: "nudge".to_string(),
146        });
147        let mut ctx = make_ctx(10, 10); // remaining = 0
148        mw.on_pre_llm(&mut ctx).await.unwrap();
149        assert!(ctx.messages.is_empty());
150    }
151}
agent_base/engine/max_turns_nudge.rs

agent_base/engine/
max_turns_nudge.rs
