You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@apisix.apache.org by "nic-6443 (via GitHub)" <gi...@apache.org> on 2026/04/17 14:48:40 UTC
[PR] feat(ai-proxy): add max_stream_duration_ms and max_response_bytes safeguards [apisix]
nic-6443 opened a new pull request, #13250:
URL: https://github.com/apache/apisix/pull/13250
### What this does
Adds two opt-in configuration knobs to `ai-proxy` and `ai-proxy-multi` to protect the gateway from a runaway upstream LLM service:
- `max_stream_duration_ms` — wall-clock cap on total streaming response duration.
- `max_response_bytes` — cap on total bytes read from the upstream for a single response (streaming or non-streaming).
Both are opt-in (no default) — existing deployments are unaffected.
### Why
The existing `timeout` field is fed to `httpc:set_timeout()`, which is a per-socket-operation timeout (connect / send / read-one-block). It does **not** bound the total duration of a streaming response. If an upstream LLM has a bug that causes it to continuously emit valid SSE tokens without ever sending a terminator (`[DONE]`, `message_stop`, `response.completed`), `parse_streaming_response` sits in an uncapped `while true` loop, pinning the worker at ~100% CPU indefinitely and degrading availability for all other traffic on that worker.
### Behavior on abort
- **Streaming, limit hit mid-stream (bytes already flushed):** stop feeding chunks and force-close the upstream httpc (`close()` + `res._httpc = nil`, so we don't pool a half-drained connection). nginx closes the downstream connection at end of content phase. The client detects truncation via the missing protocol-specific terminator. We intentionally do **not** synthesize a per-protocol "graceful error" SSE frame: we support three client protocols (OpenAI chat, Anthropic messages, OpenAI responses) with different terminators, and a missing terminator is the standard SSE way any mid-stream network failure is communicated to clients.
- **Streaming, limit hit before any output:** return `504` (duration) or `502` (size) so `on_error` / fallback / retry hooks can kick in like any other upstream failure.
- **Non-streaming, `Content-Length` exceeds cap:** pre-check the header, force-close the connection, return `502` without ever reading the body.
- **Non-streaming, chunked / no `Content-Length`:** post-read size check catches the oversized body and returns `502`.
- `ctx.var.llm_request_done = true` is set on abort so downstream filters (e.g. moderation plugins that defer work until completion) finalize their state.
- A `core.log.warn` line is emitted on every abort (`aborting AI stream: <limit> exceeded; bytes=X duration_ms=Y route_id=Z`) so log-based alerting can surface the event. No new Prometheus metric — the log line is sufficient and avoids expanding the plugin's metric surface.
### Caveat (documented)
Both limits are best-effort: they are enforced after each chunk is read from the upstream, so the byte cap can overshoot by up to one upstream chunk (≈8 KiB in practice) and the duration cap can overshoot by up to one chunk's processing time. This is acceptable for the failure mode we are defending against (runaway streams produce tens of MB/s, so a one-chunk overshoot is negligible compared to "run forever").
### Testing
New `t/plugin/ai-proxy-stream-limits.t` with a mock upstream that either streams OpenAI chat SSE chunks forever (no `[DONE]`) or returns a 100 KB body with matching `Content-Length`. Covers:
1. `max_stream_duration_ms=500` → request aborted in <5 s with the expected log line.
2. `max_response_bytes=2048` → request aborted in <5 s with the expected log line.
3. Non-streaming `max_response_bytes=1024` vs 100 KB upstream response → 502 + expected log line.
4. Schema validation rejects `max_stream_duration_ms: 0`.
`luacheck` passes on all three modified Lua files.
### Docs
Added rows to the config tables in `docs/en/latest/plugins/ai-proxy.md`, `ai-proxy-multi.md`, and their Chinese translations, with a clarifying note that `timeout` only bounds per-socket-operation timeouts and the new fields are needed to bound total stream duration / total bytes read.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@apisix.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(ai-proxy): add max_stream_duration_ms and max_response_bytes safeguards [apisix]
Posted by "Baoyuantop (via GitHub)" <gi...@apache.org>.
Baoyuantop merged PR #13250:
URL: https://github.com/apache/apisix/pull/13250
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@apisix.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(ai-proxy): add max_stream_duration_ms and max_response_bytes safeguards [apisix]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #13250:
URL: https://github.com/apache/apisix/pull/13250#discussion_r3101212290
##########
apisix/plugins/ai-providers/base.lua:
##########
@@ -196,15 +197,45 @@ end
-- using the client protocol module.
-- @param client_proto table The protocol module for the client's protocol
-- @param converter table|nil The converter module (if protocol conversion needed)
+-- @param conf table|nil Plugin configuration (used for response size limits)
-- @return table|nil Parsed and optionally converted response body
-- @return string|nil Error
-function _M.parse_response(self, ctx, res, client_proto, converter)
+function _M.parse_response(self, ctx, res, client_proto, converter, conf)
local headers = res.headers
+
+ -- Pre-check Content-Length against max_response_bytes when available.
+ -- Responses without Content-Length (chunked) fall through to read_body;
+ -- use conf.timeout as the backstop in that case.
+ local max_bytes = conf and conf.max_response_bytes
+ if max_bytes then
+ local content_length = tonumber(headers["Content-Length"])
+ if content_length and content_length > max_bytes then
+ core.log.warn("aborting AI response: Content-Length ", content_length,
+ " exceeds max_response_bytes ", max_bytes)
+ if res._httpc then
+ res._httpc:close()
+ res._httpc = nil
+ end
+ return nil, "max_response_bytes exceeded", 502
+ end
+ end
+
local raw_res_body, err = res:read_body()
Review Comment:
`max_response_bytes` is not actually enforced for responses without `Content-Length`: the code falls through to `res:read_body()` which buffers the entire (potentially unbounded) chunked body before the size check runs. This defeats the safeguard and can still hang or OOM the worker if the upstream streams indefinitely. Consider reading via `res.body_reader()` in a loop, accumulating size and aborting/closing the upstream connection as soon as the limit is exceeded (returning 502).
##########
apisix/plugins/ai-proxy/base.lua:
##########
@@ -228,12 +228,12 @@ function _M.before_proxy(conf, ctx, on_error)
return 500
end
code = ai_provider:parse_streaming_response(
Review Comment:
`parse_streaming_response` can return both a status code and an error message (e.g. when the new stream limits are hit before any bytes are flushed), but only the status code is captured here. Capture and return the second value as `body` as well so clients and `on_error` handlers can see the reason when no output was sent yet.
```suggestion
code, body = ai_provider:parse_streaming_response(
```
##########
t/plugin/ai-proxy-stream-limits.t:
##########
@@ -0,0 +1,367 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+use t::APISIX 'no_plan';
+
+log_level("info");
+repeat_each(1);
+no_long_string();
+no_root_location();
+
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!defined $block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ # Shared mock upstream: a runaway SSE server that streams OpenAI chat
+ # completion chunks indefinitely and never sends "[DONE]".
+ my $http_config = $block->http_config // <<_EOC_;
+ server {
+ server_name runaway_openai_sse;
+ listen 7740;
+
+ default_type 'text/event-stream';
+
+ location /v1/chat/completions {
+ content_by_lua_block {
+ ngx.header["Content-Type"] = "text/event-stream"
+ -- Bound by an upper limit so test-nginx never hangs even
+ -- if the plugin safeguard misfires; in practice the plugin
+ -- should abort long before this completes.
+ for i = 1, 10000 do
+ ngx.print('data: {"id":"chatcmpl-1","object":'
+ .. '"chat.completion.chunk","choices":[{"delta":'
+ .. '{"content":"token"},"index":0,'
+ .. '"finish_reason":null}],"usage":null}\\n\\n')
+ ngx.flush(true)
+ ngx.sleep(0.01)
+ end
+ -- Deliberately never send [DONE].
+ }
+ }
+
+ location /v1/oversized {
+ content_by_lua_block {
+ -- Advertise a large Content-Length to trigger the
+ -- non-streaming max_response_bytes pre-check.
+ ngx.header["Content-Type"] = "application/json"
+ ngx.header["Content-Length"] = "100000"
+ ngx.print(string.rep("x", 100000))
+ }
+ }
+ }
+_EOC_
+
+ $block->set_value("http_config", $http_config);
+});
+
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: set route with max_stream_duration_ms against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_stream_duration_ms": 500,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 2: max_stream_duration_ms aborts the stream within budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+ local core = require("apisix.core")
Review Comment:
In TEST 2's Lua snippet, `local core = require("apisix.core")` is unused. Removing it would reduce noise and avoid confusion about an intended assertion/logging.
```suggestion
```
##########
t/plugin/ai-proxy-stream-limits.t:
##########
@@ -0,0 +1,367 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+use t::APISIX 'no_plan';
+
+log_level("info");
+repeat_each(1);
+no_long_string();
+no_root_location();
+
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!defined $block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ # Shared mock upstream: a runaway SSE server that streams OpenAI chat
+ # completion chunks indefinitely and never sends "[DONE]".
+ my $http_config = $block->http_config // <<_EOC_;
+ server {
+ server_name runaway_openai_sse;
+ listen 7740;
+
+ default_type 'text/event-stream';
+
+ location /v1/chat/completions {
+ content_by_lua_block {
+ ngx.header["Content-Type"] = "text/event-stream"
+ -- Bound by an upper limit so test-nginx never hangs even
+ -- if the plugin safeguard misfires; in practice the plugin
+ -- should abort long before this completes.
+ for i = 1, 10000 do
+ ngx.print('data: {"id":"chatcmpl-1","object":'
+ .. '"chat.completion.chunk","choices":[{"delta":'
+ .. '{"content":"token"},"index":0,'
+ .. '"finish_reason":null}],"usage":null}\\n\\n')
+ ngx.flush(true)
+ ngx.sleep(0.01)
+ end
+ -- Deliberately never send [DONE].
+ }
+ }
+
+ location /v1/oversized {
+ content_by_lua_block {
+ -- Advertise a large Content-Length to trigger the
+ -- non-streaming max_response_bytes pre-check.
+ ngx.header["Content-Type"] = "application/json"
+ ngx.header["Content-Length"] = "100000"
+ ngx.print(string.rep("x", 100000))
+ }
+ }
+ }
+_EOC_
+
+ $block->set_value("http_config", $http_config);
+});
+
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: set route with max_stream_duration_ms against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_stream_duration_ms": 500,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 2: max_stream_duration_ms aborts the stream within budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+ local core = require("apisix.core")
+
+ local ok, err = httpc:connect({
+ scheme = "http",
+ host = "localhost",
+ port = ngx.var.server_port,
+ })
+ if not ok then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local start = ngx.now()
+ local res, err = httpc:request({
+ method = "POST",
+ headers = { ["Content-Type"] = "application/json" },
+ path = "/anything",
+ body = [[{"messages": [{"role": "user", "content": "hi"}]}]],
+ })
+ if not res then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ -- Drain whatever the gateway is willing to send; it must end
+ -- within a few seconds because the plugin will close the
+ -- upstream after ~500ms.
+ local chunks = 0
+ while true do
+ local chunk, rerr = res.body_reader()
+ if rerr or not chunk then break end
+ chunks = chunks + 1
+ end
+ local elapsed = ngx.now() - start
+
+ -- The test mock would run for ~100s without the safeguard;
+ -- with the safeguard the whole exchange must finish in well
+ -- under 5s.
+ if elapsed >= 5 then
+ ngx.status = 500
+ ngx.say("stream did not abort in time: ", elapsed, "s")
+ return
+ end
+ if chunks == 0 then
+ ngx.status = 500
+ ngx.say("no chunks received")
+ return
+ end
+ ngx.say("ok")
+ }
+ }
+--- response_body
+ok
+--- error_log
+aborting AI stream: max_stream_duration_ms exceeded
+
+
+
+=== TEST 3: set route with max_response_bytes against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_response_bytes": 2048,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 4: max_response_bytes aborts the stream after byte budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+
+ local ok, err = httpc:connect({
+ scheme = "http",
+ host = "localhost",
+ port = ngx.var.server_port,
+ })
+ if not ok then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local res, err = httpc:request({
+ method = "POST",
+ headers = { ["Content-Type"] = "application/json" },
+ path = "/anything",
+ body = [[{"messages": [{"role": "user", "content": "hi"}]}]],
+ })
+ if not res then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local total = 0
+ local start = ngx.now()
+ while true do
+ local chunk, rerr = res.body_reader()
+ if rerr or not chunk then break end
+ total = total + #chunk
+ end
+ local elapsed = ngx.now() - start
+
+ if elapsed >= 5 then
+ ngx.status = 500
+ ngx.say("stream did not abort in time: ", elapsed, "s")
+ return
+ end
+ if total == 0 then
+ ngx.status = 500
+ ngx.say("no bytes received")
+ return
+ end
+ ngx.say("ok")
+ }
+ }
+--- response_body
+ok
+--- error_log
+aborting AI stream: max_response_bytes exceeded
+
+
+
+=== TEST 5: set route for non-streaming oversized Content-Length
+--- config
+ location /t {
+ content_by_lua_block {
Review Comment:
This test file covers the `Content-Length` pre-check path for `max_response_bytes`, but doesn’t cover the chunked / missing-`Content-Length` path. Add a case where the upstream sends a chunked response that exceeds `max_response_bytes` to ensure the limit is enforced for that scenario too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@apisix.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(ai-proxy): add max_stream_duration_ms and max_response_bytes safeguards [apisix]
Posted by "nic-6443 (via GitHub)" <gi...@apache.org>.
nic-6443 commented on code in PR #13250:
URL: https://github.com/apache/apisix/pull/13250#discussion_r3103604972
##########
apisix/plugins/ai-providers/base.lua:
##########
@@ -196,15 +197,45 @@ end
-- using the client protocol module.
-- @param client_proto table The protocol module for the client's protocol
-- @param converter table|nil The converter module (if protocol conversion needed)
+-- @param conf table|nil Plugin configuration (used for response size limits)
-- @return table|nil Parsed and optionally converted response body
-- @return string|nil Error
-function _M.parse_response(self, ctx, res, client_proto, converter)
+function _M.parse_response(self, ctx, res, client_proto, converter, conf)
local headers = res.headers
+
+ -- Pre-check Content-Length against max_response_bytes when available.
+ -- Responses without Content-Length (chunked) fall through to read_body;
+ -- use conf.timeout as the backstop in that case.
+ local max_bytes = conf and conf.max_response_bytes
+ if max_bytes then
+ local content_length = tonumber(headers["Content-Length"])
+ if content_length and content_length > max_bytes then
+ core.log.warn("aborting AI response: Content-Length ", content_length,
+ " exceeds max_response_bytes ", max_bytes)
+ if res._httpc then
+ res._httpc:close()
+ res._httpc = nil
+ end
+ return nil, "max_response_bytes exceeded", 502
+ end
+ end
+
local raw_res_body, err = res:read_body()
Review Comment:
Already fixed in commit 61e2551b7 (the chunked enforcement follow-up). `parse_response` now reads via `res.body_reader()` in a loop and trips the cap incrementally for chunked / no-Content-Length responses; the falling-through to `res:read_body()` only happens when `max_response_bytes` is unset (opt-in).
##########
apisix/plugins/ai-proxy/base.lua:
##########
@@ -228,12 +228,12 @@ function _M.before_proxy(conf, ctx, on_error)
return 500
end
code = ai_provider:parse_streaming_response(
Review Comment:
Fixed in f1195052b — capturing both `code` and `body` so the abort reason (`max_stream_duration_ms exceeded` / `max_response_bytes exceeded`) reaches `on_error` / fallback handlers when no bytes were flushed yet.
##########
t/plugin/ai-proxy-stream-limits.t:
##########
@@ -0,0 +1,367 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+use t::APISIX 'no_plan';
+
+log_level("info");
+repeat_each(1);
+no_long_string();
+no_root_location();
+
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!defined $block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ # Shared mock upstream: a runaway SSE server that streams OpenAI chat
+ # completion chunks indefinitely and never sends "[DONE]".
+ my $http_config = $block->http_config // <<_EOC_;
+ server {
+ server_name runaway_openai_sse;
+ listen 7740;
+
+ default_type 'text/event-stream';
+
+ location /v1/chat/completions {
+ content_by_lua_block {
+ ngx.header["Content-Type"] = "text/event-stream"
+ -- Bound by an upper limit so test-nginx never hangs even
+ -- if the plugin safeguard misfires; in practice the plugin
+ -- should abort long before this completes.
+ for i = 1, 10000 do
+ ngx.print('data: {"id":"chatcmpl-1","object":'
+ .. '"chat.completion.chunk","choices":[{"delta":'
+ .. '{"content":"token"},"index":0,'
+ .. '"finish_reason":null}],"usage":null}\\n\\n')
+ ngx.flush(true)
+ ngx.sleep(0.01)
+ end
+ -- Deliberately never send [DONE].
+ }
+ }
+
+ location /v1/oversized {
+ content_by_lua_block {
+ -- Advertise a large Content-Length to trigger the
+ -- non-streaming max_response_bytes pre-check.
+ ngx.header["Content-Type"] = "application/json"
+ ngx.header["Content-Length"] = "100000"
+ ngx.print(string.rep("x", 100000))
+ }
+ }
+ }
+_EOC_
+
+ $block->set_value("http_config", $http_config);
+});
+
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: set route with max_stream_duration_ms against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_stream_duration_ms": 500,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 2: max_stream_duration_ms aborts the stream within budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+ local core = require("apisix.core")
Review Comment:
Fixed in f1195052b — removed the unused require.
##########
t/plugin/ai-proxy-stream-limits.t:
##########
@@ -0,0 +1,367 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+use t::APISIX 'no_plan';
+
+log_level("info");
+repeat_each(1);
+no_long_string();
+no_root_location();
+
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!defined $block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ # Shared mock upstream: a runaway SSE server that streams OpenAI chat
+ # completion chunks indefinitely and never sends "[DONE]".
+ my $http_config = $block->http_config // <<_EOC_;
+ server {
+ server_name runaway_openai_sse;
+ listen 7740;
+
+ default_type 'text/event-stream';
+
+ location /v1/chat/completions {
+ content_by_lua_block {
+ ngx.header["Content-Type"] = "text/event-stream"
+ -- Bound by an upper limit so test-nginx never hangs even
+ -- if the plugin safeguard misfires; in practice the plugin
+ -- should abort long before this completes.
+ for i = 1, 10000 do
+ ngx.print('data: {"id":"chatcmpl-1","object":'
+ .. '"chat.completion.chunk","choices":[{"delta":'
+ .. '{"content":"token"},"index":0,'
+ .. '"finish_reason":null}],"usage":null}\\n\\n')
+ ngx.flush(true)
+ ngx.sleep(0.01)
+ end
+ -- Deliberately never send [DONE].
+ }
+ }
+
+ location /v1/oversized {
+ content_by_lua_block {
+ -- Advertise a large Content-Length to trigger the
+ -- non-streaming max_response_bytes pre-check.
+ ngx.header["Content-Type"] = "application/json"
+ ngx.header["Content-Length"] = "100000"
+ ngx.print(string.rep("x", 100000))
+ }
+ }
+ }
+_EOC_
+
+ $block->set_value("http_config", $http_config);
+});
+
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: set route with max_stream_duration_ms against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_stream_duration_ms": 500,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 2: max_stream_duration_ms aborts the stream within budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+ local core = require("apisix.core")
+
+ local ok, err = httpc:connect({
+ scheme = "http",
+ host = "localhost",
+ port = ngx.var.server_port,
+ })
+ if not ok then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local start = ngx.now()
+ local res, err = httpc:request({
+ method = "POST",
+ headers = { ["Content-Type"] = "application/json" },
+ path = "/anything",
+ body = [[{"messages": [{"role": "user", "content": "hi"}]}]],
+ })
+ if not res then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ -- Drain whatever the gateway is willing to send; it must end
+ -- within a few seconds because the plugin will close the
+ -- upstream after ~500ms.
+ local chunks = 0
+ while true do
+ local chunk, rerr = res.body_reader()
+ if rerr or not chunk then break end
+ chunks = chunks + 1
+ end
+ local elapsed = ngx.now() - start
+
+ -- The test mock would run for ~100s without the safeguard;
+ -- with the safeguard the whole exchange must finish in well
+ -- under 5s.
+ if elapsed >= 5 then
+ ngx.status = 500
+ ngx.say("stream did not abort in time: ", elapsed, "s")
+ return
+ end
+ if chunks == 0 then
+ ngx.status = 500
+ ngx.say("no chunks received")
+ return
+ end
+ ngx.say("ok")
+ }
+ }
+--- response_body
+ok
+--- error_log
+aborting AI stream: max_stream_duration_ms exceeded
+
+
+
+=== TEST 3: set route with max_response_bytes against runaway SSE upstream
+--- config
+ location /t {
+ content_by_lua_block {
+ local t = require("lib.test_admin").test
+ local code, body = t('/apisix/admin/routes/1',
+ ngx.HTTP_PUT,
+ [[{
+ "uri": "/anything",
+ "plugins": {
+ "ai-proxy": {
+ "provider": "openai",
+ "auth": {
+ "header": {
+ "Authorization": "Bearer token"
+ }
+ },
+ "options": {
+ "model": "gpt-3.5-turbo",
+ "stream": true
+ },
+ "max_response_bytes": 2048,
+ "override": {
+ "endpoint": "http://localhost:7740"
+ },
+ "ssl_verify": false
+ }
+ }
+ }]]
+ )
+
+ if code >= 300 then
+ ngx.status = code
+ end
+ ngx.say(body)
+ }
+ }
+--- response_body
+passed
+
+
+
+=== TEST 4: max_response_bytes aborts the stream after byte budget
+--- config
+ location /t {
+ content_by_lua_block {
+ local http = require("resty.http")
+ local httpc = http.new()
+
+ local ok, err = httpc:connect({
+ scheme = "http",
+ host = "localhost",
+ port = ngx.var.server_port,
+ })
+ if not ok then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local res, err = httpc:request({
+ method = "POST",
+ headers = { ["Content-Type"] = "application/json" },
+ path = "/anything",
+ body = [[{"messages": [{"role": "user", "content": "hi"}]}]],
+ })
+ if not res then
+ ngx.status = 500
+ ngx.say(err)
+ return
+ end
+
+ local total = 0
+ local start = ngx.now()
+ while true do
+ local chunk, rerr = res.body_reader()
+ if rerr or not chunk then break end
+ total = total + #chunk
+ end
+ local elapsed = ngx.now() - start
+
+ if elapsed >= 5 then
+ ngx.status = 500
+ ngx.say("stream did not abort in time: ", elapsed, "s")
+ return
+ end
+ if total == 0 then
+ ngx.status = 500
+ ngx.say("no bytes received")
+ return
+ end
+ ngx.say("ok")
+ }
+ }
+--- response_body
+ok
+--- error_log
+aborting AI stream: max_response_bytes exceeded
+
+
+
+=== TEST 5: set route for non-streaming oversized Content-Length
+--- config
+ location /t {
+ content_by_lua_block {
Review Comment:
Already covered in commit 61e2551b7 by TEST 8 + TEST 9 (`/v1/oversized_chunked` mock upstream emits 10 × 10KB chunks with no Content-Length, configured limit 5KB; expects 502).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@apisix.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org