feat: streaming responses — server→client push over ZAP (bidirectional)

Adds multi-frame streamed responses so SSE / MCP notifications / progress / chunked bodies ride ZAP as real server→client push, not a buffered blob. A streamed response is a head frame (status+headers) + N data frames (one per handler flush) + an end frame — the analogue of HTTP/2 HEADERS+DATA+END_STREAM. - wire.go: FrameResponseHead/FrameData/FrameEnd + Marshal/Unmarshal + FrameTypeOf. - server.go: when ctx.Response.IsBodyStream(), stream via BodyWriteTo→chunkWriter (each flush = one FrameData frame on the wire immediately). Non-stream path unchanged (single FrameResponse) — fully backward-compatible. - client.go: Transport.Do dispatches on frame type; a head attaches a streaming body (streamReader) that reads data frames until end, releasing the pooled conn only when the stream drains (or closes it if abandoned mid-stream). Verified: TestStream_SSEOverZAP flushes 3 events one-at-a-time and the client receives each BEFORE the next is released (a buffering regression would deadlock/ time out — the guard proves real incremental push). All existing round-trip + keep-alive tests pass; go test -race clean; gofmt clean.
2026-10-04 03:07:19 +00:00 · 2026-07-02 00:43:58 -07:00
parent d3753d2af7
commit 5deee64a3e
4 changed files with 383 additions and 6 deletions
@@ -21,6 +21,7 @@ package zaphttp
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -103,6 +104,24 @@ func (t *Transport) Do(req *fasthttp.Request, resp *fasthttp.Response) error {
 		conn.Close()
 		return fmt.Errorf("zaphttp: read response: %w", err)
 	}
+	ft, err := FrameTypeOf(respFrame)
+	if err != nil {
+		conn.Close()
+		return fmt.Errorf("zaphttp: response frame: %w", err)
+	}
+
+	// Streamed response: apply the head, then attach a body stream that reads
+	// FrameData chunks until FrameEnd. The conn is NOT pooled until the stream
+	// is fully consumed (or the body is closed) — the reader owns it.
+	if ft == FrameResponseHead {
+		if err := UnmarshalResponseHead(respFrame, resp); err != nil {
+			conn.Close()
+			return fmt.Errorf("zaphttp: unmarshal response head: %w", err)
+		}
+		resp.SetBodyStream(&streamReader{t: t, pc: pc, br: br}, -1)
+		return nil
+	}
+
 	if err := UnmarshalResponse(respFrame, resp); err != nil {
 		conn.Close()
 		return fmt.Errorf("zaphttp: unmarshal response: %w", err)
@@ -115,6 +134,84 @@ func (t *Transport) Do(req *fasthttp.Request, resp *fasthttp.Response) error {
 	return nil
 }

+// streamReader turns the FrameData/FrameEnd sequence following a streamed
+// response head into an io.ReadCloser body: each Read drains the current chunk,
+// pulling the next FrameData frame when empty and returning io.EOF at FrameEnd.
+// fasthttp closes it when the response is released, which returns the conn to
+// the pool (fully drained) or closes it (abandoned mid-stream).
+type streamReader struct {
+	t    *Transport
+	pc   *pooledConn
+	br   *bufio.Reader
+	rest []byte // undrained tail of the current chunk
+	done bool
+	err  error
+}
+
+func (s *streamReader) Read(p []byte) (int, error) {
+	if len(s.rest) > 0 {
+		n := copy(p, s.rest)
+		s.rest = s.rest[n:]
+		return n, nil
+	}
+	if s.done {
+		return 0, io.EOF
+	}
+	if s.err != nil {
+		return 0, s.err
+	}
+	frame, err := readFrame(s.br)
+	if err != nil {
+		s.err = err
+		return 0, err
+	}
+	ft, err := FrameTypeOf(frame)
+	if err != nil {
+		s.err = err
+		return 0, err
+	}
+	switch ft {
+	case FrameData:
+		chunk, err := DataChunkOf(frame)
+		if err != nil {
+			s.err = err
+			return 0, err
+		}
+		n := copy(p, chunk)
+		if n < len(chunk) {
+			s.rest = append(s.rest, chunk[n:]...) // buffer the tail (chunk aliases frame; copy)
+		}
+		return n, nil
+	case FrameEnd:
+		s.done = true
+		s.release()
+		return 0, io.EOF
+	default:
+		s.err = fmt.Errorf("zaphttp: unexpected frame %#x mid-stream", ft)
+		return 0, s.err
+	}
+}
+
+// Close returns the conn to the pool if the stream drained cleanly, else closes
+// it (abandoned mid-stream leaves unread frames on the wire — not reusable).
+func (s *streamReader) Close() error {
+	if s.done {
+		s.release()
+	} else {
+		s.pc.c.Close()
+	}
+	return nil
+}
+
+func (s *streamReader) release() {
+	if s.pc == nil {
+		return
+	}
+	_ = s.pc.c.SetReadDeadline(time.Time{})
+	s.t.releaseConn(s.pc)
+	s.pc = nil
+}
+
 // acquireConn pops an idle conn or dials a new one.
 func (t *Transport) acquireConn() (*pooledConn, error) {
 	t.mu.Lock()
@@ -27,9 +27,9 @@ import (
 // Server is a ZAP-HTTP server. Zero-value is usable; common knobs (Addr,
 // Handler, ReadTimeout, …) mirror fasthttp.Server.
 type Server struct {
-	Addr         string                    // ":9999" if empty
-	Handler      fasthttp.RequestHandler   // required
-	ReadTimeout  time.Duration             // 0 means no timeout
+	Addr         string                  // ":9999" if empty
+	Handler      fasthttp.RequestHandler // required
+	ReadTimeout  time.Duration           // 0 means no timeout
 	WriteTimeout time.Duration
 	IdleTimeout  time.Duration
 	Logger       fasthttp.Logger // passed to each RequestCtx; nil is fine
@@ -135,14 +135,29 @@ func (s *Server) serveConn(conn net.Conn) {
 			s.Handler(ctx)
 		}()

+		if s.WriteTimeout > 0 {
+			_ = conn.SetWriteDeadline(time.Now().Add(s.WriteTimeout))
+		}
+
+		// A streamed body (SSE, MCP notifications, chunked) is delivered as a
+		// head frame + N data frames + an end frame, so the client sees each
+		// chunk as the handler flushes it — real server→client push over ZAP.
+		// A normal body takes the single-frame path (unchanged).
+		if ctx.Response.IsBodyStream() {
+			if err := s.streamResponse(conn, &ctx.Response); err != nil {
+				if !isClosedConn(err) {
+					log.Printf("zaphttp: stream response: %v", err)
+				}
+				return
+			}
+			continue
+		}
+
 		respBytes, err := MarshalResponse(&ctx.Response)
 		if err != nil {
 			log.Printf("zaphttp: marshal response: %v", err)
 			return
 		}
-		if s.WriteTimeout > 0 {
-			_ = conn.SetWriteDeadline(time.Now().Add(s.WriteTimeout))
-		}
 		if err := writeFrame(conn, respBytes); err != nil {
 			if !isClosedConn(err) {
 				log.Printf("zaphttp: write frame: %v", err)
@@ -152,6 +167,51 @@ func (s *Server) serveConn(conn net.Conn) {
 	}
 }

+// streamResponse writes a streamed response: the head (status + headers), then
+// one data frame per flush of the body stream, then an end frame. BodyWriteTo
+// runs the fasthttp body-stream writer against chunkWriter, so each handler
+// flush becomes a frame on the wire immediately (no full buffering).
+func (s *Server) streamResponse(conn net.Conn, resp *fasthttp.Response) error {
+	head, err := MarshalResponseHead(resp)
+	if err != nil {
+		return err
+	}
+	if err := writeFrame(conn, head); err != nil {
+		return err
+	}
+	cw := &chunkWriter{w: conn}
+	if err := resp.BodyWriteTo(cw); err != nil {
+		return err
+	}
+	if cw.err != nil {
+		return cw.err
+	}
+	return writeFrame(conn, MarshalEnd(nil))
+}
+
+// chunkWriter turns each Write (one flush of the body stream) into a FrameData
+// frame. The bytes are copied because fasthttp may reuse the buffer.
+type chunkWriter struct {
+	w   io.Writer
+	err error
+}
+
+func (c *chunkWriter) Write(p []byte) (int, error) {
+	if c.err != nil {
+		return 0, c.err
+	}
+	if len(p) == 0 {
+		return 0, nil
+	}
+	chunk := make([]byte, len(p))
+	copy(chunk, p)
+	if err := writeFrame(c.w, MarshalData(chunk)); err != nil {
+		c.err = err
+		return 0, err
+	}
+	return len(p), nil
+}
+
 func (s *Server) writeError(w io.Writer, status int, msg string) {
 	var resp fasthttp.Response
 	resp.SetStatusCode(status)
@@ -0,0 +1,111 @@
+package zaphttp_test
+
+import (
+	"bufio"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/valyala/fasthttp"
+
+	zaphttp "github.com/zap-proto/http"
+)
+
+// TestStream_SSEOverZAP proves server→client PUSH over ZAP: an SSE handler that
+// flushes N events incrementally is received as a live stream by the client,
+// chunk-by-chunk (each flush a FrameData frame), not one buffered blob. This is
+// the bidirectional-streaming brick that makes MCP notifications / progress /
+// chunked bodies real over the ZAP transport.
+func TestStream_SSEOverZAP(t *testing.T) {
+	const n = 3
+	release := make(chan struct{}, n) // test controls when each event flushes
+
+	addr, stop := listen(t, func(ctx *fasthttp.RequestCtx) {
+		ctx.Response.Header.Set("Content-Type", "text/event-stream")
+		ctx.SetBodyStreamWriter(func(w *bufio.Writer) {
+			for i := 0; i < n; i++ {
+				<-release // block until the test lets this event go
+				_, _ = w.WriteString("data: event-")
+				_, _ = w.WriteString(string(rune('0' + i)))
+				_, _ = w.WriteString("\n\n")
+				_ = w.Flush() // one flush -> one frame on the wire
+			}
+		})
+	})
+	defer stop()
+
+	tr := zaphttp.NewTransport(addr)
+	defer tr.CloseIdleConnections()
+	req := fasthttp.AcquireRequest()
+	resp := fasthttp.AcquireResponse()
+	defer fasthttp.ReleaseRequest(req)
+	defer fasthttp.ReleaseResponse(resp)
+	req.Header.SetMethod("GET")
+	req.SetRequestURI("/events")
+
+	if err := tr.Do(req, resp); err != nil {
+		t.Fatalf("Do: %v", err)
+	}
+	if ct := string(resp.Header.ContentType()); !strings.HasPrefix(ct, "text/event-stream") {
+		t.Fatalf("content-type = %q, want text/event-stream (head frame carried headers)", ct)
+	}
+	if !resp.IsBodyStream() {
+		t.Fatal("response body is not a stream — server sent a buffered response, not a head+chunks stream")
+	}
+
+	// Read the body stream incrementally, releasing one event at a time and
+	// asserting it arrives BEFORE the next is released — i.e. real push, not a
+	// buffered dump at end-of-handler.
+	br := bufio.NewReader(resp.BodyStream())
+	for i := 0; i < n; i++ {
+		release <- struct{}{}
+		line, err := readEvent(t, br)
+		if err != nil {
+			t.Fatalf("event %d read: %v", i, err)
+		}
+		want := "data: event-" + string(rune('0'+i))
+		if !strings.Contains(line, want) {
+			t.Fatalf("event %d = %q, want %q", i, line, want)
+		}
+	}
+}
+
+// readEvent reads until a blank-line-terminated SSE event, with a timeout so a
+// buffering regression (nothing arrives until the handler returns) fails fast.
+func readEvent(t *testing.T, br *bufio.Reader) (string, error) {
+	t.Helper()
+	type res struct {
+		s   string
+		err error
+	}
+	ch := make(chan res, 1)
+	go func() {
+		var b strings.Builder
+		for {
+			line, err := br.ReadString('\n')
+			b.WriteString(line)
+			if err != nil {
+				ch <- res{b.String(), err}
+				return
+			}
+			if line == "\n" && b.Len() > 1 {
+				ch <- res{b.String(), nil}
+				return
+			}
+		}
+	}()
+	select {
+	case r := <-ch:
+		return r.s, r.err
+	case <-time.After(2 * time.Second):
+		return "", errTimeout
+	}
+}
+
+var errTimeout = &timeoutErr{}
+
+type timeoutErr struct{}
+
+func (*timeoutErr) Error() string {
+	return "timed out waiting for streamed event (buffering regression?)"
+}
@@ -88,8 +88,117 @@ import (
 const (
 	FrameRequest  uint16 = 0x01
 	FrameResponse uint16 = 0x02
+	// Streaming response frames. A streamed response is a FrameResponseHead
+	// (status + headers, no body) followed by zero or more FrameData chunks and
+	// a terminating FrameEnd (optional trailers). This is how server→client
+	// push (SSE, MCP notifications, chunked bodies) rides ZAP — the analogue of
+	// HTTP/2 HEADERS + DATA + END_STREAM. The non-streaming FrameResponse path
+	// is unchanged, so existing peers interoperate.
+	FrameResponseHead uint16 = 0x03
+	FrameData         uint16 = 0x04
+	FrameEnd          uint16 = 0x05
 )

+// Data/End frame layout: a single bytes slot (the chunk, or the trailer JSON).
+const (
+	chunkBytes    = 0 // bytes (8)
+	chunkSlotSize = 8
+)
+
+// FrameTypeOf peeks a frame's type from its ZAP header without decoding the
+// body, so a reader can dispatch (response vs streamed head vs data vs end).
+func FrameTypeOf(frame []byte) (uint16, error) {
+	msg, err := zap.Parse(frame)
+	if err != nil {
+		return 0, err
+	}
+	return msg.Flags() >> 8, nil
+}
+
+// MarshalResponseHead serializes a response's status + headers (NO body) as the
+// opening frame of a stream. The body is delivered by subsequent FrameData
+// frames.
+func MarshalResponseHead(resp *fasthttp.Response) ([]byte, error) {
+	status := resp.StatusCode()
+	if status == 0 {
+		status = fasthttp.StatusOK
+	}
+	reason := string(resp.Header.StatusMessage())
+	if reason == "" {
+		reason = fasthttp.StatusMessage(status)
+	}
+	proto := string(resp.Header.Protocol())
+	headers, err := encodeHeaderMap(collectResponseHeaders(&resp.Header))
+	if err != nil {
+		return nil, fmt.Errorf("zaphttp: encode headers: %w", err)
+	}
+
+	b := zap.NewBuilder(respSlotSize + len(headers) + len(reason) + len(proto) + 64)
+	ob := b.StartObject(respSlotSize)
+	ob.SetUint16(respStatus, uint16(status))
+	ob.SetText(respReason, reason)
+	ob.SetText(respProto, orDefault(proto, defaultProto))
+	ob.SetBytes(respHeaders, headers)
+	ob.SetBytes(respBody, nil)
+	ob.SetBytes(respTrailer, nil)
+	ob.FinishAsRoot()
+	return b.FinishWithFlags(FrameResponseHead << 8), nil
+}
+
+// UnmarshalResponseHead applies a streamed response's status + headers to dst.
+// The body is left empty; the caller attaches a streaming body that reads the
+// following FrameData frames.
+func UnmarshalResponseHead(frame []byte, dst *fasthttp.Response) error {
+	msg, err := zap.Parse(frame)
+	if err != nil {
+		return err
+	}
+	if t := msg.Flags() >> 8; t != FrameResponseHead {
+		return fmt.Errorf("zaphttp: frame type %#x, want response-head (%#x)", t, FrameResponseHead)
+	}
+	r := msg.Root()
+	dst.Reset()
+	dst.SetStatusCode(int(r.Uint16(respStatus)))
+	if reason := r.Text(respReason); reason != "" {
+		dst.Header.SetStatusMessage([]byte(reason))
+	}
+	if proto := r.Text(respProto); proto != "" {
+		dst.Header.SetProtocol([]byte(proto))
+	}
+	return applyHeadersToResponse(&dst.Header, r.Bytes(respHeaders))
+}
+
+// MarshalData wraps one body chunk as a FrameData frame.
+func MarshalData(chunk []byte) []byte {
+	b := zap.NewBuilder(chunkSlotSize + len(chunk) + 32)
+	ob := b.StartObject(chunkSlotSize)
+	ob.SetBytes(chunkBytes, chunk)
+	ob.FinishAsRoot()
+	return b.FinishWithFlags(FrameData << 8)
+}
+
+// DataChunkOf extracts the chunk bytes from a FrameData frame. The returned
+// slice aliases the frame buffer; copy it to retain past the frame's lifetime.
+func DataChunkOf(frame []byte) ([]byte, error) {
+	msg, err := zap.Parse(frame)
+	if err != nil {
+		return nil, err
+	}
+	if t := msg.Flags() >> 8; t != FrameData {
+		return nil, fmt.Errorf("zaphttp: frame type %#x, want data (%#x)", t, FrameData)
+	}
+	return msg.Root().Bytes(chunkBytes), nil
+}
+
+// MarshalEnd terminates a stream, carrying optional encoded trailers.
+func MarshalEnd(trailer []byte) []byte {
+	b := zap.NewBuilder(chunkSlotSize + len(trailer) + 32)
+	ob := b.StartObject(chunkSlotSize)
+	ob.SetBytes(chunkBytes, trailer)
+	ob.FinishAsRoot()
+	return b.FinishWithFlags(FrameEnd << 8)
+}
+
 // Request frame field offsets within the object's fixed payload. Every
 // text/bytes slot is 8 bytes (relOffset+length); see "Offset discipline"
 // in the package doc.