codemod improvements (#2156)

packages/codemod/batch-test/repos.json

@@ -1,7 +1,7 @@
 [
     {
-        "repo": "KKonstantinov/mcp-servers-fork",
-        "ref": "feature/upgrade-zod-v4",
+        "repo": "modelcontextprotocol/servers",
+        "ref": "main",
         "packages": [
             {
                 "dir": "src/everything",

packages/codemod/src/bin/batchTest.ts

@@ -353,7 +353,7 @@ function main(): void {
                 codemodResult = run(migration, { targetDir: fullSourceDir, verbose: true });
             } catch (error) {
                 console.log(`    ERROR: codemod threw: ${error}`);
-                codemodResult = { filesChanged: 0, totalChanges: 0, diagnostics: [], fileResults: [] };
+                codemodResult = { filesChanged: 0, totalChanges: 0, diagnostics: [], fileResults: [], commentCount: 0 };
             }
             console.log(
                 `    Codemod: files=${codemodResult.filesChanged} changes=${codemodResult.totalChanges} diags=${codemodResult.diagnostics.length}`

packages/codemod/src/cli.ts

@@ -9,7 +9,7 @@ import { Command } from 'commander';
 import { listMigrations } from './migrations/index.js';
 import { run } from './runner.js';
 import { DiagnosticLevel } from './types.js';
-import { formatDiagnostic } from './utils/diagnostics.js';
+import { CODEMOD_ERROR_PREFIX, formatDiagnostic } from './utils/diagnostics.js';
 
 const require = createRequire(import.meta.url);
 const { version } = require('../package.json') as { version: string };
@@ -140,6 +140,13 @@ for (const [name, migration] of listMigrations()) {
                     console.log('');
                 }
 
+                if (result.commentCount > 0) {
+                    console.log(
+                        `${result.commentCount} location(s) marked with ${CODEMOD_ERROR_PREFIX} comments — search your code to find them:\n` +
+                            `  grep -r '${CODEMOD_ERROR_PREFIX}' "${resolvedDir}"\n`
+                    );
+                }
+
                 if (opts['dryRun']) {
                     console.log('Run without --dry-run to apply changes.\n');
                 } else {

packages/codemod/src/migrations/v1-to-v2/mappings/importMap.ts

@@ -34,14 +34,8 @@ export const IMPORT_MAP: Record<string, ImportMapping> = {
         status: 'moved'
     },
     '@modelcontextprotocol/sdk/client/stdio.js': {
-        target: '@modelcontextprotocol/client',
-        status: 'moved',
-        symbolTargetOverrides: {
-            StdioClientTransport: '@modelcontextprotocol/client/stdio',
-            DEFAULT_INHERITED_ENV_VARS: '@modelcontextprotocol/client/stdio',
-            getDefaultEnvironment: '@modelcontextprotocol/client/stdio',
-            StdioServerParameters: '@modelcontextprotocol/client/stdio'
-        }
+        target: '@modelcontextprotocol/client/stdio',
+        status: 'moved'
     },
     '@modelcontextprotocol/sdk/client/websocket.js': {
         target: '',

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

@@ -2,7 +2,8 @@ import type { SourceFile } from 'ts-morph';
 import { Node, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
-import { info, warning } from '../../../utils/diagnostics.js';
+import { isKeyPositionIdentifier } from '../../../utils/astUtils.js';
+import { actionRequired, info } from '../../../utils/diagnostics.js';
 import { hasMcpImports } from '../../../utils/importUtils.js';
 import { CONTEXT_PROPERTY_MAP, CTX_PARAM_NAME, EXTRA_PARAM_NAME } from '../mappings/contextPropertyMap.js';
 
@@ -32,9 +33,9 @@ function processCallback(
     const paramNameNode = extraParam.getNameNode();
     if (Node.isObjectBindingPattern(paramNameNode)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                extraParam.getStartLineNumber(),
+                extraParam,
                 `Destructuring of context parameter in signature: "${paramNameNode.getText()}". ` +
                     'Properties have been reorganized in v2 (e.g., signal is now ctx.mcpReq.signal). Manual refactoring required.'
             )
@@ -49,9 +50,9 @@ function processCallback(
     const otherParams = callbackNode.getParameters().filter(p => p !== extraParam);
     if (otherParams.some(p => p.getName() === CTX_PARAM_NAME)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                extraParam.getStartLineNumber(),
+                extraParam,
                 `Cannot rename '${EXTRA_PARAM_NAME}' to '${CTX_PARAM_NAME}': another parameter is already named '${CTX_PARAM_NAME}'. Manual migration required.`
             )
         );
@@ -74,9 +75,9 @@ function processCallback(
         });
         if (ctxAlreadyInScope) {
             diagnostics.push(
-                warning(
+                actionRequired(
                     sourceFile.getFilePath(),
-                    extraParam.getStartLineNumber(),
+                    extraParam,
                     `Cannot rename '${EXTRA_PARAM_NAME}' to '${CTX_PARAM_NAME}': '${CTX_PARAM_NAME}' is already referenced in this scope. Manual migration required.`
                 )
             );
@@ -98,11 +99,7 @@ function processCallback(
         body.forEachDescendant(node => {
             if (!Node.isIdentifier(node) || node.getText() !== EXTRA_PARAM_NAME) return;
             const parent = node.getParent();
-            // Skip property-name positions (e.g., meta.extra, { extra: value }, { extra }, { extra: x } = obj)
-            if (parent && Node.isPropertyAccessExpression(parent) && parent.getNameNode() === node) return;
-            if (parent && Node.isPropertyAssignment(parent) && parent.getNameNode() === node) return;
-            if (parent && Node.isShorthandPropertyAssignment(parent)) return;
-            if (parent && Node.isBindingElement(parent) && parent.getPropertyNameNode() === node) return;
+            if (parent && isKeyPositionIdentifier(node)) return;
             identifiers.push(node);
         });
 
@@ -131,6 +128,11 @@ function processCallback(
                     }
                 }
             }
+            // Shorthand property assignment: { extra } → { extra: ctx }
+            if (parent && Node.isShorthandPropertyAssignment(parent)) {
+                replacements.push({ node: parent, newText: `${EXTRA_PARAM_NAME}: ${CTX_PARAM_NAME}` });
+                continue;
+            }
             replacements.push({ node: id, newText: CTX_PARAM_NAME });
         }
 
@@ -163,9 +165,9 @@ function processCallback(
             const nameNode = node.getNameNode();
             if (!Node.isObjectBindingPattern(nameNode)) return;
             diagnostics.push(
-                warning(
+                actionRequired(
                     sourceFile.getFilePath(),
-                    node.getStartLineNumber(),
+                    node,
                     `Destructuring of context parameter detected: "const ${nameNode.getText()} = ${CTX_PARAM_NAME}". ` +
                         'Properties have been reorganized in v2 (e.g., signal is now ctx.mcpReq.signal). Manual refactoring required.'
                 )

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

@@ -2,8 +2,8 @@ import type { SourceFile } from 'ts-morph';
 import { Node, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
-import { warning } from '../../../utils/diagnostics.js';
-import { isImportedFromMcp, removeUnusedImport, resolveOriginalImportName } from '../../../utils/importUtils.js';
+import { actionRequired } from '../../../utils/diagnostics.js';
+import { hasMcpImports, isImportedFromMcp, removeUnusedImport, resolveOriginalImportName } from '../../../utils/importUtils.js';
 import { NOTIFICATION_SCHEMA_TO_METHOD, SCHEMA_TO_METHOD } from '../mappings/schemaToMethodMap.js';
 
 const ALL_SCHEMA_TO_METHOD: Record<string, string> = {
@@ -15,6 +15,10 @@ export const handlerRegistrationTransform: Transform = {
     name: 'Handler registration migration',
     id: 'handlers',
     apply(sourceFile: SourceFile, _context: TransformContext): TransformResult {
+        if (!hasMcpImports(sourceFile)) {
+            return { changesCount: 0, diagnostics: [] };
+        }
+
         let changesCount = 0;
         const diagnostics: Diagnostic[] = [];
 
@@ -40,9 +44,9 @@ export const handlerRegistrationTransform: Transform = {
             const methodString = ALL_SCHEMA_TO_METHOD[originalName];
             if (!methodString) {
                 diagnostics.push(
-                    warning(
+                    actionRequired(
                         sourceFile.getFilePath(),
-                        call.getStartLineNumber(),
+                        call,
                         `Custom method handler: ${methodName}(${schemaName}, ...). ` +
                             `In v2, use the 3-arg form: ${methodName}('method/name', { params, result? }, handler). ` +
                             `See migration.md for details.`

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

@@ -2,7 +2,7 @@ import type { SourceFile } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
 import { renameAllReferences } from '../../../utils/astUtils.js';
-import { info, v2Gap, warning } from '../../../utils/diagnostics.js';
+import { actionRequired, info, v2Gap, warning } from '../../../utils/diagnostics.js';
 import { addOrMergeImport, getSdkExports, getSdkImports, isTypeOnlyImport } from '../../../utils/importUtils.js';
 import { resolveTypesPackage } from '../../../utils/projectAnalyzer.js';
 import { IMPORT_MAP, isAuthImport } from '../mappings/importMap.js';
@@ -82,7 +82,7 @@ export const importPathsTransform: Transform = {
             }
 
             if (!mapping) {
-                diagnostics.push(warning(filePath, line, `Unknown SDK import path: ${specifier}. Manual migration required.`));
+                diagnostics.push(actionRequired(filePath, imp, `Unknown SDK import path: ${specifier}. Manual migration required.`));
                 continue;
             }
 
@@ -125,9 +125,9 @@ export const importPathsTransform: Transform = {
                         effectiveTarget = mapping.symbolTargetOverrides[namedImports[0]!.getName()]!;
                     } else if (namedImports.some(n => n.getName() in mapping.symbolTargetOverrides!)) {
                         diagnostics.push(
-                            warning(
+                            actionRequired(
                                 filePath,
-                                line,
+                                imp,
                                 `Aliased import from ${specifier} mixes symbols that belong to different v2 packages. ` +
                                     `Split the import manually so each symbol targets the correct package.`
                             )
@@ -145,9 +145,9 @@ export const importPathsTransform: Transform = {
                     }
                     if (namespaceImport) {
                         diagnostics.push(
-                            warning(
+                            actionRequired(
                                 filePath,
-                                line,
+                                imp,
                                 `Namespace import of ${specifier}: exported symbol(s) ${Object.keys(mapping.renamedSymbols).join(', ')} ` +
                                     `were renamed in ${effectiveTarget}. Update qualified accesses manually.`
                             )
@@ -234,7 +234,7 @@ function rewriteExportDeclarations(
         }
 
         if (!mapping) {
-            diagnostics.push(warning(filePath, line, `Unknown SDK export path: ${specifier}. Manual migration required.`));
+            diagnostics.push(actionRequired(filePath, exp, `Unknown SDK export path: ${specifier}. Manual migration required.`));
             continue;
         }
 
@@ -269,9 +269,9 @@ function rewriteExportDeclarations(
                 targetPackage = mapping.symbolTargetOverrides[namedExports[0]!.getName()]!;
             } else if (namedExports.some(s => s.getName() in mapping.symbolTargetOverrides!)) {
                 diagnostics.push(
-                    warning(
+                    actionRequired(
                         filePath,
-                        line,
+                        exp,
                         `Re-export from ${specifier} mixes symbols that belong to different v2 packages. ` +
                             `Split the export manually so each symbol targets the correct package.`
                     )
@@ -288,7 +288,7 @@ function rewriteExportDeclarations(
                 spec.setName(newName);
             }
             if (REEXPORT_WARNINGS[name]) {
-                diagnostics.push(warning(filePath, line, REEXPORT_WARNINGS[name]!));
+                diagnostics.push(actionRequired(filePath, exp, REEXPORT_WARNINGS[name]!));
             }
         }
         changesCount++;

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

@@ -2,8 +2,8 @@ import type { CallExpression, SourceFile } from 'ts-morph';
 import { Node, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
-import { info, warning } from '../../../utils/diagnostics.js';
-import { isOriginalNameImportedFromMcp } from '../../../utils/importUtils.js';
+import { actionRequired, info } from '../../../utils/diagnostics.js';
+import { isOriginalNameImportedFromMcp, resolveLocalImportName } from '../../../utils/importUtils.js';
 
 export const mcpServerApiTransform: Transform = {
     name: 'McpServer API migration',
@@ -64,9 +64,9 @@ export const mcpServerApiTransform: Transform = {
                 changesCount++;
             } else {
                 diagnostics.push(
-                    warning(
+                    actionRequired(
                         sourceFile.getFilePath(),
-                        call.getStartLineNumber(),
+                        call,
                         'Could not automatically migrate .tool() call. Manual migration required.'
                     )
                 );
@@ -79,9 +79,9 @@ export const mcpServerApiTransform: Transform = {
                 changesCount++;
             } else {
                 diagnostics.push(
-                    warning(
+                    actionRequired(
                         sourceFile.getFilePath(),
-                        call.getStartLineNumber(),
+                        call,
                         'Could not automatically migrate .prompt() call. Manual migration required.'
                     )
                 );
@@ -94,9 +94,9 @@ export const mcpServerApiTransform: Transform = {
                 changesCount++;
             } else {
                 diagnostics.push(
-                    warning(
+                    actionRequired(
                         sourceFile.getFilePath(),
-                        call.getStartLineNumber(),
+                        call,
                         'Could not automatically migrate .resource() call. Manual migration required.'
                     )
                 );
@@ -121,6 +121,8 @@ export const mcpServerApiTransform: Transform = {
             }
         }
 
+        changesCount += migrateConstructorTaskOptions(sourceFile, diagnostics);
+
         return { changesCount, diagnostics };
     }
 };
@@ -129,6 +131,13 @@ function isStringArg(node: Node): boolean {
     return Node.isStringLiteral(node) || Node.isNoSubstitutionTemplateLiteral(node) || Node.isTemplateExpression(node);
 }
 
+function isZodObjectCall(node: Node): boolean {
+    if (!Node.isCallExpression(node)) return false;
+    const expr = node.getExpression();
+    if (!Node.isPropertyAccessExpression(expr)) return false;
+    return expr.getName() === 'object' && expr.getExpression().getText() === 'z';
+}
+
 function wrapWithZObject(schemaText: string): string {
     return `z.object(${schemaText})`;
 }
@@ -150,6 +159,14 @@ function emitWrapDiagnostic(node: Node, sourceFile: SourceFile, call: CallExpres
                 'Raw object literal wrapped with z.object(). Verify that zod (z) is imported in this file.'
             )
         );
+    } else if (!isZodObjectCall(node)) {
+        diagnostics.push(
+            actionRequired(
+                sourceFile.getFilePath(),
+                call,
+                'Schema argument is not an object literal — verify it is a z.object() schema. V2 requires a Zod schema, not a raw object.'
+            )
+        );
     }
 }
 
@@ -177,9 +194,9 @@ function wrapSchemaInConfig(call: CallExpression, schemaPropertyName: string, so
 
     if (Node.isShorthandPropertyAssignment(schemaProp)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                call.getStartLineNumber(),
+                call,
                 `Shorthand \`{ ${schemaPropertyName} }\` in config: verify the value is a z.object() schema, not a raw object. V2 requires a Zod schema.`
             )
         );
@@ -204,13 +221,15 @@ function wrapSchemaInConfig(call: CallExpression, schemaPropertyName: string, so
         return true;
     }
 
-    diagnostics.push(
-        warning(
-            sourceFile.getFilePath(),
-            call.getStartLineNumber(),
-            `\`${schemaPropertyName}\` value is not an object literal — verify it is a z.object() schema. V2 requires a Zod schema, not a raw object.`
-        )
-    );
+    if (!isZodObjectCall(initializer)) {
+        diagnostics.push(
+            actionRequired(
+                sourceFile.getFilePath(),
+                call,
+                `\`${schemaPropertyName}\` value is not an object literal — verify it is a z.object() schema. V2 requires a Zod schema, not a raw object.`
+            )
+        );
+    }
     return false;
 }
 
@@ -392,3 +411,130 @@ function migrateResourceCall(call: CallExpression, _sourceFile: SourceFile): boo
 
     return true;
 }
+
+const TASK_OPTIONS = ['taskStore', 'taskMessageQueue'] as const;
+
+function migrateConstructorTaskOptions(sourceFile: SourceFile, diagnostics: Diagnostic[]): number {
+    const localName = resolveLocalImportName(sourceFile, 'McpServer');
+    if (!localName) return 0;
+
+    let changes = 0;
+
+    for (const node of sourceFile.getDescendantsOfKind(SyntaxKind.NewExpression)) {
+        if (node.wasForgotten()) continue;
+        const expr = node.getExpression();
+        if (!Node.isIdentifier(expr) || expr.getText() !== localName) continue;
+
+        const args = node.getArguments();
+        if (args.length < 2) continue;
+
+        const optionsArg = args[1]!;
+        if (!Node.isObjectLiteralExpression(optionsArg)) continue;
+
+        // Check if any task options are present at the top level
+        const propsToMove: string[] = [];
+        for (const propName of TASK_OPTIONS) {
+            if (optionsArg.getProperty(propName)) {
+                propsToMove.push(propName);
+            }
+        }
+        if (propsToMove.length === 0) continue;
+
+        // Find the tasks object's position within the options text using AST,
+        // then do all mutations via a single text replacement to avoid node invalidation.
+        const capabilitiesProp = optionsArg.getProperty('capabilities');
+        let tasksObjStart = -1;
+        let tasksObjEnd = -1;
+        const optionsStart = optionsArg.getStart();
+        if (capabilitiesProp && Node.isPropertyAssignment(capabilitiesProp)) {
+            const capInit = capabilitiesProp.getInitializer();
+            if (capInit && Node.isObjectLiteralExpression(capInit)) {
+                const tasksProp = capInit.getProperty('tasks');
+                if (tasksProp && Node.isPropertyAssignment(tasksProp)) {
+                    const tasksInit = tasksProp.getInitializer();
+                    if (tasksInit && Node.isObjectLiteralExpression(tasksInit)) {
+                        tasksObjStart = tasksInit.getStart() - optionsStart;
+                        tasksObjEnd = tasksInit.getEnd() - optionsStart;
+                    }
+                }
+            }
+        }
+
+        if (tasksObjStart === -1) {
+            for (const propName of propsToMove) {
+                diagnostics.push(
+                    actionRequired(
+                        sourceFile.getFilePath(),
+                        node,
+                        `Move '${propName}' from McpServer options into capabilities.tasks — v2 expects task runtime options inside the tasks capability.`
+                    )
+                );
+            }
+            continue;
+        }
+
+        // Single text replacement: remove top-level props and insert into tasks object.
+        // Use AST nodes (already located via getProperty) to get brace-balanced text and
+        // exact positions, avoiding regex truncation on values containing commas/braces.
+        // Collect all properties first, then process in reverse position order so each
+        // removal doesn't invalidate the positions of subsequent removals.
+        let optionsText = optionsArg.getText();
+        const argStart = optionsArg.getStart();
+        const propsWithPositions: { text: string; start: number; end: number }[] = [];
+        for (const propName of propsToMove) {
+            const prop = optionsArg.getProperty(propName);
+            if (!prop) continue;
+            propsWithPositions.push({
+                text: prop.getText(),
+                start: prop.getStart() - argStart,
+                end: prop.getEnd() - argStart
+            });
+        }
+        const propTexts = propsWithPositions.map(p => p.text);
+
+        // Remove in reverse position order so earlier positions remain valid
+        const sortedProps = propsWithPositions.toSorted((a, b) => b.start - a.start);
+        for (const { start, end } of sortedProps) {
+            let remStart = start;
+            let remEnd = end;
+            // Consume trailing comma and whitespace
+            const afterProp = optionsText.slice(remEnd);
+            const trailingMatch = afterProp.match(/^\s*,?\s*/);
+            if (trailingMatch) {
+                remEnd += trailingMatch[0].length;
+            }
+            // Consume leading whitespace/newline
+            const beforeProp = optionsText.slice(0, remStart);
+            const leadingMatch = beforeProp.match(/[\n\r]?\s*$/);
+            if (leadingMatch) {
+                remStart -= leadingMatch[0].length;
+            }
+            optionsText = optionsText.slice(0, remStart) + optionsText.slice(remEnd);
+            // Adjust tasks position if removal was before it
+            if (remStart < tasksObjStart) {
+                const shift = remEnd - remStart;
+                tasksObjStart -= shift;
+                tasksObjEnd -= shift;
+            }
+        }
+
+        if (propTexts.length === 0) continue;
+
+        // Insert into the tasks object (just before its closing brace)
+        const tasksText = optionsText.slice(tasksObjStart, tasksObjEnd);
+        const closingBrace = tasksText.lastIndexOf('}');
+        const before = tasksText.slice(0, closingBrace).trimEnd();
+        const sep = before.length > 1 ? ',\n' : '\n';
+        const newTasksText = before + sep + propTexts.join(',\n') + '\n' + tasksText.slice(closingBrace);
+        optionsText = optionsText.slice(0, tasksObjStart) + newTasksText + optionsText.slice(tasksObjEnd);
+
+        // Clean up double/trailing commas
+        optionsText = optionsText.replaceAll(/,(\s*,)/g, ',');
+        optionsText = optionsText.replaceAll(/,(\s*})/g, '$1');
+
+        optionsArg.replaceWithText(optionsText);
+        changes += propTexts.length;
+    }
+
+    return changes;
+}

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

@@ -2,13 +2,23 @@ import type { SourceFile } from 'ts-morph';
 import { Node, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
-import { v2Gap, warning } from '../../../utils/diagnostics.js';
+import { actionRequired, v2Gap, warning } from '../../../utils/diagnostics.js';
 import { isSdkSpecifier } from '../../../utils/importUtils.js';
 import { resolveTypesPackage } from '../../../utils/projectAnalyzer.js';
 import { IMPORT_MAP, isAuthImport } from '../mappings/importMap.js';
 import { SIMPLE_RENAMES } from '../mappings/symbolMap.js';
 
-const MOCK_METHODS = new Set(['mock', 'doMock']);
+const MOCK_METHODS = new Set([
+    'mock',
+    'doMock',
+    'unmock',
+    'dontMock',
+    'deepUnmock',
+    'requireActual',
+    'importActual',
+    'requireMock',
+    'createMockFromModule'
+]);
 const MOCK_CALLERS = new Set(['vi', 'jest']);
 
 export const mockPathsTransform: Transform = {
@@ -94,9 +104,7 @@ function rewriteMockCall(
         diagnostics
     });
     if (resolved === null) {
-        diagnostics.push(
-            warning(sourceFile.getFilePath(), call.getStartLineNumber(), `Unknown SDK mock path: ${specifier}. Manual migration required.`)
-        );
+        diagnostics.push(actionRequired(sourceFile.getFilePath(), call, `Unknown SDK mock path: ${specifier}. Manual migration required.`));
         return 0;
     }
     if ('removed' in resolved) {
@@ -122,9 +130,9 @@ function rewriteMockCall(
             effectiveTarget = resolved.symbolTargetOverrides[factorySymbols[0]!]!;
         } else if (someOverridden) {
             diagnostics.push(
-                warning(
+                actionRequired(
                     sourceFile.getFilePath(),
-                    call.getStartLineNumber(),
+                    call,
                     `Mock factory from ${specifier} mixes symbols that belong to different v2 packages. ` +
                         `Split the mock manually so each symbol targets the correct package.`
                 )
@@ -238,11 +246,7 @@ function rewriteDynamicImports(
         });
         if (resolved === null) {
             diagnostics.push(
-                warning(
-                    sourceFile.getFilePath(),
-                    node.getStartLineNumber(),
-                    `Unknown SDK dynamic import path: ${specifier}. Manual migration required.`
-                )
+                actionRequired(sourceFile.getFilePath(), node, `Unknown SDK dynamic import path: ${specifier}. Manual migration required.`)
             );
             return;
         }
@@ -313,9 +317,9 @@ function rewriteDynamicImports(
                 const moduleRenames = resolved.renamedSymbols ?? {};
                 if (!Node.isObjectBindingPattern(nameNode) && Object.keys(moduleRenames).length > 0) {
                     diagnostics.push(
-                        warning(
+                        actionRequired(
                             sourceFile.getFilePath(),
-                            node.getStartLineNumber(),
+                            node,
                             `Dynamic import assigned to variable (not destructured). Symbol renames (${Object.keys(moduleRenames).join(', ')}) were not applied. Manual update may be needed.`
                         )
                     );

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

@@ -3,7 +3,8 @@ import { Node, SyntaxKind } from 'ts-morph';
 
 import { SPEC_SCHEMA_NAMES, specSchemaToTypeName } from '../../../generated/specSchemaMap.js';
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types.js';
-import { warning } from '../../../utils/diagnostics.js';
+import { isKeyPositionIdentifier } from '../../../utils/astUtils.js';
+import { actionRequired, warning } from '../../../utils/diagnostics.js';
 import { addOrMergeImport, isAnyMcpSpecifier, removeUnusedImport } from '../../../utils/importUtils.js';
 
 export const specSchemaAccessTransform: Transform = {
@@ -70,9 +71,9 @@ function handleReference(
     // Pattern: z.infer<typeof XSchema> — type position
     if (isTypeofInTypePosition(ref)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                ref.getStartLineNumber(),
+                ref,
                 `Replace \`z.infer<typeof ${localName}>\` with the \`${typeName}\` type (already exported from the same v2 package).`
             )
         );
@@ -101,9 +102,9 @@ function handleReference(
         }
 
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                ref.getStartLineNumber(),
+                ref,
                 `${localName}.safeParse() not available in v2. Use \`isSpecType.${typeName}(value)\` for boolean validation, ` +
                     `or \`specTypeSchemas.${typeName}['~standard'].validate(value)\` for full result.`
             )
@@ -114,9 +115,9 @@ function handleReference(
     // Pattern: XSchema.parse(v) — diagnostic only
     if (isParsePattern(ref)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                ref.getStartLineNumber(),
+                ref,
                 `${localName}.parse() not available in v2. Use \`isSpecType.${typeName}(value)\` for validation, ` +
                     `or \`specTypeSchemas.${typeName}['~standard'].validate(value)\` and check for issues.`
             )
@@ -142,9 +143,9 @@ function handleReference(
 
     if (parent && Node.isExportSpecifier(parent)) {
         diagnostics.push(
-            warning(
+            actionRequired(
                 sourceFile.getFilePath(),
-                ref.getStartLineNumber(),
+                ref,
                 `Re-export of ${localName} requires manual update: replace with specTypeSchemas.${typeName} or remove.`
             )
         );
@@ -165,15 +166,7 @@ function handleReference(
         return true;
     }
 
-    if (parent && Node.isPropertyAssignment(parent) && parent.getNameNode() === ref) {
-        return false;
-    }
-
-    if (parent && Node.isBindingElement(parent) && parent.getPropertyNameNode() === ref) {
-        return false;
-    }
-
-    if (parent && Node.isPropertyAccessExpression(parent) && parent.getNameNode() === ref) {
+    if (parent && isKeyPositionIdentifier(ref)) {
         return false;
     }
 
@@ -301,9 +294,9 @@ function rewriteCapturedSafeParse(
                         replacements.push({ node: errorParent, newText: `${varName}.issues?.map(i => i.message).join(', ')` });
                     } else {
                         diagnostics.push(
-                            warning(
+                            actionRequired(
                                 sourceFile.getFilePath(),
-                                errorParent.getStartLineNumber(),
+                                errorParent,
                                 `${varName}.error.${subProp} has no StandardSchema equivalent. Manual migration required.`
                             )
                         );

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

@@ -116,7 +116,7 @@ function handleErrorCodeSplit(sourceFile: SourceFile, diagnostics: Diagnostic[])
         if (newImports.length > 0) {
             const existingImp = sourceFile
                 .getImportDeclarations()
-                .find(i => i.getModuleSpecifierValue() === targetModule && !i.isTypeOnly());
+                .find(i => i.getModuleSpecifierValue() === targetModule && !i.isTypeOnly() && !i.getNamespaceImport());
             if (existingImp) {
                 const existingNames = new Set(existingImp.getNamedImports().map(n => n.getName()));
                 const toAdd = newImports.filter(n => !existingNames.has(n));
@@ -235,14 +235,18 @@ function handleRequestHandlerExtra(sourceFile: SourceFile, context: TransformCon
         if (needsClientContext) newImports.push({ name: 'ClientContext', target: '@modelcontextprotocol/client' });
 
         for (const { name, target } of newImports) {
-            const existingImp = sourceFile.getImportDeclarations().find(i => i.getModuleSpecifierValue() === target && i.isTypeOnly());
+            const existingImp = sourceFile
+                .getImportDeclarations()
+                .find(i => i.getModuleSpecifierValue() === target && i.isTypeOnly() && !i.getNamespaceImport());
             if (existingImp) {
                 const existingNames = new Set(existingImp.getNamedImports().map(n => n.getName()));
                 if (!existingNames.has(name)) {
                     existingImp.addNamedImports([name]);
                 }
             } else {
-                const valueImp = sourceFile.getImportDeclarations().find(i => i.getModuleSpecifierValue() === target && !i.isTypeOnly());
+                const valueImp = sourceFile
+                    .getImportDeclarations()
+                    .find(i => i.getModuleSpecifierValue() === target && !i.isTypeOnly() && !i.getNamespaceImport());
                 if (valueImp) {
                     const existingNames = new Set(valueImp.getNamedImports().map(n => n.getName()));
                     if (!existingNames.has(name)) {

packages/codemod/src/runner.ts

@@ -1,10 +1,80 @@
-import { Project } from 'ts-morph';
+import type { Node } from 'ts-morph';
+import { Project, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, FileResult, Migration, RunnerOptions, RunnerResult } from './types.js';
-import { error } from './utils/diagnostics.js';
+import { CODEMOD_ERROR_PREFIX, error } from './utils/diagnostics.js';
 import { updatePackageJson } from './utils/packageJsonUpdater.js';
 import { analyzeProject } from './utils/projectAnalyzer.js';
 
+const LITERAL_NODE_KINDS = new Set([
+    SyntaxKind.NoSubstitutionTemplateLiteral,
+    SyntaxKind.TemplateHead,
+    SyntaxKind.TemplateMiddle,
+    SyntaxKind.TemplateTail,
+    SyntaxKind.StringLiteral,
+    SyntaxKind.JsxText
+]);
+
+function isInsideLiteral(node: Node | undefined): boolean {
+    let current = node;
+    while (current) {
+        if (LITERAL_NODE_KINDS.has(current.getKind())) return true;
+        current = current.getParent();
+    }
+    return false;
+}
+
+function insertDiagnosticComments(project: Project, fileResults: FileResult[]): number {
+    let insertedCount = 0;
+
+    for (const fr of fileResults) {
+        const commentDiags = fr.diagnostics.filter(d => d.insertComment).toSorted((a, b) => b.line - a.line);
+
+        if (commentDiags.length === 0) continue;
+
+        const merged: { line: number; message: string }[] = [];
+        for (const diag of commentDiags) {
+            const prev = merged.at(-1);
+            if (prev && prev.line === diag.line) {
+                prev.message += ' | ' + diag.message;
+            } else {
+                merged.push({ line: diag.line, message: diag.message });
+            }
+        }
+
+        const sf = project.getSourceFile(fr.filePath);
+        if (!sf) continue;
+
+        // `sourceText` and `lines` are computed once from the pre-insertion text.
+        // Insertions below mutate sf, but we process in descending line order, so
+        // each insertText only shifts positions above the next insertion point —
+        // prior byte offsets stay valid.
+        const sourceText = sf.getFullText();
+        const lines = sourceText.split('\n');
+        const lineEnding = sourceText.includes('\r\n') ? '\r\n' : '\n';
+
+        for (const diag of merged) {
+            const lineIndex = diag.line - 1;
+            if (lineIndex < 0 || lineIndex >= lines.length) continue;
+
+            if (lineIndex > 0 && lines[lineIndex - 1]!.includes(CODEMOD_ERROR_PREFIX)) continue;
+
+            const indent = lines[lineIndex]!.match(/^(\s*)/)?.[1] ?? '';
+            const safeMessage = diag.message.replaceAll('*/', '* /');
+            const comment = `${indent}/* ${CODEMOD_ERROR_PREFIX} ${safeMessage} */`;
+
+            const lineStart = lines.slice(0, lineIndex).reduce((sum, l) => sum + l.length + 1, 0);
+
+            if (isInsideLiteral(sf.getDescendantAtPos(lineStart))) continue;
+
+            sf.insertText(lineStart, comment + lineEnding);
+            insertedCount++;
+        }
+    }
+
+    return insertedCount;
+}
+
 function escapeGlobPath(p: string): string {
     return p.replaceAll(/[[\]{}()*?!@#]/g, String.raw`\$&`);
 }
@@ -97,6 +167,17 @@ export function run(migration: Migration, options: RunnerOptions): RunnerResult
             fileUsedPackages.clear();
         }
 
+        for (const d of fileDiagnostics) {
+            if (d.resolveCurrentLine) {
+                try {
+                    d.line = d.resolveCurrentLine();
+                } catch {
+                    // Node was removed by a later transform; keep snapshot line
+                }
+                delete d.resolveCurrentLine;
+            }
+        }
+
         if (fileChanges > 0 || fileDiagnostics.length > 0) {
             if (fileChanges > 0) {
                 filesChanged++;
@@ -118,7 +199,9 @@ export function run(migration: Migration, options: RunnerOptions): RunnerResult
 
     // Per-file mutations are atomic: if any transform fails, the file is rolled back to its
     // original state and an error diagnostic is emitted.
+    let commentCount = 0;
     if (!options.dryRun) {
+        commentCount = insertDiagnosticComments(project, fileResults);
         project.saveSync();
     }
 
@@ -127,6 +210,7 @@ export function run(migration: Migration, options: RunnerOptions): RunnerResult
         totalChanges,
         diagnostics: allDiagnostics,
         fileResults,
-        packageJsonChanges
+        packageJsonChanges,
+        commentCount
     };
 }

packages/codemod/src/types.ts

@@ -12,6 +12,8 @@ export interface Diagnostic {
     line: number;
     message: string;
     category?: 'v2-gap';
+    insertComment?: boolean;
+    resolveCurrentLine?: () => number;
 }
 
 export interface TransformResult {
@@ -62,4 +64,5 @@ export interface RunnerResult {
     diagnostics: Diagnostic[];
     fileResults: FileResult[];
     packageJsonChanges?: PackageJsonChange;
+    commentCount: number;
 }

packages/codemod/src/utils/astUtils.ts

@@ -1,6 +1,22 @@
 import type { SourceFile } from 'ts-morph';
 import { Node } from 'ts-morph';
 
+export function isKeyPositionIdentifier(node: import('ts-morph').Node): boolean {
+    const parent = node.getParent();
+    if (!parent) return false;
+    if (Node.isPropertyAssignment(parent) && parent.getNameNode() === node) return true;
+    if (Node.isPropertyAccessExpression(parent) && parent.getNameNode() === node) return true;
+    if (Node.isPropertySignature(parent) && parent.getNameNode() === node) return true;
+    if (Node.isMethodDeclaration(parent) && parent.getNameNode() === node) return true;
+    if (Node.isMethodSignature(parent) && parent.getNameNode() === node) return true;
+    if (Node.isPropertyDeclaration(parent) && parent.getNameNode() === node) return true;
+    if (Node.isEnumMember(parent) && parent.getNameNode() === node) return true;
+    if (Node.isBindingElement(parent) && parent.getPropertyNameNode() === node) return true;
+    if (Node.isGetAccessorDeclaration(parent) && parent.getNameNode() === node) return true;
+    if (Node.isSetAccessorDeclaration(parent) && parent.getNameNode() === node) return true;
+    return false;
+}
+
 export function renameAllReferences(sourceFile: SourceFile, oldName: string, newName: string): void {
     sourceFile.forEachDescendant(node => {
         if (Node.isIdentifier(node) && node.getText() === oldName) {
@@ -13,16 +29,7 @@ export function renameAllReferences(sourceFile: SourceFile, oldName: string, new
                 parent.getNameNode().replaceWithText(newName);
                 return;
             }
-            if (Node.isPropertyAssignment(parent) && parent.getNameNode() === node) return;
-            if (Node.isPropertyAccessExpression(parent) && parent.getNameNode() === node) return;
-            if (Node.isPropertySignature(parent) && parent.getNameNode() === node) return;
-            if (Node.isMethodDeclaration(parent) && parent.getNameNode() === node) return;
-            if (Node.isMethodSignature(parent) && parent.getNameNode() === node) return;
-            if (Node.isPropertyDeclaration(parent) && parent.getNameNode() === node) return;
-            if (Node.isEnumMember(parent) && parent.getNameNode() === node) return;
-            if (Node.isBindingElement(parent) && parent.getPropertyNameNode() === node) return;
-            if (Node.isGetAccessorDeclaration(parent) && parent.getNameNode() === node) return;
-            if (Node.isSetAccessorDeclaration(parent) && parent.getNameNode() === node) return;
+            if (isKeyPositionIdentifier(node)) return;
             if (Node.isShorthandPropertyAssignment(parent)) {
                 parent.replaceWithText(`${oldName}: ${newName}`);
                 return;

packages/codemod/src/utils/diagnostics.ts

@@ -1,6 +1,10 @@
+import type { Node } from 'ts-morph';
+
 import type { Diagnostic } from '../types.js';
 import { DiagnosticLevel } from '../types.js';
 
+export const CODEMOD_ERROR_PREFIX = '@mcp-codemod-error';
+
 export function error(file: string, line: number, message: string): Diagnostic {
     return { level: DiagnosticLevel.Error, file, line, message };
 }
@@ -17,6 +21,17 @@ export function v2Gap(file: string, line: number, message: string): Diagnostic {
     return { level: DiagnosticLevel.Warning, file, line, message, category: 'v2-gap' };
 }
 
+export function actionRequired(file: string, node: Node, message: string): Diagnostic {
+    return {
+        level: DiagnosticLevel.Warning,
+        file,
+        line: node.getStartLineNumber(),
+        message,
+        insertComment: true,
+        resolveCurrentLine: () => node.getStartLineNumber()
+    };
+}
+
 const LEVEL_PREFIX: Record<DiagnosticLevel, string> = {
     [DiagnosticLevel.Error]: 'ERROR',
     [DiagnosticLevel.Warning]: 'WARNING',

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';`,