Add include_usage option to Tavily MCP

This commit is contained in:
Cursor Agent
2025-12-11 18:16:37 +00:00
parent 2be989a881
commit 3169015657
5 changed files with 125 additions and 23 deletions
+1
View File
@@ -13,6 +13,7 @@ The Tavily MCP server provides:
- Intelligent data extraction from web pages via the tavily-extract tool
- Powerful web mapping tool that creates a structured map of website
- Web crawler that systematically explores websites
- Optional `include_usage` flag on every tool to request Tavily credit usage data (omit it to preserve current behavior)
### 📚 Helpful Resources
+2 -2
View File
@@ -1,12 +1,12 @@
{
"name": "tavily-mcp",
"version": "0.2.10",
"version": "0.7.15",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "tavily-mcp",
"version": "0.2.10",
"version": "0.7.15",
"license": "MIT",
"dependencies": {
"@modelcontextprotocol/sdk": "0.6.0",
+2 -1
View File
@@ -1,6 +1,6 @@
{
"name": "tavily-mcp",
"version": "0.2.10",
"version": "0.7.15",
"description": "MCP server for advanced web search using Tavily",
"type": "module",
"bin": {
@@ -11,6 +11,7 @@
],
"scripts": {
"build": "tsc && node -e \"require('fs').chmodSync('build/index.js', '755')\"",
"test": "npm run build && node --test build/__tests__/*.js",
"prepare": "npm run build",
"watch": "tsc --watch",
"inspector": "npx @modelcontextprotocol/inspector build/index.js",
+34
View File
@@ -0,0 +1,34 @@
import test from "node:test";
import assert from "node:assert/strict";
process.env.TAVILY_API_KEY = process.env.TAVILY_API_KEY ?? "test-key";
process.env.TAVILY_MCP_SKIP_RUN = "true";
const clientModulePromise = import("../index.js");
async function createMockedClient() {
const { TavilyClient } = await clientModulePromise;
const client = new TavilyClient();
const calls: any[] = [];
client.setHttpClient({
post: async (_url: string, payload: any) => {
calls.push(payload);
return { data: { query: "", results: [] } };
}
});
return { client, calls };
}
test("search omits include_usage by default", async () => {
const { client, calls } = await createMockedClient();
await client.search({ query: "tavily" });
assert.equal("include_usage" in calls[0], false);
});
test("crawl includes include_usage when requested", async () => {
const { client, calls } = await createMockedClient();
await client.crawl({ url: "https://example.com" }, true);
assert.equal(calls[0].include_usage, true);
});
+86 -20
View File
@@ -35,6 +35,7 @@ interface TavilyResponse {
raw_content?: string;
favicon?: string;
}>;
usage?: Record<string, number>;
}
interface TavilyCrawlResponse {
@@ -45,15 +46,17 @@ interface TavilyCrawlResponse {
favicon?: string;
}>;
response_time: number;
usage?: Record<string, number>;
}
interface TavilyMapResponse {
base_url: string;
results: string[];
response_time: number;
usage?: Record<string, number>;
}
class TavilyClient {
export class TavilyClient {
// Core client properties
private server: Server;
private axiosInstance;
@@ -68,7 +71,7 @@ class TavilyClient {
this.server = new Server(
{
name: "tavily-mcp",
version: "0.2.10",
version: "0.7.15",
},
{
capabilities: {
@@ -90,6 +93,13 @@ class TavilyClient {
this.setupErrorHandling();
}
/**
* Overrides the Axios instance (intended for testing).
*/
setHttpClient(client: any): void {
this.axiosInstance = client;
}
private setupErrorHandling(): void {
this.server.onerror = (error) => {
console.error("[MCP Error]", error);
@@ -204,6 +214,10 @@ class TavilyClient {
type: "boolean",
description: "Whether to include the favicon URL for each result",
default: false,
},
include_usage: {
type: "boolean",
description: "Optional flag to request credit usage details in the response. Defaults to omitted.",
}
},
required: ["query"]
@@ -242,6 +256,10 @@ class TavilyClient {
description: "Whether to include the favicon URL for each result",
default: false,
},
include_usage: {
type: "boolean",
description: "Optional flag to request credit usage details in the response. Defaults to omitted.",
}
},
required: ["urls"]
}
@@ -312,6 +330,10 @@ class TavilyClient {
description: "Whether to include the favicon URL for each result",
default: false,
},
include_usage: {
type: "boolean",
description: "Optional flag to request credit usage details in the response. Defaults to omitted.",
}
},
required: ["url"]
}
@@ -364,6 +386,10 @@ class TavilyClient {
type: "boolean",
description: "Whether to return external links in the final response",
default: true
},
include_usage: {
type: "boolean",
description: "Optional flag to request credit usage details in the response. Defaults to omitted.",
}
},
required: ["url"]
@@ -376,7 +402,7 @@ class TavilyClient {
this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
try {
let response: TavilyResponse;
const args = request.params.arguments ?? {};
const args: any = request.params.arguments ?? {};
switch (request.params.name) {
case "tavily-search":
@@ -399,7 +425,7 @@ class TavilyClient {
exclude_domains: Array.isArray(args.exclude_domains) ? args.exclude_domains : [],
country: args.country,
include_favicon: args.include_favicon
});
}, args.include_usage);
break;
case "tavily-extract":
@@ -409,7 +435,7 @@ class TavilyClient {
include_images: args.include_images,
format: args.format,
include_favicon: args.include_favicon
});
}, args.include_usage);
break;
case "tavily-crawl":
@@ -425,7 +451,7 @@ class TavilyClient {
extract_depth: args.extract_depth,
format: args.format,
include_favicon: args.include_favicon
});
}, args.include_usage);
return {
content: [{
type: "text",
@@ -443,7 +469,7 @@ class TavilyClient {
select_paths: Array.isArray(args.select_paths) ? args.select_paths : [],
select_domains: Array.isArray(args.select_domains) ? args.select_domains : [],
allow_external: args.allow_external
});
}, args.include_usage);
return {
content: [{
type: "text",
@@ -486,14 +512,22 @@ class TavilyClient {
console.error("Tavily MCP server running on stdio");
}
async search(params: any): Promise<TavilyResponse> {
/**
* Executes a Tavily search request.
* @param params Search parameters forwarded to the API.
* @param include_usage Optional flag that, when true, includes credit usage metadata in the response.
*/
async search(params: any, include_usage: boolean | null = null): Promise<TavilyResponse> {
try {
const endpoint = this.baseURLs.search;
const searchParams = {
const searchParams: Record<string, any> = {
...params,
api_key: API_KEY,
};
if (include_usage !== null && include_usage !== undefined) {
searchParams.include_usage = include_usage;
}
const response = await this.axiosInstance.post(endpoint, searchParams);
return response.data;
@@ -507,12 +541,22 @@ class TavilyClient {
}
}
async extract(params: any): Promise<TavilyResponse> {
/**
* Extracts structured content from the Tavily API.
* @param params Extraction parameters forwarded to the API.
* @param include_usage Optional flag that, when true, includes credit usage metadata in the response.
*/
async extract(params: any, include_usage: boolean | null = null): Promise<TavilyResponse> {
try {
const response = await this.axiosInstance.post(this.baseURLs.extract, {
const payload: Record<string, any> = {
...params,
api_key: API_KEY
});
};
if (include_usage !== null && include_usage !== undefined) {
payload.include_usage = include_usage;
}
const response = await this.axiosInstance.post(this.baseURLs.extract, payload);
return response.data;
} catch (error: any) {
if (error.response?.status === 401) {
@@ -524,12 +568,22 @@ class TavilyClient {
}
}
async crawl(params: any): Promise<TavilyCrawlResponse> {
/**
* Kicks off a Tavily crawl request.
* @param params Crawl parameters forwarded to the API.
* @param include_usage Optional flag that, when true, includes credit usage metadata in the response.
*/
async crawl(params: any, include_usage: boolean | null = null): Promise<TavilyCrawlResponse> {
try {
const response = await this.axiosInstance.post(this.baseURLs.crawl, {
const payload: Record<string, any> = {
...params,
api_key: API_KEY
});
};
if (include_usage !== null && include_usage !== undefined) {
payload.include_usage = include_usage;
}
const response = await this.axiosInstance.post(this.baseURLs.crawl, payload);
return response.data;
} catch (error: any) {
if (error.response?.status === 401) {
@@ -541,12 +595,22 @@ class TavilyClient {
}
}
async map(params: any): Promise<TavilyMapResponse> {
/**
* Generates a Tavily site map.
* @param params Mapping parameters forwarded to the API.
* @param include_usage Optional flag that, when true, includes credit usage metadata in the response.
*/
async map(params: any, include_usage: boolean | null = null): Promise<TavilyMapResponse> {
try {
const response = await this.axiosInstance.post(this.baseURLs.map, {
const payload: Record<string, any> = {
...params,
api_key: API_KEY
});
};
if (include_usage !== null && include_usage !== undefined) {
payload.include_usage = include_usage;
}
const response = await this.axiosInstance.post(this.baseURLs.map, payload);
return response.data;
} catch (error: any) {
if (error.response?.status === 401) {
@@ -689,5 +753,7 @@ if (argv['list-tools']) {
}
// Otherwise start the server
const server = new TavilyClient();
server.run().catch(console.error);
if (!process.env.TAVILY_MCP_SKIP_RUN) {
const server = new TavilyClient();
server.run().catch(console.error);
}