codemod improvements (#2156)

packages/codemod/test/commentInsertion.test.ts

@@ -0,0 +1,291 @@
+import { mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import path from 'node:path';
+import { afterEach, describe, expect, it } from 'vitest';
+
+import { getMigration } from '../src/migrations/index.js';
+import { run } from '../src/runner.js';
+import { CODEMOD_ERROR_PREFIX } from '../src/utils/diagnostics.js';
+
+const migration = getMigration('v1-to-v2')!;
+
+let tempDir: string;
+
+function createTempDir(): string {
+    tempDir = mkdtempSync(path.join(tmpdir(), 'mcp-codemod-comment-test-'));
+    return tempDir;
+}
+
+afterEach(() => {
+    if (tempDir) {
+        rmSync(tempDir, { recursive: true, force: true });
+    }
+});
+
+describe('comment insertion', () => {
+    it('inserts @mcp-codemod-error comment above an action-required location', () => {
+        const dir = createTempDir();
+        // handler with custom schema identifier triggers actionRequired in handlerRegistration
+        const input = [
+            `import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';`,
+            `import { MyCustomSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            ``,
+            `const server = new McpServer({ name: 'test', version: '1.0' });`,
+            `server.setRequestHandler(MyCustomSchema, async (req, extra) => {`,
+            `    return {};`,
+            `});`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(output).toContain(CODEMOD_ERROR_PREFIX);
+        expect(result.commentCount).toBeGreaterThan(0);
+    });
+
+    it('does not insert comments on dry-run', () => {
+        const dir = createTempDir();
+        const input = [
+            `import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';`,
+            `import { MyCustomSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            ``,
+            `const server = new McpServer({ name: 'test', version: '1.0' });`,
+            `server.setRequestHandler(MyCustomSchema, async (req, extra) => {`,
+            `    return {};`,
+            `});`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir, dryRun: true });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(output).not.toContain(CODEMOD_ERROR_PREFIX);
+        expect(output).toBe(input);
+        expect(result.commentCount).toBe(0);
+    });
+
+    it('does not insert comments for regular warnings (verification-type)', () => {
+        const dir = createTempDir();
+        // ErrorCode split produces verification warnings, not actionRequired
+        const input = [
+            `import { ErrorCode } from '@modelcontextprotocol/sdk/types.js';`,
+            `const code = ErrorCode.InvalidParams;`,
+            `const timeout = ErrorCode.RequestTimeout;`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(output).not.toContain(CODEMOD_ERROR_PREFIX);
+        expect(result.commentCount).toBe(0);
+    });
+
+    it('inserts multiple comments in one file in correct positions', () => {
+        const dir = createTempDir();
+        // Two .parse() calls on different schemas trigger two actionRequired diagnostics
+        const input = [
+            `import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const a = CallToolRequestSchema.parse(data1);`,
+            `const b = ListToolsRequestSchema.parse(data2);`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const commentLines = output.split('\n').filter(l => l.includes(CODEMOD_ERROR_PREFIX));
+        expect(commentLines.length).toBe(2);
+        expect(result.commentCount).toBe(2);
+    });
+
+    it('preserves indentation of the target line', () => {
+        const dir = createTempDir();
+        const input = [
+            `import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `function validate() {`,
+            `    const a = CallToolRequestSchema.parse(data);`,
+            `}`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const commentLine = output.split('\n').find(l => l.includes(CODEMOD_ERROR_PREFIX))!;
+        expect(commentLine).toMatch(/^    \/\*/);
+    });
+
+    it('does not duplicate comments on re-run (idempotency)', () => {
+        const dir = createTempDir();
+        const input = [
+            `import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const a = CallToolRequestSchema.parse(data);`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        run(migration, { targetDir: dir });
+        const afterFirst = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const firstCount = afterFirst.split('\n').filter(l => l.includes(CODEMOD_ERROR_PREFIX)).length;
+
+        // Run again on the already-transformed file
+        run(migration, { targetDir: dir });
+        const afterSecond = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const secondCount = afterSecond.split('\n').filter(l => l.includes(CODEMOD_ERROR_PREFIX)).length;
+
+        expect(firstCount).toBe(1);
+        expect(secondCount).toBe(firstCount);
+    });
+
+    it('sanitizes */ in diagnostic messages', () => {
+        const dir = createTempDir();
+        // The .parse() diagnostic message doesn't contain */, but we verify the comment is well-formed
+        const input = [
+            `import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const a = CallToolRequestSchema.parse(data);`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const commentLine = output.split('\n').find(l => l.includes(CODEMOD_ERROR_PREFIX))!;
+        // Comment must be well-formed: starts with /* and ends with */
+        expect(commentLine.trim()).toMatch(/^\/\*.*\*\/$/);
+    });
+
+    it('places comments at correct line after import-path transform shifts lines', () => {
+        const dir = createTempDir();
+        // Import rewrite adds new import lines (splitting into multiple packages),
+        // then handler transform emits actionRequired. The comment must land at the correct post-shift line.
+        const input = [
+            `import { McpServer, CallToolRequestSchema } from '@modelcontextprotocol/sdk/server/mcp.js';`,
+            ``,
+            `const server = new McpServer({ name: 'test', version: '1.0' });`,
+            `const a = CallToolRequestSchema.parse(data);`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const lines = output.split('\n');
+        const commentIdx = lines.findIndex(l => l.includes(CODEMOD_ERROR_PREFIX));
+        expect(commentIdx).toBeGreaterThan(-1);
+        // The comment should be directly above the parse() line (which may have moved)
+        const nextLine = lines[commentIdx + 1]!;
+        expect(nextLine).toContain('.parse(data)');
+    });
+
+    it('merges same-line diagnostics into a single comment', () => {
+        const dir = createTempDir();
+        const input = [
+            `import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const a = CallToolRequestSchema.parse(data1); const b = ListToolsRequestSchema.parse(data2);`,
+            ``
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        const commentLines = output.split('\n').filter(l => l.includes(CODEMOD_ERROR_PREFIX));
+        expect(commentLines.length).toBe(1);
+        expect(commentLines[0]).toContain(' | ');
+        expect(result.commentCount).toBe(1);
+    });
+
+    it('skips comment insertion when target line is inside a template literal', () => {
+        const dir = createTempDir();
+        const input = [
+            "import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';",
+            'const msg = `',
+            '  Result: ${CallToolRequestSchema.parse(data).method}',
+            '`;',
+            ''
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        // The diagnostic should still be reported
+        expect(result.diagnostics.some(d => d.insertComment)).toBe(true);
+        // But no comment should be injected inside the template literal
+        expect(result.commentCount).toBe(0);
+        // Verify the template literal is not corrupted
+        expect(output).not.toContain('/* ' + CODEMOD_ERROR_PREFIX);
+    });
+
+    it('skips comment insertion when target line is inside template text after interpolation', () => {
+        const dir = createTempDir();
+        // TemplateMiddle: text between two ${} spans
+        const input = [
+            "import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';",
+            'const msg = `${somePrefix}',
+            '  A: ${CallToolRequestSchema.parse(d1)}',
+            '  B: ${ListToolsRequestSchema.parse(d2)}',
+            '`;',
+            ''
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(result.diagnostics.filter(d => d.insertComment).length).toBeGreaterThanOrEqual(2);
+        expect(result.commentCount).toBe(0);
+        expect(output).not.toContain('/* ' + CODEMOD_ERROR_PREFIX);
+    });
+
+    it('still inserts comment when diagnostic line merely contains a template literal', () => {
+        const dir = createTempDir();
+        // The .parse() and template are on the same line, but lineStart is at "const",
+        // which is outside the template literal.
+        const input = [
+            "import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';",
+            'const a = CallToolRequestSchema.parse(`template ${data}`);',
+            ''
+        ].join('\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        const result = run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(result.diagnostics.some(d => d.insertComment)).toBe(true);
+        expect(result.commentCount).toBeGreaterThan(0);
+        const lines = output.split('\n');
+        const commentIdx = lines.findIndex(l => l.includes(CODEMOD_ERROR_PREFIX));
+        expect(commentIdx).toBeGreaterThan(-1);
+        expect(lines[commentIdx]!.trim()).toMatch(/^\/\*.*\*\/$/);
+    });
+
+    it('handles CRLF line endings without corrupting the file', () => {
+        const dir = createTempDir();
+        const input = [
+            `import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const a = CallToolRequestSchema.parse(data);`,
+            ``
+        ].join('\r\n');
+        writeFileSync(path.join(dir, 'server.ts'), input);
+
+        run(migration, { targetDir: dir });
+
+        const output = readFileSync(path.join(dir, 'server.ts'), 'utf8');
+        expect(output).toContain(CODEMOD_ERROR_PREFIX);
+        const lines = output.split(/\r?\n/);
+        const commentIdx = lines.findIndex(l => l.includes(CODEMOD_ERROR_PREFIX));
+        expect(commentIdx).toBeGreaterThan(-1);
+        expect(lines[commentIdx]!.trim()).toMatch(/^\/\*.*\*\/$/);
+        expect(lines[commentIdx + 1]).toContain('.parse(data)');
+    });
+});

packages/codemod/test/v1-to-v2/transforms/contextTypes.test.ts

@@ -325,7 +325,7 @@ describe('context-types transform', () => {
         expect(result).not.toContain('meta.ctx');
     });
 
-    it('does not rename "extra" in shorthand property assignment', () => {
+    it('expands shorthand property assignment when renaming extra to ctx', () => {
         const input = [
             `server.setRequestHandler('tools/call', async (request, extra) => {`,
             `    helper({ request, extra });`,
@@ -334,8 +334,8 @@ describe('context-types transform', () => {
             ''
         ].join('\n');
         const result = applyTransform(input);
-        expect(result).toContain('{ request, extra }');
-        expect(result).not.toContain('{ request, ctx }');
+        expect(result).toContain('extra: ctx');
+        expect(result).not.toContain('{ request, extra }');
         expect(result).toContain('(request, ctx)');
     });
 

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

@@ -162,6 +162,7 @@ describe('handler-registration transform', () => {
 
     it('emits diagnostic for custom method schema (not in spec map)', () => {
         const input = [
+            `import { Server } from '@modelcontextprotocol/sdk/server/index.js';`,
             `const AcmeSearch = z.object({ method: z.literal('acme/search'), params: z.object({ query: z.string() }) });`,
             `server.setRequestHandler(AcmeSearch, async (request) => {`,
             `    return { items: [] };`,
@@ -179,6 +180,7 @@ describe('handler-registration transform', () => {
 
     it('emits diagnostic for custom notification schema', () => {
         const input = [
+            `import { Server } from '@modelcontextprotocol/sdk/server/index.js';`,
             `const CustomNotification = z.object({ method: z.literal('acme/notify') });`,
             `server.setNotificationHandler(CustomNotification, async () => {});`,
             ''
@@ -192,6 +194,20 @@ describe('handler-registration transform', () => {
         expect(result.diagnostics[0]!.message).toContain('CustomNotification');
     });
 
+    it('skips files with no MCP imports', () => {
+        const input = [
+            `import { EventBus } from 'some-other-library';`,
+            `const CustomSchema = z.object({ method: z.literal('custom/op') });`,
+            `bus.setRequestHandler(CustomSchema, async (req) => {});`,
+            ''
+        ].join('\n');
+        const project = new Project({ useInMemoryFileSystem: true });
+        const sourceFile = project.createSourceFile('test.ts', input);
+        const result = handlerRegistrationTransform.apply(sourceFile, ctx);
+        expect(result.changesCount).toBe(0);
+        expect(result.diagnostics.length).toBe(0);
+    });
+
     it('replaces ElicitationCompleteNotificationSchema with method string', () => {
         const input = [
             `import { ElicitationCompleteNotificationSchema } from '@modelcontextprotocol/sdk/types.js';`,