diff --git a/go.mod b/go.mod
index 211ca79..5c6f58c 100644
--- a/go.mod
+++ b/go.mod
@@ -2,6 +2,6 @@ module github.com/JairusSW/pool
 
 go 1.24.0
 
-require github.com/wago-org/wago v0.0.0-20260812144524-1c58c9862d25
+require github.com/wago-org/wago v0.1.0-beta.8.0.20260916050830-c9ed2455acd9
 
-require github.com/wago-org/workers v0.0.0-20260812144603-4a90ff25d5e4
+require github.com/wago-org/workers v0.1.1-0.20260916050909-fa9222c78f9d
diff --git a/go.sum b/go.sum
index 3720533..701190e 100644
--- a/go.sum
+++ b/go.sum
@@ -1,4 +1,4 @@
-github.com/wago-org/wago v0.0.0-20260812144524-1c58c9862d25 h1:gopNHvaUaDSzmW2Hfd32r8dOwwzYUHH3ph1L/IuhJUA=
-github.com/wago-org/wago v0.0.0-20260812144524-1c58c9862d25/go.mod h1:6XmxI3S5qJ+YAzObyXjC6FmxHtbLhBBxrglzix42zl8=
-github.com/wago-org/workers v0.0.0-20260812144603-4a90ff25d5e4 h1:Z7S9Hj+JZ8OdNWLgzRjUQOvPIEQpEFNtyRSoHnUYOlI=
-github.com/wago-org/workers v0.0.0-20260812144603-4a90ff25d5e4/go.mod h1:FLuFlQ7n+DNP0VNAZuGc702UIxC/JfLuhyuacMQNFp8=
+github.com/wago-org/wago v0.1.0-beta.8.0.20260916050830-c9ed2455acd9 h1:+/8b9bvAnyUZPG+YTvGpDwpAV+jpp6+bBBWOrr+8Xco=
+github.com/wago-org/wago v0.1.0-beta.8.0.20260916050830-c9ed2455acd9/go.mod h1:q0pPiPF6W6GZpTXbWwy+GT+r19KHSzB/gRcPUVHyntI=
+github.com/wago-org/workers v0.1.1-0.20260916050909-fa9222c78f9d h1:hx2oFfErXwEJE54E98xzIj0874SUrJL2LXi8EwkQkrQ=
+github.com/wago-org/workers v0.1.1-0.20260916050909-fa9222c78f9d/go.mod h1:nFQYC2iT+rgsR2G60CF/NZm4flLcMLkYJoloAFFGztk=
diff --git a/integration_test.go b/integration_test.go
index 256d589..67a26a7 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -11,7 +11,7 @@ import (
 
 	"github.com/wago-org/wago"
 	wagoplugin "github.com/wago-org/wago/plugin"
-	"github.com/wago-org/wago/tests/wasmtest"
+	"github.com/wago-org/wago/tests/support/wasmtest"
 	"github.com/wago-org/workers"
 )
 
@@ -118,11 +118,9 @@ func driverProvider(d *driver) wago.PluginProvider {
 			if err != nil {
 				return err
 			}
-			module, err := imports.Module("env")
-			if err != nil {
-				return err
-			}
-			module.Func("run", d.run).Params(wago.ValI32, wago.ValI32, wago.ValI32).Results(wago.ValI32)
+			imports.HostFunc("env", "run", func(caller wago.Caller, call wago.HostCall) {
+				d.run(caller, call.ParamSlots(), call.ResultSlots())
+			}).Params(wago.ValI32, wago.ValI32, wago.ValI32).Results(wago.ValI32)
 			return reg.Lifecycle(wago.PluginLifecycle{
 				Start: d.start,
 				Stop: func(context.Context) error {
@@ -374,7 +372,7 @@ func TestContractGraphRoutesAndScales(t *testing.T) {
 	waitFor(t, "tasks", func() bool { return r.d.messages() == 20 })
 	r.invoke(t, "scale")
 	waitFor(t, "scale", func() bool { return r.stats(t).Live == 4 })
-	if err := r.rt.Close(); err != nil {
+	if err := r.rt.CloseContext(context.Background()); err != nil {
 		t.Fatal(err)
 	}
 	if err := r.d.poolRef.With(func(Service) error { return nil }); !errors.Is(err, wago.ErrPermissionDenied) {
diff --git a/stress_test.go b/stress_test.go
index bbdbdbb..05c451a 100644
--- a/stress_test.go
+++ b/stress_test.go
@@ -1,6 +1,7 @@
 package pool
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -294,10 +295,16 @@ func TestStressRuntimeCloseWhileBusy(t *testing.T) {
 			}
 		}()
 	}
-	done := make(chan struct{})
-	go func() { r.close(); close(done) }()
+	closeDone := make(chan error, 1)
+	go func() {
+		_ = r.in.Close()
+		closeDone <- r.rt.CloseContext(context.Background())
+	}()
 	select {
-	case <-done:
+	case err := <-closeDone:
+		if err != nil {
+			t.Fatalf("runtime close: %v", err)
+		}
 	case <-time.After(10 * time.Second):
 		stop.Store(true)
 		wg.Wait()