Skip to content

Commit 37b99b3

Browse files
committed
Resolve ref before validation
1 parent e87623c commit 37b99b3

File tree

2 files changed

+93
-1
lines changed

2 files changed

+93
-1
lines changed

client/src/lib/hooks/__tests__/useConnection.test.tsx

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import { renderHook, act } from "@testing-library/react";
22
import { useConnection } from "../useConnection";
33
import { z } from "zod";
4-
import { ClientRequest } from "@modelcontextprotocol/sdk/types.js";
4+
import {
5+
ClientRequest,
6+
JSONRPCMessage,
7+
} from "@modelcontextprotocol/sdk/types.js";
58
import { DEFAULT_INSPECTOR_CONFIG, CLIENT_IDENTITY } from "../../constants";
69
import {
710
SSEClientTransportOptions,
@@ -42,10 +45,12 @@ const mockSSETransport: {
4245
start: jest.Mock;
4346
url: URL | undefined;
4447
options: SSEClientTransportOptions | undefined;
48+
onmessage?: (message: JSONRPCMessage) => void;
4549
} = {
4650
start: jest.fn(),
4751
url: undefined,
4852
options: undefined,
53+
onmessage: undefined,
4954
};
5055

5156
const mockStreamableHTTPTransport: {
@@ -480,6 +485,57 @@ describe("useConnection", () => {
480485
});
481486
});
482487

488+
describe("Ref Resolution", () => {
489+
beforeEach(() => {
490+
jest.clearAllMocks();
491+
});
492+
493+
test("resolves $ref references in requestedSchema properties before validation", async () => {
494+
const mockProtocolOnMessage = jest.fn();
495+
496+
mockSSETransport.onmessage = mockProtocolOnMessage;
497+
498+
const { result } = renderHook(() => useConnection(defaultProps));
499+
500+
await act(async () => {
501+
await result.current.connect();
502+
});
503+
504+
const mockRequestWithRef: JSONRPCMessage = {
505+
jsonrpc: "2.0",
506+
id: 1,
507+
method: "elicitation/create",
508+
params: {
509+
message: "Please provide your information",
510+
requestedSchema: {
511+
type: "object",
512+
properties: {
513+
name: {
514+
$ref: "#/properties/nameDef",
515+
},
516+
nameDef: {
517+
type: "string",
518+
title: "Name",
519+
},
520+
},
521+
},
522+
},
523+
};
524+
525+
await act(async () => {
526+
mockSSETransport.onmessage!(mockRequestWithRef);
527+
});
528+
529+
expect(mockProtocolOnMessage).toHaveBeenCalledTimes(1);
530+
531+
const message = mockProtocolOnMessage.mock.calls[0][0];
532+
expect(message.params.requestedSchema.properties.name).toEqual({
533+
type: "string",
534+
title: "Name",
535+
});
536+
});
537+
});
538+
483539
describe("URL Port Handling", () => {
484540
const SSEClientTransport = jest.requireMock(
485541
"@modelcontextprotocol/sdk/client/sse.js",

client/src/lib/hooks/useConnection.ts

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ import {
3030
Progress,
3131
LoggingLevel,
3232
ElicitRequestSchema,
33+
JSONRPCMessage,
34+
isJSONRPCRequest,
3335
} from "@modelcontextprotocol/sdk/types.js";
3436
import { RequestOptions } from "@modelcontextprotocol/sdk/shared/protocol.js";
3537
import { useEffect, useState } from "react";
@@ -57,6 +59,8 @@ import { getMCPServerRequestTimeout } from "@/utils/configUtils";
5759
import { InspectorConfig } from "../configurationTypes";
5860
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
5961
import { CustomHeaders } from "../types/customHeaders";
62+
import { JsonSchemaType } from "@/utils/jsonUtils";
63+
import { resolveRef } from "@/utils/schemaUtils";
6064

6165
interface UseConnectionOptions {
6266
transportType: "stdio" | "sse" | "streamable-http";
@@ -681,6 +685,38 @@ export function useConnection({
681685

682686
await client.connect(transport as Transport);
683687

688+
const protocolOnMessage = transport.onmessage;
689+
if (protocolOnMessage) {
690+
transport.onmessage = (message: JSONRPCMessage) => {
691+
// Resolve $ref references in requests before validation
692+
if (isJSONRPCRequest(message) && message.params?.requestedSchema) {
693+
const requestedSchema = message.params
694+
.requestedSchema as JsonSchemaType;
695+
if (requestedSchema?.properties) {
696+
const resolvedProperties = Object.fromEntries(
697+
Object.entries(
698+
requestedSchema.properties as Record<
699+
string,
700+
JsonSchemaType
701+
>,
702+
).map(([key, propSchema]) => [
703+
key,
704+
resolveRef(propSchema, requestedSchema),
705+
]),
706+
);
707+
message.params = {
708+
...message.params,
709+
requestedSchema: {
710+
...requestedSchema,
711+
properties: resolvedProperties,
712+
},
713+
};
714+
}
715+
}
716+
protocolOnMessage(message);
717+
};
718+
}
719+
684720
setClientTransport(transport);
685721

686722
capabilities = client.getServerCapabilities();

0 commit comments

Comments
 (0)