Echo the request JSON-RPC id in session-not-found 404s

pkg/transport/proxy/transparent/backend_routing_test.go

@@ -5,6 +5,7 @@ package transparent
 
 import (
 	"context"
+	"encoding/json"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -196,6 +197,89 @@ func TestRoundTripReturns404ForUnknownSession(t *testing.T) {
 	assert.Contains(t, string(body), `"code":-32001`)
 }
 
+// TestRoundTrip404EchoesRequestID pins #5945: the unknown-session 404 must echo
+// the incoming JSON-RPC id so a client correlating responses by id can match the
+// error to the request that caused it. Before the fix the transparent proxy
+// hardcoded a null id here, while the sibling classification-error path on the
+// same RoundTrip already echoed it.
+//
+// Note TestRoundTripReturns404ForUnknownSession above cannot catch this: its body
+// carries no "id", so it renders a null id either way.
+func TestRoundTrip404EchoesRequestID(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name   string
+		body   string
+		wantID string
+	}{
+		{
+			name:   "numeric id is echoed verbatim",
+			body:   `{"jsonrpc":"2.0","id":7,"method":"tools/list"}`,
+			wantID: `"id":7`,
+		},
+		{
+			name:   "string id is echoed verbatim",
+			body:   `{"jsonrpc":"2.0","id":"abc-123","method":"tools/list"}`,
+			wantID: `"id":"abc-123"`,
+		},
+		{
+			// A notification has no id by definition, so JSON-RPC requires null.
+			name:   "notification renders a null id",
+			body:   `{"jsonrpc":"2.0","method":"notifications/initialized"}`,
+			wantID: `"id":null`,
+		},
+		{
+			// An explicit null id is not a correlatable id either.
+			name:   "explicit null id stays null",
+			body:   `{"jsonrpc":"2.0","id":null,"method":"tools/list"}`,
+			wantID: `"id":null`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				// Reaching the backend would mean the guard did not fire; the
+				// assertions below would fail on the 200 that results.
+				w.WriteHeader(http.StatusOK)
+			}))
+			defer backend.Close()
+
+			tt2 := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions(
+				"localhost", 0, backend.URL,
+				nil, nil, nil,
+				false, false, "sse",
+				nil, nil, "", false,
+				nil,
+			))
+
+			req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp", strings.NewReader(tt.body))
+			require.NoError(t, err)
+			req.Header.Set("Content-Type", "application/json")
+			req.Header.Set("Mcp-Session-Id", uuid.New().String()) // unknown to the store
+
+			resp, err := tt2.RoundTrip(req)
+			require.NoError(t, err)
+			require.Equal(t, http.StatusNotFound, resp.StatusCode)
+			body, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			_ = resp.Body.Close()
+
+			assert.Contains(t, string(body), `"code":-32001`)
+			assert.Contains(t, string(body), tt.wantID)
+
+			// The body must remain a single valid JSON-RPC error object -- echoing
+			// a raw id must not corrupt the envelope.
+			var decoded map[string]any
+			require.NoError(t, json.Unmarshal(body, &decoded))
+			assert.Equal(t, "2.0", decoded["jsonrpc"])
+		})
+	}
+}
+
 // TestRoundTripAllowsInitializeWithUnknownSession verifies that an initialize
 // call is forwarded even when its Mcp-Session-Id is not yet in the session store.
 func TestRoundTripAllowsInitializeWithUnknownSession(t *testing.T) {

pkg/transport/proxy/transparent/transparent_proxy.go

@@ -598,13 +598,20 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 	isJSON := strings.Contains(req.Header.Get("Content-Type"), "application/json")
 	sawInitialize := false
 	revision := mcp.RevisionLegacy
+	// requestID carries the incoming JSON-RPC id so an error response built below
+	// can echo it, letting a client correlate the error with its request. It is set
+	// only when parseRPCRequest reports a single request, which by construction
+	// means the id is present and non-null; it stays nil for a notification, a
+	// batch, or a bodiless GET/DELETE, where JSON-RPC requires a null id.
+	var requestID any
 
 	if len(reqBody) > 0 &&
 		((isMCP && isJSON) ||
 			t.p.transportType == types.TransportTypeStreamableHTTP.String()) {
 		method, params, id, singleRequest, isInit := t.parseRPCRequest(reqBody)
 		sawInitialize = isInit
 		if singleRequest {
+			requestID = id
 			meta := mcp.ExtractMeta(params)
 			rev, cerr := mcp.ClassifyRevision(method, meta, req.Header.Get("MCP-Protocol-Version"))
 			if cerr != nil {
@@ -628,7 +635,7 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 				slog.Error("session store lookup failed", "error", err)
 				return plainResponse(req, http.StatusServiceUnavailable, "session store unavailable"), nil
 			}
-			return session.NotFoundResponse(req), nil
+			return session.NotFoundResponse(req, requestID), nil
 		}
 	}
 

pkg/transport/session/jsonrpc_errors.go

@@ -24,8 +24,13 @@ const (
 
 // NotFoundBody returns the JSON-encoded body for a session-not-found
 // JSON-RPC error response. The requestID is the "id" from the incoming
-// JSON-RPC request; pass nil when the request ID is not available (e.g.,
-// DELETE requests, batch pre-parse, transparent proxy).
+// JSON-RPC request, echoed so a client correlating by id can match this error
+// to the request that caused it.
+//
+// Pass nil only when the request genuinely carries no id, in which case
+// JSON-RPC requires a null id: a bodiless GET (standalone SSE) or DELETE, a
+// notification, or a batch. Do NOT pass nil merely because threading the id to
+// the call site is inconvenient — that was the asymmetry fixed in #5945.
 func NotFoundBody(requestID any) []byte {
 	resp := map[string]any{
 		"jsonrpc": "2.0",
@@ -57,8 +62,12 @@ func WriteNotFound(w http.ResponseWriter, requestID any) {
 // NotFoundResponse constructs an *http.Response with HTTP 404 and a
 // JSON-RPC error body. Use this in httputil.ReverseProxy.ModifyResponse
 // (transparent proxy) where no http.ResponseWriter is available.
-func NotFoundResponse(req *http.Request) *http.Response {
-	body := NotFoundBody(nil)
+//
+// requestID follows NotFoundBody: pass the incoming request's JSON-RPC id so the
+// error echoes it, or nil when the request has none. Callers that reject a
+// request before parsing a body (GET/DELETE) pass nil.
+func NotFoundResponse(req *http.Request, requestID any) *http.Response {
+	body := NotFoundBody(requestID)
 	hdr := make(http.Header)
 	hdr.Set("Content-Type", "application/json")
 	return &http.Response{

pkg/transport/session/jsonrpc_errors_test.go

@@ -84,15 +84,56 @@ func TestWriteNotFound(t *testing.T) {
 func TestNotFoundResponse(t *testing.T) {
 	t.Parallel()
 
-	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
-	resp := NotFoundResponse(req)
+	tests := []struct {
+		name      string
+		requestID any
+		wantID    string
+	}{
+		{
+			// A request with no id of its own: bodiless GET/DELETE, or a
+			// notification. JSON-RPC requires a null id here.
+			name:      "nil id renders null",
+			requestID: nil,
+			wantID:    `"id":null`,
+		},
+		{
+			name:      "string id is echoed",
+			requestID: "req-1",
+			wantID:    `"id":"req-1"`,
+		},
+		{
+			// The shape the transparent proxy passes: the raw bytes of the
+			// incoming "id", forwarded verbatim so a numeric id stays numeric
+			// rather than being coerced to a float or a string (#5945).
+			name:      "raw numeric id is echoed verbatim",
+			requestID: json.RawMessage(`42`),
+			wantID:    `"id":42`,
+		},
+		{
+			name:      "raw string id is echoed verbatim",
+			requestID: json.RawMessage(`"abc"`),
+			wantID:    `"id":"abc"`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
 
-	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
-	assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
-	assert.Equal(t, req, resp.Request)
+			req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
+			resp := NotFoundResponse(req, tt.requestID)
 
-	body, err := io.ReadAll(resp.Body)
-	require.NoError(t, err)
-	assert.Contains(t, string(body), `"code":-32001`)
-	assert.Contains(t, string(body), `"id":null`)
+			assert.Equal(t, http.StatusNotFound, resp.StatusCode)
+			assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
+			assert.Equal(t, req, resp.Request)
+
+			body, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			assert.Contains(t, string(body), `"code":-32001`)
+			assert.Contains(t, string(body), tt.wantID)
+			// ContentLength must match the body actually written, or a client
+			// reading exactly that many bytes truncates or blocks.
+			assert.Equal(t, int64(len(body)), resp.ContentLength)
+		})
+	}
 }

test/e2e/hostile_input_proxy_test.go

@@ -101,11 +101,9 @@ var _ = Describe("Hostile Input Proxy", Label("proxy", "stateless", "hostile-inp
 			resp, err := client.Send(context.Background(), proxyURL, req)
 			Expect(err).ToNot(HaveOccurred(), tc.name)
 			Expect(resp.StatusCode).To(Equal(tc.wantStatus), tc.name)
-			if tc.idNotEchoed {
-				Expect(resp.ID).To(BeNil(), "%s: echoed id", tc.name)
-			} else {
-				Expect(resp.ID).To(Equal(json.Number(fmt.Sprintf("%d", tc.id))), "%s: echoed id", tc.name)
-			}
+			// Every rejection path here echoes the id, session-not-found
+			// included: #5945 removed the one asymmetry (see below).
+			Expect(resp.ID).To(Equal(json.Number(fmt.Sprintf("%d", tc.id))), "%s: echoed id", tc.name)
 			Expect(resp.Error).ToNot(BeNil(), tc.name)
 			Expect(resp.Error.Code).To(Equal(tc.wantCode), tc.name)
 			if tc.wantData != nil {
@@ -188,12 +186,6 @@ type hostileCase struct {
 	wantStatus int
 	wantCode   int64
 	wantData   map[string]any // nil means the error must carry no "data" field at all
-	// idNotEchoed marks a rejection path that -- unlike ClassificationErrorResponse
-	// -- doesn't echo the request id at all: session.NotFoundResponse hardcodes a
-	// nil id (see jsonrpc_errors.go's NotFoundBody(nil) call). A real minor
-	// asymmetry in the current code, left flagged rather than fixed (out of scope
-	// for a test-only step).
-	idNotEchoed bool
 }
 
 // hostileRejectionCases enumerates the classification-error and session-guard
@@ -273,10 +265,14 @@ func hostileRejectionCases() []hostileCase {
 				// in revision_guard_regression_test.go.
 				req.WithSessionID("foreign-" + uuid.NewString())
 			},
-			wantStatus:  http.StatusNotFound,
-			wantCode:    session.CodeSessionNotFound,
-			wantData:    nil,
-			idNotEchoed: true,
+			// Like the classification rejections above, this 404 echoes the
+			// request id: session.NotFoundResponse used to hardcode a nil id
+			// while ClassificationErrorResponse on the same RoundTrip echoed
+			// it, an asymmetry #5945 removed by threading the parsed id through
+			// to the guard.
+			wantStatus: http.StatusNotFound,
+			wantCode:   session.CodeSessionNotFound,
+			wantData:   nil,
 		},
 	}
 }