bamboo_engine/
message_hooks.rs

1//! Message preflight hooks.
2//!
3//! These hooks run before we forward requests upstream (proxy endpoints) and before we
4//! enter the agent loop. They operate on internal `bamboo_agent_core::Message` so the
5//! same behavior applies across OpenAI-compatible, Anthropic, Gemini, and agent routes.
6
7use bamboo_agent_core::Message;
8use bamboo_infrastructure::Config;
9
10#[derive(Debug, thiserror::Error)]
11pub enum HookError {
12    #[error("Invalid hook configuration: {0}")]
13    InvalidConfig(String),
14    #[error("Request not supported: {0}")]
15    Unsupported(String),
16}
17
18/// Apply all configured preflight hooks.
19pub async fn apply_message_preflight_hooks(
20    attachment_reader: Option<&dyn bamboo_agent_core::storage::AttachmentReader>,
21    config: &Config,
22    _model: &str,
23    messages: &mut [Message],
24) -> Result<(), HookError> {
25    apply_image_fallback_hook(attachment_reader, config, messages).await
26}
27
28async fn apply_image_fallback_hook(
29    attachment_reader: Option<&dyn bamboo_agent_core::storage::AttachmentReader>,
30    config: &Config,
31    messages: &mut [Message],
32) -> Result<(), HookError> {
33    let hook_cfg = &config.hooks.image_fallback;
34    if !hook_cfg.enabled {
35        return Ok(());
36    }
37
38    let mode = hook_cfg.mode.trim().to_ascii_lowercase();
39    let fallback_mode = match mode.as_str() {
40        "placeholder" => crate::ImageFallbackMode::Placeholder,
41        "error" => crate::ImageFallbackMode::Error,
42        "ocr" => crate::ImageFallbackMode::Ocr,
43        _ => {
44            return Err(HookError::InvalidConfig(format!(
45                "hooks.image_fallback.mode must be 'placeholder', 'error', or 'ocr' (got '{mode}')"
46            )));
47        }
48    };
49
50    let fallback = crate::ImageFallbackConfig {
51        mode: fallback_mode,
52        vision_model: None,
53    };
54
55    crate::runtime::runner::image_fallback::apply_image_fallback_to_llm_messages(
56        messages,
57        fallback,
58        attachment_reader,
59        None,
60    )
61    .await
62    .map_err(|e| HookError::Unsupported(e.to_string()))
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use bamboo_infrastructure::models::{ContentPart, ImageUrl};
69    use tempfile::TempDir;
70
71    fn base_config(mode: &str) -> Config {
72        let dir = TempDir::new().expect("tempdir");
73        let mut cfg = Config::from_data_dir(Some(dir.path().to_path_buf()));
74        cfg.hooks.image_fallback.enabled = true;
75        cfg.hooks.image_fallback.mode = mode.to_string();
76        cfg
77    }
78
79    #[tokio::test]
80    async fn image_fallback_placeholder_rewrites_images_to_text_without_leaking_data() {
81        let cfg = base_config("placeholder");
82
83        let mut messages = vec![Message::user_with_parts(
84            "What is in this image?",
85            vec![
86                ContentPart::Text {
87                    text: "What is in this image?".to_string(),
88                },
89                ContentPart::ImageUrl {
90                    image_url: ImageUrl {
91                        url: "data:image/png;base64,AAAABBBBCCCC".to_string(),
92                        detail: None,
93                    },
94                },
95            ]
96            .into_iter()
97            .map(Into::into)
98            .collect(),
99        )];
100
101        apply_message_preflight_hooks(None, &cfg, "m", &mut messages)
102            .await
103            .expect("hook ok");
104
105        assert!(messages[0].content.contains("Image omitted: image/png"));
106        assert!(!messages[0].content.contains("AAAABBBBCCCC"));
107        assert!(messages[0].content_parts.is_none());
108    }
109
110    #[tokio::test]
111    async fn image_fallback_error_rejects_requests_with_images() {
112        let cfg = base_config("error");
113
114        let mut messages = vec![Message::user_with_parts(
115            "",
116            vec![ContentPart::ImageUrl {
117                image_url: ImageUrl {
118                    url: "https://example.com/image.png".to_string(),
119                    detail: None,
120                },
121            }]
122            .into_iter()
123            .map(Into::into)
124            .collect(),
125        )];
126
127        let err = apply_message_preflight_hooks(None, &cfg, "m", &mut messages)
128            .await
129            .expect_err("should err");
130        assert!(err
131            .to_string()
132            .contains("does not currently support image inputs"));
133    }
134
135    #[tokio::test]
136    async fn image_fallback_invalid_mode_errors() {
137        let cfg = base_config("wat");
138        let mut messages = Vec::new();
139        let err = apply_message_preflight_hooks(None, &cfg, "m", &mut messages)
140            .await
141            .expect_err("should err");
142        assert!(matches!(err, HookError::InvalidConfig(_)));
143    }
144
145    #[cfg(not(windows))]
146    #[tokio::test]
147    async fn image_fallback_ocr_non_windows_leaves_images_intact() {
148        let cfg = base_config("ocr");
149
150        let mut messages = vec![Message::user_with_parts(
151            "hi",
152            vec![
153                ContentPart::Text {
154                    text: "hi".to_string(),
155                },
156                ContentPart::ImageUrl {
157                    image_url: ImageUrl {
158                        url: "data:image/png;base64,AAAABBBBCCCC".to_string(),
159                        detail: None,
160                    },
161                },
162            ]
163            .into_iter()
164            .map(Into::into)
165            .collect(),
166        )];
167
168        apply_message_preflight_hooks(None, &cfg, "m", &mut messages)
169            .await
170            .expect("hook ok");
171
172        assert!(messages[0].content_parts.is_some());
173        assert!(messages[0].content.contains("hi"));
174    }
175}
bamboo_engine/message_hooks.rs

bamboo_engine/
message_hooks.rs
