feat(plugins): enable ai plugins to read request body from buffered file

changelog/unreleased/kong/ai-plugin-read-file.yml

@@ -0,0 +1,3 @@
+message: "allow AI plugin to read request from buffered file" 
+type: feature
+scope: "Plugin"

changelog/unreleased/kong/pdk-read-file.yml

@@ -0,0 +1,3 @@
+message: "extend kong.request.get_body and kong.request.get_raw_body to read from buffered file"
+type: feature
+scope: "PDK"

kong/clustering/compat/removed_fields.lua

@@ -167,5 +167,23 @@ return {
       "traces_endpoint",
       "logs_endpoint",
     },
+    ai_proxy = {
+      "max_request_body_size",
+    },
+    ai_prompt_decorator = {
+      "max_request_body_size",
+    },
+    ai_prompt_guard = {
+      "max_request_body_size",
+    },
+    ai_prompt_template = {
+      "max_request_body_size",
+    },
+    ai_request_transformer = {
+      "max_request_body_size",
+    },
+    ai_response_transformer = {
+      "max_request_body_size",
+    },
   },
 }

kong/pdk/request.lua

@@ -10,6 +10,7 @@ local cjson = require "kong.tools.cjson"
 local multipart = require "multipart"
 local phase_checker = require "kong.pdk.private.phases"
 local normalize = require("kong.tools.uri").normalize
+local yield = require("kong.tools.yield").yield
 
 
 local ngx = ngx
@@ -693,30 +694,69 @@ local function new(self)
   --
   -- If the size of the body is greater than the Nginx buffer size (set by
   -- `client_body_buffer_size`), this function fails and returns an error
-  -- message explaining this limitation.
+  -- message explaining this limitation, unless `max_allowed_file_size`
+  -- is set and larger than the body size buffered to disk.
+  -- Use of `max_allowed_file_size` requires Kong to read data from filesystem
+  -- and has performance implications.
   --
   -- @function kong.request.get_raw_body
   -- @phases rewrite, access, response, admin_api
+  -- @max_allowed_file_size[opt] number the max allowed file size to be read from
   -- @treturn string|nil The plain request body or nil if it does not fit into
   -- the NGINX temporary buffer.
   -- @treturn nil|string An error message.
   -- @usage
   -- -- Given a body with payload "Hello, Earth!":
   --
   -- kong.request.get_raw_body():gsub("Earth", "Mars") -- "Hello, Mars!"
-  function _REQUEST.get_raw_body()
+  function _REQUEST.get_raw_body(max_allowed_file_size)
     check_phase(before_content)
 
     read_body()
 
     local body = get_body_data()
     if not body then
-      if get_body_file() then
+      local body_file = get_body_file()
+      if not body_file then
+        return ""
+      end
+
+      if not max_allowed_file_size then
         return nil, "request body did not fit into client body buffer, consider raising 'client_body_buffer_size'"
+      end
 
-      else
-        return ""
+      local file, err = io.open(body_file, "r")
+      if not file then
+        return nil, "failed to open cached request body '" .. body_file .. "': " .. err
+      end
+
+      local size = file:seek("end") or 0
+      if size > max_allowed_file_size then
+        return nil, ("request body file too big: %d > %d"):format(size, max_allowed_file_size)
       end
+
+      -- go to beginning
+      file:seek("set")
+      local chunk = {}
+      local chunk_idx = 1
+
+      while true do
+        local data, err = file:read(1048576) -- read in chunks of 1mb
+        if not data then
+          if err then
+            return nil, "failed to read cached request body '" .. body_file .. "': " .. err
+          end
+          break
+        end
+        chunk[chunk_idx] = data
+        chunk_idx = chunk_idx + 1
+
+        yield() -- yield to prevent starvation while doing blocking IO-reads
+      end
+
+      file:close()
+
+      return table.concat(chunk, "")
     end
 
     return body
@@ -767,6 +807,7 @@ local function new(self)
   -- @phases rewrite, access, response, admin_api
   -- @tparam[opt] string mimetype The MIME type.
   -- @tparam[opt] number max_args Sets a limit on the maximum number of parsed
+  -- @tparam[opt] number max_allowed_file_size the max allowed file size to be read from
   -- arguments.
   -- @treturn table|nil A table representation of the body.
   -- @treturn string|nil An error message.
@@ -775,7 +816,7 @@ local function new(self)
   -- local body, err, mimetype = kong.request.get_body()
   -- body.name -- "John Doe"
   -- body.age  -- "42"
-  function _REQUEST.get_body(mimetype, max_args)
+  function _REQUEST.get_body(mimetype, max_args, max_allowed_file_size)
     check_phase(before_content)
 
     local content_type = mimetype or _REQUEST.get_header(CONTENT_TYPE)
@@ -825,7 +866,7 @@ local function new(self)
       return pargs, nil, CONTENT_TYPE_POST
 
     elseif find(content_type_lower, CONTENT_TYPE_JSON, 1, true) == 1 then
-      local body, err = _REQUEST.get_raw_body()
+      local body, err = _REQUEST.get_raw_body(max_allowed_file_size)
       if not body then
         return nil, err, CONTENT_TYPE_JSON
       end
@@ -838,7 +879,7 @@ local function new(self)
       return json, nil, CONTENT_TYPE_JSON
 
     elseif find(content_type_lower, CONTENT_TYPE_FORM_DATA, 1, true) == 1 then
-      local body, err = _REQUEST.get_raw_body()
+      local body, err = _REQUEST.get_raw_body(max_allowed_file_size)
       if not body then
         return nil, err, CONTENT_TYPE_FORM_DATA
       end

kong/plugins/ai-prompt-decorator/handler.lua

@@ -55,7 +55,7 @@ function plugin:access(conf)
   kong.ctx.shared.ai_prompt_decorated = true  -- future use
 
   -- if plugin ordering was altered, receive the "decorated" request
-  local request = kong.request.get_body("application/json")
+  local request = kong.request.get_body("application/json", nil, conf.max_request_body_size)
   if type(request) ~= "table"  then
     return bad_request("this LLM route only supports application/json requests")
   end

kong/plugins/ai-prompt-decorator/schema.lua

@@ -39,7 +39,9 @@ return {
     { config = {
       type = "record",
       fields = {
-          { prompts = prompts_record }
+          { prompts = prompts_record },
+          { max_request_body_size = { type = "integer", default = 8 * 1024, gt = 0,
+                                    description = "max allowed body size allowed to be introspected" } },
         }
       }
     }

kong/plugins/ai-prompt-guard/handler.lua

@@ -121,7 +121,7 @@ function plugin:access(conf)
   kong.ctx.shared.ai_prompt_guarded = true -- future use
 
   -- if plugin ordering was altered, receive the "decorated" request
-  local request = kong.request.get_body("application/json")
+  local request = kong.request.get_body("application/json", nil, conf.max_request_body_size)
   if type(request) ~= "table" then
     return bad_request("this LLM route only supports application/json requests")
   end

kong/plugins/ai-prompt-guard/schema.lua

@@ -32,6 +32,12 @@ return {
               type = "boolean",
               required = true,
               default = false } },
+          { max_request_body_size = {
+              type = "integer",
+              default = 8 * 1024,
+              gt = 0,
+              description = "max allowed body size allowed to be introspected",}
+          },
         }
       }
     }

kong/plugins/ai-prompt-template/handler.lua

@@ -64,10 +64,10 @@ function AIPromptTemplateHandler:access(conf)
   kong.ctx.shared.ai_prompt_templated = true
 
   if conf.log_original_request then
-    kong.log.set_serialize_value(LOG_ENTRY_KEYS.REQUEST_BODY, kong.request.get_raw_body())
+    kong.log.set_serialize_value(LOG_ENTRY_KEYS.REQUEST_BODY, kong.request.get_raw_body(conf.max_request_body_size))
   end
 
-  local request = kong.request.get_body("application/json")
+  local request = kong.request.get_body("application/json", nil, conf.max_request_body_size)
   if type(request) ~= "table" then
     return bad_request("this LLM route only supports application/json requests")
   end

kong/plugins/ai-prompt-template/schema.lua

@@ -44,6 +44,12 @@ return {
             required = true,
             default = false,
         }},
+        { max_request_body_size = {
+            type = "integer",
+            default = 8 * 1024,
+            gt = 0,
+            description = "max allowed body size allowed to be introspected",
+        }},
       }
     }}
   },

kong/plugins/ai-proxy/handler.lua

@@ -316,7 +316,7 @@ function _M:access(conf)
     -- first, calculate the coordinates of the request
     local content_type = kong.request.get_header("Content-Type") or "application/json"
 
-    request_table = kong.request.get_body(content_type)
+    request_table = kong.request.get_body(content_type, nil, conf.max_request_body_size)
 
     if not request_table then
       if not string.find(content_type, "multipart/form-data", nil, true) then

kong/plugins/ai-proxy/schema.lua

@@ -13,6 +13,13 @@ local ai_proxy_only_config = {
         default = "allow",
         one_of = { "allow", "deny", "always" }},
     },
+    {
+      max_request_body_size = {
+      type = "integer",
+      default = 8 * 1024,
+      gt = 0,
+      description = "max allowed body size allowed to be introspected",}
+    },
 }
 
 for i, v in pairs(ai_proxy_only_config) do

kong/plugins/ai-request-transformer/handler.lua

@@ -55,7 +55,7 @@ function _M:access(conf)
   -- if asked, introspect the request before proxying
   kong.log.debug("introspecting request with LLM")
   local new_request_body, err = ai_driver:ai_introspect_body(
-    kong.request.get_raw_body(),
+    kong.request.get_raw_body(conf.max_request_body_size),
     conf.prompt,
     http_opts,
     conf.transformation_extract_pattern

kong/plugins/ai-request-transformer/schema.lua

@@ -37,6 +37,14 @@ return {
           default = true,
         }},
 
+        {
+          max_request_body_size = {
+          type = "integer",
+          default = 8 * 1024,
+          gt = 0,
+          description = "max allowed body size allowed to be introspected",}
+        },
+
         -- from forward-proxy
         { http_proxy_host = typedefs.host },
         { http_proxy_port = typedefs.port },
@@ -46,7 +54,7 @@ return {
         { llm = llm.config_schema },
       },
     }},
-    
+
   },
   entity_checks = {
     {

kong/plugins/ai-response-transformer/handler.lua

@@ -113,7 +113,9 @@ function _M:access(conf)
 
   kong.log.debug("intercepting plugin flow with one-shot request")
   local httpc = http.new()
-  local res, err = subrequest(httpc, kong.request.get_raw_body(), http_opts)
+  local res, err = subrequest(httpc,
+    kong.request.get_raw_body(conf.max_request_body_size),
+    http_opts)
   if err then
     return internal_server_error(err)
   end

kong/plugins/ai-response-transformer/schema.lua

@@ -45,6 +45,13 @@ return {
             default = true,
         }},
 
+        { max_request_body_size = {
+            type = "integer",
+            default = 8 * 1024,
+            gt = 0,
+            description = "max allowed body size allowed to be introspected",}
+        },
+
         -- from forward-proxy
         { http_proxy_host = typedefs.host },
         { http_proxy_port = typedefs.port },

spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua

@@ -543,16 +543,22 @@ describe("CP/DP config compat transformations #" .. strategy, function()
                 upstream_path = "/anywhere", -- becomes nil
               },
             },
+            max_request_body_size = 8192,
           },
         }
         -- ]]
 
-        local expected_ai_proxy_prior_37 = cycle_aware_deep_copy(ai_proxy)
-        expected_ai_proxy_prior_37.config.response_streaming = nil
-        expected_ai_proxy_prior_37.config.model.options.upstream_path = nil
-        expected_ai_proxy_prior_37.config.route_type = "llm/v1/chat"
+        local expected = cycle_aware_deep_copy(ai_proxy)
 
-        do_assert(uuid(), "3.6.0", expected_ai_proxy_prior_37)
+        expected.config.max_request_body_size = nil
+
+        do_assert(uuid(), "3.7.0", expected)
+
+        expected.config.response_streaming = nil
+        expected.config.model.options.upstream_path = nil
+        expected.config.route_type = "llm/v1/chat"
+
+        do_assert(uuid(), "3.6.0", expected)
 
         -- cleanup
         admin.plugins:remove({ id = ai_proxy.id })
@@ -584,14 +590,19 @@ describe("CP/DP config compat transformations #" .. strategy, function()
                 },
               },
             },
+            max_request_body_size = 8192,
           },
         }
         -- ]]
 
-        local expected_ai_request_transformer_prior_37 = cycle_aware_deep_copy(ai_request_transformer)
-        expected_ai_request_transformer_prior_37.config.llm.model.options.upstream_path = nil
+        local expected = cycle_aware_deep_copy(ai_request_transformer)
+        expected.config.max_request_body_size = nil
+
+        do_assert(uuid(), "3.7.0", expected)
 
-        do_assert(uuid(), "3.6.0", expected_ai_request_transformer_prior_37)
+        expected.config.llm.model.options.upstream_path = nil
+
+        do_assert(uuid(), "3.6.0", expected)
 
         -- cleanup
         admin.plugins:remove({ id = ai_request_transformer.id })
@@ -621,14 +632,19 @@ describe("CP/DP config compat transformations #" .. strategy, function()
                 },
               },
             },
+            max_request_body_size = 8192,
           },
         }
         -- ]]
 
-        local expected_ai_response_transformer_prior_37 = cycle_aware_deep_copy(ai_response_transformer)
-        expected_ai_response_transformer_prior_37.config.llm.model.options.upstream_path = nil
+        local expected = cycle_aware_deep_copy(ai_response_transformer)
+        expected.config.max_request_body_size = nil
+
+        do_assert(uuid(), "3.7.0", expected)
+
+        expected.config.llm.model.options.upstream_path = nil
 
-        do_assert(uuid(), "3.6.0", expected_ai_response_transformer_prior_37)
+        do_assert(uuid(), "3.6.0", expected)
 
         -- cleanup
         admin.plugins:remove({ id = ai_response_transformer.id })