From 7f713fe40e574e9d273c2d76524c478cb20633b7 Mon Sep 17 00:00:00 2001 From: Gavin Aboulhosn Date: Wed, 12 Feb 2025 19:05:51 -0500 Subject: [PATCH] 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 --- client/src/App.tsx | 43 +------ client/src/components/PromptsTab.tsx | 16 +-- client/src/components/ResourcesTab.tsx | 16 +-- client/src/lib/hooks/useCompletionState.ts | 76 +++++++++++++ client/src/lib/hooks/useConnection.ts | 69 ++++++++++- client/src/lib/useCompletion.ts | 126 --------------------- 6 files changed, 166 insertions(+), 180 deletions(-) create mode 100644 client/src/lib/hooks/useCompletionState.ts delete mode 100644 client/src/lib/useCompletion.ts diff --git a/client/src/App.tsx b/client/src/App.tsx index 1d6dbe4..a9adea5 100644 --- 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} diff --git a/client/src/components/PromptsTab.tsx b/client/src/components/PromptsTab.tsx index 6799290..f88b16f 100644 --- 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) => void; selectedPrompt: Prompt | null; setSelectedPrompt: (prompt: Prompt) => void; - onComplete: ( - ref: PromptReference, + handleCompletion: ( + ref: PromptReference | ResourceReference, argName: string, value: string, ) => Promise; + completionsSupported: boolean; promptContent: string; nextCursor: ListPromptsResult["nextCursor"]; error: string | null; }) => { const [promptArgs, setPromptArgs] = useState>({}); - const { completions, clearCompletions, requestCompletions } = useCompletions({ - onComplete, - }); + const { completions, clearCompletions, requestCompletions } = + useCompletionState(handleCompletion, completionsSupported); useEffect(() => { clearCompletions(); diff --git a/client/src/components/ResourcesTab.tsx b/client/src/components/ResourcesTab.tsx index f661325..93127a9 100644 --- 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; + 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(); diff --git a/client/src/lib/hooks/useCompletionState.ts b/client/src/lib/hooks/useCompletionState.ts new file mode 100644 index 0000000..1bbbf3f --- /dev/null +++ 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; + loading: Record; + error: Record; +} + +export function useCompletionState( + handleCompletion: ( + ref: ResourceReference | PromptReference, + argName: string, + value: string, + ) => Promise, + completionsSupported: boolean = true, +) { + const [state, setState] = useState({ + 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, + }; +} diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts index 6c42c3f..b72c92f 100644 --- 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 ( request: ClientRequest, schema: T, - ) => { + options?: RequestOptions, + ): Promise> => { 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; + } + 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 => { + 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, }; } diff --git a/client/src/lib/useCompletion.ts b/client/src/lib/useCompletion.ts deleted file mode 100644 index 621dd8c..0000000 --- a/client/src/lib/useCompletion.ts +++ /dev/null @@ -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 PromiseLike>( - func: T, - wait: number, -): (...args: Parameters) => void { - let timeout: ReturnType; - return function (...args: Parameters) { - clearTimeout(timeout); - timeout = setTimeout(() => func(...args), wait); - }; -} - -interface UseCompletionsOptions { - onComplete: ( - ref: T, - argName: string, - value: string, - signal?: AbortSignal, - ) => Promise; - debounceMs?: number; -} - -interface CompletionState { - completions: Record; - loading: Record; - error: Record; -} - -export function useCompletions({ - onComplete, - debounceMs = 300, -}: UseCompletionsOptions) { - const [state, setState] = useState({ - completions: {}, - loading: {}, - error: {}, - }); - - const completeRef = useRef(onComplete); - completeRef.current = onComplete; - - const abortControllerRef = useRef(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, - }; -}