feat: inject RELATED_TASK_META_KEY during type pre-processing

ochafik · claude · ochafik · commit 4ddcf2b29874 · 2025-12-13T13:27:26.000Z
Use ts-morph for pre-processing spec.types.ts to inject SDK extensions: - Transform extends clauses (JSONRPCRequest → Request, etc.) - Inject 'io.modelcontextprotocol/related-task' into RequestParams._meta This allows RequestParamsSchema and TaskAugmentedRequestParamsSchema to be re-exported from generated schemas, as they now include the SDK extension. Benefits: - SDK extension is injected at the type level, propagating to all derived schemas - Cleaner separation between spec types and SDK additions - Generated schemas are now SDK-compatible out of the box Total schemas re-exported: 33 Lines saved: ~356 (2600 → 2244) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
diff --git a/scripts/generate-schemas.ts b/scripts/generate-schemas.ts
@@ -94,12 +94,118 @@ const INTEGER_SCHEMAS = ['ProgressTokenSchema', 'RequestIdSchema'];
 // =============================================================================
 
 /**
- * Pre-process spec.types.ts to transform type hierarchy for SDK compatibility.
+ * The SDK-specific meta key for relating messages to tasks.
+ * This is added to RequestParams._meta during pre-processing.
+ */
+const RELATED_TASK_META_KEY = 'io.modelcontextprotocol/related-task';
+
+/**
+ * Pre-process spec.types.ts using ts-morph to transform for SDK compatibility.
+ *
+ * Transforms:
+ * 1. `extends JSONRPCRequest` → `extends Request`
+ * 2. `extends JSONRPCNotification` → `extends Notification`
+ * 3. Add RELATED_TASK_META_KEY to RequestParams._meta
  */
 function preProcessTypes(content: string): string {
-    content = content.replace(/\bextends\s+JSONRPCRequest\b/g, 'extends Request');
-    content = content.replace(/\bextends\s+JSONRPCNotification\b/g, 'extends Notification');
-    return content;
+    const project = new Project({ useInMemoryFileSystem: true });
+    const sourceFile = project.createSourceFile('types.ts', content);
+
+    console.log('  🔧 Pre-processing types...');
+
+    // Transform 1 & 2: Change extends clauses
+    transformExtendsClause(sourceFile, 'JSONRPCRequest', 'Request');
+    transformExtendsClause(sourceFile, 'JSONRPCNotification', 'Notification');
+
+    // Transform 3: Add RELATED_TASK_META_KEY to RequestParams._meta
+    injectRelatedTaskMetaKey(sourceFile);
+
+    return sourceFile.getFullText();
+}
+
+/**
+ * Transform extends clauses from one type to another.
+ */
+function transformExtendsClause(sourceFile: SourceFile, from: string, to: string): void {
+    for (const iface of sourceFile.getInterfaces()) {
+        for (const ext of iface.getExtends()) {
+            if (ext.getText() === from) {
+                ext.replaceWithText(to);
+                console.log(`    - ${iface.getName()}: extends ${from} → extends ${to}`);
+            }
+        }
+    }
+}
+
+/**
+ * Inject RELATED_TASK_META_KEY into RequestParams._meta interface.
+ *
+ * Before:
+ *   _meta?: {
+ *     progressToken?: ProgressToken;
+ *     [key: string]: unknown;
+ *   };
+ *
+ * After:
+ *   _meta?: {
+ *     progressToken?: ProgressToken;
+ *     'io.modelcontextprotocol/related-task'?: RelatedTaskMetadata;
+ *     [key: string]: unknown;
+ *   };
+ */
+function injectRelatedTaskMetaKey(sourceFile: SourceFile): void {
+    const requestParams = sourceFile.getInterface('RequestParams');
+    if (!requestParams) {
+        console.warn('    ⚠️  RequestParams interface not found');
+        return;
+    }
+
+    const metaProp = requestParams.getProperty('_meta');
+    if (!metaProp) {
+        console.warn('    ⚠️  _meta property not found in RequestParams');
+        return;
+    }
+
+    // Get the type of _meta (it's an inline type literal)
+    const typeNode = metaProp.getTypeNode();
+    if (!typeNode || !Node.isTypeLiteral(typeNode)) {
+        console.warn('    ⚠️  _meta is not a type literal');
+        return;
+    }
+
+    // Check if already has the key
+    const existingMember = typeNode.getMembers().find(m => {
+        if (Node.isPropertySignature(m)) {
+            const name = m.getName();
+            return name === `'${RELATED_TASK_META_KEY}'` || name === `"${RELATED_TASK_META_KEY}"`;
+        }
+        return false;
+    });
+
+    if (existingMember) {
+        console.log('    - RequestParams._meta already has RELATED_TASK_META_KEY');
+        return;
+    }
+
+    // Find the index signature ([key: string]: unknown) to insert before it
+    const members = typeNode.getMembers();
+    const indexSignatureIndex = members.findIndex(m => Node.isIndexSignatureDeclaration(m));
+
+    // Create the new property
+    const newProperty = `/**
+     * If specified, this request is related to the provided task.
+     */
+    '${RELATED_TASK_META_KEY}'?: RelatedTaskMetadata;`;
+
+    if (indexSignatureIndex >= 0) {
+        // Insert before index signature
+        typeNode.insertMember(indexSignatureIndex, newProperty);
+    } else {
+        // Add at the end
+        typeNode.addMember(newProperty);
+    }
+
+    console.log('    ✓ Injected RELATED_TASK_META_KEY into RequestParams._meta');
 }
 
 // =============================================================================
diff --git a/src/generated/sdk.schemas.ts b/src/generated/sdk.schemas.ts
@@ -18,35 +18,29 @@ export const ProgressTokenSchema = z.union([z.string(), z.number().int()]);
 export const CursorSchema = z.string();
 
 /**
- * Common params for any request.
+ * Metadata for augmenting a request with task execution.
+ * Include this in the `task` field of the request parameters.
  *
- * @internal
+ * @category `tasks`
  */
-export const RequestParamsSchema = z.object({
+export const TaskMetadataSchema = z.object({
     /**
-     * See [General fields: `_meta`](/specification/draft/basic/index#meta) for notes on `_meta` usage.
+     * Requested duration in milliseconds to retain task from creation.
      */
-    _meta: z
-        .looseObject({
-            /**
-             * If specified, the caller is requesting out-of-band progress notifications for this request (as represented by notifications/progress). The value of this parameter is an opaque token that will be attached to any subsequent notifications. The receiver is not obligated to provide these notifications.
-             */
-            progressToken: ProgressTokenSchema.optional()
-        })
-        .optional()
+    ttl: z.number().optional()
 });
 
 /**
- * Metadata for augmenting a request with task execution.
- * Include this in the `task` field of the request parameters.
+ * Metadata for associating messages with a task.
+ * Include this in the `_meta` field under the key `io.modelcontextprotocol/related-task`.
  *
  * @category `tasks`
  */
-export const TaskMetadataSchema = z.object({
+export const RelatedTaskMetadataSchema = z.object({
     /**
-     * Requested duration in milliseconds to retain task from creation.
+     * The task identifier this message is associated with.
      */
-    ttl: z.number().optional()
+    taskId: z.string()
 });
 
 /** @internal */
@@ -201,6 +195,29 @@ export const CancelledNotificationSchema = NotificationSchema.extend({
     params: CancelledNotificationParamsSchema
 });
 
+/**
+ * Common params for any request.
+ *
+ * @internal
+ */
+export const RequestParamsSchema = z.object({
+    /**
+     * See [General fields: `_meta`](/specification/draft/basic/index#meta) for notes on `_meta` usage.
+     */
+    _meta: z
+        .looseObject({
+            /**
+             * If specified, the caller is requesting out-of-band progress notifications for this request (as represented by notifications/progress). The value of this parameter is an opaque token that will be attached to any subsequent notifications. The receiver is not obligated to provide these notifications.
+             */
+            progressToken: ProgressTokenSchema.optional(),
+            /**
+             * If specified, this request is related to the provided task.
+             */
+            'io.modelcontextprotocol/related-task': RelatedTaskMetadataSchema.optional()
+        })
+        .optional()
+});
+
 /**
  * Capabilities a client may support. Known capabilities are defined here, in this schema, but this is not a closed set: any client can define its own, additional capabilities.
  *
@@ -1040,19 +1057,6 @@ export const ToolSchema = BaseMetadataSchema.extend(IconsSchema.shape).extend({
  */
 export const TaskStatusSchema = z.enum(['working', 'input_required', 'completed', 'failed', 'cancelled']);
 
-/**
- * Metadata for associating messages with a task.
- * Include this in the `_meta` field under the key `io.modelcontextprotocol/related-task`.
- *
- * @category `tasks`
- */
-export const RelatedTaskMetadataSchema = z.object({
-    /**
-     * The task identifier this message is associated with.
-     */
-    taskId: z.string()
-});
-
 /**
  * Data associated with a task.
  *
diff --git a/src/generated/sdk.types.ts b/src/generated/sdk.types.ts
@@ -68,6 +68,10 @@ export interface RequestParams {
          * If specified, the caller is requesting out-of-band progress notifications for this request (as represented by notifications/progress). The value of this parameter is an opaque token that will be attached to any subsequent notifications. The receiver is not obligated to provide these notifications.
          */
         progressToken?: ProgressToken;
+        /**
+         * If specified, this request is related to the provided task.
+         */
+        'io.modelcontextprotocol/related-task'?: RelatedTaskMetadata;
         [key: string]: unknown;
     };
 }
diff --git a/src/types.ts b/src/types.ts
@@ -51,8 +51,14 @@ import {
     ResourceContentsSchema,
     TextResourceContentsSchema,
     BlobResourceContentsSchema,
+    // Request/notification base schemas (with RELATED_TASK_META_KEY injected)
+    RequestParamsSchema,
+    TaskAugmentedRequestParamsSchema,
 } from './generated/sdk.schemas.js';
 
+// Alias RequestParamsSchema to BaseRequestParamsSchema for internal use
+const BaseRequestParamsSchema = RequestParamsSchema;
+
 export {
     ProgressTokenSchema,
     CursorSchema,
@@ -123,40 +129,17 @@ export const TaskCreationParamsSchema = z.looseObject({
     pollInterval: z.number().optional()
 });
 
-const RequestMetaSchema = z.looseObject({
-    /**
-     * If specified, the caller is requesting out-of-band progress notifications for this request (as represented by notifications/progress). The value of this parameter is an opaque token that will be attached to any subsequent notifications. The receiver is not obligated to provide these notifications.
-     */
-    progressToken: ProgressTokenSchema.optional(),
-    /**
-     * If specified, this request is related to the provided task.
-     */
-    [RELATED_TASK_META_KEY]: RelatedTaskMetadataSchema.optional()
-});
-
-/**
- * Common params for any request.
- */
-const BaseRequestParamsSchema = z.object({
-    /**
-     * See [General fields: `_meta`](/specification/draft/basic/index#meta) for notes on `_meta` usage.
-     */
-    _meta: RequestMetaSchema.optional()
-});
+// Note: RequestParamsSchema (aliased as BaseRequestParamsSchema) and
+// TaskAugmentedRequestParamsSchema are re-exported from generated.
+// They include RELATED_TASK_META_KEY in _meta, injected during pre-processing.
 
 /**
- * Common params for any task-augmented request.
+ * Request metadata schema - used in _meta field of requests, notifications, and results.
+ * This is extracted for reuse but the canonical definition is in RequestParamsSchema.
  */
-export const TaskAugmentedRequestParamsSchema = BaseRequestParamsSchema.extend({
-    /**
-     * If specified, the caller is requesting task-augmented execution for this request.
-     * The request will return a CreateTaskResult immediately, and the actual result can be
-     * retrieved later via tasks/result.
-     *
-     * Task augmentation is subject to capability negotiation - receivers MUST declare support
-     * for task augmentation of specific request types in their capabilities.
-     */
-    task: TaskMetadataSchema.optional()
+const RequestMetaSchema = z.looseObject({
+    progressToken: ProgressTokenSchema.optional(),
+    [RELATED_TASK_META_KEY]: RelatedTaskMetadataSchema.optional()
 });
 
 /**
