Skip to content

Commit 5af7015

Browse files
pilartomaspilartomas
committed
feat(tools): add Model Context Protocol tool
Signed-off-by: Tomas Pilar <[email protected]>
1 parent fb2153c commit 5af7015

File tree

6 files changed

+337
-10
lines changed

6 files changed

+337
-10
lines changed

‎examples/tools/mcp.ts

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/**
2+
* Copyright 2025 IBM Corp.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
18+
import { MCPResourceTool } from "bee-agent-framework/tools/mcp/mcpResource";
19+
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
20+
import { OllamaChatLLM } from "bee-agent-framework/adapters/ollama/chat";
21+
import { BeeAgent } from "bee-agent-framework/agents/bee/agent";
22+
import { UnconstrainedMemory } from "bee-agent-framework/memory/unconstrainedMemory";
23+
24+
const client = new Client(
25+
{
26+
name: "test-client",
27+
version: "1.0.0",
28+
},
29+
{
30+
capabilities: {},
31+
},
32+
);
33+
await client.connect(
34+
new StdioClientTransport({
35+
command: "npx",
36+
args: ["-y", "@modelcontextprotocol/server-everything"],
37+
}),
38+
);
39+
40+
const agent = new BeeAgent({
41+
llm: new OllamaChatLLM(),
42+
memory: new UnconstrainedMemory(),
43+
tools: [
44+
new MCPResourceTool({
45+
client,
46+
}),
47+
],
48+
});
49+
await agent.run({ prompt: "Read contents the first resource" }).observe((emitter) => {
50+
emitter.on("update", async ({ data, update, meta }) => {
51+
console.log(`Agent (${update.key}) 🤖 : `, update.value);
52+
});
53+
});
54+
55+
await client.close();

‎package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@
171171
"@ai-zen/node-fetch-event-source": "^2.1.4",
172172
"@connectrpc/connect": "^1.6.1",
173173
"@connectrpc/connect-node": "^1.6.1",
174+
"@modelcontextprotocol/sdk": "^1.0.4",
174175
"@opentelemetry/api": "^1.9.0",
175176
"@streamparser/json": "^0.0.21",
176177
"ajv": "^8.17.1",

‎src/tools/mcp/base.ts

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/**
2+
* Copyright 2025 IBM Corp.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Client as MCPClient } from "@modelcontextprotocol/sdk/client/index.js";
18+
19+
import { Tool, JSONToolOutput } from "@/tools/base.js";
20+
21+
export interface MCPToolInput {
22+
client: MCPClient;
23+
}
24+
25+
export class MCPToolOutput<T> extends JSONToolOutput<T> {}
26+
27+
export abstract class MCPTool<T> extends Tool<MCPToolOutput<T>> {
28+
public readonly client: MCPClient;
29+
30+
constructor({ client, ...options }: MCPToolInput) {
31+
super(options);
32+
this.client = client;
33+
}
34+
35+
protected async paginateUntilLimit<T>(
36+
fetcher: (
37+
params: { cursor?: string },
38+
options: { signal?: AbortSignal },
39+
) => Promise<{ nextCursor?: string; items: T[] }>,
40+
limit: number,
41+
{ signal }: { signal?: AbortSignal } = {},
42+
) {
43+
const items: T[] = [];
44+
let cursor: string | undefined = undefined;
45+
while (items.length < limit) {
46+
const result = await fetcher({ cursor }, { signal });
47+
items.push(...result.items);
48+
cursor = result.nextCursor;
49+
if (!cursor) {
50+
break;
51+
}
52+
}
53+
return items.slice(0, limit);
54+
}
55+
}

‎src/tools/mcp/mcpResource.test.ts

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/**
2+
* Copyright 2025 IBM Corp.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
18+
import {
19+
ListResourcesRequestSchema,
20+
ReadResourceRequestSchema,
21+
} from "@modelcontextprotocol/sdk/types.js";
22+
import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
23+
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
24+
import { MCPResourceTool } from "./mcpResource.js";
25+
import { entries } from "remeda";
26+
27+
const resources = {
28+
"file:///animals.txt": {
29+
name: "Info about animals",
30+
mimeType: "text/plain",
31+
text: "Lion, shark, bizon",
32+
},
33+
"file:///planets.txt": {
34+
name: "Info about planets",
35+
mimeType: "text/plain",
36+
text: "Earth, Mars, Venus",
37+
},
38+
} as const;
39+
40+
describe("MCPResourceTool", () => {
41+
const server = new Server(
42+
{
43+
name: "test-server",
44+
version: "1.0.0",
45+
},
46+
{
47+
capabilities: {
48+
resources: {},
49+
},
50+
},
51+
);
52+
53+
const client = new Client(
54+
{
55+
name: "test-client",
56+
version: "1.0.0",
57+
},
58+
{
59+
capabilities: {},
60+
},
61+
);
62+
63+
let instance: MCPResourceTool;
64+
65+
beforeAll(async () => {
66+
server.setRequestHandler(ListResourcesRequestSchema, async () => {
67+
return {
68+
resources: entries(resources).map(([uri, { name }]) => ({ uri, name })),
69+
};
70+
});
71+
server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
72+
const resource = resources[request.params.uri as keyof typeof resources];
73+
if (!resource) {
74+
throw new Error("Resource not found");
75+
}
76+
return {
77+
contents: [
78+
{
79+
uri: request.params.uri,
80+
mimeType: resource.mimeType,
81+
text: resource.text,
82+
},
83+
],
84+
};
85+
});
86+
87+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
88+
await server.connect(serverTransport);
89+
await client.connect(clientTransport);
90+
});
91+
92+
beforeEach(() => {
93+
instance = new MCPResourceTool({ client });
94+
});
95+
96+
describe("MCPTool", () => {
97+
it("Runs", async () => {
98+
const uri = "file:///planets.txt";
99+
const response = await instance.run({ uri });
100+
expect(response.result.contents[0].text).toBe(resources[uri].text);
101+
});
102+
});
103+
104+
afterAll(async () => {
105+
await client.close();
106+
await server.close();
107+
});
108+
});

‎src/tools/mcp/mcpResource.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/**
2+
* Copyright 2025 IBM Corp.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import { ToolEmitter, ToolInput } from "@/tools/base.js";
18+
import { z } from "zod";
19+
import { Emitter } from "@/emitter/emitter.js";
20+
import { ReadResourceResult } from "@modelcontextprotocol/sdk/types.js";
21+
import { MCPTool, MCPToolInput, MCPToolOutput } from "./base.js";
22+
23+
export interface MCPResourceToolInput extends MCPToolInput {
24+
resourceLimit?: number;
25+
}
26+
27+
export class MCPResourceTool extends MCPTool<ReadResourceResult> {
28+
name = "MCP Resource";
29+
description = `An MCP Resource tool provides ability to read resources. Use it to read contents of available resources.`;
30+
31+
public readonly emitter: ToolEmitter<ToolInput<this>, MCPToolOutput<ReadResourceResult>> =
32+
Emitter.root.child({
33+
namespace: ["tool", "mcp", "resource"],
34+
creator: this,
35+
});
36+
37+
public readonly resourceLimit: number;
38+
39+
constructor({ resourceLimit = Infinity, ...options }: MCPResourceToolInput) {
40+
super(options);
41+
this.resourceLimit = resourceLimit;
42+
}
43+
44+
async inputSchema() {
45+
const resources = await this.listResources().catch(() => []); // ignore errors, e.g. MCP server might not have resource capability
46+
return z.object({
47+
uri: z
48+
.string()
49+
.describe(
50+
`URI of the resource to read, ${resources.length > 0 ? `available resources are:\n\n${resources.map(({ uri, name, description }) => JSON.stringify({ uri, name, description })).join("\n\n")}` : "no resources available at the moment"}.`,
51+
),
52+
});
53+
}
54+
55+
protected async _run({ uri }: ToolInput<this>, { signal }: { signal?: AbortSignal }) {
56+
const result = await this.client.readResource({ uri }, { signal });
57+
return new MCPToolOutput(result);
58+
}
59+
60+
public async listResources({ signal }: { signal?: AbortSignal } = {}) {
61+
return await this.paginateUntilLimit(
62+
({ cursor }, { signal }) =>
63+
this.client
64+
.listResources({ cursor }, { signal })
65+
.then(({ resources, nextCursor }) => ({ items: resources, nextCursor })),
66+
this.resourceLimit,
67+
{ signal },
68+
);
69+
}
70+
}

0 commit comments

Comments
 (0)