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
This commit is contained in:
Gavin Aboulhosn
2025-02-12 19:05:51 -05:00
parent c66feff37d
commit 7f713fe40e
6 changed files with 166 additions and 180 deletions

View File

@@ -15,9 +15,6 @@ import {
Root, Root,
ServerNotification, ServerNotification,
Tool, Tool,
PromptReference,
ResourceReference,
CompleteResultSchema,
} from "@modelcontextprotocol/sdk/types.js"; } from "@modelcontextprotocol/sdk/types.js";
import React, { Suspense, useEffect, useRef, useState } from "react"; import React, { Suspense, useEffect, useRef, useState } from "react";
import { useConnection } from "./lib/hooks/useConnection"; import { useConnection } from "./lib/hooks/useConnection";
@@ -154,6 +151,8 @@ const App = () => {
requestHistory, requestHistory,
makeRequest: makeConnectionRequest, makeRequest: makeConnectionRequest,
sendNotification, sendNotification,
handleCompletion,
completionsSupported,
connect: connectMcpServer, connect: connectMcpServer,
} = useConnection({ } = useConnection({
transportType, 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 listResources = async () => {
const response = await makeRequest( const response = await makeRequest(
{ {
@@ -518,7 +485,8 @@ const App = () => {
clearError("resources"); clearError("resources");
setSelectedResource(resource); setSelectedResource(resource);
}} }}
onComplete={handleCompletion} handleCompletion={handleCompletion}
completionsSupported={completionsSupported}
resourceContent={resourceContent} resourceContent={resourceContent}
nextCursor={nextResourceCursor} nextCursor={nextResourceCursor}
nextTemplateCursor={nextResourceTemplateCursor} nextTemplateCursor={nextResourceTemplateCursor}
@@ -543,7 +511,8 @@ const App = () => {
clearError("prompts"); clearError("prompts");
setSelectedPrompt(prompt); setSelectedPrompt(prompt);
}} }}
onComplete={handleCompletion} handleCompletion={handleCompletion}
completionsSupported={completionsSupported}
promptContent={promptContent} promptContent={promptContent}
nextCursor={nextPromptCursor} nextCursor={nextPromptCursor}
error={errors.prompts} error={errors.prompts}

View File

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

View File

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

View File

@@ -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,
};
}

View File

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

View File

@@ -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,
};
}