cookie_agent_engine/runtime/
artifact_reads.rs

1use cookie_agent_protocol::{ArtifactReadPath, PersistedToolResult, ToolOutputManifest};
2
3use std::sync::Arc;
4
5use super::{Engine, artifacts::ArtifactRouter, blocking_io};
6use crate::ToolError;
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub struct ArtifactReadPage {
10    pub content: String,
11    pub next_offset_lines: Option<u64>,
12    pub source: String,
13}
14
15impl Engine {
16    /// Read a public artifact URI by possession, without a session ownership lookup.
17    /// This synchronous API blocks; async tool executors use `ToolExecutionContext::read_artifact`.
18    pub fn read_artifact(
19        &self,
20        path: &str,
21        offset: u64,
22        limit: u64,
23    ) -> Result<ArtifactReadPage, ToolError> {
24        read_artifact(&self.inner.artifacts, path, offset, limit)
25    }
26}
27
28pub(crate) fn read_artifact(
29    store: &ArtifactRouter,
30    path: &str,
31    offset: u64,
32    limit: u64,
33) -> Result<ArtifactReadPage, ToolError> {
34    let target = ArtifactReadPath::parse(path).map_err(ToolError::execution)?;
35    if limit == 0 {
36        return Err(ToolError::execution("artifact page limit must be positive"));
37    }
38    let (digest, source) = if let Some(name) = target.stream {
39        let page = store
40            .read_paged(target.digest.as_str(), 0, 1)
41            .map_err(read_error)?;
42        let manifest: ToolOutputManifest = serde_json::from_str(&page.content)
43            .map_err(|_| ToolError::execution("artifact is not a named-stream manifest"))?;
44        manifest.validate().map_err(ToolError::execution)?;
45        let stream = manifest
46            .streams
47            .into_iter()
48            .find(|stream| stream.name.as_deref() == Some(&name))
49            .ok_or_else(|| {
50                ToolError::execution(format!("artifact has no declared stream {name}"))
51            })?;
52        (stream.sha256, format!("artifact.{name}"))
53    } else {
54        (target.digest, "artifact".into())
55    };
56    let page = store
57        .read_paged(digest.as_str(), offset, limit.min(2_000))
58        .map_err(read_error)?;
59    if page.content.len() > PersistedToolResult::MAX_OUTPUT_BYTES {
60        return Err(ToolError::resource_limit(
61            "artifact page exceeds the 2 MiB output limit; request a smaller page",
62        ));
63    }
64    Ok(ArtifactReadPage {
65        content: page.content,
66        next_offset_lines: page.next_offset_lines,
67        source,
68    })
69}
70
71pub(crate) async fn read_artifact_async(
72    store: Arc<ArtifactRouter>,
73    path: &str,
74    offset: u64,
75    limit: u64,
76) -> Result<ArtifactReadPage, ToolError> {
77    let path = path.to_owned();
78    blocking_io::run(move || read_artifact(&store, &path, offset, limit)).await?
79}
80
81fn read_error(error: std::io::Error) -> ToolError {
82    if error.kind() == std::io::ErrorKind::FileTooLarge {
83        ToolError::resource_limit(error.to_string())
84    } else {
85        ToolError::execution(error.to_string())
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[tokio::test(flavor = "current_thread")]
94    async fn blocked_artifact_verification_leaves_async_tasks_and_other_readers_responsive() {
95        let directory = tempfile::tempdir().unwrap();
96        let store = ArtifactRouter::open_flat(directory.path().join("artifacts")).unwrap();
97        let (_, first) = store.retain(crate::test_session_id(), b"first\n").unwrap();
98        let (_, second) = store.retain(crate::test_session_id(), b"second\n").unwrap();
99        let (entered, release) =
100            blocking_io::gate(store.io_test_hook(), "read", Some(first.clone()));
101        let reader = store.clone();
102        let blocked = tokio::spawn(async move {
103            read_artifact_async(reader, &format!("artifact://{first}"), 0, 1).await
104        });
105        tokio::time::timeout(std::time::Duration::from_secs(10), entered)
106            .await
107            .unwrap()
108            .unwrap();
109        assert!(!blocked.is_finished());
110        assert_eq!(tokio::spawn(async { 7 }).await.unwrap(), 7);
111        let other = tokio::time::timeout(
112            std::time::Duration::from_secs(10),
113            read_artifact_async(store, &format!("artifact://{second}"), 0, 1),
114        )
115        .await
116        .unwrap()
117        .unwrap();
118        assert_eq!(other.content, "second\n");
119        assert!(
120            !blocked.is_finished(),
121            "first reader must remain behind its I/O gate"
122        );
123        release.send(()).unwrap();
124        assert_eq!(blocked.await.unwrap().unwrap().content, "first\n");
125    }
126
127    #[tokio::test(flavor = "current_thread")]
128    async fn in_flight_read_does_not_reinsert_a_handle_invalidated_by_gc() {
129        let directory = tempfile::tempdir().unwrap();
130        let sessions = directory.path().join("sessions");
131        std::fs::create_dir(&sessions).unwrap();
132        let store = ArtifactRouter::open_flat(directory.path().join("artifacts")).unwrap();
133        let (_, digest) = store
134            .retain(crate::test_session_id(), b"verified\n")
135            .unwrap();
136        let path = format!("artifact://{digest}");
137        read_artifact_async(store.clone(), &path, 0, 1)
138            .await
139            .unwrap();
140        let (entered, release) = blocking_io::gate(store.io_test_hook(), "read", Some(digest));
141        let reader = store.clone();
142        let reading_path = path.clone();
143        let reading =
144            tokio::spawn(async move { read_artifact_async(reader, &reading_path, 0, 1).await });
145        tokio::time::timeout(std::time::Duration::from_secs(10), entered)
146            .await
147            .unwrap()
148            .unwrap();
149        assert_eq!(
150            store
151                .collect_garbage(std::time::Duration::ZERO)
152                .unwrap()
153                .deleted,
154            1
155        );
156        release.send(()).unwrap();
157        assert_eq!(reading.await.unwrap().unwrap().content, "verified\n");
158        assert!(
159            read_artifact_async(store, &path, 0, 1)
160                .await
161                .unwrap_err()
162                .to_string()
163                .contains("artifact missing")
164        );
165    }
166}
cookie_agent_engine/runtime/artifact_reads.rs

cookie_agent_engine/runtime/
artifact_reads.rs
