supercode_harness/tools/
image_gen.rs

1//! BP-3 (catalog row "Image generation tool", cx§1's `image_gen` feature):
2//! a client-issued call to the PROVIDER's image endpoint.
3//!
4//! The row's semantics is "provider-executed image creation": the provider
5//! renders the image, and this tool is the client side of that call — the
6//! OpenAI-compatible `POST {base_url}/images/generations` every
7//! OpenAI-shaped endpoint (including the one `crate::Config::base_url`
8//! already points the chat provider at) exposes. The bytes come back
9//! base64-encoded, are written into the working directory through the same
10//! [`ToolContext::check_write`] sandbox check every other write-capable
11//! tool passes, and the model gets the path.
12//!
13//! **UnsupportedAction is a real answer.** Not every OpenAI-compatible
14//! gateway implements the images route. When the endpoint answers 404 / 405
15//! / 501, the tool says so in exactly those terms — "unsupported_action:
16//! this provider exposes no image endpoint" — rather than retrying, falling
17//! back to another vendor, or inventing a file. Nothing here reaches for a
18//! provider the session was not already configured with.
19
20use std::path::PathBuf;
21use std::time::Duration;
22
23use async_trait::async_trait;
24use serde::Deserialize;
25use serde_json::{json, Value};
26
27use crate::error::{Error, Result};
28use crate::tools::{Tool, ToolContext};
29
30/// Registered name of the image tool (Codex's own spelling).
31pub const IMAGE_GEN: &str = "image_gen";
32
33/// The image model used when neither the call nor the config names one.
34pub const DEFAULT_IMAGE_MODEL: &str = "gpt-image-1";
35
36/// Request timeout: image generation is slow compared with a chat call.
37const IMAGE_TIMEOUT_SECS: u64 = 180;
38
39/// Refuse to write an image larger than this (a runaway response must not
40/// fill the disk).
41const MAX_IMAGE_BYTES: usize = 32 * 1024 * 1024;
42
43/// `image_gen` — generate an image through the session's own provider.
44///
45/// Constructed from the resolved [`crate::Config`] by
46/// [`crate::tools::ToolRegistry::from_config`], so it inherits exactly the
47/// endpoint and credential the chat provider uses; it never reads a second
48/// vendor's configuration.
49#[derive(Debug, Clone)]
50pub struct ImageGenTool {
51    base_url: String,
52    api_key: Option<String>,
53    api_key_env: String,
54    default_model: String,
55}
56
57impl ImageGenTool {
58    /// Build the tool from the session's provider settings.
59    pub fn new(
60        base_url: impl Into<String>,
61        api_key: Option<String>,
62        api_key_env: impl Into<String>,
63    ) -> Self {
64        ImageGenTool {
65            base_url: base_url.into(),
66            api_key,
67            api_key_env: api_key_env.into(),
68            default_model: DEFAULT_IMAGE_MODEL.to_string(),
69        }
70    }
71
72    /// Override the default image model (the `model` argument still wins).
73    pub fn with_default_model(mut self, model: impl Into<String>) -> Self {
74        self.default_model = model.into();
75        self
76    }
77
78    /// The endpoint this tool posts to.
79    pub fn endpoint(&self) -> String {
80        format!("{}/images/generations", self.base_url.trim_end_matches('/'))
81    }
82
83    fn resolved_key(&self) -> Option<String> {
84        self.api_key.clone().filter(|k| !k.is_empty()).or_else(|| {
85            std::env::var(&self.api_key_env)
86                .ok()
87                .filter(|k| !k.is_empty())
88        })
89    }
90}
91
92#[derive(Debug, Deserialize)]
93struct ImageGenArgs {
94    prompt: String,
95    #[serde(default)]
96    path: Option<String>,
97    #[serde(default)]
98    size: Option<String>,
99    #[serde(default)]
100    model: Option<String>,
101}
102
103/// Standard base64 decoder — the mirror of `builtins::base64_encode`, same
104/// "small parser over a crate" precedent. `None` on any non-alphabet byte.
105fn base64_decode(input: &str) -> Option<Vec<u8>> {
106    fn digit(byte: u8) -> Option<u8> {
107        match byte {
108            b'A'..=b'Z' => Some(byte - b'A'),
109            b'a'..=b'z' => Some(byte - b'a' + 26),
110            b'0'..=b'9' => Some(byte - b'0' + 52),
111            b'+' => Some(62),
112            b'/' => Some(63),
113            _ => None,
114        }
115    }
116    let mut out = Vec::with_capacity(input.len() / 4 * 3 + 3);
117    let mut chunk = [0u8; 4];
118    let mut len = 0usize;
119    let mut padding = 0usize;
120    for byte in input.bytes().filter(|b| !b.is_ascii_whitespace()) {
121        if byte == b'=' {
122            padding += 1;
123            chunk[len] = 0;
124        } else {
125            chunk[len] = digit(byte)?;
126        }
127        len += 1;
128        if len == 4 {
129            let value = ((chunk[0] as u32) << 18)
130                | ((chunk[1] as u32) << 12)
131                | ((chunk[2] as u32) << 6)
132                | chunk[3] as u32;
133            out.push((value >> 16) as u8);
134            if padding < 2 {
135                out.push((value >> 8) as u8);
136            }
137            if padding < 1 {
138                out.push(value as u8);
139            }
140            len = 0;
141            padding = 0;
142        }
143    }
144    if len == 0 {
145        Some(out)
146    } else {
147        None
148    }
149}
150
151/// A filename-safe slug of the prompt, for the default output path.
152fn slug(prompt: &str) -> String {
153    let mut out = String::new();
154    for ch in prompt.chars() {
155        if out.len() >= 40 {
156            break;
157        }
158        if ch.is_ascii_alphanumeric() {
159            out.push(ch.to_ascii_lowercase());
160        } else if !out.ends_with('-') && !out.is_empty() {
161            out.push('-');
162        }
163    }
164    let trimmed = out.trim_matches('-').to_string();
165    if trimmed.is_empty() {
166        "image".to_string()
167    } else {
168        trimmed
169    }
170}
171
172#[async_trait]
173impl Tool for ImageGenTool {
174    fn name(&self) -> &str {
175        IMAGE_GEN
176    }
177    fn description(&self) -> &str {
178        "Generate an image from a text prompt using the session's provider and write it into \
179         the working directory. Returns the path of the written file."
180    }
181    fn parameters(&self) -> Value {
182        json!({
183            "type": "object",
184            "properties": {
185                "prompt": {"type": "string", "description": "What the image should show."},
186                "path": {
187                    "type": "string",
188                    "description": "Where to write the file (relative to the working \
189                                    directory). Defaults to a name derived from the prompt."
190                },
191                "size": {
192                    "type": "string",
193                    "description": "Requested size, e.g. \"1024x1024\". Provider default when \
194                                    omitted."
195                },
196                "model": {"type": "string", "description": "Image model to use."}
197            },
198            "required": ["prompt"],
199            "additionalProperties": false
200        })
201    }
202    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<String> {
203        let a: ImageGenArgs =
204            serde_json::from_value(args).map_err(|e| Error::InvalidArguments {
205                tool: self.name().to_string(),
206                message: e.to_string(),
207            })?;
208        if a.prompt.trim().is_empty() {
209            return Err(Error::InvalidArguments {
210                tool: self.name().to_string(),
211                message: "prompt must not be empty".to_string(),
212            });
213        }
214        let url = self.endpoint();
215        ctx.check_network(&url)?;
216
217        // Resolve the destination BEFORE spending a provider call on an
218        // image the sandbox would refuse to write anyway.
219        let rel = a
220            .path
221            .clone()
222            .unwrap_or_else(|| format!("{}.png", slug(&a.prompt)));
223        let dest: PathBuf = ctx.resolve(&rel);
224        ctx.check_write(&dest)?;
225
226        let mut body = json!({
227            "model": a.model.clone().unwrap_or_else(|| self.default_model.clone()),
228            "prompt": a.prompt,
229            "n": 1,
230            "response_format": "b64_json",
231        });
232        if let Some(size) = &a.size {
233            body["size"] = json!(size);
234        }
235        let client = reqwest::Client::builder()
236            .timeout(Duration::from_secs(IMAGE_TIMEOUT_SECS))
237            .redirect(crate::tools::network_checked_redirect_policy(
238                ctx.network_policy.clone(),
239                ctx.permission_rules.clone(),
240            ))
241            .build()
242            .map_err(|e| Error::tool(self.name(), e.to_string()))?;
243        let mut request = client.post(&url).json(&body);
244        if let Some(key) = self.resolved_key() {
245            request = request.bearer_auth(key);
246        }
247        let response = request
248            .send()
249            .await
250            .map_err(|e| Error::tool(self.name(), format!("image request failed: {e}")))?;
251        let status = response.status();
252        let text = response
253            .text()
254            .await
255            .map_err(|e| Error::tool(self.name(), format!("reading image response: {e}")))?;
256        if matches!(status.as_u16(), 404 | 405 | 501) {
257            return Err(Error::tool(
258                self.name(),
259                format!(
260                    "unsupported_action: the configured provider exposes no image endpoint \
261                     ({url} answered {status}). Image generation is unavailable in this \
262                     session — say so rather than describing an image you did not make."
263                ),
264            ));
265        }
266        if !status.is_success() {
267            let mut detail: String = text.chars().take(400).collect();
268            if detail.is_empty() {
269                detail = "(empty body)".to_string();
270            }
271            return Err(Error::tool(
272                self.name(),
273                format!("image endpoint returned {status}: {detail}"),
274            ));
275        }
276        let parsed: Value = serde_json::from_str(&text)
277            .map_err(|e| Error::tool(self.name(), format!("image response is not JSON: {e}")))?;
278        let first = parsed
279            .get("data")
280            .and_then(|d| d.as_array())
281            .and_then(|d| d.first())
282            .ok_or_else(|| {
283                Error::tool(
284                    self.name(),
285                    "image response carried no `data[0]` entry".to_string(),
286                )
287            })?;
288        let b64 = first.get("b64_json").and_then(|v| v.as_str());
289        let bytes = match b64 {
290            Some(b64) => base64_decode(b64).ok_or_else(|| {
291                Error::tool(self.name(), "image response's b64_json is not valid base64")
292            })?,
293            None => {
294                let Some(remote) = first.get("url").and_then(|v| v.as_str()) else {
295                    return Err(Error::tool(
296                        self.name(),
297                        "image response carried neither `b64_json` nor `url`",
298                    ));
299                };
300                ctx.check_network(remote)?;
301                let fetched = client.get(remote).send().await.map_err(|e| {
302                    Error::tool(self.name(), format!("downloading the image failed: {e}"))
303                })?;
304                if !fetched.status().is_success() {
305                    return Err(Error::tool(
306                        self.name(),
307                        format!("downloading the image returned {}", fetched.status()),
308                    ));
309                }
310                fetched
311                    .bytes()
312                    .await
313                    .map_err(|e| Error::tool(self.name(), format!("reading the image bytes: {e}")))?
314                    .to_vec()
315            }
316        };
317        if bytes.is_empty() {
318            return Err(Error::tool(
319                self.name(),
320                "the provider returned no image data",
321            ));
322        }
323        if bytes.len() > MAX_IMAGE_BYTES {
324            return Err(Error::tool(
325                self.name(),
326                format!(
327                    "the provider returned {} bytes, over this tool's {MAX_IMAGE_BYTES}-byte \
328                     ceiling",
329                    bytes.len()
330                ),
331            ));
332        }
333        if let Some(parent) = dest.parent() {
334            std::fs::create_dir_all(parent)
335                .map_err(|e| Error::tool(self.name(), format!("creating {parent:?}: {e}")))?;
336        }
337        // The write-path observers (checkpoint/formatters/lsp) see this
338        // write exactly like any other file-writing tool's.
339        if let Some(observer) = &ctx.write_observer {
340            observer.before_write(&dest).await;
341        }
342        std::fs::write(&dest, &bytes)
343            .map_err(|e| Error::tool(self.name(), format!("writing {}: {e}", dest.display())))?;
344        if let Some(observer) = &ctx.write_observer {
345            observer.after_write(&dest).await;
346        }
347        Ok(format!(
348            "Wrote {} ({} bytes) from the provider's image endpoint.",
349            dest.display(),
350            bytes.len()
351        ))
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358    use std::io::{Read, Write};
359    use std::net::TcpListener;
360
361    /// A single-shot HTTP server: answers ONE request with `status` and
362    /// `body`, then stops. No network leaves the loopback interface.
363    fn serve_once(status: u16, body: &'static str) -> (String, std::thread::JoinHandle<()>) {
364        let listener = TcpListener::bind("127.0.0.1:0").expect("bind loopback");
365        let port = listener.local_addr().unwrap().port();
366        let handle = std::thread::spawn(move || {
367            if let Ok((mut stream, _)) = listener.accept() {
368                let mut buf = [0u8; 8192];
369                let _ = stream.read(&mut buf);
370                let response = format!(
371                    "HTTP/1.1 {status} X\r\nContent-Type: application/json\r\nContent-Length: \
372                     {}\r\nConnection: close\r\n\r\n{body}",
373                    body.len()
374                );
375                let _ = stream.write_all(response.as_bytes());
376                let _ = stream.flush();
377            }
378        });
379        (format!("http://127.0.0.1:{port}/v1"), handle)
380    }
381
382    fn tool_for(base: &str) -> ImageGenTool {
383        ImageGenTool::new(
384            base,
385            Some("test-key".to_string()),
386            "SUPERCODE_TEST_KEY_UNSET",
387        )
388    }
389
390    #[test]
391    fn base64_round_trips_against_the_builtin_encoder() {
392        for bytes in [b"".to_vec(), b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
393            // The PNG magic bytes exercise the non-ASCII path too.
394            let mut sample = bytes.clone();
395            sample.extend_from_slice(&[0x89, 0x50, 0x4e, 0x47]);
396            let encoded = crate::tools::builtins::base64_encode(&sample);
397            assert_eq!(base64_decode(&encoded).as_deref(), Some(&sample[..]));
398        }
399        assert!(base64_decode("not base64!!").is_none());
400    }
401
402    #[test]
403    fn the_endpoint_is_the_providers_own_images_route() {
404        assert_eq!(
405            tool_for("https://example.test/v1/").endpoint(),
406            "https://example.test/v1/images/generations"
407        );
408    }
409
410    #[tokio::test]
411    async fn a_generated_image_lands_in_the_working_directory() {
412        let dir = std::env::temp_dir().join(format!("bp3-image-{}", std::process::id()));
413        std::fs::create_dir_all(&dir).unwrap();
414        // base64 of the four PNG magic bytes.
415        let (base, handle) = serve_once(200, r#"{"data":[{"b64_json":"iVBORw=="}]}"#);
416        let ctx = ToolContext::new(&dir);
417        let out = tool_for(&base)
418            .execute(json!({"prompt": "a red square", "path": "out.png"}), &ctx)
419            .await
420            .unwrap();
421        handle.join().unwrap();
422        assert!(out.contains("out.png"), "{out}");
423        let written = std::fs::read(dir.join("out.png")).unwrap();
424        assert_eq!(&written[..4], &[0x89, 0x50, 0x4e, 0x47]);
425        let _ = std::fs::remove_dir_all(&dir);
426    }
427
428    #[tokio::test]
429    async fn a_provider_without_the_route_reports_unsupported_action() {
430        let (base, handle) = serve_once(404, r#"{"error":"no such route"}"#);
431        let ctx = ToolContext::new(std::env::temp_dir());
432        let err = tool_for(&base)
433            .execute(json!({"prompt": "anything"}), &ctx)
434            .await
435            .expect_err("404 must not be treated as success");
436        handle.join().unwrap();
437        assert!(err.to_string().contains("unsupported_action"), "{err}");
438    }
439
440    #[tokio::test]
441    async fn a_read_only_sandbox_refuses_before_calling_the_provider() {
442        let mut ctx = ToolContext::new(std::env::temp_dir());
443        ctx.sandbox = crate::tools::SandboxPolicy::ReadOnly;
444        let err = tool_for("http://127.0.0.1:1/v1")
445            .execute(json!({"prompt": "a red square"}), &ctx)
446            .await
447            .expect_err("a read-only sandbox must refuse");
448        assert!(err.to_string().contains("read-only"), "{err}");
449    }
450
451    #[test]
452    fn prompt_slugs_are_filename_safe() {
453        assert_eq!(slug("A Red Square!"), "a-red-square");
454        assert_eq!(slug("***"), "image");
455    }
456}
supercode_harness/tools/image_gen.rs

supercode_harness/tools/
image_gen.rs
