Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
229 changes: 228 additions & 1 deletion client/src/lib/hooks/__tests__/useConnection.test.tsx
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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: {
Expand Down Expand Up @@ -482,6 +487,228 @@ 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: {
source: {
type: "string",
minLength: 1,
title: "A Connectable Node",
},
target: {
$ref: "#/properties/source",
},
},
},
},
};

await act(async () => {
mockSSETransport.onmessage!(mockRequestWithRef);
});

expect(mockProtocolOnMessage).toHaveBeenCalledTimes(1);

const message = mockProtocolOnMessage.mock.calls[0][0];
expect(message.params.requestedSchema.properties.target).toEqual({
type: "string",
minLength: 1,
title: "A Connectable Node",
});
});

test("resolves $ref references to $defs in requestedSchema", async () => {
const mockProtocolOnMessage = jest.fn();

mockSSETransport.onmessage = mockProtocolOnMessage;

const { result } = renderHook(() => useConnection(defaultProps));

await act(async () => {
await result.current.connect();
});

const mockRequestWithDefs: JSONRPCMessage = {
jsonrpc: "2.0",
id: 1,
method: "elicitation/create",
params: {
message: "Please provide your information",
requestedSchema: {
type: "object",
properties: {
user: {
$ref: "#/$defs/UserInput",
},
},
$defs: {
UserInput: {
type: "object",
properties: {
name: {
type: "string",
title: "Name",
},
age: {
type: "integer",
title: "Age",
minimum: 0,
},
},
required: ["name"],
},
},
},
},
};

await act(async () => {
mockSSETransport.onmessage!(mockRequestWithDefs);
});

expect(mockProtocolOnMessage).toHaveBeenCalledTimes(1);

const message = mockProtocolOnMessage.mock.calls[0][0];
// The $ref should be resolved to the actual UserInput definition
expect(message.params.requestedSchema.properties.user).toEqual({
type: "object",
properties: {
name: {
type: "string",
title: "Name",
},
age: {
type: "integer",
title: "Age",
minimum: 0,
},
},
required: ["name"],
});
});
test("resolves nested $ref references within $defs definitions", async () => {
const mockProtocolOnMessage = jest.fn();

mockSSETransport.onmessage = mockProtocolOnMessage;

const { result } = renderHook(() => useConnection(defaultProps));

await act(async () => {
await result.current.connect();
});

// This mirrors the pattern from FastMCP where a definition references another definition
const mockRequestWithNestedDefs: JSONRPCMessage = {
jsonrpc: "2.0",
id: 1,
method: "elicitation/create",
params: {
message: "Please provide your information",
requestedSchema: {
type: "object",
properties: {
person: {
$ref: "#/$defs/PersonWithAddress",
},
},
$defs: {
PersonWithAddress: {
type: "object",
properties: {
name: {
type: "string",
title: "Name",
},
address: {
$ref: "#/$defs/Address", // ← Nested ref: definition references another definition
},
},
required: ["name", "address"],
},
Address: {
type: "object",
properties: {
street: {
type: "string",
title: "Street",
},
city: {
type: "string",
title: "City",
},
zipcode: {
type: "string",
title: "Zipcode",
},
},
required: ["street", "city", "zipcode"],
},
},
},
},
};

await act(async () => {
mockSSETransport.onmessage!(mockRequestWithNestedDefs);
});

expect(mockProtocolOnMessage).toHaveBeenCalledTimes(1);

const message = mockProtocolOnMessage.mock.calls[0][0];
// The $ref should be resolved, and nested refs within $defs should also be resolved
expect(message.params.requestedSchema.properties.person).toEqual({
type: "object",
properties: {
name: {
type: "string",
title: "Name",
},
address: {
type: "object",
properties: {
street: {
type: "string",
title: "Street",
},
city: {
type: "string",
title: "City",
},
zipcode: {
type: "string",
title: "Zipcode",
},
},
required: ["street", "city", "zipcode"],
},
},
required: ["name", "address"],
});
});
});

describe("URL Port Handling", () => {
const SSEClientTransport = jest.requireMock(
"@modelcontextprotocol/sdk/client/sse.js",
Expand Down
9 changes: 9 additions & 0 deletions client/src/lib/hooks/useConnection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ import { getMCPServerRequestTimeout } from "@/utils/configUtils";
import { InspectorConfig } from "../configurationTypes";
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { CustomHeaders } from "../types/customHeaders";
import { resolveRefsInMessage } from "@/utils/schemaUtils";

interface UseConnectionOptions {
transportType: "stdio" | "sse" | "streamable-http";
Expand Down Expand Up @@ -691,6 +692,14 @@ export function useConnection({

await client.connect(transport as Transport);

const protocolOnMessage = transport.onmessage;
if (protocolOnMessage) {
transport.onmessage = (message) => {
const resolvedMessage = resolveRefsInMessage(message);
protocolOnMessage(resolvedMessage);
};
}

setClientTransport(transport);

capabilities = client.getServerCapabilities();
Expand Down
41 changes: 40 additions & 1 deletion client/src/utils/schemaUtils.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import type { JsonValue, JsonSchemaType, JsonObject } from "./jsonUtils";
import Ajv from "ajv";
import type { ValidateFunction } from "ajv";
import type { Tool } from "@modelcontextprotocol/sdk/types.js";
import type { Tool, JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
import { isJSONRPCRequest } from "@modelcontextprotocol/sdk/types.js";

const ajv = new Ajv();

Expand Down Expand Up @@ -299,3 +300,41 @@ export function formatFieldLabel(key: string): string {
.replace(/_/g, " ") // Replace underscores with spaces
.replace(/^\w/, (c) => c.toUpperCase()); // Capitalize first letter
}

/**
* Resolves `$ref` references in a JSON-RPC "elicitation/create" message's `requestedSchema` field
* @param message The JSON-RPC message that may contain $ref references
* @returns A new message with resolved $ref references, or the original message if no resolution is needed
*/
export function resolveRefsInMessage(message: JSONRPCMessage): JSONRPCMessage {
if (!isJSONRPCRequest(message) || !message.params?.requestedSchema) {
return message;
}

const requestedSchema = message.params.requestedSchema as JsonSchemaType;

if (!requestedSchema?.properties) {
return message;
}

const resolvedMessage = {
...message,
params: {
...message.params,
requestedSchema: {
...requestedSchema,
properties: Object.fromEntries(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this also resolve $refs that point to $defs or definitions? Seems like a perfect place to do so.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a test resolves $ref references to $defs in requestedSchema to verify that resolveRef already handles this. I believe it uses a generic path that resolves defs as well.

Object.entries(requestedSchema.properties).map(
([key, propSchema]) => {
const resolved = resolveRef(propSchema, requestedSchema);
const normalized = normalizeUnionType(resolved);
return [key, normalized];
},
),
),
},
},
};

return resolvedMessage;
}