ai-cache update body buffer limit size (#1644)

alibaba · Jan 6, 2025 · 2501895 · 2501895
1 parent 187a7b5
commit 2501895
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/plugins/wasm-go/extensions/ai-cache/main.go b/plugins/wasm-go/extensions/ai-cache/main.go
@@ -22,6 +22,8 @@ const (
 	STREAM_CONTEXT_KEY          = "stream"
 	SKIP_CACHE_HEADER           = "x-higress-skip-ai-cache"
 	ERROR_PARTIAL_MESSAGE_KEY   = "errorPartialMessage"
+
+	DEFAULT_MAX_BODY_BYTES uint32 = 10 * 1024 * 1024
 )
 
 func main() {
@@ -69,6 +71,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log wr
 		ctx.DontReadRequestBody()
 		return types.ActionContinue
 	}
+	ctx.SetRequestBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
 	_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
 	// The request has a body and requires delaying the header transmission until a cache miss occurs,
 	// at which point the header should be sent.
@@ -140,6 +143,8 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, c config.PluginConfig, log w
 	contentType, _ := proxywasm.GetHttpResponseHeader("content-type")
 	if strings.Contains(contentType, "text/event-stream") {
 		ctx.SetContext(STREAM_CONTEXT_KEY, struct{}{})
+	} else {
+		ctx.SetResponseBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
 	}
 
 	if ctx.GetContext(ERROR_PARTIAL_MESSAGE_KEY) != nil {

diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go
@@ -261,6 +261,7 @@ func checkStream(ctx *wrapper.HttpContext, log wrapper.Log) {
 			log.Errorf("unable to load content-type header from response: %v", err)
 		}
 		(*ctx).BufferResponseBody()
+		ctx.SetResponseBodyBufferLimit(defaultMaxBodyBytes)
 	}
 }
 

diff --git a/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go b/plugins/wasm-go/pkg/wrapper/plugin_wrapper.go
@@ -65,9 +65,9 @@ type HttpContext interface {
 	// You need to call this before making any header modification operations.
 	DisableReroute()
 	// Note that this parameter affects the gateway's memory usage！Support setting a maximum buffer size for each request body individually in request phase.
-	SetRequestBodyBufferLimit(size uint32)
+	SetRequestBodyBufferLimit(byteSize uint32)
 	// Note that this parameter affects the gateway's memory usage! Support setting a maximum buffer size for each response body individually in response phase.
-	SetResponseBodyBufferLimit(size uint32)
+	SetResponseBodyBufferLimit(byteSize uint32)
 }
 
 type ParseConfigFunc[PluginConfig any] func(json gjson.Result, config *PluginConfig, log Log) error
-Original file line number
+Diff line change
@@ Expand Up @@
     			log.Errorf("unable to load content-type header from response: %v", err)
     		}
     		(*ctx).BufferResponseBody()
+    		ctx.SetResponseBodyBufferLimit(defaultMaxBodyBytes)
     	}
     }
@@ Expand Down @@