From 2ab7d625e4d76b0afc477adb15eac1f5c3314142 Mon Sep 17 00:00:00 2001
From: Vincent <0426vincent@gmail.com>
Date: Sat, 1 Nov 2025 21:03:02 -0700
Subject: [PATCH] Resolve ref before validation

---
 .../hooks/__tests__/useConnection.test.tsx    | 58 ++++++++++++++++++-
 client/src/lib/hooks/useConnection.ts         | 36 ++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/client/src/lib/hooks/__tests__/useConnection.test.tsx b/client/src/lib/hooks/__tests__/useConnection.test.tsx
index a9ba6825..5c47767f 100644
--- a/client/src/lib/hooks/__tests__/useConnection.test.tsx
+++ b/client/src/lib/hooks/__tests__/useConnection.test.tsx
@@ -1,7 +1,10 @@
 import { renderHook, act } from "@testing-library/react";
 import { useConnection } from "../useConnection";
 import { z } from "zod";
-import { ClientRequest } from "@modelcontextprotocol/sdk/types.js";
+import {
+  ClientRequest,
+  JSONRPCMessage,
+} from "@modelcontextprotocol/sdk/types.js";
 import { DEFAULT_INSPECTOR_CONFIG, CLIENT_IDENTITY } from "../../constants";
 import {
   SSEClientTransportOptions,
@@ -42,10 +45,12 @@ const mockSSETransport: {
   start: jest.Mock;
   url: URL | undefined;
   options: SSEClientTransportOptions | undefined;
+  onmessage?: (message: JSONRPCMessage) => void;
 } = {
   start: jest.fn(),
   url: undefined,
   options: undefined,
+  onmessage: undefined,
 };
 
 const mockStreamableHTTPTransport: {
@@ -480,6 +485,57 @@ describe("useConnection", () => {
     });
   });
 
+  describe("Ref Resolution", () => {
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    test("resolves $ref references in requestedSchema properties before validation", async () => {
+      const mockProtocolOnMessage = jest.fn();
+
+      mockSSETransport.onmessage = mockProtocolOnMessage;
+
+      const { result } = renderHook(() => useConnection(defaultProps));
+
+      await act(async () => {
+        await result.current.connect();
+      });
+
+      const mockRequestWithRef: JSONRPCMessage = {
+        jsonrpc: "2.0",
+        id: 1,
+        method: "elicitation/create",
+        params: {
+          message: "Please provide your information",
+          requestedSchema: {
+            type: "object",
+            properties: {
+              name: {
+                $ref: "#/properties/nameDef",
+              },
+              nameDef: {
+                type: "string",
+                title: "Name",
+              },
+            },
+          },
+        },
+      };
+
+      await act(async () => {
+        mockSSETransport.onmessage!(mockRequestWithRef);
+      });
+
+      expect(mockProtocolOnMessage).toHaveBeenCalledTimes(1);
+
+      const message = mockProtocolOnMessage.mock.calls[0][0];
+      expect(message.params.requestedSchema.properties.name).toEqual({
+        type: "string",
+        title: "Name",
+      });
+    });
+  });
+
   describe("URL Port Handling", () => {
     const SSEClientTransport = jest.requireMock(
       "@modelcontextprotocol/sdk/client/sse.js",
diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index 8639feeb..486ba51b 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -30,6 +30,8 @@ import {
   Progress,
   LoggingLevel,
   ElicitRequestSchema,
+  JSONRPCMessage,
+  isJSONRPCRequest,
 } from "@modelcontextprotocol/sdk/types.js";
 import { RequestOptions } from "@modelcontextprotocol/sdk/shared/protocol.js";
 import { useEffect, useState } from "react";
@@ -57,6 +59,8 @@ import { getMCPServerRequestTimeout } from "@/utils/configUtils";
 import { InspectorConfig } from "../configurationTypes";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import { CustomHeaders } from "../types/customHeaders";
+import { JsonSchemaType } from "@/utils/jsonUtils";
+import { resolveRef } from "@/utils/schemaUtils";
 
 interface UseConnectionOptions {
   transportType: "stdio" | "sse" | "streamable-http";
@@ -681,6 +685,38 @@ export function useConnection({
 
         await client.connect(transport as Transport);
 
+        const protocolOnMessage = transport.onmessage;
+        if (protocolOnMessage) {
+          transport.onmessage = (message: JSONRPCMessage) => {
+            // Resolve $ref references in requests before validation
+            if (isJSONRPCRequest(message) && message.params?.requestedSchema) {
+              const requestedSchema = message.params
+                .requestedSchema as JsonSchemaType;
+              if (requestedSchema?.properties) {
+                const resolvedProperties = Object.fromEntries(
+                  Object.entries(
+                    requestedSchema.properties as Record<
+                      string,
+                      JsonSchemaType
+                    >,
+                  ).map(([key, propSchema]) => [
+                    key,
+                    resolveRef(propSchema, requestedSchema),
+                  ]),
+                );
+                message.params = {
+                  ...message.params,
+                  requestedSchema: {
+                    ...requestedSchema,
+                    properties: resolvedProperties,
+                  },
+                };
+              }
+            }
+            protocolOnMessage(message);
+          };
+        }
+
         setClientTransport(transport);
 
         capabilities = client.getServerCapabilities();