Repository navigation
proxy: replace streamBuffer with chunkedResponseWriter for capped response storage - #88
Merged
Merged
Conversation
…ponse storage
Replace the all-at-once Store call with a staging-based chunked protocol:
- POST /staging → PUT /staging/{sid}/chunks/{k} (×N) → PUT /staging/{sid}/complete
- Abort via PUT /staging/{sid}/abort when no bytes were written
Key changes:
- pkg/charon/client.go: extend Backend interface with AppendChunk, Complete, Abort
- cmd/proxy/chunk.go: new chunkedResponseWriter — byte-level splitter with
configurable MaxChunkBytes cap; splits within a single Add call
- internal/proxyconfig/proxyconfig.go: add --max-chunk-bytes flag (default 1 MiB)
- cmd/proxy/handler.go: wire chunkedResponseWriter into HandleCreate; add
WithMaxChunkBytes option
- cmd/proxy/sse.go: remove streamBuffer; accumulate output items, marshal once,
write via chunkedResponseWriter
- cmd/proxy/ws.go: same pattern as sse.go
- cmd/proxy/main.go: pass MaxChunkBytes from config to handler
Tests: unit tests in chunk_test.go; integration tests in handler_test.go,
backend_routing_test.go, and disruptive_test.go verify chunk counts, cap
splitting, store:false no-op, and mid-stream failure non-persistence.
Apply consensus findings from /simplify and /ponytail-review: - chunk.go: drop errWriterClosed sentinel. No caller or test compares against it; closedErr now stores backend.Abort's error directly. Removes the 'errors' import. - chunk.go (Abort): replace the errWriterClosed default + override with a single closedErr = err assignment. - backend_routing_test.go: rewrite the docstring above TestBufferedProxyStoreFalseNoStagingCalls to match what the test actually asserts (no staging calls on store:false). The prior comment named a non-existent test and rambled about aborts. - chunk_test.go (TestWriterCapFlush): fix 'chunk indices must be 0,1,2' comment — the loop iterates only 0,1 for the 2-element chunks slice.
The chunk tests added in this PR (TestBufferedProxySingleChunk,
TestStreamedProxySingleChunk, TestBufferedProxyMultipleChunks,
TestBufferedProxyStoreFalseNoStagingCalls, TestBufferedCapConfigurable)
all manually constructed 'rec := &routingRecorder{}; s := newTestStack(t,
withCharonMiddleware(rec.middleware()))' — the exact pattern that
newRoutingStack(t) already encapsulates.
Extend newRoutingStack to accept variadic stackOption so the cap test
can pass withMaxChunkBytes(64) alongside the recorder middleware.
Net: -5 lines per test, single source of truth for the recorder setup.
- New internal/bytesize package: ByteSize type with UnmarshalJSON, plus KiB/MiB/GiB/TiB constants. The constants are exported so call sites can write 'bytesize.MiB' instead of '1 << 20'. - charonconfig: drop ByteSize; reference bytesize.ByteSize in CharonOptions and fileStorageConfig. JSON parsing is preserved (the unmarshal method moved with the type). - internal/server/handlers.go: defaultMaxChunkBytes = bytesize.MiB, maxChunkBytes = 4 * bytesize.MiB. - cmd/proxy/handler.go: NewHandler's default cap is bytesize.MiB. - internal/proxyconfig/proxyconfig.go: defaultMaxChunkBytes const is bytesize.MiB.
The code-reviewer agent found a real concurrency bug in chunkedResponseWriter.Close: the second concurrent Close caller read w.closedErr under the lock BEFORE waiting on <-ch, racing against the first Close's deferred write of closedErr (which happens after the unlocked Complete network call). On failure, the second caller could return a stale nil, masking the real error. Reproduction (reviewer's 200-iter trial without the fix): ~3/200 failures under go test -race. Fix: in the early-return path, drop the lock, wait on <-ch, then re-acquire the lock to read closedErr. The channel close provides happens-before for closedErr, and the second lock acquisition prevents torn interface reads. Symmetric with Abort, which already reads closedErr after <-ch. Also add TestWriterCloseIdempotentOnError that exercises the failure path concurrently to lock in the contract. Pin both paths at 200 iter under -race. Drive-by: fix the misleading comment in TestBufferedProxyStoreFalseNoStagingCalls — the prior wording pointed at disruptive_test.go for the empty-response abort path, but disruptive_test.go only covers inference-truncation. The unit tests in chunk_test.go (TestWriterZeroBytesAborts) actually exercise that path.
Apply every Important and Minor finding from the code-reviewer agent: 1. Symmetric Abort + log Abort errors (handler.go/sse.go/ws.go): Extract (h *Handler) commitStoredResponse(ctx, stagingID, id, tenantKey, blob) error. The helper owns Add → Abort-on-error → Close → Abort-on-error, logging all stages. Abort failures are logged at warn but never mask the original error. Each call site collapses from ~12 lines to a single if that maps the error to its surface (HTTP / SSE / WS). This addresses the deferred altitude finding (three-way duplication) and the symmetry finding in one shot. 2. chunk.flush(): make flush failure terminal. When AppendChunk returns an error, the writer now sets closed=true and closes closedCh (gated by !w.closed so a Close-driven flush doesn't double-close the channel). Subsequent Add/Close/Abort calls return the same error without retrying the same chunk index. Removed the redundant closedErr = err in Close since flush now owns that write. 3. internal/server/handlers.go: rename defaultMaxChunkBytes to defaultChunkBodyBytes — disambiguates from proxyconfig's defaultMaxChunkBytes (different meaning, same value). Also remove the unused maxChunkBytes = 4*MiB const. 4. internal/proxyconfig/proxyconfig.go: add maxMaxChunkBytes = MiB cap and Validate() check that rejects --max-chunk-bytes values exceeding it with a clear error. Update --max-chunk-bytes help text to advertise the upper bound. Hardcoded rather than imported from internal/server to avoid a config→server import edge; comment cross-references server.defaultChunkBodyBytes. Verified: presubmit green; go test -race -count=200 -run TestWriter ./cmd/proxy/ passes (covers the new terminal-flush state path).
elevran
added a commit
that referenced
this pull request
Jul 10, 2026
All store paths (new chain and continuation) now use the staging protocol:
POST /staging → PUT chunks → PUT complete. The old GET /responses/{id}
(resolve chain) + POST /responses/{id} (buffered store) flows no longer
reflect what the proxy actually does.
Specific fixes:
- System diagram: replace wrong Charon calls with actual endpoints
- Proxy-Charon interaction: unify new chain and continuation under staging
protocol; document store:false path (GET /chain/{id}, added in PR #87)
- Charon API section: rename buffered-store section to staging protocol
(primary path), add GET /chain/{id} section, fix POST /responses body
fields (prev_id/response_id/request_blob/response_blob + X-Tenant-Key
header, not previous_response_id/tenant_key in body)
- Remove stale 'open question' note (answered by GET /chain/{id})
- ID flow diagrams: collapse to one staging-protocol diagram; remove wrong
GET /responses (chain) + POST /responses/{id} buffered diagrams
- store:false semantics: describe GET /chain/{id} for continuations
- Performance section: update endpoint names to match current API
elevran
added a commit
that referenced
this pull request
Jul 12, 2026
Closes a test gap surfaced after PR #88 merged: nothing proved that a response saved via the chunked staging protocol (POST /staging → PUT /staging/{sid}/chunks/{k} → PUT /staging/{sid}/complete) was correctly reassembled into a chain entry that the proxy could serve back as part of a subsequent turn's flat context. ## What this PR adds **`cmd/proxy/inference/mock.go`** — `MockServer` gains a `RequestBodies()` accessor that snapshots every POST /responses body in arrival order (deep-copy per element so callers may mutate without affecting each other or future snapshots). **`cmd/proxy/chunked_roundtrip_test.go`** — `TestChunkedRoundtripReassemblesIntoChain` saves an anchor turn with `maxChunkBytes=64` (forces ≥2 `AppendChunk`, pinned by `TestBufferedProxyMultipleChunks`), issues a continuation via `previous_response_id`, and verifies the continuation's inference call received the anchor's response output (`msg_ok`, role=assistant) in its input array — the strongest end-to-end proof that chunks → chain → served-context is intact. ## Existing coverage this complements | Test | Layer | Validates | |---|---|---| | `chunk_test.go` | unit | chunkedResponseWriter chunk→chunked byte round-trip | | `internal/chainstore/streaming_test.go` | server | AppendChunk/Complete/Retrieve reassembly | | `backend_routing_test.go` | proxy | wire pattern (path counts) — no byte check | | `handler_test.go::TestStoreTrueContinuation` | proxy | continuation returns 200 with different ID | | **This test** | **proxy end-to-end** | **chunked save → chain serve → inference context** | make presubmit is green.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Storecall with a staging-based chunked protocol:POST /staging→PUT /staging/{sid}/chunks/{k}(×N) →PUT /staging/{sid}/complete, orPUT /staging/{sid}/aborton empty outputchunkedResponseWriterincmd/proxy/chunk.gosplits response bytes at a configurableMaxChunkBytescap (default 1 MiB), including within a single largeAddcallstreamBufferis removedChanges
pkg/charon/client.go: extendBackendinterface withAppendChunk,Complete,Abortcmd/proxy/chunk.go:chunkedResponseWriter— byte-level splitter, mutex-protected, idempotentClose/Abortinternal/proxyconfig/proxyconfig.go:--max-chunk-bytesflag (0 → default 1 MiB)cmd/proxy/handler.go: wire writer intoHandleCreate;WithMaxChunkBytesoptioncmd/proxy/sse.go: removestreamBuffer; marshal full blob at end, write via writercmd/proxy/ws.go: same pattern assse.gocmd/proxy/main.go: passMaxChunkBytesfrom config to handlerTest Plan
chunk_test.go: single chunk, cap flush, abort-on-empty, idempotent close, cancel during flush, flush failure propagationTestBufferedCapConfigurable/TestBufferedProxyMultipleChunks: tiny cap forces ≥2AppendChunkcallsTestBufferedProxySingleChunk/TestStreamedProxySingleChunk: default cap → exactly 1 chunk + 1 completeTestBufferedProxyStoreFalseNoStagingCalls:store:falsemakes no staging callsTestStreamingInferenceFailureEmitsFailedNotCompleted: truncated stream → 0 chunks, 0 completego test ./... -raceandmake presubmitclean