allow customizing log prefixes for wasmtime serve command

Signed-off-by: Joseph Zhang <[email protected]>

src/commands/serve.rs

@@ -88,9 +88,18 @@ pub struct ServeCommand {
     run: RunCommon,
 
     /// Socket address for the web server to bind to.
-    #[arg(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR )]
+    #[arg(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR)]
     addr: SocketAddr,
 
+    /// The custom prefix of stdout logs from wasi-http handlers.
+    /// if unspecified, default log of 'stdout [{req_id}] :: ' will be used.
+    #[arg(long = "logging-prefix-stdout", value_name = "PREFIX")]
+    logging_prefix_stdout: Option<String>,
+    /// The custom prefix of stderr logs from wasi-http handlers.
+    /// if unspecified, default log of 'stderr [{req_id}] :: ' will be used.
+    #[arg(long = "logging-prefix-stderr", value_name = "PREFIX")]
+    logging_prefix_stderr: Option<String>,
+
     /// The WebAssembly component to run.
     #[arg(value_name = "WASM", required = true)]
     component: PathBuf,
@@ -153,15 +162,17 @@ impl ServeCommand {
 
         builder.env("REQUEST_ID", req_id.to_string());
 
-        builder.stdout(LogStream::new(
-            format!("stdout [{req_id}] :: "),
-            Output::Stdout,
-        ));
-
-        builder.stderr(LogStream::new(
-            format!("stderr [{req_id}] :: "),
-            Output::Stderr,
-        ));
+        let stdout_prefix = self
+            .logging_prefix_stdout
+            .clone()
+            .unwrap_or_else(|| format!("stdout [{req_id}] :: "));
+        builder.stdout(LogStream::new(stdout_prefix, Output::Stdout));
+
+        let stderr_prefix = self
+            .logging_prefix_stderr
+            .clone()
+            .unwrap_or_else(|| format!("stderr [{req_id}] :: "));
+        builder.stderr(LogStream::new(stderr_prefix, Output::Stderr));
 
         let mut host = Host {
             table: wasmtime::component::ResourceTable::new(),

tests/all/cli_tests.rs

@@ -1909,6 +1909,53 @@ stderr [1] :: after empty
         Ok(())
     }
 
+    #[tokio::test]
+    async fn cli_serve_with_print_no_prefix() -> Result<()> {
+        let server = WasmtimeServe::new(CLI_SERVE_WITH_PRINT_COMPONENT, |cmd| {
+            cmd.arg("-Scli");
+            cmd.arg("--logging-prefix-stdout=");
+            cmd.arg("--logging-prefix-stderr=");
+        })?;
+
+        for _ in 0..2 {
+            let resp = server
+                .send_request(
+                    hyper::Request::builder()
+                        .uri("http://localhost/")
+                        .body(String::new())
+                        .context("failed to make request")?,
+                )
+                .await?;
+            assert!(resp.status().is_success());
+        }
+
+        let (out, err) = server.finish()?;
+        assert_eq!(
+            out,
+            "\
+this is half a print to stdout
+\n\
+after empty
+this is half a print to stdout
+\n\
+after empty
+"
+        );
+        assert_eq!(
+            err,
+            "\
+this is half a print to stderr
+\n\
+after empty
+this is half a print to stderr
+\n\
+after empty
+"
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn cli_serve_authority_and_scheme() -> Result<()> {
         let server = WasmtimeServe::new(CLI_SERVE_AUTHORITY_AND_SCHEME_COMPONENT, |cmd| {