refactor(completions): improve completion handling and error states

- Move completion logic from App.tsx to useConnection hook
- Replace useCompletion with simpler useCompletionState hook
- Add graceful fallback for servers without completion support
- Improve error handling and state management
- Update PromptsTab and ResourcesTab to use new completion API
- Add type safety improvements across completion interfaces

Author: Gavin Aboulhosn

client/src/App.tsx

@@ -15,9 +15,6 @@ import {
   Root,
   ServerNotification,
   Tool,
-  PromptReference,
-  ResourceReference,
-  CompleteResultSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 import React, { Suspense, useEffect, useRef, useState } from "react";
 import { useConnection } from "./lib/hooks/useConnection";
@@ -154,6 +151,8 @@ const App = () => {
     requestHistory,
     makeRequest: makeConnectionRequest,
     sendNotification,
+    handleCompletion,
+    completionsSupported,
     connect: connectMcpServer,
   } = useConnection({
     transportType,
@@ -267,38 +266,6 @@ const App = () => {
     }
   };
 
-  const handleCompletion = async (
-    ref: ResourceReference | PromptReference,
-    argName: string,
-    value: string,
-    signal?: AbortSignal,
-  ) => {
-    if (!mcpClient) {
-      throw new Error("MCP client not connected");
-    }
-
-    const request: ClientRequest = {
-      method: "completion/complete",
-      params: {
-        argument: {
-          name: argName,
-          value,
-        },
-        ref,
-      },
-    };
-
-    try {
-      const result = await makeRequest(request, CompleteResultSchema);
-      return result.completion.values;
-    } catch (e: unknown) {
-      const errorMessage = e instanceof Error ? e.message : String(e);
-
-      toast.error(errorMessage);
-      throw e;
-    }
-  };
-
   const listResources = async () => {
     const response = await makeRequest(
       {
@@ -518,7 +485,8 @@ const App = () => {
                         clearError("resources");
                         setSelectedResource(resource);
                       }}
-                      onComplete={handleCompletion}
+                      handleCompletion={handleCompletion}
+                      completionsSupported={completionsSupported}
                       resourceContent={resourceContent}
                       nextCursor={nextResourceCursor}
                       nextTemplateCursor={nextResourceTemplateCursor}
@@ -543,7 +511,8 @@ const App = () => {
                         clearError("prompts");
                         setSelectedPrompt(prompt);
                       }}
-                      onComplete={handleCompletion}
+                      handleCompletion={handleCompletion}
+                      completionsSupported={completionsSupported}
                       promptContent={promptContent}
                       nextCursor={nextPromptCursor}
                       error={errors.prompts}

client/src/components/PromptsTab.tsx

@@ -7,11 +7,12 @@ import { Textarea } from "@/components/ui/textarea";
 import {
   ListPromptsResult,
   PromptReference,
+  ResourceReference,
 } from "@modelcontextprotocol/sdk/types.js";
 import { AlertCircle } from "lucide-react";
 import { useEffect, useState } from "react";
 import ListPane from "./ListPane";
-import { useCompletions } from "@/lib/useCompletion";
+import { useCompletionState } from "@/lib/hooks/useCompletionState";
 
 export type Prompt = {
   name: string;
@@ -30,7 +31,8 @@ const PromptsTab = ({
   getPrompt,
   selectedPrompt,
   setSelectedPrompt,
-  onComplete,
+  handleCompletion,
+  completionsSupported,
   promptContent,
   nextCursor,
   error,
@@ -41,19 +43,19 @@ const PromptsTab = ({
   getPrompt: (name: string, args: Record<string, string>) => void;
   selectedPrompt: Prompt | null;
   setSelectedPrompt: (prompt: Prompt) => void;
-  onComplete: (
-    ref: PromptReference,
+  handleCompletion: (
+    ref: PromptReference | ResourceReference,
     argName: string,
     value: string,
   ) => Promise<string[]>;
+  completionsSupported: boolean;
   promptContent: string;
   nextCursor: ListPromptsResult["nextCursor"];
   error: string | null;
 }) => {
   const [promptArgs, setPromptArgs] = useState<Record<string, string>>({});
-  const { completions, clearCompletions, requestCompletions } = useCompletions({
-    onComplete,
-  });
+  const { completions, clearCompletions, requestCompletions } =
+    useCompletionState(handleCompletion, completionsSupported);
 
   useEffect(() => {
     clearCompletions();

client/src/components/ResourcesTab.tsx

@@ -9,11 +9,12 @@ import {
   ResourceTemplate,
   ListResourceTemplatesResult,
   ResourceReference,
+  PromptReference,
 } from "@modelcontextprotocol/sdk/types.js";
 import { AlertCircle, ChevronRight, FileText, RefreshCw } from "lucide-react";
 import ListPane from "./ListPane";
 import { useEffect, useState } from "react";
-import { useCompletions } from "@/lib/useCompletion";
+import { useCompletionState } from "@/lib/hooks/useCompletionState";
 
 const ResourcesTab = ({
   resources,
@@ -25,7 +26,8 @@ const ResourcesTab = ({
   readResource,
   selectedResource,
   setSelectedResource,
-  onComplete,
+  handleCompletion,
+  completionsSupported,
   resourceContent,
   nextCursor,
   nextTemplateCursor,
@@ -40,11 +42,12 @@ const ResourcesTab = ({
   readResource: (uri: string) => void;
   selectedResource: Resource | null;
   setSelectedResource: (resource: Resource | null) => void;
-  onComplete: (
-    ref: ResourceReference,
+  handleCompletion: (
+    ref: ResourceReference | PromptReference,
     argName: string,
     value: string,
   ) => Promise<string[]>;
+  completionsSupported: boolean;
   resourceContent: string;
   nextCursor: ListResourcesResult["nextCursor"];
   nextTemplateCursor: ListResourceTemplatesResult["nextCursor"];
@@ -56,9 +59,8 @@ const ResourcesTab = ({
     {},
   );
 
-  const { clearCompletions, completions, requestCompletions } = useCompletions({
-    onComplete,
-  });
+  const { completions, clearCompletions, requestCompletions } =
+    useCompletionState(handleCompletion, completionsSupported);
 
   useEffect(() => {
     clearCompletions();

client/src/lib/hooks/useCompletionState.ts

@@ -0,0 +1,76 @@
+import { useState, useCallback, useEffect } from "react";
+import { ResourceReference, PromptReference } from "@modelcontextprotocol/sdk/types.js";
+
+interface CompletionState {
+  completions: Record<string, string[]>;
+  loading: Record<string, boolean>;
+  error: Record<string, string | null>;
+}
+
+export function useCompletionState(
+  handleCompletion: (
+    ref: ResourceReference | PromptReference,
+    argName: string,
+    value: string,
+  ) => Promise<string[]>,
+  completionsSupported: boolean = true,
+) {
+  const [state, setState] = useState<CompletionState>({
+    completions: {},
+    loading: {},
+    error: {},
+  });
+
+  const clearCompletions = useCallback(() => {
+    setState({
+      completions: {},
+      loading: {},
+      error: {},
+    });
+  }, []);
+
+  const requestCompletions = useCallback(
+    async (ref: ResourceReference | PromptReference, argName: string, value: string) => {
+      if (!completionsSupported) {
+        return;
+      }
+
+      setState((prev) => ({
+        ...prev,
+        loading: { ...prev.loading, [argName]: true },
+        error: { ...prev.error, [argName]: null },
+      }));
+
+      try {
+        const values = await handleCompletion(ref, argName, value);
+        setState((prev) => ({
+          ...prev,
+          completions: { ...prev.completions, [argName]: values },
+          loading: { ...prev.loading, [argName]: false },
+        }));
+      } catch (err) {
+        const error = err instanceof Error ? err.message : String(err);
+        setState((prev) => ({
+          ...prev,
+          loading: { ...prev.loading, [argName]: false },
+          error: { ...prev.error, [argName]: error },
+        }));
+      }
+    },
+    [handleCompletion, completionsSupported],
+  );
+
+  // Clear completions when support status changes
+  useEffect(() => {
+    if (!completionsSupported) {
+      clearCompletions();
+    }
+  }, [completionsSupported, clearCompletions]);
+
+  return {
+    ...state,
+    clearCompletions,
+    requestCompletions,
+    completionsSupported,
+  };
+}

client/src/lib/hooks/useConnection.ts

@@ -12,6 +12,10 @@ import {
   Request,
   Result,
   ServerCapabilities,
+  PromptReference,
+  ResourceReference,
+  McpError,
+  CompleteResultSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 import { useState } from "react";
 import { toast } from "react-toastify";
@@ -36,6 +40,11 @@ interface UseConnectionOptions {
   getRoots?: () => any[];
 }
 
+interface RequestOptions {
+  signal?: AbortSignal;
+  timeout?: number;
+}
+
 export function useConnection({
   transportType,
   command,
@@ -58,6 +67,7 @@ export function useConnection({
   const [requestHistory, setRequestHistory] = useState<
     { request: string; response?: string }[]
   >([]);
+  const [completionsSupported, setCompletionsSupported] = useState(true);
 
   const pushHistory = (request: object, response?: object) => {
     setRequestHistory((prev) => [
@@ -72,7 +82,8 @@ export function useConnection({
   const makeRequest = async <T extends z.ZodType>(
     request: ClientRequest,
     schema: T,
-  ) => {
+    options?: RequestOptions,
+  ): Promise<z.output<T>> => {
     if (!mcpClient) {
       throw new Error("MCP client not connected");
     }
@@ -81,12 +92,12 @@ export function useConnection({
       const abortController = new AbortController();
       const timeoutId = setTimeout(() => {
         abortController.abort("Request timed out");
-      }, requestTimeout);
+      }, options?.timeout ?? requestTimeout);
 
       let response;
       try {
         response = await mcpClient.request(request, schema, {
-          signal: abortController.signal,
+          signal: options?.signal ?? abortController.signal,
         });
         pushHistory(request, response);
       } catch (error) {
@@ -100,9 +111,58 @@ export function useConnection({
 
       return response;
     } catch (e: unknown) {
+      // Check for Method not found error specifically for completions
+      if (
+        request.method === "completion/complete" &&
+        e instanceof McpError &&
+        e.code === -32601
+      ) {
+        setCompletionsSupported(false);
+        return { completion: { values: [] } } as z.output<T>;
+      }
+
       const errorString = (e as Error).message ?? String(e);
       toast.error(errorString);
+      throw e;
+    }
+  };
+
+  const handleCompletion = async (
+    ref: ResourceReference | PromptReference,
+    argName: string,
+    value: string,
+    signal?: AbortSignal,
+  ): Promise<string[]> => {
+    if (!mcpClient || !completionsSupported) {
+      return [];
+    }
+
+    const request: ClientRequest = {
+      method: "completion/complete",
+      params: {
+        argument: {
+          name: argName,
+          value,
+        },
+        ref,
+      },
+    };
+
+    try {
+      const response = await makeRequest(request, CompleteResultSchema, {
+        signal,
+      });
+      return response?.completion.values || [];
+    } catch (e: unknown) {
+      const errorMessage = e instanceof Error ? e.message : String(e);
+      pushHistory(request, { error: errorMessage });
+
+      if (e instanceof McpError && e.code === -32601) {
+        setCompletionsSupported(false);
+        return [];
+      }
 
+      toast.error(errorMessage);
       throw e;
     }
   };
@@ -238,6 +298,7 @@ export function useConnection({
 
       const capabilities = client.getServerCapabilities();
       setServerCapabilities(capabilities ?? null);
+      setCompletionsSupported(true); // Reset completions support on new connection
 
       if (onPendingRequest) {
         client.setRequestHandler(CreateMessageRequestSchema, (request) => {
@@ -268,6 +329,8 @@ export function useConnection({
     requestHistory,
     makeRequest,
     sendNotification,
+    handleCompletion,
+    completionsSupported,
     connect,
   };
 }