proxy: replace streamBuffer with chunkedResponseWriter for capped response storage

cmd/proxy/backend_routing_test.go

@@ -44,10 +44,11 @@ func (r *routingRecorder) snapshot() []string {
 
 // newRoutingStack builds a testStack whose Charon mux is wrapped in a
 // routingRecorder, so tests can inspect which endpoints were called.
-func newRoutingStack(t *testing.T) (*testStack, *routingRecorder) {
+func newRoutingStack(t *testing.T, opts ...stackOption) (*testStack, *routingRecorder) {
 	t.Helper()
 	rec := &routingRecorder{}
-	s := newTestStack(t, withCharonMiddleware(rec.middleware()))
+	allOpts := append([]stackOption{withCharonMiddleware(rec.middleware())}, opts...)
+	s := newTestStack(t, allOpts...)
 	return s, rec
 }
 
@@ -222,3 +223,80 @@ func TestProxyStreamedStoreTrueNoChainFetches(t *testing.T) {
 	assert.GreaterOrEqual(t, hitsContaining(hits, "POST /staging"), 1, "store:true must commit via POST /staging")
 	assert.Equal(t, 0, hitsContaining(hits, "GET /chain/"), "store:true first turn has no prev to fetch via GET /chain")
 }
+
+// TestBufferedProxySingleChunk verifies that a small buffered response produces
+// exactly 1 chunk PUT and 1 complete (default 1 MiB cap is not exceeded).
+func TestBufferedProxySingleChunk(t *testing.T) {
+	s, rec := newRoutingStack(t)
+
+	resp := doRequest(t, s.proxyURL, "POST", "/responses", map[string]interface{}{
+		"model": "test",
+		"input": "hello",
+	})
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	hits := rec.snapshot()
+	assert.Equal(t, 1, hitsContaining(hits, "/chunks/"), "small response must produce exactly 1 AppendChunk call")
+	assert.Equal(t, 1, hitsContaining(hits, "/complete"), "exactly 1 Complete call")
+	assert.Equal(t, 0, hitsContaining(hits, "/abort"), "no abort on success")
+}
+
+// TestStreamedProxySingleChunk verifies that a small streamed response produces
+// exactly 1 chunk PUT and 1 complete.
+func TestStreamedProxySingleChunk(t *testing.T) {
+	s, rec := newRoutingStack(t)
+
+	req, _ := http.NewRequestWithContext(context.Background(), "POST", s.proxyURL+"/responses",
+		strings.NewReader(`{"model":"test","input":"hello","stream":true}`))
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	_ = readSSE(t, resp)
+
+	hits := rec.snapshot()
+	assert.Equal(t, 1, hitsContaining(hits, "/chunks/"), "small response must produce exactly 1 AppendChunk call")
+	assert.Equal(t, 1, hitsContaining(hits, "/complete"), "exactly 1 Complete call")
+	assert.Equal(t, 0, hitsContaining(hits, "/abort"), "no abort on success")
+}
+
+// TestBufferedProxyMultipleChunks verifies that a tiny chunk cap splits the
+// buffered response blob into multiple chunk PUTs.
+func TestBufferedProxyMultipleChunks(t *testing.T) {
+	s, rec := newRoutingStack(t, withMaxChunkBytes(64))
+
+	resp := doRequest(t, s.proxyURL, "POST", "/responses", map[string]interface{}{
+		"model": "test",
+		"input": "hello world, this request should produce a response blob that exceeds 64 bytes after marshaling",
+	})
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	hits := rec.snapshot()
+	assert.GreaterOrEqual(t, hitsContaining(hits, "/chunks/"), 2, "tiny cap must produce ≥2 AppendChunk calls")
+	assert.Equal(t, 1, hitsContaining(hits, "/complete"), "exactly 1 Complete call")
+	assert.Equal(t, 0, hitsContaining(hits, "/abort"), "no abort on success")
+}
+
+// TestBufferedProxyStoreFalseNoStagingCalls pins that a buffered store:false
+// turn issues no Charon staging calls at all — neither AppendChunk, Complete,
+// nor Abort. The empty-response abort path itself is exercised by the unit
+// tests in chunk_test.go (TestWriterZeroBytesAborts).
+func TestBufferedProxyStoreFalseNoStagingCalls(t *testing.T) {
+	s, rec := newRoutingStack(t)
+
+	resp := doRequest(t, s.proxyURL, "POST", "/responses", map[string]interface{}{
+		"model": "test",
+		"input": "hello",
+		"store": false,
+	})
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	hits := rec.snapshot()
+	assert.Equal(t, 0, hitsContaining(hits, "/chunks/"), "store:false must not call AppendChunk")
+	assert.Equal(t, 0, hitsContaining(hits, "/complete"), "store:false must not call Complete")
+	assert.Equal(t, 0, hitsContaining(hits, "/abort"), "store:false must not call Abort")
+}

cmd/proxy/chunk.go

@@ -0,0 +1,200 @@
+package main
+
+import (
+	"context"
+	"sync"
+
+	"github.com/elevran/charon/pkg/charon"
+)
+
+// chunkedResponseWriter buffers response payload up to maxChunkBytes and
+// flushes each chunk to Charon via PUT /staging/{sid}/chunks/{k}. The
+// terminal Close() issues PUT /staging/{sid}/complete with the running
+// total, OR — when total is 0 — PUT /staging/{sid}/abort to drop the
+// staging record without committing an empty turn.
+//
+// Chunks are raw bytes; Charon concatenates them in order to produce the
+// final stored blob. maxChunkBytes controls flush frequency only — there is
+// no JSON framing applied by this type.
+//
+// Concurrency: Add and Close may be called from separate goroutines.
+// The mutex serialises all state mutations.
+type chunkedResponseWriter struct {
+	ctx           context.Context
+	backend       charon.Backend
+	stagingID     string
+	responseID    string
+	tenantKey     string
+	maxChunkBytes int64
+
+	mu        sync.Mutex
+	pending   []byte // bytes awaiting next flush
+	k         uint32 // next chunk index
+	total     uint64 // running byte sum of flushed chunk bodies
+	closed    bool
+	closedCh  chan struct{}
+	closedErr error
+}
+
+// newChunkedResponseWriter constructs a writer targeting the given staging record.
+func newChunkedResponseWriter(
+	ctx context.Context, b charon.Backend,
+	stagingID, responseID, tenantKey string, maxChunkBytes int64,
+) *chunkedResponseWriter {
+	return &chunkedResponseWriter{
+		ctx:           ctx,
+		backend:       b,
+		stagingID:     stagingID,
+		responseID:    responseID,
+		tenantKey:     tenantKey,
+		maxChunkBytes: maxChunkBytes,
+		closedCh:      make(chan struct{}),
+	}
+}
+
+// Add appends p to the pending buffer. If appending p would push the pending
+// byte count past maxChunkBytes (and pending is non-empty), the existing
+// pending is flushed first as chunk k, then p starts a fresh chunk.
+//
+// Empty or nil p is a no-op. When maxChunkBytes <= 0, all bytes accumulate
+// without flushing until Close. When len(p) alone exceeds maxChunkBytes, p is
+// split across as many maxChunkBytes-sized flushes as needed.
+func (w *chunkedResponseWriter) Add(p []byte) error {
+	if len(p) == 0 {
+		return nil
+	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if w.closed {
+		return w.closedErr
+	}
+
+	if w.maxChunkBytes <= 0 {
+		w.pending = append(w.pending, p...)
+		return nil
+	}
+
+	for len(p) > 0 {
+		room := w.maxChunkBytes - int64(len(w.pending))
+		if room <= 0 {
+			if err := w.flush(); err != nil {
+				return err
+			}
+			room = w.maxChunkBytes
+		}
+		take := int64(len(p))
+		if take > room {
+			take = room
+		}
+		w.pending = append(w.pending, p[:take]...)
+		p = p[take:]
+		if int64(len(w.pending)) >= w.maxChunkBytes {
+			if err := w.flush(); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// Close flushes any remaining bytes and signals terminal completion:
+//   - total > 0: PUT /staging/{sid}/complete
+//   - total == 0: PUT /staging/{sid}/abort
+//
+// Subsequent calls block on closedCh and return the same error. Idempotent.
+func (w *chunkedResponseWriter) Close() error {
+	w.mu.Lock()
+
+	if w.closed {
+		ch := w.closedCh
+		w.mu.Unlock()
+		<-ch
+		// Read closedErr only after <-ch returns: the original Close writes
+		// it under mu *after* its unlocked I/O, so reading it before the
+		// channel close races with that write.
+		w.mu.Lock()
+		err := w.closedErr
+		w.mu.Unlock()
+		return err
+	}
+	w.closed = true
+	defer close(w.closedCh)
+
+	// Flush remaining bytes. flush() marks the writer closed on AppendChunk
+	// failure (so callers see the same error from any later Add/Close), but
+	// since we just set closed=true above, the gated !w.closed branch in
+	// flush will not re-close closedCh — the defer here handles that.
+	if len(w.pending) > 0 {
+		if err := w.flush(); err != nil {
+			w.mu.Unlock()
+			return err
+		}
+	}
+
+	total := w.total
+	w.mu.Unlock()
+
+	if total == 0 {
+		return w.backend.Abort(w.ctx, w.stagingID)
+	}
+	_, err := w.backend.Complete(w.ctx, w.stagingID, w.responseID, w.tenantKey, uint32(total))
+	w.mu.Lock()
+	w.closedErr = err
+	w.mu.Unlock()
+	return err
+}
+
+// Abort issues PUT /staging/{sid}/abort unconditionally (failure path).
+// Subsequent Add and Close calls return the abort error.
+func (w *chunkedResponseWriter) Abort() error {
+	w.mu.Lock()
+
+	if w.closed {
+		ch := w.closedCh
+		w.mu.Unlock()
+		<-ch
+		return w.closedErr
+	}
+	w.closed = true
+	defer close(w.closedCh)
+	w.mu.Unlock()
+
+	err := w.backend.Abort(w.ctx, w.stagingID)
+	w.mu.Lock()
+	w.closedErr = err
+	w.mu.Unlock()
+	return err
+}
+
+// flush sends pending bytes as chunk k.
+// Must be called with w.mu held. Releases and reacquires the lock around
+// the network call.
+func (w *chunkedResponseWriter) flush() error {
+	body := make([]byte, len(w.pending))
+	copy(body, w.pending)
+	k := w.k
+
+	// Release lock during I/O.
+	w.mu.Unlock()
+	err := w.backend.AppendChunk(w.ctx, w.stagingID, k, w.responseID, body)
+	w.mu.Lock()
+
+	if err != nil {
+		// Flush failure is terminal: mark the writer closed and signal any
+		// concurrent Close/Abort callers via closedCh. The !w.closed gate
+		// avoids double-closing the channel when this flush is invoked from
+		// inside Close (which sets closed=true and closes via defer).
+		if !w.closed {
+			w.closed = true
+			close(w.closedCh)
+		}
+		w.closedErr = err
+		return err
+	}
+	w.k++
+	w.total += uint64(len(body))
+	w.pending = w.pending[:0]
+	return nil
+}