This is page 5 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
--------------------------------------------------------------------------------
/src/services/ModelSelectionService.ts:
--------------------------------------------------------------------------------
```typescript
1 | import {
2 | ModelConfiguration,
3 | ModelSelectionCriteria,
4 | ModelCapabilities,
5 | ModelScore,
6 | ModelSelectionHistory,
7 | ModelPerformanceMetrics,
8 | } from "../types/index.js";
9 | import { logger } from "../utils/logger.js";
10 |
11 | export class ModelSelectionService {
12 | private modelCache: Map<string, ModelCapabilities>;
13 | private performanceMetrics: Map<string, ModelPerformanceMetrics>;
14 | private selectionHistory: ModelSelectionHistory[];
15 | private readonly maxHistorySize = 500;
16 |
17 | constructor(private config: ModelConfiguration) {
18 | this.modelCache = new Map();
19 | this.performanceMetrics = new Map();
20 | this.selectionHistory = [];
21 | this.initializeModelCache();
22 | }
23 |
24 | private initializeModelCache(): void {
25 | Object.entries(this.config.capabilities).forEach(
26 | ([model, capabilities]) => {
27 | this.modelCache.set(model, capabilities);
28 | }
29 | );
30 | }
31 |
32 | async selectOptimalModel(criteria: ModelSelectionCriteria): Promise<string> {
33 | const startTime = Date.now();
34 | try {
35 | const candidateModels = this.getCandidateModels(criteria);
36 | if (candidateModels.length === 0) {
37 | logger.warn(
38 | "[ModelSelectionService] No candidate models found for criteria",
39 | { criteria }
40 | );
41 | return criteria.fallbackModel || this.config.default;
42 | }
43 |
44 | const selectedModel = this.selectBestModel(candidateModels, criteria);
45 |
46 | const selectionTime = Date.now() - startTime;
47 | this.recordSelection(
48 | criteria,
49 | selectedModel,
50 | candidateModels,
51 | selectionTime
52 | );
53 |
54 | logger.debug("[ModelSelectionService] Model selected", {
55 | selectedModel,
56 | criteria,
57 | });
58 |
59 | return selectedModel;
60 | } catch (error) {
61 | logger.error("[ModelSelectionService] Model selection failed", {
62 | error,
63 | criteria,
64 | });
65 | return criteria.fallbackModel || this.config.default;
66 | }
67 | }
68 |
69 | private getCandidateModels(criteria: ModelSelectionCriteria): string[] {
70 | let baseModels: string[];
71 |
72 | switch (criteria.taskType) {
73 | case "text-generation":
74 | baseModels = this.config.textGeneration;
75 | break;
76 | case "image-generation":
77 | baseModels = this.config.imageGeneration;
78 | break;
79 | case "video-generation":
80 | baseModels = this.config.videoGeneration;
81 | break;
82 | case "code-review":
83 | baseModels = this.config.codeReview;
84 | break;
85 | case "reasoning":
86 | baseModels = this.config.complexReasoning;
87 | break;
88 | case "multimodal":
89 | baseModels = this.config.textGeneration.filter((model) => {
90 | const caps = this.modelCache.get(model);
91 | return (
92 | caps && (caps.imageInput || caps.videoInput || caps.audioInput)
93 | );
94 | });
95 | break;
96 | default:
97 | baseModels = this.config.textGeneration;
98 | }
99 |
100 | // Filter out models that don't have capabilities defined
101 | baseModels = baseModels.filter((model) => this.modelCache.has(model));
102 |
103 | if (criteria.requiredCapabilities) {
104 | return baseModels.filter((model) => {
105 | const capabilities = this.modelCache.get(model);
106 | if (!capabilities) return false;
107 |
108 | return criteria.requiredCapabilities!.every((capability) => {
109 | const value = capabilities[capability];
110 | if (typeof value === "boolean") return value;
111 | if (typeof value === "string") return value !== "none";
112 | return Boolean(value);
113 | });
114 | });
115 | }
116 |
117 | return baseModels;
118 | }
119 |
120 | private selectBestModel(
121 | models: string[],
122 | criteria: ModelSelectionCriteria
123 | ): string {
124 | if (models.length === 0) {
125 | return criteria.fallbackModel || this.config.default;
126 | }
127 |
128 | // Score-based selection that considers performance metrics
129 | const scoredModels = models.map((model) => ({
130 | model,
131 | score: this.calculateModelScore(model, criteria),
132 | }));
133 |
134 | scoredModels.sort((a, b) => b.score - a.score);
135 | return scoredModels[0].model;
136 | }
137 |
138 | private sortModelsByPreference(
139 | models: string[],
140 | criteria: ModelSelectionCriteria
141 | ): string[] {
142 | const preferCost =
143 | criteria.preferCost || this.config.routing.preferCostEffective;
144 | const preferSpeed = criteria.preferSpeed || this.config.routing.preferSpeed;
145 | const preferQuality =
146 | criteria.preferQuality || this.config.routing.preferQuality;
147 |
148 | return models.sort((a, b) => {
149 | const capsA = this.modelCache.get(a)!;
150 | const capsB = this.modelCache.get(b)!;
151 |
152 | // Primary sorting: cost preference
153 | if (preferCost) {
154 | const costComparison = this.compareCost(capsA.costTier, capsB.costTier);
155 | if (costComparison !== 0) return costComparison;
156 | }
157 |
158 | // Secondary sorting: speed preference
159 | if (preferSpeed) {
160 | const speedComparison = this.compareSpeed(
161 | capsA.speedTier,
162 | capsB.speedTier
163 | );
164 | if (speedComparison !== 0) return speedComparison;
165 | }
166 |
167 | // Tertiary sorting: quality preference (default)
168 | if (preferQuality) {
169 | const qualityComparison = this.compareQuality(capsA, capsB);
170 | if (qualityComparison !== 0) return qualityComparison;
171 | }
172 |
173 | return 0; // Equal preference
174 | });
175 | }
176 |
177 | private compareCost(costA: string, costB: string): number {
178 | const costOrder = { low: 0, medium: 1, high: 2 };
179 | const orderA = costOrder[costA as keyof typeof costOrder] ?? 1;
180 | const orderB = costOrder[costB as keyof typeof costOrder] ?? 1;
181 | return orderA - orderB; // Lower cost wins
182 | }
183 |
184 | private compareSpeed(speedA: string, speedB: string): number {
185 | const speedOrder = { fast: 0, medium: 1, slow: 2 };
186 | const orderA = speedOrder[speedA as keyof typeof speedOrder] ?? 1;
187 | const orderB = speedOrder[speedB as keyof typeof speedOrder] ?? 1;
188 | return orderA - orderB; // Faster wins
189 | }
190 |
191 | private compareQuality(
192 | capsA: ModelCapabilities,
193 | capsB: ModelCapabilities
194 | ): number {
195 | const reasoningOrder = { none: 0, basic: 1, good: 2, excellent: 3 };
196 | const codeOrder = { none: 0, basic: 1, good: 2, excellent: 3 };
197 |
198 | const reasoningA =
199 | reasoningOrder[capsA.complexReasoning as keyof typeof reasoningOrder] ??
200 | 0;
201 | const reasoningB =
202 | reasoningOrder[capsB.complexReasoning as keyof typeof reasoningOrder] ??
203 | 0;
204 |
205 | if (reasoningA !== reasoningB) {
206 | return reasoningB - reasoningA; // Higher reasoning wins
207 | }
208 |
209 | const codeA = codeOrder[capsA.codeExecution as keyof typeof codeOrder] ?? 0;
210 | const codeB = codeOrder[capsB.codeExecution as keyof typeof codeOrder] ?? 0;
211 |
212 | if (codeA !== codeB) {
213 | return codeB - codeA; // Higher code execution wins
214 | }
215 |
216 | // Additional quality factors
217 | if (capsA.contextWindow !== capsB.contextWindow) {
218 | return capsB.contextWindow - capsA.contextWindow; // Larger context wins
219 | }
220 |
221 | return 0;
222 | }
223 |
224 | getModelCapabilities(modelName: string): ModelCapabilities | undefined {
225 | return this.modelCache.get(modelName);
226 | }
227 |
228 | isModelAvailable(modelName: string): boolean {
229 | return this.modelCache.has(modelName);
230 | }
231 |
232 | getAvailableModels(): string[] {
233 | return Array.from(this.modelCache.keys());
234 | }
235 |
236 | validateModelForTask(
237 | modelName: string,
238 | taskType: ModelSelectionCriteria["taskType"]
239 | ): boolean {
240 | const capabilities = this.modelCache.get(modelName);
241 | if (!capabilities) return false;
242 |
243 | switch (taskType) {
244 | case "text-generation":
245 | return capabilities.textGeneration;
246 | case "image-generation":
247 | return capabilities.imageGeneration;
248 | case "video-generation":
249 | return capabilities.videoGeneration;
250 | case "code-review":
251 | return capabilities.codeExecution !== "none";
252 | case "reasoning":
253 | return capabilities.complexReasoning !== "none";
254 | case "multimodal":
255 | return (
256 | capabilities.imageInput ||
257 | capabilities.videoInput ||
258 | capabilities.audioInput
259 | );
260 | default:
261 | return capabilities.textGeneration;
262 | }
263 | }
264 |
265 | updateConfiguration(newConfig: ModelConfiguration): void {
266 | this.config = newConfig;
267 | this.modelCache.clear();
268 | this.initializeModelCache();
269 | logger.info("[ModelSelectionService] Configuration updated");
270 | }
271 |
272 | updatePerformanceMetrics(
273 | modelName: string,
274 | latency: number,
275 | success: boolean
276 | ): void {
277 | const existing = this.performanceMetrics.get(modelName) || {
278 | totalCalls: 0,
279 | avgLatency: 0,
280 | successRate: 0,
281 | lastUpdated: new Date(),
282 | };
283 |
284 | const newTotalCalls = existing.totalCalls + 1;
285 | const newAvgLatency =
286 | (existing.avgLatency * existing.totalCalls + latency) / newTotalCalls;
287 | const successCount =
288 | existing.successRate * existing.totalCalls + (success ? 1 : 0);
289 | const newSuccessRate = successCount / newTotalCalls;
290 |
291 | this.performanceMetrics.set(modelName, {
292 | totalCalls: newTotalCalls,
293 | avgLatency: newAvgLatency,
294 | successRate: newSuccessRate,
295 | lastUpdated: new Date(),
296 | });
297 | }
298 |
299 | getPerformanceMetrics(): Map<string, ModelPerformanceMetrics> {
300 | return new Map(this.performanceMetrics);
301 | }
302 |
303 | getSelectionHistory(limit?: number): ModelSelectionHistory[] {
304 | const history = [...this.selectionHistory];
305 | return limit ? history.slice(-limit) : history;
306 | }
307 |
308 | private recordSelection(
309 | criteria: ModelSelectionCriteria,
310 | selectedModel: string,
311 | candidateModels: string[],
312 | selectionTime: number
313 | ): void {
314 | const scores: ModelScore[] = candidateModels.map((model) => ({
315 | model,
316 | score: this.calculateModelScore(model, criteria),
317 | capabilities: this.modelCache.get(model)!,
318 | }));
319 |
320 | const record: ModelSelectionHistory = {
321 | timestamp: new Date(),
322 | criteria,
323 | selectedModel,
324 | candidateModels,
325 | scores,
326 | selectionTime,
327 | };
328 |
329 | this.selectionHistory.push(record);
330 |
331 | if (this.selectionHistory.length > this.maxHistorySize) {
332 | this.selectionHistory.shift();
333 | }
334 | }
335 |
336 | private calculateModelScore(
337 | model: string,
338 | criteria: ModelSelectionCriteria
339 | ): number {
340 | const capabilities = this.modelCache.get(model);
341 | if (!capabilities) return 0;
342 |
343 | let score = 0;
344 |
345 | // Base score from routing preferences
346 | if (criteria.preferCost) {
347 | const costScore =
348 | capabilities.costTier === "low"
349 | ? 3
350 | : capabilities.costTier === "medium"
351 | ? 2
352 | : 1;
353 | score += costScore * 0.4;
354 | }
355 |
356 | if (criteria.preferSpeed) {
357 | const speedScore =
358 | capabilities.speedTier === "fast"
359 | ? 3
360 | : capabilities.speedTier === "medium"
361 | ? 2
362 | : 1;
363 | score += speedScore * 0.4;
364 | }
365 |
366 | if (criteria.preferQuality) {
367 | const reasoningScore =
368 | capabilities.complexReasoning === "excellent"
369 | ? 3
370 | : capabilities.complexReasoning === "good"
371 | ? 2
372 | : 1;
373 | score += reasoningScore * 0.4;
374 | }
375 |
376 | // URL context scoring - prefer models with larger context windows for URL-heavy requests
377 | if (criteria.urlCount && criteria.urlCount > 0) {
378 | // Bonus for models with large context windows when processing URLs
379 | if (capabilities.contextWindow >= 1000000) {
380 | score += Math.min(criteria.urlCount / 5, 2.0); // Up to 2 points for many URLs
381 | } else if (capabilities.contextWindow >= 500000) {
382 | score += Math.min(criteria.urlCount / 10, 1.0); // Up to 1 point for medium context
383 | }
384 |
385 | // Bonus for estimated content size handling
386 | if (
387 | criteria.estimatedUrlContentSize &&
388 | criteria.estimatedUrlContentSize > 0
389 | ) {
390 | const sizeInTokens = criteria.estimatedUrlContentSize / 4; // Rough estimate: 4 chars per token
391 | const contextUtilization = sizeInTokens / capabilities.contextWindow;
392 |
393 | // Prefer models that won't be overwhelmed by the content size
394 | if (contextUtilization < 0.3) {
395 | score += 1.5; // Comfortable fit
396 | } else if (contextUtilization < 0.6) {
397 | score += 0.5; // Acceptable fit
398 | } else if (contextUtilization > 0.8) {
399 | score -= 2.0; // Penalize models that might struggle
400 | }
401 | }
402 |
403 | // Slight bonus for models that support URL context natively (Gemini 2.5 models)
404 | if (model.includes("gemini-2.5")) {
405 | score += 0.5;
406 | }
407 | }
408 |
409 | // Performance metrics influence (heavily weighted)
410 | const metrics = this.performanceMetrics.get(model);
411 | if (metrics && metrics.totalCalls >= 5) {
412 | // Strong preference for models with good performance history
413 | score += metrics.successRate * 2.0;
414 | // Prefer lower latency (significant impact)
415 | const latencyScore = Math.max(0, 1 - metrics.avgLatency / 2000);
416 | score += latencyScore * 1.5;
417 | }
418 |
419 | return score;
420 | }
421 | }
422 |
```
--------------------------------------------------------------------------------
/tests/unit/services/ModelSelectionService.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { ModelSelectionService } from "../../../src/services/ModelSelectionService.js";
3 | import {
4 | ModelConfiguration,
5 | ModelCapabilitiesMap,
6 | } from "../../../src/types/index.js";
7 |
8 | describe("ModelSelectionService", () => {
9 | let service: ModelSelectionService;
10 | let mockConfig: ModelConfiguration;
11 |
12 | beforeEach(() => {
13 | const capabilities: ModelCapabilitiesMap = {
14 | "gemini-2.5-pro-preview-05-06": {
15 | textGeneration: true,
16 | imageInput: true,
17 | videoInput: true,
18 | audioInput: true,
19 | imageGeneration: false,
20 | videoGeneration: false,
21 | codeExecution: "excellent",
22 | complexReasoning: "excellent",
23 | costTier: "high",
24 | speedTier: "medium",
25 | maxTokens: 65536,
26 | contextWindow: 1048576,
27 | supportsFunctionCalling: true,
28 | supportsSystemInstructions: true,
29 | supportsCaching: true,
30 | },
31 | "gemini-2.5-flash-preview-05-20": {
32 | textGeneration: true,
33 | imageInput: true,
34 | videoInput: true,
35 | audioInput: true,
36 | imageGeneration: false,
37 | videoGeneration: false,
38 | codeExecution: "excellent",
39 | complexReasoning: "excellent",
40 | costTier: "medium",
41 | speedTier: "fast",
42 | maxTokens: 65536,
43 | contextWindow: 1048576,
44 | supportsFunctionCalling: true,
45 | supportsSystemInstructions: true,
46 | supportsCaching: true,
47 | },
48 | "gemini-2.0-flash": {
49 | textGeneration: true,
50 | imageInput: true,
51 | videoInput: true,
52 | audioInput: true,
53 | imageGeneration: false,
54 | videoGeneration: false,
55 | codeExecution: "good",
56 | complexReasoning: "good",
57 | costTier: "medium",
58 | speedTier: "fast",
59 | maxTokens: 8192,
60 | contextWindow: 1048576,
61 | supportsFunctionCalling: true,
62 | supportsSystemInstructions: true,
63 | supportsCaching: true,
64 | },
65 | "gemini-1.5-flash": {
66 | textGeneration: true,
67 | imageInput: true,
68 | videoInput: true,
69 | audioInput: true,
70 | imageGeneration: false,
71 | videoGeneration: false,
72 | codeExecution: "basic",
73 | complexReasoning: "basic",
74 | costTier: "low",
75 | speedTier: "fast",
76 | maxTokens: 8192,
77 | contextWindow: 1000000,
78 | supportsFunctionCalling: true,
79 | supportsSystemInstructions: true,
80 | supportsCaching: true,
81 | },
82 | "imagen-3.0-generate-002": {
83 | textGeneration: false,
84 | imageInput: false,
85 | videoInput: false,
86 | audioInput: false,
87 | imageGeneration: true,
88 | videoGeneration: false,
89 | codeExecution: "none",
90 | complexReasoning: "none",
91 | costTier: "medium",
92 | speedTier: "medium",
93 | maxTokens: 0,
94 | contextWindow: 0,
95 | supportsFunctionCalling: false,
96 | supportsSystemInstructions: false,
97 | supportsCaching: false,
98 | },
99 | "gemini-2.0-flash-preview-image-generation": {
100 | textGeneration: true,
101 | imageInput: true,
102 | videoInput: false,
103 | audioInput: false,
104 | imageGeneration: true,
105 | videoGeneration: false,
106 | codeExecution: "basic",
107 | complexReasoning: "basic",
108 | costTier: "medium",
109 | speedTier: "medium",
110 | maxTokens: 8192,
111 | contextWindow: 32000,
112 | supportsFunctionCalling: false,
113 | supportsSystemInstructions: true,
114 | supportsCaching: false,
115 | },
116 | };
117 |
118 | mockConfig = {
119 | default: "gemini-2.5-flash-preview-05-20",
120 | textGeneration: [
121 | "gemini-2.5-pro-preview-05-06",
122 | "gemini-2.5-flash-preview-05-20",
123 | "gemini-2.0-flash",
124 | "gemini-1.5-flash",
125 | ],
126 | imageGeneration: [
127 | "imagen-3.0-generate-002",
128 | "gemini-2.0-flash-preview-image-generation",
129 | ],
130 | videoGeneration: [],
131 | codeReview: [
132 | "gemini-2.5-pro-preview-05-06",
133 | "gemini-2.5-flash-preview-05-20",
134 | "gemini-2.0-flash",
135 | ],
136 | complexReasoning: [
137 | "gemini-2.5-pro-preview-05-06",
138 | "gemini-2.5-flash-preview-05-20",
139 | ],
140 | capabilities,
141 | routing: {
142 | preferCostEffective: false,
143 | preferSpeed: false,
144 | preferQuality: true,
145 | },
146 | };
147 |
148 | service = new ModelSelectionService(mockConfig);
149 | });
150 |
151 | describe("selectOptimalModel", () => {
152 | it("should select a model for text generation", async () => {
153 | const model = await service.selectOptimalModel({
154 | taskType: "text-generation",
155 | complexityLevel: "simple",
156 | });
157 |
158 | expect(mockConfig.textGeneration).toContain(model);
159 | });
160 |
161 | it("should prefer cost-effective models when specified", async () => {
162 | const model = await service.selectOptimalModel({
163 | taskType: "text-generation",
164 | preferCost: true,
165 | });
166 |
167 | const capabilities = service.getModelCapabilities(model);
168 | expect(capabilities?.costTier).toBe("low");
169 | });
170 |
171 | it("should prefer fast models when speed is prioritized", async () => {
172 | const model = await service.selectOptimalModel({
173 | taskType: "text-generation",
174 | preferSpeed: true,
175 | });
176 |
177 | const capabilities = service.getModelCapabilities(model);
178 | expect(capabilities?.speedTier).toBe("fast");
179 | });
180 |
181 | it("should select high-quality models for complex tasks", async () => {
182 | const model = await service.selectOptimalModel({
183 | taskType: "reasoning",
184 | complexityLevel: "complex",
185 | preferQuality: true,
186 | });
187 |
188 | const capabilities = service.getModelCapabilities(model);
189 | expect(capabilities?.complexReasoning).toBe("excellent");
190 | });
191 |
192 | it("should return fallback model when no candidates match", async () => {
193 | const model = await service.selectOptimalModel({
194 | taskType: "text-generation",
195 | requiredCapabilities: ["imageGeneration"],
196 | fallbackModel: "gemini-1.5-flash",
197 | });
198 |
199 | expect(model).toBe("gemini-1.5-flash");
200 | });
201 |
202 | it("should select image generation models correctly", async () => {
203 | const model = await service.selectOptimalModel({
204 | taskType: "image-generation",
205 | });
206 |
207 | expect(mockConfig.imageGeneration).toContain(model);
208 | const capabilities = service.getModelCapabilities(model);
209 | expect(capabilities?.imageGeneration).toBe(true);
210 | });
211 |
212 | it("should filter models by required capabilities", async () => {
213 | const model = await service.selectOptimalModel({
214 | taskType: "text-generation",
215 | requiredCapabilities: ["supportsFunctionCalling", "supportsCaching"],
216 | });
217 |
218 | const capabilities = service.getModelCapabilities(model);
219 | expect(capabilities?.supportsFunctionCalling).toBe(true);
220 | expect(capabilities?.supportsCaching).toBe(true);
221 | });
222 | });
223 |
224 | describe("validateModelForTask", () => {
225 | it("should validate text generation models", () => {
226 | expect(
227 | service.validateModelForTask(
228 | "gemini-2.5-pro-preview-05-06",
229 | "text-generation"
230 | )
231 | ).toBe(true);
232 | expect(
233 | service.validateModelForTask(
234 | "imagen-3.0-generate-002",
235 | "text-generation"
236 | )
237 | ).toBe(false);
238 | });
239 |
240 | it("should validate image generation models", () => {
241 | expect(
242 | service.validateModelForTask(
243 | "imagen-3.0-generate-002",
244 | "image-generation"
245 | )
246 | ).toBe(true);
247 | expect(
248 | service.validateModelForTask(
249 | "gemini-2.5-pro-preview-05-06",
250 | "image-generation"
251 | )
252 | ).toBe(false);
253 | });
254 |
255 | it("should validate code review models", () => {
256 | expect(
257 | service.validateModelForTask(
258 | "gemini-2.5-pro-preview-05-06",
259 | "code-review"
260 | )
261 | ).toBe(true);
262 | expect(
263 | service.validateModelForTask("gemini-1.5-flash", "code-review")
264 | ).toBe(true);
265 | expect(
266 | service.validateModelForTask("imagen-3.0-generate-002", "code-review")
267 | ).toBe(false);
268 | });
269 |
270 | it("should validate multimodal models", () => {
271 | expect(
272 | service.validateModelForTask(
273 | "gemini-2.5-pro-preview-05-06",
274 | "multimodal"
275 | )
276 | ).toBe(true);
277 | expect(
278 | service.validateModelForTask("imagen-3.0-generate-002", "multimodal")
279 | ).toBe(false);
280 | });
281 | });
282 |
283 | describe("updatePerformanceMetrics", () => {
284 | it("should track performance metrics", () => {
285 | service.updatePerformanceMetrics(
286 | "gemini-2.5-pro-preview-05-06",
287 | 1000,
288 | true
289 | );
290 | service.updatePerformanceMetrics(
291 | "gemini-2.5-pro-preview-05-06",
292 | 1200,
293 | true
294 | );
295 | service.updatePerformanceMetrics(
296 | "gemini-2.5-pro-preview-05-06",
297 | 800,
298 | false
299 | );
300 |
301 | const metrics = service.getPerformanceMetrics();
302 | const proMetrics = metrics.get("gemini-2.5-pro-preview-05-06");
303 |
304 | expect(proMetrics).toBeDefined();
305 | expect(proMetrics?.totalCalls).toBe(3);
306 | expect(proMetrics?.avgLatency).toBe(1000);
307 | expect(proMetrics?.successRate).toBeCloseTo(0.667, 2);
308 | });
309 |
310 | it("should influence model selection based on performance", async () => {
311 | service.updatePerformanceMetrics(
312 | "gemini-2.5-flash-preview-05-20",
313 | 500,
314 | true
315 | );
316 | service.updatePerformanceMetrics(
317 | "gemini-2.5-flash-preview-05-20",
318 | 600,
319 | true
320 | );
321 | service.updatePerformanceMetrics(
322 | "gemini-2.5-flash-preview-05-20",
323 | 400,
324 | true
325 | );
326 | service.updatePerformanceMetrics(
327 | "gemini-2.5-flash-preview-05-20",
328 | 550,
329 | true
330 | );
331 | service.updatePerformanceMetrics(
332 | "gemini-2.5-flash-preview-05-20",
333 | 450,
334 | true
335 | );
336 |
337 | service.updatePerformanceMetrics(
338 | "gemini-2.5-pro-preview-05-06",
339 | 2000,
340 | false
341 | );
342 | service.updatePerformanceMetrics(
343 | "gemini-2.5-pro-preview-05-06",
344 | 1800,
345 | false
346 | );
347 | service.updatePerformanceMetrics(
348 | "gemini-2.5-pro-preview-05-06",
349 | 2200,
350 | false
351 | );
352 | service.updatePerformanceMetrics(
353 | "gemini-2.5-pro-preview-05-06",
354 | 1900,
355 | false
356 | );
357 | service.updatePerformanceMetrics(
358 | "gemini-2.5-pro-preview-05-06",
359 | 2100,
360 | false
361 | );
362 |
363 | const model = await service.selectOptimalModel({
364 | taskType: "text-generation",
365 | complexityLevel: "medium",
366 | });
367 |
368 | expect(model).toBe("gemini-2.5-flash-preview-05-20");
369 | });
370 | });
371 |
372 | describe("getSelectionHistory", () => {
373 | it("should track selection history", async () => {
374 | await service.selectOptimalModel({ taskType: "text-generation" });
375 | await service.selectOptimalModel({ taskType: "image-generation" });
376 |
377 | const history = service.getSelectionHistory();
378 | expect(history).toHaveLength(2);
379 | expect(history[0].criteria.taskType).toBe("text-generation");
380 | expect(history[1].criteria.taskType).toBe("image-generation");
381 | });
382 |
383 | it("should limit history size", async () => {
384 | for (let i = 0; i < 1200; i++) {
385 | await service.selectOptimalModel({ taskType: "text-generation" });
386 | }
387 |
388 | const history = service.getSelectionHistory();
389 | expect(history.length).toBeLessThanOrEqual(500);
390 | });
391 |
392 | it("should return limited history when requested", async () => {
393 | for (let i = 0; i < 10; i++) {
394 | await service.selectOptimalModel({ taskType: "text-generation" });
395 | }
396 |
397 | const limitedHistory = service.getSelectionHistory(5);
398 | expect(limitedHistory).toHaveLength(5);
399 | });
400 | });
401 |
402 | describe("isModelAvailable", () => {
403 | it("should check model availability", () => {
404 | expect(service.isModelAvailable("gemini-2.5-pro-preview-05-06")).toBe(
405 | true
406 | );
407 | expect(service.isModelAvailable("non-existent-model")).toBe(false);
408 | });
409 | });
410 |
411 | describe("getAvailableModels", () => {
412 | it("should return all available models", () => {
413 | const models = service.getAvailableModels();
414 | expect(models).toContain("gemini-2.5-pro-preview-05-06");
415 | expect(models).toContain("gemini-2.5-flash-preview-05-20");
416 | expect(models).toContain("gemini-1.5-flash");
417 | expect(models).toContain("imagen-3.0-generate-002");
418 | });
419 | });
420 |
421 | describe("updateConfiguration", () => {
422 | it("should update configuration and reinitialize cache", () => {
423 | const newConfig = {
424 | ...mockConfig,
425 | textGeneration: ["gemini-2.5-pro-preview-05-06"],
426 | };
427 |
428 | service.updateConfiguration(newConfig);
429 |
430 | const models = service.getAvailableModels();
431 | expect(models).toContain("gemini-2.5-pro-preview-05-06");
432 | });
433 | });
434 |
435 | describe("error handling", () => {
436 | it("should handle errors gracefully and return fallback", async () => {
437 | const corruptedService = new ModelSelectionService({
438 | ...mockConfig,
439 | capabilities: {},
440 | });
441 |
442 | const model = await corruptedService.selectOptimalModel({
443 | taskType: "text-generation",
444 | fallbackModel: "fallback-model",
445 | });
446 |
447 | expect(model).toBe("fallback-model");
448 | });
449 | });
450 | });
451 |
```
--------------------------------------------------------------------------------
/src/tools/geminiUrlAnalysisTool.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2 | import { z } from "zod";
3 | import { GeminiService } from "../services/index.js";
4 | import { logger } from "../utils/index.js";
5 | import { mapAnyErrorToMcpError } from "../utils/errors.js";
6 |
7 | // Tool Name and Description
8 | export const GEMINI_URL_ANALYSIS_TOOL_NAME = "gemini_url_analysis";
9 | export const GEMINI_URL_ANALYSIS_TOOL_DESCRIPTION = `
10 | Advanced URL analysis tool that fetches content from web pages and performs specialized analysis tasks.
11 | Supports various analysis types including summarization, comparison, information extraction, and Q&A.
12 | Automatically handles URL fetching, content processing, and intelligent model selection for optimal results.
13 | `;
14 |
15 | // Analysis types enum
16 | const analysisTypeSchema = z
17 | .enum([
18 | "summary",
19 | "comparison",
20 | "extraction",
21 | "qa",
22 | "sentiment",
23 | "fact-check",
24 | "content-classification",
25 | "readability",
26 | "seo-analysis",
27 | ])
28 | .describe("Type of analysis to perform on the URL content");
29 |
30 | // Extraction schema for structured data extraction
31 | const extractionSchemaSchema = z
32 | .record(z.unknown())
33 | .optional()
34 | .describe(
35 | "JSON schema or structure definition for extracting specific information from content"
36 | );
37 |
38 | // Parameters for the URL analysis tool
39 | export const GEMINI_URL_ANALYSIS_PARAMS = {
40 | urls: z
41 | .array(z.string().url())
42 | .min(1)
43 | .max(20)
44 | .describe("URLs to analyze (1-20 URLs supported)"),
45 |
46 | analysisType: analysisTypeSchema,
47 |
48 | query: z
49 | .string()
50 | .min(1)
51 | .optional()
52 | .describe("Specific query or instruction for the analysis"),
53 |
54 | extractionSchema: extractionSchemaSchema,
55 |
56 | questions: z
57 | .array(z.string())
58 | .optional()
59 | .describe("List of specific questions to answer (for Q&A analysis)"),
60 |
61 | compareBy: z
62 | .array(z.string())
63 | .optional()
64 | .describe("Specific aspects to compare when using comparison analysis"),
65 |
66 | outputFormat: z
67 | .enum(["text", "json", "markdown", "structured"])
68 | .default("text")
69 | .optional()
70 | .describe("Desired output format for the analysis results"),
71 |
72 | includeMetadata: z
73 | .boolean()
74 | .default(true)
75 | .optional()
76 | .describe(
77 | "Include URL metadata (title, description, etc.) in the analysis"
78 | ),
79 |
80 | fetchOptions: z
81 | .object({
82 | maxContentKb: z
83 | .number()
84 | .min(1)
85 | .max(1000)
86 | .default(100)
87 | .optional()
88 | .describe("Maximum content size per URL in KB"),
89 | timeoutMs: z
90 | .number()
91 | .min(1000)
92 | .max(30000)
93 | .default(10000)
94 | .optional()
95 | .describe("Fetch timeout per URL in milliseconds"),
96 | allowedDomains: z
97 | .array(z.string())
98 | .optional()
99 | .describe("Specific domains to allow for this request"),
100 | userAgent: z
101 | .string()
102 | .optional()
103 | .describe("Custom User-Agent header for URL requests"),
104 | })
105 | .optional()
106 | .describe("Advanced options for URL fetching"),
107 |
108 | modelName: z
109 | .string()
110 | .optional()
111 | .describe("Specific Gemini model to use (auto-selected if not specified)"),
112 | };
113 |
114 | /**
115 | * Registers the gemini_url_analysis tool with the MCP server.
116 | * Provides specialized URL analysis capabilities with intelligent content processing.
117 | */
118 | export const geminiUrlAnalysisTool = (
119 | server: McpServer,
120 | serviceInstance: GeminiService
121 | ): void => {
122 | const processRequest = async (args: unknown) => {
123 | // Parse and validate the arguments
124 | const parsedArgs = z.object(GEMINI_URL_ANALYSIS_PARAMS).parse(args);
125 |
126 | logger.debug(`Received ${GEMINI_URL_ANALYSIS_TOOL_NAME} request:`, {
127 | urls: parsedArgs.urls,
128 | analysisType: parsedArgs.analysisType,
129 | urlCount: parsedArgs.urls.length,
130 | });
131 |
132 | try {
133 | const {
134 | urls,
135 | analysisType,
136 | query,
137 | extractionSchema,
138 | questions,
139 | compareBy,
140 | outputFormat,
141 | includeMetadata,
142 | fetchOptions,
143 | modelName,
144 | } = parsedArgs;
145 |
146 | // Build the analysis prompt based on the analysis type
147 | const prompt = buildAnalysisPrompt({
148 | analysisType,
149 | query,
150 | extractionSchema,
151 | questions,
152 | compareBy,
153 | outputFormat,
154 | urlCount: urls.length,
155 | });
156 |
157 | // Prepare URL context for content generation
158 | const urlContext = {
159 | urls,
160 | fetchOptions: {
161 | ...fetchOptions,
162 | includeMetadata: includeMetadata ?? true,
163 | convertToMarkdown: true, // Always convert to markdown for better analysis
164 | },
165 | };
166 |
167 | // Calculate URL context metrics for optimal model selection
168 | const urlCount = urls.length;
169 | const maxContentKb = fetchOptions?.maxContentKb || 100;
170 | const estimatedUrlContentSize = urlCount * maxContentKb * 1024;
171 |
172 | // Select task type based on analysis type
173 | const taskType = getTaskTypeForAnalysis(analysisType);
174 |
175 | // Generate analysis using the service
176 | const analysisResult = await serviceInstance.generateContent({
177 | prompt,
178 | modelName,
179 | urlContext,
180 | taskType: taskType as
181 | | "text-generation"
182 | | "image-generation"
183 | | "video-generation"
184 | | "code-review"
185 | | "multimodal"
186 | | "reasoning",
187 | preferQuality: true, // Prefer quality for analysis tasks
188 | complexityHint: urlCount > 5 ? "complex" : "medium",
189 | urlCount,
190 | estimatedUrlContentSize,
191 | systemInstruction: getSystemInstructionForAnalysis(
192 | analysisType,
193 | outputFormat
194 | ),
195 | });
196 |
197 | // Format the result based on output format
198 | const formattedResult = formatAnalysisResult(
199 | analysisResult,
200 | outputFormat
201 | );
202 |
203 | return {
204 | content: [
205 | {
206 | type: "text" as const,
207 | text: formattedResult,
208 | },
209 | ],
210 | };
211 | } catch (error: unknown) {
212 | logger.error(`Error processing ${GEMINI_URL_ANALYSIS_TOOL_NAME}:`, error);
213 | throw mapAnyErrorToMcpError(error, GEMINI_URL_ANALYSIS_TOOL_NAME);
214 | }
215 | };
216 |
217 | // Register the tool with the server
218 | server.tool(
219 | GEMINI_URL_ANALYSIS_TOOL_NAME,
220 | GEMINI_URL_ANALYSIS_TOOL_DESCRIPTION,
221 | GEMINI_URL_ANALYSIS_PARAMS,
222 | processRequest
223 | );
224 |
225 | logger.info(`Tool registered: ${GEMINI_URL_ANALYSIS_TOOL_NAME}`);
226 | };
227 |
228 | /**
229 | * Builds the analysis prompt based on the requested analysis type and parameters
230 | */
231 | function buildAnalysisPrompt(params: {
232 | analysisType: string;
233 | query?: string;
234 | extractionSchema?: Record<string, unknown>;
235 | questions?: string[];
236 | compareBy?: string[];
237 | outputFormat?: string;
238 | urlCount: number;
239 | }): string {
240 | const {
241 | analysisType,
242 | query,
243 | extractionSchema,
244 | questions,
245 | compareBy,
246 | outputFormat,
247 | urlCount,
248 | } = params;
249 |
250 | let prompt = `Perform a ${analysisType} analysis on the provided URL content${urlCount > 1 ? "s" : ""}.\n\n`;
251 |
252 | switch (analysisType) {
253 | case "summary":
254 | prompt += `Provide a comprehensive summary of the main points, key information, and important insights from the content. `;
255 | if (query) {
256 | prompt += `Focus particularly on: ${query}. `;
257 | }
258 | break;
259 |
260 | case "comparison":
261 | if (urlCount < 2) {
262 | prompt += `Since only one URL is provided, analyze the different aspects or sections within the content. `;
263 | } else {
264 | prompt += `Compare and contrast the content from the different URLs, highlighting similarities, differences, and unique aspects. `;
265 | }
266 | if (compareBy && compareBy.length > 0) {
267 | prompt += `Focus your comparison on these specific aspects: ${compareBy.join(", ")}. `;
268 | }
269 | break;
270 |
271 | case "extraction":
272 | prompt += `Extract specific information from the content. `;
273 | if (extractionSchema) {
274 | prompt += `Structure the extracted information according to this schema: ${JSON.stringify(extractionSchema, null, 2)}. `;
275 | }
276 | if (query) {
277 | prompt += `Focus on extracting: ${query}. `;
278 | }
279 | break;
280 |
281 | case "qa":
282 | prompt += `Answer the following questions based on the content:\n`;
283 | if (questions && questions.length > 0) {
284 | questions.forEach((question, index) => {
285 | prompt += `${index + 1}. ${question}\n`;
286 | });
287 | } else if (query) {
288 | prompt += `Question: ${query}\n`;
289 | } else {
290 | prompt += `Provide comprehensive answers to common questions that would arise from this content.\n`;
291 | }
292 | break;
293 |
294 | case "sentiment":
295 | prompt += `Analyze the sentiment and emotional tone of the content. Identify the overall sentiment (positive, negative, neutral) and specific emotional indicators. `;
296 | if (query) {
297 | prompt += `Pay special attention to sentiment regarding: ${query}. `;
298 | }
299 | break;
300 |
301 | case "fact-check":
302 | prompt += `Evaluate the factual accuracy and credibility of claims made in the content. Identify verifiable facts, questionable claims, and potential misinformation. `;
303 | if (query) {
304 | prompt += `Focus particularly on claims about: ${query}. `;
305 | }
306 | break;
307 |
308 | case "content-classification":
309 | prompt += `Classify and categorize the content by topic, type, audience, and other relevant dimensions. `;
310 | if (query) {
311 | prompt += `Use this classification framework: ${query}. `;
312 | }
313 | break;
314 |
315 | case "readability":
316 | prompt += `Analyze the readability, writing quality, and accessibility of the content. Evaluate complexity, clarity, structure, and target audience. `;
317 | break;
318 |
319 | case "seo-analysis":
320 | prompt += `Perform an SEO analysis of the content, evaluating keyword usage, content structure, meta information, and optimization opportunities. `;
321 | break;
322 |
323 | default:
324 | if (query) {
325 | prompt += `Based on the following instruction: ${query}. `;
326 | }
327 | }
328 |
329 | // Add output format instructions
330 | if (outputFormat && outputFormat !== "text") {
331 | switch (outputFormat) {
332 | case "json":
333 | prompt += `\n\nFormat your response as valid JSON with appropriate structure and fields.`;
334 | break;
335 | case "markdown":
336 | prompt += `\n\nFormat your response in well-structured Markdown with appropriate headers, lists, and formatting.`;
337 | break;
338 | case "structured":
339 | prompt += `\n\nOrganize your response in a clear, structured format with distinct sections and subsections.`;
340 | break;
341 | }
342 | }
343 |
344 | prompt += `\n\nBe thorough, accurate, and insightful in your analysis.`;
345 |
346 | return prompt;
347 | }
348 |
349 | /**
350 | * Maps analysis types to task types for model selection
351 | */
352 | function getTaskTypeForAnalysis(analysisType: string): string {
353 | switch (analysisType) {
354 | case "comparison":
355 | case "fact-check":
356 | case "seo-analysis":
357 | return "reasoning";
358 | case "extraction":
359 | case "content-classification":
360 | return "text-generation";
361 | default:
362 | return "text-generation";
363 | }
364 | }
365 |
366 | /**
367 | * Generates system instructions based on analysis type and output format
368 | */
369 | function getSystemInstructionForAnalysis(
370 | analysisType: string,
371 | outputFormat?: string
372 | ): string {
373 | let instruction = `You are an expert content analyst specializing in ${analysisType} analysis. `;
374 |
375 | switch (analysisType) {
376 | case "summary":
377 | instruction += `Provide concise yet comprehensive summaries that capture the essence and key insights of the content.`;
378 | break;
379 | case "comparison":
380 | instruction += `Excel at identifying similarities, differences, and patterns across different content sources.`;
381 | break;
382 | case "extraction":
383 | instruction += `Focus on accurately identifying and extracting specific information while maintaining context and relevance.`;
384 | break;
385 | case "qa":
386 | instruction += `Provide clear, accurate, and well-supported answers based on the available content.`;
387 | break;
388 | case "sentiment":
389 | instruction += `Accurately identify emotional tone, sentiment indicators, and subjective language patterns.`;
390 | break;
391 | case "fact-check":
392 | instruction += `Evaluate claims critically, distinguish between facts and opinions, and identify potential misinformation.`;
393 | break;
394 | case "content-classification":
395 | instruction += `Categorize content accurately using relevant taxonomies and classification frameworks.`;
396 | break;
397 | case "readability":
398 | instruction += `Assess content accessibility, complexity, and effectiveness for target audiences.`;
399 | break;
400 | case "seo-analysis":
401 | instruction += `Evaluate content from an SEO perspective, focusing on optimization opportunities and best practices.`;
402 | break;
403 | }
404 |
405 | if (outputFormat === "json") {
406 | instruction += ` Always respond with valid, well-structured JSON.`;
407 | } else if (outputFormat === "markdown") {
408 | instruction += ` Use proper Markdown formatting with clear headers and structure.`;
409 | }
410 |
411 | instruction += ` Base your analysis strictly on the provided content and clearly distinguish between what is explicitly stated and what is inferred.`;
412 |
413 | return instruction;
414 | }
415 |
416 | /**
417 | * Formats the analysis result based on the requested output format
418 | */
419 | function formatAnalysisResult(result: string, outputFormat?: string): string {
420 | if (!outputFormat || outputFormat === "text") {
421 | return result;
422 | }
423 |
424 | // For other formats, the formatting should have been handled by the model
425 | // based on the prompt instructions, so we return the result as-is
426 | return result;
427 | }
428 |
```
--------------------------------------------------------------------------------
/tests/unit/utils/UrlSecurityService.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { UrlSecurityService } from "../../../src/utils/UrlSecurityService.js";
3 | import { ConfigurationManager } from "../../../src/config/ConfigurationManager.js";
4 | import { GeminiUrlValidationError } from "../../../src/utils/geminiErrors.js";
5 |
6 | // Mock dependencies
7 | vi.mock("../../../src/config/ConfigurationManager.js");
8 | vi.mock("../../../src/utils/logger.js");
9 |
10 | interface MockConfigManager {
11 | getUrlContextConfig: ReturnType<typeof vi.fn>;
12 | }
13 |
14 | describe("UrlSecurityService", () => {
15 | let service: UrlSecurityService;
16 | let mockConfig: MockConfigManager;
17 |
18 | beforeEach(() => {
19 | vi.clearAllMocks();
20 |
21 | mockConfig = {
22 | getUrlContextConfig: vi.fn().mockReturnValue({
23 | allowedDomains: ["*"],
24 | blocklistedDomains: [],
25 | }),
26 | };
27 |
28 | service = new UrlSecurityService(mockConfig as ConfigurationManager);
29 | });
30 |
31 | describe("URL format validation", () => {
32 | it("should accept valid HTTP URLs", async () => {
33 | await expect(
34 | service.validateUrl("http://example.com")
35 | ).resolves.not.toThrow();
36 | });
37 |
38 | it("should accept valid HTTPS URLs", async () => {
39 | await expect(
40 | service.validateUrl("https://example.com")
41 | ).resolves.not.toThrow();
42 | });
43 |
44 | it("should reject invalid URL formats", async () => {
45 | await expect(service.validateUrl("not-a-url")).rejects.toThrow(
46 | GeminiUrlValidationError
47 | );
48 | await expect(service.validateUrl("")).rejects.toThrow(
49 | GeminiUrlValidationError
50 | );
51 | await expect(service.validateUrl(" ")).rejects.toThrow(
52 | GeminiUrlValidationError
53 | );
54 | });
55 |
56 | it("should reject non-HTTP protocols", async () => {
57 | await expect(service.validateUrl("ftp://example.com")).rejects.toThrow(
58 | GeminiUrlValidationError
59 | );
60 | await expect(service.validateUrl("file:///etc/passwd")).rejects.toThrow(
61 | GeminiUrlValidationError
62 | );
63 | await expect(service.validateUrl("javascript:alert(1)")).rejects.toThrow(
64 | GeminiUrlValidationError
65 | );
66 | await expect(
67 | service.validateUrl("data:text/plain;base64,SGVsbG8=")
68 | ).rejects.toThrow(GeminiUrlValidationError);
69 | });
70 | });
71 |
72 | describe("Domain validation", () => {
73 | it("should allow domains in allowlist", async () => {
74 | mockConfig.getUrlContextConfig.mockReturnValue({
75 | allowedDomains: ["example.com", "test.org"],
76 | blocklistedDomains: [],
77 | });
78 |
79 | await expect(
80 | service.validateUrl("https://example.com")
81 | ).resolves.not.toThrow();
82 | await expect(
83 | service.validateUrl("https://test.org")
84 | ).resolves.not.toThrow();
85 | });
86 |
87 | it("should reject domains not in allowlist", async () => {
88 | mockConfig.getUrlContextConfig.mockReturnValue({
89 | allowedDomains: ["example.com"],
90 | blocklistedDomains: [],
91 | });
92 |
93 | await expect(
94 | service.validateUrl("https://malicious.com")
95 | ).rejects.toThrow(GeminiUrlValidationError);
96 | });
97 |
98 | it("should handle wildcard allowlist", async () => {
99 | mockConfig.getUrlContextConfig.mockReturnValue({
100 | allowedDomains: ["*"],
101 | blocklistedDomains: [],
102 | });
103 |
104 | await expect(
105 | service.validateUrl("https://any-domain.com")
106 | ).resolves.not.toThrow();
107 | });
108 |
109 | it("should handle subdomain patterns", async () => {
110 | mockConfig.getUrlContextConfig.mockReturnValue({
111 | allowedDomains: ["*.example.com"],
112 | blocklistedDomains: [],
113 | });
114 |
115 | await expect(
116 | service.validateUrl("https://sub.example.com")
117 | ).resolves.not.toThrow();
118 | await expect(
119 | service.validateUrl("https://deep.sub.example.com")
120 | ).resolves.not.toThrow();
121 | await expect(
122 | service.validateUrl("https://example.com")
123 | ).resolves.not.toThrow();
124 | await expect(service.validateUrl("https://other.com")).rejects.toThrow(
125 | GeminiUrlValidationError
126 | );
127 | });
128 |
129 | it("should block domains in blocklist", async () => {
130 | mockConfig.getUrlContextConfig.mockReturnValue({
131 | allowedDomains: ["*"],
132 | blocklistedDomains: ["malicious.com", "spam.net"],
133 | });
134 |
135 | await expect(
136 | service.validateUrl("https://malicious.com")
137 | ).rejects.toThrow(GeminiUrlValidationError);
138 | await expect(service.validateUrl("https://spam.net")).rejects.toThrow(
139 | GeminiUrlValidationError
140 | );
141 | await expect(
142 | service.validateUrl("https://safe.com")
143 | ).resolves.not.toThrow();
144 | });
145 |
146 | it("should block subdomains of blocklisted domains", async () => {
147 | mockConfig.getUrlContextConfig.mockReturnValue({
148 | allowedDomains: ["*"],
149 | blocklistedDomains: ["malicious.com"],
150 | });
151 |
152 | await expect(
153 | service.validateUrl("https://sub.malicious.com")
154 | ).rejects.toThrow(GeminiUrlValidationError);
155 | });
156 | });
157 |
158 | describe("Private network protection", () => {
159 | it("should block localhost addresses", async () => {
160 | await expect(service.validateUrl("http://localhost")).rejects.toThrow(
161 | GeminiUrlValidationError
162 | );
163 | await expect(service.validateUrl("http://127.0.0.1")).rejects.toThrow(
164 | GeminiUrlValidationError
165 | );
166 | await expect(service.validateUrl("http://0.0.0.0")).rejects.toThrow(
167 | GeminiUrlValidationError
168 | );
169 | });
170 |
171 | it("should block private IP ranges", async () => {
172 | await expect(service.validateUrl("http://192.168.1.1")).rejects.toThrow(
173 | GeminiUrlValidationError
174 | );
175 | await expect(service.validateUrl("http://10.0.0.1")).rejects.toThrow(
176 | GeminiUrlValidationError
177 | );
178 | await expect(service.validateUrl("http://172.16.0.1")).rejects.toThrow(
179 | GeminiUrlValidationError
180 | );
181 | });
182 |
183 | it("should block internal domain extensions", async () => {
184 | await expect(service.validateUrl("http://server.local")).rejects.toThrow(
185 | GeminiUrlValidationError
186 | );
187 | await expect(service.validateUrl("http://api.internal")).rejects.toThrow(
188 | GeminiUrlValidationError
189 | );
190 | await expect(service.validateUrl("http://db.corp")).rejects.toThrow(
191 | GeminiUrlValidationError
192 | );
193 | });
194 |
195 | it("should allow public IP addresses", async () => {
196 | await expect(
197 | service.validateUrl("http://8.8.8.8")
198 | ).resolves.not.toThrow();
199 | await expect(
200 | service.validateUrl("http://1.1.1.1")
201 | ).resolves.not.toThrow();
202 | });
203 | });
204 |
205 | describe("Suspicious pattern detection", () => {
206 | it("should detect path traversal attempts", async () => {
207 | await expect(
208 | service.validateUrl("http://example.com/../../../etc/passwd")
209 | ).rejects.toThrow(GeminiUrlValidationError);
210 | await expect(
211 | service.validateUrl("http://example.com/path/with/../dots")
212 | ).rejects.toThrow(GeminiUrlValidationError);
213 | });
214 |
215 | it("should detect dangerous characters", async () => {
216 | await expect(
217 | service.validateUrl("http://example.com/path<script>")
218 | ).rejects.toThrow(GeminiUrlValidationError);
219 | await expect(
220 | service.validateUrl("http://example.com/path{malicious}")
221 | ).rejects.toThrow(GeminiUrlValidationError);
222 | });
223 |
224 | it("should detect multiple @ symbols", async () => {
225 | await expect(
226 | service.validateUrl("http://user@[email protected]")
227 | ).rejects.toThrow(GeminiUrlValidationError);
228 | });
229 |
230 | it("should allow normal URLs with safe characters", async () => {
231 | await expect(
232 | service.validateUrl(
233 | "https://example.com/path/to/resource?param=value&other=123"
234 | )
235 | ).resolves.not.toThrow();
236 | await expect(
237 | service.validateUrl("https://api.example.com/v1/users/123")
238 | ).resolves.not.toThrow();
239 | });
240 | });
241 |
242 | describe("URL shortener detection", () => {
243 | it("should detect known URL shorteners", async () => {
244 | const shorteners = [
245 | "https://bit.ly/abc123",
246 | "https://tinyurl.com/abc123",
247 | "https://t.co/abc123",
248 | "https://goo.gl/abc123",
249 | ];
250 |
251 | // Note: These should not throw errors, but should be logged as warnings
252 | for (const url of shorteners) {
253 | await expect(service.validateUrl(url)).resolves.not.toThrow();
254 | }
255 | });
256 | });
257 |
258 | describe("IDN homograph attack detection", () => {
259 | it("should detect potentially confusing Unicode domains", async () => {
260 | // Cyrillic characters that look like Latin
261 | await expect(service.validateUrl("https://gоogle.com")).rejects.toThrow(
262 | GeminiUrlValidationError
263 | ); // 'о' is Cyrillic
264 | await expect(service.validateUrl("https://аpple.com")).rejects.toThrow(
265 | GeminiUrlValidationError
266 | ); // 'а' is Cyrillic
267 | });
268 |
269 | it("should allow legitimate Unicode domains", async () => {
270 | await expect(
271 | service.validateUrl("https://example.com")
272 | ).resolves.not.toThrow();
273 | await expect(
274 | service.validateUrl("https://测试.example.com")
275 | ).resolves.not.toThrow();
276 | });
277 | });
278 |
279 | describe("Port validation", () => {
280 | it("should allow standard HTTP/HTTPS ports", async () => {
281 | await expect(
282 | service.validateUrl("http://example.com:80")
283 | ).resolves.not.toThrow();
284 | await expect(
285 | service.validateUrl("https://example.com:443")
286 | ).resolves.not.toThrow();
287 | await expect(
288 | service.validateUrl("http://example.com:8080")
289 | ).resolves.not.toThrow();
290 | await expect(
291 | service.validateUrl("https://example.com:8443")
292 | ).resolves.not.toThrow();
293 | });
294 |
295 | it("should reject non-standard ports", async () => {
296 | await expect(
297 | service.validateUrl("http://example.com:22")
298 | ).rejects.toThrow(GeminiUrlValidationError);
299 | await expect(
300 | service.validateUrl("http://example.com:3389")
301 | ).rejects.toThrow(GeminiUrlValidationError);
302 | await expect(
303 | service.validateUrl("http://example.com:1337")
304 | ).rejects.toThrow(GeminiUrlValidationError);
305 | });
306 | });
307 |
308 | describe("URL length validation", () => {
309 | it("should reject extremely long URLs", async () => {
310 | const longPath = "a".repeat(3000);
311 | const longUrl = `https://example.com/${longPath}`;
312 |
313 | await expect(service.validateUrl(longUrl)).rejects.toThrow(
314 | GeminiUrlValidationError
315 | );
316 | });
317 |
318 | it("should accept reasonable length URLs", async () => {
319 | const normalPath = "a".repeat(100);
320 | const normalUrl = `https://example.com/${normalPath}`;
321 |
322 | await expect(service.validateUrl(normalUrl)).resolves.not.toThrow();
323 | });
324 | });
325 |
326 | describe("Random domain detection", () => {
327 | it("should flag potentially randomly generated domains", async () => {
328 | // These should log warnings but not necessarily throw errors
329 | const suspiciousDomains = [
330 | "https://xkcd123456789.com",
331 | "https://aaaaaaaaaaaa.com",
332 | "https://1234567890abcd.com",
333 | ];
334 |
335 | for (const url of suspiciousDomains) {
336 | // Should not throw, but may log warnings
337 | await expect(service.validateUrl(url)).resolves.not.toThrow();
338 | }
339 | });
340 | });
341 |
342 | describe("Security metrics", () => {
343 | it("should track validation attempts and failures", async () => {
344 | const initialMetrics = service.getSecurityMetrics();
345 | expect(initialMetrics.validationAttempts).toBe(0);
346 | expect(initialMetrics.validationFailures).toBe(0);
347 |
348 | // Valid URL
349 | await service.validateUrl("https://example.com").catch(() => {});
350 |
351 | // Invalid URL
352 | await service.validateUrl("invalid-url").catch(() => {});
353 |
354 | const updatedMetrics = service.getSecurityMetrics();
355 | expect(updatedMetrics.validationAttempts).toBe(2);
356 | expect(updatedMetrics.validationFailures).toBe(1);
357 | });
358 |
359 | it("should track blocked domains", async () => {
360 | mockConfig.getUrlContextConfig.mockReturnValue({
361 | allowedDomains: ["*"],
362 | blocklistedDomains: ["malicious.com"],
363 | });
364 |
365 | await service.validateUrl("https://malicious.com").catch(() => {});
366 |
367 | const metrics = service.getSecurityMetrics();
368 | expect(metrics.blockedDomains.has("malicious.com")).toBe(true);
369 | });
370 |
371 | it("should allow resetting metrics", () => {
372 | service.resetSecurityMetrics();
373 | const metrics = service.getSecurityMetrics();
374 |
375 | expect(metrics.validationAttempts).toBe(0);
376 | expect(metrics.validationFailures).toBe(0);
377 | expect(metrics.blockedDomains.size).toBe(0);
378 | expect(metrics.suspiciousPatterns).toHaveLength(0);
379 | });
380 | });
381 |
382 | describe("Custom domain management", () => {
383 | it("should allow adding custom malicious domains", () => {
384 | service.addMaliciousDomain("custom-malicious.com");
385 |
386 | // This should not throw immediately since domain checking happens in validateUrl
387 | expect(() => service.addMaliciousDomain("another-bad.com")).not.toThrow();
388 | });
389 | });
390 |
391 | describe("URL accessibility checking", () => {
392 | it("should check URL accessibility", async () => {
393 | // Mock fetch for accessibility check
394 | const mockFetch = vi.fn().mockResolvedValue({
395 | ok: true,
396 | status: 200,
397 | });
398 | global.fetch = mockFetch;
399 |
400 | const isAccessible = await service.checkUrlAccessibility(
401 | "https://example.com"
402 | );
403 | expect(isAccessible).toBe(true);
404 | expect(mockFetch).toHaveBeenCalledWith(
405 | "https://example.com",
406 | expect.objectContaining({
407 | method: "HEAD",
408 | })
409 | );
410 | });
411 |
412 | it("should handle inaccessible URLs", async () => {
413 | const mockFetch = vi.fn().mockRejectedValue(new Error("Network error"));
414 | global.fetch = mockFetch;
415 |
416 | const isAccessible = await service.checkUrlAccessibility(
417 | "https://unreachable.com"
418 | );
419 | expect(isAccessible).toBe(false);
420 | });
421 | });
422 | });
423 |
```
--------------------------------------------------------------------------------
/src/services/gemini/GeminiContentService.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { GoogleGenAI } from "@google/genai";
2 | import {
3 | GeminiApiError,
4 | GeminiValidationError,
5 | mapGeminiError,
6 | } from "../../utils/geminiErrors.js";
7 | import { logger } from "../../utils/logger.js";
8 | import {
9 | Content,
10 | GenerationConfig,
11 | SafetySetting,
12 | Part,
13 | ThinkingConfig,
14 | } from "./GeminiTypes.js";
15 | import { ZodError } from "zod";
16 | import { validateGenerateContentParams } from "./GeminiValidationSchemas.js";
17 | import { RetryService } from "../../utils/RetryService.js";
18 | import { GeminiUrlContextService } from "./GeminiUrlContextService.js";
19 | import { ConfigurationManager } from "../../config/ConfigurationManager.js";
20 |
21 | // Request configuration type definition for reuse
22 | interface RequestConfig {
23 | model: string;
24 | contents: Content[];
25 | generationConfig?: GenerationConfig;
26 | safetySettings?: SafetySetting[];
27 | systemInstruction?: Content;
28 | cachedContent?: string;
29 | thinkingConfig?: ThinkingConfig;
30 | }
31 |
32 | /**
33 | * Interface for URL context parameters
34 | */
35 | interface UrlContextParams {
36 | urls: string[];
37 | fetchOptions?: {
38 | maxContentKb?: number;
39 | timeoutMs?: number;
40 | includeMetadata?: boolean;
41 | convertToMarkdown?: boolean;
42 | allowedDomains?: string[];
43 | userAgent?: string;
44 | };
45 | }
46 |
47 | /**
48 | * Interface for the parameters of the generateContent method
49 | * This interface is used internally, while the parent GeminiService exports a compatible version
50 | */
51 | interface GenerateContentParams {
52 | prompt: string;
53 | modelName?: string;
54 | generationConfig?: GenerationConfig;
55 | safetySettings?: SafetySetting[];
56 | systemInstruction?: Content | string;
57 | cachedContentName?: string;
58 | urlContext?: UrlContextParams;
59 | }
60 |
61 | /**
62 | * Default retry options for Gemini API calls
63 | */
64 | const DEFAULT_RETRY_OPTIONS = {
65 | maxAttempts: 3,
66 | initialDelayMs: 500,
67 | maxDelayMs: 10000,
68 | backoffFactor: 2,
69 | jitter: true,
70 | onRetry: (error: unknown, attempt: number, delayMs: number) => {
71 | logger.warn(
72 | `Retrying Gemini API call after error (attempt ${attempt}, delay: ${delayMs}ms): ${error instanceof Error ? error.message : String(error)}`
73 | );
74 | },
75 | };
76 |
77 | /**
78 | * Service for handling content generation related operations for the Gemini service.
79 | * Manages content generation in both streaming and non-streaming modes.
80 | */
81 | export class GeminiContentService {
82 | private genAI: GoogleGenAI;
83 | private defaultModelName?: string;
84 | private defaultThinkingBudget?: number;
85 | private retryService: RetryService;
86 | private configManager: ConfigurationManager;
87 | private urlContextService: GeminiUrlContextService;
88 |
89 | /**
90 | * Creates a new instance of the GeminiContentService.
91 | * @param genAI The GoogleGenAI instance to use for API calls
92 | * @param defaultModelName Optional default model name to use if not specified in method calls
93 | * @param defaultThinkingBudget Optional default budget for reasoning (thinking) tokens
94 | */
95 | constructor(
96 | genAI: GoogleGenAI,
97 | defaultModelName?: string,
98 | defaultThinkingBudget?: number
99 | ) {
100 | this.genAI = genAI;
101 | this.defaultModelName = defaultModelName;
102 | this.defaultThinkingBudget = defaultThinkingBudget;
103 | this.retryService = new RetryService(DEFAULT_RETRY_OPTIONS);
104 | this.configManager = ConfigurationManager.getInstance();
105 | this.urlContextService = new GeminiUrlContextService(this.configManager);
106 | }
107 |
108 | /**
109 | * Streams content generation using the Gemini model.
110 | * Returns an async generator that yields text chunks as they are generated.
111 | *
112 | * @param params An object containing all necessary parameters for content generation
113 | * @returns An async generator yielding text chunks as they become available
114 | */
115 | public async *generateContentStream(
116 | params: GenerateContentParams
117 | ): AsyncGenerator<string> {
118 | // Log with truncated prompt for privacy/security
119 | logger.debug(
120 | `generateContentStream called with prompt: ${params.prompt.substring(0, 30)}...`
121 | );
122 |
123 | try {
124 | // Validate parameters using Zod schema
125 | try {
126 | // Create a proper object for validation
127 | const validationParams: Record<string, unknown> = {
128 | prompt: params.prompt,
129 | modelName: params.modelName,
130 | generationConfig: params.generationConfig,
131 | safetySettings: params.safetySettings,
132 | systemInstruction: params.systemInstruction,
133 | cachedContentName: params.cachedContentName,
134 | };
135 | validateGenerateContentParams(validationParams);
136 | } catch (validationError: unknown) {
137 | if (validationError instanceof ZodError) {
138 | const fieldErrors = validationError.errors
139 | .map((err) => `${err.path.join(".")}: ${err.message}`)
140 | .join(", ");
141 | throw new GeminiValidationError(
142 | `Invalid parameters for content generation: ${fieldErrors}`,
143 | validationError.errors[0]?.path.join(".")
144 | );
145 | }
146 | throw validationError;
147 | }
148 |
149 | // Create the request configuration using the helper method
150 | const requestConfig = await this.createRequestConfig(params);
151 |
152 | // Call generateContentStream with retry
153 | // Note: We can't use the retry service directly here because we need to handle streaming
154 | // Instead, we'll add retry logic to the initial API call, but not the streaming part
155 | let streamResult;
156 | try {
157 | streamResult = await this.retryService.execute(async () => {
158 | return this.genAI.models.generateContentStream(requestConfig);
159 | });
160 | } catch (error: unknown) {
161 | throw mapGeminiError(error, "generateContentStream");
162 | }
163 |
164 | // Stream the results (no retry for individual chunks)
165 | try {
166 | for await (const chunk of streamResult) {
167 | // Extract text from the chunk if available - text is a getter, not a method
168 | const chunkText = chunk.text;
169 | if (chunkText) {
170 | yield chunkText;
171 | }
172 | }
173 | } catch (error: unknown) {
174 | throw mapGeminiError(error, "generateContentStream");
175 | }
176 | } catch (error: unknown) {
177 | // Map to appropriate error type for any other errors
178 | throw mapGeminiError(error, "generateContentStream");
179 | }
180 | }
181 |
182 | /**
183 | * Creates the request configuration object for both content generation methods.
184 | * This helper method reduces code duplication between generateContent and generateContentStream.
185 | *
186 | * @param params The content generation parameters
187 | * @returns A properly formatted request configuration object
188 | * @throws GeminiApiError if parameters are invalid or model name is missing
189 | */
190 | private async createRequestConfig(
191 | params: GenerateContentParams
192 | ): Promise<RequestConfig> {
193 | const {
194 | prompt,
195 | modelName,
196 | generationConfig,
197 | safetySettings,
198 | systemInstruction,
199 | cachedContentName,
200 | urlContext,
201 | } = params;
202 |
203 | const effectiveModelName = modelName ?? this.defaultModelName;
204 | if (!effectiveModelName) {
205 | throw new GeminiValidationError(
206 | "Model name must be provided either as a parameter or via the GOOGLE_GEMINI_MODEL environment variable.",
207 | "modelName"
208 | );
209 | }
210 | logger.debug(`Creating request config for model: ${effectiveModelName}`);
211 |
212 | // Construct base content parts array
213 | const contentParts: Part[] = [];
214 |
215 | // Process URL context first if provided
216 | if (urlContext?.urls && urlContext.urls.length > 0) {
217 | const urlConfig = this.configManager.getUrlContextConfig();
218 |
219 | if (!urlConfig.enabled) {
220 | throw new GeminiValidationError(
221 | "URL context feature is not enabled. Set GOOGLE_GEMINI_ENABLE_URL_CONTEXT=true to enable.",
222 | "urlContext"
223 | );
224 | }
225 |
226 | try {
227 | logger.debug(`Processing ${urlContext.urls.length} URLs for context`);
228 |
229 | const urlFetchOptions = {
230 | maxContentLength:
231 | (urlContext.fetchOptions?.maxContentKb ||
232 | urlConfig.defaultMaxContentKb) * 1024,
233 | timeout:
234 | urlContext.fetchOptions?.timeoutMs || urlConfig.defaultTimeoutMs,
235 | includeMetadata:
236 | urlContext.fetchOptions?.includeMetadata ??
237 | urlConfig.includeMetadata,
238 | convertToMarkdown:
239 | urlContext.fetchOptions?.convertToMarkdown ??
240 | urlConfig.convertToMarkdown,
241 | allowedDomains:
242 | urlContext.fetchOptions?.allowedDomains || urlConfig.allowedDomains,
243 | userAgent: urlContext.fetchOptions?.userAgent || urlConfig.userAgent,
244 | };
245 |
246 | const { contents: urlContents, batchResult } =
247 | await this.urlContextService.processUrlsForContext(
248 | urlContext.urls,
249 | urlFetchOptions
250 | );
251 |
252 | // Log the batch result for monitoring
253 | logger.info("URL context processing completed", {
254 | totalUrls: batchResult.summary.totalUrls,
255 | successful: batchResult.summary.successCount,
256 | failed: batchResult.summary.failureCount,
257 | totalContentSize: batchResult.summary.totalContentSize,
258 | avgResponseTime: batchResult.summary.averageResponseTime,
259 | });
260 |
261 | // Add URL content parts to the beginning (before the user's prompt)
262 | for (const urlContent of urlContents) {
263 | if (urlContent.parts) {
264 | contentParts.push(...urlContent.parts);
265 | }
266 | }
267 |
268 | // Log any failed URLs as warnings
269 | if (batchResult.failed.length > 0) {
270 | for (const failure of batchResult.failed) {
271 | logger.warn("Failed to fetch URL for context", {
272 | url: failure.url,
273 | error: failure.error.message,
274 | errorCode: failure.errorCode,
275 | });
276 | }
277 | }
278 | } catch (error) {
279 | logger.error("URL context processing failed", { error });
280 | // Depending on configuration, we could either fail the request or continue without URL context
281 | // For now, we'll throw the error to fail fast
282 | throw mapGeminiError(error, "URL context processing");
283 | }
284 | }
285 |
286 | // Add the user's prompt after URL context
287 | contentParts.push({ text: prompt });
288 |
289 | // Process systemInstruction if it's a string
290 | let formattedSystemInstruction: Content | undefined;
291 | if (systemInstruction) {
292 | if (typeof systemInstruction === "string") {
293 | formattedSystemInstruction = {
294 | parts: [{ text: systemInstruction }],
295 | };
296 | } else {
297 | formattedSystemInstruction = systemInstruction;
298 | }
299 | }
300 |
301 | // Create the request configuration for v0.10.0
302 | const requestConfig: RequestConfig = {
303 | model: effectiveModelName,
304 | contents: [{ role: "user", parts: contentParts }],
305 | };
306 |
307 | // Add optional parameters if provided
308 | if (generationConfig) {
309 | requestConfig.generationConfig = generationConfig;
310 |
311 | // Extract thinking config if it exists within generation config
312 | if (generationConfig.thinkingConfig) {
313 | requestConfig.thinkingConfig = generationConfig.thinkingConfig;
314 | }
315 | }
316 |
317 | // Map reasoningEffort to thinkingBudget if provided
318 | if (requestConfig.thinkingConfig?.reasoningEffort) {
319 | const effortMap: Record<string, number> = {
320 | none: 0,
321 | low: 1024, // 1K tokens
322 | medium: 8192, // 8K tokens
323 | high: 24576, // 24K tokens
324 | };
325 |
326 | requestConfig.thinkingConfig.thinkingBudget =
327 | effortMap[requestConfig.thinkingConfig.reasoningEffort];
328 | logger.debug(
329 | `Mapped reasoning effort '${requestConfig.thinkingConfig.reasoningEffort}' to thinking budget: ${requestConfig.thinkingConfig.thinkingBudget} tokens`
330 | );
331 | }
332 |
333 | // Apply default thinking budget if available and not specified in request
334 | if (
335 | this.defaultThinkingBudget !== undefined &&
336 | !requestConfig.thinkingConfig
337 | ) {
338 | requestConfig.thinkingConfig = {
339 | thinkingBudget: this.defaultThinkingBudget,
340 | };
341 | logger.debug(
342 | `Applied default thinking budget: ${this.defaultThinkingBudget} tokens`
343 | );
344 | }
345 | if (safetySettings) {
346 | requestConfig.safetySettings = safetySettings;
347 | }
348 | if (formattedSystemInstruction) {
349 | requestConfig.systemInstruction = formattedSystemInstruction;
350 | }
351 | if (cachedContentName) {
352 | requestConfig.cachedContent = cachedContentName;
353 | }
354 |
355 | return requestConfig;
356 | }
357 |
358 | /**
359 | * Generates content using the Gemini model with automatic retries for transient errors.
360 | * Uses exponential backoff to avoid overwhelming the API during temporary issues.
361 | *
362 | * @param params An object containing all necessary parameters for content generation
363 | * @returns A promise resolving to the generated text content
364 | */
365 | public async generateContent(params: GenerateContentParams): Promise<string> {
366 | // Log with truncated prompt for privacy/security
367 | logger.debug(
368 | `generateContent called with prompt: ${params.prompt.substring(0, 30)}...`
369 | );
370 |
371 | try {
372 | // Validate parameters using Zod schema
373 | try {
374 | // Create a proper object for validation
375 | const validationParams: Record<string, unknown> = {
376 | prompt: params.prompt,
377 | modelName: params.modelName,
378 | generationConfig: params.generationConfig,
379 | safetySettings: params.safetySettings,
380 | systemInstruction: params.systemInstruction,
381 | cachedContentName: params.cachedContentName,
382 | };
383 | validateGenerateContentParams(validationParams);
384 | } catch (validationError: unknown) {
385 | if (validationError instanceof ZodError) {
386 | const fieldErrors = validationError.errors
387 | .map((err) => `${err.path.join(".")}: ${err.message}`)
388 | .join(", ");
389 | throw new GeminiValidationError(
390 | `Invalid parameters for content generation: ${fieldErrors}`,
391 | validationError.errors[0]?.path.join(".")
392 | );
393 | }
394 | throw validationError;
395 | }
396 |
397 | // Create the request configuration using the helper method
398 | const requestConfig = await this.createRequestConfig(params);
399 |
400 | // Call generateContent with retry logic
401 | return await this.retryService.execute(async () => {
402 | const result = await this.genAI.models.generateContent(requestConfig);
403 |
404 | // Handle potentially undefined text property
405 | if (!result.text) {
406 | throw new GeminiApiError("No text was generated in the response");
407 | }
408 |
409 | return result.text;
410 | });
411 | } catch (error: unknown) {
412 | // Map to appropriate error type
413 | throw mapGeminiError(error, "generateContent");
414 | }
415 | }
416 | }
417 |
```
--------------------------------------------------------------------------------
/tests/unit/tools/geminiCacheTool.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { geminiCacheTool } from "../../../src/tools/geminiCacheTool.js";
3 | import { GeminiApiError } from "../../../src/utils/errors.js";
4 | import { McpError } from "@modelcontextprotocol/sdk/types.js";
5 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
6 | import { GeminiService } from "../../../src/services/index.js";
7 |
8 | describe("geminiCacheTool", () => {
9 | // Mock server and service instances
10 | const mockTool = vi.fn();
11 | const mockServer = {
12 | tool: mockTool,
13 | } as unknown as McpServer;
14 |
15 | // Create mock functions for the service methods
16 | const mockCreateCache = vi.fn();
17 | const mockListCaches = vi.fn();
18 | const mockGetCache = vi.fn();
19 | const mockUpdateCache = vi.fn();
20 | const mockDeleteCache = vi.fn();
21 |
22 | // Create a minimal mock service with just the necessary methods for testing
23 | const mockService = {
24 | createCache: mockCreateCache,
25 | listCaches: mockListCaches,
26 | getCache: mockGetCache,
27 | updateCache: mockUpdateCache,
28 | deleteCache: mockDeleteCache,
29 | // Add empty implementations for required GeminiService methods
30 | generateContent: () => Promise.resolve("mock"),
31 | } as unknown as GeminiService;
32 |
33 | // Reset mocks before each test
34 | beforeEach(() => {
35 | vi.resetAllMocks();
36 | });
37 |
38 | it("should register the tool with the server", () => {
39 | // Call the tool registration function
40 | geminiCacheTool(mockServer, mockService);
41 |
42 | // Verify tool was registered
43 | expect(mockTool).toHaveBeenCalledTimes(1);
44 | const [name, description, params, handler] = mockTool.mock.calls[0];
45 |
46 | // Check tool registration parameters
47 | expect(name).toBe("gemini_cache");
48 | expect(description).toContain("Manages cached content resources");
49 | expect(params).toBeDefined();
50 | expect(typeof handler).toBe("function");
51 | });
52 |
53 | describe("create operation", () => {
54 | it("should create a cache successfully", async () => {
55 | // Register tool to get the request handler
56 | geminiCacheTool(mockServer, mockService);
57 | const [, , , handler] = mockTool.mock.calls[0];
58 |
59 | // Mock successful response
60 | const mockCacheMetadata = {
61 | name: "cachedContents/abc123xyz",
62 | displayName: "Test Cache",
63 | model: "gemini-1.5-flash",
64 | createTime: "2024-01-01T00:00:00Z",
65 | updateTime: "2024-01-01T00:00:00Z",
66 | expirationTime: "2024-01-02T00:00:00Z",
67 | state: "ACTIVE",
68 | usageMetadata: {
69 | totalTokenCount: 1000,
70 | },
71 | };
72 | mockCreateCache.mockResolvedValueOnce(mockCacheMetadata);
73 |
74 | // Prepare test request
75 | const testRequest = {
76 | operation: "create",
77 | model: "gemini-1.5-flash",
78 | contents: [
79 | {
80 | role: "user",
81 | parts: [{ text: "This is cached content" }],
82 | },
83 | ],
84 | displayName: "Test Cache",
85 | ttl: "3600s",
86 | };
87 |
88 | // Call the handler
89 | const result = await handler(testRequest);
90 |
91 | // Verify the service method was called with correct parameters
92 | expect(mockCreateCache).toHaveBeenCalledWith(
93 | "gemini-1.5-flash",
94 | testRequest.contents,
95 | {
96 | displayName: "Test Cache",
97 | ttl: "3600s",
98 | }
99 | );
100 |
101 | // Verify the result
102 | expect(result).toEqual({
103 | content: [
104 | {
105 | type: "text",
106 | text: JSON.stringify(mockCacheMetadata, null, 2),
107 | },
108 | ],
109 | });
110 | });
111 |
112 | it("should create cache with system instruction and tools", async () => {
113 | // Register tool to get the request handler
114 | geminiCacheTool(mockServer, mockService);
115 | const [, , , handler] = mockTool.mock.calls[0];
116 |
117 | // Mock successful response
118 | const mockCacheMetadata = {
119 | name: "cachedContents/def456xyz",
120 | model: "gemini-1.5-pro",
121 | createTime: "2024-01-01T00:00:00Z",
122 | updateTime: "2024-01-01T00:00:00Z",
123 | };
124 | mockCreateCache.mockResolvedValueOnce(mockCacheMetadata);
125 |
126 | // Prepare test request with optional parameters
127 | const testRequest = {
128 | operation: "create",
129 | model: "gemini-1.5-pro",
130 | contents: [
131 | {
132 | role: "user",
133 | parts: [{ text: "Cached content" }],
134 | },
135 | ],
136 | systemInstruction: {
137 | role: "system",
138 | parts: [{ text: "You are a helpful assistant" }],
139 | },
140 | tools: [
141 | {
142 | functionDeclarations: [
143 | {
144 | name: "get_weather",
145 | description: "Get weather information",
146 | parameters: {
147 | type: "OBJECT",
148 | properties: {
149 | location: {
150 | type: "STRING",
151 | description: "The location",
152 | },
153 | },
154 | },
155 | },
156 | ],
157 | },
158 | ],
159 | toolConfig: {
160 | functionCallingConfig: {
161 | mode: "AUTO",
162 | },
163 | },
164 | };
165 |
166 | // Call the handler
167 | const result = await handler(testRequest);
168 | expect(result).toBeDefined();
169 |
170 | // Verify all parameters were passed
171 | expect(mockCreateCache).toHaveBeenCalledWith(
172 | "gemini-1.5-pro",
173 | testRequest.contents,
174 | expect.objectContaining({
175 | systemInstruction: testRequest.systemInstruction,
176 | tools: testRequest.tools,
177 | toolConfig: testRequest.toolConfig,
178 | })
179 | );
180 | });
181 |
182 | it("should throw error if contents is missing", async () => {
183 | // Register tool to get the request handler
184 | geminiCacheTool(mockServer, mockService);
185 | const [, , , handler] = mockTool.mock.calls[0];
186 |
187 | // Prepare test request without contents
188 | const testRequest = {
189 | operation: "create",
190 | model: "gemini-1.5-flash",
191 | };
192 |
193 | // Call the handler and expect error
194 | await expect(handler(testRequest)).rejects.toThrow(
195 | "contents is required for operation 'create'"
196 | );
197 | });
198 | });
199 |
200 | describe("list operation", () => {
201 | it("should list caches successfully", async () => {
202 | // Register tool to get the request handler
203 | geminiCacheTool(mockServer, mockService);
204 | const [, , , handler] = mockTool.mock.calls[0];
205 |
206 | // Mock successful response
207 | const mockListResult = {
208 | cachedContents: [
209 | {
210 | name: "cachedContents/cache1",
211 | displayName: "Cache 1",
212 | model: "gemini-1.5-flash",
213 | state: "ACTIVE",
214 | },
215 | {
216 | name: "cachedContents/cache2",
217 | displayName: "Cache 2",
218 | model: "gemini-1.5-pro",
219 | state: "ACTIVE",
220 | },
221 | ],
222 | nextPageToken: "token123",
223 | };
224 | mockListCaches.mockResolvedValueOnce(mockListResult);
225 |
226 | // Prepare test request
227 | const testRequest = {
228 | operation: "list",
229 | pageSize: 50,
230 | pageToken: "previousToken",
231 | };
232 |
233 | // Call the handler
234 | const result = await handler(testRequest);
235 |
236 | // Verify the service method was called
237 | expect(mockListCaches).toHaveBeenCalledWith(50, "previousToken");
238 |
239 | // Verify the result
240 | expect(result).toEqual({
241 | content: [
242 | {
243 | type: "text",
244 | text: JSON.stringify(mockListResult, null, 2),
245 | },
246 | ],
247 | });
248 | });
249 | });
250 |
251 | describe("get operation", () => {
252 | it("should get cache metadata successfully", async () => {
253 | // Register tool to get the request handler
254 | geminiCacheTool(mockServer, mockService);
255 | const [, , , handler] = mockTool.mock.calls[0];
256 |
257 | // Mock successful response
258 | const mockCacheMetadata = {
259 | name: "cachedContents/abc123xyz",
260 | displayName: "Test Cache",
261 | model: "gemini-1.5-flash",
262 | createTime: "2024-01-01T00:00:00Z",
263 | updateTime: "2024-01-01T00:00:00Z",
264 | expirationTime: "2024-01-02T00:00:00Z",
265 | state: "ACTIVE",
266 | };
267 | mockGetCache.mockResolvedValueOnce(mockCacheMetadata);
268 |
269 | // Prepare test request
270 | const testRequest = {
271 | operation: "get",
272 | cacheName: "cachedContents/abc123xyz",
273 | };
274 |
275 | // Call the handler
276 | const result = await handler(testRequest);
277 |
278 | // Verify the service method was called
279 | expect(mockGetCache).toHaveBeenCalledWith("cachedContents/abc123xyz");
280 |
281 | // Verify the result
282 | expect(result).toEqual({
283 | content: [
284 | {
285 | type: "text",
286 | text: JSON.stringify(mockCacheMetadata, null, 2),
287 | },
288 | ],
289 | });
290 | });
291 |
292 | it("should throw error if cacheName is missing", async () => {
293 | // Register tool to get the request handler
294 | geminiCacheTool(mockServer, mockService);
295 | const [, , , handler] = mockTool.mock.calls[0];
296 |
297 | // Prepare test request without cacheName
298 | const testRequest = {
299 | operation: "get",
300 | };
301 |
302 | // Call the handler and expect error
303 | await expect(handler(testRequest)).rejects.toThrow(
304 | "cacheName is required for operation 'get'"
305 | );
306 | });
307 |
308 | it("should throw error if cacheName format is invalid", async () => {
309 | // Register tool to get the request handler
310 | geminiCacheTool(mockServer, mockService);
311 | const [, , , handler] = mockTool.mock.calls[0];
312 |
313 | // Prepare test request with invalid cacheName
314 | const testRequest = {
315 | operation: "get",
316 | cacheName: "invalid-format",
317 | };
318 |
319 | // Call the handler and expect error
320 | await expect(handler(testRequest)).rejects.toThrow(
321 | "cacheName must start with 'cachedContents/'"
322 | );
323 | });
324 | });
325 |
326 | describe("update operation", () => {
327 | it("should update cache with TTL successfully", async () => {
328 | // Register tool to get the request handler
329 | geminiCacheTool(mockServer, mockService);
330 | const [, , , handler] = mockTool.mock.calls[0];
331 |
332 | // Mock successful response
333 | const mockUpdatedMetadata = {
334 | name: "cachedContents/abc123xyz",
335 | displayName: "Test Cache",
336 | model: "gemini-1.5-flash",
337 | updateTime: "2024-01-01T01:00:00Z",
338 | expirationTime: "2024-01-03T00:00:00Z",
339 | };
340 | mockUpdateCache.mockResolvedValueOnce(mockUpdatedMetadata);
341 |
342 | // Prepare test request
343 | const testRequest = {
344 | operation: "update",
345 | cacheName: "cachedContents/abc123xyz",
346 | ttl: "7200s",
347 | };
348 |
349 | // Call the handler
350 | const result = await handler(testRequest);
351 |
352 | // Verify the service method was called
353 | expect(mockUpdateCache).toHaveBeenCalledWith("cachedContents/abc123xyz", {
354 | ttl: "7200s",
355 | });
356 |
357 | // Verify the result
358 | expect(result).toEqual({
359 | content: [
360 | {
361 | type: "text",
362 | text: JSON.stringify(mockUpdatedMetadata, null, 2),
363 | },
364 | ],
365 | });
366 | });
367 |
368 | it("should update cache with displayName successfully", async () => {
369 | // Register tool to get the request handler
370 | geminiCacheTool(mockServer, mockService);
371 | const [, , , handler] = mockTool.mock.calls[0];
372 |
373 | // Mock successful response
374 | const mockUpdatedMetadata = {
375 | name: "cachedContents/abc123xyz",
376 | displayName: "Updated Cache Name",
377 | model: "gemini-1.5-flash",
378 | updateTime: "2024-01-01T01:00:00Z",
379 | };
380 | mockUpdateCache.mockResolvedValueOnce(mockUpdatedMetadata);
381 |
382 | // Prepare test request
383 | const testRequest = {
384 | operation: "update",
385 | cacheName: "cachedContents/abc123xyz",
386 | displayName: "Updated Cache Name",
387 | };
388 |
389 | // Call the handler
390 | const result = await handler(testRequest);
391 | expect(result).toBeDefined();
392 |
393 | // Verify the service method was called
394 | expect(mockUpdateCache).toHaveBeenCalledWith("cachedContents/abc123xyz", {
395 | displayName: "Updated Cache Name",
396 | });
397 | });
398 |
399 | it("should throw error if neither ttl nor displayName is provided", async () => {
400 | // Register tool to get the request handler
401 | geminiCacheTool(mockServer, mockService);
402 | const [, , , handler] = mockTool.mock.calls[0];
403 |
404 | // Prepare test request without update fields
405 | const testRequest = {
406 | operation: "update",
407 | cacheName: "cachedContents/abc123xyz",
408 | };
409 |
410 | // Call the handler and expect error
411 | await expect(handler(testRequest)).rejects.toThrow(
412 | "At least one of 'ttl' or 'displayName' must be provided for update operation"
413 | );
414 | });
415 | });
416 |
417 | describe("delete operation", () => {
418 | it("should delete cache successfully", async () => {
419 | // Register tool to get the request handler
420 | geminiCacheTool(mockServer, mockService);
421 | const [, , , handler] = mockTool.mock.calls[0];
422 |
423 | // Mock successful response
424 | mockDeleteCache.mockResolvedValueOnce({ success: true });
425 |
426 | // Prepare test request
427 | const testRequest = {
428 | operation: "delete",
429 | cacheName: "cachedContents/abc123xyz",
430 | };
431 |
432 | // Call the handler
433 | const result = await handler(testRequest);
434 |
435 | // Verify the service method was called
436 | expect(mockDeleteCache).toHaveBeenCalledWith("cachedContents/abc123xyz");
437 |
438 | // Verify the result
439 | expect(result).toEqual({
440 | content: [
441 | {
442 | type: "text",
443 | text: JSON.stringify({
444 | success: true,
445 | message: "Cache cachedContents/abc123xyz deleted successfully",
446 | }),
447 | },
448 | ],
449 | });
450 | });
451 | });
452 |
453 | describe("error handling", () => {
454 | it("should map GeminiApiError to McpError", async () => {
455 | // Register tool to get the request handler
456 | geminiCacheTool(mockServer, mockService);
457 | const [, , , handler] = mockTool.mock.calls[0];
458 |
459 | // Mock service to throw GeminiApiError
460 | const geminiError = new GeminiApiError("API error occurred");
461 | mockListCaches.mockRejectedValueOnce(geminiError);
462 |
463 | // Prepare test request
464 | const testRequest = {
465 | operation: "list",
466 | };
467 |
468 | // Call the handler and expect McpError
469 | try {
470 | await handler(testRequest);
471 | expect.fail("Should have thrown an error");
472 | } catch (error) {
473 | expect(error).toBeInstanceOf(McpError);
474 | expect((error as McpError).message).toContain("API error occurred");
475 | }
476 | });
477 |
478 | it("should handle invalid operation", async () => {
479 | // Register tool to get the request handler
480 | geminiCacheTool(mockServer, mockService);
481 | const [, , , handler] = mockTool.mock.calls[0];
482 |
483 | // Prepare test request with invalid operation
484 | const testRequest = {
485 | operation: "invalid_operation",
486 | };
487 |
488 | // Call the handler and expect error
489 | await expect(handler(testRequest)).rejects.toThrow(
490 | "Invalid operation: invalid_operation"
491 | );
492 | });
493 | });
494 | });
495 |
```
--------------------------------------------------------------------------------
/tests/integration/urlContextIntegration.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import { GeminiService } from "../../src/services/GeminiService.js";
3 | import { ConfigurationManager } from "../../src/config/ConfigurationManager.js";
4 |
5 | // Mock external dependencies
6 | vi.mock("../../src/config/ConfigurationManager.js");
7 | vi.mock("@google/genai");
8 |
9 | // Mock fetch globally for URL fetching tests
10 | const mockFetch = vi.fn();
11 | global.fetch = mockFetch;
12 |
13 | interface MockConfigInstance {
14 | getGeminiServiceConfig: ReturnType<typeof vi.fn>;
15 | getUrlContextConfig: ReturnType<typeof vi.fn>;
16 | }
17 |
18 | interface MockConfig {
19 | getInstance: ReturnType<typeof vi.fn<[], MockConfigInstance>>;
20 | }
21 |
22 | describe("URL Context Integration Tests", () => {
23 | let geminiService: GeminiService;
24 | let mockConfig: MockConfig;
25 |
26 | beforeEach(async () => {
27 | vi.clearAllMocks();
28 |
29 | // Mock configuration with URL context enabled
30 | mockConfig = {
31 | getInstance: vi.fn().mockReturnValue({
32 | getGeminiServiceConfig: vi.fn().mockReturnValue({
33 | apiKey: "test-api-key",
34 | defaultModel: "gemini-2.5-flash-preview-05-20",
35 | }),
36 | getUrlContextConfig: vi.fn().mockReturnValue({
37 | enabled: true,
38 | maxUrlsPerRequest: 20,
39 | defaultMaxContentKb: 100,
40 | defaultTimeoutMs: 10000,
41 | allowedDomains: ["*"],
42 | blocklistedDomains: [],
43 | convertToMarkdown: true,
44 | includeMetadata: true,
45 | enableCaching: true,
46 | cacheExpiryMinutes: 15,
47 | maxCacheSize: 1000,
48 | rateLimitPerDomainPerMinute: 10,
49 | userAgent: "MCP-Gemini-Server/1.0",
50 | }),
51 | }),
52 | };
53 |
54 | ConfigurationManager.getInstance = mockConfig.getInstance;
55 |
56 | // Mock Gemini API
57 | const mockGenAI = {
58 | models: {
59 | generateContent: vi.fn().mockResolvedValue({
60 | text: "Generated response based on URL content",
61 | }),
62 | generateContentStream: vi.fn().mockImplementation(async function* () {
63 | yield "Generated ";
64 | yield "response ";
65 | yield "based on ";
66 | yield "URL content";
67 | }),
68 | },
69 | };
70 |
71 | const { GoogleGenAI } = await import("@google/genai");
72 | vi.mocked(GoogleGenAI).mockImplementation(() => mockGenAI);
73 |
74 | geminiService = new GeminiService();
75 | });
76 |
77 | afterEach(() => {
78 | vi.resetAllMocks();
79 | });
80 |
81 | describe("URL Context with Content Generation", () => {
82 | it("should successfully generate content with single URL context", async () => {
83 | const mockHtmlContent = `
84 | <!DOCTYPE html>
85 | <html>
86 | <head>
87 | <title>Test Article</title>
88 | <meta name="description" content="A comprehensive guide to testing">
89 | </head>
90 | <body>
91 | <h1>Introduction to Testing</h1>
92 | <p>Testing is essential for software quality assurance.</p>
93 | <h2>Types of Testing</h2>
94 | <ul>
95 | <li>Unit Testing</li>
96 | <li>Integration Testing</li>
97 | <li>End-to-End Testing</li>
98 | </ul>
99 | </body>
100 | </html>
101 | `;
102 |
103 | mockFetch.mockResolvedValueOnce({
104 | ok: true,
105 | status: 200,
106 | statusText: "OK",
107 | url: "https://example.com/testing-guide",
108 | headers: new Map([
109 | ["content-type", "text/html; charset=utf-8"],
110 | ["content-length", mockHtmlContent.length.toString()],
111 | ]),
112 | text: () => Promise.resolve(mockHtmlContent),
113 | });
114 |
115 | const result = await geminiService.generateContent({
116 | prompt: "Summarize the main points from the provided article",
117 | urlContext: {
118 | urls: ["https://example.com/testing-guide"],
119 | fetchOptions: {
120 | maxContentKb: 50,
121 | includeMetadata: true,
122 | },
123 | },
124 | });
125 |
126 | expect(result).toBeDefined();
127 | expect(result).toBe("Generated response based on URL content");
128 | expect(mockFetch).toHaveBeenCalledTimes(1);
129 | });
130 |
131 | it("should handle multiple URLs in context", async () => {
132 | const mockContent1 = `
133 | <html>
134 | <head><title>Article 1</title></head>
135 | <body><p>Content from first article about React development.</p></body>
136 | </html>
137 | `;
138 |
139 | const mockContent2 = `
140 | <html>
141 | <head><title>Article 2</title></head>
142 | <body><p>Content from second article about Vue.js development.</p></body>
143 | </html>
144 | `;
145 |
146 | mockFetch
147 | .mockResolvedValueOnce({
148 | ok: true,
149 | status: 200,
150 | url: "https://example1.com/react",
151 | headers: new Map([["content-type", "text/html"]]),
152 | text: () => Promise.resolve(mockContent1),
153 | })
154 | .mockResolvedValueOnce({
155 | ok: true,
156 | status: 200,
157 | url: "https://example2.com/vue",
158 | headers: new Map([["content-type", "text/html"]]),
159 | text: () => Promise.resolve(mockContent2),
160 | });
161 |
162 | const result = await geminiService.generateContent({
163 | prompt:
164 | "Compare the development approaches mentioned in these articles",
165 | urlContext: {
166 | urls: ["https://example1.com/react", "https://example2.com/vue"],
167 | },
168 | });
169 |
170 | expect(result).toBeDefined();
171 | expect(mockFetch).toHaveBeenCalledTimes(2);
172 | });
173 |
174 | it("should work with streaming content generation", async () => {
175 | const mockJsonContent = JSON.stringify({
176 | title: "API Documentation",
177 | endpoints: [
178 | { path: "/users", method: "GET" },
179 | { path: "/users", method: "POST" },
180 | ],
181 | });
182 |
183 | mockFetch.mockResolvedValueOnce({
184 | ok: true,
185 | status: 200,
186 | url: "https://api.example.com/docs",
187 | headers: new Map([["content-type", "application/json"]]),
188 | text: () => Promise.resolve(mockJsonContent),
189 | });
190 |
191 | const chunks: string[] = [];
192 | for await (const chunk of geminiService.generateContentStream({
193 | prompt: "Explain the API endpoints described in the documentation",
194 | urlContext: {
195 | urls: ["https://api.example.com/docs"],
196 | fetchOptions: {
197 | convertToMarkdown: false, // Keep JSON as-is
198 | },
199 | },
200 | })) {
201 | chunks.push(chunk);
202 | }
203 |
204 | const fullResponse = chunks.join("");
205 | expect(fullResponse).toBe("Generated response based on URL content");
206 | expect(mockFetch).toHaveBeenCalledTimes(1);
207 | });
208 | });
209 |
210 | describe("URL Context Error Handling", () => {
211 | it("should handle URL fetch failures gracefully", async () => {
212 | mockFetch.mockRejectedValueOnce(new Error("Network error"));
213 |
214 | await expect(
215 | geminiService.generateContent({
216 | prompt: "Analyze the content from this URL",
217 | urlContext: {
218 | urls: ["https://unreachable.com"],
219 | },
220 | })
221 | ).rejects.toThrow();
222 | });
223 |
224 | it("should handle mixed success/failure scenarios", async () => {
225 | const mockSuccessContent =
226 | "<html><body><p>Successful content</p></body></html>";
227 |
228 | mockFetch
229 | .mockResolvedValueOnce({
230 | ok: true,
231 | status: 200,
232 | url: "https://success.com",
233 | headers: new Map([["content-type", "text/html"]]),
234 | text: () => Promise.resolve(mockSuccessContent),
235 | })
236 | .mockRejectedValueOnce(new Error("Failed to fetch"))
237 | .mockResolvedValueOnce({
238 | ok: true,
239 | status: 200,
240 | url: "https://success2.com",
241 | headers: new Map([["content-type", "text/html"]]),
242 | text: () => Promise.resolve(mockSuccessContent),
243 | });
244 |
245 | // This should continue processing successful URLs despite some failures
246 | const result = await geminiService.generateContent({
247 | prompt: "Summarize the available content",
248 | urlContext: {
249 | urls: [
250 | "https://success.com",
251 | "https://failed.com",
252 | "https://success2.com",
253 | ],
254 | },
255 | });
256 |
257 | expect(result).toBeDefined();
258 | expect(mockFetch).toHaveBeenCalledTimes(3);
259 | });
260 |
261 | it("should respect URL context disabled configuration", async () => {
262 | mockConfig.getInstance().getUrlContextConfig.mockReturnValue({
263 | enabled: false,
264 | maxUrlsPerRequest: 20,
265 | defaultMaxContentKb: 100,
266 | defaultTimeoutMs: 10000,
267 | allowedDomains: ["*"],
268 | blocklistedDomains: [],
269 | });
270 |
271 | await expect(
272 | geminiService.generateContent({
273 | prompt: "Analyze this content",
274 | urlContext: {
275 | urls: ["https://example.com"],
276 | },
277 | })
278 | ).rejects.toThrow("URL context feature is not enabled");
279 |
280 | expect(mockFetch).not.toHaveBeenCalled();
281 | });
282 | });
283 |
284 | describe("URL Security Integration", () => {
285 | it("should block access to private networks", async () => {
286 | await expect(
287 | geminiService.generateContent({
288 | prompt: "Analyze the content",
289 | urlContext: {
290 | urls: ["http://192.168.1.1/admin"],
291 | },
292 | })
293 | ).rejects.toThrow();
294 |
295 | expect(mockFetch).not.toHaveBeenCalled();
296 | });
297 |
298 | it("should respect domain restrictions", async () => {
299 | mockConfig.getInstance().getUrlContextConfig.mockReturnValue({
300 | enabled: true,
301 | maxUrlsPerRequest: 20,
302 | defaultMaxContentKb: 100,
303 | defaultTimeoutMs: 10000,
304 | allowedDomains: ["example.com"],
305 | blocklistedDomains: [],
306 | });
307 |
308 | // Allowed domain should work
309 | mockFetch.mockResolvedValueOnce({
310 | ok: true,
311 | status: 200,
312 | url: "https://example.com",
313 | headers: new Map([["content-type", "text/html"]]),
314 | text: () => Promise.resolve("<html><body>Content</body></html>"),
315 | });
316 |
317 | await geminiService.generateContent({
318 | prompt: "Analyze this content",
319 | urlContext: {
320 | urls: ["https://example.com"],
321 | },
322 | });
323 |
324 | expect(mockFetch).toHaveBeenCalledTimes(1);
325 |
326 | // Disallowed domain should fail
327 | await expect(
328 | geminiService.generateContent({
329 | prompt: "Analyze this content",
330 | urlContext: {
331 | urls: ["https://other.com"],
332 | },
333 | })
334 | ).rejects.toThrow();
335 | });
336 |
337 | it("should enforce URL count limits", async () => {
338 | const manyUrls = Array.from(
339 | { length: 25 },
340 | (_, i) => `https://example${i}.com`
341 | );
342 |
343 | await expect(
344 | geminiService.generateContent({
345 | prompt: "Analyze all these URLs",
346 | urlContext: {
347 | urls: manyUrls,
348 | },
349 | })
350 | ).rejects.toThrow("Too many URLs");
351 |
352 | expect(mockFetch).not.toHaveBeenCalled();
353 | });
354 | });
355 |
356 | describe("Content Processing Integration", () => {
357 | it("should correctly convert HTML to Markdown", async () => {
358 | const complexHtml = `
359 | <html>
360 | <head><title>Complex Document</title></head>
361 | <body>
362 | <h1>Main Title</h1>
363 | <p>Paragraph with <strong>bold</strong> and <em>italic</em> text.</p>
364 | <ul>
365 | <li>List item 1</li>
366 | <li>List item 2 with <a href="https://example.com">link</a></li>
367 | </ul>
368 | <blockquote>This is a quote</blockquote>
369 | <code>inline code</code>
370 | <pre>code block</pre>
371 | </body>
372 | </html>
373 | `;
374 |
375 | mockFetch.mockResolvedValueOnce({
376 | ok: true,
377 | status: 200,
378 | url: "https://example.com/complex",
379 | headers: new Map([["content-type", "text/html"]]),
380 | text: () => Promise.resolve(complexHtml),
381 | });
382 |
383 | await geminiService.generateContent({
384 | prompt: "Process this complex document",
385 | urlContext: {
386 | urls: ["https://example.com/complex"],
387 | fetchOptions: {
388 | convertToMarkdown: true,
389 | includeMetadata: true,
390 | },
391 | },
392 | });
393 |
394 | expect(mockFetch).toHaveBeenCalledTimes(1);
395 | // The actual content processing is tested in unit tests
396 | });
397 |
398 | it("should handle large content with truncation", async () => {
399 | const largeContent = "x".repeat(500 * 1024); // 500KB content
400 |
401 | mockFetch.mockResolvedValueOnce({
402 | ok: true,
403 | status: 200,
404 | url: "https://example.com/large",
405 | headers: new Map([
406 | ["content-type", "text/html"],
407 | ["content-length", largeContent.length.toString()],
408 | ]),
409 | text: () => Promise.resolve(largeContent),
410 | });
411 |
412 | await geminiService.generateContent({
413 | prompt: "Summarize this large document",
414 | urlContext: {
415 | urls: ["https://example.com/large"],
416 | fetchOptions: {
417 | maxContentKb: 100, // Limit to 100KB
418 | },
419 | },
420 | });
421 |
422 | expect(mockFetch).toHaveBeenCalledTimes(1);
423 | });
424 | });
425 |
426 | describe("Model Selection Integration", () => {
427 | it("should prefer models with larger context windows for URL-heavy requests", async () => {
428 | const urls = Array.from(
429 | { length: 15 },
430 | (_, i) => `https://example${i}.com`
431 | );
432 |
433 | // Mock multiple successful fetches
434 | for (let i = 0; i < 15; i++) {
435 | mockFetch.mockResolvedValueOnce({
436 | ok: true,
437 | status: 200,
438 | url: urls[i],
439 | headers: new Map([["content-type", "text/html"]]),
440 | text: () => Promise.resolve(`<html><body>Content ${i}</body></html>`),
441 | });
442 | }
443 |
444 | const result = await geminiService.generateContent({
445 | prompt: "Analyze and compare all these sources",
446 | urlContext: {
447 | urls,
448 | },
449 | // Don't specify a model - let the service choose based on URL count
450 | taskType: "reasoning",
451 | complexityHint: "complex",
452 | });
453 |
454 | expect(result).toBeDefined();
455 | expect(mockFetch).toHaveBeenCalledTimes(15);
456 | });
457 | });
458 |
459 | describe("Caching Integration", () => {
460 | it("should cache URL content between requests", async () => {
461 | const mockContent = "<html><body><p>Cached content</p></body></html>";
462 |
463 | mockFetch.mockResolvedValue({
464 | ok: true,
465 | status: 200,
466 | url: "https://example.com/cached",
467 | headers: new Map([["content-type", "text/html"]]),
468 | text: () => Promise.resolve(mockContent),
469 | });
470 |
471 | // First request
472 | await geminiService.generateContent({
473 | prompt: "Analyze this content",
474 | urlContext: {
475 | urls: ["https://example.com/cached"],
476 | },
477 | });
478 |
479 | // Second request with same URL - should use cache
480 | await geminiService.generateContent({
481 | prompt: "Different analysis of the same content",
482 | urlContext: {
483 | urls: ["https://example.com/cached"],
484 | },
485 | });
486 |
487 | // Should only fetch once due to caching
488 | expect(mockFetch).toHaveBeenCalledTimes(1);
489 | });
490 | });
491 |
492 | describe("Rate Limiting Integration", () => {
493 | it("should enforce rate limits per domain", async () => {
494 | const baseUrl = "https://example.com/page";
495 |
496 | // Mock successful responses for rate limit testing
497 | for (let i = 0; i < 12; i++) {
498 | mockFetch.mockResolvedValueOnce({
499 | ok: true,
500 | status: 200,
501 | url: `${baseUrl}${i}`,
502 | headers: new Map([["content-type", "text/html"]]),
503 | text: () => Promise.resolve("<html><body>Content</body></html>"),
504 | });
505 | }
506 |
507 | // First 10 requests should succeed
508 | for (let i = 0; i < 10; i++) {
509 | await geminiService.generateContent({
510 | prompt: `Analyze page ${i}`,
511 | urlContext: {
512 | urls: [`${baseUrl}${i}`],
513 | },
514 | });
515 | }
516 |
517 | // 11th request should fail due to rate limiting
518 | await expect(
519 | geminiService.generateContent({
520 | prompt: "Analyze page 11",
521 | urlContext: {
522 | urls: [`${baseUrl}11`],
523 | },
524 | })
525 | ).rejects.toThrow();
526 | });
527 | });
528 | });
529 |
```
--------------------------------------------------------------------------------
/src/utils/errors.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Base custom error class for application-specific errors.
3 | */
4 | export class BaseError extends Error {
5 | public code: string;
6 | public readonly status: number; // HTTP status code equivalent
7 | public readonly details?: unknown; // Additional details
8 |
9 | constructor(
10 | message: string,
11 | code: string,
12 | status: number,
13 | details?: unknown
14 | ) {
15 | super(message);
16 | this.name = this.constructor.name; // Set the error name to the class name
17 | this.code = code;
18 | this.status = status;
19 | this.details = details;
20 | // Capture stack trace (excluding constructor)
21 | Error.captureStackTrace(this, this.constructor);
22 | }
23 | }
24 |
25 | /**
26 | * Error for validation failures (e.g., invalid input).
27 | * Maps typically to a 400 Bad Request or MCP InvalidParams.
28 | */
29 | export class ValidationError extends BaseError {
30 | constructor(message: string, details?: unknown) {
31 | super(message, "VALIDATION_ERROR", 400, details);
32 | }
33 | }
34 |
35 | /**
36 | * Error when an expected entity or resource is not found.
37 | * Maps typically to a 404 Not Found.
38 | */
39 | export class NotFoundError extends BaseError {
40 | constructor(message: string = "Resource not found") {
41 | super(message, "NOT_FOUND", 404);
42 | }
43 | }
44 |
45 | /**
46 | * Error for configuration problems.
47 | */
48 | export class ConfigurationError extends BaseError {
49 | constructor(message: string) {
50 | super(message, "CONFIG_ERROR", 500);
51 | }
52 | }
53 |
54 | /**
55 | * Error for issues during service processing unrelated to input validation.
56 | * Maps typically to a 500 Internal Server Error or MCP InternalError.
57 | */
58 | export class ServiceError extends BaseError {
59 | constructor(message: string, details?: unknown) {
60 | super(message, "SERVICE_ERROR", 500, details);
61 | }
62 | }
63 |
64 | /**
65 | * Error specifically for issues encountered when interacting with the Google Gemini API.
66 | * Extends ServiceError as it relates to an external service failure.
67 | */
68 | export class GeminiApiError extends ServiceError {
69 | constructor(message: string, details?: unknown) {
70 | // Call ServiceError constructor with only message and details
71 | super(`Gemini API Error: ${message}`, details);
72 | // Optionally add a specific code property if needed for finer-grained handling
73 | // this.code = 'GEMINI_API_ERROR'; // Overrides the 'SERVICE_ERROR' code from BaseError via ServiceError
74 | }
75 | }
76 |
77 | /**
78 | * Error specifically for when a file or resource is not found in the Gemini API.
79 | * Extends GeminiApiError to maintain the error hierarchy.
80 | */
81 | export class GeminiResourceNotFoundError extends GeminiApiError {
82 | constructor(resourceType: string, resourceId: string, details?: unknown) {
83 | super(`${resourceType} not found: ${resourceId}`, details);
84 | this.code = "GEMINI_RESOURCE_NOT_FOUND";
85 | }
86 | }
87 |
88 | /**
89 | * Error for invalid parameters when calling the Gemini API.
90 | * Extends GeminiApiError to maintain the error hierarchy.
91 | */
92 | export class GeminiInvalidParameterError extends GeminiApiError {
93 | constructor(message: string, details?: unknown) {
94 | super(`Invalid parameter: ${message}`, details);
95 | this.code = "GEMINI_INVALID_PARAMETER";
96 | }
97 | }
98 |
99 | /**
100 | * Error for authentication failures with the Gemini API.
101 | * Extends GeminiApiError to maintain the error hierarchy.
102 | */
103 | export class GeminiAuthenticationError extends GeminiApiError {
104 | constructor(message: string, details?: unknown) {
105 | super(`Authentication error: ${message}`, details);
106 | this.code = "GEMINI_AUTHENTICATION_ERROR";
107 | }
108 | }
109 |
110 | /**
111 | * Error for when Gemini API quota is exceeded or rate limits are hit.
112 | * Extends GeminiApiError to maintain the error hierarchy.
113 | */
114 | export class GeminiQuotaExceededError extends GeminiApiError {
115 | constructor(message: string, details?: unknown) {
116 | super(`Quota exceeded: ${message}`, details);
117 | this.code = "GEMINI_QUOTA_EXCEEDED";
118 | }
119 | }
120 |
121 | /**
122 | * Error for when content is blocked by Gemini's safety settings.
123 | * Extends GeminiApiError to maintain the error hierarchy.
124 | */
125 | export class GeminiSafetyError extends GeminiApiError {
126 | constructor(message: string, details?: unknown) {
127 | super(`Content blocked by safety settings: ${message}`, details);
128 | this.code = "GEMINI_SAFETY_ERROR";
129 | }
130 | }
131 |
132 | // Import the McpError and ErrorCode from the MCP SDK for use in the mapping function
133 | import { McpError, ErrorCode } from "@modelcontextprotocol/sdk/types.js";
134 | import { ToolError } from "./ToolError.js";
135 |
136 | // Re-export ToolError for use by tools
137 | export { ToolError };
138 |
139 | /**
140 | * Maps internal application errors to standardized MCP errors.
141 | * This function ensures consistent error mapping across all tool handlers.
142 | *
143 | * @param error - The error to be mapped to an MCP error
144 | * @param toolName - The name of the tool where the error occurred (for better error messages)
145 | * @returns McpError - A properly mapped MCP error
146 | */
147 | export function mapToMcpError(error: unknown, toolName: string): McpError {
148 | // If error is already an McpError, return it directly
149 | if (error instanceof McpError) {
150 | return error;
151 | }
152 |
153 | // Default error message if error is not an Error instance
154 | let errorMessage = "An unknown error occurred";
155 | let errorDetails: unknown = undefined;
156 |
157 | // Extract error message and details if error is an Error instance
158 | if (error instanceof Error) {
159 | errorMessage = error.message;
160 |
161 | // Extract details from BaseError instances
162 | if (error instanceof BaseError && error.details) {
163 | errorDetails = error.details;
164 | }
165 | } else if (typeof error === "string") {
166 | errorMessage = error;
167 | } else if (error !== null && typeof error === "object") {
168 | // Try to extract information from unknown object errors
169 | try {
170 | errorMessage = JSON.stringify(error);
171 | } catch {
172 | // If JSON stringification fails, use default message
173 | }
174 | }
175 |
176 | // ValidationError mapping
177 | if (error instanceof ValidationError) {
178 | return new McpError(
179 | ErrorCode.InvalidParams,
180 | `Validation error: ${errorMessage}`,
181 | errorDetails
182 | );
183 | }
184 |
185 | // NotFoundError mapping
186 | if (error instanceof NotFoundError) {
187 | return new McpError(
188 | ErrorCode.InvalidRequest,
189 | `Resource not found: ${errorMessage}`,
190 | errorDetails
191 | );
192 | }
193 |
194 | // ConfigurationError mapping
195 | if (error instanceof ConfigurationError) {
196 | return new McpError(
197 | ErrorCode.InternalError, // Changed from FailedPrecondition which is not in MCP SDK
198 | `Configuration error: ${errorMessage}`,
199 | errorDetails
200 | );
201 | }
202 |
203 | // Handle more specific Gemini API error subtypes first
204 | if (error instanceof GeminiResourceNotFoundError) {
205 | return new McpError(
206 | ErrorCode.InvalidRequest, // MCP SDK lacks NotFound, mapping to InvalidRequest
207 | `Resource not found: ${errorMessage}`,
208 | errorDetails
209 | );
210 | }
211 |
212 | if (error instanceof GeminiInvalidParameterError) {
213 | return new McpError(
214 | ErrorCode.InvalidParams,
215 | `Invalid parameters: ${errorMessage}`,
216 | errorDetails
217 | );
218 | }
219 |
220 | if (error instanceof GeminiAuthenticationError) {
221 | return new McpError(
222 | ErrorCode.InvalidRequest, // Changed from PermissionDenied which is not in MCP SDK
223 | `Authentication failed: ${errorMessage}`,
224 | errorDetails
225 | );
226 | }
227 |
228 | if (error instanceof GeminiQuotaExceededError) {
229 | return new McpError(
230 | ErrorCode.InternalError, // Changed from ResourceExhausted which is not in MCP SDK
231 | `Quota exceeded or rate limit hit: ${errorMessage}`,
232 | errorDetails
233 | );
234 | }
235 |
236 | if (error instanceof GeminiSafetyError) {
237 | return new McpError(
238 | ErrorCode.InvalidRequest,
239 | `Content blocked by safety settings: ${errorMessage}`,
240 | errorDetails
241 | );
242 | }
243 |
244 | // Generic GeminiApiError mapping with enhanced pattern detection
245 | if (error instanceof GeminiApiError) {
246 | // Convert message to lowercase for case-insensitive pattern matching
247 | const lowerCaseMessage = errorMessage.toLowerCase();
248 |
249 | // Handle rate limiting and quota errors
250 | if (
251 | lowerCaseMessage.includes("quota") ||
252 | lowerCaseMessage.includes("rate limit") ||
253 | lowerCaseMessage.includes("resource has been exhausted") ||
254 | lowerCaseMessage.includes("resource exhausted") ||
255 | lowerCaseMessage.includes("429") ||
256 | lowerCaseMessage.includes("too many requests")
257 | ) {
258 | return new McpError(
259 | ErrorCode.InternalError, // Changed from ResourceExhausted which is not in MCP SDK
260 | `Quota exceeded or rate limit hit: ${errorMessage}`,
261 | errorDetails
262 | );
263 | }
264 |
265 | // Handle permission and authorization errors
266 | if (
267 | lowerCaseMessage.includes("permission") ||
268 | lowerCaseMessage.includes("not authorized") ||
269 | lowerCaseMessage.includes("unauthorized") ||
270 | lowerCaseMessage.includes("forbidden") ||
271 | lowerCaseMessage.includes("403") ||
272 | lowerCaseMessage.includes("access denied")
273 | ) {
274 | return new McpError(
275 | ErrorCode.InvalidRequest, // Changed from PermissionDenied which is not in MCP SDK
276 | `Permission denied: ${errorMessage}`,
277 | errorDetails
278 | );
279 | }
280 |
281 | // Handle not found errors
282 | if (
283 | lowerCaseMessage.includes("not found") ||
284 | lowerCaseMessage.includes("does not exist") ||
285 | lowerCaseMessage.includes("404") ||
286 | lowerCaseMessage.includes("could not find") ||
287 | lowerCaseMessage.includes("no such file")
288 | ) {
289 | return new McpError(
290 | ErrorCode.InvalidRequest, // MCP SDK lacks NotFound, mapping to InvalidRequest
291 | `Resource not found: ${errorMessage}`,
292 | errorDetails
293 | );
294 | }
295 |
296 | // Handle invalid argument/parameter errors
297 | if (
298 | lowerCaseMessage.includes("invalid argument") ||
299 | lowerCaseMessage.includes("invalid parameter") ||
300 | lowerCaseMessage.includes("invalid request") ||
301 | lowerCaseMessage.includes("failed precondition") ||
302 | lowerCaseMessage.includes("400") ||
303 | lowerCaseMessage.includes("bad request") ||
304 | lowerCaseMessage.includes("malformed")
305 | ) {
306 | return new McpError(
307 | ErrorCode.InvalidParams,
308 | `Invalid parameters: ${errorMessage}`,
309 | errorDetails
310 | );
311 | }
312 |
313 | // Handle safety-related errors
314 | if (
315 | lowerCaseMessage.includes("safety") ||
316 | lowerCaseMessage.includes("blocked") ||
317 | lowerCaseMessage.includes("content policy") ||
318 | lowerCaseMessage.includes("harmful") ||
319 | lowerCaseMessage.includes("inappropriate") ||
320 | lowerCaseMessage.includes("offensive")
321 | ) {
322 | return new McpError(
323 | ErrorCode.InvalidRequest,
324 | `Content blocked by safety settings: ${errorMessage}`,
325 | errorDetails
326 | );
327 | }
328 |
329 | // Handle File API and other unsupported feature errors
330 | if (
331 | lowerCaseMessage.includes("file api is not supported") ||
332 | lowerCaseMessage.includes("not supported") ||
333 | lowerCaseMessage.includes("unsupported") ||
334 | lowerCaseMessage.includes("not implemented")
335 | ) {
336 | return new McpError(
337 | ErrorCode.InvalidRequest, // Changed from FailedPrecondition which is not in MCP SDK
338 | `Operation not supported: ${errorMessage}`,
339 | errorDetails
340 | );
341 | }
342 |
343 | // Default case for GeminiApiError - map to internal error
344 | return new McpError(
345 | ErrorCode.InternalError,
346 | `Gemini API Error: ${errorMessage}`,
347 | errorDetails
348 | );
349 | }
350 |
351 | // Generic ServiceError mapping
352 | if (error instanceof ServiceError) {
353 | return new McpError(
354 | ErrorCode.InternalError,
355 | `Service error: ${errorMessage}`,
356 | errorDetails
357 | );
358 | }
359 |
360 | // Default case for all other errors
361 | return new McpError(
362 | ErrorCode.InternalError,
363 | `[${toolName}] Failed: ${errorMessage}`
364 | );
365 | }
366 |
367 | /**
368 | * Combined error mapping function that handles both standard errors and ToolError instances.
369 | * This function accommodates the different error types used across different tool implementations.
370 | *
371 | * @param error - Any error type, including McpError, BaseError, ToolError, or standard Error
372 | * @param toolName - The name of the tool where the error occurred
373 | * @returns McpError - A consistently mapped MCP error
374 | */
375 | export function mapAnyErrorToMcpError(
376 | error: unknown,
377 | toolName: string
378 | ): McpError {
379 | // Check if error is a ToolError from image feature tools
380 | if (
381 | error !== null &&
382 | typeof error === "object" &&
383 | "code" in error &&
384 | typeof (error as ToolErrorLike).code === "string"
385 | ) {
386 | // For objects that match the ToolError interface
387 | return mapToolErrorToMcpError(error as ToolErrorLike, toolName);
388 | }
389 |
390 | // For standard errors and BaseError types
391 | return mapToMcpError(error, toolName);
392 | }
393 |
394 | /**
395 | * Interface for objects that conform to the ToolError structure
396 | * This provides type safety for objects that have a similar structure to ToolError
397 | * but may not be actual instances of the ToolError class.
398 | */
399 | export interface ToolErrorLike {
400 | code?: string;
401 | message?: string;
402 | details?: unknown;
403 | [key: string]: unknown; // Allow additional properties for flexibility
404 | }
405 |
406 | // These tools use a different error structure than the rest of the application
407 | // but need to maintain consistent error mapping to McpError
408 |
409 | /**
410 | * Maps ToolError instances used in some image feature tools to McpError.
411 | * This is a compatibility layer for tools that use a different error structure.
412 | *
413 | * @param toolError - The ToolError instance or object with code/details properties
414 | * @param toolName - The name of the tool for better error messages
415 | * @returns McpError - A consistent MCP error
416 | */
417 | export function mapToolErrorToMcpError(
418 | toolError: ToolErrorLike | unknown,
419 | toolName: string
420 | ): McpError {
421 | // Default message if more specific extraction fails
422 | let errorMessage = `Error in ${toolName}`;
423 | let errorDetails: unknown = undefined;
424 |
425 | // Extract error message and details if possible
426 | if (toolError && typeof toolError === "object") {
427 | const errorObj = toolError as ToolErrorLike;
428 |
429 | // Extract message
430 | if ("message" in errorObj && typeof errorObj.message === "string") {
431 | errorMessage = errorObj.message;
432 | }
433 |
434 | // Extract details
435 | if ("details" in errorObj) {
436 | errorDetails = errorObj.details;
437 | }
438 |
439 | // Extract code for mapping
440 | if ("code" in errorObj && typeof errorObj.code === "string") {
441 | const code = errorObj.code.toUpperCase();
442 |
443 | // Map common ToolError codes to appropriate ErrorCode values
444 | if (code.includes("SAFETY") || code.includes("BLOCKED")) {
445 | return new McpError(
446 | ErrorCode.InvalidRequest,
447 | `Content blocked by safety settings: ${errorMessage}`,
448 | errorDetails
449 | );
450 | }
451 |
452 | if (code.includes("QUOTA") || code.includes("RATE_LIMIT")) {
453 | return new McpError(
454 | ErrorCode.InternalError, // Changed from ResourceExhausted which is not in MCP SDK
455 | `API quota or rate limit exceeded: ${errorMessage}`,
456 | errorDetails
457 | );
458 | }
459 |
460 | if (code.includes("PERMISSION") || code.includes("AUTH")) {
461 | return new McpError(
462 | ErrorCode.InvalidRequest, // Changed from PermissionDenied which is not in MCP SDK
463 | `Permission denied: ${errorMessage}`,
464 | errorDetails
465 | );
466 | }
467 |
468 | if (code.includes("NOT_FOUND")) {
469 | return new McpError(
470 | ErrorCode.InvalidRequest,
471 | `Resource not found: ${errorMessage}`,
472 | errorDetails
473 | );
474 | }
475 |
476 | if (code.includes("INVALID") || code.includes("ARGUMENT")) {
477 | return new McpError(
478 | ErrorCode.InvalidParams,
479 | `Invalid parameters: ${errorMessage}`,
480 | errorDetails
481 | );
482 | }
483 |
484 | if (code.includes("UNSUPPORTED") || code.includes("NOT_SUPPORTED")) {
485 | return new McpError(
486 | ErrorCode.InvalidRequest, // Changed from FailedPrecondition which is not in MCP SDK
487 | `Operation not supported: ${errorMessage}`,
488 | errorDetails
489 | );
490 | }
491 | }
492 | }
493 |
494 | // Default to internal error for any other case
495 | return new McpError(
496 | ErrorCode.InternalError,
497 | `[${toolName}] Error: ${errorMessage}`,
498 | errorDetails
499 | );
500 | }
501 |
```
--------------------------------------------------------------------------------
/tests/unit/tools/schemas/ToolSchemas.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import {
3 | ToolSchema,
4 | ToolResponseSchema,
5 | FunctionParameterSchema,
6 | FunctionDeclarationSchema,
7 | } from "../../../../src/tools/schemas/ToolSchemas.js";
8 |
9 | import {
10 | HarmCategorySchema,
11 | SafetySettingSchema,
12 | ThinkingConfigSchema,
13 | GenerationConfigSchema,
14 | FilePathSchema,
15 | FileOverwriteSchema,
16 | EncodingSchema,
17 | ModelNameSchema,
18 | PromptSchema,
19 | } from "../../../../src/tools/schemas/CommonSchemas.js";
20 |
21 | describe("Tool Schemas Validation", () => {
22 | describe("ToolSchema", () => {
23 | it("should validate a valid tool definition with function declarations", () => {
24 | const validTool = {
25 | functionDeclarations: [
26 | {
27 | name: "testFunction",
28 | description: "A test function",
29 | parameters: {
30 | type: "OBJECT",
31 | properties: {
32 | name: {
33 | type: "STRING",
34 | description: "The name parameter",
35 | },
36 | },
37 | required: ["name"],
38 | },
39 | },
40 | ],
41 | };
42 |
43 | const result = ToolSchema.safeParse(validTool);
44 | expect(result.success).toBe(true);
45 | });
46 |
47 | it("should validate a tool with no function declarations", () => {
48 | const emptyTool = {};
49 | const result = ToolSchema.safeParse(emptyTool);
50 | expect(result.success).toBe(true);
51 | });
52 |
53 | it("should reject invalid function declarations", () => {
54 | const invalidTool = {
55 | functionDeclarations: [
56 | {
57 | // Missing required name field
58 | description: "A test function",
59 | parameters: {
60 | type: "OBJECT",
61 | properties: {},
62 | },
63 | },
64 | ],
65 | };
66 |
67 | const result = ToolSchema.safeParse(invalidTool);
68 | expect(result.success).toBe(false);
69 | });
70 | });
71 |
72 | describe("ToolResponseSchema", () => {
73 | it("should validate a valid tool response", () => {
74 | const validResponse = {
75 | name: "testTool",
76 | response: { result: "success" },
77 | };
78 |
79 | const result = ToolResponseSchema.safeParse(validResponse);
80 | expect(result.success).toBe(true);
81 | });
82 |
83 | it("should reject response with missing name", () => {
84 | const invalidResponse = {
85 | response: { result: "success" },
86 | };
87 |
88 | const result = ToolResponseSchema.safeParse(invalidResponse);
89 | expect(result.success).toBe(false);
90 | });
91 | });
92 |
93 | describe("FunctionParameterSchema", () => {
94 | it("should validate primitive parameter types", () => {
95 | const stringParam = {
96 | type: "STRING",
97 | description: "A string parameter",
98 | };
99 |
100 | const numberParam = {
101 | type: "NUMBER",
102 | description: "A number parameter",
103 | };
104 |
105 | const booleanParam = {
106 | type: "BOOLEAN",
107 | };
108 |
109 | expect(FunctionParameterSchema.safeParse(stringParam).success).toBe(true);
110 | expect(FunctionParameterSchema.safeParse(numberParam).success).toBe(true);
111 | expect(FunctionParameterSchema.safeParse(booleanParam).success).toBe(
112 | true
113 | );
114 | });
115 |
116 | it("should validate object parameter with nested properties", () => {
117 | const objectParam = {
118 | type: "OBJECT",
119 | description: "An object parameter",
120 | properties: {
121 | name: {
122 | type: "STRING",
123 | },
124 | age: {
125 | type: "INTEGER",
126 | },
127 | details: {
128 | type: "OBJECT",
129 | properties: {
130 | address: {
131 | type: "STRING",
132 | },
133 | },
134 | },
135 | },
136 | required: ["name"],
137 | };
138 |
139 | const result = FunctionParameterSchema.safeParse(objectParam);
140 | expect(result.success).toBe(true);
141 | });
142 |
143 | it("should validate array parameter with items", () => {
144 | const arrayParam = {
145 | type: "ARRAY",
146 | description: "An array parameter",
147 | items: {
148 | type: "STRING",
149 | },
150 | };
151 |
152 | const result = FunctionParameterSchema.safeParse(arrayParam);
153 | expect(result.success).toBe(true);
154 | });
155 |
156 | it("should reject parameter with invalid type", () => {
157 | const invalidParam = {
158 | type: "INVALID_TYPE", // Not a valid type
159 | description: "An invalid parameter",
160 | };
161 |
162 | const result = FunctionParameterSchema.safeParse(invalidParam);
163 | expect(result.success).toBe(false);
164 | });
165 | });
166 |
167 | describe("FunctionDeclarationSchema", () => {
168 | it("should validate a valid function declaration", () => {
169 | const validFunction = {
170 | name: "testFunction",
171 | description: "A test function",
172 | parameters: {
173 | type: "OBJECT",
174 | properties: {
175 | name: {
176 | type: "STRING",
177 | description: "The name parameter",
178 | },
179 | age: {
180 | type: "INTEGER",
181 | },
182 | },
183 | required: ["name"],
184 | },
185 | };
186 |
187 | const result = FunctionDeclarationSchema.safeParse(validFunction);
188 | expect(result.success).toBe(true);
189 | });
190 |
191 | it("should reject function declaration with missing required fields", () => {
192 | const invalidFunction = {
193 | // Missing name
194 | description: "A test function",
195 | parameters: {
196 | type: "OBJECT",
197 | properties: {},
198 | },
199 | };
200 |
201 | const result = FunctionDeclarationSchema.safeParse(invalidFunction);
202 | expect(result.success).toBe(false);
203 | });
204 |
205 | it("should reject function declaration with invalid parameters type", () => {
206 | const invalidFunction = {
207 | name: "testFunction",
208 | description: "A test function",
209 | parameters: {
210 | type: "STRING", // Should be "OBJECT"
211 | properties: {},
212 | },
213 | };
214 |
215 | const result = FunctionDeclarationSchema.safeParse(invalidFunction);
216 | expect(result.success).toBe(false);
217 | });
218 | });
219 |
220 | describe("CommonSchemas", () => {
221 | describe("HarmCategorySchema", () => {
222 | it("should validate valid harm categories", () => {
223 | const validCategories = [
224 | "HARM_CATEGORY_UNSPECIFIED",
225 | "HARM_CATEGORY_HATE_SPEECH",
226 | "HARM_CATEGORY_SEXUALLY_EXPLICIT",
227 | "HARM_CATEGORY_HARASSMENT",
228 | "HARM_CATEGORY_DANGEROUS_CONTENT",
229 | ];
230 |
231 | validCategories.forEach((category) => {
232 | expect(HarmCategorySchema.safeParse(category).success).toBe(true);
233 | });
234 | });
235 |
236 | it("should reject invalid harm categories", () => {
237 | expect(HarmCategorySchema.safeParse("INVALID_CATEGORY").success).toBe(
238 | false
239 | );
240 | });
241 | });
242 |
243 | describe("SafetySettingSchema", () => {
244 | it("should validate a valid safety setting", () => {
245 | const validSetting = {
246 | category: "HARM_CATEGORY_HATE_SPEECH",
247 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
248 | };
249 |
250 | const result = SafetySettingSchema.safeParse(validSetting);
251 | expect(result.success).toBe(true);
252 | });
253 |
254 | it("should validate all valid combinations of categories and thresholds", () => {
255 | const validCategories = [
256 | "HARM_CATEGORY_UNSPECIFIED",
257 | "HARM_CATEGORY_HATE_SPEECH",
258 | "HARM_CATEGORY_SEXUALLY_EXPLICIT",
259 | "HARM_CATEGORY_HARASSMENT",
260 | "HARM_CATEGORY_DANGEROUS_CONTENT",
261 | ];
262 |
263 | const validThresholds = [
264 | "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
265 | "BLOCK_LOW_AND_ABOVE",
266 | "BLOCK_MEDIUM_AND_ABOVE",
267 | "BLOCK_ONLY_HIGH",
268 | "BLOCK_NONE",
269 | ];
270 |
271 | // Test a sampling of combinations
272 | for (const category of validCategories) {
273 | for (const threshold of validThresholds) {
274 | const setting = { category, threshold };
275 | expect(SafetySettingSchema.safeParse(setting).success).toBe(true);
276 | }
277 | }
278 | });
279 |
280 | it("should reject setting with valid structure but invalid category", () => {
281 | const invalidSetting = {
282 | category: "INVALID_CATEGORY",
283 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
284 | };
285 |
286 | const result = SafetySettingSchema.safeParse(invalidSetting);
287 | expect(result.success).toBe(false);
288 | });
289 |
290 | it("should reject setting with valid structure but invalid threshold", () => {
291 | const invalidSetting = {
292 | category: "HARM_CATEGORY_HATE_SPEECH",
293 | threshold: "INVALID_THRESHOLD",
294 | };
295 |
296 | const result = SafetySettingSchema.safeParse(invalidSetting);
297 | expect(result.success).toBe(false);
298 | });
299 |
300 | it("should reject setting with missing required fields", () => {
301 | const missingCategory = {
302 | threshold: "BLOCK_MEDIUM_AND_ABOVE",
303 | };
304 |
305 | const missingThreshold = {
306 | category: "HARM_CATEGORY_HATE_SPEECH",
307 | };
308 |
309 | expect(SafetySettingSchema.safeParse(missingCategory).success).toBe(
310 | false
311 | );
312 | expect(SafetySettingSchema.safeParse(missingThreshold).success).toBe(
313 | false
314 | );
315 | });
316 | });
317 |
318 | describe("GenerationConfigSchema", () => {
319 | it("should validate a valid generation config", () => {
320 | const validConfig = {
321 | temperature: 0.7,
322 | topP: 0.9,
323 | topK: 40,
324 | maxOutputTokens: 1024,
325 | stopSequences: ["STOP", "END"],
326 | thinkingConfig: {
327 | thinkingBudget: 1000,
328 | reasoningEffort: "medium",
329 | },
330 | };
331 |
332 | const result = GenerationConfigSchema.safeParse(validConfig);
333 | expect(result.success).toBe(true);
334 | });
335 |
336 | it("should validate minimal generation config", () => {
337 | const minimalConfig = {};
338 | const result = GenerationConfigSchema.safeParse(minimalConfig);
339 | expect(result.success).toBe(true);
340 | });
341 |
342 | describe("temperature parameter boundary values", () => {
343 | it("should validate minimum valid temperature (0)", () => {
344 | const config = { temperature: 0 };
345 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
346 | });
347 |
348 | it("should validate maximum valid temperature (1)", () => {
349 | const config = { temperature: 1 };
350 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
351 | });
352 |
353 | it("should reject temperature below minimum (-0.1)", () => {
354 | const config = { temperature: -0.1 };
355 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
356 | });
357 |
358 | it("should reject temperature above maximum (1.01)", () => {
359 | const config = { temperature: 1.01 };
360 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
361 | });
362 | });
363 |
364 | describe("topP parameter boundary values", () => {
365 | it("should validate minimum valid topP (0)", () => {
366 | const config = { topP: 0 };
367 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
368 | });
369 |
370 | it("should validate maximum valid topP (1)", () => {
371 | const config = { topP: 1 };
372 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
373 | });
374 |
375 | it("should reject topP below minimum (-0.1)", () => {
376 | const config = { topP: -0.1 };
377 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
378 | });
379 |
380 | it("should reject topP above maximum (1.01)", () => {
381 | const config = { topP: 1.01 };
382 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
383 | });
384 | });
385 |
386 | describe("topK parameter boundary values", () => {
387 | it("should validate minimum valid topK (1)", () => {
388 | const config = { topK: 1 };
389 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
390 | });
391 |
392 | it("should reject topK below minimum (0)", () => {
393 | const config = { topK: 0 };
394 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
395 | });
396 |
397 | it("should validate large topK values", () => {
398 | const config = { topK: 1000 };
399 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
400 | });
401 | });
402 |
403 | describe("maxOutputTokens parameter boundary values", () => {
404 | it("should validate minimum valid maxOutputTokens (1)", () => {
405 | const config = { maxOutputTokens: 1 };
406 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
407 | });
408 |
409 | it("should reject maxOutputTokens below minimum (0)", () => {
410 | const config = { maxOutputTokens: 0 };
411 | expect(GenerationConfigSchema.safeParse(config).success).toBe(false);
412 | });
413 |
414 | it("should validate large maxOutputTokens values", () => {
415 | const config = { maxOutputTokens: 10000 };
416 | expect(GenerationConfigSchema.safeParse(config).success).toBe(true);
417 | });
418 | });
419 | });
420 |
421 | describe("ThinkingConfigSchema", () => {
422 | it("should validate valid thinking configs", () => {
423 | const validConfigs = [
424 | { thinkingBudget: 1000 },
425 | { reasoningEffort: "medium" },
426 | { thinkingBudget: 5000, reasoningEffort: "high" },
427 | {}, // Empty config is valid
428 | ];
429 |
430 | validConfigs.forEach((config) => {
431 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(true);
432 | });
433 | });
434 |
435 | describe("thinkingBudget parameter boundary values", () => {
436 | it("should validate minimum valid thinkingBudget (0)", () => {
437 | const config = { thinkingBudget: 0 };
438 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(true);
439 | });
440 |
441 | it("should validate maximum valid thinkingBudget (24576)", () => {
442 | const config = { thinkingBudget: 24576 };
443 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(true);
444 | });
445 |
446 | it("should reject thinkingBudget below minimum (-1)", () => {
447 | const config = { thinkingBudget: -1 };
448 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(false);
449 | });
450 |
451 | it("should reject thinkingBudget above maximum (24577)", () => {
452 | const config = { thinkingBudget: 24577 };
453 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(false);
454 | });
455 |
456 | it("should reject non-integer thinkingBudget (1000.5)", () => {
457 | const config = { thinkingBudget: 1000.5 };
458 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(false);
459 | });
460 | });
461 |
462 | describe("reasoningEffort parameter values", () => {
463 | it("should validate all valid reasoningEffort options", () => {
464 | const validOptions = ["none", "low", "medium", "high"];
465 |
466 | validOptions.forEach((option) => {
467 | const config = { reasoningEffort: option };
468 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(true);
469 | });
470 | });
471 |
472 | it("should reject invalid reasoningEffort options", () => {
473 | const invalidOptions = ["maximum", "minimal", "very-high", ""];
474 |
475 | invalidOptions.forEach((option) => {
476 | const config = { reasoningEffort: option };
477 | expect(ThinkingConfigSchema.safeParse(config).success).toBe(false);
478 | });
479 | });
480 | });
481 | });
482 |
483 | describe("File Operation Schemas", () => {
484 | it("should validate valid file paths", () => {
485 | const validPaths = [
486 | "/path/to/file.txt",
487 | "C:\\Windows\\System32\\file.exe",
488 | ];
489 |
490 | validPaths.forEach((path) => {
491 | expect(FilePathSchema.safeParse(path).success).toBe(true);
492 | });
493 | });
494 |
495 | it("should reject empty file paths", () => {
496 | expect(FilePathSchema.safeParse("").success).toBe(false);
497 | });
498 |
499 | it("should validate file overwrite options", () => {
500 | expect(FileOverwriteSchema.safeParse(true).success).toBe(true);
501 | expect(FileOverwriteSchema.safeParse(false).success).toBe(true);
502 | expect(FileOverwriteSchema.safeParse(undefined).success).toBe(true);
503 | });
504 |
505 | it("should validate encoding options", () => {
506 | expect(EncodingSchema.safeParse("utf8").success).toBe(true);
507 | expect(EncodingSchema.safeParse("base64").success).toBe(true);
508 | expect(EncodingSchema.safeParse(undefined).success).toBe(true);
509 | expect(EncodingSchema.safeParse("binary").success).toBe(false);
510 | });
511 | });
512 |
513 | describe("Other Common Schemas", () => {
514 | it("should validate model names", () => {
515 | expect(ModelNameSchema.safeParse("gemini-pro").success).toBe(true);
516 | expect(ModelNameSchema.safeParse("").success).toBe(false);
517 | });
518 |
519 | it("should validate prompts", () => {
520 | expect(PromptSchema.safeParse("Tell me a story").success).toBe(true);
521 | expect(PromptSchema.safeParse("").success).toBe(false);
522 | });
523 | });
524 | });
525 | });
526 |
```
--------------------------------------------------------------------------------
/tests/unit/utils/FileSecurityService.test.vitest.ts:
--------------------------------------------------------------------------------
```typescript
1 | // Using vitest globals - see vitest.config.ts globals: true
2 | import * as path from "path";
3 | import * as fs from "fs/promises";
4 | import * as fsSync from "fs";
5 |
6 | // Import the code to test
7 | import { FileSecurityService } from "../../../src/utils/FileSecurityService.js";
8 | import { ValidationError } from "../../../src/utils/errors.js";
9 | import { logger } from "../../../src/utils/logger.js";
10 |
11 | describe("FileSecurityService", () => {
12 | // Mock logger
13 | const loggerMock = {
14 | info: vi.fn(),
15 | warn: vi.fn(),
16 | error: vi.fn(),
17 | debug: vi.fn(),
18 | };
19 |
20 | // Define test constants for all tests
21 | const TEST_CONTENT = "Test file content";
22 | const TEST_DIR = path.resolve("./test-security-dir");
23 | const OUTSIDE_DIR = path.resolve("./outside-security-dir");
24 |
25 | // Setup before each test
26 | beforeEach(() => {
27 | // Reset mocks and create test directories
28 | vi.clearAllMocks();
29 |
30 | // Replace logger with mock
31 | vi.spyOn(logger, "info").mockImplementation(loggerMock.info);
32 | vi.spyOn(logger, "warn").mockImplementation(loggerMock.warn);
33 | vi.spyOn(logger, "error").mockImplementation(loggerMock.error);
34 | vi.spyOn(logger, "debug").mockImplementation(loggerMock.debug);
35 |
36 | // Create test directories
37 | fsSync.mkdirSync(TEST_DIR, { recursive: true });
38 | fsSync.mkdirSync(OUTSIDE_DIR, { recursive: true });
39 | });
40 |
41 | // Cleanup after each test
42 | afterEach(() => {
43 | // Restore original logger
44 | vi.restoreAllMocks();
45 |
46 | // Clean up test directories
47 | try {
48 | fsSync.rmSync(TEST_DIR, { recursive: true, force: true });
49 | fsSync.rmSync(OUTSIDE_DIR, { recursive: true, force: true });
50 | } catch (err) {
51 | // Ignore cleanup errors
52 | }
53 | });
54 |
55 | describe("Constructor and Configuration", () => {
56 | it("should initialize with default allowed directories", () => {
57 | const service = new FileSecurityService();
58 | const allowedDirs = service.getAllowedDirectories();
59 |
60 | expect(allowedDirs.length).toBeGreaterThan(0);
61 | expect(allowedDirs).toContain(path.resolve(process.cwd()));
62 | });
63 |
64 | it("should initialize with custom allowed directories", () => {
65 | const customDirs = [TEST_DIR, OUTSIDE_DIR];
66 | const service = new FileSecurityService(customDirs);
67 | const allowedDirs = service.getAllowedDirectories();
68 |
69 | expect(allowedDirs.length).toBe(2);
70 | expect(allowedDirs).toContain(path.resolve(TEST_DIR));
71 | expect(allowedDirs).toContain(path.resolve(OUTSIDE_DIR));
72 | });
73 |
74 | it("should initialize with a secure base path", () => {
75 | const service = new FileSecurityService([], TEST_DIR);
76 | const basePath = service.getSecureBasePath();
77 |
78 | expect(basePath).toBe(path.normalize(TEST_DIR));
79 |
80 | // Verify allowed directories includes the base path
81 | const allowedDirs = service.getAllowedDirectories();
82 | expect(allowedDirs).toContain(path.normalize(TEST_DIR));
83 | });
84 |
85 | it("should set allowed directories", () => {
86 | const service = new FileSecurityService();
87 | const newDirs = [TEST_DIR, OUTSIDE_DIR];
88 |
89 | service.setAllowedDirectories(newDirs);
90 | const allowedDirs = service.getAllowedDirectories();
91 |
92 | expect(allowedDirs.length).toBe(2);
93 | expect(allowedDirs).toContain(path.normalize(TEST_DIR));
94 | expect(allowedDirs).toContain(path.normalize(OUTSIDE_DIR));
95 | });
96 |
97 | it("should throw error when setting empty allowed directories", () => {
98 | const service = new FileSecurityService();
99 |
100 | expect(() => service.setAllowedDirectories([])).toThrow(ValidationError);
101 | expect(() => service.setAllowedDirectories([])).toThrow(
102 | /At least one allowed directory/
103 | );
104 | });
105 |
106 | it("should throw error when setting non-absolute allowed directories", () => {
107 | const service = new FileSecurityService();
108 |
109 | expect(() => service.setAllowedDirectories(["./relative/path"])).toThrow(
110 | ValidationError
111 | );
112 | expect(() => service.setAllowedDirectories(["./relative/path"])).toThrow(
113 | /Directory path must be absolute/
114 | );
115 | });
116 |
117 | it("should set and get secure base path", () => {
118 | const service = new FileSecurityService();
119 | service.setSecureBasePath(TEST_DIR);
120 |
121 | const basePath = service.getSecureBasePath();
122 | expect(basePath).toBe(path.normalize(TEST_DIR));
123 | });
124 |
125 | it("should throw error when setting non-absolute secure base path", () => {
126 | const service = new FileSecurityService();
127 |
128 | expect(() => service.setSecureBasePath("./relative/path")).toThrow(
129 | ValidationError
130 | );
131 | expect(() => service.setSecureBasePath("./relative/path")).toThrow(
132 | /Base path must be absolute/
133 | );
134 | });
135 |
136 | it("should configure from environment", () => {
137 | // Save original env var
138 | const originalEnvVar = process.env.GEMINI_SAFE_FILE_BASE_DIR;
139 |
140 | // Set env var for test
141 | process.env.GEMINI_SAFE_FILE_BASE_DIR = TEST_DIR;
142 |
143 | const service = FileSecurityService.configureFromEnvironment();
144 | const allowedDirs = service.getAllowedDirectories();
145 |
146 | expect(allowedDirs).toContain(path.normalize(TEST_DIR));
147 |
148 | // Restore original env var
149 | if (originalEnvVar) {
150 | process.env.GEMINI_SAFE_FILE_BASE_DIR = originalEnvVar;
151 | } else {
152 | delete process.env.GEMINI_SAFE_FILE_BASE_DIR;
153 | }
154 | });
155 | });
156 |
157 | describe("Path Validation", () => {
158 | let service: FileSecurityService;
159 |
160 | beforeEach(() => {
161 | service = new FileSecurityService([TEST_DIR]);
162 | });
163 |
164 | it("should validate path within allowed directory", () => {
165 | const testFilePath = path.join(TEST_DIR, "test-file.txt");
166 | const validatedPath = service.validateAndResolvePath(testFilePath);
167 |
168 | expect(validatedPath).toBe(path.normalize(testFilePath));
169 | });
170 |
171 | it("should validate paths with relative components", () => {
172 | const complexPath = path.join(
173 | TEST_DIR,
174 | ".",
175 | "subdir",
176 | "..",
177 | "test-file.txt"
178 | );
179 | const validatedPath = service.validateAndResolvePath(complexPath);
180 |
181 | // Should normalize to TEST_DIR/test-file.txt
182 | const expectedPath = path.normalize(path.join(TEST_DIR, "test-file.txt"));
183 | expect(validatedPath).toBe(expectedPath);
184 | });
185 |
186 | it("should reject paths outside allowed directories", () => {
187 | const outsidePath = path.join(OUTSIDE_DIR, "test-file.txt");
188 |
189 | expect(() => service.validateAndResolvePath(outsidePath)).toThrow(
190 | ValidationError
191 | );
192 | expect(() => service.validateAndResolvePath(outsidePath)).toThrow(
193 | /Access denied/
194 | );
195 | });
196 |
197 | it("should reject paths with directory traversal", () => {
198 | const traversalPath = path.join(
199 | TEST_DIR,
200 | "..",
201 | "outside",
202 | "test-file.txt"
203 | );
204 |
205 | expect(() => service.validateAndResolvePath(traversalPath)).toThrow(
206 | ValidationError
207 | );
208 | expect(() => service.validateAndResolvePath(traversalPath)).toThrow(
209 | /Access denied/
210 | );
211 | });
212 |
213 | it("should check file existence with mustExist option", () => {
214 | const nonExistentPath = path.join(TEST_DIR, "non-existent.txt");
215 |
216 | expect(() =>
217 | service.validateAndResolvePath(nonExistentPath, { mustExist: true })
218 | ).toThrow(ValidationError);
219 | expect(() =>
220 | service.validateAndResolvePath(nonExistentPath, { mustExist: true })
221 | ).toThrow(/File not found/);
222 | });
223 |
224 | it("should use custom allowed directories when provided", () => {
225 | // Path is outside the service's configured directory but inside custom allowed dir
226 | const customAllowedPath = path.join(OUTSIDE_DIR, "custom-allowed.txt");
227 |
228 | const validatedPath = service.validateAndResolvePath(customAllowedPath, {
229 | allowedDirs: [OUTSIDE_DIR],
230 | });
231 |
232 | expect(validatedPath).toBe(path.normalize(customAllowedPath));
233 | });
234 | });
235 |
236 | describe("isPathWithinAllowedDirs", () => {
237 | let service: FileSecurityService;
238 |
239 | beforeEach(() => {
240 | service = new FileSecurityService([TEST_DIR]);
241 | });
242 |
243 | it("should return true for paths within allowed directories", () => {
244 | const insidePath = path.join(TEST_DIR, "test-file.txt");
245 | const result = service.isPathWithinAllowedDirs(insidePath);
246 |
247 | expect(result).toBe(true);
248 | });
249 |
250 | it("should return true for exact match with allowed directory", () => {
251 | const result = service.isPathWithinAllowedDirs(TEST_DIR);
252 |
253 | expect(result).toBe(true);
254 | });
255 |
256 | it("should return false for paths outside allowed directories", () => {
257 | const outsidePath = path.join(OUTSIDE_DIR, "test-file.txt");
258 | const result = service.isPathWithinAllowedDirs(outsidePath);
259 |
260 | expect(result).toBe(false);
261 | });
262 |
263 | it("should return false for paths with directory traversal", () => {
264 | const traversalPath = path.join(
265 | TEST_DIR,
266 | "..",
267 | "outside",
268 | "test-file.txt"
269 | );
270 | const result = service.isPathWithinAllowedDirs(traversalPath);
271 |
272 | expect(result).toBe(false);
273 | });
274 |
275 | it("should use custom allowed directories when provided", () => {
276 | const outsidePath = path.join(OUTSIDE_DIR, "test-file.txt");
277 |
278 | // Should be false with default allowed dirs
279 | expect(service.isPathWithinAllowedDirs(outsidePath)).toBe(false);
280 |
281 | // Should be true with custom allowed dirs
282 | expect(service.isPathWithinAllowedDirs(outsidePath, [OUTSIDE_DIR])).toBe(
283 | true
284 | );
285 | });
286 |
287 | it("should return false when no allowed directories exist", () => {
288 | const result = service.isPathWithinAllowedDirs(TEST_DIR, []);
289 |
290 | expect(result).toBe(false);
291 | });
292 | });
293 |
294 | describe("fullyResolvePath", () => {
295 | let service: FileSecurityService;
296 |
297 | beforeEach(() => {
298 | service = new FileSecurityService([TEST_DIR, OUTSIDE_DIR]);
299 | });
300 |
301 | it("should resolve a normal file path", async () => {
302 | const testPath = path.join(TEST_DIR, "test-file.txt");
303 | const resolvedPath = await service.fullyResolvePath(testPath);
304 |
305 | expect(resolvedPath).toBe(path.normalize(testPath));
306 | });
307 |
308 | it("should handle non-existent paths", async () => {
309 | const nonExistentPath = path.join(
310 | TEST_DIR,
311 | "non-existent",
312 | "test-file.txt"
313 | );
314 | const resolvedPath = await service.fullyResolvePath(nonExistentPath);
315 |
316 | expect(resolvedPath).toBe(path.normalize(nonExistentPath));
317 | });
318 |
319 | it("should resolve and validate a symlink to a file", async () => {
320 | // Create target file
321 | const targetPath = path.join(TEST_DIR, "target.txt");
322 | await fs.writeFile(targetPath, TEST_CONTENT, "utf8");
323 |
324 | // Create symlink
325 | const symlinkPath = path.join(TEST_DIR, "symlink.txt");
326 | await fs.symlink(targetPath, symlinkPath);
327 |
328 | // Resolve the symlink
329 | const resolvedPath = await service.fullyResolvePath(symlinkPath);
330 |
331 | // Should resolve to the target path
332 | expect(resolvedPath).toBe(path.normalize(targetPath));
333 | });
334 |
335 | it("should reject symlinks pointing outside allowed directories", async () => {
336 | // Create target file in outside (non-allowed) directory
337 | const targetPath = path.join(OUTSIDE_DIR, "target.txt");
338 | await fs.writeFile(targetPath, TEST_CONTENT, "utf8");
339 |
340 | // Create symlink in test (allowed) directory pointing to outside
341 | const symlinkPath = path.join(TEST_DIR, "bad-symlink.txt");
342 |
343 | // Setup service with only TEST_DIR allowed (not OUTSIDE_DIR)
344 | const restrictedService = new FileSecurityService([TEST_DIR]);
345 |
346 | await fs.symlink(targetPath, symlinkPath);
347 |
348 | // Try to resolve the symlink
349 | await expect(
350 | restrictedService.fullyResolvePath(symlinkPath)
351 | ).rejects.toThrow(ValidationError);
352 | await expect(
353 | restrictedService.fullyResolvePath(symlinkPath)
354 | ).rejects.toThrow(/Security error/);
355 | await expect(
356 | restrictedService.fullyResolvePath(symlinkPath)
357 | ).rejects.toThrow(/outside allowed directories/);
358 | });
359 |
360 | it("should detect and validate symlinked parent directories", async () => {
361 | // Create target directory in allowed location
362 | const targetDir = path.join(TEST_DIR, "target-dir");
363 | await fs.mkdir(targetDir, { recursive: true });
364 |
365 | // Create symlink to directory
366 | const symlinkDir = path.join(TEST_DIR, "symlink-dir");
367 | await fs.symlink(targetDir, symlinkDir);
368 |
369 | // Create a file path inside the symlinked directory
370 | const filePath = path.join(symlinkDir, "test-file.txt");
371 |
372 | // Resolve the path
373 | const resolvedPath = await service.fullyResolvePath(filePath);
374 |
375 | // Should resolve to actual path in target directory
376 | const expectedPath = path.join(targetDir, "test-file.txt");
377 | expect(resolvedPath).toBe(path.normalize(expectedPath));
378 | });
379 |
380 | it("should reject symlinked parent directories pointing outside allowed directories", async () => {
381 | // Create target directory in outside (not allowed) directory
382 | const targetDir = path.join(OUTSIDE_DIR, "target-dir");
383 | await fs.mkdir(targetDir, { recursive: true });
384 |
385 | // Create symlink in test directory pointing to outside directory
386 | const symlinkDir = path.join(TEST_DIR, "bad-symlink-dir");
387 | await fs.symlink(targetDir, symlinkDir);
388 |
389 | // Create a file path inside the symlinked directory
390 | const filePath = path.join(symlinkDir, "test-file.txt");
391 |
392 | // Setup service with only TEST_DIR allowed
393 | const restrictedService = new FileSecurityService([TEST_DIR]);
394 |
395 | // Try to resolve the path
396 | await expect(
397 | restrictedService.fullyResolvePath(filePath)
398 | ).rejects.toThrow(ValidationError);
399 | await expect(
400 | restrictedService.fullyResolvePath(filePath)
401 | ).rejects.toThrow(/Security error/);
402 | });
403 | });
404 |
405 | describe("secureWriteFile", () => {
406 | let service: FileSecurityService;
407 |
408 | beforeEach(() => {
409 | service = new FileSecurityService([TEST_DIR]);
410 | });
411 |
412 | it("should write file to an allowed directory", async () => {
413 | const filePath = path.join(TEST_DIR, "test-file.txt");
414 |
415 | await service.secureWriteFile(filePath, TEST_CONTENT);
416 |
417 | // Verify file was written
418 | const content = await fs.readFile(filePath, "utf8");
419 | expect(content).toBe(TEST_CONTENT);
420 | });
421 |
422 | it("should create directories if they don't exist", async () => {
423 | const nestedFilePath = path.join(
424 | TEST_DIR,
425 | "nested",
426 | "deep",
427 | "test-file.txt"
428 | );
429 |
430 | await service.secureWriteFile(nestedFilePath, TEST_CONTENT);
431 |
432 | // Verify directories were created and file exists
433 | const content = await fs.readFile(nestedFilePath, "utf8");
434 | expect(content).toBe(TEST_CONTENT);
435 | });
436 |
437 | it("should reject writing outside allowed directories", async () => {
438 | const outsidePath = path.join(OUTSIDE_DIR, "test-file.txt");
439 |
440 | await expect(
441 | service.secureWriteFile(outsidePath, TEST_CONTENT)
442 | ).rejects.toThrow(ValidationError);
443 | await expect(
444 | service.secureWriteFile(outsidePath, TEST_CONTENT)
445 | ).rejects.toThrow(/Access denied/);
446 |
447 | // Verify file was not created
448 | await expect(fs.access(outsidePath)).rejects.toThrow();
449 | });
450 |
451 | it("should reject overwriting existing files by default", async () => {
452 | const filePath = path.join(TEST_DIR, "existing-file.txt");
453 |
454 | // Create the file first
455 | await fs.writeFile(filePath, "Original content", "utf8");
456 |
457 | // Try to overwrite without setting overwrite flag
458 | await expect(
459 | service.secureWriteFile(filePath, TEST_CONTENT)
460 | ).rejects.toThrow(ValidationError);
461 | await expect(
462 | service.secureWriteFile(filePath, TEST_CONTENT)
463 | ).rejects.toThrow(/File already exists/);
464 |
465 | // Verify file wasn't changed
466 | const content = await fs.readFile(filePath, "utf8");
467 | expect(content).toBe("Original content");
468 | });
469 |
470 | it("should allow overwriting existing files with overwrite flag", async () => {
471 | const filePath = path.join(TEST_DIR, "existing-file.txt");
472 |
473 | // Create the file first
474 | await fs.writeFile(filePath, "Original content", "utf8");
475 |
476 | // Overwrite with overwrite flag
477 | await service.secureWriteFile(filePath, TEST_CONTENT, {
478 | overwrite: true,
479 | });
480 |
481 | // Verify file was overwritten
482 | const content = await fs.readFile(filePath, "utf8");
483 | expect(content).toBe(TEST_CONTENT);
484 | });
485 |
486 | it("should support custom allowed directories", async () => {
487 | // Path is outside the service's configured directories
488 | const customAllowedPath = path.join(OUTSIDE_DIR, "custom-allowed.txt");
489 |
490 | // Use explicit allowedDirs
491 | await service.secureWriteFile(customAllowedPath, TEST_CONTENT, {
492 | allowedDirs: [OUTSIDE_DIR],
493 | });
494 |
495 | // Verify file was written
496 | const content = await fs.readFile(customAllowedPath, "utf8");
497 | expect(content).toBe(TEST_CONTENT);
498 | });
499 | });
500 | });
501 |
```