codemod iterations 5 (#2398)

packages/codemod/src/bin/batchTest.ts

@@ -157,19 +157,30 @@ function detectPm(repoRoot: string): string {
     return 'npm';
 }
 
-function installCommand(pm: string): string {
+export function installCommand(pm: string, opts: { hasOwnPnpmWorkspace: boolean; packageDirs: string[] }): string {
     if (pm !== 'pnpm') return `${pm} install --ignore-scripts`;
-    // pnpm walks up to find a workspace; clones live inside this SDK's pnpm workspace, so a plain
-    // `pnpm install` targets the OUTER workspace and never populates the clone's node_modules — every
-    // downstream check (tsc base config, tsup, vitest) then fails identically at baseline and post,
-    // masking real codemod signal.
-    //   --ignore-workspace:    treat the clone as a standalone project (not part of the SDK workspace).
-    //   --no-frozen-lockfile:  the codemod rewrites package.json to swap v1 → v2 deps, so the lockfile
-    //                          must be allowed to change. CI=true (set in shell()) otherwise defaults
-    //                          pnpm to a frozen lockfile and the post-codemod reinstall silently skips
-    //                          the new v2 deps, leaving the clone on v1.
-    // npm/yarn/bun key off a `workspaces` field in package.json (absent at this repo root), so they
-    // need no equivalent flags.
+    // --no-frozen-lockfile: the codemod rewrites package.json to swap v1 → v2 deps, so the lockfile must
+    //   be allowed to change. CI=true (set in shell()) otherwise defaults pnpm to a frozen lockfile and the
+    //   post-codemod reinstall silently skips the new v2 deps, leaving the clone on v1.
+    if (opts.hasOwnPnpmWorkspace) {
+        // The clone is its OWN pnpm workspace — pnpm-workspace.yaml defines catalog:/workspace: deps
+        // (e.g. mastra). `--ignore-workspace` would discard that file and pnpm would fail to resolve them
+        // (ERR_PNPM_CATALOG_ENTRY_NOT_FOUND_FOR_SPEC; unresolved workspace: links → repo skipped). We don't
+        // need it: the SDK workspace excludes the clones (`!packages/codemod/batch-test/**`) and pnpm uses
+        // the clone's own pnpm-workspace.yaml as the nearest root. Scope the install to the target packages
+        // and their dependencies (`{./<dir>}...`) so a monorepo target installs only what the checks need
+        // (e.g. 2 of mastra's 161 projects) instead of the whole tree. The braces are load-bearing: pnpm
+        // silently ignores the trailing `...` on a bare `./<dir>...` path selector (installing the package
+        // without its workspace deps); the `{./<dir>}...` form honors it.
+        const filters = opts.packageDirs
+            .filter(dir => dir !== '.')
+            .map(dir => `--filter ${JSON.stringify(`{./${dir}}...`)}`)
+            .join(' ');
+        return `pnpm install --ignore-scripts --no-frozen-lockfile${filters ? ` ${filters}` : ''}`;
+    }
+    // Single-package clone with no workspace of its own: pnpm would walk up to the (clone-excluding) SDK
+    // workspace and never populate the clone's node_modules. `--ignore-workspace` treats it as standalone.
+    // npm/yarn/bun key off a `workspaces` field in package.json (absent at this repo root).
     return 'pnpm install --ignore-scripts --ignore-workspace --no-frozen-lockfile';
 }
 
@@ -643,18 +654,23 @@ function main(): void {
         const pm = detectPm(clonePath);
         console.log(`  Package manager: ${pm}`);
 
-        // Step 3: Install
+        // Process packages
+        const packages: PackageEntry[] = entry.packages ?? [{ dir: '.', sourceDir: 'src' }];
+        const repoPkgResults: PackageReport[] = [];
+
+        // Step 3: Install. A clone that is its own pnpm workspace (catalog:/workspace: deps) must keep its
+        // pnpm-workspace.yaml — see installCommand. Computed once and reused for the post-codemod reinstall.
+        const installCmd = installCommand(pm, {
+            hasOwnPnpmWorkspace: existsSync(path.join(clonePath, 'pnpm-workspace.yaml')),
+            packageDirs: packages.map(p => p.dir)
+        });
         console.log('  Installing dependencies...');
-        const installResult = shell(installCommand(pm), clonePath);
+        const installResult = shell(installCmd, clonePath);
         if (installResult.exitCode !== 0) {
             console.log(`  ERROR: install failed, skipping\n  ${installResult.stderr.split('\n')[0]}`);
             continue;
         }
 
-        // Process packages
-        const packages: PackageEntry[] = entry.packages ?? [{ dir: '.', sourceDir: 'src' }];
-        const repoPkgResults: PackageReport[] = [];
-
         for (const pkg of packages) {
             const sourceDir = pkg.sourceDir ?? 'src';
             const fullPkgDir = path.join(clonePath, pkg.dir);
@@ -697,7 +713,7 @@ function main(): void {
                 if (rewrites > 0) console.log(`    Pinned ${rewrites} deps to resolved published versions`);
             }
             console.log('    Re-installing dependencies...');
-            shell(installCommand(pm), clonePath);
+            shell(installCmd, clonePath);
 
             // Step 7: Post-codemod checks
             console.log('    Running post-codemod checks...');

packages/codemod/src/generated/versions.ts

@@ -1,9 +1,9 @@
 // AUTO-GENERATED — do not edit. Run `pnpm run generate:versions` to regenerate.
 export const V2_PACKAGE_VERSIONS: Record<string, string> = {
-    '@modelcontextprotocol/client': '^2.0.0-alpha.3',
-    '@modelcontextprotocol/server': '^2.0.0-alpha.3',
-    '@modelcontextprotocol/node': '^2.0.0-alpha.3',
-    '@modelcontextprotocol/express': '^2.0.0-alpha.3',
-    '@modelcontextprotocol/server-legacy': '^2.0.0-alpha.3',
-    '@modelcontextprotocol/core': '^2.0.0-alpha.1'
+    '@modelcontextprotocol/client': '^2.0.0-beta.1',
+    '@modelcontextprotocol/server': '^2.0.0-beta.1',
+    '@modelcontextprotocol/node': '^2.0.0-beta.1',
+    '@modelcontextprotocol/express': '^2.0.0-beta.1',
+    '@modelcontextprotocol/server-legacy': '^2.0.0-beta.1',
+    '@modelcontextprotocol/core': '^2.0.0-beta.1'
 };

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

@@ -1,4 +1,4 @@
-import type { SourceFile } from 'ts-morph';
+import type { ObjectLiteralExpression, SourceFile } from 'ts-morph';
 import { Node, SyntaxKind } from 'ts-morph';
 
 import type { Diagnostic, Transform, TransformContext, TransformResult } from '../../../types';
@@ -9,8 +9,33 @@ import { CONTEXT_PROPERTY_MAP, CTX_PARAM_NAME, EXTRA_PARAM_NAME } from '../mappi
 
 const CONTEXT_LIKE_KEYS = new Set(CONTEXT_PROPERTY_MAP.map(mapping => mapping.from.slice(1)));
 
+/**
+ * v1 context keys distinctive enough that a single one on an object literal is a strong
+ * signal it's a hand-built handler-context mock (vs. generic keys like `signal`/`sessionId`/
+ * `requestId` — a bare correlation-ID literal such as `logger.info(msg, { requestId })` is
+ * not a context mock — which appear on unrelated objects and only count in aggregate).
+ */
+const DISTINCTIVE_CONTEXT_KEYS = new Set([
+    'sendRequest',
+    'sendNotification',
+    'requestInfo',
+    'authInfo',
+    'closeSSEStream',
+    'closeStandaloneSSEStream'
+]);
+
+/** A literal already carrying one of these is in the v2 nested shape — not a stale v1 mock. */
+const V2_SHAPE_KEYS = new Set(['mcpReq', 'http', 'task']);
+
 const HANDLER_METHODS = new Set(['setRequestHandler', 'setNotificationHandler']);
 
+/**
+ * Transport ingestion methods whose second argument is a flat `MessageExtraInfo`
+ * (authInfo/request/closeSSEStream/… stay top-level in v2), NOT a handler context —
+ * so a literal handed to them must never get handler-context reshape guidance.
+ */
+const TRANSPORT_MESSAGE_METHODS = new Set(['onmessage', 'handleMessage']);
+
 const REGISTER_METHODS = new Set(['registerTool', 'registerPrompt', 'registerResource', 'tool', 'prompt', 'resource']);
 
 /**
@@ -303,6 +328,7 @@ export const contextTypesTransform: Transform = {
 
         changesCount += processFallbackHandlerAssignments(sourceFile, diagnostics);
         changesCount += remapAnnotatedContextParams(sourceFile, diagnostics);
+        flagV1MockContextLiterals(sourceFile, diagnostics);
 
         return { changesCount, diagnostics };
     }
@@ -329,6 +355,108 @@ function processFallbackHandlerAssignments(sourceFile: SourceFile, diagnostics:
     return changes;
 }
 
+/**
+ * Render a v1 context key as a reshape hint, e.g. `sendRequest` → `mcpReq.send`. Returns
+ * undefined for `sessionId` (a no-op — it stays top-level in v2) and for non-context keys.
+ */
+function contextKeyReshapeHint(key: string): string | undefined {
+    const mapping = CONTEXT_PROPERTY_MAP.find(m => m.from === '.' + key);
+    if (mapping === undefined || mapping.from === mapping.to) return undefined;
+    // Render the target as a plain object path: '.http?.authInfo' → 'http.authInfo'.
+    return `${key} → ${mapping.to.replace(/^\./, '').replaceAll('?', '')}`;
+}
+
+/**
+ * Flag hand-built mocks of the handler context (common in tests). The call-site scan above
+ * only reshapes `extra.X` inside handler definitions it can anchor on (registerTool,
+ * setRequestHandler, fallback handlers, annotated params). A test hands its mock to a bare
+ * `handler(args, mockCtx)` invocation, so the object literal is never reached and keeps the
+ * flat v1 shape — at runtime the migrated handler reads `ctx.mcpReq.send` / `.id` / … against
+ * it and throws "Cannot read properties of undefined (reading 'send')". Advisory only: an
+ * untyped literal that merely shares a key name might not be a context mock, so never rewrite.
+ */
+function flagV1MockContextLiterals(sourceFile: SourceFile, diagnostics: Diagnostic[]): void {
+    for (const obj of sourceFile.getDescendantsOfKind(SyntaxKind.ObjectLiteralExpression)) {
+        // Only literals in a mock-context position: a call argument or a variable initializer.
+        // Covers both `handler({}, { sendRequest: fn })` and `const extra = { … }; handler(a, extra)`.
+        // Unwrap casts/parens first so typed mocks — `{ … } as unknown as RequestHandlerExtra`,
+        // `({ … })`, `{ … } satisfies X` — are anchored by what encloses the cast, not the
+        // AsExpression that directly wraps the literal.
+        let expr: Node = obj;
+        let parent = expr.getParent();
+        while (
+            parent !== undefined &&
+            (Node.isAsExpression(parent) || Node.isSatisfiesExpression(parent) || Node.isParenthesizedExpression(parent))
+        ) {
+            expr = parent;
+            parent = parent.getParent();
+        }
+        const isCallArg = parent !== undefined && Node.isCallExpression(parent) && parent.getArguments().includes(expr);
+        const isVarInit = parent !== undefined && Node.isVariableDeclaration(parent) && parent.getInitializer() === expr;
+        if (!isCallArg && !isVarInit) continue;
+
+        // A literal handed to a transport ingestion method (`transport.onmessage(msg, { … })`,
+        // `transport.handleMessage(msg, { … })`) is a flat MessageExtraInfo, not a handler-context
+        // mock — reshaping its authInfo/request/closeSSEStream/… under http/mcpReq would be wrong.
+        if (isCallArg && Node.isCallExpression(parent)) {
+            const callee = parent.getExpression();
+            if (Node.isPropertyAccessExpression(callee) && TRANSPORT_MESSAGE_METHODS.has(callee.getName())) continue;
+        }
+
+        // The codemod's OWN output: an options object inside a just-rewritten handler whose values
+        // now read `ctx.mcpReq.*` / `ctx.http.*`. The keys keep their v1 names — so V2_SHAPE_KEYS
+        // below, which only inspects key names, won't catch it — but the values are already v2, not
+        // a stale mock. Flagging it would insert a comment above code the codemod itself produced.
+        if (readsFromMigratedContext(obj)) continue;
+
+        // Collect named property keys (skip spreads and computed names).
+        const keys: string[] = [];
+        for (const prop of obj.getProperties()) {
+            if (!Node.isPropertyAssignment(prop) && !Node.isShorthandPropertyAssignment(prop) && !Node.isMethodDeclaration(prop)) continue;
+            const nameNode = prop.getNameNode();
+            if (Node.isIdentifier(nameNode)) keys.push(nameNode.getText());
+            else if (Node.isStringLiteral(nameNode)) keys.push(nameNode.getLiteralText());
+        }
+        if (keys.some(key => V2_SHAPE_KEYS.has(key))) continue; // already v2-shaped
+
+        const contextKeys = keys.filter(key => CONTEXT_LIKE_KEYS.has(key));
+        const hasDistinctive = contextKeys.some(key => DISTINCTIVE_CONTEXT_KEYS.has(key));
+        if (!hasDistinctive && contextKeys.length < 2) continue;
+
+        const reshapes = contextKeys.map(key => contextKeyReshapeHint(key)).filter((hint): hint is string => hint !== undefined);
+        const sessionNote = contextKeys.includes('sessionId') ? '; sessionId stays top-level' : '';
+        diagnostics.push({
+            ...actionRequired(
+                sourceFile.getFilePath(),
+                obj,
+                `This object looks like a v1 handler-context mock (${contextKeys.join(', ')}). v2 nests the context — ` +
+                    `reshape it (${reshapes.join('; ')}${sessionNote}), e.g. { sendRequest: fn } → { mcpReq: { send: fn } }. ` +
+                    `Passed as-is to a migrated handler that reads ctx.mcpReq.*, the v1 shape throws ` +
+                    `"Cannot read properties of undefined".`
+            ),
+            advisoryOnly: true
+        });
+    }
+}
+
+/**
+ * True when any property value already reads from the v2-nested context — a `.mcpReq` or `.http`
+ * property access (e.g. `ctx.mcpReq.signal`, `ctx.http?.authInfo`). Such a literal is migrated
+ * output, not a hand-built v1 mock: a real v1 mock supplies raw values (`fn`, `ac.signal`, a
+ * literal), never the nested v2 shape the codemod itself just wrote.
+ */
+function readsFromMigratedContext(obj: ObjectLiteralExpression): boolean {
+    for (const prop of obj.getProperties()) {
+        if (!Node.isPropertyAssignment(prop)) continue;
+        const init = prop.getInitializer();
+        if (init === undefined) continue;
+        const accesses = init.getDescendantsOfKind(SyntaxKind.PropertyAccessExpression);
+        if (Node.isPropertyAccessExpression(init)) accesses.push(init);
+        if (accesses.some(access => access.getName() === 'mcpReq' || access.getName() === 'http')) return true;
+    }
+    return false;
+}
+
 const CONTEXT_TYPE_NAMES = new Set(['RequestHandlerExtra', 'ServerContext', 'ClientContext']);
 /**
  * Split a type's text on top-level `|` only (angle brackets tracked). Returns

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

@@ -173,6 +173,54 @@ export const handlerRegistrationTransform: Transform = {
             removeUnusedImport(sourceFile, schemaName, true);
         }
 
+        // Flag surviving registration-schema references. A method-mapped *RequestSchema/*NotificationSchema
+        // whose import survived the conversion above, used as a VALUE (a call argument or an `===`/`!==`
+        // operand), is almost always a stale v1 registration assertion/lookup — in v2 the registration key is
+        // the method string, not the schema. Advisory, not rewritten: the same constant is still valid for
+        // `.parse()` etc., so we flag rather than risk breaking a legitimate schema use.
+        //
+        // Gate on a registration MOCK/lookup, not on any setRequestHandler access. A file that only *calls*
+        // setRequestHandler/setNotificationHandler (real registrations, which the pass above rewrites) tells us
+        // nothing about its other schema references — those may be plain schema uses (`validateSchema(S)`,
+        // `zodToJsonSchema(S)`, `ajv.compile(S)`). Only a NON-call reference — `inner.setRequestHandler.mock`,
+        // `expect(x.setRequestHandler)…`, a comparison operand — signals the file makes registration
+        // assertions, where a surviving schema key is likely stale.
+        const usesRegistrationMock = sourceFile.getDescendantsOfKind(SyntaxKind.PropertyAccessExpression).some(pa => {
+            if (pa.getName() !== 'setRequestHandler' && pa.getName() !== 'setNotificationHandler') return false;
+            const paParent = pa.getParent();
+            // Exclude real registration calls `x.setRequestHandler(…)`, where pa is the callee.
+            return !(Node.isCallExpression(paParent) && paParent.getExpression() === pa);
+        });
+        if (usesRegistrationMock) {
+            const flagged = new Set<string>();
+            for (const id of sourceFile.getDescendantsOfKind(SyntaxKind.Identifier)) {
+                const local = id.getText();
+                if (flagged.has(local)) continue;
+                const original = resolveOriginalImportName(sourceFile, local) ?? local;
+                const method = ALL_SCHEMA_TO_METHOD[original];
+                if (method === undefined || !isImportedFromMcp(sourceFile, local)) continue;
+                const parent = id.getParent();
+                if (parent === undefined) continue;
+                const isCallArg = Node.isCallExpression(parent) && parent.getArguments().includes(id);
+                const isComparand =
+                    Node.isBinaryExpression(parent) &&
+                    ['===', '!==', '==', '!='].includes(parent.getOperatorToken().getText()) &&
+                    (parent.getLeft() === id || parent.getRight() === id);
+                if (!isCallArg && !isComparand) continue;
+                flagged.add(local);
+                diagnostics.push({
+                    ...actionRequired(
+                        sourceFile.getFilePath(),
+                        id,
+                        `${local} is no longer the setRequestHandler/setNotificationHandler key in v2 — handlers ` +
+                            `register by the method string '${method}'. Update registration assertions/lookups ` +
+                            `(e.g. against a setRequestHandler mock) to compare against '${method}'.`
+                    ),
+                    advisoryOnly: true
+                });
+            }
+        }
+
         return { changesCount, diagnostics };
     }
 };

packages/codemod/test/batchTest.test.ts

@@ -3,7 +3,7 @@ import { tmpdir } from 'node:os';
 import path from 'node:path';
 import { describe, it, expect, vi, afterEach } from 'vitest';
 
-import { parseArgs, parseCodemodCliOutput, runCodemod } from '../src/bin/batchTest';
+import { parseArgs, parseCodemodCliOutput, runCodemod, installCommand } from '../src/bin/batchTest';
 import { computeResultsDirName, type ResolvedConfig } from '../src/bin/batchTest';
 import { parseNpmViewVersion, rewriteToPublishedVersion, cleanSubprocessEnv } from '../src/bin/batchTest';
 import { getMigration } from '../src/migrations/index';
@@ -12,6 +12,35 @@ afterEach(() => {
     vi.restoreAllMocks();
 });
 
+describe('installCommand', () => {
+    it('uses plain --ignore-scripts for non-pnpm managers', () => {
+        expect(installCommand('npm', { hasOwnPnpmWorkspace: false, packageDirs: ['.'] })).toBe('npm install --ignore-scripts');
+        expect(installCommand('yarn', { hasOwnPnpmWorkspace: true, packageDirs: ['packages/a'] })).toBe('yarn install --ignore-scripts');
+    });
+
+    it('isolates a pnpm clone with no workspace of its own via --ignore-workspace', () => {
+        expect(installCommand('pnpm', { hasOwnPnpmWorkspace: false, packageDirs: ['.'] })).toBe(
+            'pnpm install --ignore-scripts --ignore-workspace --no-frozen-lockfile'
+        );
+    });
+
+    it('respects a pnpm clone that is its own workspace and scopes the install to the target packages', () => {
+        const cmd = installCommand('pnpm', { hasOwnPnpmWorkspace: true, packageDirs: ['packages/core', 'packages/mcp'] });
+        expect(cmd).not.toContain('--ignore-workspace');
+        // braces required: pnpm ignores the `...` on a bare `./dir...` selector (drops workspace deps);
+        // `{./dir}...` includes them.
+        expect(cmd).toBe(
+            'pnpm install --ignore-scripts --no-frozen-lockfile --filter "{./packages/core}..." --filter "{./packages/mcp}..."'
+        );
+    });
+
+    it('installs the whole workspace (no --filter) when the only target is the clone root', () => {
+        expect(installCommand('pnpm', { hasOwnPnpmWorkspace: true, packageDirs: ['.'] })).toBe(
+            'pnpm install --ignore-scripts --no-frozen-lockfile'
+        );
+    });
+});
+
 describe('parseArgs', () => {
     it('defaults to local/local with latest codemod version and unset sdk version', () => {
         const opts = parseArgs([]);

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

@@ -2,7 +2,7 @@ import { describe, it, expect } from 'vitest';
 import { Project } from 'ts-morph';
 
 import { contextTypesTransform } from '../../../src/migrations/v1-to-v2/transforms/contextTypes';
-import type { TransformContext } from '../../../src/types';
+import type { Diagnostic, TransformContext } from '../../../src/types';
 
 const ctx: TransformContext = { projectType: 'server' };
 
@@ -637,3 +637,121 @@ describe('fallback handler and annotated context params (B5, #152)', () => {
         expect(diag?.advisoryOnly).toBe(true);
     });
 });
+
+describe('v1 mock context literals in tests (test doubles)', () => {
+    function apply(code: string) {
+        const project = new Project({ useInMemoryFileSystem: true });
+        const sourceFile = project.createSourceFile('test.ts', code);
+        const result = contextTypesTransform.apply(sourceFile, { projectType: 'server' });
+        return { text: sourceFile.getFullText(), result };
+    }
+    const IMPORT = `import { McpServer } from '@modelcontextprotocol/server';\n`;
+    const findMockDiag = (diags: Diagnostic[]) => diags.find(d => d.message.includes('v1 handler-context mock'));
+
+    it('advises (without rewriting) on a v1 context mock passed inline as a call argument', () => {
+        const code = IMPORT + [`const r = await handler({}, { sendRequest: mockSendRequest });`, ''].join('\n');
+        const { text, result } = apply(code);
+        // Advisory only — the literal is left untouched.
+        expect(text).toContain('{ sendRequest: mockSendRequest }');
+        const diag = findMockDiag(result.diagnostics);
+        expect(diag).toBeDefined();
+        expect(diag?.advisoryOnly).toBe(true);
+        expect(diag?.message).toContain('sendRequest → mcpReq.send');
+    });
+
+    it('advises on a v1 context mock bound to a variable, and notes sessionId stays top-level', () => {
+        const code = IMPORT + [`const extra = { sessionId: 'session-1', sendRequest: fn };`, `await handler(args, extra);`, ''].join('\n');
+        const { result } = apply(code);
+        const diag = findMockDiag(result.diagnostics);
+        expect(diag).toBeDefined();
+        expect(diag?.advisoryOnly).toBe(true);
+        expect(diag?.message).toContain('sessionId stays top-level');
+    });
+
+    it('advises on the _meta/requestId progress-notification mock shape', () => {
+        const code = IMPORT + [`await handler({ duration: 1 }, { _meta: {}, requestId: 'test-123' });`, ''].join('\n');
+        const { result } = apply(code);
+        const diag = findMockDiag(result.diagnostics);
+        expect(diag).toBeDefined();
+        expect(diag?.message).toContain('requestId → mcpReq.id');
+    });
+
+    it('advises when two generic context keys appear together (no distinctive key)', () => {
+        const code = IMPORT + [`await handler(args, { signal: ac.signal, sessionId: 's' });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeDefined();
+    });
+
+    it('does not advise on an already-migrated v2 context mock', () => {
+        const code = IMPORT + [`await handler({}, { mcpReq: { send: fn } });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('does not advise on a non-context object literal', () => {
+        const code = IMPORT + [`const config = { title: 'x', description: 'y' };`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('does not advise on a lone generic key (needs a distinctive key or >=2 context keys)', () => {
+        const code = IMPORT + [`doThing({ signal: ac.signal });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('does not advise on a lone requestId (a bare correlation-ID literal, not a context mock)', () => {
+        const code = IMPORT + [`logger.info('handling lookup', { requestId: args.id });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('advises on a cast-wrapped inline mock (`{ … } as unknown as RequestHandlerExtra`)', () => {
+        const code = IMPORT + [`await handler(args, { sendRequest: fn } as unknown as RequestHandlerExtra);`, ''].join('\n');
+        const { result } = apply(code);
+        const diag = findMockDiag(result.diagnostics);
+        expect(diag).toBeDefined();
+        expect(diag?.advisoryOnly).toBe(true);
+    });
+
+    it('advises on a cast-wrapped variable mock (`const ctx = { … } as any`)', () => {
+        const code = IMPORT + [`const mockCtx = { requestId: 1, sendRequest: fn } as any;`, `await handler(args, mockCtx);`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeDefined();
+    });
+
+    it('does not advise on already-v2 values (an options object reading ctx.mcpReq / ctx.sessionId)', () => {
+        const code = IMPORT + [`doThing({ signal: ctx.mcpReq.signal, sessionId: ctx.sessionId });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it("does not re-flag the codemod's own migrated handler output", () => {
+        const code =
+            IMPORT +
+            [
+                `server.setRequestHandler('x', async (req, extra) => {`,
+                `    await loadData(req.uri, { signal: extra.signal, sessionId: extra.sessionId });`,
+                `    return {};`,
+                `});`,
+                ''
+            ].join('\n');
+        const { text, result } = apply(code);
+        // sanity: the handler body WAS migrated to the v2 nested shape …
+        expect(text).toContain('ctx.mcpReq.signal');
+        // … and the resulting options object is not mistaken for a stale v1 mock.
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('does not advise on a flat MessageExtraInfo literal passed to onmessage', () => {
+        const code = IMPORT + [`transport.onmessage(msg, { requestInfo: req, authInfo });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+
+    it('does not advise on a flat MessageExtraInfo literal passed to handleMessage', () => {
+        const code = IMPORT + [`await transport.handleMessage(msg, { authInfo, closeSSEStream: fn });`, ''].join('\n');
+        const { result } = apply(code);
+        expect(findMockDiag(result.diagnostics)).toBeUndefined();
+    });
+});

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

@@ -394,3 +394,72 @@ describe('variable-hop guard rails (B6 review)', () => {
         expect(result.diagnostics.some(d => d.message.includes('Custom method handler'))).toBe(true);
     });
 });
+
+describe('stale registration-schema references (analysis: mcp-servers/memory)', () => {
+    function applyWithDiagnostics(code: string) {
+        const project = new Project({ useInMemoryFileSystem: true });
+        const sourceFile = project.createSourceFile('test.ts', code);
+        const result = handlerRegistrationTransform.apply(sourceFile, { projectType: 'server' });
+        return { text: sourceFile.getFullText(), result };
+    }
+
+    it('flags a request schema used as a setRequestHandler-mock assertion arg with the method string', () => {
+        const code = [
+            `import { SubscribeRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const schemas = inner.setRequestHandler.mock.calls.map((c) => c[0]);`,
+            `expect(schemas).toContain(SubscribeRequestSchema);`,
+            ''
+        ].join('\n');
+        const { result } = applyWithDiagnostics(code);
+        const diag = result.diagnostics.find(
+            d => d.message.includes('SubscribeRequestSchema') && d.message.includes("'resources/subscribe'")
+        );
+        expect(diag).toBeDefined();
+    });
+
+    it('flags a schema passed to a registration-lookup helper', () => {
+        const code = [
+            `import { UnsubscribeRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `const c = inner.setRequestHandler.mock.calls.find((x) => x[0] === s);`,
+            `handlerFor(inner, UnsubscribeRequestSchema);`,
+            ''
+        ].join('\n');
+        const { result } = applyWithDiagnostics(code);
+        expect(result.diagnostics.some(d => d.message.includes("'resources/unsubscribe'"))).toBe(true);
+    });
+
+    it('does NOT flag a request schema used as a schema (.parse) — property-access base', () => {
+        const code = [
+            `import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `server.setRequestHandler(CallToolRequestSchema, async () => ({ content: [] }));`,
+            `const parsed = CallToolRequestSchema.parse(x);`,
+            ''
+        ].join('\n');
+        const { result } = applyWithDiagnostics(code);
+        expect(result.diagnostics.some(d => d.message.includes('no longer the setRequestHandler'))).toBe(false);
+    });
+
+    it('does NOT flag in a file that performs no handler registration', () => {
+        const code = [
+            `import { SubscribeRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `someValidator(SubscribeRequestSchema);`,
+            ''
+        ].join('\n');
+        const { result } = applyWithDiagnostics(code);
+        expect(result.diagnostics.some(d => d.message.includes('no longer the setRequestHandler'))).toBe(false);
+    });
+
+    it('does NOT flag a schema consumed by a validator when setRequestHandler is only a real registration', () => {
+        // The file registers a handler for real (a call the pass rewrites); an unrelated schema-consuming
+        // call like validateSchema(S) / zodToJsonSchema(S) is not a stale registration key — only a
+        // setRequestHandler MOCK/lookup reference signals that surviving schema refs are stale.
+        const code = [
+            `import { CallToolRequestSchema, SubscribeRequestSchema } from '@modelcontextprotocol/sdk/types.js';`,
+            `server.setRequestHandler(CallToolRequestSchema, async () => ({ content: [] }));`,
+            `validateSchema(SubscribeRequestSchema);`,
+            ''
+        ].join('\n');
+        const { result } = applyWithDiagnostics(code);
+        expect(result.diagnostics.some(d => d.message.includes('no longer the setRequestHandler'))).toBe(false);
+    });
+});