supercode-harness 0.4.20

//! BP-3 (catalog row "Image generation tool", cx§1's `image_gen` feature):
//! a client-issued call to the PROVIDER's image endpoint.
//!
//! The row's semantics is "provider-executed image creation": the provider
//! renders the image, and this tool is the client side of that call — the
//! OpenAI-compatible `POST {base_url}/images/generations` every
//! OpenAI-shaped endpoint (including the one `crate::Config::base_url`
//! already points the chat provider at) exposes. The bytes come back
//! base64-encoded, are written into the working directory through the same
//! [`ToolContext::check_write`] sandbox check every other write-capable
//! tool passes, and the model gets the path.
//!
//! **UnsupportedAction is a real answer.** Not every OpenAI-compatible
//! gateway implements the images route. When the endpoint answers 404 / 405
//! / 501, the tool says so in exactly those terms — "unsupported_action:
//! this provider exposes no image endpoint" — rather than retrying, falling
//! back to another vendor, or inventing a file. Nothing here reaches for a
//! provider the session was not already configured with.

use std::path::PathBuf;
use std::time::Duration;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::{json, Value};

use crate::error::{Error, Result};
use crate::tools::{Tool, ToolContext};

/// Registered name of the image tool (Codex's own spelling).
pub const IMAGE_GEN: &str = "image_gen";

/// The image model used when neither the call nor the config names one.
pub const DEFAULT_IMAGE_MODEL: &str = "gpt-image-1";

/// Request timeout: image generation is slow compared with a chat call.
const IMAGE_TIMEOUT_SECS: u64 = 180;

/// Refuse to write an image larger than this (a runaway response must not
/// fill the disk).
const MAX_IMAGE_BYTES: usize = 32 * 1024 * 1024;

/// `image_gen` — generate an image through the session's own provider.
///
/// Constructed from the resolved [`crate::Config`] by
/// [`crate::tools::ToolRegistry::from_config`], so it inherits exactly the
/// endpoint and credential the chat provider uses; it never reads a second
/// vendor's configuration.
#[derive(Debug, Clone)]
pub struct ImageGenTool {
    base_url: String,
    api_key: Option<String>,
    api_key_env: String,
    default_model: String,
}

impl ImageGenTool {
    /// Build the tool from the session's provider settings.
    pub fn new(
        base_url: impl Into<String>,
        api_key: Option<String>,
        api_key_env: impl Into<String>,
    ) -> Self {
        ImageGenTool {
            base_url: base_url.into(),
            api_key,
            api_key_env: api_key_env.into(),
            default_model: DEFAULT_IMAGE_MODEL.to_string(),
        }
    }

    /// Override the default image model (the `model` argument still wins).
    pub fn with_default_model(mut self, model: impl Into<String>) -> Self {
        self.default_model = model.into();
        self
    }

    /// The endpoint this tool posts to.
    pub fn endpoint(&self) -> String {
        format!("{}/images/generations", self.base_url.trim_end_matches('/'))
    }

    fn resolved_key(&self) -> Option<String> {
        self.api_key.clone().filter(|k| !k.is_empty()).or_else(|| {
            std::env::var(&self.api_key_env)
                .ok()
                .filter(|k| !k.is_empty())
        })
    }
}

#[derive(Debug, Deserialize)]
struct ImageGenArgs {
    prompt: String,
    #[serde(default)]
    path: Option<String>,
    #[serde(default)]
    size: Option<String>,
    #[serde(default)]
    model: Option<String>,
}

/// Standard base64 decoder — the mirror of `builtins::base64_encode`, same
/// "small parser over a crate" precedent. `None` on any non-alphabet byte.
fn base64_decode(input: &str) -> Option<Vec<u8>> {
    fn digit(byte: u8) -> Option<u8> {
        match byte {
            b'A'..=b'Z' => Some(byte - b'A'),
            b'a'..=b'z' => Some(byte - b'a' + 26),
            b'0'..=b'9' => Some(byte - b'0' + 52),
            b'+' => Some(62),
            b'/' => Some(63),
            _ => None,
        }
    }
    let mut out = Vec::with_capacity(input.len() / 4 * 3 + 3);
    let mut chunk = [0u8; 4];
    let mut len = 0usize;
    let mut padding = 0usize;
    for byte in input.bytes().filter(|b| !b.is_ascii_whitespace()) {
        if byte == b'=' {
            padding += 1;
            chunk[len] = 0;
        } else {
            chunk[len] = digit(byte)?;
        }
        len += 1;
        if len == 4 {
            let value = ((chunk[0] as u32) << 18)
                | ((chunk[1] as u32) << 12)
                | ((chunk[2] as u32) << 6)
                | chunk[3] as u32;
            out.push((value >> 16) as u8);
            if padding < 2 {
                out.push((value >> 8) as u8);
            }
            if padding < 1 {
                out.push(value as u8);
            }
            len = 0;
            padding = 0;
        }
    }
    if len == 0 {
        Some(out)
    } else {
        None
    }
}

/// A filename-safe slug of the prompt, for the default output path.
fn slug(prompt: &str) -> String {
    let mut out = String::new();
    for ch in prompt.chars() {
        if out.len() >= 40 {
            break;
        }
        if ch.is_ascii_alphanumeric() {
            out.push(ch.to_ascii_lowercase());
        } else if !out.ends_with('-') && !out.is_empty() {
            out.push('-');
        }
    }
    let trimmed = out.trim_matches('-').to_string();
    if trimmed.is_empty() {
        "image".to_string()
    } else {
        trimmed
    }
}

#[async_trait]
impl Tool for ImageGenTool {
    fn name(&self) -> &str {
        IMAGE_GEN
    }
    fn description(&self) -> &str {
        "Generate an image from a text prompt using the session's provider and write it into \
         the working directory. Returns the path of the written file."
    }
    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "prompt": {"type": "string", "description": "What the image should show."},
                "path": {
                    "type": "string",
                    "description": "Where to write the file (relative to the working \
                                    directory). Defaults to a name derived from the prompt."
                },
                "size": {
                    "type": "string",
                    "description": "Requested size, e.g. \"1024x1024\". Provider default when \
                                    omitted."
                },
                "model": {"type": "string", "description": "Image model to use."}
            },
            "required": ["prompt"],
            "additionalProperties": false
        })
    }
    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<String> {
        let a: ImageGenArgs =
            serde_json::from_value(args).map_err(|e| Error::InvalidArguments {
                tool: self.name().to_string(),
                message: e.to_string(),
            })?;
        if a.prompt.trim().is_empty() {
            return Err(Error::InvalidArguments {
                tool: self.name().to_string(),
                message: "prompt must not be empty".to_string(),
            });
        }
        let url = self.endpoint();
        ctx.check_network(&url)?;

        // Resolve the destination BEFORE spending a provider call on an
        // image the sandbox would refuse to write anyway.
        let rel = a
            .path
            .clone()
            .unwrap_or_else(|| format!("{}.png", slug(&a.prompt)));
        let dest: PathBuf = ctx.resolve(&rel);
        ctx.check_write(&dest)?;

        let mut body = json!({
            "model": a.model.clone().unwrap_or_else(|| self.default_model.clone()),
            "prompt": a.prompt,
            "n": 1,
            "response_format": "b64_json",
        });
        if let Some(size) = &a.size {
            body["size"] = json!(size);
        }
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(IMAGE_TIMEOUT_SECS))
            .redirect(crate::tools::network_checked_redirect_policy(
                ctx.network_policy.clone(),
                ctx.permission_rules.clone(),
            ))
            .build()
            .map_err(|e| Error::tool(self.name(), e.to_string()))?;
        let mut request = client.post(&url).json(&body);
        if let Some(key) = self.resolved_key() {
            request = request.bearer_auth(key);
        }
        let response = request
            .send()
            .await
            .map_err(|e| Error::tool(self.name(), format!("image request failed: {e}")))?;
        let status = response.status();
        let text = response
            .text()
            .await
            .map_err(|e| Error::tool(self.name(), format!("reading image response: {e}")))?;
        if matches!(status.as_u16(), 404 | 405 | 501) {
            return Err(Error::tool(
                self.name(),
                format!(
                    "unsupported_action: the configured provider exposes no image endpoint \
                     ({url} answered {status}). Image generation is unavailable in this \
                     session — say so rather than describing an image you did not make."
                ),
            ));
        }
        if !status.is_success() {
            let mut detail: String = text.chars().take(400).collect();
            if detail.is_empty() {
                detail = "(empty body)".to_string();
            }
            return Err(Error::tool(
                self.name(),
                format!("image endpoint returned {status}: {detail}"),
            ));
        }
        let parsed: Value = serde_json::from_str(&text)
            .map_err(|e| Error::tool(self.name(), format!("image response is not JSON: {e}")))?;
        let first = parsed
            .get("data")
            .and_then(|d| d.as_array())
            .and_then(|d| d.first())
            .ok_or_else(|| {
                Error::tool(
                    self.name(),
                    "image response carried no `data[0]` entry".to_string(),
                )
            })?;
        let b64 = first.get("b64_json").and_then(|v| v.as_str());
        let bytes = match b64 {
            Some(b64) => base64_decode(b64).ok_or_else(|| {
                Error::tool(self.name(), "image response's b64_json is not valid base64")
            })?,
            None => {
                let Some(remote) = first.get("url").and_then(|v| v.as_str()) else {
                    return Err(Error::tool(
                        self.name(),
                        "image response carried neither `b64_json` nor `url`",
                    ));
                };
                ctx.check_network(remote)?;
                let fetched = client.get(remote).send().await.map_err(|e| {
                    Error::tool(self.name(), format!("downloading the image failed: {e}"))
                })?;
                if !fetched.status().is_success() {
                    return Err(Error::tool(
                        self.name(),
                        format!("downloading the image returned {}", fetched.status()),
                    ));
                }
                fetched
                    .bytes()
                    .await
                    .map_err(|e| Error::tool(self.name(), format!("reading the image bytes: {e}")))?
                    .to_vec()
            }
        };
        if bytes.is_empty() {
            return Err(Error::tool(
                self.name(),
                "the provider returned no image data",
            ));
        }
        if bytes.len() > MAX_IMAGE_BYTES {
            return Err(Error::tool(
                self.name(),
                format!(
                    "the provider returned {} bytes, over this tool's {MAX_IMAGE_BYTES}-byte \
                     ceiling",
                    bytes.len()
                ),
            ));
        }
        if let Some(parent) = dest.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| Error::tool(self.name(), format!("creating {parent:?}: {e}")))?;
        }
        // The write-path observers (checkpoint/formatters/lsp) see this
        // write exactly like any other file-writing tool's.
        if let Some(observer) = &ctx.write_observer {
            observer.before_write(&dest).await;
        }
        std::fs::write(&dest, &bytes)
            .map_err(|e| Error::tool(self.name(), format!("writing {}: {e}", dest.display())))?;
        if let Some(observer) = &ctx.write_observer {
            observer.after_write(&dest).await;
        }
        Ok(format!(
            "Wrote {} ({} bytes) from the provider's image endpoint.",
            dest.display(),
            bytes.len()
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpListener;

    /// A single-shot HTTP server: answers ONE request with `status` and
    /// `body`, then stops. No network leaves the loopback interface.
    fn serve_once(status: u16, body: &'static str) -> (String, std::thread::JoinHandle<()>) {
        let listener = TcpListener::bind("127.0.0.1:0").expect("bind loopback");
        let port = listener.local_addr().unwrap().port();
        let handle = std::thread::spawn(move || {
            if let Ok((mut stream, _)) = listener.accept() {
                let mut buf = [0u8; 8192];
                let _ = stream.read(&mut buf);
                let response = format!(
                    "HTTP/1.1 {status} X\r\nContent-Type: application/json\r\nContent-Length: \
                     {}\r\nConnection: close\r\n\r\n{body}",
                    body.len()
                );
                let _ = stream.write_all(response.as_bytes());
                let _ = stream.flush();
            }
        });
        (format!("http://127.0.0.1:{port}/v1"), handle)
    }

    fn tool_for(base: &str) -> ImageGenTool {
        ImageGenTool::new(
            base,
            Some("test-key".to_string()),
            "SUPERCODE_TEST_KEY_UNSET",
        )
    }

    #[test]
    fn base64_round_trips_against_the_builtin_encoder() {
        for bytes in [b"".to_vec(), b"a".to_vec(), b"ab".to_vec(), b"abc".to_vec()] {
            // The PNG magic bytes exercise the non-ASCII path too.
            let mut sample = bytes.clone();
            sample.extend_from_slice(&[0x89, 0x50, 0x4e, 0x47]);
            let encoded = crate::tools::builtins::base64_encode(&sample);
            assert_eq!(base64_decode(&encoded).as_deref(), Some(&sample[..]));
        }
        assert!(base64_decode("not base64!!").is_none());
    }

    #[test]
    fn the_endpoint_is_the_providers_own_images_route() {
        assert_eq!(
            tool_for("https://example.test/v1/").endpoint(),
            "https://example.test/v1/images/generations"
        );
    }

    #[tokio::test]
    async fn a_generated_image_lands_in_the_working_directory() {
        let dir = std::env::temp_dir().join(format!("bp3-image-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        // base64 of the four PNG magic bytes.
        let (base, handle) = serve_once(200, r#"{"data":[{"b64_json":"iVBORw=="}]}"#);
        let ctx = ToolContext::new(&dir);
        let out = tool_for(&base)
            .execute(json!({"prompt": "a red square", "path": "out.png"}), &ctx)
            .await
            .unwrap();
        handle.join().unwrap();
        assert!(out.contains("out.png"), "{out}");
        let written = std::fs::read(dir.join("out.png")).unwrap();
        assert_eq!(&written[..4], &[0x89, 0x50, 0x4e, 0x47]);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn a_provider_without_the_route_reports_unsupported_action() {
        let (base, handle) = serve_once(404, r#"{"error":"no such route"}"#);
        let ctx = ToolContext::new(std::env::temp_dir());
        let err = tool_for(&base)
            .execute(json!({"prompt": "anything"}), &ctx)
            .await
            .expect_err("404 must not be treated as success");
        handle.join().unwrap();
        assert!(err.to_string().contains("unsupported_action"), "{err}");
    }

    #[tokio::test]
    async fn a_read_only_sandbox_refuses_before_calling_the_provider() {
        let mut ctx = ToolContext::new(std::env::temp_dir());
        ctx.sandbox = crate::tools::SandboxPolicy::ReadOnly;
        let err = tool_for("http://127.0.0.1:1/v1")
            .execute(json!({"prompt": "a red square"}), &ctx)
            .await
            .expect_err("a read-only sandbox must refuse");
        assert!(err.to_string().contains("read-only"), "{err}");
    }

    #[test]
    fn prompt_slugs_are_filename_safe() {
        assert_eq!(slug("A Red Square!"), "a-red-square");
        assert_eq!(slug("***"), "image");
    }
}