forge_runtime/jobs/
registry.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use forge_core::Result;
7use forge_core::job::{ForgeJob, JobContext, JobInfo};
8use serde_json::Value;
9
10/// Normalize args for deserialization.
11/// - Converts `null` to `{}` so both unit `()` and empty structs deserialize correctly.
12/// - Unwraps `{"args": ...}` or `{"input": ...}` wrapper if present (callers may use either format).
13fn normalize_args(args: Value) -> Value {
14    let unwrapped = match &args {
15        Value::Object(map) if map.len() == 1 => {
16            if map.contains_key("args") {
17                map.get("args").cloned().unwrap_or(Value::Null)
18            } else if map.contains_key("input") {
19                map.get("input").cloned().unwrap_or(Value::Null)
20            } else {
21                args
22            }
23        }
24        _ => args,
25    };
26
27    match &unwrapped {
28        Value::Null => Value::Object(serde_json::Map::new()),
29        _ => unwrapped,
30    }
31}
32
33/// Type alias for boxed job handler function.
34pub type BoxedJobHandler = Arc<
35    dyn Fn(&JobContext, Value) -> Pin<Box<dyn Future<Output = Result<Value>> + Send + '_>>
36        + Send
37        + Sync,
38>;
39
40/// Entry in the job registry.
41pub struct JobEntry {
42    /// Job metadata.
43    pub info: JobInfo,
44    /// Job handler function.
45    pub handler: BoxedJobHandler,
46}
47
48/// Registry of all FORGE jobs.
49#[derive(Clone, Default)]
50pub struct JobRegistry {
51    jobs: HashMap<String, Arc<JobEntry>>,
52}
53
54impl JobRegistry {
55    /// Create a new empty registry.
56    pub fn new() -> Self {
57        Self {
58            jobs: HashMap::new(),
59        }
60    }
61
62    /// Register a job type.
63    pub fn register<J: ForgeJob>(&mut self)
64    where
65        J::Args: serde::de::DeserializeOwned + Send + 'static,
66        J::Output: serde::Serialize + Send + 'static,
67    {
68        let info = J::info();
69        let name = info.name.to_string();
70
71        let handler: BoxedJobHandler = Arc::new(move |ctx, args| {
72            Box::pin(async move {
73                let parsed_args: J::Args = serde_json::from_value(normalize_args(args))
74                    .map_err(|e| forge_core::ForgeError::Validation(e.to_string()))?;
75                let result = J::execute(ctx, parsed_args).await?;
76                serde_json::to_value(result)
77                    .map_err(|e| forge_core::ForgeError::Internal(e.to_string()))
78            })
79        });
80
81        self.jobs.insert(name, Arc::new(JobEntry { info, handler }));
82    }
83
84    /// Get a job entry by name.
85    pub fn get(&self, name: &str) -> Option<Arc<JobEntry>> {
86        self.jobs.get(name).cloned()
87    }
88
89    /// Get job info by name.
90    pub fn info(&self, name: &str) -> Option<&JobInfo> {
91        self.jobs.get(name).map(|e| &e.info)
92    }
93
94    /// Check if a job exists.
95    pub fn exists(&self, name: &str) -> bool {
96        self.jobs.contains_key(name)
97    }
98
99    /// Get all job names.
100    pub fn job_names(&self) -> impl Iterator<Item = &str> {
101        self.jobs.keys().map(|s| s.as_str())
102    }
103
104    /// Get all jobs.
105    pub fn jobs(&self) -> impl Iterator<Item = (&str, &Arc<JobEntry>)> {
106        self.jobs.iter().map(|(k, v)| (k.as_str(), v))
107    }
108
109    /// Get the number of registered jobs.
110    pub fn len(&self) -> usize {
111        self.jobs.len()
112    }
113
114    /// Check if registry is empty.
115    pub fn is_empty(&self) -> bool {
116        self.jobs.is_empty()
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_empty_registry() {
126        let registry = JobRegistry::new();
127        assert!(registry.is_empty());
128        assert_eq!(registry.len(), 0);
129        assert!(registry.get("nonexistent").is_none());
130    }
131}
forge_runtime/jobs/registry.rs

forge_runtime/jobs/
registry.rs
