mobius/middleware/
image_generation.rs

1//! Native image generation, publication, and transcript presentation.
2
3use std::sync::{Arc, LazyLock};
4
5use serde::Deserialize;
6use serde_json::Value;
7
8use super::tools::{ApprovalRequirement, Catalog, Tool, ToolContext, ToolSpec};
9use super::{Middleware, RuntimeContext};
10use crate::backend::model::{
11    ImageAspect, ImageGenerationReference, ImageGenerationRequest, ModelRouter, ToolDefinition,
12};
13use crate::backend::session_files::SessionFileStore;
14use crate::protocol::{
15    ContentPart, EventMsg, FrontendBlock, FrontendBlockFormat, FrontendBlockRole,
16    FrontendBlockUpdate, FrontendContribution, FrontendSettingValue, ToolContent, ToolResponse,
17};
18use crate::{BoxFuture, Error, Result};
19
20#[derive(Deserialize)]
21#[serde(deny_unknown_fields)]
22struct Definition {
23    default_enabled: bool,
24    manifest_label: String,
25    manifest_description: String,
26    render_ready: String,
27    render_failed: String,
28    #[serde(deserialize_with = "super::manifest::deserialize_settings")]
29    settings: Vec<super::manifest::MiddlewareSettingManifest>,
30    generate_image: ToolSpec,
31}
32
33const MAX_REFERENCE_IMAGES: usize = 4;
34
35static DEFINITION: LazyLock<Definition> =
36    LazyLock::new(|| crate::config::embedded(include_str!("image_generation.toml")));
37
38super::manifest::middleware_manifest! {
39/// Configuration metadata for native image generation.
40    "image_generation", DEFINITION, required: false, settings: &DEFINITION.settings
41}
42
43/// Generates a session image and publishes it as a chat artifact.
44pub struct ImageGeneration {
45    models: Arc<ModelRouter>,
46    store: SessionFileStore,
47    route: Option<Arc<str>>,
48}
49
50impl ImageGeneration {
51    /// Creates the image capability with the configured model routes and session files.
52    /// `route` selects an image model route; `None` uses the first configured one.
53    #[must_use]
54    pub fn new(models: Arc<ModelRouter>, store: SessionFileStore, route: Option<Arc<str>>) -> Self {
55        Self {
56            models,
57            store,
58            route,
59        }
60    }
61
62    /// Creates the image capability from its manifest settings, read through `setting`.
63    /// # Errors
64    ///
65    /// Returns an error if the `model` setting is not a string.
66    pub fn from_settings<'a>(
67        models: Arc<ModelRouter>,
68        store: SessionFileStore,
69        setting: impl FnOnce(&str) -> Option<&'a FrontendSettingValue>,
70    ) -> Result<Self> {
71        let route = match setting("model") {
72            Some(FrontendSettingValue::String(route)) => Some(Arc::from(route.as_str())),
73            Some(FrontendSettingValue::Integer(_)) => {
74                return Err(Error::Config(format!(
75                    "middleware setting `{}.model` must be string",
76                    MANIFEST.id
77                )));
78            }
79            None => None,
80        };
81        Ok(Self::new(models, store, route))
82    }
83}
84
85impl Middleware for ImageGeneration {
86    fn name(&self) -> &'static str {
87        MANIFEST.id
88    }
89
90    fn register(&self, catalog: &mut Catalog, runtime: &RuntimeContext) -> Result<()> {
91        if self.models.image_choices().len() == 0 {
92            return Ok(());
93        }
94        catalog.register(Arc::new(GenerateImage {
95            models: Arc::clone(&self.models),
96            store: self.store.clone(),
97            session_id: runtime.session_id.clone(),
98            route: self.route.as_ref().map(Arc::clone),
99        }))
100    }
101
102    fn frontend(&self, _session_id: &str) -> FrontendContribution {
103        FrontendContribution {
104            capability: MANIFEST.id.into(),
105            ..FrontendContribution::default()
106        }
107    }
108
109    fn render(&self, event: &EventMsg, _session_id: &str) -> Option<FrontendBlock> {
110        let mut block = DEFINITION.generate_image.render(event)?;
111        block.role = FrontendBlockRole::Artifact;
112        block.format = FrontendBlockFormat::Image;
113        block.update = FrontendBlockUpdate::Replace;
114        match event {
115            EventMsg::ToolCallBegin(_) => block.text.clear(),
116            EventMsg::ToolCallEnd(result) if !result.is_error => {
117                block.title = DEFINITION.render_ready.clone();
118                block.text.clear();
119                block.files = result.output.files().cloned().collect();
120                block.content = ToolContent(
121                    result
122                        .output
123                        .0
124                        .iter()
125                        .filter(|part| matches!(part, ContentPart::Image { .. }))
126                        .cloned()
127                        .collect(),
128                );
129            }
130            EventMsg::ToolCallEnd(_) => block.title = DEFINITION.render_failed.clone(),
131            _ => {}
132        }
133        Some(block)
134    }
135}
136
137#[derive(Deserialize)]
138#[serde(deny_unknown_fields)]
139struct GenerateImageArgs {
140    prompt: String,
141    #[serde(default)]
142    image_aspect: ImageAspect,
143    #[serde(default)]
144    reference_file_ids: Vec<String>,
145}
146
147struct GenerateImage {
148    models: Arc<ModelRouter>,
149    store: SessionFileStore,
150    session_id: String,
151    route: Option<Arc<str>>,
152}
153
154impl Tool for GenerateImage {
155    fn definition(&self) -> ToolDefinition {
156        let mut tool = DEFINITION.generate_image.tool.clone();
157        let properties = &mut tool.parameters["properties"];
158        properties["prompt"]["maxLength"] = crate::backend::model::MAX_IMAGE_PROMPT_CHARS.into();
159        properties["reference_file_ids"]["maxItems"] = MAX_REFERENCE_IMAGES.into();
160        tool
161    }
162
163    fn approval(&self) -> ApprovalRequirement {
164        ApprovalRequirement::Always
165    }
166
167    fn call<'a>(
168        &'a self,
169        context: ToolContext,
170        arguments: Value,
171    ) -> BoxFuture<'a, Result<ToolResponse>> {
172        Box::pin(async move {
173            let arguments: GenerateImageArgs = serde_json::from_value(arguments)?;
174            if arguments.reference_file_ids.len() > MAX_REFERENCE_IMAGES {
175                return Err(Error::Tool("too many reference images".into()));
176            }
177            let mut owned_references = Vec::with_capacity(arguments.reference_file_ids.len());
178            for id in &arguments.reference_file_ids {
179                let file = self.store.file_reference(&self.session_id, id).await?;
180                let bytes = self.store.read_image(&self.session_id, &file).await?;
181                owned_references.push((file.media_type, bytes));
182            }
183            let references = owned_references
184                .iter()
185                .map(|(media_type, bytes)| ImageGenerationReference { media_type, bytes })
186                .collect::<Vec<_>>();
187            let image = self.models.image_model(self.route.as_deref())?;
188            let generated = self
189                .models
190                .generate_image(
191                    image.route,
192                    ImageGenerationRequest {
193                        model: image.model,
194                        quality: image.quality,
195                        prompt: &arguments.prompt,
196                        image_aspect: arguments.image_aspect,
197                        references: &references,
198                    },
199                )
200                .await?;
201            if let Some(usage) = generated.usage {
202                context.report_usage(usage)?;
203            }
204            let extension = image::ImageFormat::from_mime_type(&generated.media_type)
205                .filter(image::ImageFormat::reading_enabled)
206                .and_then(|format| format.extensions_str().first())
207                .ok_or_else(|| {
208                    Error::Tool("provider returned an unsupported image format".into())
209                })?;
210            let image = self
211                .store
212                .publish_image(
213                    &self.session_id,
214                    format!("generated-image.{extension}"),
215                    &generated.media_type,
216                    generated.bytes,
217                )
218                .await?;
219            // Models without image tool results receive the file; frontends still get its pixels.
220            let part = if self
221                .models
222                .supports_tool_image_input(context.model_route())
223                .unwrap_or(false)
224            {
225                ContentPart::Image { image }
226            } else {
227                ContentPart::File { file: image.file }
228            };
229            Ok(ToolResponse {
230                content: ToolContent(vec![part]),
231                is_error: false,
232            })
233        })
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::agent::AgentRole;
241    use crate::backend::checkpoint::sqlite::SqliteCheckpoint;
242    use crate::backend::model::{GeneratedImage, Model, ModelEventSink, ModelOutput, ModelRequest};
243    use crate::backend::sandbox::local::LocalSandbox;
244    use crate::backend::sandbox::{
245        ApprovalPolicy, NetworkAccess, Sandbox, SandboxMode, SandboxPermissions,
246    };
247    use crate::middleware::tools::ToolExposure;
248    use crate::protocol::{
249        FrontendBlockState, SessionContext, TokenUsage, ToolCallBeginEvent, ToolCallEndEvent,
250    };
251
252    struct ImageModel(Vec<u8>, bool);
253
254    impl Model for ImageModel {
255        fn supports_image_generation(&self) -> bool {
256            self.1
257        }
258
259        fn supports_image_input(&self) -> bool {
260            self.1
261        }
262
263        fn supports_tool_image_input(&self) -> bool {
264            self.1
265        }
266
267        fn generate_image<'a>(
268            &'a self,
269            request: ImageGenerationRequest<'a>,
270        ) -> BoxFuture<'a, Result<GeneratedImage>> {
271            assert_eq!(request.model, "gpt-image-2.5-sunburst");
272            assert_eq!(request.image_aspect, ImageAspect::Landscape);
273            let bytes = self.0.clone();
274            Box::pin(async move {
275                Ok(GeneratedImage {
276                    bytes,
277                    media_type: "image/png".into(),
278                    usage: Some(TokenUsage {
279                        output_tokens: 1,
280                        total_tokens: 1,
281                        ..TokenUsage::default()
282                    }),
283                })
284            })
285        }
286
287        fn respond<'a>(
288            &'a self,
289            _request: ModelRequest<'a>,
290            _events: ModelEventSink,
291        ) -> BoxFuture<'a, Result<ModelOutput>> {
292            Box::pin(async { unreachable!("image tool never starts a conversation response") })
293        }
294    }
295
296    #[tokio::test]
297    async fn generates_one_artifact_and_replaces_its_pending_transcript_block() {
298        let state = tempfile::tempdir().expect("state");
299        let mut bytes = std::io::Cursor::new(Vec::new());
300        image::DynamicImage::ImageRgba8(image::RgbaImage::new(3, 2))
301            .write_to(&mut bytes, image::ImageFormat::Png)
302            .expect("test image");
303        let png = bytes.into_inner();
304        let store = SessionFileStore::new(state.path(), None);
305        let mut router = ModelRouter::new("image", Arc::new(ImageModel(png.clone(), true)));
306        router
307            .register_image(
308                Arc::new(ImageModel(png.clone(), true)),
309                crate::protocol::ModelChoice {
310                    route: "openai::gpt-image-2.5-sunburst".into(),
311                    group: "OpenAI".into(),
312                    model: "gpt-image-2.5-sunburst".into(),
313                    reasoning_effort: None,
314                    variant_label: None,
315                    context_window: None,
316                    supports_image_input: false,
317                    supports_image_generation: true,
318                    supports_realtime_voice: false,
319                    tool_discovery: crate::protocol::ToolDiscoveryMode::Rebuild,
320                },
321                Default::default(),
322            )
323            .expect("image route");
324        router
325            .register_image(
326                Arc::new(ImageModel(png.clone(), true)),
327                crate::protocol::ModelChoice {
328                    route: "openai::gpt-image-2.5-flare::high".into(),
329                    group: "OpenAI".into(),
330                    model: "gpt-image-2.5-flare".into(),
331                    reasoning_effort: Some("high".into()),
332                    variant_label: None,
333                    context_window: None,
334                    supports_image_input: false,
335                    supports_image_generation: true,
336                    supports_realtime_voice: false,
337                    tool_discovery: crate::protocol::ToolDiscoveryMode::Rebuild,
338                },
339                Default::default(),
340            )
341            .expect("image tier route");
342        let flare = router
343            .image_model(Some("openai::gpt-image-2.5-flare::high"))
344            .expect("quality route");
345        assert_eq!(
346            (flare.model, flare.quality),
347            ("gpt-image-2.5-flare", Some("high"))
348        );
349        let middleware = ImageGeneration::new(Arc::new(router), store.clone(), None);
350        let runtime = RuntimeContext {
351            children: crate::agent::ChildAgents::default(),
352            sender: crate::agent::test_sender(),
353            checkpoints: Arc::new(
354                SqliteCheckpoint::new(state.path().join("checkpoints.sqlite3"))
355                    .expect("checkpoint store"),
356            ),
357            session_id: "session".into(),
358            model_route: "image".into(),
359            model: "image".into(),
360            approval_policy: ApprovalPolicy::Ask,
361            session_context: SessionContext::default(),
362            metadata: Default::default(),
363            role: AgentRole::Main,
364            frontend: Arc::new(|_| Ok(())),
365        };
366        let mut catalog = Catalog::default();
367        middleware
368            .register(&mut catalog, &runtime)
369            .expect("supported route registers tool");
370        let unsupported = ImageGeneration::new(
371            Arc::new(ModelRouter::new(
372                "unsupported",
373                Arc::new(ImageModel(png.clone(), false)),
374            )),
375            store.clone(),
376            None,
377        );
378        let mut unsupported_runtime = runtime.clone();
379        unsupported_runtime.model_route = "unsupported".into();
380        unsupported
381            .register(&mut Catalog::default(), &unsupported_runtime)
382            .expect("a router without image models registers no tool");
383        let source = store
384            .publish_artifact("session", "source.png".into(), "image/png".into(), &png)
385            .await
386            .expect("source image");
387        let tool = GenerateImage {
388            models: Arc::clone(&middleware.models),
389            store: store.clone(),
390            session_id: "session".into(),
391            route: None,
392        };
393        assert_eq!(tool.exposure(), ToolExposure::Deferred);
394        assert_eq!(tool.approval(), ApprovalRequirement::Always);
395        let sandbox = Arc::new(Sandbox::new(
396            Arc::new(LocalSandbox::new(state.path()).expect("sandbox")),
397            ApprovalPolicy::Ask,
398        ));
399        let permissions = SandboxPermissions::restore(
400            "session",
401            SandboxMode::WorkspaceWrite,
402            NetworkAccess::Allowed,
403            ["image-call".into()],
404        );
405        let context = ToolContext::new(
406            Arc::clone(&sandbox),
407            permissions.for_call("image-call"),
408            "turn",
409        )
410        .with_model_route("image");
411        let response = tool
412            .call(
413                context,
414                serde_json::json!({"prompt": "a red landscape", "image_aspect": "landscape", "reference_file_ids": [source.id]}),
415            )
416            .await
417            .expect("generate image");
418        let Some(ContentPart::Image { image }) = response.content.0.first() else {
419            panic!("generated image part");
420        };
421        assert!(image.width > 0 && image.height > 0);
422        let file = response
423            .content
424            .files()
425            .next()
426            .expect("published file")
427            .clone();
428        assert_eq!(store.read_file("session", &file).await.expect("image"), png);
429        let artifacts = store
430            .list_files(
431                "session",
432                &[crate::backend::session_files::SessionFileOrigin::Artifact],
433            )
434            .await
435            .expect("artifacts");
436        assert_eq!(artifacts.len(), 2);
437        assert!(artifacts.iter().any(|(_, stored)| stored == &file));
438        assert_eq!(
439            store
440                .list_files(
441                    "session",
442                    &[
443                        crate::backend::session_files::SessionFileOrigin::Upload,
444                        crate::backend::session_files::SessionFileOrigin::Artifact
445                    ]
446                )
447                .await
448                .expect("stored files")
449                .len(),
450            2
451        );
452
453        let begin = middleware
454            .render(
455                &EventMsg::ToolCallBegin(ToolCallBeginEvent {
456                    turn_id: "turn".into(),
457                    call_id: "image-call".into(),
458                    name: "generate_image".into(),
459                    arguments: serde_json::json!({"prompt": "a red landscape", "image_aspect": "landscape"}),
460                }),
461                "session",
462            )
463            .expect("pending block");
464        let end = middleware
465            .render(
466                &EventMsg::ToolCallEnd(ToolCallEndEvent {
467                    turn_id: "turn".into(),
468                    call_id: "image-call".into(),
469                    name: "generate_image".into(),
470                    output: response.content,
471                    is_error: false,
472                }),
473                "session",
474            )
475            .expect("completed block");
476        assert_eq!(begin.id, end.id);
477        assert_eq!(begin.state, FrontendBlockState::Pending);
478        assert_eq!(
479            (begin.title.as_str(), begin.text.as_str()),
480            ("Generating image", "")
481        );
482        assert_eq!(end.title, "Image ready");
483        assert_eq!(begin.format, FrontendBlockFormat::Image);
484        assert_eq!(end.state, FrontendBlockState::Complete);
485        assert_eq!(end.files, vec![file.clone()]);
486        assert!(matches!(
487            end.content.0.as_slice(),
488            [ContentPart::Image { image }] if image.file == file && image.width > 0
489        ));
490
491        let default: GenerateImageArgs =
492            serde_json::from_value(serde_json::json!({"prompt": "x"})).expect("default shape");
493        assert_eq!(default.image_aspect, ImageAspect::Square);
494        assert!(
495            serde_json::from_value::<GenerateImageArgs>(
496                serde_json::json!({"prompt": "x", "image_aspect": "auto"})
497            )
498            .is_err()
499        );
500
501        let unsupported_tool = GenerateImage {
502            models: Arc::clone(&unsupported.models),
503            store,
504            session_id: "session".into(),
505            route: None,
506        };
507        let schema = unsupported_tool.definition().parameters;
508        assert_eq!(schema["properties"]["prompt"]["maxLength"], 32_000);
509        assert_eq!(schema["properties"]["reference_file_ids"]["maxItems"], 4);
510        let unsupported_context =
511            ToolContext::new(sandbox, permissions.for_call("image-call"), "turn")
512                .with_model_route("unsupported");
513        assert!(matches!(
514            unsupported_tool
515                .call(
516                    unsupported_context,
517                    serde_json::json!({"prompt": "a red square"})
518                )
519                .await,
520            Err(Error::Config(_))
521        ));
522    }
523}
mobius/middleware/image_generation.rs

mobius/middleware/
image_generation.rs
