From 0784fa661b29b2d496efae23133116ca978353e8 Mon Sep 17 00:00:00 2001
From: Jason <me@jasonharrison.us>
Date: Mon, 11 Dec 2023 01:13:06 -0700
Subject: [PATCH] feat(core): add --host option to specify the host Tabby
 server will listen on. (#1021)

* Add --host option to specify the host Tabby server will listen on.

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 crates/tabby/src/routes/mod.rs | 7 +++----
 crates/tabby/src/serve.rs      | 9 ++++++---
 crates/tabby/src/worker.rs     | 9 ++++++---
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/crates/tabby/src/routes/mod.rs b/crates/tabby/src/routes/mod.rs
index 1ef1542..f8311e8 100644
--- a/crates/tabby/src/routes/mod.rs
+++ b/crates/tabby/src/routes/mod.rs
@@ -1,7 +1,7 @@
 mod metrics;
 
 use std::{
-    net::{Ipv4Addr, SocketAddr},
+    net::{IpAddr, SocketAddr},
     sync::Arc,
 };
 
@@ -14,7 +14,7 @@ use tracing::info;
 
 use crate::fatal;
 
-pub async fn run_app(api: Router, ui: Option<Router>, port: u16) {
+pub async fn run_app(api: Router, ui: Option<Router>, host: IpAddr, port: u16) {
     let (prometheus_layer, prometheus_handle) = PrometheusMetricLayer::pair();
     let app = api
         .layer(CorsLayer::permissive())
@@ -31,9 +31,8 @@ pub async fn run_app(api: Router, ui: Option<Router>, port: u16) {
         app
     };
 
-    let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, port));
+    let address = SocketAddr::from((host, port));
     info!("Listening at {}", address);
-
     Server::bind(&address)
         .serve(app.into_make_service_with_connect_info::<SocketAddr>())
         .await
diff --git a/crates/tabby/src/serve.rs b/crates/tabby/src/serve.rs
index 9df6c52..bd7a425 100644
--- a/crates/tabby/src/serve.rs
+++ b/crates/tabby/src/serve.rs
@@ -1,4 +1,4 @@
-use std::{sync::Arc, time::Duration};
+use std::{net::IpAddr, sync::Arc, time::Duration};
 
 use axum::{routing, Router};
 use clap::Args;
@@ -81,6 +81,9 @@ pub struct ServeArgs {
     #[clap(long)]
     chat_model: Option<String>,
 
+    #[clap(long, default_value = "0.0.0.0")]
+    host: IpAddr,
+
     #[clap(long, default_value_t = 8080)]
     port: u16,
 
@@ -108,7 +111,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
     #[cfg(not(feature = "experimental-http"))]
     load_model(args).await;
 
-    info!("Starting server, this might takes a few minutes...");
+    info!("Starting server, this might take a few minutes...");
 
     let logger = Arc::new(create_logger());
     let code = Arc::new(create_code_search());
@@ -129,7 +132,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
     let ui = ui.fallback(|| async { axum::response::Redirect::permanent("/swagger-ui") });
 
     start_heartbeat(args);
-    run_app(api, Some(ui), args.port).await
+    run_app(api, Some(ui), args.host, args.port).await
 }
 
 async fn load_model(args: &ServeArgs) {
diff --git a/crates/tabby/src/worker.rs b/crates/tabby/src/worker.rs
index eedfafe..417d110 100644
--- a/crates/tabby/src/worker.rs
+++ b/crates/tabby/src/worker.rs
@@ -1,4 +1,4 @@
-use std::{env::consts::ARCH, sync::Arc};
+use std::{env::consts::ARCH, net::IpAddr, sync::Arc};
 
 use anyhow::Result;
 use axum::{routing, Router};
@@ -23,6 +23,9 @@ pub struct WorkerArgs {
     #[clap(long)]
     url: String,
 
+    #[clap(long, default_value = "0.0.0.0")]
+    host: IpAddr,
+
     #[clap(long, default_value_t = 8080)]
     port: u16,
 
@@ -74,7 +77,7 @@ async fn make_completion_route(context: WorkerContext, args: &WorkerArgs) -> Rou
 pub async fn main(kind: WorkerKind, args: &WorkerArgs) {
     download_model_if_needed(&args.model).await;
 
-    info!("Starting worker, this might takes a few minutes...");
+    info!("Starting worker, this might take a few minutes...");
 
     let context = WorkerContext::new(&args.url).await;
 
@@ -83,7 +86,7 @@ pub async fn main(kind: WorkerKind, args: &WorkerArgs) {
         WorkerKind::Chat => make_chat_route(context, args).await,
     };
 
-    run_app(app, None, args.port).await
+    run_app(app, None, args.host, args.port).await
 }
 
 struct WorkerContext {