generate wasmtime-wasi export bindings with task_exit option (#12185)

* generate `wasmtime-wasi` export bindings with `task_exit` option

Previously, `wasmtime run` only ran `wasi:cli/[email protected]` tasks until they
returned a result, then exited.  That meant any post-return computation done by
the task might not have a chance to run.

This commit aligns `wasmtime-wasi` with what `wasmtime-wasi-http` was already
doing: specifying the `task_exit` option for export bindings to
`wasmtime-wit-bindgen` so the host can `.await` both the return value _and_ the
task completion.  It also modifies the `wasmtime run` implementation to await
both.

* address review feedback

Author: dicej

crates/test-programs/src/bin/p3_cli_hello_stdout_post_return.rs

@@ -0,0 +1,33 @@
+use test_programs::p3::*;
+
+struct Component;
+
+export!(Component);
+
+impl exports::wasi::cli::run::Guest for Component {
+    async fn run() -> Result<(), ()> {
+        let (mut tx, rx) = wit_stream::new();
+        futures::join!(
+            async {
+                wasi::cli::stdout::write_via_stream(rx).await.unwrap();
+            },
+            async {
+                tx.write_all(b"hello, world\n".to_vec()).await;
+                wit_bindgen::spawn(async move {
+                    // Yield a few times to allow the host to accept and process
+                    // the `run` result.
+                    for _ in 0..10 {
+                        wit_bindgen::yield_async().await;
+                    }
+                    tx.write_all(b"hello again, after return\n".to_vec()).await;
+                    drop(tx);
+                });
+            },
+        );
+        Ok(())
+    }
+}
+
+fn main() {
+    panic!("should call p3 entrypoint");
+}

crates/wasi-http/tests/all/p3/mod.rs

@@ -144,6 +144,7 @@ async fn run_cli(path: &str, server: &Server) -> anyhow::Result<()> {
         .await
         .context("failed to call `wasi:cli/run#run`")?
         .context("guest trapped")?
+        .0
         .map_err(|()| anyhow!("`wasi:cli/run#run` failed"))
 }
 

crates/wasi/src/p3/bindings.rs

@@ -90,7 +90,7 @@ mod generated {
             "wasi:sockets/types.[method]udp-socket.connect": async | tracing | trappable,
             default: tracing | trappable,
         },
-        exports: { default: async | store },
+        exports: { default: async | store | task_exit },
         with: {
             "wasi:cli/terminal-input.terminal-input": crate::p3::cli::TerminalInput,
             "wasi:cli/terminal-output.terminal-output": crate::p3::cli::TerminalOutput,
@@ -159,7 +159,7 @@ pub use self::generated::wasi::*;
 ///     let command = Command::instantiate_async(&mut store, &component, &linker).await?;
 ///     let program_result = store.run_concurrent(async move |store| {
 ///         command.wasi_cli_run().call_run(store).await
-///     }).await??;
+///     }).await??.0;
 ///     match program_result {
 ///         Ok(()) => Ok(()),
 ///         Err(()) => std::process::exit(1),

crates/wasi/tests/all/p3/mod.rs

@@ -41,6 +41,7 @@ async fn run_allow_blocking_current_thread(
         .await
         .context("failed to call `wasi:cli/run#run`")?
         .context("guest trapped")?
+        .0
         .map_err(|()| anyhow!("`wasi:cli/run#run` failed"))
 }
 
@@ -157,6 +158,12 @@ async fn p3_file_write_blocking() -> anyhow::Result<()> {
 )]
 fn p3_cli_hello_stdout() {}
 
+#[expect(
+    dead_code,
+    reason = "tested in the wasi-cli crate, satisfying foreach_api! macro"
+)]
+fn p3_cli_hello_stdout_post_return() {}
+
 #[expect(
     dead_code,
     reason = "tested in the wasi-cli crate, satisfying foreach_api! macro"

src/commands/run.rs

@@ -645,7 +645,11 @@ impl RunCommand {
             if let Ok(command) = wasmtime_wasi::p3::bindings::Command::new(&mut *store, &instance) {
                 result = Some(
                     store
-                        .run_concurrent(async |store| command.wasi_cli_run().call_run(store).await)
+                        .run_concurrent(async |store| {
+                            let (result, task) = command.wasi_cli_run().call_run(store).await?;
+                            task.block(store).await;
+                            Ok(result)
+                        })
                         .await?,
                 );
             }

tests/all/cli_tests.rs

@@ -2165,6 +2165,23 @@ start a print 1234
         Ok(())
     }
 
+    #[test]
+    fn p2_cli_p3_hello_stdout_post_return() -> Result<()> {
+        let output = run_wasmtime(&[
+            "run",
+            "-Wcomponent-model-async",
+            "-Sp3",
+            P3_CLI_HELLO_STDOUT_POST_RETURN_COMPONENT,
+        ]);
+        if cfg!(feature = "component-model-async") {
+            let output = output?;
+            assert_eq!(output, "hello, world\nhello again, after return\n");
+        } else {
+            assert!(output.is_err());
+        }
+        Ok(())
+    }
+
     mod invoke {
         use super::*;