diff --git a/worker-build/src/js/shim.js b/worker-build/src/js/shim.js
index 19d55198..b520d77f 100644
--- a/worker-build/src/js/shim.js
+++ b/worker-build/src/js/shim.js
@@ -10,7 +10,9 @@ export { wasmModule };
 
 class Entrypoint extends WorkerEntrypoint {
     async fetch(request) {
-        return await imports.fetch(request, this.env, this.ctx)
+        let response = imports.fetch(request, this.env, this.ctx);
+        $WAIT_UNTIL_RESPONSE
+        return await response;
     }
 
     async queue(batch) {
diff --git a/worker-build/src/main.rs b/worker-build/src/main.rs
index 807f3694..18310906 100644
--- a/worker-build/src/main.rs
+++ b/worker-build/src/main.rs
@@ -18,6 +18,8 @@ const OUT_DIR: &str = "build";
 const OUT_NAME: &str = "index";
 const WORKER_SUBDIR: &str = "worker";
 
+const SHIM_TEMPLATE: &str = include_str!("./js/shim.js");
+
 const WASM_IMPORT: &str = r#"let wasm;
 export function __wbg_set_wasm(val) {
     wasm = val;
@@ -54,7 +56,13 @@ pub fn main() -> Result<()> {
     use_glue_import()?;
 
     write_string_to_file(worker_path("glue.js"), include_str!("./js/glue.js"))?;
-    write_string_to_file(worker_path("shim.js"), include_str!("./js/shim.js"))?;
+    let shim = if env::var("RUN_TO_COMPLETION").is_ok() {
+        SHIM_TEMPLATE.replace("$WAIT_UNTIL_RESPONSE", "this.ctx.waitUntil(response);")
+    } else {
+        SHIM_TEMPLATE.replace("$WAIT_UNTIL_RESPONSE", "")
+    };
+
+    write_string_to_file(worker_path("shim.js"), shim)?;
 
     bundle(&esbuild_path)?;