package mcphttp_test import ( "bufio" "io" "net/http" "net/http/httptest" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "sourcecraft.dev/bigbes/sr-ht-ecore/mcphttp" ) const ( wantCacheControl = "private, no-store, no-transform" wantVary = "Cookie, Authorization" ) // TestHeadersLandOnAnExplicitWriteHeader is the ordinary path: a handler that // commits with WriteHeader. func TestHeadersLandOnAnExplicitWriteHeader(t *testing.T) { h := mcphttp.PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusAccepted) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil)) assert.Equal(t, http.StatusAccepted, rec.Code) assert.Equal(t, wantCacheControl, rec.Header().Get("Cache-Control")) assert.Equal(t, wantVary, rec.Header().Get("Vary")) } // TestHeadersLandOnAnImplicitCommit covers the handler that never calls // WriteHeader at all. net/http commits on the first Write, and headers set after // that point are dropped silently — so the wrapper has to catch Write too. func TestHeadersLandOnAnImplicitCommit(t *testing.T) { h := mcphttp.PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(`{"jsonrpc":"2.0"}`)) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil)) assert.Equal(t, http.StatusOK, rec.Code) assert.Equal(t, wantCacheControl, rec.Header().Get("Cache-Control")) assert.Equal(t, wantVary, rec.Header().Get("Vary")) assert.Equal(t, `{"jsonrpc":"2.0"}`, rec.Body.String()) } // TestTheSDKsOwnDirectivesAreOverridden is the reason this wrapper exists rather // than middleware.PrivateCache. The streamable transport sets Cache-Control with // Set from inside the handler, so anything written on the way in loses. What // arrives here is a handler doing exactly what the SDK does. func TestTheSDKsOwnDirectivesAreOverridden(t *testing.T) { h := mcphttp.PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Cache-Control", "no-cache, no-transform") w.WriteHeader(http.StatusOK) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil)) assert.Equal(t, wantCacheControl, rec.Header().Get("Cache-Control"), "no-cache permits a cache to store the body and revalidate, which is what no-store forbids") assert.Equal(t, wantVary, rec.Header().Get("Vary")) } // TestOnlyTheFirstCommitWrites checks the committed flag: a handler that writes // its own header after committing is doing something net/http would ignore // anyway, and the wrapper must not undo a status that is already on the wire by // re-running its Set on every Write. func TestOnlyTheFirstCommitWrites(t *testing.T) { h := mcphttp.PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("a")) _, _ = w.Write([]byte("b")) })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", nil)) assert.Equal(t, []string{wantCacheControl}, rec.Header().Values("Cache-Control")) assert.Equal(t, []string{wantVary}, rec.Header().Values("Vary")) assert.Equal(t, "ab", rec.Body.String()) } // TestAFlushBeforeAnyWriteStillCarriesTheHeaders is what the Flush method earns. // // A handler that flushes before writing anything commits the response through // the flush, not through Write — so the Write hook never runs, and without a // Flush of its own the wrapper is bypassed entirely and the response leaves with // no Cache-Control and no Vary. Every other test here writes first, which sets // the headers on the way in and leaves the flush path unmeasured: deleting Flush // kept all of them green. // // It has to run against a real server. httptest.ResponseRecorder's Header() // hands back the live map, so a recorder reports the headers as set no matter // when — or whether — the wrapper committed them, and reports a pass for exactly // this bug. func TestAFlushBeforeAnyWriteStillCarriesTheHeaders(t *testing.T) { release := make(chan struct{}) srv := httptest.NewServer(mcphttp.PrivateCache(http.HandlerFunc( func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") if err := http.NewResponseController(w).Flush(); err != nil { t.Errorf("flush: %v", err) } <-release _, _ = w.Write([]byte(": keepalive\n\n")) }))) t.Cleanup(func() { close(release) srv.Close() }) resp, err := (&http.Client{Timeout: 5 * time.Second}).Get(srv.URL) require.NoError(t, err) t.Cleanup(func() { _ = resp.Body.Close() }) assert.Equal(t, wantCacheControl, resp.Header.Get("Cache-Control"), "the response committed through Flush, so only Flush could have set this") assert.Equal(t, wantVary, resp.Header.Get("Vary")) } // TestFlushOnTheWrapperStreams covers the flush path, which cacheWriter answers // itself rather than delegating. // // It is deliberately NOT named for Unwrap. An earlier version of this test was // named that and claimed to be what the Unwrap method exists for — and it passed // with Unwrap deleted, because cacheWriter has a Flush method and // http.NewResponseController prefers a method on the writer it is handed over // one reached by unwrapping. It never got that far. See // TestUnwrapReachesTheWriterBelow for the property that does need Unwrap. func TestFlushOnTheWrapperStreams(t *testing.T) { var flushErr error h := mcphttp.PrivateCache(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("event: message\n")) flushErr = http.NewResponseController(w).Flush() })) rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/mcp", nil)) require.NoError(t, flushErr, "an SSE stream that cannot be flushed is a response no client sees") assert.True(t, rec.Flushed) assert.Equal(t, wantCacheControl, rec.Header().Get("Cache-Control")) } // TestUnwrapReachesTheWriterBelow is the test the Unwrap method actually earns. // // cacheWriter embeds the http.ResponseWriter *interface*, so it promotes nothing // but the four methods that interface declares. Flush it answers itself; every // other thing http.ResponseController offers — the deadlines, Hijack — can only // be reached by unwrapping. A long-lived MCP stream is exactly the response that // wants a write deadline pushed out, so this is not a hypothetical. // // It has to run against a real server: httptest.ResponseRecorder supports no // deadlines at all, so a recorder would report ErrNotSupported whether Unwrap // were there or not — the same vacuity the old test had. func TestUnwrapReachesTheWriterBelow(t *testing.T) { var deadlineErr error srv := httptest.NewServer(mcphttp.PrivateCache(http.HandlerFunc( func(w http.ResponseWriter, _ *http.Request) { deadlineErr = http.NewResponseController(w).SetWriteDeadline(time.Now().Add(time.Minute)) _, _ = w.Write([]byte("ok")) }))) t.Cleanup(srv.Close) resp, err := (&http.Client{Timeout: 5 * time.Second}).Get(srv.URL) require.NoError(t, err) t.Cleanup(func() { _ = resp.Body.Close() }) _, _ = io.Copy(io.Discard, resp.Body) require.NoError(t, deadlineErr, "without Unwrap the controller cannot reach the real writer and this is ErrNotSupported") } // TestAStreamReachesTheClientBeforeTheHandlerReturns is the same property // measured rather than asserted: over a real connection, with a real client, // bytes written and flushed inside the handler have to arrive while the handler // is still running. A recorder cannot tell a flush that worked from one that was // buffered until the end. func TestAStreamReachesTheClientBeforeTheHandlerReturns(t *testing.T) { release := make(chan struct{}) srv := httptest.NewServer(mcphttp.PrivateCache(http.HandlerFunc( func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: first\n\n")) if err := http.NewResponseController(w).Flush(); err != nil { t.Errorf("flush: %v", err) } <-release }))) t.Cleanup(func() { close(release) srv.Close() }) req, err := http.NewRequest(http.MethodGet, srv.URL, nil) require.NoError(t, err) resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) require.NoError(t, err) t.Cleanup(func() { _ = resp.Body.Close() }) assert.Equal(t, wantCacheControl, resp.Header.Get("Cache-Control")) assert.Equal(t, wantVary, resp.Header.Get("Vary")) line, err := bufio.NewReader(resp.Body).ReadString('\n') require.NoError(t, err, "the first event must arrive while the handler is still blocked") assert.Equal(t, "data: first\n", line) }