refactor(completions): restore debouncing and improve MCP error handling

This commit is contained in:
Gavin Aboulhosn
2025-02-12 20:19:07 -05:00
parent 5b451a7cfe
commit d9df5ff860
2 changed files with 101 additions and 47 deletions

View File

@@ -1,10 +1,24 @@
import { useState, useCallback, useEffect } from "react"; import { useState, useCallback, useEffect, useRef } from "react";
import { ResourceReference, PromptReference } from "@modelcontextprotocol/sdk/types.js"; import {
ResourceReference,
PromptReference,
} from "@modelcontextprotocol/sdk/types.js";
interface CompletionState { interface CompletionState {
completions: Record<string, string[]>; completions: Record<string, string[]>;
loading: Record<string, boolean>; loading: Record<string, boolean>;
error: Record<string, string | null>; }
// 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);
};
} }
export function useCompletionState( export function useCompletionState(
@@ -12,52 +26,90 @@ export function useCompletionState(
ref: ResourceReference | PromptReference, ref: ResourceReference | PromptReference,
argName: string, argName: string,
value: string, value: string,
signal?: AbortSignal,
) => Promise<string[]>, ) => Promise<string[]>,
completionsSupported: boolean = true, completionsSupported: boolean = true,
debounceMs: number = 300,
) { ) {
const [state, setState] = useState<CompletionState>({ const [state, setState] = useState<CompletionState>({
completions: {}, completions: {},
loading: {}, loading: {},
error: {},
}); });
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 clearCompletions = useCallback(() => { const clearCompletions = useCallback(() => {
cleanup();
setState({ setState({
completions: {}, completions: {},
loading: {}, loading: {},
error: {},
}); });
}, []); }, [cleanup]);
const requestCompletions = useCallback( const requestCompletions = useCallback(
async (ref: ResourceReference | PromptReference, argName: string, value: string) => { debounce(
if (!completionsSupported) { async (
return; ref: ResourceReference | PromptReference,
} argName: string,
value: string,
) => {
if (!completionsSupported) {
return;
}
setState((prev) => ({ cleanup();
...prev,
loading: { ...prev.loading, [argName]: true }, const abortController = new AbortController();
error: { ...prev.error, [argName]: null }, abortControllerRef.current = abortController;
}));
try {
const values = await handleCompletion(ref, argName, value);
setState((prev) => ({ setState((prev) => ({
...prev, ...prev,
completions: { ...prev.completions, [argName]: values }, loading: { ...prev.loading, [argName]: true },
loading: { ...prev.loading, [argName]: false },
})); }));
} catch (err) {
const error = err instanceof Error ? err.message : String(err); try {
setState((prev) => ({ const values = await handleCompletion(
...prev, ref,
loading: { ...prev.loading, [argName]: false }, argName,
error: { ...prev.error, [argName]: error }, value,
})); abortController.signal,
} );
},
[handleCompletion, completionsSupported], if (!abortController.signal.aborted) {
setState((prev) => ({
...prev,
completions: { ...prev.completions, [argName]: values },
loading: { ...prev.loading, [argName]: false },
}));
}
} catch (err) {
if (!abortController.signal.aborted) {
setState((prev) => ({
...prev,
loading: { ...prev.loading, [argName]: false },
}));
}
} finally {
if (abortControllerRef.current === abortController) {
abortControllerRef.current = null;
}
}
},
debounceMs,
),
[handleCompletion, completionsSupported, cleanup, debounceMs],
); );
// Clear completions when support status changes // Clear completions when support status changes

View File

@@ -16,6 +16,7 @@ import {
ResourceReference, ResourceReference,
McpError, McpError,
CompleteResultSchema, CompleteResultSchema,
ErrorCode,
} from "@modelcontextprotocol/sdk/types.js"; } from "@modelcontextprotocol/sdk/types.js";
import { useState } from "react"; import { useState } from "react";
import { toast } from "react-toastify"; import { toast } from "react-toastify";
@@ -43,6 +44,7 @@ interface UseConnectionOptions {
interface RequestOptions { interface RequestOptions {
signal?: AbortSignal; signal?: AbortSignal;
timeout?: number; timeout?: number;
suppressToast?: boolean;
} }
export function useConnection({ export function useConnection({
@@ -111,18 +113,10 @@ export function useConnection({
return response; return response;
} catch (e: unknown) { } catch (e: unknown) {
// Check for Method not found error specifically for completions if (!options?.suppressToast) {
if ( const errorString = (e as Error).message ?? String(e);
request.method === "completion/complete" && toast.error(errorString);
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; throw e;
} }
}; };
@@ -151,32 +145,40 @@ export function useConnection({
try { try {
const response = await makeRequest(request, CompleteResultSchema, { const response = await makeRequest(request, CompleteResultSchema, {
signal, signal,
suppressToast: true,
}); });
return response?.completion.values || []; return response?.completion.values || [];
} catch (e: unknown) { } catch (e: unknown) {
const errorMessage = e instanceof Error ? e.message : String(e); // Disable completions silently if the server doesn't support them.
pushHistory(request, { error: errorMessage }); // See https://github.com/modelcontextprotocol/specification/discussions/122
if (e instanceof McpError && e.code === ErrorCode.MethodNotFound) {
if (e instanceof McpError && e.code === -32601) {
setCompletionsSupported(false); setCompletionsSupported(false);
return []; return [];
} }
toast.error(errorMessage); // Unexpected errors - show toast and rethrow
toast.error(e instanceof Error ? e.message : String(e));
throw e; throw e;
} }
}; };
const sendNotification = async (notification: ClientNotification) => { const sendNotification = async (notification: ClientNotification) => {
if (!mcpClient) { if (!mcpClient) {
throw new Error("MCP client not connected"); const error = new Error("MCP client not connected");
toast.error(error.message);
throw error;
} }
try { try {
await mcpClient.notification(notification); await mcpClient.notification(notification);
// Log successful notifications
pushHistory(notification); pushHistory(notification);
} catch (e: unknown) { } catch (e: unknown) {
toast.error((e as Error).message ?? String(e)); if (e instanceof McpError) {
// Log MCP protocol errors
pushHistory(notification, { error: e.message });
}
toast.error(e instanceof Error ? e.message : String(e));
throw e; throw e;
} }
}; };