1 files changed, 13 insertions, 3 deletions
diff --git a/internal/tag/tag.go b/internal/tag/tag.go
index 197d2d6..b4b627d 100644
--- a/internal/tag/tag.go
+++ b/internal/tag/tag.go
@@ -1,12 +1,22 @@
 package tag
 
+import "context"
+
 type Tagged[V, T any] struct {
 	Val V
 	Tag T
 }
 
-func Tag[V, T any](out chan<- Tagged[V, T], in <-chan V, f func(V) T) {
-	for val := range in {
-		out <- Tagged[V, T]{val, f(val)}
+func Tag[V, T any](ctx context.Context, out chan<- Tagged[V, T], in <-chan V, f func(V) T) {
+	for {
+		select {
+		case val, ok := <-in:
+			if !ok {
+				return
+			}
+			out <- Tagged[V, T]{val, f(val)}
+		case <-ctx.Done():
+			return
+		}
 	}
 }