diff --git a/cmd/proxy/chunked_roundtrip_test.go b/cmd/proxy/chunked_roundtrip_test.go new file mode 100644 index 0000000..e496ebe --- /dev/null +++ b/cmd/proxy/chunked_roundtrip_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// inputItem is the subset of fields we assert on for an item in the chained +// context. The full shape is open-ended (user messages, assistant messages, +// tool calls, ...) so we only decode what we check. +type inputItem struct { + ID string `json:"id,omitempty"` + Role string `json:"role,omitempty"` + Type string `json:"type,omitempty"` + Text json.RawMessage `json:"content,omitempty"` // string for user, []item for assistant +} + +// TestChunkedRoundtripReassemblesIntoChain closes the test gap between the +// chunked-staging wire pattern (asserted by backend_routing_test.go, in +// particular TestBufferedProxyMultipleChunks) and the downstream chain-serve +// path: a response saved via PUT /staging/{sid}/chunks/{k} → +// PUT /staging/{sid}/complete must be reassembled by Charon into a single +// entry, and that entry must be served back as part of the chained context +// for any subsequent turn that references the anchor via previous_response_id. +// +// The proxy builds the chained context in turnsToFlatCtx and prepends it to +// the current turn's input before forwarding to the inference backend (see +// buildInferenceMap in assemble.go). So the strongest end-to-end assertion is +// that the inference call for the continuation sees the anchor's response +// output (msg_ok) somewhere in its input array — proving the chunks were +// reassembled into a chain entry that the proxy can serve. +// +// Pairs with: +// - cmd/proxy/chunk_test.go: unit-level chunk → chunked byte round-trip. +// - internal/chainstore/streaming_test.go: server-side chunk reassembly. +// - cmd/proxy/backend_routing_test.go: wire-pattern pinning (no byte check). +func TestChunkedRoundtripReassemblesIntoChain(t *testing.T) { + // maxChunkBytes=64 forces ≥2 AppendChunk calls for a non-trivial stored + // blob. The wire pattern itself is pinned by TestBufferedProxyMultipleChunks + // — this test focuses on the chain-serve property. + s := newTestStack(t, withMaxChunkBytes(64)) + + // Anchor turn — store:true, multi-chunk save under the hood. + anchorHTTP := doRequest(t, s.proxyURL, "POST", "/responses", map[string]interface{}{ + "model": "test", + "input": "anchor turn input", + }) + require.Equal(t, http.StatusOK, anchorHTTP.StatusCode, + "anchor turn must succeed before we can chain off it") + anchor := decodeJSON[ResponseResource](t, anchorHTTP) + + // Continuation — references the anchor via previous_response_id. + contHTTP := doRequest(t, s.proxyURL, "POST", "/responses", map[string]interface{}{ + "model": "test", + "input": "follow up input", + "previous_response_id": anchor.ID, + }) + require.Equal(t, http.StatusOK, contHTTP.StatusCode, + "continuation must succeed") + follow := decodeJSON[ResponseResource](t, contHTTP) + assert.Equal(t, "completed", follow.Status) + + // Capture the inference backend's request bodies in arrival order: + // call 1 = anchor, call 2 = continuation. + bodies := s.mockInf.RequestBodies() + require.GreaterOrEqual(t, len(bodies), 2, + "both anchor and continuation should have hit the inference backend") + + // Decode the continuation's inference request and walk its input array. + var infReq struct { + Input json.RawMessage `json:"input"` + } + require.NoError(t, json.Unmarshal(bodies[1], &infReq)) + + var inputItems []inputItem + require.NoError(t, json.Unmarshal(infReq.Input, &inputItems), + "continuation input must be a JSON array (stringified for stateless inference)") + + // Strongest assertion: the anchor's response output (msg_ok, role=assistant) + // appears in the chained context. Proves chunks → chain → served-context. + var foundAssistant *inputItem + for i := range inputItems { + if inputItems[i].ID == "msg_ok" { + foundAssistant = &inputItems[i] + break + } + } + require.NotNil(t, foundAssistant, + "continuation's inference call must include the anchor's response output "+ + "(msg_ok) in its input — proves the chunked save was reassembled "+ + "into a chain entry the proxy can serve back") + assert.Equal(t, "assistant", foundAssistant.Role, + "anchor's msg_ok must appear as an assistant item in the chained context") + + // Sanity: at least three items — anchor input, anchor output, continuation input. + assert.GreaterOrEqual(t, len(inputItems), 3, + "chained input must contain anchor input + anchor output + continuation input") +} diff --git a/cmd/proxy/inference/mock.go b/cmd/proxy/inference/mock.go index c1ee930..4b72dbd 100644 --- a/cmd/proxy/inference/mock.go +++ b/cmd/proxy/inference/mock.go @@ -1,10 +1,13 @@ package inference import ( + "bytes" "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" + "sync" "sync/atomic" ) @@ -19,6 +22,9 @@ import ( type MockServer struct { *httptest.Server counter atomic.Int64 + + mu sync.Mutex + bodies [][]byte // captured request bodies, in arrival order } // NewMockServer starts a mock inference server. The caller must call Close(). @@ -33,6 +39,21 @@ func NewMockServer() *MockServer { // Calls returns the total number of inference requests handled. func (m *MockServer) Calls() int64 { return m.counter.Load() } +// RequestBodies returns a snapshot of every request body the mock has +// received, in arrival order. Each returned byte slice is a deep copy of +// the stored body, so callers may mutate the slices without affecting +// future snapshots or each other. Tests use this to assert what the +// proxy actually forwarded to the inference backend. +func (m *MockServer) RequestBodies() [][]byte { + m.mu.Lock() + defer m.mu.Unlock() + out := make([][]byte, len(m.bodies)) + for i, b := range m.bodies { + out[i] = append([]byte(nil), b...) + } + return out +} + // BaseURL returns the mock server's base URL (no trailing slash), satisfying // the inference.Backend interface. func (m *MockServer) BaseURL() string { return m.URL } @@ -43,11 +64,17 @@ func (m *MockServer) nextID() string { } func (m *MockServer) handle(w http.ResponseWriter, r *http.Request) { + // recordAndReplaceBody MUST run before any JSON decode of r.Body — it + // drains r.Body and re-installs a fresh reader so subsequent decoding + // observes the same bytes we captured. If a future change decodes + // r.Body directly before this call, stream=false would be observed + // for every request, regardless of what the client sent. + recordAndReplaceBody(r, m) + var req struct { Stream bool `json:"stream"` } _ = json.NewDecoder(r.Body).Decode(&req) - _ = r.Body.Close() id := m.nextID() outputItem := json.RawMessage(`{"type":"message","id":"msg_ok","role":"assistant","status":"completed","content":[{"type":"output_text","text":"OK."}]}`) @@ -60,6 +87,23 @@ func (m *MockServer) handle(w http.ResponseWriter, r *http.Request) { m.writeComplete(w, id, outputItem, usage) } +// recordAndReplaceBody reads r.Body fully, stores a copy in m.bodies, and +// re-installs a fresh reader on r.Body so subsequent decoders in the handler +// see the same bytes. On read error the body may be partial or empty — +// either way we still replace r.Body so the handler decodes a deterministic +// (possibly empty) payload rather than re-draining an unknown state. +func recordAndReplaceBody(r *http.Request, m *MockServer) { + body, _ := io.ReadAll(r.Body) + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(body)) + + cp := make([]byte, len(body)) + copy(cp, body) + m.mu.Lock() + m.bodies = append(m.bodies, cp) + m.mu.Unlock() +} + func (m *MockServer) writeComplete(w http.ResponseWriter, id string, item json.RawMessage, usage *UsageInfo) { w.Header().Set("Content-Type", "application/json") resp := Response{