agent_base/engine/
auto_continue.rs

1use async_trait::async_trait;
2
3use crate::types::AgentResult;
4
5use super::middleware::{Middleware, PostLlmCtx};
6
7/// Middleware that automatically injects a continuation message when the LLM
8/// response is truncated by the token limit (e.g. `max_tokens` / `length`)
9/// and no tool calls are present.
10///
11/// This turns a silent truncation into an automatic retry: the react loop
12/// will push the follow-up as a user message and call the LLM again.
13///
14/// # When it triggers
15/// - `finish_reason` is `Truncated` **and**
16/// - `tool_calls` is empty (tool-call truncation is already handled by the
17///   react-loop's built-in truncation guard, which injects error results
18///   instead of executing incomplete tool calls).
19///
20/// # Usage
21/// ```ignore
22/// use agent_base::engine::{AgentBuilder, auto_continue::AutoContinueMiddleware};
23///
24/// let runtime = AgentBuilder::new(client)
25///     .middleware(AutoContinueMiddleware::new())
26///     .build()?;
27/// ```
28pub struct AutoContinueMiddleware {
29    /// The message to inject as a user turn when truncation is detected.
30    prompt: String,
31}
32
33impl AutoContinueMiddleware {
34    /// Create a new middleware with the default continuation prompt.
35    pub fn new() -> Self {
36        Self {
37            prompt: "Please continue.".to_string(),
38        }
39    }
40
41    /// Create a new middleware with a custom continuation prompt.
42    pub fn with_prompt(prompt: impl Into<String>) -> Self {
43        Self {
44            prompt: prompt.into(),
45        }
46    }
47}
48
49impl Default for AutoContinueMiddleware {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55#[async_trait]
56impl Middleware for AutoContinueMiddleware {
57    async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
58        if ctx.finish_reason.is_truncated() && ctx.tool_calls.is_empty() {
59            tracing::info!(
60                session_id = ctx.session_id.id,
61                turn = ctx.turn_count,
62                "text-only response truncated — injecting auto-continue prompt"
63            );
64            ctx.follow_up_message = Some(self.prompt.clone());
65        }
66        Ok(())
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use crate::types::{FinishReason, SessionId};
74
75    fn ctx(finish_reason: FinishReason, tool_calls: Vec<(String, String, String)>) -> PostLlmCtx {
76        PostLlmCtx {
77            session_id: SessionId::new(1),
78            full_text: "partial answer...".to_string(),
79            is_tool_call: !tool_calls.is_empty(),
80            tool_calls,
81            available_tools: vec![],
82            turn_count: 1,
83            total_tool_calls: 0,
84            nudge_count: 0,
85            turn_tool_calls: 0,
86            skip_push: false,
87            follow_up_message: None,
88            finish_reason,
89        }
90    }
91
92    #[tokio::test]
93    async fn triggers_on_truncated_text_only() {
94        let mw = AutoContinueMiddleware::new();
95        let mut c = ctx(
96            FinishReason::Truncated {
97                reason: Some("max_tokens".into()),
98            },
99            vec![],
100        );
101        mw.on_post_llm(&mut c).await.unwrap();
102        assert_eq!(c.follow_up_message, Some("Please continue.".to_string()));
103    }
104
105    #[tokio::test]
106    async fn triggers_on_truncated_no_reason() {
107        let mw = AutoContinueMiddleware::new();
108        let mut c = ctx(FinishReason::Truncated { reason: None }, vec![]);
109        mw.on_post_llm(&mut c).await.unwrap();
110        assert_eq!(c.follow_up_message, Some("Please continue.".to_string()));
111    }
112
113    #[tokio::test]
114    async fn skips_when_tool_calls_present() {
115        let mw = AutoContinueMiddleware::new();
116        let mut c = ctx(
117            FinishReason::Truncated {
118                reason: Some("length".into()),
119            },
120            vec![("id".into(), "shell".into(), "{}".into())],
121        );
122        mw.on_post_llm(&mut c).await.unwrap();
123        assert!(c.follow_up_message.is_none());
124    }
125
126    #[tokio::test]
127    async fn skips_when_stop() {
128        let mw = AutoContinueMiddleware::new();
129        let mut c = ctx(FinishReason::Stop, vec![]);
130        mw.on_post_llm(&mut c).await.unwrap();
131        assert!(c.follow_up_message.is_none());
132    }
133
134    #[tokio::test]
135    async fn skips_when_tool_use() {
136        let mw = AutoContinueMiddleware::new();
137        let mut c = ctx(FinishReason::ToolUse, vec![]);
138        mw.on_post_llm(&mut c).await.unwrap();
139        assert!(c.follow_up_message.is_none());
140    }
141
142    #[tokio::test]
143    async fn custom_prompt() {
144        let mw = AutoContinueMiddleware::with_prompt("Continue please.");
145        let mut c = ctx(
146            FinishReason::Truncated {
147                reason: Some("max_tokens".into()),
148            },
149            vec![],
150        );
151        mw.on_post_llm(&mut c).await.unwrap();
152        assert_eq!(c.follow_up_message, Some("Continue please.".to_string()));
153    }
154}
agent_base/engine/auto_continue.rs

agent_base/engine/
auto_continue.rs
