fix: ctx in stream callback handlers (#88)

Change-Id: Ic4a3d2e0a22493ba389bd9db5a841977986b0eef
cloudwego · Feb 26, 2025 · b769ad2 · b769ad2
1 parent 706ccb8
commit b769ad2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/compose/utils.go b/compose/utils.go
@@ -112,7 +112,7 @@ func genericOnStartWithStreamInputHandle(ctx context.Context, input streamReader
 
 	cpy := input.copy
 
-	handle := func(handler icb.Handler, in streamReader) context.Context {
+	handle := func(ctx context.Context, handler icb.Handler, in streamReader) context.Context {
 		in_, ok := unpackStreamReader[icb.CallbackInput](in)
 		if !ok {
 			panic("impossible")
@@ -139,7 +139,7 @@ func genericOnEndWithStreamOutputHandle(ctx context.Context, output streamReader
 
 	cpy := output.copy
 
-	handle := func(handler icb.Handler, out streamReader) context.Context {
+	handle := func(ctx context.Context, handler icb.Handler, out streamReader) context.Context {
 		out_, ok := unpackStreamReader[icb.CallbackOutput](out)
 		if !ok {
 			panic("impossible")

diff --git a/internal/callbacks/inject.go b/internal/callbacks/inject.go
@@ -95,7 +95,7 @@ func OnWithStreamHandle[S any](
 	inOut S,
 	handlers []Handler,
 	cpy func(int) []S,
-	handle func(Handler, S) context.Context) (context.Context, S) {
+	handle func(context.Context, Handler, S) context.Context) (context.Context, S) {
 
 	if len(handlers) == 0 {
 		return ctx, inOut
@@ -104,7 +104,7 @@ func OnWithStreamHandle[S any](
 	inOuts := cpy(len(handlers) + 1)
 
 	for i, handler := range handlers {
-		ctx = handle(handler, inOuts[i])
+		ctx = handle(ctx, handler, inOuts[i])
 	}
 
 	return ctx, inOuts[len(inOuts)-1]
@@ -117,7 +117,7 @@ func OnStartWithStreamInputHandle[T any](ctx context.Context, input *schema.Stre
 
 	cpy := input.Copy
 
-	handle := func(handler Handler, in *schema.StreamReader[T]) context.Context {
+	handle := func(ctx context.Context, handler Handler, in *schema.StreamReader[T]) context.Context {
 		in_ := schema.StreamReaderWithConvert(in, func(i T) (CallbackInput, error) {
 			return i, nil
 		})
@@ -132,7 +132,7 @@ func OnEndWithStreamOutputHandle[T any](ctx context.Context, output *schema.Stre
 
 	cpy := output.Copy
 
-	handle := func(handler Handler, out *schema.StreamReader[T]) context.Context {
+	handle := func(ctx context.Context, handler Handler, out *schema.StreamReader[T]) context.Context {
 		out_ := schema.StreamReaderWithConvert(out, func(i T) (CallbackOutput, error) {
 			return i, nil
 		})