diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 75892d6..239c4c2 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -88,7 +88,7 @@ export abstract class BaseMCPServer extends Server { * Create a standardized error response */ protected createErrorResponse(message: string, statusMessage?: string): ErrorResponse { - const span = getActiveSpan(); + const span = this.initSpanWithCommonAttributes(); span?.setStatus({ code: SpanStatusCode.ERROR, message: statusMessage || message }); return { isError: true, @@ -100,7 +100,7 @@ export abstract class BaseMCPServer extends Server { * Create a standardized success response with a single message */ protected createSuccessResponse(message: string, statusMessage?: string): SuccessResponse { - const span = getActiveSpan(); + const span = this.initSpanWithCommonAttributes(); span?.setStatus({ code: SpanStatusCode.OK, message: statusMessage || message }); return { content: [{ type: "text", text: message }], @@ -111,7 +111,7 @@ export abstract class BaseMCPServer extends Server { * Create a standardized success response with multiple content items */ protected createMultiContentSuccessResponse(content: ContentItem[], statusMessage: string): SuccessResponse { - const span = getActiveSpan(); + const span = this.initSpanWithCommonAttributes(); span?.setStatus({ code: SpanStatusCode.OK, message: statusMessage }); return { content, @@ -122,7 +122,7 @@ export abstract class BaseMCPServer extends Server { * Create a standardized success response with an array of text items */ protected createArraySuccessResponse(texts: string[], statusMessage: string): SuccessResponse { - const span = getActiveSpan(); + const span = this.initSpanWithCommonAttributes(); span?.setStatus({ code: SpanStatusCode.OK, message: statusMessage }); return { content: texts.map(text => ({ type: "text", text })), @@ -130,7 +130,7 @@ export abstract class BaseMCPServer extends Server { } protected createStructuredContentSuccessResponse(structuredContent: T, statusMessage: string): SuccessResponse { - const span = getActiveSpan(); + const span = this.initSpanWithCommonAttributes(); span?.setStatus({ code: SpanStatusCode.OK, message: statusMessage }); return { content: [{ diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index b4e6335..85d267a 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -11,7 +11,8 @@ import type { DataSource } from "../thoughtspot/thoughtspot-service"; import { TrackEvent } from "../metrics"; -import { WithSpan } from "../metrics/tracing/tracing-utils"; +import { getActiveSpan, WithSpan } from "../metrics/tracing/tracing-utils"; +import { SpanStatusCode } from "@opentelemetry/api"; import { BaseMCPServer, type Context, type ToolResponse } from "./mcp-server-base"; @@ -250,7 +251,20 @@ export class MCPServer extends BaseMCPServer { ); if (relevantQuestions.error) { - return this.createErrorResponse(relevantQuestions.error.message, `Error getting relevant questions ${relevantQuestions.error.message}`); + console.error("Error getting relevant questions: ", relevantQuestions.error); + + const structuredContent = { questions: [{ question: query, datasourceId: sourceIds?.[0] ?? '' }] }; + const span = this.initSpanWithCommonAttributes(); + span?.setStatus({ code: SpanStatusCode.ERROR, message: "Relevant questions failed, sending back the query as it is" }); + span?.setAttribute("datasource_ids", sourceIds?.join(",") ?? ""); + span?.setAttribute("error", relevantQuestions.error.message); + return { + content: [{ + type: "text", + text: JSON.stringify(structuredContent), + }], + structuredContent, + }; } if (relevantQuestions.questions.length === 0) { diff --git a/test/servers/mcp-server-base.spec.ts b/test/servers/mcp-server-base.spec.ts new file mode 100644 index 0000000..4bbd9bf --- /dev/null +++ b/test/servers/mcp-server-base.spec.ts @@ -0,0 +1,270 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { MCPServer } from "../../src/servers/mcp-server"; +import * as thoughtspotClient from "../../src/thoughtspot/thoughtspot-client"; +import { MixpanelTracker } from "../../src/metrics/mixpanel/mixpanel"; + +// Mock the MixpanelTracker +vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ + MixpanelTracker: vi.fn().mockImplementation(() => ({ + track: vi.fn(), + })), +})); + +// Test subclass to expose protected methods +class TestMCPServer extends MCPServer { + public testCreateMultiContentSuccessResponse(content: any[], statusMessage: string) { + return this.createMultiContentSuccessResponse(content, statusMessage); + } + + public testCreateArraySuccessResponse(texts: string[], statusMessage: string) { + return this.createArraySuccessResponse(texts, statusMessage); + } + + public testCreateErrorResponse(message: string, statusMessage?: string) { + return this.createErrorResponse(message, statusMessage); + } + + public testCreateSuccessResponse(message: string, statusMessage?: string) { + return this.createSuccessResponse(message, statusMessage); + } + + public testCreateStructuredContentSuccessResponse(structuredContent: T, statusMessage: string) { + return this.createStructuredContentSuccessResponse(structuredContent, statusMessage); + } + + public testIsDatasourceDiscoveryAvailable() { + return this.isDatasourceDiscoveryAvailable(); + } +} + +describe("MCP Server Base", () => { + let server: TestMCPServer; + let mockProps: any; + + beforeEach(() => { + vi.clearAllMocks(); + + // Mock getThoughtSpotClient + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "10.13.0.cl-110", + userGUID: "test-user-123", + configInfo: { + mixpanelConfig: { + devSdkKey: "test-dev-token", + prodSdkKey: "test-prod-token", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + enableSpotterDataSourceDiscovery: true, + }), + searchMetadata: vi.fn().mockResolvedValue([]), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + mockProps = { + instanceUrl: "https://test.thoughtspot.cloud", + accessToken: "test-token", + clientName: { + clientId: "test-client-id", + clientName: "test-client", + registrationDate: Date.now(), + }, + }; + + server = new TestMCPServer({ props: mockProps }); + }); + + describe("Response Helper Methods", () => { + beforeEach(async () => { + await server.init(); + }); + + it("should create multi-content success response", () => { + const content = [ + { type: "text" as const, text: "First message" }, + { type: "text" as const, text: "Second message" }, + { type: "text" as const, text: "Third message" }, + ]; + + const result = server.testCreateMultiContentSuccessResponse(content, "Multiple messages"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(3); + expect(result.content[0].text).toBe("First message"); + expect(result.content[1].text).toBe("Second message"); + expect(result.content[2].text).toBe("Third message"); + }); + + it("should create multi-content success response with empty array", () => { + const result = server.testCreateMultiContentSuccessResponse([], "No messages"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(0); + }); + + it("should create array success response from string array", () => { + const texts = ["Item 1", "Item 2", "Item 3"]; + + const result = server.testCreateArraySuccessResponse(texts, "Array response"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(3); + expect(result.content[0]).toEqual({ type: "text", text: "Item 1" }); + expect(result.content[1]).toEqual({ type: "text", text: "Item 2" }); + expect(result.content[2]).toEqual({ type: "text", text: "Item 3" }); + }); + + it("should create array success response with empty array", () => { + const result = server.testCreateArraySuccessResponse([], "Empty array"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(0); + }); + + it("should create array success response with single item", () => { + const texts = ["Single item"]; + + const result = server.testCreateArraySuccessResponse(texts, "Single item response"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(1); + expect(result.content[0]).toEqual({ type: "text", text: "Single item" }); + }); + + it("should create error response with message", () => { + const result = server.testCreateErrorResponse("Something went wrong"); + + expect(result.isError).toBe(true); + expect(result.content).toHaveLength(1); + expect(result.content[0].text).toBe("ERROR: Something went wrong"); + }); + + it("should create error response with custom status message", () => { + const result = server.testCreateErrorResponse("Error occurred", "Custom status"); + + expect(result.isError).toBe(true); + expect(result.content).toHaveLength(1); + expect(result.content[0].text).toBe("ERROR: Error occurred"); + }); + + it("should create success response with message", () => { + const result = server.testCreateSuccessResponse("Operation successful"); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(1); + expect(result.content[0].text).toBe("Operation successful"); + }); + + it("should create structured content success response", () => { + const structuredContent = { + items: ["a", "b", "c"], + count: 3, + }; + + const result = server.testCreateStructuredContentSuccessResponse( + structuredContent, + "Structured response" + ); + + expect(result.isError).toBeUndefined(); + expect(result.content).toHaveLength(1); + expect(result.content[0].text).toBe(JSON.stringify(structuredContent)); + expect(result.structuredContent).toEqual(structuredContent); + }); + }); + + describe("Datasource Discovery Check", () => { + it("should return true when enableSpotterDataSourceDiscovery is enabled", async () => { + await server.init(); + const result = server.testIsDatasourceDiscoveryAvailable(); + expect(result).toBe(true); + }); + + it("should return false when enableSpotterDataSourceDiscovery is disabled", async () => { + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "10.13.0.cl-110", + userGUID: "test-user-123", + configInfo: { + mixpanelConfig: { + devSdkKey: "test-dev-token", + prodSdkKey: "test-prod-token", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + enableSpotterDataSourceDiscovery: false, + }), + searchMetadata: vi.fn().mockResolvedValue([]), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + const testServer = new TestMCPServer({ props: mockProps }); + await testServer.init(); + const result = testServer.testIsDatasourceDiscoveryAvailable(); + expect(result).toBe(false); + }); + + it("should return false when enableSpotterDataSourceDiscovery is undefined", async () => { + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "10.13.0.cl-110", + userGUID: "test-user-123", + configInfo: { + mixpanelConfig: { + devSdkKey: "test-dev-token", + prodSdkKey: "test-prod-token", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + enableSpotterDataSourceDiscovery: undefined, + }), + searchMetadata: vi.fn().mockResolvedValue([]), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + const testServer = new TestMCPServer({ props: mockProps }); + await testServer.init(); + const result = testServer.testIsDatasourceDiscoveryAvailable(); + expect(result).toBe(false); + }); + }); + + describe("Server Initialization", () => { + it("should initialize with custom server name and version", () => { + const customServer = new TestMCPServer( + { props: mockProps }, + "CustomServer", + "2.0.0" + ); + + expect(customServer).toBeDefined(); + }); + + it("should initialize with default server name and version", () => { + expect(server).toBeDefined(); + }); + }); +}); + diff --git a/test/servers/mcp-server.spec.ts b/test/servers/mcp-server.spec.ts index 7db3aa1..85c0924 100644 --- a/test/servers/mcp-server.spec.ts +++ b/test/servers/mcp-server.spec.ts @@ -322,8 +322,98 @@ describe("MCP Server", () => { datasourceIds: ["ds-123"], }); - expect(result.isError).toBe(true); - expect((result.content as any[])[0].text).toBe("ERROR: Service unavailable"); + // When error occurs, the code returns the query as fallback (graceful degradation) + expect(result.isError).toBeUndefined(); + const resultText = JSON.parse((result.content as any[])[0].text); + expect(resultText.questions).toBeInstanceOf(Array); + expect(resultText.questions[0].question).toBe("Show me revenue data"); + expect(resultText.questions[0].datasourceId).toBe("ds-123"); + }); + + it("should handle error from service with multiple datasource IDs", async () => { + // Mock client to return error + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "1.0.0", + userGUID: "test-user-123", + configInfo: { + mixpanelConfig: { + devSdkKey: "test-dev-token", + prodSdkKey: "test-prod-token", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + enableSpotterDataSourceDiscovery: true, + }), + queryGetDecomposedQuery: vi.fn().mockRejectedValue(new Error("Service unavailable")), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("getRelevantQuestions", { + query: "What are the sales?", + datasourceIds: ["ds-123", "ds-456", "ds-789"], + }); + + // Should use the first datasourceId for the fallback + expect(result.isError).toBeUndefined(); + const resultText = JSON.parse((result.content as any[])[0].text); + expect(resultText.questions).toBeInstanceOf(Array); + expect(resultText.questions[0].question).toBe("What are the sales?"); + expect(resultText.questions[0].datasourceId).toBe("ds-123"); + // Verify structuredContent is also set + expect((result.structuredContent as any).questions).toEqual(resultText.questions); + }); + + it("should handle error from service with empty datasourceIds array", async () => { + // Mock client to return error + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockResolvedValue({ + clusterId: "test-cluster-123", + clusterName: "test-cluster", + releaseVersion: "1.0.0", + userGUID: "test-user-123", + configInfo: { + mixpanelConfig: { + devSdkKey: "test-dev-token", + prodSdkKey: "test-prod-token", + production: false, + }, + selfClusterName: "test-cluster", + selfClusterId: "test-cluster-123", + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + enableSpotterDataSourceDiscovery: true, + }), + queryGetDecomposedQuery: vi.fn().mockRejectedValue(new Error("Network error")), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("getRelevantQuestions", { + query: "Show me data", + datasourceIds: [], + }); + + // Should handle empty array gracefully (fallback to empty string for datasourceId) + expect(result.isError).toBeUndefined(); + const resultText = JSON.parse((result.content as any[])[0].text); + expect(resultText.questions).toBeInstanceOf(Array); + expect(resultText.questions[0].question).toBe("Show me data"); + expect(resultText.questions[0].datasourceId).toBe(""); }); it("should handle empty questions response", async () => { @@ -642,7 +732,7 @@ describe("MCP Server", () => { })).rejects.toThrow("Invalid datasource uri"); }); - it("should use cached datasources for resource lookup", async () => { + it("should use cached datasources for resource lookup", async () => { await server.init(); // First call should fetch from service