tenant.openagents/omega

1mod check_permissions;
2mod check_run_patterns;
3
4use std::{
5    fs,
6    path::{Path, PathBuf},
7};
8
9use annotate_snippets::{Group, Renderer};
10use anyhow::{Result, anyhow};
11use clap::Parser;
12use itertools::{Either, Itertools};
13use serde_yaml::Value;
14use strum::IntoEnumIterator;
15
16use crate::tasks::workflows::WorkflowType;
17
18use check_permissions::PermissionsError;
19use check_run_patterns::RunValidationError;
20
21pub use check_run_patterns::validate_run_command;
22
23#[derive(Default, Parser)]
24pub struct WorkflowValidationArgs {}
25
26pub fn validate(_: WorkflowValidationArgs) -> Result<()> {
27    let (parsing_errors, file_errors): (Vec<_>, Vec<_>) = get_all_workflow_files()
28        .map(check_workflow)
29        .flat_map(Result::err)
30        .partition_map(|error| match error {
31            WorkflowError::ParseError(error) => Either::Left(error),
32            WorkflowError::ValidationError(error) => Either::Right(error),
33        });
34
35    if !parsing_errors.is_empty() {
36        Err(anyhow!(
37            "Failed to read or parse some workflow files: {}",
38            parsing_errors.into_iter().join("\n")
39        ))
40    } else if !file_errors.is_empty() {
41        let groups: Vec<_> = file_errors
42            .iter()
43            .flat_map(|error| error.annotation_groups())
44            .collect();
45
46        let renderer =
47            Renderer::styled().decor_style(annotate_snippets::renderer::DecorStyle::Ascii);
48        println!("{}", renderer.render(groups.as_slice()));
49
50        Err(anyhow!("Workflow checks failed!"))
51    } else {
52        Ok(())
53    }
54}
55
56struct WorkflowFile {
57    raw_content: String,
58    parsed_content: Value,
59}
60
61impl WorkflowFile {
62    fn load(workflow_file_path: &Path) -> Result<Self> {
63        fs::read_to_string(workflow_file_path)
64            .map_err(|_| {
65                anyhow!(
66                    "Could not read workflow file at {}",
67                    workflow_file_path.display()
68                )
69            })
70            .and_then(|file_content| {
71                serde_yaml::from_str(&file_content)
72                    .map(|parsed_content| Self {
73                        raw_content: file_content,
74                        parsed_content,
75                    })
76                    .map_err(|e| anyhow!("Failed to parse workflow file: {e:?}"))
77            })
78    }
79}
80
81/// A single kind of validation failure found within a workflow file.
82enum ValidationError {
83    RunInjection(RunValidationError),
84    Permissions(PermissionsError),
85}
86
87impl ValidationError {
88    fn annotation_group<'a>(&self, file_path: &Path, raw_content: &'a str) -> Group<'a> {
89        match self {
90            ValidationError::RunInjection(error) => error.annotation_group(file_path, raw_content),
91            ValidationError::Permissions(error) => error.annotation_group(file_path, raw_content),
92        }
93    }
94}
95
96struct WorkflowValidationError {
97    file_path: PathBuf,
98    contents: WorkflowFile,
99    errors: Vec<ValidationError>,
100}
101
102impl WorkflowValidationError {
103    fn annotation_groups(&self) -> Vec<Group<'_>> {
104        self.errors
105            .iter()
106            .map(|error| error.annotation_group(&self.file_path, &self.contents.raw_content))
107            .collect()
108    }
109}
110
111enum WorkflowError {
112    ParseError(anyhow::Error),
113    ValidationError(Box<WorkflowValidationError>),
114}
115
116fn get_all_workflow_files() -> impl Iterator<Item = PathBuf> {
117    WorkflowType::iter()
118        .map(|workflow_type| workflow_type.folder_path())
119        .flat_map(|folder_path| {
120            fs::read_dir(folder_path).into_iter().flat_map(|entries| {
121                entries
122                    .flat_map(Result::ok)
123                    .map(|entry| entry.path())
124                    .filter(|path| {
125                        path.extension()
126                            .is_some_and(|ext| ext == "yaml" || ext == "yml")
127                    })
128            })
129        })
130}
131
132fn check_workflow(workflow_file_path: PathBuf) -> Result<(), WorkflowError> {
133    let file_content =
134        WorkflowFile::load(&workflow_file_path).map_err(WorkflowError::ParseError)?;
135
136    let mut errors = Vec::new();
137
138    if let Err(error) = check_permissions::validate_permissions(&file_content.parsed_content) {
139        errors.push(ValidationError::Permissions(error));
140    }
141
142    errors.extend(
143        collect_run_injection_errors(&Value::Null, &file_content.parsed_content)
144            .into_iter()
145            .map(ValidationError::RunInjection),
146    );
147
148    if errors.is_empty() {
149        Ok(())
150    } else {
151        Err(WorkflowError::ValidationError(Box::new(
152            WorkflowValidationError {
153                file_path: workflow_file_path,
154                contents: file_content,
155                errors,
156            },
157        )))
158    }
159}
160
161fn collect_run_injection_errors(key: &Value, value: &Value) -> Vec<RunValidationError> {
162    match value {
163        Value::Mapping(mapping) => mapping
164            .iter()
165            .flat_map(|(key, value)| collect_run_injection_errors(key, value))
166            .collect(),
167        Value::Sequence(sequence) => sequence
168            .iter()
169            .flat_map(|value| collect_run_injection_errors(key, value))
170            .collect(),
171        Value::String(string) => check_string(key, string).err().into_iter().collect(),
172        Value::Null | Value::Bool(_) | Value::Number(_) | Value::Tagged(_) => Vec::new(),
173    }
174}
175
176fn check_string(key: &Value, value: &str) -> Result<(), RunValidationError> {
177    match key {
178        Value::String(key) if key == "run" => validate_run_command(value),
179        _ => Ok(()),
180    }
181}
182