Autocomplete on resource template variables
This commit is contained in:
Justin Spahr-Summers
parent
8635c6384f
commit
189cc84113
2 changed files with 294 additions and 74 deletions
|
|
@ -560,32 +560,34 @@ test("should handle request timeout", async () => {
|
|||
|
||||
describe("ResourceTemplate", () => {
|
||||
test("should create ResourceTemplate with string pattern", () => {
|
||||
const template = new ResourceTemplate("test://{category}/{id}", undefined);
|
||||
const template = new ResourceTemplate("test://{category}/{id}", {
|
||||
list: undefined,
|
||||
});
|
||||
expect(template.uriTemplate.toString()).toBe("test://{category}/{id}");
|
||||
expect(template.listCallback).toBeUndefined();
|
||||
});
|
||||
|
||||
test("should create ResourceTemplate with UriTemplate", () => {
|
||||
const uriTemplate = new UriTemplate("test://{category}/{id}");
|
||||
const template = new ResourceTemplate(uriTemplate, undefined);
|
||||
const template = new ResourceTemplate(uriTemplate, { list: undefined });
|
||||
expect(template.uriTemplate).toBe(uriTemplate);
|
||||
expect(template.listCallback).toBeUndefined();
|
||||
});
|
||||
|
||||
test("should create ResourceTemplate with list callback", async () => {
|
||||
const listCallback = jest.fn().mockResolvedValue({
|
||||
const list = jest.fn().mockResolvedValue({
|
||||
resources: [{ name: "Test", uri: "test://example" }],
|
||||
});
|
||||
|
||||
const template = new ResourceTemplate("test://{id}", listCallback);
|
||||
expect(template.listCallback).toBe(listCallback);
|
||||
const template = new ResourceTemplate("test://{id}", { list });
|
||||
expect(template.listCallback).toBe(list);
|
||||
|
||||
const abortController = new AbortController();
|
||||
const result = await template.listCallback?.({
|
||||
signal: abortController.signal,
|
||||
});
|
||||
expect(result?.resources).toHaveLength(1);
|
||||
expect(listCallback).toHaveBeenCalled();
|
||||
expect(list).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -1068,7 +1070,7 @@ describe("Server.resource", () => {
|
|||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{id}", undefined),
|
||||
new ResourceTemplate("test://resource/{id}", { list: undefined }),
|
||||
async () => ({
|
||||
contents: [
|
||||
{
|
||||
|
|
@ -1113,18 +1115,20 @@ describe("Server.resource", () => {
|
|||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{id}", async () => ({
|
||||
resources: [
|
||||
{
|
||||
name: "Resource 1",
|
||||
uri: "test://resource/1",
|
||||
},
|
||||
{
|
||||
name: "Resource 2",
|
||||
uri: "test://resource/2",
|
||||
},
|
||||
],
|
||||
})),
|
||||
new ResourceTemplate("test://resource/{id}", {
|
||||
list: async () => ({
|
||||
resources: [
|
||||
{
|
||||
name: "Resource 1",
|
||||
uri: "test://resource/1",
|
||||
},
|
||||
{
|
||||
name: "Resource 2",
|
||||
uri: "test://resource/2",
|
||||
},
|
||||
],
|
||||
}),
|
||||
}),
|
||||
async (uri) => ({
|
||||
contents: [
|
||||
{
|
||||
|
|
@ -1169,7 +1173,9 @@ describe("Server.resource", () => {
|
|||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{category}/{id}", undefined),
|
||||
new ResourceTemplate("test://resource/{category}/{id}", {
|
||||
list: undefined,
|
||||
}),
|
||||
async (uri, { category, id }) => ({
|
||||
contents: [
|
||||
{
|
||||
|
|
@ -1236,7 +1242,7 @@ describe("Server.resource", () => {
|
|||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{id}", undefined),
|
||||
new ResourceTemplate("test://resource/{id}", { list: undefined }),
|
||||
async () => ({
|
||||
contents: [
|
||||
{
|
||||
|
|
@ -1250,7 +1256,7 @@ describe("Server.resource", () => {
|
|||
expect(() => {
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{id}", undefined),
|
||||
new ResourceTemplate("test://resource/{id}", { list: undefined }),
|
||||
async () => ({
|
||||
contents: [
|
||||
{
|
||||
|
|
@ -1337,6 +1343,139 @@ describe("Server.resource", () => {
|
|||
),
|
||||
).rejects.toThrow(/Resource test:\/\/nonexistent not found/);
|
||||
});
|
||||
|
||||
test("should support completion of resource template parameters", async () => {
|
||||
const server = new Server({
|
||||
name: "test server",
|
||||
version: "1.0",
|
||||
});
|
||||
|
||||
const client = new Client(
|
||||
{
|
||||
name: "test client",
|
||||
version: "1.0",
|
||||
},
|
||||
{
|
||||
capabilities: {
|
||||
resources: {},
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{category}", {
|
||||
list: undefined,
|
||||
complete: {
|
||||
category: () => ["books", "movies", "music"],
|
||||
},
|
||||
}),
|
||||
async () => ({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://resource/test",
|
||||
text: "Test content",
|
||||
},
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
const [clientTransport, serverTransport] =
|
||||
InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
const result = await client.request(
|
||||
{
|
||||
method: "completion/complete",
|
||||
params: {
|
||||
ref: {
|
||||
type: "ref/resource",
|
||||
uri: "test://resource/{category}",
|
||||
},
|
||||
argument: {
|
||||
name: "category",
|
||||
value: "",
|
||||
},
|
||||
},
|
||||
},
|
||||
CompleteResultSchema,
|
||||
);
|
||||
|
||||
expect(result.completion.values).toEqual(["books", "movies", "music"]);
|
||||
expect(result.completion.total).toBe(3);
|
||||
});
|
||||
|
||||
test("should support filtered completion of resource template parameters", async () => {
|
||||
const server = new Server({
|
||||
name: "test server",
|
||||
version: "1.0",
|
||||
});
|
||||
|
||||
const client = new Client(
|
||||
{
|
||||
name: "test client",
|
||||
version: "1.0",
|
||||
},
|
||||
{
|
||||
capabilities: {
|
||||
resources: {},
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
server.resource(
|
||||
"test",
|
||||
new ResourceTemplate("test://resource/{category}", {
|
||||
list: undefined,
|
||||
complete: {
|
||||
category: (test) =>
|
||||
["books", "movies", "music"].filter((value) =>
|
||||
value.startsWith(test),
|
||||
),
|
||||
},
|
||||
}),
|
||||
async () => ({
|
||||
contents: [
|
||||
{
|
||||
uri: "test://resource/test",
|
||||
text: "Test content",
|
||||
},
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
const [clientTransport, serverTransport] =
|
||||
InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
const result = await client.request(
|
||||
{
|
||||
method: "completion/complete",
|
||||
params: {
|
||||
ref: {
|
||||
type: "ref/resource",
|
||||
uri: "test://resource/{category}",
|
||||
},
|
||||
argument: {
|
||||
name: "category",
|
||||
value: "m",
|
||||
},
|
||||
},
|
||||
},
|
||||
CompleteResultSchema,
|
||||
);
|
||||
|
||||
expect(result.completion.values).toEqual(["movies", "music"]);
|
||||
expect(result.completion.total).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Server.prompt", () => {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import {
|
|||
CallToolRequestSchema,
|
||||
CallToolResult,
|
||||
ClientCapabilities,
|
||||
CompleteRequest,
|
||||
CompleteRequestSchema,
|
||||
CompleteResult,
|
||||
CreateMessageRequest,
|
||||
|
|
@ -48,10 +49,12 @@ import {
|
|||
Notification,
|
||||
Prompt,
|
||||
PromptArgument,
|
||||
PromptReference,
|
||||
ReadResourceRequestSchema,
|
||||
ReadResourceResult,
|
||||
Request,
|
||||
Resource,
|
||||
ResourceReference,
|
||||
ResourceUpdatedNotification,
|
||||
Result,
|
||||
ServerCapabilities,
|
||||
|
|
@ -500,6 +503,86 @@ export class Server<
|
|||
this.setToolRequestHandlers();
|
||||
}
|
||||
|
||||
private setCompletionRequestHandler() {
|
||||
this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value);
|
||||
|
||||
this.setRequestHandler(
|
||||
CompleteRequestSchema,
|
||||
async (request): Promise<CompleteResult> => {
|
||||
switch (request.params.ref.type) {
|
||||
case "ref/prompt":
|
||||
return this.handlePromptCompletion(request, request.params.ref);
|
||||
|
||||
case "ref/resource":
|
||||
return this.handleResourceCompletion(request, request.params.ref);
|
||||
|
||||
default:
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Invalid completion reference: ${request.params.ref}`,
|
||||
);
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
private async handlePromptCompletion(
|
||||
request: CompleteRequest,
|
||||
ref: PromptReference,
|
||||
): Promise<CompleteResult> {
|
||||
const prompt = this._registeredPrompts[ref.name];
|
||||
if (!prompt) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Prompt ${request.params.ref.name} not found`,
|
||||
);
|
||||
}
|
||||
|
||||
if (!prompt.argsSchema) {
|
||||
return EMPTY_COMPLETION_RESULT;
|
||||
}
|
||||
|
||||
const field = prompt.argsSchema.shape[request.params.argument.name];
|
||||
if (!(field instanceof Completable)) {
|
||||
return EMPTY_COMPLETION_RESULT;
|
||||
}
|
||||
|
||||
const def: CompletableDef<ZodString> = field._def;
|
||||
const suggestions = await def.complete(request.params.argument.value);
|
||||
return createCompletionResult(suggestions);
|
||||
}
|
||||
|
||||
private async handleResourceCompletion(
|
||||
request: CompleteRequest,
|
||||
ref: ResourceReference,
|
||||
): Promise<CompleteResult> {
|
||||
const template = Object.values(this._registeredResourceTemplates).find(
|
||||
(t) => t.resourceTemplate.uriTemplate.toString() === ref.uri,
|
||||
);
|
||||
|
||||
if (!template) {
|
||||
if (this._registeredResources[ref.uri]) {
|
||||
// Attempting to autocomplete a fixed resource URI is not an error in the spec (but probably should be).
|
||||
return EMPTY_COMPLETION_RESULT;
|
||||
}
|
||||
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Resource template ${request.params.ref.uri} not found`,
|
||||
);
|
||||
}
|
||||
|
||||
const completer = template.resourceTemplate.completeCallback(
|
||||
request.params.argument.name,
|
||||
);
|
||||
if (!completer) {
|
||||
return EMPTY_COMPLETION_RESULT;
|
||||
}
|
||||
|
||||
const suggestions = await completer(request.params.argument.value);
|
||||
return createCompletionResult(suggestions);
|
||||
}
|
||||
|
||||
private setResourceRequestHandlers() {
|
||||
this.assertCanSetRequestHandler(
|
||||
ListResourcesRequestSchema.shape.method.value,
|
||||
|
|
@ -588,6 +671,8 @@ export class Server<
|
|||
);
|
||||
},
|
||||
);
|
||||
|
||||
this.setCompletionRequestHandler();
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -708,7 +793,6 @@ export class Server<
|
|||
ListPromptsRequestSchema.shape.method.value,
|
||||
);
|
||||
this.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value);
|
||||
this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value);
|
||||
|
||||
this.registerCapabilities({
|
||||
prompts: {},
|
||||
|
|
@ -763,55 +847,7 @@ export class Server<
|
|||
},
|
||||
);
|
||||
|
||||
this.setRequestHandler(
|
||||
CompleteRequestSchema,
|
||||
async (request): Promise<CompleteResult> => {
|
||||
if (request.params.ref.type !== "ref/prompt") {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
"Only prompt completions are supported",
|
||||
);
|
||||
}
|
||||
|
||||
const prompt = this._registeredPrompts[request.params.ref.name];
|
||||
if (!prompt) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Prompt ${request.params.ref.name} not found`,
|
||||
);
|
||||
}
|
||||
|
||||
if (!prompt.argsSchema) {
|
||||
return {
|
||||
completion: {
|
||||
values: [],
|
||||
hasMore: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const field = prompt.argsSchema.shape[request.params.argument.name];
|
||||
if (!(field instanceof Completable)) {
|
||||
return {
|
||||
completion: {
|
||||
values: [],
|
||||
hasMore: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const def: CompletableDef<ZodString> = field._def;
|
||||
const completer = def.complete;
|
||||
const suggestions = await completer(request.params.argument.value);
|
||||
return {
|
||||
completion: {
|
||||
values: suggestions.slice(0, 100),
|
||||
total: suggestions.length,
|
||||
hasMore: suggestions.length > 100,
|
||||
},
|
||||
};
|
||||
},
|
||||
);
|
||||
this.setCompletionRequestHandler();
|
||||
}
|
||||
|
||||
prompt(name: string, cb: PromptCallback): void;
|
||||
|
|
@ -854,6 +890,13 @@ export class Server<
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A callback to complete one variable within a resource template's URI template.
|
||||
*/
|
||||
export type CompleteResourceTemplateCallback = (
|
||||
value: string,
|
||||
) => string[] | Promise<string[]>;
|
||||
|
||||
/**
|
||||
* A resource template combines a URI pattern with optional functionality to enumerate
|
||||
* all resources matching that pattern.
|
||||
|
|
@ -863,7 +906,19 @@ export class ResourceTemplate {
|
|||
|
||||
constructor(
|
||||
uriTemplate: string | UriTemplate,
|
||||
private _listCallback: ListResourcesCallback | undefined,
|
||||
private _callbacks: {
|
||||
/**
|
||||
* A callback to list all resources matching this template. This is required to specified, even if `undefined`, to avoid accidentally forgetting resource listing.
|
||||
*/
|
||||
list: ListResourcesCallback | undefined;
|
||||
|
||||
/**
|
||||
* An optional callback to autocomplete variables within the URI template. Useful for clients and users to discover possible values.
|
||||
*/
|
||||
complete?: {
|
||||
[variable: string]: CompleteResourceTemplateCallback;
|
||||
};
|
||||
},
|
||||
) {
|
||||
this._uriTemplate =
|
||||
typeof uriTemplate === "string"
|
||||
|
|
@ -882,7 +937,16 @@ export class ResourceTemplate {
|
|||
* Gets the list callback, if one was provided.
|
||||
*/
|
||||
get listCallback(): ListResourcesCallback | undefined {
|
||||
return this._listCallback;
|
||||
return this._callbacks.list;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the callback for completing a specific URI template variable, if one was provided.
|
||||
*/
|
||||
completeCallback(
|
||||
variable: string,
|
||||
): CompleteResourceTemplateCallback | undefined {
|
||||
return this._callbacks.complete?.[variable];
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -982,3 +1046,20 @@ function promptArgumentsFromSchema(
|
|||
}),
|
||||
);
|
||||
}
|
||||
|
||||
function createCompletionResult(suggestions: string[]): CompleteResult {
|
||||
return {
|
||||
completion: {
|
||||
values: suggestions.slice(0, 100),
|
||||
total: suggestions.length,
|
||||
hasMore: suggestions.length > 100,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const EMPTY_COMPLETION_RESULT: CompleteResult = {
|
||||
completion: {
|
||||
values: [],
|
||||
hasMore: false,
|
||||
},
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue