1- import type { SourceFile } from 'ts-morph' ;
1+ import type { ObjectLiteralExpression , SourceFile } from 'ts-morph' ;
22import { Node , SyntaxKind } from 'ts-morph' ;
33
44import type { Diagnostic , Transform , TransformContext , TransformResult } from '../../../types' ;
@@ -9,8 +9,33 @@ import { CONTEXT_PROPERTY_MAP, CTX_PARAM_NAME, EXTRA_PARAM_NAME } from '../mappi
99
1010const CONTEXT_LIKE_KEYS = new Set ( CONTEXT_PROPERTY_MAP . map ( mapping => mapping . from . slice ( 1 ) ) ) ;
1111
12+ /**
13+ * v1 context keys distinctive enough that a single one on an object literal is a strong
14+ * signal it's a hand-built handler-context mock (vs. generic keys like `signal`/`sessionId`/
15+ * `requestId` — a bare correlation-ID literal such as `logger.info(msg, { requestId })` is
16+ * not a context mock — which appear on unrelated objects and only count in aggregate).
17+ */
18+ const DISTINCTIVE_CONTEXT_KEYS = new Set ( [
19+ 'sendRequest' ,
20+ 'sendNotification' ,
21+ 'requestInfo' ,
22+ 'authInfo' ,
23+ 'closeSSEStream' ,
24+ 'closeStandaloneSSEStream'
25+ ] ) ;
26+
27+ /** A literal already carrying one of these is in the v2 nested shape — not a stale v1 mock. */
28+ const V2_SHAPE_KEYS = new Set ( [ 'mcpReq' , 'http' , 'task' ] ) ;
29+
1230const HANDLER_METHODS = new Set ( [ 'setRequestHandler' , 'setNotificationHandler' ] ) ;
1331
32+ /**
33+ * Transport ingestion methods whose second argument is a flat `MessageExtraInfo`
34+ * (authInfo/request/closeSSEStream/… stay top-level in v2), NOT a handler context —
35+ * so a literal handed to them must never get handler-context reshape guidance.
36+ */
37+ const TRANSPORT_MESSAGE_METHODS = new Set ( [ 'onmessage' , 'handleMessage' ] ) ;
38+
1439const REGISTER_METHODS = new Set ( [ 'registerTool' , 'registerPrompt' , 'registerResource' , 'tool' , 'prompt' , 'resource' ] ) ;
1540
1641/**
@@ -303,6 +328,7 @@ export const contextTypesTransform: Transform = {
303328
304329 changesCount += processFallbackHandlerAssignments ( sourceFile , diagnostics ) ;
305330 changesCount += remapAnnotatedContextParams ( sourceFile , diagnostics ) ;
331+ flagV1MockContextLiterals ( sourceFile , diagnostics ) ;
306332
307333 return { changesCount, diagnostics } ;
308334 }
@@ -329,6 +355,108 @@ function processFallbackHandlerAssignments(sourceFile: SourceFile, diagnostics:
329355 return changes ;
330356}
331357
358+ /**
359+ * Render a v1 context key as a reshape hint, e.g. `sendRequest` → `mcpReq.send`. Returns
360+ * undefined for `sessionId` (a no-op — it stays top-level in v2) and for non-context keys.
361+ */
362+ function contextKeyReshapeHint ( key : string ) : string | undefined {
363+ const mapping = CONTEXT_PROPERTY_MAP . find ( m => m . from === '.' + key ) ;
364+ if ( mapping === undefined || mapping . from === mapping . to ) return undefined ;
365+ // Render the target as a plain object path: '.http?.authInfo' → 'http.authInfo'.
366+ return `${ key } ${ mapping . to . replace ( / ^ \. / , '' ) . replaceAll ( '?' , '' ) } ;
367+ }
368+
369+ /**
370+ * Flag hand-built mocks of the handler context (common in tests). The call-site scan above
371+ * only reshapes `extra.X` inside handler definitions it can anchor on (registerTool,
372+ * setRequestHandler, fallback handlers, annotated params). A test hands its mock to a bare
373+ * `handler(args, mockCtx)` invocation, so the object literal is never reached and keeps the
374+ * flat v1 shape — at runtime the migrated handler reads `ctx.mcpReq.send` / `.id` / … against
375+ * it and throws "Cannot read properties of undefined (reading 'send')". Advisory only: an
376+ * untyped literal that merely shares a key name might not be a context mock, so never rewrite.
377+ */
378+ function flagV1MockContextLiterals ( sourceFile : SourceFile , diagnostics : Diagnostic [ ] ) : void {
379+ for ( const obj of sourceFile . getDescendantsOfKind ( SyntaxKind . ObjectLiteralExpression ) ) {
380+ // Only literals in a mock-context position: a call argument or a variable initializer.
381+ // Covers both `handler({}, { sendRequest: fn })` and `const extra = { … }; handler(a, extra)`.
382+ // Unwrap casts/parens first so typed mocks — `{ … } as unknown as RequestHandlerExtra`,
383+ // `({ … })`, `{ … } satisfies X` — are anchored by what encloses the cast, not the
384+ // AsExpression that directly wraps the literal.
385+ let expr : Node = obj ;
386+ let parent = expr . getParent ( ) ;
387+ while (
388+ parent !== undefined &&
389+ ( Node . isAsExpression ( parent ) || Node . isSatisfiesExpression ( parent ) || Node . isParenthesizedExpression ( parent ) )
390+ ) {
391+ expr = parent ;
392+ parent = parent . getParent ( ) ;
393+ }
394+ const isCallArg = parent !== undefined && Node . isCallExpression ( parent ) && parent . getArguments ( ) . includes ( expr ) ;
395+ const isVarInit = parent !== undefined && Node . isVariableDeclaration ( parent ) && parent . getInitializer ( ) === expr ;
396+ if ( ! isCallArg && ! isVarInit ) continue ;
397+
398+ // A literal handed to a transport ingestion method (`transport.onmessage(msg, { … })`,
399+ // `transport.handleMessage(msg, { … })`) is a flat MessageExtraInfo, not a handler-context
400+ // mock — reshaping its authInfo/request/closeSSEStream/… under http/mcpReq would be wrong.
401+ if ( isCallArg && Node . isCallExpression ( parent ) ) {
402+ const callee = parent . getExpression ( ) ;
403+ if ( Node . isPropertyAccessExpression ( callee ) && TRANSPORT_MESSAGE_METHODS . has ( callee . getName ( ) ) ) continue ;
404+ }
405+
406+ // The codemod's OWN output: an options object inside a just-rewritten handler whose values
407+ // now read `ctx.mcpReq.*` / `ctx.http.*`. The keys keep their v1 names — so V2_SHAPE_KEYS
408+ // below, which only inspects key names, won't catch it — but the values are already v2, not
409+ // a stale mock. Flagging it would insert a comment above code the codemod itself produced.
410+ if ( readsFromMigratedContext ( obj ) ) continue ;
411+
412+ // Collect named property keys (skip spreads and computed names).
413+ const keys : string [ ] = [ ] ;
414+ for ( const prop of obj . getProperties ( ) ) {
415+ if ( ! Node . isPropertyAssignment ( prop ) && ! Node . isShorthandPropertyAssignment ( prop ) && ! Node . isMethodDeclaration ( prop ) ) continue ;
416+ const nameNode = prop . getNameNode ( ) ;
417+ if ( Node . isIdentifier ( nameNode ) ) keys . push ( nameNode . getText ( ) ) ;
418+ else if ( Node . isStringLiteral ( nameNode ) ) keys . push ( nameNode . getLiteralText ( ) ) ;
419+ }
420+ if ( keys . some ( key => V2_SHAPE_KEYS . has ( key ) ) ) continue ; // already v2-shaped
421+
422+ const contextKeys = keys . filter ( key => CONTEXT_LIKE_KEYS . has ( key ) ) ;
423+ const hasDistinctive = contextKeys . some ( key => DISTINCTIVE_CONTEXT_KEYS . has ( key ) ) ;
424+ if ( ! hasDistinctive && contextKeys . length < 2 ) continue ;
425+
426+ const reshapes = contextKeys . map ( key => contextKeyReshapeHint ( key ) ) . filter ( ( hint ) : hint is string => hint !== undefined ) ;
427+ const sessionNote = contextKeys . includes ( 'sessionId' ) ? '; sessionId stays top-level' : '' ;
428+ diagnostics . push ( {
429+ ...actionRequired (
430+ sourceFile . getFilePath ( ) ,
431+ obj ,
432+ `This object looks like a v1 handler-context mock (${ contextKeys . join ( ', ' ) } +
433+ `reshape it (${ reshapes . join ( '; ' ) } ${ sessionNote } +
434+ `Passed as-is to a migrated handler that reads ctx.mcpReq.*, the v1 shape throws ` +
435+ `"Cannot read properties of undefined".`
436+ ) ,
437+ advisoryOnly : true
438+ } ) ;
439+ }
440+ }
441+
442+ /**
443+ * True when any property value already reads from the v2-nested context — a `.mcpReq` or `.http`
444+ * property access (e.g. `ctx.mcpReq.signal`, `ctx.http?.authInfo`). Such a literal is migrated
445+ * output, not a hand-built v1 mock: a real v1 mock supplies raw values (`fn`, `ac.signal`, a
446+ * literal), never the nested v2 shape the codemod itself just wrote.
447+ */
448+ function readsFromMigratedContext ( obj : ObjectLiteralExpression ) : boolean {
449+ for ( const prop of obj . getProperties ( ) ) {
450+ if ( ! Node . isPropertyAssignment ( prop ) ) continue ;
451+ const init = prop . getInitializer ( ) ;
452+ if ( init === undefined ) continue ;
453+ const accesses = init . getDescendantsOfKind ( SyntaxKind . PropertyAccessExpression ) ;
454+ if ( Node . isPropertyAccessExpression ( init ) ) accesses . push ( init ) ;
455+ if ( accesses . some ( access => access . getName ( ) === 'mcpReq' || access . getName ( ) === 'http' ) ) return true ;
456+ }
457+ return false ;
458+ }
459+
332460const CONTEXT_TYPE_NAMES = new Set ( [ 'RequestHandlerExtra' , 'ServerContext' , 'ClientContext' ] ) ;
333461/**
334462 * Split a type's text on top-level `|` only (angle brackets tracked). Returns
0 commit comments