Skip to content

Commit 1a46c4c

Browse files
KKonstantinovKKonstantinov
committed
fixes & edge case handling
1 parent b8e397a commit 1a46c4c

11 files changed

Lines changed: 193 additions & 41 deletions

File tree

‎packages/codemod/src/migrations/v1-to-v2/transforms/handlerRegistration.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { SourceFile } from 'ts-morph';
22
import { Node, SyntaxKind } from 'ts-morph';
33

44
import type { Transform, TransformContext, TransformResult } from '../../../types.js';
5-
import { isImportedFromMcp, removeUnusedImport } from '../../../utils/importUtils.js';
5+
import { isImportedFromMcp, removeUnusedImport, resolveOriginalImportName } from '../../../utils/importUtils.js';
66
import { NOTIFICATION_SCHEMA_TO_METHOD, SCHEMA_TO_METHOD } from '../mappings/schemaToMethodMap.js';
77

88
const ALL_SCHEMA_TO_METHOD: Record<string, string> = {
@@ -34,7 +34,8 @@ export const handlerRegistrationTransform: Transform = {
3434
if (!Node.isIdentifier(firstArg)) continue;
3535

3636
const schemaName = firstArg.getText();
37-
const methodString = ALL_SCHEMA_TO_METHOD[schemaName];
37+
const originalName = resolveOriginalImportName(sourceFile, schemaName) ?? schemaName;
38+
const methodString = ALL_SCHEMA_TO_METHOD[originalName];
3839
if (!methodString) continue;
3940

4041
if (!isImportedFromMcp(sourceFile, schemaName)) continue;

‎packages/codemod/src/migrations/v1-to-v2/transforms/importPaths.ts‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,16 @@ import { warning } from '../../../utils/diagnostics.js';
66
import { addOrMergeImport, getSdkExports, getSdkImports, isTypeOnlyImport } from '../../../utils/importUtils.js';
77
import { resolveTypesPackage } from '../../../utils/projectAnalyzer.js';
88
import { IMPORT_MAP, isAuthImport } from '../mappings/importMap.js';
9+
import { SIMPLE_RENAMES } from '../mappings/symbolMap.js';
10+
11+
const REEXPORT_WARNINGS: Record<string, string> = {
12+
ErrorCode: 'Re-exported ErrorCode was split into ProtocolErrorCode and SdkErrorCode in v2. Update this re-export manually.',
13+
RequestHandlerExtra:
14+
'Re-exported RequestHandlerExtra was renamed to ServerContext/ClientContext in v2. Update this re-export manually.',
15+
IsomorphicHeaders: 'Re-exported IsomorphicHeaders was removed in v2 (replaced by standard Headers API). Remove this re-export.',
16+
StreamableHTTPError:
17+
'Re-exported StreamableHTTPError was renamed to SdkError in v2 with different constructor. Update this re-export manually.'
18+
};
919

1020
export const importPathsTransform: Transform = {
1121
name: 'Import path rewrites',
@@ -204,13 +214,15 @@ function rewriteExportDeclarations(
204214
}
205215

206216
exp.setModuleSpecifier(targetPackage);
207-
if (mapping.renamedSymbols) {
208-
for (const spec of exp.getNamedExports()) {
209-
const newName = mapping.renamedSymbols[spec.getName()];
210-
if (newName) {
211-
if (!spec.getAliasNode()) spec.setAlias(spec.getName());
212-
spec.setName(newName);
213-
}
217+
for (const spec of exp.getNamedExports()) {
218+
const name = spec.getName();
219+
const newName = mapping.renamedSymbols?.[name] ?? SIMPLE_RENAMES[name];
220+
if (newName) {
221+
if (!spec.getAliasNode()) spec.setAlias(name);
222+
spec.setName(newName);
223+
}
224+
if (REEXPORT_WARNINGS[name]) {
225+
diagnostics.push(warning(filePath, line, REEXPORT_WARNINGS[name]!));
214226
}
215227
}
216228
changesCount++;

‎packages/codemod/src/migrations/v1-to-v2/transforms/mockPaths.ts‎

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { warning } from '../../../utils/diagnostics.js';
66
import { isSdkSpecifier } from '../../../utils/importUtils.js';
77
import { resolveTypesPackage } from '../../../utils/projectAnalyzer.js';
88
import { IMPORT_MAP, isAuthImport } from '../mappings/importMap.js';
9+
import { SIMPLE_RENAMES } from '../mappings/symbolMap.js';
910

1011
const MOCK_METHODS = new Set(['mock', 'doMock']);
1112
const MOCK_CALLERS = new Set(['vi', 'jest']);
@@ -101,8 +102,9 @@ function rewriteMockCall(
101102
firstArg.setLiteralValue(resolved.target);
102103
changes++;
103104

104-
if (resolved.renamedSymbols && args.length >= 2) {
105-
changes += renameSymbolsInFactory(args[1]!, resolved.renamedSymbols);
105+
const allRenames: Record<string, string> = { ...SIMPLE_RENAMES, ...resolved.renamedSymbols };
106+
if (args.length >= 2) {
107+
changes += renameSymbolsInFactory(args[1]!, allRenames);
106108
}
107109

108110
return changes;
@@ -177,24 +179,25 @@ function rewriteDynamicImports(sourceFile: SourceFile, context: TransformContext
177179
firstArg.setLiteralValue(resolved.target);
178180
changes++;
179181

180-
if (resolved.renamedSymbols) {
181-
const parent = node.getParent();
182-
if (parent && Node.isAwaitExpression(parent)) {
183-
const grandParent = parent.getParent();
184-
if (grandParent && Node.isVariableDeclaration(grandParent)) {
185-
const nameNode = grandParent.getNameNode();
186-
if (Node.isObjectBindingPattern(nameNode)) {
187-
for (const element of nameNode.getElements()) {
188-
const bindingName = element.getName();
189-
const newName = resolved.renamedSymbols[bindingName];
190-
if (newName) {
191-
if (element.getPropertyNameNode()) {
192-
element.getPropertyNameNode()!.replaceWithText(newName);
193-
} else {
194-
element.replaceWithText(`${newName}: ${bindingName}`);
195-
}
196-
changes++;
182+
const allRenames: Record<string, string> = { ...SIMPLE_RENAMES, ...resolved.renamedSymbols };
183+
const parent = node.getParent();
184+
if (parent && Node.isAwaitExpression(parent)) {
185+
const grandParent = parent.getParent();
186+
if (grandParent && Node.isVariableDeclaration(grandParent)) {
187+
const nameNode = grandParent.getNameNode();
188+
if (Node.isObjectBindingPattern(nameNode)) {
189+
for (const element of nameNode.getElements()) {
190+
const propertyName = element.getPropertyNameNode()?.getText();
191+
const bindingName = element.getName();
192+
const lookupKey = propertyName ?? bindingName;
193+
const newName = allRenames[lookupKey];
194+
if (newName) {
195+
if (propertyName) {
196+
element.getPropertyNameNode()!.replaceWithText(newName);
197+
} else {
198+
element.replaceWithText(`${newName}: ${bindingName}`);
197199
}
200+
changes++;
198201
}
199202
}
200203
}

‎packages/codemod/src/migrations/v1-to-v2/transforms/schemaParamRemoval.ts‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { SourceFile } from 'ts-morph';
22
import { Node, SyntaxKind } from 'ts-morph';
33

44
import type { Transform, TransformContext, TransformResult } from '../../../types.js';
5-
import { isImportedFromMcp, removeUnusedImport } from '../../../utils/importUtils.js';
5+
import { isImportedFromMcp, removeUnusedImport, resolveOriginalImportName } from '../../../utils/importUtils.js';
66

77
const TARGET_METHODS = new Set(['request', 'callTool', 'send', 'sendRequest']);
88

@@ -25,9 +25,11 @@ export const schemaParamRemovalTransform: Transform = {
2525
if (args.length < 2) continue;
2626

2727
const secondArg = args[1]!;
28-
if (!isSchemaReference(secondArg)) continue;
28+
if (!Node.isIdentifier(secondArg)) continue;
2929

3030
const schemaName = secondArg.getText();
31+
const originalName = resolveOriginalImportName(sourceFile, schemaName) ?? schemaName;
32+
if (!originalName.endsWith('Schema')) continue;
3133
if (!isImportedFromMcp(sourceFile, schemaName)) continue;
3234

3335
call.removeArgument(1);
@@ -39,9 +41,3 @@ export const schemaParamRemovalTransform: Transform = {
3941
return { changesCount, diagnostics: [] };
4042
}
4143
};
42-
43-
function isSchemaReference(node: Node): boolean {
44-
if (!Node.isIdentifier(node)) return false;
45-
const text = node.getText();
46-
return text.endsWith('Schema');
47-
}

‎packages/codemod/src/utils/astUtils.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export function renameAllReferences(sourceFile: SourceFile, oldName: string, new
1919
if (Node.isMethodSignature(parent) && parent.getNameNode() === node) return;
2020
if (Node.isPropertyDeclaration(parent) && parent.getNameNode() === node) return;
2121
if (Node.isEnumMember(parent) && parent.getNameNode() === node) return;
22+
if (Node.isBindingElement(parent) && parent.getPropertyNameNode() === node) return;
2223
if (Node.isGetAccessorDeclaration(parent) && parent.getNameNode() === node) return;
2324
if (Node.isSetAccessorDeclaration(parent) && parent.getNameNode() === node) return;
2425
if (Node.isShorthandPropertyAssignment(parent)) {

‎packages/codemod/src/utils/importUtils.ts‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,24 @@ export function hasMcpImports(sourceFile: SourceFile): boolean {
7676
export function isImportedFromMcp(sourceFile: SourceFile, symbolName: string): boolean {
7777
return sourceFile.getImportDeclarations().some(imp => {
7878
if (!isAnyMcpSpecifier(imp.getModuleSpecifierValue())) return false;
79-
return imp.getNamedImports().some(n => n.getName() === symbolName);
79+
return imp.getNamedImports().some(n => {
80+
const localName = n.getAliasNode()?.getText() ?? n.getName();
81+
return localName === symbolName;
82+
});
8083
});
8184
}
8285

86+
export function resolveOriginalImportName(sourceFile: SourceFile, localName: string): string | undefined {
87+
for (const imp of sourceFile.getImportDeclarations()) {
88+
for (const n of imp.getNamedImports()) {
89+
const alias = n.getAliasNode()?.getText();
90+
if (alias === localName) return n.getName();
91+
if (!alias && n.getName() === localName) return localName;
92+
}
93+
}
94+
return undefined;
95+
}
96+
8397
export function removeUnusedImport(sourceFile: SourceFile, symbolName: string, onlyMcpImports?: boolean): void {
8498
let referenceCount = 0;
8599
sourceFile.forEachDescendant(node => {
@@ -95,7 +109,7 @@ export function removeUnusedImport(sourceFile: SourceFile, symbolName: string, o
95109
for (const imp of sourceFile.getImportDeclarations()) {
96110
if (onlyMcpImports && !isAnyMcpSpecifier(imp.getModuleSpecifierValue())) continue;
97111
for (const namedImport of imp.getNamedImports()) {
98-
if (namedImport.getName() === symbolName) {
112+
if ((namedImport.getAliasNode()?.getText() ?? namedImport.getName()) === symbolName) {
99113
namedImport.remove();
100114
if (imp.getNamedImports().length === 0 && !imp.getDefaultImport() && !imp.getNamespaceImport()) {
101115
imp.remove();

‎packages/codemod/test/v1-to-v2/transforms/handlerRegistration.test.ts‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ describe('handler-registration transform', () => {
104104
expect(result).not.toContain("'tools/call'");
105105
});
106106

107-
it('does not remove non-MCP import when MCP import of same name is consumed', () => {
107+
it('does not rewrite local import when aliased MCP import has same export name', () => {
108108
const input = [
109109
`import { CallToolRequestSchema } from './local-schemas.js';`,
110110
`import { CallToolRequestSchema as McpSchema } from '@modelcontextprotocol/sdk/types.js';`,
@@ -114,8 +114,8 @@ describe('handler-registration transform', () => {
114114
].join('\n');
115115
const result = applyTransform(input);
116116
expect(result).toContain("from './local-schemas.js'");
117-
expect(result).toContain("'tools/call'");
118-
expect(result).not.toMatch(/setRequestHandler\(CallToolRequestSchema/);
117+
expect(result).toContain('setRequestHandler(CallToolRequestSchema');
118+
expect(result).not.toContain("'tools/call'");
119119
});
120120

121121
it('replaces ListRootsRequestSchema with method string', () => {
@@ -139,4 +139,17 @@ describe('handler-registration transform', () => {
139139
expect(result).toContain("'notifications/roots/list_changed'");
140140
expect(result).not.toContain('RootsListChangedNotificationSchema');
141141
});
142+
143+
it('handles aliased schema imports', () => {
144+
const input = [
145+
`import { CallToolRequestSchema as CTRS } from '@modelcontextprotocol/sdk/types.js';`,
146+
`server.setRequestHandler(CTRS, async (request) => {`,
147+
` return { content: [] };`,
148+
`});`,
149+
''
150+
].join('\n');
151+
const result = applyTransform(input);
152+
expect(result).toContain("'tools/call'");
153+
expect(result).not.toContain('CTRS');
154+
});
142155
});

‎packages/codemod/test/v1-to-v2/transforms/importPaths.test.ts‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,4 +233,41 @@ describe('import-paths transform', () => {
233233
expect(output).toContain('@modelcontextprotocol/client');
234234
expect(output).not.toContain('@modelcontextprotocol/sdk');
235235
});
236+
237+
it('applies SIMPLE_RENAMES to re-export specifiers', () => {
238+
const input = `export { McpError, ResourceReference } from '@modelcontextprotocol/sdk/types.js';\n`;
239+
const result = applyTransform(input);
240+
expect(result).toContain('ProtocolError as McpError');
241+
expect(result).toContain('ResourceTemplateReference as ResourceReference');
242+
expect(result).toContain('@modelcontextprotocol/server');
243+
});
244+
245+
it('emits warning for re-exported ErrorCode', () => {
246+
const input = `export { ErrorCode } from '@modelcontextprotocol/sdk/types.js';\n`;
247+
const project = new Project({ useInMemoryFileSystem: true });
248+
const sourceFile = project.createSourceFile('test.ts', input);
249+
const result = importPathsTransform.apply(sourceFile, { projectType: 'server' });
250+
expect(result.diagnostics.length).toBeGreaterThan(0);
251+
expect(result.diagnostics[0]!.message).toContain('ErrorCode');
252+
expect(result.diagnostics[0]!.message).toContain('split');
253+
});
254+
255+
it('emits warning for re-exported RequestHandlerExtra', () => {
256+
const input = `export { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';\n`;
257+
const project = new Project({ useInMemoryFileSystem: true });
258+
const sourceFile = project.createSourceFile('test.ts', input);
259+
const result = importPathsTransform.apply(sourceFile, { projectType: 'server' });
260+
expect(result.diagnostics.length).toBeGreaterThan(0);
261+
expect(result.diagnostics[0]!.message).toContain('RequestHandlerExtra');
262+
});
263+
264+
it('emits warning for re-exported IsomorphicHeaders', () => {
265+
const input = `export { IsomorphicHeaders } from '@modelcontextprotocol/sdk/types.js';\n`;
266+
const project = new Project({ useInMemoryFileSystem: true });
267+
const sourceFile = project.createSourceFile('test.ts', input);
268+
const result = importPathsTransform.apply(sourceFile, { projectType: 'server' });
269+
expect(result.diagnostics.length).toBeGreaterThan(0);
270+
expect(result.diagnostics[0]!.message).toContain('IsomorphicHeaders');
271+
expect(result.diagnostics[0]!.message).toContain('removed');
272+
});
236273
});

‎packages/codemod/test/v1-to-v2/transforms/mockPaths.test.ts‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,18 @@ describe('mock-paths transform', () => {
118118
expect(result).toContain('new StreamableHTTPServerTransport({})');
119119
});
120120

121+
it('handles aliased dynamic import destructuring', () => {
122+
const input = [
123+
`const { StreamableHTTPServerTransport: MyTransport } = await import('@modelcontextprotocol/sdk/server/streamableHttp.js');`,
124+
`const t = new MyTransport({});`,
125+
''
126+
].join('\n');
127+
const result = applyTransform(input);
128+
expect(result).toContain(`import('@modelcontextprotocol/node')`);
129+
expect(result).toContain('NodeStreamableHTTPServerTransport: MyTransport');
130+
expect(result).toContain('new MyTransport({})');
131+
});
132+
121133
it('rewrites dynamic import for server/mcp.js', () => {
122134
const input = [`const { McpServer } = await import('@modelcontextprotocol/sdk/server/mcp.js');`, ''].join('\n');
123135
const result = applyTransform(input);
@@ -174,5 +186,44 @@ describe('mock-paths transform', () => {
174186
expect(result.diagnostics.length).toBeGreaterThan(0);
175187
expect(result.diagnostics[0]!.message).toContain('Unknown SDK mock path');
176188
});
189+
190+
it('renames SIMPLE_RENAMES symbols in mock factory', () => {
191+
const input = [
192+
`vi.mock('@modelcontextprotocol/sdk/types.js', () => ({`,
193+
` McpError: vi.fn(),`,
194+
` ResourceReference: { type: 'resource' },`,
195+
`}));`,
196+
''
197+
].join('\n');
198+
const result = applyTransform(input);
199+
expect(result).toContain('@modelcontextprotocol/server');
200+
expect(result).toContain('ProtocolError');
201+
expect(result).toContain('ResourceTemplateReference');
202+
expect(result).not.toMatch(/(?<!Protocol)\bMcpError\b/);
203+
});
204+
205+
it('renames SIMPLE_RENAMES symbols in dynamic import destructuring', () => {
206+
const input = [
207+
`const { McpError, ResourceReference } = await import('@modelcontextprotocol/sdk/types.js');`,
208+
`const err = new McpError('fail');`,
209+
''
210+
].join('\n');
211+
const result = applyTransform(input);
212+
expect(result).toContain('@modelcontextprotocol/server');
213+
expect(result).toContain('ProtocolError: McpError');
214+
expect(result).toContain('ResourceTemplateReference: ResourceReference');
215+
expect(result).toContain('new McpError(');
216+
});
217+
218+
it('renames SIMPLE_RENAMES symbols in aliased dynamic import destructuring', () => {
219+
const input = [
220+
`const { McpError: MyError } = await import('@modelcontextprotocol/sdk/types.js');`,
221+
`throw new MyError('fail');`,
222+
''
223+
].join('\n');
224+
const result = applyTransform(input);
225+
expect(result).toContain('ProtocolError: MyError');
226+
expect(result).toContain('new MyError(');
227+
});
177228
});
178229
});

‎packages/codemod/test/v1-to-v2/transforms/schemaParamRemoval.test.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,16 @@ describe('schema-param-removal transform', () => {
8585
expect(result).not.toContain('CreateMessageResultSchema');
8686
expect(result).toContain('extra.sendRequest(');
8787
});
88+
89+
it('removes aliased schema from sendRequest calls', () => {
90+
const input = [
91+
`import { CreateMessageResultSchema as CMRS } from '@modelcontextprotocol/sdk/types.js';`,
92+
`const result = await extra.sendRequest({ method: 'sampling/createMessage', params }, CMRS);`,
93+
''
94+
].join('\n');
95+
const result = applyTransform(input);
96+
expect(result).not.toContain('CMRS');
97+
expect(result).toContain('extra.sendRequest(');
98+
expect(result).not.toContain('CreateMessageResultSchema');
99+
});
88100
});

0 commit comments

Comments
 (0)