pkg/models/http_stream_body.go
Rama main · réplica de lectura del repositorio original.
package models
import (
"bufio"
"encoding/base64"
"encoding/json"
"strings"
"time"
"unicode/utf8"
yamlLib "gopkg.in/yaml.v3"
)
// maxModelStreamTokenSize caps the maximum size of a single streaming token (10 MB)
// to prevent runaway memory allocation when scanning large streaming responses.
const maxModelStreamTokenSize = 10 * 1024 * 1024
// HTTPStreamDataField represents a single key-value field within a streaming chunk.
//
// For SSE responses, each field maps to an SSE line:
// - key="data", value=<event payload>
// - key="event", value=<event type name>
// - key="id", value=<event ID>
// - key="retry", value=<reconnect delay (ms)>
// - key="comment", value=<comment text after ':'>
//
// For other stream formats (NDJSON, plain text, binary), every chunk
// has exactly one field with key="raw" and value=<the raw line or payload>.
type HTTPStreamDataField struct {
Key string
Value string
}
// HTTPStreamChunk represents one logical unit of a streaming response body.
//
// For SSE this is one complete SSE "frame" (aka event, delimited by blank lines).
// For NDJSON and plain-text streams this is one newline-delimited line.
// For binary streams this is the full captured body.
//
// TS is when the chunk reached the application during recording: for a mock
// recorded with marks (HTTPResp.Marks), the time of the write that completed
// it. Replay keeps no timing between chunks: it writes each part of an
// answer once the replay has reached or passed its recorded place, which it
// reads from the marks.
type HTTPStreamChunk struct {
TS time.Time // When this chunk reached the application during recording
Data []HTTPStreamDataField // Key-value fields that make up this chunk
}
// ValueByKey looks up the value of the field with the given key (case-insensitive).
// Returns ("", false) if no such field exists.
func (c HTTPStreamChunk) ValueByKey(key string) (string, bool) {
key = strings.ToLower(strings.TrimSpace(key))
for _, field := range c.Data {
if strings.ToLower(strings.TrimSpace(field.Key)) == key {
return field.Value, true
}
}
return "", false
}
// ── YAML Serialization ─────────────────────────────────────────────────────────
//
// HTTPResp can be serialized to YAML in two modes:
//
// 1. FLAT (plain string): used for normal, non-streaming HTTP responses.
// body: '{"id":1}'
//
// 2. STRUCTURED (YAML sequence): used for streaming responses (SSE, NDJSON, etc.).
// Each element in the sequence is one chunk (frame/line), storing the
// timestamp and the key-value fields that make up that chunk.
//
// body:
// - ts: "2024-01-01T00:00:01Z"
// data:
// event: message
// data: '{"id":1}'
// - ts: "2024-01-01T00:00:02Z"
// data:
// event: message
// data: '{"id":2}'
//
// MarshalYAML and UnmarshalYAML implement this two-mode logic.
// MarshalYAML serializes an HTTPResp to YAML.
// If the response has streaming chunks (StreamBody), the body field is written
// as a YAML sequence of chunks. Otherwise, it is written as a plain scalar string.
func (h HTTPResp) MarshalYAML() (interface{}, error) {
bodyNode := &yamlLib.Node{}
marks := h.Marks
if chunks, ok := streamChunksForYAML(h); ok {
// Streaming response — encode as a YAML sequence of chunks. A fresh
// recording's frames are split from its body here: each takes the
// time of the write that completed it, and the marks move to the
// body the frames read back as. A response read back from frames
// (StreamBody) has both already.
if marks != nil && len(h.StreamBody) == 0 {
chunks, marks = marksOnFrames(h, chunks)
}
bodyNode = encodeStreamChunksNode(chunks)
} else {
// Non-streaming response, OR a stream whose bytes cannot survive the
// chunked representation — encode the body as a plain string scalar
// (yaml.v3 auto-promotes an untagged non-UTF-8 scalar to !!binary).
//
// Reconstruct the flat body from the chunks when the chunk form was
// rejected and Body is empty: StreamBody travels independently of Body
// over the transport, so a producer that populated only the chunk form
// would otherwise have its entire body silently written as "".
body := h.Body
if body == "" && len(h.StreamBody) > 0 {
body = streamChunksToLegacyBody(h.StreamBody, detectStreamBodyKind(h.Header))
}
if YAMLBlockScalarUnsafe(body) {
// Node.Encode writes it as a block scalar and parses that
// back, which fails (YAMLBlockScalarUnsafe): the mock was lost.
bodyNode = &yamlLib.Node{Kind: yamlLib.ScalarNode, Tag: "!!str", Value: body, Style: yamlLib.DoubleQuotedStyle}
} else if err := bodyNode.Encode(body); err != nil {
return nil, err
}
}
type httpRespYAML struct {
StatusCode int `yaml:"status_code"`
Header map[string]string `yaml:"header"`
Body *yamlLib.Node `yaml:"body"`
BodySkipped bool `yaml:"body_skipped,omitempty"`
BodySize int64 `yaml:"body_size,omitempty"`
StatusMessage string `yaml:"status_message"`
ProtoMajor int `yaml:"proto_major"`
ProtoMinor int `yaml:"proto_minor"`
Binary string `yaml:"binary,omitempty"`
Timestamp time.Time `yaml:"timestamp"`
HeaderLineLengths HeaderLineLengths `yaml:"header_line_lengths,omitempty"`
Marks *HTTPAnswerMarks `yaml:"marks,omitempty"`
}
return httpRespYAML{
StatusCode: h.StatusCode,
Header: h.Header,
Body: bodyNode,
BodySkipped: h.BodySkipped,
BodySize: h.BodySize,
StatusMessage: h.StatusMessage,
ProtoMajor: h.ProtoMajor,
ProtoMinor: h.ProtoMinor,
Binary: h.Binary,
Timestamp: h.Timestamp,
HeaderLineLengths: h.HeaderLineLengths,
Marks: marks,
}, nil
}
// UnmarshalYAML deserializes an HTTPResp from YAML.
// The body field can be either a plain string (non-streaming)
// or a YAML sequence of chunks (streaming). Both formats are handled.
func (h *HTTPResp) UnmarshalYAML(node *yamlLib.Node) error {
type httpRespYAML struct {
StatusCode int `yaml:"status_code"`
Header map[string]string `yaml:"header"`
Body yamlLib.Node `yaml:"body"`
BodySkipped bool `yaml:"body_skipped,omitempty"`
BodySize int64 `yaml:"body_size,omitempty"`
StatusMessage string `yaml:"status_message"`
ProtoMajor int `yaml:"proto_major"`
ProtoMinor int `yaml:"proto_minor"`
Binary string `yaml:"binary,omitempty"`
Timestamp time.Time `yaml:"timestamp"`
HeaderLineLengths HeaderLineLengths `yaml:"header_line_lengths,omitempty"`
Marks *HTTPAnswerMarks `yaml:"marks,omitempty"`
}
var raw httpRespYAML
if err := node.Decode(&raw); err != nil {
return err
}
h.StatusCode = raw.StatusCode
h.Header = raw.Header
h.BodySkipped = raw.BodySkipped
h.BodySize = raw.BodySize
h.StatusMessage = raw.StatusMessage
h.ProtoMajor = raw.ProtoMajor
h.ProtoMinor = raw.ProtoMinor
h.Binary = raw.Binary
h.Timestamp = raw.Timestamp
h.HeaderLineLengths = raw.HeaderLineLengths
h.Marks = raw.Marks
// Decode body — may produce a plain string body, streaming chunks, or both.
// StreamBody is populated only when the YAML body field was a sequence node.
body, streamChunks, err := decodeStreamBody(raw.Body, raw.Header, raw.Timestamp)
if err != nil {
return err
}
h.Body = body
h.StreamBody = streamChunks
return nil
}
// ── Stream body kind detection ─────────────────────────────────────────────────
// streamBodyKind identifies what format a streaming response body uses.
// This is used to decide how to reconstruct the flat Body string from
// structured StreamBody chunks (see streamChunksToLegacyBody).
type streamBodyKind string
const (
streamBodyKindUnknown streamBodyKind = "" // Unknown / non-streaming
streamBodyKindSSE streamBodyKind = "sse" // Server-Sent Events (text/event-stream)
streamBodyKindRaw streamBodyKind = "raw" // All other streaming formats (NDJSON, plain text, binary)
)
// streamChunksAreUTF8 reports whether every key and value in an already-parsed
// chunk list is valid UTF-8. The chunked representation puts arbitrary body
// bytes into YAML scalars, and an SSE field NAME additionally goes through
// strings.ToLower during parsing, which rewrites every invalid byte to U+FFFD
// — so a body that fails this check must fall back to the flat scalar form
// rather than be chunked.
func streamChunksAreUTF8(chunks []HTTPStreamChunk) bool {
for _, c := range chunks {
for _, f := range c.Data {
if !utf8.ValidString(f.Key) || !utf8.ValidString(f.Value) {
return false
}
}
}
return true
}
// streamChunksForYAML returns the chunks to encode if this response has a
// streaming body, and (nil, false) for a non-streaming response or one whose
// bytes cannot survive the chunked representation.
func streamChunksForYAML(resp HTTPResp) ([]HTTPStreamChunk, bool) {
if len(resp.StreamBody) > 0 {
if !streamChunksAreUTF8(resp.StreamBody) {
return nil, false
}
return cloneStreamChunks(resp.StreamBody), true
}
if resp.Body != "" && !resp.bodyOnWire {
// A fault's body is the bytes the upstream sent before it
// (bodyOnWire): kept as they came, never split.
//
// A non-UTF-8 body is not a text stream, and must never reach the
// chunkers below.
//
// Two things go wrong if it does. The chunk encoder emits every key and
// value through stringNode, which stamps an explicit `!!str` tag, and
// yaml.v3 refuses that outright: "cannot marshal invalid UTF-8 data as
// !!str". That error is returned all the way out of InsertMock, and the
// recorder treats a failed insert as fatal — one gzip response body
// ended a 46-hour production recording of prod/api-server. Second, even
// when it does not abort, parseSSEBodyToChunks lower-cases the SSE field
// name, and strings.ToLower rewrites every invalid byte to U+FFFD, so
// the mock is silently corrupted and replays the wrong bytes.
//
// Falling through to the flat scalar body costs nothing and fixes both:
// yaml.v3 auto-promotes an untagged non-UTF-8 scalar to `!!binary` with
// base64, which is exactly how non-streaming bodies (application/json
// and friends) have always been stored. No schema change, and the
// decoder already handles it — decodeStreamBody dispatches on node kind
// and yaml.v3 base64-decodes `!!binary` scalars back into the string.
if !utf8.ValidString(resp.Body) {
return nil, false
}
// A body is kept as frames only when the response is a stream by
// keploy's one rule (HTTPStreamOf) on its Content-Type, never on how
// its length was framed: a text/plain body, chunked or with no
// Content-Length (Go's JSON written with no Content-Type), is one
// document, kept whole, and replayed as one (recordedFramesOf).
format := recordedFramesOf(getHeaderValueCaseInsensitiveModel(resp.Header, "Content-Type"))
if format == HTTPStreamSSE && !sseEndsItsLastEvent(resp.Body) {
// An SSE stream that ended inside its last event, before the
// blank line that ends it, is kept as its bytes. Its frames would
// read back with that event ended (streamChunksToLegacyBody), and
// replay would send an event its client never dispatched at
// record: a close-delimited body the upstream closed mid-event,
// or a server that never ended its last event. Kept so, replay
// sends the bytes as recorded, and the client drops the
// unfinished event as it did.
return nil, false
}
if chunks := splitStreamBody(format, resp.Body, resp.Timestamp); len(chunks) > 0 {
return chunks, true
}
}
return nil, false
}
// StreamChunksOf splits the body of a response that is a stream (HTTPStreamOf)
// into frames by the format of its Content-Type, as a recording of the stream
// keeps them (StreamBody): SSE frames by blank lines, NDJSON and text by
// lines, an octet stream whole. It is nil for a body that does not split
// (multipart, a Content-Type of no stream format), and for one that is empty
// or not UTF-8, which a recording keeps whole. A replay splits what a stream
// delivered with it, so frame i of a recording and frame i of an answer are
// read alike. It does not say whether a response is a stream.
func StreamChunksOf(header map[string]string, body string) []HTTPStreamChunk {
if !utf8.ValidString(body) {
return nil
}
return SplitStreamBody(httpStreamFormatOf(getHeaderValueCaseInsensitiveModel(header, "Content-Type")), body)
}
// SplitStreamBody splits a stream's body into the frames of a format, as
// StreamChunksOf does by a Content-Type's, whatever its bytes: it is for
// comparing a stream, which reads the answer's bytes as they come, not for
// keeping one (a recording keeps a body that is not UTF-8 whole).
func SplitStreamBody(format HTTPStreamFormat, body string) []HTTPStreamChunk {
if body == "" {
return nil
}
return splitStreamBody(format, body, time.Time{})
}
// detectStreamBodyKind is the format a recording's frames were kept in, by
// its Content-Type: how a flat body is reconstructed from them. It says
// nothing of whether the response is a stream (HTTPStreamOf). It reads frames
// as every keploy kept them: as SSE events under a Content-Type that names
// text/event-stream, and as lines under the rest (recordedFramesOf).
func detectStreamBodyKind(headers map[string]string) streamBodyKind {
contentType := strings.ToLower(getHeaderValueCaseInsensitiveModel(headers, "Content-Type"))
if strings.Contains(contentType, "text/event-stream") {
return streamBodyKindSSE
}
if strings.Contains(contentType, "application/x-ndjson") ||
strings.Contains(contentType, "application/ndjson") ||
strings.Contains(contentType, "text/plain") ||
strings.Contains(contentType, "application/octet-stream") {
return streamBodyKindRaw
}
return streamBodyKindUnknown
}
// ── YAML body decode/encode ────────────────────────────────────────────────────
// decodeStreamBody parses the YAML body node and returns:
// - body: the flat string representation (used by the HTTP matcher).
// - streamChunks: the structured chunk list (used for chunk-by-chunk comparison
// and for re-encoding back to YAML). Nil for non-streaming responses.
//
// If the YAML node is a sequence, it's a structured streaming body.
// If it's a scalar, it's a plain string body.
func decodeStreamBody(node yamlLib.Node, headers map[string]string, respTimestamp time.Time) (string, []HTTPStreamChunk, error) {
if node.Kind == 0 {
// Empty body node — nothing to decode.
return "", nil, nil
}
if node.Kind == yamlLib.SequenceNode {
// Structured streaming body: decode the sequence into chunks,
// then reconstruct the legacy flat Body string so that the rest
// of the code that reads h.Body still works correctly.
chunks, err := decodeStreamChunks(node)
if err != nil {
return "", nil, err
}
flatBody := streamChunksToLegacyBody(chunks, detectStreamBodyKind(headers))
return flatBody, chunks, nil
}
// Plain string body (non-streaming).
var body string
if err := node.Decode(&body); err != nil {
return "", nil, err
}
return body, nil, nil
}
// decodeStreamChunks parses a YAML sequence node into a list of HTTPStreamChunk.
// Each element of the sequence is a YAML mapping with "ts" and "data" keys.
func decodeStreamChunks(seqNode yamlLib.Node) ([]HTTPStreamChunk, error) {
chunks := make([]HTTPStreamChunk, 0, len(seqNode.Content))
for _, chunkNode := range seqNode.Content {
if chunkNode == nil || chunkNode.Kind != yamlLib.MappingNode {
continue
}
chunk := HTTPStreamChunk{}
// Iterate over the mapping key-value pairs (YAML mapping nodes store
// alternating key nodes and value nodes in Content).
for i := 0; i+1 < len(chunkNode.Content); i += 2 {
keyNode := chunkNode.Content[i]
valueNode := chunkNode.Content[i+1]
if keyNode == nil || valueNode == nil {
continue
}
switch strings.ToLower(strings.TrimSpace(keyNode.Value)) {
case "ts":
if ts := parseStreamChunkTime(valueNode.Value); !ts.IsZero() {
chunk.TS = ts
}
case "data":
fields, err := decodeStreamDataFields(valueNode)
if err != nil {
return nil, err
}
chunk.Data = fields
}
}
chunks = append(chunks, chunk)
}
return chunks, nil
}
// decodeStreamDataFields parses the "data" value node of a chunk into key-value fields.
//
// If the data node is a YAML mapping, each entry becomes one HTTPStreamDataField.
// If it's a scalar (e.g. a plain line of text), it's wrapped in a single
// field with key="raw".
func decodeStreamDataFields(dataNode *yamlLib.Node) ([]HTTPStreamDataField, error) {
if dataNode == nil {
return nil, nil
}
if dataNode.Kind != yamlLib.MappingNode {
// Scalar or sequence — wrap as a raw field.
return []HTTPStreamDataField{{
Key: "raw",
Value: yamlNodeToString(dataNode),
}}, nil
}
fields := make([]HTTPStreamDataField, 0, len(dataNode.Content)/2)
for i := 0; i+1 < len(dataNode.Content); i += 2 {
keyNode := dataNode.Content[i]
valueNode := dataNode.Content[i+1]
if keyNode == nil || valueNode == nil {
continue
}
// Keys decode the same way values do. stringNode can emit a key as
// !!binary, and yaml.v3 leaves the BASE64 TEXT in node.Value for those
// — reading Value raw would substitute the base64 for the real field
// name. TrimSpace stays off the binary branch: for decoded bytes it is
// not cosmetic, it would silently eat payload.
key := decodeYAMLScalarKey(keyNode)
if keyNode.Tag != "!!binary" {
key = strings.TrimSpace(key)
}
fields = append(fields, HTTPStreamDataField{
Key: key,
Value: yamlNodeToString(valueNode),
})
}
return fields, nil
}
// encodeStreamChunksNode converts a list of HTTPStreamChunk into a YAML sequence node
// for serialization. Each chunk becomes a YAML mapping with "ts" and "data" keys.
func encodeStreamChunksNode(chunks []HTTPStreamChunk) *yamlLib.Node {
root := &yamlLib.Node{Kind: yamlLib.SequenceNode, Tag: "!!seq"}
for _, chunk := range chunks {
chunkNode := &yamlLib.Node{Kind: yamlLib.MappingNode, Tag: "!!map"}
if !chunk.TS.IsZero() {
chunkNode.Content = append(chunkNode.Content,
stringNode("ts"),
stringNode(chunk.TS.UTC().Format(time.RFC3339Nano)),
)
}
// Encode each key-value field under the "data" mapping.
dataNode := &yamlLib.Node{Kind: yamlLib.MappingNode, Tag: "!!map"}
for _, field := range chunk.Data {
dataNode.Content = append(dataNode.Content, stringNode(field.Key), stringNode(field.Value))
}
chunkNode.Content = append(chunkNode.Content, stringNode("data"), dataNode)
root.Content = append(root.Content, chunkNode)
}
return root
}
// ── Flat body reconstruction ──────────────────────────────────────────────────
// streamChunksToLegacyBody reconstructs the flat Body string from structured chunks.
//
// This is used when loading a test case from YAML so that the rest of the code
// (which reads HTTPResp.Body as a single string) still works as expected,
// even though the YAML file now stores structured streaming data.
//
// For SSE bodies, each chunk becomes an SSE frame (fields joined by \n,
// each frame ended by a blank line, the last one too, matching the SSE wire
// format).
//
// For all other streaming formats (NDJSON, plain text, binary), each chunk's
// "raw" field value becomes one line.
//
// The body is the one form of the stream that JSON carries (StreamBody is not
// JSON): a test case read from YAML and sent on as JSON (an upload to Keploy
// Cloud, which keeps and serves it so) reaches the replay with this body
// alone, which reads it as it reads the app's stream to compare the two, and
// a recording of it splits it again (MarshalYAML). So the body is written as
// a wire that splits back into the frames: an SSE value of more than one line
// (an event's data lines, which a read joins) is one line of its field each,
// a value that starts with a space gets the one space a read takes off, and a
// last line that is blank ends in its newline.
func streamChunksToLegacyBody(chunks []HTTPStreamChunk, kind streamBodyKind) string {
if len(chunks) == 0 {
return ""
}
switch kind {
case streamBodyKindSSE:
// Reconstruct SSE wire format: each frame is a block of "key:value" lines
// separated by \n, and each frame ends in a blank line.
frames := make([]string, 0, len(chunks))
for _, chunk := range chunks {
if frame := sseFrame(chunk); frame != "" {
frames = append(frames, frame)
}
}
body := strings.Join(frames, "\n\n")
if len(frames) > 0 {
// The last event ends in its blank line as every other does: a
// client dispatches an event only at the blank line after it, and
// drops the one a stream ends inside (the WHATWG event-stream
// parse), so a body that stopped at its last field would lose its
// last event on replay (an LLM stream's "data: [DONE]"). A frame
// never ends in a blank line of its own (every line carries its
// field's name or the comment's ':'), so this is the one blank
// line. A recording keeps a stream that ended inside its last
// event as its bytes, never as frames (streamChunksForYAML), so
// frames are of a stream that ended its last event. Frames an
// earlier keploy kept of one that did not carry no word of it,
// and read back with that event ended.
body += "\n\n"
}
return body
default:
// NDJSON, plain text, binary — each chunk contributes one line via its "raw" field.
lines := make([]string, 0, len(chunks))
for _, chunk := range chunks {
if raw, ok := chunk.ValueByKey("raw"); ok {
lines = append(lines, raw)
continue
}
// Fallback: use the first field value if no "raw" key exists.
if len(chunk.Data) > 0 {
lines = append(lines, chunk.Data[0].Value)
}
}
body := strings.Join(lines, "\n")
if len(lines) > 0 && lines[len(lines)-1] == "" {
// A last line that is blank (a text stream's) ends in the
// newline a read of lines reads it by: without it the read
// ends one line short.
body += "\n"
}
return body
}
}
// sseEndsItsLastEvent reports whether an SSE body ends in a blank line: two
// line ends in a row, each "\r\n", "\n" or "\r" as the WHATWG event-stream
// parse reads them ("\r\n\r\n", "\n\r", "\r\r" too). A client dispatches an
// event at the blank line that ends it, and drops the one a stream ends
// inside. Only the line ends the body ends in are read: no "\r\n" spans the
// last other byte and them.
func sseEndsItsLastEvent(body string) bool {
ends := body[len(strings.TrimRight(body, "\r\n")):]
return len(normalizeModelLineEndings(ends)) >= 2
}
// sseFrame is chunk as an SSE event's lines, joined by "\n", without the blank
// line that ends the event; "" for a chunk of no line, which the body leaves
// out.
func sseFrame(chunk HTTPStreamChunk) string {
lines := make([]string, 0, len(chunk.Data))
for _, field := range chunk.Data {
key := strings.TrimSpace(field.Key)
if strings.EqualFold(key, "comment") {
// SSE comments start with ':' and have no key prefix.
lines = appendSSEFieldLines(lines, ":", field.Value, false)
continue
}
if key == "" {
continue
}
// strings.ToLower rewrites every invalid byte to U+FFFD, so it
// must not touch a field name that is not valid UTF-8. This is
// reachable precisely because MarshalYAML now routes non-UTF-8
// chunks through here as its flat-body fallback — lower-casing
// them would reintroduce the silent corruption the fallback
// exists to avoid. SSE field names are ASCII in practice, so
// skipping the fold for the pathological case costs nothing.
name := key
if utf8.ValidString(name) {
name = strings.ToLower(name)
}
lines = appendSSEFieldLines(lines, name+":", field.Value, true)
}
return strings.Join(lines, "\n")
}
// appendSSEFieldLines appends an SSE field's lines to lines: prefix (its
// name and ':', or ':' alone for a comment) and each line of value, so a
// value of more than one line never reads back as a field of its own, or as
// the end of its frame (an empty line). A field the parse takes a space off
// (spaced) keeps a value that starts with one: it is written after one more.
// Only data lines join when read: a value of more than one line in a field of
// another name, which no SSE wire carries, reads back as a field per line.
func appendSSEFieldLines(lines []string, prefix, value string, spaced bool) []string {
for _, line := range strings.Split(value, "\n") {
if spaced && strings.HasPrefix(line, " ") {
line = " " + line
}
lines = append(lines, prefix+line)
}
return lines
}
// ── Body → Chunk parsers ───────────────────────────────────────────────────────
// These are used during recording to convert a raw HTTP response body string into
// structured HTTPStreamChunk slices for storage in the test case YAML.
// parseSSEBodyToChunks splits a raw SSE body string into HTTPStreamChunk slices.
//
// SSE frames are delimited by blank lines (\n\n). Within each frame, each line
// is a "key: value" pair. Consecutive "data" lines are merged per the SSE spec.
// Comment lines (starting with ':') are stored with key="comment".
func parseSSEBodyToChunks(body string, ts time.Time) []HTTPStreamChunk {
body = normalizeModelLineEndings(body)
frames := splitModelSSEFrames(body)
chunks := make([]HTTPStreamChunk, 0, len(frames))
for _, frame := range frames {
lines := strings.Split(frame, "\n")
fields := make([]HTTPStreamDataField, 0, len(lines))
for _, line := range lines {
line = strings.TrimRight(line, "\r")
if strings.TrimSpace(line) == "" {
continue
}
// SSE comment line: starts with ':'
if strings.HasPrefix(line, ":") {
fields = append(fields, HTTPStreamDataField{
Key: "comment",
Value: strings.TrimPrefix(line, ":"),
})
continue
}
// Regular SSE field: "key: value" or just "key" (value is empty string).
key, value, ok := strings.Cut(line, ":")
if !ok {
key = strings.TrimSpace(line)
value = ""
} else if strings.HasPrefix(value, " ") {
// Per SSE spec, one optional leading space after ':' is stripped.
value = value[1:]
}
if key == "" {
continue
}
normalizedKey := strings.ToLower(strings.TrimSpace(key))
// Per SSE spec, consecutive "data" lines in the same frame
// are concatenated with a '\n' separator.
if normalizedKey == "data" && len(fields) > 0 && strings.EqualFold(fields[len(fields)-1].Key, "data") {
fields[len(fields)-1].Value = fields[len(fields)-1].Value + "\n" + value
continue
}
fields = append(fields, HTTPStreamDataField{
Key: normalizedKey,
Value: value,
})
}
if len(fields) == 0 {
continue
}
chunks = append(chunks, HTTPStreamChunk{
TS: ts,
Data: fields,
})
}
return chunks
}
// parseRawBodyToChunks splits a raw body string into HTTPStreamChunk slices
// where each line becomes one chunk with a single "raw" field.
//
// Used for NDJSON, plain text, and binary streaming formats.
// If ignoreEmpty is true, blank lines are skipped (useful for NDJSON which
// can have empty separator lines between JSON objects).
func parseRawBodyToChunks(body string, ts time.Time, ignoreEmpty bool) []HTTPStreamChunk {
body = normalizeModelLineEndings(body)
if body == "" {
return nil
}
lines := make([]string, 0)
scanner := bufio.NewScanner(strings.NewReader(body))
scanner.Buffer(make([]byte, 0, 64*1024), maxModelStreamTokenSize)
for scanner.Scan() {
line := scanner.Text()
if ignoreEmpty && strings.TrimSpace(line) == "" {
continue
}
lines = append(lines, line)
}
if err := scanner.Err(); err != nil {
// Scanner error — fall back to treating the whole body as one chunk.
lines = nil
}
if len(lines) == 0 {
// No lines found — treat the entire body as a single raw chunk.
lines = append(lines, body)
}
chunks := make([]HTTPStreamChunk, 0, len(lines))
for _, line := range lines {
chunks = append(chunks, HTTPStreamChunk{
TS: ts,
Data: []HTTPStreamDataField{{
Key: "raw",
Value: line,
}},
})
}
return chunks
}
// ── Helpers ───────────────────────────────────────────────────────────────────
// splitModelSSEFrames splits an SSE body string into individual frames
// by splitting on blank lines (\n\n). Empty frames are discarded.
func splitModelSSEFrames(body string) []string {
body = normalizeModelLineEndings(body)
frames := make([]string, 0)
for _, part := range strings.Split(body, "\n\n") {
part = strings.Trim(part, "\n")
if strings.TrimSpace(part) == "" {
continue
}
frames = append(frames, part)
}
return frames
}
// normalizeModelLineEndings converts all line ending variants (\r\n, \r)
// to the canonical Unix form (\n) so that parsing is line-ending agnostic.
func normalizeModelLineEndings(value string) string {
value = strings.ReplaceAll(value, "\r\n", "\n")
value = strings.ReplaceAll(value, "\r", "\n")
return value
}
// getHeaderValueCaseInsensitiveModel looks up a header value by key, ignoring case.
// Returns "" if the key is not found.
func getHeaderValueCaseInsensitiveModel(headers map[string]string, key string) string {
key = strings.ToLower(strings.TrimSpace(key))
for k, v := range headers {
if strings.ToLower(strings.TrimSpace(k)) == key {
return v
}
}
return ""
}
// parseStreamChunkTime tries to parse a timestamp string in RFC3339Nano or RFC3339 format.
// Returns the zero time.Time if the string is empty or cannot be parsed.
func parseStreamChunkTime(raw string) time.Time {
raw = strings.TrimSpace(raw)
if raw == "" {
return time.Time{}
}
// Try nanosecond precision first, then second precision.
if parsed, err := time.Parse(time.RFC3339Nano, raw); err == nil {
return parsed
}
if parsed, err := time.Parse(time.RFC3339, raw); err == nil {
return parsed
}
return time.Time{}
}
// yamlNodeToString converts a YAML node to its string representation.
// Scalar nodes return their value directly.
// For complex nodes, the decoded value is re-marshalled to compact JSON.
func yamlNodeToString(node *yamlLib.Node) string {
if node == nil {
return ""
}
if node.Kind == yamlLib.ScalarNode {
// yaml.v3 leaves the BASE64 TEXT in node.Value for a !!binary scalar
// rather than the decoded bytes, so returning it raw would hand replay
// the base64 string and silently mismatch. Decode it here, mirroring
// stringNode. Whitespace strip: yaml.v3 folds long base64 scalars.
if node.Tag == "!!binary" {
if raw, err := base64.StdEncoding.DecodeString(stripBase64Whitespace(node.Value)); err == nil {
return string(raw)
}
}
return node.Value
}
// Complex node (mapping or sequence) — decode then re-encode as JSON.
var decoded interface{}
if err := node.Decode(&decoded); err != nil {
return strings.TrimSpace(node.Value)
}
b, err := json.Marshal(decoded)
if err != nil {
return strings.TrimSpace(node.Value)
}
return string(b)
}
// stringNode creates a YAML scalar string node with the given value.
// Used as a helper when constructing YAML trees programmatically.
func stringNode(value string) *yamlLib.Node {
// Defence in depth behind streamChunksForYAML's guard. yaml.v3 hard-fails
// on invalid UTF-8 under an EXPLICIT tag ("cannot marshal invalid UTF-8
// data as !!str") but silently base64s it under `!!binary`, so stamping
// !!str on arbitrary bytes is what turns one bad chunk into a failed mock
// insert — and a failed insert used to end the whole recording.
// yamlNodeToString decodes this back, so the round-trip is lossless.
if !utf8.ValidString(value) {
return &yamlLib.Node{
Kind: yamlLib.ScalarNode,
Tag: "!!binary",
Value: base64.StdEncoding.EncodeToString([]byte(value)),
}
}
n := &yamlLib.Node{
Kind: yamlLib.ScalarNode,
Tag: "!!str",
Value: value,
}
if YAMLBlockScalarUnsafe(value) {
n.Style = yamlLib.DoubleQuotedStyle // not a block scalar that reads back as another string
}
return n
}
// cloneStreamChunks returns a deep copy of a slice of HTTPStreamChunk.
// This ensures that mutations to one copy do not affect the original.
func cloneStreamChunks(in []HTTPStreamChunk) []HTTPStreamChunk {
out := make([]HTTPStreamChunk, len(in))
for i := range in {
out[i].TS = in[i].TS
if len(in[i].Data) == 0 {
continue
}
out[i].Data = make([]HTTPStreamDataField, len(in[i].Data))
copy(out[i].Data, in[i].Data)
}
return out
}