fix(agent): remove unwraps

src/agent/src/agents/mod.rs

@@ -1,4 +1,4 @@
-use crate::AgentResult;
+use crate::{AgentError, AgentResult};
 use serde::Deserialize;
 
 #[cfg(feature = "debug-agent")]
@@ -36,14 +36,17 @@ impl std::fmt::Display for Language {
 }
 
 impl TryFrom<&str> for Language {
-    type Error = String;
+    type Error = AgentError;
 
-    fn try_from(value: &str) -> Result<Self, Self::Error> {
+    fn try_from(value: &str) -> Result<Self, AgentError> {
         match value {
             "rust" => Ok(Language::Rust),
             #[cfg(feature = "debug-agent")]
             "debug" => Ok(Language::Debug),
-            _ => Err(format!("Invalid language: {}", value)),
+            _ => Err(AgentError::InvalidLanguage(format!(
+                "Invalid language: {}",
+                value
+            ))),
         }
     }
 }

src/agent/src/lib.rs

@@ -8,6 +8,7 @@ pub mod workload;
 pub enum AgentError {
     OpenConfigFileError(std::io::Error),
     ParseConfigError(toml::de::Error),
+    InvalidLanguage(String),
     BuildFailed(AgentOutput),
 }
 
@@ -17,6 +18,7 @@ impl fmt::Display for AgentError {
             AgentError::OpenConfigFileError(e) => write!(f, "Failed to open config file: {}", e),
             AgentError::ParseConfigError(e) => write!(f, "Failed to parse config file: {}", e),
             AgentError::BuildFailed(output) => write!(f, "Build failed: {:?}", output),
+            AgentError::InvalidLanguage(e) => write!(f, "Invalid language: {}", e),
         }
     }
 }

src/agent/src/workload/config.rs

@@ -35,15 +35,15 @@ impl Config {
         Ok(config)
     }
 
-    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Self {
-        Self {
+    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Result<Self, AgentError> {
+        Ok(Self {
             workload_name: execute_request.workload_name.clone(),
             // TODO: Fix this unwrap
-            language: Language::try_from(execute_request.language.clone().as_str()).unwrap(),
+            language: Language::try_from(execute_request.language.clone().as_str())?,
             action: execute_request.action().into(),
             config_string: execute_request.config_str,
             code: execute_request.code,
-        }
+        })
     }
 }
 

src/agent/src/workload/runner.rs

@@ -2,7 +2,7 @@ use crate::{
     agent::ExecuteRequest,
     agents::{rust, Agent, AgentOutput, Language},
     workload::config::Action,
-    AgentResult,
+    AgentError, AgentResult,
 };
 
 #[cfg(feature = "debug-agent")]
@@ -28,9 +28,9 @@ impl Runner {
         Runner { config, agent }
     }
 
-    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Self {
-        let config = Config::new_from_execute_request(execute_request);
-        Self::new(config)
+    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Result<Self, AgentError> {
+        let config = Config::new_from_execute_request(execute_request)?;
+        Ok(Self::new(config))
     }
 
     pub fn run(&self) -> AgentResult<AgentOutput> {

src/agent/src/workload/service.rs

@@ -18,7 +18,8 @@ impl WorkloadRunner for WorkloadRunnerService {
 
         let execute_request = req.into_inner();
 
-        let runner = Runner::new_from_execute_request(execute_request);
+        let runner = Runner::new_from_execute_request(execute_request)
+            .map_err(|e| tonic::Status::internal(e.to_string()))?;
 
         let res = runner
             .run()