This is page 6 of 6. Use http://codebase.md/bsmi021/mcp-gemini-server?page={x} to view the full context.
# Directory Structure
```
├── .env.example
├── .eslintignore
├── .eslintrc.json
├── .gitignore
├── .prettierrc.json
├── Dockerfile
├── LICENSE
├── package-lock.json
├── package.json
├── README.md
├── review-prompt.txt
├── scripts
│ ├── gemini-review.sh
│ └── run-with-health-check.sh
├── smithery.yaml
├── src
│ ├── config
│ │ └── ConfigurationManager.ts
│ ├── createServer.ts
│ ├── index.ts
│ ├── resources
│ │ └── system-prompt.md
│ ├── server.ts
│ ├── services
│ │ ├── ExampleService.ts
│ │ ├── gemini
│ │ │ ├── GeminiCacheService.ts
│ │ │ ├── GeminiChatService.ts
│ │ │ ├── GeminiContentService.ts
│ │ │ ├── GeminiGitDiffService.ts
│ │ │ ├── GeminiPromptTemplates.ts
│ │ │ ├── GeminiTypes.ts
│ │ │ ├── GeminiUrlContextService.ts
│ │ │ ├── GeminiValidationSchemas.ts
│ │ │ ├── GitHubApiService.ts
│ │ │ ├── GitHubUrlParser.ts
│ │ │ └── ModelMigrationService.ts
│ │ ├── GeminiService.ts
│ │ ├── index.ts
│ │ ├── mcp
│ │ │ ├── index.ts
│ │ │ └── McpClientService.ts
│ │ ├── ModelSelectionService.ts
│ │ ├── session
│ │ │ ├── index.ts
│ │ │ ├── InMemorySessionStore.ts
│ │ │ ├── SessionStore.ts
│ │ │ └── SQLiteSessionStore.ts
│ │ └── SessionService.ts
│ ├── tools
│ │ ├── exampleToolParams.ts
│ │ ├── geminiCacheParams.ts
│ │ ├── geminiCacheTool.ts
│ │ ├── geminiChatParams.ts
│ │ ├── geminiChatTool.ts
│ │ ├── geminiCodeReviewParams.ts
│ │ ├── geminiCodeReviewTool.ts
│ │ ├── geminiGenerateContentConsolidatedParams.ts
│ │ ├── geminiGenerateContentConsolidatedTool.ts
│ │ ├── geminiGenerateImageParams.ts
│ │ ├── geminiGenerateImageTool.ts
│ │ ├── geminiGenericParamSchemas.ts
│ │ ├── geminiRouteMessageParams.ts
│ │ ├── geminiRouteMessageTool.ts
│ │ ├── geminiUrlAnalysisTool.ts
│ │ ├── index.ts
│ │ ├── mcpClientParams.ts
│ │ ├── mcpClientTool.ts
│ │ ├── registration
│ │ │ ├── index.ts
│ │ │ ├── registerAllTools.ts
│ │ │ ├── ToolAdapter.ts
│ │ │ └── ToolRegistry.ts
│ │ ├── schemas
│ │ │ ├── BaseToolSchema.ts
│ │ │ ├── CommonSchemas.ts
│ │ │ ├── index.ts
│ │ │ ├── ToolSchemas.ts
│ │ │ └── writeToFileParams.ts
│ │ └── writeToFileTool.ts
│ ├── types
│ │ ├── exampleServiceTypes.ts
│ │ ├── geminiServiceTypes.ts
│ │ ├── gitdiff-parser.d.ts
│ │ ├── googleGenAI.d.ts
│ │ ├── googleGenAITypes.ts
│ │ ├── index.ts
│ │ ├── micromatch.d.ts
│ │ ├── modelcontextprotocol-sdk.d.ts
│ │ ├── node-fetch.d.ts
│ │ └── serverTypes.ts
│ └── utils
│ ├── errors.ts
│ ├── filePathSecurity.ts
│ ├── FileSecurityService.ts
│ ├── geminiErrors.ts
│ ├── healthCheck.ts
│ ├── index.ts
│ ├── logger.ts
│ ├── RetryService.ts
│ ├── ToolError.ts
│ └── UrlSecurityService.ts
├── tests
│ ├── .env.test.example
│ ├── basic-router.test.vitest.ts
│ ├── e2e
│ │ ├── clients
│ │ │ └── mcp-test-client.ts
│ │ ├── README.md
│ │ └── streamableHttpTransport.test.vitest.ts
│ ├── integration
│ │ ├── dummyMcpServerSse.ts
│ │ ├── dummyMcpServerStdio.ts
│ │ ├── geminiRouterIntegration.test.vitest.ts
│ │ ├── mcpClientIntegration.test.vitest.ts
│ │ ├── multiModelIntegration.test.vitest.ts
│ │ └── urlContextIntegration.test.vitest.ts
│ ├── tsconfig.test.json
│ ├── unit
│ │ ├── config
│ │ │ └── ConfigurationManager.multimodel.test.vitest.ts
│ │ ├── server
│ │ │ └── transportLogic.test.vitest.ts
│ │ ├── services
│ │ │ ├── gemini
│ │ │ │ ├── GeminiChatService.test.vitest.ts
│ │ │ │ ├── GeminiGitDiffService.test.vitest.ts
│ │ │ │ ├── geminiImageGeneration.test.vitest.ts
│ │ │ │ ├── GeminiPromptTemplates.test.vitest.ts
│ │ │ │ ├── GeminiUrlContextService.test.vitest.ts
│ │ │ │ ├── GeminiValidationSchemas.test.vitest.ts
│ │ │ │ ├── GitHubApiService.test.vitest.ts
│ │ │ │ ├── GitHubUrlParser.test.vitest.ts
│ │ │ │ └── ThinkingBudget.test.vitest.ts
│ │ │ ├── mcp
│ │ │ │ └── McpClientService.test.vitest.ts
│ │ │ ├── ModelSelectionService.test.vitest.ts
│ │ │ └── session
│ │ │ └── SQLiteSessionStore.test.vitest.ts
│ │ ├── tools
│ │ │ ├── geminiCacheTool.test.vitest.ts
│ │ │ ├── geminiChatTool.test.vitest.ts
│ │ │ ├── geminiCodeReviewTool.test.vitest.ts
│ │ │ ├── geminiGenerateContentConsolidatedTool.test.vitest.ts
│ │ │ ├── geminiGenerateImageTool.test.vitest.ts
│ │ │ ├── geminiRouteMessageTool.test.vitest.ts
│ │ │ ├── mcpClientTool.test.vitest.ts
│ │ │ ├── mcpToolsTests.test.vitest.ts
│ │ │ └── schemas
│ │ │ ├── BaseToolSchema.test.vitest.ts
│ │ │ ├── ToolParamSchemas.test.vitest.ts
│ │ │ └── ToolSchemas.test.vitest.ts
│ │ └── utils
│ │ ├── errors.test.vitest.ts
│ │ ├── FileSecurityService.test.vitest.ts
│ │ ├── FileSecurityService.vitest.ts
│ │ ├── FileSecurityServiceBasics.test.vitest.ts
│ │ ├── healthCheck.test.vitest.ts
│ │ ├── RetryService.test.vitest.ts
│ │ └── UrlSecurityService.test.vitest.ts
│ └── utils
│ ├── assertions.ts
│ ├── debug-error.ts
│ ├── env-check.ts
│ ├── environment.ts
│ ├── error-helpers.ts
│ ├── express-mocks.ts
│ ├── integration-types.ts
│ ├── mock-types.ts
│ ├── test-fixtures.ts
│ ├── test-generators.ts
│ ├── test-setup.ts
│ └── vitest.d.ts
├── tsconfig.json
├── tsconfig.test.json
├── vitest-globals.d.ts
├── vitest.config.ts
└── vitest.setup.ts
```
# Files
--------------------------------------------------------------------------------
/src/config/ConfigurationManager.ts:
--------------------------------------------------------------------------------
```typescript
import * as path from "path";
import {
ExampleServiceConfig,
GeminiServiceConfig,
ModelConfiguration,
ModelCapabilitiesMap,
} from "../types/index.js";
import { FileSecurityService } from "../utils/FileSecurityService.js";
import { ModelMigrationService } from "../services/gemini/ModelMigrationService.js";
import { logger } from "../utils/logger.js";
// Define the structure for all configurations managed
interface ManagedConfigs {
exampleService: Required<ExampleServiceConfig>;
geminiService: GeminiServiceConfig;
github: {
apiToken: string;
};
allowedOutputPaths: string[];
mcpConfig: {
host: string;
port: number;
connectionToken: string;
clientId: string;
logLevel?: "debug" | "info" | "warn" | "error";
transport?: "stdio" | "sse";
enableStreaming?: boolean;
sessionTimeoutSeconds?: number;
};
urlContext: {
enabled: boolean;
maxUrlsPerRequest: number;
defaultMaxContentKb: number;
defaultTimeoutMs: number;
allowedDomains: string[];
blocklistedDomains: string[];
convertToMarkdown: boolean;
includeMetadata: boolean;
enableCaching: boolean;
cacheExpiryMinutes: number;
maxCacheSize: number;
rateLimitPerDomainPerMinute: number;
userAgent: string;
};
modelConfiguration: ModelConfiguration;
}
/**
* Centralized configuration management for all services.
* Implements singleton pattern to ensure consistent configuration.
*/
export class ConfigurationManager {
private static instance: ConfigurationManager | null = null;
private static instanceLock = false;
private config: ManagedConfigs;
private constructor() {
// Initialize with default configurations
this.config = {
exampleService: {
// Define defaults for ExampleService
greeting: "Hello",
enableDetailedLogs: false,
},
geminiService: {
apiKey: "",
defaultModel: undefined,
defaultImageResolution: "1024x1024",
maxImageSizeMB: 10,
supportedImageFormats: ["image/jpeg", "image/png", "image/webp"],
defaultThinkingBudget: undefined,
},
modelConfiguration: this.buildDefaultModelConfiguration(),
github: {
// Default GitHub API token is empty; will be loaded from environment variable
apiToken: "",
},
allowedOutputPaths: [],
mcpConfig: {
// Initialize MCP config
host: "localhost",
port: 8080,
connectionToken: "", // Must be set via env
clientId: "gemini-sdk-client",
logLevel: "info",
transport: "stdio",
},
urlContext: {
// Initialize URL context config with secure defaults
enabled: false, // Disabled by default for security
maxUrlsPerRequest: 20,
defaultMaxContentKb: 100,
defaultTimeoutMs: 10000,
allowedDomains: ["*"], // Allow all by default (can be restricted)
blocklistedDomains: [], // Empty by default
convertToMarkdown: true,
includeMetadata: true,
enableCaching: true,
cacheExpiryMinutes: 15,
maxCacheSize: 1000,
rateLimitPerDomainPerMinute: 10,
userAgent:
"MCP-Gemini-Server/1.0 (+https://github.com/bsmi021/mcp-gemini-server)",
},
// Initialize other service configs with defaults:
// yourService: {
// someSetting: 'default value',
// retryCount: 3,
// },
};
const migrationService = ModelMigrationService.getInstance();
migrationService.migrateEnvironmentVariables();
const validation = migrationService.validateConfiguration();
if (!validation.isValid) {
logger.error("[ConfigurationManager] Configuration validation failed", {
errors: validation.errors,
});
}
const deprecated = migrationService.getDeprecatedFeatures();
if (deprecated.length > 0) {
logger.warn("[ConfigurationManager] Deprecated features detected", {
deprecated,
});
}
this.validateRequiredEnvVars();
this.loadEnvironmentOverrides();
this.config.modelConfiguration = this.parseModelConfiguration();
FileSecurityService.configureFromEnvironment();
}
private validateRequiredEnvVars(): void {
// Skip validation in test environment
if (process.env.NODE_ENV === "test") {
logger.info(
"Skipping environment variable validation in test environment"
);
return;
}
// Always require Gemini API key
const requiredVars = ["GOOGLE_GEMINI_API_KEY"];
// Check transport type to determine if MCP server variables are required
const transportType =
process.env.MCP_TRANSPORT || process.env.MCP_TRANSPORT_TYPE || "stdio";
// Only require MCP server variables for HTTP/SSE transport modes
// Note: MCP_CLIENT_ID is not required as it's optional with a default value
if (
transportType === "http" ||
transportType === "sse" ||
transportType === "streamable"
) {
requiredVars.push(
"MCP_SERVER_HOST",
"MCP_SERVER_PORT",
"MCP_CONNECTION_TOKEN"
);
}
const missingVars = requiredVars.filter((varName) => !process.env[varName]);
if (missingVars.length > 0) {
throw new Error(
`Missing required environment variables: ${missingVars.join(", ")}`
);
}
}
/**
* Get the singleton instance of ConfigurationManager.
* Basic lock to prevent race conditions during initial creation.
*/
public static getInstance(): ConfigurationManager {
if (!ConfigurationManager.instance) {
if (!ConfigurationManager.instanceLock) {
ConfigurationManager.instanceLock = true; // Lock
try {
ConfigurationManager.instance = new ConfigurationManager();
} finally {
ConfigurationManager.instanceLock = false; // Unlock
}
} else {
// Basic busy wait if locked (consider a more robust async lock if high contention is expected)
while (ConfigurationManager.instanceLock) {
// Small delay to prevent tight loop
const now = Date.now();
while (Date.now() - now < 10) {
// Intentional minimal delay
}
}
// Re-check instance after wait
if (!ConfigurationManager.instance) {
// This path is less likely but handles edge cases if lock logic needs refinement
return ConfigurationManager.getInstance();
}
}
}
return ConfigurationManager.instance;
}
// --- Getters for specific configurations ---
public getExampleServiceConfig(): Required<ExampleServiceConfig> {
// Return a copy to prevent accidental modification of the internal state
return { ...this.config.exampleService };
}
public getGeminiServiceConfig(): GeminiServiceConfig {
// Return a copy to prevent accidental modification
return { ...this.config.geminiService };
}
// Getter for MCP Configuration
public getMcpConfig(): Required<ManagedConfigs["mcpConfig"]> {
// Return a copy to ensure type safety and prevent modification
// Cast to Required because we validate essential fields are set from env vars.
// Optional fields will have their defaults.
return { ...this.config.mcpConfig } as Required<
ManagedConfigs["mcpConfig"]
>;
}
// Getter specifically for the default model name
public getDefaultModelName(): string | undefined {
return this.config.geminiService.defaultModel;
}
public getModelConfiguration(): ModelConfiguration {
return { ...this.config.modelConfiguration };
}
/**
* Returns the GitHub API token for GitHub API requests
* @returns The configured GitHub API token or undefined if not set
*/
public getGitHubApiToken(): string | undefined {
return this.config.github.apiToken || undefined;
}
/**
* Returns the list of allowed output paths for file writing
* @returns A copy of the configured allowed output paths array
*/
public getAllowedOutputPaths(): string[] {
// Return a copy to prevent accidental modification
return [...this.config.allowedOutputPaths];
}
/**
* Returns the URL context configuration
* @returns A copy of the URL context configuration
*/
public getUrlContextConfig(): Required<ManagedConfigs["urlContext"]> {
return { ...this.config.urlContext };
}
// Add getters for other service configs:
// public getYourServiceConfig(): Required<YourServiceConfig> {
// return { ...this.config.yourService };
// }
// --- Updaters for specific configurations (if runtime updates are needed) ---
public updateExampleServiceConfig(
update: Partial<ExampleServiceConfig>
): void {
this.config.exampleService = {
...this.config.exampleService,
...update,
};
// Optional: Notify relevant services about the config change
}
// Add updaters for other service configs:
// public updateYourServiceConfig(update: Partial<YourServiceConfig>): void {
// this.config.yourService = {
// ...this.config.yourService,
// ...update,
// };
// }
/**
* Example method to load configuration overrides from environment variables.
* Call this in the constructor.
*/
private loadEnvironmentOverrides(): void {
// Example for ExampleService
if (process.env.EXAMPLE_GREETING) {
this.config.exampleService.greeting = process.env.EXAMPLE_GREETING;
}
if (process.env.EXAMPLE_ENABLE_LOGS) {
this.config.exampleService.enableDetailedLogs =
process.env.EXAMPLE_ENABLE_LOGS.toLowerCase() === "true";
}
// Load GitHub API token if provided
if (process.env.GITHUB_API_TOKEN) {
this.config.github.apiToken = process.env.GITHUB_API_TOKEN;
logger.info("[ConfigurationManager] GitHub API token configured");
} else {
logger.warn(
"[ConfigurationManager] GITHUB_API_TOKEN environment variable not set. GitHub code review features may not work properly."
);
}
// Add logic for other services based on their environment variables
// if (process.env.YOUR_SERVICE_RETRY_COUNT) {
// const retryCount = parseInt(process.env.YOUR_SERVICE_RETRY_COUNT, 10);
// if (!isNaN(retryCount)) {
// this.config.yourService.retryCount = retryCount;
// }
// }
// Load Gemini API Key (using the name from .env)
if (process.env.GOOGLE_GEMINI_API_KEY) {
this.config.geminiService.apiKey = process.env.GOOGLE_GEMINI_API_KEY;
} else {
// Log a warning if the key is missing, the service constructor will throw
logger.warn(
"[ConfigurationManager] WARNING: GOOGLE_GEMINI_API_KEY environment variable not set."
);
}
// Load Default Gemini Model Name
if (process.env.GOOGLE_GEMINI_MODEL) {
this.config.geminiService.defaultModel = process.env.GOOGLE_GEMINI_MODEL;
logger.info(
`[ConfigurationManager] Default Gemini model set to: ${this.config.geminiService.defaultModel}`
);
} else {
logger.info(
"[ConfigurationManager] GOOGLE_GEMINI_MODEL environment variable not set. No default model configured."
);
}
// Load image-specific settings if provided
if (process.env.GOOGLE_GEMINI_IMAGE_RESOLUTION) {
const resolution = process.env.GOOGLE_GEMINI_IMAGE_RESOLUTION;
if (["512x512", "1024x1024", "1536x1536"].includes(resolution)) {
this.config.geminiService.defaultImageResolution = resolution as
| "512x512"
| "1024x1024"
| "1536x1536";
logger.info(
`[ConfigurationManager] Default image resolution set to: ${resolution}`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid image resolution '${resolution}' specified in GOOGLE_GEMINI_IMAGE_RESOLUTION. Using default.`
);
}
}
if (process.env.GOOGLE_GEMINI_MAX_IMAGE_SIZE_MB) {
const sizeMB = parseInt(process.env.GOOGLE_GEMINI_MAX_IMAGE_SIZE_MB, 10);
if (!isNaN(sizeMB) && sizeMB > 0) {
this.config.geminiService.maxImageSizeMB = sizeMB;
logger.info(
`[ConfigurationManager] Maximum image size set to: ${sizeMB}MB`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid max image size '${process.env.GOOGLE_GEMINI_MAX_IMAGE_SIZE_MB}' specified. Using default.`
);
}
}
if (process.env.GOOGLE_GEMINI_SUPPORTED_IMAGE_FORMATS) {
try {
const formats = JSON.parse(
process.env.GOOGLE_GEMINI_SUPPORTED_IMAGE_FORMATS
);
if (
Array.isArray(formats) &&
formats.every((f) => typeof f === "string")
) {
this.config.geminiService.supportedImageFormats = formats;
logger.info(
`[ConfigurationManager] Supported image formats set to: ${formats.join(", ")}`
);
} else {
throw new Error("Invalid format array");
}
} catch (error) {
logger.warn(
`[ConfigurationManager] Invalid image formats specified in GOOGLE_GEMINI_SUPPORTED_IMAGE_FORMATS: '${process.env.GOOGLE_GEMINI_SUPPORTED_IMAGE_FORMATS}'. Using default.`
);
}
}
// Load default thinking budget if provided
if (process.env.GOOGLE_GEMINI_DEFAULT_THINKING_BUDGET) {
const budget = parseInt(
process.env.GOOGLE_GEMINI_DEFAULT_THINKING_BUDGET,
10
);
if (!isNaN(budget) && budget >= 0 && budget <= 24576) {
this.config.geminiService.defaultThinkingBudget = budget;
logger.info(
`[ConfigurationManager] Default thinking budget set to: ${budget} tokens`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid thinking budget '${process.env.GOOGLE_GEMINI_DEFAULT_THINKING_BUDGET}' specified. Must be between 0 and 24576. Not using default thinking budget.`
);
}
}
// Load MCP Configuration
if (process.env.MCP_SERVER_HOST) {
this.config.mcpConfig.host = process.env.MCP_SERVER_HOST;
}
if (process.env.MCP_SERVER_PORT) {
const port = parseInt(process.env.MCP_SERVER_PORT, 10);
if (!isNaN(port) && port > 0 && port < 65536) {
this.config.mcpConfig.port = port;
} else {
logger.warn(
`[ConfigurationManager] Invalid MCP_SERVER_PORT: '${process.env.MCP_SERVER_PORT}'. Using default ${this.config.mcpConfig.port}.`
);
}
}
if (process.env.MCP_CONNECTION_TOKEN) {
this.config.mcpConfig.connectionToken = process.env.MCP_CONNECTION_TOKEN;
}
if (process.env.MCP_CLIENT_ID) {
this.config.mcpConfig.clientId = process.env.MCP_CLIENT_ID;
}
if (process.env.MCP_LOG_LEVEL) {
const logLevel = process.env.MCP_LOG_LEVEL.toLowerCase();
if (["debug", "info", "warn", "error"].includes(logLevel)) {
this.config.mcpConfig.logLevel = logLevel as
| "debug"
| "info"
| "warn"
| "error";
} else {
logger.warn(
`[ConfigurationManager] Invalid MCP_LOG_LEVEL: '${process.env.MCP_LOG_LEVEL}'. Using default '${this.config.mcpConfig.logLevel}'.`
);
}
}
if (process.env.MCP_TRANSPORT) {
const transport = process.env.MCP_TRANSPORT.toLowerCase();
if (["stdio", "sse"].includes(transport)) {
this.config.mcpConfig.transport = transport as "stdio" | "sse";
} else {
logger.warn(
`[ConfigurationManager] Invalid MCP_TRANSPORT: '${process.env.MCP_TRANSPORT}'. Using default '${this.config.mcpConfig.transport}'.`
);
}
}
if (process.env.MCP_ENABLE_STREAMING) {
this.config.mcpConfig.enableStreaming =
process.env.MCP_ENABLE_STREAMING.toLowerCase() === "true";
logger.info(
`[ConfigurationManager] MCP streaming enabled: ${this.config.mcpConfig.enableStreaming}`
);
}
if (process.env.MCP_SESSION_TIMEOUT) {
const timeout = parseInt(process.env.MCP_SESSION_TIMEOUT, 10);
if (!isNaN(timeout) && timeout > 0) {
this.config.mcpConfig.sessionTimeoutSeconds = timeout;
logger.info(
`[ConfigurationManager] MCP session timeout set to: ${timeout} seconds`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid MCP_SESSION_TIMEOUT: '${process.env.MCP_SESSION_TIMEOUT}'. Using default.`
);
}
}
logger.info("[ConfigurationManager] MCP configuration loaded.");
// Load URL Context Configuration
if (process.env.GOOGLE_GEMINI_ENABLE_URL_CONTEXT) {
this.config.urlContext.enabled =
process.env.GOOGLE_GEMINI_ENABLE_URL_CONTEXT.toLowerCase() === "true";
logger.info(
`[ConfigurationManager] URL context feature enabled: ${this.config.urlContext.enabled}`
);
}
if (process.env.GOOGLE_GEMINI_URL_MAX_COUNT) {
const maxCount = parseInt(process.env.GOOGLE_GEMINI_URL_MAX_COUNT, 10);
if (!isNaN(maxCount) && maxCount > 0 && maxCount <= 20) {
this.config.urlContext.maxUrlsPerRequest = maxCount;
logger.info(`[ConfigurationManager] URL max count set to: ${maxCount}`);
} else {
logger.warn(
`[ConfigurationManager] Invalid URL max count '${process.env.GOOGLE_GEMINI_URL_MAX_COUNT}'. Must be between 1 and 20.`
);
}
}
if (process.env.GOOGLE_GEMINI_URL_MAX_CONTENT_KB) {
const maxKb = parseInt(process.env.GOOGLE_GEMINI_URL_MAX_CONTENT_KB, 10);
if (!isNaN(maxKb) && maxKb > 0 && maxKb <= 1000) {
this.config.urlContext.defaultMaxContentKb = maxKb;
logger.info(
`[ConfigurationManager] URL max content size set to: ${maxKb}KB`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid URL max content size '${process.env.GOOGLE_GEMINI_URL_MAX_CONTENT_KB}'. Must be between 1 and 1000 KB.`
);
}
}
if (process.env.GOOGLE_GEMINI_URL_FETCH_TIMEOUT_MS) {
const timeout = parseInt(
process.env.GOOGLE_GEMINI_URL_FETCH_TIMEOUT_MS,
10
);
if (!isNaN(timeout) && timeout >= 1000 && timeout <= 30000) {
this.config.urlContext.defaultTimeoutMs = timeout;
logger.info(
`[ConfigurationManager] URL fetch timeout set to: ${timeout}ms`
);
} else {
logger.warn(
`[ConfigurationManager] Invalid URL fetch timeout '${process.env.GOOGLE_GEMINI_URL_FETCH_TIMEOUT_MS}'. Must be between 1000 and 30000 ms.`
);
}
}
if (process.env.GOOGLE_GEMINI_URL_ALLOWED_DOMAINS) {
try {
const domains = this.parseStringArray(
process.env.GOOGLE_GEMINI_URL_ALLOWED_DOMAINS
);
this.config.urlContext.allowedDomains = domains;
logger.info(
`[ConfigurationManager] URL allowed domains set to: ${domains.join(", ")}`
);
} catch (error) {
logger.warn(
`[ConfigurationManager] Invalid URL allowed domains format: ${error}`
);
}
}
if (process.env.GOOGLE_GEMINI_URL_BLOCKLIST) {
try {
const domains = this.parseStringArray(
process.env.GOOGLE_GEMINI_URL_BLOCKLIST
);
this.config.urlContext.blocklistedDomains = domains;
logger.info(
`[ConfigurationManager] URL blocklisted domains set to: ${domains.join(", ")}`
);
} catch (error) {
logger.warn(
`[ConfigurationManager] Invalid URL blocklist format: ${error}`
);
}
}
if (process.env.GOOGLE_GEMINI_URL_CONVERT_TO_MARKDOWN) {
this.config.urlContext.convertToMarkdown =
process.env.GOOGLE_GEMINI_URL_CONVERT_TO_MARKDOWN.toLowerCase() ===
"true";
logger.info(
`[ConfigurationManager] URL markdown conversion enabled: ${this.config.urlContext.convertToMarkdown}`
);
}
if (process.env.GOOGLE_GEMINI_URL_INCLUDE_METADATA) {
this.config.urlContext.includeMetadata =
process.env.GOOGLE_GEMINI_URL_INCLUDE_METADATA.toLowerCase() === "true";
logger.info(
`[ConfigurationManager] URL metadata inclusion enabled: ${this.config.urlContext.includeMetadata}`
);
}
if (process.env.GOOGLE_GEMINI_URL_ENABLE_CACHING) {
this.config.urlContext.enableCaching =
process.env.GOOGLE_GEMINI_URL_ENABLE_CACHING.toLowerCase() === "true";
logger.info(
`[ConfigurationManager] URL caching enabled: ${this.config.urlContext.enableCaching}`
);
}
if (process.env.GOOGLE_GEMINI_URL_USER_AGENT) {
this.config.urlContext.userAgent =
process.env.GOOGLE_GEMINI_URL_USER_AGENT;
logger.info(
`[ConfigurationManager] URL user agent set to: ${this.config.urlContext.userAgent}`
);
}
logger.info("[ConfigurationManager] URL context configuration loaded.");
this.config.allowedOutputPaths = [];
const allowedOutputPathsEnv = process.env.ALLOWED_OUTPUT_PATHS;
if (allowedOutputPathsEnv && allowedOutputPathsEnv.trim().length > 0) {
const pathsArray = allowedOutputPathsEnv
.split(",")
.map((p) => p.trim()) // Trim whitespace from each path
.filter((p) => p.length > 0); // Filter out any empty strings resulting from split
if (pathsArray.length > 0) {
this.config.allowedOutputPaths = pathsArray.map((p) => path.resolve(p)); // Resolve to absolute paths
logger.info(
`[ConfigurationManager] Allowed output paths configured: ${this.config.allowedOutputPaths.join(
", "
)}`
);
} else {
// This case handles if ALLOWED_OUTPUT_PATHS was something like ",," or " , "
logger.warn(
"[ConfigurationManager] ALLOWED_OUTPUT_PATHS environment variable was provided but contained no valid paths after trimming. File writing might be restricted."
);
}
} else {
logger.warn(
"[ConfigurationManager] ALLOWED_OUTPUT_PATHS environment variable not set or is empty. File writing might be restricted or disabled."
);
}
}
private buildDefaultModelConfiguration(): ModelConfiguration {
return {
default: "gemini-2.5-flash-preview-05-20",
textGeneration: [
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-flash-preview-05-20",
"gemini-2.0-flash",
"gemini-1.5-pro",
"gemini-1.5-flash",
],
imageGeneration: [
"imagen-3.0-generate-002",
"gemini-2.0-flash-preview-image-generation",
],
videoGeneration: ["veo-2.0-generate-001"],
codeReview: [
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-flash-preview-05-20",
"gemini-2.0-flash",
],
complexReasoning: [
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-flash-preview-05-20",
],
capabilities: this.buildCapabilitiesMap(),
routing: {
preferCostEffective: false,
preferSpeed: false,
preferQuality: true,
},
};
}
private buildCapabilitiesMap(): ModelCapabilitiesMap {
return {
"gemini-2.5-pro-preview-05-06": {
textGeneration: true,
imageInput: true,
videoInput: true,
audioInput: true,
imageGeneration: false,
videoGeneration: false,
codeExecution: "excellent",
complexReasoning: "excellent",
costTier: "high",
speedTier: "medium",
maxTokens: 65536,
contextWindow: 1048576,
supportsFunctionCalling: true,
supportsSystemInstructions: true,
supportsCaching: true,
},
"gemini-2.5-flash-preview-05-20": {
textGeneration: true,
imageInput: true,
videoInput: true,
audioInput: true,
imageGeneration: false,
videoGeneration: false,
codeExecution: "excellent",
complexReasoning: "excellent",
costTier: "medium",
speedTier: "fast",
maxTokens: 65536,
contextWindow: 1048576,
supportsFunctionCalling: true,
supportsSystemInstructions: true,
supportsCaching: true,
},
"gemini-2.0-flash": {
textGeneration: true,
imageInput: true,
videoInput: true,
audioInput: true,
imageGeneration: false,
videoGeneration: false,
codeExecution: "good",
complexReasoning: "good",
costTier: "medium",
speedTier: "fast",
maxTokens: 8192,
contextWindow: 1048576,
supportsFunctionCalling: true,
supportsSystemInstructions: true,
supportsCaching: true,
},
"gemini-2.0-flash-preview-image-generation": {
textGeneration: true,
imageInput: true,
videoInput: false,
audioInput: false,
imageGeneration: true,
videoGeneration: false,
codeExecution: "basic",
complexReasoning: "basic",
costTier: "medium",
speedTier: "medium",
maxTokens: 8192,
contextWindow: 32000,
supportsFunctionCalling: false,
supportsSystemInstructions: true,
supportsCaching: false,
},
"gemini-1.5-pro": {
textGeneration: true,
imageInput: true,
videoInput: true,
audioInput: true,
imageGeneration: false,
videoGeneration: false,
codeExecution: "good",
complexReasoning: "good",
costTier: "high",
speedTier: "medium",
maxTokens: 8192,
contextWindow: 2000000,
supportsFunctionCalling: true,
supportsSystemInstructions: true,
supportsCaching: true,
},
"gemini-1.5-flash": {
textGeneration: true,
imageInput: true,
videoInput: true,
audioInput: true,
imageGeneration: false,
videoGeneration: false,
codeExecution: "basic",
complexReasoning: "basic",
costTier: "low",
speedTier: "fast",
maxTokens: 8192,
contextWindow: 1000000,
supportsFunctionCalling: true,
supportsSystemInstructions: true,
supportsCaching: true,
},
"imagen-3.0-generate-002": {
textGeneration: false,
imageInput: false,
videoInput: false,
audioInput: false,
imageGeneration: true,
videoGeneration: false,
codeExecution: "none",
complexReasoning: "none",
costTier: "medium",
speedTier: "medium",
maxTokens: 0,
contextWindow: 0,
supportsFunctionCalling: false,
supportsSystemInstructions: false,
supportsCaching: false,
},
"veo-2.0-generate-001": {
textGeneration: false,
imageInput: true,
videoInput: false,
audioInput: false,
imageGeneration: false,
videoGeneration: true,
codeExecution: "none",
complexReasoning: "none",
costTier: "high",
speedTier: "slow",
maxTokens: 0,
contextWindow: 0,
supportsFunctionCalling: false,
supportsSystemInstructions: true,
supportsCaching: false,
},
};
}
private parseModelConfiguration(): ModelConfiguration {
const textModels = this.parseModelArray("GOOGLE_GEMINI_MODELS") ||
this.parseModelArray("GOOGLE_GEMINI_TEXT_MODELS") || [
process.env.GOOGLE_GEMINI_MODEL || "gemini-2.5-flash-preview-05-20",
];
const imageModels = this.parseModelArray("GOOGLE_GEMINI_IMAGE_MODELS") || [
"imagen-3.0-generate-002",
"gemini-2.0-flash-preview-image-generation",
];
const videoModels = this.parseModelArray("GOOGLE_GEMINI_VIDEO_MODELS") || [
"veo-2.0-generate-001",
];
const codeModels = this.parseModelArray("GOOGLE_GEMINI_CODE_MODELS") || [
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-flash-preview-05-20",
"gemini-2.0-flash",
];
return {
default: process.env.GOOGLE_GEMINI_DEFAULT_MODEL || textModels[0],
textGeneration: textModels,
imageGeneration: imageModels,
videoGeneration: videoModels,
codeReview: codeModels,
complexReasoning: textModels.filter((m) => this.isHighReasoningModel(m)),
capabilities: this.buildCapabilitiesMap(),
routing: this.parseRoutingPreferences(),
};
}
private parseModelArray(envVarName: string): string[] | null {
const envValue = process.env[envVarName];
if (!envValue) return null;
try {
const parsed = JSON.parse(envValue);
if (
Array.isArray(parsed) &&
parsed.every((item) => typeof item === "string")
) {
return parsed;
}
logger.warn(
`[ConfigurationManager] Invalid ${envVarName} format: expected JSON array of strings`
);
return null;
} catch (error) {
logger.warn(
`[ConfigurationManager] Failed to parse ${envVarName}: ${error}`
);
return null;
}
}
private isHighReasoningModel(modelName: string): boolean {
const highReasoningModels = [
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-flash-preview-05-20",
"gemini-1.5-pro",
];
return highReasoningModels.includes(modelName);
}
private parseRoutingPreferences(): ModelConfiguration["routing"] {
return {
preferCostEffective:
process.env.GOOGLE_GEMINI_ROUTING_PREFER_COST?.toLowerCase() === "true",
preferSpeed:
process.env.GOOGLE_GEMINI_ROUTING_PREFER_SPEED?.toLowerCase() ===
"true",
preferQuality:
process.env.GOOGLE_GEMINI_ROUTING_PREFER_QUALITY?.toLowerCase() ===
"true" ||
(!process.env.GOOGLE_GEMINI_ROUTING_PREFER_COST &&
!process.env.GOOGLE_GEMINI_ROUTING_PREFER_SPEED),
};
}
/**
* Parse a comma-separated string or JSON array into a string array
*/
private parseStringArray(value: string): string[] {
if (!value || value.trim() === "") {
return [];
}
// Try to parse as JSON first
if (value.trim().startsWith("[")) {
try {
const parsed = JSON.parse(value);
if (
Array.isArray(parsed) &&
parsed.every((item) => typeof item === "string")
) {
return parsed;
}
throw new Error("Not a string array");
} catch (error) {
throw new Error(`Invalid JSON array format: ${error}`);
}
}
// Parse as comma-separated string
return value
.split(",")
.map((item) => item.trim())
.filter((item) => item.length > 0);
}
}
```
--------------------------------------------------------------------------------
/tests/unit/tools/geminiGenerateContentConsolidatedTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
// Using vitest globals - see vitest.config.ts globals: true
import { geminiGenerateContentConsolidatedTool } from "../../../src/tools/geminiGenerateContentConsolidatedTool.js";
import { GeminiApiError } from "../../../src/utils/errors.js";
import { McpError } from "@modelcontextprotocol/sdk/types.js";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { GeminiService } from "../../../src/services/index.js";
describe("geminiGenerateContentConsolidatedTool", () => {
// Mock server and service instances
const mockTool = vi.fn();
const mockServer = {
tool: mockTool,
} as unknown as McpServer;
// Create mock functions for the service methods
const mockGenerateContent = vi.fn();
const mockGenerateContentStream = vi.fn();
// Create a minimal mock service with just the necessary methods for testing
const mockService = {
generateContent: mockGenerateContent,
generateContentStream: mockGenerateContentStream,
// Add empty implementations for required GeminiService methods
} as unknown as GeminiService;
// Reset mocks before each test
beforeEach(() => {
vi.resetAllMocks();
});
it("should register the tool with the server", () => {
// Call the tool registration function
geminiGenerateContentConsolidatedTool(mockServer, mockService);
// Verify tool was registered
expect(mockTool).toHaveBeenCalledTimes(1);
const [name, description, params, handler] = mockTool.mock.calls[0];
// Check tool registration parameters
expect(name).toBe("gemini_generate_content");
expect(description).toContain("Generates text content");
expect(params).toBeDefined();
expect(typeof handler).toBe("function");
});
it("should handle standard content generation", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "This is a test response";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "What is the capital of France?",
stream: false,
};
// Call the handler
const result = await handler(testRequest);
// Verify the service method was called with correct parameters
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-1.5-flash",
prompt: "What is the capital of France?",
})
);
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: "This is a test response",
},
],
});
});
it("should handle streaming content generation", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Create an async generator mock for streaming
async function* mockStreamGenerator() {
yield "This is ";
yield "a streaming ";
yield "response";
}
mockGenerateContentStream.mockReturnValueOnce(mockStreamGenerator());
// Prepare test request with streaming enabled
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Tell me a story",
stream: true,
};
// Call the handler
const result = await handler(testRequest);
// Verify the streaming service method was called
expect(mockGenerateContentStream).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-1.5-flash",
prompt: "Tell me a story",
})
);
// Verify the result contains the concatenated stream
expect(result).toEqual({
content: [
{
type: "text",
text: "This is a streaming response",
},
],
});
});
it("should handle function calling with function declarations", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock function call response
const mockFunctionCallResponse = {
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
};
mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
// Prepare test request with function declarations
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "What's the weather in Paris?",
stream: false,
functionDeclarations: [
{
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "OBJECT" as const,
properties: {
location: {
type: "STRING" as const,
description: "The location to get weather for",
},
},
required: ["location"],
},
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the service method was called with function declarations
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-1.5-flash",
prompt: "What's the weather in Paris?",
functionDeclarations: expect.arrayContaining([
expect.objectContaining({
name: "get_weather",
}),
]),
})
);
// Verify the result contains the serialized function call
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({
name: "get_weather",
args: { location: "Paris" },
}),
},
],
});
});
it("should handle optional parameters", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Generated with parameters";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with all optional parameters
const testRequest = {
modelName: "gemini-1.5-pro",
prompt: "Generate creative content",
stream: false,
generationConfig: {
temperature: 0.8,
topK: 40,
topP: 0.95,
maxOutputTokens: 1024,
stopSequences: ["END"],
thinkingConfig: {
thinkingBudget: 8192,
reasoningEffort: "medium",
},
},
safetySettings: [
{
category: "HARM_CATEGORY_HATE_SPEECH",
threshold: "BLOCK_MEDIUM_AND_ABOVE",
},
],
systemInstruction: "You are a helpful assistant",
cachedContentName: "cachedContents/12345",
urlContext: {
urls: ["https://example.com"],
fetchOptions: {
maxContentKb: 100,
timeoutMs: 10000,
},
},
modelPreferences: {
preferQuality: true,
preferSpeed: false,
taskType: "creative_writing",
},
};
// Call the handler
const result = await handler(testRequest);
// Verify all parameters were passed to the service
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-1.5-pro",
prompt: "Generate creative content",
generationConfig: expect.objectContaining({
temperature: 0.8,
topK: 40,
topP: 0.95,
maxOutputTokens: 1024,
}),
systemInstruction: "You are a helpful assistant",
cachedContentName: "cachedContents/12345",
urlContext: expect.objectContaining({
urls: ["https://example.com"],
}),
})
);
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: "Generated with parameters",
},
],
});
});
it("should handle errors and map them to MCP errors", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock an API error
const apiError = new GeminiApiError("API rate limit exceeded", {
statusCode: 429,
statusText: "Too Many Requests",
});
mockGenerateContent.mockRejectedValueOnce(apiError);
// Prepare test request
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Test error handling",
stream: false,
};
// Call the handler and expect it to throw
await expect(handler(testRequest)).rejects.toThrow(McpError);
});
it("should handle URL context metrics calculation", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
mockGenerateContent.mockResolvedValueOnce("Response with URL context");
// Prepare test request with URL context
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Analyze these URLs",
stream: false,
urlContext: {
urls: [
"https://example1.com",
"https://example2.com",
"https://example3.com",
],
fetchOptions: {
maxContentKb: 200,
},
},
};
// Call the handler
await handler(testRequest);
// Verify URL metrics were calculated and passed
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
urlCount: 3,
estimatedUrlContentSize: 3 * 200 * 1024, // 3 URLs * 200KB * 1024 bytes/KB
})
);
});
it("should handle function call response with text fallback", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock function call response with text fallback
const mockFunctionCallResponse = {
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
text: "I'll get the weather for Paris.",
};
mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
// Prepare test request with function declarations
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "What's the weather in Paris?",
stream: false,
functionDeclarations: [
{
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "OBJECT" as const,
properties: {
location: {
type: "STRING" as const,
description: "The location to get weather for",
},
},
required: ["location"],
},
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the result contains the serialized function call (not the text)
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({
name: "get_weather",
args: { location: "Paris" },
}),
},
],
});
});
it("should handle function call response with only text", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with only text (no function call)
const mockTextResponse = {
text: "The weather in Paris is sunny and 22°C.",
};
mockGenerateContent.mockResolvedValueOnce(mockTextResponse);
// Prepare test request with function declarations
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "What's the weather in Paris?",
stream: false,
functionDeclarations: [
{
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "OBJECT" as const,
properties: {
location: {
type: "STRING" as const,
description: "The location to get weather for",
},
},
required: ["location"],
},
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the result contains the text response
expect(result).toEqual({
content: [
{
type: "text",
text: "The weather in Paris is sunny and 22°C.",
},
],
});
});
it("should handle toolConfig parameter", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Response with tool config";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with toolConfig
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Test with tool config",
stream: false,
functionDeclarations: [
{
name: "test_function",
description: "A test function",
parameters: {
type: "OBJECT" as const,
properties: {
param: {
type: "STRING" as const,
description: "A test parameter",
},
},
required: ["param"],
},
},
],
toolConfig: {
functionCallingConfig: {
mode: "AUTO",
allowedFunctionNames: ["test_function"],
},
},
};
// Call the handler
await handler(testRequest);
// Verify toolConfig was passed to the service
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
toolConfig: {
functionCallingConfig: {
mode: "AUTO",
allowedFunctionNames: ["test_function"],
},
},
})
);
});
it("should handle thinking configuration parameters", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Response with thinking config";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with thinking configuration
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Complex reasoning task",
stream: false,
generationConfig: {
temperature: 0.7,
thinkingConfig: {
thinkingBudget: 16384,
reasoningEffort: "high",
},
},
};
// Call the handler
await handler(testRequest);
// Verify thinking config was passed to the service
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
generationConfig: expect.objectContaining({
thinkingConfig: {
thinkingBudget: 16384,
reasoningEffort: "high",
},
}),
})
);
});
it("should handle model preferences for task optimization", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Optimized response";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with comprehensive model preferences
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Generate creative content",
stream: false,
modelPreferences: {
preferQuality: true,
preferSpeed: false,
preferCost: false,
complexityHint: "high",
taskType: "creative_writing",
},
};
// Call the handler
await handler(testRequest);
// Verify model preferences were passed to the service
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
preferQuality: true,
preferSpeed: false,
preferCost: false,
complexityHint: "high",
taskType: "creative_writing",
})
);
});
it("should handle comprehensive safety settings", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Safe response";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with comprehensive safety settings
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Generate content with safety controls",
stream: false,
safetySettings: [
{
category: "HARM_CATEGORY_HATE_SPEECH",
threshold: "BLOCK_MEDIUM_AND_ABOVE",
},
{
category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
threshold: "BLOCK_LOW_AND_ABOVE",
},
{
category: "HARM_CATEGORY_HARASSMENT",
threshold: "BLOCK_ONLY_HIGH",
},
{
category: "HARM_CATEGORY_DANGEROUS_CONTENT",
threshold: "BLOCK_NONE",
},
],
};
// Call the handler
await handler(testRequest);
// Verify safety settings were properly mapped and passed
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
safetySettings: [
{
category: "HARM_CATEGORY_HATE_SPEECH",
threshold: "BLOCK_MEDIUM_AND_ABOVE",
},
{
category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
threshold: "BLOCK_LOW_AND_ABOVE",
},
{
category: "HARM_CATEGORY_HARASSMENT",
threshold: "BLOCK_ONLY_HIGH",
},
{
category: "HARM_CATEGORY_DANGEROUS_CONTENT",
threshold: "BLOCK_NONE",
},
],
})
);
});
it("should handle URL context with comprehensive fetch options", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
mockGenerateContent.mockResolvedValueOnce("Response with URL context");
// Prepare test request with comprehensive URL context options
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Analyze these web pages",
stream: false,
urlContext: {
urls: ["https://example1.com/article", "https://example2.com/blog"],
fetchOptions: {
maxContentKb: 150,
timeoutMs: 15000,
includeMetadata: true,
convertToMarkdown: true,
allowedDomains: ["example1.com", "example2.com"],
userAgent: "Custom-Agent/1.0",
},
},
};
// Call the handler
await handler(testRequest);
// Verify URL context was passed with all options
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
urlContext: {
urls: ["https://example1.com/article", "https://example2.com/blog"],
fetchOptions: {
maxContentKb: 150,
timeoutMs: 15000,
includeMetadata: true,
convertToMarkdown: true,
allowedDomains: ["example1.com", "example2.com"],
userAgent: "Custom-Agent/1.0",
},
},
urlCount: 2,
estimatedUrlContentSize: 2 * 150 * 1024, // 2 URLs * 150KB * 1024 bytes/KB
})
);
});
it("should handle URL context with default fetch options", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
mockGenerateContent.mockResolvedValueOnce(
"Response with default URL context"
);
// Prepare test request with minimal URL context (using defaults)
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Analyze this web page",
stream: false,
urlContext: {
urls: ["https://example.com"],
},
};
// Call the handler
await handler(testRequest);
// Verify URL context was passed with default maxContentKb
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
urlContext: {
urls: ["https://example.com"],
},
urlCount: 1,
estimatedUrlContentSize: 1 * 100 * 1024, // 1 URL * 100KB default * 1024 bytes/KB
})
);
});
it("should handle unexpected response structure from service", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock unexpected response structure
const unexpectedResponse = { unexpected: "structure" };
mockGenerateContent.mockResolvedValueOnce(unexpectedResponse);
// Prepare test request
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Test unexpected response",
stream: false,
};
// Call the handler and expect it to throw
await expect(handler(testRequest)).rejects.toThrow(
"Invalid response structure received from Gemini service."
);
});
it("should handle streaming with empty chunks gracefully", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Create an async generator mock with empty chunks
async function* mockStreamGenerator() {
yield "Start ";
yield ""; // Empty chunk
yield "middle ";
yield ""; // Another empty chunk
yield "end";
}
mockGenerateContentStream.mockReturnValueOnce(mockStreamGenerator());
// Prepare test request with streaming enabled
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Stream with empty chunks",
stream: true,
};
// Call the handler
const result = await handler(testRequest);
// Verify the result contains the concatenated stream (empty chunks should be included)
expect(result).toEqual({
content: [
{
type: "text",
text: "Start middle end",
},
],
});
});
it("should handle complex generation config with all parameters", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Complex config response";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with all generation config parameters
const testRequest = {
modelName: "gemini-1.5-pro",
prompt: "Complex generation task",
stream: false,
generationConfig: {
temperature: 0.9,
topP: 0.8,
topK: 50,
maxOutputTokens: 2048,
stopSequences: ["STOP", "END", "FINISH"],
thinkingConfig: {
thinkingBudget: 12288,
reasoningEffort: "medium",
},
},
};
// Call the handler
await handler(testRequest);
// Verify all generation config parameters were passed
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
generationConfig: {
temperature: 0.9,
topP: 0.8,
topK: 50,
maxOutputTokens: 2048,
stopSequences: ["STOP", "END", "FINISH"],
thinkingConfig: {
thinkingBudget: 12288,
reasoningEffort: "medium",
},
},
})
);
});
it("should handle cached content parameter", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Response using cached content";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with cached content
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Use cached context for this request",
stream: false,
cachedContentName: "cachedContents/abc123def456",
};
// Call the handler
await handler(testRequest);
// Verify cached content name was passed
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
cachedContentName: "cachedContents/abc123def456",
})
);
});
it("should handle system instruction parameter", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = "Response following system instruction";
mockGenerateContent.mockResolvedValueOnce(mockResponse);
// Prepare test request with system instruction
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Generate a response",
stream: false,
systemInstruction:
"You are a helpful assistant that always responds in a friendly tone.",
};
// Call the handler
await handler(testRequest);
// Verify system instruction was passed
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
systemInstruction:
"You are a helpful assistant that always responds in a friendly tone.",
})
);
});
it("should handle streaming errors gracefully", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Create an async generator that throws an error
async function* mockStreamGeneratorWithError() {
yield "Start ";
throw new Error("Streaming error occurred");
}
mockGenerateContentStream.mockReturnValueOnce(
mockStreamGeneratorWithError()
);
// Prepare test request with streaming enabled
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Stream that will error",
stream: true,
};
// Call the handler and expect it to throw
await expect(handler(testRequest)).rejects.toThrow();
});
it("should handle function declarations with complex parameter schemas", async () => {
// Register tool to get the request handler
geminiGenerateContentConsolidatedTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock function call response
const mockFunctionCallResponse = {
functionCall: {
name: "complex_function",
args: {
user: { name: "John", age: 30 },
preferences: ["option1", "option2"],
settings: { theme: "dark", notifications: true },
},
},
};
mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
// Prepare test request with complex function declarations
const testRequest = {
modelName: "gemini-1.5-flash",
prompt: "Call a complex function",
stream: false,
functionDeclarations: [
{
name: "complex_function",
description: "A function with complex nested parameters",
parameters: {
type: "OBJECT" as const,
properties: {
user: {
type: "OBJECT" as const,
properties: {
name: {
type: "STRING" as const,
description: "User's name",
},
age: {
type: "INTEGER" as const,
description: "User's age",
},
},
required: ["name", "age"],
},
preferences: {
type: "ARRAY" as const,
items: {
type: "STRING" as const,
description: "User preference",
},
description: "List of user preferences",
},
settings: {
type: "OBJECT" as const,
properties: {
theme: {
type: "STRING" as const,
enum: ["light", "dark"],
description: "UI theme preference",
},
notifications: {
type: "BOOLEAN" as const,
description: "Enable notifications",
},
},
},
},
required: ["user"],
},
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the service was called with complex function declarations
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
functionDeclarations: expect.arrayContaining([
expect.objectContaining({
name: "complex_function",
parameters: expect.objectContaining({
type: "OBJECT",
properties: expect.objectContaining({
user: expect.objectContaining({
type: "OBJECT",
properties: expect.any(Object),
}),
preferences: expect.objectContaining({
type: "ARRAY",
}),
settings: expect.objectContaining({
type: "OBJECT",
}),
}),
}),
}),
]),
})
);
// Verify the result contains the serialized function call
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({
name: "complex_function",
args: {
user: { name: "John", age: 30 },
preferences: ["option1", "option2"],
settings: { theme: "dark", notifications: true },
},
}),
},
],
});
});
});
```
--------------------------------------------------------------------------------
/tests/unit/tools/geminiChatTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
// Using vitest globals - see vitest.config.ts globals: true
import { geminiChatTool } from "../../../src/tools/geminiChatTool.js";
import { GeminiApiError } from "../../../src/utils/errors.js";
import { McpError } from "@modelcontextprotocol/sdk/types.js";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { GeminiService } from "../../../src/services/index.js";
import { BlockedReason, FinishReason } from "@google/genai";
describe("geminiChatTool", () => {
// Mock server and service instances
const mockTool = vi.fn();
const mockServer = {
tool: mockTool,
} as unknown as McpServer;
// Create mock functions for the service methods
const mockStartChatSession = vi.fn();
const mockSendMessageToSession = vi.fn();
const mockSendFunctionResultToSession = vi.fn();
// Create a minimal mock service with just the necessary methods for testing
const mockService = {
startChatSession: mockStartChatSession,
sendMessageToSession: mockSendMessageToSession,
sendFunctionResultToSession: mockSendFunctionResultToSession,
// Add empty implementations for required GeminiService methods
generateContent: () => Promise.resolve("mock"),
} as unknown as GeminiService;
// Reset mocks before each test
beforeEach(() => {
vi.resetAllMocks();
});
it("should register the tool with the server", () => {
// Call the tool registration function
geminiChatTool(mockServer, mockService);
// Verify tool was registered
expect(mockTool).toHaveBeenCalledTimes(1);
const [name, description, params, handler] = mockTool.mock.calls[0];
// Check tool registration parameters
expect(name).toBe("gemini_chat");
expect(description).toContain("Manages stateful chat sessions");
expect(params).toBeDefined();
expect(typeof handler).toBe("function");
});
describe("start operation", () => {
it("should start a new chat session", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-123";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request
const testRequest = {
operation: "start",
modelName: "gemini-1.5-flash",
history: [
{
role: "user",
parts: [{ text: "Hello" }],
},
{
role: "model",
parts: [{ text: "Hi there!" }],
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the service method was called with correct parameters
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-1.5-flash",
history: testRequest.history,
})
);
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
it("should start chat session with optional parameters", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-456";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request with optional parameters
const testRequest = {
operation: "start",
modelName: "gemini-2.0-flash",
systemInstruction: {
parts: [{ text: "You are a helpful assistant" }],
},
generationConfig: {
temperature: 0.7,
maxOutputTokens: 1000,
},
safetySettings: [
{
category: "HARM_CATEGORY_HATE_SPEECH",
threshold: "BLOCK_MEDIUM_AND_ABOVE",
},
],
cachedContentName: "cachedContents/abc123",
};
// Call the handler
const result = await handler(testRequest);
// Verify all parameters were passed
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
modelName: "gemini-2.0-flash",
systemInstruction: testRequest.systemInstruction,
generationConfig: testRequest.generationConfig,
safetySettings: testRequest.safetySettings,
cachedContentName: "cachedContents/abc123",
})
);
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
});
describe("send_message operation", () => {
it("should send a message to an existing session", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "The capital of France is Paris." }],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "What is the capital of France?",
};
// Call the handler
const result = await handler(testRequest);
// Verify the service method was called
expect(mockSendMessageToSession).toHaveBeenCalledWith({
sessionId: "test-session-123",
message: "What is the capital of France?",
generationConfig: undefined,
safetySettings: undefined,
tools: undefined,
toolConfig: undefined,
cachedContentName: undefined,
});
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: "The capital of France is Paris.",
},
],
});
});
it("should handle function call responses", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock function call response
const mockResponse = {
candidates: [
{
content: {
parts: [
{
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
},
],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "What's the weather in Paris?",
tools: [
{
functionDeclarations: [
{
name: "get_weather",
description: "Get weather information",
parameters: {
type: "OBJECT",
properties: {
location: {
type: "STRING",
description: "The location",
},
},
},
},
],
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the result contains function call
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
}),
},
],
});
});
it("should handle safety blocked responses", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock safety blocked response
const mockResponse = {
promptFeedback: {
blockReason: BlockedReason.SAFETY,
},
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify error response
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Prompt blocked due to safety settings . Reason: SAFETY",
},
],
isError: true,
});
});
it("should throw error if sessionId is missing", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Prepare test request without sessionId
const testRequest = {
operation: "send_message",
message: "Test message",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow(
"sessionId is required for operation 'send_message'"
);
});
it("should throw error if message is missing", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Prepare test request without message
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow(
"message is required for operation 'send_message'"
);
});
});
describe("send_function_result operation", () => {
it("should send function results to session", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "The weather in Paris is sunny and 22°C." }],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
functionResponses: [
{
name: "get_weather",
response: {
temperature: 22,
condition: "sunny",
location: "Paris",
},
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify the service method was called
expect(mockSendFunctionResultToSession).toHaveBeenCalledWith({
sessionId: "test-session-123",
functionResponse: JSON.stringify(testRequest.functionResponses),
functionCall: undefined,
});
// Verify the result
expect(result).toEqual({
content: [
{
type: "text",
text: "The weather in Paris is sunny and 22°C.",
},
],
});
});
it("should handle empty candidates", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with no candidates
const mockResponse = {
candidates: [],
};
mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
functionResponses: [
{
name: "test_function",
response: { result: "test" },
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify error response
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: No response candidates returned by the model after function result.",
},
],
isError: true,
});
});
it("should throw error if sessionId is missing", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Prepare test request without sessionId
const testRequest = {
operation: "send_function_result",
functionResponses: [
{
name: "test_function",
response: { result: "test" },
},
],
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow(
"sessionId is required for operation 'send_function_result'"
);
});
it("should throw error if functionResponses is missing", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Prepare test request without functionResponses
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow(
"functionResponses is required for operation 'send_function_result'"
);
});
});
describe("error handling", () => {
it("should map GeminiApiError to McpError", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock service to throw GeminiApiError
const geminiError = new GeminiApiError("API error occurred");
mockStartChatSession.mockImplementationOnce(() => {
throw geminiError;
});
// Prepare test request
const testRequest = {
operation: "start",
modelName: "gemini-1.5-flash",
};
// Call the handler and expect McpError
await expect(handler(testRequest)).rejects.toThrow();
// Verify the error was caught and mapped
try {
await handler(testRequest);
} catch (error) {
expect(error).toBeInstanceOf(McpError);
expect((error as McpError).message).toContain("API error occurred");
}
});
it("should handle invalid operation", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Prepare test request with invalid operation
const testRequest = {
operation: "invalid_operation",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow(
"Invalid operation: invalid_operation"
);
});
});
describe("response processing edge cases", () => {
it("should handle candidate with SAFETY finish reason", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with SAFETY finish reason
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Partial response..." }],
},
finishReason: FinishReason.SAFETY,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify error response for SAFETY finish reason
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Response generation stopped due to safety settings . FinishReason: SAFETY",
},
],
isError: true,
});
});
it("should handle candidate with MAX_TOKENS finish reason", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with MAX_TOKENS finish reason
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Response cut off due to token limit..." }],
},
finishReason: FinishReason.MAX_TOKENS,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify successful response even with MAX_TOKENS (this is acceptable)
expect(result).toEqual({
content: [
{
type: "text",
text: "Response cut off due to token limit...",
},
],
});
});
it("should handle candidate with OTHER finish reason", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with OTHER finish reason
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Some response..." }],
},
finishReason: FinishReason.OTHER,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify response is still returned but with warning logged
expect(result).toEqual({
content: [
{
type: "text",
text: "Some response...",
},
],
});
});
it("should handle empty content parts", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with empty content parts
const mockResponse = {
candidates: [
{
content: {
parts: [],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify error response for empty content
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Empty response from the model .",
},
],
isError: true,
});
});
it("should handle missing content in candidate", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with missing content
const mockResponse = {
candidates: [
{
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify error response for missing content
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Empty response from the model .",
},
],
isError: true,
});
});
it("should handle mixed content parts (text and function call)", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with both text and function call
const mockResponse = {
candidates: [
{
content: {
parts: [
{ text: "I'll help you with that. " },
{
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
},
{ text: " Let me check the weather for you." },
],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "What's the weather in Paris?",
};
// Call the handler
const result = await handler(testRequest);
// Verify function call is returned (function call takes precedence)
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({
functionCall: {
name: "get_weather",
args: { location: "Paris" },
},
}),
},
],
});
});
it("should handle unexpected response structure", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock response with unexpected structure (no text or function call)
const mockResponse = {
candidates: [
{
content: {
parts: [{ someOtherField: "unexpected data" }],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler
const result = await handler(testRequest);
// Verify error response for unexpected structure
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Unexpected response structure from the model .",
},
],
isError: true,
});
});
});
describe("advanced parameter combinations", () => {
it("should handle start operation with tools and toolConfig", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-tools";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request with tools
const testRequest = {
operation: "start",
modelName: "gemini-1.5-pro",
tools: [
{
functionDeclarations: [
{
name: "calculate",
description: "Perform calculations",
parameters: {
type: "OBJECT",
properties: {
expression: {
type: "STRING",
description: "Mathematical expression",
},
},
required: ["expression"],
},
},
],
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify tools were passed
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
tools: testRequest.tools,
})
);
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
it("should handle send_message with all optional parameters", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Response with all parameters" }],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request with all optional parameters
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message with all params",
generationConfig: {
temperature: 0.5,
topP: 0.9,
topK: 40,
maxOutputTokens: 2048,
stopSequences: ["END"],
thinkingConfig: {
thinkingBudget: 1024,
},
},
safetySettings: [
{
category: "HARM_CATEGORY_HARASSMENT",
threshold: "BLOCK_MEDIUM_AND_ABOVE",
},
],
tools: [
{
functionDeclarations: [
{
name: "test_function",
description: "Test function",
parameters: {
type: "OBJECT",
properties: {},
},
},
],
},
],
toolConfig: {
functionCallingConfig: {
mode: "AUTO",
},
},
cachedContentName: "cachedContents/test123",
};
// Call the handler
const result = await handler(testRequest);
// Verify all parameters were passed
expect(mockSendMessageToSession).toHaveBeenCalledWith({
sessionId: "test-session-123",
message: "Test message with all params",
generationConfig: testRequest.generationConfig,
safetySettings: testRequest.safetySettings,
tools: testRequest.tools,
toolConfig: testRequest.toolConfig,
cachedContentName: "cachedContents/test123",
});
expect(result).toEqual({
content: [
{
type: "text",
text: "Response with all parameters",
},
],
});
});
it("should handle start operation with string systemInstruction", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-string-instruction";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request with string system instruction
const testRequest = {
operation: "start",
systemInstruction:
"You are a helpful assistant specialized in mathematics.",
};
// Call the handler
const result = await handler(testRequest);
// Verify string system instruction was passed
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
systemInstruction:
"You are a helpful assistant specialized in mathematics.",
})
);
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
it("should handle start operation with complex history", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-complex-history";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request with complex history
const testRequest = {
operation: "start",
history: [
{
role: "user",
parts: [{ text: "Hello, I need help with math." }],
},
{
role: "model",
parts: [
{
text: "I'd be happy to help you with mathematics! What specific topic or problem would you like assistance with?",
},
],
},
{
role: "user",
parts: [{ text: "Can you solve quadratic equations?" }],
},
{
role: "model",
parts: [
{
text: "Yes, I can help you solve quadratic equations. The general form is ax² + bx + c = 0.",
},
],
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify complex history was passed
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
history: testRequest.history,
})
);
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
});
describe("function result processing", () => {
it("should handle function result with safety blocked response", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock safety blocked response after function result
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Partial response..." }],
},
finishReason: FinishReason.SAFETY,
},
],
};
mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
functionResponses: [
{
name: "test_function",
response: { result: "test result" },
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify error response includes "after function result" context
expect(result).toEqual({
content: [
{
type: "text",
text: "Error: Response generation stopped due to safety settings after function result. FinishReason: SAFETY",
},
],
isError: true,
});
});
it("should handle multiple function responses", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = {
candidates: [
{
content: {
parts: [
{
text: "Based on the function results, here's the summary...",
},
],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request with multiple function responses
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
functionResponses: [
{
name: "get_weather",
response: { temperature: 22, condition: "sunny" },
},
{
name: "get_time",
response: { time: "14:30", timezone: "UTC" },
},
],
};
// Call the handler
const result = await handler(testRequest);
// Verify multiple function responses were serialized correctly
expect(mockSendFunctionResultToSession).toHaveBeenCalledWith({
sessionId: "test-session-123",
functionResponse: JSON.stringify(testRequest.functionResponses),
functionCall: undefined,
});
expect(result).toEqual({
content: [
{
type: "text",
text: "Based on the function results, here's the summary...",
},
],
});
});
});
describe("service error handling", () => {
it("should handle service errors during start operation", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock service to throw error
const serviceError = new Error("Service unavailable");
mockStartChatSession.mockImplementationOnce(() => {
throw serviceError;
});
// Prepare test request
const testRequest = {
operation: "start",
modelName: "gemini-1.5-flash",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow();
});
it("should handle service errors during send_message operation", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock service to throw error
const serviceError = new Error("Network error");
mockSendMessageToSession.mockRejectedValueOnce(serviceError);
// Prepare test request
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Test message",
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow();
});
it("should handle service errors during send_function_result operation", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock service to throw error
const serviceError = new Error("Function processing error");
mockSendFunctionResultToSession.mockRejectedValueOnce(serviceError);
// Prepare test request
const testRequest = {
operation: "send_function_result",
sessionId: "test-session-123",
functionResponses: [
{
name: "test_function",
response: { result: "test" },
},
],
};
// Call the handler and expect error
await expect(handler(testRequest)).rejects.toThrow();
});
});
describe("thinking configuration", () => {
it("should handle thinkingConfig in generationConfig for start operation", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockSessionId = "test-session-thinking";
mockStartChatSession.mockReturnValueOnce(mockSessionId);
// Prepare test request with thinking configuration
const testRequest = {
operation: "start",
generationConfig: {
temperature: 0.7,
thinkingConfig: {
thinkingBudget: 2048,
reasoningEffort: "medium",
},
},
};
// Call the handler
const result = await handler(testRequest);
// Verify thinking config was passed
expect(mockStartChatSession).toHaveBeenCalledWith(
expect.objectContaining({
generationConfig: testRequest.generationConfig,
})
);
expect(result).toEqual({
content: [
{
type: "text",
text: JSON.stringify({ sessionId: mockSessionId }),
},
],
});
});
it("should handle reasoningEffort in thinkingConfig", async () => {
// Register tool to get the request handler
geminiChatTool(mockServer, mockService);
const [, , , handler] = mockTool.mock.calls[0];
// Mock successful response
const mockResponse = {
candidates: [
{
content: {
parts: [{ text: "Response with reasoning effort" }],
},
finishReason: FinishReason.STOP,
},
],
};
mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
// Prepare test request with reasoning effort
const testRequest = {
operation: "send_message",
sessionId: "test-session-123",
message: "Complex reasoning task",
generationConfig: {
thinkingConfig: {
reasoningEffort: "high",
},
},
};
// Call the handler
const result = await handler(testRequest);
// Verify reasoning effort was passed
expect(mockSendMessageToSession).toHaveBeenCalledWith(
expect.objectContaining({
generationConfig: testRequest.generationConfig,
})
);
expect(result).toEqual({
content: [
{
type: "text",
text: "Response with reasoning effort",
},
],
});
});
});
});
```
--------------------------------------------------------------------------------
/src/services/mcp/McpClientService.ts:
--------------------------------------------------------------------------------
```typescript
import { logger } from "../../utils/index.js";
import { spawn, ChildProcess } from "child_process";
import EventSource from "eventsource";
import {
McpError as SdkMcpError,
ErrorCode,
} from "@modelcontextprotocol/sdk/types.js";
import { v4 as uuidv4 } from "uuid";
// Import node-fetch types only
// We'll dynamically import the actual implementation later to handle CJS/ESM compatibility
import type { Response, RequestInit } from "node-fetch";
// Define custom types for EventSource events since the eventsource package
// doesn't export its own types
interface ESMessageEvent {
data: string;
type: string;
lastEventId: string;
origin: string;
}
interface ESErrorEvent {
type: string;
message?: string;
error?: Error;
}
// Add appropriate error handler typings
type ESErrorHandler = (event: ESErrorEvent) => void;
type ESMessageHandler = (event: ESMessageEvent) => void;
// Extended EventSource interface to properly type the handlers
interface ExtendedEventSource extends EventSource {
onopen: (this: EventSource, ev: MessageEvent<unknown>) => unknown;
onmessage: (this: EventSource, ev: MessageEvent) => unknown;
onerror: (this: EventSource, ev: Event) => unknown;
}
export interface McpRequest {
id: string;
method: "listTools" | "callTool" | "initialize";
params?: Record<string, unknown>;
}
export interface McpResponseError {
code: number;
message: string;
data?: Record<string, unknown>;
}
export interface McpResponse {
id: string;
result?: Record<string, unknown> | Array<unknown>;
error?: McpResponseError;
}
export interface ToolDefinition {
name: string;
description: string;
parametersSchema: Record<string, unknown>; // JSON Schema object
}
export interface ConnectionDetails {
type: "sse" | "stdio";
sseUrl?: string;
stdioCommand?: string;
stdioArgs?: string[];
connectionToken?: string;
}
// Re-export SDK McpError under the local name used throughout this file
type McpError = SdkMcpError;
const McpError = SdkMcpError;
/**
* Service for connecting to external Model Context Protocol (MCP) servers.
* Provides methods to establish different types of connections (SSE, stdio).
*/
export class McpClientService {
// Maps to store active connections
private activeSseConnections: Map<
string,
{
eventSource: EventSource;
baseUrl: string;
lastActivityTimestamp: number; // Track when connection was last used
}
>;
private activeStdioConnections: Map<
string,
{
process: ChildProcess;
lastActivityTimestamp: number; // Track when connection was last used
}
>;
private pendingStdioRequests: Map<
string, // connectionId
Map<
string,
{
resolve: (value: Record<string, unknown> | Array<unknown>) => void;
reject: (reason: Error | McpError) => void;
}
> // requestId -> handlers
> = new Map();
// Configuration values
private static readonly DEFAULT_REQUEST_TIMEOUT_MS = 30000; // 30 seconds
private static readonly DEFAULT_CONNECTION_MAX_IDLE_MS = 600000; // 10 minutes
private static readonly CONNECTION_CLEANUP_INTERVAL_MS = 300000; // Check every 5 minutes
/**
* Helper method to fetch with timeout
* @param url - The URL to fetch
* @param options - Fetch options
* @param timeoutMs - Timeout in milliseconds
* @param timeoutMessage - Message to include in timeout error
* @returns The fetch response
* @throws {SdkMcpError} - If the request times out
*/
private async fetchWithTimeout(
url: string,
options: RequestInit,
timeoutMs = McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
timeoutMessage = "Request timed out"
): Promise<Response> {
// Create controller for aborting the fetch
const controller = new AbortController();
const id = setTimeout(() => controller.abort(), timeoutMs);
try {
// Add the signal to the options
const fetchOptions = {
...options,
signal: controller.signal,
};
// Dynamically import node-fetch (v2 is CommonJS)
const nodeFetch = await import("node-fetch");
const fetch = nodeFetch.default;
// Make the fetch request
const response = await fetch(url, fetchOptions);
clearTimeout(id);
return response;
} catch (error) {
clearTimeout(id);
throw new SdkMcpError(
ErrorCode.InternalError,
`${timeoutMessage} after ${timeoutMs}ms`
);
}
}
// Cleanup timer reference
private cleanupIntervalId?: NodeJS.Timeout;
constructor() {
this.activeSseConnections = new Map();
this.activeStdioConnections = new Map();
logger.info("McpClientService initialized.");
// Start the connection cleanup interval
this.cleanupIntervalId = setInterval(
() => this.cleanupStaleConnections(),
McpClientService.CONNECTION_CLEANUP_INTERVAL_MS
);
}
/**
* Cleans up stale connections that haven't been used for a while
* @private
*/
private cleanupStaleConnections(): void {
const now = Date.now();
const maxIdleTime = McpClientService.DEFAULT_CONNECTION_MAX_IDLE_MS;
let closedCount = 0;
// Check SSE connections
for (const [
connectionId,
{ lastActivityTimestamp },
] of this.activeSseConnections.entries()) {
if (now - lastActivityTimestamp > maxIdleTime) {
logger.info(
`Closing stale SSE connection ${connectionId} (idle for ${Math.floor((now - lastActivityTimestamp) / 1000)} seconds)`
);
this.closeSseConnection(connectionId);
closedCount++;
}
}
// Check stdio connections
for (const [
connectionId,
{ lastActivityTimestamp },
] of this.activeStdioConnections.entries()) {
if (now - lastActivityTimestamp > maxIdleTime) {
logger.info(
`Closing stale stdio connection ${connectionId} (idle for ${Math.floor((now - lastActivityTimestamp) / 1000)} seconds)`
);
this.closeStdioConnection(connectionId);
closedCount++;
}
}
if (closedCount > 0) {
logger.info(`Cleaned up ${closedCount} stale connections`);
}
}
/**
* Validates a server ID.
* @param serverId - The server ID to validate.
* @throws {McpError} Throws an error if the server ID is invalid.
*/
private validateServerId(serverId: string): void {
if (!serverId || typeof serverId !== "string" || serverId.trim() === "") {
throw new McpError(
ErrorCode.InvalidParams,
"Server ID must be a non-empty string"
);
}
}
/**
* Checks if a connection exists for the given server ID.
* @param serverId - The server ID to check.
* @throws {McpError} Throws an error if the connection doesn't exist.
*/
private validateConnectionExists(serverId: string): void {
if (
!this.activeSseConnections.has(serverId) &&
!this.activeStdioConnections.has(serverId)
) {
throw new McpError(
ErrorCode.InvalidRequest,
`Connection not found for serverId: ${serverId}`
);
}
}
/**
* Establishes a connection to an MCP server.
* @param serverId - A unique identifier provided by the caller to reference this server connection.
* Note: This is NOT used as the internal connection tracking ID.
* @param connectionDetails - The details for establishing the connection.
* @param messageHandler - Optional callback for handling received messages.
* @returns A promise that resolves to a connection ID (different from serverId) when the connection is established.
* This returned connectionId should be used for all subsequent interactions with this connection.
* @throws {McpError} Throws an error if the parameters are invalid.
*/
public async connect(
serverId: string,
connectionDetails: ConnectionDetails,
messageHandler?: (data: unknown) => void
): Promise<string> {
// Validate serverId
this.validateServerId(serverId);
// Validate connectionDetails
if (!connectionDetails || typeof connectionDetails !== "object") {
throw new McpError(
ErrorCode.InvalidParams,
"Connection details must be an object"
);
}
// Validate connection type
if (
connectionDetails.type !== "sse" &&
connectionDetails.type !== "stdio"
) {
throw new McpError(
ErrorCode.InvalidParams,
"Connection type must be 'sse' or 'stdio'"
);
}
// Validate SSE connection details
if (connectionDetails.type === "sse") {
if (
!connectionDetails.sseUrl ||
typeof connectionDetails.sseUrl !== "string" ||
connectionDetails.sseUrl.trim() === ""
) {
throw new McpError(
ErrorCode.InvalidParams,
"For SSE connections, sseUrl must be a non-empty string"
);
}
// Basic URL format validation
if (
!connectionDetails.sseUrl.startsWith("http://") &&
!connectionDetails.sseUrl.startsWith("https://")
) {
throw new McpError(
ErrorCode.InvalidParams,
"sseUrl must be a valid URL format starting with http:// or https://"
);
}
return this.connectSse(
connectionDetails.sseUrl,
connectionDetails.connectionToken,
messageHandler
);
}
// Validate stdio connection details
else if (connectionDetails.type === "stdio") {
if (
!connectionDetails.stdioCommand ||
typeof connectionDetails.stdioCommand !== "string" ||
connectionDetails.stdioCommand.trim() === ""
) {
throw new McpError(
ErrorCode.InvalidParams,
"For stdio connections, stdioCommand must be a non-empty string"
);
}
return this.connectStdio(
connectionDetails.stdioCommand,
connectionDetails.stdioArgs || [],
connectionDetails.connectionToken,
messageHandler
);
}
// This should never be reached due to the type check above
throw new McpError(
ErrorCode.InvalidParams,
"Invalid connection type specified"
);
}
/**
* Establishes an SSE connection to the specified MCP server.
* @param url - The URL of the MCP server to connect to.
* @param connectionToken - Optional token for authentication with the server.
* @param messageHandler - Optional callback for handling received messages.
* @returns A promise that resolves to a connection ID when the connection is established.
*/
private connectSse(
url: string,
connectionToken?: string,
messageHandler?: (data: unknown) => void
): Promise<string> {
return new Promise((resolve, reject) => {
logger.info(`Connecting to MCP server via SSE: ${url}`);
try {
// Generate a unique connectionId for internal tracking
// This will be different from the serverId passed to the connect() method
const connectionId = uuidv4();
// Create a timeout for the connection attempt
const connectionTimeout = setTimeout(() => {
reject(
new SdkMcpError(
ErrorCode.InternalError,
`Connection timeout while attempting to connect to ${url}`
)
);
}, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
// Add connectionToken to headers if provided
const options: EventSource.EventSourceInitDict = {};
if (connectionToken) {
logger.debug(`Adding connection token to SSE request`);
options.headers = {
Authorization: `Bearer ${connectionToken}`,
};
}
// Create EventSource for SSE connection with options
const eventSource = new EventSource(url, options);
// Handler functions to store for proper cleanup
const onOpen = () => {
// Clear the connection timeout
clearTimeout(connectionTimeout);
logger.info(`SSE connection established to ${url}`);
this.activeSseConnections.set(connectionId, {
eventSource,
baseUrl: url,
lastActivityTimestamp: Date.now(),
});
resolve(connectionId);
};
const onMessage = ((event: ESMessageEvent) => {
logger.debug(`SSE message received from ${url}:`, event.data);
// Update the last activity timestamp
const connection = this.activeSseConnections.get(connectionId);
if (connection) {
connection.lastActivityTimestamp = Date.now();
}
if (messageHandler) {
try {
const parsedData = JSON.parse(event.data);
messageHandler(parsedData);
} catch (error) {
logger.error(`Error parsing SSE message:`, error);
messageHandler(event.data);
}
}
}) as ESMessageHandler;
const onError = ((error: ESErrorEvent) => {
// Clear the connection timeout if it's still pending
clearTimeout(connectionTimeout);
logger.error(
`SSE connection error for ${url}:`,
error.message || "Unknown error"
);
if (!this.activeSseConnections.has(connectionId)) {
// If we haven't resolved yet, this is a connection failure
reject(
new SdkMcpError(
ErrorCode.InternalError,
`Failed to establish SSE connection to ${url}: ${error.message || "Unknown error"}`
)
);
} else if (eventSource.readyState === EventSource.CLOSED) {
// Connection was established but is now closed
logger.info(`SSE connection ${connectionId} closed due to error.`);
this.activeSseConnections.delete(connectionId);
} else {
// Connection is still open but had an error
logger.warn(
`SSE connection ${connectionId} had an error but is still open. Monitoring for further issues.`
);
}
}) as ESErrorHandler;
// Set up event handlers
eventSource.onopen = onOpen;
eventSource.onmessage = onMessage;
eventSource.onerror = onError;
} catch (error) {
logger.error(`Error creating SSE connection to ${url}:`, error);
reject(error);
}
});
}
/**
* Closes an SSE connection.
* @param connectionId - The ID of the connection to close.
* @returns True if the connection was closed, false if it wasn't found.
*/
public closeSseConnection(connectionId: string): boolean {
const connection = this.activeSseConnections.get(connectionId);
if (connection) {
// Close the EventSource and remove listeners
const eventSource = connection.eventSource;
// Clean up event listeners by setting handlers to empty functions
// (EventSource doesn't support removeEventListener)
(eventSource as ExtendedEventSource).onopen = () => {};
(eventSource as ExtendedEventSource).onmessage = () => {};
(eventSource as ExtendedEventSource).onerror = () => {};
// Close the connection
eventSource.close();
// Remove from active connections
this.activeSseConnections.delete(connectionId);
// Clean up any pending requests for this connection (this shouldn't generally happen for SSE)
this.cleanupPendingRequestsForConnection(connectionId);
logger.info(`SSE connection ${connectionId} closed.`);
return true;
}
logger.warn(
`Attempted to close non-existent SSE connection: ${connectionId}`
);
return false;
}
/**
* Helper method to clean up pending requests for a connection
* @param connectionId - The ID of the connection to clean up pending requests for
*/
private cleanupPendingRequestsForConnection(connectionId: string): void {
// If there are any pending requests for this connection, reject them all
if (this.pendingStdioRequests.has(connectionId)) {
const pendingRequests = this.pendingStdioRequests.get(connectionId)!;
for (const [
requestId,
{ reject: rejectRequest },
] of pendingRequests.entries()) {
logger.warn(
`Rejecting pending request ${requestId} due to connection cleanup`
);
rejectRequest(
new McpError(
ErrorCode.InternalError,
`Connection closed during cleanup before response was received`
)
);
}
// Clean up the map entry
this.pendingStdioRequests.delete(connectionId);
}
}
/**
* Establishes a stdio connection using the specified command.
* @param command - The command to execute for stdio connection.
* @param args - Arguments to pass to the command.
* @param connectionToken - Optional token for authentication with the server.
* @param messageHandler - Optional callback for handling stdout data.
* @returns A promise that resolves to a connection ID when the process is spawned.
*/
private connectStdio(
command: string,
args: string[] = [],
connectionToken?: string,
messageHandler?: (data: unknown) => void
): Promise<string> {
return new Promise((resolve, reject) => {
logger.info(
`Connecting to MCP server via stdio using command: ${command} ${args.join(" ")}`
);
try {
// Generate a unique connectionId for internal tracking
// This will be different from the serverId passed to the connect() method
const connectionId = uuidv4();
// Create a timeout for the connection establishment
const connectionTimeout = setTimeout(() => {
reject(
new SdkMcpError(
ErrorCode.InternalError,
`Timeout while establishing stdio connection for command: ${command}`
)
);
}, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
// Prepare the environment for the child process
const env = { ...process.env };
// Add connectionToken to environment if provided
if (connectionToken) {
logger.debug("Adding connection token to stdio environment");
env.MCP_CONNECTION_TOKEN = connectionToken;
}
// Spawn the child process with environment
const childProcess = spawn(command, args, {
stdio: "pipe",
env: env,
});
// Store the connection with timestamp
this.activeStdioConnections.set(connectionId, {
process: childProcess,
lastActivityTimestamp: Date.now(),
});
// Buffer to accumulate data chunks
let buffer = "";
// We'll mark connection as established when the process is ready
const connectionEstablished = () => {
clearTimeout(connectionTimeout);
logger.info(`Stdio connection established for ${command}`);
resolve(connectionId);
};
// Data handler function for stdout
const onStdoutData = (data: Buffer) => {
// Update the last activity timestamp to prevent cleanup
const connection = this.activeStdioConnections.get(connectionId);
if (connection) {
connection.lastActivityTimestamp = Date.now();
}
// Append the new data to our buffer
buffer += data.toString();
logger.debug(
`Stdio stdout from ${command} (raw chunk):`,
data.toString()
);
// Process complete lines in the buffer
let newlineIndex;
while ((newlineIndex = buffer.indexOf("\n")) !== -1) {
// Extract the line (excluding the newline)
const line = buffer.substring(0, newlineIndex);
// Remove the processed line from the buffer (including the newline)
buffer = buffer.substring(newlineIndex + 1);
// Skip empty lines
if (!line.trim()) continue;
try {
// Try to parse the line as JSON
const parsedData = JSON.parse(line);
logger.debug(`Parsed JSON message:`, parsedData);
// Check if this is a response to a pending request
if (parsedData.id) {
// Only check pending requests for the current connection (connectionId)
const requestsMap = this.pendingStdioRequests.get(connectionId);
let foundRequest = false;
if (requestsMap && requestsMap.has(parsedData.id)) {
const { resolve, reject } = requestsMap.get(parsedData.id)!;
requestsMap.delete(parsedData.id);
// If this was the last pending request, clean up the connection map
if (requestsMap.size === 0) {
this.pendingStdioRequests.delete(connectionId);
}
foundRequest = true;
if (parsedData.error) {
reject(
new SdkMcpError(
(parsedData.error.code as ErrorCode) ||
ErrorCode.InternalError,
parsedData.error.message || "Tool execution error",
parsedData.error.data
)
);
} else {
// Verify the result is an object or array
if (
parsedData.result === null ||
parsedData.result === undefined
) {
reject(
new McpError(
ErrorCode.InternalError,
"Received null or undefined result from tool",
{ responseId: parsedData.id }
)
);
} else if (
typeof parsedData.result !== "object" &&
!Array.isArray(parsedData.result)
) {
reject(
new McpError(
ErrorCode.InternalError,
"Expected object or array result from tool",
{
responseId: parsedData.id,
receivedType: typeof parsedData.result,
}
)
);
} else {
resolve(
parsedData.result as
| Record<string, unknown>
| Array<unknown>
);
}
}
}
// Only log if we didn't find the request
if (!foundRequest && messageHandler) {
logger.debug(
`Received message with ID ${parsedData.id} but no matching pending request found for this connection`
);
messageHandler(parsedData);
}
} else if (messageHandler) {
// If not a response to a pending request, pass to the message handler
messageHandler(parsedData);
}
} catch (error) {
logger.warn(`Error parsing JSON from stdio:`, error);
// If not valid JSON and we have a message handler, pass the raw line
if (messageHandler) {
messageHandler(line);
}
}
}
};
// Error handler for stderr
const onStderrData = (data: Buffer) => {
// Update the last activity timestamp to prevent cleanup
const connection = this.activeStdioConnections.get(connectionId);
if (connection) {
connection.lastActivityTimestamp = Date.now();
}
logger.warn(`Stdio stderr from ${command}:`, data.toString());
};
// Error handler for the process
const onError = (error: Error) => {
// Clear the connection timeout
clearTimeout(connectionTimeout);
logger.error(`Stdio error for ${command}:`, error);
if (this.activeStdioConnections.has(connectionId)) {
this.activeStdioConnections.delete(connectionId);
// Reject all pending requests for this connection
if (this.pendingStdioRequests.has(connectionId)) {
const pendingRequests =
this.pendingStdioRequests.get(connectionId)!;
for (const [
requestId,
{ reject: rejectRequest },
] of pendingRequests.entries()) {
logger.warn(
`Rejecting pending request ${requestId} due to connection error`
);
rejectRequest(
new SdkMcpError(
ErrorCode.InternalError,
`Connection error occurred before response: ${error.message}`
)
);
}
this.pendingStdioRequests.delete(connectionId);
}
}
reject(error);
};
// Close handler for the process
const onClose = (
code: number | null,
signal: NodeJS.Signals | null
) => {
// Clear the connection timeout if process closes before we establish connection
clearTimeout(connectionTimeout);
logger.info(
`Stdio process ${command} closed with code ${code} and signal ${signal}`
);
if (this.activeStdioConnections.has(connectionId)) {
this.activeStdioConnections.delete(connectionId);
// Reject all pending requests for this connection
if (this.pendingStdioRequests.has(connectionId)) {
const pendingRequests =
this.pendingStdioRequests.get(connectionId)!;
for (const [
requestId,
{ reject: rejectRequest },
] of pendingRequests.entries()) {
logger.warn(
`Rejecting pending request ${requestId} due to connection closure`
);
rejectRequest(
new McpError(
ErrorCode.InternalError,
`Connection closed before response (code: ${code}, signal: ${signal})`
)
);
}
this.pendingStdioRequests.delete(connectionId);
}
}
};
// Set up event handlers
childProcess.stdout.on("data", onStdoutData);
childProcess.stderr.on("data", onStderrData);
childProcess.on("error", onError);
childProcess.on("close", onClose);
// The connection is established immediately after we set up event handlers
connectionEstablished();
} catch (error) {
logger.error(`Error creating stdio connection for ${command}:`, error);
reject(error);
}
});
}
/**
* Sends data to a stdio connection.
* @param connectionId - The ID of the connection to send data to.
* @param data - The data to send.
* @returns True if the data was sent, false if the connection wasn't found.
*/
private sendToStdio(connectionId: string, data: string | object): boolean {
const connection = this.activeStdioConnections.get(connectionId);
if (connection) {
const childProcess = connection.process;
// Update the last activity timestamp
connection.lastActivityTimestamp = Date.now();
// Safety check for data size to prevent buffer overflow
const dataStr = typeof data === "string" ? data : JSON.stringify(data);
// Limit data size to 1MB to prevent abuse
const MAX_DATA_SIZE = 1024 * 1024; // 1MB
if (dataStr.length > MAX_DATA_SIZE) {
logger.error(
`Data to send to stdio connection ${connectionId} exceeds size limit (${dataStr.length} > ${MAX_DATA_SIZE})`
);
return false;
}
if (childProcess.stdin) {
try {
childProcess.stdin.write(dataStr + "\n");
} catch (error) {
logger.error(
`Error writing to stdin for connection ${connectionId}:`,
error
);
return false;
}
} else {
logger.error(`Stdio connection ${connectionId} has no stdin`);
return false;
}
logger.debug(`Sent data to stdio connection ${connectionId}`);
return true;
}
logger.warn(
`Attempted to send data to non-existent stdio connection: ${connectionId}`
);
return false;
}
/**
* Closes a stdio connection.
* @param connectionId - The ID of the connection to close.
* @param signal - Optional signal to send to the process. Default is 'SIGTERM'.
* @returns True if the connection was closed, false if it wasn't found.
*/
public closeStdioConnection(
connectionId: string,
signal: NodeJS.Signals = "SIGTERM"
): boolean {
const connection = this.activeStdioConnections.get(connectionId);
if (connection) {
const childProcess = connection.process;
// Remove all listeners to prevent memory leaks
childProcess.stdout?.removeAllListeners();
childProcess.stderr?.removeAllListeners();
childProcess.removeAllListeners();
// Kill the process
childProcess.kill(signal);
// Remove from active connections
this.activeStdioConnections.delete(connectionId);
// Clean up any pending requests for this connection
this.cleanupPendingRequestsForConnection(connectionId);
logger.info(
`Stdio connection ${connectionId} closed with signal ${signal}.`
);
return true;
}
logger.warn(
`Attempted to close non-existent stdio connection: ${connectionId}`
);
return false;
}
/**
* Gets all active SSE connection IDs.
* @returns Array of active SSE connection IDs.
*/
public getActiveSseConnectionIds(): string[] {
return Array.from(this.activeSseConnections.keys());
}
/**
* Gets all active stdio connection IDs.
* @returns Array of active stdio connection IDs.
*/
public getActiveStdioConnectionIds(): string[] {
return Array.from(this.activeStdioConnections.keys());
}
/**
* Gets the last activity timestamp for a connection
* @param connectionId - The ID of the connection to check
* @returns The last activity timestamp in milliseconds since the epoch, or undefined if the connection doesn't exist
*/
public getLastActivityTimestamp(connectionId: string): number | undefined {
const sseConnection = this.activeSseConnections.get(connectionId);
if (sseConnection) {
return sseConnection.lastActivityTimestamp;
}
const stdioConnection = this.activeStdioConnections.get(connectionId);
if (stdioConnection) {
return stdioConnection.lastActivityTimestamp;
}
return undefined;
}
/**
* Lists all available tools from an MCP server.
* @param serverId - The ID of the connection to query.
* @returns A promise that resolves to an array of tool definitions.
* @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
*/
public async listTools(serverId: string): Promise<ToolDefinition[]> {
// Validate serverId
this.validateServerId(serverId);
// Validate connection exists
this.validateConnectionExists(serverId);
logger.info(`Listing tools for connection ${serverId}`);
// Check if this is an SSE connection
if (this.activeSseConnections.has(serverId)) {
const connection = this.activeSseConnections.get(serverId)!;
const requestId = uuidv4();
const request: McpRequest = { id: requestId, method: "listTools" };
try {
// Create URL for the MCP request
const mcpRequestUrl = new URL(connection.baseUrl);
// Update the connection's last activity timestamp
connection.lastActivityTimestamp = Date.now();
// Make the request with timeout
const response = await this.fetchWithTimeout(
mcpRequestUrl.toString(),
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(request),
},
McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
"Request timed out"
);
if (!response.ok) {
throw new McpError(
ErrorCode.InternalError,
`HTTP error from MCP server: ${response.status} ${response.statusText}`
);
}
const mcpResponse = (await response.json()) as McpResponse;
if (mcpResponse.error) {
throw new SdkMcpError(
ErrorCode.InternalError,
`MCP error: ${mcpResponse.error.message} (code: ${mcpResponse.error.code})`,
mcpResponse.error.data
);
}
// Type assertion with verification to ensure we have an array of ToolDefinition
const result = mcpResponse.result;
if (!Array.isArray(result)) {
throw new McpError(
ErrorCode.InternalError,
"Expected array of tools in response",
{ receivedType: typeof result }
);
}
return result as ToolDefinition[];
} catch (error) {
logger.error(
`Error listing tools for SSE connection ${serverId}:`,
error
);
// Wrap non-McpError instances
if (!(error instanceof McpError)) {
throw new McpError(
ErrorCode.InternalError,
`Failed to list tools for connection ${serverId}: ${error instanceof Error ? error.message : String(error)}`
);
}
throw error;
}
}
// Check if this is a stdio connection
else if (this.activeStdioConnections.has(serverId)) {
const requestId = uuidv4();
const request: McpRequest = { id: requestId, method: "listTools" };
return new Promise<ToolDefinition[]>((resolve, reject) => {
// Initialize the map for this connection if it doesn't exist
if (!this.pendingStdioRequests.has(serverId)) {
this.pendingStdioRequests.set(serverId, new Map());
}
// Store the promise resolution functions
this.pendingStdioRequests.get(serverId)!.set(requestId, {
resolve: (value) => {
// Type-safe resolution for tool definitions
resolve(value as ToolDefinition[]);
},
reject: reject,
});
// Set up a timeout to automatically reject this request if it takes too long
setTimeout(() => {
// If the request is still pending, reject it
if (
this.pendingStdioRequests.has(serverId) &&
this.pendingStdioRequests.get(serverId)!.has(requestId)
) {
// Get the reject function
const { reject: rejectRequest } = this.pendingStdioRequests
.get(serverId)!
.get(requestId)!;
// Delete the request
this.pendingStdioRequests.get(serverId)!.delete(requestId);
// If this was the last pending request, clean up the connection map
if (this.pendingStdioRequests.get(serverId)!.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
// Reject the request with a timeout error
rejectRequest(
new SdkMcpError(
ErrorCode.InternalError,
"Request timed out waiting for response"
)
);
}
}, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
// Send the request
const sent = this.sendToStdio(serverId, request);
if (!sent) {
// Clean up the pending request if sending fails
this.pendingStdioRequests.get(serverId)!.delete(requestId);
// If this was the last pending request, clean up the connection map
if (this.pendingStdioRequests.get(serverId)!.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
reject(
new Error(`Failed to send request to stdio connection ${serverId}`)
);
}
});
}
// This should never be reached due to the validateConnectionExists check above
throw new McpError(
ErrorCode.InvalidRequest,
`No connection found with ID ${serverId}`
);
}
/**
* Gets server information from an MCP server.
* @param serverId - The ID of the connection to use.
* @returns A promise that resolves to the server information.
* @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
*/
public async getServerInfo(
serverId: string
): Promise<Record<string, unknown>> {
// Validate serverId
this.validateServerId(serverId);
// Check if connection exists
this.validateConnectionExists(serverId);
logger.debug(`Getting server info for connection: ${serverId}`);
// Check if this is an SSE connection
if (this.activeSseConnections.has(serverId)) {
const connection = this.activeSseConnections.get(serverId)!;
connection.lastActivityTimestamp = Date.now();
// For SSE connections, we'll make an HTTP request to get server info
const baseUrl = connection.baseUrl;
const infoUrl = `${baseUrl}/info`;
try {
const response = await this.fetchWithTimeout(
infoUrl,
{
method: "GET",
headers: {
"Content-Type": "application/json",
},
},
McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
`Server info request timed out for ${serverId}`
);
if (!response.ok) {
throw new McpError(
ErrorCode.InternalError,
`Server info request failed with status ${response.status}: ${response.statusText}`
);
}
const serverInfo = await response.json();
return serverInfo as Record<string, unknown>;
} catch (error) {
logger.error(
`Error getting server info for SSE connection ${serverId}:`,
error
);
throw new McpError(
ErrorCode.InternalError,
`Failed to get server info: ${error instanceof Error ? error.message : "Unknown error"}`
);
}
}
// Check if this is a stdio connection
if (this.activeStdioConnections.has(serverId)) {
const connection = this.activeStdioConnections.get(serverId)!;
connection.lastActivityTimestamp = Date.now();
// For stdio connections, send an initialize request
const requestId = uuidv4();
const request: McpRequest = {
id: requestId,
method: "initialize",
params: {
protocolVersion: "2024-11-05",
capabilities: {},
clientInfo: {
name: "mcp-gemini-server",
version: "1.0.0",
},
},
};
return new Promise<Record<string, unknown>>((resolve, reject) => {
// Set up timeout for the request
const timeout = setTimeout(() => {
// Clean up the pending request
const pendingRequests = this.pendingStdioRequests.get(serverId);
if (pendingRequests) {
pendingRequests.delete(requestId);
if (pendingRequests.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
}
reject(
new McpError(
ErrorCode.InternalError,
`Server info request timed out for ${serverId} after ${McpClientService.DEFAULT_REQUEST_TIMEOUT_MS}ms`
)
);
}, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
// Store the request handlers
if (!this.pendingStdioRequests.has(serverId)) {
this.pendingStdioRequests.set(serverId, new Map());
}
this.pendingStdioRequests.get(serverId)!.set(requestId, {
resolve: (result) => {
clearTimeout(timeout);
resolve(result as Record<string, unknown>);
},
reject: (error) => {
clearTimeout(timeout);
reject(error);
},
});
// Send the request
const success = this.sendToStdio(serverId, request);
if (!success) {
// Clean up the pending request
const pendingRequests = this.pendingStdioRequests.get(serverId);
if (pendingRequests) {
pendingRequests.delete(requestId);
if (pendingRequests.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
}
clearTimeout(timeout);
reject(
new McpError(
ErrorCode.InternalError,
`Failed to send server info request to ${serverId}`
)
);
}
});
}
// This should never be reached due to the validateConnectionExists check above
throw new McpError(
ErrorCode.InvalidRequest,
`No connection found with ID ${serverId}`
);
}
/**
* Calls a tool on an MCP server.
* @param serverId - The ID of the connection to use.
* @param toolName - The name of the tool to call.
* @param toolArgs - The arguments to pass to the tool.
* @returns A promise that resolves to the tool's result.
* @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
*/
public async callTool(
serverId: string,
toolName: string,
toolArgs: Record<string, unknown> | null | undefined
): Promise<Record<string, unknown> | Array<unknown>> {
// Validate serverId
this.validateServerId(serverId);
// Validate connection exists
this.validateConnectionExists(serverId);
// Validate toolName
if (!toolName || typeof toolName !== "string" || toolName.trim() === "") {
throw new McpError(
ErrorCode.InvalidParams,
"Tool name must be a non-empty string"
);
}
// Validate toolArgs (ensure it's an object if provided)
if (
toolArgs !== null &&
toolArgs !== undefined &&
typeof toolArgs !== "object"
) {
throw new McpError(
ErrorCode.InvalidParams,
"Tool arguments must be an object, null, or undefined"
);
}
// Normalize toolArgs to an empty object if null or undefined
const normalizedToolArgs: Record<string, unknown> = toolArgs || {};
logger.info(`Calling tool ${toolName} on connection ${serverId}`);
// Check if this is an SSE connection
if (this.activeSseConnections.has(serverId)) {
const connection = this.activeSseConnections.get(serverId)!;
const requestId = uuidv4();
const request: McpRequest = {
id: requestId,
method: "callTool",
params: { toolName, arguments: normalizedToolArgs },
};
try {
// Create URL for the MCP request
const mcpRequestUrl = new URL(connection.baseUrl);
// Update the connection's last activity timestamp
connection.lastActivityTimestamp = Date.now();
// Make the request with timeout
const response = await this.fetchWithTimeout(
mcpRequestUrl.toString(),
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(request),
},
McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
"Request timed out"
);
if (!response.ok) {
throw new McpError(
ErrorCode.InternalError,
`HTTP error from MCP server: ${response.status} ${response.statusText}`
);
}
const mcpResponse = (await response.json()) as McpResponse;
if (mcpResponse.error) {
throw new SdkMcpError(
ErrorCode.InternalError,
`MCP error: ${mcpResponse.error.message} (code: ${mcpResponse.error.code})`,
mcpResponse.error.data
);
}
// Ensure result is either an object or array
if (
!mcpResponse.result ||
(typeof mcpResponse.result !== "object" &&
!Array.isArray(mcpResponse.result))
) {
throw new McpError(
ErrorCode.InternalError,
"Expected object or array result from tool call",
{ receivedType: typeof mcpResponse.result }
);
}
return mcpResponse.result;
} catch (error) {
logger.error(
`Error calling tool ${toolName} on SSE connection ${serverId}:`,
error
);
// Wrap non-McpError instances
if (!(error instanceof McpError)) {
throw new McpError(
ErrorCode.InternalError,
`Failed to call tool ${toolName} on connection ${serverId}: ${error instanceof Error ? error.message : String(error)}`
);
}
throw error;
}
}
// Check if this is a stdio connection
else if (this.activeStdioConnections.has(serverId)) {
const requestId = uuidv4();
const request: McpRequest = {
id: requestId,
method: "callTool",
params: { toolName, arguments: normalizedToolArgs },
};
return new Promise((resolve, reject) => {
// Initialize the map for this connection if it doesn't exist
if (!this.pendingStdioRequests.has(serverId)) {
this.pendingStdioRequests.set(serverId, new Map());
}
// Store the promise resolution functions
this.pendingStdioRequests.get(serverId)!.set(requestId, {
resolve: (value) => {
// Type-safe resolution for tool call results
resolve(value as Record<string, unknown>);
},
reject: reject,
});
// Set up a timeout to automatically reject this request if it takes too long
setTimeout(() => {
// If the request is still pending, reject it
if (
this.pendingStdioRequests.has(serverId) &&
this.pendingStdioRequests.get(serverId)!.has(requestId)
) {
// Get the reject function
const { reject: rejectRequest } = this.pendingStdioRequests
.get(serverId)!
.get(requestId)!;
// Delete the request
this.pendingStdioRequests.get(serverId)!.delete(requestId);
// If this was the last pending request, clean up the connection map
if (this.pendingStdioRequests.get(serverId)!.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
// Reject the request with a timeout error
rejectRequest(
new SdkMcpError(
ErrorCode.InternalError,
"Request timed out waiting for response"
)
);
}
}, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
// Send the request
const sent = this.sendToStdio(serverId, request);
if (!sent) {
// Clean up the pending request if sending fails
this.pendingStdioRequests.get(serverId)!.delete(requestId);
// If this was the last pending request, clean up the connection map
if (this.pendingStdioRequests.get(serverId)!.size === 0) {
this.pendingStdioRequests.delete(serverId);
}
reject(
new Error(`Failed to send request to stdio connection ${serverId}`)
);
}
});
}
// This should never be reached due to the validateConnectionExists check above
throw new McpError(
ErrorCode.InvalidRequest,
`No connection found with ID ${serverId}`
);
}
/**
* Disconnects from an MCP server.
* @param serverId - The ID of the connection to close.
* @returns True if the connection was closed, false if it wasn't found.
* @throws {McpError} Throws an error if the parameters are invalid.
*/
public disconnect(serverId: string): boolean {
// Validate serverId
this.validateServerId(serverId);
// Check if this is an SSE connection
if (this.activeSseConnections.has(serverId)) {
return this.closeSseConnection(serverId);
}
// Check if this is a stdio connection
else if (this.activeStdioConnections.has(serverId)) {
return this.closeStdioConnection(serverId);
}
// Connection not found
throw new McpError(
ErrorCode.InvalidRequest,
`Connection not found for serverId: ${serverId}`
);
}
/**
* Closes all active connections.
*/
public closeAllConnections(): void {
// Close all SSE connections
for (const id of this.activeSseConnections.keys()) {
this.closeSseConnection(id);
}
// Close all stdio connections
for (const id of this.activeStdioConnections.keys()) {
this.closeStdioConnection(id);
}
// Clean up all pending requests
for (const [, requestsMap] of this.pendingStdioRequests.entries()) {
for (const [requestId, { reject }] of requestsMap.entries()) {
logger.warn(
`Rejecting pending request ${requestId} due to service shutdown`
);
reject(new Error("Connection closed due to service shutdown"));
}
}
// Clear the pending requests map
this.pendingStdioRequests.clear();
// Clear the cleanup interval
if (this.cleanupIntervalId) {
clearInterval(this.cleanupIntervalId);
this.cleanupIntervalId = undefined;
}
logger.info("All MCP connections closed.");
}
}
```