refactor: Add generic type safety for completions

Improve type safety for completions in both prompts and resources:
- Add generic type parameters to PromptProtocol and ResourceProtocol
- Add type-safe argument completion methods with proper inference
- Update MCPServer to handle typed protocols
- Add schema validation in completion methods

This ensures completion arguments and values maintain proper typing
throughout the completion flow.

Author: Gavin Aboulhosn

src/core/MCPServer.ts

@@ -255,7 +255,9 @@ export class MCPServer {
       const { ref, argument } = request.params;
 
       if (ref.type === "ref/prompt") {
-        const prompt = this.promptsMap.get(ref.name);
+        const prompt = this.promptsMap.get(ref.name) as
+          | PromptProtocol<any>
+          | undefined;
         if (!prompt?.complete) {
           return { completion: { values: [] } };
         }
@@ -265,7 +267,9 @@ export class MCPServer {
       }
 
       if (ref.type === "ref/resource") {
-        const resource = this.resourcesMap.get(ref.uri);
+        const resource = this.resourcesMap.get(ref.uri) as
+          | ResourceProtocol<any>
+          | undefined;
         if (!resource?.complete) {
           return { completion: { values: [] } };
         }

src/prompts/BasePrompt.ts

@@ -18,7 +18,7 @@ export type PromptCompletion = {
   hasMore?: boolean;
 };
 
-export interface PromptProtocol {
+export interface PromptProtocol<TArgs extends Record<string, any> = {}> {
   name: string;
   description: string;
   promptDefinition: {
@@ -30,7 +30,7 @@ export interface PromptProtocol {
       required?: boolean;
     }>;
   };
-  getMessages(args?: Record<string, unknown>): Promise<
+  getMessages(args?: Partial<TArgs>): Promise<
     Array<{
       role: string;
       content: {
@@ -44,11 +44,14 @@ export interface PromptProtocol {
       };
     }>
   >;
-  complete?(argument: string, name: string): Promise<PromptCompletion>;
+  complete?<K extends keyof TArgs & string>(
+    argumentName: K,
+    value: string,
+  ): Promise<PromptCompletion>;
 }
 
 export abstract class MCPPrompt<TArgs extends Record<string, any> = {}>
-  implements PromptProtocol
+  implements PromptProtocol<TArgs>
 {
   abstract name: string;
   abstract description: string;
@@ -92,10 +95,14 @@ export abstract class MCPPrompt<TArgs extends Record<string, any> = {}>
     return this.generateMessages(validatedArgs);
   }
 
-  async complete?(
-    argumentName: string,
+  async complete<K extends keyof TArgs & string>(
+    argumentName: K,
     value: string,
   ): Promise<PromptCompletion> {
+    if (!this.schema[argumentName].type) {
+      throw new Error(`No schema found for argument: ${argumentName}`);
+    }
+
     return {
       values: [],
       total: 0,

src/resources/BaseResource.ts

@@ -25,7 +25,7 @@ export type ResourceCompletion = {
   hasMore?: boolean;
 };
 
-export interface ResourceProtocol {
+export interface ResourceProtocol<TArgs extends Record<string, any> = {}> {
   uri: string;
   name: string;
   description?: string;
@@ -35,10 +35,15 @@ export interface ResourceProtocol {
   read(): Promise<ResourceContent[]>;
   subscribe?(): Promise<void>;
   unsubscribe?(): Promise<void>;
-  complete?(argument: string, name: string): Promise<ResourceCompletion>;
+  complete?<K extends keyof TArgs & string>(
+    argumentName: K,
+    value: TArgs[K],
+  ): Promise<ResourceCompletion>;
 }
 
-export abstract class MCPResource implements ResourceProtocol {
+export abstract class MCPResource<TArgs extends Record<string, any> = {}>
+  implements ResourceProtocol<TArgs>
+{
   abstract uri: string;
   abstract name: string;
   description?: string;