Skip to main content

atman_runtime/tools/
image.rs

1use crate::error::RuntimeError;
2use crate::message::{Message, MessageOrigin, MessagePart, MessageRole};
3use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
4use crate::value::Value;
5
6pub struct ImageRead;
7
8impl Tool for ImageRead {
9    fn name(&self) -> &str {
10        "image.read"
11    }
12
13    fn tier(&self) -> Tier {
14        Tier::Zero
15    }
16
17    fn description(&self) -> Option<&str> {
18        Some(
19            "Read a local PNG, JPEG, GIF, or WebP image and attach it as visual input for the next model call. The image must be at most 20 MiB.",
20        )
21    }
22
23    fn input_schema(&self) -> serde_json::Value {
24        serde_json::json!({
25            "type": "object",
26            "properties": {
27                "path": {
28                    "type": "string",
29                    "description": "Local image path. Relative paths resolve inside the active workspace."
30                }
31            },
32            "required": ["path"],
33            "additionalProperties": false
34        })
35    }
36
37    fn invocation_provenance(
38        &self,
39        args: &ToolArgs,
40        ctx: &ToolCtx,
41    ) -> Result<crate::permission::ResourceProvenance, RuntimeError> {
42        crate::permission::ResourceProvenance::for_ctx(ctx).with_path(ctx, &extract_path(args)?)
43    }
44
45    fn model_followups(&self, result: &Value, _ctx: &ToolCtx) -> Vec<Message> {
46        match result {
47            Value::Message(message)
48                if message
49                    .parts
50                    .iter()
51                    .any(|part| matches!(part, MessagePart::Image { .. })) =>
52            {
53                vec![message.clone()]
54            }
55            _ => Vec::new(),
56        }
57    }
58
59    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
60        Box::pin(async move {
61            let path = ctx.resolve_path(&extract_path(&args)?)?;
62            let byte_len = tokio::fs::metadata(&path)
63                .await
64                .map_err(|error| {
65                    RuntimeError::ToolFailed(format!("image.read({}): {error}", path.display()))
66                })?
67                .len();
68            let source = match ctx.session_runtime.as_ref() {
69                Some(session) => session.import_image_path(&path)?,
70                None => crate::attachment_store::AttachmentStore::at(
71                    ctx.session_dir
72                        .as_deref()
73                        .unwrap_or_else(|| std::path::Path::new("")),
74                )
75                .import_path(&path)?,
76            };
77            let canonical = std::fs::canonicalize(&path).unwrap_or_else(|_| path.clone());
78            ctx.note_read(&canonical);
79            let label = format!(
80                "Loaded image {} ({}; {} bytes) as visual input for the next model call.",
81                crate::attachment_store::display_name(&source),
82                source.media_type,
83                byte_len
84            );
85            Ok(Value::Message(Message {
86                role: MessageRole::User,
87                parts: vec![
88                    MessagePart::Text { text: label },
89                    MessagePart::Image { source },
90                ],
91                turn_id: ctx
92                    .turn_id
93                    .clone()
94                    .unwrap_or_else(crate::event::TurnId::now),
95                origin: MessageOrigin::Internal,
96            }))
97        })
98    }
99}
100
101fn extract_path(args: &ToolArgs) -> Result<std::path::PathBuf, RuntimeError> {
102    let value = args.named("path").or_else(|| args.positional.first());
103    match value {
104        Some(Value::Str(path)) if !path.is_empty() => Ok(path.into()),
105        Some(Value::Path(path)) => Ok(path.clone()),
106        Some(other) => Err(RuntimeError::TypeMismatch {
107            expected: "non-empty image path".into(),
108            actual: other.kind_name().into(),
109        }),
110        None => Err(RuntimeError::MissingArg("image.read.path".into())),
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    const PNG_HEADER: &[u8] = b"\x89PNG\r\n\x1a\n\0\0\0\r";
119
120    #[tokio::test]
121    async fn read_returns_internal_model_image_without_base64_text() {
122        let dir = tempfile::tempdir().unwrap();
123        let path = dir.path().join("pixel.png");
124        std::fs::write(&path, PNG_HEADER).unwrap();
125        let reads = std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
126        let ctx = ToolCtx::new().with_read_files(reads.clone());
127
128        let value = ImageRead
129            .call(
130                ToolArgs {
131                    positional: Vec::new(),
132                    named: vec![("path".into(), Value::Str(path.display().to_string()))],
133                },
134                &ctx,
135            )
136            .await
137            .unwrap();
138
139        let Value::Message(message) = &value else {
140            panic!("expected image message")
141        };
142        assert_eq!(message.role, MessageRole::User);
143        assert_eq!(message.origin, MessageOrigin::Internal);
144        assert!(
145            matches!(message.parts.as_slice(), [MessagePart::Text { text }, MessagePart::Image { source }] if !text.contains("base64") && source.media_type == "image/png")
146        );
147        assert_eq!(
148            ImageRead.model_followups(&value, &ctx),
149            vec![message.clone()]
150        );
151        assert!(
152            reads
153                .lock()
154                .unwrap()
155                .contains(&std::fs::canonicalize(path).unwrap())
156        );
157    }
158
159    #[tokio::test]
160    async fn read_rejects_non_image_content() {
161        let dir = tempfile::tempdir().unwrap();
162        let path = dir.path().join("notes.txt");
163        std::fs::write(&path, b"not an image").unwrap();
164        let ctx = ToolCtx::new();
165
166        let error = ImageRead
167            .call(
168                ToolArgs {
169                    positional: vec![Value::Str(path.display().to_string())],
170                    named: Vec::new(),
171                },
172                &ctx,
173            )
174            .await
175            .unwrap_err();
176
177        assert!(error.to_string().contains("unsupported image format"));
178    }
179}