fix: add max bytes request limit to aibridge (#26164) (#26305) · coder/coder@481857f
11package aibridge_test
2233import (
4+"bytes"
5+"fmt"
6+"io"
47"net/http"
58"net/http/httptest"
9+"strings"
610"testing"
711812"github.com/stretchr/testify/assert"
@@ -205,3 +209,74 @@ func TestPassthroughRoutesForProviders(t *testing.T) {
205209 })
206210 }
207211}
212+213+func TestRequestBodySizeLimit(t *testing.T) {
214+t.Parallel()
215+216+newOpenAI := func(baseURL string) provider.Provider {
217+return aibridge.NewOpenAIProvider(config.OpenAI{Name: "openai", BaseURL: baseURL})
218+ }
219+newAnthropic := func(baseURL string) provider.Provider {
220+return aibridge.NewAnthropicProvider(config.Anthropic{Name: "anthropic", BaseURL: baseURL}, nil)
221+ }
222+newCopilot := func(baseURL string) provider.Provider {
223+return aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: baseURL})
224+ }
225+226+// Each body is a well-formed, schema-valid request for its provider, with
227+// an oversized message content that pushes it past the 32 MiB limit.
228+filler := strings.Repeat("A", 32<<20)
229+chatCompletionsBody := fmt.Appendf(nil, `{"model":"gpt-4","messages":[{"role":"user","content":"%s"}]}`, filler)
230+responsesBody := fmt.Appendf(nil, `{"model":"gpt-4","input":"%s"}`, filler)
231+messagesBody := fmt.Appendf(nil, `{"model":"claude-3-5-sonnet-latest","max_tokens":1024,"messages":[{"role":"user","content":"%s"}]}`, filler)
232+233+tests := []struct {
234+name string
235+provider func(baseURL string) provider.Provider
236+path string
237+body []byte
238+ }{
239+ {name: "openai_passthrough", provider: newOpenAI, path: "/openai/v1/models", body: chatCompletionsBody},
240+ {name: "openai_chat_completions", provider: newOpenAI, path: "/openai/v1/chat/completions", body: chatCompletionsBody},
241+ {name: "openai_responses", provider: newOpenAI, path: "/openai/v1/responses", body: responsesBody},
242+ {name: "anthropic_passthrough", provider: newAnthropic, path: "/anthropic/v1/models", body: messagesBody},
243+ {name: "anthropic_messages", provider: newAnthropic, path: "/anthropic/v1/messages", body: messagesBody},
244+ {name: "copilot_passthrough", provider: newCopilot, path: "/copilot/models", body: chatCompletionsBody},
245+ {name: "copilot_chat_completions", provider: newCopilot, path: "/copilot/chat/completions", body: chatCompletionsBody},
246+ {name: "copilot_responses", provider: newCopilot, path: "/copilot/responses", body: responsesBody},
247+ }
248+249+for _, tc := range tests {
250+t.Run(tc.name, func(t *testing.T) {
251+t.Parallel()
252+253+logger := slogtest.Make(t, nil)
254+255+upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
256+_, _ = io.ReadAll(r.Body)
257+w.WriteHeader(http.StatusOK)
258+ }))
259+t.Cleanup(upstream.Close)
260+261+prov := tc.provider(upstream.URL)
262+bridge, err := aibridge.NewRequestBridge(
263+t.Context(),
264+ []provider.Provider{prov},
265+nil, nil, logger, nil, bridgeTestTracer,
266+ )
267+require.NoError(t, err)
268+269+req := httptest.NewRequest(http.MethodPost, tc.path, bytes.NewReader(tc.body))
270+// Unknown Content-Length
271+req.ContentLength = -1
272+// Copilot's bridged route checks Authorization before reading the
273+// body, so provide a token to reach the read path.
274+req.Header.Set("Authorization", "Bearer test-key")
275+resp := httptest.NewRecorder()
276+bridge.ServeHTTP(resp, req)
277+278+assert.Equal(t, http.StatusRequestEntityTooLarge, resp.Code)
279+assert.Contains(t, resp.Body.String(), "Request body too large")
280+ })
281+ }
282+}