From 75d19433be36b89de8de7b0b1eb7661094856df5 Mon Sep 17 00:00:00 2001
From: Tristan Druyen <tristan@vault81.de>
Date: Tue, 18 Mar 2025 00:55:37 +0100
Subject: [PATCH] Remove panics

---
 llama_forge_rs/src/main.rs  |  2 +-
 llama_proxy_man/src/lib.rs  | 14 +++++++-------
 llama_proxy_man/src/main.rs |  7 +++++--
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/llama_forge_rs/src/main.rs b/llama_forge_rs/src/main.rs
index f34e639..69bd62f 100644
--- a/llama_forge_rs/src/main.rs
+++ b/llama_forge_rs/src/main.rs
@@ -98,7 +98,7 @@ async fn main() -> anyhow::Result<()> {
     let proxy_man_fut = async move {
         use llama_proxy_man::{config::AppConfig, start_server};
         let config = AppConfig::default_figment();
-        start_server(config).await;
+        start_server(config).await?;
 
         Ok::<(), anyhow::Error>(())
     };
diff --git a/llama_proxy_man/src/lib.rs b/llama_proxy_man/src/lib.rs
index 9d1a58c..2bbbd25 100644
--- a/llama_proxy_man/src/lib.rs
+++ b/llama_proxy_man/src/lib.rs
@@ -51,7 +51,7 @@ pub fn axum_router(spec: &ModelSpec, state: AppState) -> Router {
 }
 
 /// Starts an inference server for each model defined in the config.
-pub async fn start_server(config: AppConfig) {
+pub async fn start_server(config: AppConfig) -> anyhow::Result<()> {
     let state = AppState::from_config(config.clone());
 
     let mut handles = Vec::new();
@@ -59,17 +59,17 @@ pub async fn start_server(config: AppConfig) {
         let state = state.clone();
         let spec = spec.clone();
 
-        let handle = tokio::spawn(async move {
+        let handle: tokio::task::JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
             let app = axum_router(&spec, state);
             let addr = SocketAddr::from(([0, 0, 0, 0], spec.port));
             tracing::info!(msg = "Listening", ?spec);
-            let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
-            axum::serve(listener, app.into_make_service())
-                .await
-                .unwrap();
+            let listener = tokio::net::TcpListener::bind(&addr).await?;
+            axum::serve(listener, app.into_make_service()).await?;
+            Ok(())
         });
         handles.push(handle);
     }
 
-    futures::future::join_all(handles).await;
+    let _ = futures::future::try_join_all(handles).await?;
+    Ok(())
 }
diff --git a/llama_proxy_man/src/main.rs b/llama_proxy_man/src/main.rs
index f7cd786..b5bc200 100644
--- a/llama_proxy_man/src/main.rs
+++ b/llama_proxy_man/src/main.rs
@@ -1,10 +1,13 @@
+use anyhow;
 use llama_proxy_man::{config::AppConfig, logging, start_server};
 
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
     logging::initialize_logger();
 
     let config = AppConfig::default_from_pwd_yml();
 
-    start_server(config).await;
+    start_server(config).await?;
+
+    Ok(())
 }