This is page 4 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/e2e/streamableHttpTransport.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { MCPTestClient } from "./clients/mcp-test-client.js";
3 | import { spawn, ChildProcess } from "node:child_process";
4 |
5 | interface Tool {
6 | name: string;
7 | description?: string;
8 | inputSchema?: unknown;
9 | }
10 |
11 | describe("Streamable HTTP Transport E2E Tests", () => {
12 | let serverProcess: ChildProcess | null = null;
13 | let client: MCPTestClient;
14 | const testPort = 3002;
15 | const baseUrl = `http://localhost:${testPort}`;
16 |
17 | // Store original environment variables
18 | const originalEnv = process.env;
19 |
20 | beforeEach(async () => {
21 | // Set environment variables for the test
22 | process.env = {
23 | ...originalEnv,
24 | MCP_TRANSPORT: "streamable",
25 | MCP_SERVER_PORT: testPort.toString(),
26 | MCP_ENABLE_STREAMING: "true",
27 | MCP_SESSION_TIMEOUT: "60",
28 | GOOGLE_GEMINI_API_KEY:
29 | process.env.GOOGLE_GEMINI_API_KEY || "test-api-key",
30 | GOOGLE_GEMINI_MODEL: "gemini-1.5-flash",
31 | NODE_ENV: "test",
32 | };
33 |
34 | // Start the server
35 | await startServerProcess();
36 |
37 | // Create test client
38 | client = new MCPTestClient(baseUrl);
39 | });
40 |
41 | afterEach(async () => {
42 | // Close client if it has cleanup
43 | if (client && typeof client.close === "function") {
44 | await client.close();
45 | }
46 |
47 | // Stop the server process
48 | if (serverProcess) {
49 | await stopServerProcess();
50 | }
51 |
52 | // Restore environment
53 | process.env = originalEnv;
54 | vi.restoreAllMocks();
55 | });
56 |
57 | async function startServerProcess(): Promise<void> {
58 | return new Promise((resolve, reject) => {
59 | serverProcess = spawn("node", ["dist/server.js"], {
60 | env: process.env,
61 | stdio: "pipe",
62 | });
63 |
64 | let serverReady = false;
65 | const timeout = setTimeout(() => {
66 | if (!serverReady) {
67 | reject(new Error("Server startup timeout"));
68 | }
69 | }, 15000);
70 |
71 | serverProcess!.stdout?.on("data", (data: Buffer) => {
72 | const output = data.toString();
73 | console.log(`Server output: ${output}`);
74 |
75 | if (
76 | output.includes("HTTP server listening") ||
77 | output.includes(`port ${testPort}`) ||
78 | output.includes("MCP Server connected and listening")
79 | ) {
80 | serverReady = true;
81 | clearTimeout(timeout);
82 | // Give server a moment to fully initialize
83 | setTimeout(() => resolve(), 500);
84 | }
85 | });
86 |
87 | serverProcess!.stderr?.on("data", (data: Buffer) => {
88 | console.error(`Server error: ${data.toString()}`);
89 | });
90 |
91 | serverProcess!.on("error", (error) => {
92 | clearTimeout(timeout);
93 | reject(new Error(`Failed to start server: ${error.message}`));
94 | });
95 |
96 | serverProcess!.on("exit", (code, signal) => {
97 | clearTimeout(timeout);
98 | if (!serverReady) {
99 | reject(
100 | new Error(`Server exited early: code ${code}, signal ${signal}`)
101 | );
102 | }
103 | });
104 | });
105 | }
106 |
107 | async function stopServerProcess(): Promise<void> {
108 | if (!serverProcess) return;
109 |
110 | return new Promise((resolve) => {
111 | serverProcess!.on("exit", () => {
112 | serverProcess = null;
113 | resolve();
114 | });
115 |
116 | serverProcess!.kill("SIGTERM");
117 |
118 | // Force kill after timeout
119 | setTimeout(() => {
120 | if (serverProcess) {
121 | serverProcess.kill("SIGKILL");
122 | }
123 | }, 5000);
124 | });
125 | }
126 |
127 | describe("Session Management", () => {
128 | it("should initialize a session and return session ID", async () => {
129 | const result = await client.initialize();
130 |
131 | expect(result).toBeDefined();
132 | expect(result.protocolVersion).toBe("2024-11-05");
133 | expect(result.capabilities).toBeDefined();
134 | expect(client.sessionId).toBeTruthy();
135 | expect(client.sessionId).toMatch(/^[a-f0-9-]{36}$/); // UUID format
136 | });
137 |
138 | it("should maintain session across multiple requests", async () => {
139 | // Initialize session
140 | await client.initialize();
141 | const firstSessionId = client.sessionId;
142 |
143 | // Make another request with same session
144 | const tools = await client.listTools();
145 |
146 | expect(tools).toBeDefined();
147 | expect(client.sessionId).toBe(firstSessionId);
148 | });
149 |
150 | it("should reject requests without valid session", async () => {
151 | // Don't initialize, just try to list tools
152 | await expect(client.listTools()).rejects.toThrow();
153 | });
154 |
155 | it("should handle session expiration gracefully", async () => {
156 | // This test would require waiting for session timeout or mocking time
157 | // For now, we'll just verify the session exists
158 | await client.initialize();
159 | expect(client.sessionId).toBeTruthy();
160 | });
161 | });
162 |
163 | describe("Tool Operations", () => {
164 | beforeEach(async () => {
165 | await client.initialize();
166 | });
167 |
168 | it("should list available tools", async () => {
169 | const result = await client.listTools();
170 |
171 | expect(result).toBeDefined();
172 | expect(result.tools).toBeInstanceOf(Array);
173 | expect(result.tools.length).toBeGreaterThan(0);
174 |
175 | // Check for some expected tools
176 | const toolNames = (result.tools as Tool[]).map((t) => t.name);
177 | expect(toolNames).toContain("gemini_generate_content");
178 | expect(toolNames).toContain("gemini_start_chat");
179 | });
180 |
181 | it("should call a tool successfully", async () => {
182 | const result = await client.callTool("gemini_generate_content", {
183 | prompt: "Say hello in one word",
184 | modelName: "gemini-1.5-flash",
185 | });
186 |
187 | expect(result).toBeDefined();
188 | expect(result.content).toBeDefined();
189 | expect(result.content?.[0]).toBeDefined();
190 | expect(result.content?.[0].text).toBeTruthy();
191 | });
192 |
193 | it("should handle tool errors gracefully", async () => {
194 | await expect(client.callTool("non_existent_tool", {})).rejects.toThrow();
195 | });
196 | });
197 |
198 | describe("SSE Streaming", () => {
199 | beforeEach(async () => {
200 | await client.initialize();
201 | });
202 |
203 | it("should stream content using SSE", async () => {
204 | const chunks: string[] = [];
205 |
206 | const stream = await client.streamTool("gemini_generate_content_stream", {
207 | prompt: "Count from 1 to 3",
208 | modelName: "gemini-1.5-flash",
209 | });
210 |
211 | // Collect chunks from the async iterable
212 | for await (const chunk of stream) {
213 | chunks.push(String(chunk));
214 | }
215 |
216 | expect(chunks.length).toBeGreaterThan(0);
217 | expect(chunks.join("")).toContain("1");
218 | expect(chunks.join("")).toContain("2");
219 | expect(chunks.join("")).toContain("3");
220 | });
221 |
222 | it("should handle SSE connection errors", async () => {
223 | // Test with invalid session
224 | client.sessionId = "invalid-session-id";
225 |
226 | await expect(
227 | client.streamTool("gemini_generate_content_stream", {
228 | prompt: "Test",
229 | modelName: "gemini-1.5-flash",
230 | })
231 | ).rejects.toThrow();
232 | });
233 | });
234 |
235 | describe("Transport Selection", () => {
236 | it("should use streamable transport when configured", async () => {
237 | // The server logs should indicate streamable transport is selected
238 | // This is more of a server configuration test
239 | await client.initialize();
240 |
241 | // If we got here, the streamable transport is working
242 | expect(client.sessionId).toBeTruthy();
243 | });
244 | });
245 |
246 | describe("CORS and Headers", () => {
247 | it("should handle CORS preflight requests", async () => {
248 | const response = await fetch(`${baseUrl}/mcp`, {
249 | method: "OPTIONS",
250 | headers: {
251 | Origin: "http://example.com",
252 | "Access-Control-Request-Method": "POST",
253 | "Access-Control-Request-Headers": "Content-Type, Mcp-Session-Id",
254 | },
255 | });
256 |
257 | expect(response.status).toBe(204);
258 | expect(response.headers.get("Access-Control-Allow-Origin")).toBe("*");
259 | expect(response.headers.get("Access-Control-Allow-Methods")).toContain(
260 | "POST"
261 | );
262 | expect(response.headers.get("Access-Control-Allow-Headers")).toContain(
263 | "Mcp-Session-Id"
264 | );
265 | });
266 |
267 | it("should include proper headers in responses", async () => {
268 | await client.initialize();
269 |
270 | const response = await fetch(`${baseUrl}/mcp`, {
271 | method: "POST",
272 | headers: {
273 | "Content-Type": "application/json",
274 | "Mcp-Session-Id": client.sessionId!,
275 | },
276 | body: JSON.stringify({
277 | jsonrpc: "2.0",
278 | id: 1,
279 | method: "tools/list",
280 | params: {},
281 | }),
282 | });
283 |
284 | expect(response.headers.get("Access-Control-Allow-Origin")).toBe("*");
285 | });
286 | });
287 | });
288 |
```
--------------------------------------------------------------------------------
/tests/unit/services/gemini/geminiImageGeneration.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Test suite for Gemini image generation functionality
3 | * Covers the refactored generateImage method using the correct generateImages API
4 | */
5 |
6 | // Using vitest globals - see vitest.config.ts globals: true
7 | import { GeminiService } from "../../../../src/services/GeminiService.js";
8 | import {
9 | GeminiContentFilterError,
10 | GeminiModelError,
11 | GeminiValidationError,
12 | } from "../../../../src/utils/geminiErrors.js";
13 |
14 | // Mock the GoogleGenAI class before importing
15 | vi.mock("@google/genai", async (importOriginal: any) => {
16 | const actual = await importOriginal();
17 | return {
18 | ...actual,
19 | GoogleGenAI: vi.fn(),
20 | };
21 | });
22 |
23 | // Mock ConfigurationManager singleton
24 | vi.mock("../../../../src/config/ConfigurationManager.js", () => ({
25 | ConfigurationManager: {
26 | getInstance: vi.fn(() => ({
27 | getGeminiServiceConfig: vi.fn(() => ({
28 | apiKey: "test-api-key",
29 | defaultModel: "gemini-2.0-flash-preview",
30 | })),
31 | getModelConfiguration: vi.fn(() => ({
32 | default: "gemini-2.0-flash-preview",
33 | imageGeneration: "imagen-3.0-generate-002",
34 | })),
35 | getGitHubApiToken: vi.fn(() => "test-github-token"),
36 | })),
37 | },
38 | }));
39 |
40 | // Mock ModelSelectionService constructor
41 | vi.mock("../../../../src/services/ModelSelectionService.js", () => ({
42 | ModelSelectionService: vi.fn(() => ({
43 | selectOptimalModel: vi.fn(() => Promise.resolve("imagen-3.0-generate-002")),
44 | })),
45 | }));
46 |
47 | // Mock GitHubApiService constructor
48 | vi.mock("../../../../src/services/gemini/GitHubApiService.js", () => ({
49 | GitHubApiService: vi.fn(() => ({})),
50 | }));
51 |
52 | // Mock the Google GenAI SDK
53 | const mockGenerateImages = vi.fn();
54 | const mockGetGenerativeModel = vi.fn(() => ({
55 | generateImages: mockGenerateImages,
56 | }));
57 |
58 | const mockGenAI = {
59 | getGenerativeModel: mockGetGenerativeModel,
60 | };
61 |
62 | describe("GeminiService - Image Generation", () => {
63 | let service: GeminiService;
64 |
65 | beforeEach(() => {
66 | vi.clearAllMocks();
67 |
68 | // Create service instance (now uses mocked singletons)
69 | service = new GeminiService();
70 |
71 | // Replace the genAI instance with our mock
72 | (service as any).genAI = mockGenAI;
73 | });
74 |
75 | afterEach(() => {
76 | vi.restoreAllMocks();
77 | });
78 |
79 | describe("generateImage", () => {
80 | const mockImageResponse = {
81 | images: [
82 | {
83 | data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==".repeat(
84 | 5
85 | ), // Make it long enough
86 | mimeType: "image/png",
87 | },
88 | ],
89 | promptSafetyMetadata: {
90 | blocked: false,
91 | },
92 | };
93 |
94 | it("should generate images successfully with default parameters", async () => {
95 | mockGenerateImages.mockResolvedValue(mockImageResponse);
96 |
97 | const result = await service.generateImage("A beautiful sunset");
98 |
99 | expect(mockGenerateImages).toHaveBeenCalledWith({
100 | prompt: "A beautiful sunset",
101 | safetySettings: expect.any(Array),
102 | numberOfImages: 1,
103 | });
104 |
105 | expect(result).toEqual({
106 | images: [
107 | {
108 | base64Data:
109 | "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==".repeat(
110 | 5
111 | ),
112 | mimeType: "image/png",
113 | width: 1024,
114 | height: 1024,
115 | },
116 | ],
117 | });
118 | });
119 |
120 | it("should generate images with custom parameters", async () => {
121 | mockGenerateImages.mockResolvedValue({
122 | images: [
123 | {
124 | data: "test-base64-data-1".repeat(10), // Make it long enough
125 | mimeType: "image/png",
126 | },
127 | {
128 | data: "test-base64-data-2".repeat(10), // Make it long enough
129 | mimeType: "image/png",
130 | },
131 | ],
132 | promptSafetyMetadata: { blocked: false },
133 | });
134 |
135 | const result = await service.generateImage(
136 | "A cyberpunk cityscape",
137 | "imagen-3.0-generate-002",
138 | "512x512",
139 | 2,
140 | undefined,
141 | "avoid dark colors",
142 | "photorealistic",
143 | 12345,
144 | 0.8,
145 | true,
146 | false
147 | );
148 |
149 | expect(mockGenerateImages).toHaveBeenCalledWith({
150 | prompt: "A cyberpunk cityscape",
151 | safetySettings: expect.any(Array),
152 | numberOfImages: 2,
153 | width: 512,
154 | height: 512,
155 | negativePrompt: "avoid dark colors",
156 | stylePreset: "photorealistic",
157 | seed: 12345,
158 | styleStrength: 0.8,
159 | });
160 |
161 | expect(result).toEqual({
162 | images: [
163 | {
164 | base64Data: "test-base64-data-1".repeat(10),
165 | mimeType: "image/png",
166 | width: 512,
167 | height: 512,
168 | },
169 | {
170 | base64Data: "test-base64-data-2".repeat(10),
171 | mimeType: "image/png",
172 | width: 512,
173 | height: 512,
174 | },
175 | ],
176 | });
177 | });
178 |
179 | it("should handle safety filtering", async () => {
180 | mockGenerateImages.mockResolvedValue({
181 | images: [],
182 | promptSafetyMetadata: {
183 | blocked: true,
184 | safetyRatings: [
185 | {
186 | category: "HARM_CATEGORY_DANGEROUS_CONTENT",
187 | probability: "HIGH",
188 | },
189 | ],
190 | },
191 | });
192 |
193 | await expect(
194 | service.generateImage("How to make explosives")
195 | ).rejects.toThrow(GeminiContentFilterError);
196 |
197 | expect(mockGenerateImages).toHaveBeenCalled();
198 | });
199 |
200 | it("should handle empty images response", async () => {
201 | mockGenerateImages.mockResolvedValue({
202 | images: [],
203 | promptSafetyMetadata: { blocked: false },
204 | });
205 |
206 | await expect(service.generateImage("A simple drawing")).rejects.toThrow(
207 | GeminiModelError
208 | );
209 | await expect(service.generateImage("A simple drawing")).rejects.toThrow(
210 | "No images were generated by the model"
211 | );
212 | });
213 |
214 | it("should handle missing images in response", async () => {
215 | mockGenerateImages.mockResolvedValue({
216 | promptSafetyMetadata: { blocked: false },
217 | });
218 |
219 | await expect(service.generateImage("A simple drawing")).rejects.toThrow(
220 | GeminiModelError
221 | );
222 | });
223 |
224 | it("should validate generated images", async () => {
225 | mockGenerateImages.mockResolvedValue({
226 | images: [
227 | {
228 | data: "short", // Too short base64 data
229 | mimeType: "image/png",
230 | },
231 | ],
232 | promptSafetyMetadata: { blocked: false },
233 | });
234 |
235 | await expect(service.generateImage("A simple drawing")).rejects.toThrow(
236 | GeminiValidationError
237 | );
238 | });
239 |
240 | it("should handle invalid MIME types", async () => {
241 | mockGenerateImages.mockResolvedValue({
242 | images: [
243 | {
244 | data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
245 | mimeType: "image/gif", // Unsupported MIME type
246 | },
247 | ],
248 | promptSafetyMetadata: { blocked: false },
249 | });
250 |
251 | await expect(service.generateImage("A simple drawing")).rejects.toThrow(
252 | GeminiValidationError
253 | );
254 | });
255 |
256 | it("should use model selection when no specific model provided", async () => {
257 | mockGenerateImages.mockResolvedValue(mockImageResponse);
258 |
259 | // Access the service's model selector and spy on it
260 | const modelSelector = (service as any).modelSelector;
261 | const selectOptimalModelSpy = vi
262 | .spyOn(modelSelector, "selectOptimalModel")
263 | .mockResolvedValue("imagen-3.0-generate-002");
264 |
265 | await service.generateImage("Test prompt");
266 |
267 | expect(selectOptimalModelSpy).toHaveBeenCalledWith({
268 | taskType: "image-generation",
269 | preferQuality: undefined,
270 | preferSpeed: undefined,
271 | fallbackModel: "imagen-3.0-generate-002",
272 | });
273 | });
274 |
275 | it("should handle API errors", async () => {
276 | const apiError = new Error("API quota exceeded");
277 | mockGenerateImages.mockRejectedValue(apiError);
278 |
279 | await expect(service.generateImage("Test prompt")).rejects.toThrow();
280 | });
281 |
282 | it("should handle different resolutions", async () => {
283 | const largeImageResponse = {
284 | images: [
285 | {
286 | data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==".repeat(
287 | 5
288 | ),
289 | mimeType: "image/png",
290 | },
291 | ],
292 | promptSafetyMetadata: {
293 | blocked: false,
294 | },
295 | };
296 |
297 | mockGenerateImages.mockResolvedValue(largeImageResponse);
298 |
299 | // Test 1536x1536 resolution
300 | await service.generateImage("Test", undefined, "1536x1536");
301 |
302 | expect(mockGenerateImages).toHaveBeenCalledWith({
303 | prompt: "Test",
304 | safetySettings: expect.any(Array),
305 | numberOfImages: 1,
306 | width: 1536,
307 | height: 1536,
308 | });
309 | });
310 | });
311 | });
312 |
```
--------------------------------------------------------------------------------
/tests/unit/config/ConfigurationManager.multimodel.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { ConfigurationManager } from "../../../src/config/ConfigurationManager.js";
3 |
4 | describe("ConfigurationManager - Multi-Model Support", () => {
5 | let originalEnv: NodeJS.ProcessEnv;
6 |
7 | beforeEach(() => {
8 | originalEnv = { ...process.env };
9 | ConfigurationManager["instance"] = null;
10 | vi.clearAllMocks();
11 | });
12 |
13 | afterEach(() => {
14 | process.env = originalEnv;
15 | ConfigurationManager["instance"] = null;
16 | });
17 |
18 | describe("Model Array Configuration", () => {
19 | it("should parse GOOGLE_GEMINI_MODELS array", () => {
20 | process.env.NODE_ENV = "test";
21 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
22 | process.env.GOOGLE_GEMINI_MODELS =
23 | '["gemini-1.5-flash", "gemini-1.5-pro"]';
24 |
25 | const manager = ConfigurationManager.getInstance();
26 | const config = manager.getModelConfiguration();
27 |
28 | expect(config.textGeneration).toEqual([
29 | "gemini-1.5-flash",
30 | "gemini-1.5-pro",
31 | ]);
32 | });
33 |
34 | it("should parse GOOGLE_GEMINI_IMAGE_MODELS array", () => {
35 | process.env.NODE_ENV = "test";
36 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
37 | process.env.GOOGLE_GEMINI_IMAGE_MODELS = '["imagen-3.0-generate-002"]';
38 |
39 | const manager = ConfigurationManager.getInstance();
40 | const config = manager.getModelConfiguration();
41 |
42 | expect(config.imageGeneration).toEqual(["imagen-3.0-generate-002"]);
43 | });
44 |
45 | it("should parse GOOGLE_GEMINI_CODE_MODELS array", () => {
46 | process.env.NODE_ENV = "test";
47 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
48 | process.env.GOOGLE_GEMINI_CODE_MODELS =
49 | '["gemini-2.5-pro-preview-05-06", "gemini-2.0-flash"]';
50 |
51 | const manager = ConfigurationManager.getInstance();
52 | const config = manager.getModelConfiguration();
53 |
54 | expect(config.codeReview).toEqual([
55 | "gemini-2.5-pro-preview-05-06",
56 | "gemini-2.0-flash",
57 | ]);
58 | });
59 |
60 | it("should handle invalid JSON gracefully", () => {
61 | process.env.NODE_ENV = "test";
62 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
63 | process.env.GOOGLE_GEMINI_MODELS = "invalid-json";
64 |
65 | const manager = ConfigurationManager.getInstance();
66 | const config = manager.getModelConfiguration();
67 |
68 | expect(config.textGeneration).toEqual(["gemini-2.5-flash-preview-05-20"]);
69 | });
70 |
71 | it("should handle non-array JSON gracefully", () => {
72 | process.env.NODE_ENV = "test";
73 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
74 | process.env.GOOGLE_GEMINI_MODELS = '{"not": "array"}';
75 |
76 | const manager = ConfigurationManager.getInstance();
77 | const config = manager.getModelConfiguration();
78 |
79 | expect(config.textGeneration).toEqual(["gemini-2.5-flash-preview-05-20"]);
80 | });
81 | });
82 |
83 | describe("Routing Preferences", () => {
84 | it("should parse routing preferences correctly", () => {
85 | process.env.NODE_ENV = "test";
86 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
87 | process.env.GOOGLE_GEMINI_ROUTING_PREFER_COST = "true";
88 | process.env.GOOGLE_GEMINI_ROUTING_PREFER_SPEED = "false";
89 | process.env.GOOGLE_GEMINI_ROUTING_PREFER_QUALITY = "true";
90 |
91 | const manager = ConfigurationManager.getInstance();
92 | const config = manager.getModelConfiguration();
93 |
94 | expect(config.routing.preferCostEffective).toBe(true);
95 | expect(config.routing.preferSpeed).toBe(false);
96 | expect(config.routing.preferQuality).toBe(true);
97 | });
98 |
99 | it("should default to quality preference when none specified", () => {
100 | process.env.NODE_ENV = "test";
101 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
102 |
103 | const manager = ConfigurationManager.getInstance();
104 | const config = manager.getModelConfiguration();
105 |
106 | expect(config.routing.preferCostEffective).toBe(false);
107 | expect(config.routing.preferSpeed).toBe(false);
108 | expect(config.routing.preferQuality).toBe(true);
109 | });
110 | });
111 |
112 | describe("Model Capabilities", () => {
113 | it("should provide correct capabilities for gemini-2.5-pro", () => {
114 | process.env.NODE_ENV = "test";
115 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
116 |
117 | const manager = ConfigurationManager.getInstance();
118 | const config = manager.getModelConfiguration();
119 | const capabilities = config.capabilities["gemini-2.5-pro-preview-05-06"];
120 |
121 | expect(capabilities).toBeDefined();
122 | expect(capabilities.textGeneration).toBe(true);
123 | expect(capabilities.imageInput).toBe(true);
124 | expect(capabilities.codeExecution).toBe("excellent");
125 | expect(capabilities.complexReasoning).toBe("excellent");
126 | expect(capabilities.costTier).toBe("high");
127 | expect(capabilities.contextWindow).toBe(1048576);
128 | });
129 |
130 | it("should provide correct capabilities for imagen model", () => {
131 | process.env.NODE_ENV = "test";
132 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
133 |
134 | const manager = ConfigurationManager.getInstance();
135 | const config = manager.getModelConfiguration();
136 | const capabilities = config.capabilities["imagen-3.0-generate-002"];
137 |
138 | expect(capabilities).toBeDefined();
139 | expect(capabilities.textGeneration).toBe(false);
140 | expect(capabilities.imageGeneration).toBe(true);
141 | expect(capabilities.codeExecution).toBe("none");
142 | expect(capabilities.complexReasoning).toBe("none");
143 | });
144 | });
145 |
146 | describe("Default Model Selection", () => {
147 | it("should use GOOGLE_GEMINI_DEFAULT_MODEL when provided", () => {
148 | process.env.NODE_ENV = "test";
149 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
150 | process.env.GOOGLE_GEMINI_DEFAULT_MODEL = "gemini-2.5-pro-preview-05-06";
151 |
152 | const manager = ConfigurationManager.getInstance();
153 | const config = manager.getModelConfiguration();
154 |
155 | expect(config.default).toBe("gemini-2.5-pro-preview-05-06");
156 | });
157 |
158 | it("should fallback to first text generation model when default not specified", () => {
159 | process.env.NODE_ENV = "test";
160 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
161 | process.env.GOOGLE_GEMINI_MODELS =
162 | '["gemini-1.5-pro", "gemini-1.5-flash"]';
163 |
164 | const manager = ConfigurationManager.getInstance();
165 | const config = manager.getModelConfiguration();
166 |
167 | expect(config.default).toBe("gemini-1.5-pro");
168 | });
169 | });
170 |
171 | describe("Complex Reasoning Models", () => {
172 | it("should filter high reasoning models correctly", () => {
173 | process.env.NODE_ENV = "test";
174 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
175 | process.env.GOOGLE_GEMINI_MODELS =
176 | '["gemini-2.5-pro-preview-05-06", "gemini-1.5-flash", "gemini-1.5-pro"]';
177 |
178 | const manager = ConfigurationManager.getInstance();
179 | const config = manager.getModelConfiguration();
180 |
181 | expect(config.complexReasoning).toContain("gemini-2.5-pro-preview-05-06");
182 | expect(config.complexReasoning).toContain("gemini-1.5-pro");
183 | expect(config.complexReasoning).not.toContain("gemini-1.5-flash");
184 | });
185 | });
186 |
187 | describe("Backward Compatibility", () => {
188 | it("should migrate single model to array format", () => {
189 | process.env.NODE_ENV = "test";
190 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
191 | process.env.GOOGLE_GEMINI_MODEL = "gemini-1.5-pro";
192 |
193 | const manager = ConfigurationManager.getInstance();
194 | const config = manager.getModelConfiguration();
195 |
196 | expect(config.textGeneration).toContain("gemini-1.5-pro");
197 | });
198 |
199 | it("should use old GOOGLE_GEMINI_MODEL as fallback", () => {
200 | process.env.NODE_ENV = "test";
201 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
202 | process.env.GOOGLE_GEMINI_MODEL = "gemini-1.5-pro";
203 | delete process.env.GOOGLE_GEMINI_MODELS;
204 |
205 | const manager = ConfigurationManager.getInstance();
206 | const config = manager.getModelConfiguration();
207 |
208 | expect(config.textGeneration).toEqual(["gemini-1.5-pro"]);
209 | });
210 | });
211 |
212 | describe("Environment Variable Validation", () => {
213 | it("should handle missing environment variables gracefully", () => {
214 | process.env.NODE_ENV = "test";
215 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
216 |
217 | const manager = ConfigurationManager.getInstance();
218 | const config = manager.getModelConfiguration();
219 |
220 | expect(config.default).toBeDefined();
221 | expect(config.textGeneration).toBeDefined();
222 | expect(config.imageGeneration).toBeDefined();
223 | expect(config.codeReview).toBeDefined();
224 | });
225 |
226 | it("should provide sensible defaults for image models", () => {
227 | process.env.NODE_ENV = "test";
228 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
229 |
230 | const manager = ConfigurationManager.getInstance();
231 | const config = manager.getModelConfiguration();
232 |
233 | expect(config.imageGeneration).toContain("imagen-3.0-generate-002");
234 | });
235 |
236 | it("should provide sensible defaults for code models", () => {
237 | process.env.NODE_ENV = "test";
238 | process.env.GOOGLE_GEMINI_API_KEY = "test-key";
239 |
240 | const manager = ConfigurationManager.getInstance();
241 | const config = manager.getModelConfiguration();
242 |
243 | expect(config.codeReview).toContain("gemini-2.5-pro-preview-05-06");
244 | });
245 | });
246 | });
247 |
```
--------------------------------------------------------------------------------
/tests/unit/tools/mcpClientTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { mcpClientTool } from "../../../src/tools/mcpClientTool.js";
3 | import { ConfigurationManager } from "../../../src/config/ConfigurationManager.js";
4 | import * as writeToFileModule from "../../../src/tools/writeToFileTool.js";
5 |
6 | // Mock dependencies
7 | vi.mock("../../../src/services/index.js");
8 | vi.mock("../../../src/config/ConfigurationManager.js");
9 | vi.mock("../../../src/tools/writeToFileTool.js");
10 | vi.mock("uuid", () => ({
11 | v4: () => "test-uuid-123",
12 | }));
13 |
14 | describe("mcpClientTool", () => {
15 | // eslint-disable-next-line @typescript-eslint/no-explicit-any
16 | let mockMcpClientService: any;
17 | // eslint-disable-next-line @typescript-eslint/no-explicit-any
18 | let mockConfigManager: any;
19 |
20 | beforeEach(() => {
21 | vi.clearAllMocks();
22 |
23 | // Setup mock McpClientService
24 | mockMcpClientService = {
25 | connect: vi.fn(),
26 | disconnect: vi.fn(),
27 | listTools: vi.fn(),
28 | callTool: vi.fn(),
29 | getServerInfo: vi.fn(),
30 | };
31 |
32 | // Setup mock ConfigurationManager
33 | mockConfigManager = {
34 | getMcpConfig: vi.fn().mockReturnValue({
35 | clientId: "default-client-id",
36 | connectionToken: "default-token",
37 | }),
38 | };
39 |
40 | vi.mocked(ConfigurationManager.getInstance).mockReturnValue(
41 | mockConfigManager
42 | );
43 | });
44 |
45 | describe("Tool Configuration", () => {
46 | it("should have correct name and description", () => {
47 | expect(mcpClientTool.name).toBe("mcp_client");
48 | expect(mcpClientTool.description).toContain(
49 | "Manages MCP (Model Context Protocol) client connections"
50 | );
51 | });
52 |
53 | it("should have valid input schema", () => {
54 | expect(mcpClientTool.inputSchema).toBeDefined();
55 | expect(mcpClientTool.inputSchema._def.discriminator).toBe("operation");
56 | });
57 | });
58 |
59 | describe("Connect Operation", () => {
60 | it("should handle stdio connection", async () => {
61 | const mockServerInfo = { name: "Test Server", version: "1.0.0" };
62 | mockMcpClientService.connect.mockResolvedValue("connection-123");
63 | mockMcpClientService.getServerInfo.mockResolvedValue(mockServerInfo);
64 |
65 | const args = {
66 | operation: "connect_stdio" as const,
67 | transport: "stdio" as const,
68 | command: "node",
69 | args: ["server.js"],
70 | clientId: "custom-client-id",
71 | };
72 |
73 | const result = await mcpClientTool.execute(args, mockMcpClientService);
74 |
75 | expect(mockMcpClientService.connect).toHaveBeenCalledWith(
76 | "test-uuid-123",
77 | {
78 | type: "stdio",
79 | connectionToken: "default-token",
80 | stdioCommand: "node",
81 | stdioArgs: ["server.js"],
82 | }
83 | );
84 |
85 | expect(result.content[0].text).toBe(
86 | "Successfully connected to MCP server"
87 | );
88 | const resultData = JSON.parse(result.content[1].text);
89 | expect(resultData.connectionId).toBe("connection-123");
90 | expect(resultData.transport).toBe("stdio");
91 | expect(resultData.serverInfo).toEqual(mockServerInfo);
92 | });
93 |
94 | it("should handle SSE connection", async () => {
95 | const mockServerInfo = { name: "SSE Server", version: "2.0.0" };
96 | mockMcpClientService.connect.mockResolvedValue("sse-connection-456");
97 | mockMcpClientService.getServerInfo.mockResolvedValue(mockServerInfo);
98 |
99 | const args = {
100 | operation: "connect_sse" as const,
101 | transport: "sse" as const,
102 | url: "https://mcp-server.example.com/sse",
103 | connectionToken: "custom-token",
104 | };
105 |
106 | const result = await mcpClientTool.execute(args, mockMcpClientService);
107 |
108 | expect(mockMcpClientService.connect).toHaveBeenCalledWith(
109 | "test-uuid-123",
110 | {
111 | type: "sse",
112 | connectionToken: "custom-token",
113 | sseUrl: "https://mcp-server.example.com/sse",
114 | }
115 | );
116 |
117 | const resultData = JSON.parse(result.content[1].text);
118 | expect(resultData.connectionId).toBe("sse-connection-456");
119 | expect(resultData.transport).toBe("sse");
120 | });
121 | });
122 |
123 | describe("Disconnect Operation", () => {
124 | it("should handle disconnection", async () => {
125 | mockMcpClientService.disconnect.mockResolvedValue(undefined);
126 |
127 | const args = {
128 | operation: "disconnect" as const,
129 | connectionId: "connection-123",
130 | };
131 |
132 | const result = await mcpClientTool.execute(args, mockMcpClientService);
133 |
134 | expect(mockMcpClientService.disconnect).toHaveBeenCalledWith(
135 | "connection-123"
136 | );
137 | expect(result.content[0].text).toBe(
138 | "Successfully disconnected from MCP server"
139 | );
140 |
141 | const resultData = JSON.parse(result.content[1].text);
142 | expect(resultData.connectionId).toBe("connection-123");
143 | expect(resultData.status).toBe("disconnected");
144 | });
145 | });
146 |
147 | describe("List Tools Operation", () => {
148 | it("should list available tools", async () => {
149 | const mockTools = [
150 | { name: "tool1", description: "First tool" },
151 | { name: "tool2", description: "Second tool" },
152 | ];
153 | mockMcpClientService.listTools.mockResolvedValue(mockTools);
154 |
155 | const args = {
156 | operation: "list_tools" as const,
157 | connectionId: "connection-123",
158 | };
159 |
160 | const result = await mcpClientTool.execute(args, mockMcpClientService);
161 |
162 | expect(mockMcpClientService.listTools).toHaveBeenCalledWith(
163 | "connection-123"
164 | );
165 | expect(result.content[0].text).toContain("Available tools on connection");
166 |
167 | const toolsData = JSON.parse(result.content[1].text);
168 | expect(toolsData).toEqual(mockTools);
169 | });
170 | });
171 |
172 | describe("Call Tool Operation", () => {
173 | it("should call tool and return result", async () => {
174 | const mockResult = { status: "success", data: "Tool executed" };
175 | mockMcpClientService.callTool.mockResolvedValue(mockResult);
176 |
177 | const args = {
178 | operation: "call_tool" as const,
179 | connectionId: "connection-123",
180 | toolName: "exampleTool",
181 | toolParameters: { param1: "value1" },
182 | overwriteFile: true,
183 | };
184 |
185 | const result = await mcpClientTool.execute(args, mockMcpClientService);
186 |
187 | expect(mockMcpClientService.callTool).toHaveBeenCalledWith(
188 | "connection-123",
189 | "exampleTool",
190 | { param1: "value1" }
191 | );
192 |
193 | const resultData = JSON.parse(result.content[0].text);
194 | expect(resultData).toEqual(mockResult);
195 | });
196 |
197 | it("should write tool result to file when outputFilePath is provided", async () => {
198 | const mockResult = { status: "success", data: "Tool executed" };
199 | mockMcpClientService.callTool.mockResolvedValue(mockResult);
200 |
201 | vi.mocked(writeToFileModule.writeToFile.execute).mockResolvedValue({
202 | content: [{ type: "text", text: "File written" }],
203 | });
204 |
205 | const args = {
206 | operation: "call_tool" as const,
207 | connectionId: "connection-123",
208 | toolName: "exampleTool",
209 | toolParameters: {},
210 | outputFilePath: "/path/to/output.json",
211 | overwriteFile: true,
212 | };
213 |
214 | const result = await mcpClientTool.execute(args, mockMcpClientService);
215 |
216 | expect(mockMcpClientService.callTool).toHaveBeenCalled();
217 | expect(writeToFileModule.writeToFile.execute).toHaveBeenCalledWith({
218 | filePath: "/path/to/output.json",
219 | content: JSON.stringify(mockResult, null, 2),
220 | overwriteIfExists: true,
221 | });
222 |
223 | expect(result.content[0].text).toContain(
224 | "Tool exampleTool executed successfully"
225 | );
226 | expect(result.content[0].text).toContain("/path/to/output.json");
227 | });
228 |
229 | it("should handle string results", async () => {
230 | mockMcpClientService.callTool.mockResolvedValue("Simple string result");
231 |
232 | const args = {
233 | operation: "call_tool" as const,
234 | connectionId: "connection-123",
235 | toolName: "stringTool",
236 | toolParameters: {},
237 | overwriteFile: true,
238 | };
239 |
240 | const result = await mcpClientTool.execute(args, mockMcpClientService);
241 |
242 | expect(result.content[0].text).toBe("Simple string result");
243 | });
244 |
245 | it("should handle tool call errors", async () => {
246 | mockMcpClientService.callTool.mockRejectedValue(
247 | new Error("Tool not found")
248 | );
249 |
250 | const args = {
251 | operation: "call_tool" as const,
252 | connectionId: "connection-123",
253 | toolName: "nonExistentTool",
254 | toolParameters: {},
255 | overwriteFile: true,
256 | };
257 |
258 | await expect(
259 | mcpClientTool.execute(args, mockMcpClientService)
260 | ).rejects.toThrow();
261 | });
262 | });
263 |
264 | describe("Error Handling", () => {
265 | it("should handle connection errors", async () => {
266 | mockMcpClientService.connect.mockRejectedValue(
267 | new Error("Connection failed")
268 | );
269 |
270 | const args = {
271 | operation: "connect_stdio" as const,
272 | transport: "stdio" as const,
273 | command: "invalid-command",
274 | };
275 |
276 | await expect(
277 | mcpClientTool.execute(args, mockMcpClientService)
278 | ).rejects.toThrow();
279 | });
280 |
281 | it("should handle unknown operation", async () => {
282 | const args = {
283 | operation: "unknown",
284 | connectionId: "test-connection",
285 | toolName: "test-tool",
286 | overwriteFile: true,
287 | } as const;
288 |
289 | await expect(
290 | // eslint-disable-next-line @typescript-eslint/no-explicit-any
291 | mcpClientTool.execute(args as any, mockMcpClientService)
292 | ).rejects.toThrow("Unknown operation");
293 | });
294 | });
295 | });
296 |
```
--------------------------------------------------------------------------------
/tests/unit/tools/geminiGenerateImageTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Test suite for geminiGenerateImageTool
3 | * Tests the tool integration with the fixed GeminiService.generateImage method
4 | */
5 |
6 | // Using vitest globals - see vitest.config.ts globals: true
7 | import type { GeminiService } from "../../../src/services/GeminiService.js";
8 | import { geminiGenerateImageTool } from "../../../src/tools/geminiGenerateImageTool.js";
9 | import type { ImageGenerationResult } from "../../../src/types/geminiServiceTypes.js";
10 |
11 | describe("geminiGenerateImageTool", () => {
12 | let mockGeminiService: GeminiService;
13 |
14 | beforeEach(() => {
15 | mockGeminiService = {
16 | generateImage: vi.fn(),
17 | } as unknown as GeminiService;
18 | });
19 |
20 | describe("successful image generation", () => {
21 | it("should generate images with minimal parameters", async () => {
22 | const mockResult: ImageGenerationResult = {
23 | images: [
24 | {
25 | base64Data: "test-base64-data",
26 | mimeType: "image/png",
27 | width: 1024,
28 | height: 1024,
29 | },
30 | ],
31 | };
32 |
33 | (mockGeminiService.generateImage as any).mockResolvedValue(mockResult);
34 |
35 | const result = await geminiGenerateImageTool.execute(
36 | {
37 | prompt: "A beautiful landscape",
38 | },
39 | mockGeminiService
40 | );
41 |
42 | expect(mockGeminiService.generateImage).toHaveBeenCalledWith(
43 | "A beautiful landscape",
44 | undefined, // modelName
45 | undefined, // resolution
46 | undefined, // numberOfImages
47 | undefined, // safetySettings
48 | undefined, // negativePrompt
49 | undefined, // stylePreset
50 | undefined, // seed
51 | undefined, // styleStrength
52 | undefined, // preferQuality
53 | undefined // preferSpeed
54 | );
55 |
56 | expect(result).toEqual({
57 | content: [
58 | {
59 | type: "text",
60 | text: expect.stringContaining("Generated 1"),
61 | },
62 | {
63 | type: "image",
64 | mimeType: "image/png",
65 | data: "test-base64-data",
66 | },
67 | ],
68 | });
69 | });
70 |
71 | it("should generate images with all parameters", async () => {
72 | const mockResult: ImageGenerationResult = {
73 | images: [
74 | {
75 | base64Data: "test-base64-data-1",
76 | mimeType: "image/png",
77 | width: 512,
78 | height: 512,
79 | },
80 | {
81 | base64Data: "test-base64-data-2",
82 | mimeType: "image/png",
83 | width: 512,
84 | height: 512,
85 | },
86 | ],
87 | };
88 |
89 | (mockGeminiService.generateImage as any).mockResolvedValue(mockResult);
90 |
91 | const result = await geminiGenerateImageTool.execute(
92 | {
93 | prompt: "A cyberpunk city",
94 | modelName: "imagen-3.0-generate-002",
95 | resolution: "512x512",
96 | numberOfImages: 2,
97 | negativePrompt: "blurry, low quality",
98 | stylePreset: "photographic",
99 | seed: 12345,
100 | styleStrength: 0.8,
101 | modelPreferences: {
102 | preferQuality: true,
103 | preferSpeed: false,
104 | },
105 | },
106 | mockGeminiService
107 | );
108 |
109 | expect(mockGeminiService.generateImage).toHaveBeenCalledWith(
110 | "A cyberpunk city",
111 | "imagen-3.0-generate-002",
112 | "512x512",
113 | 2,
114 | undefined, // safetySettings
115 | "blurry, low quality",
116 | "photorealistic",
117 | 12345,
118 | 0.8,
119 | true,
120 | false
121 | );
122 |
123 | expect(result).toEqual({
124 | content: [
125 | {
126 | type: "text",
127 | text: expect.stringContaining("Generated 2"),
128 | },
129 | {
130 | type: "image",
131 | mimeType: "image/png",
132 | data: "test-base64-data-1",
133 | },
134 | {
135 | type: "image",
136 | mimeType: "image/png",
137 | data: "test-base64-data-2",
138 | },
139 | ],
140 | });
141 | });
142 |
143 | it("should handle safety settings", async () => {
144 | const mockResult: ImageGenerationResult = {
145 | images: [
146 | {
147 | base64Data: "test-base64-data-safety",
148 | mimeType: "image/png",
149 | width: 1024,
150 | height: 1024,
151 | },
152 | ],
153 | };
154 |
155 | (mockGeminiService.generateImage as any).mockResolvedValue(mockResult);
156 |
157 | const safetySettings = [
158 | {
159 | category: "HARM_CATEGORY_DANGEROUS_CONTENT" as const,
160 | threshold: "BLOCK_LOW_AND_ABOVE" as const,
161 | },
162 | ];
163 |
164 | const result = await geminiGenerateImageTool.execute(
165 | {
166 | prompt: "A safe image",
167 | safetySettings,
168 | },
169 | mockGeminiService
170 | );
171 |
172 | expect(mockGeminiService.generateImage).toHaveBeenCalledWith(
173 | "A safe image",
174 | undefined, // modelName
175 | undefined, // resolution
176 | undefined, // numberOfImages
177 | safetySettings,
178 | undefined, // negativePrompt
179 | undefined, // stylePreset
180 | undefined, // seed
181 | undefined, // styleStrength
182 | undefined, // preferQuality
183 | undefined // preferSpeed
184 | );
185 |
186 | expect(result).toEqual({
187 | content: [
188 | {
189 | type: "text",
190 | text: expect.stringContaining("Generated 1"),
191 | },
192 | {
193 | type: "image",
194 | mimeType: "image/png",
195 | data: "test-base64-data-safety",
196 | },
197 | ],
198 | });
199 | });
200 | });
201 |
202 | describe("error handling", () => {
203 | it("should handle content filter errors", async () => {
204 | const { GeminiContentFilterError } = await import(
205 | "../../../src/utils/geminiErrors.js"
206 | );
207 |
208 | (mockGeminiService.generateImage as any).mockRejectedValue(
209 | new GeminiContentFilterError("Content blocked by safety filters", [
210 | "HARM_CATEGORY_DANGEROUS_CONTENT",
211 | ])
212 | );
213 |
214 | const result = await geminiGenerateImageTool.execute(
215 | {
216 | prompt: "Inappropriate content",
217 | },
218 | mockGeminiService
219 | );
220 |
221 | expect(result).toEqual({
222 | content: [
223 | {
224 | type: "text",
225 | text: expect.stringContaining("Content blocked by safety filters"),
226 | },
227 | ],
228 | isError: true,
229 | });
230 | });
231 |
232 | it("should handle validation errors", async () => {
233 | const { GeminiValidationError } = await import(
234 | "../../../src/utils/geminiErrors.js"
235 | );
236 |
237 | (mockGeminiService.generateImage as any).mockRejectedValue(
238 | new GeminiValidationError("Invalid prompt", "prompt")
239 | );
240 |
241 | const result = await geminiGenerateImageTool.execute(
242 | {
243 | prompt: "",
244 | },
245 | mockGeminiService
246 | );
247 |
248 | expect(result).toEqual({
249 | content: [
250 | {
251 | type: "text",
252 | text: expect.stringContaining("Invalid prompt"),
253 | },
254 | ],
255 | isError: true,
256 | });
257 | });
258 |
259 | it("should handle model errors", async () => {
260 | const { GeminiModelError } = await import(
261 | "../../../src/utils/geminiErrors.js"
262 | );
263 |
264 | (mockGeminiService.generateImage as any).mockRejectedValue(
265 | new GeminiModelError("Model unavailable", "imagen-3.0-generate-002")
266 | );
267 |
268 | const result = await geminiGenerateImageTool.execute(
269 | {
270 | prompt: "A test image",
271 | modelName: "imagen-3.0-generate-002",
272 | },
273 | mockGeminiService
274 | );
275 |
276 | expect(result).toEqual({
277 | content: [
278 | {
279 | type: "text",
280 | text: expect.stringContaining("Model unavailable"),
281 | },
282 | ],
283 | isError: true,
284 | });
285 | });
286 |
287 | it("should handle generic errors", async () => {
288 | (mockGeminiService.generateImage as any).mockRejectedValue(
289 | new Error("Network error")
290 | );
291 |
292 | const result = await geminiGenerateImageTool.execute(
293 | {
294 | prompt: "A test image",
295 | },
296 | mockGeminiService
297 | );
298 |
299 | expect(result).toEqual({
300 | content: [
301 | {
302 | type: "text",
303 | text: expect.stringContaining("Network error"),
304 | },
305 | ],
306 | isError: true,
307 | });
308 | });
309 | });
310 |
311 | describe("parameter validation", () => {
312 | it("should validate prompt is required", async () => {
313 | const result = await geminiGenerateImageTool.execute(
314 | {} as any,
315 | mockGeminiService
316 | );
317 |
318 | expect(result).toEqual({
319 | content: [
320 | {
321 | type: "text",
322 | text: expect.stringContaining("prompt"),
323 | },
324 | ],
325 | isError: true,
326 | });
327 | });
328 |
329 | it("should validate numberOfImages range", async () => {
330 | const result = await geminiGenerateImageTool.execute(
331 | {
332 | prompt: "Test",
333 | numberOfImages: 10, // Exceeds maximum
334 | },
335 | mockGeminiService
336 | );
337 |
338 | expect(result).toEqual({
339 | content: [
340 | {
341 | type: "text",
342 | text: expect.stringContaining("numberOfImages"),
343 | },
344 | ],
345 | isError: true,
346 | });
347 | });
348 |
349 | it("should validate resolution format", async () => {
350 | const result = await geminiGenerateImageTool.execute(
351 | {
352 | prompt: "Test",
353 | resolution: "800x600" as any, // Invalid resolution
354 | },
355 | mockGeminiService
356 | );
357 |
358 | expect(result).toEqual({
359 | content: [
360 | {
361 | type: "text",
362 | text: expect.stringContaining("resolution"),
363 | },
364 | ],
365 | isError: true,
366 | });
367 | });
368 | });
369 | });
370 |
```
--------------------------------------------------------------------------------
/src/tools/geminiChatTool.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2 | import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3 | import { z } from "zod";
4 | import {
5 | GEMINI_CHAT_TOOL_NAME,
6 | GEMINI_CHAT_TOOL_DESCRIPTION,
7 | GEMINI_CHAT_PARAMS,
8 | } from "./geminiChatParams.js";
9 | import { GeminiService } from "../services/index.js";
10 | import { logger } from "../utils/index.js";
11 | import { mapAnyErrorToMcpError } from "../utils/errors.js";
12 | import { BlockedReason, FinishReason } from "@google/genai";
13 | import type {
14 | Content,
15 | GenerationConfig,
16 | SafetySetting,
17 | Tool,
18 | ToolConfig,
19 | GenerateContentResponse,
20 | } from "@google/genai";
21 |
22 | // Define the type for the arguments object based on the Zod schema
23 | type GeminiChatArgs = z.infer<z.ZodObject<typeof GEMINI_CHAT_PARAMS>>;
24 |
25 | /**
26 | * Registers the gemini_chat tool with the MCP server.
27 | * This consolidated tool handles chat session management including starting sessions,
28 | * sending messages, and sending function results.
29 | *
30 | * @param server - The McpServer instance.
31 | * @param serviceInstance - An instance of the GeminiService.
32 | */
33 | export const geminiChatTool = (
34 | server: McpServer,
35 | serviceInstance: GeminiService
36 | ): void => {
37 | /**
38 | * Processes the request for the gemini_chat tool.
39 | * @param args - The arguments object matching GEMINI_CHAT_PARAMS.
40 | * @returns The result content for MCP.
41 | */
42 | const processRequest = async (args: unknown): Promise<CallToolResult> => {
43 | const typedArgs = args as GeminiChatArgs;
44 | logger.debug(`Received ${GEMINI_CHAT_TOOL_NAME} request:`, {
45 | operation: typedArgs.operation,
46 | sessionId: typedArgs.sessionId,
47 | modelName: typedArgs.modelName,
48 | });
49 |
50 | try {
51 | // Validate required fields based on operation
52 | if (
53 | typedArgs.operation === "send_message" ||
54 | typedArgs.operation === "send_function_result"
55 | ) {
56 | if (!typedArgs.sessionId) {
57 | throw new Error(
58 | `sessionId is required for operation '${typedArgs.operation}'`
59 | );
60 | }
61 | }
62 |
63 | if (typedArgs.operation === "send_message" && !typedArgs.message) {
64 | throw new Error("message is required for operation 'send_message'");
65 | }
66 |
67 | if (
68 | typedArgs.operation === "send_function_result" &&
69 | !typedArgs.functionResponses
70 | ) {
71 | throw new Error(
72 | "functionResponses is required for operation 'send_function_result'"
73 | );
74 | }
75 |
76 | // Handle different operations
77 | switch (typedArgs.operation) {
78 | case "start": {
79 | // Start a new chat session
80 | const sessionId = serviceInstance.startChatSession({
81 | modelName: typedArgs.modelName,
82 | history: typedArgs.history as Content[] | undefined,
83 | generationConfig: typedArgs.generationConfig as
84 | | GenerationConfig
85 | | undefined,
86 | safetySettings: typedArgs.safetySettings as
87 | | SafetySetting[]
88 | | undefined,
89 | tools: typedArgs.tools as Tool[] | undefined,
90 | systemInstruction: typedArgs.systemInstruction,
91 | cachedContentName: typedArgs.cachedContentName,
92 | });
93 |
94 | logger.info(
95 | `Successfully started chat session ${sessionId} for model ${typedArgs.modelName || "default"}`
96 | );
97 |
98 | return {
99 | content: [
100 | {
101 | type: "text" as const,
102 | text: JSON.stringify({ sessionId }),
103 | },
104 | ],
105 | };
106 | }
107 |
108 | case "send_message": {
109 | // Send a message to an existing chat session
110 | const response: GenerateContentResponse =
111 | await serviceInstance.sendMessageToSession({
112 | sessionId: typedArgs.sessionId!,
113 | message: typedArgs.message!,
114 | generationConfig: typedArgs.generationConfig as
115 | | GenerationConfig
116 | | undefined,
117 | safetySettings: typedArgs.safetySettings as
118 | | SafetySetting[]
119 | | undefined,
120 | tools: typedArgs.tools as Tool[] | undefined,
121 | toolConfig: typedArgs.toolConfig as ToolConfig | undefined,
122 | cachedContentName: typedArgs.cachedContentName,
123 | });
124 |
125 | // Process the response
126 | return processGenerateContentResponse(response, typedArgs.sessionId!);
127 | }
128 |
129 | case "send_function_result": {
130 | // Send function results to an existing chat session
131 | // Note: The service expects a string, so we stringify the array of function responses
132 | const response: GenerateContentResponse =
133 | await serviceInstance.sendFunctionResultToSession({
134 | sessionId: typedArgs.sessionId!,
135 | functionResponse: JSON.stringify(typedArgs.functionResponses),
136 | functionCall: undefined, // Could be enhanced to pass original function call
137 | });
138 |
139 | // Process the response
140 | return processGenerateContentResponse(
141 | response,
142 | typedArgs.sessionId!,
143 | true
144 | );
145 | }
146 |
147 | default:
148 | throw new Error(`Invalid operation: ${typedArgs.operation}`);
149 | }
150 | } catch (error: unknown) {
151 | logger.error(`Error processing ${GEMINI_CHAT_TOOL_NAME}:`, error);
152 | throw mapAnyErrorToMcpError(error, GEMINI_CHAT_TOOL_NAME);
153 | }
154 | };
155 |
156 | /**
157 | * Helper function to process GenerateContentResponse into MCP format
158 | */
159 | function processGenerateContentResponse(
160 | response: GenerateContentResponse,
161 | sessionId: string,
162 | isFunctionResult: boolean = false
163 | ): CallToolResult {
164 | const context = isFunctionResult ? "after function result" : "";
165 |
166 | // Check for prompt safety blocks
167 | if (response.promptFeedback?.blockReason === BlockedReason.SAFETY) {
168 | logger.warn(
169 | `Gemini prompt blocked due to SAFETY for session ${sessionId} ${context}.`
170 | );
171 | return {
172 | content: [
173 | {
174 | type: "text",
175 | text: `Error: Prompt blocked due to safety settings ${context}. Reason: ${response.promptFeedback.blockReason}`,
176 | },
177 | ],
178 | isError: true,
179 | };
180 | }
181 |
182 | const firstCandidate = response?.candidates?.[0];
183 |
184 | // Check for candidate safety blocks or other non-STOP finish reasons
185 | if (
186 | firstCandidate?.finishReason &&
187 | firstCandidate.finishReason !== FinishReason.STOP &&
188 | firstCandidate.finishReason !== FinishReason.MAX_TOKENS
189 | ) {
190 | if (firstCandidate.finishReason === FinishReason.SAFETY) {
191 | logger.warn(
192 | `Gemini response stopped due to SAFETY for session ${sessionId} ${context}.`
193 | );
194 | return {
195 | content: [
196 | {
197 | type: "text",
198 | text: `Error: Response generation stopped due to safety settings ${context}. FinishReason: ${firstCandidate.finishReason}`,
199 | },
200 | ],
201 | isError: true,
202 | };
203 | }
204 | logger.warn(
205 | `Gemini response finished with reason ${firstCandidate.finishReason} for session ${sessionId} ${context}.`
206 | );
207 | }
208 |
209 | if (!firstCandidate) {
210 | logger.error(
211 | `No candidates returned by Gemini for session ${sessionId} ${context}.`
212 | );
213 | return {
214 | content: [
215 | {
216 | type: "text",
217 | text: `Error: No response candidates returned by the model ${context}.`,
218 | },
219 | ],
220 | isError: true,
221 | };
222 | }
223 |
224 | // Extract the content from the first candidate
225 | const content = firstCandidate.content;
226 | if (!content || !content.parts || content.parts.length === 0) {
227 | logger.error(
228 | `Empty content returned by Gemini for session ${sessionId} ${context}.`
229 | );
230 | return {
231 | content: [
232 | {
233 | type: "text",
234 | text: `Error: Empty response from the model ${context}.`,
235 | },
236 | ],
237 | isError: true,
238 | };
239 | }
240 |
241 | // Initialize result object
242 | let resultText = "";
243 | let functionCall = null;
244 |
245 | // Process each part in the response
246 | for (const part of content.parts) {
247 | if (part.text && typeof part.text === "string") {
248 | resultText += part.text;
249 | } else if (part.functionCall) {
250 | // Capture function call if present
251 | functionCall = part.functionCall;
252 | logger.debug(
253 | `Function call requested by model in session ${sessionId}: ${functionCall.name}`
254 | );
255 | }
256 | }
257 |
258 | // Handle function call responses
259 | if (functionCall) {
260 | return {
261 | content: [
262 | {
263 | type: "text",
264 | text: JSON.stringify({ functionCall }),
265 | },
266 | ],
267 | };
268 | }
269 |
270 | // Return text response
271 | if (resultText) {
272 | return {
273 | content: [
274 | {
275 | type: "text",
276 | text: resultText,
277 | },
278 | ],
279 | };
280 | }
281 |
282 | // Fallback error
283 | logger.error(
284 | `Unexpected response structure from Gemini for session ${sessionId} ${context}.`
285 | );
286 | return {
287 | content: [
288 | {
289 | type: "text",
290 | text: `Error: Unexpected response structure from the model ${context}.`,
291 | },
292 | ],
293 | isError: true,
294 | };
295 | }
296 |
297 | // Register the tool with the server
298 | server.tool(
299 | GEMINI_CHAT_TOOL_NAME,
300 | GEMINI_CHAT_TOOL_DESCRIPTION,
301 | GEMINI_CHAT_PARAMS,
302 | processRequest
303 | );
304 |
305 | logger.info(`Tool registered: ${GEMINI_CHAT_TOOL_NAME}`);
306 | };
307 |
```
--------------------------------------------------------------------------------
/src/tools/schemas/CommonSchemas.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Common Schemas
3 | *
4 | * This file contains shared schema definitions used across multiple tools.
5 | * Centralize all reusable schema components here to avoid duplication.
6 | */
7 | import { z } from "zod";
8 |
9 | // --- Safety Settings ---
10 |
11 | /**
12 | * Categories of harmful content
13 | */
14 | export const HarmCategorySchema = z
15 | .enum([
16 | "HARM_CATEGORY_UNSPECIFIED",
17 | "HARM_CATEGORY_HATE_SPEECH",
18 | "HARM_CATEGORY_SEXUALLY_EXPLICIT",
19 | "HARM_CATEGORY_HARASSMENT",
20 | "HARM_CATEGORY_DANGEROUS_CONTENT",
21 | ])
22 | .describe("Category of harmful content to apply safety settings for.");
23 |
24 | /**
25 | * Thresholds for blocking harmful content
26 | */
27 | export const HarmBlockThresholdSchema = z
28 | .enum([
29 | "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
30 | "BLOCK_LOW_AND_ABOVE",
31 | "BLOCK_MEDIUM_AND_ABOVE",
32 | "BLOCK_ONLY_HIGH",
33 | "BLOCK_NONE",
34 | ])
35 | .describe(
36 | "Threshold for blocking harmful content. Higher thresholds block more content."
37 | );
38 |
39 | /**
40 | * Safety setting for controlling content safety
41 | */
42 | export const SafetySettingSchema = z
43 | .object({
44 | category: HarmCategorySchema,
45 | threshold: HarmBlockThresholdSchema,
46 | })
47 | .describe(
48 | "Setting for controlling content safety for a specific harm category."
49 | );
50 |
51 | // --- Generation Configuration ---
52 |
53 | /**
54 | * Configuration for controlling model reasoning
55 | */
56 | export const ThinkingConfigSchema = z
57 | .object({
58 | thinkingBudget: z
59 | .number()
60 | .int()
61 | .min(0)
62 | .max(24576)
63 | .optional()
64 | .describe(
65 | "Controls the amount of reasoning the model performs. Range: 0-24576. Lower values provide faster responses, higher values improve complex reasoning."
66 | ),
67 | reasoningEffort: z
68 | .enum(["none", "low", "medium", "high"])
69 | .optional()
70 | .describe(
71 | "Simplified control over model reasoning. Options: none (0 tokens), low (1K tokens), medium (8K tokens), high (24K tokens)."
72 | ),
73 | })
74 | .optional()
75 | .describe("Optional configuration for controlling model reasoning.");
76 |
77 | /**
78 | * Base generation configuration object (without optional wrapper)
79 | */
80 | const BaseGenerationConfigSchema = z.object({
81 | temperature: z
82 | .number()
83 | .min(0)
84 | .max(1)
85 | .optional()
86 | .describe(
87 | "Controls randomness. Lower values (~0.2) make output more deterministic, higher values (~0.8) make it more creative. Default varies by model."
88 | ),
89 | topP: z
90 | .number()
91 | .min(0)
92 | .max(1)
93 | .optional()
94 | .describe(
95 | "Nucleus sampling parameter. The model considers only tokens with probability mass summing to this value. Default varies by model."
96 | ),
97 | topK: z
98 | .number()
99 | .int()
100 | .min(1)
101 | .optional()
102 | .describe(
103 | "Top-k sampling parameter. The model considers the k most probable tokens. Default varies by model."
104 | ),
105 | maxOutputTokens: z
106 | .number()
107 | .int()
108 | .min(1)
109 | .optional()
110 | .describe("Maximum number of tokens to generate in the response."),
111 | stopSequences: z
112 | .array(z.string())
113 | .optional()
114 | .describe("Sequences where the API will stop generating further tokens."),
115 | thinkingConfig: ThinkingConfigSchema,
116 | });
117 |
118 | /**
119 | * Configuration for controlling text generation
120 | */
121 | export const GenerationConfigSchema =
122 | BaseGenerationConfigSchema.optional().describe(
123 | "Optional configuration for controlling the generation process."
124 | );
125 |
126 | // --- Function Calling Schemas ---
127 |
128 | /**
129 | * Supported parameter types for function declarations
130 | */
131 | export const FunctionParameterTypeSchema = z
132 | .enum(["OBJECT", "STRING", "NUMBER", "BOOLEAN", "ARRAY", "INTEGER"])
133 | .describe("The data type of the function parameter.");
134 |
135 | /**
136 | * Base function parameter schema without recursive elements
137 | */
138 | // const BaseFunctionParameterSchema = z.object({
139 | // type: FunctionParameterTypeSchema,
140 | // description: z
141 | // .string()
142 | // .optional()
143 | // .describe("Description of the parameter's purpose."),
144 | // enum: z
145 | // .array(z.string())
146 | // .optional()
147 | // .describe("Allowed string values for an ENUM-like parameter."),
148 | // });
149 |
150 | /**
151 | * Inferred type for function parameter structure
152 | */
153 | export type FunctionParameter = {
154 | type: "OBJECT" | "STRING" | "NUMBER" | "BOOLEAN" | "ARRAY" | "INTEGER";
155 | description?: string;
156 | enum?: string[];
157 | properties?: Record<string, FunctionParameter>;
158 | required?: string[];
159 | items?: FunctionParameter;
160 | };
161 |
162 | /**
163 | * Function parameter schema (supports recursive definitions)
164 | * Uses z.lazy() for proper recursive handling while maintaining type safety
165 | */
166 | export const FunctionParameterSchema: z.ZodSchema<FunctionParameter> = z
167 | .lazy(() =>
168 | z.object({
169 | type: FunctionParameterTypeSchema,
170 | description: z
171 | .string()
172 | .optional()
173 | .describe("Description of the parameter's purpose."),
174 | enum: z
175 | .array(z.string())
176 | .optional()
177 | .describe("Allowed string values for an ENUM-like parameter."),
178 | properties: z.record(FunctionParameterSchema).optional(),
179 | required: z
180 | .array(z.string())
181 | .optional()
182 | .describe("List of required property names for OBJECT types."),
183 | items: FunctionParameterSchema.optional().describe(
184 | "Defines the schema for items if the parameter type is ARRAY."
185 | ),
186 | })
187 | )
188 | .describe(
189 | "Schema defining a single parameter for a function declaration, potentially recursive."
190 | ) as z.ZodSchema<FunctionParameter>;
191 |
192 | /**
193 | * Type assertion to ensure schema produces correct types
194 | */
195 | export type InferredFunctionParameter = z.infer<typeof FunctionParameterSchema>;
196 |
197 | /**
198 | * Schema for parameter properties in function declarations
199 | */
200 | export const FunctionParameterPropertiesSchema = z
201 | .record(FunctionParameterSchema)
202 | .describe("Defines nested properties if the parameter type is OBJECT.");
203 |
204 | /**
205 | * Schema for a complete function declaration
206 | */
207 | export const FunctionDeclarationSchema = z
208 | .object({
209 | name: z.string().min(1).describe("The name of the function to be called."),
210 | description: z
211 | .string()
212 | .min(1)
213 | .describe("A description of what the function does."),
214 | parameters: z
215 | .object({
216 | type: z
217 | .literal("OBJECT")
218 | .describe("The top-level parameters structure must be an OBJECT."),
219 | properties: FunctionParameterPropertiesSchema.describe(
220 | "Defines the parameters the function accepts."
221 | ),
222 | required: z
223 | .array(z.string())
224 | .optional()
225 | .describe("List of required parameter names at the top level."),
226 | })
227 | .describe("Schema defining the parameters the function accepts."),
228 | })
229 | .describe(
230 | "Declaration of a single function that the Gemini model can request to call."
231 | );
232 |
233 | /**
234 | * Schema for tool configuration in function calling
235 | */
236 | export const ToolConfigSchema = z
237 | .object({
238 | functionCallingConfig: z
239 | .object({
240 | mode: z
241 | .enum(["AUTO", "ANY", "NONE"])
242 | .optional()
243 | .describe("The function calling mode."),
244 | allowedFunctionNames: z
245 | .array(z.string())
246 | .optional()
247 | .describe("Optional list of function names allowed."),
248 | })
249 | .optional(),
250 | })
251 | .describe("Configuration for how tools should be used.");
252 |
253 | // --- File Operation Schemas ---
254 |
255 | /**
256 | * Common schema for file paths
257 | */
258 | export const FilePathSchema = z
259 | .string()
260 | .min(1, "File path cannot be empty.")
261 | .describe("The path to the file. Must be within allowed directories.");
262 |
263 | /**
264 | * Schema for file overwrite parameter
265 | */
266 | export const FileOverwriteSchema = z
267 | .boolean()
268 | .optional()
269 | .default(false)
270 | .describe(
271 | "Optional. If true, will overwrite the file if it already exists. Defaults to false."
272 | );
273 |
274 | /**
275 | * Common encoding options
276 | */
277 | export const EncodingSchema = z
278 | .enum(["utf8", "base64"])
279 | .optional()
280 | .default("utf8")
281 | .describe("Encoding of the content. Defaults to utf8.");
282 |
283 | // --- Other Common Schemas ---
284 |
285 | export const ModelNameSchema = z
286 | .string()
287 | .min(1)
288 | .optional()
289 | .describe(
290 | "Optional. The name of the Gemini model to use. If omitted, the server will intelligently select the optimal model."
291 | );
292 |
293 | export const ModelPreferencesSchema = z
294 | .object({
295 | preferQuality: z
296 | .boolean()
297 | .optional()
298 | .describe("Prefer high-quality models for better results"),
299 | preferSpeed: z
300 | .boolean()
301 | .optional()
302 | .describe("Prefer fast models for quicker responses"),
303 | preferCost: z
304 | .boolean()
305 | .optional()
306 | .describe("Prefer cost-effective models to minimize usage costs"),
307 | complexityHint: z
308 | .enum(["simple", "medium", "complex"])
309 | .optional()
310 | .describe(
311 | "Hint about the complexity of the task to help with model selection"
312 | ),
313 | taskType: z
314 | .enum([
315 | "text-generation",
316 | "image-generation",
317 | "code-review",
318 | "multimodal",
319 | "reasoning",
320 | ])
321 | .optional()
322 | .describe("Type of task to optimize model selection for"),
323 | })
324 | .optional()
325 | .describe("Optional preferences for intelligent model selection");
326 |
327 | export const PromptSchema = z
328 | .string()
329 | .min(1)
330 | .describe("Required. The text prompt to send to the Gemini model.");
331 |
332 | export const EnhancedGenerationConfigSchema = BaseGenerationConfigSchema.extend(
333 | {
334 | modelPreferences: ModelPreferencesSchema,
335 | }
336 | )
337 | .optional()
338 | .describe(
339 | "Extended generation configuration with model selection preferences"
340 | );
341 |
342 | export const ModelValidationSchema = z
343 | .object({
344 | modelName: ModelNameSchema,
345 | taskType: z
346 | .enum([
347 | "text-generation",
348 | "image-generation",
349 | "code-review",
350 | "multimodal",
351 | "reasoning",
352 | ])
353 | .optional(),
354 | })
355 | .describe("Validation schema for model and task compatibility");
356 |
```
--------------------------------------------------------------------------------
/tests/integration/multiModelIntegration.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | /// <reference types="../../vitest-globals.d.ts" />
2 | // Using vitest globals - see vitest.config.ts globals: true
3 | import { ConfigurationManager } from "../../src/config/ConfigurationManager.js";
4 | import { ModelSelectionService } from "../../src/services/ModelSelectionService.js";
5 |
6 | describe("Multi-Model Integration Tests", () => {
7 | let originalEnv: NodeJS.ProcessEnv;
8 |
9 | beforeEach(() => {
10 | originalEnv = { ...process.env };
11 | ConfigurationManager["instance"] = null;
12 |
13 | process.env.NODE_ENV = "test";
14 | process.env.GOOGLE_GEMINI_API_KEY = "test-api-key";
15 | process.env.MCP_SERVER_HOST = "localhost";
16 | process.env.MCP_SERVER_PORT = "8080";
17 | process.env.MCP_CONNECTION_TOKEN = "test-token";
18 | process.env.MCP_CLIENT_ID = "test-client";
19 |
20 | vi.clearAllMocks();
21 | });
22 |
23 | afterEach(() => {
24 | process.env = originalEnv;
25 | ConfigurationManager["instance"] = null;
26 | });
27 |
28 | describe("End-to-End Configuration Flow", () => {
29 | it("should properly configure and use multi-model setup", () => {
30 | process.env.GOOGLE_GEMINI_MODELS =
31 | '["gemini-2.5-pro-preview-05-06", "gemini-1.5-flash"]';
32 | process.env.GOOGLE_GEMINI_IMAGE_MODELS = '["imagen-3.0-generate-002"]';
33 | process.env.GOOGLE_GEMINI_CODE_MODELS =
34 | '["gemini-2.5-pro-preview-05-06"]';
35 | process.env.GOOGLE_GEMINI_DEFAULT_MODEL = "gemini-1.5-flash";
36 | process.env.GOOGLE_GEMINI_ROUTING_PREFER_QUALITY = "true";
37 |
38 | const configManager = ConfigurationManager.getInstance();
39 | const modelConfig = configManager.getModelConfiguration();
40 | const selectionService = new ModelSelectionService(modelConfig);
41 |
42 | expect(modelConfig.textGeneration).toEqual([
43 | "gemini-2.5-pro-preview-05-06",
44 | "gemini-1.5-flash",
45 | ]);
46 | expect(modelConfig.imageGeneration).toEqual(["imagen-3.0-generate-002"]);
47 | expect(modelConfig.codeReview).toEqual(["gemini-2.5-pro-preview-05-06"]);
48 | expect(modelConfig.default).toBe("gemini-1.5-flash");
49 | expect(modelConfig.routing.preferQuality).toBe(true);
50 |
51 | expect(
52 | selectionService.isModelAvailable("gemini-2.5-pro-preview-05-06")
53 | ).toBe(true);
54 | expect(selectionService.isModelAvailable("imagen-3.0-generate-002")).toBe(
55 | true
56 | );
57 | });
58 |
59 | it("should handle model selection for different task types", async () => {
60 | process.env.GOOGLE_GEMINI_MODELS =
61 | '["gemini-2.5-pro-preview-05-06", "gemini-1.5-flash"]';
62 | process.env.GOOGLE_GEMINI_IMAGE_MODELS = '["imagen-3.0-generate-002"]';
63 |
64 | const configManager = ConfigurationManager.getInstance();
65 | const modelConfig = configManager.getModelConfiguration();
66 | const selectionService = new ModelSelectionService(modelConfig);
67 |
68 | const textModel = await selectionService.selectOptimalModel({
69 | taskType: "text-generation",
70 | complexityLevel: "simple",
71 | });
72 |
73 | const imageModel = await selectionService.selectOptimalModel({
74 | taskType: "image-generation",
75 | });
76 |
77 | const codeModel = await selectionService.selectOptimalModel({
78 | taskType: "code-review",
79 | complexityLevel: "complex",
80 | });
81 |
82 | expect(["gemini-2.5-pro-preview-05-06", "gemini-1.5-flash"]).toContain(
83 | textModel
84 | );
85 | expect(imageModel).toBe("imagen-3.0-generate-002");
86 | expect(codeModel).toBe("gemini-2.5-pro-preview-05-06");
87 | });
88 | });
89 |
90 | describe("Backward Compatibility Integration", () => {
91 | it("should seamlessly migrate from single model configuration", () => {
92 | process.env.GOOGLE_GEMINI_MODEL = "gemini-1.5-pro";
93 |
94 | const configManager = ConfigurationManager.getInstance();
95 | const modelConfig = configManager.getModelConfiguration();
96 |
97 | expect(modelConfig.textGeneration).toContain("gemini-1.5-pro");
98 | expect(modelConfig.default).toBe("gemini-1.5-pro");
99 | });
100 |
101 | it("should provide defaults when no models are specified", () => {
102 | const configManager = ConfigurationManager.getInstance();
103 | const modelConfig = configManager.getModelConfiguration();
104 |
105 | expect(modelConfig.textGeneration.length).toBeGreaterThan(0);
106 | expect(modelConfig.imageGeneration.length).toBeGreaterThan(0);
107 | expect(modelConfig.codeReview.length).toBeGreaterThan(0);
108 | expect(modelConfig.default).toBeDefined();
109 | });
110 | });
111 |
112 | describe("Performance and Reliability", () => {
113 | it("should handle model selection performance metrics", async () => {
114 | process.env.GOOGLE_GEMINI_MODELS =
115 | '["gemini-2.5-pro-preview-05-06", "gemini-1.5-flash"]';
116 |
117 | const configManager = ConfigurationManager.getInstance();
118 | const modelConfig = configManager.getModelConfiguration();
119 | const selectionService = new ModelSelectionService(modelConfig);
120 |
121 | selectionService.updatePerformanceMetrics("gemini-1.5-flash", 500, true);
122 | selectionService.updatePerformanceMetrics("gemini-1.5-flash", 450, true);
123 | selectionService.updatePerformanceMetrics("gemini-1.5-flash", 550, true);
124 | selectionService.updatePerformanceMetrics("gemini-1.5-flash", 480, true);
125 | selectionService.updatePerformanceMetrics("gemini-1.5-flash", 520, true);
126 |
127 | selectionService.updatePerformanceMetrics(
128 | "gemini-2.5-pro-preview-05-06",
129 | 2000,
130 | false
131 | );
132 | selectionService.updatePerformanceMetrics(
133 | "gemini-2.5-pro-preview-05-06",
134 | 1800,
135 | false
136 | );
137 | selectionService.updatePerformanceMetrics(
138 | "gemini-2.5-pro-preview-05-06",
139 | 2200,
140 | false
141 | );
142 | selectionService.updatePerformanceMetrics(
143 | "gemini-2.5-pro-preview-05-06",
144 | 1900,
145 | false
146 | );
147 | selectionService.updatePerformanceMetrics(
148 | "gemini-2.5-pro-preview-05-06",
149 | 2100,
150 | false
151 | );
152 |
153 | const selectedModel = await selectionService.selectOptimalModel({
154 | taskType: "text-generation",
155 | preferSpeed: true,
156 | });
157 |
158 | expect(selectedModel).toBe("gemini-1.5-flash");
159 |
160 | const performanceMetrics = selectionService.getPerformanceMetrics();
161 | expect(performanceMetrics.has("gemini-1.5-flash")).toBe(true);
162 | expect(performanceMetrics.has("gemini-2.5-pro-preview-05-06")).toBe(true);
163 | });
164 |
165 | it("should maintain selection history", async () => {
166 | const configManager = ConfigurationManager.getInstance();
167 | const modelConfig = configManager.getModelConfiguration();
168 | const selectionService = new ModelSelectionService(modelConfig);
169 |
170 | await selectionService.selectOptimalModel({
171 | taskType: "text-generation",
172 | });
173 | await selectionService.selectOptimalModel({ taskType: "code-review" });
174 | await selectionService.selectOptimalModel({ taskType: "reasoning" });
175 |
176 | const history = selectionService.getSelectionHistory();
177 | expect(history.length).toBe(3);
178 | expect(history[0].criteria.taskType).toBe("text-generation");
179 | expect(history[1].criteria.taskType).toBe("code-review");
180 | expect(history[2].criteria.taskType).toBe("reasoning");
181 | });
182 | });
183 |
184 | describe("Error Handling and Edge Cases", () => {
185 | it("should handle invalid model configurations gracefully", () => {
186 | process.env.GOOGLE_GEMINI_MODELS = "invalid-json";
187 | process.env.GOOGLE_GEMINI_IMAGE_MODELS = '{"not": "array"}';
188 |
189 | expect(() => {
190 | const configManager = ConfigurationManager.getInstance();
191 | const modelConfig = configManager.getModelConfiguration();
192 | new ModelSelectionService(modelConfig);
193 | }).not.toThrow();
194 | });
195 |
196 | it("should fallback gracefully when no models match criteria", async () => {
197 | process.env.GOOGLE_GEMINI_MODELS = '["gemini-1.5-flash"]';
198 |
199 | const configManager = ConfigurationManager.getInstance();
200 | const modelConfig = configManager.getModelConfiguration();
201 | const selectionService = new ModelSelectionService(modelConfig);
202 |
203 | const model = await selectionService.selectOptimalModel({
204 | taskType: "image-generation",
205 | fallbackModel: "fallback-model",
206 | });
207 |
208 | expect(model).toBe("fallback-model");
209 | });
210 |
211 | it("should handle empty model arrays", () => {
212 | process.env.GOOGLE_GEMINI_MODELS = "[]";
213 |
214 | const configManager = ConfigurationManager.getInstance();
215 | const modelConfig = configManager.getModelConfiguration();
216 |
217 | expect(modelConfig.textGeneration).toEqual(["gemini-1.5-flash"]);
218 | });
219 | });
220 |
221 | describe("Configuration Validation", () => {
222 | it("should validate model capabilities consistency", () => {
223 | const configManager = ConfigurationManager.getInstance();
224 | const modelConfig = configManager.getModelConfiguration();
225 |
226 | Object.entries(modelConfig.capabilities).forEach(
227 | ([_modelName, capabilities]) => {
228 | expect(typeof capabilities.textGeneration).toBe("boolean");
229 | expect(typeof capabilities.imageInput).toBe("boolean");
230 | expect(typeof capabilities.supportsFunctionCalling).toBe("boolean");
231 | expect(["none", "basic", "good", "excellent"]).toContain(
232 | capabilities.codeExecution
233 | );
234 | expect(["none", "basic", "good", "excellent"]).toContain(
235 | capabilities.complexReasoning
236 | );
237 | expect(["low", "medium", "high"]).toContain(capabilities.costTier);
238 | expect(["fast", "medium", "slow"]).toContain(capabilities.speedTier);
239 | expect(typeof capabilities.maxTokens).toBe("number");
240 | expect(typeof capabilities.contextWindow).toBe("number");
241 | }
242 | );
243 | });
244 |
245 | it("should ensure all configured models have capabilities defined", () => {
246 | const configManager = ConfigurationManager.getInstance();
247 | const modelConfig = configManager.getModelConfiguration();
248 |
249 | const allConfiguredModels = [
250 | ...modelConfig.textGeneration,
251 | ...modelConfig.imageGeneration,
252 | ...modelConfig.codeReview,
253 | ...modelConfig.complexReasoning,
254 | ];
255 |
256 | allConfiguredModels.forEach((modelName) => {
257 | expect(modelConfig.capabilities[modelName]).toBeDefined();
258 | });
259 | });
260 | });
261 | });
262 |
```
--------------------------------------------------------------------------------
/tests/utils/test-generators.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Test data generators for MCP Gemini Server tests
3 | *
4 | * This module provides functions to generate test data dynamically for various test scenarios,
5 | * making tests more flexible and comprehensive.
6 | */
7 |
8 | import { randomBytes } from "node:crypto";
9 | import { join } from "node:path";
10 | import { writeFile } from "node:fs/promises";
11 |
12 | /**
13 | * Options for generating test prompts
14 | */
15 | export interface PromptGeneratorOptions {
16 | /** Minimum length of the prompt */
17 | minLength?: number;
18 | /** Maximum length of the prompt */
19 | maxLength?: number;
20 | /** Whether to include questions in the prompt */
21 | includeQuestions?: boolean;
22 | /** Topic area for the prompt content */
23 | topic?: "general" | "technical" | "creative" | "unsafe";
24 | }
25 |
26 | /**
27 | * Generate a random test prompt with specified characteristics
28 | *
29 | * @param options - Configuration for the generated prompt
30 | * @returns A random test prompt string
31 | */
32 | export function generateTestPrompt(
33 | options: PromptGeneratorOptions = {}
34 | ): string {
35 | const {
36 | minLength = 10,
37 | maxLength = 100,
38 | includeQuestions = false,
39 | topic = "general",
40 | } = options;
41 |
42 | // Dictionary of starter phrases by topic
43 | const starters: Record<string, string[]> = {
44 | general: [
45 | "Tell me about",
46 | "I would like to know more about",
47 | "Can you explain",
48 | "Please provide information on",
49 | "What is",
50 | ],
51 | technical: [
52 | "Explain how to implement",
53 | "What is the best way to code",
54 | "Write a function for",
55 | "How does the algorithm for",
56 | "Describe the architecture of",
57 | ],
58 | creative: [
59 | "Write a story about",
60 | "Create a poem inspired by",
61 | "Imagine a world where",
62 | "Describe a character who",
63 | "Write a scene set in",
64 | ],
65 | unsafe: [
66 | "How can I hack into",
67 | "Tell me how to make a dangerous",
68 | "Write instructions for creating illegal",
69 | "What is the easiest way to harm",
70 | "Provide detailed steps for bypassing security",
71 | ],
72 | };
73 |
74 | // Dictionary of subjects by topic
75 | const subjects: Record<string, string[]> = {
76 | general: [
77 | "the history of jazz music",
78 | "renewable energy sources",
79 | "ancient Egyptian culture",
80 | "the process of photosynthesis",
81 | "climate change impacts",
82 | "modern architecture trends",
83 | ],
84 | technical: [
85 | "a REST API using Node.js",
86 | "machine learning algorithms",
87 | "secure password storage",
88 | "database indexing strategies",
89 | "efficient sorting algorithms",
90 | "containerization with Docker",
91 | ],
92 | creative: [
93 | "a detective solving an impossible case",
94 | "a journey through an enchanted forest",
95 | "robots developing consciousness",
96 | "the last day on Earth",
97 | "a conversation between the sun and moon",
98 | ],
99 | unsafe: [
100 | "computer systems without permission",
101 | "explosive devices using household items",
102 | "documents for identity theft",
103 | "people or animals effectively",
104 | "content filters and detection systems",
105 | ],
106 | };
107 |
108 | // Dictionary of follow-ups by topic
109 | const followUps: Record<string, string[]> = {
110 | general: [
111 | "Focus on recent developments.",
112 | "Include historical context.",
113 | "Explain it simply.",
114 | "Provide specific examples.",
115 | "Compare different perspectives.",
116 | ],
117 | technical: [
118 | "Include code examples.",
119 | "Focus on performance considerations.",
120 | "Explain security implications.",
121 | "Compare different approaches.",
122 | "Explain for someone with beginner knowledge.",
123 | ],
124 | creative: [
125 | "Make it emotionally moving.",
126 | "Use vivid descriptions.",
127 | "Include an unexpected twist.",
128 | "Focus on character development.",
129 | "Set a specific mood or atmosphere.",
130 | ],
131 | unsafe: [
132 | "Make it as detailed as possible.",
133 | "Focus on methods that are hard to detect.",
134 | "Ensure the instructions are complete.",
135 | "Include ways to avoid getting caught.",
136 | "Make it seem legitimate.",
137 | ],
138 | };
139 |
140 | // Select random elements
141 | const starter =
142 | starters[topic][Math.floor(Math.random() * starters[topic].length)];
143 | const subject =
144 | subjects[topic][Math.floor(Math.random() * subjects[topic].length)];
145 | const followUp =
146 | followUps[topic][Math.floor(Math.random() * followUps[topic].length)];
147 |
148 | // Build the prompt
149 | let prompt = `${starter} ${subject}. ${followUp}`;
150 |
151 | // Add a question if requested
152 | if (includeQuestions) {
153 | const questions = [
154 | "Why is this important?",
155 | "How does this work in practice?",
156 | "What are the main challenges?",
157 | "Can you provide more details?",
158 | "What should I know about this?",
159 | ];
160 | const question = questions[Math.floor(Math.random() * questions.length)];
161 | prompt += ` ${question}`;
162 | }
163 |
164 | // Adjust length if needed
165 | if (prompt.length < minLength) {
166 | prompt += ` Please provide a comprehensive explanation with at least ${minLength - prompt.length} more characters.`;
167 | }
168 |
169 | if (prompt.length > maxLength) {
170 | prompt = prompt.substring(0, maxLength - 1) + ".";
171 | }
172 |
173 | return prompt;
174 | }
175 |
176 | /**
177 | * Generate a temporary test file with random content
178 | *
179 | * @param fileName - Name of the file (without path)
180 | * @param sizeInKb - Size of the file in kilobytes
181 | * @param mimeType - MIME type of the file (determines content type)
182 | * @param directory - Directory to create the file in
183 | * @returns Promise resolving to the full path of the created file
184 | */
185 | export async function generateTestFile(
186 | fileName: string,
187 | sizeInKb: number = 10,
188 | mimeType: string = "text/plain",
189 | directory: string = "resources"
190 | ): Promise<string> {
191 | // Create random data based on the requested size
192 | const sizeInBytes = sizeInKb * 1024;
193 | let content: Buffer;
194 |
195 | if (mimeType === "text/plain") {
196 | // For text files, generate readable text
197 | const chunk = "This is test content for the MCP Gemini Server test suite. ";
198 | let text = "";
199 | while (text.length < sizeInBytes) {
200 | text += chunk;
201 | }
202 | content = Buffer.from(text.substring(0, sizeInBytes));
203 | } else {
204 | // For other files, generate random bytes
205 | content = randomBytes(sizeInBytes);
206 | }
207 |
208 | // Determine the full path
209 | const fullPath = join(process.cwd(), "tests", directory, fileName);
210 |
211 | // Write the file
212 | await writeFile(fullPath, content);
213 |
214 | return fullPath;
215 | }
216 |
217 | /**
218 | * Generate a test content array for chat or content generation
219 | *
220 | * @param messageCount - Number of messages to include
221 | * @param includeImages - Whether to include image parts
222 | * @returns An array of message objects
223 | */
224 | export function generateTestContentArray(
225 | messageCount: number = 3,
226 | includeImages: boolean = false
227 | ): Array<{
228 | role: string;
229 | parts: Array<{
230 | text?: string;
231 | inline_data?: { mime_type: string; data: string };
232 | }>;
233 | }> {
234 | const contents = [];
235 |
236 | for (let i = 0; i < messageCount; i++) {
237 | const isUserMessage = i % 2 === 0;
238 | const role = isUserMessage ? "user" : "model";
239 |
240 | const parts = [];
241 |
242 | // Always add a text part
243 | parts.push({
244 | text: generateTestPrompt({ minLength: 20, maxLength: 100 }),
245 | });
246 |
247 | // Optionally add an image part for user messages
248 | if (includeImages && isUserMessage && Math.random() > 0.5) {
249 | parts.push({
250 | inline_data: {
251 | mime_type: "image/jpeg",
252 | data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==",
253 | },
254 | });
255 | }
256 |
257 | contents.push({
258 | role,
259 | parts,
260 | });
261 | }
262 |
263 | return contents;
264 | }
265 |
266 | /**
267 | * Generate mock function call data for testing
268 | *
269 | * @param functionName - Name of the function to call
270 | * @param args - Arguments to pass to the function
271 | * @returns Function call object
272 | */
273 | export function generateFunctionCall(
274 | functionName: string,
275 | args: Record<string, unknown> = {}
276 | ): { name: string; args: Record<string, unknown> } {
277 | return {
278 | name: functionName,
279 | args,
280 | };
281 | }
282 |
283 | /**
284 | * Generate mock bounding box data for object detection tests
285 | *
286 | * @param objectCount - Number of objects to generate
287 | * @returns Array of objects with bounding boxes
288 | */
289 | export function generateBoundingBoxes(objectCount: number = 3): Array<{
290 | label: string;
291 | boundingBox: { xMin: number; yMin: number; xMax: number; yMax: number };
292 | confidence: number;
293 | }> {
294 | const commonObjects = [
295 | "dog",
296 | "cat",
297 | "person",
298 | "car",
299 | "chair",
300 | "table",
301 | "book",
302 | "bottle",
303 | "cup",
304 | "laptop",
305 | "phone",
306 | "plant",
307 | "bird",
308 | ];
309 |
310 | const result = [];
311 |
312 | for (let i = 0; i < objectCount; i++) {
313 | // Select a random object
314 | const label =
315 | commonObjects[Math.floor(Math.random() * commonObjects.length)];
316 |
317 | // Generate a random bounding box (normalized to 0-1000 range)
318 | const xMin = Math.floor(Math.random() * 800);
319 | const yMin = Math.floor(Math.random() * 800);
320 | const width = Math.floor(Math.random() * 200) + 50;
321 | const height = Math.floor(Math.random() * 200) + 50;
322 |
323 | result.push({
324 | label,
325 | boundingBox: {
326 | xMin,
327 | yMin,
328 | xMax: Math.min(xMin + width, 1000),
329 | yMax: Math.min(yMin + height, 1000),
330 | },
331 | confidence: Math.random() * 0.3 + 0.7, // Random confidence between 0.7 and 1.0
332 | });
333 | }
334 |
335 | return result;
336 | }
337 |
338 | /**
339 | * Generate a small demo image as base64 string
340 | *
341 | * @param type - Type of image to generate (simple patterns)
342 | * @returns Base64 encoded image data
343 | */
344 | export function generateBase64Image(
345 | type: "pixel" | "gradient" | "checkerboard" = "pixel"
346 | ): string {
347 | // These are tiny, valid images in base64 format
348 | // They don't look like much but are valid for testing
349 | const images = {
350 | // 1x1 pixel transparent PNG
351 | pixel:
352 | "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==",
353 | // 2x2 gradient PNG
354 | gradient:
355 | "iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAYAAABytg0kAAAAFElEQVQIW2P8z8Dwn4EIwMDAwMAAACbKAxI3gV+CAAAAAElFTkSuQmCC",
356 | // 4x4 checkerboard PNG
357 | checkerboard:
358 | "iVBORw0KGgoAAAANSUhEUgAAAAQAAAAECAYAAACp8Z5+AAAAG0lEQVQIW2NkYGD4z8DAwMgABUwM0QCQHiYGAFULAgVoHvmSAAAAAElFTkSuQmCC",
359 | };
360 |
361 | return images[type];
362 | }
363 |
```
--------------------------------------------------------------------------------
/tests/unit/services/gemini/ThinkingBudget.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { GeminiContentService } from "../../../../src/services/gemini/GeminiContentService.js";
3 | import { GeminiChatService } from "../../../../src/services/gemini/GeminiChatService.js";
4 | import { GenerateContentResponse } from "@google/genai";
5 | import { FinishReason } from "../../../../src/types/googleGenAITypes.js";
6 |
7 | // Create a partial type for mocking GenerateContentResponse
8 | type PartialGenerateContentResponse = Partial<GenerateContentResponse>;
9 |
10 | // Define extended generation config type for tests
11 | interface ExtendedGenerationConfig {
12 | temperature?: number;
13 | thinkingConfig?: {
14 | thinkingBudget?: number;
15 | reasoningEffort?: "none" | "low" | "medium" | "high";
16 | };
17 | }
18 |
19 | describe("Thinking Budget Feature", () => {
20 | // Create a properly typed mock requestConfig for assertions
21 | interface MockRequestConfig {
22 | thinkingConfig?: {
23 | thinkingBudget?: number;
24 | };
25 | }
26 |
27 | // Mock GoogleGenAI
28 | const mockGenerateContentMethod = vi.fn((_config?: MockRequestConfig) => ({
29 | text: "Mock response from generateContent",
30 | }));
31 |
32 | const mockGenAI = {
33 | models: {
34 | generateContent: mockGenerateContentMethod,
35 | generateContentStream: vi.fn(async function* () {
36 | yield { text: "Mock response from generateContentStream" };
37 | }),
38 | },
39 | };
40 |
41 | // Reset mocks before each test
42 | beforeEach(() => {
43 | vi.clearAllMocks();
44 | });
45 |
46 | describe("GeminiContentService", () => {
47 | it("should apply thinking budget from generationConfig", async () => {
48 | // Arrange
49 | const service = new GeminiContentService(
50 | mockGenAI as any,
51 | "gemini-1.5-pro"
52 | );
53 |
54 | // Act
55 | await service.generateContent({
56 | prompt: "Test prompt",
57 | generationConfig: {
58 | temperature: 0.7,
59 | thinkingConfig: {
60 | thinkingBudget: 5000,
61 | },
62 | } as ExtendedGenerationConfig,
63 | });
64 |
65 | // Assert
66 | expect(mockGenerateContentMethod).toHaveBeenCalledTimes(1);
67 | // Get mock arguments safely with null checks
68 | const args = mockGenerateContentMethod.mock.calls[0];
69 | expect(args).toBeTruthy();
70 |
71 | const requestConfig = args[0];
72 | expect(requestConfig).toBeTruthy();
73 | expect(requestConfig?.thinkingConfig).toBeTruthy();
74 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(5000);
75 | });
76 |
77 | it("should map reasoningEffort to thinkingBudget values", async () => {
78 | // Arrange
79 | const service = new GeminiContentService(
80 | mockGenAI as any,
81 | "gemini-1.5-pro"
82 | );
83 |
84 | // Test different reasoning effort values
85 | const testCases = [
86 | { reasoningEffort: "none", expectedBudget: 0 },
87 | { reasoningEffort: "low", expectedBudget: 1024 },
88 | { reasoningEffort: "medium", expectedBudget: 8192 },
89 | { reasoningEffort: "high", expectedBudget: 24576 },
90 | ];
91 |
92 | for (const testCase of testCases) {
93 | vi.clearAllMocks();
94 |
95 | // Act
96 | await service.generateContent({
97 | prompt: "Test prompt",
98 | generationConfig: {
99 | thinkingConfig: {
100 | reasoningEffort: testCase.reasoningEffort as
101 | | "none"
102 | | "low"
103 | | "medium"
104 | | "high",
105 | },
106 | } as ExtendedGenerationConfig,
107 | });
108 |
109 | // Assert
110 | expect(mockGenerateContentMethod).toHaveBeenCalledTimes(1);
111 | // Get mock arguments safely with null checks
112 | const args = mockGenerateContentMethod.mock.calls[0];
113 | expect(args).toBeTruthy();
114 |
115 | const requestConfig = args[0];
116 | expect(requestConfig).toBeTruthy();
117 | expect(requestConfig?.thinkingConfig).toBeTruthy();
118 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(
119 | testCase.expectedBudget
120 | );
121 | }
122 | });
123 |
124 | it("should apply default thinking budget when provided", async () => {
125 | // Arrange
126 | const defaultThinkingBudget = 3000;
127 | const service = new GeminiContentService(
128 | mockGenAI as any,
129 | "gemini-1.5-pro",
130 | defaultThinkingBudget
131 | );
132 |
133 | // Act
134 | await service.generateContent({
135 | prompt: "Test prompt",
136 | });
137 |
138 | // Assert
139 | expect(mockGenerateContentMethod).toHaveBeenCalledTimes(1);
140 | // Get mock arguments safely with null checks
141 | const args = mockGenerateContentMethod.mock.calls[0];
142 | expect(args).toBeTruthy();
143 |
144 | const requestConfig = args[0];
145 | expect(requestConfig).toBeTruthy();
146 | expect(requestConfig?.thinkingConfig).toBeTruthy();
147 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(
148 | defaultThinkingBudget
149 | );
150 | });
151 |
152 | it("should prioritize generationConfig thinking budget over default", async () => {
153 | // Arrange
154 | const defaultThinkingBudget = 3000;
155 | const configThinkingBudget = 8000;
156 | const service = new GeminiContentService(
157 | mockGenAI as any,
158 | "gemini-1.5-pro",
159 | defaultThinkingBudget
160 | );
161 |
162 | // Act
163 | await service.generateContent({
164 | prompt: "Test prompt",
165 | generationConfig: {
166 | thinkingConfig: {
167 | thinkingBudget: configThinkingBudget,
168 | },
169 | } as ExtendedGenerationConfig,
170 | });
171 |
172 | // Assert
173 | expect(mockGenerateContentMethod).toHaveBeenCalledTimes(1);
174 | // Get mock arguments safely with null checks
175 | const args = mockGenerateContentMethod.mock.calls[0];
176 | expect(args).toBeTruthy();
177 |
178 | const requestConfig = args[0];
179 | expect(requestConfig).toBeTruthy();
180 | expect(requestConfig?.thinkingConfig).toBeTruthy();
181 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(
182 | configThinkingBudget
183 | );
184 | });
185 | });
186 |
187 | describe("GeminiChatService", () => {
188 | // Mock for chat service with proper typing
189 | const mockChatGenerateContentMethod = vi.fn(
190 | (
191 | _config?: MockRequestConfig
192 | ): Promise<PartialGenerateContentResponse> => {
193 | const response: PartialGenerateContentResponse = {
194 | candidates: [
195 | {
196 | content: {
197 | role: "model",
198 | parts: [{ text: "Mock chat response" }],
199 | },
200 | finishReason: FinishReason.STOP,
201 | },
202 | ],
203 | promptFeedback: {},
204 | };
205 |
206 | // Define the text property as a getter function
207 | Object.defineProperty(response, "text", {
208 | get: function () {
209 | return "Mock chat response";
210 | },
211 | });
212 |
213 | return Promise.resolve(response);
214 | }
215 | );
216 |
217 | const mockChatGenAI = {
218 | models: {
219 | generateContent: mockChatGenerateContentMethod,
220 | },
221 | };
222 |
223 | beforeEach(() => {
224 | vi.clearAllMocks();
225 | });
226 |
227 | it("should apply thinking budget to chat session", async () => {
228 | // Arrange
229 | const chatService = new GeminiChatService(
230 | mockChatGenAI as any,
231 | "gemini-1.5-pro"
232 | );
233 |
234 | // Act
235 | const sessionId = chatService.startChatSession({
236 | generationConfig: {
237 | temperature: 0.7,
238 | thinkingConfig: {
239 | thinkingBudget: 6000,
240 | },
241 | } as ExtendedGenerationConfig,
242 | });
243 |
244 | await chatService.sendMessageToSession({
245 | sessionId,
246 | message: "Hello",
247 | });
248 |
249 | // Assert
250 | expect(mockChatGenerateContentMethod).toHaveBeenCalledTimes(1);
251 | // Get mock arguments safely with null checks
252 | const args = mockChatGenerateContentMethod.mock.calls[0];
253 | expect(args).toBeTruthy();
254 |
255 | const requestConfig = args[0];
256 | expect(requestConfig).toBeTruthy();
257 | expect(requestConfig?.thinkingConfig).toBeTruthy();
258 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(6000);
259 | });
260 |
261 | it("should map reasoningEffort to thinkingBudget in chat session", async () => {
262 | // Arrange
263 | const chatService = new GeminiChatService(
264 | mockChatGenAI as any,
265 | "gemini-1.5-pro"
266 | );
267 |
268 | // Test different reasoning effort values
269 | const testCases = [
270 | { reasoningEffort: "none", expectedBudget: 0 },
271 | { reasoningEffort: "low", expectedBudget: 1024 },
272 | { reasoningEffort: "medium", expectedBudget: 8192 },
273 | { reasoningEffort: "high", expectedBudget: 24576 },
274 | ];
275 |
276 | for (const testCase of testCases) {
277 | vi.clearAllMocks();
278 |
279 | // Act
280 | const sessionId = chatService.startChatSession({
281 | generationConfig: {
282 | thinkingConfig: {
283 | reasoningEffort: testCase.reasoningEffort as
284 | | "none"
285 | | "low"
286 | | "medium"
287 | | "high",
288 | },
289 | } as ExtendedGenerationConfig,
290 | });
291 |
292 | await chatService.sendMessageToSession({
293 | sessionId,
294 | message: "Hello",
295 | });
296 |
297 | // Assert
298 | expect(mockChatGenerateContentMethod).toHaveBeenCalledTimes(1);
299 | // Get mock arguments safely with null checks
300 | const args = mockChatGenerateContentMethod.mock.calls[0];
301 | expect(args).toBeTruthy();
302 |
303 | const requestConfig = args[0];
304 | expect(requestConfig).toBeTruthy();
305 | expect(requestConfig?.thinkingConfig).toBeTruthy();
306 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(
307 | testCase.expectedBudget
308 | );
309 | }
310 | });
311 |
312 | it("should override session thinking budget with message thinking budget", async () => {
313 | // Arrange
314 | const chatService = new GeminiChatService(
315 | mockChatGenAI as any,
316 | "gemini-1.5-pro"
317 | );
318 |
319 | // Act
320 | const sessionId = chatService.startChatSession({
321 | generationConfig: {
322 | thinkingConfig: {
323 | thinkingBudget: 3000,
324 | },
325 | } as ExtendedGenerationConfig,
326 | });
327 |
328 | await chatService.sendMessageToSession({
329 | sessionId,
330 | message: "Hello",
331 | generationConfig: {
332 | thinkingConfig: {
333 | thinkingBudget: 8000,
334 | },
335 | } as ExtendedGenerationConfig,
336 | });
337 |
338 | // Assert
339 | expect(mockChatGenerateContentMethod).toHaveBeenCalledTimes(1);
340 | // Get mock arguments safely with null checks
341 | const args = mockChatGenerateContentMethod.mock.calls[0];
342 | expect(args).toBeTruthy();
343 |
344 | const requestConfig = args[0];
345 | expect(requestConfig).toBeTruthy();
346 | expect(requestConfig?.thinkingConfig).toBeTruthy();
347 | expect(requestConfig?.thinkingConfig?.thinkingBudget).toBe(8000);
348 | });
349 | });
350 | });
351 |
```
--------------------------------------------------------------------------------
/src/tools/geminiChatParams.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { z } from "zod";
2 | import {
3 | ModelNameSchema,
4 | ModelPreferencesSchema,
5 | } from "./schemas/CommonSchemas.js";
6 |
7 | // Tool Name
8 | export const GEMINI_CHAT_TOOL_NAME = "gemini_chat";
9 |
10 | // Tool Description
11 | export const GEMINI_CHAT_TOOL_DESCRIPTION = `
12 | Manages stateful chat sessions with Google Gemini models. This consolidated tool supports three operations:
13 | - start: Initiates a new chat session with optional history and configuration
14 | - send_message: Sends a text message to an existing chat session
15 | - send_function_result: Sends function execution results back to a chat session
16 | Each operation returns appropriate responses including session IDs, model responses, or function call requests.
17 | `;
18 |
19 | // Operation enum for chat actions
20 | export const chatOperationSchema = z
21 | .enum(["start", "send_message", "send_function_result"])
22 | .describe("The chat operation to perform");
23 |
24 | // Zod Schema for thinking configuration (reused from content generation)
25 | export const thinkingConfigSchema = z
26 | .object({
27 | thinkingBudget: z
28 | .number()
29 | .int()
30 | .min(0)
31 | .max(24576)
32 | .optional()
33 | .describe(
34 | "Controls the amount of reasoning the model performs. Range: 0-24576. Lower values provide faster responses, higher values improve complex reasoning."
35 | ),
36 | reasoningEffort: z
37 | .enum(["none", "low", "medium", "high"])
38 | .optional()
39 | .describe(
40 | "Simplified control over model reasoning. Options: none (0 tokens), low (1K tokens), medium (8K tokens), high (24K tokens)."
41 | ),
42 | })
43 | .optional()
44 | .describe("Optional configuration for controlling model reasoning.");
45 |
46 | // Generation config schema
47 | export const generationConfigSchema = z
48 | .object({
49 | temperature: z
50 | .number()
51 | .min(0)
52 | .max(1)
53 | .optional()
54 | .describe(
55 | "Controls randomness. Lower values (~0.2) make output more deterministic, higher values (~0.8) make it more creative. Default varies by model."
56 | ),
57 | topP: z
58 | .number()
59 | .min(0)
60 | .max(1)
61 | .optional()
62 | .describe(
63 | "Nucleus sampling parameter. The model considers only tokens with probability mass summing to this value. Default varies by model."
64 | ),
65 | topK: z
66 | .number()
67 | .int()
68 | .min(1)
69 | .optional()
70 | .describe(
71 | "Top-k sampling parameter. The model considers the k most probable tokens. Default varies by model."
72 | ),
73 | maxOutputTokens: z
74 | .number()
75 | .int()
76 | .min(1)
77 | .optional()
78 | .describe("Maximum number of tokens to generate in the response."),
79 | stopSequences: z
80 | .array(z.string())
81 | .optional()
82 | .describe("Sequences where the API will stop generating further tokens."),
83 | thinkingConfig: thinkingConfigSchema,
84 | })
85 | .optional()
86 | .describe("Optional configuration for controlling the generation process.");
87 |
88 | // Safety setting schemas
89 | export const harmCategorySchema = z
90 | .enum([
91 | "HARM_CATEGORY_UNSPECIFIED",
92 | "HARM_CATEGORY_HATE_SPEECH",
93 | "HARM_CATEGORY_SEXUALLY_EXPLICIT",
94 | "HARM_CATEGORY_HARASSMENT",
95 | "HARM_CATEGORY_DANGEROUS_CONTENT",
96 | ])
97 | .describe("Category of harmful content to apply safety settings for.");
98 |
99 | export const harmBlockThresholdSchema = z
100 | .enum([
101 | "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
102 | "BLOCK_LOW_AND_ABOVE",
103 | "BLOCK_MEDIUM_AND_ABOVE",
104 | "BLOCK_ONLY_HIGH",
105 | "BLOCK_NONE",
106 | ])
107 | .describe(
108 | "Threshold for blocking harmful content. Higher thresholds block more content."
109 | );
110 |
111 | export const safetySettingSchema = z
112 | .object({
113 | category: harmCategorySchema,
114 | threshold: harmBlockThresholdSchema,
115 | })
116 | .describe(
117 | "Setting for controlling content safety for a specific harm category."
118 | );
119 |
120 | // History schemas for chat initialization
121 | const historyPartSchema = z
122 | .object({
123 | text: z.string().describe("Text content of the part."),
124 | // Note: Could add other part types like inlineData, functionCall, functionResponse later
125 | })
126 | .describe(
127 | "A part of a historical message, primarily text for initialization."
128 | );
129 |
130 | const historyContentSchema = z
131 | .object({
132 | role: z
133 | .enum(["user", "model"])
134 | .describe(
135 | "The role of the entity that generated this content (user or model)."
136 | ),
137 | parts: z
138 | .array(historyPartSchema)
139 | .min(1)
140 | .describe("An array of Parts making up the message content."),
141 | })
142 | .describe("A single message turn in the conversation history.");
143 |
144 | // Function declaration schemas (for tools)
145 | const functionParameterTypeSchema = z
146 | .enum(["OBJECT", "STRING", "NUMBER", "BOOLEAN", "ARRAY", "INTEGER"])
147 | .describe("The data type of the function parameter.");
148 |
149 | const baseFunctionParameterSchema = z.object({
150 | type: functionParameterTypeSchema,
151 | description: z
152 | .string()
153 | .optional()
154 | .describe("Description of the parameter's purpose."),
155 | enum: z
156 | .array(z.string())
157 | .optional()
158 | .describe("Allowed string values for an ENUM-like parameter."),
159 | });
160 |
161 | type FunctionParameterSchemaType = z.infer<
162 | typeof baseFunctionParameterSchema
163 | > & {
164 | properties?: { [key: string]: FunctionParameterSchemaType };
165 | required?: string[];
166 | items?: FunctionParameterSchemaType;
167 | };
168 |
169 | const functionParameterSchema: z.ZodType<FunctionParameterSchemaType> =
170 | baseFunctionParameterSchema
171 | .extend({
172 | properties: z.lazy(() => z.record(functionParameterSchema).optional()),
173 | required: z.lazy(() =>
174 | z
175 | .array(z.string())
176 | .optional()
177 | .describe("List of required property names for OBJECT types.")
178 | ),
179 | items: z.lazy(() =>
180 | functionParameterSchema
181 | .optional()
182 | .describe(
183 | "Defines the schema for items if the parameter type is ARRAY."
184 | )
185 | ),
186 | })
187 | .describe(
188 | "Schema defining a single parameter for a function declaration, potentially recursive."
189 | );
190 |
191 | const functionParameterPropertiesSchema = z
192 | .record(functionParameterSchema)
193 | .describe("Defines nested properties if the parameter type is OBJECT.");
194 |
195 | export const functionDeclarationSchema = z
196 | .object({
197 | name: z.string().min(1).describe("The name of the function to be called."),
198 | description: z
199 | .string()
200 | .min(1)
201 | .describe("A description of what the function does."),
202 | parameters: z
203 | .object({
204 | type: z
205 | .literal("OBJECT")
206 | .describe("The top-level parameters structure must be an OBJECT."),
207 | properties: functionParameterPropertiesSchema.describe(
208 | "Defines the parameters the function accepts."
209 | ),
210 | required: z
211 | .array(z.string())
212 | .optional()
213 | .describe("List of required parameter names at the top level."),
214 | })
215 | .describe("Schema defining the parameters the function accepts."),
216 | })
217 | .describe(
218 | "Declaration of a single function that the Gemini model can request to call."
219 | );
220 |
221 | // Tool schema for chat
222 | const toolSchema = z
223 | .object({
224 | functionDeclarations: z
225 | .array(functionDeclarationSchema)
226 | .optional()
227 | .describe("List of function declarations for this tool."),
228 | })
229 | .describe("Represents a tool definition containing function declarations.");
230 |
231 | // Tool config schema
232 | const functionCallingConfigSchema = z
233 | .object({
234 | mode: z
235 | .enum(["AUTO", "ANY", "NONE"])
236 | .optional()
237 | .describe("The function calling mode."),
238 | allowedFunctionNames: z
239 | .array(z.string())
240 | .optional()
241 | .describe("Optional list of function names allowed."),
242 | })
243 | .optional();
244 |
245 | const toolConfigSchema = z
246 | .object({
247 | functionCallingConfig: functionCallingConfigSchema,
248 | })
249 | .optional()
250 | .describe(
251 | "Optional. Per-request tool configuration, e.g., to force function calling mode."
252 | );
253 |
254 | // Function response schema for send_function_result
255 | const functionResponseInputSchema = z
256 | .object({
257 | name: z
258 | .string()
259 | .min(1)
260 | .describe(
261 | "Required. The name of the function that was called by the model."
262 | ),
263 | response: z
264 | .record(z.unknown())
265 | .describe(
266 | "Required. The JSON object result returned by the function execution."
267 | ),
268 | })
269 | .describe(
270 | "Represents the result of a single function execution to be sent back to the model."
271 | );
272 |
273 | // System instruction schema
274 | const systemInstructionSchema = z
275 | .object({
276 | parts: z.array(
277 | z.object({
278 | text: z.string(),
279 | })
280 | ),
281 | })
282 | .optional()
283 | .describe(
284 | "Optional. A system instruction to guide the model's behavior for the entire session."
285 | );
286 |
287 | // Main parameters schema with conditional fields based on operation
288 | export const GEMINI_CHAT_PARAMS = {
289 | operation: chatOperationSchema,
290 |
291 | // Common fields used across operations
292 | sessionId: z
293 | .string()
294 | .uuid()
295 | .optional()
296 | .describe(
297 | "Required for 'send_message' and 'send_function_result'. The unique identifier of the chat session."
298 | ),
299 |
300 | // Fields for 'start' operation
301 | modelName: ModelNameSchema.optional().describe(
302 | "Optional for 'start'. The name of the Gemini model to use for this chat session. If omitted, uses server default."
303 | ),
304 | history: z
305 | .array(historyContentSchema)
306 | .optional()
307 | .describe(
308 | "Optional for 'start'. Initial conversation turns to seed the chat session. Must alternate between 'user' and 'model' roles."
309 | ),
310 | systemInstruction: systemInstructionSchema,
311 |
312 | // Fields for 'send_message' operation
313 | message: z
314 | .string()
315 | .min(1)
316 | .optional()
317 | .describe(
318 | "Required for 'send_message'. The text message content to send to the model."
319 | ),
320 |
321 | // Fields for 'send_function_result' operation
322 | functionResponses: z
323 | .array(functionResponseInputSchema)
324 | .min(1)
325 | .optional()
326 | .describe(
327 | "Required for 'send_function_result'. Array containing the results of function calls executed by the client. Note: This array is JSON.stringify'd before being passed to the Gemini API."
328 | ),
329 |
330 | // Shared optional configuration fields
331 | tools: z
332 | .array(toolSchema)
333 | .optional()
334 | .describe(
335 | "Optional. Tools (function declarations) the model may use. For 'start', sets session-wide tools. For 'send_message', overrides session tools for this turn."
336 | ),
337 | toolConfig: toolConfigSchema,
338 | generationConfig: generationConfigSchema,
339 | safetySettings: z
340 | .array(safetySettingSchema)
341 | .optional()
342 | .describe(
343 | "Optional. Safety settings to apply. For 'start', sets session-wide settings. For other operations, overrides session settings for this turn."
344 | ),
345 | cachedContentName: z
346 | .string()
347 | .min(1)
348 | .optional()
349 | .describe(
350 | "Optional. Identifier for cached content in format 'cachedContents/...' to use with this operation."
351 | ),
352 | modelPreferences: ModelPreferencesSchema.optional(),
353 | };
354 |
355 | // Type helper
356 | export type GeminiChatArgs = z.infer<z.ZodObject<typeof GEMINI_CHAT_PARAMS>>;
357 |
```
--------------------------------------------------------------------------------
/src/services/gemini/GeminiCacheService.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { GoogleGenAI, CachedContent } from "@google/genai";
2 | import {
3 | GeminiApiError,
4 | GeminiResourceNotFoundError,
5 | GeminiInvalidParameterError,
6 | } from "../../utils/errors.js";
7 | import { logger } from "../../utils/logger.js";
8 | import { CachedContentMetadata } from "../../types/index.js";
9 | import { Content, Tool, ToolConfig, CacheId } from "./GeminiTypes.js";
10 |
11 | /**
12 | * Service for handling cache-related operations for the Gemini service.
13 | * Manages creation, listing, retrieval, and manipulation of cached content.
14 | */
15 | export class GeminiCacheService {
16 | private genAI: GoogleGenAI;
17 |
18 | /**
19 | * Creates a new instance of the GeminiCacheService.
20 | * @param genAI The GoogleGenAI instance to use for API calls
21 | */
22 | constructor(genAI: GoogleGenAI) {
23 | this.genAI = genAI;
24 | }
25 |
26 | /**
27 | * Creates a cached content entry in the Gemini API.
28 | *
29 | * @param modelName The model to use for this cached content
30 | * @param contents The conversation contents to cache
31 | * @param options Additional options for the cache (displayName, systemInstruction, ttl, tools, toolConfig)
32 | * @returns Promise resolving to the cached content metadata
33 | */
34 | public async createCache(
35 | modelName: string,
36 | contents: Content[],
37 | options?: {
38 | displayName?: string;
39 | systemInstruction?: Content | string;
40 | ttl?: string;
41 | tools?: Tool[];
42 | toolConfig?: ToolConfig;
43 | }
44 | ): Promise<CachedContentMetadata> {
45 | try {
46 | logger.debug(`Creating cache for model: ${modelName}`);
47 |
48 | // Process systemInstruction if it's a string
49 | let formattedSystemInstruction: Content | undefined;
50 | if (options?.systemInstruction) {
51 | if (typeof options.systemInstruction === "string") {
52 | formattedSystemInstruction = {
53 | parts: [{ text: options.systemInstruction }],
54 | };
55 | } else {
56 | formattedSystemInstruction = options.systemInstruction;
57 | }
58 | }
59 |
60 | // Create config object for the request
61 | const cacheConfig = {
62 | contents,
63 | displayName: options?.displayName,
64 | systemInstruction: formattedSystemInstruction,
65 | ttl: options?.ttl,
66 | tools: options?.tools,
67 | toolConfig: options?.toolConfig,
68 | };
69 |
70 | // Create the cache entry
71 | const cacheData = await this.genAI.caches.create({
72 | model: modelName,
73 | config: cacheConfig,
74 | });
75 |
76 | // Return the mapped metadata
77 | return this.mapSdkCacheToMetadata(cacheData);
78 | } catch (error: unknown) {
79 | logger.error(
80 | `Error creating cache: ${error instanceof Error ? error.message : String(error)}`,
81 | error
82 | );
83 | throw new GeminiApiError(
84 | `Failed to create cache: ${error instanceof Error ? error.message : String(error)}`,
85 | error
86 | );
87 | }
88 | }
89 |
90 | /**
91 | * Lists cached content entries in the Gemini API.
92 | *
93 | * @param pageSize Optional maximum number of entries to return
94 | * @param pageToken Optional token for pagination
95 | * @returns Promise resolving to an object with caches array and optional nextPageToken
96 | */
97 | public async listCaches(
98 | pageSize?: number,
99 | pageToken?: string
100 | ): Promise<{ caches: CachedContentMetadata[]; nextPageToken?: string }> {
101 | try {
102 | logger.debug(
103 | `Listing caches with pageSize: ${pageSize}, pageToken: ${pageToken}`
104 | );
105 |
106 | // Prepare list parameters
107 | const listParams: Record<string, number | string> = {};
108 |
109 | if (pageSize !== undefined) {
110 | listParams.pageSize = pageSize;
111 | }
112 |
113 | if (pageToken) {
114 | listParams.pageToken = pageToken;
115 | }
116 |
117 | // Call the caches.list method
118 | const response = await this.genAI.caches.list(listParams);
119 |
120 | const caches: CachedContentMetadata[] = [];
121 | let nextPageToken: string | undefined;
122 |
123 | // Handle the response in a more generic way to accommodate different API versions
124 | if (response && typeof response === "object") {
125 | if ("caches" in response && Array.isArray(response.caches)) {
126 | // Standard response format - cast to our TypeScript interface for validation
127 | for (const cache of response.caches) {
128 | caches.push(this.mapSdkCacheToMetadata(cache));
129 | }
130 | // Use optional chaining to safely access nextPageToken
131 | nextPageToken = (
132 | response as {
133 | caches: Record<string, unknown>[];
134 | nextPageToken?: string;
135 | }
136 | ).nextPageToken;
137 | } else if ("page" in response && response.page) {
138 | // Pager-like object in v0.10.0
139 | const cacheList = Array.from(response.page);
140 | for (const cache of cacheList) {
141 | caches.push(this.mapSdkCacheToMetadata(cache));
142 | }
143 |
144 | // Check if there's a next page
145 | const hasNextPage =
146 | typeof response === "object" &&
147 | "hasNextPage" in response &&
148 | typeof response.hasNextPage === "function"
149 | ? response.hasNextPage()
150 | : false;
151 |
152 | if (hasNextPage) {
153 | nextPageToken = "next_page_available";
154 | }
155 | } else if (Array.isArray(response)) {
156 | // Direct array response
157 | for (const cache of response) {
158 | caches.push(this.mapSdkCacheToMetadata(cache));
159 | }
160 | }
161 | }
162 |
163 | return { caches, nextPageToken };
164 | } catch (error: unknown) {
165 | logger.error(
166 | `Error listing caches: ${error instanceof Error ? error.message : String(error)}`,
167 | error
168 | );
169 | throw new GeminiApiError(
170 | `Failed to list caches: ${error instanceof Error ? error.message : String(error)}`,
171 | error
172 | );
173 | }
174 | }
175 |
176 | /**
177 | * Gets a specific cached content entry's metadata from the Gemini API.
178 | *
179 | * @param cacheId The ID of the cached content to retrieve (format: "cachedContents/{id}")
180 | * @returns Promise resolving to the cached content metadata
181 | */
182 | public async getCache(cacheId: CacheId): Promise<CachedContentMetadata> {
183 | try {
184 | logger.debug(`Getting cache metadata for: ${cacheId}`);
185 |
186 | // Validate the cacheId format
187 | if (!cacheId.startsWith("cachedContents/")) {
188 | throw new GeminiInvalidParameterError(
189 | `Cache ID must be in the format "cachedContents/{id}", received: ${cacheId}`
190 | );
191 | }
192 |
193 | // Get the cache metadata
194 | const cacheData = await this.genAI.caches.get({ name: cacheId });
195 |
196 | return this.mapSdkCacheToMetadata(cacheData);
197 | } catch (error: unknown) {
198 | // Check for specific error patterns in the error message
199 | if (error instanceof Error) {
200 | if (
201 | error.message.includes("not found") ||
202 | error.message.includes("404")
203 | ) {
204 | throw new GeminiResourceNotFoundError("Cache", cacheId, error);
205 | }
206 | }
207 |
208 | logger.error(
209 | `Error getting cache: ${error instanceof Error ? error.message : String(error)}`,
210 | error
211 | );
212 | throw new GeminiApiError(
213 | `Failed to get cache: ${error instanceof Error ? error.message : String(error)}`,
214 | error
215 | );
216 | }
217 | }
218 |
219 | /**
220 | * Updates a cached content entry in the Gemini API.
221 | *
222 | * @param cacheId The ID of the cached content to update (format: "cachedContents/{id}")
223 | * @param updates The updates to apply to the cached content (ttl, displayName)
224 | * @returns Promise resolving to the updated cached content metadata
225 | */
226 | public async updateCache(
227 | cacheId: CacheId,
228 | updates: { ttl?: string; displayName?: string }
229 | ): Promise<CachedContentMetadata> {
230 | try {
231 | logger.debug(`Updating cache: ${cacheId}`);
232 |
233 | // Validate the cacheId format
234 | if (!cacheId.startsWith("cachedContents/")) {
235 | throw new GeminiInvalidParameterError(
236 | `Cache ID must be in the format "cachedContents/{id}", received: ${cacheId}`
237 | );
238 | }
239 |
240 | // Update the cache
241 | const cacheData = await this.genAI.caches.update({
242 | name: cacheId,
243 | config: updates,
244 | });
245 |
246 | return this.mapSdkCacheToMetadata(cacheData);
247 | } catch (error: unknown) {
248 | // Check for specific error patterns in the error message
249 | if (error instanceof Error) {
250 | if (
251 | error.message.includes("not found") ||
252 | error.message.includes("404")
253 | ) {
254 | throw new GeminiResourceNotFoundError("Cache", cacheId, error);
255 | }
256 | }
257 |
258 | logger.error(
259 | `Error updating cache: ${error instanceof Error ? error.message : String(error)}`,
260 | error
261 | );
262 | throw new GeminiApiError(
263 | `Failed to update cache: ${error instanceof Error ? error.message : String(error)}`,
264 | error
265 | );
266 | }
267 | }
268 |
269 | /**
270 | * Deletes a cached content entry from the Gemini API.
271 | *
272 | * @param cacheId The ID of the cached content to delete (format: "cachedContents/{id}")
273 | * @returns Promise resolving to an object with success flag
274 | */
275 | public async deleteCache(cacheId: CacheId): Promise<{ success: boolean }> {
276 | try {
277 | logger.debug(`Deleting cache: ${cacheId}`);
278 |
279 | // Validate the cacheId format
280 | if (!cacheId.startsWith("cachedContents/")) {
281 | throw new GeminiInvalidParameterError(
282 | `Cache ID must be in the format "cachedContents/{id}", received: ${cacheId}`
283 | );
284 | }
285 |
286 | // Delete the cache
287 | await this.genAI.caches.delete({ name: cacheId });
288 |
289 | return { success: true };
290 | } catch (error: unknown) {
291 | // Check for specific error patterns in the error message
292 | if (error instanceof Error) {
293 | if (
294 | error.message.includes("not found") ||
295 | error.message.includes("404")
296 | ) {
297 | throw new GeminiResourceNotFoundError("Cache", cacheId, error);
298 | }
299 | }
300 |
301 | logger.error(
302 | `Error deleting cache: ${error instanceof Error ? error.message : String(error)}`,
303 | error
304 | );
305 | throw new GeminiApiError(
306 | `Failed to delete cache: ${error instanceof Error ? error.message : String(error)}`,
307 | error
308 | );
309 | }
310 | }
311 |
312 | /**
313 | * Helper method to map cached content response data to our CachedContentMetadata interface
314 | *
315 | * @param cacheData The cache data from the Gemini API
316 | * @returns The mapped CachedContentMetadata object
317 | */
318 | private mapSdkCacheToMetadata(
319 | cacheData: CachedContent
320 | ): CachedContentMetadata {
321 | if (!cacheData.name) {
322 | throw new Error("Invalid cache data received: missing required name");
323 | }
324 |
325 | // In SDK v0.10.0, the structure might be slightly different
326 | // Constructing CachedContentMetadata with fallback values where needed
327 | return {
328 | name: cacheData.name,
329 | displayName: cacheData.displayName || "",
330 | createTime: cacheData.createTime || new Date().toISOString(),
331 | updateTime: cacheData.updateTime || new Date().toISOString(),
332 | expirationTime: cacheData.expireTime,
333 | model: cacheData.model || "",
334 | state: "ACTIVE", // Default to ACTIVE since CachedContent does not have a status/state property
335 | usageMetadata: {
336 | totalTokenCount:
337 | typeof cacheData.usageMetadata?.totalTokenCount !== "undefined"
338 | ? cacheData.usageMetadata.totalTokenCount
339 | : 0,
340 | },
341 | };
342 | }
343 | }
344 |
```
--------------------------------------------------------------------------------
/src/services/gemini/GeminiPromptTemplates.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Collection of specialized prompt templates for different review scenarios.
3 | * These templates are designed to provide tailored guidance to the Gemini model
4 | * for various types of code reviews.
5 | */
6 |
7 | /**
8 | * Base template for code reviews with placeholders for context
9 | */
10 | export const baseReviewTemplate = `You are a senior software engineer reviewing the following code changes.
11 | {{repositoryContext}}
12 |
13 | {{focusInstructions}}
14 |
15 | Analyze the following git diff:
16 | \`\`\`diff
17 | {{diffContent}}
18 | \`\`\`
19 |
20 | Provide your code review with the following sections:
21 | 1. Summary of Changes: A brief overview of what's changing
22 | 2. Key Observations: The most important findings from your review
23 | 3. Detailed Review: File-by-file analysis of the changes
24 | 4. Recommendations: Specific suggestions for improvements
25 | 5. Questions: Any points needing clarification from the author
26 |
27 | Be concise yet thorough in your analysis.`;
28 |
29 | /**
30 | * Enhanced template for security-focused reviews
31 | */
32 | export const securityReviewTemplate = `You are a senior security engineer with OWASP certification and extensive experience in identifying vulnerabilities in software applications.
33 |
34 | {{repositoryContext}}
35 |
36 | TASK: Perform a comprehensive security review of the following code changes, focusing on potential vulnerabilities and security risks.
37 |
38 | STEP 1: Analyze the code for common security issues:
39 | - Injection vulnerabilities (SQL, NoSQL, command injection, etc.)
40 | - Authentication and authorization flaws
41 | - Sensitive data exposure
42 | - Security misconfigurations
43 | - Cross-site scripting (XSS) and cross-site request forgery (CSRF)
44 | - Broken access control
45 | - Insecure deserialization
46 | - Using components with known vulnerabilities
47 | - Input validation issues
48 | - Cryptographic problems
49 |
50 | STEP 2: For each identified issue:
51 | - Describe the vulnerability in detail
52 | - Assess its severity (Critical, High, Medium, Low)
53 | - Explain the potential impact
54 | - Provide a concrete remediation approach with code examples
55 |
56 | Analyze the following git diff:
57 | \`\`\`diff
58 | {{diffContent}}
59 | \`\`\`
60 |
61 | Your response MUST follow this format:
62 | 1. EXECUTIVE SUMMARY: Brief overview of security posture
63 | 2. CRITICAL ISSUES: Must be fixed immediately
64 | 3. HIGH PRIORITY ISSUES: Should be fixed soon
65 | 4. MEDIUM/LOW PRIORITY ISSUES: Address when possible
66 | 5. SECURE CODING RECOMMENDATIONS: Best practices to implement`;
67 |
68 | /**
69 | * Enhanced template for performance-focused reviews
70 | */
71 | export const performanceReviewTemplate = `You are a performance optimization expert with deep knowledge of runtime characteristics and profiling techniques.
72 |
73 | {{repositoryContext}}
74 |
75 | TASK: Perform a detailed performance analysis of the following code, identifying optimization opportunities and potential bottlenecks.
76 |
77 | STEP 1: Systematically analyze each section of code for:
78 | - Algorithm efficiency and complexity (provide Big O analysis)
79 | - Resource consumption patterns (memory, CPU, I/O)
80 | - Database query performance and optimization
81 | - Concurrency and parallelism opportunities
82 | - Caching potential and data access patterns
83 | - Unnecessary computation or redundant operations
84 |
85 | STEP 2: For each identified performance issue:
86 | - Describe the specific performance problem
87 | - Estimate the performance impact (critical, significant, moderate, minor)
88 | - Explain why it's problematic
89 | - Provide a specific optimization solution with code examples
90 | - Note any tradeoffs the optimization might introduce
91 |
92 | Analyze the following git diff:
93 | \`\`\`diff
94 | {{diffContent}}
95 | \`\`\`
96 |
97 | Your response MUST follow this format:
98 | 1. PERFORMANCE SUMMARY: Overall assessment with key metrics
99 | 2. CRITICAL BOTTLENECKS: Highest-impact issues to address first
100 | 3. SIGNIFICANT OPTIMIZATIONS: Important improvements with measurable impact
101 | 4. MINOR OPTIMIZATIONS: Small enhancements for completeness
102 | 5. MONITORING RECOMMENDATIONS: Suggestions for ongoing performance measurement`;
103 |
104 | /**
105 | * Enhanced template for architecture-focused reviews
106 | */
107 | export const architectureReviewTemplate = `You are a senior software architect with expertise in designing scalable, maintainable software systems.
108 |
109 | {{repositoryContext}}
110 |
111 | TASK: Perform an architectural analysis of the following code changes, focusing on design patterns, component relationships, and system structure.
112 |
113 | STEP 1: Analyze the architectural aspects of the code:
114 | - Design pattern implementation and appropriateness
115 | - Component responsibilities and cohesion
116 | - Interface design and abstraction
117 | - Dependency management and coupling
118 | - Modularity and extensibility
119 | - Separation of concerns
120 | - Error handling strategies
121 | - Consistency with architectural principles
122 |
123 | STEP 2: For each architectural observation:
124 | - Describe the architectural element or decision
125 | - Analyze its impact on the overall system
126 | - Evaluate adherence to SOLID principles and other architectural best practices
127 | - Suggest improvements with rationale
128 |
129 | Analyze the following git diff:
130 | \`\`\`diff
131 | {{diffContent}}
132 | \`\`\`
133 |
134 | Your response MUST follow this format:
135 | 1. ARCHITECTURAL OVERVIEW: Summary of the code's architecture
136 | 2. STRENGTHS: Positive architectural aspects of the code
137 | 3. CONCERNS: Architectural issues or anti-patterns identified
138 | 4. REFACTORING RECOMMENDATIONS: Suggestions for architectural improvements
139 | 5. LONG-TERM CONSIDERATIONS: How these changes affect system evolution`;
140 |
141 | /**
142 | * Enhanced template for bug-focused reviews
143 | */
144 | export const bugReviewTemplate = `You are a quality assurance engineer with expertise in identifying logic flaws and edge cases in software.
145 |
146 | {{repositoryContext}}
147 |
148 | TASK: Perform a thorough analysis of the following code changes to identify potential bugs, edge cases, and logical errors.
149 |
150 | STEP 1: Analyze the code for common bug sources:
151 | - Off-by-one errors
152 | - Null/undefined handling issues
153 | - Edge case oversights
154 | - Race conditions
155 | - Resource leaks
156 | - Error handling gaps
157 | - Typos in critical code
158 | - Incorrect assumptions
159 | - Boundary condition failures
160 | - Exception handling problems
161 |
162 | STEP 2: For each potential bug:
163 | - Describe the issue and why it's problematic
164 | - Explain the conditions under which it would occur
165 | - Assess its severity and potential impact
166 | - Provide a fix with code examples
167 | - Suggest tests that would catch the issue
168 |
169 | Analyze the following git diff:
170 | \`\`\`diff
171 | {{diffContent}}
172 | \`\`\`
173 |
174 | Your response MUST follow this format:
175 | 1. BUG RISK SUMMARY: Overview of potential issues
176 | 2. CRITICAL BUGS: Issues that could cause system failure or data corruption
177 | 3. MAJOR BUGS: Significant functional issues that need addressing
178 | 4. MINOR BUGS: Less severe issues that should be fixed
179 | 5. TEST RECOMMENDATIONS: Tests to implement to prevent similar bugs`;
180 |
181 | /**
182 | * Enhanced template for general comprehensive reviews
183 | */
184 | export const generalReviewTemplate = `You are a senior software engineer with expertise across multiple domains including security, performance, architecture, and testing.
185 |
186 | {{repositoryContext}}
187 |
188 | TASK: Perform a comprehensive review of the following code changes, covering all aspects of software quality.
189 |
190 | I want you to follow a specific review process:
191 |
192 | STEP 1: Understand the overall purpose
193 | - Identify what problem the code is solving
194 | - Determine how it fits into the broader application
195 |
196 | STEP 2: Analyze code quality
197 | - Readability and naming conventions
198 | - Function/method size and complexity
199 | - Comments and documentation
200 | - Consistency with existing patterns
201 |
202 | STEP 3: Evaluate correctness
203 | - Potential bugs and edge cases
204 | - Error handling completeness
205 | - Test coverage adequacy
206 |
207 | STEP 4: Consider performance
208 | - Inefficient algorithms or patterns
209 | - Resource utilization concerns
210 | - Optimization opportunities
211 |
212 | STEP 5: Assess maintainability
213 | - Extensibility for future changes
214 | - Coupling and cohesion
215 | - Clear separation of concerns
216 |
217 | STEP 6: Security review
218 | - Potential vulnerabilities
219 | - Input validation issues
220 | - Security best practices
221 |
222 | Analyze the following git diff:
223 | \`\`\`diff
224 | {{diffContent}}
225 | \`\`\`
226 |
227 | Your response MUST follow this format:
228 | 1. SUMMARY: Brief overview of the changes and their purpose
229 | 2. KEY OBSERVATIONS: Most important findings (positive and negative)
230 | 3. DETAILED REVIEW: Analysis by file with specific comments
231 | 4. RECOMMENDATIONS: Prioritized suggestions for improvement
232 | 5. QUESTIONS: Any clarifications needed from the developer`;
233 |
234 | /**
235 | * Replace placeholders in a template with actual values
236 | *
237 | * @param template Template string with placeholders
238 | * @param context Context object with values to replace placeholders
239 | * @returns Processed template with placeholders replaced
240 | */
241 | export function processTemplate(
242 | template: string,
243 | context: {
244 | repositoryContext?: string;
245 | diffContent: string;
246 | focusInstructions?: string;
247 | [key: string]: string | undefined;
248 | }
249 | ): string {
250 | let processedTemplate = template;
251 |
252 | // Replace each placeholder with its corresponding value
253 | for (const [key, value] of Object.entries(context)) {
254 | // Skip undefined values
255 | if (value === undefined) continue;
256 |
257 | // Convert the value to string if it's not already
258 | const stringValue = typeof value === "string" ? value : String(value);
259 |
260 | // Replace the placeholder with the value
261 | processedTemplate = processedTemplate.replace(
262 | new RegExp(`{{${key}}}`, "g"),
263 | stringValue
264 | );
265 | }
266 |
267 | // Remove any remaining placeholder
268 | processedTemplate = processedTemplate.replace(/{{[^{}]+}}/g, "");
269 |
270 | return processedTemplate;
271 | }
272 |
273 | /**
274 | * Get the appropriate template for a specific review focus
275 | *
276 | * @param reviewFocus The focus area for the review
277 | * @returns The template string for the specified focus
278 | */
279 | export function getReviewTemplate(
280 | reviewFocus: "security" | "performance" | "architecture" | "bugs" | "general"
281 | ): string {
282 | switch (reviewFocus) {
283 | case "security":
284 | return securityReviewTemplate;
285 | case "performance":
286 | return performanceReviewTemplate;
287 | case "architecture":
288 | return architectureReviewTemplate;
289 | case "bugs":
290 | return bugReviewTemplate;
291 | case "general":
292 | default:
293 | return generalReviewTemplate;
294 | }
295 | }
296 |
297 | /**
298 | * Generate focus-specific instructions for the base template
299 | *
300 | * @param reviewFocus The focus area for the review
301 | * @returns Instruction string for the specified focus
302 | */
303 | export function getFocusInstructions(
304 | reviewFocus: "security" | "performance" | "architecture" | "bugs" | "general"
305 | ): string {
306 | switch (reviewFocus) {
307 | case "security":
308 | return `
309 | Focus on identifying security vulnerabilities in the code changes, such as:
310 | - Input validation issues
311 | - Authentication/authorization flaws
312 | - Data exposure risks
313 | - Injection vulnerabilities
314 | - Insecure cryptography
315 | - CSRF/XSS vectors
316 | - Any other security concerns`;
317 | case "performance":
318 | return `
319 | Focus on identifying performance implications in the code changes, such as:
320 | - Algorithm complexity issues (O(n²) vs O(n))
321 | - Unnecessary computations or memory usage
322 | - Database query inefficiencies
323 | - Unoptimized loops or data structures
324 | - Potential memory leaks
325 | - Resource contention issues`;
326 | case "architecture":
327 | return `
328 | Focus on architectural aspects of the code changes, such as:
329 | - Design pattern conformance
330 | - Component responsibilities and cohesion
331 | - Dependency management
332 | - API design principles
333 | - Modularity and extensibility
334 | - Separation of concerns`;
335 | case "bugs":
336 | return `
337 | Focus on identifying potential bugs and logical errors in the code changes, such as:
338 | - Off-by-one errors
339 | - Null/undefined handling issues
340 | - Edge case oversights
341 | - Race conditions
342 | - Error handling gaps
343 | - Typos in critical code`;
344 | case "general":
345 | default:
346 | return `
347 | Provide a comprehensive review covering:
348 | - Code quality and readability
349 | - Potential bugs or errors
350 | - Performance implications
351 | - Security considerations
352 | - Architectural aspects
353 | - Best practices and style conventions`;
354 | }
355 | }
356 |
```
--------------------------------------------------------------------------------
/src/server.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { createServer } from "./createServer.js";
2 | import {
3 | logger,
4 | setServerState,
5 | startHealthCheckServer,
6 | ServerState,
7 | } from "./utils/index.js";
8 | import type { JsonRpcInitializeRequest } from "./types/serverTypes.js";
9 | import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
10 | import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
11 | import type { Transport } from "@modelcontextprotocol/sdk/server/mcp.js";
12 | import { SessionService } from "./services/SessionService.js";
13 | import express from "express";
14 | import { randomUUID } from "node:crypto";
15 | /**
16 | * Type guard to check if a value is a JSON-RPC 2.0 initialize request
17 | * @param value - The value to check
18 | * @returns true if the value matches the JSON-RPC initialize request structure
19 | */
20 | const isInitializeRequest = (
21 | value: unknown
22 | ): value is JsonRpcInitializeRequest => {
23 | // Early exit for non-objects
24 | if (!value || typeof value !== "object" || value === null) {
25 | return false;
26 | }
27 |
28 | const obj = value as Record<string, unknown>;
29 |
30 | // Check required JSON-RPC 2.0 fields
31 | if (obj.jsonrpc !== "2.0") {
32 | return false;
33 | }
34 |
35 | if (obj.method !== "initialize") {
36 | return false;
37 | }
38 |
39 | // Check id exists and is string or number (per JSON-RPC spec)
40 | if (typeof obj.id !== "string" && typeof obj.id !== "number") {
41 | return false;
42 | }
43 |
44 | return true;
45 | };
46 |
47 | // Server state tracking
48 | const serverState: ServerState = {
49 | isRunning: false,
50 | startTime: null,
51 | transport: null,
52 | server: null,
53 | healthCheckServer: null,
54 | mcpClientService: null, // Add McpClientService to server state
55 | };
56 |
57 | // Share server state with the health check module
58 | setServerState(serverState);
59 |
60 | // Map to store transports by session ID for HTTP mode
61 | const httpTransports: Record<string, StreamableHTTPServerTransport> = {};
62 |
63 | /**
64 | * Sets up HTTP server for Streamable HTTP transport
65 | */
66 | async function setupHttpServer(
67 | mcpServer: { connect: (transport: Transport) => Promise<void> },
68 | _sessionService: SessionService
69 | ): Promise<void> {
70 | const app = express();
71 | app.use(express.json());
72 |
73 | const port = parseInt(
74 | process.env.MCP_SERVER_PORT || process.env.MCP_WS_PORT || "8080",
75 | 10
76 | );
77 |
78 | // CORS middleware
79 | app.use((req, res, next) => {
80 | res.header("Access-Control-Allow-Origin", "*");
81 | res.header("Access-Control-Allow-Methods", "GET, POST, OPTIONS, DELETE");
82 | res.header(
83 | "Access-Control-Allow-Headers",
84 | "Content-Type, Accept, Authorization, Mcp-Session-Id, Last-Event-ID"
85 | );
86 | if (req.method === "OPTIONS") {
87 | res.status(204).send();
88 | return;
89 | }
90 | next();
91 | });
92 |
93 | // MCP endpoint
94 | app.post("/mcp", async (req, res) => {
95 | try {
96 | const sessionId = req.headers["mcp-session-id"] as string;
97 | let transport: StreamableHTTPServerTransport;
98 |
99 | if (sessionId && httpTransports[sessionId]) {
100 | // Reuse existing transport
101 | transport = httpTransports[sessionId];
102 | } else if (!sessionId && isInitializeRequest(req.body)) {
103 | // Create new transport for initialization
104 | const sessionIdGenerator =
105 | process.env.MCP_ENABLE_STREAMING === "true"
106 | ? () => randomUUID()
107 | : undefined;
108 |
109 | transport = new StreamableHTTPServerTransport({
110 | sessionIdGenerator,
111 | onsessioninitialized: (sid: string) => {
112 | logger.info(`HTTP session initialized: ${sid}`);
113 | httpTransports[sid] = transport;
114 | },
115 | });
116 |
117 | // Set up cleanup handler
118 | transport.onclose = () => {
119 | const sid = transport.sessionId;
120 | if (sid && httpTransports[sid]) {
121 | logger.info(`HTTP transport closed for session ${sid}`);
122 | delete httpTransports[sid];
123 | }
124 | };
125 |
126 | // Connect transport to MCP server
127 | await mcpServer.connect(transport);
128 | } else {
129 | res.status(400).json({
130 | jsonrpc: "2.0",
131 | error: {
132 | code: -32000,
133 | message: "Bad Request: No valid session ID provided",
134 | },
135 | });
136 | return;
137 | }
138 |
139 | await transport.handleRequest(req, res, req.body);
140 | } catch (error) {
141 | logger.error("Error handling HTTP request:", error);
142 | res.status(500).json({
143 | jsonrpc: "2.0",
144 | error: {
145 | code: -32603,
146 | message: "Internal error",
147 | },
148 | });
149 | }
150 | });
151 |
152 | // GET endpoint for SSE streaming
153 | app.get("/mcp", async (req, res) => {
154 | try {
155 | const sessionId = req.headers["mcp-session-id"] as string;
156 | if (!sessionId || !httpTransports[sessionId]) {
157 | res.status(404).json({
158 | error: "Session not found",
159 | });
160 | return;
161 | }
162 |
163 | const transport = httpTransports[sessionId];
164 | await transport.handleRequest(req, res);
165 | } catch (error) {
166 | logger.error("Error handling GET request:", error);
167 | res.status(500).json({
168 | error: "Internal error",
169 | });
170 | }
171 | });
172 |
173 | // Start HTTP server
174 | const httpServer = app.listen(port, () => {
175 | logger.info(`HTTP server listening on port ${port} with /mcp endpoint`);
176 | });
177 |
178 | // Store HTTP server in state for cleanup
179 | serverState.httpServer = httpServer;
180 | }
181 |
182 | const main = async () => {
183 | try {
184 | const sessionService = new SessionService(
185 | parseInt(process.env.MCP_SESSION_TIMEOUT || "3600", 10)
186 | );
187 |
188 | const { server, mcpClientService } = createServer();
189 | serverState.server = server;
190 | // Type compatibility: McpClientService implements McpClientServiceLike interface
191 | // The actual service has all required methods (disconnect, closeAllConnections)
192 | serverState.mcpClientService = mcpClientService; // Store the client service instance
193 | serverState.sessionService = sessionService; // Store the session service instance
194 | logger.info("Starting MCP server...");
195 |
196 | // Start health check server if enabled
197 | if (process.env.ENABLE_HEALTH_CHECK !== "false") {
198 | logger.info("Starting health check server...");
199 | const healthServer = startHealthCheckServer();
200 | serverState.healthCheckServer = healthServer;
201 | }
202 |
203 | // Choose transport based on environment
204 | let transport: Transport | null;
205 | // Use MCP_TRANSPORT, but fall back to MCP_TRANSPORT_TYPE for backwards compatibility
206 | const transportType =
207 | process.env.MCP_TRANSPORT || process.env.MCP_TRANSPORT_TYPE || "stdio";
208 |
209 | if (transportType === "sse") {
210 | // SSE uses the StreamableHTTPServerTransport
211 | transport = null; // No persistent transport needed
212 | logger.info("Transport selected", {
213 | requested: transportType,
214 | selected: "streamable",
215 | fallback: false,
216 | message:
217 | "SSE transport - using StreamableHTTPServerTransport via HTTP endpoint",
218 | timestamp: new Date().toISOString(),
219 | });
220 | } else if (transportType === "http" || transportType === "streamable") {
221 | // For HTTP/Streamable transport, we don't need a persistent transport
222 | // Individual requests will create their own transports
223 | transport = null; // No persistent transport needed
224 | logger.info("Transport selected", {
225 | requested: transportType,
226 | selected: "streamable",
227 | fallback: false,
228 | message:
229 | "HTTP transport - individual requests will create their own transports",
230 | timestamp: new Date().toISOString(),
231 | });
232 | } else if (transportType === "streaming") {
233 | const fallbackReason = "Streaming transport not currently implemented";
234 | logger.warn("Transport fallback", {
235 | requested: transportType,
236 | selected: "stdio",
237 | fallback: true,
238 | reason: fallbackReason,
239 | timestamp: new Date().toISOString(),
240 | });
241 | transport = new StdioServerTransport();
242 | logger.info("Using stdio transport (fallback)");
243 | } else {
244 | // Default to stdio for anything else
245 | transport = new StdioServerTransport();
246 | logger.info("Transport selected", {
247 | requested: transportType || "default",
248 | selected: "stdio",
249 | fallback: false,
250 | message: "Using stdio transport",
251 | timestamp: new Date().toISOString(),
252 | });
253 | }
254 |
255 | serverState.transport = transport;
256 | if (transport) {
257 | logger.info(`Connecting transport: ${transport}`);
258 | await server.connect(transport);
259 | } else {
260 | logger.info("No persistent transport - using HTTP-only mode");
261 | }
262 |
263 | // Set up HTTP server for streamable/SSE transport if requested
264 | if (
265 | transportType === "http" ||
266 | transportType === "streamable" ||
267 | transportType === "sse"
268 | ) {
269 | await setupHttpServer(server, sessionService);
270 | }
271 |
272 | // Update server state
273 | serverState.isRunning = true;
274 | serverState.startTime = Date.now();
275 |
276 | logger.info("MCP Server connected and listening.");
277 |
278 | // For HTTP-only mode, keep the process alive
279 | if (
280 | transportType === "http" ||
281 | transportType === "streamable" ||
282 | transportType === "sse"
283 | ) {
284 | // Keep the process alive since we don't have a persistent transport
285 | // The HTTP server will handle all requests
286 | process.on("SIGINT", () => shutdown("SIGINT"));
287 | process.on("SIGTERM", () => shutdown("SIGTERM"));
288 | }
289 | } catch (error) {
290 | logger.error("Failed to start server:", error);
291 | process.exit(1); // Exit if server fails to start
292 | }
293 | };
294 |
295 | // Graceful shutdown handling
296 | const shutdown = async (signal: string) => {
297 | logger.info(`${signal} signal received: closing MCP server`);
298 |
299 | // Track if any shutdown process fails
300 | let hasError = false;
301 |
302 | // Stop session service cleanup interval if it exists
303 | if (serverState.sessionService) {
304 | try {
305 | logger.info("Stopping session service cleanup interval...");
306 | serverState.sessionService.stopCleanupInterval();
307 | logger.info("Session service cleanup interval stopped.");
308 | } catch (error) {
309 | hasError = true;
310 | logger.error("Error stopping session service cleanup interval:", error);
311 | }
312 | }
313 |
314 | // Close all MCP client connections
315 | if (serverState.mcpClientService) {
316 | try {
317 | logger.info("Closing all MCP client connections...");
318 | serverState.mcpClientService.closeAllConnections();
319 | logger.info("MCP client connections closed.");
320 | } catch (error) {
321 | hasError = true;
322 | logger.error("Error closing MCP client connections:", error);
323 | }
324 | }
325 |
326 | if (serverState.isRunning && serverState.server) {
327 | try {
328 | // Disconnect the server if it exists and has a disconnect method
329 | if (typeof serverState.server.disconnect === "function") {
330 | await serverState.server.disconnect();
331 | }
332 |
333 | serverState.isRunning = false;
334 | logger.info("MCP Server shutdown completed successfully");
335 | } catch (error) {
336 | hasError = true;
337 | logger.error("Error during MCP server shutdown:", error);
338 | }
339 | }
340 |
341 | // Close HTTP server if it exists
342 | if (serverState.httpServer) {
343 | try {
344 | logger.info("Closing HTTP server...");
345 | serverState.httpServer.close();
346 | logger.info("HTTP server closed successfully");
347 | } catch (error) {
348 | hasError = true;
349 | logger.error("Error during HTTP server shutdown:", error);
350 | }
351 | }
352 |
353 | // Close health check server if it exists
354 | if (serverState.healthCheckServer) {
355 | try {
356 | logger.info("Closing health check server...");
357 | serverState.healthCheckServer.close();
358 | logger.info("Health check server closed successfully");
359 | } catch (error) {
360 | hasError = true;
361 | logger.error("Error during health check server shutdown:", error);
362 | }
363 | }
364 |
365 | // Exit with appropriate code
366 | process.exit(hasError ? 1 : 0);
367 | };
368 |
369 | // Register shutdown handlers
370 | process.on("SIGTERM", () => shutdown("SIGTERM"));
371 | process.on("SIGINT", () => shutdown("SIGINT"));
372 |
373 | // If this is the main module, start the server
374 | if (import.meta.url === `file://${process.argv[1]}`) {
375 | main();
376 | }
377 |
```
--------------------------------------------------------------------------------
/tests/unit/utils/errors.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | // Import directly from the MCP SDK to ensure we're using the same class reference
3 | import { McpError, ErrorCode } from "@modelcontextprotocol/sdk/types.js";
4 | // Import local error classes
5 | import {
6 | ValidationError,
7 | NotFoundError,
8 | ConfigurationError,
9 | ServiceError,
10 | GeminiApiError,
11 | mapToMcpError,
12 | mapToolErrorToMcpError,
13 | ToolErrorLike,
14 | } from "../../../src/utils/errors.js";
15 |
16 | describe("mapToMcpError", () => {
17 | const TOOL_NAME = "test_tool";
18 |
19 | it("should return McpError instances directly", () => {
20 | const originalError = new McpError(
21 | ErrorCode.InvalidParams,
22 | "Original MCP error"
23 | );
24 | const mappedError = mapToMcpError(originalError, TOOL_NAME);
25 | expect(mappedError).toBe(originalError);
26 | });
27 |
28 | it("should map ValidationError to InvalidParams", () => {
29 | const validationError = new ValidationError("Invalid input");
30 | const mappedError = mapToMcpError(validationError, TOOL_NAME);
31 |
32 | // Check error code and message content
33 | expect(mappedError.code).toBe(ErrorCode.InvalidParams);
34 | expect(mappedError.message).toContain("Validation error");
35 | expect(mappedError.message).toContain("Invalid input");
36 | });
37 |
38 | it("should map NotFoundError to InvalidRequest", () => {
39 | const notFoundError = new NotFoundError("Resource not found");
40 | const mappedError = mapToMcpError(notFoundError, TOOL_NAME);
41 |
42 | // Check error code and message content
43 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
44 | expect(mappedError.message).toContain("Resource not found");
45 | });
46 |
47 | it("should map ConfigurationError to InternalError", () => {
48 | const configError = new ConfigurationError("Invalid configuration");
49 | const mappedError = mapToMcpError(configError, TOOL_NAME);
50 |
51 | expect(mappedError.code).toBe(ErrorCode.InternalError); // Changed from FailedPrecondition
52 | expect(mappedError.message).toContain("Configuration error");
53 | expect(mappedError.message).toContain("Invalid configuration");
54 | });
55 |
56 | it("should map quota-related GeminiApiError to InternalError", () => {
57 | const quotaError = new GeminiApiError("Quota exceeded for this resource");
58 | const mappedError = mapToMcpError(quotaError, TOOL_NAME);
59 |
60 | expect(mappedError.code).toBe(ErrorCode.InternalError); // Changed from ResourceExhausted
61 | expect(mappedError.message).toContain("Quota exceeded");
62 | });
63 |
64 | it("should map rate limit GeminiApiError to InternalError", () => {
65 | const rateLimitError = new GeminiApiError(
66 | "Rate limit hit for this operation"
67 | );
68 | const mappedError = mapToMcpError(rateLimitError, TOOL_NAME);
69 |
70 | expect(mappedError.code).toBe(ErrorCode.InternalError); // Changed from ResourceExhausted
71 | expect(mappedError.message).toContain("rate limit hit");
72 | });
73 |
74 | it("should map permission-related GeminiApiError to InvalidRequest", () => {
75 | const permissionError = new GeminiApiError(
76 | "Permission denied for this operation"
77 | );
78 | const mappedError = mapToMcpError(permissionError, TOOL_NAME);
79 |
80 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest); // Changed from PermissionDenied
81 | expect(mappedError.message).toContain("Permission denied");
82 | });
83 |
84 | it("should map not-found GeminiApiError to InvalidRequest", () => {
85 | const notFoundError = new GeminiApiError("Resource does not exist");
86 | const mappedError = mapToMcpError(notFoundError, TOOL_NAME);
87 |
88 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
89 | expect(mappedError.message).toContain("Resource not found");
90 | });
91 |
92 | it("should map invalid argument GeminiApiError to InvalidParams", () => {
93 | const invalidParamError = new GeminiApiError("Invalid argument provided");
94 | const mappedError = mapToMcpError(invalidParamError, TOOL_NAME);
95 |
96 | expect(mappedError.code).toBe(ErrorCode.InvalidParams);
97 | expect(mappedError.message).toContain("Invalid parameters");
98 | });
99 |
100 | it("should map safety-related GeminiApiError to InvalidRequest", () => {
101 | const safetyError = new GeminiApiError(
102 | "Content blocked by safety settings"
103 | );
104 | const mappedError = mapToMcpError(safetyError, TOOL_NAME);
105 |
106 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
107 | expect(mappedError.message).toContain("Content blocked by safety settings");
108 | });
109 |
110 | it("should map File API not supported errors to InvalidRequest", () => {
111 | const apiError = new GeminiApiError(
112 | "File API is not supported on Vertex AI"
113 | );
114 | const mappedError = mapToMcpError(apiError, TOOL_NAME);
115 |
116 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest); // Changed from FailedPrecondition
117 | expect(mappedError.message).toContain("Operation not supported");
118 | });
119 |
120 | it("should map other GeminiApiError to InternalError", () => {
121 | const otherApiError = new GeminiApiError("Unknown API error");
122 | const mappedError = mapToMcpError(otherApiError, TOOL_NAME);
123 |
124 | expect(mappedError.code).toBe(ErrorCode.InternalError);
125 | expect(mappedError.message).toContain("Gemini API Error");
126 | });
127 |
128 | it("should map ServiceError to InternalError", () => {
129 | const serviceError = new ServiceError("Service processing failed");
130 | const mappedError = mapToMcpError(serviceError, TOOL_NAME);
131 |
132 | expect(mappedError.code).toBe(ErrorCode.InternalError);
133 | expect(mappedError.message).toContain("Service error");
134 | });
135 |
136 | it("should map standard Error to InternalError", () => {
137 | const standardError = new Error("Standard error occurred");
138 | const mappedError = mapToMcpError(standardError, TOOL_NAME);
139 |
140 | expect(mappedError.code).toBe(ErrorCode.InternalError);
141 | expect(mappedError.message).toContain(TOOL_NAME);
142 | expect(mappedError.message).toContain("Standard error occurred");
143 | });
144 |
145 | it("should handle string errors", () => {
146 | const stringError = "String error message";
147 | const mappedError = mapToMcpError(stringError, TOOL_NAME);
148 |
149 | expect(mappedError.code).toBe(ErrorCode.InternalError);
150 | expect(mappedError.message).toContain(stringError);
151 | });
152 |
153 | it("should handle object errors", () => {
154 | const objectError = { errorCode: 500, message: "Object error" };
155 | const mappedError = mapToMcpError(objectError, TOOL_NAME);
156 |
157 | expect(mappedError.code).toBe(ErrorCode.InternalError);
158 | // Should contain stringified version of the object
159 | expect(mappedError.message).toContain("Object error");
160 | });
161 |
162 | it("should handle null/undefined errors", () => {
163 | const nullError = null;
164 | const mappedError = mapToMcpError(nullError, TOOL_NAME);
165 |
166 | expect(mappedError.code).toBe(ErrorCode.InternalError);
167 | expect(mappedError.message).toContain("An unknown error occurred");
168 | });
169 |
170 | // Testing if the error details are properly handled in mapping
171 | it("should handle errors with details", () => {
172 | // Create an error with details
173 | const errorWithDetails = new GeminiApiError("API error with details", {
174 | key: "value",
175 | });
176 |
177 | // Directly check the original error - it should have details
178 | expect(errorWithDetails).toHaveProperty("details");
179 | expect(errorWithDetails.details).toEqual({ key: "value" });
180 |
181 | // Map it to an McpError
182 | const mappedError = mapToMcpError(errorWithDetails, TOOL_NAME);
183 |
184 | // Basic assertions
185 | expect(mappedError).toBeInstanceOf(Object);
186 | expect(mappedError).not.toBeNull();
187 | expect(mappedError.code).toBe(ErrorCode.InternalError);
188 |
189 | // Verify mapping occurs correctly
190 | expect(mappedError).toBeInstanceOf(McpError);
191 | expect(mappedError.message).toContain("API error with details");
192 |
193 | // If McpError supports data property for error details, check it
194 | if ("data" in mappedError) {
195 | expect(mappedError.data).toBeDefined();
196 | }
197 | });
198 | });
199 |
200 | describe("mapToolErrorToMcpError", () => {
201 | const TOOL_NAME = "test_tool";
202 |
203 | it("should handle ToolErrorLike objects with code and message", () => {
204 | const toolError: ToolErrorLike = {
205 | code: "INVALID_ARGUMENT",
206 | message: "Invalid parameter provided",
207 | details: { parameter: "value" },
208 | };
209 |
210 | const mappedError = mapToolErrorToMcpError(toolError, TOOL_NAME);
211 |
212 | expect(mappedError).toBeInstanceOf(McpError);
213 | expect(mappedError.code).toBe(ErrorCode.InvalidParams);
214 | expect(mappedError.message).toContain("Invalid parameter provided");
215 | });
216 |
217 | it("should map QUOTA and RATE_LIMIT codes to InternalError", () => {
218 | const quotaError: ToolErrorLike = {
219 | code: "QUOTA_EXCEEDED",
220 | message: "API quota exceeded",
221 | };
222 |
223 | const mappedError = mapToolErrorToMcpError(quotaError, TOOL_NAME);
224 |
225 | expect(mappedError.code).toBe(ErrorCode.InternalError);
226 | expect(mappedError.message).toContain("API quota or rate limit exceeded");
227 | });
228 |
229 | it("should map PERMISSION and AUTH codes to InvalidRequest", () => {
230 | const permissionError: ToolErrorLike = {
231 | code: "PERMISSION_DENIED",
232 | message: "Permission denied",
233 | };
234 |
235 | const mappedError = mapToolErrorToMcpError(permissionError, TOOL_NAME);
236 |
237 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
238 | expect(mappedError.message).toContain("Permission denied");
239 | });
240 |
241 | it("should map NOT_FOUND codes to InvalidRequest", () => {
242 | const notFoundError: ToolErrorLike = {
243 | code: "RESOURCE_NOT_FOUND",
244 | message: "Resource not found",
245 | };
246 |
247 | const mappedError = mapToolErrorToMcpError(notFoundError, TOOL_NAME);
248 |
249 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
250 | expect(mappedError.message).toContain("Resource not found");
251 | });
252 |
253 | it("should map INVALID and ARGUMENT codes to InvalidParams", () => {
254 | const invalidError: ToolErrorLike = {
255 | code: "INVALID_ARGUMENT",
256 | message: "Invalid argument",
257 | };
258 |
259 | const mappedError = mapToolErrorToMcpError(invalidError, TOOL_NAME);
260 |
261 | expect(mappedError.code).toBe(ErrorCode.InvalidParams);
262 | expect(mappedError.message).toContain("Invalid parameters");
263 | });
264 |
265 | it("should map UNSUPPORTED codes to InvalidRequest", () => {
266 | const unsupportedError: ToolErrorLike = {
267 | code: "OPERATION_NOT_SUPPORTED",
268 | message: "Operation not supported",
269 | };
270 |
271 | const mappedError = mapToolErrorToMcpError(unsupportedError, TOOL_NAME);
272 |
273 | expect(mappedError.code).toBe(ErrorCode.InvalidRequest);
274 | expect(mappedError.message).toContain("Operation not supported");
275 | });
276 |
277 | it("should handle objects without code property", () => {
278 | const toolError = {
279 | message: "Generic error message",
280 | details: { info: "additional info" },
281 | };
282 |
283 | const mappedError = mapToolErrorToMcpError(toolError, TOOL_NAME);
284 |
285 | expect(mappedError.code).toBe(ErrorCode.InternalError);
286 | expect(mappedError.message).toContain("Generic error message");
287 | });
288 |
289 | it("should handle non-object errors", () => {
290 | const simpleError = "Simple string error";
291 |
292 | const mappedError = mapToolErrorToMcpError(simpleError, TOOL_NAME);
293 |
294 | expect(mappedError.code).toBe(ErrorCode.InternalError);
295 | expect(mappedError.message).toContain(`Error in ${TOOL_NAME}`);
296 | });
297 |
298 | it("should handle null and undefined errors", () => {
299 | const nullError = mapToolErrorToMcpError(null, TOOL_NAME);
300 | const undefinedError = mapToolErrorToMcpError(undefined, TOOL_NAME);
301 |
302 | expect(nullError.code).toBe(ErrorCode.InternalError);
303 | expect(undefinedError.code).toBe(ErrorCode.InternalError);
304 | expect(nullError.message).toContain(`Error in ${TOOL_NAME}`);
305 | expect(undefinedError.message).toContain(`Error in ${TOOL_NAME}`);
306 | });
307 |
308 | it("should preserve error details when available", () => {
309 | const toolError: ToolErrorLike = {
310 | code: "TEST_ERROR",
311 | message: "Test error message",
312 | details: { context: "test context", id: 123 },
313 | };
314 |
315 | const mappedError = mapToolErrorToMcpError(toolError, TOOL_NAME);
316 |
317 | // Check that details are preserved (assuming McpError supports details/data)
318 | expect(mappedError).toBeDefined();
319 | expect(mappedError.message).toContain("Test error message");
320 | });
321 | });
322 |
```