package mcphttp_test import ( "bufio" "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()) } // TestUnwrapReachesTheUnderlyingFlusher is the test the Unwrap method exists // for. // // cacheWriter embeds the http.ResponseWriter *interface*, so it promotes no // Flush of its own; http.NewResponseController can only reach the real writer's // through Unwrap. Delete the method and this fails with ErrNotSupported while // every status-code and header test above stays green — which is precisely why // it is written as a test and not as a comment. func TestUnwrapReachesTheUnderlyingFlusher(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")) } // 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) }