diff --git a/cmd/proxy/backend_routing_test.go b/cmd/proxy/backend_routing_test.go index 39aec10..fa0019c 100644 --- a/cmd/proxy/backend_routing_test.go +++ b/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") +} diff --git a/cmd/proxy/chunk.go b/cmd/proxy/chunk.go new file mode 100644 index 0000000..4bfdac8 --- /dev/null +++ b/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 +} diff --git a/cmd/proxy/chunk_test.go b/cmd/proxy/chunk_test.go new file mode 100644 index 0000000..54ab181 --- /dev/null +++ b/cmd/proxy/chunk_test.go @@ -0,0 +1,312 @@ +package main + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/elevran/charon/pkg/charon" +) + +// --------------------------------------------------------------------------- +// recordingBackend — minimal Backend mock for chunkedResponseWriter tests +// --------------------------------------------------------------------------- + +type chunkCall struct { + k uint32 + body []byte +} + +type completeCall struct { + stagingID string + responseID string + tenantKey string + total uint32 +} + +type recordingBackend struct { + mu sync.Mutex + chunks []chunkCall + completes []completeCall + aborts []string + chunkErr error + completeErr error +} + +func (r *recordingBackend) AppendChunk(_ context.Context, _ string, k uint32, _ string, body []byte) error { + r.mu.Lock() + defer r.mu.Unlock() + if r.chunkErr != nil { + return r.chunkErr + } + cp := make([]byte, len(body)) + copy(cp, body) + r.chunks = append(r.chunks, chunkCall{k: k, body: cp}) + return nil +} + +func (r *recordingBackend) Complete(_ context.Context, stagingID, responseID, tenantKey string, total uint32) (string, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.completeErr != nil { + return "", r.completeErr + } + r.completes = append(r.completes, completeCall{stagingID: stagingID, responseID: responseID, tenantKey: tenantKey, total: total}) + return responseID, nil +} + +func (r *recordingBackend) Abort(_ context.Context, stagingID string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.aborts = append(r.aborts, stagingID) + return nil +} + +func (r *recordingBackend) snapshot() (chunks []chunkCall, completes []completeCall, aborts []string) { + r.mu.Lock() + defer r.mu.Unlock() + return append([]chunkCall(nil), r.chunks...), append([]completeCall(nil), r.completes...), append([]string(nil), r.aborts...) +} + +// Unused interface methods — must exist to satisfy charon.Backend. +func (r *recordingBackend) Resolve(_ context.Context, _, _ string, _ []byte) (string, []charon.ResolveTurn, error) { + return "", nil, nil +} +func (r *recordingBackend) GetChain(_ context.Context, _, _ string) ([]charon.ResolveTurn, error) { + return nil, nil +} +func (r *recordingBackend) Store(_ context.Context, _, _, _ string, _ []byte) error { return nil } +func (r *recordingBackend) Retrieve(_ context.Context, _, _ string) (*charon.RetrieveResponse, error) { + return nil, nil +} +func (r *recordingBackend) Delete(_ context.Context, _, _ string) error { return nil } + +var _ charon.Backend = (*recordingBackend)(nil) + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +func newWriter(b *recordingBackend, maxChunkBytes int64) *chunkedResponseWriter { + return newChunkedResponseWriter(context.Background(), b, "sid1", "resp1", "tenant1", maxChunkBytes) +} + +// --------------------------------------------------------------------------- +// tests +// --------------------------------------------------------------------------- + +func TestWriterSingleSmallChunk(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + payload := []byte(`{"id":"r1","output":[]}`) + require.NoError(t, w.Add(payload)) + require.NoError(t, w.Close()) + + chunks, completes, aborts := b.snapshot() + assert.Len(t, chunks, 1) + assert.Equal(t, uint32(0), chunks[0].k) + assert.Equal(t, payload, chunks[0].body) + assert.Len(t, completes, 1) + assert.Equal(t, uint32(len(payload)), completes[0].total) + assert.Empty(t, aborts) +} + +func TestWriterCapFlush(t *testing.T) { + // cap = 10 bytes; 20 bytes total → 2 full chunks (intra-Add splitting fills + // each chunk to cap before moving on, so bytes span chunk boundaries freely). + b := &recordingBackend{} + w := newWriter(b, 10) + p1 := []byte("12345678") // 8 bytes + p2 := []byte("abcdefgh") // 8 bytes + p3 := []byte("ABCD") // 4 bytes → 20 total → 2 full chunks of 10, no remainder at Close + require.NoError(t, w.Add(p1)) + require.NoError(t, w.Add(p2)) + require.NoError(t, w.Add(p3)) + require.NoError(t, w.Close()) + + chunks, completes, aborts := b.snapshot() + assert.GreaterOrEqual(t, len(chunks), 2, "20 bytes at 10-byte cap must produce ≥2 chunks") + assert.Len(t, completes, 1) + assert.Empty(t, aborts) + + // chunk indices must be 0,1 + for i, c := range chunks { + assert.Equal(t, uint32(i), c.k) + } + + // total must equal sum of chunk body lengths + var expectedTotal uint32 + for _, c := range chunks { + expectedTotal += uint32(len(c.body)) + } + assert.Equal(t, expectedTotal, completes[0].total) + + // concatenation must reproduce the original bytes in order + var concat []byte + for _, c := range chunks { + concat = append(concat, c.body...) + } + assert.Equal(t, append(append(p1, p2...), p3...), concat) +} + +func TestWriterZeroBytesAborts(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + require.NoError(t, w.Close()) + + chunks, completes, aborts := b.snapshot() + assert.Empty(t, chunks) + assert.Empty(t, completes) + assert.Len(t, aborts, 1) + assert.Equal(t, "sid1", aborts[0]) +} + +func TestWriterZeroItemAborts(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + require.NoError(t, w.Add(nil)) + require.NoError(t, w.Add([]byte{})) + require.NoError(t, w.Close()) + + chunks, completes, aborts := b.snapshot() + assert.Empty(t, chunks) + assert.Empty(t, completes) + assert.Len(t, aborts, 1) +} + +func TestWriterNonZeroCloses(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + require.NoError(t, w.Add([]byte(`{"id":"r1"}`))) + require.NoError(t, w.Close()) + + _, completes, aborts := b.snapshot() + assert.Len(t, completes, 1) + assert.Empty(t, aborts) +} + +func TestWriterCloseIdempotent(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + require.NoError(t, w.Add([]byte(`{"id":"r1"}`))) + + var wg sync.WaitGroup + errs := make([]error, 2) + for i := range 2 { + wg.Add(1) + go func(idx int) { + defer wg.Done() + errs[idx] = w.Close() + }(i) + } + wg.Wait() + + assert.Equal(t, errs[0], errs[1], "both Close calls must return the same error") + + _, completes, _ := b.snapshot() + assert.Len(t, completes, 1, "terminal flush must execute exactly once") +} + +// TestWriterCloseIdempotentOnError pins that concurrent Close calls return the +// same non-nil error when Complete fails. Without the closedCh-aware read in +// Close, a second caller can read a stale nil closedErr before the first +// caller's deferred close(closedCh) fires, masking the failure. +func TestWriterCloseIdempotentOnError(t *testing.T) { + sentinel := errors.New("complete failed") + b := &recordingBackend{completeErr: sentinel} + w := newWriter(b, 1024) + require.NoError(t, w.Add([]byte(`{"id":"r1"}`))) + + var wg sync.WaitGroup + errs := make([]error, 2) + for i := range 2 { + wg.Add(1) + go func(idx int) { + defer wg.Done() + errs[idx] = w.Close() + }(i) + } + wg.Wait() + + assert.ErrorIs(t, errs[0], sentinel, "first Close must propagate Complete error") + assert.ErrorIs(t, errs[1], sentinel, "second Close must also propagate the same error, not a stale nil") +} + +func TestWriterAbortSkipsComplete(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 1024) + require.NoError(t, w.Add([]byte(`{"id":"r1"}`))) + require.NoError(t, w.Abort()) + + _, completes, aborts := b.snapshot() + assert.Empty(t, completes) + assert.Len(t, aborts, 1) +} + +func TestWriterCancelDuringFlush(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + b := &recordingBackend{} + w := newChunkedResponseWriter(ctx, b, "sid1", "resp1", "tenant1", 4) + + cancel() + b.chunkErr = context.Canceled + + // Two adds: first buffers (4 bytes exactly == cap, no flush yet since not exceeding), + // second Add triggers flush because combined would exceed cap. + err := w.Add([]byte("abcd")) // 4 bytes == cap, no flush (not > cap) + if err == nil { + err = w.Add([]byte("ef")) // 4+2=6 > 4, triggers flush → error + } + if err == nil { + err = w.Close() + } + assert.ErrorIs(t, err, context.Canceled) + + // Subsequent Add returns the same error. + err2 := w.Add([]byte("z")) + assert.ErrorIs(t, err2, context.Canceled) +} + +func TestWriterFlushFailurePropagates(t *testing.T) { + sentinel := errors.New("chunk write failed") + b := &recordingBackend{chunkErr: sentinel} + w := newWriter(b, 5) + + // cap=5: first Add (4 bytes) buffers without flushing (4 < 5). + // Second Add (2 bytes) fills pending to 6 > 5, triggers flush → error. + require.NoError(t, w.Add([]byte("abcd"))) + err := w.Add([]byte("ef")) + if err == nil { + err = w.Close() + } + assert.ErrorIs(t, err, sentinel) + + // Subsequent calls return the same error. + assert.ErrorIs(t, w.Add([]byte("x")), sentinel) +} + +func TestWriterBodyShape(t *testing.T) { + b := &recordingBackend{} + w := newWriter(b, 10) // small cap to force multiple chunks + + // Simulate storing a large blob in two parts. + part1 := []byte(`{"id":"r1","output":[`) + part2 := []byte(`{"type":"msg"}]}`) + require.NoError(t, w.Add(part1)) + require.NoError(t, w.Add(part2)) + require.NoError(t, w.Close()) + + chunks, _, _ := b.snapshot() + require.NotEmpty(t, chunks) + + // Concatenated chunks must reproduce the original bytes. + var concat []byte + for _, c := range chunks { + concat = append(concat, c.body...) + } + assert.Equal(t, append(part1, part2...), concat) +} diff --git a/cmd/proxy/disruptive_test.go b/cmd/proxy/disruptive_test.go index 97754fb..3df83f4 100644 --- a/cmd/proxy/disruptive_test.go +++ b/cmd/proxy/disruptive_test.go @@ -151,3 +151,36 @@ func TestStreamingStoreFailureIsFatal(t *testing.T) { assert.Equal(t, http.StatusNotFound, getResp.StatusCode, "response must not be accessible via GET after a store failure") } + +// TestStreamingInferenceFailureEmitsFailedNotCompleted verifies that a +// mid-stream inference failure results in no chunk/complete calls reaching +// Charon — the staging record must remain uncommitted (abort via TTL). +func TestStreamingInferenceFailureEmitsFailedNotCompleted(t *testing.T) { + partialSrv := inference.NewPartialMockServer() + t.Cleanup(partialSrv.Close) + + rec := &routingRecorder{} + s := newTestStack(t, withInferenceURL(partialSrv.URL), withCharonMiddleware(rec.middleware())) + + body, _ := json.Marshal(map[string]any{ + "model": "mock", + "stream": true, + "input": []map[string]string{{"type": "message", "role": "user", "content": "hi"}}, + }) + resp, err := http.Post(s.proxyURL+"/responses", "application/json", bytes.NewReader(body)) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + sse := readSSE(t, resp) + canonicalID := sse.createdID() + + assert.NotEmpty(t, canonicalID, "proxy must emit response.created") + assert.NotContains(t, sse.EventTypes, "response.completed", + "proxy must NOT emit response.completed when inference truncated") + + hits := rec.snapshot() + assert.Equal(t, 0, hitsContaining(hits, "/chunks/"), + "no AppendChunk calls when inference stream was truncated before response.completed") + assert.Equal(t, 0, hitsContaining(hits, "/complete"), + "no Complete call when inference stream was truncated") +} diff --git a/cmd/proxy/handler.go b/cmd/proxy/handler.go index cb906f4..b4f4afc 100644 --- a/cmd/proxy/handler.go +++ b/cmd/proxy/handler.go @@ -10,20 +10,55 @@ import ( "time" "github.com/elevran/charon/cmd/proxy/inference" + "github.com/elevran/charon/internal/bytesize" "github.com/elevran/charon/internal/server" "github.com/elevran/charon/pkg/charon" ) // Handler is the client-facing Responses API proxy handler. type Handler struct { - charon charon.Backend - inf inference.Backend - log *slog.Logger + charon charon.Backend + inf inference.Backend + log *slog.Logger + maxChunkBytes int64 } -// NewHandler creates a Handler. +// commitStoredResponse persists blob to Charon via the staging protocol +// (POST /staging → PUT /staging/{sid}/chunks/{k} → PUT /staging/{sid}/complete). +// On either Add or Close failure it explicitly aborts the staging record so +// partial chunks are not left orphaned; abort failures are logged at warn +// level but do not mask the original error. Returns the underlying storage +// error so the caller can map it to its surface (HTTP / SSE / WS). +func (h *Handler) commitStoredResponse(ctx context.Context, stagingID, id, tenantKey string, blob []byte) error { + cw := newChunkedResponseWriter(ctx, h.charon, stagingID, id, tenantKey, h.maxChunkBytes) + if err := cw.Add(blob); err != nil { + h.log.Error("chunk add", "id", id, "err", err) + if abortErr := cw.Abort(); abortErr != nil { + h.log.Warn("chunk abort", "id", id, "err", abortErr) + } + return err + } + if err := cw.Close(); err != nil { + h.log.Error("chunk close", "id", id, "err", err) + if abortErr := cw.Abort(); abortErr != nil { + h.log.Warn("chunk abort", "id", id, "err", abortErr) + } + return err + } + return nil +} + +// NewHandler creates a Handler with the default MaxChunkBytes cap (1 MiB). func NewHandler(ch charon.Backend, inf inference.Backend, log *slog.Logger) *Handler { - return &Handler{charon: ch, inf: inf, log: log} + return &Handler{charon: ch, inf: inf, log: log, maxChunkBytes: bytesize.MiB} +} + +// WithMaxChunkBytes sets the per-response chunk cap and returns h for chaining. +func (h *Handler) WithMaxChunkBytes(n int64) *Handler { + if n > 0 { + h.maxChunkBytes = n + } + return h } // RegisterHandlers mounts Responses API routes on mux. @@ -109,8 +144,7 @@ func (h *Handler) HandleCreate(w http.ResponseWriter, r *http.Request) { if req.ShouldStore() { responseBlob := marshalStoredResponse(infResp, req.PreviousResponseID, req.Instructions, req.Background) - if err := h.charon.Store(ctx, infResp.ID, stagingID, tenantKey, responseBlob); err != nil { - h.log.Error("charon store", "id", infResp.ID, "err", err) + if err := h.commitStoredResponse(ctx, stagingID, infResp.ID, tenantKey, responseBlob); err != nil { server.WriteError(w, http.StatusInternalServerError, "storage error") return } diff --git a/cmd/proxy/handler_test.go b/cmd/proxy/handler_test.go index 94a2c4d..4bd3181 100644 --- a/cmd/proxy/handler_test.go +++ b/cmd/proxy/handler_test.go @@ -159,6 +159,24 @@ func TestStoreTrueContinuation(t *testing.T) { assert.NotEqual(t, anchor.ID, follow.ID) } +// TestBufferedCapConfigurable verifies that a tiny MaxChunkBytes forces the +// buffered path to split a response blob into multiple AppendChunk calls. +func TestBufferedCapConfigurable(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 is a test request that should produce a response blob exceeding 64 bytes", + }) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + hits := rec.snapshot() + assert.GreaterOrEqual(t, hitsContaining(hits, "/chunks/"), 2, + "tiny cap must force ≥2 AppendChunk calls for a non-trivial response blob") + assert.Equal(t, 1, hitsContaining(hits, "/complete"), "exactly 1 Complete call") +} + // TestCreateStoreFalseProduces200NoCommit ensures that the proxy // doesn't 5xx on a store:false turn — the response is returned to the // client and no Charon state is committed. diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go index 2071646..3ccf09f 100644 --- a/cmd/proxy/main.go +++ b/cmd/proxy/main.go @@ -56,7 +56,7 @@ func run() error { infClient := inference.New(opts.Backend, opts.APIKey, timeout) charonClient := charon.New(opts.CharonURL, timeout) - h := NewHandler(charonClient, infClient, log) + h := NewHandler(charonClient, infClient, log).WithMaxChunkBytes(opts.MaxChunkBytes) mux := http.NewServeMux() RegisterHandlers(mux, h) srv := server.NewServerFromMux(opts.Listen, mux, log, tp) diff --git a/cmd/proxy/sse.go b/cmd/proxy/sse.go index 9e281bb..f354ed0 100644 --- a/cmd/proxy/sse.go +++ b/cmd/proxy/sse.go @@ -29,21 +29,6 @@ func writeSSE(w http.ResponseWriter, evt sseEvent) { } } -// streamBuffer accumulates bare output item JSON (SSE framing already stripped). -type streamBuffer struct { - items []json.RawMessage -} - -func (b *streamBuffer) add(item json.RawMessage) { - b.items = append(b.items, item) -} - -func (b *streamBuffer) drain() []json.RawMessage { - items := b.items - b.items = nil - return items -} - // handleStream implements POST /responses with stream:true. // // Output items are extracted from response.output_item.done events (bare item @@ -101,8 +86,7 @@ func (h *Handler) handleStream(w http.ResponseWriter, r *http.Request, req Creat var canonicalID string var created bool var finalInfResp *inference.Response - - buf := &streamBuffer{} + var outputItems []json.RawMessage // accumulate output items for Charon store for evt := range ch { if evt.Response != nil && evt.Response.ID != "" && canonicalID == "" { @@ -130,7 +114,7 @@ func (h *Handler) handleStream(w http.ResponseWriter, r *http.Request, req Creat // Forward to client immediately (no buffering of client events). writeSSE(w, sseEvent{Type: evt.Type, SequenceNumber: seq, OutputIndex: &outIdx, Item: evt.Item}) seq++ - buf.add(evt.Item) + outputItems = append(outputItems, evt.Item) case "response.completed": finalInfResp = evt.Response @@ -161,11 +145,10 @@ func (h *Handler) handleStream(w http.ResponseWriter, r *http.Request, req Creat } if req.ShouldStore() && canonicalID != "" { - finalInfResp.Output = buf.drain() + finalInfResp.Output = outputItems responseBlob := marshalStoredResponse(finalInfResp, req.PreviousResponseID, req.Instructions, req.Background) - if err := h.charon.Store(ctx, canonicalID, stagingID, tenantKey, responseBlob); err != nil { - h.log.Error("charon store after stream", "id", canonicalID, "err", err) - return // do not emit response.completed — client must not believe the response was persisted + if err := h.commitStoredResponse(ctx, stagingID, canonicalID, tenantKey, responseBlob); err != nil { + return // do not emit response.completed — staging was aborted } } diff --git a/cmd/proxy/testhelpers_test.go b/cmd/proxy/testhelpers_test.go index 5b65461..94fff7b 100644 --- a/cmd/proxy/testhelpers_test.go +++ b/cmd/proxy/testhelpers_test.go @@ -47,6 +47,7 @@ type stackConfig struct { infURL string // non-empty → use this URL instead of a fresh MockServer realListeners bool timeout time.Duration + maxChunkBytes int64 // 0 → use Handler default (1 MiB) } // stackOption is a functional option for newTestStack. @@ -64,6 +65,11 @@ func withInferenceURL(u string) stackOption { return func(c *stackConfig) { c.infURL = u } } +// withMaxChunkBytes sets the per-chunk byte cap on the proxy Handler. +func withMaxChunkBytes(n int64) stackOption { + return func(c *stackConfig) { c.maxChunkBytes = n } +} + // withRealListeners uses real OS-assigned TCP ports instead of httptest servers. // Required when an out-of-process client (e.g. bun) needs to connect. func withRealListeners() stackOption { @@ -117,6 +123,9 @@ func newTestStack(t testing.TB, opts ...stackOption) *testStack { buildProxy := func(charonURL string) http.Handler { proxyH := NewHandler(charon.New(charonURL, cfg.timeout), infClient, log) + if cfg.maxChunkBytes > 0 { + proxyH.WithMaxChunkBytes(cfg.maxChunkBytes) + } mux := http.NewServeMux() RegisterHandlers(mux, proxyH) return mux diff --git a/cmd/proxy/ws.go b/cmd/proxy/ws.go index 8f82068..fefcecc 100644 --- a/cmd/proxy/ws.go +++ b/cmd/proxy/ws.go @@ -206,8 +206,7 @@ func (h *Handler) wsTurn(ctx context.Context, conn *websocket.Conn, cache *wsCac var canonicalID string var sentCreated bool var finalInfResp *inference.Response - - buf := &streamBuffer{} + var outputItems []json.RawMessage // accumulate output items for store:true and wsCache for evt := range ch { if evt.Response != nil && evt.Response.ID != "" && canonicalID == "" { @@ -233,7 +232,7 @@ func (h *Handler) wsTurn(ctx context.Context, conn *websocket.Conn, cache *wsCac case "response.output_item.done": h.wsSend(conn, sseEvent{Type: evt.Type, SequenceNumber: seq, OutputIndex: &outIdx, Item: evt.Item}) seq++ - buf.add(evt.Item) + outputItems = append(outputItems, evt.Item) case "response.completed": finalInfResp = evt.Response @@ -262,12 +261,11 @@ func (h *Handler) wsTurn(ctx context.Context, conn *websocket.Conn, cache *wsCac } if msg.ShouldStore() { - finalInfResp.Output = buf.drain() + finalInfResp.Output = outputItems responseBlob := marshalStoredResponse(finalInfResp, msg.PreviousResponseID, msg.Instructions, msg.Background) - if err := h.charon.Store(ctx, canonicalID, stagingID, tenantKey, responseBlob); err != nil { - h.log.Error("ws charon store", "id", canonicalID, "err", err) + if err := h.commitStoredResponse(ctx, stagingID, canonicalID, tenantKey, responseBlob); err != nil { h.wsSendError(conn, 500, "storage_error", "response not persisted") - return // do not emit response.completed — client must not believe the response was persisted + return } } else { // store:false — cache assembled flat_context for subsequent turns. diff --git a/internal/charonconfig/bytesize.go b/internal/bytesize/bytesize.go similarity index 71% rename from internal/charonconfig/bytesize.go rename to internal/bytesize/bytesize.go index 265ea2b..a50c029 100644 --- a/internal/charonconfig/bytesize.go +++ b/internal/bytesize/bytesize.go @@ -1,4 +1,8 @@ -package charonconfig +// Package bytesize provides a configurable byte-size type that unmarshals +// from either a plain integer (bytes) or a string with an optional unit suffix +// (B, KB, MB, GB), plus named binary multipliers (KiB, MiB, GiB) so call sites +// can stop writing "1 << 20". +package bytesize import ( "encoding/json" @@ -7,6 +11,14 @@ import ( "strings" ) +// Binary multipliers. K = 1024 to match the unit parsing below. +const ( + KiB int64 = 1 << 10 + MiB int64 = 1 << 20 + GiB int64 = 1 << 30 + TiB int64 = 1 << 40 +) + // ByteSize is an int64 that unmarshals from either a plain integer (bytes) or // a string with an optional unit suffix: B, KB, MB, GB. K=1024. type ByteSize int64 @@ -18,6 +30,8 @@ var unitMultipliers = map[string]int64{ "gb": 1024 * 1024 * 1024, } +// UnmarshalJSON accepts a JSON number (raw bytes) or string (number with +// optional unit suffix). func (b *ByteSize) UnmarshalJSON(data []byte) error { var n int64 if err := json.Unmarshal(data, &n); err == nil { diff --git a/internal/charonconfig/bytesize_test.go b/internal/bytesize/bytesize_test.go similarity index 65% rename from internal/charonconfig/bytesize_test.go rename to internal/bytesize/bytesize_test.go index ff31622..440ed13 100644 --- a/internal/charonconfig/bytesize_test.go +++ b/internal/bytesize/bytesize_test.go @@ -1,14 +1,13 @@ -package charonconfig_test +package bytesize_test import ( "encoding/json" - "flag" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/elevran/charon/internal/charonconfig" + "github.com/elevran/charon/internal/bytesize" ) func TestByteSizeUnmarshal(t *testing.T) { @@ -36,9 +35,9 @@ func TestByteSizeUnmarshal(t *testing.T) { for _, tc := range tests { t.Run(tc.input, func(t *testing.T) { - var b charonconfig.ByteSize + var b bytesize.ByteSize require.NoError(t, json.Unmarshal([]byte(tc.input), &b)) - assert.Equal(t, charonconfig.ByteSize(tc.want), b) + assert.Equal(t, bytesize.ByteSize(tc.want), b) }) } } @@ -57,21 +56,10 @@ func TestByteSizeUnmarshalErrors(t *testing.T) { for _, tc := range tests { t.Run(tc.input, func(t *testing.T) { - var b charonconfig.ByteSize + var b bytesize.ByteSize err := json.Unmarshal([]byte(tc.input), &b) require.Error(t, err) assert.Contains(t, err.Error(), tc.errFrag) }) } } - -func TestByteSizeInConfig(t *testing.T) { - yaml := []byte("charon:\n storage:\n max_payload: \"10MB\"\n") - cfg := configFromBytes(t, yaml) - opts := charonconfig.NewOptions() - fs := flag.NewFlagSet("test", flag.ContinueOnError) - opts.AddFlags(fs) - require.NoError(t, fs.Parse([]string{"--config", cfg})) - require.NoError(t, opts.Complete(fs)) - assert.Equal(t, charonconfig.ByteSize(10*1024*1024), opts.MaxPayload) -} diff --git a/internal/charonconfig/charonconfig.go b/internal/charonconfig/charonconfig.go index 8d941df..90fa340 100644 --- a/internal/charonconfig/charonconfig.go +++ b/internal/charonconfig/charonconfig.go @@ -8,6 +8,7 @@ import ( "sigs.k8s.io/yaml" + "github.com/elevran/charon/internal/bytesize" "github.com/elevran/charon/internal/telemetry" ) @@ -24,9 +25,9 @@ type CharonOptions struct { // TTLDays is the maximum age of a stored response. TTLDays int MaxResponses int64 - MaxPayload ByteSize + MaxPayload bytesize.ByteSize MaxChainDepth int - MaxContextBytes ByteSize + MaxContextBytes bytesize.ByteSize // WorkerTTLInterval is how often the background TTL reaper runs (not the TTL itself). WorkerTTLInterval time.Duration @@ -116,12 +117,12 @@ type fileCharonConfig struct { } type fileStorageConfig struct { - DataDir string `json:"data_dir"` - TTLDays int `json:"ttl_days"` - MaxResponses int64 `json:"max_responses"` - MaxPayload ByteSize `json:"max_payload"` - MaxChainDepth int `json:"max_chain_depth"` - MaxContextBytes ByteSize `json:"max_context_bytes"` + DataDir string `json:"data_dir"` + TTLDays int `json:"ttl_days"` + MaxResponses int64 `json:"max_responses"` + MaxPayload bytesize.ByteSize `json:"max_payload"` + MaxChainDepth int `json:"max_chain_depth"` + MaxContextBytes bytesize.ByteSize `json:"max_context_bytes"` } type fileWorkerConfig struct { diff --git a/internal/proxyconfig/proxyconfig.go b/internal/proxyconfig/proxyconfig.go index 9e0b1a1..bc3da9e 100644 --- a/internal/proxyconfig/proxyconfig.go +++ b/internal/proxyconfig/proxyconfig.go @@ -8,9 +8,21 @@ import ( "sigs.k8s.io/yaml" + "github.com/elevran/charon/internal/bytesize" "github.com/elevran/charon/internal/telemetry" ) +// defaultMaxChunkBytes is the default cap for chunkedResponseWriter when +// --max-chunk-bytes is unset (1 MiB). +const defaultMaxChunkBytes int64 = bytesize.MiB + +// maxMaxChunkBytes is the upper bound the proxy accepts for --max-chunk-bytes. +// Hardcoded rather than imported from internal/server to avoid a config→server +// import edge — keep this in sync with server.defaultChunkBodyBytes +// (internal/server/handlers.go). Charon enforces the body cap on incoming +// chunk PUTs and rejects oversize bodies with 400. +const maxMaxChunkBytes int64 = bytesize.MiB + // ProxyOptions holds configuration for the proxy server. type ProxyOptions struct { // Config file path — set by --config flag. @@ -28,6 +40,10 @@ type ProxyOptions struct { // Auto-derived from config file proxy.charon_url or Charon's listen address. CharonURL string + // MaxChunkBytes caps the in-memory response buffer before flushing to + // Charon as a chunk. 0 or negative applies the default (1 MiB). + MaxChunkBytes int64 + Telemetry telemetry.Options } @@ -48,6 +64,8 @@ func (o *ProxyOptions) AddFlags(fs *flag.FlagSet) { fs.StringVar(&o.Listen, "listen", o.Listen, "proxy server listen address") fs.StringVar(&o.Backend, "backend", o.Backend, "inference backend base URL") fs.StringVar(&o.CharonURL, "charon-url", o.CharonURL, "charon internal API base URL") + fs.Int64Var(&o.MaxChunkBytes, "max-chunk-bytes", 0, + "max response bytes buffered before flushing to Charon as a chunk (0 = default 1 MiB, max 1 MiB)") o.Telemetry.AddFlags(fs) } @@ -88,6 +106,10 @@ func (o *ProxyOptions) Complete(fs *flag.FlagSet) error { } o.Telemetry.ServiceName = fc.Telemetry.ServiceName + if o.MaxChunkBytes <= 0 { + o.MaxChunkBytes = defaultMaxChunkBytes + } + return nil } @@ -96,6 +118,9 @@ func (o *ProxyOptions) Validate() error { if o.Backend == "" { return fmt.Errorf("proxy backend (inference base URL) is empty") } + if o.MaxChunkBytes > maxMaxChunkBytes { + return fmt.Errorf("--max-chunk-bytes=%d exceeds server chunk body cap %d", o.MaxChunkBytes, maxMaxChunkBytes) + } return nil } diff --git a/internal/server/handlers.go b/internal/server/handlers.go index 6337e5c..caedd2a 100644 --- a/internal/server/handlers.go +++ b/internal/server/handlers.go @@ -10,17 +10,17 @@ import ( "github.com/google/uuid" + "github.com/elevran/charon/internal/bytesize" "github.com/elevran/charon/internal/chainstore" ) -// maxChunkBytes is the hard upper bound on a streaming-chunk PUT body. -// 1 MB default / 4 MB cap keeps the proxy→Charon ingest path bounded: small -// enough to stay under typical 256 MB per-process caps with 50+ concurrent -// inferences, large enough that per-batch HTTP framing overhead is small. -const ( - defaultMaxChunkBytes = 1 << 20 - maxChunkBytes = 4 << 20 -) +// defaultChunkBodyBytes caps the HTTP body of a streaming-chunk PUT +// (PUT /staging/{sid}/chunks/{k}). Distinct from defaultMaxChunkBytes +// in internal/proxyconfig — that const is the proxy's *flush* size; +// this const is the server's *body-read* cap. The two must stay in sync +// because the proxy emits chunks sized at the flush cap, and Charon +// rejects anything larger. +const defaultChunkBodyBytes = bytesize.MiB // Handler wires chainstore.Store to HTTP endpoints. type Handler struct { @@ -351,7 +351,7 @@ func (h *Handler) HandleAppendChunk(w http.ResponseWriter, r *http.Request) { // Cap per-chunk reads independently of the global maxBodyBytes. Default // 1 MB; configurable up to 4 MB via WithMaxChunkBytes. - r.Body = http.MaxBytesReader(w, r.Body, defaultMaxChunkBytes) + r.Body = http.MaxBytesReader(w, r.Body, defaultChunkBodyBytes) chunkBody, err := io.ReadAll(r.Body) if err != nil { WriteError(w, http.StatusBadRequest, "failed to read chunk body") diff --git a/pkg/charon/client.go b/pkg/charon/client.go index 26dad99..01c6fbf 100644 --- a/pkg/charon/client.go +++ b/pkg/charon/client.go @@ -349,6 +349,9 @@ type Backend interface { Resolve(ctx context.Context, previousID, tenantKey string, requestBlob []byte) (string, []ResolveTurn, error) GetChain(ctx context.Context, id, tenantKey string) ([]ResolveTurn, error) Store(ctx context.Context, id, stagingID, tenantKey string, responseBlob []byte) error + AppendChunk(ctx context.Context, stagingID string, k uint32, responseID string, body []byte) error + Complete(ctx context.Context, stagingID, responseID, tenantKey string, total uint32) (string, error) + Abort(ctx context.Context, stagingID string) error Retrieve(ctx context.Context, id, tenantKey string) (*RetrieveResponse, error) Delete(ctx context.Context, id, tenantKey string) error }