From 7fcb608fce4aafab8ff491d49a6a6e8033987cc1 Mon Sep 17 00:00:00 2001
From: Kent Dong <ch3cho@qq.com>
Date: Thu, 23 May 2024 09:35:57 +0800
Subject: [PATCH] fix: Fix the incorrect usage of DisableReroute (#991)

---
 plugins/wasm-go/extensions/ai-proxy/main.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go
index f4d67de42..3c5676129 100644
--- a/plugins/wasm-go/extensions/ai-proxy/main.go
+++ b/plugins/wasm-go/extensions/ai-proxy/main.go
@@ -70,6 +70,9 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf
 	ctx.SetContext(ctxKeyApiName, apiName)
 
 	if handler, ok := activeProvider.(provider.RequestHeadersHandler); ok {
+		// Disable the route re-calculation since the plugin may modify some headers related to  the chosen route.
+		ctx.DisableReroute()
+
 		action, err := handler.OnRequestHeaders(ctx, apiName, log)
 		if err == nil {
 			return action
@@ -153,9 +156,6 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginCo
 		ctx.BufferResponseBody()
 	}
 
-	// Disable the route re-calculation since the plugin may modify some headers related to  the chosen route.
-	ctx.DisableReroute()
-
 	return types.ActionContinue
 }