@pagopa/dx-mcpserver 0.0.11 → 0.1.1

package/dist/server.js

@@ -0,0 +1,109 @@
+import { getLogger } from "@logtape/logtape";
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import crypto from "node:crypto";
+import * as http from "node:http";
+import { createBedrockRuntimeClient } from "./config/aws.js";
+import { handleAskEndpoint } from "./handlers/ask.js";
+import { handleSearchEndpoint } from "./handlers/search.js";
+import { createServer } from "./mcp/server.js";
+import { sessionStorage } from "./session.js";
+import { createToolDefinitions } from "./tools/registry.js";
+import { parseJsonBody, sendErrorResponse, setCorsHeaders, } from "./utils/http.js";
+export async function startHttpServer(config, enabledPrompts) {
+    const logger = getLogger(["mcpserver"]);
+    const awsLogger = getLogger(["mcpserver", "aws-config"]);
+    const kbRuntimeClient = createBedrockRuntimeClient(config.aws.region, awsLogger);
+    const toolDefinitions = createToolDefinitions({
+        aws: config.aws,
+        githubSearchOrg: config.github.searchOrg,
+        kbRuntimeClient,
+    });
+    const httpServer = http.createServer(async (req, res) => {
+        // Log incoming request for debugging
+        logger.debug("Incoming request", {
+            headers: req.headers,
+            method: req.method,
+            url: req.url,
+        });
+        // Configure CORS headers
+        setCorsHeaders(res);
+        // Handle OPTIONS for CORS preflight
+        if (req.method === "OPTIONS") {
+            res.writeHead(200);
+            res.end();
+            return;
+        }
+        // Handle /ask endpoint for Bedrock Knowledge Base queries
+        if (req.url?.includes("/ask") && req.method === "POST") {
+            await handleAskEndpoint(req, res, config, kbRuntimeClient);
+            return;
+        }
+        // Handle /search endpoint for documentation search
+        if (req.url?.includes("/search") && req.method === "POST") {
+            await handleSearchEndpoint(req, res, config, kbRuntimeClient);
+            return;
+        }
+        // Only allow POST and DELETE for MCP endpoints
+        if (req.method !== "POST" && req.method !== "DELETE") {
+            return sendErrorResponse(res, 405, "Method not allowed");
+        }
+        try {
+            // Parse request body
+            let jsonBody;
+            try {
+                jsonBody = await parseJsonBody(req);
+            }
+            catch {
+                return sendErrorResponse(res, 400, "Invalid JSON");
+            }
+            // Create session for AsyncLocalStorage context (stateless per-request)
+            // Extract request ID from AWS Lambda trace (undefined if not in AWS)
+            const traceHeader = req.headers["x-amzn-trace-id"];
+            const requestId = typeof traceHeader === "string"
+                ? traceHeader.split(";")[0].replace("Root=", "")
+                : undefined;
+            const session = {
+                id: crypto.randomUUID(),
+                requestId,
+            };
+            // Execute request in isolated session context
+            await sessionStorage.run(session, async () => {
+                // Create new server and transport for this request
+                // This ensures complete isolation between concurrent requests
+                const server = createServer({
+                    enabledPrompts,
+                    requestId,
+                    toolDefinitions,
+                });
+                const transport = new StreamableHTTPServerTransport({
+                    // This MCP server is stateless at the transport layer: each HTTP request
+                    // runs in its own AsyncLocalStorage-backed session context via
+                    // `sessionStorage.run(...)`. Because we do not multiplex logical sessions
+                    // over a shared connection, transport-level session IDs are unnecessary,
+                    // so `sessionIdGenerator` is explicitly set to `undefined`.
+                    sessionIdGenerator: undefined,
+                });
+                await server.connect(transport);
+                await transport.handleRequest(req, res, jsonBody);
+                // Clean up after response
+                res.on("close", () => {
+                    transport.close();
+                    server.close();
+                });
+            });
+        }
+        catch (error) {
+            logger.error("Error handling request", { error });
+            if (!res.headersSent) {
+                sendErrorResponse(res, 500, "Internal server error");
+            }
+        }
+    });
+    await new Promise((resolve) => {
+        httpServer.listen(config.port, () => {
+            logger.info(`MCP Server started on port ${config.port}`);
+            resolve();
+        });
+    });
+    return httpServer;
+}

package/dist/services/__tests__/bedrock-retrieve-and-generate.test.js

@@ -0,0 +1,116 @@
+import { getLogger } from "@logtape/logtape";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { retrieveAndGenerate } from "../bedrock-retrieve-and-generate.js";
+// Helper to create mock Bedrock client
+function createMockBedrockClient(sendMock) {
+    return {
+        config: {
+            apiVersion: "2023-11-20",
+            region: async () => "eu-central-1",
+            requestHandler: { handle: vi.fn() },
+        },
+        destroy: vi.fn(),
+        middlewareStack: {},
+        send: sendMock,
+    };
+}
+describe("retrieveAndGenerate", () => {
+    let loggerSpy;
+    beforeEach(() => {
+        vi.clearAllMocks();
+        const logger = getLogger(["mcpserver", "bedrock", "retrieve-and-generate"]);
+        loggerSpy = {
+            debug: vi.spyOn(logger, "debug"),
+            error: vi.spyOn(logger, "error"),
+            info: vi.spyOn(logger, "info"),
+        };
+    });
+    it("should successfully retrieve and generate with valid parameters", async () => {
+        const mockResponse = {
+            citations: [
+                {
+                    retrievedReferences: [
+                        {
+                            content: { text: "Reference content" },
+                            location: {
+                                s3Location: { uri: "s3://bucket/doc.md" },
+                                type: "S3",
+                            },
+                        },
+                    ],
+                },
+            ],
+            output: {
+                text: "This is the AI-generated answer based on the knowledge base.",
+            },
+            sessionId: "test-session-123",
+        };
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue(mockResponse));
+        const result = await retrieveAndGenerate("kb-test-id", "arn:aws:bedrock:eu-central-1::foundation-model/amazon.nova-2-lite-v1:0", "How do I setup Terraform?", mockClient, 5);
+        expect(result).toEqual(mockResponse);
+        expect(result.output?.text).toBe("This is the AI-generated answer based on the knowledge base.");
+        expect(result.sessionId).toBe("test-session-123");
+        expect(loggerSpy.debug).toHaveBeenCalledWith("Calling RetrieveAndGenerate", {
+            knowledgeBaseId: "kb-test-id",
+            modelArn: "arn:aws:bedrock:eu-central-1::foundation-model/amazon.nova-2-lite-v1:0",
+            numberOfResults: 5,
+            query: "How do I setup Terraform?",
+        });
+        expect(loggerSpy.info).toHaveBeenCalledWith("RetrieveAndGenerate successful", {
+            sessionId: "test-session-123",
+        });
+    });
+    it("should use default numberOfResults when not provided", async () => {
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue({
+            output: { text: "Answer" },
+            sessionId: "session-456",
+        }));
+        await retrieveAndGenerate("kb-id", "model-arn", "test query", mockClient);
+        expect(loggerSpy.debug).toHaveBeenCalledWith("Calling RetrieveAndGenerate", expect.objectContaining({
+            numberOfResults: 5,
+        }));
+    });
+    it("should log error details and rethrow on failure", async () => {
+        const mockError = {
+            $metadata: {
+                httpStatusCode: 403,
+            },
+            __type: "AccessDeniedException",
+            message: "User is not authorized to perform: bedrock:RetrieveAndGenerate",
+            name: "AccessDeniedException",
+        };
+        const mockClient = createMockBedrockClient(vi.fn().mockRejectedValue(mockError));
+        await expect(retrieveAndGenerate("kb-id", "model-arn", "test query", mockClient)).rejects.toThrow();
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("RetrieveAndGenerate failed"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("AccessDeniedException"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("kb-id"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("model-arn"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("403"));
+    });
+    it("should handle errors without $metadata or __type", async () => {
+        const mockError = new Error("Network timeout");
+        const mockClient = createMockBedrockClient(vi.fn().mockRejectedValue(mockError));
+        await expect(retrieveAndGenerate("kb-id", "model-arn", "test query", mockClient)).rejects.toThrow("Network timeout");
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("RetrieveAndGenerate failed"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("Error"));
+        expect(loggerSpy.error).toHaveBeenCalledWith(expect.stringContaining("Network timeout"));
+    });
+    it("should send correct RetrieveAndGenerateCommand configuration", async () => {
+        const sendSpy = vi.fn().mockResolvedValue({
+            output: { text: "Test answer" },
+            sessionId: "test-session",
+        });
+        const mockClient = createMockBedrockClient(sendSpy);
+        await retrieveAndGenerate("kb-test-id", "arn:aws:bedrock:eu-central-1::foundation-model/test-model", "What is the capital of France?", mockClient, 10);
+        const commandArg = sendSpy.mock.calls[0][0];
+        expect(commandArg.input.input.text).toBe("What is the capital of France?");
+        expect(commandArg.input.retrieveAndGenerateConfiguration.type).toBe("KNOWLEDGE_BASE");
+        expect(commandArg.input.retrieveAndGenerateConfiguration
+            .knowledgeBaseConfiguration.knowledgeBaseId).toBe("kb-test-id");
+        expect(commandArg.input.retrieveAndGenerateConfiguration
+            .knowledgeBaseConfiguration.modelArn).toBe("arn:aws:bedrock:eu-central-1::foundation-model/test-model");
+        expect(commandArg.input.retrieveAndGenerateConfiguration
+            .knowledgeBaseConfiguration.retrievalConfiguration
+            .vectorSearchConfiguration.numberOfResults).toBe(10);
+    });
+});

package/dist/services/__tests__/bedrock.test.js

@@ -1,6 +1,19 @@
 import { getLogger } from "@logtape/logtape";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { queryKnowledgeBase } from "../bedrock.js";
+import { queryKnowledgeBase, queryKnowledgeBaseStructured, } from "../bedrock.js";
+// Helper to create mock Bedrock client
+function createMockBedrockClient(sendMock) {
+    return {
+        config: {
+            apiVersion: "2023-11-20",
+            region: async () => "eu-central-1",
+            requestHandler: { handle: vi.fn() },
+        },
+        destroy: vi.fn(),
+        middlewareStack: {},
+        send: sendMock,
+    };
+}
 describe("queryKnowledgeBase", () => {
     let loggerSpy;
     beforeEach(() => {
@@ -47,3 +60,149 @@ describe("queryKnowledgeBase", () => {
         expect(loggerSpy.warn).toHaveBeenCalledWith("Reranking is not supported in region unsupported-region");
     });
 });
+describe("queryKnowledgeBaseStructured - Basic Functionality", () => {
+    let loggerSpy;
+    beforeEach(() => {
+        vi.clearAllMocks();
+        const logger = getLogger(["mcpserver", "bedrock"]);
+        loggerSpy = {
+            warn: vi.spyOn(logger, "warn"),
+        };
+    });
+    it("should return structured QueryKnowledgeBasesOutput[] instead of a string", async () => {
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue({
+            retrievalResults: [
+                {
+                    content: { text: "Documentation content", type: "TEXT" },
+                    location: {
+                        s3Location: { uri: "s3://bucket/docs/guide.md" },
+                        type: "S3",
+                    },
+                    score: 0.95,
+                },
+            ],
+        }));
+        const result = await queryKnowledgeBaseStructured("kbId", "query", mockClient, 5, false);
+        expect(Array.isArray(result)).toBe(true);
+        expect(result.length).toBe(1);
+        expect(typeof result[0]).toBe("object");
+    });
+    it("should return result objects with content, location, and score properties", async () => {
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue({
+            retrievalResults: [
+                {
+                    content: { text: "Test content", type: "TEXT" },
+                    location: {
+                        s3Location: { uri: "s3://bucket/path/file.md" },
+                        type: "S3",
+                    },
+                    score: 0.87,
+                },
+            ],
+        }));
+        const result = await queryKnowledgeBaseStructured("kbId", "test query", mockClient);
+        expect(result[0]).toHaveProperty("content");
+        expect(result[0]).toHaveProperty("location");
+        expect(result[0]).toHaveProperty("score");
+        expect(result[0].content).toBe("Test content");
+        expect(result[0].score).toBe(0.87);
+    });
+    it("should resolve S3 locations to website URLs via resolveToWebsiteUrl", async () => {
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue({
+            retrievalResults: [
+                {
+                    content: { text: "Content", type: "TEXT" },
+                    location: {
+                        s3Location: {
+                            uri: "s3://pagopa-dx-documentation/docs/azure/index.md",
+                        },
+                        type: "S3",
+                    },
+                    score: 0.9,
+                },
+            ],
+        }));
+        const result = await queryKnowledgeBaseStructured("kbId", "query", mockClient);
+        // resolveToWebsiteUrl should transform S3 URI to website URL
+        expect(result[0].location).toBeDefined();
+        if (result[0].location?.webLocation?.url) {
+            expect(result[0].location.webLocation.url).toContain("https://");
+            // URL should have /index removed
+            expect(result[0].location.webLocation.url).not.toContain("/index");
+        }
+    });
+    it("should skip image content with warning logs", async () => {
+        const mockClient = createMockBedrockClient(vi.fn().mockResolvedValue({
+            retrievalResults: [
+                {
+                    content: { text: "Valid text content", type: "TEXT" },
+                    location: {
+                        s3Location: { uri: "s3://bucket/file.md" },
+                        type: "S3",
+                    },
+                    score: 0.9,
+                },
+                {
+                    content: { type: "IMAGE" },
+                    location: {
+                        s3Location: { uri: "s3://bucket/image.png" },
+                        type: "S3",
+                    },
+                    score: 0.85,
+                },
+            ],
+        }));
+        const result = await queryKnowledgeBaseStructured("kbId", "query", mockClient);
+        expect(result.length).toBe(1);
+        expect(result[0].content).toBe("Valid text content");
+        expect(loggerSpy.warn).toHaveBeenCalledWith("Images are not supported at this time. Skipping...");
+    });
+});
+describe("queryKnowledgeBaseStructured - Reranking", () => {
+    let loggerSpy;
+    beforeEach(() => {
+        vi.clearAllMocks();
+        const logger = getLogger(["mcpserver", "bedrock"]);
+        loggerSpy = {
+            warn: vi.spyOn(logger, "warn"),
+        };
+    });
+    it("should apply reranking configuration when enabled in supported region", async () => {
+        const sendSpy = vi.fn().mockResolvedValue({
+            retrievalResults: [
+                {
+                    content: { text: "Reranked content", type: "TEXT" },
+                    location: { s3Location: { uri: "s3://bucket/doc.md" }, type: "S3" },
+                    score: 0.98,
+                },
+            ],
+        });
+        const mockClient = createMockBedrockClient(sendSpy);
+        // Override region for this test
+        mockClient.config.region = async () => "us-east-1"; // Supported region
+        await queryKnowledgeBaseStructured("kbId", "query", mockClient, 5, true, // Enable reranking
+        "AMAZON");
+        // Verify that send was called with reranking configuration
+        const commandArg = sendSpy.mock.calls[0][0];
+        expect(commandArg.input.retrievalConfiguration).toBeDefined();
+        expect(commandArg.input.retrievalConfiguration.vectorSearchConfiguration
+            .rerankingConfiguration).toBeDefined();
+        expect(commandArg.input.retrievalConfiguration.vectorSearchConfiguration
+            .rerankingConfiguration.bedrockRerankingConfiguration.modelConfiguration
+            .modelArn).toContain("amazon.rerank-v1:0");
+    });
+    it("should disable reranking in unsupported region with warning", async () => {
+        const sendSpy = vi.fn().mockResolvedValue({
+            retrievalResults: [],
+        });
+        const mockClient = createMockBedrockClient(sendSpy);
+        // Override region for this test
+        mockClient.config.region = async () => "ap-southeast-1"; // Unsupported region
+        await queryKnowledgeBaseStructured("kbId", "query", mockClient, 5, true);
+        expect(loggerSpy.warn).toHaveBeenCalledWith("Reranking is not supported in region ap-southeast-1");
+        // Verify that send was called WITHOUT reranking configuration
+        const commandArg = sendSpy.mock.calls[0][0];
+        expect(commandArg.input.retrievalConfiguration.vectorSearchConfiguration
+            .rerankingConfiguration).toBeUndefined();
+    });
+});

package/dist/services/bedrock-retrieve-and-generate.js

@@ -0,0 +1,55 @@
+import { RetrieveAndGenerateCommand, } from "@aws-sdk/client-bedrock-agent-runtime";
+import { getLogger } from "@logtape/logtape";
+/**
+ * Calls Bedrock RetrieveAndGenerate API to get an AI-generated answer
+ * based on documents from a knowledge base.
+ *
+ * @param knowledgeBaseId The ID of the knowledge base to query
+ * @param modelArn The ARN of the Bedrock model to use for generation
+ * @param query The user's natural language query
+ * @param kbAgentClient The Bedrock Agent Runtime client
+ * @param numberOfResults The maximum number of documents to retrieve (default: 5)
+ * @returns The complete RetrieveAndGenerate response including answer and citations
+ */
+export async function retrieveAndGenerate(knowledgeBaseId, modelArn, query, kbAgentClient, numberOfResults = 5) {
+    const logger = getLogger(["mcpserver", "bedrock", "retrieve-and-generate"]);
+    try {
+        const command = new RetrieveAndGenerateCommand({
+            input: {
+                text: query,
+            },
+            retrieveAndGenerateConfiguration: {
+                knowledgeBaseConfiguration: {
+                    knowledgeBaseId,
+                    modelArn,
+                    retrievalConfiguration: {
+                        vectorSearchConfiguration: {
+                            numberOfResults,
+                        },
+                    },
+                },
+                type: "KNOWLEDGE_BASE",
+            },
+        });
+        logger.debug("Calling RetrieveAndGenerate", {
+            knowledgeBaseId,
+            modelArn,
+            numberOfResults,
+            query,
+        });
+        const response = await kbAgentClient.send(command);
+        logger.info("RetrieveAndGenerate successful", {
+            sessionId: response.sessionId,
+        });
+        return response;
+    }
+    catch (error) {
+        const errorMessage = error instanceof Error ? error.message : String(error);
+        const errorName = error instanceof Error ? error.name : "Unknown";
+        const errorCode = error
+            .$metadata?.httpStatusCode;
+        const errorType = error.__type;
+        logger.error(`RetrieveAndGenerate failed: ${errorName} - ${errorMessage}. KB: ${knowledgeBaseId}, Model: ${modelArn}, HTTP Status: ${errorCode}, Type: ${errorType}`);
+        throw error;
+    }
+}

package/dist/services/bedrock.js

@@ -65,6 +65,60 @@ export async function queryKnowledgeBase(knowledgeBaseId, query, kbAgentClient,
     }
     return serializeResults(documents);
 }
+/**
+ * Returns structured results from knowledge base query without serialization.
+ * Same as queryKnowledgeBase but returns structured data instead of string.
+ */
+export async function queryKnowledgeBaseStructured(knowledgeBaseId, query, kbAgentClient, numberOfResults = 5, reranking = false, rerankingModelName = "AMAZON") {
+    const logger = getLogger(["mcpserver", "bedrock"]);
+    const clientRegion = await kbAgentClient.config.region();
+    let rerankingEnabled = reranking;
+    if (reranking && !rerankingSupportedRegions.includes(clientRegion)) {
+        logger.warn(`Reranking is not supported in region ${clientRegion}`);
+        rerankingEnabled = false;
+    }
+    const retrieveRequest = {
+        vectorSearchConfiguration: {
+            numberOfResults,
+        },
+    };
+    if (rerankingEnabled && retrieveRequest.vectorSearchConfiguration) {
+        const modelNameMapping = {
+            AMAZON: "amazon.rerank-v1:0",
+            COHERE: "cohere.rerank-v3-5:0",
+        };
+        retrieveRequest.vectorSearchConfiguration.rerankingConfiguration = {
+            bedrockRerankingConfiguration: {
+                modelConfiguration: {
+                    modelArn: `arn:aws:bedrock:${clientRegion}::foundation-model/${modelNameMapping[rerankingModelName]}`,
+                },
+            },
+            type: "BEDROCK_RERANKING_MODEL",
+        };
+    }
+    const command = new RetrieveCommand({
+        knowledgeBaseId,
+        retrievalConfiguration: retrieveRequest,
+        retrievalQuery: { text: query },
+    });
+    const response = await kbAgentClient.send(command);
+    const results = response.retrievalResults || [];
+    const documents = [];
+    for (const result of results) {
+        if (result.content?.type === "IMAGE") {
+            logger.warn("Images are not supported at this time. Skipping...");
+            continue;
+        }
+        else if (result.content?.text) {
+            documents.push({
+                content: result.content.text,
+                location: resolveToWebsiteUrl(result.location),
+                score: result.score || -1.0,
+            });
+        }
+    }
+    return documents;
+}
 /**
  * Resolves an S3 location from a knowledge base result to a public website URL.
  * This method converts S3 URIs to publicly accessible URLs based on specific rules.
@@ -102,6 +156,8 @@ export function resolveToWebsiteUrl(location) {
             else {
                 url = `https://dx.pagopa.it/docs/${key.replace(/\.md$/, "")}`;
             }
+            // Post-process: remove trailing /index from URLs
+            url = url.replace(/\/index$/, "/");
             // Return a RetrievalResultLocation object of type WEB with the computed URL
             return {
                 type: "WEB",

package/dist/session.js

@@ -0,0 +1,12 @@
+/**
+ * Session storage for stateless MCP server operations
+ *
+ * Uses AsyncLocalStorage to provide request-scoped context without global state.
+ * This enables the server to run in stateless environments (e.g., AWS Lambda)
+ * while still passing session data to tool handlers.
+ *
+ * Each HTTP request creates its own isolated session context that is automatically
+ * cleaned up after the request completes.
+ */
+import { AsyncLocalStorage } from "node:async_hooks";
+export const sessionStorage = new AsyncLocalStorage();

package/dist/tools/__tests__/QueryValidation.test.js

@@ -0,0 +1,83 @@
+import { describe, expect, it } from "vitest";
+import { QueryPagoPADXDocumentationInputSchema } from "../query-pagopa-dx-documentation.js";
+// Test the Zod schema validation used in the tools
+describe("Query Validation", () => {
+    // Use the actual schema from QueryPagoPADXDocumentationTool
+    const queryDocSchema = QueryPagoPADXDocumentationInputSchema;
+    describe("QueryPagoPADXDocumentationTool validation", () => {
+        it("should accept valid queries", () => {
+            const validQueries = [
+                { query: "How do I set up the project?" },
+                { query: "abc" }, // minimum 3 characters
+                { query: "What is the DX CLI?" },
+                { query: "Terraform Azure provider setup" },
+            ];
+            validQueries.forEach((input) => {
+                const result = queryDocSchema.safeParse(input);
+                expect(result.success).toBe(true);
+            });
+        });
+        it("should reject queries shorter than 3 characters", () => {
+            const shortQueries = [{ query: "" }, { query: "a" }, { query: "ab" }];
+            shortQueries.forEach((input) => {
+                const result = queryDocSchema.safeParse(input);
+                expect(result.success).toBe(false);
+                if (!result.success) {
+                    expect(result.error.errors[0].message).toBe("Query must be at least 3 characters");
+                }
+            });
+        });
+        it("should reject queries exceeding 500 characters", () => {
+            const longQuery = { query: "a".repeat(501) };
+            const result = queryDocSchema.safeParse(longQuery);
+            expect(result.success).toBe(false);
+            if (!result.success) {
+                expect(result.error.errors[0].message).toBe("Query must not exceed 500 characters");
+            }
+        });
+        it("should reject missing query field", () => {
+            const missingQuery = {};
+            const result = queryDocSchema.safeParse(missingQuery);
+            expect(result.success).toBe(false);
+        });
+        it("should reject non-string query values", () => {
+            const invalidTypes = [
+                { query: 123 },
+                { query: null },
+                { query: undefined },
+                { query: [] },
+                { query: {} },
+            ];
+            invalidTypes.forEach((input) => {
+                const result = queryDocSchema.safeParse(input);
+                expect(result.success).toBe(false);
+            });
+        });
+    });
+    describe("Edge cases", () => {
+        it("should handle queries with special characters", () => {
+            const specialCharQueries = [
+                { query: "api/management@v1" },
+                { query: "terraform-module (azure-app-service)" },
+                { query: "query with 'quotes' and \"double quotes\"" },
+            ];
+            specialCharQueries.forEach((input) => {
+                const result = queryDocSchema.safeParse(input);
+                expect(result.success).toBe(true);
+            });
+        });
+        it("should handle very long query strings within limit", () => {
+            const longQuery = { query: "a".repeat(500) }; // exactly at limit
+            const result = queryDocSchema.safeParse(longQuery);
+            expect(result.success).toBe(true);
+        });
+        it("should accept whitespace-only strings if they meet min length", () => {
+            // Note: The Zod schema doesn't call .trim(), so whitespace-only strings will be accepted
+            // if they meet the minimum length requirement (3 characters for queryDoc)
+            const whitespaceOnly = { query: "   " }; // 3 spaces = 3 characters
+            const result = queryDocSchema.safeParse(whitespaceOnly);
+            // Whitespace-only string has length >= 3, so it passes
+            expect(result.success).toBe(true);
+        });
+    });
+});

package/dist/tools/__tests__/query-pago-pa-dx-documentation.test.js

@@ -0,0 +1,47 @@
+import { describe, expect, it, vi } from "vitest";
+vi.mock("../../services/bedrock", () => ({
+    queryKnowledgeBase: vi.fn().mockResolvedValue("mocked result"),
+}));
+import { BedrockAgentRuntimeClient } from "@aws-sdk/client-bedrock-agent-runtime";
+import { createQueryPagoPADXDocumentationTool } from "../query-pagopa-dx-documentation.js";
+const QueryPagoPADXDocumentationTool = createQueryPagoPADXDocumentationTool({
+    kbRuntimeClient: new BedrockAgentRuntimeClient({ region: "eu-central-1" }),
+    knowledgeBaseId: "mockKbId",
+    rerankingEnabled: false,
+});
+describe("QueryPagoPADXDocumentationTool", () => {
+    it("should return results from the knowledge base", async () => {
+        const args = { format: "markdown", query: "test query" };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toBe("mocked result");
+    });
+    it("should use default format if not provided", async () => {
+        const args = { query: "test query" };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toBe("mocked result");
+    });
+    it("should reject unknown parameters due to strict schema", async () => {
+        const args = { query: "test query", unknownParam: "value" };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toContain("Error: Invalid input");
+        expect(result).toContain("Unrecognized key");
+    });
+    it("should reject queries shorter than 3 characters", async () => {
+        const args = { query: "ab" };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toContain("Error: Invalid input");
+        expect(result).toContain("Query must be at least 3 characters");
+    });
+    it("should reject empty queries", async () => {
+        const args = { query: "" };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toContain("Error: Invalid input");
+        expect(result).toContain("Query must be at least 3 characters");
+    });
+    it("should reject queries exceeding 500 characters", async () => {
+        const args = { query: "a".repeat(501) };
+        const result = await QueryPagoPADXDocumentationTool.execute(args);
+        expect(result).toContain("Error: Invalid input");
+        expect(result).toContain("Query must not exceed 500 characters");
+    });
+});