mobius/middleware/
artifacts.rs

1//! Agent-published session files and their frontend presentation.
2
3use std::path::Path;
4use std::sync::Arc;
5
6use serde::Deserialize;
7use serde_json::Value;
8
9use super::tools::{ApprovalRequirement, Catalog, Tool, ToolContext};
10use super::{Middleware, PromptSection, RuntimeContext};
11use crate::backend::model::ToolDefinition;
12use crate::backend::sandbox::MAX_BINARY_FILE_BYTES;
13use crate::backend::session_files::{SessionFileStore, media_type};
14use crate::protocol::{EventMsg, FrontendBlock, FrontendContribution};
15use crate::{BoxFuture, Error, Result};
16
17mod text {
18    #[derive(serde::Deserialize)]
19    #[serde(deny_unknown_fields)]
20    pub(super) struct Definition {
21        pub(super) send_artifact: crate::middleware::tools::ToolSpec,
22        pub(super) default_enabled: bool,
23        pub(super) manifest_description: String,
24        pub(super) manifest_label: String,
25        pub(super) prompt_main: String,
26        pub(super) render_sent_prefix: String,
27        pub(super) error_path_not_file: String,
28        pub(super) error_path_or_file_id: String,
29    }
30    crate::embedded_config! { pub(super) static DEFINITION: Definition = include_str!("artifacts.toml"); }
31}
32super::manifest::middleware_manifest! {
33/// Configuration metadata for agent-published files.
34    "artifacts", text::DEFINITION, required: false, settings: &[]
35}
36
37/// Publishes workspace files as session-bound assistant files.
38pub struct Artifacts {
39    store: Arc<SessionFileStore>,
40}
41
42impl Artifacts {
43    #[must_use]
44    /// Creates a new instance.
45    pub fn new(store: SessionFileStore) -> Self {
46        Self {
47            store: Arc::new(store),
48        }
49    }
50}
51
52impl Middleware for Artifacts {
53    fn name(&self) -> &'static str {
54        MANIFEST.id
55    }
56
57    fn register(&self, catalog: &mut Catalog, runtime: &RuntimeContext) -> Result<()> {
58        catalog.register(Arc::new(SendArtifact {
59            store: Arc::clone(&self.store),
60            session_id: runtime.session_id.clone(),
61        }))
62    }
63
64    fn prompt_section(&self, _runtime: &RuntimeContext) -> Result<Option<PromptSection>> {
65        Ok(Some(PromptSection::new(
66            text::DEFINITION.prompt_main.as_str(),
67        )))
68    }
69
70    fn frontend(&self, _session_id: &str) -> FrontendContribution {
71        FrontendContribution {
72            capability: MANIFEST.id.into(),
73            ..FrontendContribution::default()
74        }
75    }
76
77    fn render(&self, event: &EventMsg, _session_id: &str) -> Option<FrontendBlock> {
78        let mut block = text::DEFINITION.send_artifact.render(event)?;
79        let EventMsg::ToolCallEnd(result) = event else {
80            return Some(block);
81        };
82        if result.is_error {
83            return Some(block);
84        }
85        let Some(file) = result.output.files().next().cloned() else {
86            return Some(block);
87        };
88        block.update = crate::protocol::FrontendBlockUpdate::Replace;
89        block.state = crate::protocol::FrontendBlockState::Complete;
90        block.role = crate::protocol::FrontendBlockRole::Artifact;
91        block.title = format!(
92            "{}{}",
93            text::DEFINITION.render_sent_prefix.as_str(),
94            file.name
95        );
96        block.text.clear();
97        block.content = Default::default();
98        block.files = vec![file];
99        Some(block)
100    }
101}
102
103#[derive(Deserialize)]
104#[serde(deny_unknown_fields)]
105struct SendArtifactArgs {
106    path: Option<String>,
107    file_id: Option<String>,
108}
109
110struct SendArtifact {
111    store: Arc<SessionFileStore>,
112    session_id: String,
113}
114
115impl Tool for SendArtifact {
116    fn definition(&self) -> ToolDefinition {
117        text::DEFINITION.send_artifact.tool.clone()
118    }
119
120    fn approval(&self) -> ApprovalRequirement {
121        ApprovalRequirement::Always
122    }
123
124    fn call<'a>(
125        &'a self,
126        context: ToolContext,
127        arguments: Value,
128    ) -> BoxFuture<'a, Result<crate::protocol::ToolResponse>> {
129        Box::pin(async move {
130            let arguments: SendArtifactArgs = serde_json::from_value(arguments)?;
131            let file = match (arguments.path, arguments.file_id) {
132                (Some(path), None) => {
133                    let name = Path::new(&path)
134                        .file_name()
135                        .and_then(|name| name.to_str())
136                        .ok_or_else(|| {
137                            Error::Tool(text::DEFINITION.error_path_not_file.as_str().into())
138                        })?
139                        .to_string();
140                    let bytes = context
141                        .sandbox
142                        .read_bytes(&path, MAX_BINARY_FILE_BYTES, &context.permissions)
143                        .await?;
144                    let media_type = media_type(&name).into();
145                    self.store
146                        .publish_artifact(&self.session_id, name, media_type, &bytes)
147                        .await?
148                }
149                (None, Some(id)) => {
150                    let file = self.store.file_reference(&self.session_id, &id).await?;
151                    self.store
152                        .publish_reference(&self.session_id, &file)
153                        .await?
154                }
155                _ => {
156                    return Err(Error::Tool(
157                        text::DEFINITION.error_path_or_file_id.as_str().into(),
158                    ));
159                }
160            };
161            Ok(crate::protocol::ToolResponse {
162                content: crate::protocol::ToolContent(vec![crate::protocol::ContentPart::File {
163                    file,
164                }]),
165                is_error: false,
166            })
167        })
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use crate::protocol::SessionFileReference;
175    use std::collections::BTreeMap;
176
177    use crate::backend::checkpoint::{CheckpointStore, sqlite::SqliteCheckpoint};
178    use crate::backend::sandbox::local::LocalSandbox;
179    use crate::backend::sandbox::{ApprovalPolicy, NetworkAccess, Sandbox, SandboxPermissions};
180    use crate::protocol::{SessionContext, ToolCallEndEvent};
181
182    #[test]
183    fn one_middleware_registers_tools_for_multiple_sessions() {
184        let state = tempfile::tempdir().expect("state");
185        let checkpoints: Arc<dyn CheckpointStore> = Arc::new(
186            SqliteCheckpoint::new(state.path().join("checkpoints.sqlite3"))
187                .expect("checkpoint store"),
188        );
189        let middleware = Artifacts::new(SessionFileStore::new(state.path(), None));
190        let runtime = RuntimeContext {
191            children: crate::agent::ChildAgents::default(),
192            sender: crate::agent::test_sender(),
193            checkpoints,
194            session_id: "session-a".into(),
195            model_route: "model".into(),
196            model: "model".into(),
197            approval_policy: crate::backend::sandbox::ApprovalPolicy::Ask,
198            session_context: SessionContext::default(),
199            metadata: BTreeMap::new(),
200            role: crate::agent::AgentRole::Main,
201            frontend: Arc::new(|_| Ok(())),
202        };
203        let mut second_runtime = runtime.clone();
204        second_runtime.session_id = "session-b".into();
205
206        let mut first = Catalog::default();
207        middleware
208            .register(&mut first, &runtime)
209            .expect("first session catalog");
210        let mut second = Catalog::default();
211        middleware
212            .register(&mut second, &second_runtime)
213            .expect("second session catalog");
214
215        assert_eq!(
216            first.registered_definitions(),
217            second.registered_definitions()
218        );
219    }
220
221    #[tokio::test]
222    async fn send_artifact_publishes_binary_outside_the_upload_list() {
223        let workspace = tempfile::tempdir().expect("workspace");
224        std::fs::write(workspace.path().join("report.xlsx"), [0, 255, 1, 254]).expect("file");
225        let state = tempfile::tempdir().expect("state");
226        let store = SessionFileStore::new(state.path(), None);
227        let tool = SendArtifact {
228            store: Arc::new(store.clone()),
229            session_id: "session-a".into(),
230        };
231        let sandbox = Arc::new(Sandbox::new(
232            Arc::new(LocalSandbox::new(workspace.path()).expect("local sandbox")),
233            ApprovalPolicy::Ask,
234        ));
235        let permissions = SandboxPermissions::restore(
236            "session-a",
237            crate::backend::sandbox::SandboxMode::WorkspaceWrite,
238            NetworkAccess::Denied,
239            ["call".into()],
240        )
241        .for_call("call");
242
243        let output = tool
244            .call(
245                ToolContext::new(sandbox, permissions, "turn"),
246                serde_json::json!({"path": "report.xlsx"}),
247            )
248            .await
249            .expect("send artifact");
250        let file = output.content.files().next().expect("published file");
251
252        assert_eq!(tool.approval(), ApprovalRequirement::Always);
253        assert!(
254            store
255                .list_files(
256                    "session-a",
257                    &[crate::backend::session_files::SessionFileOrigin::Upload]
258                )
259                .await
260                .map(|files| files.into_iter().map(|(_, file)| file).collect::<Vec<_>>())
261                .expect("uploads")
262                .is_empty()
263        );
264        assert_eq!(
265            store
266                .read_chunk("session-a", &file.id, 0, 16)
267                .await
268                .expect("read")
269                .data,
270            [0, 255, 1, 254]
271        );
272    }
273
274    #[test]
275    fn successful_send_replaces_the_pending_block_with_a_file() {
276        let state = tempfile::tempdir().expect("state");
277        let middleware = Artifacts::new(SessionFileStore::new(state.path(), None));
278        let file = SessionFileReference {
279            id: uuid::Uuid::new_v4().to_string(),
280            name: "diagram.svg".into(),
281            size: 4,
282            media_type: "image/svg+xml".into(),
283        };
284        let output = crate::protocol::ToolContent(vec![crate::protocol::ContentPart::File {
285            file: file.clone(),
286        }]);
287
288        let block = middleware
289            .render(
290                &EventMsg::ToolCallEnd(ToolCallEndEvent {
291                    turn_id: "turn".into(),
292                    call_id: "call".into(),
293                    name: "send_artifact".into(),
294                    output,
295                    is_error: false,
296                }),
297                "session-a",
298            )
299            .expect("block");
300
301        assert_eq!(block.update, crate::protocol::FrontendBlockUpdate::Replace);
302        assert_eq!(block.title, "Sent diagram.svg");
303        assert!(block.text.is_empty());
304        assert_eq!(block.files, [file]);
305    }
306}
mobius/middleware/artifacts.rs

mobius/middleware/
artifacts.rs
