fix(llm): read a bounded response body, 8 MiB on success and the error bound on failure (#53)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-10-03 21:22:17 +02:00
co-authored by Claude Sonnet 5.5
parent 9cc954df40
commit 61d1ff2255
+24 -3
View File
@@ -37,6 +37,21 @@ type Client struct {
temp float64 temp float64
httpClient *http.Client httpClient *http.Client
usageHook func(model string, prompt, completion int) usageHook func(model string, prompt, completion int)
maxBody int64
}
// defaultMaxResponseBytes bounds a success body. A completion is a few KB of text; this
// leaves room for a long one and none for a gateway that streams forever (#53).
const defaultMaxResponseBytes = 8 << 20
// WithMaxResponseBytes overrides the bound on a success response body. A non-positive n is
// ignored (keeps the default).
func WithMaxResponseBytes(n int64) Option {
return func(c *Client) {
if n > 0 {
c.maxBody = n
}
}
} }
// Option configures a Client at construction. Variadic so the existing 4-arg // Option configures a Client at construction. Variadic so the existing 4-arg
@@ -101,6 +116,7 @@ func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *
maxTokens: defaultMaxTokens, maxTokens: defaultMaxTokens,
temp: defaultTemperature, temp: defaultTemperature,
httpClient: &http.Client{Timeout: timeout}, httpClient: &http.Client{Timeout: timeout},
maxBody: defaultMaxResponseBytes,
} }
for _, opt := range opts { for _, opt := range opts {
opt(c) opt(c)
@@ -184,12 +200,17 @@ func (c *Client) Complete(ctx context.Context, system, user string) (string, err
} }
defer resp.Body.Close() //nolint:errcheck defer resp.Body.Close() //nolint:errcheck
out, err := io.ReadAll(resp.Body) if resp.StatusCode != http.StatusOK {
// Only the bounded head is ever used (HTTPError.Error truncates to it), so read no more.
head, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody+1))
return "", &HTTPError{Code: resp.StatusCode, Body: string(head)}
}
out, err := io.ReadAll(io.LimitReader(resp.Body, c.maxBody+1))
if err != nil { if err != nil {
return "", fmt.Errorf("read response: %w", err) return "", fmt.Errorf("read response: %w", err)
} }
if resp.StatusCode != http.StatusOK { if int64(len(out)) > c.maxBody {
return "", &HTTPError{Code: resp.StatusCode, Body: string(out)} return "", fmt.Errorf("LLM response exceeds %d bytes", c.maxBody)
} }
var cr chatResponse var cr chatResponse