diff --git a/pool.go b/pool.go
index 429e403..c9fb68e 100644
--- a/pool.go
+++ b/pool.go
@@ -270,6 +270,12 @@ type contextTask struct {
 	task Task
 }
 
+// UpdateTask wraps task with ctx so the task can observe cancellation and
+// deadlines while it is waiting in or executing from the pool.
+func UpdateTask(ctx context.Context, task Task) *contextTask {
+	return &contextTask{ctx: ctx, task: task}
+}
+
 func (t *contextTask) Process() {
 	if t.ctx.Err() != nil {
 		return
diff --git a/test/count.go b/test/count.go
index b8c51b9..117cd41 100644
--- a/test/count.go
+++ b/test/count.go
@@ -6,7 +6,7 @@ import (
 	"runtime"
 	"time"
 
-	agilepool "github.com/Yiming1997/agilePool"
+	agilepool "github.com/Yiming1997/agilePool/v2"
 	"github.com/shirou/gopsutil/v3/cpu"
 )
 
diff --git a/test/flag.go b/test/flag.go
index cf7be9d..183ff26 100644
--- a/test/flag.go
+++ b/test/flag.go
@@ -13,7 +13,7 @@ import (
 	"sync/atomic"
 	"time"
 
-	agilepool "github.com/Yiming1997/agilePool"
+	agilepool "github.com/Yiming1997/agilePool/v2"
 )
 
 // TaskType controls how simulated task durations are generated.
diff --git a/test/go.mod b/test/go.mod
index ed6a666..76f5a71 100644
--- a/test/go.mod
+++ b/test/go.mod
@@ -3,7 +3,7 @@ module agilepool-test
 go 1.23.4
 
 require (
-	github.com/Yiming1997/agilePool v0.0.0
+	github.com/Yiming1997/agilePool/v2 v2.0.0
 	github.com/shirou/gopsutil/v3 v3.24.5
 )
 
@@ -18,4 +18,4 @@ require (
 	golang.org/x/sys v0.20.0 // indirect
 )
 
-replace github.com/Yiming1997/agilePool => ../
+replace github.com/Yiming1997/agilePool/v2 => ../
diff --git a/test/main.go b/test/main.go
index 099355f..85fb530 100644
--- a/test/main.go
+++ b/test/main.go
@@ -8,7 +8,7 @@ import (
 	"runtime/pprof"
 	"time"
 
-	agilepool "github.com/Yiming1997/agilePool"
+	agilepool "github.com/Yiming1997/agilePool/v2"
 )
 
 func main() {