Skip to content

Commit f4e1a38

Browse files
dclark27dclark27
committed
chore: working tests
1 parent d39b4bd commit f4e1a38

File tree

12 files changed

+167
-75
lines changed

12 files changed

+167
-75
lines changed

‎src/client/auth.ts‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import {
2727
UnauthorizedClientError
2828
} from '../server/auth/errors.js';
2929
import { FetchLike } from '../shared/transport.js';
30+
import type { AnySchema, SchemaOutput } from '../server/zod-compat.js';
31+
import { safeParse } from '../server/zod-compat.js';
3032

3133
/**
3234
* Implements an end-to-end OAuth client to be used with one MCP server.
@@ -271,7 +273,8 @@ export async function parseErrorResponse(input: Response | string): Promise<OAut
271273
const body = input instanceof Response ? await input.text() : input;
272274

273275
try {
274-
const result = OAuthErrorResponseSchema.parse(JSON.parse(body));
276+
const parsedJson = JSON.parse(body);
277+
const result = parseSchemaOrThrow(OAuthErrorResponseSchema, parsedJson);
275278
const { error, error_description, error_uri } = result;
276279
const errorClass = OAUTH_ERRORS[error] || ServerError;
277280
return new errorClass(error_description || '', error_uri);
@@ -515,7 +518,7 @@ export async function discoverOAuthProtectedResourceMetadata(
515518
if (!response.ok) {
516519
throw new Error(`HTTP ${response.status} trying to load well-known OAuth protected resource metadata.`);
517520
}
518-
return OAuthProtectedResourceMetadataSchema.parse(await response.json());
521+
return parseSchemaOrThrow(OAuthProtectedResourceMetadataSchema, await response.json());
519522
}
520523

521524
/**
@@ -647,7 +650,7 @@ export async function discoverOAuthMetadata(
647650
throw new Error(`HTTP ${response.status} trying to load well-known OAuth metadata`);
648651
}
649652

650-
return OAuthMetadataSchema.parse(await response.json());
653+
return parseSchemaOrThrow(OAuthMetadataSchema, await response.json());
651654
}
652655

653656
/**
@@ -770,9 +773,9 @@ export async function discoverAuthorizationServerMetadata(
770773

771774
// Parse and validate based on type
772775
if (type === 'oauth') {
773-
return OAuthMetadataSchema.parse(await response.json());
776+
return parseSchemaOrThrow(OAuthMetadataSchema, await response.json());
774777
} else {
775-
return OpenIdProviderDiscoveryMetadataSchema.parse(await response.json());
778+
return parseSchemaOrThrow(OpenIdProviderDiscoveryMetadataSchema, await response.json());
776779
}
777780
}
778781

@@ -929,7 +932,7 @@ export async function exchangeAuthorization(
929932
throw await parseErrorResponse(response);
930933
}
931934

932-
return OAuthTokensSchema.parse(await response.json());
935+
return parseSchemaOrThrow(OAuthTokensSchema, await response.json());
933936
}
934937

935938
/**
@@ -1007,7 +1010,7 @@ export async function refreshAuthorization(
10071010
throw await parseErrorResponse(response);
10081011
}
10091012

1010-
return OAuthTokensSchema.parse({ refresh_token: refreshToken, ...(await response.json()) });
1013+
return parseSchemaOrThrow(OAuthTokensSchema, { refresh_token: refreshToken, ...(await response.json()) });
10111014
}
10121015

10131016
/**
@@ -1049,5 +1052,14 @@ export async function registerClient(
10491052
throw await parseErrorResponse(response);
10501053
}
10511054

1052-
return OAuthClientInformationFullSchema.parse(await response.json());
1055+
return parseSchemaOrThrow(OAuthClientInformationFullSchema, await response.json());
1056+
}
1057+
1058+
function parseSchemaOrThrow<S extends AnySchema>(schema: S, value: unknown): SchemaOutput<S> {
1059+
const result = safeParse(schema, value);
1060+
if (!result.success) {
1061+
throw result.error instanceof Error ? result.error : new Error(String(result.error));
1062+
}
1063+
1064+
return result.data;
10531065
}

‎src/client/index.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
/* eslint-disable no-constant-binary-expression */
33
/* eslint-disable @typescript-eslint/no-unused-expressions */
44
import { Client } from './index.js';
5-
import { z } from 'zod';
5+
import * as z from 'zod/v4';
66
import {
77
RequestSchema,
88
NotificationSchema,

‎src/server/completable.test.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import { z } from 'zod';
2-
import { completable } from './completable.js';
1+
import * as z from 'zod/v3';
2+
import { completable, getCompleter } from './completable.js';
33

44
describe('completable', () => {
55
it('preserves types and values of underlying schema', () => {
@@ -14,27 +14,27 @@ describe('completable', () => {
1414
const completions = ['foo', 'bar', 'baz'];
1515
const schema = completable(z.string(), () => completions);
1616

17-
expect(await schema._def.complete('')).toEqual(completions);
17+
expect(await getCompleter(schema)!('')).toEqual(completions);
1818
});
1919

2020
it('allows async completion functions', async () => {
2121
const completions = ['foo', 'bar', 'baz'];
2222
const schema = completable(z.string(), async () => completions);
2323

24-
expect(await schema._def.complete('')).toEqual(completions);
24+
expect(await getCompleter(schema)!('')).toEqual(completions);
2525
});
2626

2727
it('passes current value to completion function', async () => {
2828
const schema = completable(z.string(), value => [value + '!']);
2929

30-
expect(await schema._def.complete('test')).toEqual(['test!']);
30+
expect(await getCompleter(schema)!('test')).toEqual(['test!']);
3131
});
3232

3333
it('works with number schemas', async () => {
3434
const schema = completable(z.number(), () => [1, 2, 3]);
3535

3636
expect(schema.parse(1)).toBe(1);
37-
expect(await schema._def.complete(0)).toEqual([1, 2, 3]);
37+
expect(await getCompleter(schema)!(0)).toEqual([1, 2, 3]);
3838
});
3939

4040
it('preserves schema description', () => {

‎src/server/index.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
/* eslint-disable no-constant-binary-expression */
33
/* eslint-disable @typescript-eslint/no-unused-expressions */
44
import { Server } from './index.js';
5-
import { z } from 'zod';
5+
import * as z from 'zod/v4';
66
import {
77
RequestSchema,
88
NotificationSchema,

‎src/server/mcp.test.ts‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import { McpServer } from './mcp.js';
22
import { Client } from '../client/index.js';
33
import { InMemoryTransport } from '../inMemory.js';
4-
import { z } from 'zod';
4+
import * as z from 'zod/v4';
5+
import * as zMini from 'zod/v4-mini';
56
import {
67
ListToolsResultSchema,
78
CallToolResultSchema,
@@ -522,8 +523,8 @@ describe('tool()', () => {
522523
mcpServer.tool(
523524
'test',
524525
{
525-
name: z.string(),
526-
value: z.number()
526+
name: zMini.string(),
527+
value: zMini.number()
527528
},
528529
async ({ name, value }) => ({
529530
content: [
@@ -695,14 +696,14 @@ describe('tool()', () => {
695696
version: '1.0'
696697
});
697698

698-
mcpServer.tool('test', { name: z.string() }, { title: 'Test Tool', readOnlyHint: true }, async ({ name }) => ({
699+
mcpServer.tool('test', { name: zMini.string() }, { title: 'Test Tool', readOnlyHint: true }, async ({ name }) => ({
699700
content: [{ type: 'text', text: `Hello, ${name}!` }]
700701
}));
701702

702703
mcpServer.registerTool(
703704
'test (new api)',
704705
{
705-
inputSchema: { name: z.string() },
706+
inputSchema: { name: zMini.string() },
706707
annotations: { title: 'Test Tool', readOnlyHint: true }
707708
},
708709
async ({ name }) => ({
@@ -879,8 +880,8 @@ describe('tool()', () => {
879880
'test (new api)',
880881
{
881882
inputSchema: {
882-
name: z.string(),
883-
value: z.number()
883+
name: zMini.string(),
884+
value: zMini.number()
884885
}
885886
},
886887
async ({ name, value }) => ({
@@ -3081,7 +3082,7 @@ describe('prompt()', () => {
30813082
);
30823083

30833084
// Register a prompt with completion
3084-
mcpServer.prompt('echo', { message: completable(z.string(), () => ['hello', 'world']) }, ({ message }) => ({
3085+
mcpServer.prompt('echo', { message: completable(zMini.string(), () => ['hello', 'world']) }, ({ message }) => ({
30853086
messages: [
30863087
{
30873088
role: 'user',
@@ -3153,7 +3154,7 @@ describe('prompt()', () => {
31533154
mcpServer.prompt(
31543155
'test-prompt',
31553156
{
3156-
name: completable(z.string(), () => ['Alice', 'Bob', 'Charlie'])
3157+
name: completable(zMini.string(), () => ['Alice', 'Bob', 'Charlie'])
31573158
},
31583159
async ({ name }) => ({
31593160
messages: [
@@ -3192,7 +3193,7 @@ describe('prompt()', () => {
31923193
mcpServer.prompt(
31933194
'test-prompt',
31943195
{
3195-
name: completable(z.string(), () => ['Alice', 'Bob', 'Charlie'])
3196+
name: completable(zMini.string(), () => ['Alice', 'Bob', 'Charlie'])
31963197
},
31973198
async ({ name }) => ({
31983199
messages: [
@@ -3249,7 +3250,7 @@ describe('prompt()', () => {
32493250
mcpServer.prompt(
32503251
'test-prompt',
32513252
{
3252-
name: completable(z.string(), test => ['Alice', 'Bob', 'Charlie'].filter(value => value.startsWith(test)))
3253+
name: completable(zMini.string(), test => ['Alice', 'Bob', 'Charlie'].filter(value => value.startsWith(test)))
32533254
},
32543255
async ({ name }) => ({
32553256
messages: [
@@ -3741,10 +3742,10 @@ describe('Tool title precedence', () => {
37413742
title: 'Team Greeting',
37423743
description: 'Generate a greeting for team members',
37433744
argsSchema: {
3744-
department: completable(z.string(), value => {
3745+
department: completable(zMini.string(), value => {
37453746
return ['engineering', 'sales', 'marketing', 'support'].filter(d => d.startsWith(value));
37463747
}),
3747-
name: completable(z.string(), (value, context) => {
3748+
name: completable(zMini.string(), (value, context) => {
37483749
const department = context?.arguments?.['department'];
37493750
if (department === 'engineering') {
37503751
return ['Alice', 'Bob', 'Charlie'].filter(n => n.startsWith(value));

‎src/server/mcp.ts‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -268,9 +268,8 @@ export class McpServer {
268268
return EMPTY_COMPLETION_RESULT;
269269
}
270270

271-
// Access field schema by name; works for v3 and v4 Mini objects
272-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
273-
const field = (prompt.argsSchema as any).shape?.[request.params.argument.name];
271+
const promptShape = getObjectShape(prompt.argsSchema);
272+
const field = promptShape?.[request.params.argument.name];
274273
if (!field) {
275274
return EMPTY_COMPLETION_RESULT;
276275
}
@@ -1135,6 +1134,25 @@ function isZodTypeLike(value: unknown): value is AnySchema {
11351134
return false;
11361135
}
11371136

1137+
function getObjectShape(schema: AnyObjectSchema | undefined): Record<string, AnySchema> | undefined {
1138+
if (!schema) return undefined;
1139+
1140+
// Zod v3 exposes `.shape`; Zod v4 keeps the shape on `_zod.def.shape`
1141+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1142+
const rawShape = (schema as any).shape ?? (isZ4Schema(schema as any) ? (schema as any)._zod?.def?.shape : undefined);
1143+
if (!rawShape) return undefined;
1144+
1145+
if (typeof rawShape === 'function') {
1146+
try {
1147+
return rawShape();
1148+
} catch {
1149+
return undefined;
1150+
}
1151+
}
1152+
1153+
return rawShape as Record<string, AnySchema>;
1154+
}
1155+
11381156
/**
11391157
* Additional, optional information for annotating a resource.
11401158
*/
@@ -1233,12 +1251,7 @@ export type RegisteredPrompt = {
12331251
};
12341252

12351253
function promptArgumentsFromSchema(schema: AnyObjectSchema): PromptArgument[] {
1236-
// v3 exposes .shape; v4 keeps the shape on _zod.def.shape
1237-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1238-
const shape: Record<string, unknown> | undefined =
1239-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1240-
(schema as any).shape ?? (isZ4Schema(schema as any) ? (schema as any)._zod?.def?.shape : undefined);
1241-
1254+
const shape = getObjectShape(schema);
12421255
if (!shape) return [];
12431256

12441257
return Object.entries(shape).map(
@@ -1247,8 +1260,7 @@ function promptArgumentsFromSchema(schema: AnyObjectSchema): PromptArgument[] {
12471260
// Determine required by attempting to parse `undefined`.
12481261
// If undefined fails, the arg is required.
12491262
const test = safeParse(field, undefined);
1250-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1251-
const required = !(test as any).success;
1263+
const required = !test.success;
12521264

12531265
return {
12541266
name,

‎src/server/sse.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { SSEServerTransport } from './sse.js';
44
import { McpServer } from './mcp.js';
55
import { createServer, type Server } from 'node:http';
66
import { AddressInfo } from 'node:net';
7-
import { z } from 'zod';
7+
import * as z from 'zod/v3';
88
import { CallToolResult, JSONRPCMessage } from 'src/types.js';
99

1010
const createMockResponse = () => {

‎src/server/title.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Server } from './index.js';
22
import { Client } from '../client/index.js';
33
import { InMemoryTransport } from '../inMemory.js';
4-
import { z } from 'zod';
4+
import * as z from 'zod/v3';
55
import { McpServer, ResourceTemplate } from './mcp.js';
66

77
describe('Title field backwards compatibility', () => {

‎src/server/zod-json-schema-compat.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import type * as z4c from 'zod/v4/core';
1010
import * as z4mini from 'zod/v4-mini';
1111

1212
import { AnyObjectSchema, isZ4Schema } from './zod-compat.js';
13-
import { zodToJsonSchema } from 'src/_vendor/zodToJsonSchema.js';
13+
import { zodToJsonSchema } from '../_vendor/zodToJsonSchema.js';
1414

1515
type JsonSchema = Record<string, unknown>;
1616

‎src/shared/protocol-transport-handling.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { describe, expect, test, beforeEach } from '@jest/globals';
22
import { Protocol } from './protocol.js';
33
import { Transport } from './transport.js';
44
import { Request, Notification, Result, JSONRPCMessage } from '../types.js';
5-
import { z } from 'zod';
5+
import * as z from 'zod/v3';
66

77
// Mock Transport class
88
class MockTransport implements Transport {

0 commit comments

Comments
 (0)