agents_runtime/
middleware.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock};
3
4use agents_core::agent::{AgentHandle, ToolHandle, ToolResponse};
5use agents_core::messaging::{
6    AgentMessage, CacheControl, MessageContent, MessageMetadata, MessageRole, ToolInvocation,
7};
8use agents_core::prompts::{
9    BASE_AGENT_PROMPT, FILESYSTEM_SYSTEM_PROMPT, TASK_SYSTEM_PROMPT, TASK_TOOL_DESCRIPTION,
10    WRITE_TODOS_SYSTEM_PROMPT,
11};
12use agents_core::state::AgentStateSnapshot;
13use agents_toolkit::{EditFileTool, LsTool, ReadFileTool, WriteFileTool, WriteTodosTool};
14use async_trait::async_trait;
15use serde::Deserialize;
16
17/// Request sent to the underlying language model. Middlewares can augment
18/// the system prompt or mutate the pending message list before the model call.
19#[derive(Debug, Clone)]
20pub struct ModelRequest {
21    pub system_prompt: String,
22    pub messages: Vec<AgentMessage>,
23}
24
25impl ModelRequest {
26    pub fn new(system_prompt: impl Into<String>, messages: Vec<AgentMessage>) -> Self {
27        Self {
28            system_prompt: system_prompt.into(),
29            messages,
30        }
31    }
32
33    pub fn append_prompt(&mut self, fragment: &str) {
34        if !fragment.is_empty() {
35            self.system_prompt.push_str("\n\n");
36            self.system_prompt.push_str(fragment);
37        }
38    }
39}
40
41/// Read/write state handle exposed to middleware implementations.
42pub struct MiddlewareContext<'a> {
43    pub request: &'a mut ModelRequest,
44    pub state: Arc<RwLock<AgentStateSnapshot>>,
45}
46
47impl<'a> MiddlewareContext<'a> {
48    pub fn with_request(
49        request: &'a mut ModelRequest,
50        state: Arc<RwLock<AgentStateSnapshot>>,
51    ) -> Self {
52        Self { request, state }
53    }
54}
55
56/// Middleware hook that can register additional tools and mutate the model request
57/// prior to execution. Mirrors the Python AgentMiddleware contracts but keeps the
58/// interface async-first for future network calls.
59#[async_trait]
60pub trait AgentMiddleware: Send + Sync {
61    /// Unique identifier for logging and diagnostics.
62    fn id(&self) -> &'static str;
63
64    /// Tools to expose when this middleware is active.
65    fn tools(&self) -> Vec<Arc<dyn ToolHandle>> {
66        Vec::new()
67    }
68
69    /// Apply middleware-specific mutations to the pending model request.
70    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()>;
71}
72
73pub struct SummarizationMiddleware {
74    pub messages_to_keep: usize,
75    pub summary_note: String,
76}
77
78impl SummarizationMiddleware {
79    pub fn new(messages_to_keep: usize, summary_note: impl Into<String>) -> Self {
80        Self {
81            messages_to_keep,
82            summary_note: summary_note.into(),
83        }
84    }
85}
86
87#[async_trait]
88impl AgentMiddleware for SummarizationMiddleware {
89    fn id(&self) -> &'static str {
90        "summarization"
91    }
92
93    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
94        if ctx.request.messages.len() > self.messages_to_keep {
95            let dropped = ctx.request.messages.len() - self.messages_to_keep;
96            let mut truncated = ctx
97                .request
98                .messages
99                .split_off(ctx.request.messages.len() - self.messages_to_keep);
100            truncated.insert(
101                0,
102                AgentMessage {
103                    role: MessageRole::System,
104                    content: MessageContent::Text(format!(
105                        "{} ({} earlier messages summarized)",
106                        self.summary_note, dropped
107                    )),
108                    metadata: None,
109                },
110            );
111            ctx.request.messages = truncated;
112        }
113        Ok(())
114    }
115}
116
117pub struct PlanningMiddleware {
118    state: Arc<RwLock<AgentStateSnapshot>>,
119}
120
121impl PlanningMiddleware {
122    pub fn new(state: Arc<RwLock<AgentStateSnapshot>>) -> Self {
123        Self { state }
124    }
125}
126
127#[async_trait]
128impl AgentMiddleware for PlanningMiddleware {
129    fn id(&self) -> &'static str {
130        "planning"
131    }
132
133    fn tools(&self) -> Vec<Arc<dyn ToolHandle>> {
134        vec![Arc::new(WriteTodosTool {
135            name: "write_todos".into(),
136            state: self.state.clone(),
137        })]
138    }
139
140    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
141        ctx.request.append_prompt(WRITE_TODOS_SYSTEM_PROMPT);
142        Ok(())
143    }
144}
145
146pub struct FilesystemMiddleware {
147    state: Arc<RwLock<AgentStateSnapshot>>,
148}
149
150impl FilesystemMiddleware {
151    pub fn new(state: Arc<RwLock<AgentStateSnapshot>>) -> Self {
152        Self { state }
153    }
154}
155
156#[async_trait]
157impl AgentMiddleware for FilesystemMiddleware {
158    fn id(&self) -> &'static str {
159        "filesystem"
160    }
161
162    fn tools(&self) -> Vec<Arc<dyn ToolHandle>> {
163        vec![
164            Arc::new(LsTool {
165                name: "ls".into(),
166                state: self.state.clone(),
167            }),
168            Arc::new(ReadFileTool {
169                name: "read_file".into(),
170                state: self.state.clone(),
171            }),
172            Arc::new(WriteFileTool {
173                name: "write_file".into(),
174                state: self.state.clone(),
175            }),
176            Arc::new(EditFileTool {
177                name: "edit_file".into(),
178                state: self.state.clone(),
179            }),
180        ]
181    }
182
183    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
184        ctx.request.append_prompt(FILESYSTEM_SYSTEM_PROMPT);
185        Ok(())
186    }
187}
188
189#[derive(Clone)]
190pub struct SubAgentRegistration {
191    pub descriptor: SubAgentDescriptor,
192    pub agent: Arc<dyn AgentHandle>,
193}
194
195struct SubAgentRegistry {
196    agents: HashMap<String, Arc<dyn AgentHandle>>,
197}
198
199impl SubAgentRegistry {
200    fn new(registrations: Vec<SubAgentRegistration>) -> Self {
201        let mut agents = HashMap::new();
202        for reg in registrations {
203            agents.insert(reg.descriptor.name.clone(), reg.agent.clone());
204        }
205        Self { agents }
206    }
207
208    fn available_names(&self) -> Vec<String> {
209        self.agents.keys().cloned().collect()
210    }
211
212    fn get(&self, name: &str) -> Option<Arc<dyn AgentHandle>> {
213        self.agents.get(name).cloned()
214    }
215}
216
217pub struct SubAgentMiddleware {
218    task_tool: Arc<dyn ToolHandle>,
219    descriptors: Vec<SubAgentDescriptor>,
220    _registry: Arc<SubAgentRegistry>,
221}
222
223impl SubAgentMiddleware {
224    pub fn new(registrations: Vec<SubAgentRegistration>) -> Self {
225        let descriptors = registrations.iter().map(|r| r.descriptor.clone()).collect();
226        let registry = Arc::new(SubAgentRegistry::new(registrations));
227        let task_tool: Arc<dyn ToolHandle> = Arc::new(TaskRouterTool::new(registry.clone()));
228        Self {
229            task_tool,
230            descriptors,
231            _registry: registry,
232        }
233    }
234
235    fn prompt_fragment(&self) -> String {
236        let descriptions: Vec<String> = if self.descriptors.is_empty() {
237            vec![String::from("- general-purpose: Default reasoning agent")]
238        } else {
239            self.descriptors
240                .iter()
241                .map(|agent| format!("- {}: {}", agent.name, agent.description))
242                .collect()
243        };
244
245        TASK_TOOL_DESCRIPTION.replace("{other_agents}", &descriptions.join("\n"))
246    }
247}
248
249#[async_trait]
250impl AgentMiddleware for SubAgentMiddleware {
251    fn id(&self) -> &'static str {
252        "subagent"
253    }
254
255    fn tools(&self) -> Vec<Arc<dyn ToolHandle>> {
256        vec![self.task_tool.clone()]
257    }
258
259    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
260        ctx.request.append_prompt(TASK_SYSTEM_PROMPT);
261        ctx.request.append_prompt(&self.prompt_fragment());
262        Ok(())
263    }
264}
265
266#[derive(Clone, Debug)]
267pub struct HitlPolicy {
268    pub allow_auto: bool,
269    pub note: Option<String>,
270}
271
272pub struct HumanInLoopMiddleware {
273    policies: HashMap<String, HitlPolicy>,
274}
275
276impl HumanInLoopMiddleware {
277    pub fn new(policies: HashMap<String, HitlPolicy>) -> Self {
278        Self { policies }
279    }
280
281    pub fn requires_approval(&self, tool_name: &str) -> Option<&HitlPolicy> {
282        self.policies
283            .get(tool_name)
284            .filter(|policy| !policy.allow_auto)
285    }
286
287    fn prompt_fragment(&self) -> Option<String> {
288        let pending: Vec<String> = self
289            .policies
290            .iter()
291            .filter(|(_, policy)| !policy.allow_auto)
292            .map(|(tool, policy)| match &policy.note {
293                Some(note) => format!("- {tool}: {note}"),
294                None => format!("- {tool}: Requires approval"),
295            })
296            .collect();
297        if pending.is_empty() {
298            None
299        } else {
300            Some(format!(
301                "The following tools require human approval before execution:\n{}",
302                pending.join("\n")
303            ))
304        }
305    }
306}
307
308#[async_trait]
309impl AgentMiddleware for HumanInLoopMiddleware {
310    fn id(&self) -> &'static str {
311        "human-in-loop"
312    }
313
314    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
315        if let Some(fragment) = self.prompt_fragment() {
316            ctx.request.append_prompt(&fragment);
317        }
318        ctx.request.messages.push(AgentMessage {
319            role: MessageRole::System,
320            content: MessageContent::Text(
321                "Tools marked for human approval will emit interrupts requiring external resolution."
322                    .into(),
323            ),
324            metadata: None,
325        });
326        Ok(())
327    }
328}
329
330pub struct BaseSystemPromptMiddleware;
331
332#[async_trait]
333impl AgentMiddleware for BaseSystemPromptMiddleware {
334    fn id(&self) -> &'static str {
335        "base-system-prompt"
336    }
337
338    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
339        ctx.request.append_prompt(BASE_AGENT_PROMPT);
340        Ok(())
341    }
342}
343
344/// Anthropic-specific prompt caching middleware. Marks system prompts for caching
345/// to reduce latency on subsequent requests with the same base prompt.
346pub struct AnthropicPromptCachingMiddleware {
347    pub ttl: String,
348    pub unsupported_model_behavior: String,
349}
350
351impl AnthropicPromptCachingMiddleware {
352    pub fn new(ttl: impl Into<String>, unsupported_model_behavior: impl Into<String>) -> Self {
353        Self {
354            ttl: ttl.into(),
355            unsupported_model_behavior: unsupported_model_behavior.into(),
356        }
357    }
358
359    pub fn with_defaults() -> Self {
360        Self::new("5m", "ignore")
361    }
362
363    /// Parse TTL string like "5m" to detect if caching is requested.
364    /// For now, any non-empty TTL enables ephemeral caching.
365    fn should_enable_caching(&self) -> bool {
366        !self.ttl.is_empty() && self.ttl != "0" && self.ttl != "0s"
367    }
368}
369
370#[async_trait]
371impl AgentMiddleware for AnthropicPromptCachingMiddleware {
372    fn id(&self) -> &'static str {
373        "anthropic-prompt-caching"
374    }
375
376    async fn modify_model_request(&self, ctx: &mut MiddlewareContext<'_>) -> anyhow::Result<()> {
377        if !self.should_enable_caching() {
378            return Ok(());
379        }
380
381        // Mark system prompt for caching by converting it to a system message with cache control
382        if !ctx.request.system_prompt.is_empty() {
383            let system_message = AgentMessage {
384                role: MessageRole::System,
385                content: MessageContent::Text(ctx.request.system_prompt.clone()),
386                metadata: Some(MessageMetadata {
387                    tool_call_id: None,
388                    cache_control: Some(CacheControl {
389                        cache_type: "ephemeral".to_string(),
390                    }),
391                }),
392            };
393
394            // Insert system message at the beginning of the messages
395            ctx.request.messages.insert(0, system_message);
396
397            // Clear the system_prompt since it's now in messages
398            ctx.request.system_prompt.clear();
399
400            tracing::debug!(
401                ttl = %self.ttl,
402                behavior = %self.unsupported_model_behavior,
403                "Applied Anthropic prompt caching to system message"
404            );
405        }
406
407        Ok(())
408    }
409}
410
411pub struct TaskRouterTool {
412    registry: Arc<SubAgentRegistry>,
413}
414
415impl TaskRouterTool {
416    fn new(registry: Arc<SubAgentRegistry>) -> Self {
417        Self { registry }
418    }
419
420    fn available_subagents(&self) -> Vec<String> {
421        self.registry.available_names()
422    }
423}
424
425#[derive(Debug, Clone, Deserialize)]
426struct TaskInvocationArgs {
427    description: String,
428    subagent_type: String,
429}
430
431#[async_trait]
432impl ToolHandle for TaskRouterTool {
433    fn name(&self) -> &str {
434        "task"
435    }
436
437    async fn invoke(&self, invocation: ToolInvocation) -> anyhow::Result<ToolResponse> {
438        let args: TaskInvocationArgs = serde_json::from_value(invocation.args.clone())?;
439        let available = self.available_subagents();
440        if let Some(agent) = self.registry.get(&args.subagent_type) {
441            let user_message = AgentMessage {
442                role: MessageRole::User,
443                content: MessageContent::Text(args.description.clone()),
444                metadata: None,
445            };
446            let response = agent
447                .handle_message(user_message, Arc::new(AgentStateSnapshot::default()))
448                .await?;
449
450            return Ok(ToolResponse::Message(AgentMessage {
451                role: MessageRole::Tool,
452                content: response.content,
453                metadata: invocation.tool_call_id.map(|id| MessageMetadata {
454                    tool_call_id: Some(id),
455                    cache_control: None,
456                }),
457            }));
458        }
459
460        Ok(ToolResponse::Message(AgentMessage {
461            role: MessageRole::Tool,
462            content: MessageContent::Text(format!(
463                "Unknown subagent '{subagent}'. Available: {available:?}",
464                subagent = args.subagent_type,
465                available = available
466            )),
467            metadata: invocation.tool_call_id.map(|id| MessageMetadata {
468                tool_call_id: Some(id),
469                cache_control: None,
470            }),
471        }))
472    }
473}
474
475#[derive(Debug, Clone)]
476pub struct SubAgentDescriptor {
477    pub name: String,
478    pub description: String,
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484    use agents_core::agent::{AgentDescriptor, AgentHandle};
485    use agents_core::messaging::{MessageContent, MessageRole};
486    use serde_json::json;
487
488    struct AppendPromptMiddleware;
489
490    #[async_trait]
491    impl AgentMiddleware for AppendPromptMiddleware {
492        fn id(&self) -> &'static str {
493            "append-prompt"
494        }
495
496        async fn modify_model_request(
497            &self,
498            ctx: &mut MiddlewareContext<'_>,
499        ) -> anyhow::Result<()> {
500            ctx.request.system_prompt.push_str("\nExtra directives.");
501            Ok(())
502        }
503    }
504
505    #[tokio::test]
506    async fn middleware_mutates_prompt() {
507        let mut request = ModelRequest::new(
508            "System",
509            vec![AgentMessage {
510                role: MessageRole::User,
511                content: MessageContent::Text("Hi".into()),
512                metadata: None,
513            }],
514        );
515        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
516        let mut ctx = MiddlewareContext::with_request(&mut request, state);
517        let middleware = AppendPromptMiddleware;
518        middleware.modify_model_request(&mut ctx).await.unwrap();
519        assert!(ctx.request.system_prompt.contains("Extra directives"));
520    }
521
522    #[tokio::test]
523    async fn planning_middleware_registers_write_todos() {
524        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
525        let middleware = PlanningMiddleware::new(state);
526        let tool_names: Vec<_> = middleware
527            .tools()
528            .iter()
529            .map(|t| t.name().to_string())
530            .collect();
531        assert!(tool_names.contains(&"write_todos".to_string()));
532
533        let mut request = ModelRequest::new("System", vec![]);
534        let mut ctx = MiddlewareContext::with_request(
535            &mut request,
536            Arc::new(RwLock::new(AgentStateSnapshot::default())),
537        );
538        middleware.modify_model_request(&mut ctx).await.unwrap();
539        assert!(ctx.request.system_prompt.contains("write_todos"));
540    }
541
542    #[tokio::test]
543    async fn filesystem_middleware_registers_tools() {
544        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
545        let middleware = FilesystemMiddleware::new(state);
546        let tool_names: Vec<_> = middleware
547            .tools()
548            .iter()
549            .map(|t| t.name().to_string())
550            .collect();
551        for expected in ["ls", "read_file", "write_file", "edit_file"] {
552            assert!(tool_names.contains(&expected.to_string()));
553        }
554    }
555
556    #[tokio::test]
557    async fn summarization_middleware_trims_messages() {
558        let middleware = SummarizationMiddleware::new(2, "Summary note");
559        let mut request = ModelRequest::new(
560            "System",
561            vec![
562                AgentMessage {
563                    role: MessageRole::User,
564                    content: MessageContent::Text("one".into()),
565                    metadata: None,
566                },
567                AgentMessage {
568                    role: MessageRole::Agent,
569                    content: MessageContent::Text("two".into()),
570                    metadata: None,
571                },
572                AgentMessage {
573                    role: MessageRole::User,
574                    content: MessageContent::Text("three".into()),
575                    metadata: None,
576                },
577            ],
578        );
579        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
580        let mut ctx = MiddlewareContext::with_request(&mut request, state);
581        middleware.modify_model_request(&mut ctx).await.unwrap();
582        assert_eq!(ctx.request.messages.len(), 3);
583        match &ctx.request.messages[0].content {
584            MessageContent::Text(text) => assert!(text.contains("Summary note")),
585            other => panic!("expected text, got {other:?}"),
586        }
587    }
588
589    struct StubAgent;
590
591    #[async_trait]
592    impl AgentHandle for StubAgent {
593        async fn describe(&self) -> AgentDescriptor {
594            AgentDescriptor {
595                name: "stub".into(),
596                version: "0.0.1".into(),
597                description: None,
598            }
599        }
600
601        async fn handle_message(
602            &self,
603            _input: AgentMessage,
604            _state: Arc<AgentStateSnapshot>,
605        ) -> anyhow::Result<AgentMessage> {
606            Ok(AgentMessage {
607                role: MessageRole::Agent,
608                content: MessageContent::Text("stub-response".into()),
609                metadata: None,
610            })
611        }
612    }
613
614    #[tokio::test]
615    async fn task_router_reports_unknown_subagent() {
616        let registry = Arc::new(SubAgentRegistry::new(vec![]));
617        let task_tool = TaskRouterTool::new(registry.clone());
618
619        let response = task_tool
620            .invoke(ToolInvocation {
621                tool_name: "task".into(),
622                args: json!({
623                    "description": "Do something",
624                    "subagent_type": "unknown"
625                }),
626                tool_call_id: None,
627            })
628            .await
629            .unwrap();
630
631        match response {
632            ToolResponse::Message(msg) => match msg.content {
633                MessageContent::Text(text) => assert!(text.contains("Unknown subagent")),
634                other => panic!("expected text, got {other:?}"),
635            },
636            _ => panic!("expected message"),
637        }
638    }
639
640    #[tokio::test]
641    async fn subagent_middleware_appends_prompt() {
642        let subagents = vec![SubAgentRegistration {
643            descriptor: SubAgentDescriptor {
644                name: "research-agent".into(),
645                description: "Deep research specialist".into(),
646            },
647            agent: Arc::new(StubAgent),
648        }];
649        let middleware = SubAgentMiddleware::new(subagents);
650
651        let mut request = ModelRequest::new("System", vec![]);
652        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
653        let mut ctx = MiddlewareContext::with_request(&mut request, state);
654        middleware.modify_model_request(&mut ctx).await.unwrap();
655
656        assert!(ctx.request.system_prompt.contains("research-agent"));
657        let tool_names: Vec<_> = middleware
658            .tools()
659            .iter()
660            .map(|t| t.name().to_string())
661            .collect();
662        assert!(tool_names.contains(&"task".to_string()));
663    }
664
665    #[tokio::test]
666    async fn task_router_invokes_registered_subagent() {
667        let registry = Arc::new(SubAgentRegistry::new(vec![SubAgentRegistration {
668            descriptor: SubAgentDescriptor {
669                name: "stub-agent".into(),
670                description: "Stub".into(),
671            },
672            agent: Arc::new(StubAgent),
673        }]));
674        let task_tool = TaskRouterTool::new(registry.clone());
675        let response = task_tool
676            .invoke(ToolInvocation {
677                tool_name: "task".into(),
678                args: json!({
679                    "description": "do work",
680                    "subagent_type": "stub-agent"
681                }),
682                tool_call_id: Some("call-42".into()),
683            })
684            .await
685            .unwrap();
686
687        match response {
688            ToolResponse::Message(msg) => {
689                assert_eq!(msg.metadata.unwrap().tool_call_id.unwrap(), "call-42");
690                match msg.content {
691                    MessageContent::Text(text) => assert_eq!(text, "stub-response"),
692                    other => panic!("expected text, got {other:?}"),
693                }
694            }
695            _ => panic!("expected message"),
696        }
697    }
698
699    #[tokio::test]
700    async fn human_in_loop_appends_prompt() {
701        let middleware = HumanInLoopMiddleware::new(HashMap::from([(
702            "danger-tool".into(),
703            HitlPolicy {
704                allow_auto: false,
705                note: Some("Requires security review".into()),
706            },
707        )]));
708        let mut request = ModelRequest::new("System", vec![]);
709        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
710        let mut ctx = MiddlewareContext::with_request(&mut request, state);
711        middleware.modify_model_request(&mut ctx).await.unwrap();
712        assert!(ctx
713            .request
714            .system_prompt
715            .contains("danger-tool: Requires security review"));
716    }
717
718    #[tokio::test]
719    async fn anthropic_prompt_caching_moves_system_prompt_to_messages() {
720        let middleware = AnthropicPromptCachingMiddleware::new("5m", "ignore");
721        let mut request = ModelRequest::new(
722            "This is the system prompt",
723            vec![AgentMessage {
724                role: MessageRole::User,
725                content: MessageContent::Text("Hello".into()),
726                metadata: None,
727            }],
728        );
729        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
730        let mut ctx = MiddlewareContext::with_request(&mut request, state);
731
732        // Apply the middleware
733        middleware.modify_model_request(&mut ctx).await.unwrap();
734
735        // System prompt should be cleared
736        assert!(ctx.request.system_prompt.is_empty());
737
738        // Should have added a system message with cache control at the beginning
739        assert_eq!(ctx.request.messages.len(), 2);
740
741        let system_message = &ctx.request.messages[0];
742        assert!(matches!(system_message.role, MessageRole::System));
743        assert_eq!(
744            system_message.content.as_text().unwrap(),
745            "This is the system prompt"
746        );
747
748        // Check cache control metadata
749        let metadata = system_message.metadata.as_ref().unwrap();
750        let cache_control = metadata.cache_control.as_ref().unwrap();
751        assert_eq!(cache_control.cache_type, "ephemeral");
752
753        // Original user message should still be there
754        let user_message = &ctx.request.messages[1];
755        assert!(matches!(user_message.role, MessageRole::User));
756        assert_eq!(user_message.content.as_text().unwrap(), "Hello");
757    }
758
759    #[tokio::test]
760    async fn anthropic_prompt_caching_disabled_with_zero_ttl() {
761        let middleware = AnthropicPromptCachingMiddleware::new("0", "ignore");
762        let mut request = ModelRequest::new("This is the system prompt", vec![]);
763        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
764        let mut ctx = MiddlewareContext::with_request(&mut request, state);
765
766        // Apply the middleware
767        middleware.modify_model_request(&mut ctx).await.unwrap();
768
769        // System prompt should be unchanged
770        assert_eq!(ctx.request.system_prompt, "This is the system prompt");
771        assert_eq!(ctx.request.messages.len(), 0);
772    }
773
774    #[tokio::test]
775    async fn anthropic_prompt_caching_no_op_with_empty_system_prompt() {
776        let middleware = AnthropicPromptCachingMiddleware::new("5m", "ignore");
777        let mut request = ModelRequest::new(
778            "",
779            vec![AgentMessage {
780                role: MessageRole::User,
781                content: MessageContent::Text("Hello".into()),
782                metadata: None,
783            }],
784        );
785        let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
786        let mut ctx = MiddlewareContext::with_request(&mut request, state);
787
788        // Apply the middleware
789        middleware.modify_model_request(&mut ctx).await.unwrap();
790
791        // Should be unchanged
792        assert!(ctx.request.system_prompt.is_empty());
793        assert_eq!(ctx.request.messages.len(), 1);
794        assert!(matches!(ctx.request.messages[0].role, MessageRole::User));
795    }
796}
agents_runtime/middleware.rs

agents_runtime/
middleware.rs
