This is page 8 of 8. Use http://codebase.md/bsmi021/mcp-gemini-server?lines=true&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
--------------------------------------------------------------------------------
/tests/unit/tools/geminiGenerateContentConsolidatedTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { geminiGenerateContentConsolidatedTool } from "../../../src/tools/geminiGenerateContentConsolidatedTool.js";
3 | import { GeminiApiError } from "../../../src/utils/errors.js";
4 | import { McpError } from "@modelcontextprotocol/sdk/types.js";
5 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
6 | import { GeminiService } from "../../../src/services/index.js";
7 |
8 | describe("geminiGenerateContentConsolidatedTool", () => {
9 | // Mock server and service instances
10 | const mockTool = vi.fn();
11 | const mockServer = {
12 | tool: mockTool,
13 | } as unknown as McpServer;
14 |
15 | // Create mock functions for the service methods
16 | const mockGenerateContent = vi.fn();
17 | const mockGenerateContentStream = vi.fn();
18 |
19 | // Create a minimal mock service with just the necessary methods for testing
20 | const mockService = {
21 | generateContent: mockGenerateContent,
22 | generateContentStream: mockGenerateContentStream,
23 | // Add empty implementations for required GeminiService methods
24 | } as unknown as GeminiService;
25 |
26 | // Reset mocks before each test
27 | beforeEach(() => {
28 | vi.resetAllMocks();
29 | });
30 |
31 | it("should register the tool with the server", () => {
32 | // Call the tool registration function
33 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
34 |
35 | // Verify tool was registered
36 | expect(mockTool).toHaveBeenCalledTimes(1);
37 | const [name, description, params, handler] = mockTool.mock.calls[0];
38 |
39 | // Check tool registration parameters
40 | expect(name).toBe("gemini_generate_content");
41 | expect(description).toContain("Generates text content");
42 | expect(params).toBeDefined();
43 | expect(typeof handler).toBe("function");
44 | });
45 |
46 | it("should handle standard content generation", async () => {
47 | // Register tool to get the request handler
48 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
49 | const [, , , handler] = mockTool.mock.calls[0];
50 |
51 | // Mock successful response
52 | const mockResponse = "This is a test response";
53 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
54 |
55 | // Prepare test request
56 | const testRequest = {
57 | modelName: "gemini-1.5-flash",
58 | prompt: "What is the capital of France?",
59 | stream: false,
60 | };
61 |
62 | // Call the handler
63 | const result = await handler(testRequest);
64 |
65 | // Verify the service method was called with correct parameters
66 | expect(mockGenerateContent).toHaveBeenCalledWith(
67 | expect.objectContaining({
68 | modelName: "gemini-1.5-flash",
69 | prompt: "What is the capital of France?",
70 | })
71 | );
72 |
73 | // Verify the result
74 | expect(result).toEqual({
75 | content: [
76 | {
77 | type: "text",
78 | text: "This is a test response",
79 | },
80 | ],
81 | });
82 | });
83 |
84 | it("should handle streaming content generation", async () => {
85 | // Register tool to get the request handler
86 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
87 | const [, , , handler] = mockTool.mock.calls[0];
88 |
89 | // Create an async generator mock for streaming
90 | async function* mockStreamGenerator() {
91 | yield "This is ";
92 | yield "a streaming ";
93 | yield "response";
94 | }
95 | mockGenerateContentStream.mockReturnValueOnce(mockStreamGenerator());
96 |
97 | // Prepare test request with streaming enabled
98 | const testRequest = {
99 | modelName: "gemini-1.5-flash",
100 | prompt: "Tell me a story",
101 | stream: true,
102 | };
103 |
104 | // Call the handler
105 | const result = await handler(testRequest);
106 |
107 | // Verify the streaming service method was called
108 | expect(mockGenerateContentStream).toHaveBeenCalledWith(
109 | expect.objectContaining({
110 | modelName: "gemini-1.5-flash",
111 | prompt: "Tell me a story",
112 | })
113 | );
114 |
115 | // Verify the result contains the concatenated stream
116 | expect(result).toEqual({
117 | content: [
118 | {
119 | type: "text",
120 | text: "This is a streaming response",
121 | },
122 | ],
123 | });
124 | });
125 |
126 | it("should handle function calling with function declarations", async () => {
127 | // Register tool to get the request handler
128 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
129 | const [, , , handler] = mockTool.mock.calls[0];
130 |
131 | // Mock function call response
132 | const mockFunctionCallResponse = {
133 | functionCall: {
134 | name: "get_weather",
135 | args: { location: "Paris" },
136 | },
137 | };
138 | mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
139 |
140 | // Prepare test request with function declarations
141 | const testRequest = {
142 | modelName: "gemini-1.5-flash",
143 | prompt: "What's the weather in Paris?",
144 | stream: false,
145 | functionDeclarations: [
146 | {
147 | name: "get_weather",
148 | description: "Get the weather for a location",
149 | parameters: {
150 | type: "OBJECT" as const,
151 | properties: {
152 | location: {
153 | type: "STRING" as const,
154 | description: "The location to get weather for",
155 | },
156 | },
157 | required: ["location"],
158 | },
159 | },
160 | ],
161 | };
162 |
163 | // Call the handler
164 | const result = await handler(testRequest);
165 |
166 | // Verify the service method was called with function declarations
167 | expect(mockGenerateContent).toHaveBeenCalledWith(
168 | expect.objectContaining({
169 | modelName: "gemini-1.5-flash",
170 | prompt: "What's the weather in Paris?",
171 | functionDeclarations: expect.arrayContaining([
172 | expect.objectContaining({
173 | name: "get_weather",
174 | }),
175 | ]),
176 | })
177 | );
178 |
179 | // Verify the result contains the serialized function call
180 | expect(result).toEqual({
181 | content: [
182 | {
183 | type: "text",
184 | text: JSON.stringify({
185 | name: "get_weather",
186 | args: { location: "Paris" },
187 | }),
188 | },
189 | ],
190 | });
191 | });
192 |
193 | it("should handle optional parameters", async () => {
194 | // Register tool to get the request handler
195 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
196 | const [, , , handler] = mockTool.mock.calls[0];
197 |
198 | // Mock successful response
199 | const mockResponse = "Generated with parameters";
200 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
201 |
202 | // Prepare test request with all optional parameters
203 | const testRequest = {
204 | modelName: "gemini-1.5-pro",
205 | prompt: "Generate creative content",
206 | stream: false,
207 | generationConfig: {
208 | temperature: 0.8,
209 | topK: 40,
210 | topP: 0.95,
211 | maxOutputTokens: 1024,
212 | stopSequences: ["END"],
213 | thinkingConfig: {
214 | thinkingBudget: 8192,
215 | reasoningEffort: "medium",
216 | },
217 | },
218 | safetySettings: [
219 | {
220 | category: "HARM_CATEGORY_HATE_SPEECH",
221 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
222 | },
223 | ],
224 | systemInstruction: "You are a helpful assistant",
225 | cachedContentName: "cachedContents/12345",
226 | urlContext: {
227 | urls: ["https://example.com"],
228 | fetchOptions: {
229 | maxContentKb: 100,
230 | timeoutMs: 10000,
231 | },
232 | },
233 | modelPreferences: {
234 | preferQuality: true,
235 | preferSpeed: false,
236 | taskType: "creative_writing",
237 | },
238 | };
239 |
240 | // Call the handler
241 | const result = await handler(testRequest);
242 |
243 | // Verify all parameters were passed to the service
244 | expect(mockGenerateContent).toHaveBeenCalledWith(
245 | expect.objectContaining({
246 | modelName: "gemini-1.5-pro",
247 | prompt: "Generate creative content",
248 | generationConfig: expect.objectContaining({
249 | temperature: 0.8,
250 | topK: 40,
251 | topP: 0.95,
252 | maxOutputTokens: 1024,
253 | }),
254 | systemInstruction: "You are a helpful assistant",
255 | cachedContentName: "cachedContents/12345",
256 | urlContext: expect.objectContaining({
257 | urls: ["https://example.com"],
258 | }),
259 | })
260 | );
261 |
262 | // Verify the result
263 | expect(result).toEqual({
264 | content: [
265 | {
266 | type: "text",
267 | text: "Generated with parameters",
268 | },
269 | ],
270 | });
271 | });
272 |
273 | it("should handle errors and map them to MCP errors", async () => {
274 | // Register tool to get the request handler
275 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
276 | const [, , , handler] = mockTool.mock.calls[0];
277 |
278 | // Mock an API error
279 | const apiError = new GeminiApiError("API rate limit exceeded", {
280 | statusCode: 429,
281 | statusText: "Too Many Requests",
282 | });
283 | mockGenerateContent.mockRejectedValueOnce(apiError);
284 |
285 | // Prepare test request
286 | const testRequest = {
287 | modelName: "gemini-1.5-flash",
288 | prompt: "Test error handling",
289 | stream: false,
290 | };
291 |
292 | // Call the handler and expect it to throw
293 | await expect(handler(testRequest)).rejects.toThrow(McpError);
294 | });
295 |
296 | it("should handle URL context metrics calculation", async () => {
297 | // Register tool to get the request handler
298 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
299 | const [, , , handler] = mockTool.mock.calls[0];
300 |
301 | // Mock successful response
302 | mockGenerateContent.mockResolvedValueOnce("Response with URL context");
303 |
304 | // Prepare test request with URL context
305 | const testRequest = {
306 | modelName: "gemini-1.5-flash",
307 | prompt: "Analyze these URLs",
308 | stream: false,
309 | urlContext: {
310 | urls: [
311 | "https://example1.com",
312 | "https://example2.com",
313 | "https://example3.com",
314 | ],
315 | fetchOptions: {
316 | maxContentKb: 200,
317 | },
318 | },
319 | };
320 |
321 | // Call the handler
322 | await handler(testRequest);
323 |
324 | // Verify URL metrics were calculated and passed
325 | expect(mockGenerateContent).toHaveBeenCalledWith(
326 | expect.objectContaining({
327 | urlCount: 3,
328 | estimatedUrlContentSize: 3 * 200 * 1024, // 3 URLs * 200KB * 1024 bytes/KB
329 | })
330 | );
331 | });
332 | it("should handle function call response with text fallback", async () => {
333 | // Register tool to get the request handler
334 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
335 | const [, , , handler] = mockTool.mock.calls[0];
336 |
337 | // Mock function call response with text fallback
338 | const mockFunctionCallResponse = {
339 | functionCall: {
340 | name: "get_weather",
341 | args: { location: "Paris" },
342 | },
343 | text: "I'll get the weather for Paris.",
344 | };
345 | mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
346 |
347 | // Prepare test request with function declarations
348 | const testRequest = {
349 | modelName: "gemini-1.5-flash",
350 | prompt: "What's the weather in Paris?",
351 | stream: false,
352 | functionDeclarations: [
353 | {
354 | name: "get_weather",
355 | description: "Get the weather for a location",
356 | parameters: {
357 | type: "OBJECT" as const,
358 | properties: {
359 | location: {
360 | type: "STRING" as const,
361 | description: "The location to get weather for",
362 | },
363 | },
364 | required: ["location"],
365 | },
366 | },
367 | ],
368 | };
369 |
370 | // Call the handler
371 | const result = await handler(testRequest);
372 |
373 | // Verify the result contains the serialized function call (not the text)
374 | expect(result).toEqual({
375 | content: [
376 | {
377 | type: "text",
378 | text: JSON.stringify({
379 | name: "get_weather",
380 | args: { location: "Paris" },
381 | }),
382 | },
383 | ],
384 | });
385 | });
386 |
387 | it("should handle function call response with only text", async () => {
388 | // Register tool to get the request handler
389 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
390 | const [, , , handler] = mockTool.mock.calls[0];
391 |
392 | // Mock response with only text (no function call)
393 | const mockTextResponse = {
394 | text: "The weather in Paris is sunny and 22°C.",
395 | };
396 | mockGenerateContent.mockResolvedValueOnce(mockTextResponse);
397 |
398 | // Prepare test request with function declarations
399 | const testRequest = {
400 | modelName: "gemini-1.5-flash",
401 | prompt: "What's the weather in Paris?",
402 | stream: false,
403 | functionDeclarations: [
404 | {
405 | name: "get_weather",
406 | description: "Get the weather for a location",
407 | parameters: {
408 | type: "OBJECT" as const,
409 | properties: {
410 | location: {
411 | type: "STRING" as const,
412 | description: "The location to get weather for",
413 | },
414 | },
415 | required: ["location"],
416 | },
417 | },
418 | ],
419 | };
420 |
421 | // Call the handler
422 | const result = await handler(testRequest);
423 |
424 | // Verify the result contains the text response
425 | expect(result).toEqual({
426 | content: [
427 | {
428 | type: "text",
429 | text: "The weather in Paris is sunny and 22°C.",
430 | },
431 | ],
432 | });
433 | });
434 |
435 | it("should handle toolConfig parameter", async () => {
436 | // Register tool to get the request handler
437 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
438 | const [, , , handler] = mockTool.mock.calls[0];
439 |
440 | // Mock successful response
441 | const mockResponse = "Response with tool config";
442 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
443 |
444 | // Prepare test request with toolConfig
445 | const testRequest = {
446 | modelName: "gemini-1.5-flash",
447 | prompt: "Test with tool config",
448 | stream: false,
449 | functionDeclarations: [
450 | {
451 | name: "test_function",
452 | description: "A test function",
453 | parameters: {
454 | type: "OBJECT" as const,
455 | properties: {
456 | param: {
457 | type: "STRING" as const,
458 | description: "A test parameter",
459 | },
460 | },
461 | required: ["param"],
462 | },
463 | },
464 | ],
465 | toolConfig: {
466 | functionCallingConfig: {
467 | mode: "AUTO",
468 | allowedFunctionNames: ["test_function"],
469 | },
470 | },
471 | };
472 |
473 | // Call the handler
474 | await handler(testRequest);
475 |
476 | // Verify toolConfig was passed to the service
477 | expect(mockGenerateContent).toHaveBeenCalledWith(
478 | expect.objectContaining({
479 | toolConfig: {
480 | functionCallingConfig: {
481 | mode: "AUTO",
482 | allowedFunctionNames: ["test_function"],
483 | },
484 | },
485 | })
486 | );
487 | });
488 |
489 | it("should handle thinking configuration parameters", async () => {
490 | // Register tool to get the request handler
491 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
492 | const [, , , handler] = mockTool.mock.calls[0];
493 |
494 | // Mock successful response
495 | const mockResponse = "Response with thinking config";
496 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
497 |
498 | // Prepare test request with thinking configuration
499 | const testRequest = {
500 | modelName: "gemini-1.5-flash",
501 | prompt: "Complex reasoning task",
502 | stream: false,
503 | generationConfig: {
504 | temperature: 0.7,
505 | thinkingConfig: {
506 | thinkingBudget: 16384,
507 | reasoningEffort: "high",
508 | },
509 | },
510 | };
511 |
512 | // Call the handler
513 | await handler(testRequest);
514 |
515 | // Verify thinking config was passed to the service
516 | expect(mockGenerateContent).toHaveBeenCalledWith(
517 | expect.objectContaining({
518 | generationConfig: expect.objectContaining({
519 | thinkingConfig: {
520 | thinkingBudget: 16384,
521 | reasoningEffort: "high",
522 | },
523 | }),
524 | })
525 | );
526 | });
527 |
528 | it("should handle model preferences for task optimization", async () => {
529 | // Register tool to get the request handler
530 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
531 | const [, , , handler] = mockTool.mock.calls[0];
532 |
533 | // Mock successful response
534 | const mockResponse = "Optimized response";
535 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
536 |
537 | // Prepare test request with comprehensive model preferences
538 | const testRequest = {
539 | modelName: "gemini-1.5-flash",
540 | prompt: "Generate creative content",
541 | stream: false,
542 | modelPreferences: {
543 | preferQuality: true,
544 | preferSpeed: false,
545 | preferCost: false,
546 | complexityHint: "high",
547 | taskType: "creative_writing",
548 | },
549 | };
550 |
551 | // Call the handler
552 | await handler(testRequest);
553 |
554 | // Verify model preferences were passed to the service
555 | expect(mockGenerateContent).toHaveBeenCalledWith(
556 | expect.objectContaining({
557 | preferQuality: true,
558 | preferSpeed: false,
559 | preferCost: false,
560 | complexityHint: "high",
561 | taskType: "creative_writing",
562 | })
563 | );
564 | });
565 |
566 | it("should handle comprehensive safety settings", async () => {
567 | // Register tool to get the request handler
568 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
569 | const [, , , handler] = mockTool.mock.calls[0];
570 |
571 | // Mock successful response
572 | const mockResponse = "Safe response";
573 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
574 |
575 | // Prepare test request with comprehensive safety settings
576 | const testRequest = {
577 | modelName: "gemini-1.5-flash",
578 | prompt: "Generate content with safety controls",
579 | stream: false,
580 | safetySettings: [
581 | {
582 | category: "HARM_CATEGORY_HATE_SPEECH",
583 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
584 | },
585 | {
586 | category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
587 | threshold: "BLOCK_LOW_AND_ABOVE",
588 | },
589 | {
590 | category: "HARM_CATEGORY_HARASSMENT",
591 | threshold: "BLOCK_ONLY_HIGH",
592 | },
593 | {
594 | category: "HARM_CATEGORY_DANGEROUS_CONTENT",
595 | threshold: "BLOCK_NONE",
596 | },
597 | ],
598 | };
599 |
600 | // Call the handler
601 | await handler(testRequest);
602 |
603 | // Verify safety settings were properly mapped and passed
604 | expect(mockGenerateContent).toHaveBeenCalledWith(
605 | expect.objectContaining({
606 | safetySettings: [
607 | {
608 | category: "HARM_CATEGORY_HATE_SPEECH",
609 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
610 | },
611 | {
612 | category: "HARM_CATEGORY_SEXUALLY_EXPLICIT",
613 | threshold: "BLOCK_LOW_AND_ABOVE",
614 | },
615 | {
616 | category: "HARM_CATEGORY_HARASSMENT",
617 | threshold: "BLOCK_ONLY_HIGH",
618 | },
619 | {
620 | category: "HARM_CATEGORY_DANGEROUS_CONTENT",
621 | threshold: "BLOCK_NONE",
622 | },
623 | ],
624 | })
625 | );
626 | });
627 |
628 | it("should handle URL context with comprehensive fetch options", async () => {
629 | // Register tool to get the request handler
630 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
631 | const [, , , handler] = mockTool.mock.calls[0];
632 |
633 | // Mock successful response
634 | mockGenerateContent.mockResolvedValueOnce("Response with URL context");
635 |
636 | // Prepare test request with comprehensive URL context options
637 | const testRequest = {
638 | modelName: "gemini-1.5-flash",
639 | prompt: "Analyze these web pages",
640 | stream: false,
641 | urlContext: {
642 | urls: ["https://example1.com/article", "https://example2.com/blog"],
643 | fetchOptions: {
644 | maxContentKb: 150,
645 | timeoutMs: 15000,
646 | includeMetadata: true,
647 | convertToMarkdown: true,
648 | allowedDomains: ["example1.com", "example2.com"],
649 | userAgent: "Custom-Agent/1.0",
650 | },
651 | },
652 | };
653 |
654 | // Call the handler
655 | await handler(testRequest);
656 |
657 | // Verify URL context was passed with all options
658 | expect(mockGenerateContent).toHaveBeenCalledWith(
659 | expect.objectContaining({
660 | urlContext: {
661 | urls: ["https://example1.com/article", "https://example2.com/blog"],
662 | fetchOptions: {
663 | maxContentKb: 150,
664 | timeoutMs: 15000,
665 | includeMetadata: true,
666 | convertToMarkdown: true,
667 | allowedDomains: ["example1.com", "example2.com"],
668 | userAgent: "Custom-Agent/1.0",
669 | },
670 | },
671 | urlCount: 2,
672 | estimatedUrlContentSize: 2 * 150 * 1024, // 2 URLs * 150KB * 1024 bytes/KB
673 | })
674 | );
675 | });
676 |
677 | it("should handle URL context with default fetch options", async () => {
678 | // Register tool to get the request handler
679 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
680 | const [, , , handler] = mockTool.mock.calls[0];
681 |
682 | // Mock successful response
683 | mockGenerateContent.mockResolvedValueOnce(
684 | "Response with default URL context"
685 | );
686 |
687 | // Prepare test request with minimal URL context (using defaults)
688 | const testRequest = {
689 | modelName: "gemini-1.5-flash",
690 | prompt: "Analyze this web page",
691 | stream: false,
692 | urlContext: {
693 | urls: ["https://example.com"],
694 | },
695 | };
696 |
697 | // Call the handler
698 | await handler(testRequest);
699 |
700 | // Verify URL context was passed with default maxContentKb
701 | expect(mockGenerateContent).toHaveBeenCalledWith(
702 | expect.objectContaining({
703 | urlContext: {
704 | urls: ["https://example.com"],
705 | },
706 | urlCount: 1,
707 | estimatedUrlContentSize: 1 * 100 * 1024, // 1 URL * 100KB default * 1024 bytes/KB
708 | })
709 | );
710 | });
711 |
712 | it("should handle unexpected response structure from service", async () => {
713 | // Register tool to get the request handler
714 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
715 | const [, , , handler] = mockTool.mock.calls[0];
716 |
717 | // Mock unexpected response structure
718 | const unexpectedResponse = { unexpected: "structure" };
719 | mockGenerateContent.mockResolvedValueOnce(unexpectedResponse);
720 |
721 | // Prepare test request
722 | const testRequest = {
723 | modelName: "gemini-1.5-flash",
724 | prompt: "Test unexpected response",
725 | stream: false,
726 | };
727 |
728 | // Call the handler and expect it to throw
729 | await expect(handler(testRequest)).rejects.toThrow(
730 | "Invalid response structure received from Gemini service."
731 | );
732 | });
733 |
734 | it("should handle streaming with empty chunks gracefully", async () => {
735 | // Register tool to get the request handler
736 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
737 | const [, , , handler] = mockTool.mock.calls[0];
738 |
739 | // Create an async generator mock with empty chunks
740 | async function* mockStreamGenerator() {
741 | yield "Start ";
742 | yield ""; // Empty chunk
743 | yield "middle ";
744 | yield ""; // Another empty chunk
745 | yield "end";
746 | }
747 | mockGenerateContentStream.mockReturnValueOnce(mockStreamGenerator());
748 |
749 | // Prepare test request with streaming enabled
750 | const testRequest = {
751 | modelName: "gemini-1.5-flash",
752 | prompt: "Stream with empty chunks",
753 | stream: true,
754 | };
755 |
756 | // Call the handler
757 | const result = await handler(testRequest);
758 |
759 | // Verify the result contains the concatenated stream (empty chunks should be included)
760 | expect(result).toEqual({
761 | content: [
762 | {
763 | type: "text",
764 | text: "Start middle end",
765 | },
766 | ],
767 | });
768 | });
769 |
770 | it("should handle complex generation config with all parameters", async () => {
771 | // Register tool to get the request handler
772 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
773 | const [, , , handler] = mockTool.mock.calls[0];
774 |
775 | // Mock successful response
776 | const mockResponse = "Complex config response";
777 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
778 |
779 | // Prepare test request with all generation config parameters
780 | const testRequest = {
781 | modelName: "gemini-1.5-pro",
782 | prompt: "Complex generation task",
783 | stream: false,
784 | generationConfig: {
785 | temperature: 0.9,
786 | topP: 0.8,
787 | topK: 50,
788 | maxOutputTokens: 2048,
789 | stopSequences: ["STOP", "END", "FINISH"],
790 | thinkingConfig: {
791 | thinkingBudget: 12288,
792 | reasoningEffort: "medium",
793 | },
794 | },
795 | };
796 |
797 | // Call the handler
798 | await handler(testRequest);
799 |
800 | // Verify all generation config parameters were passed
801 | expect(mockGenerateContent).toHaveBeenCalledWith(
802 | expect.objectContaining({
803 | generationConfig: {
804 | temperature: 0.9,
805 | topP: 0.8,
806 | topK: 50,
807 | maxOutputTokens: 2048,
808 | stopSequences: ["STOP", "END", "FINISH"],
809 | thinkingConfig: {
810 | thinkingBudget: 12288,
811 | reasoningEffort: "medium",
812 | },
813 | },
814 | })
815 | );
816 | });
817 |
818 | it("should handle cached content parameter", async () => {
819 | // Register tool to get the request handler
820 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
821 | const [, , , handler] = mockTool.mock.calls[0];
822 |
823 | // Mock successful response
824 | const mockResponse = "Response using cached content";
825 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
826 |
827 | // Prepare test request with cached content
828 | const testRequest = {
829 | modelName: "gemini-1.5-flash",
830 | prompt: "Use cached context for this request",
831 | stream: false,
832 | cachedContentName: "cachedContents/abc123def456",
833 | };
834 |
835 | // Call the handler
836 | await handler(testRequest);
837 |
838 | // Verify cached content name was passed
839 | expect(mockGenerateContent).toHaveBeenCalledWith(
840 | expect.objectContaining({
841 | cachedContentName: "cachedContents/abc123def456",
842 | })
843 | );
844 | });
845 |
846 | it("should handle system instruction parameter", async () => {
847 | // Register tool to get the request handler
848 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
849 | const [, , , handler] = mockTool.mock.calls[0];
850 |
851 | // Mock successful response
852 | const mockResponse = "Response following system instruction";
853 | mockGenerateContent.mockResolvedValueOnce(mockResponse);
854 |
855 | // Prepare test request with system instruction
856 | const testRequest = {
857 | modelName: "gemini-1.5-flash",
858 | prompt: "Generate a response",
859 | stream: false,
860 | systemInstruction:
861 | "You are a helpful assistant that always responds in a friendly tone.",
862 | };
863 |
864 | // Call the handler
865 | await handler(testRequest);
866 |
867 | // Verify system instruction was passed
868 | expect(mockGenerateContent).toHaveBeenCalledWith(
869 | expect.objectContaining({
870 | systemInstruction:
871 | "You are a helpful assistant that always responds in a friendly tone.",
872 | })
873 | );
874 | });
875 |
876 | it("should handle streaming errors gracefully", async () => {
877 | // Register tool to get the request handler
878 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
879 | const [, , , handler] = mockTool.mock.calls[0];
880 |
881 | // Create an async generator that throws an error
882 | async function* mockStreamGeneratorWithError() {
883 | yield "Start ";
884 | throw new Error("Streaming error occurred");
885 | }
886 | mockGenerateContentStream.mockReturnValueOnce(
887 | mockStreamGeneratorWithError()
888 | );
889 |
890 | // Prepare test request with streaming enabled
891 | const testRequest = {
892 | modelName: "gemini-1.5-flash",
893 | prompt: "Stream that will error",
894 | stream: true,
895 | };
896 |
897 | // Call the handler and expect it to throw
898 | await expect(handler(testRequest)).rejects.toThrow();
899 | });
900 |
901 | it("should handle function declarations with complex parameter schemas", async () => {
902 | // Register tool to get the request handler
903 | geminiGenerateContentConsolidatedTool(mockServer, mockService);
904 | const [, , , handler] = mockTool.mock.calls[0];
905 |
906 | // Mock function call response
907 | const mockFunctionCallResponse = {
908 | functionCall: {
909 | name: "complex_function",
910 | args: {
911 | user: { name: "John", age: 30 },
912 | preferences: ["option1", "option2"],
913 | settings: { theme: "dark", notifications: true },
914 | },
915 | },
916 | };
917 | mockGenerateContent.mockResolvedValueOnce(mockFunctionCallResponse);
918 |
919 | // Prepare test request with complex function declarations
920 | const testRequest = {
921 | modelName: "gemini-1.5-flash",
922 | prompt: "Call a complex function",
923 | stream: false,
924 | functionDeclarations: [
925 | {
926 | name: "complex_function",
927 | description: "A function with complex nested parameters",
928 | parameters: {
929 | type: "OBJECT" as const,
930 | properties: {
931 | user: {
932 | type: "OBJECT" as const,
933 | properties: {
934 | name: {
935 | type: "STRING" as const,
936 | description: "User's name",
937 | },
938 | age: {
939 | type: "INTEGER" as const,
940 | description: "User's age",
941 | },
942 | },
943 | required: ["name", "age"],
944 | },
945 | preferences: {
946 | type: "ARRAY" as const,
947 | items: {
948 | type: "STRING" as const,
949 | description: "User preference",
950 | },
951 | description: "List of user preferences",
952 | },
953 | settings: {
954 | type: "OBJECT" as const,
955 | properties: {
956 | theme: {
957 | type: "STRING" as const,
958 | enum: ["light", "dark"],
959 | description: "UI theme preference",
960 | },
961 | notifications: {
962 | type: "BOOLEAN" as const,
963 | description: "Enable notifications",
964 | },
965 | },
966 | },
967 | },
968 | required: ["user"],
969 | },
970 | },
971 | ],
972 | };
973 |
974 | // Call the handler
975 | const result = await handler(testRequest);
976 |
977 | // Verify the service was called with complex function declarations
978 | expect(mockGenerateContent).toHaveBeenCalledWith(
979 | expect.objectContaining({
980 | functionDeclarations: expect.arrayContaining([
981 | expect.objectContaining({
982 | name: "complex_function",
983 | parameters: expect.objectContaining({
984 | type: "OBJECT",
985 | properties: expect.objectContaining({
986 | user: expect.objectContaining({
987 | type: "OBJECT",
988 | properties: expect.any(Object),
989 | }),
990 | preferences: expect.objectContaining({
991 | type: "ARRAY",
992 | }),
993 | settings: expect.objectContaining({
994 | type: "OBJECT",
995 | }),
996 | }),
997 | }),
998 | }),
999 | ]),
1000 | })
1001 | );
1002 |
1003 | // Verify the result contains the serialized function call
1004 | expect(result).toEqual({
1005 | content: [
1006 | {
1007 | type: "text",
1008 | text: JSON.stringify({
1009 | name: "complex_function",
1010 | args: {
1011 | user: { name: "John", age: 30 },
1012 | preferences: ["option1", "option2"],
1013 | settings: { theme: "dark", notifications: true },
1014 | },
1015 | }),
1016 | },
1017 | ],
1018 | });
1019 | });
1020 | });
1021 |
```
--------------------------------------------------------------------------------
/tests/unit/tools/geminiChatTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { geminiChatTool } from "../../../src/tools/geminiChatTool.js";
3 | import { GeminiApiError } from "../../../src/utils/errors.js";
4 | import { McpError } from "@modelcontextprotocol/sdk/types.js";
5 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
6 | import { GeminiService } from "../../../src/services/index.js";
7 | import { BlockedReason, FinishReason } from "@google/genai";
8 |
9 | describe("geminiChatTool", () => {
10 | // Mock server and service instances
11 | const mockTool = vi.fn();
12 | const mockServer = {
13 | tool: mockTool,
14 | } as unknown as McpServer;
15 |
16 | // Create mock functions for the service methods
17 | const mockStartChatSession = vi.fn();
18 | const mockSendMessageToSession = vi.fn();
19 | const mockSendFunctionResultToSession = vi.fn();
20 |
21 | // Create a minimal mock service with just the necessary methods for testing
22 | const mockService = {
23 | startChatSession: mockStartChatSession,
24 | sendMessageToSession: mockSendMessageToSession,
25 | sendFunctionResultToSession: mockSendFunctionResultToSession,
26 | // Add empty implementations for required GeminiService methods
27 | generateContent: () => Promise.resolve("mock"),
28 | } as unknown as GeminiService;
29 |
30 | // Reset mocks before each test
31 | beforeEach(() => {
32 | vi.resetAllMocks();
33 | });
34 |
35 | it("should register the tool with the server", () => {
36 | // Call the tool registration function
37 | geminiChatTool(mockServer, mockService);
38 |
39 | // Verify tool was registered
40 | expect(mockTool).toHaveBeenCalledTimes(1);
41 | const [name, description, params, handler] = mockTool.mock.calls[0];
42 |
43 | // Check tool registration parameters
44 | expect(name).toBe("gemini_chat");
45 | expect(description).toContain("Manages stateful chat sessions");
46 | expect(params).toBeDefined();
47 | expect(typeof handler).toBe("function");
48 | });
49 |
50 | describe("start operation", () => {
51 | it("should start a new chat session", async () => {
52 | // Register tool to get the request handler
53 | geminiChatTool(mockServer, mockService);
54 | const [, , , handler] = mockTool.mock.calls[0];
55 |
56 | // Mock successful response
57 | const mockSessionId = "test-session-123";
58 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
59 |
60 | // Prepare test request
61 | const testRequest = {
62 | operation: "start",
63 | modelName: "gemini-1.5-flash",
64 | history: [
65 | {
66 | role: "user",
67 | parts: [{ text: "Hello" }],
68 | },
69 | {
70 | role: "model",
71 | parts: [{ text: "Hi there!" }],
72 | },
73 | ],
74 | };
75 |
76 | // Call the handler
77 | const result = await handler(testRequest);
78 |
79 | // Verify the service method was called with correct parameters
80 | expect(mockStartChatSession).toHaveBeenCalledWith(
81 | expect.objectContaining({
82 | modelName: "gemini-1.5-flash",
83 | history: testRequest.history,
84 | })
85 | );
86 |
87 | // Verify the result
88 | expect(result).toEqual({
89 | content: [
90 | {
91 | type: "text",
92 | text: JSON.stringify({ sessionId: mockSessionId }),
93 | },
94 | ],
95 | });
96 | });
97 |
98 | it("should start chat session with optional parameters", async () => {
99 | // Register tool to get the request handler
100 | geminiChatTool(mockServer, mockService);
101 | const [, , , handler] = mockTool.mock.calls[0];
102 |
103 | // Mock successful response
104 | const mockSessionId = "test-session-456";
105 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
106 |
107 | // Prepare test request with optional parameters
108 | const testRequest = {
109 | operation: "start",
110 | modelName: "gemini-2.0-flash",
111 | systemInstruction: {
112 | parts: [{ text: "You are a helpful assistant" }],
113 | },
114 | generationConfig: {
115 | temperature: 0.7,
116 | maxOutputTokens: 1000,
117 | },
118 | safetySettings: [
119 | {
120 | category: "HARM_CATEGORY_HATE_SPEECH",
121 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
122 | },
123 | ],
124 | cachedContentName: "cachedContents/abc123",
125 | };
126 |
127 | // Call the handler
128 | const result = await handler(testRequest);
129 |
130 | // Verify all parameters were passed
131 | expect(mockStartChatSession).toHaveBeenCalledWith(
132 | expect.objectContaining({
133 | modelName: "gemini-2.0-flash",
134 | systemInstruction: testRequest.systemInstruction,
135 | generationConfig: testRequest.generationConfig,
136 | safetySettings: testRequest.safetySettings,
137 | cachedContentName: "cachedContents/abc123",
138 | })
139 | );
140 |
141 | // Verify the result
142 | expect(result).toEqual({
143 | content: [
144 | {
145 | type: "text",
146 | text: JSON.stringify({ sessionId: mockSessionId }),
147 | },
148 | ],
149 | });
150 | });
151 | });
152 |
153 | describe("send_message operation", () => {
154 | it("should send a message to an existing session", async () => {
155 | // Register tool to get the request handler
156 | geminiChatTool(mockServer, mockService);
157 | const [, , , handler] = mockTool.mock.calls[0];
158 |
159 | // Mock successful response
160 | const mockResponse = {
161 | candidates: [
162 | {
163 | content: {
164 | parts: [{ text: "The capital of France is Paris." }],
165 | },
166 | finishReason: FinishReason.STOP,
167 | },
168 | ],
169 | };
170 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
171 |
172 | // Prepare test request
173 | const testRequest = {
174 | operation: "send_message",
175 | sessionId: "test-session-123",
176 | message: "What is the capital of France?",
177 | };
178 |
179 | // Call the handler
180 | const result = await handler(testRequest);
181 |
182 | // Verify the service method was called
183 | expect(mockSendMessageToSession).toHaveBeenCalledWith({
184 | sessionId: "test-session-123",
185 | message: "What is the capital of France?",
186 | generationConfig: undefined,
187 | safetySettings: undefined,
188 | tools: undefined,
189 | toolConfig: undefined,
190 | cachedContentName: undefined,
191 | });
192 |
193 | // Verify the result
194 | expect(result).toEqual({
195 | content: [
196 | {
197 | type: "text",
198 | text: "The capital of France is Paris.",
199 | },
200 | ],
201 | });
202 | });
203 |
204 | it("should handle function call responses", async () => {
205 | // Register tool to get the request handler
206 | geminiChatTool(mockServer, mockService);
207 | const [, , , handler] = mockTool.mock.calls[0];
208 |
209 | // Mock function call response
210 | const mockResponse = {
211 | candidates: [
212 | {
213 | content: {
214 | parts: [
215 | {
216 | functionCall: {
217 | name: "get_weather",
218 | args: { location: "Paris" },
219 | },
220 | },
221 | ],
222 | },
223 | finishReason: FinishReason.STOP,
224 | },
225 | ],
226 | };
227 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
228 |
229 | // Prepare test request
230 | const testRequest = {
231 | operation: "send_message",
232 | sessionId: "test-session-123",
233 | message: "What's the weather in Paris?",
234 | tools: [
235 | {
236 | functionDeclarations: [
237 | {
238 | name: "get_weather",
239 | description: "Get weather information",
240 | parameters: {
241 | type: "OBJECT",
242 | properties: {
243 | location: {
244 | type: "STRING",
245 | description: "The location",
246 | },
247 | },
248 | },
249 | },
250 | ],
251 | },
252 | ],
253 | };
254 |
255 | // Call the handler
256 | const result = await handler(testRequest);
257 |
258 | // Verify the result contains function call
259 | expect(result).toEqual({
260 | content: [
261 | {
262 | type: "text",
263 | text: JSON.stringify({
264 | functionCall: {
265 | name: "get_weather",
266 | args: { location: "Paris" },
267 | },
268 | }),
269 | },
270 | ],
271 | });
272 | });
273 |
274 | it("should handle safety blocked responses", async () => {
275 | // Register tool to get the request handler
276 | geminiChatTool(mockServer, mockService);
277 | const [, , , handler] = mockTool.mock.calls[0];
278 |
279 | // Mock safety blocked response
280 | const mockResponse = {
281 | promptFeedback: {
282 | blockReason: BlockedReason.SAFETY,
283 | },
284 | };
285 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
286 |
287 | // Prepare test request
288 | const testRequest = {
289 | operation: "send_message",
290 | sessionId: "test-session-123",
291 | message: "Test message",
292 | };
293 |
294 | // Call the handler
295 | const result = await handler(testRequest);
296 |
297 | // Verify error response
298 | expect(result).toEqual({
299 | content: [
300 | {
301 | type: "text",
302 | text: "Error: Prompt blocked due to safety settings . Reason: SAFETY",
303 | },
304 | ],
305 | isError: true,
306 | });
307 | });
308 |
309 | it("should throw error if sessionId is missing", async () => {
310 | // Register tool to get the request handler
311 | geminiChatTool(mockServer, mockService);
312 | const [, , , handler] = mockTool.mock.calls[0];
313 |
314 | // Prepare test request without sessionId
315 | const testRequest = {
316 | operation: "send_message",
317 | message: "Test message",
318 | };
319 |
320 | // Call the handler and expect error
321 | await expect(handler(testRequest)).rejects.toThrow(
322 | "sessionId is required for operation 'send_message'"
323 | );
324 | });
325 |
326 | it("should throw error if message is missing", async () => {
327 | // Register tool to get the request handler
328 | geminiChatTool(mockServer, mockService);
329 | const [, , , handler] = mockTool.mock.calls[0];
330 |
331 | // Prepare test request without message
332 | const testRequest = {
333 | operation: "send_message",
334 | sessionId: "test-session-123",
335 | };
336 |
337 | // Call the handler and expect error
338 | await expect(handler(testRequest)).rejects.toThrow(
339 | "message is required for operation 'send_message'"
340 | );
341 | });
342 | });
343 |
344 | describe("send_function_result operation", () => {
345 | it("should send function results to session", async () => {
346 | // Register tool to get the request handler
347 | geminiChatTool(mockServer, mockService);
348 | const [, , , handler] = mockTool.mock.calls[0];
349 |
350 | // Mock successful response
351 | const mockResponse = {
352 | candidates: [
353 | {
354 | content: {
355 | parts: [{ text: "The weather in Paris is sunny and 22°C." }],
356 | },
357 | finishReason: FinishReason.STOP,
358 | },
359 | ],
360 | };
361 | mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
362 |
363 | // Prepare test request
364 | const testRequest = {
365 | operation: "send_function_result",
366 | sessionId: "test-session-123",
367 | functionResponses: [
368 | {
369 | name: "get_weather",
370 | response: {
371 | temperature: 22,
372 | condition: "sunny",
373 | location: "Paris",
374 | },
375 | },
376 | ],
377 | };
378 |
379 | // Call the handler
380 | const result = await handler(testRequest);
381 |
382 | // Verify the service method was called
383 | expect(mockSendFunctionResultToSession).toHaveBeenCalledWith({
384 | sessionId: "test-session-123",
385 | functionResponse: JSON.stringify(testRequest.functionResponses),
386 | functionCall: undefined,
387 | });
388 |
389 | // Verify the result
390 | expect(result).toEqual({
391 | content: [
392 | {
393 | type: "text",
394 | text: "The weather in Paris is sunny and 22°C.",
395 | },
396 | ],
397 | });
398 | });
399 |
400 | it("should handle empty candidates", async () => {
401 | // Register tool to get the request handler
402 | geminiChatTool(mockServer, mockService);
403 | const [, , , handler] = mockTool.mock.calls[0];
404 |
405 | // Mock response with no candidates
406 | const mockResponse = {
407 | candidates: [],
408 | };
409 | mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
410 |
411 | // Prepare test request
412 | const testRequest = {
413 | operation: "send_function_result",
414 | sessionId: "test-session-123",
415 | functionResponses: [
416 | {
417 | name: "test_function",
418 | response: { result: "test" },
419 | },
420 | ],
421 | };
422 |
423 | // Call the handler
424 | const result = await handler(testRequest);
425 |
426 | // Verify error response
427 | expect(result).toEqual({
428 | content: [
429 | {
430 | type: "text",
431 | text: "Error: No response candidates returned by the model after function result.",
432 | },
433 | ],
434 | isError: true,
435 | });
436 | });
437 |
438 | it("should throw error if sessionId is missing", async () => {
439 | // Register tool to get the request handler
440 | geminiChatTool(mockServer, mockService);
441 | const [, , , handler] = mockTool.mock.calls[0];
442 |
443 | // Prepare test request without sessionId
444 | const testRequest = {
445 | operation: "send_function_result",
446 | functionResponses: [
447 | {
448 | name: "test_function",
449 | response: { result: "test" },
450 | },
451 | ],
452 | };
453 |
454 | // Call the handler and expect error
455 | await expect(handler(testRequest)).rejects.toThrow(
456 | "sessionId is required for operation 'send_function_result'"
457 | );
458 | });
459 |
460 | it("should throw error if functionResponses is missing", async () => {
461 | // Register tool to get the request handler
462 | geminiChatTool(mockServer, mockService);
463 | const [, , , handler] = mockTool.mock.calls[0];
464 |
465 | // Prepare test request without functionResponses
466 | const testRequest = {
467 | operation: "send_function_result",
468 | sessionId: "test-session-123",
469 | };
470 |
471 | // Call the handler and expect error
472 | await expect(handler(testRequest)).rejects.toThrow(
473 | "functionResponses is required for operation 'send_function_result'"
474 | );
475 | });
476 | });
477 |
478 | describe("error handling", () => {
479 | it("should map GeminiApiError to McpError", async () => {
480 | // Register tool to get the request handler
481 | geminiChatTool(mockServer, mockService);
482 | const [, , , handler] = mockTool.mock.calls[0];
483 |
484 | // Mock service to throw GeminiApiError
485 | const geminiError = new GeminiApiError("API error occurred");
486 | mockStartChatSession.mockImplementationOnce(() => {
487 | throw geminiError;
488 | });
489 |
490 | // Prepare test request
491 | const testRequest = {
492 | operation: "start",
493 | modelName: "gemini-1.5-flash",
494 | };
495 |
496 | // Call the handler and expect McpError
497 | await expect(handler(testRequest)).rejects.toThrow();
498 |
499 | // Verify the error was caught and mapped
500 | try {
501 | await handler(testRequest);
502 | } catch (error) {
503 | expect(error).toBeInstanceOf(McpError);
504 | expect((error as McpError).message).toContain("API error occurred");
505 | }
506 | });
507 |
508 | it("should handle invalid operation", async () => {
509 | // Register tool to get the request handler
510 | geminiChatTool(mockServer, mockService);
511 | const [, , , handler] = mockTool.mock.calls[0];
512 |
513 | // Prepare test request with invalid operation
514 | const testRequest = {
515 | operation: "invalid_operation",
516 | };
517 |
518 | // Call the handler and expect error
519 | await expect(handler(testRequest)).rejects.toThrow(
520 | "Invalid operation: invalid_operation"
521 | );
522 | });
523 | });
524 |
525 | describe("response processing edge cases", () => {
526 | it("should handle candidate with SAFETY finish reason", async () => {
527 | // Register tool to get the request handler
528 | geminiChatTool(mockServer, mockService);
529 | const [, , , handler] = mockTool.mock.calls[0];
530 |
531 | // Mock response with SAFETY finish reason
532 | const mockResponse = {
533 | candidates: [
534 | {
535 | content: {
536 | parts: [{ text: "Partial response..." }],
537 | },
538 | finishReason: FinishReason.SAFETY,
539 | },
540 | ],
541 | };
542 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
543 |
544 | // Prepare test request
545 | const testRequest = {
546 | operation: "send_message",
547 | sessionId: "test-session-123",
548 | message: "Test message",
549 | };
550 |
551 | // Call the handler
552 | const result = await handler(testRequest);
553 |
554 | // Verify error response for SAFETY finish reason
555 | expect(result).toEqual({
556 | content: [
557 | {
558 | type: "text",
559 | text: "Error: Response generation stopped due to safety settings . FinishReason: SAFETY",
560 | },
561 | ],
562 | isError: true,
563 | });
564 | });
565 |
566 | it("should handle candidate with MAX_TOKENS finish reason", async () => {
567 | // Register tool to get the request handler
568 | geminiChatTool(mockServer, mockService);
569 | const [, , , handler] = mockTool.mock.calls[0];
570 |
571 | // Mock response with MAX_TOKENS finish reason
572 | const mockResponse = {
573 | candidates: [
574 | {
575 | content: {
576 | parts: [{ text: "Response cut off due to token limit..." }],
577 | },
578 | finishReason: FinishReason.MAX_TOKENS,
579 | },
580 | ],
581 | };
582 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
583 |
584 | // Prepare test request
585 | const testRequest = {
586 | operation: "send_message",
587 | sessionId: "test-session-123",
588 | message: "Test message",
589 | };
590 |
591 | // Call the handler
592 | const result = await handler(testRequest);
593 |
594 | // Verify successful response even with MAX_TOKENS (this is acceptable)
595 | expect(result).toEqual({
596 | content: [
597 | {
598 | type: "text",
599 | text: "Response cut off due to token limit...",
600 | },
601 | ],
602 | });
603 | });
604 |
605 | it("should handle candidate with OTHER finish reason", async () => {
606 | // Register tool to get the request handler
607 | geminiChatTool(mockServer, mockService);
608 | const [, , , handler] = mockTool.mock.calls[0];
609 |
610 | // Mock response with OTHER finish reason
611 | const mockResponse = {
612 | candidates: [
613 | {
614 | content: {
615 | parts: [{ text: "Some response..." }],
616 | },
617 | finishReason: FinishReason.OTHER,
618 | },
619 | ],
620 | };
621 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
622 |
623 | // Prepare test request
624 | const testRequest = {
625 | operation: "send_message",
626 | sessionId: "test-session-123",
627 | message: "Test message",
628 | };
629 |
630 | // Call the handler
631 | const result = await handler(testRequest);
632 |
633 | // Verify response is still returned but with warning logged
634 | expect(result).toEqual({
635 | content: [
636 | {
637 | type: "text",
638 | text: "Some response...",
639 | },
640 | ],
641 | });
642 | });
643 |
644 | it("should handle empty content parts", async () => {
645 | // Register tool to get the request handler
646 | geminiChatTool(mockServer, mockService);
647 | const [, , , handler] = mockTool.mock.calls[0];
648 |
649 | // Mock response with empty content parts
650 | const mockResponse = {
651 | candidates: [
652 | {
653 | content: {
654 | parts: [],
655 | },
656 | finishReason: FinishReason.STOP,
657 | },
658 | ],
659 | };
660 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
661 |
662 | // Prepare test request
663 | const testRequest = {
664 | operation: "send_message",
665 | sessionId: "test-session-123",
666 | message: "Test message",
667 | };
668 |
669 | // Call the handler
670 | const result = await handler(testRequest);
671 |
672 | // Verify error response for empty content
673 | expect(result).toEqual({
674 | content: [
675 | {
676 | type: "text",
677 | text: "Error: Empty response from the model .",
678 | },
679 | ],
680 | isError: true,
681 | });
682 | });
683 |
684 | it("should handle missing content in candidate", async () => {
685 | // Register tool to get the request handler
686 | geminiChatTool(mockServer, mockService);
687 | const [, , , handler] = mockTool.mock.calls[0];
688 |
689 | // Mock response with missing content
690 | const mockResponse = {
691 | candidates: [
692 | {
693 | finishReason: FinishReason.STOP,
694 | },
695 | ],
696 | };
697 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
698 |
699 | // Prepare test request
700 | const testRequest = {
701 | operation: "send_message",
702 | sessionId: "test-session-123",
703 | message: "Test message",
704 | };
705 |
706 | // Call the handler
707 | const result = await handler(testRequest);
708 |
709 | // Verify error response for missing content
710 | expect(result).toEqual({
711 | content: [
712 | {
713 | type: "text",
714 | text: "Error: Empty response from the model .",
715 | },
716 | ],
717 | isError: true,
718 | });
719 | });
720 |
721 | it("should handle mixed content parts (text and function call)", async () => {
722 | // Register tool to get the request handler
723 | geminiChatTool(mockServer, mockService);
724 | const [, , , handler] = mockTool.mock.calls[0];
725 |
726 | // Mock response with both text and function call
727 | const mockResponse = {
728 | candidates: [
729 | {
730 | content: {
731 | parts: [
732 | { text: "I'll help you with that. " },
733 | {
734 | functionCall: {
735 | name: "get_weather",
736 | args: { location: "Paris" },
737 | },
738 | },
739 | { text: " Let me check the weather for you." },
740 | ],
741 | },
742 | finishReason: FinishReason.STOP,
743 | },
744 | ],
745 | };
746 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
747 |
748 | // Prepare test request
749 | const testRequest = {
750 | operation: "send_message",
751 | sessionId: "test-session-123",
752 | message: "What's the weather in Paris?",
753 | };
754 |
755 | // Call the handler
756 | const result = await handler(testRequest);
757 |
758 | // Verify function call is returned (function call takes precedence)
759 | expect(result).toEqual({
760 | content: [
761 | {
762 | type: "text",
763 | text: JSON.stringify({
764 | functionCall: {
765 | name: "get_weather",
766 | args: { location: "Paris" },
767 | },
768 | }),
769 | },
770 | ],
771 | });
772 | });
773 |
774 | it("should handle unexpected response structure", async () => {
775 | // Register tool to get the request handler
776 | geminiChatTool(mockServer, mockService);
777 | const [, , , handler] = mockTool.mock.calls[0];
778 |
779 | // Mock response with unexpected structure (no text or function call)
780 | const mockResponse = {
781 | candidates: [
782 | {
783 | content: {
784 | parts: [{ someOtherField: "unexpected data" }],
785 | },
786 | finishReason: FinishReason.STOP,
787 | },
788 | ],
789 | };
790 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
791 |
792 | // Prepare test request
793 | const testRequest = {
794 | operation: "send_message",
795 | sessionId: "test-session-123",
796 | message: "Test message",
797 | };
798 |
799 | // Call the handler
800 | const result = await handler(testRequest);
801 |
802 | // Verify error response for unexpected structure
803 | expect(result).toEqual({
804 | content: [
805 | {
806 | type: "text",
807 | text: "Error: Unexpected response structure from the model .",
808 | },
809 | ],
810 | isError: true,
811 | });
812 | });
813 | });
814 |
815 | describe("advanced parameter combinations", () => {
816 | it("should handle start operation with tools and toolConfig", async () => {
817 | // Register tool to get the request handler
818 | geminiChatTool(mockServer, mockService);
819 | const [, , , handler] = mockTool.mock.calls[0];
820 |
821 | // Mock successful response
822 | const mockSessionId = "test-session-tools";
823 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
824 |
825 | // Prepare test request with tools
826 | const testRequest = {
827 | operation: "start",
828 | modelName: "gemini-1.5-pro",
829 | tools: [
830 | {
831 | functionDeclarations: [
832 | {
833 | name: "calculate",
834 | description: "Perform calculations",
835 | parameters: {
836 | type: "OBJECT",
837 | properties: {
838 | expression: {
839 | type: "STRING",
840 | description: "Mathematical expression",
841 | },
842 | },
843 | required: ["expression"],
844 | },
845 | },
846 | ],
847 | },
848 | ],
849 | };
850 |
851 | // Call the handler
852 | const result = await handler(testRequest);
853 |
854 | // Verify tools were passed
855 | expect(mockStartChatSession).toHaveBeenCalledWith(
856 | expect.objectContaining({
857 | tools: testRequest.tools,
858 | })
859 | );
860 |
861 | expect(result).toEqual({
862 | content: [
863 | {
864 | type: "text",
865 | text: JSON.stringify({ sessionId: mockSessionId }),
866 | },
867 | ],
868 | });
869 | });
870 |
871 | it("should handle send_message with all optional parameters", async () => {
872 | // Register tool to get the request handler
873 | geminiChatTool(mockServer, mockService);
874 | const [, , , handler] = mockTool.mock.calls[0];
875 |
876 | // Mock successful response
877 | const mockResponse = {
878 | candidates: [
879 | {
880 | content: {
881 | parts: [{ text: "Response with all parameters" }],
882 | },
883 | finishReason: FinishReason.STOP,
884 | },
885 | ],
886 | };
887 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
888 |
889 | // Prepare test request with all optional parameters
890 | const testRequest = {
891 | operation: "send_message",
892 | sessionId: "test-session-123",
893 | message: "Test message with all params",
894 | generationConfig: {
895 | temperature: 0.5,
896 | topP: 0.9,
897 | topK: 40,
898 | maxOutputTokens: 2048,
899 | stopSequences: ["END"],
900 | thinkingConfig: {
901 | thinkingBudget: 1024,
902 | },
903 | },
904 | safetySettings: [
905 | {
906 | category: "HARM_CATEGORY_HARASSMENT",
907 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
908 | },
909 | ],
910 | tools: [
911 | {
912 | functionDeclarations: [
913 | {
914 | name: "test_function",
915 | description: "Test function",
916 | parameters: {
917 | type: "OBJECT",
918 | properties: {},
919 | },
920 | },
921 | ],
922 | },
923 | ],
924 | toolConfig: {
925 | functionCallingConfig: {
926 | mode: "AUTO",
927 | },
928 | },
929 | cachedContentName: "cachedContents/test123",
930 | };
931 |
932 | // Call the handler
933 | const result = await handler(testRequest);
934 |
935 | // Verify all parameters were passed
936 | expect(mockSendMessageToSession).toHaveBeenCalledWith({
937 | sessionId: "test-session-123",
938 | message: "Test message with all params",
939 | generationConfig: testRequest.generationConfig,
940 | safetySettings: testRequest.safetySettings,
941 | tools: testRequest.tools,
942 | toolConfig: testRequest.toolConfig,
943 | cachedContentName: "cachedContents/test123",
944 | });
945 |
946 | expect(result).toEqual({
947 | content: [
948 | {
949 | type: "text",
950 | text: "Response with all parameters",
951 | },
952 | ],
953 | });
954 | });
955 |
956 | it("should handle start operation with string systemInstruction", async () => {
957 | // Register tool to get the request handler
958 | geminiChatTool(mockServer, mockService);
959 | const [, , , handler] = mockTool.mock.calls[0];
960 |
961 | // Mock successful response
962 | const mockSessionId = "test-session-string-instruction";
963 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
964 |
965 | // Prepare test request with string system instruction
966 | const testRequest = {
967 | operation: "start",
968 | systemInstruction:
969 | "You are a helpful assistant specialized in mathematics.",
970 | };
971 |
972 | // Call the handler
973 | const result = await handler(testRequest);
974 |
975 | // Verify string system instruction was passed
976 | expect(mockStartChatSession).toHaveBeenCalledWith(
977 | expect.objectContaining({
978 | systemInstruction:
979 | "You are a helpful assistant specialized in mathematics.",
980 | })
981 | );
982 |
983 | expect(result).toEqual({
984 | content: [
985 | {
986 | type: "text",
987 | text: JSON.stringify({ sessionId: mockSessionId }),
988 | },
989 | ],
990 | });
991 | });
992 |
993 | it("should handle start operation with complex history", async () => {
994 | // Register tool to get the request handler
995 | geminiChatTool(mockServer, mockService);
996 | const [, , , handler] = mockTool.mock.calls[0];
997 |
998 | // Mock successful response
999 | const mockSessionId = "test-session-complex-history";
1000 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
1001 |
1002 | // Prepare test request with complex history
1003 | const testRequest = {
1004 | operation: "start",
1005 | history: [
1006 | {
1007 | role: "user",
1008 | parts: [{ text: "Hello, I need help with math." }],
1009 | },
1010 | {
1011 | role: "model",
1012 | parts: [
1013 | {
1014 | text: "I'd be happy to help you with mathematics! What specific topic or problem would you like assistance with?",
1015 | },
1016 | ],
1017 | },
1018 | {
1019 | role: "user",
1020 | parts: [{ text: "Can you solve quadratic equations?" }],
1021 | },
1022 | {
1023 | role: "model",
1024 | parts: [
1025 | {
1026 | text: "Yes, I can help you solve quadratic equations. The general form is ax² + bx + c = 0.",
1027 | },
1028 | ],
1029 | },
1030 | ],
1031 | };
1032 |
1033 | // Call the handler
1034 | const result = await handler(testRequest);
1035 |
1036 | // Verify complex history was passed
1037 | expect(mockStartChatSession).toHaveBeenCalledWith(
1038 | expect.objectContaining({
1039 | history: testRequest.history,
1040 | })
1041 | );
1042 |
1043 | expect(result).toEqual({
1044 | content: [
1045 | {
1046 | type: "text",
1047 | text: JSON.stringify({ sessionId: mockSessionId }),
1048 | },
1049 | ],
1050 | });
1051 | });
1052 | });
1053 |
1054 | describe("function result processing", () => {
1055 | it("should handle function result with safety blocked response", async () => {
1056 | // Register tool to get the request handler
1057 | geminiChatTool(mockServer, mockService);
1058 | const [, , , handler] = mockTool.mock.calls[0];
1059 |
1060 | // Mock safety blocked response after function result
1061 | const mockResponse = {
1062 | candidates: [
1063 | {
1064 | content: {
1065 | parts: [{ text: "Partial response..." }],
1066 | },
1067 | finishReason: FinishReason.SAFETY,
1068 | },
1069 | ],
1070 | };
1071 | mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
1072 |
1073 | // Prepare test request
1074 | const testRequest = {
1075 | operation: "send_function_result",
1076 | sessionId: "test-session-123",
1077 | functionResponses: [
1078 | {
1079 | name: "test_function",
1080 | response: { result: "test result" },
1081 | },
1082 | ],
1083 | };
1084 |
1085 | // Call the handler
1086 | const result = await handler(testRequest);
1087 |
1088 | // Verify error response includes "after function result" context
1089 | expect(result).toEqual({
1090 | content: [
1091 | {
1092 | type: "text",
1093 | text: "Error: Response generation stopped due to safety settings after function result. FinishReason: SAFETY",
1094 | },
1095 | ],
1096 | isError: true,
1097 | });
1098 | });
1099 |
1100 | it("should handle multiple function responses", async () => {
1101 | // Register tool to get the request handler
1102 | geminiChatTool(mockServer, mockService);
1103 | const [, , , handler] = mockTool.mock.calls[0];
1104 |
1105 | // Mock successful response
1106 | const mockResponse = {
1107 | candidates: [
1108 | {
1109 | content: {
1110 | parts: [
1111 | {
1112 | text: "Based on the function results, here's the summary...",
1113 | },
1114 | ],
1115 | },
1116 | finishReason: FinishReason.STOP,
1117 | },
1118 | ],
1119 | };
1120 | mockSendFunctionResultToSession.mockResolvedValueOnce(mockResponse);
1121 |
1122 | // Prepare test request with multiple function responses
1123 | const testRequest = {
1124 | operation: "send_function_result",
1125 | sessionId: "test-session-123",
1126 | functionResponses: [
1127 | {
1128 | name: "get_weather",
1129 | response: { temperature: 22, condition: "sunny" },
1130 | },
1131 | {
1132 | name: "get_time",
1133 | response: { time: "14:30", timezone: "UTC" },
1134 | },
1135 | ],
1136 | };
1137 |
1138 | // Call the handler
1139 | const result = await handler(testRequest);
1140 |
1141 | // Verify multiple function responses were serialized correctly
1142 | expect(mockSendFunctionResultToSession).toHaveBeenCalledWith({
1143 | sessionId: "test-session-123",
1144 | functionResponse: JSON.stringify(testRequest.functionResponses),
1145 | functionCall: undefined,
1146 | });
1147 |
1148 | expect(result).toEqual({
1149 | content: [
1150 | {
1151 | type: "text",
1152 | text: "Based on the function results, here's the summary...",
1153 | },
1154 | ],
1155 | });
1156 | });
1157 | });
1158 |
1159 | describe("service error handling", () => {
1160 | it("should handle service errors during start operation", async () => {
1161 | // Register tool to get the request handler
1162 | geminiChatTool(mockServer, mockService);
1163 | const [, , , handler] = mockTool.mock.calls[0];
1164 |
1165 | // Mock service to throw error
1166 | const serviceError = new Error("Service unavailable");
1167 | mockStartChatSession.mockImplementationOnce(() => {
1168 | throw serviceError;
1169 | });
1170 |
1171 | // Prepare test request
1172 | const testRequest = {
1173 | operation: "start",
1174 | modelName: "gemini-1.5-flash",
1175 | };
1176 |
1177 | // Call the handler and expect error
1178 | await expect(handler(testRequest)).rejects.toThrow();
1179 | });
1180 |
1181 | it("should handle service errors during send_message operation", async () => {
1182 | // Register tool to get the request handler
1183 | geminiChatTool(mockServer, mockService);
1184 | const [, , , handler] = mockTool.mock.calls[0];
1185 |
1186 | // Mock service to throw error
1187 | const serviceError = new Error("Network error");
1188 | mockSendMessageToSession.mockRejectedValueOnce(serviceError);
1189 |
1190 | // Prepare test request
1191 | const testRequest = {
1192 | operation: "send_message",
1193 | sessionId: "test-session-123",
1194 | message: "Test message",
1195 | };
1196 |
1197 | // Call the handler and expect error
1198 | await expect(handler(testRequest)).rejects.toThrow();
1199 | });
1200 |
1201 | it("should handle service errors during send_function_result operation", async () => {
1202 | // Register tool to get the request handler
1203 | geminiChatTool(mockServer, mockService);
1204 | const [, , , handler] = mockTool.mock.calls[0];
1205 |
1206 | // Mock service to throw error
1207 | const serviceError = new Error("Function processing error");
1208 | mockSendFunctionResultToSession.mockRejectedValueOnce(serviceError);
1209 |
1210 | // Prepare test request
1211 | const testRequest = {
1212 | operation: "send_function_result",
1213 | sessionId: "test-session-123",
1214 | functionResponses: [
1215 | {
1216 | name: "test_function",
1217 | response: { result: "test" },
1218 | },
1219 | ],
1220 | };
1221 |
1222 | // Call the handler and expect error
1223 | await expect(handler(testRequest)).rejects.toThrow();
1224 | });
1225 | });
1226 |
1227 | describe("thinking configuration", () => {
1228 | it("should handle thinkingConfig in generationConfig for start operation", async () => {
1229 | // Register tool to get the request handler
1230 | geminiChatTool(mockServer, mockService);
1231 | const [, , , handler] = mockTool.mock.calls[0];
1232 |
1233 | // Mock successful response
1234 | const mockSessionId = "test-session-thinking";
1235 | mockStartChatSession.mockReturnValueOnce(mockSessionId);
1236 |
1237 | // Prepare test request with thinking configuration
1238 | const testRequest = {
1239 | operation: "start",
1240 | generationConfig: {
1241 | temperature: 0.7,
1242 | thinkingConfig: {
1243 | thinkingBudget: 2048,
1244 | reasoningEffort: "medium",
1245 | },
1246 | },
1247 | };
1248 |
1249 | // Call the handler
1250 | const result = await handler(testRequest);
1251 |
1252 | // Verify thinking config was passed
1253 | expect(mockStartChatSession).toHaveBeenCalledWith(
1254 | expect.objectContaining({
1255 | generationConfig: testRequest.generationConfig,
1256 | })
1257 | );
1258 |
1259 | expect(result).toEqual({
1260 | content: [
1261 | {
1262 | type: "text",
1263 | text: JSON.stringify({ sessionId: mockSessionId }),
1264 | },
1265 | ],
1266 | });
1267 | });
1268 |
1269 | it("should handle reasoningEffort in thinkingConfig", async () => {
1270 | // Register tool to get the request handler
1271 | geminiChatTool(mockServer, mockService);
1272 | const [, , , handler] = mockTool.mock.calls[0];
1273 |
1274 | // Mock successful response
1275 | const mockResponse = {
1276 | candidates: [
1277 | {
1278 | content: {
1279 | parts: [{ text: "Response with reasoning effort" }],
1280 | },
1281 | finishReason: FinishReason.STOP,
1282 | },
1283 | ],
1284 | };
1285 | mockSendMessageToSession.mockResolvedValueOnce(mockResponse);
1286 |
1287 | // Prepare test request with reasoning effort
1288 | const testRequest = {
1289 | operation: "send_message",
1290 | sessionId: "test-session-123",
1291 | message: "Complex reasoning task",
1292 | generationConfig: {
1293 | thinkingConfig: {
1294 | reasoningEffort: "high",
1295 | },
1296 | },
1297 | };
1298 |
1299 | // Call the handler
1300 | const result = await handler(testRequest);
1301 |
1302 | // Verify reasoning effort was passed
1303 | expect(mockSendMessageToSession).toHaveBeenCalledWith(
1304 | expect.objectContaining({
1305 | generationConfig: testRequest.generationConfig,
1306 | })
1307 | );
1308 |
1309 | expect(result).toEqual({
1310 | content: [
1311 | {
1312 | type: "text",
1313 | text: "Response with reasoning effort",
1314 | },
1315 | ],
1316 | });
1317 | });
1318 | });
1319 | });
1320 |
```
--------------------------------------------------------------------------------
/src/services/mcp/McpClientService.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { logger } from "../../utils/index.js";
2 | import { spawn, ChildProcess } from "child_process";
3 | import EventSource from "eventsource";
4 | import {
5 | McpError as SdkMcpError,
6 | ErrorCode,
7 | } from "@modelcontextprotocol/sdk/types.js";
8 | import { v4 as uuidv4 } from "uuid";
9 | // Import node-fetch types only
10 | // We'll dynamically import the actual implementation later to handle CJS/ESM compatibility
11 | import type { Response, RequestInit } from "node-fetch";
12 |
13 | // Define custom types for EventSource events since the eventsource package
14 | // doesn't export its own types
15 | interface ESMessageEvent {
16 | data: string;
17 | type: string;
18 | lastEventId: string;
19 | origin: string;
20 | }
21 |
22 | interface ESErrorEvent {
23 | type: string;
24 | message?: string;
25 | error?: Error;
26 | }
27 |
28 | // Add appropriate error handler typings
29 | type ESErrorHandler = (event: ESErrorEvent) => void;
30 | type ESMessageHandler = (event: ESMessageEvent) => void;
31 |
32 | // Extended EventSource interface to properly type the handlers
33 | interface ExtendedEventSource extends EventSource {
34 | onopen: (this: EventSource, ev: MessageEvent<unknown>) => unknown;
35 | onmessage: (this: EventSource, ev: MessageEvent) => unknown;
36 | onerror: (this: EventSource, ev: Event) => unknown;
37 | }
38 |
39 | export interface McpRequest {
40 | id: string;
41 | method: "listTools" | "callTool" | "initialize";
42 | params?: Record<string, unknown>;
43 | }
44 |
45 | export interface McpResponseError {
46 | code: number;
47 | message: string;
48 | data?: Record<string, unknown>;
49 | }
50 |
51 | export interface McpResponse {
52 | id: string;
53 | result?: Record<string, unknown> | Array<unknown>;
54 | error?: McpResponseError;
55 | }
56 |
57 | export interface ToolDefinition {
58 | name: string;
59 | description: string;
60 | parametersSchema: Record<string, unknown>; // JSON Schema object
61 | }
62 |
63 | export interface ConnectionDetails {
64 | type: "sse" | "stdio";
65 | sseUrl?: string;
66 | stdioCommand?: string;
67 | stdioArgs?: string[];
68 | connectionToken?: string;
69 | }
70 |
71 | // Re-export SDK McpError under the local name used throughout this file
72 | type McpError = SdkMcpError;
73 | const McpError = SdkMcpError;
74 |
75 | /**
76 | * Service for connecting to external Model Context Protocol (MCP) servers.
77 | * Provides methods to establish different types of connections (SSE, stdio).
78 | */
79 | export class McpClientService {
80 | // Maps to store active connections
81 | private activeSseConnections: Map<
82 | string,
83 | {
84 | eventSource: EventSource;
85 | baseUrl: string;
86 | lastActivityTimestamp: number; // Track when connection was last used
87 | }
88 | >;
89 | private activeStdioConnections: Map<
90 | string,
91 | {
92 | process: ChildProcess;
93 | lastActivityTimestamp: number; // Track when connection was last used
94 | }
95 | >;
96 | private pendingStdioRequests: Map<
97 | string, // connectionId
98 | Map<
99 | string,
100 | {
101 | resolve: (value: Record<string, unknown> | Array<unknown>) => void;
102 | reject: (reason: Error | McpError) => void;
103 | }
104 | > // requestId -> handlers
105 | > = new Map();
106 |
107 | // Configuration values
108 | private static readonly DEFAULT_REQUEST_TIMEOUT_MS = 30000; // 30 seconds
109 | private static readonly DEFAULT_CONNECTION_MAX_IDLE_MS = 600000; // 10 minutes
110 | private static readonly CONNECTION_CLEANUP_INTERVAL_MS = 300000; // Check every 5 minutes
111 |
112 | /**
113 | * Helper method to fetch with timeout
114 | * @param url - The URL to fetch
115 | * @param options - Fetch options
116 | * @param timeoutMs - Timeout in milliseconds
117 | * @param timeoutMessage - Message to include in timeout error
118 | * @returns The fetch response
119 | * @throws {SdkMcpError} - If the request times out
120 | */
121 | private async fetchWithTimeout(
122 | url: string,
123 | options: RequestInit,
124 | timeoutMs = McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
125 | timeoutMessage = "Request timed out"
126 | ): Promise<Response> {
127 | // Create controller for aborting the fetch
128 | const controller = new AbortController();
129 | const id = setTimeout(() => controller.abort(), timeoutMs);
130 |
131 | try {
132 | // Add the signal to the options
133 | const fetchOptions = {
134 | ...options,
135 | signal: controller.signal,
136 | };
137 |
138 | // Dynamically import node-fetch (v2 is CommonJS)
139 | const nodeFetch = await import("node-fetch");
140 | const fetch = nodeFetch.default;
141 |
142 | // Make the fetch request
143 | const response = await fetch(url, fetchOptions);
144 | clearTimeout(id);
145 | return response;
146 | } catch (error) {
147 | clearTimeout(id);
148 | throw new SdkMcpError(
149 | ErrorCode.InternalError,
150 | `${timeoutMessage} after ${timeoutMs}ms`
151 | );
152 | }
153 | }
154 |
155 | // Cleanup timer reference
156 | private cleanupIntervalId?: NodeJS.Timeout;
157 |
158 | constructor() {
159 | this.activeSseConnections = new Map();
160 | this.activeStdioConnections = new Map();
161 | logger.info("McpClientService initialized.");
162 |
163 | // Start the connection cleanup interval
164 | this.cleanupIntervalId = setInterval(
165 | () => this.cleanupStaleConnections(),
166 | McpClientService.CONNECTION_CLEANUP_INTERVAL_MS
167 | );
168 | }
169 |
170 | /**
171 | * Cleans up stale connections that haven't been used for a while
172 | * @private
173 | */
174 | private cleanupStaleConnections(): void {
175 | const now = Date.now();
176 | const maxIdleTime = McpClientService.DEFAULT_CONNECTION_MAX_IDLE_MS;
177 | let closedCount = 0;
178 |
179 | // Check SSE connections
180 | for (const [
181 | connectionId,
182 | { lastActivityTimestamp },
183 | ] of this.activeSseConnections.entries()) {
184 | if (now - lastActivityTimestamp > maxIdleTime) {
185 | logger.info(
186 | `Closing stale SSE connection ${connectionId} (idle for ${Math.floor((now - lastActivityTimestamp) / 1000)} seconds)`
187 | );
188 | this.closeSseConnection(connectionId);
189 | closedCount++;
190 | }
191 | }
192 |
193 | // Check stdio connections
194 | for (const [
195 | connectionId,
196 | { lastActivityTimestamp },
197 | ] of this.activeStdioConnections.entries()) {
198 | if (now - lastActivityTimestamp > maxIdleTime) {
199 | logger.info(
200 | `Closing stale stdio connection ${connectionId} (idle for ${Math.floor((now - lastActivityTimestamp) / 1000)} seconds)`
201 | );
202 | this.closeStdioConnection(connectionId);
203 | closedCount++;
204 | }
205 | }
206 |
207 | if (closedCount > 0) {
208 | logger.info(`Cleaned up ${closedCount} stale connections`);
209 | }
210 | }
211 |
212 | /**
213 | * Validates a server ID.
214 | * @param serverId - The server ID to validate.
215 | * @throws {McpError} Throws an error if the server ID is invalid.
216 | */
217 | private validateServerId(serverId: string): void {
218 | if (!serverId || typeof serverId !== "string" || serverId.trim() === "") {
219 | throw new McpError(
220 | ErrorCode.InvalidParams,
221 | "Server ID must be a non-empty string"
222 | );
223 | }
224 | }
225 |
226 | /**
227 | * Checks if a connection exists for the given server ID.
228 | * @param serverId - The server ID to check.
229 | * @throws {McpError} Throws an error if the connection doesn't exist.
230 | */
231 | private validateConnectionExists(serverId: string): void {
232 | if (
233 | !this.activeSseConnections.has(serverId) &&
234 | !this.activeStdioConnections.has(serverId)
235 | ) {
236 | throw new McpError(
237 | ErrorCode.InvalidRequest,
238 | `Connection not found for serverId: ${serverId}`
239 | );
240 | }
241 | }
242 |
243 | /**
244 | * Establishes a connection to an MCP server.
245 | * @param serverId - A unique identifier provided by the caller to reference this server connection.
246 | * Note: This is NOT used as the internal connection tracking ID.
247 | * @param connectionDetails - The details for establishing the connection.
248 | * @param messageHandler - Optional callback for handling received messages.
249 | * @returns A promise that resolves to a connection ID (different from serverId) when the connection is established.
250 | * This returned connectionId should be used for all subsequent interactions with this connection.
251 | * @throws {McpError} Throws an error if the parameters are invalid.
252 | */
253 | public async connect(
254 | serverId: string,
255 | connectionDetails: ConnectionDetails,
256 | messageHandler?: (data: unknown) => void
257 | ): Promise<string> {
258 | // Validate serverId
259 | this.validateServerId(serverId);
260 |
261 | // Validate connectionDetails
262 | if (!connectionDetails || typeof connectionDetails !== "object") {
263 | throw new McpError(
264 | ErrorCode.InvalidParams,
265 | "Connection details must be an object"
266 | );
267 | }
268 |
269 | // Validate connection type
270 | if (
271 | connectionDetails.type !== "sse" &&
272 | connectionDetails.type !== "stdio"
273 | ) {
274 | throw new McpError(
275 | ErrorCode.InvalidParams,
276 | "Connection type must be 'sse' or 'stdio'"
277 | );
278 | }
279 |
280 | // Validate SSE connection details
281 | if (connectionDetails.type === "sse") {
282 | if (
283 | !connectionDetails.sseUrl ||
284 | typeof connectionDetails.sseUrl !== "string" ||
285 | connectionDetails.sseUrl.trim() === ""
286 | ) {
287 | throw new McpError(
288 | ErrorCode.InvalidParams,
289 | "For SSE connections, sseUrl must be a non-empty string"
290 | );
291 | }
292 |
293 | // Basic URL format validation
294 | if (
295 | !connectionDetails.sseUrl.startsWith("http://") &&
296 | !connectionDetails.sseUrl.startsWith("https://")
297 | ) {
298 | throw new McpError(
299 | ErrorCode.InvalidParams,
300 | "sseUrl must be a valid URL format starting with http:// or https://"
301 | );
302 | }
303 |
304 | return this.connectSse(
305 | connectionDetails.sseUrl,
306 | connectionDetails.connectionToken,
307 | messageHandler
308 | );
309 | }
310 | // Validate stdio connection details
311 | else if (connectionDetails.type === "stdio") {
312 | if (
313 | !connectionDetails.stdioCommand ||
314 | typeof connectionDetails.stdioCommand !== "string" ||
315 | connectionDetails.stdioCommand.trim() === ""
316 | ) {
317 | throw new McpError(
318 | ErrorCode.InvalidParams,
319 | "For stdio connections, stdioCommand must be a non-empty string"
320 | );
321 | }
322 |
323 | return this.connectStdio(
324 | connectionDetails.stdioCommand,
325 | connectionDetails.stdioArgs || [],
326 | connectionDetails.connectionToken,
327 | messageHandler
328 | );
329 | }
330 |
331 | // This should never be reached due to the type check above
332 | throw new McpError(
333 | ErrorCode.InvalidParams,
334 | "Invalid connection type specified"
335 | );
336 | }
337 |
338 | /**
339 | * Establishes an SSE connection to the specified MCP server.
340 | * @param url - The URL of the MCP server to connect to.
341 | * @param connectionToken - Optional token for authentication with the server.
342 | * @param messageHandler - Optional callback for handling received messages.
343 | * @returns A promise that resolves to a connection ID when the connection is established.
344 | */
345 | private connectSse(
346 | url: string,
347 | connectionToken?: string,
348 | messageHandler?: (data: unknown) => void
349 | ): Promise<string> {
350 | return new Promise((resolve, reject) => {
351 | logger.info(`Connecting to MCP server via SSE: ${url}`);
352 |
353 | try {
354 | // Generate a unique connectionId for internal tracking
355 | // This will be different from the serverId passed to the connect() method
356 | const connectionId = uuidv4();
357 |
358 | // Create a timeout for the connection attempt
359 | const connectionTimeout = setTimeout(() => {
360 | reject(
361 | new SdkMcpError(
362 | ErrorCode.InternalError,
363 | `Connection timeout while attempting to connect to ${url}`
364 | )
365 | );
366 | }, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
367 |
368 | // Add connectionToken to headers if provided
369 | const options: EventSource.EventSourceInitDict = {};
370 | if (connectionToken) {
371 | logger.debug(`Adding connection token to SSE request`);
372 | options.headers = {
373 | Authorization: `Bearer ${connectionToken}`,
374 | };
375 | }
376 |
377 | // Create EventSource for SSE connection with options
378 | const eventSource = new EventSource(url, options);
379 |
380 | // Handler functions to store for proper cleanup
381 | const onOpen = () => {
382 | // Clear the connection timeout
383 | clearTimeout(connectionTimeout);
384 |
385 | logger.info(`SSE connection established to ${url}`);
386 | this.activeSseConnections.set(connectionId, {
387 | eventSource,
388 | baseUrl: url,
389 | lastActivityTimestamp: Date.now(),
390 | });
391 | resolve(connectionId);
392 | };
393 |
394 | const onMessage = ((event: ESMessageEvent) => {
395 | logger.debug(`SSE message received from ${url}:`, event.data);
396 |
397 | // Update the last activity timestamp
398 | const connection = this.activeSseConnections.get(connectionId);
399 | if (connection) {
400 | connection.lastActivityTimestamp = Date.now();
401 | }
402 |
403 | if (messageHandler) {
404 | try {
405 | const parsedData = JSON.parse(event.data);
406 | messageHandler(parsedData);
407 | } catch (error) {
408 | logger.error(`Error parsing SSE message:`, error);
409 | messageHandler(event.data);
410 | }
411 | }
412 | }) as ESMessageHandler;
413 |
414 | const onError = ((error: ESErrorEvent) => {
415 | // Clear the connection timeout if it's still pending
416 | clearTimeout(connectionTimeout);
417 |
418 | logger.error(
419 | `SSE connection error for ${url}:`,
420 | error.message || "Unknown error"
421 | );
422 |
423 | if (!this.activeSseConnections.has(connectionId)) {
424 | // If we haven't resolved yet, this is a connection failure
425 | reject(
426 | new SdkMcpError(
427 | ErrorCode.InternalError,
428 | `Failed to establish SSE connection to ${url}: ${error.message || "Unknown error"}`
429 | )
430 | );
431 | } else if (eventSource.readyState === EventSource.CLOSED) {
432 | // Connection was established but is now closed
433 | logger.info(`SSE connection ${connectionId} closed due to error.`);
434 | this.activeSseConnections.delete(connectionId);
435 | } else {
436 | // Connection is still open but had an error
437 | logger.warn(
438 | `SSE connection ${connectionId} had an error but is still open. Monitoring for further issues.`
439 | );
440 | }
441 | }) as ESErrorHandler;
442 |
443 | // Set up event handlers
444 | eventSource.onopen = onOpen;
445 | eventSource.onmessage = onMessage;
446 | eventSource.onerror = onError;
447 | } catch (error) {
448 | logger.error(`Error creating SSE connection to ${url}:`, error);
449 | reject(error);
450 | }
451 | });
452 | }
453 |
454 | /**
455 | * Closes an SSE connection.
456 | * @param connectionId - The ID of the connection to close.
457 | * @returns True if the connection was closed, false if it wasn't found.
458 | */
459 | public closeSseConnection(connectionId: string): boolean {
460 | const connection = this.activeSseConnections.get(connectionId);
461 | if (connection) {
462 | // Close the EventSource and remove listeners
463 | const eventSource = connection.eventSource;
464 |
465 | // Clean up event listeners by setting handlers to empty functions
466 | // (EventSource doesn't support removeEventListener)
467 | (eventSource as ExtendedEventSource).onopen = () => {};
468 | (eventSource as ExtendedEventSource).onmessage = () => {};
469 | (eventSource as ExtendedEventSource).onerror = () => {};
470 |
471 | // Close the connection
472 | eventSource.close();
473 |
474 | // Remove from active connections
475 | this.activeSseConnections.delete(connectionId);
476 |
477 | // Clean up any pending requests for this connection (this shouldn't generally happen for SSE)
478 | this.cleanupPendingRequestsForConnection(connectionId);
479 |
480 | logger.info(`SSE connection ${connectionId} closed.`);
481 | return true;
482 | }
483 | logger.warn(
484 | `Attempted to close non-existent SSE connection: ${connectionId}`
485 | );
486 | return false;
487 | }
488 |
489 | /**
490 | * Helper method to clean up pending requests for a connection
491 | * @param connectionId - The ID of the connection to clean up pending requests for
492 | */
493 | private cleanupPendingRequestsForConnection(connectionId: string): void {
494 | // If there are any pending requests for this connection, reject them all
495 | if (this.pendingStdioRequests.has(connectionId)) {
496 | const pendingRequests = this.pendingStdioRequests.get(connectionId)!;
497 | for (const [
498 | requestId,
499 | { reject: rejectRequest },
500 | ] of pendingRequests.entries()) {
501 | logger.warn(
502 | `Rejecting pending request ${requestId} due to connection cleanup`
503 | );
504 | rejectRequest(
505 | new McpError(
506 | ErrorCode.InternalError,
507 | `Connection closed during cleanup before response was received`
508 | )
509 | );
510 | }
511 | // Clean up the map entry
512 | this.pendingStdioRequests.delete(connectionId);
513 | }
514 | }
515 |
516 | /**
517 | * Establishes a stdio connection using the specified command.
518 | * @param command - The command to execute for stdio connection.
519 | * @param args - Arguments to pass to the command.
520 | * @param connectionToken - Optional token for authentication with the server.
521 | * @param messageHandler - Optional callback for handling stdout data.
522 | * @returns A promise that resolves to a connection ID when the process is spawned.
523 | */
524 | private connectStdio(
525 | command: string,
526 | args: string[] = [],
527 | connectionToken?: string,
528 | messageHandler?: (data: unknown) => void
529 | ): Promise<string> {
530 | return new Promise((resolve, reject) => {
531 | logger.info(
532 | `Connecting to MCP server via stdio using command: ${command} ${args.join(" ")}`
533 | );
534 |
535 | try {
536 | // Generate a unique connectionId for internal tracking
537 | // This will be different from the serverId passed to the connect() method
538 | const connectionId = uuidv4();
539 |
540 | // Create a timeout for the connection establishment
541 | const connectionTimeout = setTimeout(() => {
542 | reject(
543 | new SdkMcpError(
544 | ErrorCode.InternalError,
545 | `Timeout while establishing stdio connection for command: ${command}`
546 | )
547 | );
548 | }, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
549 |
550 | // Prepare the environment for the child process
551 | const env = { ...process.env };
552 |
553 | // Add connectionToken to environment if provided
554 | if (connectionToken) {
555 | logger.debug("Adding connection token to stdio environment");
556 | env.MCP_CONNECTION_TOKEN = connectionToken;
557 | }
558 |
559 | // Spawn the child process with environment
560 | const childProcess = spawn(command, args, {
561 | stdio: "pipe",
562 | env: env,
563 | });
564 |
565 | // Store the connection with timestamp
566 | this.activeStdioConnections.set(connectionId, {
567 | process: childProcess,
568 | lastActivityTimestamp: Date.now(),
569 | });
570 |
571 | // Buffer to accumulate data chunks
572 | let buffer = "";
573 |
574 | // We'll mark connection as established when the process is ready
575 | const connectionEstablished = () => {
576 | clearTimeout(connectionTimeout);
577 | logger.info(`Stdio connection established for ${command}`);
578 | resolve(connectionId);
579 | };
580 |
581 | // Data handler function for stdout
582 | const onStdoutData = (data: Buffer) => {
583 | // Update the last activity timestamp to prevent cleanup
584 | const connection = this.activeStdioConnections.get(connectionId);
585 | if (connection) {
586 | connection.lastActivityTimestamp = Date.now();
587 | }
588 |
589 | // Append the new data to our buffer
590 | buffer += data.toString();
591 | logger.debug(
592 | `Stdio stdout from ${command} (raw chunk):`,
593 | data.toString()
594 | );
595 |
596 | // Process complete lines in the buffer
597 | let newlineIndex;
598 | while ((newlineIndex = buffer.indexOf("\n")) !== -1) {
599 | // Extract the line (excluding the newline)
600 | const line = buffer.substring(0, newlineIndex);
601 | // Remove the processed line from the buffer (including the newline)
602 | buffer = buffer.substring(newlineIndex + 1);
603 |
604 | // Skip empty lines
605 | if (!line.trim()) continue;
606 |
607 | try {
608 | // Try to parse the line as JSON
609 | const parsedData = JSON.parse(line);
610 | logger.debug(`Parsed JSON message:`, parsedData);
611 |
612 | // Check if this is a response to a pending request
613 | if (parsedData.id) {
614 | // Only check pending requests for the current connection (connectionId)
615 | const requestsMap = this.pendingStdioRequests.get(connectionId);
616 | let foundRequest = false;
617 |
618 | if (requestsMap && requestsMap.has(parsedData.id)) {
619 | const { resolve, reject } = requestsMap.get(parsedData.id)!;
620 |
621 | requestsMap.delete(parsedData.id);
622 |
623 | // If this was the last pending request, clean up the connection map
624 | if (requestsMap.size === 0) {
625 | this.pendingStdioRequests.delete(connectionId);
626 | }
627 |
628 | foundRequest = true;
629 |
630 | if (parsedData.error) {
631 | reject(
632 | new SdkMcpError(
633 | (parsedData.error.code as ErrorCode) ||
634 | ErrorCode.InternalError,
635 | parsedData.error.message || "Tool execution error",
636 | parsedData.error.data
637 | )
638 | );
639 | } else {
640 | // Verify the result is an object or array
641 | if (
642 | parsedData.result === null ||
643 | parsedData.result === undefined
644 | ) {
645 | reject(
646 | new McpError(
647 | ErrorCode.InternalError,
648 | "Received null or undefined result from tool",
649 | { responseId: parsedData.id }
650 | )
651 | );
652 | } else if (
653 | typeof parsedData.result !== "object" &&
654 | !Array.isArray(parsedData.result)
655 | ) {
656 | reject(
657 | new McpError(
658 | ErrorCode.InternalError,
659 | "Expected object or array result from tool",
660 | {
661 | responseId: parsedData.id,
662 | receivedType: typeof parsedData.result,
663 | }
664 | )
665 | );
666 | } else {
667 | resolve(
668 | parsedData.result as
669 | | Record<string, unknown>
670 | | Array<unknown>
671 | );
672 | }
673 | }
674 | }
675 |
676 | // Only log if we didn't find the request
677 | if (!foundRequest && messageHandler) {
678 | logger.debug(
679 | `Received message with ID ${parsedData.id} but no matching pending request found for this connection`
680 | );
681 | messageHandler(parsedData);
682 | }
683 | } else if (messageHandler) {
684 | // If not a response to a pending request, pass to the message handler
685 | messageHandler(parsedData);
686 | }
687 | } catch (error) {
688 | logger.warn(`Error parsing JSON from stdio:`, error);
689 | // If not valid JSON and we have a message handler, pass the raw line
690 | if (messageHandler) {
691 | messageHandler(line);
692 | }
693 | }
694 | }
695 | };
696 |
697 | // Error handler for stderr
698 | const onStderrData = (data: Buffer) => {
699 | // Update the last activity timestamp to prevent cleanup
700 | const connection = this.activeStdioConnections.get(connectionId);
701 | if (connection) {
702 | connection.lastActivityTimestamp = Date.now();
703 | }
704 |
705 | logger.warn(`Stdio stderr from ${command}:`, data.toString());
706 | };
707 |
708 | // Error handler for the process
709 | const onError = (error: Error) => {
710 | // Clear the connection timeout
711 | clearTimeout(connectionTimeout);
712 |
713 | logger.error(`Stdio error for ${command}:`, error);
714 | if (this.activeStdioConnections.has(connectionId)) {
715 | this.activeStdioConnections.delete(connectionId);
716 |
717 | // Reject all pending requests for this connection
718 | if (this.pendingStdioRequests.has(connectionId)) {
719 | const pendingRequests =
720 | this.pendingStdioRequests.get(connectionId)!;
721 | for (const [
722 | requestId,
723 | { reject: rejectRequest },
724 | ] of pendingRequests.entries()) {
725 | logger.warn(
726 | `Rejecting pending request ${requestId} due to connection error`
727 | );
728 | rejectRequest(
729 | new SdkMcpError(
730 | ErrorCode.InternalError,
731 | `Connection error occurred before response: ${error.message}`
732 | )
733 | );
734 | }
735 | this.pendingStdioRequests.delete(connectionId);
736 | }
737 | }
738 | reject(error);
739 | };
740 |
741 | // Close handler for the process
742 | const onClose = (
743 | code: number | null,
744 | signal: NodeJS.Signals | null
745 | ) => {
746 | // Clear the connection timeout if process closes before we establish connection
747 | clearTimeout(connectionTimeout);
748 |
749 | logger.info(
750 | `Stdio process ${command} closed with code ${code} and signal ${signal}`
751 | );
752 | if (this.activeStdioConnections.has(connectionId)) {
753 | this.activeStdioConnections.delete(connectionId);
754 |
755 | // Reject all pending requests for this connection
756 | if (this.pendingStdioRequests.has(connectionId)) {
757 | const pendingRequests =
758 | this.pendingStdioRequests.get(connectionId)!;
759 | for (const [
760 | requestId,
761 | { reject: rejectRequest },
762 | ] of pendingRequests.entries()) {
763 | logger.warn(
764 | `Rejecting pending request ${requestId} due to connection closure`
765 | );
766 | rejectRequest(
767 | new McpError(
768 | ErrorCode.InternalError,
769 | `Connection closed before response (code: ${code}, signal: ${signal})`
770 | )
771 | );
772 | }
773 | this.pendingStdioRequests.delete(connectionId);
774 | }
775 | }
776 | };
777 |
778 | // Set up event handlers
779 | childProcess.stdout.on("data", onStdoutData);
780 | childProcess.stderr.on("data", onStderrData);
781 | childProcess.on("error", onError);
782 | childProcess.on("close", onClose);
783 |
784 | // The connection is established immediately after we set up event handlers
785 | connectionEstablished();
786 | } catch (error) {
787 | logger.error(`Error creating stdio connection for ${command}:`, error);
788 | reject(error);
789 | }
790 | });
791 | }
792 |
793 | /**
794 | * Sends data to a stdio connection.
795 | * @param connectionId - The ID of the connection to send data to.
796 | * @param data - The data to send.
797 | * @returns True if the data was sent, false if the connection wasn't found.
798 | */
799 | private sendToStdio(connectionId: string, data: string | object): boolean {
800 | const connection = this.activeStdioConnections.get(connectionId);
801 | if (connection) {
802 | const childProcess = connection.process;
803 |
804 | // Update the last activity timestamp
805 | connection.lastActivityTimestamp = Date.now();
806 |
807 | // Safety check for data size to prevent buffer overflow
808 | const dataStr = typeof data === "string" ? data : JSON.stringify(data);
809 |
810 | // Limit data size to 1MB to prevent abuse
811 | const MAX_DATA_SIZE = 1024 * 1024; // 1MB
812 | if (dataStr.length > MAX_DATA_SIZE) {
813 | logger.error(
814 | `Data to send to stdio connection ${connectionId} exceeds size limit (${dataStr.length} > ${MAX_DATA_SIZE})`
815 | );
816 | return false;
817 | }
818 |
819 | if (childProcess.stdin) {
820 | try {
821 | childProcess.stdin.write(dataStr + "\n");
822 | } catch (error) {
823 | logger.error(
824 | `Error writing to stdin for connection ${connectionId}:`,
825 | error
826 | );
827 | return false;
828 | }
829 | } else {
830 | logger.error(`Stdio connection ${connectionId} has no stdin`);
831 | return false;
832 | }
833 | logger.debug(`Sent data to stdio connection ${connectionId}`);
834 | return true;
835 | }
836 | logger.warn(
837 | `Attempted to send data to non-existent stdio connection: ${connectionId}`
838 | );
839 | return false;
840 | }
841 |
842 | /**
843 | * Closes a stdio connection.
844 | * @param connectionId - The ID of the connection to close.
845 | * @param signal - Optional signal to send to the process. Default is 'SIGTERM'.
846 | * @returns True if the connection was closed, false if it wasn't found.
847 | */
848 | public closeStdioConnection(
849 | connectionId: string,
850 | signal: NodeJS.Signals = "SIGTERM"
851 | ): boolean {
852 | const connection = this.activeStdioConnections.get(connectionId);
853 | if (connection) {
854 | const childProcess = connection.process;
855 |
856 | // Remove all listeners to prevent memory leaks
857 | childProcess.stdout?.removeAllListeners();
858 | childProcess.stderr?.removeAllListeners();
859 | childProcess.removeAllListeners();
860 |
861 | // Kill the process
862 | childProcess.kill(signal);
863 |
864 | // Remove from active connections
865 | this.activeStdioConnections.delete(connectionId);
866 |
867 | // Clean up any pending requests for this connection
868 | this.cleanupPendingRequestsForConnection(connectionId);
869 |
870 | logger.info(
871 | `Stdio connection ${connectionId} closed with signal ${signal}.`
872 | );
873 | return true;
874 | }
875 | logger.warn(
876 | `Attempted to close non-existent stdio connection: ${connectionId}`
877 | );
878 | return false;
879 | }
880 |
881 | /**
882 | * Gets all active SSE connection IDs.
883 | * @returns Array of active SSE connection IDs.
884 | */
885 | public getActiveSseConnectionIds(): string[] {
886 | return Array.from(this.activeSseConnections.keys());
887 | }
888 |
889 | /**
890 | * Gets all active stdio connection IDs.
891 | * @returns Array of active stdio connection IDs.
892 | */
893 | public getActiveStdioConnectionIds(): string[] {
894 | return Array.from(this.activeStdioConnections.keys());
895 | }
896 |
897 | /**
898 | * Gets the last activity timestamp for a connection
899 | * @param connectionId - The ID of the connection to check
900 | * @returns The last activity timestamp in milliseconds since the epoch, or undefined if the connection doesn't exist
901 | */
902 | public getLastActivityTimestamp(connectionId: string): number | undefined {
903 | const sseConnection = this.activeSseConnections.get(connectionId);
904 | if (sseConnection) {
905 | return sseConnection.lastActivityTimestamp;
906 | }
907 |
908 | const stdioConnection = this.activeStdioConnections.get(connectionId);
909 | if (stdioConnection) {
910 | return stdioConnection.lastActivityTimestamp;
911 | }
912 |
913 | return undefined;
914 | }
915 |
916 | /**
917 | * Lists all available tools from an MCP server.
918 | * @param serverId - The ID of the connection to query.
919 | * @returns A promise that resolves to an array of tool definitions.
920 | * @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
921 | */
922 | public async listTools(serverId: string): Promise<ToolDefinition[]> {
923 | // Validate serverId
924 | this.validateServerId(serverId);
925 |
926 | // Validate connection exists
927 | this.validateConnectionExists(serverId);
928 |
929 | logger.info(`Listing tools for connection ${serverId}`);
930 |
931 | // Check if this is an SSE connection
932 | if (this.activeSseConnections.has(serverId)) {
933 | const connection = this.activeSseConnections.get(serverId)!;
934 | const requestId = uuidv4();
935 | const request: McpRequest = { id: requestId, method: "listTools" };
936 |
937 | try {
938 | // Create URL for the MCP request
939 | const mcpRequestUrl = new URL(connection.baseUrl);
940 |
941 | // Update the connection's last activity timestamp
942 | connection.lastActivityTimestamp = Date.now();
943 |
944 | // Make the request with timeout
945 | const response = await this.fetchWithTimeout(
946 | mcpRequestUrl.toString(),
947 | {
948 | method: "POST",
949 | headers: {
950 | "Content-Type": "application/json",
951 | },
952 | body: JSON.stringify(request),
953 | },
954 | McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
955 | "Request timed out"
956 | );
957 |
958 | if (!response.ok) {
959 | throw new McpError(
960 | ErrorCode.InternalError,
961 | `HTTP error from MCP server: ${response.status} ${response.statusText}`
962 | );
963 | }
964 |
965 | const mcpResponse = (await response.json()) as McpResponse;
966 |
967 | if (mcpResponse.error) {
968 | throw new SdkMcpError(
969 | ErrorCode.InternalError,
970 | `MCP error: ${mcpResponse.error.message} (code: ${mcpResponse.error.code})`,
971 | mcpResponse.error.data
972 | );
973 | }
974 |
975 | // Type assertion with verification to ensure we have an array of ToolDefinition
976 | const result = mcpResponse.result;
977 | if (!Array.isArray(result)) {
978 | throw new McpError(
979 | ErrorCode.InternalError,
980 | "Expected array of tools in response",
981 | { receivedType: typeof result }
982 | );
983 | }
984 |
985 | return result as ToolDefinition[];
986 | } catch (error) {
987 | logger.error(
988 | `Error listing tools for SSE connection ${serverId}:`,
989 | error
990 | );
991 |
992 | // Wrap non-McpError instances
993 | if (!(error instanceof McpError)) {
994 | throw new McpError(
995 | ErrorCode.InternalError,
996 | `Failed to list tools for connection ${serverId}: ${error instanceof Error ? error.message : String(error)}`
997 | );
998 | }
999 | throw error;
1000 | }
1001 | }
1002 |
1003 | // Check if this is a stdio connection
1004 | else if (this.activeStdioConnections.has(serverId)) {
1005 | const requestId = uuidv4();
1006 | const request: McpRequest = { id: requestId, method: "listTools" };
1007 |
1008 | return new Promise<ToolDefinition[]>((resolve, reject) => {
1009 | // Initialize the map for this connection if it doesn't exist
1010 | if (!this.pendingStdioRequests.has(serverId)) {
1011 | this.pendingStdioRequests.set(serverId, new Map());
1012 | }
1013 |
1014 | // Store the promise resolution functions
1015 | this.pendingStdioRequests.get(serverId)!.set(requestId, {
1016 | resolve: (value) => {
1017 | // Type-safe resolution for tool definitions
1018 | resolve(value as ToolDefinition[]);
1019 | },
1020 | reject: reject,
1021 | });
1022 |
1023 | // Set up a timeout to automatically reject this request if it takes too long
1024 | setTimeout(() => {
1025 | // If the request is still pending, reject it
1026 | if (
1027 | this.pendingStdioRequests.has(serverId) &&
1028 | this.pendingStdioRequests.get(serverId)!.has(requestId)
1029 | ) {
1030 | // Get the reject function
1031 | const { reject: rejectRequest } = this.pendingStdioRequests
1032 | .get(serverId)!
1033 | .get(requestId)!;
1034 |
1035 | // Delete the request
1036 | this.pendingStdioRequests.get(serverId)!.delete(requestId);
1037 |
1038 | // If this was the last pending request, clean up the connection map
1039 | if (this.pendingStdioRequests.get(serverId)!.size === 0) {
1040 | this.pendingStdioRequests.delete(serverId);
1041 | }
1042 |
1043 | // Reject the request with a timeout error
1044 | rejectRequest(
1045 | new SdkMcpError(
1046 | ErrorCode.InternalError,
1047 | "Request timed out waiting for response"
1048 | )
1049 | );
1050 | }
1051 | }, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
1052 |
1053 | // Send the request
1054 | const sent = this.sendToStdio(serverId, request);
1055 |
1056 | if (!sent) {
1057 | // Clean up the pending request if sending fails
1058 | this.pendingStdioRequests.get(serverId)!.delete(requestId);
1059 |
1060 | // If this was the last pending request, clean up the connection map
1061 | if (this.pendingStdioRequests.get(serverId)!.size === 0) {
1062 | this.pendingStdioRequests.delete(serverId);
1063 | }
1064 |
1065 | reject(
1066 | new Error(`Failed to send request to stdio connection ${serverId}`)
1067 | );
1068 | }
1069 | });
1070 | }
1071 |
1072 | // This should never be reached due to the validateConnectionExists check above
1073 | throw new McpError(
1074 | ErrorCode.InvalidRequest,
1075 | `No connection found with ID ${serverId}`
1076 | );
1077 | }
1078 | /**
1079 | * Gets server information from an MCP server.
1080 | * @param serverId - The ID of the connection to use.
1081 | * @returns A promise that resolves to the server information.
1082 | * @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
1083 | */
1084 | public async getServerInfo(
1085 | serverId: string
1086 | ): Promise<Record<string, unknown>> {
1087 | // Validate serverId
1088 | this.validateServerId(serverId);
1089 |
1090 | // Check if connection exists
1091 | this.validateConnectionExists(serverId);
1092 |
1093 | logger.debug(`Getting server info for connection: ${serverId}`);
1094 |
1095 | // Check if this is an SSE connection
1096 | if (this.activeSseConnections.has(serverId)) {
1097 | const connection = this.activeSseConnections.get(serverId)!;
1098 | connection.lastActivityTimestamp = Date.now();
1099 |
1100 | // For SSE connections, we'll make an HTTP request to get server info
1101 | const baseUrl = connection.baseUrl;
1102 | const infoUrl = `${baseUrl}/info`;
1103 |
1104 | try {
1105 | const response = await this.fetchWithTimeout(
1106 | infoUrl,
1107 | {
1108 | method: "GET",
1109 | headers: {
1110 | "Content-Type": "application/json",
1111 | },
1112 | },
1113 | McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
1114 | `Server info request timed out for ${serverId}`
1115 | );
1116 |
1117 | if (!response.ok) {
1118 | throw new McpError(
1119 | ErrorCode.InternalError,
1120 | `Server info request failed with status ${response.status}: ${response.statusText}`
1121 | );
1122 | }
1123 |
1124 | const serverInfo = await response.json();
1125 | return serverInfo as Record<string, unknown>;
1126 | } catch (error) {
1127 | logger.error(
1128 | `Error getting server info for SSE connection ${serverId}:`,
1129 | error
1130 | );
1131 | throw new McpError(
1132 | ErrorCode.InternalError,
1133 | `Failed to get server info: ${error instanceof Error ? error.message : "Unknown error"}`
1134 | );
1135 | }
1136 | }
1137 |
1138 | // Check if this is a stdio connection
1139 | if (this.activeStdioConnections.has(serverId)) {
1140 | const connection = this.activeStdioConnections.get(serverId)!;
1141 | connection.lastActivityTimestamp = Date.now();
1142 |
1143 | // For stdio connections, send an initialize request
1144 | const requestId = uuidv4();
1145 | const request: McpRequest = {
1146 | id: requestId,
1147 | method: "initialize",
1148 | params: {
1149 | protocolVersion: "2024-11-05",
1150 | capabilities: {},
1151 | clientInfo: {
1152 | name: "mcp-gemini-server",
1153 | version: "1.0.0",
1154 | },
1155 | },
1156 | };
1157 |
1158 | return new Promise<Record<string, unknown>>((resolve, reject) => {
1159 | // Set up timeout for the request
1160 | const timeout = setTimeout(() => {
1161 | // Clean up the pending request
1162 | const pendingRequests = this.pendingStdioRequests.get(serverId);
1163 | if (pendingRequests) {
1164 | pendingRequests.delete(requestId);
1165 | if (pendingRequests.size === 0) {
1166 | this.pendingStdioRequests.delete(serverId);
1167 | }
1168 | }
1169 |
1170 | reject(
1171 | new McpError(
1172 | ErrorCode.InternalError,
1173 | `Server info request timed out for ${serverId} after ${McpClientService.DEFAULT_REQUEST_TIMEOUT_MS}ms`
1174 | )
1175 | );
1176 | }, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
1177 |
1178 | // Store the request handlers
1179 | if (!this.pendingStdioRequests.has(serverId)) {
1180 | this.pendingStdioRequests.set(serverId, new Map());
1181 | }
1182 | this.pendingStdioRequests.get(serverId)!.set(requestId, {
1183 | resolve: (result) => {
1184 | clearTimeout(timeout);
1185 | resolve(result as Record<string, unknown>);
1186 | },
1187 | reject: (error) => {
1188 | clearTimeout(timeout);
1189 | reject(error);
1190 | },
1191 | });
1192 |
1193 | // Send the request
1194 | const success = this.sendToStdio(serverId, request);
1195 | if (!success) {
1196 | // Clean up the pending request
1197 | const pendingRequests = this.pendingStdioRequests.get(serverId);
1198 | if (pendingRequests) {
1199 | pendingRequests.delete(requestId);
1200 | if (pendingRequests.size === 0) {
1201 | this.pendingStdioRequests.delete(serverId);
1202 | }
1203 | }
1204 | clearTimeout(timeout);
1205 | reject(
1206 | new McpError(
1207 | ErrorCode.InternalError,
1208 | `Failed to send server info request to ${serverId}`
1209 | )
1210 | );
1211 | }
1212 | });
1213 | }
1214 |
1215 | // This should never be reached due to the validateConnectionExists check above
1216 | throw new McpError(
1217 | ErrorCode.InvalidRequest,
1218 | `No connection found with ID ${serverId}`
1219 | );
1220 | }
1221 |
1222 | /**
1223 | * Calls a tool on an MCP server.
1224 | * @param serverId - The ID of the connection to use.
1225 | * @param toolName - The name of the tool to call.
1226 | * @param toolArgs - The arguments to pass to the tool.
1227 | * @returns A promise that resolves to the tool's result.
1228 | * @throws {McpError} Throws an error if the parameters are invalid or the connection doesn't exist.
1229 | */
1230 | public async callTool(
1231 | serverId: string,
1232 | toolName: string,
1233 | toolArgs: Record<string, unknown> | null | undefined
1234 | ): Promise<Record<string, unknown> | Array<unknown>> {
1235 | // Validate serverId
1236 | this.validateServerId(serverId);
1237 |
1238 | // Validate connection exists
1239 | this.validateConnectionExists(serverId);
1240 |
1241 | // Validate toolName
1242 | if (!toolName || typeof toolName !== "string" || toolName.trim() === "") {
1243 | throw new McpError(
1244 | ErrorCode.InvalidParams,
1245 | "Tool name must be a non-empty string"
1246 | );
1247 | }
1248 |
1249 | // Validate toolArgs (ensure it's an object if provided)
1250 | if (
1251 | toolArgs !== null &&
1252 | toolArgs !== undefined &&
1253 | typeof toolArgs !== "object"
1254 | ) {
1255 | throw new McpError(
1256 | ErrorCode.InvalidParams,
1257 | "Tool arguments must be an object, null, or undefined"
1258 | );
1259 | }
1260 |
1261 | // Normalize toolArgs to an empty object if null or undefined
1262 | const normalizedToolArgs: Record<string, unknown> = toolArgs || {};
1263 |
1264 | logger.info(`Calling tool ${toolName} on connection ${serverId}`);
1265 |
1266 | // Check if this is an SSE connection
1267 | if (this.activeSseConnections.has(serverId)) {
1268 | const connection = this.activeSseConnections.get(serverId)!;
1269 | const requestId = uuidv4();
1270 | const request: McpRequest = {
1271 | id: requestId,
1272 | method: "callTool",
1273 | params: { toolName, arguments: normalizedToolArgs },
1274 | };
1275 |
1276 | try {
1277 | // Create URL for the MCP request
1278 | const mcpRequestUrl = new URL(connection.baseUrl);
1279 |
1280 | // Update the connection's last activity timestamp
1281 | connection.lastActivityTimestamp = Date.now();
1282 |
1283 | // Make the request with timeout
1284 | const response = await this.fetchWithTimeout(
1285 | mcpRequestUrl.toString(),
1286 | {
1287 | method: "POST",
1288 | headers: {
1289 | "Content-Type": "application/json",
1290 | },
1291 | body: JSON.stringify(request),
1292 | },
1293 | McpClientService.DEFAULT_REQUEST_TIMEOUT_MS,
1294 | "Request timed out"
1295 | );
1296 |
1297 | if (!response.ok) {
1298 | throw new McpError(
1299 | ErrorCode.InternalError,
1300 | `HTTP error from MCP server: ${response.status} ${response.statusText}`
1301 | );
1302 | }
1303 |
1304 | const mcpResponse = (await response.json()) as McpResponse;
1305 |
1306 | if (mcpResponse.error) {
1307 | throw new SdkMcpError(
1308 | ErrorCode.InternalError,
1309 | `MCP error: ${mcpResponse.error.message} (code: ${mcpResponse.error.code})`,
1310 | mcpResponse.error.data
1311 | );
1312 | }
1313 |
1314 | // Ensure result is either an object or array
1315 | if (
1316 | !mcpResponse.result ||
1317 | (typeof mcpResponse.result !== "object" &&
1318 | !Array.isArray(mcpResponse.result))
1319 | ) {
1320 | throw new McpError(
1321 | ErrorCode.InternalError,
1322 | "Expected object or array result from tool call",
1323 | { receivedType: typeof mcpResponse.result }
1324 | );
1325 | }
1326 |
1327 | return mcpResponse.result;
1328 | } catch (error) {
1329 | logger.error(
1330 | `Error calling tool ${toolName} on SSE connection ${serverId}:`,
1331 | error
1332 | );
1333 |
1334 | // Wrap non-McpError instances
1335 | if (!(error instanceof McpError)) {
1336 | throw new McpError(
1337 | ErrorCode.InternalError,
1338 | `Failed to call tool ${toolName} on connection ${serverId}: ${error instanceof Error ? error.message : String(error)}`
1339 | );
1340 | }
1341 | throw error;
1342 | }
1343 | }
1344 |
1345 | // Check if this is a stdio connection
1346 | else if (this.activeStdioConnections.has(serverId)) {
1347 | const requestId = uuidv4();
1348 | const request: McpRequest = {
1349 | id: requestId,
1350 | method: "callTool",
1351 | params: { toolName, arguments: normalizedToolArgs },
1352 | };
1353 |
1354 | return new Promise((resolve, reject) => {
1355 | // Initialize the map for this connection if it doesn't exist
1356 | if (!this.pendingStdioRequests.has(serverId)) {
1357 | this.pendingStdioRequests.set(serverId, new Map());
1358 | }
1359 |
1360 | // Store the promise resolution functions
1361 | this.pendingStdioRequests.get(serverId)!.set(requestId, {
1362 | resolve: (value) => {
1363 | // Type-safe resolution for tool call results
1364 | resolve(value as Record<string, unknown>);
1365 | },
1366 | reject: reject,
1367 | });
1368 |
1369 | // Set up a timeout to automatically reject this request if it takes too long
1370 | setTimeout(() => {
1371 | // If the request is still pending, reject it
1372 | if (
1373 | this.pendingStdioRequests.has(serverId) &&
1374 | this.pendingStdioRequests.get(serverId)!.has(requestId)
1375 | ) {
1376 | // Get the reject function
1377 | const { reject: rejectRequest } = this.pendingStdioRequests
1378 | .get(serverId)!
1379 | .get(requestId)!;
1380 |
1381 | // Delete the request
1382 | this.pendingStdioRequests.get(serverId)!.delete(requestId);
1383 |
1384 | // If this was the last pending request, clean up the connection map
1385 | if (this.pendingStdioRequests.get(serverId)!.size === 0) {
1386 | this.pendingStdioRequests.delete(serverId);
1387 | }
1388 |
1389 | // Reject the request with a timeout error
1390 | rejectRequest(
1391 | new SdkMcpError(
1392 | ErrorCode.InternalError,
1393 | "Request timed out waiting for response"
1394 | )
1395 | );
1396 | }
1397 | }, McpClientService.DEFAULT_REQUEST_TIMEOUT_MS);
1398 |
1399 | // Send the request
1400 | const sent = this.sendToStdio(serverId, request);
1401 |
1402 | if (!sent) {
1403 | // Clean up the pending request if sending fails
1404 | this.pendingStdioRequests.get(serverId)!.delete(requestId);
1405 |
1406 | // If this was the last pending request, clean up the connection map
1407 | if (this.pendingStdioRequests.get(serverId)!.size === 0) {
1408 | this.pendingStdioRequests.delete(serverId);
1409 | }
1410 |
1411 | reject(
1412 | new Error(`Failed to send request to stdio connection ${serverId}`)
1413 | );
1414 | }
1415 | });
1416 | }
1417 |
1418 | // This should never be reached due to the validateConnectionExists check above
1419 | throw new McpError(
1420 | ErrorCode.InvalidRequest,
1421 | `No connection found with ID ${serverId}`
1422 | );
1423 | }
1424 |
1425 | /**
1426 | * Disconnects from an MCP server.
1427 | * @param serverId - The ID of the connection to close.
1428 | * @returns True if the connection was closed, false if it wasn't found.
1429 | * @throws {McpError} Throws an error if the parameters are invalid.
1430 | */
1431 | public disconnect(serverId: string): boolean {
1432 | // Validate serverId
1433 | this.validateServerId(serverId);
1434 |
1435 | // Check if this is an SSE connection
1436 | if (this.activeSseConnections.has(serverId)) {
1437 | return this.closeSseConnection(serverId);
1438 | }
1439 |
1440 | // Check if this is a stdio connection
1441 | else if (this.activeStdioConnections.has(serverId)) {
1442 | return this.closeStdioConnection(serverId);
1443 | }
1444 |
1445 | // Connection not found
1446 | throw new McpError(
1447 | ErrorCode.InvalidRequest,
1448 | `Connection not found for serverId: ${serverId}`
1449 | );
1450 | }
1451 |
1452 | /**
1453 | * Closes all active connections.
1454 | */
1455 | public closeAllConnections(): void {
1456 | // Close all SSE connections
1457 | for (const id of this.activeSseConnections.keys()) {
1458 | this.closeSseConnection(id);
1459 | }
1460 |
1461 | // Close all stdio connections
1462 | for (const id of this.activeStdioConnections.keys()) {
1463 | this.closeStdioConnection(id);
1464 | }
1465 |
1466 | // Clean up all pending requests
1467 | for (const [, requestsMap] of this.pendingStdioRequests.entries()) {
1468 | for (const [requestId, { reject }] of requestsMap.entries()) {
1469 | logger.warn(
1470 | `Rejecting pending request ${requestId} due to service shutdown`
1471 | );
1472 | reject(new Error("Connection closed due to service shutdown"));
1473 | }
1474 | }
1475 |
1476 | // Clear the pending requests map
1477 | this.pendingStdioRequests.clear();
1478 |
1479 | // Clear the cleanup interval
1480 | if (this.cleanupIntervalId) {
1481 | clearInterval(this.cleanupIntervalId);
1482 | this.cleanupIntervalId = undefined;
1483 | }
1484 |
1485 | logger.info("All MCP connections closed.");
1486 | }
1487 | }
1488 |
```