81 lines
2.0 KiB
Go
81 lines
2.0 KiB
Go
package llm
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"slices"
|
|
)
|
|
|
|
const maxObservedBodyBytes = 32*1024*1024 + 1
|
|
|
|
// bodyObserver counts the bytes a reader returns from an HTTP response body
|
|
// and whether it reached io.EOF. It never copies or keeps the bytes.
|
|
type bodyObserver struct {
|
|
src io.Reader
|
|
protocol string
|
|
framing string
|
|
declared int64
|
|
uncompressed bool
|
|
bytes int64
|
|
saturated bool
|
|
eof bool
|
|
}
|
|
|
|
// newBodyObserver classifies only what net/http exposes on the response;
|
|
// anything it cannot see reliably is recorded as unknown.
|
|
func newBodyObserver(resp *http.Response) *bodyObserver {
|
|
o := &bodyObserver{src: resp.Body, protocol: "other", framing: "unknown",
|
|
declared: resp.ContentLength, uncompressed: resp.Uncompressed}
|
|
switch resp.ProtoMajor {
|
|
case 1:
|
|
o.protocol = "http1"
|
|
switch {
|
|
case slices.Contains(resp.TransferEncoding, "chunked"):
|
|
o.framing = "chunked"
|
|
case resp.Uncompressed:
|
|
// Transparent decompression hides the wire framing and length.
|
|
case resp.ContentLength >= 0:
|
|
o.framing = "known_length"
|
|
case resp.Close:
|
|
o.framing = "close_delimited"
|
|
}
|
|
case 2:
|
|
o.protocol = "http2"
|
|
o.framing = "http2_stream"
|
|
}
|
|
return o
|
|
}
|
|
|
|
func (o *bodyObserver) Read(p []byte) (int, error) {
|
|
n, err := o.src.Read(p)
|
|
if n > 0 && !o.saturated {
|
|
if int64(n) >= maxObservedBodyBytes-o.bytes {
|
|
o.bytes = maxObservedBodyBytes
|
|
o.saturated = true
|
|
} else {
|
|
o.bytes += int64(n)
|
|
}
|
|
}
|
|
if err == io.EOF {
|
|
o.eof = true
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
// snapshot must be taken after the last Read has returned.
|
|
func (o *bodyObserver) snapshot() *BodyBoundaryObservation {
|
|
relation := "unknown"
|
|
if o.declared >= 0 && !o.uncompressed && !o.saturated && o.eof {
|
|
relation = "mismatch"
|
|
if o.bytes == o.declared {
|
|
relation = "match"
|
|
}
|
|
}
|
|
return &BodyBoundaryObservation{
|
|
ProtocolClass: o.protocol, FramingClass: o.framing,
|
|
DeclaredLengthKnown: o.declared >= 0, BytesRead: o.bytes,
|
|
BytesSaturated: o.saturated, EOFSeen: o.eof,
|
|
LengthRelation: relation, Uncompressed: o.uncompressed,
|
|
}
|
|
}
|