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
2025-02-12 19:05:51 -05:00
parent c66feff37d
commit 7f713fe40e
6 changed files with 166 additions and 180 deletions
--- a/client/src/App.tsx
+++ b/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}
--- a/client/src/components/PromptsTab.tsx
+++ b/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();
--- a/client/src/components/ResourcesTab.tsx
+++ b/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();
--- a/client/src/lib/hooks/useCompletionState.ts
+++ b/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,
+  };
+}
--- a/client/src/lib/hooks/useConnection.ts
+++ b/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,
  };
 }
--- a/client/src/lib/useCompletion.ts
+++ b/client/src/lib/useCompletion.ts
@@ -1,126 +0,0 @@
-import { useState, useCallback, useRef, useEffect } from "react";
-import {
-  ResourceReference,
-  PromptReference,
-} from "@modelcontextprotocol/sdk/types.js";
-
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-function debounce<T extends (...args: any[]) => PromiseLike<void>>(
-  func: T,
-  wait: number,
-): (...args: Parameters<T>) => void {
-  let timeout: ReturnType<typeof setTimeout>;
-  return function (...args: Parameters<T>) {
-    clearTimeout(timeout);
-    timeout = setTimeout(() => func(...args), wait);
-  };
-}
-
-interface UseCompletionsOptions<T extends ResourceReference | PromptReference> {
-  onComplete: (
-    ref: T,
-    argName: string,
-    value: string,
-    signal?: AbortSignal,
-  ) => Promise<string[]>;
-  debounceMs?: number;
-}
-
-interface CompletionState {
-  completions: Record<string, string[]>;
-  loading: Record<string, boolean>;
-  error: Record<string, string | null>;
-}
-
-export function useCompletions<T extends ResourceReference | PromptReference>({
-  onComplete,
-  debounceMs = 300,
-}: UseCompletionsOptions<T>) {
-  const [state, setState] = useState<CompletionState>({
-    completions: {},
-    loading: {},
-    error: {},
-  });
-
-  const completeRef = useRef(onComplete);
-  completeRef.current = onComplete;
-
-  const abortControllerRef = useRef<AbortController | null>(null);
-
-  const cleanup = useCallback(() => {
-    if (abortControllerRef.current) {
-      abortControllerRef.current.abort();
-      abortControllerRef.current = null;
-    }
-  }, []);
-
-  // Cleanup on unmount
-  useEffect(() => {
-    return cleanup;
-  }, [cleanup]);
-
-  const requestCompletions = useCallback(
-    debounce(async (ref: T, argName: string, value: string) => {
-      // Abort any pending request
-      cleanup();
-
-      const abortController = new AbortController();
-      abortControllerRef.current = abortController;
-
-      setState((prev) => ({
-        ...prev,
-        loading: { ...prev.loading, [argName]: true },
-        error: { ...prev.error, [argName]: null },
-      }));
-
-      try {
-        const values = await completeRef.current(
-          ref,
-          argName,
-          value,
-          abortController.signal,
-        );
-
-        // Check if this request was aborted
-        if (!abortController.signal.aborted) {
-          setState((prev) => ({
-            ...prev,
-            completions: { ...prev.completions, [argName]: values },
-            loading: { ...prev.loading, [argName]: false },
-          }));
-        }
-      } catch (err) {
-        // Only update state if the request wasn't aborted
-        if (!abortController.signal.aborted) {
-          const error = err instanceof Error ? err.message : String(err);
-          setState((prev) => ({
-            ...prev,
-            loading: { ...prev.loading, [argName]: false },
-            error: { ...prev.error, [argName]: error },
-          }));
-        }
-      } finally {
-        // Clear the abort controller if it's still the current one
-        if (abortControllerRef.current === abortController) {
-          abortControllerRef.current = null;
-        }
-      }
-    }, debounceMs),
-    [cleanup, debounceMs],
-  );
-
-  const clearCompletions = useCallback(() => {
-    cleanup();
-    setState({
-      completions: {},
-      loading: {},
-      error: {},
-    });
-  }, [cleanup]);
-
-  return {
-    requestCompletions,
-    clearCompletions,
-    ...state,
-  };
-}