batch_mode_batch_workspace/
load_seed_by_custom_id.rs

1crate::ix!();
2
3use serde::Deserialize;
4use tokio::{fs, io::{self, AsyncBufReadExt, BufReader}};
5
6#[derive(Deserialize)]
7struct ManifestLine {
8    custom_id: String,
9    // future‑proof: allow an optional header/name field
10    #[serde(default)]
11    header: Option<String>,
12}
13
14#[async_trait]
15impl LoadSeedByCustomId for BatchWorkspace {
16    /// Locate the seed (header) that corresponds to `custom_id`
17    /// by scanning seed‑manifest files (.json **or** .jsonl).
18    #[instrument(
19        level = "debug",
20        skip(self),
21        fields(%custom_id, workdir = %self.workdir().display())
22    )]
23    async fn load_seed_by_custom_id(
24        &self,
25        custom_id: &CustomRequestId,
26    ) -> Result<Box<dyn Named + Send + Sync>, BatchWorkspaceError> {
27        let mut dir = fs::read_dir(self.workdir()).await?;
28
29        while let Some(entry) = dir.next_entry().await? {
30            let p = entry.path();
31
32            if !p.file_name()
33                .and_then(|s| s.to_str())
34                .map_or(false, |s| s.starts_with("batch_seed_manifest_"))
35            {
36                continue;
37            }
38
39            // ---- Fast path: single‑JSON object (current expectation) ----
40            match fs::read_to_string(&p).await {
41                Ok(raw) => {
42                    if let Ok(map) =
43                        serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(&raw)
44                    {
45                        if let Some(serde_json::Value::String(name)) =
46                            map.get(custom_id.as_str())
47                        {
48                            debug!(file=%p.display(), "found via single‑JSON manifest");
49                            return Ok(Box::new(SimpleSeed(name.clone())));
50                        }
51                    } else {
52                        // fall through to line‑by‑line parsing
53                    }
54                }
55                Err(e) => return Err(e.into()),
56            }
57
58            // ---- Slow path: JSON‑Lines (one object per line) ------------
59            let file = fs::File::open(&p).await?;
60            let mut reader = BufReader::new(file);
61            let mut line = String::new();
62
63            loop {
64                line.clear();
65                let bytes = reader.read_line(&mut line).await?;
66                if bytes == 0 {
67                    break; // EOF
68                }
69
70                let parsed: ManifestLine = match serde_json::from_str(&line) {
71                    Ok(v) => v,
72                    Err(_) => continue, // skip malformed lines
73                };
74
75                if parsed.custom_id == custom_id.as_str() {
76                    let name = parsed
77                        .header
78                        .unwrap_or_else(|| parsed.custom_id.clone());
79                    debug!(file=%p.display(), "found via JSON‑Lines manifest");
80                    return Ok(Box::new(SimpleSeed(name)));
81                }
82            }
83        }
84
85        Err(BatchWorkspaceError::NoBatchFileTripleAtIndex {
86            index: BatchIndex::new(),
87        })
88    }
89}
90
91/// Small helper so we can return something that implements `Named`
92struct SimpleSeed(String);
93impl Named for SimpleSeed {
94    fn name(&self) -> std::borrow::Cow<'_, str> {
95        std::borrow::Cow::Borrowed(&self.0)
96    }
97}