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

forge_runtime/jobs/
registry.rs
