From 72e5bad2e923ce424a9b9ec7e099a9f315e40448 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Mon, 27 Apr 2026 17:28:10 -0700 Subject: [PATCH 01/25] Add conversation storage server (#121) - Add new DO for conversation storage, and serve it on a new /storage/conversation-id route - Add test coverage - This DO is not yet used, we will switch to this in a future PR Co-authored-by: Rifdhan Nazeer --- src/index.ts | 19 + src/servers/conversation-storage-server.ts | 129 ++++++ .../conversation-storage-server.spec.ts | 375 ++++++++++++++++++ worker-configuration.d.ts | 1 + wrangler.jsonc | 10 + 5 files changed, 534 insertions(+) create mode 100644 src/servers/conversation-storage-server.ts create mode 100644 test/servers/conversation-storage-server.spec.ts diff --git a/src/index.ts b/src/index.ts index 1510805..6b0b693 100644 --- a/src/index.ts +++ b/src/index.ts @@ -12,6 +12,9 @@ import { MCPServer } from "./servers/mcp-server"; import { apiServer } from "./servers/api-server"; import { withBearerHandler } from "./bearer"; import { OpenAIDeepResearchMCPServer } from "./servers/openai-mcp-server"; +import { ConversationStorageServer } from "./servers/conversation-storage-server"; + +export { ConversationStorageServer }; // OTEL configuration function const config: ResolveConfigFn = (env: Env, _trigger) => { @@ -63,6 +66,21 @@ function createMCPRouter( }; } +const conversationStorageHandler = { + async fetch(request: Request, env: Env): Promise { + const url = new URL(request.url); + // Path format: /storage/[/] + const parts = url.pathname.split("/"); + const conversationId = parts[2]; + if (!conversationId) { + return new Response("Missing conversation ID", { status: 400 }); + } + const id = env.CONVERSATION_STORAGE_OBJECT.idFromName(conversationId); + const stub = env.CONVERSATION_STORAGE_OBJECT.get(id); + return stub.fetch(request); + }, +}; + // Create the OAuth provider instance const oauthProvider = new OAuthProvider({ apiHandlers: { @@ -75,6 +93,7 @@ const oauthProvider = new OAuthProvider({ binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", }) as any, // TODO: Remove 'any' "/api": apiServer as any, // TODO: Remove 'any' + "/storage": conversationStorageHandler as any, // TODO: Remove 'any' }, defaultHandler: withBearerHandler(handler, ThoughtSpotMCP) as any, // TODO: Remove 'any' authorizeEndpoint: "/authorize", diff --git a/src/servers/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts new file mode 100644 index 0000000..431191d --- /dev/null +++ b/src/servers/conversation-storage-server.ts @@ -0,0 +1,129 @@ +import type { Message, StreamingMessagesState } from "../thoughtspot/types"; + +const DEFAULT_TTL_MS = 30 * 60 * 1000; // 30 minutes + +const STATE_KEY = "streaming-messages-state"; +const BOOKMARK_KEY = "streaming-messages-bookmark"; + +/** + * A Durable Object that stores streaming conversation messages and exposes them over HTTP. + * + * Each instance corresponds to a single conversation. This means we don't need to use the + * conversationId internally, instead it is used to route to a unique instance per conversation. + * The parent DurableObject routes requests here via /storage/, and this DO + * handles the following sub-routes: + * + * POST /storage//initialize —> initializeConversation + * POST /storage//append —> appendMessagesAndRestartTtl + * GET /storage//messages —> getNewMessagesAndUpdateBookmark + */ +export class ConversationStorageServer { + constructor( + private state: DurableObjectState, + private env: Env, + ) {} + + async fetch(request: Request): Promise { + const url = new URL(request.url); + // Strip the /storage/ prefix; remaining path is the operation + // e.g. /storage/abc123/initialize -> /initialize + const parts = url.pathname.split("/"); + // parts: ["", "storage", "", ""] + const operation = parts[3] ?? ""; + + try { + switch (`${request.method} /${operation}`) { + case "POST /initialize": { + await this.initializeConversation(); + return Response.json({ ok: true }); + } + + case "POST /append": { + const body = (await request.json()) as StreamingMessagesState; + await this.appendMessagesAndRestartTtl(body.messages, body.isDone); + return Response.json({ ok: true }); + } + + case "GET /messages": { + const state = await this.getNewMessagesAndUpdateBookmark(); + return Response.json(state); + } + + default: + return new Response("Not Found", { status: 404 }); + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + console.error("Error handling conversation storage request:", message); + return Response.json({ error: "Something went wrong" }, { status: 500 }); + } + } + + /* + * Initialize the conversation. This can be a brand new conversation, or it can be priming an + * existing conversation which is already marked done for a followup message. + */ + private async initializeConversation(): Promise { + const existing = + await this.state.storage.get(STATE_KEY); + if (existing && !existing.isDone) { + throw new Error("Conversation already exists and is not marked done"); + } + + await this.setStateAndRestartTtl({ messages: [], isDone: false }); + await this.state.storage.put(BOOKMARK_KEY, 0); + } + + private async appendMessagesAndRestartTtl( + newMessages: Message[], + isDone = false, + ): Promise { + const oldState = + await this.state.storage.get(STATE_KEY); + if (!oldState) { + throw new Error("Conversation not found"); + } + if (oldState.isDone) { + throw new Error("Cannot append messages to a conversation marked done"); + } + + await this.setStateAndRestartTtl({ + messages: [...oldState.messages, ...newMessages], + isDone, + }); + } + + private async getNewMessagesAndUpdateBookmark(): Promise { + const bookmark = (await this.state.storage.get(BOOKMARK_KEY)) ?? 0; + + const conversationState = + await this.state.storage.get(STATE_KEY); + if (!conversationState) { + throw new Error("Conversation not found"); + } + + await this.state.storage.put( + BOOKMARK_KEY, + conversationState.messages.length, + ); + + return { + messages: conversationState.messages.slice(bookmark), + isDone: conversationState.isDone, + }; + } + + private async setStateAndRestartTtl( + newState: StreamingMessagesState, + ): Promise { + // Cancel any existing alarm and schedule a fresh one + await this.state.storage.deleteAlarm(); + await this.state.storage.setAlarm(Date.now() + DEFAULT_TTL_MS); + + await this.state.storage.put(STATE_KEY, newState); + } + + async alarm(): Promise { + await this.state.storage.delete([STATE_KEY, BOOKMARK_KEY]); + } +} diff --git a/test/servers/conversation-storage-server.spec.ts b/test/servers/conversation-storage-server.spec.ts new file mode 100644 index 0000000..e86999c --- /dev/null +++ b/test/servers/conversation-storage-server.spec.ts @@ -0,0 +1,375 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { ConversationStorageServer } from "../../src/servers/conversation-storage-server"; +import type { + Message, + StreamingMessagesState, +} from "../../src/thoughtspot/types"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function createMockStorage() { + const store = new Map(); + let alarm: number | null = null; + + return { + store, + get alarm() { + return alarm; + }, + storage: { + get: vi.fn(async (key: string): Promise => { + return store.get(key) as T | undefined; + }), + put: vi.fn(async (key: string, value: unknown): Promise => { + store.set(key, value); + }), + delete: vi.fn(async (keys: string[]): Promise => { + for (const key of keys) { + store.delete(key); + } + }), + setAlarm: vi.fn(async (scheduledTime: number): Promise => { + alarm = scheduledTime; + }), + deleteAlarm: vi.fn(async (): Promise => { + alarm = null; + }), + }, + }; +} + +function createServer(mock: ReturnType) { + const state = { storage: mock.storage } as unknown as DurableObjectState; + return new ConversationStorageServer(state, {} as Env); +} + +function makeRequest( + method: string, + operation: string, + body?: unknown, +): Request { + const url = `https://example.com/storage/conv-1/${operation}`; + return new Request(url, { + method, + headers: body ? { "Content-Type": "application/json" } : {}, + body: body ? JSON.stringify(body) : undefined, + }); +} + +// Sample messages +const textMessage: Message = { type: "text", text: "Hello" }; +const chunkMessage: Message = { type: "text_chunk", text: " world" }; +const answerMessage: Message = { + type: "answer", + answer_id: "ans-1", + answer_title: "My Answer", + answer_query: "SELECT 1", + iframe_url: "https://example.com/answer/1", +}; + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("ConversationStorageServer", () => { + let mock: ReturnType; + let server: ConversationStorageServer; + + beforeEach(() => { + mock = createMockStorage(); + server = createServer(mock); + }); + + // ------------------------------------------------------------------------- + // Routing + // ------------------------------------------------------------------------- + + describe("routing", () => { + it("returns 404 for an unknown route", async () => { + const res = await server.fetch(makeRequest("GET", "unknown")); + expect(res.status).toBe(404); + }); + + it("returns 404 for a valid operation with the wrong HTTP method", async () => { + const res = await server.fetch(makeRequest("GET", "initialize")); + expect(res.status).toBe(404); + }); + }); + + // ------------------------------------------------------------------------- + // POST /initialize + // ------------------------------------------------------------------------- + + describe("POST /initialize", () => { + it("responds with { ok: true } on success", async () => { + const res = await server.fetch(makeRequest("POST", "initialize")); + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ ok: true }); + }); + + it("stores empty messages and isDone=false", async () => { + await server.fetch(makeRequest("POST", "initialize")); + + const state = mock.store.get( + "streaming-messages-state", + ) as StreamingMessagesState; + expect(state).toMatchObject({ messages: [], isDone: false }); + }); + + it("sets bookmark to 0", async () => { + await server.fetch(makeRequest("POST", "initialize")); + + expect(mock.store.get("streaming-messages-bookmark")).toBe(0); + }); + + it("schedules a TTL alarm", async () => { + const before = Date.now(); + await server.fetch(makeRequest("POST", "initialize")); + + expect(mock.storage.setAlarm).toHaveBeenCalledOnce(); + const scheduledTime = mock.storage.setAlarm.mock.calls[0][0] as number; + expect(scheduledTime).toBeGreaterThanOrEqual(before + 30 * 60 * 1000); + }); + + it("returns 500 when conversation already exists and is not done", async () => { + await server.fetch(makeRequest("POST", "initialize")); + // Second init while not done should fail + const res = await server.fetch(makeRequest("POST", "initialize")); + expect(res.status).toBe(500); + }); + + it("allows re-initialization after the conversation is marked done", async () => { + await server.fetch(makeRequest("POST", "initialize")); + await server.fetch( + makeRequest("POST", "append", { + messages: [textMessage], + isDone: true, + }), + ); + + const res = await server.fetch(makeRequest("POST", "initialize")); + expect(res.status).toBe(200); + + const state = mock.store.get( + "streaming-messages-state", + ) as StreamingMessagesState; + expect(state).toMatchObject({ messages: [], isDone: false }); + }); + }); + + // ------------------------------------------------------------------------- + // POST /append + // ------------------------------------------------------------------------- + + describe("POST /append", () => { + beforeEach(async () => { + await server.fetch(makeRequest("POST", "initialize")); + vi.clearAllMocks(); + }); + + it("responds with { ok: true } on success", async () => { + const res = await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ ok: true }); + }); + + it("appends messages to storage", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + + const state = mock.store.get( + "streaming-messages-state", + ) as StreamingMessagesState; + expect(state.messages).toEqual([textMessage]); + }); + + it("accumulates messages across multiple calls", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + await server.fetch( + makeRequest("POST", "append", { + messages: [chunkMessage, answerMessage], + }), + ); + + const state = mock.store.get( + "streaming-messages-state", + ) as StreamingMessagesState; + expect(state.messages).toEqual([ + textMessage, + chunkMessage, + answerMessage, + ]); + }); + + it("marks the conversation done when isDone is true", async () => { + await server.fetch( + makeRequest("POST", "append", { + messages: [textMessage], + isDone: true, + }), + ); + + const state = mock.store.get( + "streaming-messages-state", + ) as StreamingMessagesState; + expect(state.isDone).toBe(true); + }); + + it("restarts the TTL alarm on each call", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + + expect(mock.storage.deleteAlarm).toHaveBeenCalledOnce(); + expect(mock.storage.setAlarm).toHaveBeenCalledOnce(); + }); + + it("returns 500 when the conversation does not exist", async () => { + // Wipe the state so the conversation is gone + mock.store.clear(); + + const res = await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + expect(res.status).toBe(500); + }); + + it("returns 500 when the conversation is already marked done", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [], isDone: true }), + ); + + const res = await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + expect(res.status).toBe(500); + }); + }); + + // ------------------------------------------------------------------------- + // GET /messages + // ------------------------------------------------------------------------- + + describe("GET /messages", () => { + beforeEach(async () => { + await server.fetch(makeRequest("POST", "initialize")); + }); + + it("returns empty messages and isDone=false on a fresh conversation", async () => { + const res = await server.fetch(makeRequest("GET", "messages")); + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ messages: [], isDone: false }); + }); + + it("returns all messages appended since the last call", async () => { + await server.fetch( + makeRequest("POST", "append", { + messages: [textMessage, chunkMessage], + }), + ); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + expect(body.messages).toEqual([textMessage, chunkMessage]); + }); + + it("advances the bookmark so subsequent calls only return new messages", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + // First poll — consumes textMessage + await server.fetch(makeRequest("GET", "messages")); + + // Append another message + await server.fetch( + makeRequest("POST", "append", { messages: [chunkMessage] }), + ); + + // Second poll — should only see chunkMessage + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + expect(body.messages).toEqual([chunkMessage]); + }); + + it("returns empty messages when polled again with no new messages", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + await server.fetch(makeRequest("GET", "messages")); // advances bookmark + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + expect(body.messages).toHaveLength(0); + }); + + it("reflects isDone=true when the conversation has been completed", async () => { + await server.fetch( + makeRequest("POST", "append", { + messages: [textMessage], + isDone: true, + }), + ); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + expect(body.isDone).toBe(true); + }); + + it("returns 500 when the conversation does not exist", async () => { + mock.store.clear(); + + const res = await server.fetch(makeRequest("GET", "messages")); + expect(res.status).toBe(500); + }); + }); + + // ------------------------------------------------------------------------- + // alarm() + // ------------------------------------------------------------------------- + + describe("alarm()", () => { + beforeEach(async () => { + await server.fetch(makeRequest("POST", "initialize")); + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + }); + + it("deletes the conversation state and bookmark", async () => { + await server.alarm(); + + expect(mock.store.has("streaming-messages-state")).toBe(false); + expect(mock.store.has("streaming-messages-bookmark")).toBe(false); + }); + + it("causes subsequent append to return 500", async () => { + await server.alarm(); + + const res = await server.fetch( + makeRequest("POST", "append", { messages: [chunkMessage] }), + ); + expect(res.status).toBe(500); + }); + + it("causes subsequent GET /messages to return 500", async () => { + await server.alarm(); + + const res = await server.fetch(makeRequest("GET", "messages")); + expect(res.status).toBe(500); + }); + + it("allows re-initialization after the alarm fires", async () => { + await server.alarm(); + + const res = await server.fetch(makeRequest("POST", "initialize")); + expect(res.status).toBe(200); + }); + }); +}); diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 4e97ea3..0b4c91b 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -10,6 +10,7 @@ declare namespace Cloudflare { OAUTH_KV: KVNamespace; MCP_OBJECT: DurableObjectNamespace; OPENAI_DEEP_RESEARCH_MCP_OBJECT: DurableObjectNamespace; + CONVERSATION_STORAGE_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; ASSETS: Fetcher; } diff --git a/wrangler.jsonc b/wrangler.jsonc index bf4897f..0d51a86 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -20,6 +20,10 @@ { "class_name": "ThoughtSpotOpenAIDeepResearchMCP", "name": "OPENAI_DEEP_RESEARCH_MCP_OBJECT" + }, + { + "class_name": "ConversationStorageServer", + "name": "CONVERSATION_STORAGE_OBJECT" } ] }, @@ -35,6 +39,12 @@ "new_sqlite_classes": [ "ThoughtSpotOpenAIDeepResearchMCP" ] + }, + { + "tag": "v4", + "new_classes": [ + "ConversationStorageServer" + ] } ], "kv_namespaces": [{ From 462edfbe3a32a620f838e6bfa8e39397b27f5e1b Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Mon, 27 Apr 2026 17:28:48 -0700 Subject: [PATCH 02/25] Add storage service (#122) - Create a storage-service client for the new storage server DO - Add test coverage - Not used yet, will be used in a future PR Co-authored-by: Rifdhan Nazeer --- src/storage-service/storage-service.ts | 93 ++++++ test/storage-service/storage-service.spec.ts | 280 +++++++++++++++++++ 2 files changed, 373 insertions(+) create mode 100644 src/storage-service/storage-service.ts create mode 100644 test/storage-service/storage-service.spec.ts diff --git a/src/storage-service/storage-service.ts b/src/storage-service/storage-service.ts new file mode 100644 index 0000000..b0af916 --- /dev/null +++ b/src/storage-service/storage-service.ts @@ -0,0 +1,93 @@ +import type { Message, StreamingMessagesState } from "../thoughtspot/types"; + +/** + * Client for the ConversationStorageServer Durable Object. + * + * Provides typed methods for each HTTP endpoint exposed by the server: + * POST /storage//initialize —> initializeConversation + * POST /storage//append —> appendMessagesAndRestartTtl + * GET /storage//messages —> getNewMessagesAndUpdateBookmark + */ +export class StorageServiceClient { + constructor( + private readonly baseUrl: string, + private readonly authToken: string, + ) {} + + private headers(): HeadersInit { + return { + "Content-Type": "application/json", + Accept: "application/json", + Authorization: `Bearer ${this.authToken}`, + }; + } + + private url(conversationId: string, operation: string): string { + return `${this.baseUrl}/storage/${encodeURIComponent(conversationId)}/${operation}`; + } + + /** + * Initialize a conversation. Must be called before appending messages. + * Can also be called on an existing conversation that is already marked done, + * to prime it for a follow-up message. + */ + async initializeConversation(conversationId: string): Promise { + const response = await fetch(this.url(conversationId, "initialize"), { + method: "POST", + headers: this.headers(), + }); + + if (!response.ok) { + const body = await response.text(); + throw new Error( + `Failed to initialize conversation (${response.status}): ${body}`, + ); + } + } + + /** + * Append new messages to a conversation and restart its TTL. + * Optionally mark the conversation as done. + */ + async appendMessages( + conversationId: string, + messages: Message[], + isDone = false, + ): Promise { + const body: StreamingMessagesState = { messages, isDone }; + + const response = await fetch(this.url(conversationId, "append"), { + method: "POST", + headers: this.headers(), + body: JSON.stringify(body), + }); + + if (!response.ok) { + const text = await response.text(); + throw new Error( + `Failed to append messages (${response.status}): ${text}`, + ); + } + } + + /** + * Retrieve all messages that have been added since the last call to this method + * (tracked via a per-conversation bookmark) and advance the bookmark. + * Also returns whether the conversation has been marked done. + */ + async getNewMessages( + conversationId: string, + ): Promise { + const response = await fetch(this.url(conversationId, "messages"), { + method: "GET", + headers: this.headers(), + }); + + if (!response.ok) { + const text = await response.text(); + throw new Error(`Failed to get messages (${response.status}): ${text}`); + } + + return response.json() as Promise; + } +} diff --git a/test/storage-service/storage-service.spec.ts b/test/storage-service/storage-service.spec.ts new file mode 100644 index 0000000..781ea1d --- /dev/null +++ b/test/storage-service/storage-service.spec.ts @@ -0,0 +1,280 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { StorageServiceClient } from "../../src/storage-service/storage-service"; +import type { + Message, + StreamingMessagesState, +} from "../../src/thoughtspot/types"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +const BASE_URL = "https://example.com"; +const AUTH_TOKEN = "test-token"; +const CONVERSATION_ID = "conv-abc123"; + +const textMessage: Message = { type: "text", text: "Hello" }; +const chunkMessage: Message = { type: "text_chunk", text: " world" }; +const answerMessage: Message = { + type: "answer", + answer_id: "ans-1", + answer_title: "My Answer", + answer_query: "SELECT 1", + iframe_url: "https://example.com/answer/1", +}; + +function mockFetchOk(body: unknown = { ok: true }): void { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + new Response(JSON.stringify(body), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ), + ); +} + +function mockFetchError(status: number, body: string): void { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(new Response(body, { status })), + ); +} + +function lastFetchCall(): { url: string; init: RequestInit } { + const mockFn = vi.mocked(fetch); + const [url, init] = mockFn.mock.calls[mockFn.mock.calls.length - 1] as [ + string, + RequestInit, + ]; + return { url, init }; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("StorageServiceClient", () => { + let client: StorageServiceClient; + + beforeEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + client = new StorageServiceClient(BASE_URL, AUTH_TOKEN); + }); + + // ------------------------------------------------------------------------- + // initializeConversation + // ------------------------------------------------------------------------- + + describe("initializeConversation", () => { + it("sends POST to /storage//initialize", async () => { + mockFetchOk(); + + await client.initializeConversation(CONVERSATION_ID); + + const { url, init } = lastFetchCall(); + expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/initialize`); + expect(init.method).toBe("POST"); + }); + + it("sends the Authorization header", async () => { + mockFetchOk(); + + await client.initializeConversation(CONVERSATION_ID); + + const { init } = lastFetchCall(); + expect((init.headers as Record).Authorization).toBe( + `Bearer ${AUTH_TOKEN}`, + ); + }); + + it("URL-encodes the conversation ID", async () => { + mockFetchOk(); + + await client.initializeConversation("conv with spaces/and-slash"); + + const { url } = lastFetchCall(); + expect(url).toBe( + `${BASE_URL}/storage/conv%20with%20spaces%2Fand-slash/initialize`, + ); + }); + + it("resolves without error on a 200 response", async () => { + mockFetchOk(); + await expect( + client.initializeConversation(CONVERSATION_ID), + ).resolves.toBeUndefined(); + }); + + it("throws when the server returns a non-ok status", async () => { + mockFetchError(500, "Something went wrong"); + + await expect( + client.initializeConversation(CONVERSATION_ID), + ).rejects.toThrow("Failed to initialize conversation (500)"); + }); + + it("includes the error body in the thrown error message", async () => { + mockFetchError(400, "Conversation already exists and is not marked done"); + + await expect( + client.initializeConversation(CONVERSATION_ID), + ).rejects.toThrow("Conversation already exists and is not marked done"); + }); + }); + + // ------------------------------------------------------------------------- + // appendMessages + // ------------------------------------------------------------------------- + + describe("appendMessages", () => { + it("sends POST to /storage//append", async () => { + mockFetchOk(); + + await client.appendMessages(CONVERSATION_ID, [textMessage]); + + const { url, init } = lastFetchCall(); + expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/append`); + expect(init.method).toBe("POST"); + }); + + it("sends messages and isDone=false in the request body by default", async () => { + mockFetchOk(); + + await client.appendMessages(CONVERSATION_ID, [textMessage, chunkMessage]); + + const { init } = lastFetchCall(); + const body = JSON.parse(init.body as string) as StreamingMessagesState; + expect(body.messages).toEqual([textMessage, chunkMessage]); + expect(body.isDone).toBe(false); + }); + + it("sends isDone=true when specified", async () => { + mockFetchOk(); + + await client.appendMessages(CONVERSATION_ID, [answerMessage], true); + + const { init } = lastFetchCall(); + const body = JSON.parse(init.body as string) as StreamingMessagesState; + expect(body.isDone).toBe(true); + }); + + it("sends the Authorization header", async () => { + mockFetchOk(); + + await client.appendMessages(CONVERSATION_ID, []); + + const { init } = lastFetchCall(); + expect((init.headers as Record).Authorization).toBe( + `Bearer ${AUTH_TOKEN}`, + ); + }); + + it("sends Content-Type: application/json", async () => { + mockFetchOk(); + + await client.appendMessages(CONVERSATION_ID, []); + + const { init } = lastFetchCall(); + expect((init.headers as Record)["Content-Type"]).toBe( + "application/json", + ); + }); + + it("resolves without error on a 200 response", async () => { + mockFetchOk(); + await expect( + client.appendMessages(CONVERSATION_ID, [textMessage]), + ).resolves.toBeUndefined(); + }); + + it("throws when the server returns a non-ok status", async () => { + mockFetchError(500, "Conversation not found"); + + await expect( + client.appendMessages(CONVERSATION_ID, [textMessage]), + ).rejects.toThrow("Failed to append messages (500)"); + }); + + it("includes the error body in the thrown error message", async () => { + mockFetchError( + 400, + "Cannot append messages to a conversation marked done", + ); + + await expect( + client.appendMessages(CONVERSATION_ID, [textMessage]), + ).rejects.toThrow("Cannot append messages to a conversation marked done"); + }); + }); + + // ------------------------------------------------------------------------- + // getNewMessages + // ------------------------------------------------------------------------- + + describe("getNewMessages", () => { + it("sends GET to /storage//messages", async () => { + const state: StreamingMessagesState = { + messages: [textMessage], + isDone: false, + }; + mockFetchOk(state); + + await client.getNewMessages(CONVERSATION_ID); + + const { url, init } = lastFetchCall(); + expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/messages`); + expect(init.method).toBe("GET"); + }); + + it("sends the Authorization header", async () => { + mockFetchOk({ messages: [], isDone: false }); + + await client.getNewMessages(CONVERSATION_ID); + + const { init } = lastFetchCall(); + expect((init.headers as Record).Authorization).toBe( + `Bearer ${AUTH_TOKEN}`, + ); + }); + + it("returns the parsed StreamingMessagesState", async () => { + const state: StreamingMessagesState = { + messages: [textMessage, answerMessage], + isDone: true, + }; + mockFetchOk(state); + + const result = await client.getNewMessages(CONVERSATION_ID); + + expect(result).toEqual(state); + }); + + it("returns an empty messages array when there are no new messages", async () => { + mockFetchOk({ messages: [], isDone: false }); + + const result = await client.getNewMessages(CONVERSATION_ID); + + expect(result.messages).toHaveLength(0); + expect(result.isDone).toBe(false); + }); + + it("throws when the server returns a non-ok status", async () => { + mockFetchError(404, "Conversation not found"); + + await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( + "Failed to get messages (404)", + ); + }); + + it("includes the error body in the thrown error message", async () => { + mockFetchError(500, "Internal error"); + + await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( + "Internal error", + ); + }); + }); +}); From 5dfe6c6b0ed39660a36cb681fb702b5b6ee6aec4 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Tue, 28 Apr 2026 11:28:19 -0700 Subject: [PATCH 03/25] Add is_thinking flag to analytical session updates (#123) - Each message will have a new is_thinking flag, indicating whether the message is part of Spotter's thinking process or final results - Add test coverage - Fix some minor biome formatter violations Co-authored-by: Rifdhan Nazeer --- src/servers/tool-definitions.ts | 5 + src/streaming-utils.ts | 4 + src/thoughtspot/types.ts | 8 +- test/handlers.spec.ts | 12 +- ...streaming-message-storage-with-ttl.spec.ts | 13 ++- test/streaming-utils.spec.ts | 105 +++++++++++++++--- ...ughtspot-client-nanoid-integration.spec.ts | 4 +- 7 files changed, 119 insertions(+), 32 deletions(-) diff --git a/src/servers/tool-definitions.ts b/src/servers/tool-definitions.ts index 7143dd5..c69a97e 100644 --- a/src/servers/tool-definitions.ts +++ b/src/servers/tool-definitions.ts @@ -138,6 +138,11 @@ export const ConversationUpdateSchema = z.object({ .describe( "The type of update: `text` or `text_chunk` for a natural language message from the Analytics Agent, or `answer` for a data visualization with query results. Determines which other fields are populated.", ), + is_thinking: z + .boolean() + .describe( + "Whether this update is part of the Analytics Agent's thinking process. Use this to separate thinking updates from final result updates. You can use the thinking updates to show intermediate status or progress updates to the user.", + ), text: z .string() .optional() diff --git a/src/streaming-utils.ts b/src/streaming-utils.ts index 8efb208..63b5ee2 100644 --- a/src/streaming-utils.ts +++ b/src/streaming-utils.ts @@ -73,12 +73,14 @@ export const processSendAgentConversationMessageStreamingResponse = async ( if (item.type === "text") { nTextMessagesParsed++; newMessages.push({ + is_thinking: item.metadata?.type === "thinking", type: "text", text: item.content, }); } else if (item.type === "text-chunk") { nTextMessagesParsed++; newMessages.push({ + is_thinking: item.metadata?.type === "thinking", type: "text_chunk", text: item.content, }); @@ -86,6 +88,7 @@ export const processSendAgentConversationMessageStreamingResponse = async ( nAnswerMessagesParsed++; const iframeUrl = `${instanceUrl}/?tsmcp=true#/embed/conv-assist-answer?sessionId=${item.metadata?.session_id}&genNo=${item.metadata?.gen_no}&acSessionId=${item.metadata?.transaction_id}&acGenNo=${item.metadata?.generation_number}`; newMessages.push({ + is_thinking: item.metadata?.type === "thinking", type: "answer", answer_id: JSON.stringify({ session_id: item.metadata?.session_id, @@ -113,6 +116,7 @@ export const processSendAgentConversationMessageStreamingResponse = async ( message: item, }); newMessages.push({ + is_thinking: false, type: "text", text: item.display_message || "Something went wrong", }); diff --git a/src/thoughtspot/types.ts b/src/thoughtspot/types.ts index 1b6f499..b4d6a60 100644 --- a/src/thoughtspot/types.ts +++ b/src/thoughtspot/types.ts @@ -30,12 +30,16 @@ export interface SessionInfo { enableSpotterDataSourceDiscovery?: boolean; } -export interface TextMessage { +export interface BaseMessage { + is_thinking: boolean; +} + +export interface TextMessage extends BaseMessage { type: "text" | "text_chunk"; text: string; } -export interface AnswerMessage { +export interface AnswerMessage extends BaseMessage { type: "answer"; answer_id: string; answer_title: string; diff --git a/test/handlers.spec.ts b/test/handlers.spec.ts index da9b144..0ac34e6 100644 --- a/test/handlers.spec.ts +++ b/test/handlers.spec.ts @@ -105,13 +105,11 @@ describe("Handlers", () => { // Mock the OAUTH_PROVIDER to return valid client info const mockOAuthProvider = { - parseAuthRequest: vi - .fn() - .mockResolvedValue({ - clientId: "test-client", - codeChallenge: "test-code-challenge", - codeChallengeMethod: "S256", - }), + parseAuthRequest: vi.fn().mockResolvedValue({ + clientId: "test-client", + codeChallenge: "test-code-challenge", + codeChallengeMethod: "S256", + }), lookupClient: vi.fn().mockResolvedValue({ clientId: "test-client", clientName: "Test Client", diff --git a/test/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl.spec.ts b/test/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl.spec.ts index f0fd448..7074853 100644 --- a/test/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl.spec.ts +++ b/test/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl.spec.ts @@ -29,9 +29,18 @@ function createMockStorage() { } // Sample messages used across tests -const textMessage: Message = { type: "text", text: "Hello" }; -const chunkMessage: Message = { type: "text_chunk", text: " world" }; +const textMessage: Message = { + is_thinking: false, + type: "text", + text: "Hello", +}; +const chunkMessage: Message = { + is_thinking: false, + type: "text_chunk", + text: " world", +}; const answerMessage: Message = { + is_thinking: false, type: "answer", answer_id: "ans-1", answer_title: "My Answer", diff --git a/test/streaming-utils.spec.ts b/test/streaming-utils.spec.ts index 9511fb8..1b31349 100644 --- a/test/streaming-utils.spec.ts +++ b/test/streaming-utils.spec.ts @@ -62,7 +62,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { it("parses a text event and stores a text message", async () => { const storage = makeMockStorage(); - const line = `data: ${JSON.stringify([{ type: "text", content: "Hello world" }])}\n`; + const line = `data: ${JSON.stringify([{ type: "text", content: "Hello world", metadata: {} }])}\n`; const reader = makeReader([line]); await processSendAgentConversationMessageStreamingResponse( @@ -73,7 +73,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "Hello world" }, + { is_thinking: false, type: "text", text: "Hello world" }, ]); // Final done call expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith( @@ -83,9 +83,26 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); }); + it("sets is_thinking=true on a text event when metadata.type is 'thinking'", async () => { + const storage = makeMockStorage(); + const line = `data: ${JSON.stringify([{ type: "text", content: "Reasoning...", metadata: { type: "thinking" } }])}\n`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage as any, + INSTANCE_URL, + ); + + expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ + { is_thinking: true, type: "text", text: "Reasoning..." }, + ]); + }); + it("parses a text-chunk event and stores a text_chunk message", async () => { const storage = makeMockStorage(); - const line = `data: ${JSON.stringify([{ type: "text-chunk", content: "chunk content" }])}\n`; + const line = `data: ${JSON.stringify([{ type: "text-chunk", content: "chunk content", metadata: {} }])}\n`; const reader = makeReader([line]); await processSendAgentConversationMessageStreamingResponse( @@ -96,7 +113,24 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text_chunk", text: "chunk content" }, + { is_thinking: false, type: "text_chunk", text: "chunk content" }, + ]); + }); + + it("sets is_thinking=true on a text-chunk event when metadata.type is 'thinking'", async () => { + const storage = makeMockStorage(); + const line = `data: ${JSON.stringify([{ type: "text-chunk", content: "thinking chunk", metadata: { type: "thinking" } }])}\n`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage as any, + INSTANCE_URL, + ); + + expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ + { is_thinking: true, type: "text_chunk", text: "thinking chunk" }, ]); }); @@ -123,6 +157,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { const expectedIframeUrl = `${INSTANCE_URL}/?tsmcp=true#/embed/conv-assist-answer?sessionId=sess-1&genNo=42&acSessionId=txn-1&acGenNo=7`; expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ { + is_thinking: false, type: "answer", answer_id: JSON.stringify({ session_id: "sess-1", gen_no: 42 }), answer_title: "My Answer", @@ -132,6 +167,40 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ]); }); + it("sets is_thinking=true on an answer event when metadata.type is 'thinking'", async () => { + const storage = makeMockStorage(); + const metadata = { + type: "thinking", + session_id: "sess-2", + gen_no: 1, + transaction_id: "txn-2", + generation_number: 2, + title: "Thinking Answer", + sage_query: "show revenue", + }; + const line = `data: ${JSON.stringify([{ type: "answer", metadata }])}\n`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage as any, + INSTANCE_URL, + ); + + const expectedIframeUrl = `${INSTANCE_URL}/?tsmcp=true#/embed/conv-assist-answer?sessionId=sess-2&genNo=1&acSessionId=txn-2&acGenNo=2`; + expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ + { + is_thinking: true, + type: "answer", + answer_id: JSON.stringify({ session_id: "sess-2", gen_no: 1 }), + answer_title: "Thinking Answer", + answer_query: "show revenue", + iframe_url: expectedIframeUrl, + }, + ]); + }); + it("parses an error event and stores a text message with the display_message", async () => { const storage = makeMockStorage(); const line = `data: ${JSON.stringify([{ type: "error", code: "ERR_001", message: "internal", display_message: "Something went wrong" }])}\n`; @@ -145,7 +214,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "Something went wrong" }, + { is_thinking: false, type: "text", text: "Something went wrong" }, ]); }); @@ -162,7 +231,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "Something went wrong" }, + { is_thinking: false, type: "text", text: "Something went wrong" }, ]); }); @@ -198,7 +267,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { it("ignores blank lines and heartbeat lines", async () => { const storage = makeMockStorage(); // Blank line and heartbeat, then a real message - const chunk = `\n: heartbeat\ndata: ${JSON.stringify([{ type: "text", content: "hi" }])}\n`; + const chunk = `\n: heartbeat\ndata: ${JSON.stringify([{ type: "text", content: "hi", metadata: {} }])}\n`; const reader = makeReader([chunk]); await processSendAgentConversationMessageStreamingResponse( @@ -209,7 +278,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "hi" }, + { is_thinking: false, type: "text", text: "hi" }, ]); }); @@ -273,7 +342,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { it("handles multiple chunks and assembles partial lines across reads correctly", async () => { const storage = makeMockStorage(); // Split the data line across two chunks - const fullLine = `data: ${JSON.stringify([{ type: "text", content: "split message" }])}`; + const fullLine = `data: ${JSON.stringify([{ type: "text", content: "split message", metadata: {} }])}`; const part1 = fullLine.slice(0, 20); const part2 = `${fullLine.slice(20)}\n`; const reader = makeReader([part1, part2]); @@ -286,14 +355,14 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "split message" }, + { is_thinking: false, type: "text", text: "split message" }, ]); }); it("processes multiple messages from multiple chunks and stores them in order", async () => { const storage = makeMockStorage(); - const chunk1 = `data: ${JSON.stringify([{ type: "text", content: "first" }])}\n`; - const chunk2 = `data: ${JSON.stringify([{ type: "text-chunk", content: "second" }])}\n`; + const chunk1 = `data: ${JSON.stringify([{ type: "text", content: "first", metadata: {} }])}\n`; + const chunk2 = `data: ${JSON.stringify([{ type: "text-chunk", content: "second", metadata: {} }])}\n`; const reader = makeReader([chunk1, chunk2]); await processSendAgentConversationMessageStreamingResponse( @@ -306,12 +375,12 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { expect(storage.appendMessagesAndRestartTtl).toHaveBeenNthCalledWith( 1, CONV_ID, - [{ type: "text", text: "first" }], + [{ is_thinking: false, type: "text", text: "first" }], ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenNthCalledWith( 2, CONV_ID, - [{ type: "text_chunk", text: "second" }], + [{ is_thinking: false, type: "text_chunk", text: "second" }], ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenNthCalledWith( 3, @@ -324,8 +393,8 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { it("processes multiple events in the same line as a batch", async () => { const storage = makeMockStorage(); const items = [ - { type: "text", content: "one" }, - { type: "text-chunk", content: "two" }, + { type: "text", content: "one", metadata: {} }, + { type: "text-chunk", content: "two", metadata: {} }, ]; const chunk = `data: ${JSON.stringify(items)}\n`; const reader = makeReader([chunk]); @@ -338,8 +407,8 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "one" }, - { type: "text_chunk", text: "two" }, + { is_thinking: false, type: "text", text: "one" }, + { is_thinking: false, type: "text_chunk", text: "two" }, ]); }); }); diff --git a/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts b/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts index 759b952..41d262e 100644 --- a/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts +++ b/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts @@ -165,9 +165,7 @@ describe("sendAgentConversationMessageStreaming — nano ID integration", () => const body = JSON.parse(options.body); // Endpoint construction - expect(url).toBe( - `${INSTANCE_URL}/conversation/v2/${conversationId}/query`, - ); + expect(url).toBe(`${INSTANCE_URL}/conversation/v2/${conversationId}/query`); // The id must be present, valid length, and from the correct alphabet expect(body.id).toHaveLength(NANO_ID_SIZE); From 65db06fc74d5ab1935313482e9a138b0eacb4139 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Wed, 29 Apr 2026 11:32:22 -0700 Subject: [PATCH 04/25] Use new Storage Server for streaming message storage (#124) - Remove /storage public route and use internal binding to access Storage Server (no longer need to expose the endpoint publicly) - Wire up Storage Client to use the new Storage Server in place of existing storage class - Update tests Co-authored-by: Rifdhan Nazeer --- src/cloudflare-utils.ts | 5 +- src/index.ts | 16 -- src/servers/mcp-server-base.ts | 11 +- src/servers/mcp-server.ts | 19 +- src/storage-service/storage-service.ts | 42 ++--- src/streaming-utils.ts | 18 +- src/thoughtspot/thoughtspot-service.ts | 9 +- test/storage-service/storage-service.spec.ts | 176 ++++++++----------- test/streaming-utils.spec.ts | 38 ++-- 9 files changed, 152 insertions(+), 182 deletions(-) diff --git a/src/cloudflare-utils.ts b/src/cloudflare-utils.ts index 6c656b2..00c866b 100644 --- a/src/cloudflare-utils.ts +++ b/src/cloudflare-utils.ts @@ -18,7 +18,10 @@ export function instrumentedMCPServer( this.scheduleTimer.bind(this), this.cancelTimer.bind(this), ); - server = new MCPServer(this as Context, this.streamingMessageStorage); + server = new MCPServer( + this as unknown as Context, + this.streamingMessageStorage, + ); // Argument of type 'typeof ThoughtSpotMCPWrapper' is not assignable to parameter of type 'DOClass'. // Cannot assign a 'protected' constructor type to a 'public' constructor type. diff --git a/src/index.ts b/src/index.ts index 6b0b693..123ffb5 100644 --- a/src/index.ts +++ b/src/index.ts @@ -66,21 +66,6 @@ function createMCPRouter( }; } -const conversationStorageHandler = { - async fetch(request: Request, env: Env): Promise { - const url = new URL(request.url); - // Path format: /storage/[/] - const parts = url.pathname.split("/"); - const conversationId = parts[2]; - if (!conversationId) { - return new Response("Missing conversation ID", { status: 400 }); - } - const id = env.CONVERSATION_STORAGE_OBJECT.idFromName(conversationId); - const stub = env.CONVERSATION_STORAGE_OBJECT.get(id); - return stub.fetch(request); - }, -}; - // Create the OAuth provider instance const oauthProvider = new OAuthProvider({ apiHandlers: { @@ -93,7 +78,6 @@ const oauthProvider = new OAuthProvider({ binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", }) as any, // TODO: Remove 'any' "/api": apiServer as any, // TODO: Remove 'any' - "/storage": conversationStorageHandler as any, // TODO: Remove 'any' }, defaultHandler: withBearerHandler(handler, ThoughtSpotMCP) as any, // TODO: Remove 'any' authorizeEndpoint: "/authorize", diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index c8c0ddf..106e3ef 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -8,13 +8,14 @@ import { type ListToolsResult, } from "@modelcontextprotocol/sdk/types.js"; import type { z } from "zod"; -import { context, type Span, SpanStatusCode } from "@opentelemetry/api"; +import { type Span, SpanStatusCode } from "@opentelemetry/api"; import { getActiveSpan, withSpan } from "../metrics/tracing/tracing-utils"; import { Trackers, type Tracker, TrackEvent } from "../metrics"; import type { Props } from "../utils"; import { MixpanelTracker } from "../metrics/mixpanel/mixpanel"; import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client"; import { ThoughtSpotService } from "../thoughtspot/thoughtspot-service"; +import { StorageServiceClient } from "../storage-service/storage-service"; const ToolInputSchema = ToolSchema.shape.inputSchema; type ToolInput = z.infer; @@ -39,6 +40,7 @@ export type ToolResponse = SuccessResponse | ErrorResponse; export interface Context { props: Props; + env: Env; } export abstract class BaseMCPServer extends Server { @@ -172,6 +174,13 @@ export abstract class BaseMCPServer extends Server { }; } + protected getStorageService(): StorageServiceClient { + return new StorageServiceClient( + this.ctx.env + .CONVERSATION_STORAGE_OBJECT as unknown as DurableObjectNamespace, + ); + } + protected getThoughtSpotService() { return new ThoughtSpotService( getThoughtSpotClient( diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 5b5a57b..cd56861 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -317,11 +317,14 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; has_additional_context: !!additional_context, }); + const storageService = this.getStorageService(); try { - await this.streamingMessageStorage.initializeConversation( - analytical_session_id, - ); + await storageService.initializeConversation(analytical_session_id); } catch (error) { + console.error( + "Error initializing conversation in storage service:", + error, + ); return this.createErrorResponse( "The analytical session has an ongoing response to the previous message. Please continue to call `get_session_updates` until `is_done` is true before sending a followup message.", `Error sending message to conversation ${analytical_session_id}: ${error}`, @@ -331,7 +334,7 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; await this.getThoughtSpotService().sendAgentConversationMessageStreaming( analytical_session_id, message, - this.streamingMessageStorage, + storageService.appendMessages.bind(storageService), additional_context, ); @@ -356,6 +359,7 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; // returning too quickly, which leads to too many get updates tool calls. // 4. If there are no updates after waiting for 10 seconds, return an empty response. We // want to avoid waiting indefinitely in case of errors or unexpected problems. + const storageService = this.getStorageService(); const messagesState: StreamingMessagesState = { messages: [], isDone: false, @@ -363,10 +367,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; let i = 0; for (; i < 20; i++) { // Get latest updates - const newMessagesState = - await this.streamingMessageStorage.getNewMessagesAndUpdateBookmark( - analytical_session_id, - ); + const newMessagesState = await storageService.getNewMessages( + analytical_session_id, + ); messagesState.messages.push(...newMessagesState.messages); messagesState.isDone = newMessagesState.isDone; diff --git a/src/storage-service/storage-service.ts b/src/storage-service/storage-service.ts index b0af916..805698d 100644 --- a/src/storage-service/storage-service.ts +++ b/src/storage-service/storage-service.ts @@ -3,27 +3,30 @@ import type { Message, StreamingMessagesState } from "../thoughtspot/types"; /** * Client for the ConversationStorageServer Durable Object. * - * Provides typed methods for each HTTP endpoint exposed by the server: + * Communicates directly with the DO via its stub (bypassing the OAuth layer), mapping to the + * following HTTP endpoints exposed by the server: * POST /storage//initialize —> initializeConversation * POST /storage//append —> appendMessagesAndRestartTtl * GET /storage//messages —> getNewMessagesAndUpdateBookmark */ export class StorageServiceClient { - constructor( - private readonly baseUrl: string, - private readonly authToken: string, - ) {} + constructor(private readonly namespace: DurableObjectNamespace) {} private headers(): HeadersInit { return { "Content-Type": "application/json", Accept: "application/json", - Authorization: `Bearer ${this.authToken}`, }; } + private stubFor(conversationId: string): DurableObjectStub { + const id = this.namespace.idFromName(conversationId); + return this.namespace.get(id); + } + + // DO stubs ignore the hostname; we use a placeholder so the path is parsed correctly. private url(conversationId: string, operation: string): string { - return `${this.baseUrl}/storage/${encodeURIComponent(conversationId)}/${operation}`; + return `https://internal/storage/${encodeURIComponent(conversationId)}/${operation}`; } /** @@ -32,10 +35,10 @@ export class StorageServiceClient { * to prime it for a follow-up message. */ async initializeConversation(conversationId: string): Promise { - const response = await fetch(this.url(conversationId, "initialize"), { - method: "POST", - headers: this.headers(), - }); + const response = await this.stubFor(conversationId).fetch( + this.url(conversationId, "initialize"), + { method: "POST", headers: this.headers() }, + ); if (!response.ok) { const body = await response.text(); @@ -56,11 +59,10 @@ export class StorageServiceClient { ): Promise { const body: StreamingMessagesState = { messages, isDone }; - const response = await fetch(this.url(conversationId, "append"), { - method: "POST", - headers: this.headers(), - body: JSON.stringify(body), - }); + const response = await this.stubFor(conversationId).fetch( + this.url(conversationId, "append"), + { method: "POST", headers: this.headers(), body: JSON.stringify(body) }, + ); if (!response.ok) { const text = await response.text(); @@ -78,10 +80,10 @@ export class StorageServiceClient { async getNewMessages( conversationId: string, ): Promise { - const response = await fetch(this.url(conversationId, "messages"), { - method: "GET", - headers: this.headers(), - }); + const response = await this.stubFor(conversationId).fetch( + this.url(conversationId, "messages"), + { method: "GET", headers: this.headers() }, + ); if (!response.ok) { const text = await response.text(); diff --git a/src/streaming-utils.ts b/src/streaming-utils.ts index 63b5ee2..ed3eee9 100644 --- a/src/streaming-utils.ts +++ b/src/streaming-utils.ts @@ -1,5 +1,4 @@ import type { Message } from "./thoughtspot/types"; -import type { StreamingMessagesStorageWithTtl } from "./streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; import { withSpan } from "./metrics/tracing/tracing-utils"; import { type Span, SpanStatusCode } from "@opentelemetry/api"; @@ -11,7 +10,11 @@ import { type Span, SpanStatusCode } from "@opentelemetry/api"; export const processSendAgentConversationMessageStreamingResponse = async ( conversationId: string, streamingResponseReader: ReadableStreamDefaultReader, - streamingMessageStorage: StreamingMessagesStorageWithTtl, + appendStoredMessages: ( + conversationId: string, + messages: Message[], + isDone?: boolean, + ) => Promise, instanceUrl: string, ) => { return await withSpan( @@ -33,11 +36,7 @@ export const processSendAgentConversationMessageStreamingResponse = async ( // If stream is marked done, mark the conversation as done and exit if (done) { - await streamingMessageStorage.appendMessagesAndRestartTtl( - conversationId, - [], - true, - ); + await appendStoredMessages(conversationId, [], true); break; } @@ -129,10 +128,7 @@ export const processSendAgentConversationMessageStreamingResponse = async ( // If we parsed any new messages, store them in the storage if (newMessages.length > 0) { - await streamingMessageStorage.appendMessagesAndRestartTtl( - conversationId, - newMessages, - ); + await appendStoredMessages(conversationId, newMessages); } } } catch (error) { diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index c44e304..424331b 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -11,7 +11,6 @@ import type { Message, Answer, } from "./types"; -import type { StreamingMessagesStorageWithTtl } from "../streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; import { processSendAgentConversationMessageStreamingResponse } from "../streaming-utils"; /** @@ -281,7 +280,11 @@ export class ThoughtSpotService { async sendAgentConversationMessageStreaming( conversationId: string, message: string, - streamingMessageStorage: StreamingMessagesStorageWithTtl, + appendStoredMessages: ( + conversationId: string, + messages: Message[], + isDone?: boolean, + ) => Promise, additionalContext?: string | undefined, ): Promise { const span = trace.getSpan(context.active()); @@ -315,7 +318,7 @@ export class ThoughtSpotService { processSendAgentConversationMessageStreamingResponse( conversationId, reader, - streamingMessageStorage, + appendStoredMessages, (this.client as any).instanceUrl, ); diff --git a/test/storage-service/storage-service.spec.ts b/test/storage-service/storage-service.spec.ts index 781ea1d..67194c4 100644 --- a/test/storage-service/storage-service.spec.ts +++ b/test/storage-service/storage-service.spec.ts @@ -9,8 +9,6 @@ import type { // Helpers // --------------------------------------------------------------------------- -const BASE_URL = "https://example.com"; -const AUTH_TOKEN = "test-token"; const CONVERSATION_ID = "conv-abc123"; const textMessage: Message = { type: "text", text: "Hello" }; @@ -23,32 +21,37 @@ const answerMessage: Message = { iframe_url: "https://example.com/answer/1", }; -function mockFetchOk(body: unknown = { ok: true }): void { - vi.stubGlobal( - "fetch", - vi.fn().mockResolvedValue( - new Response(JSON.stringify(body), { - status: 200, +// Captured request from the stub's last fetch call +let lastStubRequest: Request | undefined; + +function makeNamespaceMock( + responseBody: unknown = { ok: true }, + status = 200, +): DurableObjectNamespace { + lastStubRequest = undefined; + const stub = { + fetch: vi.fn(async (input: RequestInfo, init?: RequestInit) => { + lastStubRequest = new Request(input, init); + const body = + typeof responseBody === "string" + ? responseBody + : JSON.stringify(responseBody); + return new Response(body, { + status, headers: { "Content-Type": "application/json" }, - }), - ), - ); + }); + }), + } as unknown as DurableObjectStub; + + return { + idFromName: vi.fn(() => ({ toString: () => "stub-id" }) as DurableObjectId), + get: vi.fn(() => stub), + } as unknown as DurableObjectNamespace; } -function mockFetchError(status: number, body: string): void { - vi.stubGlobal( - "fetch", - vi.fn().mockResolvedValue(new Response(body, { status })), - ); -} - -function lastFetchCall(): { url: string; init: RequestInit } { - const mockFn = vi.mocked(fetch); - const [url, init] = mockFn.mock.calls[mockFn.mock.calls.length - 1] as [ - string, - RequestInit, - ]; - return { url, init }; +function lastRequest(): Request { + if (!lastStubRequest) throw new Error("No stub request recorded"); + return lastStubRequest; } // --------------------------------------------------------------------------- @@ -57,11 +60,12 @@ function lastFetchCall(): { url: string; init: RequestInit } { describe("StorageServiceClient", () => { let client: StorageServiceClient; + let namespaceMock: DurableObjectNamespace; beforeEach(() => { vi.restoreAllMocks(); - vi.unstubAllGlobals(); - client = new StorageServiceClient(BASE_URL, AUTH_TOKEN); + namespaceMock = makeNamespaceMock(); + client = new StorageServiceClient(namespaceMock); }); // ------------------------------------------------------------------------- @@ -70,46 +74,33 @@ describe("StorageServiceClient", () => { describe("initializeConversation", () => { it("sends POST to /storage//initialize", async () => { - mockFetchOk(); - - await client.initializeConversation(CONVERSATION_ID); - - const { url, init } = lastFetchCall(); - expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/initialize`); - expect(init.method).toBe("POST"); - }); - - it("sends the Authorization header", async () => { - mockFetchOk(); - await client.initializeConversation(CONVERSATION_ID); - const { init } = lastFetchCall(); - expect((init.headers as Record).Authorization).toBe( - `Bearer ${AUTH_TOKEN}`, + const req = lastRequest(); + expect(req.url).toBe( + `https://internal/storage/${CONVERSATION_ID}/initialize`, ); + expect(req.method).toBe("POST"); }); it("URL-encodes the conversation ID", async () => { - mockFetchOk(); - await client.initializeConversation("conv with spaces/and-slash"); - const { url } = lastFetchCall(); - expect(url).toBe( - `${BASE_URL}/storage/conv%20with%20spaces%2Fand-slash/initialize`, + const req = lastRequest(); + expect(req.url).toBe( + "https://internal/storage/conv%20with%20spaces%2Fand-slash/initialize", ); }); it("resolves without error on a 200 response", async () => { - mockFetchOk(); await expect( client.initializeConversation(CONVERSATION_ID), ).resolves.toBeUndefined(); }); it("throws when the server returns a non-ok status", async () => { - mockFetchError(500, "Something went wrong"); + namespaceMock = makeNamespaceMock("Something went wrong", 500); + client = new StorageServiceClient(namespaceMock); await expect( client.initializeConversation(CONVERSATION_ID), @@ -117,7 +108,11 @@ describe("StorageServiceClient", () => { }); it("includes the error body in the thrown error message", async () => { - mockFetchError(400, "Conversation already exists and is not marked done"); + namespaceMock = makeNamespaceMock( + "Conversation already exists and is not marked done", + 400, + ); + client = new StorageServiceClient(namespaceMock); await expect( client.initializeConversation(CONVERSATION_ID), @@ -131,67 +126,47 @@ describe("StorageServiceClient", () => { describe("appendMessages", () => { it("sends POST to /storage//append", async () => { - mockFetchOk(); - await client.appendMessages(CONVERSATION_ID, [textMessage]); - const { url, init } = lastFetchCall(); - expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/append`); - expect(init.method).toBe("POST"); + const req = lastRequest(); + expect(req.url).toBe( + `https://internal/storage/${CONVERSATION_ID}/append`, + ); + expect(req.method).toBe("POST"); }); it("sends messages and isDone=false in the request body by default", async () => { - mockFetchOk(); - await client.appendMessages(CONVERSATION_ID, [textMessage, chunkMessage]); - const { init } = lastFetchCall(); - const body = JSON.parse(init.body as string) as StreamingMessagesState; + const body = (await lastRequest().json()) as StreamingMessagesState; expect(body.messages).toEqual([textMessage, chunkMessage]); expect(body.isDone).toBe(false); }); it("sends isDone=true when specified", async () => { - mockFetchOk(); - await client.appendMessages(CONVERSATION_ID, [answerMessage], true); - const { init } = lastFetchCall(); - const body = JSON.parse(init.body as string) as StreamingMessagesState; + const body = (await lastRequest().json()) as StreamingMessagesState; expect(body.isDone).toBe(true); }); - it("sends the Authorization header", async () => { - mockFetchOk(); - - await client.appendMessages(CONVERSATION_ID, []); - - const { init } = lastFetchCall(); - expect((init.headers as Record).Authorization).toBe( - `Bearer ${AUTH_TOKEN}`, - ); - }); - it("sends Content-Type: application/json", async () => { - mockFetchOk(); - await client.appendMessages(CONVERSATION_ID, []); - const { init } = lastFetchCall(); - expect((init.headers as Record)["Content-Type"]).toBe( + expect(lastRequest().headers.get("Content-Type")).toBe( "application/json", ); }); it("resolves without error on a 200 response", async () => { - mockFetchOk(); await expect( client.appendMessages(CONVERSATION_ID, [textMessage]), ).resolves.toBeUndefined(); }); it("throws when the server returns a non-ok status", async () => { - mockFetchError(500, "Conversation not found"); + namespaceMock = makeNamespaceMock("Conversation not found", 500); + client = new StorageServiceClient(namespaceMock); await expect( client.appendMessages(CONVERSATION_ID, [textMessage]), @@ -199,10 +174,11 @@ describe("StorageServiceClient", () => { }); it("includes the error body in the thrown error message", async () => { - mockFetchError( - 400, + namespaceMock = makeNamespaceMock( "Cannot append messages to a conversation marked done", + 400, ); + client = new StorageServiceClient(namespaceMock); await expect( client.appendMessages(CONVERSATION_ID, [textMessage]), @@ -216,28 +192,16 @@ describe("StorageServiceClient", () => { describe("getNewMessages", () => { it("sends GET to /storage//messages", async () => { - const state: StreamingMessagesState = { - messages: [textMessage], - isDone: false, - }; - mockFetchOk(state); - - await client.getNewMessages(CONVERSATION_ID); - - const { url, init } = lastFetchCall(); - expect(url).toBe(`${BASE_URL}/storage/${CONVERSATION_ID}/messages`); - expect(init.method).toBe("GET"); - }); - - it("sends the Authorization header", async () => { - mockFetchOk({ messages: [], isDone: false }); + namespaceMock = makeNamespaceMock({ messages: [textMessage], isDone: false }); + client = new StorageServiceClient(namespaceMock); await client.getNewMessages(CONVERSATION_ID); - const { init } = lastFetchCall(); - expect((init.headers as Record).Authorization).toBe( - `Bearer ${AUTH_TOKEN}`, + const req = lastRequest(); + expect(req.url).toBe( + `https://internal/storage/${CONVERSATION_ID}/messages`, ); + expect(req.method).toBe("GET"); }); it("returns the parsed StreamingMessagesState", async () => { @@ -245,7 +209,8 @@ describe("StorageServiceClient", () => { messages: [textMessage, answerMessage], isDone: true, }; - mockFetchOk(state); + namespaceMock = makeNamespaceMock(state); + client = new StorageServiceClient(namespaceMock); const result = await client.getNewMessages(CONVERSATION_ID); @@ -253,7 +218,8 @@ describe("StorageServiceClient", () => { }); it("returns an empty messages array when there are no new messages", async () => { - mockFetchOk({ messages: [], isDone: false }); + namespaceMock = makeNamespaceMock({ messages: [], isDone: false }); + client = new StorageServiceClient(namespaceMock); const result = await client.getNewMessages(CONVERSATION_ID); @@ -262,7 +228,8 @@ describe("StorageServiceClient", () => { }); it("throws when the server returns a non-ok status", async () => { - mockFetchError(404, "Conversation not found"); + namespaceMock = makeNamespaceMock("Conversation not found", 404); + client = new StorageServiceClient(namespaceMock); await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( "Failed to get messages (404)", @@ -270,7 +237,8 @@ describe("StorageServiceClient", () => { }); it("includes the error body in the thrown error message", async () => { - mockFetchError(500, "Internal error"); + namespaceMock = makeNamespaceMock("Internal error", 500); + client = new StorageServiceClient(namespaceMock); await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( "Internal error", diff --git a/test/streaming-utils.spec.ts b/test/streaming-utils.spec.ts index 1b31349..92d3dfa 100644 --- a/test/streaming-utils.spec.ts +++ b/test/streaming-utils.spec.ts @@ -19,8 +19,10 @@ function makeReader(chunks: string[]): ReadableStreamDefaultReader { // Mock storage function makeMockStorage() { + const fn = vi.fn(async () => {}); return { - appendMessagesAndRestartTtl: vi.fn(async () => {}), + appendMessages: fn, + appendMessagesAndRestartTtl: fn, }; } @@ -48,7 +50,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -68,7 +70,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -91,7 +93,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -108,7 +110,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -125,7 +127,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -150,7 +152,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -184,7 +186,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -209,7 +211,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -226,7 +228,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -251,7 +253,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -273,7 +275,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -290,7 +292,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -310,7 +312,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -328,7 +330,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -350,7 +352,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -368,7 +370,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -402,7 +404,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); From 575dc94b6caffd7759b32c2270e86aaf528ee6a9 Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:56:56 -0700 Subject: [PATCH 05/25] SCAL-308277 Add metrics recorder core (#119) * SCAL-308277 Add metrics recorder core ## Summary Introduce the request-scoped metrics runtime core for the Worker and add focused tests so the new runtime metrics files do not regress repo coverage. ## What Changed - add metric definitions, label normalization, and runtime config for sink selection - add the request-scoped recorder, sink interfaces, composite/noop sinks, and context helpers - wrap the Worker fetch entrypoint so every request gets a recorder and non-blocking flush behavior - add targeted tests for metric context, sink behavior, recorder branches, runtime config, and request-scoped recorder helpers ## Notes - this PR only adds the runtime core, request wrapper plumbing, and tests - external Analytics Engine and Grafana emission lands in follow-up PRs ## Validation - `npm test -- --coverage.enabled=false test/metrics/runtime` - `npm test` * SCAL-308277 Fix recorder flush race ## Summary Close the recorder as soon as flush begins, add coverage for the in-flight flush case, and address the follow-up lint and label-policy feedback. ## What Changed - mark the recorder flushed before awaiting the sink and snapshot observations once - add a test that rejects metrics recorded after flush has started - document why forbidden metric label keys stay explicit - remove the redundant `both` switch case flagged by Biome ## Validation - `npm run lint` - `npm test` * Refine request metrics routing context ## Summary Strip tracing headers before metrics setup and centralize route metric classification so request metrics are easier to maintain as routes evolve. ## What Changed - move tracing header stripping ahead of `withRequestMetrics(...)` in the Worker entrypoint - replace duplicated metric path classification helpers with a shared explicit route-context table and resolver - expand focused metric-context tests to cover the shared route metadata and fallback behavior ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/metric-context.spec.ts test/index.header-stripping.spec.ts` * SCAL-308277 Add production route metrics registry Move known public routes into production-owned constants so metrics context coverage stays tied to the routes served by the Worker. - add shared route constants for exact routes and grouped prefixes - use route constants in Worker, bearer, and Hono route registration - key explicit metrics route contexts from production route constants - add storage, OpenAI challenge, token/register, hello, and OpenAPI route groups - allow the low-cardinality is_thinking stream update label - cover route context registration and label normalization in tests - `npm run lint` - `npm test` * SCAL-308277 Classify OpenAPI spec subroutes ## Summary Classify mounted OpenAPI spec subroutes with the same route context as the base `/openapi-spec` endpoint. ## What Changed - add an OpenAPI spec route prefix constant - map `/openapi-spec/*` to `route_group=openapi_spec` - keep OpenAPI spec routes on the static, unauthenticated API surface - cover generated tool spec paths in route context tests ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/metric-context.spec.ts` - `npm test` --- src/bearer.ts | 13 +- src/handlers.ts | 33 +- src/index.ts | 63 ++-- src/metrics/runtime/composite-sink.ts | 20 ++ src/metrics/runtime/metric-context.ts | 296 ++++++++++++++++++ src/metrics/runtime/metric-types.ts | 203 ++++++++++++ src/metrics/runtime/metrics-recorder.ts | 112 +++++++ src/metrics/runtime/metrics-sink.ts | 30 ++ src/metrics/runtime/noop-sink.ts | 5 + src/metrics/runtime/request-metrics.ts | 65 ++++ src/metrics/runtime/runtime-config.ts | 145 +++++++++ src/routes.ts | 46 +++ test/metrics/runtime/metric-context.spec.ts | 108 +++++++ test/metrics/runtime/metric-types.spec.ts | 36 +++ test/metrics/runtime/metrics-recorder.spec.ts | 186 +++++++++++ test/metrics/runtime/request-metrics.spec.ts | 93 ++++++ test/metrics/runtime/runtime-config.spec.ts | 129 ++++++++ test/metrics/runtime/sinks.spec.ts | 42 +++ 18 files changed, 1583 insertions(+), 42 deletions(-) create mode 100644 src/metrics/runtime/composite-sink.ts create mode 100644 src/metrics/runtime/metric-context.ts create mode 100644 src/metrics/runtime/metric-types.ts create mode 100644 src/metrics/runtime/metrics-recorder.ts create mode 100644 src/metrics/runtime/metrics-sink.ts create mode 100644 src/metrics/runtime/noop-sink.ts create mode 100644 src/metrics/runtime/request-metrics.ts create mode 100644 src/metrics/runtime/runtime-config.ts create mode 100644 src/routes.ts create mode 100644 test/metrics/runtime/metric-context.spec.ts create mode 100644 test/metrics/runtime/metric-types.spec.ts create mode 100644 test/metrics/runtime/metrics-recorder.spec.ts create mode 100644 test/metrics/runtime/request-metrics.spec.ts create mode 100644 test/metrics/runtime/runtime-config.spec.ts create mode 100644 test/metrics/runtime/sinks.spec.ts diff --git a/src/bearer.ts b/src/bearer.ts index 5a18de1..fdae0e9 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -1,6 +1,7 @@ import type { ThoughtSpotMCP } from "."; import type honoApp from "./handlers"; import { validateAndSanitizeUrl } from "./oauth-manager/oauth-utils"; +import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; /** * Handler function for bearer/token authentication endpoints @@ -62,12 +63,12 @@ function handleTokenAuth( // Route to appropriate handler const pathname = url.pathname; - if (pathname.endsWith("/mcp")) { - return MCPServer.serve("/mcp").fetch(req, env, ctx); + if (pathname.endsWith(PUBLIC_ROUTES.mcp)) { + return MCPServer.serve(PUBLIC_ROUTES.mcp).fetch(req, env, ctx); } - if (pathname.endsWith("/sse")) { - return MCPServer.serveSSE("/sse").fetch(req, env, ctx); + if (pathname.endsWith(PUBLIC_ROUTES.sse)) { + return MCPServer.serveSSE(PUBLIC_ROUTES.sse).fetch(req, env, ctx); } return new Response("Not found", { status: 404 }); @@ -79,13 +80,13 @@ export function withBearerHandler( ) { // These endpoints do NOT support api-version query params (will be removed in future) // Use /token endpoints instead for new implementations - app.mount("/bearer", (req, env, ctx) => { + app.mount(PUBLIC_ROUTE_PREFIXES.bearer, (req, env, ctx) => { return handleTokenAuth(req, env, ctx, MCPServer, false); }); // NEW: /token endpoints - supports api-version query params // Recommended for all new implementations - app.mount("/token", (req, env, ctx) => { + app.mount(PUBLIC_ROUTE_PREFIXES.token, (req, env, ctx) => { return handleTokenAuth(req, env, ctx, MCPServer, true); }); diff --git a/src/handlers.ts b/src/handlers.ts index 95d847d..dbaa2a5 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -2,20 +2,21 @@ import type { AuthRequest, OAuthHelpers, } from "@cloudflare/workers-oauth-provider"; +import { type Span, SpanStatusCode, context, trace } from "@opentelemetry/api"; import { Hono } from "hono"; -import type { Props } from "./utils"; -import { McpServerError } from "./utils"; +import { decodeBase64Url, encodeBase64Url } from "hono/utils/encode"; +import { any } from "zod"; +import { openApiSpecHandler } from "./api-schemas/open-api-spec"; +import { WithSpan, getActiveSpan } from "./metrics/tracing/tracing-utils"; import { + buildSamlRedirectUrl, parseRedirectApproval, renderApprovalDialog, - buildSamlRedirectUrl, } from "./oauth-manager/oauth-utils"; import { renderTokenCallback } from "./oauth-manager/token-utils"; -import { any } from "zod"; -import { encodeBase64Url, decodeBase64Url } from "hono/utils/encode"; -import { getActiveSpan, WithSpan } from "./metrics/tracing/tracing-utils"; -import { context, type Span, SpanStatusCode, trace } from "@opentelemetry/api"; -import { openApiSpecHandler } from "./api-schemas/open-api-spec"; +import { PUBLIC_ROUTES } from "./routes"; +import type { Props } from "./utils"; +import { McpServerError } from "./utils"; const app = new Hono<{ Bindings: Env & { OAUTH_PROVIDER: OAuthHelpers } }>(); @@ -235,17 +236,17 @@ class Handler { const handler = new Handler(); -app.get("/", async (c) => { +app.get(PUBLIC_ROUTES.root, async (c) => { const response = await handler.serveIndex(c.env); return response; }); -app.get("/hello", async (c) => { +app.get(PUBLIC_ROUTES.hello, async (c) => { const result = await handler.helloWorld(); return c.json(result); }); -app.get("/authorize", async (c) => { +app.get(PUBLIC_ROUTES.authorize, async (c) => { try { const response = await handler.getAuthorize( c.req.raw, @@ -257,7 +258,7 @@ app.get("/authorize", async (c) => { } }); -app.post("/authorize", async (c) => { +app.post(PUBLIC_ROUTES.authorize, async (c) => { try { const redirectUrl = await handler.postAuthorize(c.req.raw, c.req.url); return Response.redirect(redirectUrl); @@ -272,7 +273,7 @@ app.post("/authorize", async (c) => { } }); -app.get("/callback", async (c) => { +app.get(PUBLIC_ROUTES.callback, async (c) => { try { const htmlContent = await handler.handleCallback( c.req.raw, @@ -300,7 +301,7 @@ app.get("/callback", async (c) => { } }); -app.post("/store-token", async (c) => { +app.post(PUBLIC_ROUTES.storeToken, async (c) => { try { const result = await handler.storeToken(c.req.raw, c.env.OAUTH_PROVIDER); return new Response(JSON.stringify(result), { @@ -329,10 +330,10 @@ app.post("/store-token", async (c) => { } }); -app.get("/.well-known/openai-apps-challenge", (c) => { +app.get(PUBLIC_ROUTES.openaiAppsChallenge, (c) => { return c.text(process.env.OPEN_AI_TOKEN); }); -app.route("/openapi-spec", openApiSpecHandler); +app.route(PUBLIC_ROUTES.openapiSpec, openApiSpecHandler); export default app; diff --git a/src/index.ts b/src/index.ts index 123ffb5..7ea4c88 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,18 +1,20 @@ -import { trace } from "@opentelemetry/api"; +import OAuthProvider from "@cloudflare/workers-oauth-provider"; import { - instrument, type ResolveConfigFn, type TraceConfig, + instrument, } from "@microlabs/otel-cf-workers"; -import OAuthProvider from "@cloudflare/workers-oauth-provider"; +import { trace } from "@opentelemetry/api"; -import handler from "./handlers"; +import { withBearerHandler } from "./bearer"; import { instrumentedMCPServer } from "./cloudflare-utils"; -import { MCPServer } from "./servers/mcp-server"; +import handler from "./handlers"; +import { withRequestMetrics } from "./metrics/runtime/request-metrics"; +import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; import { apiServer } from "./servers/api-server"; -import { withBearerHandler } from "./bearer"; -import { OpenAIDeepResearchMCPServer } from "./servers/openai-mcp-server"; import { ConversationStorageServer } from "./servers/conversation-storage-server"; +import { MCPServer } from "./servers/mcp-server"; +import { OpenAIDeepResearchMCPServer } from "./servers/openai-mcp-server"; export { ConversationStorageServer }; @@ -69,20 +71,34 @@ function createMCPRouter( // Create the OAuth provider instance const oauthProvider = new OAuthProvider({ apiHandlers: { - "/mcp": createMCPRouter("/mcp", ThoughtSpotMCP, "serve") as any, - "/sse": createMCPRouter("/sse", ThoughtSpotMCP, "serveSSE") as any, - "/openai/mcp": ThoughtSpotOpenAIDeepResearchMCP.serve("/openai/mcp", { - binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", - }) as any, // TODO: Remove 'any' - "/openai/sse": ThoughtSpotOpenAIDeepResearchMCP.serveSSE("/openai/sse", { - binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", - }) as any, // TODO: Remove 'any' - "/api": apiServer as any, // TODO: Remove 'any' + [PUBLIC_ROUTES.mcp]: createMCPRouter( + PUBLIC_ROUTES.mcp, + ThoughtSpotMCP, + "serve", + ) as any, + [PUBLIC_ROUTES.sse]: createMCPRouter( + PUBLIC_ROUTES.sse, + ThoughtSpotMCP, + "serveSSE", + ) as any, + [PUBLIC_ROUTES.openaiMcp]: ThoughtSpotOpenAIDeepResearchMCP.serve( + PUBLIC_ROUTES.openaiMcp, + { + binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", + }, + ) as any, // TODO: Remove 'any' + [PUBLIC_ROUTES.openaiSse]: ThoughtSpotOpenAIDeepResearchMCP.serveSSE( + PUBLIC_ROUTES.openaiSse, + { + binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", + }, + ) as any, // TODO: Remove 'any' + [PUBLIC_ROUTE_PREFIXES.api]: apiServer as any, // TODO: Remove 'any' }, defaultHandler: withBearerHandler(handler, ThoughtSpotMCP) as any, // TODO: Remove 'any' - authorizeEndpoint: "/authorize", - tokenEndpoint: "/token", - clientRegistrationEndpoint: "/register", + authorizeEndpoint: PUBLIC_ROUTES.authorize, + tokenEndpoint: PUBLIC_ROUTES.oauthToken, + clientRegistrationEndpoint: PUBLIC_ROUTES.register, }); // Wrap the OAuth provider with a handler that includes tracing @@ -123,6 +139,13 @@ export default { HEADERS_TO_STRIP.forEach((header) => headers.delete(header)); request = new Request(request, { headers }); } - return instrumentedOAuthHandler.fetch!(request, env, ctx); + + return withRequestMetrics( + env as unknown as Record, + ctx, + async () => { + return instrumentedOAuthHandler.fetch!(request, env, ctx); + }, + ); }, }; diff --git a/src/metrics/runtime/composite-sink.ts b/src/metrics/runtime/composite-sink.ts new file mode 100644 index 0000000..81b982b --- /dev/null +++ b/src/metrics/runtime/composite-sink.ts @@ -0,0 +1,20 @@ +import type { MetricsFlushPayload, MetricsSink } from "./metrics-sink"; + +export class CompositeMetricsSink implements MetricsSink { + constructor(private readonly sinks: readonly MetricsSink[]) {} + + async flush(payload: MetricsFlushPayload): Promise { + const results = await Promise.allSettled( + this.sinks.map((sink) => sink.flush(payload)), + ); + + for (const [index, result] of results.entries()) { + if (result.status === "rejected") { + console.error( + `[metrics] Sink at index ${index} failed during flush`, + result.reason, + ); + } + } + } +} diff --git a/src/metrics/runtime/metric-context.ts b/src/metrics/runtime/metric-context.ts new file mode 100644 index 0000000..5b704e8 --- /dev/null +++ b/src/metrics/runtime/metric-context.ts @@ -0,0 +1,296 @@ +import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "../../routes"; +import type { + ApiSurface, + AuthMode, + RouteGroup, + StatusClass, + Transport, +} from "./metric-types"; + +export type RequestMetricContext = { + routeGroup: RouteGroup; + transport: Transport; + apiSurface: ApiSurface; + authMode: AuthMode; +}; + +// Keep explicit route classifications here so adding a new public path only +// requires a single metadata update instead of touching multiple helper +// functions. +export const EXPLICIT_ROUTE_CONTEXTS = { + [PUBLIC_ROUTES.root]: { + routeGroup: "root", + transport: "http", + apiSurface: "static", + authMode: "none", + }, + [PUBLIC_ROUTES.hello]: { + routeGroup: "hello", + transport: "http", + apiSurface: "static", + authMode: "none", + }, + [PUBLIC_ROUTES.authorize]: { + routeGroup: "authorize", + transport: "http", + apiSurface: "oauth", + authMode: "none", + }, + [PUBLIC_ROUTES.callback]: { + routeGroup: "callback", + transport: "http", + apiSurface: "oauth", + authMode: "none", + }, + [PUBLIC_ROUTES.storeToken]: { + routeGroup: "store_token", + transport: "http", + apiSurface: "oauth", + authMode: "none", + }, + [PUBLIC_ROUTES.oauthToken]: { + routeGroup: "oauth_token", + transport: "http", + apiSurface: "oauth", + authMode: "none", + }, + [PUBLIC_ROUTES.register]: { + routeGroup: "register", + transport: "http", + apiSurface: "oauth", + authMode: "none", + }, + [PUBLIC_ROUTES.mcp]: { + routeGroup: "mcp", + transport: "mcp", + apiSurface: "mcp", + authMode: "oauth", + }, + [PUBLIC_ROUTES.sse]: { + routeGroup: "sse", + transport: "sse", + apiSurface: "mcp", + authMode: "oauth", + }, + [PUBLIC_ROUTES.openaiMcp]: { + routeGroup: "openai_mcp", + transport: "mcp", + apiSurface: "openai_mcp", + authMode: "oauth", + }, + [PUBLIC_ROUTES.openaiSse]: { + routeGroup: "openai_sse", + transport: "sse", + apiSurface: "openai_mcp", + authMode: "oauth", + }, + [PUBLIC_ROUTES.bearerMcp]: { + routeGroup: "bearer_mcp", + transport: "mcp", + apiSurface: "mcp", + authMode: "bearer", + }, + [PUBLIC_ROUTES.bearerSse]: { + routeGroup: "bearer_sse", + transport: "sse", + apiSurface: "mcp", + authMode: "bearer", + }, + [PUBLIC_ROUTES.tokenMcp]: { + routeGroup: "token_mcp", + transport: "mcp", + apiSurface: "mcp", + authMode: "token", + }, + [PUBLIC_ROUTES.tokenSse]: { + routeGroup: "token_sse", + transport: "sse", + apiSurface: "mcp", + authMode: "token", + }, + [PUBLIC_ROUTES.openaiAppsChallenge]: { + routeGroup: "openai_apps_challenge", + transport: "http", + apiSurface: "static", + authMode: "none", + }, + [PUBLIC_ROUTES.openapiSpec]: { + routeGroup: "openapi_spec", + transport: "http", + apiSurface: "static", + authMode: "none", + }, +} as const satisfies Record; + +const API_ROUTE_CONTEXT: RequestMetricContext = { + routeGroup: "api", + transport: "http", + apiSurface: "api", + authMode: "oauth", +}; + +const OPENAPI_SPEC_ROUTE_CONTEXT: RequestMetricContext = { + routeGroup: "openapi_spec", + transport: "http", + apiSurface: "static", + authMode: "none", +}; + +const UNKNOWN_ROUTE_CONTEXT: RequestMetricContext = { + routeGroup: "unknown", + transport: "http", + apiSurface: "unknown", + authMode: "unknown", +}; + +function getExplicitRouteContext( + pathname: string, +): RequestMetricContext | undefined { + return EXPLICIT_ROUTE_CONTEXTS[ + pathname as keyof typeof EXPLICIT_ROUTE_CONTEXTS + ]; +} + +function matchesRoutePrefix(pathname: string, prefix: string): boolean { + return pathname === prefix || pathname.startsWith(`${prefix}/`); +} + +function inferTransport(pathname: string): Transport { + if (pathname.endsWith(PUBLIC_ROUTES.mcp) || pathname === PUBLIC_ROUTES.mcp) { + return "mcp"; + } + if (pathname.endsWith(PUBLIC_ROUTES.sse) || pathname === PUBLIC_ROUTES.sse) { + return "sse"; + } + return "http"; +} + +function inferApiSurface(pathname: string): ApiSurface { + if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api)) { + return "api"; + } + if (pathname.startsWith("/openai/")) { + return "openai_mcp"; + } + if ( + pathname === PUBLIC_ROUTES.mcp || + pathname === PUBLIC_ROUTES.sse || + matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.bearer) || + matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.token) + ) { + return "mcp"; + } + if ( + pathname === PUBLIC_ROUTES.root || + pathname === PUBLIC_ROUTES.hello || + pathname === PUBLIC_ROUTES.openaiAppsChallenge || + matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec) + ) { + return "static"; + } + if ( + pathname === PUBLIC_ROUTES.authorize || + pathname === PUBLIC_ROUTES.callback || + pathname === PUBLIC_ROUTES.storeToken || + pathname === PUBLIC_ROUTES.oauthToken || + pathname === PUBLIC_ROUTES.register + ) { + return "oauth"; + } + return "unknown"; +} + +function inferAuthMode(pathname: string): AuthMode { + if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.bearer)) { + return "bearer"; + } + if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.token)) { + return "token"; + } + if ( + pathname === PUBLIC_ROUTES.mcp || + pathname === PUBLIC_ROUTES.sse || + pathname.startsWith("/openai/") || + matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api) + ) { + return "oauth"; + } + if ( + pathname === PUBLIC_ROUTES.root || + pathname === PUBLIC_ROUTES.hello || + pathname === PUBLIC_ROUTES.authorize || + pathname === PUBLIC_ROUTES.callback || + pathname === PUBLIC_ROUTES.storeToken || + pathname === PUBLIC_ROUTES.oauthToken || + pathname === PUBLIC_ROUTES.register || + pathname === PUBLIC_ROUTES.openaiAppsChallenge || + matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec) + ) { + return "none"; + } + return "unknown"; +} + +export function resolvePathMetricContext( + pathname: string, +): RequestMetricContext { + const explicitContext = getExplicitRouteContext(pathname); + if (explicitContext) { + return explicitContext; + } + + if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api)) { + return API_ROUTE_CONTEXT; + } + + if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec)) { + return OPENAPI_SPEC_ROUTE_CONTEXT; + } + + return { + ...UNKNOWN_ROUTE_CONTEXT, + transport: inferTransport(pathname), + apiSurface: inferApiSurface(pathname), + authMode: inferAuthMode(pathname), + }; +} + +export function getRouteGroup(pathname: string): RouteGroup { + return resolvePathMetricContext(pathname).routeGroup; +} + +export function getTransport(pathname: string): Transport { + return resolvePathMetricContext(pathname).transport; +} + +export function getApiSurface(pathname: string): ApiSurface { + return resolvePathMetricContext(pathname).apiSurface; +} + +export function getAuthMode(pathname: string): AuthMode { + return resolvePathMetricContext(pathname).authMode; +} + +export function getStatusClass(status: number): StatusClass { + if (status >= 100 && status < 200) { + return "1xx"; + } + if (status >= 200 && status < 300) { + return "2xx"; + } + if (status >= 300 && status < 400) { + return "3xx"; + } + if (status >= 400 && status < 500) { + return "4xx"; + } + if (status >= 500 && status < 600) { + return "5xx"; + } + return "unknown"; +} + +export function resolveRequestMetricContext(request: Request) { + const pathname = new URL(request.url).pathname; + return resolvePathMetricContext(pathname); +} diff --git a/src/metrics/runtime/metric-types.ts b/src/metrics/runtime/metric-types.ts new file mode 100644 index 0000000..58a72ac --- /dev/null +++ b/src/metrics/runtime/metric-types.ts @@ -0,0 +1,203 @@ +export const HISTOGRAM_BUCKETS_MS = [ + 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000, +] as const; + +export const METRIC_NAMES = { + httpRequestsTotal: "ts_mcp_http_requests_total", + httpRequestDurationMs: "ts_mcp_http_request_duration_ms", + httpInflightRequests: "ts_mcp_http_inflight_requests", + sessionsStartedTotal: "ts_mcp_sessions_started_total", + toolCallsTotal: "ts_mcp_tool_calls_total", + toolDurationMs: "ts_mcp_tool_duration_ms", + resourceReadsTotal: "ts_mcp_resource_reads_total", + oauthAuthorizeRequestsTotal: "ts_mcp_oauth_authorize_requests_total", + oauthAuthorizeSubmitTotal: "ts_mcp_oauth_authorize_submit_total", + oauthCallbackTotal: "ts_mcp_oauth_callback_total", + oauthStoreTokenTotal: "ts_mcp_oauth_store_token_total", + bearerAuthRequestsTotal: "ts_mcp_bearer_auth_requests_total", + upstreamCallsTotal: "ts_mcp_upstream_calls_total", + upstreamDurationMs: "ts_mcp_upstream_duration_ms", + upstreamStreamsStartedTotal: "ts_mcp_upstream_streams_started_total", + upstreamStreamMessagesTotal: "ts_mcp_upstream_stream_messages_total", + analysisSessionsCreatedTotal: "ts_mcp_analysis_sessions_created_total", + analysisMessagesSentTotal: "ts_mcp_analysis_messages_sent_total", + analysisUpdatesPolledTotal: "ts_mcp_analysis_updates_polled_total", + analysisPollWaitMs: "ts_mcp_analysis_poll_wait_ms", + analysisFirstBufferedUpdateMs: "ts_mcp_analysis_first_buffered_update_ms", + analysisFirstPollDelayMs: "ts_mcp_analysis_first_poll_delay_ms", + analysisFirstNonEmptyResponseMs: + "ts_mcp_analysis_first_non_empty_response_ms", + analysisSessionsNeverPolledTotal: + "ts_mcp_analysis_sessions_never_polled_total", + streamStorageErrorsTotal: "ts_mcp_stream_storage_errors_total", + dashboardsCreatedTotal: "ts_mcp_dashboards_created_total", + dashboardTilesCount: "ts_mcp_dashboard_tiles_count", +} as const; + +export type MetricName = (typeof METRIC_NAMES)[keyof typeof METRIC_NAMES]; +export type MetricKind = "counter" | "histogram" | "gauge"; + +const COUNTER_METRIC_NAMES = new Set([ + METRIC_NAMES.httpRequestsTotal, + METRIC_NAMES.sessionsStartedTotal, + METRIC_NAMES.toolCallsTotal, + METRIC_NAMES.resourceReadsTotal, + METRIC_NAMES.oauthAuthorizeRequestsTotal, + METRIC_NAMES.oauthAuthorizeSubmitTotal, + METRIC_NAMES.oauthCallbackTotal, + METRIC_NAMES.oauthStoreTokenTotal, + METRIC_NAMES.bearerAuthRequestsTotal, + METRIC_NAMES.upstreamCallsTotal, + METRIC_NAMES.upstreamStreamsStartedTotal, + METRIC_NAMES.upstreamStreamMessagesTotal, + METRIC_NAMES.analysisSessionsCreatedTotal, + METRIC_NAMES.analysisMessagesSentTotal, + METRIC_NAMES.analysisUpdatesPolledTotal, + METRIC_NAMES.analysisSessionsNeverPolledTotal, + METRIC_NAMES.streamStorageErrorsTotal, + METRIC_NAMES.dashboardsCreatedTotal, +]); + +const HISTOGRAM_METRIC_NAMES = new Set([ + METRIC_NAMES.httpRequestDurationMs, + METRIC_NAMES.toolDurationMs, + METRIC_NAMES.upstreamDurationMs, + METRIC_NAMES.analysisPollWaitMs, + METRIC_NAMES.analysisFirstBufferedUpdateMs, + METRIC_NAMES.analysisFirstPollDelayMs, + METRIC_NAMES.analysisFirstNonEmptyResponseMs, + METRIC_NAMES.dashboardTilesCount, +]); + +const GAUGE_METRIC_NAMES = new Set([ + METRIC_NAMES.httpInflightRequests, +]); + +export function getMetricKind(name: MetricName): MetricKind { + if (COUNTER_METRIC_NAMES.has(name)) { + return "counter"; + } + if (HISTOGRAM_METRIC_NAMES.has(name)) { + return "histogram"; + } + if (GAUGE_METRIC_NAMES.has(name)) { + return "gauge"; + } + throw new Error(`Unknown metric kind for metric: ${name}`); +} + +export const APPROVED_METRIC_LABEL_KEYS = [ + "route_group", + "transport", + "auth_mode", + "api_surface", + "api_version", + "outcome", + "status_class", + "tool_name", + "upstream_operation", + "message_type", + "is_thinking", + "is_done", + "operation", +] as const; + +export const FORBIDDEN_METRIC_LABEL_KEYS = [ + // These are repo-known high-cardinality or sensitive fields that must never + // be promoted into metric labels, even if someone tries to pass them. + "instanceUrl", + "userGUID", + "userName", + "clientId", + "datasourceId", + "conversationId", + "question", + "query", + "redirectUrl", + "frameUrl", + "authorization", + "x-ts-host", +] as const; + +const APPROVED_METRIC_LABEL_KEYS_SET = new Set( + APPROVED_METRIC_LABEL_KEYS, +); +const FORBIDDEN_METRIC_LABEL_KEYS_SET = new Set( + FORBIDDEN_METRIC_LABEL_KEYS, +); + +export type MetricLabelKey = (typeof APPROVED_METRIC_LABEL_KEYS)[number]; +export type MetricLabelValue = string | number | boolean; +export type MetricLabels = Partial>; +export type MetricLabelInput = Partial< + Record +> & + Record; + +export type MetricOutcome = + | "success" + | "error" + | "client_error" + | "upstream_error" + | "validation_error"; + +export type RouteGroup = + | "root" + | "hello" + | "authorize" + | "callback" + | "store_token" + | "oauth_token" + | "register" + | "mcp" + | "sse" + | "openai_mcp" + | "openai_sse" + | "openai_apps_challenge" + | "openapi_spec" + | "api" + | "bearer_mcp" + | "bearer_sse" + | "token_mcp" + | "token_sse" + | "unknown"; + +export type Transport = "mcp" | "sse" | "http" | "unknown"; +export type AuthMode = "oauth" | "bearer" | "token" | "none" | "unknown"; +export type ApiSurface = + | "mcp" + | "openai_mcp" + | "api" + | "oauth" + | "static" + | "unknown"; +export type StatusClass = "1xx" | "2xx" | "3xx" | "4xx" | "5xx" | "unknown"; + +function warnOnInvalidMetricLabel(key: string, reason: string) { + console.warn(`[metrics] Dropping label "${key}": ${reason}`); +} + +export function normalizeMetricLabels(labels?: MetricLabelInput): MetricLabels { + if (!labels) { + return {}; + } + + const normalized: MetricLabels = {}; + for (const key of Object.keys(labels).sort()) { + const rawValue = labels[key]; + if (rawValue === undefined || rawValue === null || rawValue === "") { + continue; + } + if (FORBIDDEN_METRIC_LABEL_KEYS_SET.has(key)) { + warnOnInvalidMetricLabel(key, "forbidden by cardinality guardrail"); + continue; + } + if (!APPROVED_METRIC_LABEL_KEYS_SET.has(key)) { + warnOnInvalidMetricLabel(key, "not in approved label set"); + continue; + } + normalized[key as MetricLabelKey] = String(rawValue); + } + + return normalized; +} diff --git a/src/metrics/runtime/metrics-recorder.ts b/src/metrics/runtime/metrics-recorder.ts new file mode 100644 index 0000000..e07a2dc --- /dev/null +++ b/src/metrics/runtime/metrics-recorder.ts @@ -0,0 +1,112 @@ +import { + getMetricKind, + normalizeMetricLabels, + type MetricKind, + type MetricLabelInput, + type MetricName, +} from "./metric-types"; +import type { + MetricObservation, + MetricResourceAttributes, + MetricsSink, +} from "./metrics-sink"; + +type MetricsRecorderOptions = { + sink: MetricsSink; + resourceAttributes?: MetricResourceAttributes; + now?: () => number; +}; + +export interface MetricsRecorder { + count(name: MetricName, value?: number, labels?: MetricLabelInput): void; + histogram(name: MetricName, value: number, labels?: MetricLabelInput): void; + gauge(name: MetricName, value: number, labels?: MetricLabelInput): void; + flush(ctx?: Pick): Promise; + snapshot(): readonly MetricObservation[]; +} + +export class RequestMetricsRecorder implements MetricsRecorder { + private readonly observations: MetricObservation[] = []; + private flushPromise?: Promise; + private flushed = false; + + constructor(private readonly options: MetricsRecorderOptions) {} + + count(name: MetricName, value = 1, labels?: MetricLabelInput): void { + this.record("counter", name, value, labels); + } + + histogram(name: MetricName, value: number, labels?: MetricLabelInput): void { + this.record("histogram", name, value, labels); + } + + gauge(name: MetricName, value: number, labels?: MetricLabelInput): void { + this.record("gauge", name, value, labels); + } + + snapshot(): readonly MetricObservation[] { + return [...this.observations]; + } + + async flush(ctx?: Pick): Promise { + if (!this.flushPromise) { + this.flushPromise = this.flushInternal(); + } + + if (ctx) { + ctx.waitUntil(this.flushPromise); + } + + return this.flushPromise; + } + + private record( + expectedKind: MetricKind, + name: MetricName, + value: number, + labels?: MetricLabelInput, + ): void { + if (this.flushed) { + console.warn(`[metrics] Ignoring metric recorded after flush: ${name}`); + return; + } + if (!Number.isFinite(value)) { + console.warn(`[metrics] Ignoring non-finite metric value for ${name}`); + return; + } + + const actualKind = getMetricKind(name); + if (actualKind !== expectedKind) { + console.warn( + `[metrics] Ignoring ${expectedKind} write for ${name}; metric is defined as ${actualKind}`, + ); + return; + } + + this.observations.push({ + kind: actualKind, + name, + value, + labels: normalizeMetricLabels(labels), + timestampMs: this.options.now?.() ?? Date.now(), + }); + } + + private async flushInternal(): Promise { + this.flushed = true; + const observations = this.snapshot(); + + try { + if (observations.length === 0) { + return; + } + + await this.options.sink.flush({ + observations, + resourceAttributes: { ...this.options.resourceAttributes }, + }); + } catch (error) { + console.error("[metrics] Flush failed", error); + } + } +} diff --git a/src/metrics/runtime/metrics-sink.ts b/src/metrics/runtime/metrics-sink.ts new file mode 100644 index 0000000..edc9c47 --- /dev/null +++ b/src/metrics/runtime/metrics-sink.ts @@ -0,0 +1,30 @@ +import type { MetricKind, MetricLabels, MetricName } from "./metric-types"; + +export type MetricObservation = { + kind: MetricKind; + name: MetricName; + value: number; + labels: MetricLabels; + timestampMs: number; +}; + +export type MetricResourceAttributes = Partial< + Record< + | "service.name" + | "service.namespace" + | "service.version" + | "deployment.environment" + | "cloud.provider" + | "cloud.platform", + string + > +>; + +export type MetricsFlushPayload = { + observations: readonly MetricObservation[]; + resourceAttributes: MetricResourceAttributes; +}; + +export interface MetricsSink { + flush(payload: MetricsFlushPayload): Promise; +} diff --git a/src/metrics/runtime/noop-sink.ts b/src/metrics/runtime/noop-sink.ts new file mode 100644 index 0000000..fa7bfa7 --- /dev/null +++ b/src/metrics/runtime/noop-sink.ts @@ -0,0 +1,5 @@ +import type { MetricsFlushPayload, MetricsSink } from "./metrics-sink"; + +export class NoopMetricsSink implements MetricsSink { + async flush(_payload: MetricsFlushPayload): Promise {} +} diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts new file mode 100644 index 0000000..9aaaf3e --- /dev/null +++ b/src/metrics/runtime/request-metrics.ts @@ -0,0 +1,65 @@ +import { RequestMetricsRecorder, type MetricsRecorder } from "./metrics-recorder"; +import { + createConfiguredMetricsSink, + resolveMetricsRuntimeConfig, + type ConfiguredMetricsSinks, + type MetricsEnvLike, +} from "./runtime-config"; + +const METRICS_RECORDER_SYMBOL = Symbol.for( + "thoughtspot.mcp.metrics.requestRecorder", +); + +type MetricsExecutionContext = ExecutionContext & { + [METRICS_RECORDER_SYMBOL]?: MetricsRecorder; +}; + +export function setMetricsRecorderOnExecutionContext( + ctx: ExecutionContext, + recorder: MetricsRecorder, +): MetricsRecorder { + (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL] = recorder; + return recorder; +} + +export function getMetricsRecorderFromExecutionContext( + ctx: ExecutionContext, +): MetricsRecorder | undefined { + return (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL]; +} + +export function clearMetricsRecorderFromExecutionContext( + ctx: ExecutionContext, +): void { + delete (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL]; +} + +export function createRequestMetricsRecorder( + env?: MetricsEnvLike, + sinks: ConfiguredMetricsSinks = {}, +): RequestMetricsRecorder { + const config = resolveMetricsRuntimeConfig(env); + const sink = createConfiguredMetricsSink(config, sinks); + + return new RequestMetricsRecorder({ + sink, + resourceAttributes: config.resourceAttributes, + }); +} + +export async function withRequestMetrics( + env: MetricsEnvLike | undefined, + ctx: ExecutionContext, + handler: (recorder: MetricsRecorder) => Promise, + sinks: ConfiguredMetricsSinks = {}, +): Promise { + const recorder = createRequestMetricsRecorder(env, sinks); + setMetricsRecorderOnExecutionContext(ctx, recorder); + + try { + return await handler(recorder); + } finally { + await recorder.flush(ctx); + clearMetricsRecorderFromExecutionContext(ctx); + } +} diff --git a/src/metrics/runtime/runtime-config.ts b/src/metrics/runtime/runtime-config.ts new file mode 100644 index 0000000..8d63e3e --- /dev/null +++ b/src/metrics/runtime/runtime-config.ts @@ -0,0 +1,145 @@ +import { CompositeMetricsSink } from "./composite-sink"; +import { NoopMetricsSink } from "./noop-sink"; +import type { MetricResourceAttributes, MetricsSink } from "./metrics-sink"; + +export type MetricsSinkMode = "none" | "analytics_engine" | "grafana" | "both"; +export type MetricsDeploymentEnvironment = "production" | "local"; +export type MetricsEnvLike = Partial>; + +export type MetricsRuntimeConfig = { + sinkMode: MetricsSinkMode; + deploymentEnvironment: MetricsDeploymentEnvironment; + resourceAttributes: MetricResourceAttributes; +}; + +export type ConfiguredMetricsSinks = { + analyticsEngineSink?: MetricsSink; + grafanaSink?: MetricsSink; +}; + +function getProcessEnvValue(name: string): string | undefined { + if (typeof process === "undefined") { + return undefined; + } + return process.env?.[name]; +} + +function readConfigValue( + env: MetricsEnvLike | undefined, + ...keys: string[] +): string | undefined { + for (const key of keys) { + const envValue = env?.[key]; + if (typeof envValue === "string" && envValue.length > 0) { + return envValue; + } + + const processEnvValue = getProcessEnvValue(key); + if (processEnvValue && processEnvValue.length > 0) { + return processEnvValue; + } + } + + return undefined; +} + +export function resolveMetricsSinkMode(rawValue?: string): MetricsSinkMode { + switch (rawValue?.trim().toLowerCase()) { + case "none": + return "none"; + case "analytics-engine": + case "analytics_engine": + case "analytics": + return "analytics_engine"; + case "grafana": + return "grafana"; + case "both": + case undefined: + return "both"; + default: + console.warn( + `[metrics] Unknown METRICS_SINK_MODE "${rawValue}", defaulting to "both"`, + ); + return "both"; + } +} + +export function resolveMetricsDeploymentEnvironment( + rawValue?: string, +): MetricsDeploymentEnvironment { + switch (rawValue?.trim().toLowerCase()) { + case "local": + return "local"; + case "production": + case undefined: + return "production"; + default: + console.warn( + `[metrics] Unknown metrics environment "${rawValue}", defaulting to "production"`, + ); + return "production"; + } +} + +export function resolveMetricResourceAttributes( + deploymentEnvironment: MetricsDeploymentEnvironment, + serviceVersion?: string, +): MetricResourceAttributes { + const resourceAttributes: MetricResourceAttributes = { + "service.name": "thoughtspot-mcp-server", + "service.namespace": "thoughtspot", + "deployment.environment": deploymentEnvironment, + "cloud.provider": "cloudflare", + "cloud.platform": "cloudflare_workers", + }; + + if (serviceVersion) { + resourceAttributes["service.version"] = serviceVersion; + } + + return resourceAttributes; +} + +export function resolveMetricsRuntimeConfig( + env?: MetricsEnvLike, +): MetricsRuntimeConfig { + const sinkMode = resolveMetricsSinkMode( + readConfigValue(env, "METRICS_SINK_MODE"), + ); + const deploymentEnvironment = resolveMetricsDeploymentEnvironment( + readConfigValue( + env, + "METRICS_DEPLOYMENT_ENVIRONMENT", + "DEPLOYMENT_ENVIRONMENT", + ), + ); + const serviceVersion = readConfigValue(env, "SERVICE_VERSION", "npm_package_version"); + + return { + sinkMode, + deploymentEnvironment, + resourceAttributes: resolveMetricResourceAttributes( + deploymentEnvironment, + serviceVersion, + ), + }; +} + +export function createConfiguredMetricsSink( + config: Pick, + sinks: ConfiguredMetricsSinks = {}, +): MetricsSink { + switch (config.sinkMode) { + case "none": + return new NoopMetricsSink(); + case "analytics_engine": + return sinks.analyticsEngineSink ?? new NoopMetricsSink(); + case "grafana": + return sinks.grafanaSink ?? new NoopMetricsSink(); + default: + return new CompositeMetricsSink([ + sinks.analyticsEngineSink ?? new NoopMetricsSink(), + sinks.grafanaSink ?? new NoopMetricsSink(), + ]); + } +} diff --git a/src/routes.ts b/src/routes.ts new file mode 100644 index 0000000..0f768e5 --- /dev/null +++ b/src/routes.ts @@ -0,0 +1,46 @@ +export const PUBLIC_ROUTES = { + root: "/", + hello: "/hello", + authorize: "/authorize", + callback: "/callback", + storeToken: "/store-token", + oauthToken: "/token", + register: "/register", + mcp: "/mcp", + sse: "/sse", + openaiMcp: "/openai/mcp", + openaiSse: "/openai/sse", + bearerMcp: "/bearer/mcp", + bearerSse: "/bearer/sse", + tokenMcp: "/token/mcp", + tokenSse: "/token/sse", + openaiAppsChallenge: "/.well-known/openai-apps-challenge", + openapiSpec: "/openapi-spec", +} as const; + +export const PUBLIC_ROUTE_PREFIXES = { + api: "/api", + bearer: "/bearer", + openapiSpec: PUBLIC_ROUTES.openapiSpec, + token: "/token", +} as const; + +export const EXACT_PUBLIC_ROUTES_REQUIRING_METRICS = [ + PUBLIC_ROUTES.root, + PUBLIC_ROUTES.hello, + PUBLIC_ROUTES.authorize, + PUBLIC_ROUTES.callback, + PUBLIC_ROUTES.storeToken, + PUBLIC_ROUTES.oauthToken, + PUBLIC_ROUTES.register, + PUBLIC_ROUTES.mcp, + PUBLIC_ROUTES.sse, + PUBLIC_ROUTES.openaiMcp, + PUBLIC_ROUTES.openaiSse, + PUBLIC_ROUTES.bearerMcp, + PUBLIC_ROUTES.bearerSse, + PUBLIC_ROUTES.tokenMcp, + PUBLIC_ROUTES.tokenSse, + PUBLIC_ROUTES.openaiAppsChallenge, + PUBLIC_ROUTES.openapiSpec, +] as const; diff --git a/test/metrics/runtime/metric-context.spec.ts b/test/metrics/runtime/metric-context.spec.ts new file mode 100644 index 0000000..7fee837 --- /dev/null +++ b/test/metrics/runtime/metric-context.spec.ts @@ -0,0 +1,108 @@ +import { describe, expect, it } from "vitest"; +import { + EXPLICIT_ROUTE_CONTEXTS, + getApiSurface, + getAuthMode, + getRouteGroup, + getStatusClass, + getTransport, + resolvePathMetricContext, + resolveRequestMetricContext, +} from "../../../src/metrics/runtime/metric-context"; +import { + EXACT_PUBLIC_ROUTES_REQUIRING_METRICS, + PUBLIC_ROUTES, + PUBLIC_ROUTE_PREFIXES, +} from "../../../src/routes"; + +describe("metric-context", () => { + it("requires exact public routes to have explicit metric context entries", () => { + for (const pathname of EXACT_PUBLIC_ROUTES_REQUIRING_METRICS) { + expect(EXPLICIT_ROUTE_CONTEXTS).toHaveProperty(pathname); + expect(resolvePathMetricContext(pathname).routeGroup).not.toBe("unknown"); + } + }); + + it("maps explicit request paths through the shared route context table", () => { + for (const [pathname, context] of Object.entries(EXPLICIT_ROUTE_CONTEXTS)) { + expect(resolvePathMetricContext(pathname)).toEqual(context); + expect(getRouteGroup(pathname)).toBe(context.routeGroup); + expect(getTransport(pathname)).toBe(context.transport); + expect(getApiSurface(pathname)).toBe(context.apiSurface); + expect(getAuthMode(pathname)).toBe(context.authMode); + } + }); + + it("maps known grouped request paths to route groups", () => { + expect( + getRouteGroup(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), + ).toBe("api"); + expect( + getRouteGroup(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`), + ).toBe("openapi_spec"); + expect(getRouteGroup("/not-a-route")).toBe("unknown"); + }); + + it("derives transport from fallback request paths", () => { + expect(getTransport("/future/mcp")).toBe("mcp"); + expect(getTransport("/future/sse")).toBe("sse"); + expect(getTransport(PUBLIC_ROUTES.authorize)).toBe("http"); + }); + + it("derives API surface from fallback request paths", () => { + expect(getApiSurface("/openai/future-endpoint")).toBe("openai_mcp"); + expect( + getApiSurface(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), + ).toBe("api"); + expect( + getApiSurface(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`), + ).toBe("static"); + expect(getApiSurface("/bearer/future-endpoint")).toBe("mcp"); + expect(getApiSurface("/token/future-endpoint")).toBe("mcp"); + expect(getApiSurface(PUBLIC_ROUTES.root)).toBe("static"); + expect(getApiSurface(PUBLIC_ROUTES.authorize)).toBe("oauth"); + expect(getApiSurface(PUBLIC_ROUTES.callback)).toBe("oauth"); + expect(getApiSurface(PUBLIC_ROUTES.storeToken)).toBe("oauth"); + expect(getApiSurface("/mystery")).toBe("unknown"); + }); + + it("derives auth mode from fallback request paths", () => { + expect(getAuthMode("/bearer/future-endpoint")).toBe("bearer"); + expect(getAuthMode("/token/future-endpoint")).toBe("token"); + expect(getAuthMode(PUBLIC_ROUTES.openaiMcp)).toBe("oauth"); + expect( + getAuthMode(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), + ).toBe("oauth"); + expect(getAuthMode(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`)).toBe( + "none", + ); + expect(getAuthMode(PUBLIC_ROUTES.root)).toBe("none"); + expect(getAuthMode(PUBLIC_ROUTES.authorize)).toBe("none"); + expect(getAuthMode(PUBLIC_ROUTES.callback)).toBe("none"); + expect(getAuthMode(PUBLIC_ROUTES.storeToken)).toBe("none"); + expect(getAuthMode("/unknown")).toBe("unknown"); + }); + + it("maps response status codes into status classes", () => { + expect(getStatusClass(101)).toBe("1xx"); + expect(getStatusClass(204)).toBe("2xx"); + expect(getStatusClass(302)).toBe("3xx"); + expect(getStatusClass(404)).toBe("4xx"); + expect(getStatusClass(503)).toBe("5xx"); + expect(getStatusClass(99)).toBe("unknown"); + expect(getStatusClass(600)).toBe("unknown"); + }); + + it("resolves the full request metric context from a Request", () => { + const request = new Request( + `https://example.com${PUBLIC_ROUTES.bearerMcp}?api-version=2026-04-23`, + ); + + expect(resolveRequestMetricContext(request)).toEqual({ + routeGroup: "bearer_mcp", + transport: "mcp", + apiSurface: "mcp", + authMode: "bearer", + }); + }); +}); diff --git a/test/metrics/runtime/metric-types.spec.ts b/test/metrics/runtime/metric-types.spec.ts new file mode 100644 index 0000000..a037375 --- /dev/null +++ b/test/metrics/runtime/metric-types.spec.ts @@ -0,0 +1,36 @@ +import { describe, expect, it, vi } from "vitest"; +import { + METRIC_NAMES, + getMetricKind, + normalizeMetricLabels, +} from "../../../src/metrics/runtime/metric-types"; + +describe("metric-types", () => { + it("returns the configured metric kind for known metrics", () => { + expect(getMetricKind(METRIC_NAMES.httpRequestsTotal)).toBe("counter"); + expect(getMetricKind(METRIC_NAMES.httpRequestDurationMs)).toBe("histogram"); + expect(getMetricKind(METRIC_NAMES.httpInflightRequests)).toBe("gauge"); + }); + + it("normalizes labels and drops forbidden or unknown keys", () => { + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); + + const labels = normalizeMetricLabels({ + route_group: "mcp", + outcome: "success", + is_thinking: true, + instanceUrl: "https://tenant.thoughtspot.cloud", + unexpected_key: "value", + tool_name: undefined, + }); + + expect(labels).toEqual({ + is_thinking: "true", + outcome: "success", + route_group: "mcp", + }); + expect(warnSpy).toHaveBeenCalledTimes(2); + }); +}); diff --git a/test/metrics/runtime/metrics-recorder.spec.ts b/test/metrics/runtime/metrics-recorder.spec.ts new file mode 100644 index 0000000..e33c0b9 --- /dev/null +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -0,0 +1,186 @@ +import { describe, expect, it, vi } from "vitest"; +import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; +import { RequestMetricsRecorder } from "../../../src/metrics/runtime/metrics-recorder"; + +describe("RequestMetricsRecorder", () => { + it("records normalized observations and flushes them once", async () => { + const flushSpy = vi.fn().mockResolvedValue(undefined); + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + resourceAttributes: { "service.name": "thoughtspot-mcp-server" }, + now: () => 123, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + instanceUrl: "forbidden", + }); + recorder.histogram(METRIC_NAMES.httpRequestDurationMs, 50, { + outcome: "success", + }); + recorder.gauge(METRIC_NAMES.httpInflightRequests, 2, { + transport: "mcp", + }); + + expect(recorder.snapshot()).toEqual([ + { + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 1, + labels: { route_group: "mcp" }, + timestampMs: 123, + }, + { + kind: "histogram", + name: METRIC_NAMES.httpRequestDurationMs, + value: 50, + labels: { outcome: "success" }, + timestampMs: 123, + }, + { + kind: "gauge", + name: METRIC_NAMES.httpInflightRequests, + value: 2, + labels: { transport: "mcp" }, + timestampMs: 123, + }, + ]); + + await recorder.flush(); + await recorder.flush(); + + expect(warnSpy).toHaveBeenCalledTimes(1); + expect(flushSpy).toHaveBeenCalledTimes(1); + expect(flushSpy).toHaveBeenCalledWith({ + observations: recorder.snapshot(), + resourceAttributes: { "service.name": "thoughtspot-mcp-server" }, + }); + }); + + it("schedules the flush with waitUntil when an execution context is provided", async () => { + const flushSpy = vi.fn().mockResolvedValue(undefined); + const waitUntil = vi.fn(); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + const flushPromise = recorder.flush({ waitUntil } as ExecutionContext); + + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(waitUntil).toHaveBeenCalledWith(flushPromise); + await flushPromise; + expect(flushSpy).toHaveBeenCalledTimes(1); + }); + + it("does not emit observations for the wrong metric kind", () => { + const recorder = new RequestMetricsRecorder({ + sink: { flush: vi.fn().mockResolvedValue(undefined) }, + }); + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + + recorder.histogram(METRIC_NAMES.httpRequestsTotal, 10); + + expect(recorder.snapshot()).toEqual([]); + expect(warnSpy).toHaveBeenCalledOnce(); + }); + + it("swallows sink flush failures", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => undefined); + const recorder = new RequestMetricsRecorder({ + sink: { + flush: vi.fn().mockRejectedValue(new Error("flush failed")), + }, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + + await expect(recorder.flush()).resolves.toBeUndefined(); + expect(errorSpy).toHaveBeenCalled(); + }); + + it("rejects new metrics once a flush has started", async () => { + let resolveFlush!: () => void; + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const flushSpy = vi.fn().mockImplementation( + () => + new Promise((resolve) => { + resolveFlush = resolve; + }), + ); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + const flushPromise = recorder.flush(); + recorder.histogram(METRIC_NAMES.httpRequestDurationMs, 25); + + expect(recorder.snapshot()).toHaveLength(1); + expect(warnSpy).toHaveBeenCalledWith( + `[metrics] Ignoring metric recorded after flush: ${METRIC_NAMES.httpRequestDurationMs}`, + ); + + resolveFlush(); + await flushPromise; + + expect(flushSpy).toHaveBeenCalledWith({ + observations: [ + expect.objectContaining({ + name: METRIC_NAMES.httpRequestsTotal, + kind: "counter", + }), + ], + resourceAttributes: {}, + }); + }); + + it("ignores metrics recorded after the recorder has been flushed", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const flushSpy = vi.fn().mockResolvedValue(undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + await recorder.flush(); + recorder.count(METRIC_NAMES.httpRequestsTotal, 5); + + expect(flushSpy).toHaveBeenCalledTimes(1); + expect(recorder.snapshot()).toHaveLength(1); + expect(warnSpy).toHaveBeenCalledWith( + `[metrics] Ignoring metric recorded after flush: ${METRIC_NAMES.httpRequestsTotal}`, + ); + }); + + it("ignores non-finite metric values", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: vi.fn().mockResolvedValue(undefined) }, + }); + + recorder.histogram(METRIC_NAMES.httpRequestDurationMs, Number.NaN); + recorder.gauge(METRIC_NAMES.httpInflightRequests, Number.POSITIVE_INFINITY); + + expect(recorder.snapshot()).toEqual([]); + expect(warnSpy).toHaveBeenCalledTimes(2); + }); + + it("does not flush empty observations but still closes the recorder", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const flushSpy = vi.fn().mockResolvedValue(undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + await recorder.flush(); + recorder.count(METRIC_NAMES.httpRequestsTotal); + + expect(flushSpy).not.toHaveBeenCalled(); + expect(recorder.snapshot()).toEqual([]); + expect(warnSpy).toHaveBeenCalledWith( + `[metrics] Ignoring metric recorded after flush: ${METRIC_NAMES.httpRequestsTotal}`, + ); + }); +}); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts new file mode 100644 index 0000000..37fca42 --- /dev/null +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -0,0 +1,93 @@ +import { describe, expect, it, vi } from "vitest"; +import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; +import { + clearMetricsRecorderFromExecutionContext, + createRequestMetricsRecorder, + getMetricsRecorderFromExecutionContext, + setMetricsRecorderOnExecutionContext, + withRequestMetrics, +} from "../../../src/metrics/runtime/request-metrics"; + +describe("withRequestMetrics", () => { + it("exposes a request-scoped recorder during handler execution and clears it afterwards", async () => { + const waitUntil = vi.fn(); + const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const ctx = { waitUntil } as ExecutionContext; + + await withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + expect(getMetricsRecorderFromExecutionContext(ctx)).toBe(recorder); + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + }, + { analyticsEngineSink }, + ); + + expect(getMetricsRecorderFromExecutionContext(ctx)).toBeUndefined(); + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); + }); + + it("flushes and clears the request-scoped recorder when the handler throws", async () => { + const waitUntil = vi.fn(); + const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const ctx = { waitUntil } as ExecutionContext; + + await expect( + withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + throw new Error("boom"); + }, + { analyticsEngineSink }, + ), + ).rejects.toThrow("boom"); + + expect(getMetricsRecorderFromExecutionContext(ctx)).toBeUndefined(); + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); + }); + + it("creates a recorder with resolved resource attributes", async () => { + const grafanaSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const recorder = createRequestMetricsRecorder( + { + METRICS_SINK_MODE: "grafana", + METRICS_DEPLOYMENT_ENVIRONMENT: "local", + SERVICE_VERSION: "1.2.3", + }, + { grafanaSink }, + ); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + await recorder.flush(); + + expect(grafanaSink.flush).toHaveBeenCalledWith( + expect.objectContaining({ + resourceAttributes: expect.objectContaining({ + "deployment.environment": "local", + "service.version": "1.2.3", + }), + }), + ); + }); + + it("supports setting and clearing the recorder on the execution context", () => { + const ctx = {} as ExecutionContext; + const recorder = createRequestMetricsRecorder(); + + expect(setMetricsRecorderOnExecutionContext(ctx, recorder)).toBe(recorder); + expect(getMetricsRecorderFromExecutionContext(ctx)).toBe(recorder); + + clearMetricsRecorderFromExecutionContext(ctx); + + expect(getMetricsRecorderFromExecutionContext(ctx)).toBeUndefined(); + }); +}); diff --git a/test/metrics/runtime/runtime-config.spec.ts b/test/metrics/runtime/runtime-config.spec.ts new file mode 100644 index 0000000..f558ace --- /dev/null +++ b/test/metrics/runtime/runtime-config.spec.ts @@ -0,0 +1,129 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { CompositeMetricsSink } from "../../../src/metrics/runtime/composite-sink"; +import { NoopMetricsSink } from "../../../src/metrics/runtime/noop-sink"; +import { + createConfiguredMetricsSink, + resolveMetricResourceAttributes, + resolveMetricsDeploymentEnvironment, + resolveMetricsRuntimeConfig, + resolveMetricsSinkMode, +} from "../../../src/metrics/runtime/runtime-config"; + +describe("runtime-config", () => { + const baseEnv = { ...process.env }; + + afterEach(() => { + vi.stubGlobal("process", { env: { ...baseEnv } }); + vi.restoreAllMocks(); + }); + + it("defaults to both sinks and production attributes", () => { + const config = resolveMetricsRuntimeConfig(); + + expect(config.sinkMode).toBe("both"); + expect(config.deploymentEnvironment).toBe("production"); + expect(config.resourceAttributes["service.name"]).toBe( + "thoughtspot-mcp-server", + ); + expect(config.resourceAttributes["deployment.environment"]).toBe( + "production", + ); + }); + + it("parses supported sink mode aliases", () => { + expect(resolveMetricsSinkMode("analytics-engine")).toBe("analytics_engine"); + expect(resolveMetricsSinkMode("analytics")).toBe("analytics_engine"); + expect(resolveMetricsSinkMode("grafana")).toBe("grafana"); + expect(resolveMetricsSinkMode("none")).toBe("none"); + expect(resolveMetricsSinkMode(" both ")).toBe("both"); + }); + + it("warns and defaults for unknown sink mode and deployment environment", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + + expect(resolveMetricsSinkMode("mystery")).toBe("both"); + expect(resolveMetricsDeploymentEnvironment("qa")).toBe("production"); + expect(warnSpy).toHaveBeenCalledTimes(2); + }); + + it("resolves runtime config from explicit env values", () => { + const config = resolveMetricsRuntimeConfig({ + METRICS_SINK_MODE: "analytics", + METRICS_DEPLOYMENT_ENVIRONMENT: "local", + SERVICE_VERSION: "1.2.3", + }); + + expect(config.sinkMode).toBe("analytics_engine"); + expect(config.deploymentEnvironment).toBe("local"); + expect(config.resourceAttributes).toMatchObject({ + "deployment.environment": "local", + "service.version": "1.2.3", + }); + }); + + it("falls back to alternate env key names", () => { + const config = resolveMetricsRuntimeConfig({ + METRICS_SINK_MODE: "grafana", + DEPLOYMENT_ENVIRONMENT: "local", + npm_package_version: "9.9.9", + }); + + expect(config.sinkMode).toBe("grafana"); + expect(config.deploymentEnvironment).toBe("local"); + expect(config.resourceAttributes["service.version"]).toBe("9.9.9"); + }); + + it("includes service.version only when provided", () => { + expect(resolveMetricResourceAttributes("production")).not.toHaveProperty( + "service.version", + ); + expect( + resolveMetricResourceAttributes("local", "2026.04.23")["service.version"], + ).toBe("2026.04.23"); + }); + + it("creates noop sinks for none and missing single-sink modes", async () => { + const noneSink = createConfiguredMetricsSink({ sinkMode: "none" }); + const analyticsSink = createConfiguredMetricsSink({ + sinkMode: "analytics_engine", + }); + const grafanaSink = createConfiguredMetricsSink({ sinkMode: "grafana" }); + + expect(noneSink).toBeInstanceOf(NoopMetricsSink); + expect(analyticsSink).toBeInstanceOf(NoopMetricsSink); + expect(grafanaSink).toBeInstanceOf(NoopMetricsSink); + await expect( + noneSink.flush({ observations: [], resourceAttributes: {} }), + ).resolves.toBeUndefined(); + }); + + it("returns the provided single sink when configured", () => { + const analyticsEngineSink = { flush: vi.fn() }; + const grafanaSink = { flush: vi.fn() }; + + expect( + createConfiguredMetricsSink( + { sinkMode: "analytics_engine" }, + { analyticsEngineSink }, + ), + ).toBe(analyticsEngineSink); + expect( + createConfiguredMetricsSink({ sinkMode: "grafana" }, { grafanaSink }), + ).toBe(grafanaSink); + }); + + it("creates a composite sink for both mode and tolerates missing sinks", async () => { + const analyticsEngineSink = { flush: vi.fn() }; + const sink = createConfiguredMetricsSink( + { sinkMode: "both" }, + { analyticsEngineSink }, + ); + + await expect( + sink.flush({ observations: [], resourceAttributes: {} }), + ).resolves.toBeUndefined(); + + expect(sink).toBeInstanceOf(CompositeMetricsSink); + expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); + }); +}); diff --git a/test/metrics/runtime/sinks.spec.ts b/test/metrics/runtime/sinks.spec.ts new file mode 100644 index 0000000..53bf090 --- /dev/null +++ b/test/metrics/runtime/sinks.spec.ts @@ -0,0 +1,42 @@ +import { describe, expect, it, vi } from "vitest"; +import { CompositeMetricsSink } from "../../../src/metrics/runtime/composite-sink"; +import { NoopMetricsSink } from "../../../src/metrics/runtime/noop-sink"; + +describe("runtime sinks", () => { + it("flushes every sink in the composite", async () => { + const payload = { observations: [], resourceAttributes: {} }; + const firstSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const secondSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const sink = new CompositeMetricsSink([firstSink, secondSink]); + + await expect(sink.flush(payload)).resolves.toBeUndefined(); + + expect(firstSink.flush).toHaveBeenCalledWith(payload); + expect(secondSink.flush).toHaveBeenCalledWith(payload); + }); + + it("logs rejected sinks but still resolves the composite flush", async () => { + const payload = { observations: [], resourceAttributes: {} }; + const failure = new Error("grafana unavailable"); + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => undefined); + const sink = new CompositeMetricsSink([ + { flush: vi.fn().mockRejectedValue(failure) }, + { flush: vi.fn().mockResolvedValue(undefined) }, + ]); + + await expect(sink.flush(payload)).resolves.toBeUndefined(); + + expect(errorSpy).toHaveBeenCalledWith( + "[metrics] Sink at index 0 failed during flush", + failure, + ); + }); + + it("treats the noop sink as a successful no-op", async () => { + const sink = new NoopMetricsSink(); + + await expect( + sink.flush({ observations: [], resourceAttributes: {} }), + ).resolves.toBeUndefined(); + }); +}); From 0f63d8820c4c98f3d097b36f2e7cb7f9b84cbffc Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Wed, 29 Apr 2026 17:46:37 -0700 Subject: [PATCH 06/25] Format metrics files with Biome (#128) --- src/metrics/mixpanel/mixpanel.ts | 2 +- src/metrics/runtime/metrics-recorder.ts | 4 +-- src/metrics/runtime/request-metrics.ts | 9 ++++-- src/metrics/runtime/runtime-config.ts | 8 ++++-- src/metrics/tracing/tracing-utils.ts | 2 +- test/metrics/mixpanel/integration.spec.ts | 2 +- test/metrics/mixpanel/mixpanel-client.spec.ts | 2 +- .../metrics/mixpanel/mixpanel-tracker.spec.ts | 2 +- test/metrics/runtime/metrics-recorder.spec.ts | 28 ++++++++++++++----- test/metrics/runtime/runtime-config.spec.ts | 4 ++- test/metrics/runtime/sinks.spec.ts | 4 ++- test/metrics/tracing/tracing-utils.spec.ts | 6 ++-- 12 files changed, 49 insertions(+), 24 deletions(-) diff --git a/src/metrics/mixpanel/mixpanel.ts b/src/metrics/mixpanel/mixpanel.ts index e62e654..8e92f09 100644 --- a/src/metrics/mixpanel/mixpanel.ts +++ b/src/metrics/mixpanel/mixpanel.ts @@ -1,6 +1,6 @@ -import { MixpanelClient } from "./mixpanel-client"; import type { SessionInfo } from "../../thoughtspot/types"; import type { Tracker } from "../index"; +import { MixpanelClient } from "./mixpanel-client"; export class MixpanelTracker implements Tracker { private mixpanel: MixpanelClient; diff --git a/src/metrics/runtime/metrics-recorder.ts b/src/metrics/runtime/metrics-recorder.ts index e07a2dc..05e236c 100644 --- a/src/metrics/runtime/metrics-recorder.ts +++ b/src/metrics/runtime/metrics-recorder.ts @@ -1,9 +1,9 @@ import { - getMetricKind, - normalizeMetricLabels, type MetricKind, type MetricLabelInput, type MetricName, + getMetricKind, + normalizeMetricLabels, } from "./metric-types"; import type { MetricObservation, diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 9aaaf3e..9109886 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,9 +1,12 @@ -import { RequestMetricsRecorder, type MetricsRecorder } from "./metrics-recorder"; import { - createConfiguredMetricsSink, - resolveMetricsRuntimeConfig, + type MetricsRecorder, + RequestMetricsRecorder, +} from "./metrics-recorder"; +import { type ConfiguredMetricsSinks, type MetricsEnvLike, + createConfiguredMetricsSink, + resolveMetricsRuntimeConfig, } from "./runtime-config"; const METRICS_RECORDER_SYMBOL = Symbol.for( diff --git a/src/metrics/runtime/runtime-config.ts b/src/metrics/runtime/runtime-config.ts index 8d63e3e..ec4b932 100644 --- a/src/metrics/runtime/runtime-config.ts +++ b/src/metrics/runtime/runtime-config.ts @@ -1,6 +1,6 @@ import { CompositeMetricsSink } from "./composite-sink"; -import { NoopMetricsSink } from "./noop-sink"; import type { MetricResourceAttributes, MetricsSink } from "./metrics-sink"; +import { NoopMetricsSink } from "./noop-sink"; export type MetricsSinkMode = "none" | "analytics_engine" | "grafana" | "both"; export type MetricsDeploymentEnvironment = "production" | "local"; @@ -113,7 +113,11 @@ export function resolveMetricsRuntimeConfig( "DEPLOYMENT_ENVIRONMENT", ), ); - const serviceVersion = readConfigValue(env, "SERVICE_VERSION", "npm_package_version"); + const serviceVersion = readConfigValue( + env, + "SERVICE_VERSION", + "npm_package_version", + ); return { sinkMode, diff --git a/src/metrics/tracing/tracing-utils.ts b/src/metrics/tracing/tracing-utils.ts index 58f5fd6..18f14d9 100644 --- a/src/metrics/tracing/tracing-utils.ts +++ b/src/metrics/tracing/tracing-utils.ts @@ -1,5 +1,5 @@ // tracing-utils.ts -import { type Span, trace, context } from "@opentelemetry/api"; +import { type Span, context, trace } from "@opentelemetry/api"; export function getActiveSpan(spanOverride?: Span): Span | undefined { return spanOverride ?? trace.getSpan(context.active()); diff --git a/test/metrics/mixpanel/integration.spec.ts b/test/metrics/mixpanel/integration.spec.ts index d7503ff..c412838 100644 --- a/test/metrics/mixpanel/integration.spec.ts +++ b/test/metrics/mixpanel/integration.spec.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { MixpanelTracker } from "../../../src/metrics/mixpanel/mixpanel"; import type { SessionInfo } from "../../../src/thoughtspot/types"; diff --git a/test/metrics/mixpanel/mixpanel-client.spec.ts b/test/metrics/mixpanel/mixpanel-client.spec.ts index a33e7b9..9b65146 100644 --- a/test/metrics/mixpanel/mixpanel-client.spec.ts +++ b/test/metrics/mixpanel/mixpanel-client.spec.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { MixpanelClient } from "../../../src/metrics/mixpanel/mixpanel-client"; // Mock fetch globally diff --git a/test/metrics/mixpanel/mixpanel-tracker.spec.ts b/test/metrics/mixpanel/mixpanel-tracker.spec.ts index 823af97..b1467d9 100644 --- a/test/metrics/mixpanel/mixpanel-tracker.spec.ts +++ b/test/metrics/mixpanel/mixpanel-tracker.spec.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { MixpanelTracker } from "../../../src/metrics/mixpanel/mixpanel"; import { MixpanelClient } from "../../../src/metrics/mixpanel/mixpanel-client"; import type { SessionInfo } from "../../../src/thoughtspot/types"; diff --git a/test/metrics/runtime/metrics-recorder.spec.ts b/test/metrics/runtime/metrics-recorder.spec.ts index e33c0b9..183de96 100644 --- a/test/metrics/runtime/metrics-recorder.spec.ts +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -5,7 +5,9 @@ import { RequestMetricsRecorder } from "../../../src/metrics/runtime/metrics-rec describe("RequestMetricsRecorder", () => { it("records normalized observations and flushes them once", async () => { const flushSpy = vi.fn().mockResolvedValue(undefined); - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const recorder = new RequestMetricsRecorder({ sink: { flush: flushSpy }, resourceAttributes: { "service.name": "thoughtspot-mcp-server" }, @@ -78,7 +80,9 @@ describe("RequestMetricsRecorder", () => { const recorder = new RequestMetricsRecorder({ sink: { flush: vi.fn().mockResolvedValue(undefined) }, }); - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); recorder.histogram(METRIC_NAMES.httpRequestsTotal, 10); @@ -87,7 +91,9 @@ describe("RequestMetricsRecorder", () => { }); it("swallows sink flush failures", async () => { - const errorSpy = vi.spyOn(console, "error").mockImplementation(() => undefined); + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); const recorder = new RequestMetricsRecorder({ sink: { flush: vi.fn().mockRejectedValue(new Error("flush failed")), @@ -102,7 +108,9 @@ describe("RequestMetricsRecorder", () => { it("rejects new metrics once a flush has started", async () => { let resolveFlush!: () => void; - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const flushSpy = vi.fn().mockImplementation( () => new Promise((resolve) => { @@ -137,7 +145,9 @@ describe("RequestMetricsRecorder", () => { }); it("ignores metrics recorded after the recorder has been flushed", async () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const flushSpy = vi.fn().mockResolvedValue(undefined); const recorder = new RequestMetricsRecorder({ sink: { flush: flushSpy }, @@ -155,7 +165,9 @@ describe("RequestMetricsRecorder", () => { }); it("ignores non-finite metric values", () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const recorder = new RequestMetricsRecorder({ sink: { flush: vi.fn().mockResolvedValue(undefined) }, }); @@ -168,7 +180,9 @@ describe("RequestMetricsRecorder", () => { }); it("does not flush empty observations but still closes the recorder", async () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const flushSpy = vi.fn().mockResolvedValue(undefined); const recorder = new RequestMetricsRecorder({ sink: { flush: flushSpy }, diff --git a/test/metrics/runtime/runtime-config.spec.ts b/test/metrics/runtime/runtime-config.spec.ts index f558ace..308968b 100644 --- a/test/metrics/runtime/runtime-config.spec.ts +++ b/test/metrics/runtime/runtime-config.spec.ts @@ -39,7 +39,9 @@ describe("runtime-config", () => { }); it("warns and defaults for unknown sink mode and deployment environment", () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const warnSpy = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); expect(resolveMetricsSinkMode("mystery")).toBe("both"); expect(resolveMetricsDeploymentEnvironment("qa")).toBe("production"); diff --git a/test/metrics/runtime/sinks.spec.ts b/test/metrics/runtime/sinks.spec.ts index 53bf090..5ad8e00 100644 --- a/test/metrics/runtime/sinks.spec.ts +++ b/test/metrics/runtime/sinks.spec.ts @@ -18,7 +18,9 @@ describe("runtime sinks", () => { it("logs rejected sinks but still resolves the composite flush", async () => { const payload = { observations: [], resourceAttributes: {} }; const failure = new Error("grafana unavailable"); - const errorSpy = vi.spyOn(console, "error").mockImplementation(() => undefined); + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); const sink = new CompositeMetricsSink([ { flush: vi.fn().mockRejectedValue(failure) }, { flush: vi.fn().mockResolvedValue(undefined) }, diff --git a/test/metrics/tracing/tracing-utils.spec.ts b/test/metrics/tracing/tracing-utils.spec.ts index 8a9058d..e071e4b 100644 --- a/test/metrics/tracing/tracing-utils.spec.ts +++ b/test/metrics/tracing/tracing-utils.spec.ts @@ -1,9 +1,9 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { trace, context } from "@opentelemetry/api"; +import { context, trace } from "@opentelemetry/api"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { + WithSpan, getActiveSpan, withSpan, - WithSpan, withSpanNamed, } from "../../../src/metrics/tracing/tracing-utils"; From a9a4ef3605d3b09010fb9242613bceb74ad5e723 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 30 Apr 2026 10:11:56 -0700 Subject: [PATCH 07/25] Store messages in separate keys in DurableObject KV store (#127) - Each value has a maximum of 131072 bytes, and we are running into that limit in real world scenarios - Instead of storing the full conversation in a single key, split each message into its own key, giving us much higher storage - Update tests Co-authored-by: Rifdhan Nazeer --- src/servers/conversation-storage-server.ts | 145 +++++++++++++----- .../conversation-storage-server.spec.ts | 61 ++++---- 2 files changed, 139 insertions(+), 67 deletions(-) diff --git a/src/servers/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts index 431191d..5b4621c 100644 --- a/src/servers/conversation-storage-server.ts +++ b/src/servers/conversation-storage-server.ts @@ -1,9 +1,12 @@ +import { isBoolean, isNumber } from "lodash"; import type { Message, StreamingMessagesState } from "../thoughtspot/types"; const DEFAULT_TTL_MS = 30 * 60 * 1000; // 30 minutes -const STATE_KEY = "streaming-messages-state"; -const BOOKMARK_KEY = "streaming-messages-bookmark"; +const MESSAGE_KEY_PREFIX = "message-"; +const IS_DONE_KEY = "is-done"; +const WRITE_BOOKMARK_KEY = "write-bookmark"; +const READ_BOOKMARK_KEY = "read-bookmark"; /** * A Durable Object that stores streaming conversation messages and exposes them over HTTP. @@ -18,6 +21,8 @@ const BOOKMARK_KEY = "streaming-messages-bookmark"; * GET /storage//messages —> getNewMessagesAndUpdateBookmark */ export class ConversationStorageServer { + private conversationId = ""; + constructor( private state: DurableObjectState, private env: Env, @@ -29,6 +34,7 @@ export class ConversationStorageServer { // e.g. /storage/abc123/initialize -> /initialize const parts = url.pathname.split("/"); // parts: ["", "storage", "", ""] + this.conversationId = parts[2]; const operation = parts[3] ?? ""; try { @@ -54,76 +60,139 @@ export class ConversationStorageServer { } } catch (err) { const message = err instanceof Error ? err.message : String(err); - console.error("Error handling conversation storage request:", message); + console.error( + `Error handling conversation storage request for conversation ${this.conversationId}:`, + message, + ); return Response.json({ error: "Something went wrong" }, { status: 500 }); } } /* * Initialize the conversation. This can be a brand new conversation, or it can be priming an - * existing conversation which is already marked done for a followup message. + * existing conversation which is already marked done for a followup message. We never delete + * messages in the conversation, instead the next messages begin at the existing bookmark. */ private async initializeConversation(): Promise { - const existing = - await this.state.storage.get(STATE_KEY); - if (existing && !existing.isDone) { - throw new Error("Conversation already exists and is not marked done"); + const existingIsDone = await this.state.storage.get(IS_DONE_KEY); + if (isBoolean(existingIsDone) && !existingIsDone) { + throw new Error( + `Conversation ${this.conversationId} already exists and is not marked done`, + ); } - await this.setStateAndRestartTtl({ messages: [], isDone: false }); - await this.state.storage.put(BOOKMARK_KEY, 0); + await this.state.storage.put(IS_DONE_KEY, false); + await this.restartTtl(); } + /* + * Append new messages to the conversation, starting at the current state of WRITE_BOOKMARK and + * saving the new state of WRITE_BOOKMARK after. We write the isDone flag state after writing + * all messages, so that if a reader is executing concurrently, they will never think the + * conversation is done without having already seen all messages. We also restart the TTL for + * the conversation after all writes are done. + */ private async appendMessagesAndRestartTtl( newMessages: Message[], isDone = false, ): Promise { - const oldState = - await this.state.storage.get(STATE_KEY); - if (!oldState) { - throw new Error("Conversation not found"); + const existingIsDone = await this.state.storage.get(IS_DONE_KEY); + if (!isBoolean(existingIsDone)) { + throw new Error(`Conversation ${this.conversationId} not found`); } - if (oldState.isDone) { - throw new Error("Cannot append messages to a conversation marked done"); + if (existingIsDone) { + throw new Error( + `Cannot append messages to conversation ${this.conversationId} marked done`, + ); } - await this.setStateAndRestartTtl({ - messages: [...oldState.messages, ...newMessages], - isDone, - }); + let idx = (await this.state.storage.get(WRITE_BOOKMARK_KEY)) ?? 0; + for (const message of newMessages) { + await this.state.storage.put( + `${MESSAGE_KEY_PREFIX}${idx}`, + message, + ); + idx++; + } + await this.state.storage.put(WRITE_BOOKMARK_KEY, idx); + + if (isDone) { + await this.state.storage.put(IS_DONE_KEY, true); + } + + await this.restartTtl(); } + /* + * Retrieve all new messages since the last time this was called. We use a READ_BOOKMARK to + * track the index of the last returned message, and update it when returning new messages. We + * read the isDone flag state before reading messages, so that if a writer is executing + * concurrently, we will only see isDone=true if all messages have already been written. Note + * that we don't restart the TTL here, since it is only meant to be based on writes. + */ private async getNewMessagesAndUpdateBookmark(): Promise { - const bookmark = (await this.state.storage.get(BOOKMARK_KEY)) ?? 0; - - const conversationState = - await this.state.storage.get(STATE_KEY); - if (!conversationState) { - throw new Error("Conversation not found"); + const isDone = await this.state.storage.get(IS_DONE_KEY); + if (!isBoolean(isDone)) { + throw new Error(`Conversation ${this.conversationId} not found`); } - await this.state.storage.put( - BOOKMARK_KEY, - conversationState.messages.length, - ); + let bookmark = + (await this.state.storage.get(READ_BOOKMARK_KEY)) ?? 0; + const newMessages: Message[] = []; + while (true) { + const message = await this.state.storage.get( + `${MESSAGE_KEY_PREFIX}${bookmark}`, + ); + if (!message) { + break; + } + newMessages.push(message); + bookmark++; + } + await this.state.storage.put(READ_BOOKMARK_KEY, bookmark); return { - messages: conversationState.messages.slice(bookmark), - isDone: conversationState.isDone, + messages: newMessages, + isDone, }; } - private async setStateAndRestartTtl( - newState: StreamingMessagesState, - ): Promise { + private async restartTtl(): Promise { // Cancel any existing alarm and schedule a fresh one await this.state.storage.deleteAlarm(); await this.state.storage.setAlarm(Date.now() + DEFAULT_TTL_MS); - - await this.state.storage.put(STATE_KEY, newState); } async alarm(): Promise { - await this.state.storage.delete([STATE_KEY, BOOKMARK_KEY]); + // Check for any abnormalities in the state prior to deleting + const isDone = await this.state.storage.get(IS_DONE_KEY); + if (!isBoolean(isDone) || !isDone) { + console.warn( + `Conversation ${this.conversationId} expired without being marked done`, + { + isDone, + }, + ); + } + const writeBookmark = + await this.state.storage.get(WRITE_BOOKMARK_KEY); + const readBookmark = + await this.state.storage.get(READ_BOOKMARK_KEY); + if (!isNumber(writeBookmark)) { + console.warn( + `Conversation ${this.conversationId} expired without any messages written`, + ); + } else if (!isNumber(readBookmark) || writeBookmark !== readBookmark) { + console.warn( + `Conversation ${this.conversationId} expired with unread messages`, + { + writeBookmark, + readBookmark, + }, + ); + } + + // Delete everything in storage + await this.state.storage.deleteAll(); } } diff --git a/test/servers/conversation-storage-server.spec.ts b/test/servers/conversation-storage-server.spec.ts index e86999c..34b25b6 100644 --- a/test/servers/conversation-storage-server.spec.ts +++ b/test/servers/conversation-storage-server.spec.ts @@ -36,6 +36,9 @@ function createMockStorage() { deleteAlarm: vi.fn(async (): Promise => { alarm = null; }), + deleteAll: vi.fn(async (): Promise => { + store.clear(); + }), }, }; } @@ -59,14 +62,23 @@ function makeRequest( } // Sample messages -const textMessage: Message = { type: "text", text: "Hello" }; -const chunkMessage: Message = { type: "text_chunk", text: " world" }; +const textMessage: Message = { + type: "text", + text: "Hello", + is_thinking: false, +}; +const chunkMessage: Message = { + type: "text_chunk", + text: " world", + is_thinking: false, +}; const answerMessage: Message = { type: "answer", answer_id: "ans-1", answer_title: "My Answer", answer_query: "SELECT 1", iframe_url: "https://example.com/answer/1", + is_thinking: false, }; // --------------------------------------------------------------------------- @@ -112,16 +124,17 @@ describe("ConversationStorageServer", () => { it("stores empty messages and isDone=false", async () => { await server.fetch(makeRequest("POST", "initialize")); - const state = mock.store.get( - "streaming-messages-state", - ) as StreamingMessagesState; - expect(state).toMatchObject({ messages: [], isDone: false }); + expect(mock.store.get("is-done")).toBe(false); + // No message keys should exist yet + expect(mock.store.has("message-0")).toBe(false); }); it("sets bookmark to 0", async () => { await server.fetch(makeRequest("POST", "initialize")); - expect(mock.store.get("streaming-messages-bookmark")).toBe(0); + // write-bookmark and read-bookmark are lazily initialised to 0 + expect(mock.store.get("write-bookmark") ?? 0).toBe(0); + expect(mock.store.get("read-bookmark") ?? 0).toBe(0); }); it("schedules a TTL alarm", async () => { @@ -152,10 +165,7 @@ describe("ConversationStorageServer", () => { const res = await server.fetch(makeRequest("POST", "initialize")); expect(res.status).toBe(200); - const state = mock.store.get( - "streaming-messages-state", - ) as StreamingMessagesState; - expect(state).toMatchObject({ messages: [], isDone: false }); + expect(mock.store.get("is-done")).toBe(false); }); }); @@ -182,10 +192,8 @@ describe("ConversationStorageServer", () => { makeRequest("POST", "append", { messages: [textMessage] }), ); - const state = mock.store.get( - "streaming-messages-state", - ) as StreamingMessagesState; - expect(state.messages).toEqual([textMessage]); + expect(mock.store.get("message-0")).toEqual(textMessage); + expect(mock.store.get("write-bookmark")).toBe(1); }); it("accumulates messages across multiple calls", async () => { @@ -198,14 +206,10 @@ describe("ConversationStorageServer", () => { }), ); - const state = mock.store.get( - "streaming-messages-state", - ) as StreamingMessagesState; - expect(state.messages).toEqual([ - textMessage, - chunkMessage, - answerMessage, - ]); + expect(mock.store.get("message-0")).toEqual(textMessage); + expect(mock.store.get("message-1")).toEqual(chunkMessage); + expect(mock.store.get("message-2")).toEqual(answerMessage); + expect(mock.store.get("write-bookmark")).toBe(3); }); it("marks the conversation done when isDone is true", async () => { @@ -216,10 +220,7 @@ describe("ConversationStorageServer", () => { }), ); - const state = mock.store.get( - "streaming-messages-state", - ) as StreamingMessagesState; - expect(state.isDone).toBe(true); + expect(mock.store.get("is-done")).toBe(true); }); it("restarts the TTL alarm on each call", async () => { @@ -345,8 +346,10 @@ describe("ConversationStorageServer", () => { it("deletes the conversation state and bookmark", async () => { await server.alarm(); - expect(mock.store.has("streaming-messages-state")).toBe(false); - expect(mock.store.has("streaming-messages-bookmark")).toBe(false); + expect(mock.store.has("is-done")).toBe(false); + expect(mock.store.has("message-0")).toBe(false); + expect(mock.store.has("write-bookmark")).toBe(false); + expect(mock.store.has("read-bookmark")).toBe(false); }); it("causes subsequent append to return 500", async () => { From 13ee175e2cd7f57449f70adb1b70a892a3c47f2b Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Thu, 30 Apr 2026 15:01:21 -0700 Subject: [PATCH 08/25] SCAL-309983 Harden metrics request isolation (#129) * SCAL-309983 Harden metrics request isolation ## Summary Harden the metrics request wrapper so metrics setup, scheduling, and flush failures cannot block MCP business logic or become visible to MCP clients. ## What Changed - fall back to a noop sink if metrics config or sink construction throws - keep `recorder.flush()` focused on metrics delivery, not Worker scheduling - schedule request metrics flush through `ctx.waitUntil(...)` without awaiting delivery - log waitUntil scheduling failures separately from metrics execution failures - add tests for setup fallback, slow flushes, handler failures, waitUntil failures, and delivery failures ## Semantics Direct `recorder.flush()` calls still await delivery. `withRequestMetrics(...)` only schedules the request-scoped flush so slow, failing, or unavailable sinks do not delay responses. Metrics failures remain best-effort side effects and are never propagated to MCP clients. ## Validation - `npm run lint` - focused runtime metrics tests - `npm test` * SCAL-309983 Return noop recorder on init fallback ## Summary Use a stateless noop metrics recorder for initialization fallback so the fail-open path does not construct a second request recorder. ## What Changed - add a shared `NOOP_METRICS_RECORDER` implementation - return the noop recorder when metrics config or sink setup throws - update tests and log wording to reflect the recorder-level fallback --- src/metrics/runtime/metrics-recorder.ts | 23 ++- src/metrics/runtime/request-metrics.ts | 46 ++++- test/metrics/runtime/metrics-recorder.spec.ts | 21 ++- test/metrics/runtime/request-metrics.spec.ts | 178 +++++++++++++++++- 4 files changed, 247 insertions(+), 21 deletions(-) diff --git a/src/metrics/runtime/metrics-recorder.ts b/src/metrics/runtime/metrics-recorder.ts index 05e236c..ff4421e 100644 --- a/src/metrics/runtime/metrics-recorder.ts +++ b/src/metrics/runtime/metrics-recorder.ts @@ -21,10 +21,25 @@ export interface MetricsRecorder { count(name: MetricName, value?: number, labels?: MetricLabelInput): void; histogram(name: MetricName, value: number, labels?: MetricLabelInput): void; gauge(name: MetricName, value: number, labels?: MetricLabelInput): void; - flush(ctx?: Pick): Promise; + flush(): Promise; snapshot(): readonly MetricObservation[]; } +const NOOP_FLUSH_PROMISE: Promise = Promise.resolve(); +const NOOP_METRIC_OBSERVATIONS: readonly MetricObservation[] = []; + +export const NOOP_METRICS_RECORDER: MetricsRecorder = { + count(_name, _value, _labels): void {}, + histogram(_name, _value, _labels): void {}, + gauge(_name, _value, _labels): void {}, + flush(): Promise { + return NOOP_FLUSH_PROMISE; + }, + snapshot(): readonly MetricObservation[] { + return NOOP_METRIC_OBSERVATIONS; + }, +}; + export class RequestMetricsRecorder implements MetricsRecorder { private readonly observations: MetricObservation[] = []; private flushPromise?: Promise; @@ -48,15 +63,11 @@ export class RequestMetricsRecorder implements MetricsRecorder { return [...this.observations]; } - async flush(ctx?: Pick): Promise { + flush(): Promise { if (!this.flushPromise) { this.flushPromise = this.flushInternal(); } - if (ctx) { - ctx.waitUntil(this.flushPromise); - } - return this.flushPromise; } diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 9109886..000f1f4 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,5 +1,6 @@ import { type MetricsRecorder, + NOOP_METRICS_RECORDER, RequestMetricsRecorder, } from "./metrics-recorder"; import { @@ -40,14 +41,43 @@ export function clearMetricsRecorderFromExecutionContext( export function createRequestMetricsRecorder( env?: MetricsEnvLike, sinks: ConfiguredMetricsSinks = {}, -): RequestMetricsRecorder { - const config = resolveMetricsRuntimeConfig(env); - const sink = createConfiguredMetricsSink(config, sinks); +): MetricsRecorder { + try { + const config = resolveMetricsRuntimeConfig(env); + const sink = createConfiguredMetricsSink(config, sinks); - return new RequestMetricsRecorder({ - sink, - resourceAttributes: config.resourceAttributes, - }); + return new RequestMetricsRecorder({ + sink, + resourceAttributes: config.resourceAttributes, + }); + } catch (error) { + console.error( + "[metrics] Failed to initialize request metrics recorder; using noop recorder", + error, + ); + return NOOP_METRICS_RECORDER; + } +} + +function scheduleRequestMetricsFlush( + recorder: MetricsRecorder, + ctx: ExecutionContext, +): void { + let flushPromise: Promise; + try { + flushPromise = recorder.flush().catch((error) => { + console.error("[metrics] Failed to execute request metrics flush", error); + }); + } catch (error) { + console.error("[metrics] Failed to execute request metrics flush", error); + return; + } + + try { + ctx.waitUntil(flushPromise); + } catch (error) { + console.error("[metrics] Failed to schedule request metrics flush", error); + } } export async function withRequestMetrics( @@ -62,7 +92,7 @@ export async function withRequestMetrics( try { return await handler(recorder); } finally { - await recorder.flush(ctx); + scheduleRequestMetricsFlush(recorder, ctx); clearMetricsRecorderFromExecutionContext(ctx); } } diff --git a/test/metrics/runtime/metrics-recorder.spec.ts b/test/metrics/runtime/metrics-recorder.spec.ts index 183de96..41c4e4c 100644 --- a/test/metrics/runtime/metrics-recorder.spec.ts +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -60,19 +60,28 @@ describe("RequestMetricsRecorder", () => { }); }); - it("schedules the flush with waitUntil when an execution context is provided", async () => { + it("returns the same in-flight flush promise when called repeatedly", async () => { + let resolveFlush!: () => void; const flushSpy = vi.fn().mockResolvedValue(undefined); - const waitUntil = vi.fn(); + flushSpy.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFlush = resolve; + }), + ); const recorder = new RequestMetricsRecorder({ sink: { flush: flushSpy }, }); recorder.count(METRIC_NAMES.httpRequestsTotal); - const flushPromise = recorder.flush({ waitUntil } as ExecutionContext); + const firstFlushPromise = recorder.flush(); + const secondFlushPromise = recorder.flush(); - expect(waitUntil).toHaveBeenCalledTimes(1); - expect(waitUntil).toHaveBeenCalledWith(flushPromise); - await flushPromise; + expect(secondFlushPromise).toBe(firstFlushPromise); + expect(flushSpy).toHaveBeenCalledTimes(1); + + resolveFlush(); + await firstFlushPromise; expect(flushSpy).toHaveBeenCalledTimes(1); }); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index 37fca42..530ec58 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; import { clearMetricsRecorderFromExecutionContext, @@ -9,6 +9,10 @@ import { } from "../../../src/metrics/runtime/request-metrics"; describe("withRequestMetrics", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + it("exposes a request-scoped recorder during handler execution and clears it afterwards", async () => { const waitUntil = vi.fn(); const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; @@ -31,6 +35,40 @@ describe("withRequestMetrics", () => { expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); }); + it("does not await slow metrics flushes on the request path", async () => { + let resolveFlush!: () => void; + const waitUntil = vi.fn(); + const analyticsEngineSink = { + flush: vi.fn( + () => + new Promise((resolve) => { + resolveFlush = resolve; + }), + ), + }; + const ctx = { waitUntil } as unknown as ExecutionContext; + + const result = await withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + return "handler-result"; + }, + { analyticsEngineSink }, + ); + + expect(result).toBe("handler-result"); + expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); + expect(waitUntil).toHaveBeenCalledTimes(1); + expect(getMetricsRecorderFromExecutionContext(ctx)).toBeUndefined(); + + resolveFlush(); + await waitUntil.mock.calls[0][0]; + }); + it("flushes and clears the request-scoped recorder when the handler throws", async () => { const waitUntil = vi.fn(); const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; @@ -55,6 +93,38 @@ describe("withRequestMetrics", () => { expect(analyticsEngineSink.flush).toHaveBeenCalledTimes(1); }); + it("does not mask handler failures with slow metrics flushes", async () => { + let resolveFlush!: () => void; + const waitUntil = vi.fn(); + const analyticsEngineSink = { + flush: vi.fn( + () => + new Promise((resolve) => { + resolveFlush = resolve; + }), + ), + }; + const ctx = { waitUntil } as unknown as ExecutionContext; + + await expect( + withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + throw new Error("boom"); + }, + { analyticsEngineSink }, + ), + ).rejects.toThrow("boom"); + + expect(waitUntil).toHaveBeenCalledTimes(1); + resolveFlush(); + await waitUntil.mock.calls[0][0]; + }); + it("creates a recorder with resolved resource attributes", async () => { const grafanaSink = { flush: vi.fn().mockResolvedValue(undefined) }; const recorder = createRequestMetricsRecorder( @@ -79,6 +149,112 @@ describe("withRequestMetrics", () => { ); }); + it("falls back to a noop recorder when metrics config resolution throws", async () => { + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const env = Object.defineProperty({}, "METRICS_SINK_MODE", { + get() { + throw new Error("env unavailable"); + }, + }); + const recorder = createRequestMetricsRecorder(env); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + await expect(recorder.flush()).resolves.toBeUndefined(); + expect(recorder.snapshot()).toEqual([]); + + expect(errorSpy).toHaveBeenCalledWith( + "[metrics] Failed to initialize request metrics recorder; using noop recorder", + expect.any(Error), + ); + }); + + it("falls back to a noop recorder when sink construction throws", async () => { + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const sinks = Object.defineProperty({}, "analyticsEngineSink", { + get() { + throw new Error("sink unavailable"); + }, + }); + const recorder = createRequestMetricsRecorder( + { METRICS_SINK_MODE: "analytics_engine" }, + sinks, + ); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + await expect(recorder.flush()).resolves.toBeUndefined(); + expect(recorder.snapshot()).toEqual([]); + + expect(errorSpy).toHaveBeenCalledWith( + "[metrics] Failed to initialize request metrics recorder; using noop recorder", + expect.any(Error), + ); + }); + + it("does not fail the request when waitUntil rejects scheduling", async () => { + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const ctx = { + waitUntil: vi.fn(() => { + throw new Error("waitUntil unavailable"); + }), + } as unknown as ExecutionContext; + const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; + + const result = await withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + return "handler-result"; + }, + { analyticsEngineSink }, + ); + + expect(result).toBe("handler-result"); + expect(errorSpy).toHaveBeenCalledWith( + "[metrics] Failed to schedule request metrics flush", + expect.any(Error), + ); + }); + + it("does not fail the request when metrics delivery fails", async () => { + const errorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const waitUntil = vi.fn(); + const ctx = { waitUntil } as unknown as ExecutionContext; + const analyticsEngineSink = { + flush: vi.fn().mockRejectedValue(new Error("metrics unavailable")), + }; + + const result = await withRequestMetrics( + { METRICS_SINK_MODE: "analytics_engine" }, + ctx, + async (recorder) => { + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + return "handler-result"; + }, + { analyticsEngineSink }, + ); + + expect(result).toBe("handler-result"); + expect(waitUntil).toHaveBeenCalledTimes(1); + await waitUntil.mock.calls[0][0]; + expect(errorSpy).toHaveBeenCalledWith( + "[metrics] Flush failed", + expect.any(Error), + ); + }); + it("supports setting and clearing the recorder on the execution context", () => { const ctx = {} as ExecutionContext; const recorder = createRequestMetricsRecorder(); From cc80fcc9ad2b978be40f63ca3510d5ebb9388d2a Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 30 Apr 2026 17:34:56 -0700 Subject: [PATCH 09/25] Improve error handling for getSessionInfo (#130) - If getSessionInfo call fails, handle the exception and log error - Add test coverage Co-authored-by: Rifdhan Nazeer --- src/servers/mcp-server-base.ts | 16 +++++++++------ test/servers/mcp-server-base.spec.ts | 29 ++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 106e3ef..bd6b8af 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -191,12 +191,16 @@ export abstract class BaseMCPServer extends Server { } protected async initializeService(): Promise { - this.sessionInfo = await this.getThoughtSpotService().getSessionInfo(); - const mixpanel = new MixpanelTracker( - this.sessionInfo, - this.ctx.props.clientName, - ); - this.addTracker(mixpanel); + try { + this.sessionInfo = await this.getThoughtSpotService().getSessionInfo(); + const mixpanel = new MixpanelTracker( + this.sessionInfo, + this.ctx.props.clientName, + ); + this.addTracker(mixpanel); + } catch (error) { + console.error("Error initializing session info:", error); + } } /** diff --git a/test/servers/mcp-server-base.spec.ts b/test/servers/mcp-server-base.spec.ts index d7c5fea..20b9c9e 100644 --- a/test/servers/mcp-server-base.spec.ts +++ b/test/servers/mcp-server-base.spec.ts @@ -278,6 +278,35 @@ describe("MCP Server Base", () => { }); }); + describe("initializeService", () => { + it("should catch and log error if getSessionInfo throws", async () => { + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi + .fn() + .mockRejectedValue(new Error("Session info fetch failed")), + searchMetadata: vi.fn().mockResolvedValue([]), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); + + const consoleErrorSpy = vi + .spyOn(console, "error") + .mockImplementation(() => {}); + + const testServer = new TestMCPServer( + { props: mockProps } as any, + null as any, + ); + await expect(testServer.init()).resolves.not.toThrow(); + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Error initializing session info:", + expect.any(Error), + ); + + consoleErrorSpy.mockRestore(); + }); + }); + describe("Server Initialization", () => { it("should initialize with custom server name and version", () => { const customServer = new TestMCPServer( From 9e49d1e6ca568a63a20cecd8cf72fe298c506784 Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Fri, 1 May 2026 10:15:55 -0700 Subject: [PATCH 10/25] Stabilize nanoid integration fetch mock (#131) ## Summary Make the nanoid integration tests independent of global `fetch` state so the full Vitest run is stable. ## What Changed - replace the module-level `global.fetch = vi.fn()` assignment with a per-test fetch mock - stub and unstub `fetch` in `beforeEach` and `afterEach` - keep the existing nanoid assertions unchanged while removing order dependence ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts` - `npm test` --- ...ughtspot-client-nanoid-integration.spec.ts | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts b/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts index 41d262e..4cff999 100644 --- a/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts +++ b/test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts @@ -7,12 +7,12 @@ * carry a unique, well-formed ID. */ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { getThoughtSpotClient } from "../../src/thoughtspot/thoughtspot-client"; import { - createBearerAuthenticationConfig, ThoughtSpotRestApi, + createBearerAuthenticationConfig, } from "@thoughtspot/rest-api-sdk"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { getThoughtSpotClient } from "../../src/thoughtspot/thoughtspot-client"; // Mock only the SDK plumbing — nanoid is intentionally NOT mocked so that the // real customAlphabet implementation is exercised. @@ -21,8 +21,6 @@ vi.mock("@thoughtspot/rest-api-sdk", () => ({ ThoughtSpotRestApi: vi.fn(), })); -global.fetch = vi.fn(); - const CUSTOM_ALPHABET = "_-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; const NANO_ID_SIZE = 12; @@ -31,6 +29,8 @@ const ALLOWED_CHARS_RE = /^[_\-0-9a-zA-Z]+$/; const INSTANCE_URL = "https://integration-test.thoughtspot.com"; const BEARER_TOKEN = "integration-test-token"; +let fetchMock: ReturnType; + function buildClient() { const mockConfig = { middleware: [] }; const mockClient: Record = { instanceUrl: INSTANCE_URL }; @@ -42,13 +42,11 @@ function buildClient() { } function mockOkFetch() { - (fetch as any).mockResolvedValue({ ok: true }); + fetchMock.mockResolvedValue({ ok: true }); } function parsedBodies(): any[] { - return (fetch as any).mock.calls.map((call: any[]) => - JSON.parse(call[1].body), - ); + return fetchMock.mock.calls.map((call: any[]) => JSON.parse(call[1].body)); } describe("sendAgentConversationMessageStreaming — nano ID integration", () => { @@ -56,10 +54,13 @@ describe("sendAgentConversationMessageStreaming — nano ID integration", () => beforeEach(() => { vi.clearAllMocks(); + fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); client = buildClient(); }); afterEach(() => { + vi.unstubAllGlobals(); vi.restoreAllMocks(); }); @@ -161,7 +162,7 @@ describe("sendAgentConversationMessageStreaming — nano ID integration", () => message: userMessage, }); - const [url, options] = (fetch as any).mock.calls[0]; + const [url, options] = fetchMock.mock.calls[0]; const body = JSON.parse(options.body); // Endpoint construction From a82d9f78778d90e27a6a115ff12346eb499020ac Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Fri, 1 May 2026 16:48:50 -0700 Subject: [PATCH 11/25] Release 2026-05-01 version (#132) - Add new entry in version registry for 2026-05-01 - Add new "latest" version option, mapping to newest available non-beta version - Remove "default" version, default will now be latest - Add "backwards-compatibility-default" version which points to older version for now, will be removed in the future - Add metrics tracking of versioning - Update tests Co-authored-by: Rifdhan Nazeer --- src/bearer.ts | 26 +++--- src/index.ts | 20 +++-- src/servers/mcp-server.ts | 23 +++++- src/servers/version-registry.ts | 97 ++++++++++------------ test/bearer.spec.ts | 16 ++-- test/servers/mcp-server.spec.ts | 60 ++++++-------- test/servers/version-registry.spec.ts | 112 ++++++++++++++------------ 7 files changed, 186 insertions(+), 168 deletions(-) diff --git a/src/bearer.ts b/src/bearer.ts index fdae0e9..e1c28db 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -9,15 +9,15 @@ import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; * @param env - Environment bindings * @param ctx - Execution context * @param MCPServer - MCP server instance - * @param supportApiVersion - Whether to support api-version query param + * @param apiVersionOverride - Optional API version override (ignore value in request) */ function handleTokenAuth( req: Request, env: Env, ctx: ExecutionContext, MCPServer: typeof ThoughtSpotMCP, - supportApiVersion: boolean, -): Response { + apiVersionOverride?: string, +): Response | Promise { const authHeader = req.headers.get("authorization"); if (!authHeader) { return new Response("Bearer token is required", { status: 400 }); @@ -51,12 +51,10 @@ function handleTokenAuth( clientName, }; - // Add api-version support only for /token endpoints (supports "beta" or "YYYY-MM-DD" format) - if (supportApiVersion) { - const apiVersion = url.searchParams.get("api-version"); - if (apiVersion) { - props.apiVersion = apiVersion; - } + // Resolve API version to use + const apiVersion = apiVersionOverride ?? url.searchParams.get("api-version"); + if (apiVersion) { + props.apiVersion = apiVersion; } (ctx as any).props = props; @@ -81,13 +79,19 @@ export function withBearerHandler( // These endpoints do NOT support api-version query params (will be removed in future) // Use /token endpoints instead for new implementations app.mount(PUBLIC_ROUTE_PREFIXES.bearer, (req, env, ctx) => { - return handleTokenAuth(req, env, ctx, MCPServer, false); + return handleTokenAuth( + req, + env, + ctx, + MCPServer, + "backwards-compatibility-default", + ); }); // NEW: /token endpoints - supports api-version query params // Recommended for all new implementations app.mount(PUBLIC_ROUTE_PREFIXES.token, (req, env, ctx) => { - return handleTokenAuth(req, env, ctx, MCPServer, true); + return handleTokenAuth(req, env, ctx, MCPServer); }); return app; diff --git a/src/index.ts b/src/index.ts index 7ea4c88..9b9e550 100644 --- a/src/index.ts +++ b/src/index.ts @@ -51,17 +51,21 @@ function createMCPRouter( ctx: ExecutionContext, ): Promise { const url = new URL(request.url); - const apiVersion = url.searchParams.get("api-version"); + let apiVersion = url.searchParams.get("api-version"); - // Inject apiVersion into props if provided (supports "beta" or "YYYY-MM-DD" format) - if (apiVersion) { - const originalProps = (ctx as any).props || {}; - (ctx as any).props = { - ...originalProps, - apiVersion, - }; + // TODO(Rifdhan): this is a temporary backwards compatibility measure. In the future + // we will use latest by default. + if (!apiVersion) { + apiVersion = "backwards-compatibility-default"; } + // Inject apiVersion into props + const originalProps = (ctx as any).props || {}; + (ctx as any).props = { + ...originalProps, + apiVersion, + }; + // Route to the appropriate serve method return serverClass[serveMethod](path, options).fetch(request, env, ctx); }, diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index cd56861..9684f8b 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -9,7 +9,7 @@ import { TrackEvent } from "../metrics"; import { WithSpan } from "../metrics/tracing/tracing-utils"; import { context, SpanStatusCode, trace } from "@opentelemetry/api"; import { BaseMCPServer, type Context } from "./mcp-server-base"; -import { resolveApiVersion } from "./version-registry"; +import { resolveApiVersion, type VersionConfig } from "./version-registry"; import { GetRelevantQuestionsSchema, GetAnswerSchema, @@ -32,9 +32,28 @@ export class MCPServer extends BaseMCPServer { super(ctx, "ThoughtSpot", "2.0.0"); } + @WithSpan("call-list-tools") protected async listTools() { + const span = this.initSpanWithCommonAttributes(); + span?.setAttribute( + "api_version_requested", + this.ctx.props.apiVersion ?? "(not passed)", + ); + // Resolve the API version to get the appropriate tool configuration - const versionConfig = resolveApiVersion(this.ctx.props.apiVersion); + let versionConfig: VersionConfig; + try { + versionConfig = resolveApiVersion(this.ctx.props.apiVersion); + } catch (error) { + console.error("Error resolving API version, using default:", error); + span?.recordException(error as Error); + versionConfig = resolveApiVersion(); + } + span?.setAttribute( + "api_version_resolved", + // The plain date will be the last entry if multiple labels + versionConfig.version[versionConfig.version.length - 1], + ); // Get base tools from version config let tools = [...versionConfig.tools]; diff --git a/src/servers/version-registry.ts b/src/servers/version-registry.ts index 0be8760..5374f98 100644 --- a/src/servers/version-registry.ts +++ b/src/servers/version-registry.ts @@ -1,5 +1,7 @@ import { toolDefinitionsV1, toolDefinitionsV2 } from "./tool-definitions"; +const YYYY_MM_DD_DATE_REGEX = /^\d{4}-\d{2}-\d{2}$/; + /** * Version configuration interface */ @@ -13,26 +15,23 @@ export interface VersionConfig { } /** - * Parses the release date from a version's date identifiers - * @param versions - The version identifiers array (e.g., ["beta", "2026-03-25"]) - * @returns The parsed date from the first valid YYYY-MM-DD string in the array, or null if none found + * Given a list of version names, returns a date if there is a valid date-based version identifier. + * This should always be the last entry in the list if present. */ function getReleaseDateFromVersion(versions: string[]): Date | null { - const dateVersion = versions.find((v) => /^\d{4}-\d{2}-\d{2}$/.test(v)); - if (dateVersion) { - return new Date(dateVersion); + const possibleDateVersion = versions[versions.length - 1]; + const isDate = YYYY_MM_DD_DATE_REGEX.test(possibleDateVersion); + if (!isDate) { + return null; } - return null; + return new Date(possibleDateVersion); } /** - * Version registry mapping version identifiers to their configurations - * - * IMPORTANT: Versions MUST be ordered by release date (newest first). - * When adding new versions, ensure they are inserted in the correct position - * to maintain this ordering, as the resolution logic depends on it. - * There should always be a "default" stable version that serves as the fallback for - * any requests without a specified version or with a date before all versions. + * Version registry, with respective tools to expose by API version. The "version" field can + * contain multiple identifiers for the same version. Important ordering rules: + * - Entries in the registry must be in chronological order, with newest first + * - Entries in the "version" field must include the plain date (if present) as the last entry */ export const VERSION_REGISTRY: VersionConfig[] = [ { @@ -41,39 +40,22 @@ export const VERSION_REGISTRY: VersionConfig[] = [ description: "Spotter3 agent conversation tools released", }, { - version: ["default", "2025-01-01"], + version: ["latest", "2026-05-01"], + tools: [...toolDefinitionsV2], + description: "Spotter3 agent conversation tools released", + }, + { + version: ["backwards-compatibility-default", "2025-01-01"], tools: [...toolDefinitionsV1], description: "Base version with getRelevantQuestions and getAnswer tools", }, ]; -function getDefaultVersionConfig() { - const defaultVersion = VERSION_REGISTRY.find((v) => - v.version.includes("default"), - ); - if (defaultVersion) { - return defaultVersion; - } - if (VERSION_REGISTRY.length > 0) { - return VERSION_REGISTRY[0]; - } - throw new Error("No available API versions in registry."); -} - /** - * Resolves an API version string to a version configuration - * @param apiVersion - The API version string (e.g., "beta", "2025-03-01", or null for default) - * @returns The resolved version configuration + * Resolves an API version string to a version configuration, defaulting to latest */ -export function resolveApiVersion( - apiVersion: string | null | undefined, -): VersionConfig { - // No version specified - return the entry marked as "default" - if (!apiVersion) { - return getDefaultVersionConfig(); - } - - // Check for exact match (including "beta") +export function resolveApiVersion(apiVersion = "latest"): VersionConfig { + // Check for exact match (including non-dates like "beta", "latest", etc) const exactMatch = VERSION_REGISTRY.find((v) => v.version.includes(apiVersion), ); @@ -81,31 +63,36 @@ export function resolveApiVersion( return exactMatch; } - // Try to parse as date (YYYY-MM-DD format) - const dateMatch = apiVersion.match(/^(\d{4})-(\d{2})-(\d{2})$/); - if (!dateMatch) { - return getDefaultVersionConfig(); + // Try to parse as date + const isDate = YYYY_MM_DD_DATE_REGEX.test(apiVersion); + if (!isDate) { + throw new Error( + `Invalid date format in API version, expected YYYY-MM-DD: ${apiVersion}`, + ); } const requestedDate = new Date(apiVersion); - - // Validate the date is valid if (Number.isNaN(requestedDate.getTime())) { - return getDefaultVersionConfig(); + throw new Error( + `Invalid date format in API version, expected YYYY-MM-DD: ${apiVersion}`, + ); } - // Find the latest version released on or before the requested date - // Entries without a date identifier are excluded from date-based resolution - // Note: No sort needed as VERSION_REGISTRY is already ordered by release date (newest first) - const matchingVersion = VERSION_REGISTRY.filter((v) => { + // Find the newest version on or before the requested date. Note that the version registry is + // already ordered from newest to oldest. We ignore any entries without a date-based version. + const matchingVersion = VERSION_REGISTRY.find((v) => { const releaseDate = getReleaseDateFromVersion(v.version); return releaseDate !== null && releaseDate <= requestedDate; - })[0]; - + }); if (matchingVersion) { return matchingVersion; } - // If no version found on or before the date, return the default entry - return getDefaultVersionConfig(); + // If requesting an API version older than the oldest available version, return the oldest + // available version + console.warn( + "Requested API version is older than all available versions, defaulting to oldest available version", + apiVersion, + ); + return VERSION_REGISTRY[VERSION_REGISTRY.length - 1]; } diff --git a/test/bearer.spec.ts b/test/bearer.spec.ts index b64276a..d351f91 100644 --- a/test/bearer.spec.ts +++ b/test/bearer.spec.ts @@ -177,6 +177,7 @@ describe("Bearer Handler", () => { accessToken: "my-access-token", instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Custom Test Client", + apiVersion: "backwards-compatibility-default", }); // Verify the response @@ -201,6 +202,7 @@ describe("Bearer Handler", () => { accessToken: "my-access-token", instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Bearer Token client", + apiVersion: "backwards-compatibility-default", }); // Verify the response @@ -479,8 +481,8 @@ describe("Bearer Handler", () => { }); }); - describe("DEPRECATED: /bearer endpoints - No API Version Support", () => { - it("should NOT inject apiVersion even when query param is present on /bearer/mcp", async () => { + describe("DEPRECATED: /bearer endpoints - Fixed API Version Override", () => { + it("should use backwards-compatibility-default apiVersion and ignore query param on /bearer/mcp", async () => { const appWithBearer = withBearerHandler(app, ThoughtSpotMCP); const request = new Request( @@ -494,15 +496,15 @@ describe("Bearer Handler", () => { await appWithBearer.fetch(request, mockEnv, mockCtx); - // LEGACY: /bearer endpoints do NOT support api-version for backward compatibility + // LEGACY: /bearer endpoints always use backwards-compatibility-default, ignoring any query param expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", }); - expect(mockCtx.props.apiVersion).toBeUndefined(); + expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); }); - it("should NOT inject apiVersion even when query param is present on /bearer/sse", async () => { + it("should use backwards-compatibility-default apiVersion and ignore query param on /bearer/sse", async () => { const appWithBearer = withBearerHandler(app, ThoughtSpotMCP); const request = new Request( @@ -516,12 +518,12 @@ describe("Bearer Handler", () => { await appWithBearer.fetch(request, mockEnv, mockCtx); - // LEGACY: /bearer endpoints do NOT support api-version + // LEGACY: /bearer endpoints always use backwards-compatibility-default, ignoring any query param expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", }); - expect(mockCtx.props.apiVersion).toBeUndefined(); + expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); }); }); diff --git a/test/servers/mcp-server.spec.ts b/test/servers/mcp-server.spec.ts index f1e8ac5..39249bd 100644 --- a/test/servers/mcp-server.spec.ts +++ b/test/servers/mcp-server.spec.ts @@ -178,14 +178,14 @@ describe("MCP Server", () => { const result = await listTools(); - // Common tools (2) + Standard tools (2) + DataSourceDiscovery (1) = 5 + // V2 tools (latest version): 5 tools expect(result.tools).toHaveLength(5); expect(result.tools?.map((t) => t.name)).toEqual([ - "ping", - "createLiveboard", - "getDataSourceSuggestions", - "getRelevantQuestions", - "getAnswer", + "check_connectivity", + "create_analysis_session", + "send_session_message", + "get_session_updates", + "create_dashboard", ]); }); @@ -195,32 +195,27 @@ describe("MCP Server", () => { const result = await listTools(); - const pingTool = result.tools?.find((t) => t.name === "ping"); - expect(pingTool?.description).toBe( - "Simple ping tool to test connectivity and Auth", + const connectivityTool = result.tools?.find( + (t) => t.name === "check_connectivity", ); - - const questionsTool = result.tools?.find( - (t) => t.name === "getRelevantQuestions", - ); - expect(questionsTool?.description).toBe( - "Get relevant data questions from ThoughtSpot database", + expect(connectivityTool?.description).toBe( + "Ping tool to test connectivity and authentication. This can be used if other tool calls are failing to verify if the connection is working.", ); - const answerTool = result.tools?.find((t) => t.name === "getAnswer"); - expect(answerTool?.description).toBe( - "Get the answer to a question from ThoughtSpot database", + const sessionTool = result.tools?.find( + (t) => t.name === "create_analysis_session", ); + expect(sessionTool).toBeDefined(); - const liveboardTool = result.tools?.find( - (t) => t.name === "createLiveboard", + const dashboardTool = result.tools?.find( + (t) => t.name === "create_dashboard", ); - expect(liveboardTool?.description).toBe( - "Create a liveboard from a list of answers", + expect(dashboardTool?.description).toBe( + "Create a dashboard from a list of answers, allowing the user to revisit the results later. Use this if the user asks for a dashboard, or asks to save the results from the analysis.", ); }); - it("should return 4 tools when enableSpotterDataSourceDiscovery is false", async () => { + it("should return 5 tools regardless of enableSpotterDataSourceDiscovery when using latest (V2)", async () => { // Mock getThoughtSpotClient with enableSpotterDataSourceDiscovery set to false vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ getSessionInfo: vi.fn().mockResolvedValue({ @@ -256,20 +251,15 @@ describe("MCP Server", () => { const result = await listTools(); - // Common tools (2) + Standard tools (2) = 4 (no datasource discovery) - expect(result.tools).toHaveLength(4); + // V2 tools don't have a datasource discovery tool, so filtering has no effect + expect(result.tools).toHaveLength(5); expect(result.tools?.map((t) => t.name)).toEqual([ - "ping", - "createLiveboard", - "getRelevantQuestions", - "getAnswer", + "check_connectivity", + "create_analysis_session", + "send_session_message", + "get_session_updates", + "create_dashboard", ]); - - // Verify that getDataSourceSuggestions is not included - const dataSourceTool = result.tools?.find( - (t) => t.name === "getDataSourceSuggestions", - ); - expect(dataSourceTool).toBeUndefined(); }); }); diff --git a/test/servers/version-registry.spec.ts b/test/servers/version-registry.spec.ts index f370b07..fa60659 100644 --- a/test/servers/version-registry.spec.ts +++ b/test/servers/version-registry.spec.ts @@ -16,14 +16,13 @@ function isValidApiVersion(apiVersion: string): boolean { describe("Version Registry", () => { describe("resolveApiVersion", () => { - it("should return default version when no apiVersion is provided", () => { - const result = resolveApiVersion(null); - expect(result.version).toContain("default"); + it("should throw for null apiVersion", () => { + expect(() => resolveApiVersion(null as any)).toThrow(); }); - it("should return default version when undefined is provided", () => { + it("should return latest version when undefined is provided", () => { const result = resolveApiVersion(undefined); - expect(result.version).toContain("default"); + expect(result.version).toContain("latest"); }); it("should return beta version when 'beta' is specified", () => { @@ -31,6 +30,12 @@ describe("Version Registry", () => { expect(result.version).toContain("beta"); }); + it("should return latest stable version when 'latest' is specified", () => { + const result = resolveApiVersion("latest"); + expect(result.version).toContain("latest"); + expect(result.version).not.toContain("beta"); + }); + it("should resolve exact date match", () => { const result = resolveApiVersion("2025-01-01"); expect(result.version).toContain("2025-01-01"); @@ -46,10 +51,10 @@ describe("Version Registry", () => { expect(result.version).toContain("2025-01-01"); }); - it("should return default version when date is before all versions", () => { + it("should return oldest version when date is before all versions", () => { // Use a date before the earliest version in VERSION_REGISTRY const result = resolveApiVersion("2020-01-01"); - expect(result.version).toContain("default"); + expect(result.version).toContain("backwards-compatibility-default"); }); it("should exclude beta from date-based resolution", () => { @@ -58,25 +63,23 @@ describe("Version Registry", () => { expect(result.version).not.toContain("beta"); }); - it("should return default version for invalid date format", () => { - const result = resolveApiVersion("invalid-date"); - expect(result.version).toContain("default"); + it("should throw for invalid date format", () => { + expect(() => resolveApiVersion("invalid-date")).toThrow(); }); - it("should return default version for malformed date", () => { - const result = resolveApiVersion("2025-13-01"); - expect(result.version).toContain("default"); + it("should throw for malformed date", () => { + expect(() => resolveApiVersion("2025-13-01")).toThrow(); }); - it("should return default version for partial date", () => { - const result = resolveApiVersion("2025-03"); - expect(result.version).toContain("default"); + it("should throw for partial date", () => { + expect(() => resolveApiVersion("2025-03")).toThrow(); }); it("should handle future dates", () => { const result = resolveApiVersion("2030-01-01"); - // Only dated entry is 2025-01-01, so future dates resolve to it - expect(result.version).toContain("2025-01-01"); + // Resolves to the latest dated stable entry + expect(result.version).toContain("latest"); + expect(result.version).not.toContain("beta"); }); }); @@ -89,16 +92,16 @@ describe("Version Registry", () => { expect(isValidApiVersion("2025-03-01")).toBe(true); }); - it("should return true for invalid format (falls back to default)", () => { - expect(isValidApiVersion("invalid")).toBe(true); + it("should return false for invalid format (throws)", () => { + expect(isValidApiVersion("invalid")).toBe(false); }); - it("should return true for malformed date (falls back to default)", () => { - expect(isValidApiVersion("2025-13-45")).toBe(true); + it("should return false for malformed date (throws)", () => { + expect(isValidApiVersion("2025-13-45")).toBe(false); }); - it("should return true for partial date (falls back to default)", () => { - expect(isValidApiVersion("2025-03")).toBe(true); + it("should return false for partial date (throws)", () => { + expect(isValidApiVersion("2025-03")).toBe(false); }); }); @@ -106,7 +109,8 @@ describe("Version Registry", () => { it("should return all version identifiers", () => { const versions = VERSION_REGISTRY.flatMap((v) => v.version); expect(versions).toContain("beta"); - expect(versions).toContain("default"); + expect(versions).toContain("backwards-compatibility-default"); + expect(versions).toContain("latest"); expect(versions).toContain("2025-01-01"); expect(versions.length).toBeGreaterThan(0); }); @@ -170,9 +174,10 @@ describe("Version Registry", () => { describe("Date range resolution", () => { // Registry: // ["beta"] → Spotter3 tools (no date, only reachable via "beta") - // ["default", "2025-01-01"] → base MCP tools + // ["latest", "2026-05-01"] → newest stable entry + // ["backwards-compatibility-default", "2025-01-01"] → base MCP tools it.each([ - // Before all dated versions → falls back to default (2025-01-01) + // Before all dated versions → falls back to oldest (2025-01-01) { apiVersion: "2020-01-01", expectedVersionDate: "2025-01-01", @@ -183,7 +188,7 @@ describe("Version Registry", () => { expectedVersionDate: "2025-01-01", label: "day before earliest version", }, - // On or after 2025-01-01 → resolves to 2025-01-01 (only dated entry) + // On or after 2025-01-01 but before 2026-05-01 → resolves to 2025-01-01 { apiVersion: "2025-01-01", expectedVersionDate: "2025-01-01", @@ -204,14 +209,20 @@ describe("Version Registry", () => { expectedVersionDate: "2025-01-01", label: "end of 2025", }, + // On or after 2026-05-01 → resolves to latest stable (2026-05-01) { - apiVersion: "2026-01-01", - expectedVersionDate: "2025-01-01", + apiVersion: "2026-05-01", + expectedVersionDate: "2026-05-01", + label: "exact match for latest version", + }, + { + apiVersion: "2026-06-01", + expectedVersionDate: "2026-05-01", label: "start of 2026", }, { apiVersion: "2030-01-01", - expectedVersionDate: "2025-01-01", + expectedVersionDate: "2026-05-01", label: "far future date", }, ])( @@ -230,9 +241,14 @@ describe("Version Registry", () => { label: "beta identifier", }, { - apiVersion: "default", - expectedIdentifier: "default", - label: "default identifier", + apiVersion: "backwards-compatibility-default", + expectedIdentifier: "backwards-compatibility-default", + label: "backwards-compatibility-default identifier", + }, + { + apiVersion: "latest", + expectedIdentifier: "latest", + label: "latest identifier", }, ])( "$label: apiVersion=$apiVersion → contains $expectedIdentifier", @@ -242,27 +258,23 @@ describe("Version Registry", () => { }, ); - // Invalid versions → fall back to default + // Invalid versions → throw it.each([ { apiVersion: "beta2", label: "unknown identifier" }, { apiVersion: "invalid", label: "invalid string" }, { apiVersion: "2025-13-01", label: "malformed date" }, - ])( - "$label: apiVersion=$apiVersion → falls back to default", - ({ apiVersion }) => { - const result = resolveApiVersion(apiVersion); - expect(result.version).toContain("default"); - }, - ); + ])("$label: apiVersion=$apiVersion → throws", ({ apiVersion }) => { + expect(() => resolveApiVersion(apiVersion)).toThrow(); + }); - // Null/undefined → default - it.each([ - { apiVersion: null, label: "null" }, - { apiVersion: undefined, label: "undefined" }, - ])("$label → resolves to default (2025-01-01)", ({ apiVersion }) => { - const result = resolveApiVersion(apiVersion); - expect(result.version).toContain("default"); - expect(result.version).toContain("2025-01-01"); + // Null → throws; undefined → resolves to latest + it("null → throws", () => { + expect(() => resolveApiVersion(null as any)).toThrow(); + }); + + it("undefined → resolves to latest", () => { + const result = resolveApiVersion(undefined); + expect(result.version).toContain("latest"); }); }); }); From 866f98685219bfc606d726302032a0be2dc1ce6a Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Fri, 1 May 2026 17:44:29 -0700 Subject: [PATCH 12/25] Update README for 2026-05-01 release (#133) Co-authored-by: Rifdhan Nazeer --- README.md | 38 +++++++++++++++++--------------------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 0682c87..a1c91ad 100644 --- a/README.md +++ b/README.md @@ -40,19 +40,18 @@ See our privacy statement [here](https://www.thoughtspot.com/privacy-statement). ## Connect -If using a client which supports remote MCPs natively (Claude.ai etc) then just enter: +As of May 1, 2026, the ThoughtSpot MCP Server supports Spotter 3, enabling advanced analytics, forecasting, multi-step reasoning, and deep research. The new MCP tools support real-time streaming and session-based conversations. -MCP Server URL: +The server uses date-based versioning via the `?api-version=YYYY-MM-DD` query parameter. Appending this to your MCP server URL pins your integration to a specific version (or the closest earlier version if an exact match doesn't exist). Each new version may introduce new tools, enhancements, or bug fixes. -``` -https://agent.thoughtspot.app/mcp -``` -Preferred Auth method: Oauth + OAuth apps (plug-and-play integrations for Claude, ChatGPT, or custom OAuth apps): +- https://agent.thoughtspot.app/mcp?api-version=latest — Always points to the latest version. +- https://agent.thoughtspot.app/mcp?api-version=YYYY-MM-DD — Pins to a specific version. Example: `?api-version=2026-04-30`. +- https://agent.thoughtspot.app/mcp — Not recommended. Currently points to the baseline version with legacy MCP tools. -- For OpenAI ChatGPT Deep Research, add the URL as: -```js -https://agent.thoughtspot.app/openai/mcp -``` + Bearer token apps (custom apps using bearer tokens): +- https://agent.thoughtspot.app/token/mcp — Always points to the latest version. +- https://agent.thoughtspot.app/token/mcp?api-version=YYYY-MM-DD — Pins to a specific version. Example: `?api-version=2026-04-30`. To configure this MCP server in your MCP client (such as Claude Desktop, Windsurf, Cursor, etc.) which do not support remote MCPs, add the following configuration to your MCP client settings: @@ -63,7 +62,7 @@ To configure this MCP server in your MCP client (such as Claude Desktop, Windsur "command": "npx", "args": [ "mcp-remote", - "https://agent.thoughtspot.app/mcp" + "https://agent.thoughtspot.app/mcp?api-version=latest" ] } } @@ -391,28 +390,25 @@ The ThoughtSpot MCP Server supports API versioning to access different tool sets **Version Formats:** - **Beta version**: `?api-version=beta` - Access the latest beta features - **Date-based version**: `?api-version=YYYY-MM-DD` - Access tools from a specific release date or the latest version on or before that date -- **Default** (no parameter): Returns the stable default tool set +- **Default** (no parameter): Returns the latest stable tool set **Examples:** ```bash # Beta version (latest experimental features) https://agent.thoughtspot.app/token/mcp?api-version=beta -# Specific date version (Spotter3 agent tools) -https://agent.thoughtspot.app/token/mcp?api-version=2025-03-01 +# Specific date version (Spotter 3 agent tools) +https://agent.thoughtspot.app/token/mcp?api-version=2026-05-01 -# Date range resolution (returns latest version ≤ specified date) +# Date range resolution (returns newest version ≤ specified date) https://agent.thoughtspot.app/token/mcp?api-version=2025-03-15 - -# Default version (stable tools) -https://agent.thoughtspot.app/token/mcp ``` **Available Versions:** - `beta`: Latest beta features with Spotter3 agent conversation tools -- `2025-03-01`: Spotter3 agent conversation tools (`createConversation`, `sendConversationMessage`, `getConversationUpdates`) -- `2024-12-01`: Base MCP tools (`getRelevantQuestions`, `getAnswer`, `getDataSourceSuggestions`) -- `default`: Stable base tools (same as `2024-12-01`) +- `latest`: Most recent non-beta version +- `2026-05-01`: Spotter 3 agent conversation tools (`create_analysis_session`, `send_session_message`, `get_session_updates`) +- `2025-01-01`: Base MCP tools (`getRelevantQuestions`, `getAnswer`, `getDataSourceSuggestions`) **Note:** The `/bearer/*` endpoints always return the default stable tool set and ignore the `api-version` parameter for backward compatibility. From 5306811109c411c6c33211b9f4cfe4cb1a7db85e Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Mon, 4 May 2026 15:14:11 -0700 Subject: [PATCH 13/25] SCAL-308279 Add Grafana OTLP metrics sink (#126) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * SCAL-308279 Add Grafana OTLP metrics sink Add the direct Grafana OTLP/HTTP sink for low-cardinality MCP service metrics. This keeps PR 3 independent from the Analytics Engine sink so it can be reviewed and merged in parallel. - add a Grafana OTLP metrics sink that emits OTLP JSON over HTTP - convert counters, gauges, and histograms into OTLP metric datapoints - normalize OTLP endpoints to `/v1/metrics` - support Grafana Cloud Basic auth and explicit OTLP auth headers - wire request metrics to build the Grafana sink from env config - test payload mapping, endpoint config, auth headers, and failure behavior The sink sends delta sums and histograms because request-scoped observations are emitted as per-flush deltas. Export failures are surfaced to the recorder, where the existing metrics flush isolation logs and drops them without changing business responses. - `npm run lint` - focused runtime metrics tests for Grafana sink and request metrics * SCAL-308279 Aggregate OTLP metric datapoints ## Summary Address review feedback for the Grafana OTLP payload shape and timestamp precision. OTLP requires each metric to contain at most one datapoint for a given attribute set. ## What Changed - aggregate same-metric and same-label observations before creating OTLP datapoints - sum counters and histogram count, sum, and bucket counts per attribute set - keep the latest gauge value per attribute set - convert millisecond timestamps to nanoseconds with `BigInt` precision - add regression coverage for duplicate attributes and fractional millisecond timestamps ## Validation - `npm run lint` - focused runtime metrics tests - `npm test` * SCAL-308279 Avoid counter bucket allocation ## Summary Avoid allocating histogram bucket storage for non-histogram OTLP datapoints. ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/grafana-otlp-sink.spec.ts` * SCAL-308279 Harden OTLP auth and timestamps ## Summary Address Grafana OTLP review feedback for Basic auth encoding, timestamp conversion, and export error messages. ## What Changed - encode Basic auth input as UTF-8 before base64 conversion when Buffer is unavailable - split millisecond timestamps into integer and fractional parts before BigInt conversion - keep timestamp export at microsecond precision to avoid unsafe epoch-scale float math - cap Grafana OTLP export error bodies included in exception messages - add regression coverage for non-ASCII credentials and large export error bodies ## Notes The attribute aggregation key remains JSON-based because request-scoped flushes are small and attributes are already sorted, making the key stable. A cheaper custom key can be added later if this sink is reused for high-volume background aggregation. ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308279 Cache Grafana OTLP sink per env Avoid rebuilding the default Grafana OTLP sink for every request in a warm Worker isolate. - add a module-level WeakMap keyed by the Worker env object - reuse the Grafana sink for repeated recorder creation with the same env - keep caller-provided sinks unchanged for tests and explicit injection - add coverage for same-env cache reuse The cache is opportunistic and per Worker isolate. It does not provide correctness guarantees, and cold starts or new isolates will rebuild the sink. - `npm run lint` - focused runtime metrics tests * SCAL-308279 Avoid histogram bucket copies ## Summary Avoid per-observation histogram bucket array allocation during OTLP aggregation. ## What Changed - replace per-observation bucket count arrays with a bucket-index helper - increment the target aggregate bucket directly - preserve existing histogram aggregation semantics ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308279 Refine OTLP timestamp and helper naming ## Summary Tighten the Grafana OTLP sink helpers around timestamp conversion, naming clarity, and error-body truncation. ## What Changed - rename OTLP attribute and authorization helper functions for clearer intent - remove redundant string unions from OTLP attribute helper types - convert fractional millisecond timestamps directly to nanoseconds with carry handling - enforce the export error-body length limit including the trailing ellipsis - add regression coverage for fractional-millisecond precision and truncation behavior ## Validation - > @thoughtspot/mcp-server@0.5.0 lint > biome lint Checked 72 files in 44ms. No fixes applied. - > @thoughtspot/mcp-server@0.5.0 test > vitest run --coverage.enabled=false test/metrics/runtime/grafana-otlp-sink.spec.ts test/metrics/runtime/request-metrics.spec.ts test/metrics/runtime/runtime-config.spec.ts test/metrics/runtime/sinks.spec.ts RUN v3.2.4 /Users/kanwarbir.bajwa/mcp-server [vpw:info] Starting single runtime for vitest.config.ts... ✓ test/metrics/runtime/request-metrics.spec.ts (12 tests) 8ms ✓ test/metrics/runtime/runtime-config.spec.ts (9 tests) 2ms ✓ test/metrics/runtime/grafana-otlp-sink.spec.ts (12 tests) 3ms ✓ test/metrics/runtime/sinks.spec.ts (3 tests) 1ms Test Files 4 passed (4) Tests 36 passed (36) Start at 09:19:17 Duration 1.43s (transform 193ms, setup 124ms, collect 85ms, tests 14ms, environment 0ms, prepare 253ms) [vpw:debug] Shutting down runtimes... - > @thoughtspot/mcp-server@0.5.0 test > vitest run RUN v3.2.4 /Users/kanwarbir.bajwa/mcp-server Coverage enabled with istanbul [vpw:info] Starting single runtime for vitest.config.ts... ✓ test/index.header-stripping.spec.ts (5 tests) 2473ms ✓ Header stripping > strips traceparent before the request reaches the OAuth provider 527ms ✓ Header stripping > strips tracestate before the request reaches the OAuth provider 499ms ✓ Header stripping > strips all tracing headers while preserving others 487ms ✓ Header stripping > does not mutate the original request object 478ms ✓ Header stripping > passes the request through unchanged when no tracing headers are present 479ms ✓ test/index.spec.ts (4 tests) 1947ms ✓ The ThoughtSpot MCP Worker: Auth handler > responds with Hello World! on '/hello' 488ms ✓ MCP Router with API Version > should create router with correct serve method for /mcp 478ms ✓ MCP Router with API Version > should create router with correct serve method for /sse 496ms ✓ MCP Router with API Version > should handle query parameters in router paths 484ms stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should return relevant questions for a query [DEBUG] Getting relevant questions for datasource: [ 'ds-123', 'ds-456' ] [DEBUG] Getting relevant questions with datasource: [ 'ds-123', 'ds-456' ] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should handle error from service [DEBUG] Getting relevant questions for datasource: [ 'ds-123' ] [DEBUG] Getting relevant questions with datasource: [ 'ds-123' ] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should handle error from service with multiple datasource IDs [DEBUG] Getting relevant questions for datasource: [ 'ds-123', 'ds-456', 'ds-789' ] [DEBUG] Getting relevant questions with datasource: [ 'ds-123', 'ds-456', 'ds-789' ] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should handle error from service with empty datasourceIds array [DEBUG] Getting relevant questions for datasource: [] [DEBUG] Getting relevant questions with datasource: [] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should handle empty questions response [DEBUG] Getting relevant questions for datasource: [ 'ds-123' ] [DEBUG] Getting relevant questions with datasource: [ 'ds-123' ] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Relevant Questions Tool > should handle optional additional context [DEBUG] Getting relevant questions for datasource: [ 'ds-123' ] [DEBUG] Getting relevant questions with datasource: [ 'ds-123' ] stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Answer Tool > should return answer for a question [DEBUG] Getting answer for sourceId: ds-123 shouldGetTML: false stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Answer Tool > should return answer for a question [DEBUG] Getting Data for session_identifier: session-123 generation_number: 1 instanceUrl: https://test.thoughtspot.cloud stdout | test/servers/mcp-server.spec.ts > MCP Server > Get Answer Tool > should handle error from service [DEBUG] Getting answer for sourceId: ds-123 shouldGetTML: false stdout | test/servers/mcp-server.spec.ts > MCP Server > Create Liveboard Tool > should create liveboard successfully [DEBUG] Getting TML for answer: What is the total revenue? stdout | test/servers/mcp-server.spec.ts > MCP Server > Create Liveboard Tool > should handle error from service [DEBUG] Getting TML for answer: What is the total revenue? ✓ test/servers/mcp-server.spec.ts (36 tests) 79ms stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://company.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fcompany.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://team.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fteam.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://my.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fmy.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://teamabc.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fteamabc.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://myabc.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fmyabc.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://team1test.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fteam1test.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://my1test.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fmy1test.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://test-team1.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Ftest-team1.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://test-my1.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Ftest-my1.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://team1.test.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fteam1.test.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://my1.test.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fmy1.test.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://team123.thoughtspot.com/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fteam123.thoughtspot.com%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl https://my123.thoughtspot.com/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttps%253A%252F%252Fmy123.thoughtspot.com%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl http://team1.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttp%253A%252F%252Fteam1.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D stdout | test/handlers.spec.ts > Handlers > POST /authorize > Instance URL regex pattern matching > should NOT redirect to callback for URLs that don't match the pattern redirectUrl http://my1.thoughtspot.cloud/callosum/v1/saml/login?targetURLPath=https%3A%2F%2Fexample.com%2Fcallback%3FinstanceUrl%3Dhttp%253A%252F%252Fmy1.thoughtspot.cloud%26oauthReqInfo%3DeyJjbGllbnRJZCI6InRlc3QifQ%253D%253D ✓ test/handlers.spec.ts (34 tests | 5 skipped) 45ms ✓ test/thoughtspot/thoughtspot-client.spec.ts (37 tests) 21ms ✓ test/metrics/mixpanel/mixpanel-tracker.spec.ts (25 tests) 27ms ✓ test/utils.spec.ts (41 tests) 20ms ✓ test/streaming-utils.spec.ts (17 tests) 11ms ✓ test/metrics/mixpanel/integration.spec.ts (13 tests) 10ms ✓ test/servers/openai-mcp-server.spec.ts (19 tests) 10ms ✓ test/storage-service/storage-service.spec.ts (17 tests) 10ms ✓ test/metrics/mixpanel/mixpanel-client.spec.ts (21 tests) 9ms ✓ test/thoughtspot/thoughtspot-client-nanoid-integration.spec.ts (6 tests) 6ms ✓ test/metrics/tracing/tracing-utils.spec.ts (17 tests) 9ms ✓ test/servers/conversation-storage-server.spec.ts (25 tests) 7ms ✓ test/servers/mcp-server-base.spec.ts (15 tests) 8ms ✓ test/metrics/runtime/request-metrics.spec.ts (12 tests) 11ms ✓ test/bearer.spec.ts (36 tests) 7ms ✓ test/servers/api-server.spec.ts (13 tests) 10ms stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getRelevantQuestions > should return relevant questions successfully [DEBUG] Getting relevant questions with datasource: [ 'ws1', 'ws2' ] stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getRelevantQuestions > should handle empty additional context [DEBUG] Getting relevant questions with datasource: [ 'ws1' ] stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getRelevantQuestions > should handle null additional context [DEBUG] Getting relevant questions with datasource: [ 'ws1' ] stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getRelevantQuestions > should handle missing decomposedQueryResponse [DEBUG] Getting relevant questions with datasource: [ 'ws1' ] stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getRelevantQuestions > should handle API errors [DEBUG] Getting relevant questions with datasource: [ 'ws1' ] stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should return answer data successfully without TML [DEBUG] Getting answer for sourceId: ws1 shouldGetTML: false stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should return answer data successfully without TML [DEBUG] Getting Data for session_identifier: session123 generation_number: 1 instanceUrl: https://test.thoughtspot.com stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should return answer data with TML when requested [DEBUG] Getting answer for sourceId: ws1 shouldGetTML: true stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should return answer data with TML when requested [DEBUG] Getting Data for session_identifier: session123 generation_number: 1 instanceUrl: https://test.thoughtspot.com [DEBUG] Getting TML for answer: What is the revenue? stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should limit CSV data to 100 lines [DEBUG] Getting answer for sourceId: ws1 shouldGetTML: false stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should limit CSV data to 100 lines [DEBUG] Getting Data for session_identifier: session123 generation_number: 1 instanceUrl: https://test.thoughtspot.com stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should handle TML export errors gracefully [DEBUG] Getting answer for sourceId: ws1 shouldGetTML: true stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should handle TML export errors gracefully [DEBUG] Getting Data for session_identifier: session123 generation_number: 1 instanceUrl: https://test.thoughtspot.com [DEBUG] Getting TML for answer: What is the revenue? stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > getAnswerForQuestion > should handle API errors [DEBUG] Getting answer for sourceId: ws1 shouldGetTML: false stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > fetchTMLAndCreateLiveboard > should fetch TML and create liveboard successfully [DEBUG] Getting TML for answer: undefined [DEBUG] Getting TML for answer: undefined stdout | test/thoughtspot/thoughtspot-service.spec.ts > thoughtspot-service > fetchTMLAndCreateLiveboard > should handle TML fetch errors [DEBUG] Getting TML for answer: undefined ✓ test/thoughtspot/thoughtspot-service.spec.ts (46 tests) 14ms ✓ test/oauth-manager/oauth-utils.spec.ts (20 tests) 5ms ✓ test/oauth-manager/token-utils.integration.spec.ts (12 tests) 6ms ✓ test/oauth-manager/token-utils.spec.ts (16 tests) 4ms ✓ test/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl.spec.ts (27 tests) 4ms ✓ test/metrics/runtime/grafana-otlp-sink.spec.ts (12 tests) 2ms ✓ test/servers/version-registry.spec.ts (39 tests) 2ms ✓ test/metrics/runtime/metric-context.spec.ts (8 tests) 1ms ✓ test/metrics/runtime/runtime-config.spec.ts (9 tests) 9ms ✓ test/metrics/runtime/metric-types.spec.ts (2 tests) ✓ test/metrics/runtime/metrics-recorder.spec.ts (8 tests) 1ms ✓ test/metrics/runtime/sinks.spec.ts (3 tests) 1ms Test Files 31 passed (31) Tests 590 passed | 5 skipped (595) Start at 09:19:19 Duration 13.73s (transform 1.06s, setup 87ms, collect 6.35s, tests 4.77s, environment 0ms, prepare 100ms) % Coverage report from istanbul -------------------|---------|----------|---------|---------|------------------- File | % Stmts | % Branch | % Funcs | % Lines | Uncovered Line #s -------------------|---------|----------|---------|---------|------------------- All files | 89.24 | 84.47 | 90.24 | 89.34 | src | 83.76 | 82.51 | 73.33 | 83.93 | bearer.ts | 100 | 100 | 100 | 100 | ...lare-utils.ts | 35.29 | 0 | 37.5 | 35.29 | 39-84 handlers.ts | 93.22 | 79.59 | 92.3 | 93.22 | ...68,300,329,334 index.ts | 74.07 | 50 | 75 | 72 | 23,53-66 routes.ts | 100 | 100 | 100 | 100 | stdio.ts | 0 | 0 | 0 | 0 | 9-55 ...ming-utils.ts | 100 | 96.96 | 100 | 100 | 139 utils.ts | 90 | 96.29 | 100 | 92.3 | 110,128-129 src/metrics | 100 | 100 | 100 | 100 | index.ts | 100 | 100 | 100 | 100 | ...trics/mixpanel | 100 | 100 | 100 | 100 | ...nel-client.ts | 100 | 100 | 100 | 100 | mixpanel.ts | 100 | 100 | 100 | 100 | ...etrics/runtime | 91.97 | 89.61 | 96 | 91.95 | ...osite-sink.ts | 100 | 100 | 100 | 100 | ...-otlp-sink.ts | 87.2 | 78.57 | 100 | 87.09 | ...76,181,184,197 ...ic-context.ts | 92.85 | 95 | 100 | 92.85 | 170,189,198,229 metric-types.ts | 96.87 | 94.11 | 100 | 96.87 | 86 ...s-recorder.ts | 100 | 100 | 84.61 | 100 | noop-sink.ts | 0 | 0 | 0 | 0 | ...st-metrics.ts | 91.66 | 100 | 88.88 | 91.66 | 104-108 ...ime-config.ts | 94.44 | 89.18 | 100 | 94.44 | 22,39 ...etrics/tracing | 100 | 100 | 100 | 100 | tracing-utils.ts | 100 | 100 | 100 | 100 | src/oauth-manager | 98.59 | 90 | 100 | 98.59 | oauth-utils.ts | 98 | 85 | 100 | 98 | 451 token-utils.ts | 100 | 100 | 100 | 100 | src/servers | 84.4 | 74.19 | 87.95 | 84.61 | api-server.ts | 92.18 | 37.5 | 93.33 | 92.18 | 158-165 ...age-server.ts | 98.52 | 83.78 | 100 | 98.52 | 182 ...erver-base.ts | 88.88 | 66.66 | 86.36 | 88.88 | 75-78,178,255-257 mcp-server.ts | 65.94 | 72.13 | 77.27 | 65.92 | ...49-395,406-446 ...mcp-server.ts | 86.95 | 70.58 | 88.88 | 88.63 | 151-167 ...efinitions.ts | 100 | 100 | 100 | 100 | ...n-registry.ts | 90.32 | 83.33 | 100 | 90 | 57-60 ...torage-service | 100 | 100 | 100 | 100 | ...ge-service.ts | 100 | 100 | 100 | 100 | ...orage-with-ttl | 100 | 94.44 | 100 | 100 | ...e-with-ttl.ts | 100 | 94.44 | 100 | 100 | 96 src/thoughtspot | 94.26 | 85.41 | 93.61 | 94.21 | ...pot-client.ts | 98.38 | 100 | 92.85 | 98.38 | 23 ...ot-service.ts | 92.85 | 79.41 | 93.93 | 92.77 | ...98-208,647-652 -------------------|---------|----------|---------|---------|------------------- [vpw:debug] Shutting down runtimes... * SCAL-308279 Simplify Grafana OTLP env config ## Summary Reduce the Grafana OTLP sink to one canonical env name per setting so the config surface is explicit and easier to maintain. ## What Changed - collapse Grafana OTLP config resolution to `GRAFANA_OTLP_*` env names only - remove alias-based precedence across Grafana Cloud and OTLP-style env names - drop `OTEL_EXPORTER_OTLP_HEADERS` parsing from the sink config path - update Grafana sink tests to cover only the canonical env names ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/grafana-otlp-sink.spec.ts test/metrics/runtime/request-metrics.spec.ts test/metrics/runtime/runtime-config.spec.ts test/metrics/runtime/sinks.spec.ts` --- src/metrics/runtime/grafana-otlp-sink.ts | 413 ++++++++++++++++++ src/metrics/runtime/request-metrics.ts | 37 +- .../metrics/runtime/grafana-otlp-sink.spec.ts | 411 +++++++++++++++++ test/metrics/runtime/request-metrics.spec.ts | 68 +++ 4 files changed, 928 insertions(+), 1 deletion(-) create mode 100644 src/metrics/runtime/grafana-otlp-sink.ts create mode 100644 test/metrics/runtime/grafana-otlp-sink.spec.ts diff --git a/src/metrics/runtime/grafana-otlp-sink.ts b/src/metrics/runtime/grafana-otlp-sink.ts new file mode 100644 index 0000000..eb10c43 --- /dev/null +++ b/src/metrics/runtime/grafana-otlp-sink.ts @@ -0,0 +1,413 @@ +import { + HISTOGRAM_BUCKETS_MS, + type MetricKind, + type MetricLabelValue, + type MetricName, +} from "./metric-types"; +import type { + MetricObservation, + MetricResourceAttributes, + MetricsFlushPayload, + MetricsSink, +} from "./metrics-sink"; + +type OtlpStringValue = { stringValue: string }; +type OtlpBoolValue = { boolValue: boolean }; +type OtlpDoubleValue = { doubleValue: number }; +type OtlpAttributeValue = OtlpStringValue | OtlpBoolValue | OtlpDoubleValue; + +type OtlpAttribute = { + key: string; + value: OtlpAttributeValue; +}; + +type OtlpNumberDataPoint = { + attributes?: OtlpAttribute[]; + asDouble: number; + timeUnixNano: string; +}; + +type OtlpHistogramDataPoint = { + attributes?: OtlpAttribute[]; + count: string; + sum: number; + bucketCounts: string[]; + explicitBounds: readonly number[]; + timeUnixNano: string; +}; + +type OtlpMetric = + | { + name: string; + sum: { + aggregationTemporality: typeof OTLP_AGGREGATION_TEMPORALITY_DELTA; + isMonotonic: true; + dataPoints: OtlpNumberDataPoint[]; + }; + } + | { + name: string; + gauge: { + dataPoints: OtlpNumberDataPoint[]; + }; + } + | { + name: string; + histogram: { + aggregationTemporality: typeof OTLP_AGGREGATION_TEMPORALITY_DELTA; + dataPoints: OtlpHistogramDataPoint[]; + }; + }; + +export type OtlpMetricsPayload = { + resourceMetrics: Array<{ + resource: { + attributes: OtlpAttribute[]; + }; + scopeMetrics: Array<{ + scope: { + name: string; + }; + metrics: OtlpMetric[]; + }>; + }>; +}; + +export type GrafanaOtlpEnvLike = Partial>; + +export type GrafanaOtlpSinkConfig = { + endpoint: string; + username?: string; + apiToken?: string; + authHeader?: string; +}; + +type GrafanaOtlpMetricsSinkOptions = GrafanaOtlpSinkConfig & { + fetchFn?: typeof fetch; +}; + +type AggregatedMetricDataPoint = { + attributes: OtlpAttribute[]; + bucketCounts: number[]; + count: number; + timestampMs: number; + value: number; +}; + +const OTLP_SCOPE_NAME = "thoughtspot.mcp.metrics.runtime"; +const OTLP_AGGREGATION_TEMPORALITY_DELTA = 1; +const MAX_EXPORT_ERROR_BODY_LENGTH = 1_000; +const JSON_HEADERS = { + "Content-Type": "application/json", +}; + +function getProcessEnvValue(name: string): string | undefined { + if (typeof process === "undefined") { + return undefined; + } + return process.env?.[name]; +} + +function readConfigValue( + env: GrafanaOtlpEnvLike | undefined, + key: string, +): string | undefined { + const envValue = env?.[key]; + if (typeof envValue === "string" && envValue.length > 0) { + return envValue; + } + + const processEnvValue = getProcessEnvValue(key); + if (processEnvValue && processEnvValue.length > 0) { + return processEnvValue; + } + + return undefined; +} + +function encodeBase64(value: string): string { + if (typeof Buffer !== "undefined") { + return Buffer.from(value, "utf8").toString("base64"); + } + if (typeof TextEncoder !== "undefined" && typeof btoa === "function") { + const bytes = new TextEncoder().encode(value); + let binary = ""; + for (const byte of bytes) { + binary += String.fromCharCode(byte); + } + return btoa(binary); + } + throw new Error("No base64 encoder is available for Grafana OTLP auth"); +} + +function toOtlpValue(value: MetricLabelValue): OtlpAttributeValue { + if (typeof value === "boolean") { + return { boolValue: value }; + } + if (typeof value === "number") { + return { doubleValue: value }; + } + return { stringValue: value }; +} + +function toOtlpAttributes( + attributes: Record, +): OtlpAttribute[] { + return Object.keys(attributes) + .sort() + .flatMap((key) => { + const value = attributes[key]; + if (value === undefined || value === "") { + return []; + } + return [{ key, value: toOtlpValue(value) }]; + }); +} + +function toTimeUnixNano(timestampMs: number): string { + const integerMs = Math.trunc(timestampMs); + const remainderNs = Math.round((timestampMs - integerMs) * 1_000_000); + const carryMs = Math.trunc(remainderNs / 1_000_000); + const normalizedRemainderNs = remainderNs - carryMs * 1_000_000; + + return ( + (BigInt(integerMs) + BigInt(carryMs)) * 1_000_000n + + BigInt(normalizedRemainderNs) + ).toString(); +} + +function getHistogramBucketIndex(value: number): number { + const bucketIndex = HISTOGRAM_BUCKETS_MS.findIndex((bound) => value <= bound); + return bucketIndex === -1 ? HISTOGRAM_BUCKETS_MS.length : bucketIndex; +} + +function groupObservationsByMetric( + observations: readonly MetricObservation[], +): Map { + const grouped = new Map(); + for (const observation of observations) { + const metricObservations = grouped.get(observation.name) ?? []; + metricObservations.push(observation); + grouped.set(observation.name, metricObservations); + } + return grouped; +} + +function getAttributeSetKey(attributes: readonly OtlpAttribute[]): string { + return JSON.stringify(attributes); +} + +function toNumberDataPoint( + observation: AggregatedMetricDataPoint, +): OtlpNumberDataPoint { + return { + attributes: observation.attributes, + asDouble: observation.value, + timeUnixNano: toTimeUnixNano(observation.timestampMs), + }; +} + +function toHistogramDataPoint( + observation: AggregatedMetricDataPoint, +): OtlpHistogramDataPoint { + return { + attributes: observation.attributes, + count: String(observation.count), + sum: observation.value, + bucketCounts: observation.bucketCounts.map(String), + explicitBounds: HISTOGRAM_BUCKETS_MS, + timeUnixNano: toTimeUnixNano(observation.timestampMs), + }; +} + +function aggregateObservations( + kind: MetricKind, + observations: readonly MetricObservation[], +): AggregatedMetricDataPoint[] { + const aggregated = new Map(); + + for (const observation of observations) { + const attributes = toOtlpAttributes(observation.labels); + const attributeSetKey = getAttributeSetKey(attributes); + const dataPoint = aggregated.get(attributeSetKey) ?? { + attributes, + bucketCounts: + kind === "histogram" + ? (new Array(HISTOGRAM_BUCKETS_MS.length + 1).fill(0) as number[]) + : [], + count: 0, + timestampMs: observation.timestampMs, + value: 0, + }; + + switch (kind) { + case "counter": + dataPoint.value += observation.value; + dataPoint.count += 1; + dataPoint.timestampMs = observation.timestampMs; + break; + case "gauge": + dataPoint.value = observation.value; + dataPoint.count = 1; + dataPoint.timestampMs = observation.timestampMs; + break; + case "histogram": { + const bucketIndex = getHistogramBucketIndex(observation.value); + dataPoint.bucketCounts[bucketIndex] += 1; + dataPoint.value += observation.value; + dataPoint.count += 1; + dataPoint.timestampMs = observation.timestampMs; + break; + } + } + + aggregated.set(attributeSetKey, dataPoint); + } + + return [...aggregated.values()]; +} + +function toOtlpMetric( + name: MetricName, + kind: MetricKind, + observations: readonly MetricObservation[], +): OtlpMetric { + const dataPoints = aggregateObservations(kind, observations); + + switch (kind) { + case "counter": + return { + name, + sum: { + aggregationTemporality: OTLP_AGGREGATION_TEMPORALITY_DELTA, + isMonotonic: true, + dataPoints: dataPoints.map(toNumberDataPoint), + }, + }; + case "gauge": + return { + name, + gauge: { + dataPoints: dataPoints.map(toNumberDataPoint), + }, + }; + case "histogram": + return { + name, + histogram: { + aggregationTemporality: OTLP_AGGREGATION_TEMPORALITY_DELTA, + dataPoints: dataPoints.map(toHistogramDataPoint), + }, + }; + } +} + +function normalizeOtlpMetricsEndpoint(endpoint: string): string { + const trimmed = endpoint.trim().replace(/\/+$/, ""); + return trimmed.endsWith("/v1/metrics") ? trimmed : `${trimmed}/v1/metrics`; +} + +function buildAuthorizationHeaderValue( + config: Pick, +): string | undefined { + if (config.authHeader) { + return config.authHeader; + } + if (config.username && config.apiToken) { + return `Basic ${encodeBase64(`${config.username}:${config.apiToken}`)}`; + } + return undefined; +} + +function buildRequestHeaders(config: GrafanaOtlpSinkConfig): HeadersInit { + const authorization = buildAuthorizationHeaderValue(config); + return authorization + ? { ...JSON_HEADERS, Authorization: authorization } + : JSON_HEADERS; +} + +async function getExportErrorBody(response: Response): Promise { + const text = await response.text(); + if (text.length <= MAX_EXPORT_ERROR_BODY_LENGTH) { + return text; + } + const truncatedLength = Math.max(0, MAX_EXPORT_ERROR_BODY_LENGTH - 3); + return `${text.slice(0, truncatedLength)}...`; +} + +export function toOtlpMetricsPayload( + payload: MetricsFlushPayload, +): OtlpMetricsPayload { + const grouped = groupObservationsByMetric(payload.observations); + + return { + resourceMetrics: [ + { + resource: { + attributes: toOtlpAttributes(payload.resourceAttributes), + }, + scopeMetrics: [ + { + scope: { name: OTLP_SCOPE_NAME }, + metrics: [...grouped.entries()].map(([name, observations]) => + toOtlpMetric(name, observations[0].kind, observations), + ), + }, + ], + }, + ], + }; +} + +export function resolveGrafanaOtlpSinkConfig( + env?: GrafanaOtlpEnvLike, +): GrafanaOtlpSinkConfig | undefined { + const endpoint = readConfigValue(env, "GRAFANA_OTLP_ENDPOINT"); + if (!endpoint) { + return undefined; + } + + return { + endpoint: normalizeOtlpMetricsEndpoint(endpoint), + username: readConfigValue(env, "GRAFANA_OTLP_USERNAME"), + apiToken: readConfigValue(env, "GRAFANA_OTLP_API_TOKEN"), + authHeader: readConfigValue(env, "GRAFANA_OTLP_AUTH_HEADER"), + }; +} + +export class GrafanaOtlpMetricsSink implements MetricsSink { + private readonly endpoint: string; + private readonly fetchFn: typeof fetch; + private readonly headers: HeadersInit; + + constructor(options: GrafanaOtlpMetricsSinkOptions) { + this.endpoint = normalizeOtlpMetricsEndpoint(options.endpoint); + this.fetchFn = options.fetchFn ?? fetch; + this.headers = buildRequestHeaders(options); + } + + async flush(payload: MetricsFlushPayload): Promise { + if (payload.observations.length === 0) { + return; + } + + const response = await this.fetchFn(this.endpoint, { + method: "POST", + headers: this.headers, + body: JSON.stringify(toOtlpMetricsPayload(payload)), + }); + + if (!response.ok) { + throw new Error( + `Grafana OTLP metrics export failed with status ${response.status}: ${await getExportErrorBody(response)}`, + ); + } + } +} + +export function createGrafanaOtlpMetricsSink( + env?: GrafanaOtlpEnvLike, +): GrafanaOtlpMetricsSink | undefined { + const config = resolveGrafanaOtlpSinkConfig(env); + return config ? new GrafanaOtlpMetricsSink(config) : undefined; +} diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 000f1f4..0a2abe2 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,8 +1,10 @@ +import { createGrafanaOtlpMetricsSink } from "./grafana-otlp-sink"; import { type MetricsRecorder, NOOP_METRICS_RECORDER, RequestMetricsRecorder, } from "./metrics-recorder"; +import type { MetricsSink } from "./metrics-sink"; import { type ConfiguredMetricsSinks, type MetricsEnvLike, @@ -13,6 +15,7 @@ import { const METRICS_RECORDER_SYMBOL = Symbol.for( "thoughtspot.mcp.metrics.requestRecorder", ); +const GRAFANA_SINK_CACHE = new WeakMap(); type MetricsExecutionContext = ExecutionContext & { [METRICS_RECORDER_SYMBOL]?: MetricsRecorder; @@ -38,13 +41,45 @@ export function clearMetricsRecorderFromExecutionContext( delete (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL]; } +function getCachedGrafanaSink( + env: MetricsEnvLike | undefined, +): MetricsSink | undefined { + if (!env || typeof env !== "object") { + return createGrafanaOtlpMetricsSink(env); + } + + const cachedSink = GRAFANA_SINK_CACHE.get(env); + if (cachedSink) { + return cachedSink; + } + + const sink = createGrafanaOtlpMetricsSink(env); + if (sink) { + GRAFANA_SINK_CACHE.set(env, sink); + } + return sink; +} + +function createDefaultConfiguredMetricsSinks( + env: MetricsEnvLike | undefined, + sinks: ConfiguredMetricsSinks, +): ConfiguredMetricsSinks { + return { + analyticsEngineSink: sinks.analyticsEngineSink, + grafanaSink: sinks.grafanaSink ?? getCachedGrafanaSink(env), + }; +} + export function createRequestMetricsRecorder( env?: MetricsEnvLike, sinks: ConfiguredMetricsSinks = {}, ): MetricsRecorder { try { const config = resolveMetricsRuntimeConfig(env); - const sink = createConfiguredMetricsSink(config, sinks); + const sink = createConfiguredMetricsSink( + config, + createDefaultConfiguredMetricsSinks(env, sinks), + ); return new RequestMetricsRecorder({ sink, diff --git a/test/metrics/runtime/grafana-otlp-sink.spec.ts b/test/metrics/runtime/grafana-otlp-sink.spec.ts new file mode 100644 index 0000000..e572f18 --- /dev/null +++ b/test/metrics/runtime/grafana-otlp-sink.spec.ts @@ -0,0 +1,411 @@ +import { describe, expect, it, vi } from "vitest"; +import { + GrafanaOtlpMetricsSink, + createGrafanaOtlpMetricsSink, + resolveGrafanaOtlpSinkConfig, + toOtlpMetricsPayload, +} from "../../../src/metrics/runtime/grafana-otlp-sink"; +import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; +import type { MetricsFlushPayload } from "../../../src/metrics/runtime/metrics-sink"; + +describe("GrafanaOtlpMetricsSink", () => { + const payload = { + observations: [ + { + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 1, + labels: { + route_group: "mcp", + status_class: "2xx", + }, + timestampMs: 1_714_000_000_000, + }, + { + kind: "histogram", + name: METRIC_NAMES.httpRequestDurationMs, + value: 123, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_123, + }, + { + kind: "gauge", + name: METRIC_NAMES.httpInflightRequests, + value: 3, + labels: {}, + timestampMs: 1_714_000_000_456, + }, + ], + resourceAttributes: { + "deployment.environment": "production", + "service.name": "thoughtspot-mcp-server", + "service.namespace": "thoughtspot", + }, + } satisfies MetricsFlushPayload; + + it("maps observations into OTLP JSON metrics payloads", () => { + const otlpPayload = toOtlpMetricsPayload(payload); + const metrics = otlpPayload.resourceMetrics[0].scopeMetrics[0].metrics; + + expect(otlpPayload.resourceMetrics[0].resource.attributes).toEqual([ + { + key: "deployment.environment", + value: { stringValue: "production" }, + }, + { + key: "service.name", + value: { stringValue: "thoughtspot-mcp-server" }, + }, + { key: "service.namespace", value: { stringValue: "thoughtspot" } }, + ]); + expect(metrics).toHaveLength(3); + expect(metrics[0]).toMatchObject({ + name: METRIC_NAMES.httpRequestsTotal, + sum: { + aggregationTemporality: 1, + isMonotonic: true, + dataPoints: [ + { + asDouble: 1, + timeUnixNano: "1714000000000000000", + attributes: [ + { key: "route_group", value: { stringValue: "mcp" } }, + { key: "status_class", value: { stringValue: "2xx" } }, + ], + }, + ], + }, + }); + expect(metrics[1]).toMatchObject({ + name: METRIC_NAMES.httpRequestDurationMs, + histogram: { + aggregationTemporality: 1, + dataPoints: [ + { + count: "1", + sum: 123, + bucketCounts: [ + "0", + "0", + "0", + "1", + "0", + "0", + "0", + "0", + "0", + "0", + "0", + ], + }, + ], + }, + }); + expect(metrics[2]).toMatchObject({ + name: METRIC_NAMES.httpInflightRequests, + gauge: { + dataPoints: [{ asDouble: 3 }], + }, + }); + }); + + it("aggregates observations with identical attributes before export", () => { + const otlpPayload = toOtlpMetricsPayload({ + observations: [ + { + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 1, + labels: { + route_group: "mcp", + status_class: "2xx", + }, + timestampMs: 1_714_000_000_100, + }, + { + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 2, + labels: { + status_class: "2xx", + route_group: "mcp", + }, + timestampMs: 1_714_000_000_123.5, + }, + { + kind: "histogram", + name: METRIC_NAMES.httpRequestDurationMs, + value: 25, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_200, + }, + { + kind: "histogram", + name: METRIC_NAMES.httpRequestDurationMs, + value: 75, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_201, + }, + { + kind: "gauge", + name: METRIC_NAMES.httpInflightRequests, + value: 1, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_300, + }, + { + kind: "gauge", + name: METRIC_NAMES.httpInflightRequests, + value: 4, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_301, + }, + ], + resourceAttributes: {}, + }); + const [counterMetric, histogramMetric, gaugeMetric] = + otlpPayload.resourceMetrics[0].scopeMetrics[0].metrics; + + expect(counterMetric).toMatchObject({ + sum: { + dataPoints: [ + { + asDouble: 3, + timeUnixNano: "1714000000123500000", + }, + ], + }, + }); + expect(counterMetric.sum.dataPoints).toHaveLength(1); + expect(histogramMetric).toMatchObject({ + histogram: { + dataPoints: [ + { + count: "2", + sum: 100, + bucketCounts: [ + "1", + "0", + "1", + "0", + "0", + "0", + "0", + "0", + "0", + "0", + "0", + ], + }, + ], + }, + }); + expect(histogramMetric.histogram.dataPoints).toHaveLength(1); + expect(gaugeMetric).toMatchObject({ + gauge: { + dataPoints: [ + { + asDouble: 4, + timeUnixNano: "1714000000301000000", + }, + ], + }, + }); + expect(gaugeMetric.gauge.dataPoints).toHaveLength(1); + }); + + it("keeps nanosecond precision when fractional milliseconds round up a microsecond", () => { + const otlpPayload = toOtlpMetricsPayload({ + observations: [ + { + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 1, + labels: { + route_group: "mcp", + }, + timestampMs: 1_714_000_000_123.9995, + }, + ], + resourceAttributes: {}, + }); + const [metric] = otlpPayload.resourceMetrics[0].scopeMetrics[0].metrics; + + expect(metric).toMatchObject({ + sum: { + dataPoints: [ + { + timeUnixNano: "1714000000123999512", + }, + ], + }, + }); + }); + + it("posts OTLP JSON to the normalized metrics endpoint", async () => { + const fetchFn = vi + .fn() + .mockResolvedValue(new Response(null, { status: 200 })); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com/otlp", + username: "12345", + apiToken: "secret", + fetchFn, + }); + + await sink.flush(payload); + + expect(fetchFn).toHaveBeenCalledTimes(1); + expect(fetchFn).toHaveBeenCalledWith( + "https://otlp.example.com/otlp/v1/metrics", + expect.objectContaining({ + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: "Basic MTIzNDU6c2VjcmV0", + }, + body: expect.any(String), + }), + ); + }); + + it("supports full metrics endpoints and explicit auth headers", async () => { + const fetchFn = vi + .fn() + .mockResolvedValue(new Response(null, { status: 200 })); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com/v1/metrics", + authHeader: "Bearer token", + fetchFn, + }); + + await sink.flush(payload); + + expect(fetchFn).toHaveBeenCalledWith( + "https://otlp.example.com/v1/metrics", + expect.objectContaining({ + headers: { + "Content-Type": "application/json", + Authorization: "Bearer token", + }, + }), + ); + }); + + it("encodes Basic auth credentials as UTF-8", async () => { + const fetchFn = vi + .fn() + .mockResolvedValue(new Response(null, { status: 200 })); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com/otlp", + username: "üser", + apiToken: "päss", + fetchFn, + }); + + await sink.flush(payload); + + expect(fetchFn).toHaveBeenCalledWith( + "https://otlp.example.com/otlp/v1/metrics", + expect.objectContaining({ + headers: { + "Content-Type": "application/json", + Authorization: "Basic w7xzZXI6cMOkc3M=", + }, + }), + ); + }); + + it("skips export when there are no observations", async () => { + const fetchFn = vi.fn(); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com", + fetchFn, + }); + + await sink.flush({ observations: [], resourceAttributes: {} }); + + expect(fetchFn).not.toHaveBeenCalled(); + }); + + it("throws on non-2xx responses so the recorder can isolate failures", async () => { + const fetchFn = vi + .fn() + .mockResolvedValue(new Response("bad token", { status: 401 })); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com", + fetchFn, + }); + + await expect(sink.flush(payload)).rejects.toThrow( + "Grafana OTLP metrics export failed with status 401: bad token", + ); + }); + + it("truncates large export error bodies", async () => { + const fetchFn = vi + .fn() + .mockResolvedValue(new Response("x".repeat(1_010), { status: 500 })); + const sink = new GrafanaOtlpMetricsSink({ + endpoint: "https://otlp.example.com", + fetchFn, + }); + + await expect(sink.flush(payload)).rejects.toThrow( + `Grafana OTLP metrics export failed with status 500: ${"x".repeat(997)}...`, + ); + }); + + it("resolves config from canonical Grafana environment names", () => { + expect( + resolveGrafanaOtlpSinkConfig({ + GRAFANA_OTLP_ENDPOINT: "https://otlp.example.com/otlp", + GRAFANA_OTLP_USERNAME: "12345", + GRAFANA_OTLP_API_TOKEN: "test-api-token", + }), + ).toEqual({ + endpoint: "https://otlp.example.com/otlp/v1/metrics", + username: "12345", + apiToken: "test-api-token", + authHeader: undefined, + }); + expect( + resolveGrafanaOtlpSinkConfig({ + GRAFANA_OTLP_ENDPOINT: "https://collector.example.com", + GRAFANA_OTLP_AUTH_HEADER: "Bearer test", + }), + ).toMatchObject({ + endpoint: "https://collector.example.com/v1/metrics", + authHeader: "Bearer test", + }); + expect(resolveGrafanaOtlpSinkConfig()).toBeUndefined(); + }); + + it("ignores non-canonical OTLP environment names", () => { + expect( + resolveGrafanaOtlpSinkConfig({ + OTEL_EXPORTER_OTLP_ENDPOINT: "https://otlp.example.com/otlp", + OTEL_EXPORTER_OTLP_HEADERS: "Authorization=Basic%20abc123", + }), + ).toBeUndefined(); + }); + + it("only creates a sink when an OTLP endpoint is configured", () => { + expect( + createGrafanaOtlpMetricsSink({ + GRAFANA_OTLP_ENDPOINT: "https://otlp.example.com", + }), + ).toBeInstanceOf(GrafanaOtlpMetricsSink); + expect(createGrafanaOtlpMetricsSink({})).toBeUndefined(); + }); +}); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index 530ec58..89b80ca 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -194,6 +194,34 @@ describe("withRequestMetrics", () => { ); }); + it("uses Grafana OTLP config as the default grafana sink", async () => { + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue(new Response(null, { status: 200 })); + const recorder = createRequestMetricsRecorder({ + METRICS_SINK_MODE: "grafana", + GRAFANA_OTLP_ENDPOINT: "https://otlp.example.com/otlp", + GRAFANA_OTLP_AUTH_HEADER: "Bearer test", + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + await recorder.flush(); + + expect(fetchSpy).toHaveBeenCalledWith( + "https://otlp.example.com/otlp/v1/metrics", + expect.objectContaining({ + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: "Bearer test", + }, + body: expect.any(String), + }), + ); + }); + it("does not fail the request when waitUntil rejects scheduling", async () => { const errorSpy = vi .spyOn(console, "error") @@ -255,6 +283,46 @@ describe("withRequestMetrics", () => { ); }); + it("reuses the default Grafana sink for the same env object", async () => { + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue(new Response(null, { status: 200 })); + const env = { + METRICS_SINK_MODE: "grafana", + GRAFANA_OTLP_ENDPOINT: "https://otlp.example.com/first", + GRAFANA_OTLP_AUTH_HEADER: "Bearer first", + }; + const firstRecorder = createRequestMetricsRecorder(env); + + env.GRAFANA_OTLP_ENDPOINT = "https://otlp.example.com/second"; + env.GRAFANA_OTLP_AUTH_HEADER = "Bearer second"; + const secondRecorder = createRequestMetricsRecorder(env); + + firstRecorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + secondRecorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + await firstRecorder.flush(); + await secondRecorder.flush(); + + expect(fetchSpy).toHaveBeenCalledTimes(2); + expect(fetchSpy).toHaveBeenCalledWith( + "https://otlp.example.com/first/v1/metrics", + expect.objectContaining({ + headers: { + "Content-Type": "application/json", + Authorization: "Bearer first", + }, + }), + ); + expect(fetchSpy).not.toHaveBeenCalledWith( + "https://otlp.example.com/second/v1/metrics", + expect.anything(), + ); + }); + it("supports setting and clearing the recorder on the execution context", () => { const ctx = {} as ExecutionContext; const recorder = createRequestMetricsRecorder(); From 28a9f18ee143a7fe83ed9faf7a56e59e718688dc Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Mon, 4 May 2026 15:24:45 -0700 Subject: [PATCH 14/25] SCAL-308278 Add Analytics Engine metrics sink (#125) * SCAL-308278 Add Analytics Engine metrics sink Add the Analytics Engine sink for the metrics runtime so PM and tenant event observations can be written to the existing Cloudflare `ANALYTICS` binding. - add an Analytics Engine sink that writes one datapoint per metric observation - define the positional Analytics Engine schema for indexes, blobs, and doubles - classify metric names into stable event families for PM-oriented querying - wire request metrics to use `env.ANALYTICS` when the analytics sink is enabled - keep missing bindings as noop through the existing sink selection path - cover datapoint mapping, binding detection, and default sink wiring in tests The schema reserves index positions for tenant and user identity. This PR wires the sink and raw event shape; later instrumentation PRs will populate those identity fields from ThoughtSpot session context where available. - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/analytics-engine-sink.spec.ts test/metrics/runtime/request-metrics.spec.ts test/metrics/runtime/runtime-config.spec.ts test/metrics/runtime/sinks.spec.ts` - `npm test` * SCAL-308278 Cover all Analytics Engine metric families ## Summary Make Analytics Engine metric family classification exhaustive and cover all current metric names in tests. ## What Changed - map `sessionsStartedTotal` into the auth event family - add an exhaustive check for future metric additions - verify every `METRIC_NAMES` value resolves to a valid Analytics Engine family ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/analytics-engine-sink.spec.ts` * SCAL-308278 Isolate Analytics Engine write failures ## Summary Keep one failed Analytics Engine write from aborting the rest of a metrics flush. ## What Changed - allow Analytics Engine bindings to return either sync or async write results - await each data point write during flush - catch and warn per failed data point before continuing with the batch - add regression coverage for continuing after a rejected write ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308278 Parallelize Analytics Engine writes ## Summary Write Analytics Engine metric datapoints in parallel during flush while keeping per-point failure isolation. ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308278 Add Analytics Engine event identity ## Summary Make the future tenant/user enrichment path explicit in the metrics flush contract. ## What Changed - add optional event identity to `MetricsFlushPayload` - pass event identity into Analytics Engine data point conversion - map tenant and user identity into Analytics Engine indexes when present - add regression coverage for identity-backed Analytics Engine indexes ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308278 Use sync Analytics Engine writes ## Summary Match the Analytics Engine sink to Cloudflare's synchronous `writeDataPoint` API. ## What Changed - narrow the Analytics Engine dataset contract back to synchronous `void` - use a simple `for...of` loop with per-point error isolation - update the failure-isolation test to throw synchronously ## Validation - `npm run lint` - focused runtime metrics tests * SCAL-308278 Centralize Analytics Engine schema guards ## Summary Tighten the Analytics Engine sink schema wiring so the dataset boundary is typed explicitly and schema drift is caught by tests. ## What Changed - add a dedicated `isAnalyticsEngineDatasetLike(...)` type guard for the binding - derive persisted Analytics Engine label fields from `APPROVED_METRIC_LABEL_KEYS` - centralize the persisted resource attribute field mapping used by the AE sink - add a guard test that locks the AE schema arrays to the approved labels and resource attributes ## Validation - `npm run lint` - `npm test -- --coverage.enabled=false test/metrics/runtime/analytics-engine-sink.spec.ts test/metrics/runtime/request-metrics.spec.ts test/metrics/runtime/runtime-config.spec.ts test/metrics/runtime/sinks.spec.ts` --- src/metrics/runtime/analytics-engine-sink.ts | 206 +++++++++++++++++ src/metrics/runtime/metrics-sink.ts | 6 + src/metrics/runtime/request-metrics.ts | 23 +- .../runtime/analytics-engine-sink.spec.ts | 216 ++++++++++++++++++ test/metrics/runtime/request-metrics.spec.ts | 20 ++ 5 files changed, 461 insertions(+), 10 deletions(-) create mode 100644 src/metrics/runtime/analytics-engine-sink.ts create mode 100644 test/metrics/runtime/analytics-engine-sink.spec.ts diff --git a/src/metrics/runtime/analytics-engine-sink.ts b/src/metrics/runtime/analytics-engine-sink.ts new file mode 100644 index 0000000..fd1bc57 --- /dev/null +++ b/src/metrics/runtime/analytics-engine-sink.ts @@ -0,0 +1,206 @@ +import { + APPROVED_METRIC_LABEL_KEYS, + METRIC_NAMES, + type MetricLabels, + type MetricName, +} from "./metric-types"; +import type { + MetricObservation, + MetricResourceAttributes, + MetricsFlushPayload, + MetricsSink, +} from "./metrics-sink"; + +export type AnalyticsEngineDataPointLike = { + indexes?: ((ArrayBuffer | string) | null)[]; + blobs?: ((ArrayBuffer | string) | null)[]; + doubles?: number[]; +}; + +export type AnalyticsEngineDatasetLike = { + writeDataPoint(event?: AnalyticsEngineDataPointLike): void; +}; + +export const ANALYTICS_ENGINE_SCHEMA_VERSION = "mcp_metrics_v1"; + +export const ANALYTICS_ENGINE_INDEX_FIELDS = [ + "schema_version", + "event_family", + "metric_name", + "tenant_id", + "user_id", +] as const; + +export const ANALYTICS_ENGINE_LABEL_FIELDS = APPROVED_METRIC_LABEL_KEYS; + +export const ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS = [ + ["deployment_environment", "deployment.environment"], + ["service_name", "service.name"], + ["service_namespace", "service.namespace"], + ["service_version", "service.version"], +] as const satisfies readonly (readonly [ + string, + keyof MetricResourceAttributes, +])[]; + +export const ANALYTICS_ENGINE_BLOB_FIELDS = [ + "metric_kind", + ...ANALYTICS_ENGINE_LABEL_FIELDS, + ...ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS.map(([field]) => field), +] as const; + +export const ANALYTICS_ENGINE_DOUBLE_FIELDS = [ + "metric_value", + "timestamp_ms", +] as const; + +type AnalyticsEngineIdentity = { + tenantId?: string; + userId?: string; +}; + +type AnalyticsEngineMetricFamily = + | "analysis" + | "auth" + | "dashboard" + | "http" + | "resource" + | "stream_storage" + | "tool" + | "upstream"; + +function nullableString(value: string | undefined): string | null { + return value && value.length > 0 ? value : null; +} + +function getLabel( + labels: MetricLabels, + key: keyof MetricLabels, +): string | null { + return nullableString(labels[key]); +} + +function getResourceAttribute( + resourceAttributes: MetricResourceAttributes, + key: keyof MetricResourceAttributes, +): string | null { + return nullableString(resourceAttributes[key]); +} + +export function getAnalyticsEngineMetricFamily( + name: MetricName, +): AnalyticsEngineMetricFamily { + switch (name) { + case METRIC_NAMES.httpRequestsTotal: + case METRIC_NAMES.httpRequestDurationMs: + case METRIC_NAMES.httpInflightRequests: + return "http"; + case METRIC_NAMES.toolCallsTotal: + case METRIC_NAMES.toolDurationMs: + return "tool"; + case METRIC_NAMES.resourceReadsTotal: + return "resource"; + case METRIC_NAMES.sessionsStartedTotal: + case METRIC_NAMES.oauthAuthorizeRequestsTotal: + case METRIC_NAMES.oauthAuthorizeSubmitTotal: + case METRIC_NAMES.oauthCallbackTotal: + case METRIC_NAMES.oauthStoreTokenTotal: + case METRIC_NAMES.bearerAuthRequestsTotal: + return "auth"; + case METRIC_NAMES.upstreamCallsTotal: + case METRIC_NAMES.upstreamDurationMs: + case METRIC_NAMES.upstreamStreamsStartedTotal: + case METRIC_NAMES.upstreamStreamMessagesTotal: + return "upstream"; + case METRIC_NAMES.analysisSessionsCreatedTotal: + case METRIC_NAMES.analysisMessagesSentTotal: + case METRIC_NAMES.analysisUpdatesPolledTotal: + case METRIC_NAMES.analysisPollWaitMs: + case METRIC_NAMES.analysisFirstBufferedUpdateMs: + case METRIC_NAMES.analysisFirstPollDelayMs: + case METRIC_NAMES.analysisFirstNonEmptyResponseMs: + case METRIC_NAMES.analysisSessionsNeverPolledTotal: + return "analysis"; + case METRIC_NAMES.streamStorageErrorsTotal: + return "stream_storage"; + case METRIC_NAMES.dashboardsCreatedTotal: + case METRIC_NAMES.dashboardTilesCount: + return "dashboard"; + default: { + const _exhaustiveCheck: never = name; + throw new Error( + `Unhandled Analytics Engine metric family: ${_exhaustiveCheck}`, + ); + } + } +} + +export function toAnalyticsEngineDataPoint( + observation: MetricObservation, + resourceAttributes: MetricResourceAttributes, + identity: AnalyticsEngineIdentity = {}, +): AnalyticsEngineDataPointLike { + return { + indexes: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + getAnalyticsEngineMetricFamily(observation.name), + observation.name, + nullableString(identity.tenantId), + nullableString(identity.userId), + ], + blobs: [ + observation.kind, + ...ANALYTICS_ENGINE_LABEL_FIELDS.map((key) => + getLabel(observation.labels, key), + ), + ...ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS.map(([, key]) => + getResourceAttribute(resourceAttributes, key), + ), + ], + doubles: [observation.value, observation.timestampMs], + }; +} + +export class AnalyticsEngineMetricsSink implements MetricsSink { + constructor(private readonly dataset: AnalyticsEngineDatasetLike) {} + + async flush(payload: MetricsFlushPayload): Promise { + for (const observation of payload.observations) { + try { + this.dataset.writeDataPoint( + toAnalyticsEngineDataPoint( + observation, + payload.resourceAttributes, + payload.eventIdentity, + ), + ); + } catch (error) { + console.warn( + `[metrics] Failed to write Analytics Engine data point for ${observation.name}`, + error, + ); + } + } + } +} + +function isAnalyticsEngineDatasetLike( + dataset: unknown, +): dataset is AnalyticsEngineDatasetLike { + return ( + typeof dataset === "object" && + dataset !== null && + "writeDataPoint" in dataset && + typeof dataset.writeDataPoint === "function" + ); +} + +export function createAnalyticsEngineMetricsSink( + dataset: unknown, +): AnalyticsEngineMetricsSink | undefined { + if (isAnalyticsEngineDatasetLike(dataset)) { + return new AnalyticsEngineMetricsSink(dataset); + } + + return undefined; +} diff --git a/src/metrics/runtime/metrics-sink.ts b/src/metrics/runtime/metrics-sink.ts index edc9c47..33dacdc 100644 --- a/src/metrics/runtime/metrics-sink.ts +++ b/src/metrics/runtime/metrics-sink.ts @@ -20,9 +20,15 @@ export type MetricResourceAttributes = Partial< > >; +export type MetricEventIdentity = { + tenantId?: string; + userId?: string; +}; + export type MetricsFlushPayload = { observations: readonly MetricObservation[]; resourceAttributes: MetricResourceAttributes; + eventIdentity?: MetricEventIdentity; }; export interface MetricsSink { diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 0a2abe2..617f8a6 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,4 +1,5 @@ import { createGrafanaOtlpMetricsSink } from "./grafana-otlp-sink"; +import { createAnalyticsEngineMetricsSink } from "./analytics-engine-sink"; import { type MetricsRecorder, NOOP_METRICS_RECORDER, @@ -29,6 +30,18 @@ export function setMetricsRecorderOnExecutionContext( return recorder; } +function createDefaultConfiguredMetricsSinks( + env: MetricsEnvLike | undefined, + sinks: ConfiguredMetricsSinks, +): ConfiguredMetricsSinks { + return { + analyticsEngineSink: + sinks.analyticsEngineSink ?? + createAnalyticsEngineMetricsSink(env?.ANALYTICS), + grafanaSink: sinks.grafanaSink ?? getCachedGrafanaSink(env), + }; +} + export function getMetricsRecorderFromExecutionContext( ctx: ExecutionContext, ): MetricsRecorder | undefined { @@ -60,16 +73,6 @@ function getCachedGrafanaSink( return sink; } -function createDefaultConfiguredMetricsSinks( - env: MetricsEnvLike | undefined, - sinks: ConfiguredMetricsSinks, -): ConfiguredMetricsSinks { - return { - analyticsEngineSink: sinks.analyticsEngineSink, - grafanaSink: sinks.grafanaSink ?? getCachedGrafanaSink(env), - }; -} - export function createRequestMetricsRecorder( env?: MetricsEnvLike, sinks: ConfiguredMetricsSinks = {}, diff --git a/test/metrics/runtime/analytics-engine-sink.spec.ts b/test/metrics/runtime/analytics-engine-sink.spec.ts new file mode 100644 index 0000000..7538866 --- /dev/null +++ b/test/metrics/runtime/analytics-engine-sink.spec.ts @@ -0,0 +1,216 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + ANALYTICS_ENGINE_BLOB_FIELDS, + ANALYTICS_ENGINE_DOUBLE_FIELDS, + ANALYTICS_ENGINE_INDEX_FIELDS, + ANALYTICS_ENGINE_LABEL_FIELDS, + ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS, + ANALYTICS_ENGINE_SCHEMA_VERSION, + AnalyticsEngineMetricsSink, + createAnalyticsEngineMetricsSink, + getAnalyticsEngineMetricFamily, + toAnalyticsEngineDataPoint, +} from "../../../src/metrics/runtime/analytics-engine-sink"; +import { + APPROVED_METRIC_LABEL_KEYS, + METRIC_NAMES, +} from "../../../src/metrics/runtime/metric-types"; +import type { MetricObservation } from "../../../src/metrics/runtime/metrics-sink"; + +describe("AnalyticsEngineMetricsSink", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + const observation = { + kind: "counter", + name: METRIC_NAMES.toolCallsTotal, + value: 2, + labels: { + tool_name: "create_liveboard", + outcome: "success", + route_group: "mcp", + transport: "mcp", + api_surface: "mcp", + auth_mode: "oauth", + }, + timestampMs: 1_714_000_000_000, + } satisfies MetricObservation; + + it("maps metric observations into the Analytics Engine schema", () => { + const dataPoint = toAnalyticsEngineDataPoint(observation, { + "deployment.environment": "production", + "service.name": "thoughtspot-mcp-server", + "service.namespace": "thoughtspot", + "service.version": "0.5.0", + }); + + expect(dataPoint.indexes).toEqual([ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + null, + null, + ]); + expect(dataPoint.blobs).toEqual([ + "counter", + "mcp", + "mcp", + "oauth", + "mcp", + null, + "success", + null, + "create_liveboard", + null, + null, + null, + null, + null, + "production", + "thoughtspot-mcp-server", + "thoughtspot", + "0.5.0", + ]); + expect(dataPoint.doubles).toEqual([2, 1_714_000_000_000]); + }); + + it("classifies every metric name into a stable event family", () => { + const validFamilies = [ + "analysis", + "auth", + "dashboard", + "http", + "resource", + "stream_storage", + "tool", + "upstream", + ]; + + for (const name of Object.values(METRIC_NAMES)) { + expect(validFamilies).toContain(getAnalyticsEngineMetricFamily(name)); + } + expect( + getAnalyticsEngineMetricFamily(METRIC_NAMES.sessionsStartedTotal), + ).toBe("auth"); + }); + + it("keeps the Analytics Engine schema aligned with approved labels and resource attributes", () => { + expect(ANALYTICS_ENGINE_LABEL_FIELDS).toEqual(APPROVED_METRIC_LABEL_KEYS); + expect(ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS).toEqual([ + ["deployment_environment", "deployment.environment"], + ["service_name", "service.name"], + ["service_namespace", "service.namespace"], + ["service_version", "service.version"], + ]); + expect(ANALYTICS_ENGINE_INDEX_FIELDS).toEqual([ + "schema_version", + "event_family", + "metric_name", + "tenant_id", + "user_id", + ]); + expect(ANALYTICS_ENGINE_BLOB_FIELDS).toEqual([ + "metric_kind", + ...APPROVED_METRIC_LABEL_KEYS, + "deployment_environment", + "service_name", + "service_namespace", + "service_version", + ]); + expect(ANALYTICS_ENGINE_DOUBLE_FIELDS).toEqual([ + "metric_value", + "timestamp_ms", + ]); + }); + + it("writes one data point per observation", async () => { + const dataset = { writeDataPoint: vi.fn() }; + const sink = new AnalyticsEngineMetricsSink(dataset); + + await sink.flush({ + observations: [observation], + resourceAttributes: { + "deployment.environment": "local", + }, + }); + + expect(dataset.writeDataPoint).toHaveBeenCalledTimes(1); + expect(dataset.writeDataPoint).toHaveBeenCalledWith( + expect.objectContaining({ + indexes: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + null, + null, + ], + doubles: [2, 1_714_000_000_000], + }), + ); + }); + + it("maps request-scoped event identity into Analytics Engine indexes", async () => { + const dataset = { writeDataPoint: vi.fn() }; + const sink = new AnalyticsEngineMetricsSink(dataset); + + await sink.flush({ + observations: [observation], + resourceAttributes: {}, + eventIdentity: { + tenantId: "tenant-123", + userId: "user-456", + }, + }); + + expect(dataset.writeDataPoint).toHaveBeenCalledWith( + expect.objectContaining({ + indexes: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + "tenant-123", + "user-456", + ], + }), + ); + }); + + it("continues writing remaining data points when one write fails", async () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const secondObservation = { + ...observation, + name: METRIC_NAMES.httpRequestsTotal, + } satisfies MetricObservation; + const dataset = { + writeDataPoint: vi.fn((dataPoint) => { + if (dataPoint?.indexes?.[2] === METRIC_NAMES.toolCallsTotal) { + throw new Error("write failed"); + } + }), + }; + const sink = new AnalyticsEngineMetricsSink(dataset); + + await sink.flush({ + observations: [observation, secondObservation], + resourceAttributes: {}, + }); + + expect(dataset.writeDataPoint).toHaveBeenCalledTimes(2); + expect(warnSpy).toHaveBeenCalledWith( + `[metrics] Failed to write Analytics Engine data point for ${METRIC_NAMES.toolCallsTotal}`, + expect.any(Error), + ); + }); + + it("only creates a sink when the Analytics Engine binding is present", () => { + expect( + createAnalyticsEngineMetricsSink({ writeDataPoint: vi.fn() }), + ).toBeInstanceOf(AnalyticsEngineMetricsSink); + expect(createAnalyticsEngineMetricsSink(undefined)).toBeUndefined(); + expect(createAnalyticsEngineMetricsSink({})).toBeUndefined(); + expect( + createAnalyticsEngineMetricsSink({ writeDataPoint: "not-a-function" }), + ).toBeUndefined(); + }); +}); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index 89b80ca..2d27f8c 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -222,6 +222,26 @@ describe("withRequestMetrics", () => { ); }); + it("uses the Analytics Engine binding as the default analytics sink", async () => { + const analyticsDataset = { writeDataPoint: vi.fn() }; + const recorder = createRequestMetricsRecorder({ + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + route_group: "mcp", + }); + await recorder.flush(); + + expect(analyticsDataset.writeDataPoint).toHaveBeenCalledTimes(1); + expect(analyticsDataset.writeDataPoint).toHaveBeenCalledWith( + expect.objectContaining({ + indexes: expect.arrayContaining([METRIC_NAMES.httpRequestsTotal]), + }), + ); + }); + it("does not fail the request when waitUntil rejects scheduling", async () => { const errorSpy = vi .spyOn(console, "error") From 19c20f9973d02a3113f9b5431455b4621e5bd14f Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Tue, 5 May 2026 12:34:28 -0700 Subject: [PATCH 15/25] Batch storage operations in Storage Server (#135) - Use a single transaction where possible for reading and writing multiple fields, to improve performance and consistency - Update tests Co-authored-by: Rifdhan Nazeer --- src/servers/conversation-storage-server.ts | 118 ++++++++---- test/metrics/mixpanel/integration.spec.ts | 5 +- .../conversation-storage-server.spec.ts | 170 +++++++++++++++++- 3 files changed, 249 insertions(+), 44 deletions(-) diff --git a/src/servers/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts index 5b4621c..19ebd75 100644 --- a/src/servers/conversation-storage-server.ts +++ b/src/servers/conversation-storage-server.ts @@ -2,6 +2,7 @@ import { isBoolean, isNumber } from "lodash"; import type { Message, StreamingMessagesState } from "../thoughtspot/types"; const DEFAULT_TTL_MS = 30 * 60 * 1000; // 30 minutes +const STORAGE_BATCH_SIZE = 127; // Cloudflare DO bulk get/put limit is 128, we use 127 to be safe const MESSAGE_KEY_PREFIX = "message-"; const IS_DONE_KEY = "is-done"; @@ -87,10 +88,9 @@ export class ConversationStorageServer { /* * Append new messages to the conversation, starting at the current state of WRITE_BOOKMARK and - * saving the new state of WRITE_BOOKMARK after. We write the isDone flag state after writing - * all messages, so that if a reader is executing concurrently, they will never think the - * conversation is done without having already seen all messages. We also restart the TTL for - * the conversation after all writes are done. + * saving the new state of WRITE_BOOKMARK after. Writes are done un bulk, but batched if there + * are too many operations. The isDone flag is always in the last batch, so that any reader + * will never think the conversation is done before all messages have been written. */ private async appendMessagesAndRestartTtl( newMessages: Message[], @@ -107,49 +107,54 @@ export class ConversationStorageServer { } let idx = (await this.state.storage.get(WRITE_BOOKMARK_KEY)) ?? 0; + const entriesToStore = {} as Record; for (const message of newMessages) { - await this.state.storage.put( - `${MESSAGE_KEY_PREFIX}${idx}`, - message, - ); + entriesToStore[`${MESSAGE_KEY_PREFIX}${idx}`] = message; idx++; } - await this.state.storage.put(WRITE_BOOKMARK_KEY, idx); + entriesToStore[WRITE_BOOKMARK_KEY] = idx; if (isDone) { - await this.state.storage.put(IS_DONE_KEY, true); + entriesToStore[IS_DONE_KEY] = true; } + // Perform all writes in batches, then restart TTL + await this.putInBatches(entriesToStore); await this.restartTtl(); } /* * Retrieve all new messages since the last time this was called. We use a READ_BOOKMARK to * track the index of the last returned message, and update it when returning new messages. We - * read the isDone flag state before reading messages, so that if a writer is executing - * concurrently, we will only see isDone=true if all messages have already been written. Note - * that we don't restart the TTL here, since it is only meant to be based on writes. + * use WRITE_BOOKMARK to know up to which index new messages have been written. */ private async getNewMessagesAndUpdateBookmark(): Promise { - const isDone = await this.state.storage.get(IS_DONE_KEY); + const [isDone, readBookmark, writeBookmark] = + await this.getIsDoneAndReadWriteBookmarks(); if (!isBoolean(isDone)) { throw new Error(`Conversation ${this.conversationId} not found`); } - let bookmark = - (await this.state.storage.get(READ_BOOKMARK_KEY)) ?? 0; + const keys = []; + for (let i = readBookmark; i < writeBookmark; i++) { + keys.push(MESSAGE_KEY_PREFIX + i); + } + const newMessages: Message[] = []; - while (true) { - const message = await this.state.storage.get( - `${MESSAGE_KEY_PREFIX}${bookmark}`, - ); + const messagesMap = await this.getInBatches(keys); + for (let i = readBookmark; i < writeBookmark; i++) { + const message = messagesMap.get(MESSAGE_KEY_PREFIX + i); if (!message) { - break; + console.warn( + `Expected message at index ${i} for conversation ${this.conversationId} not found`, + { readBookmark, writeBookmark }, + ); + continue; } newMessages.push(message); - bookmark++; } - await this.state.storage.put(READ_BOOKMARK_KEY, bookmark); + + await this.state.storage.put(READ_BOOKMARK_KEY, writeBookmark); return { messages: newMessages, @@ -157,37 +162,62 @@ export class ConversationStorageServer { }; } + /* + * Perform bulk get operations in batches up to STORAGE_BATCH_SIZE + */ + private async getInBatches(keys: string[]): Promise> { + const result = new Map(); + for (let i = 0; i < keys.length; i += STORAGE_BATCH_SIZE) { + const batch = keys.slice(i, i + STORAGE_BATCH_SIZE); + const batchResult = await this.state.storage.get(batch); + for (const [k, v] of batchResult) { + result.set(k, v); + } + } + return result; + } + + /* + * Perform bulk put operations in batches up to STORAGE_BATCH_SIZE + */ + private async putInBatches(entries: Record): Promise { + const keys = Object.keys(entries); + for (let i = 0; i < keys.length; i += STORAGE_BATCH_SIZE) { + const batchKeys = keys.slice(i, i + STORAGE_BATCH_SIZE); + const batch = Object.fromEntries(batchKeys.map((k) => [k, entries[k]])); + await this.state.storage.put(batch); + } + } + + /* + * Restart TTL timer by canceling any old alarm and scheduling a new one for DEFAULT_TTL_MS + */ private async restartTtl(): Promise { - // Cancel any existing alarm and schedule a fresh one await this.state.storage.deleteAlarm(); await this.state.storage.setAlarm(Date.now() + DEFAULT_TTL_MS); } async alarm(): Promise { // Check for any abnormalities in the state prior to deleting - const isDone = await this.state.storage.get(IS_DONE_KEY); + const [isDone, readBookmark, writeBookmark] = + await this.getIsDoneAndReadWriteBookmarks(); if (!isBoolean(isDone) || !isDone) { console.warn( `Conversation ${this.conversationId} expired without being marked done`, { isDone, + readBookmark, + writeBookmark, }, ); } - const writeBookmark = - await this.state.storage.get(WRITE_BOOKMARK_KEY); - const readBookmark = - await this.state.storage.get(READ_BOOKMARK_KEY); - if (!isNumber(writeBookmark)) { - console.warn( - `Conversation ${this.conversationId} expired without any messages written`, - ); - } else if (!isNumber(readBookmark) || writeBookmark !== readBookmark) { + if (writeBookmark !== readBookmark) { console.warn( `Conversation ${this.conversationId} expired with unread messages`, { - writeBookmark, + isDone, readBookmark, + writeBookmark, }, ); } @@ -195,4 +225,22 @@ export class ConversationStorageServer { // Delete everything in storage await this.state.storage.deleteAll(); } + + /* + * Retrieve 3 fields from storage in one transaction: isDone, readBookmark, and writeBookmark + */ + async getIsDoneAndReadWriteBookmarks(): Promise< + [boolean | undefined, number, number] + > { + const result = await this.state.storage.get([ + IS_DONE_KEY, + READ_BOOKMARK_KEY, + WRITE_BOOKMARK_KEY, + ]); + return [ + result.get(IS_DONE_KEY) as boolean, + (result.get(READ_BOOKMARK_KEY) as number) ?? 0, + (result.get(WRITE_BOOKMARK_KEY) as number) ?? 0, + ]; + } } diff --git a/test/metrics/mixpanel/integration.spec.ts b/test/metrics/mixpanel/integration.spec.ts index c412838..d3810f8 100644 --- a/test/metrics/mixpanel/integration.spec.ts +++ b/test/metrics/mixpanel/integration.spec.ts @@ -2,9 +2,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { MixpanelTracker } from "../../../src/metrics/mixpanel/mixpanel"; import type { SessionInfo } from "../../../src/thoughtspot/types"; -// Mock fetch globally for integration tests -global.fetch = vi.fn(); - describe("Mixpanel Integration Tests", () => { const mockSessionInfo: SessionInfo = { mixpanelToken: "test-mixpanel-token", @@ -28,6 +25,7 @@ describe("Mixpanel Integration Tests", () => { let consoleDebugSpy: any; beforeEach(() => { + vi.stubGlobal("fetch", vi.fn()); vi.clearAllMocks(); // Mock console methods properly @@ -37,6 +35,7 @@ describe("Mixpanel Integration Tests", () => { afterEach(() => { vi.restoreAllMocks(); + vi.unstubAllGlobals(); }); describe("end-to-end tracking", () => { diff --git a/test/servers/conversation-storage-server.spec.ts b/test/servers/conversation-storage-server.spec.ts index 34b25b6..3f022f4 100644 --- a/test/servers/conversation-storage-server.spec.ts +++ b/test/servers/conversation-storage-server.spec.ts @@ -19,12 +19,36 @@ function createMockStorage() { return alarm; }, storage: { - get: vi.fn(async (key: string): Promise => { - return store.get(key) as T | undefined; - }), - put: vi.fn(async (key: string, value: unknown): Promise => { - store.set(key, value); - }), + get: vi.fn( + async ( + keyOrKeys: string | string[], + ): Promise> => { + if (Array.isArray(keyOrKeys)) { + const result = new Map(); + for (const key of keyOrKeys) { + if (store.has(key)) { + result.set(key, store.get(key) as T); + } + } + return result; + } + return store.get(keyOrKeys) as T | undefined; + }, + ), + put: vi.fn( + async ( + keyOrEntries: string | Record, + value?: unknown, + ): Promise => { + if (typeof keyOrEntries === "string") { + store.set(keyOrEntries, value); + } else { + for (const [k, v] of Object.entries(keyOrEntries)) { + store.set(k, v); + } + } + }, + ), delete: vi.fn(async (keys: string[]): Promise => { for (const key of keys) { store.delete(key); @@ -81,6 +105,18 @@ const answerMessage: Message = { is_thinking: false, }; +// Generate an array of N simple text messages +function generateMessages(n: number): Message[] { + return Array.from({ length: n }, (_, i) => ({ + type: "text", + text: `Message ${i}`, + is_thinking: false, + })); +} + +// The storage batch size used by ConversationStorageServer (must match the constant in the source) +const STORAGE_BATCH_SIZE = 127; + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -252,6 +288,61 @@ describe("ConversationStorageServer", () => { ); expect(res.status).toBe(500); }); + + // ------------------------------------------------------------------- + // Batching + // ------------------------------------------------------------------- + + it("stores exactly STORAGE_BATCH_SIZE messages in a single put call", async () => { + // STORAGE_BATCH_SIZE - 1 messages + write-bookmark = STORAGE_BATCH_SIZE entries → 1 batch + const messages = generateMessages(STORAGE_BATCH_SIZE - 1); + await server.fetch(makeRequest("POST", "append", { messages })); + + expect(mock.storage.put).toHaveBeenCalledOnce(); + expect(mock.store.get("write-bookmark")).toBe(STORAGE_BATCH_SIZE - 1); + expect(mock.store.get("message-0")).toEqual(messages[0]); + expect(mock.store.get(`message-${STORAGE_BATCH_SIZE - 2}`)).toEqual( + messages[STORAGE_BATCH_SIZE - 2], + ); + }); + + it("splits into two put calls when messages exceed STORAGE_BATCH_SIZE", async () => { + // STORAGE_BATCH_SIZE messages + write-bookmark = STORAGE_BATCH_SIZE + 1 entries → 2 batches + const messages = generateMessages(STORAGE_BATCH_SIZE); + await server.fetch(makeRequest("POST", "append", { messages })); + + expect(mock.storage.put).toHaveBeenCalledTimes(2); + expect(mock.store.get("write-bookmark")).toBe(STORAGE_BATCH_SIZE); + for (let i = 0; i < STORAGE_BATCH_SIZE; i++) { + expect(mock.store.get(`message-${i}`)).toEqual(messages[i]); + } + }); + + it("splits into two put calls when isDone adds an extra entry over the batch limit", async () => { + // STORAGE_BATCH_SIZE messages + write-bookmark + is-done = STORAGE_BATCH_SIZE + 2 entries → 2 batches + const messages = generateMessages(STORAGE_BATCH_SIZE); + await server.fetch( + makeRequest("POST", "append", { messages, isDone: true }), + ); + + expect(mock.storage.put).toHaveBeenCalledTimes(2); + expect(mock.store.get("write-bookmark")).toBe(STORAGE_BATCH_SIZE); + expect(mock.store.get("is-done")).toBe(true); + for (let i = 0; i < STORAGE_BATCH_SIZE; i++) { + expect(mock.store.get(`message-${i}`)).toEqual(messages[i]); + } + }); + + it("correctly stores messages across three or more batches", async () => { + const count = STORAGE_BATCH_SIZE * 2 + 10; + const messages = generateMessages(count); + await server.fetch(makeRequest("POST", "append", { messages })); + + expect(mock.store.get("write-bookmark")).toBe(count); + for (let i = 0; i < count; i++) { + expect(mock.store.get(`message-${i}`)).toEqual(messages[i]); + } + }); }); // ------------------------------------------------------------------------- @@ -329,6 +420,73 @@ describe("ConversationStorageServer", () => { const res = await server.fetch(makeRequest("GET", "messages")); expect(res.status).toBe(500); }); + + // ------------------------------------------------------------------- + // Batching + // ------------------------------------------------------------------- + + it("retrieves exactly STORAGE_BATCH_SIZE messages in a single get call", async () => { + const messages = generateMessages(STORAGE_BATCH_SIZE); + await server.fetch(makeRequest("POST", "append", { messages })); + vi.clearAllMocks(); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + + // 1 get for getIsDoneAndReadWriteBookmarks + 1 get for the STORAGE_BATCH_SIZE message keys + expect(mock.storage.get).toHaveBeenCalledTimes(2); + expect(body.messages).toHaveLength(STORAGE_BATCH_SIZE); + expect(body.messages).toEqual(messages); + }); + + it("splits into two get calls when fetching more than STORAGE_BATCH_SIZE messages", async () => { + const count = STORAGE_BATCH_SIZE + 1; + const messages = generateMessages(count); + await server.fetch(makeRequest("POST", "append", { messages })); + vi.clearAllMocks(); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + + // 1 get for getIsDoneAndReadWriteBookmarks + 2 get batches for the message keys + expect(mock.storage.get).toHaveBeenCalledTimes(3); + expect(body.messages).toHaveLength(count); + expect(body.messages).toEqual(messages); + }); + + it("correctly retrieves messages across three or more get batches", async () => { + const count = STORAGE_BATCH_SIZE * 2 + 10; + const messages = generateMessages(count); + await server.fetch(makeRequest("POST", "append", { messages })); + vi.clearAllMocks(); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + + expect(body.messages).toHaveLength(count); + expect(body.messages).toEqual(messages); + }); + + it("only fetches new messages across batches after bookmark advances", async () => { + const firstBatch = generateMessages(STORAGE_BATCH_SIZE + 5); + await server.fetch( + makeRequest("POST", "append", { messages: firstBatch }), + ); + // Consume first batch + await server.fetch(makeRequest("GET", "messages")); + + const secondBatch = generateMessages(STORAGE_BATCH_SIZE + 3); + await server.fetch( + makeRequest("POST", "append", { messages: secondBatch }), + ); + vi.clearAllMocks(); + + const res = await server.fetch(makeRequest("GET", "messages")); + const body = (await res.json()) as StreamingMessagesState; + + expect(body.messages).toHaveLength(secondBatch.length); + expect(body.messages).toEqual(secondBatch); + }); }); // ------------------------------------------------------------------------- From d9ed7b7984911e1aeb62c3ecdf5205e601a321d8 Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Wed, 6 May 2026 00:02:32 -0700 Subject: [PATCH 16/25] SCAL-308280 Instrument request and auth health metrics (#134) * SCAL-308280 Instrument request and auth health metrics ## Summary Add the request-layer metrics for SCAL-308280 so MCP traffic and auth flows emit bounded counters and latency through the shared runtime recorder. ## What Changed - record ts_mcp_http_requests_total and ts_mcp_http_request_duration_ms at the worker edge with route, transport, auth mode, api surface, outcome, status class, and canonical api version labels - add helpers to resolve bounded api_version values through the version registry instead of copying raw query params - emit auth outcome counters for /authorize, /callback, /store-token, and /bearer/* / /token/* requests - add focused tests for request metric labeling and auth metric emission - include the small metrics-runtime review nits from PR #126 while touching the same files ## Notes - ts_mcp_http_inflight_requests remains intentionally unimplemented because the current request-scoped recorder cannot represent a real inflight gauge safely ## Validation - npm run lint -- src/index.ts src/bearer.ts src/handlers.ts src/metrics/runtime/request-metrics.ts src/metrics/runtime/grafana-otlp-sink.ts test/metrics/runtime/request-metrics.spec.ts test/bearer.spec.ts test/handlers.spec.ts - npm test -- --coverage.enabled false test/metrics/runtime/request-metrics.spec.ts test/bearer.spec.ts test/handlers.spec.ts * SCAL-308280 Fix API version labels in request metrics ## Summary - map resolved API versions to stable low-cardinality labels - reuse the parsed request URL in bearer auth handling - add coverage for date-based API version label mapping --- src/bearer.ts | 146 +++++++++++------ src/handlers.ts | 134 ++++++++++++++-- src/index.ts | 34 +++- src/metrics/runtime/grafana-otlp-sink.ts | 22 +-- src/metrics/runtime/request-metrics.ts | 159 ++++++++++++++++++- test/bearer.spec.ts | 39 ++++- test/handlers.spec.ts | 45 +++++- test/metrics/runtime/request-metrics.spec.ts | 135 ++++++++++++++++ 8 files changed, 636 insertions(+), 78 deletions(-) diff --git a/src/bearer.ts b/src/bearer.ts index e1c28db..3755d3a 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -1,8 +1,25 @@ import type { ThoughtSpotMCP } from "."; import type honoApp from "./handlers"; +import { + getMetricsRecorderFromExecutionContext, + recordBearerAuthRequestMetric, +} from "./metrics/runtime/request-metrics"; import { validateAndSanitizeUrl } from "./oauth-manager/oauth-utils"; import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; +type AuthRouteFamily = "bearer" | "token"; + +function getAuthMetricRouteGroup( + pathname: string, + authRouteFamily: AuthRouteFamily, +): "bearer_mcp" | "bearer_sse" | "token_mcp" | "token_sse" { + if (pathname.endsWith(PUBLIC_ROUTES.sse)) { + return authRouteFamily === "bearer" ? "bearer_sse" : "token_sse"; + } + + return authRouteFamily === "bearer" ? "bearer_mcp" : "token_mcp"; +} + /** * Handler function for bearer/token authentication endpoints * @param req - Incoming request @@ -11,65 +28,103 @@ import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; * @param MCPServer - MCP server instance * @param apiVersionOverride - Optional API version override (ignore value in request) */ -function handleTokenAuth( +async function handleTokenAuth( req: Request, env: Env, ctx: ExecutionContext, MCPServer: typeof ThoughtSpotMCP, apiVersionOverride?: string, -): Response | Promise { - const authHeader = req.headers.get("authorization"); - if (!authHeader) { - return new Response("Bearer token is required", { status: 400 }); - } + authRouteFamily: AuthRouteFamily = "token", +): Promise { + const recorder = getMetricsRecorderFromExecutionContext(ctx); + const url = new URL(req.url); + const authMetricRouteGroup = getAuthMetricRouteGroup( + url.pathname, + authRouteFamily, + ); - let accessToken = authHeader.split(" ")[1]; - let tsHost: string | null; + try { + const authHeader = req.headers.get("authorization"); + if (!authHeader) { + const response = new Response("Bearer token is required", { + status: 400, + }); + recordBearerAuthRequestMetric( + recorder, + req, + response.status, + authMetricRouteGroup, + ); + return response; + } - if (accessToken.includes("@")) { - [accessToken, tsHost] = accessToken.split("@"); - } else { - tsHost = req.headers.get("x-ts-host"); - } + let accessToken = authHeader.split(" ")[1]; + let tsHost: string | null; - if (!tsHost) { - return new Response( - "TS Host is required, either in the authorization header as 'token@ts-host' or as a separate 'x-ts-host' header", - { status: 400 }, - ); - } + if (accessToken.includes("@")) { + [accessToken, tsHost] = accessToken.split("@"); + } else { + tsHost = req.headers.get("x-ts-host"); + } - const clientName = - req.headers.get("x-ts-client-name") || "Bearer Token client"; + if (!tsHost) { + const response = new Response( + "TS Host is required, either in the authorization header as 'token@ts-host' or as a separate 'x-ts-host' header", + { status: 400 }, + ); + recordBearerAuthRequestMetric( + recorder, + req, + response.status, + authMetricRouteGroup, + ); + return response; + } - const url = new URL(req.url); + const clientName = + req.headers.get("x-ts-client-name") || "Bearer Token client"; - // Build props object - const props: any = { - accessToken: accessToken, - instanceUrl: validateAndSanitizeUrl(tsHost), - clientName, - }; - - // Resolve API version to use - const apiVersion = apiVersionOverride ?? url.searchParams.get("api-version"); - if (apiVersion) { - props.apiVersion = apiVersion; - } + // Build props object + const props: any = { + accessToken: accessToken, + instanceUrl: validateAndSanitizeUrl(tsHost), + clientName, + }; - (ctx as any).props = props; + // Resolve API version to use + const apiVersion = + apiVersionOverride ?? url.searchParams.get("api-version"); + if (apiVersion) { + props.apiVersion = apiVersion; + } - // Route to appropriate handler - const pathname = url.pathname; - if (pathname.endsWith(PUBLIC_ROUTES.mcp)) { - return MCPServer.serve(PUBLIC_ROUTES.mcp).fetch(req, env, ctx); - } + (ctx as any).props = props; - if (pathname.endsWith(PUBLIC_ROUTES.sse)) { - return MCPServer.serveSSE(PUBLIC_ROUTES.sse).fetch(req, env, ctx); - } + let response: Response; + const pathname = url.pathname; + if (pathname.endsWith(PUBLIC_ROUTES.mcp)) { + response = await MCPServer.serve(PUBLIC_ROUTES.mcp).fetch(req, env, ctx); + } else if (pathname.endsWith(PUBLIC_ROUTES.sse)) { + response = await MCPServer.serveSSE(PUBLIC_ROUTES.sse).fetch( + req, + env, + ctx, + ); + } else { + response = new Response("Not found", { status: 404 }); + } - return new Response("Not found", { status: 404 }); + recordBearerAuthRequestMetric( + recorder, + req, + response.status, + authMetricRouteGroup, + ); + return response; + } catch (error) { + recordBearerAuthRequestMetric(recorder, req, 500, authMetricRouteGroup); + throw error; + } } export function withBearerHandler( @@ -85,13 +140,14 @@ export function withBearerHandler( ctx, MCPServer, "backwards-compatibility-default", + "bearer", ); }); // NEW: /token endpoints - supports api-version query params // Recommended for all new implementations app.mount(PUBLIC_ROUTE_PREFIXES.token, (req, env, ctx) => { - return handleTokenAuth(req, env, ctx, MCPServer); + return handleTokenAuth(req, env, ctx, MCPServer, undefined, "token"); }); return app; diff --git a/src/handlers.ts b/src/handlers.ts index dbaa2a5..0bcd02a 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -7,6 +7,11 @@ import { Hono } from "hono"; import { decodeBase64Url, encodeBase64Url } from "hono/utils/encode"; import { any } from "zod"; import { openApiSpecHandler } from "./api-schemas/open-api-spec"; +import { METRIC_NAMES } from "./metrics/runtime/metric-types"; +import { + getMetricsRecorderFromExecutionContext, + recordStatusMetric, +} from "./metrics/runtime/request-metrics"; import { WithSpan, getActiveSpan } from "./metrics/tracing/tracing-utils"; import { buildSamlRedirectUrl, @@ -20,6 +25,37 @@ import { McpServerError } from "./utils"; const app = new Hono<{ Bindings: Env & { OAUTH_PROVIDER: OAuthHelpers } }>(); +function getExecutionContextOrUndefined(context: { + executionCtx: ExecutionContext; +}): ExecutionContext | undefined { + try { + return context.executionCtx; + } catch { + return undefined; + } +} + +function recordAuthFlowMetric( + context: { executionCtx: ExecutionContext }, + name: + | typeof METRIC_NAMES.oauthAuthorizeRequestsTotal + | typeof METRIC_NAMES.oauthAuthorizeSubmitTotal + | typeof METRIC_NAMES.oauthCallbackTotal + | typeof METRIC_NAMES.oauthStoreTokenTotal, + status: number, +): void { + const executionContext = getExecutionContextOrUndefined(context); + if (!executionContext) { + return; + } + + recordStatusMetric( + getMetricsRecorderFromExecutionContext(executionContext), + name, + status, + ); +} + class Handler { @WithSpan("serve-index") async serveIndex(env: Env) { @@ -252,24 +288,55 @@ app.get(PUBLIC_ROUTES.authorize, async (c) => { c.req.raw, c.env.OAUTH_PROVIDER, ); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthAuthorizeRequestsTotal, + response.status, + ); return response; } catch (error) { - return c.text(`Internal Server Error ${error}`, 500); + const response = c.text(`Internal Server Error ${error}`, 500); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthAuthorizeRequestsTotal, + response.status, + ); + return response; } }); app.post(PUBLIC_ROUTES.authorize, async (c) => { try { const redirectUrl = await handler.postAuthorize(c.req.raw, c.req.url); - return Response.redirect(redirectUrl); + const response = Response.redirect(redirectUrl); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthAuthorizeSubmitTotal, + response.status, + ); + return response; } catch (error) { if ( error instanceof Error && error.message.includes("Missing instance URL") ) { - return new Response("Missing instance URL", { status: 400 }); + const response = new Response("Missing instance URL", { status: 400 }); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthAuthorizeSubmitTotal, + response.status, + ); + return response; } - return new Response(`Internal Server Error ${error}`, { status: 500 }); + const response = new Response(`Internal Server Error ${error}`, { + status: 500, + }); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthAuthorizeSubmitTotal, + response.status, + ); + return response; } }); @@ -280,53 +347,94 @@ app.get(PUBLIC_ROUTES.callback, async (c) => { c.env.ASSETS, c.req.url, ); - return new Response(htmlContent, { + const response = new Response(htmlContent, { headers: { "Content-Type": "text/html", }, }); + recordAuthFlowMetric(c, METRIC_NAMES.oauthCallbackTotal, response.status); + return response; } catch (error) { if (error instanceof Error) { if (error.message.includes("Missing instance URL")) { - return c.text(`Missing instance URL ${error}`, 400); + const response = c.text(`Missing instance URL ${error}`, 400); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthCallbackTotal, + response.status, + ); + return response; } if (error.message.includes("Missing OAuth request info")) { - return c.text(`Missing OAuth request info ${error}`, 400); + const response = c.text(`Missing OAuth request info ${error}`, 400); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthCallbackTotal, + response.status, + ); + return response; } if (error.message.includes("Invalid OAuth request info format")) { - return c.text(`Invalid OAuth request info format ${error}`, 400); + const response = c.text( + `Invalid OAuth request info format ${error}`, + 400, + ); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthCallbackTotal, + response.status, + ); + return response; } } - return c.text(`Internal server error ${error}`, 500); + const response = c.text(`Internal server error ${error}`, 500); + recordAuthFlowMetric(c, METRIC_NAMES.oauthCallbackTotal, response.status); + return response; } }); app.post(PUBLIC_ROUTES.storeToken, async (c) => { try { const result = await handler.storeToken(c.req.raw, c.env.OAUTH_PROVIDER); - return new Response(JSON.stringify(result), { + const response = new Response(JSON.stringify(result), { status: 200, headers: { "Content-Type": "application/json", }, }); + recordAuthFlowMetric(c, METRIC_NAMES.oauthStoreTokenTotal, response.status); + return response; } catch (error) { if (error instanceof Error) { if (error.message.includes("Invalid JSON format")) { - return c.text(`Invalid JSON format ${error}`, 400); + const response = c.text(`Invalid JSON format ${error}`, 400); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthStoreTokenTotal, + response.status, + ); + return response; } if ( error.message.includes( "Missing token or OAuth request info or instanceUrl", ) ) { - return c.text( + const response = c.text( `Missing token or OAuth request info or instanceUrl ${error}`, 400, ); + recordAuthFlowMetric( + c, + METRIC_NAMES.oauthStoreTokenTotal, + response.status, + ); + return response; } } - return c.text(`Internal server error ${error}`, 500); + const response = c.text(`Internal server error ${error}`, 500); + recordAuthFlowMetric(c, METRIC_NAMES.oauthStoreTokenTotal, response.status); + return response; } }); diff --git a/src/index.ts b/src/index.ts index 9b9e550..e5ea47f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -9,7 +9,10 @@ import { trace } from "@opentelemetry/api"; import { withBearerHandler } from "./bearer"; import { instrumentedMCPServer } from "./cloudflare-utils"; import handler from "./handlers"; -import { withRequestMetrics } from "./metrics/runtime/request-metrics"; +import { + recordHttpRequestMetrics, + withRequestMetrics, +} from "./metrics/runtime/request-metrics"; import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; import { apiServer } from "./servers/api-server"; import { ConversationStorageServer } from "./servers/conversation-storage-server"; @@ -147,8 +150,33 @@ export default { return withRequestMetrics( env as unknown as Record, ctx, - async () => { - return instrumentedOAuthHandler.fetch!(request, env, ctx); + async (recorder) => { + const requestStartMs = Date.now(); + + try { + const response = await instrumentedOAuthHandler.fetch!( + request, + env, + ctx, + ); + recordHttpRequestMetrics( + recorder, + request, + response, + ctx, + Date.now() - requestStartMs, + ); + return response; + } catch (error) { + recordHttpRequestMetrics( + recorder, + request, + new Response(null, { status: 500 }), + ctx, + Date.now() - requestStartMs, + ); + throw error; + } }, ); }, diff --git a/src/metrics/runtime/grafana-otlp-sink.ts b/src/metrics/runtime/grafana-otlp-sink.ts index eb10c43..14fe1bc 100644 --- a/src/metrics/runtime/grafana-otlp-sink.ts +++ b/src/metrics/runtime/grafana-otlp-sink.ts @@ -153,15 +153,19 @@ function toOtlpValue(value: MetricLabelValue): OtlpAttributeValue { function toOtlpAttributes( attributes: Record, ): OtlpAttribute[] { - return Object.keys(attributes) - .sort() - .flatMap((key) => { - const value = attributes[key]; - if (value === undefined || value === "") { - return []; - } - return [{ key, value: toOtlpValue(value) }]; - }); + return ( + Object.keys(attributes) + // Keep attribute ordering deterministic so identical metric series + // aggregate the same way regardless of insertion order upstream. + .sort() + .flatMap((key) => { + const value = attributes[key]; + if (value === undefined || value === "") { + return []; + } + return [{ key, value: toOtlpValue(value) }]; + }) + ); } function toTimeUnixNano(timestampMs: number): string { diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 617f8a6..5fbda7c 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,5 +1,13 @@ -import { createGrafanaOtlpMetricsSink } from "./grafana-otlp-sink"; +import { resolveApiVersion } from "../../servers/version-registry"; import { createAnalyticsEngineMetricsSink } from "./analytics-engine-sink"; +import { createGrafanaOtlpMetricsSink } from "./grafana-otlp-sink"; +import { getStatusClass, resolveRequestMetricContext } from "./metric-context"; +import { + METRIC_NAMES, + type MetricLabelInput, + type MetricName, + type MetricOutcome, +} from "./metric-types"; import { type MetricsRecorder, NOOP_METRICS_RECORDER, @@ -17,9 +25,25 @@ const METRICS_RECORDER_SYMBOL = Symbol.for( "thoughtspot.mcp.metrics.requestRecorder", ); const GRAFANA_SINK_CACHE = new WeakMap(); +const VERSIONED_REQUEST_ROUTE_GROUPS = new Set([ + "mcp", + "sse", + "bearer_mcp", + "bearer_sse", + "token_mcp", + "token_sse", +] as const); +type BearerAuthRouteGroup = + | "bearer_mcp" + | "bearer_sse" + | "token_mcp" + | "token_sse"; type MetricsExecutionContext = ExecutionContext & { [METRICS_RECORDER_SYMBOL]?: MetricsRecorder; + props?: { + apiVersion?: unknown; + }; }; export function setMetricsRecorderOnExecutionContext( @@ -54,6 +78,137 @@ export function clearMetricsRecorderFromExecutionContext( delete (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL]; } +export function getMetricOutcomeForStatus(status: number): MetricOutcome { + if (status >= 400 && status < 500) { + return "client_error"; + } + if (status >= 500) { + return "error"; + } + return "success"; +} + +function getCanonicalResolvedApiVersion(apiVersion: string): string { + const versionConfig = resolveApiVersion(apiVersion); + if (versionConfig.version.includes("beta")) { + return "beta"; + } + if (versionConfig.version.includes("backwards-compatibility-default")) { + return "default"; + } + if (versionConfig.version.includes("latest")) { + return "latest"; + } + + return "unknown"; +} + +export function resolveCanonicalApiVersionLabel( + request: Request, + ctx: ExecutionContext, +): string | undefined { + const requestContext = resolveRequestMetricContext(request); + if (!VERSIONED_REQUEST_ROUTE_GROUPS.has(requestContext.routeGroup)) { + return undefined; + } + + const requestedApiVersion = new URL(request.url).searchParams.get( + "api-version", + ); + const effectiveApiVersion = (ctx as MetricsExecutionContext).props + ?.apiVersion; + if ( + typeof effectiveApiVersion === "string" && + effectiveApiVersion.length > 0 + ) { + try { + return getCanonicalResolvedApiVersion(effectiveApiVersion); + } catch { + return "unknown"; + } + } + + if (!requestedApiVersion) { + return "default"; + } + + try { + return getCanonicalResolvedApiVersion(requestedApiVersion); + } catch { + return "unknown"; + } +} + +export function recordStatusMetric( + recorder: MetricsRecorder | undefined, + name: MetricName, + status: number, + labels: MetricLabelInput = {}, +): void { + if (!recorder) { + return; + } + + recorder.count(name, 1, { + ...labels, + outcome: getMetricOutcomeForStatus(status), + }); +} + +export function recordBearerAuthRequestMetric( + recorder: MetricsRecorder | undefined, + request: Request, + status: number, + routeGroupOverride?: BearerAuthRouteGroup, +): void { + if (!recorder) { + return; + } + + const requestContext = resolveRequestMetricContext(request); + recordStatusMetric(recorder, METRIC_NAMES.bearerAuthRequestsTotal, status, { + route_group: routeGroupOverride ?? requestContext.routeGroup, + transport: routeGroupOverride?.endsWith("_sse") + ? "sse" + : routeGroupOverride?.endsWith("_mcp") + ? "mcp" + : requestContext.transport, + }); +} + +export function recordHttpRequestMetrics( + recorder: MetricsRecorder, + request: Request, + response: Response, + ctx: ExecutionContext, + durationMs: number, +): void { + const requestContext = resolveRequestMetricContext(request); + const outcome = getMetricOutcomeForStatus(response.status); + const apiVersion = resolveCanonicalApiVersionLabel(request, ctx); + const baseLabels: MetricLabelInput = { + route_group: requestContext.routeGroup, + transport: requestContext.transport, + auth_mode: requestContext.authMode, + api_surface: requestContext.apiSurface, + outcome, + }; + + if (apiVersion) { + baseLabels.api_version = apiVersion; + } + + recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { + ...baseLabels, + status_class: getStatusClass(response.status), + }); + recorder.histogram( + METRIC_NAMES.httpRequestDurationMs, + durationMs, + baseLabels, + ); +} + function getCachedGrafanaSink( env: MetricsEnvLike | undefined, ): MetricsSink | undefined { @@ -61,6 +216,8 @@ function getCachedGrafanaSink( return createGrafanaOtlpMetricsSink(env); } + // Reuse the sink for the same env object so repeated request recorders do not + // rebuild identical Grafana exporter configuration within one Worker runtime. const cachedSink = GRAFANA_SINK_CACHE.get(env); if (cachedSink) { return cachedSink; diff --git a/test/bearer.spec.ts b/test/bearer.spec.ts index d351f91..f7ecfc4 100644 --- a/test/bearer.spec.ts +++ b/test/bearer.spec.ts @@ -1,8 +1,13 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { withBearerHandler } from "../src/bearer"; -import { ThoughtSpotMCP } from "../src"; import { Hono } from "hono"; -import { encodeBase64Url, decodeBase64Url } from "hono/utils/encode"; +import { decodeBase64Url, encodeBase64Url } from "hono/utils/encode"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ThoughtSpotMCP } from "../src"; +import { withBearerHandler } from "../src/bearer"; +import { METRIC_NAMES } from "../src/metrics/runtime/metric-types"; +import { + createRequestMetricsRecorder, + setMetricsRecorderOnExecutionContext, +} from "../src/metrics/runtime/request-metrics"; // For correctly-typed Request const IncomingRequest = Request; @@ -212,6 +217,32 @@ describe("Bearer Handler", () => { }); describe("Authorization Header Parsing", () => { + it("records bearer auth metrics for rejected requests", async () => { + const appWithBearer = withBearerHandler(app, ThoughtSpotMCP); + const recorder = createRequestMetricsRecorder(); + setMetricsRecorderOnExecutionContext( + mockCtx as ExecutionContext, + recorder, + ); + + const request = new Request("https://example.com/bearer/mcp"); + const result = await appWithBearer.fetch(request, mockEnv, mockCtx); + + expect(result.status).toBe(400); + expect(recorder.snapshot()).toContainEqual( + expect.objectContaining({ + kind: "counter", + name: METRIC_NAMES.bearerAuthRequestsTotal, + value: 1, + labels: { + outcome: "client_error", + route_group: "bearer_mcp", + transport: "mcp", + }, + }), + ); + }); + it("should return 400 when authorization header is missing", async () => { const appWithBearer = withBearerHandler(app, ThoughtSpotMCP); diff --git a/test/handlers.spec.ts b/test/handlers.spec.ts index 0ac34e6..3eed907 100644 --- a/test/handlers.spec.ts +++ b/test/handlers.spec.ts @@ -1,18 +1,23 @@ import { + createExecutionContext, env, runInDurableObject, - createExecutionContext, waitOnExecutionContext, } from "cloudflare:test"; -import { describe, it, expect, vi, beforeEach } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import worker, { ThoughtSpotMCP } from "../src"; import app from "../src/handlers"; +import { METRIC_NAMES } from "../src/metrics/runtime/metric-types"; +import { + createRequestMetricsRecorder, + setMetricsRecorderOnExecutionContext, +} from "../src/metrics/runtime/request-metrics"; // Type assertion for worker to have fetch method const typedWorker = worker as { fetch: (request: Request, env: any, ctx: any) => Promise; }; -import { encodeBase64Url, decodeBase64Url } from "hono/utils/encode"; +import { decodeBase64Url, encodeBase64Url } from "hono/utils/encode"; // For correctly-typed Request const IncomingRequest = Request; @@ -137,6 +142,40 @@ describe("Handlers", () => { }); describe("POST /authorize", () => { + it("records authorize submit metrics for client errors", async () => { + const recorder = createRequestMetricsRecorder(); + setMetricsRecorderOnExecutionContext( + mockCtx as ExecutionContext, + recorder, + ); + const formData = new FormData(); + formData.append( + "state", + btoa(JSON.stringify({ oauthReqInfo: { clientId: "test" } })), + ); + + const result = await app.fetch( + new Request("https://example.com/authorize", { + method: "POST", + body: formData, + }), + mockEnv, + mockCtx, + ); + + expect(result.status).toBe(400); + expect(recorder.snapshot()).toContainEqual( + expect.objectContaining({ + kind: "counter", + name: METRIC_NAMES.oauthAuthorizeSubmitTotal, + value: 1, + labels: { + outcome: "client_error", + }, + }), + ); + }); + it("should return 400 for missing instance URL", async () => { const id = env.MCP_OBJECT.idFromName("test"); const object = env.MCP_OBJECT.get(id); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index 2d27f8c..38e83d7 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -4,6 +4,10 @@ import { clearMetricsRecorderFromExecutionContext, createRequestMetricsRecorder, getMetricsRecorderFromExecutionContext, + recordBearerAuthRequestMetric, + recordHttpRequestMetrics, + recordStatusMetric, + resolveCanonicalApiVersionLabel, setMetricsRecorderOnExecutionContext, withRequestMetrics, } from "../../../src/metrics/runtime/request-metrics"; @@ -354,4 +358,135 @@ describe("withRequestMetrics", () => { expect(getMetricsRecorderFromExecutionContext(ctx)).toBeUndefined(); }); + + it("records HTTP request metrics with canonical route and version labels", () => { + const recorder = createRequestMetricsRecorder(); + const ctx = { + props: { + apiVersion: "beta", + }, + } as unknown as ExecutionContext; + const request = new Request("https://example.com/mcp?api-version=beta"); + const response = new Response("ok", { status: 200 }); + + recordHttpRequestMetrics(recorder, request, response, ctx, 123); + + expect(recorder.snapshot()).toEqual([ + expect.objectContaining({ + kind: "counter", + name: METRIC_NAMES.httpRequestsTotal, + value: 1, + labels: { + api_surface: "mcp", + api_version: "beta", + auth_mode: "oauth", + outcome: "success", + route_group: "mcp", + status_class: "2xx", + transport: "mcp", + }, + }), + expect.objectContaining({ + kind: "histogram", + name: METRIC_NAMES.httpRequestDurationMs, + value: 123, + labels: { + api_surface: "mcp", + api_version: "beta", + auth_mode: "oauth", + outcome: "success", + route_group: "mcp", + transport: "mcp", + }, + }), + ]); + }); + + it("labels versioned MCP routes as default when no explicit API version is requested", () => { + const ctx = {} as ExecutionContext; + const request = new Request("https://example.com/token/mcp"); + + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + }); + + it("uses the effective default surface when bearer routes ignore an api-version query", () => { + const ctx = { + props: { + apiVersion: "backwards-compatibility-default", + }, + } as unknown as ExecutionContext; + const request = new Request( + "https://example.com/bearer/mcp?api-version=beta", + ); + + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + }); + + it("maps stable date-based versions onto the latest label", () => { + const ctx = {} as ExecutionContext; + const request = new Request( + "https://example.com/token/mcp?api-version=2026-05-01", + ); + + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("latest"); + }); + + it("maps older date-based versions onto the default label", () => { + const ctx = {} as ExecutionContext; + const request = new Request( + "https://example.com/token/mcp?api-version=2025-12-01", + ); + + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + }); + + it("labels unresolved api-version values as unknown", () => { + const ctx = { + props: { + apiVersion: "garbage", + }, + } as unknown as ExecutionContext; + const request = new Request( + "https://example.com/token/mcp?api-version=garbage", + ); + + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("unknown"); + }); + + it("records auth outcome counters from response status", () => { + const recorder = createRequestMetricsRecorder(); + + recordStatusMetric(recorder, METRIC_NAMES.oauthAuthorizeSubmitTotal, 302); + + expect(recorder.snapshot()).toEqual([ + expect.objectContaining({ + kind: "counter", + name: METRIC_NAMES.oauthAuthorizeSubmitTotal, + value: 1, + labels: { + outcome: "success", + }, + }), + ]); + }); + + it("records bearer auth traffic with route and transport labels", () => { + const recorder = createRequestMetricsRecorder(); + const request = new Request("https://example.com/token/sse"); + + recordBearerAuthRequestMetric(recorder, request, 401); + + expect(recorder.snapshot()).toEqual([ + expect.objectContaining({ + kind: "counter", + name: METRIC_NAMES.bearerAuthRequestsTotal, + value: 1, + labels: { + outcome: "client_error", + route_group: "token_sse", + transport: "sse", + }, + }), + ]); + }); }); From b95de51396499f3448f6a34e41279ec1359200a8 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Wed, 6 May 2026 13:47:31 -0700 Subject: [PATCH 17/25] Remove deprecated servers: API Server, OpenAI MCP Server (#139) - Delete relevant code and update tests Co-authored-by: Rifdhan Nazeer --- src/api-schemas/open-api-spec.ts | 170 ------ src/handlers.ts | 3 - src/index.ts | 22 +- src/metrics/runtime/metric-context.ts | 59 +- src/routes.ts | 8 - src/servers/api-server.ts | 191 ------ src/servers/openai-mcp-server.ts | 199 ------- test/metrics/runtime/metric-context.spec.ts | 21 - test/servers/api-server.spec.ts | 535 ----------------- test/servers/openai-mcp-server.spec.ts | 606 -------------------- worker-configuration.d.ts | 3 +- wrangler.jsonc | 4 - 12 files changed, 5 insertions(+), 1816 deletions(-) delete mode 100644 src/api-schemas/open-api-spec.ts delete mode 100644 src/servers/api-server.ts delete mode 100644 src/servers/openai-mcp-server.ts delete mode 100644 test/servers/api-server.spec.ts delete mode 100644 test/servers/openai-mcp-server.spec.ts diff --git a/src/api-schemas/open-api-spec.ts b/src/api-schemas/open-api-spec.ts deleted file mode 100644 index 9a0cf81..0000000 --- a/src/api-schemas/open-api-spec.ts +++ /dev/null @@ -1,170 +0,0 @@ -import { Hono } from "hono"; -import { toolDefinitionsV1 } from "../servers/tool-definitions"; -import { capitalize } from "../utils"; - -export const openApiSpecHandler = new Hono(); - -// Helper function to generate tool schema -const generateToolSchema = (tool: (typeof toolDefinitionsV1)[0]) => { - const schemaName = `${capitalize(tool.name)}Request`; - const generatedSchema = { ...tool.inputSchema } as any; - generatedSchema.$schema = undefined; - return { schemaName, schema: generatedSchema }; -}; - -// Helper function to generate response schema -const generateResponseSchema = () => { - return { - type: "object", - description: "Response from the API endpoint", - }; -}; - -// Create individual endpoints for each tool -for (const tool of toolDefinitionsV1) { - const { schemaName, schema } = generateToolSchema(tool); - const responseSchema = generateResponseSchema(); - - openApiSpecHandler.get(`/tools/${tool.name}`, async (c) => { - const toolSpec = { - openapi: "3.0.0", - info: { - title: "ThoughtSpot API", - version: "1.0.0", - description: "API for interacting with ThoughtSpot services", - }, - servers: [ - { - url: "", - description: "ThoughtSpot agent url", - }, - ], - paths: { - [`/api/tools/${tool.name}`]: { - post: { - summary: tool.description, - description: tool.description, - operationId: tool.name, - tags: ["Tools"], - requestBody: { - required: true, - content: { - "application/json": { - schema: { - $ref: `#/components/schemas/${schemaName}`, - }, - }, - }, - }, - responses: { - "200": { - description: "Successful response", - content: { - "application/json": { - schema: { - $ref: `#/components/schemas/${capitalize(tool.name)}Response`, - }, - }, - }, - }, - "400": { - description: "Bad request - Invalid input parameters", - }, - "401": { - description: "Unauthorized - Invalid or missing authentication", - }, - "500": { - description: "Internal server error", - }, - }, - }, - }, - }, - components: { - schemas: { - [schemaName]: schema, - [`${capitalize(tool.name)}Response`]: responseSchema, - }, - }, - }; - - return c.json(toolSpec); - }); -} - -// Main OpenAPI spec endpoint that combines all tools -openApiSpecHandler.get("/", async (c) => { - const paths: Record = {}; - const schemas: Record = {}; - - // any tool added to the toolDefinitionsMCPServer will be added to the openapi spec automatically - // the api server path should be /api/tools/ - for (const tool of toolDefinitionsV1) { - const { schemaName, schema } = generateToolSchema(tool); - const responseSchema = generateResponseSchema(); - - schemas[schemaName] = schema; - schemas[`${capitalize(tool.name)}Response`] = responseSchema; - - paths[`/api/tools/${tool.name}`] = { - post: { - summary: tool.description, - description: tool.description, - operationId: tool.name, - tags: ["Tools"], - requestBody: { - required: true, - content: { - "application/json": { - schema: { - $ref: `#/components/schemas/${schemaName}`, - }, - }, - }, - }, - responses: { - "200": { - description: "Successful response", - content: { - "application/json": { - schema: { - $ref: `#/components/schemas/${capitalize(tool.name)}Response`, - }, - }, - }, - }, - "400": { - description: "Bad request - Invalid input parameters", - }, - "401": { - description: "Unauthorized - Invalid or missing authentication", - }, - "500": { - description: "Internal server error", - }, - }, - }, - }; - } - - const openApiDocument = { - openapi: "3.0.0", - info: { - title: "ThoughtSpot API", - version: "1.0.0", - description: "API for interacting with ThoughtSpot services", - }, - servers: [ - { - url: "", - description: "ThoughtSpot agent url", - }, - ], - paths: paths, - components: { - schemas: schemas, - }, - }; - - return c.json(openApiDocument); -}); diff --git a/src/handlers.ts b/src/handlers.ts index 0bcd02a..0ade3fa 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -6,7 +6,6 @@ import { type Span, SpanStatusCode, context, trace } from "@opentelemetry/api"; import { Hono } from "hono"; import { decodeBase64Url, encodeBase64Url } from "hono/utils/encode"; import { any } from "zod"; -import { openApiSpecHandler } from "./api-schemas/open-api-spec"; import { METRIC_NAMES } from "./metrics/runtime/metric-types"; import { getMetricsRecorderFromExecutionContext, @@ -442,6 +441,4 @@ app.get(PUBLIC_ROUTES.openaiAppsChallenge, (c) => { return c.text(process.env.OPEN_AI_TOKEN); }); -app.route(PUBLIC_ROUTES.openapiSpec, openApiSpecHandler); - export default app; diff --git a/src/index.ts b/src/index.ts index e5ea47f..f3af551 100644 --- a/src/index.ts +++ b/src/index.ts @@ -13,11 +13,9 @@ import { recordHttpRequestMetrics, withRequestMetrics, } from "./metrics/runtime/request-metrics"; -import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; -import { apiServer } from "./servers/api-server"; +import { PUBLIC_ROUTES } from "./routes"; import { ConversationStorageServer } from "./servers/conversation-storage-server"; import { MCPServer } from "./servers/mcp-server"; -import { OpenAIDeepResearchMCPServer } from "./servers/openai-mcp-server"; export { ConversationStorageServer }; @@ -35,11 +33,6 @@ const config: ResolveConfigFn = (env: Env, _trigger) => { // Create the instrumented ThoughtSpotMCP for the main export export const ThoughtSpotMCP = instrumentedMCPServer(MCPServer, config); -export const ThoughtSpotOpenAIDeepResearchMCP = instrumentedMCPServer( - OpenAIDeepResearchMCPServer, - config, -); - // Router function to handle query params and inject apiVersion into props function createMCPRouter( path: string, @@ -88,19 +81,6 @@ const oauthProvider = new OAuthProvider({ ThoughtSpotMCP, "serveSSE", ) as any, - [PUBLIC_ROUTES.openaiMcp]: ThoughtSpotOpenAIDeepResearchMCP.serve( - PUBLIC_ROUTES.openaiMcp, - { - binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", - }, - ) as any, // TODO: Remove 'any' - [PUBLIC_ROUTES.openaiSse]: ThoughtSpotOpenAIDeepResearchMCP.serveSSE( - PUBLIC_ROUTES.openaiSse, - { - binding: "OPENAI_DEEP_RESEARCH_MCP_OBJECT", - }, - ) as any, // TODO: Remove 'any' - [PUBLIC_ROUTE_PREFIXES.api]: apiServer as any, // TODO: Remove 'any' }, defaultHandler: withBearerHandler(handler, ThoughtSpotMCP) as any, // TODO: Remove 'any' authorizeEndpoint: PUBLIC_ROUTES.authorize, diff --git a/src/metrics/runtime/metric-context.ts b/src/metrics/runtime/metric-context.ts index 5b704e8..b3fb4d0 100644 --- a/src/metrics/runtime/metric-context.ts +++ b/src/metrics/runtime/metric-context.ts @@ -72,18 +72,6 @@ export const EXPLICIT_ROUTE_CONTEXTS = { apiSurface: "mcp", authMode: "oauth", }, - [PUBLIC_ROUTES.openaiMcp]: { - routeGroup: "openai_mcp", - transport: "mcp", - apiSurface: "openai_mcp", - authMode: "oauth", - }, - [PUBLIC_ROUTES.openaiSse]: { - routeGroup: "openai_sse", - transport: "sse", - apiSurface: "openai_mcp", - authMode: "oauth", - }, [PUBLIC_ROUTES.bearerMcp]: { routeGroup: "bearer_mcp", transport: "mcp", @@ -114,28 +102,8 @@ export const EXPLICIT_ROUTE_CONTEXTS = { apiSurface: "static", authMode: "none", }, - [PUBLIC_ROUTES.openapiSpec]: { - routeGroup: "openapi_spec", - transport: "http", - apiSurface: "static", - authMode: "none", - }, } as const satisfies Record; -const API_ROUTE_CONTEXT: RequestMetricContext = { - routeGroup: "api", - transport: "http", - apiSurface: "api", - authMode: "oauth", -}; - -const OPENAPI_SPEC_ROUTE_CONTEXT: RequestMetricContext = { - routeGroup: "openapi_spec", - transport: "http", - apiSurface: "static", - authMode: "none", -}; - const UNKNOWN_ROUTE_CONTEXT: RequestMetricContext = { routeGroup: "unknown", transport: "http", @@ -166,12 +134,6 @@ function inferTransport(pathname: string): Transport { } function inferApiSurface(pathname: string): ApiSurface { - if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api)) { - return "api"; - } - if (pathname.startsWith("/openai/")) { - return "openai_mcp"; - } if ( pathname === PUBLIC_ROUTES.mcp || pathname === PUBLIC_ROUTES.sse || @@ -183,8 +145,7 @@ function inferApiSurface(pathname: string): ApiSurface { if ( pathname === PUBLIC_ROUTES.root || pathname === PUBLIC_ROUTES.hello || - pathname === PUBLIC_ROUTES.openaiAppsChallenge || - matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec) + pathname === PUBLIC_ROUTES.openaiAppsChallenge ) { return "static"; } @@ -207,12 +168,7 @@ function inferAuthMode(pathname: string): AuthMode { if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.token)) { return "token"; } - if ( - pathname === PUBLIC_ROUTES.mcp || - pathname === PUBLIC_ROUTES.sse || - pathname.startsWith("/openai/") || - matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api) - ) { + if (pathname === PUBLIC_ROUTES.mcp || pathname === PUBLIC_ROUTES.sse) { return "oauth"; } if ( @@ -223,8 +179,7 @@ function inferAuthMode(pathname: string): AuthMode { pathname === PUBLIC_ROUTES.storeToken || pathname === PUBLIC_ROUTES.oauthToken || pathname === PUBLIC_ROUTES.register || - pathname === PUBLIC_ROUTES.openaiAppsChallenge || - matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec) + pathname === PUBLIC_ROUTES.openaiAppsChallenge ) { return "none"; } @@ -239,14 +194,6 @@ export function resolvePathMetricContext( return explicitContext; } - if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.api)) { - return API_ROUTE_CONTEXT; - } - - if (matchesRoutePrefix(pathname, PUBLIC_ROUTE_PREFIXES.openapiSpec)) { - return OPENAPI_SPEC_ROUTE_CONTEXT; - } - return { ...UNKNOWN_ROUTE_CONTEXT, transport: inferTransport(pathname), diff --git a/src/routes.ts b/src/routes.ts index 0f768e5..7952a2f 100644 --- a/src/routes.ts +++ b/src/routes.ts @@ -8,20 +8,15 @@ export const PUBLIC_ROUTES = { register: "/register", mcp: "/mcp", sse: "/sse", - openaiMcp: "/openai/mcp", - openaiSse: "/openai/sse", bearerMcp: "/bearer/mcp", bearerSse: "/bearer/sse", tokenMcp: "/token/mcp", tokenSse: "/token/sse", openaiAppsChallenge: "/.well-known/openai-apps-challenge", - openapiSpec: "/openapi-spec", } as const; export const PUBLIC_ROUTE_PREFIXES = { - api: "/api", bearer: "/bearer", - openapiSpec: PUBLIC_ROUTES.openapiSpec, token: "/token", } as const; @@ -35,12 +30,9 @@ export const EXACT_PUBLIC_ROUTES_REQUIRING_METRICS = [ PUBLIC_ROUTES.register, PUBLIC_ROUTES.mcp, PUBLIC_ROUTES.sse, - PUBLIC_ROUTES.openaiMcp, - PUBLIC_ROUTES.openaiSse, PUBLIC_ROUTES.bearerMcp, PUBLIC_ROUTES.bearerSse, PUBLIC_ROUTES.tokenMcp, PUBLIC_ROUTES.tokenSse, PUBLIC_ROUTES.openaiAppsChallenge, - PUBLIC_ROUTES.openapiSpec, ] as const; diff --git a/src/servers/api-server.ts b/src/servers/api-server.ts deleted file mode 100644 index c621f50..0000000 --- a/src/servers/api-server.ts +++ /dev/null @@ -1,191 +0,0 @@ -import { Hono } from "hono"; -import { zValidator } from "@hono/zod-validator"; -import type { Props } from "../utils"; -import { ThoughtSpotService } from "../thoughtspot/thoughtspot-service"; -import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client"; -import { getActiveSpan, WithSpan } from "../metrics/tracing/tracing-utils"; -import { - CreateLiveboardSchema, - GetAnswerSchema, - GetRelevantQuestionsSchema, -} from "./tool-definitions"; - -const apiServer = new Hono<{ Bindings: Env & { props: Props } }>(); - -class ApiHandler { - private initSpan(props: Props) { - const span = getActiveSpan(); - span?.setAttributes({ - instance_url: props.instanceUrl, - }); - } - - private getThoughtSpotService(props: Props): ThoughtSpotService { - this.initSpan(props); - return new ThoughtSpotService( - getThoughtSpotClient(props.instanceUrl, props.accessToken), - ); - } - - @WithSpan("api-relevant-questions") - async getRelevantQuestions( - props: Props, - query: string, - datasourceIds: string[], - additionalContext?: string, - ) { - const service = this.getThoughtSpotService(props); - return await service.getRelevantQuestions( - query, - datasourceIds, - additionalContext || "", - ); - } - - @WithSpan("api-get-answer") - async getAnswer(props: Props, question: string, datasourceId: string) { - const service = this.getThoughtSpotService(props); - return await service.getAnswerForQuestion(question, datasourceId, false); - } - - @WithSpan("api-create-liveboard") - async createLiveboard( - props: Props, - name: string, - answers: any[], - noteTileParsedHtml: string, - ) { - const service = this.getThoughtSpotService(props); - const result = await service.fetchTMLAndCreateLiveboard( - name, - answers, - noteTileParsedHtml, - ); - return result.url || ""; - } - - @WithSpan("api-get-datasources") - async getDataSources(props: Props) { - const service = this.getThoughtSpotService(props); - return await service.getDataSources(); - } - - @WithSpan("api-proxy-post") - async proxyPost(props: Props, path: string, body: any) { - const span = getActiveSpan(); - span?.setAttributes({ - instance_url: props.instanceUrl, - path: path, - }); - span?.addEvent("proxy-post"); - return fetch(props.instanceUrl + path, { - method: "POST", - headers: { - Authorization: `Bearer ${props.accessToken}`, - Accept: "application/json", - "Content-Type": "application/json", - "User-Agent": "ThoughtSpot-ts-client", - }, - body: JSON.stringify(body), - }); - } - - @WithSpan("api-proxy-get") - async proxyGet(props: Props, path: string) { - const span = getActiveSpan(); - span?.setAttributes({ - instance_url: props.instanceUrl, - path: path, - }); - span?.addEvent("proxy-get"); - return fetch(props.instanceUrl + path, { - method: "GET", - headers: { - Authorization: `Bearer ${props.accessToken}`, - Accept: "application/json", - "User-Agent": "ThoughtSpot-ts-client", - }, - }); - } -} - -const handler = new ApiHandler(); - -apiServer.post( - "/api/tools/relevant-questions", - zValidator("json", GetRelevantQuestionsSchema), - async (c) => { - const { props } = c.executionCtx; - const { query, datasourceIds, additionalContext } = c.req.valid("json"); - const questions = await handler.getRelevantQuestions( - props, - query, - datasourceIds, - additionalContext, - ); - return c.json(questions); - }, -); - -apiServer.post( - "/api/tools/get-answer", - zValidator("json", GetAnswerSchema), - async (c) => { - const { props } = c.executionCtx; - const { question, datasourceId } = c.req.valid("json"); - const answer = await handler.getAnswer(props, question, datasourceId); - return c.json(answer); - }, -); - -apiServer.post( - "/api/tools/create-liveboard", - zValidator("json", CreateLiveboardSchema), - async (c) => { - const { props } = c.executionCtx; - const { name, answers, noteTile } = c.req.valid("json"); - const liveboardUrl = await handler.createLiveboard( - props, - name, - answers, - noteTile, - ); - return c.text(liveboardUrl); - }, -); - -apiServer.get("/api/tools/ping", async (c) => { - const { props } = c.executionCtx; - console.log("Received Ping request"); - if (props.accessToken && props.instanceUrl) { - return c.json({ - content: [{ type: "text", text: "Pong" }], - }); - } - return c.json({ - isError: true, - content: [{ type: "text", text: "ERROR: Not authenticated" }], - }); -}); - -apiServer.get("/api/resources/datasources", async (c) => { - const { props } = c.executionCtx; - const datasources = await handler.getDataSources(props); - return c.json(datasources); -}); - -apiServer.post("/api/rest/2.0/*", async (c) => { - const { props } = c.executionCtx; - const path = c.req.path; - const method = c.req.method; - const body = await c.req.json(); - return handler.proxyPost(props, path, body); -}); - -apiServer.get("/api/rest/2.0/*", async (c) => { - const { props } = c.executionCtx; - const path = c.req.path; - return handler.proxyGet(props, path); -}); - -export { apiServer }; diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts deleted file mode 100644 index 8883882..0000000 --- a/src/servers/openai-mcp-server.ts +++ /dev/null @@ -1,199 +0,0 @@ -import type { - CallToolRequestSchema, - ReadResourceRequestSchema, -} from "@modelcontextprotocol/sdk/types.js"; -import { BaseMCPServer, type Context } from "./mcp-server-base"; -import { z } from "zod"; -import { WithSpan } from "../metrics/tracing/tracing-utils"; -import zodToJsonSchema from "zod-to-json-schema"; -import { ToolSchema } from "@modelcontextprotocol/sdk/types.js"; - -const ToolInputSchema = ToolSchema.shape.inputSchema; -export type ToolInput = z.infer; - -const ToolOutputSchema = ToolSchema.shape.outputSchema; -export type ToolOutput = z.infer; - -export const SearchInputSchema = z.object({ - query: z - .string() - .describe(`The question/task to search for relevant data queries to answer. Use the fetch tool to retrieve the data for individual queries. The datasource id should be passed as part of the query. With the syntax - datasource: . The search-query can be any textual question. - - For example: - datasource:asdhshd-123123-12dd How to reduce customer churn? - datasource:abc-123123-12dd How to increase sales? - - If the datasource id is not available, ask the user to supply one explicitly.`), -}); - -export const SearchOutputSchema = z.object({ - results: z.array( - z.object({ - id: z.string().describe("The id of the search result."), - title: z.string().describe("The title of the search result."), - text: z.string().describe("The text of the search result."), - url: z.string().describe("The url of the search result."), - }), - ), -}); - -export const FetchInputSchema = z.object({ - id: z.string().describe("The id of the search result to fetch."), -}); - -export const FetchOutputSchema = z.object({ - id: z.string().describe("The id of the search result."), - title: z.string().describe("The title of the search result."), - text: z.string().describe("The text of the search result."), - url: z.string().describe("The url of the search result."), -}); - -export const toolDefinitionsOpenAIMCPServer = [ - { - name: "search", - description: - "Tool to search for relevant data queries to answer the given question based on the datasource passed to this tool, which is a datasource id, see the query description for the syntax. The datasource id is mandatory and should be passed as part of the query. Any textual question can be passed to this tool, and it will do its best to find relevant data queries to answer the question.", - inputSchema: zodToJsonSchema(SearchInputSchema) as ToolInput, - outputSchema: zodToJsonSchema(SearchOutputSchema) as ToolOutput, - }, - { - name: "fetch", - description: - "Tool to retrieve data from the retail sales dataset for a given query.", - inputSchema: zodToJsonSchema(FetchInputSchema) as ToolInput, - outputSchema: zodToJsonSchema(FetchOutputSchema) as ToolOutput, - }, -]; - -export class OpenAIDeepResearchMCPServer extends BaseMCPServer { - constructor(ctx: Context) { - super(ctx, "ThoughtSpot", "1.0.0"); - } - - protected async listTools() { - return { - tools: [...toolDefinitionsOpenAIMCPServer], - }; - } - - protected async listResources() { - return { - resources: [], - }; - } - - protected async readResource( - request: z.infer, - ) { - return { - contents: [], - }; - } - - protected async callTool(request: z.infer) { - const { name } = request.params; - switch (name) { - case "search": - return this.callSearch(request); - case "fetch": - return this.callFetch(request); - } - } - - @WithSpan("call-search") - protected async callSearch(request: z.infer) { - const { query } = SearchInputSchema.parse(request.params.arguments); - // query could be of the form "datasource: " or just "" - // First check if the query is of the form "datasource: . The id is a string of numbers, letters, and hyphens." - const re = /^(?:datasource:(?[A-Za-z0-9-]+)\s+)?(.+)$/; - const match = re.exec(query); - const datasourceId = match?.groups?.id; - const queryWithoutDatasourceId = match![2]; - if (datasourceId) { - const relevantQuestions = - await this.getThoughtSpotService().getRelevantQuestions( - queryWithoutDatasourceId, - [datasourceId], - "", - ); - if (relevantQuestions.error) { - return this.createErrorResponse( - relevantQuestions.error.message, - `Error getting relevant questions ${relevantQuestions.error.message}`, - ); - } - - if (relevantQuestions.questions.length === 0) { - return this.createSuccessResponse("No relevant questions found"); - } - - const results = relevantQuestions.questions.map((q) => ({ - id: `${datasourceId}: ${q.question}`, - title: q.question, - text: q.question, - url: "", - })); - - return this.createStructuredContentSuccessResponse( - { results }, - "Relevant questions found", - ); - } - // Search for datasources in case the query is not of the form "datasource: " - if (!this.isDatasourceDiscoveryAvailable()) { - return this.createStructuredContentSuccessResponse( - { results: [] }, - "No relevant questions found", - ); - } - const dataSources = - await this.getThoughtSpotService().getDataSourceSuggestions( - queryWithoutDatasourceId, - ); - if (!dataSources || dataSources.length === 0) { - return this.createSuccessResponse( - "No relevant data sources found, please provide a datasource id in the query", - ); - } - const results = dataSources.map((d) => ({ - id: `datasource:///${d.header.guid}`, - title: d.header.displayName, - text: `Datasource Description: ${d.header.description}. Confidence that this datasource is relevant to the query: ${d.confidence}. Reasoning for the confidence: ${d.llmReasoning}. - Use this datasource to search for relevant questions and to get answers for the questions. - Use the search tool to search for relevant questions with the format "datasource: " and the fetch tool to get answers for the questions.`, - })); - - return this.createStructuredContentSuccessResponse( - { results }, - "Relevant questions found", - ); - } - - @WithSpan("call-fetch") - protected async callFetch(request: z.infer) { - const { id } = FetchInputSchema.parse(request.params.arguments); - // id is of the form ":" - const [datasourceId, question = ""] = id.split(":"); - const answer = await this.getThoughtSpotService().getAnswerForQuestion( - question, - datasourceId, - false, - ); - if (answer.error) { - return this.createErrorResponse( - answer.error.message, - `Error getting answer ${answer.error.message}`, - ); - } - - const result = { - id, - title: question, - text: answer.data, - url: `${this.ctx.props.instanceUrl}/#/insights/conv-assist?query=${question.trim()}&worksheet=${datasourceId}&executeSearch=true`, - }; - - return this.createStructuredContentSuccessResponse(result, "Answer found"); - } -} diff --git a/test/metrics/runtime/metric-context.spec.ts b/test/metrics/runtime/metric-context.spec.ts index 7fee837..d7e69b9 100644 --- a/test/metrics/runtime/metric-context.spec.ts +++ b/test/metrics/runtime/metric-context.spec.ts @@ -12,7 +12,6 @@ import { import { EXACT_PUBLIC_ROUTES_REQUIRING_METRICS, PUBLIC_ROUTES, - PUBLIC_ROUTE_PREFIXES, } from "../../../src/routes"; describe("metric-context", () => { @@ -34,12 +33,6 @@ describe("metric-context", () => { }); it("maps known grouped request paths to route groups", () => { - expect( - getRouteGroup(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), - ).toBe("api"); - expect( - getRouteGroup(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`), - ).toBe("openapi_spec"); expect(getRouteGroup("/not-a-route")).toBe("unknown"); }); @@ -50,13 +43,6 @@ describe("metric-context", () => { }); it("derives API surface from fallback request paths", () => { - expect(getApiSurface("/openai/future-endpoint")).toBe("openai_mcp"); - expect( - getApiSurface(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), - ).toBe("api"); - expect( - getApiSurface(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`), - ).toBe("static"); expect(getApiSurface("/bearer/future-endpoint")).toBe("mcp"); expect(getApiSurface("/token/future-endpoint")).toBe("mcp"); expect(getApiSurface(PUBLIC_ROUTES.root)).toBe("static"); @@ -69,13 +55,6 @@ describe("metric-context", () => { it("derives auth mode from fallback request paths", () => { expect(getAuthMode("/bearer/future-endpoint")).toBe("bearer"); expect(getAuthMode("/token/future-endpoint")).toBe("token"); - expect(getAuthMode(PUBLIC_ROUTES.openaiMcp)).toBe("oauth"); - expect( - getAuthMode(`${PUBLIC_ROUTE_PREFIXES.api}/resources/datasources`), - ).toBe("oauth"); - expect(getAuthMode(`${PUBLIC_ROUTE_PREFIXES.openapiSpec}/tools/ping`)).toBe( - "none", - ); expect(getAuthMode(PUBLIC_ROUTES.root)).toBe("none"); expect(getAuthMode(PUBLIC_ROUTES.authorize)).toBe("none"); expect(getAuthMode(PUBLIC_ROUTES.callback)).toBe("none"); diff --git a/test/servers/api-server.spec.ts b/test/servers/api-server.spec.ts deleted file mode 100644 index fc44aee..0000000 --- a/test/servers/api-server.spec.ts +++ /dev/null @@ -1,535 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { apiServer } from "../../src/servers/api-server"; -import { ThoughtSpotService } from "../../src/thoughtspot/thoughtspot-service"; -import * as thoughtspotClient from "../../src/thoughtspot/thoughtspot-client"; - -// Mock the ThoughtSpot service and client -vi.mock("../../src/thoughtspot/thoughtspot-service"); -vi.mock("../../src/thoughtspot/thoughtspot-client"); - -describe("API Server", () => { - let mockClient: any; - let mockProps: any; - let mockServiceInstance: any; - - beforeEach(() => { - // Reset all mocks - vi.clearAllMocks(); - - // Mock the ThoughtSpot client - mockClient = { - instanceUrl: "https://test.thoughtspot.cloud", - }; - vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue( - mockClient, - ); - - // Mock ThoughtSpotService instance methods - mockServiceInstance = { - getRelevantQuestions: vi.fn(), - getAnswerForQuestion: vi.fn(), - fetchTMLAndCreateLiveboard: vi.fn(), - getDataSources: vi.fn(), - }; - - // Mock the ThoughtSpotService constructor - vi.mocked(ThoughtSpotService).mockImplementation(() => mockServiceInstance); - - // Mock props - mockProps = { - instanceUrl: "https://test.thoughtspot.cloud", - accessToken: "test-access-token", - }; - }); - - // Helper function to create a mock execution context - const createMockExecutionContext = (props: any) => ({ - props, - waitUntil: vi.fn(), - passThroughOnException: vi.fn(), - }); - - describe("POST /api/tools/relevant-questions", () => { - it("should return relevant questions successfully", async () => { - const mockQuestions = { - questions: [ - { question: "What is the total revenue?", datasourceId: "ds-123" }, - { question: "How many customers?", datasourceId: "ds-456" }, - ], - error: null, - }; - - mockServiceInstance.getRelevantQuestions.mockResolvedValue(mockQuestions); - - const requestBody = { - query: "Show me revenue data", - datasourceIds: ["ds-123", "ds-456"], - additionalContext: "Previous analysis", - }; - - const request = new Request( - "http://localhost/api/tools/relevant-questions", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - const data = await response.json(); - expect(data).toEqual(mockQuestions); - expect(thoughtspotClient.getThoughtSpotClient).toHaveBeenCalledWith( - mockProps.instanceUrl, - mockProps.accessToken, - ); - expect(mockServiceInstance.getRelevantQuestions).toHaveBeenCalledWith( - requestBody.query, - requestBody.datasourceIds, - requestBody.additionalContext, - ); - }); - - it("should handle missing additionalContext", async () => { - const mockQuestions = { - questions: [ - { question: "What is the total revenue?", datasourceId: "ds-123" }, - ], - error: null, - }; - - mockServiceInstance.getRelevantQuestions.mockResolvedValue(mockQuestions); - - const requestBody = { - query: "Show me revenue data", - datasourceIds: ["ds-123"], - }; - - const request = new Request( - "http://localhost/api/tools/relevant-questions", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - const data = await response.json(); - expect(data).toEqual(mockQuestions); - expect(mockServiceInstance.getRelevantQuestions).toHaveBeenCalledWith( - requestBody.query, - requestBody.datasourceIds, - "", - ); - }); - }); - - describe("POST /api/tools/get-answer", () => { - it("should return answer successfully", async () => { - const mockAnswer = { - question: "What is the total revenue?", - data: "The total revenue is $1,000,000", - session_identifier: "session-123", - generation_number: 1, - tml: null, - error: null, - message_type: "TSAnswer", - } as any; - - mockServiceInstance.getAnswerForQuestion.mockResolvedValue(mockAnswer); - - const requestBody = { - question: "What is the total revenue?", - datasourceId: "ds-123", - }; - - const request = new Request("http://localhost/api/tools/get-answer", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - const data = await response.json(); - expect(data).toEqual(mockAnswer); - expect(thoughtspotClient.getThoughtSpotClient).toHaveBeenCalledWith( - mockProps.instanceUrl, - mockProps.accessToken, - ); - expect(mockServiceInstance.getAnswerForQuestion).toHaveBeenCalledWith( - requestBody.question, - requestBody.datasourceId, - false, - ); - }); - }); - - describe("POST /api/tools/create-liveboard", () => { - it("should create liveboard successfully", async () => { - const mockLiveboardUrl = - "https://test.thoughtspot.cloud/#/pinboard/liveboard-123"; - - mockServiceInstance.fetchTMLAndCreateLiveboard.mockResolvedValue({ - url: mockLiveboardUrl, - }); - - const requestBody = { - name: "Revenue Dashboard", - answers: [ - { - question: "What is the total revenue?", - session_identifier: "session-123", - generation_number: 1, - }, - ], - noteTile: - "

Revenue Dashboard

This is a revenue dashboard

", - }; - - const request = new Request( - "http://localhost/api/tools/create-liveboard", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - const data = await response.text(); - expect(data).toBe(mockLiveboardUrl); - expect(thoughtspotClient.getThoughtSpotClient).toHaveBeenCalledWith( - mockProps.instanceUrl, - mockProps.accessToken, - ); - expect( - mockServiceInstance.fetchTMLAndCreateLiveboard, - ).toHaveBeenCalledWith( - requestBody.name, - requestBody.answers, - requestBody.noteTile, - ); - }); - - it("should handle service errors", async () => { - const mockError = new Error("Failed to create liveboard"); - - mockServiceInstance.fetchTMLAndCreateLiveboard.mockRejectedValue( - mockError, - ); - - const requestBody = { - name: "Revenue Dashboard", - answers: [ - { - question: "What is the total revenue?", - session_identifier: "session-123", - generation_number: 1, - }, - ], - }; - - const request = new Request( - "http://localhost/api/tools/create-liveboard", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The endpoint should return a 500 error when the service throws - expect(response.status).toBe(400); - }); - }); - - describe("GET /api/resources/datasources", () => { - it("should return datasources successfully", async () => { - const mockDatasources = [ - { - name: "Sales Data", - id: "ds-123", - description: "Sales data for analysis", - }, - { - name: "Customer Data", - id: "ds-456", - description: "Customer information", - }, - ]; - - mockServiceInstance.getDataSources.mockResolvedValue(mockDatasources); - - const request = new Request( - "http://localhost/api/resources/datasources", - { - method: "GET", - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - const data = await response.json(); - expect(data).toEqual(mockDatasources); - expect(thoughtspotClient.getThoughtSpotClient).toHaveBeenCalledWith( - mockProps.instanceUrl, - mockProps.accessToken, - ); - expect(mockServiceInstance.getDataSources).toHaveBeenCalledWith(); - }); - - it("should handle service errors", async () => { - const mockError = new Error("Failed to fetch datasources"); - - mockServiceInstance.getDataSources.mockRejectedValue(mockError); - - const request = new Request( - "http://localhost/api/resources/datasources", - { - method: "GET", - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The endpoint should return a 500 error when the service throws - expect(response.status).toBe(500); - }); - }); - - describe("POST /api/rest/2.0/*", () => { - it("should proxy POST requests to ThoughtSpot API", async () => { - const mockFetchResponse = { - status: 200, - json: () => Promise.resolve({ success: true }), - }; - - global.fetch = vi.fn().mockResolvedValue(mockFetchResponse); - - const requestBody = { test: "data" }; - - const request = new Request( - "http://localhost/api/rest/2.0/test-endpoint", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - expect(global.fetch).toHaveBeenCalledWith( - `${mockProps.instanceUrl}/api/rest/2.0/test-endpoint`, - { - method: "POST", - headers: { - Authorization: `Bearer ${mockProps.accessToken}`, - Accept: "application/json", - "Content-Type": "application/json", - "User-Agent": "ThoughtSpot-ts-client", - }, - body: JSON.stringify(requestBody), - }, - ); - }); - - it("should handle fetch errors", async () => { - global.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const requestBody = { test: "data" }; - - const request = new Request( - "http://localhost/api/rest/2.0/test-endpoint", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The endpoint should return a 500 error when fetch throws - expect(response.status).toBe(500); - }); - }); - - describe("GET /api/rest/2.0/*", () => { - it("should proxy GET requests to ThoughtSpot API", async () => { - const mockFetchResponse = { - status: 200, - json: () => Promise.resolve({ success: true }), - }; - - global.fetch = vi.fn().mockResolvedValue(mockFetchResponse); - - const request = new Request( - "http://localhost/api/rest/2.0/test-endpoint", - { - method: "GET", - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - expect(response.status).toBe(200); - expect(global.fetch).toHaveBeenCalledWith( - `${mockProps.instanceUrl}/api/rest/2.0/test-endpoint`, - { - method: "GET", - headers: { - Authorization: `Bearer ${mockProps.accessToken}`, - Accept: "application/json", - "User-Agent": "ThoughtSpot-ts-client", - }, - }, - ); - }); - - it("should handle fetch errors", async () => { - global.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const request = new Request( - "http://localhost/api/rest/2.0/test-endpoint", - { - method: "GET", - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The endpoint should return a 500 error when fetch throws - expect(response.status).toBe(500); - }); - }); - - describe("Error handling", () => { - it("should handle malformed JSON in request body", async () => { - const request = new Request( - "http://localhost/api/tools/relevant-questions", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: "invalid json", - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The API server returns 400 for JSON parsing errors - expect(response.status).toBe(400); - }); - - it("should handle missing required fields", async () => { - const requestBody = { - // Missing required fields - }; - - const request = new Request( - "http://localhost/api/tools/relevant-questions", - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(requestBody), - }, - ); - - const response = await apiServer.fetch( - request, - { - props: mockProps, - }, - createMockExecutionContext(mockProps), - ); - - // The endpoint should return an error when required fields are missing - expect(response.status).toBe(400); - }); - }); -}); diff --git a/test/servers/openai-mcp-server.spec.ts b/test/servers/openai-mcp-server.spec.ts deleted file mode 100644 index f740860..0000000 --- a/test/servers/openai-mcp-server.spec.ts +++ /dev/null @@ -1,606 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { connect } from "mcp-testing-kit"; -import { OpenAIDeepResearchMCPServer } from "../../src/servers/openai-mcp-server"; -import * as thoughtspotService from "../../src/thoughtspot/thoughtspot-service"; -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(), - })), -})); - -describe("OpenAI Deep Research MCP Server", () => { - let server: OpenAIDeepResearchMCPServer; - let mockProps: any; - - beforeEach(() => { - // Reset all mocks - 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-10", - userGUID: "test-user-123", - configInfo: { - mixpanelConfig: { - devSdkKey: "test-dev-token", - prodSdkKey: "test-prod-token", - production: false, - }, - selfClusterName: "test-cluster", - selfClusterId: "test-cluster-123", - enableSpotterDataSourceDiscovery: true, - }, - userName: "test-user", - currentOrgId: "test-org", - privileges: [], - }), - singleAnswer: vi.fn().mockResolvedValue({ - session_identifier: "session-123", - generation_number: 1, - }), - exportAnswerReport: vi.fn().mockResolvedValue({ - text: vi.fn().mockResolvedValue("The total revenue is $1,000,000"), - }), - instanceUrl: "https://test.thoughtspot.cloud", - } as any); - - // Mock props with correct structure - mockProps = { - instanceUrl: "https://test.thoughtspot.cloud", - accessToken: "test-access-token", - clientName: { - clientId: "test-client-id", - clientName: "test-client", - registrationDate: Date.now(), - }, - }; - - server = new OpenAIDeepResearchMCPServer({ - props: mockProps, - }); - }); - - describe("Initialization", () => { - it("should initialize successfully with valid props", async () => { - await expect(server.init()).resolves.not.toThrow(); - }); - - it("should track initialization event", async () => { - await server.init(); - expect(MixpanelTracker).toHaveBeenCalledWith( - { - clusterId: "test-cluster-123", - clusterName: "test-cluster", - releaseVersion: "10.13.0.cl-10", - userGUID: "test-user-123", - mixpanelToken: "test-dev-token", - userName: "test-user", - currentOrgId: "test-org", - privileges: [], - enableSpotterDataSourceDiscovery: true, - }, - { - clientId: "test-client-id", - clientName: "test-client", - registrationDate: expect.any(Number), - }, - ); - }); - }); - - describe("List Tools", () => { - it("should return all available tools", async () => { - await server.init(); - const { listTools } = connect(server); - - const result = await listTools(); - - expect(result.tools).toHaveLength(2); - expect(result.tools?.map((t) => t.name)).toEqual(["search", "fetch"]); - }); - - it("should include correct tool descriptions", async () => { - await server.init(); - const { listTools } = connect(server); - - const result = await listTools(); - - const searchTool = result.tools?.find((t) => t.name === "search"); - expect(searchTool?.description).toBe( - "Tool to search for relevant data queries to answer the given question based on the datasource passed to this tool, which is a datasource id, see the query description for the syntax. The datasource id is mandatory and should be passed as part of the query. Any textual question can be passed to this tool, and it will do its best to find relevant data queries to answer the question.", - ); - - const fetchTool = result.tools?.find((t) => t.name === "fetch"); - expect(fetchTool?.description).toBe( - "Tool to retrieve data from the retail sales dataset for a given query.", - ); - }); - - it("should include correct input schemas", async () => { - await server.init(); - const { listTools } = connect(server); - - const result = await listTools(); - - const searchTool = result.tools?.find((t) => t.name === "search"); - expect(searchTool?.inputSchema).toMatchObject({ - type: "object", - properties: { - query: { - type: "string", - description: expect.stringContaining( - "The question/task to search for relevant data queries", - ), - }, - }, - required: ["query"], - }); - - const fetchTool = result.tools?.find((t) => t.name === "fetch"); - expect(fetchTool?.inputSchema).toMatchObject({ - type: "object", - properties: { - id: { - type: "string", - description: "The id of the search result to fetch.", - }, - }, - required: ["id"], - }); - }); - - it("should include correct output schemas", async () => { - await server.init(); - const { listTools } = connect(server); - - const result = await listTools(); - - const searchTool = result.tools?.find((t) => t.name === "search"); - expect(searchTool?.outputSchema).toMatchObject({ - type: "object", - properties: { - results: { - type: "array", - items: { - type: "object", - properties: { - id: { - type: "string", - description: "The id of the search result.", - }, - title: { - type: "string", - description: "The title of the search result.", - }, - text: { - type: "string", - description: "The text of the search result.", - }, - url: { - type: "string", - description: "The url of the search result.", - }, - }, - required: ["id", "title", "text", "url"], - }, - }, - }, - required: ["results"], - }); - - const fetchTool = result.tools?.find((t) => t.name === "fetch"); - expect(fetchTool?.outputSchema).toMatchObject({ - type: "object", - properties: { - id: { - type: "string", - description: "The id of the search result.", - }, - title: { - type: "string", - description: "The title of the search result.", - }, - text: { - type: "string", - description: "The text of the search result.", - }, - url: { - type: "string", - description: "The url of the search result.", - }, - }, - required: ["id", "title", "text", "url"], - }); - }); - }); - - describe("List Resources", () => { - it("should return empty resources list", async () => { - await server.init(); - const { listResources } = connect(server); - - const result = await listResources(); - - expect(result.resources).toHaveLength(0); - }); - }); - - describe("Read Resource", () => { - it("should return empty contents", async () => { - await server.init(); - - // Test the protected method directly since it's abstract - const result = await (server as any).readResource({ - method: "resources/read", - params: { uri: "datasource:///test-id" }, - }); - - expect(result.contents).toHaveLength(0); - }); - }); - - describe("Search Tool", () => { - it("should return relevant questions for query with datasource ID", async () => { - // Mock the ThoughtSpot service to return relevant questions - const mockGetRelevantQuestions = vi.fn().mockResolvedValue({ - questions: [ - { question: "What is the total revenue?" }, - { question: "How many customers do we have?" }, - ], - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getRelevantQuestions", - ).mockImplementation(mockGetRelevantQuestions); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("search", { - query: "datasource:asdhshd-123123-12dd How to reduce customer churn?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - results: [ - { - id: "asdhshd-123123-12dd: What is the total revenue?", - title: "What is the total revenue?", - text: "What is the total revenue?", - url: "", - }, - { - id: "asdhshd-123123-12dd: How many customers do we have?", - title: "How many customers do we have?", - text: "How many customers do we have?", - url: "", - }, - ], - }); - // The text field contains the JSON stringified structured content - expect((result.content as any[])[0].text).toContain('"results"'); - expect((result.content as any[])[0].text).toContain( - '"What is the total revenue?"', - ); - }); - - it("should handle error from ThoughtSpot service", async () => { - // Mock the ThoughtSpot service to return error - const mockGetRelevantQuestions = vi.fn().mockResolvedValue({ - questions: [], - error: { message: "Service unavailable" }, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getRelevantQuestions", - ).mockImplementation(mockGetRelevantQuestions); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("search", { - query: "datasource:asdhshd-123123-12dd How to reduce customer churn?", - }); - - expect(result.isError).toBe(true); - expect((result.content as any[])[0].text).toBe( - "ERROR: Service unavailable", - ); - }); - - it("should handle empty questions response", async () => { - // Mock the ThoughtSpot service to return empty questions - const mockGetRelevantQuestions = vi.fn().mockResolvedValue({ - questions: [], - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getRelevantQuestions", - ).mockImplementation(mockGetRelevantQuestions); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("search", { - query: "datasource:asdhshd-123123-12dd How to reduce customer churn?", - }); - - expect(result.isError).toBeUndefined(); - // When no questions found, it uses createSuccessResponse, not createStructuredContentSuccessResponse - expect((result.content as any[])[0].text).toBe( - "No relevant questions found", - ); - }); - - it("should handle query with complex datasource ID", async () => { - // Mock the ThoughtSpot service to return relevant questions - const mockGetRelevantQuestions = vi.fn().mockResolvedValue({ - questions: [{ question: "What is the total revenue?" }], - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getRelevantQuestions", - ).mockImplementation(mockGetRelevantQuestions); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("search", { - query: "datasource:abc-123-def-456 How to increase sales?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - results: [ - { - id: "abc-123-def-456: What is the total revenue?", - title: "What is the total revenue?", - text: "What is the total revenue?", - url: "", - }, - ], - }); - }); - - it("should handle query with mixed case datasource ID", async () => { - // Mock the ThoughtSpot service to return relevant questions - const mockGetRelevantQuestions = vi.fn().mockResolvedValue({ - questions: [{ question: "What is the total revenue?" }], - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getRelevantQuestions", - ).mockImplementation(mockGetRelevantQuestions); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("search", { - query: "datasource:ABC123def How to increase sales?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - results: [ - { - id: "ABC123def: What is the total revenue?", - title: "What is the total revenue?", - text: "What is the total revenue?", - url: "", - }, - ], - }); - }); - - it("should handle query without datasource ID for version 10.12 (no data source suggestions)", async () => { - // Mock version to be less than 10.13 - vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ - getSessionInfo: vi.fn().mockResolvedValue({ - clusterId: "test-cluster-123", - clusterName: "test-cluster", - releaseVersion: "10.12.0.cl-144", // Version < 10.13 - 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: [], - }), - } as any); - - const versionSpecificServer = new OpenAIDeepResearchMCPServer({ - props: { - instanceUrl: "https://test.thoughtspot.cloud", - accessToken: "test-access-token", - clientName: { - clientId: "test-client-id", - clientName: "test-client", - registrationDate: Date.now(), - }, - }, - }); - - await versionSpecificServer.init(); - const { callTool } = connect(versionSpecificServer); - - const result = await callTool("search", { - query: "How to reduce customer churn?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ results: [] }); - expect((result.content as any[])[0].text).toContain('"results"'); - expect((result.content as any[])[0].text).toContain("[]"); - }); - }); - - describe("Fetch Tool", () => { - it("should return answer for a valid question ID", async () => { - // Mock the ThoughtSpot service to return answer - const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ - data: "The total revenue is $1,000,000", - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getAnswerForQuestion", - ).mockImplementation(mockGetAnswerForQuestion); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("fetch", { - id: "asdhshd-123123-12dd: What is the total revenue?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - id: "asdhshd-123123-12dd: What is the total revenue?", - title: " What is the total revenue?", - text: "The total revenue is $1,000,000", - url: "https://test.thoughtspot.cloud/#/insights/conv-assist?query=What is the total revenue?&worksheet=asdhshd-123123-12dd&executeSearch=true", - }); - // The text field contains the JSON stringified structured content - expect((result.content as any[])[0].text).toContain('"id"'); - expect((result.content as any[])[0].text).toContain( - '"The total revenue is $1,000,000"', - ); - }); - - it("should handle error from ThoughtSpot service", async () => { - // Mock the ThoughtSpot service to return error - const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ - data: null, - error: { message: "Question not found" }, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getAnswerForQuestion", - ).mockImplementation(mockGetAnswerForQuestion); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("fetch", { - id: "asdhshd-123123-12dd: What is the total revenue?", - }); - - expect(result.isError).toBe(true); - expect((result.content as any[])[0].text).toBe( - "ERROR: Question not found", - ); - }); - - it("should handle ID with complex datasource ID", async () => { - // Mock the ThoughtSpot service to return answer - const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ - data: "The total revenue is $1,000,000", - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getAnswerForQuestion", - ).mockImplementation(mockGetAnswerForQuestion); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("fetch", { - id: "abc-123-def-456: What is the total revenue?", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - id: "abc-123-def-456: What is the total revenue?", - title: " What is the total revenue?", - text: "The total revenue is $1,000,000", - url: "https://test.thoughtspot.cloud/#/insights/conv-assist?query=What is the total revenue?&worksheet=abc-123-def-456&executeSearch=true", - }); - }); - - it("should handle ID with question containing special characters", async () => { - // Mock the ThoughtSpot service to return answer - const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ - data: "The revenue increased by 15%", - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getAnswerForQuestion", - ).mockImplementation(mockGetAnswerForQuestion); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("fetch", { - id: "ds-123: How much did revenue increase? (in %)", - }); - - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - id: "ds-123: How much did revenue increase? (in %)", - title: " How much did revenue increase? (in %)", - text: "The revenue increased by 15%", - url: "https://test.thoughtspot.cloud/#/insights/conv-assist?query=How much did revenue increase? (in %)&worksheet=ds-123&executeSearch=true", - }); - }); - }); - - describe("Error Handling", () => { - it("should handle empty fetch ID", async () => { - // Mock the ThoughtSpot service to return answer for empty ID test - const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ - data: "The total revenue is $1,000,000", - error: null, - }); - - vi.spyOn( - thoughtspotService.ThoughtSpotService.prototype, - "getAnswerForQuestion", - ).mockImplementation(mockGetAnswerForQuestion); - - await server.init(); - const { callTool } = connect(server); - - const result = await callTool("fetch", { - id: "", // Empty ID - }); - - // Empty ID will cause the split to return ["", ""], which results in empty datasourceId and undefined question - expect(result.isError).toBeUndefined(); - expect(result.structuredContent).toEqual({ - id: "", - title: "", - text: "The total revenue is $1,000,000", - url: "https://test.thoughtspot.cloud/#/insights/conv-assist?query=&worksheet=&executeSearch=true", - }); - }); - }); -}); diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 0b4c91b..8175aa3 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -4,12 +4,11 @@ declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/index"); - durableNamespaces: "ThoughtSpotMCP" | "ThoughtSpotOpenAIDeepResearchMCP"; + durableNamespaces: "ThoughtSpotMCP"; } interface Env { OAUTH_KV: KVNamespace; MCP_OBJECT: DurableObjectNamespace; - OPENAI_DEEP_RESEARCH_MCP_OBJECT: DurableObjectNamespace; CONVERSATION_STORAGE_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; ASSETS: Fetcher; diff --git a/wrangler.jsonc b/wrangler.jsonc index 0d51a86..7659219 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -17,10 +17,6 @@ "class_name": "ThoughtSpotMCP", "name": "MCP_OBJECT" }, - { - "class_name": "ThoughtSpotOpenAIDeepResearchMCP", - "name": "OPENAI_DEEP_RESEARCH_MCP_OBJECT" - }, { "class_name": "ConversationStorageServer", "name": "CONVERSATION_STORAGE_OBJECT" From 3872a621f665577c637359e7d7ab6dc989f176ee Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Wed, 6 May 2026 15:55:16 -0700 Subject: [PATCH 18/25] Add Cloudflare migration for OpenAI MCP Server deletion (#140) - Clean up a few leftover types Co-authored-by: Rifdhan Nazeer --- README.md | 1 - public/docs/glean-actions.md | 19 ------------------- src/metrics/runtime/metric-types.ts | 12 +----------- wrangler.jsonc | 6 ++++++ 4 files changed, 7 insertions(+), 31 deletions(-) delete mode 100644 public/docs/glean-actions.md diff --git a/README.md b/README.md index a1c91ad..ab2f46e 100644 --- a/README.md +++ b/README.md @@ -373,7 +373,6 @@ When adding new MCP tools to the server: **OAuth-based endpoints:** - `/mcp`: MCP HTTP Streaming endpoint (supports `?api-version`) - `/sse`: Server-sent events for MCP (supports `?api-version`) -- `/api`: MCP tools exposed as HTTP endpoints - `/authorize`, `/token`, `/register`: OAuth endpoints **Token-based endpoints (Recommended for APIs):** diff --git a/public/docs/glean-actions.md b/public/docs/glean-actions.md deleted file mode 100644 index 29bed54..0000000 --- a/public/docs/glean-actions.md +++ /dev/null @@ -1,19 +0,0 @@ -- Glean actions - -To configure the MCP server to be used via glean actions follow these steps: - -1. Create a separate action for each the tools exposed through the MCP server -2. Add open api spec for the tool that you are adding in the functionality section. The openapi spec for each tool is available on this url : https://agent.thoughtspot.app/openapi-spec/tools/{tool_name}. Tool name here would be the name given in the tool set below in [Features](#features). For example, Get relevant data questions tool has name getRelevantQuestions. Note: getDataSourceSuggestions is not yet available as openapi-spec as the feature is not yet Generally Available in ThoughtSpot for all customers. We will add this once it is available. -3. Select authentication type as Oauth User while configuring the action -4. Register the glean oauth server with TS MCP server - -```bash - curl 'https://agent.thoughtspot.app/register' \ - -H 'accept: */*' \ - -H 'accept-language: en-US,en;q=0.9' \ - --data-raw '{"redirect_uris":["${glean_callback_url}"],"token_endpoint_auth_method":"client_secret_basic","grant_types":["authorization_code","refresh_token"],"response_types":["code"],"client_name":"{company_glean_name}","client_uri":"${company_glean_uri}"}' -``` -5. Add the client_id and client secret obtained from to the glean action auth section. Along with this add https://agent.thoughtspot.app/authorize in client url and https://agent.thoughtspot.app/token in authorize url. -6. Save the spec and reload the action. -7. Once the action is saved, we will get an option to run API test. It is recommended to run it one time to make sure everything is setup. - diff --git a/src/metrics/runtime/metric-types.ts b/src/metrics/runtime/metric-types.ts index 58a72ac..4525c0e 100644 --- a/src/metrics/runtime/metric-types.ts +++ b/src/metrics/runtime/metric-types.ts @@ -151,11 +151,7 @@ export type RouteGroup = | "register" | "mcp" | "sse" - | "openai_mcp" - | "openai_sse" | "openai_apps_challenge" - | "openapi_spec" - | "api" | "bearer_mcp" | "bearer_sse" | "token_mcp" @@ -164,13 +160,7 @@ export type RouteGroup = export type Transport = "mcp" | "sse" | "http" | "unknown"; export type AuthMode = "oauth" | "bearer" | "token" | "none" | "unknown"; -export type ApiSurface = - | "mcp" - | "openai_mcp" - | "api" - | "oauth" - | "static" - | "unknown"; +export type ApiSurface = "mcp" | "oauth" | "static" | "unknown"; export type StatusClass = "1xx" | "2xx" | "3xx" | "4xx" | "5xx" | "unknown"; function warnOnInvalidMetricLabel(key: string, reason: string) { diff --git a/wrangler.jsonc b/wrangler.jsonc index 7659219..9b8d436 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -41,6 +41,12 @@ "new_classes": [ "ConversationStorageServer" ] + }, + { + "tag": "v5", + "deleted_classes": [ + "ThoughtSpotOpenAIDeepResearchMCP" + ] } ], "kv_namespaces": [{ From aa258740627dcb3fac4fb97f2f215701cf2c4ae6 Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Wed, 6 May 2026 17:58:30 -0700 Subject: [PATCH 19/25] SCAL-308281 Add tool and upstream health metrics (#138) --- src/bearer.ts | 31 +- src/index.ts | 14 +- src/metrics/runtime/analytics-engine-sink.ts | 23 +- src/metrics/runtime/metric-types.ts | 10 + src/metrics/runtime/metrics-recorder.ts | 69 ++++ src/metrics/runtime/metrics-sink.ts | 10 + src/metrics/runtime/request-metrics.ts | 198 +++++++++-- src/metrics/runtime/tool-metrics.ts | 163 +++++++++ src/servers/mcp-server-base.ts | 162 ++++++++- src/servers/mcp-server.ts | 192 +++++++---- src/servers/version-registry.ts | 51 ++- src/streaming-utils.ts | 39 ++- src/thoughtspot/thoughtspot-service.ts | 310 +++++++++++++----- src/utils.ts | 5 +- test/bearer.spec.ts | 26 +- .../runtime/analytics-engine-sink.spec.ts | 27 ++ test/metrics/runtime/metrics-recorder.spec.ts | 46 +++ test/metrics/runtime/request-metrics.spec.ts | 107 +++++- test/servers/mcp-server.spec.ts | 86 ++++- test/servers/version-registry.spec.ts | 36 +- test/streaming-utils.spec.ts | 78 ++++- test/thoughtspot/thoughtspot-service.spec.ts | 163 ++++++++- 22 files changed, 1624 insertions(+), 222 deletions(-) create mode 100644 src/metrics/runtime/tool-metrics.ts diff --git a/src/bearer.ts b/src/bearer.ts index 3755d3a..9b9478f 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -2,7 +2,9 @@ import type { ThoughtSpotMCP } from "."; import type honoApp from "./handlers"; import { getMetricsRecorderFromExecutionContext, + normalizeRequestedApiVersionForAnalytics, recordBearerAuthRequestMetric, + resolveRequestedApiVersionMode, } from "./metrics/runtime/request-metrics"; import { validateAndSanitizeUrl } from "./oauth-manager/oauth-utils"; import { PUBLIC_ROUTES, PUBLIC_ROUTE_PREFIXES } from "./routes"; @@ -90,13 +92,36 @@ async function handleTokenAuth( instanceUrl: validateAndSanitizeUrl(tsHost), clientName, }; - - // Resolve API version to use + const requestedApiVersion = url.searchParams.get("api-version"); + + // Stamp the effective served surface into props so downstream request/tool metrics + // can distinguish: + // - `api_version=backwards-compatibility-default` => tenants still on the legacy/v1 surface + // - `api_version_mode` => implicit route defaults vs explicit/latest/pinned selectors + // - `api_requested_version` (stored only in Analytics Engine) => the exact selector + // the client sent, which helps debug version-resolution confusion and future-dated pins + // - `api_release_date` (derived later from the registry) => which exact dated release is served const apiVersion = - apiVersionOverride ?? url.searchParams.get("api-version"); + apiVersionOverride ?? + requestedApiVersion ?? + (authRouteFamily === "token" ? "latest" : undefined); + const apiVersionMode = apiVersionOverride + ? "implicit_legacy" + : requestedApiVersion + ? resolveRequestedApiVersionMode(requestedApiVersion) + : authRouteFamily === "token" + ? "implicit_latest" + : undefined; + if (requestedApiVersion) { + props.apiRequestedVersion = + normalizeRequestedApiVersionForAnalytics(requestedApiVersion); + } if (apiVersion) { props.apiVersion = apiVersion; } + if (apiVersionMode) { + props.apiVersionMode = apiVersionMode; + } (ctx as any).props = props; diff --git a/src/index.ts b/src/index.ts index f3af551..a7873e3 100644 --- a/src/index.ts +++ b/src/index.ts @@ -9,8 +9,11 @@ import { trace } from "@opentelemetry/api"; import { withBearerHandler } from "./bearer"; import { instrumentedMCPServer } from "./cloudflare-utils"; import handler from "./handlers"; +import type { ApiVersionMode } from "./metrics/runtime/metric-types"; import { + normalizeRequestedApiVersionForAnalytics, recordHttpRequestMetrics, + resolveRequestedApiVersionMode, withRequestMetrics, } from "./metrics/runtime/request-metrics"; import { PUBLIC_ROUTES } from "./routes"; @@ -47,12 +50,17 @@ function createMCPRouter( ctx: ExecutionContext, ): Promise { const url = new URL(request.url); - let apiVersion = url.searchParams.get("api-version"); + const requestedApiVersion = url.searchParams.get("api-version"); + let apiVersion = requestedApiVersion; + let apiVersionMode: ApiVersionMode; // TODO(Rifdhan): this is a temporary backwards compatibility measure. In the future // we will use latest by default. if (!apiVersion) { apiVersion = "backwards-compatibility-default"; + apiVersionMode = "implicit_legacy"; + } else { + apiVersionMode = resolveRequestedApiVersionMode(apiVersion); } // Inject apiVersion into props @@ -60,6 +68,10 @@ function createMCPRouter( (ctx as any).props = { ...originalProps, apiVersion, + apiRequestedVersion: requestedApiVersion + ? normalizeRequestedApiVersionForAnalytics(requestedApiVersion) + : undefined, + apiVersionMode, }; // Route to the appropriate serve method diff --git a/src/metrics/runtime/analytics-engine-sink.ts b/src/metrics/runtime/analytics-engine-sink.ts index fd1bc57..78b7b99 100644 --- a/src/metrics/runtime/analytics-engine-sink.ts +++ b/src/metrics/runtime/analytics-engine-sink.ts @@ -5,6 +5,7 @@ import { type MetricName, } from "./metric-types"; import type { + MetricAnalyticsContext, MetricObservation, MetricResourceAttributes, MetricsFlushPayload, @@ -21,7 +22,7 @@ export type AnalyticsEngineDatasetLike = { writeDataPoint(event?: AnalyticsEngineDataPointLike): void; }; -export const ANALYTICS_ENGINE_SCHEMA_VERSION = "mcp_metrics_v1"; +export const ANALYTICS_ENGINE_SCHEMA_VERSION = "mcp_metrics_v2"; export const ANALYTICS_ENGINE_INDEX_FIELDS = [ "schema_version", @@ -33,6 +34,13 @@ export const ANALYTICS_ENGINE_INDEX_FIELDS = [ export const ANALYTICS_ENGINE_LABEL_FIELDS = APPROVED_METRIC_LABEL_KEYS; +export const ANALYTICS_ENGINE_CONTEXT_FIELDS = [ + ["api_requested_version", "apiRequestedVersion"], +] as const satisfies readonly (readonly [ + string, + keyof MetricAnalyticsContext, +])[]; + export const ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS = [ ["deployment_environment", "deployment.environment"], ["service_name", "service.name"], @@ -46,6 +54,7 @@ export const ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS = [ export const ANALYTICS_ENGINE_BLOB_FIELDS = [ "metric_kind", ...ANALYTICS_ENGINE_LABEL_FIELDS, + ...ANALYTICS_ENGINE_CONTEXT_FIELDS.map(([field]) => field), ...ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS.map(([field]) => field), ] as const; @@ -87,6 +96,13 @@ function getResourceAttribute( return nullableString(resourceAttributes[key]); } +function getAnalyticsContextField( + analyticsContext: MetricAnalyticsContext, + key: keyof MetricAnalyticsContext, +): string | null { + return nullableString(analyticsContext[key]); +} + export function getAnalyticsEngineMetricFamily( name: MetricName, ): AnalyticsEngineMetricFamily { @@ -138,6 +154,7 @@ export function getAnalyticsEngineMetricFamily( export function toAnalyticsEngineDataPoint( observation: MetricObservation, resourceAttributes: MetricResourceAttributes, + analyticsContext: MetricAnalyticsContext = {}, identity: AnalyticsEngineIdentity = {}, ): AnalyticsEngineDataPointLike { return { @@ -153,6 +170,9 @@ export function toAnalyticsEngineDataPoint( ...ANALYTICS_ENGINE_LABEL_FIELDS.map((key) => getLabel(observation.labels, key), ), + ...ANALYTICS_ENGINE_CONTEXT_FIELDS.map(([, key]) => + getAnalyticsContextField(analyticsContext, key), + ), ...ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS.map(([, key]) => getResourceAttribute(resourceAttributes, key), ), @@ -171,6 +191,7 @@ export class AnalyticsEngineMetricsSink implements MetricsSink { toAnalyticsEngineDataPoint( observation, payload.resourceAttributes, + payload.analyticsContext, payload.eventIdentity, ), ); diff --git a/src/metrics/runtime/metric-types.ts b/src/metrics/runtime/metric-types.ts index 4525c0e..2fb4110 100644 --- a/src/metrics/runtime/metric-types.ts +++ b/src/metrics/runtime/metric-types.ts @@ -92,6 +92,8 @@ export const APPROVED_METRIC_LABEL_KEYS = [ "auth_mode", "api_surface", "api_version", + "api_version_mode", + "api_release_date", "outcome", "status_class", "tool_name", @@ -161,6 +163,14 @@ export type RouteGroup = export type Transport = "mcp" | "sse" | "http" | "unknown"; export type AuthMode = "oauth" | "bearer" | "token" | "none" | "unknown"; export type ApiSurface = "mcp" | "oauth" | "static" | "unknown"; +export type ApiVersionMode = + | "implicit_legacy" + | "implicit_latest" + | "explicit_legacy" + | "explicit_latest" + | "pinned" + | "beta" + | "unknown"; export type StatusClass = "1xx" | "2xx" | "3xx" | "4xx" | "5xx" | "unknown"; function warnOnInvalidMetricLabel(key: string, reason: string) { diff --git a/src/metrics/runtime/metrics-recorder.ts b/src/metrics/runtime/metrics-recorder.ts index ff4421e..7cb5398 100644 --- a/src/metrics/runtime/metrics-recorder.ts +++ b/src/metrics/runtime/metrics-recorder.ts @@ -6,6 +6,8 @@ import { normalizeMetricLabels, } from "./metric-types"; import type { + MetricAnalyticsContext, + MetricEventIdentity, MetricObservation, MetricResourceAttributes, MetricsSink, @@ -21,6 +23,8 @@ export interface MetricsRecorder { count(name: MetricName, value?: number, labels?: MetricLabelInput): void; histogram(name: MetricName, value: number, labels?: MetricLabelInput): void; gauge(name: MetricName, value: number, labels?: MetricLabelInput): void; + setAnalyticsContext(context?: MetricAnalyticsContext): void; + setEventIdentity(identity?: MetricEventIdentity): void; flush(): Promise; snapshot(): readonly MetricObservation[]; } @@ -32,6 +36,8 @@ export const NOOP_METRICS_RECORDER: MetricsRecorder = { count(_name, _value, _labels): void {}, histogram(_name, _value, _labels): void {}, gauge(_name, _value, _labels): void {}, + setAnalyticsContext(_context): void {}, + setEventIdentity(_identity): void {}, flush(): Promise { return NOOP_FLUSH_PROMISE; }, @@ -42,6 +48,8 @@ export const NOOP_METRICS_RECORDER: MetricsRecorder = { export class RequestMetricsRecorder implements MetricsRecorder { private readonly observations: MetricObservation[] = []; + private analyticsContext?: MetricAnalyticsContext; + private eventIdentity?: MetricEventIdentity; private flushPromise?: Promise; private flushed = false; @@ -59,6 +67,37 @@ export class RequestMetricsRecorder implements MetricsRecorder { this.record("gauge", name, value, labels); } + setAnalyticsContext(context?: MetricAnalyticsContext): void { + if (!context) { + return; + } + + const nextContext: MetricAnalyticsContext = { + ...this.analyticsContext, + }; + if (context.apiRequestedVersion) { + nextContext.apiRequestedVersion = context.apiRequestedVersion; + } + this.analyticsContext = nextContext; + } + + setEventIdentity(identity?: MetricEventIdentity): void { + if (!identity) { + return; + } + + const nextIdentity: MetricEventIdentity = { + ...this.eventIdentity, + }; + if (identity.tenantId) { + nextIdentity.tenantId = identity.tenantId; + } + if (identity.userId) { + nextIdentity.userId = identity.userId; + } + this.eventIdentity = nextIdentity; + } + snapshot(): readonly MetricObservation[] { return [...this.observations]; } @@ -115,9 +154,39 @@ export class RequestMetricsRecorder implements MetricsRecorder { await this.options.sink.flush({ observations, resourceAttributes: { ...this.options.resourceAttributes }, + analyticsContext: this.analyticsContext + ? { ...this.analyticsContext } + : undefined, + eventIdentity: this.eventIdentity + ? { ...this.eventIdentity } + : undefined, }); } catch (error) { console.error("[metrics] Flush failed", error); } } } + +export type MetricsWaitUntil = (promise: Promise) => void; + +export function scheduleMetricsFlush( + recorder: MetricsRecorder, + waitUntil?: MetricsWaitUntil, +): void { + const flushPromise = recorder.flush().catch((error) => { + console.error("[metrics] Failed to execute metrics flush", error); + }); + + if (!waitUntil) { + // Non-Worker runtimes and tests may not provide waitUntil. Start the guarded + // flush anyway and intentionally detach from it. + void flushPromise; + return; + } + + try { + waitUntil(flushPromise); + } catch (error) { + console.error("[metrics] Failed to schedule metrics flush", error); + } +} diff --git a/src/metrics/runtime/metrics-sink.ts b/src/metrics/runtime/metrics-sink.ts index 33dacdc..4d5fe40 100644 --- a/src/metrics/runtime/metrics-sink.ts +++ b/src/metrics/runtime/metrics-sink.ts @@ -25,10 +25,20 @@ export type MetricEventIdentity = { userId?: string; }; +// Sink-specific context for dimensions that should not become generic metric labels. +// `api_version`, `api_version_mode`, and `api_release_date` stay in `MetricObservation.labels` +// because both Grafana and Analytics Engine should receive them as low-cardinality dimensions. +// `apiRequestedVersion` lives here instead because we only want it in Analytics Engine as +// request-debug context, not as an extra Grafana label. +export type MetricAnalyticsContext = { + apiRequestedVersion?: string; +}; + export type MetricsFlushPayload = { observations: readonly MetricObservation[]; resourceAttributes: MetricResourceAttributes; eventIdentity?: MetricEventIdentity; + analyticsContext?: MetricAnalyticsContext; }; export interface MetricsSink { diff --git a/src/metrics/runtime/request-metrics.ts b/src/metrics/runtime/request-metrics.ts index 5fbda7c..b50e254 100644 --- a/src/metrics/runtime/request-metrics.ts +++ b/src/metrics/runtime/request-metrics.ts @@ -1,8 +1,12 @@ -import { resolveApiVersion } from "../../servers/version-registry"; +import { + YYYY_MM_DD_DATE_REGEX, + resolveApiVersionMetrics, +} from "../../servers/version-registry"; import { createAnalyticsEngineMetricsSink } from "./analytics-engine-sink"; import { createGrafanaOtlpMetricsSink } from "./grafana-otlp-sink"; import { getStatusClass, resolveRequestMetricContext } from "./metric-context"; import { + type ApiVersionMode, METRIC_NAMES, type MetricLabelInput, type MetricName, @@ -12,6 +16,7 @@ import { type MetricsRecorder, NOOP_METRICS_RECORDER, RequestMetricsRecorder, + scheduleMetricsFlush, } from "./metrics-recorder"; import type { MetricsSink } from "./metrics-sink"; import { @@ -43,6 +48,8 @@ type MetricsExecutionContext = ExecutionContext & { [METRICS_RECORDER_SYMBOL]?: MetricsRecorder; props?: { apiVersion?: unknown; + apiRequestedVersion?: unknown; + apiVersionMode?: unknown; }; }; @@ -88,28 +95,77 @@ export function getMetricOutcomeForStatus(status: number): MetricOutcome { return "success"; } -function getCanonicalResolvedApiVersion(apiVersion: string): string { - const versionConfig = resolveApiVersion(apiVersion); - if (versionConfig.version.includes("beta")) { +type ApiVersionLabels = { + apiRequestedVersion?: string; + apiReleaseDate?: string; + apiVersion?: string; + apiVersionMode?: ApiVersionMode; +}; + +export function normalizeRequestedApiVersionForAnalytics( + requestedApiVersion: string, +): string { + if ( + requestedApiVersion === "beta" || + requestedApiVersion === "backwards-compatibility-default" || + requestedApiVersion === "latest" || + YYYY_MM_DD_DATE_REGEX.test(requestedApiVersion) + ) { + return requestedApiVersion; + } + + return "invalid"; +} + +export function resolveRequestedApiVersionMode( + requestedApiVersion: string, +): ApiVersionMode { + if (requestedApiVersion === "beta") { return "beta"; } - if (versionConfig.version.includes("backwards-compatibility-default")) { - return "default"; + if (requestedApiVersion === "backwards-compatibility-default") { + return "explicit_legacy"; } - if (versionConfig.version.includes("latest")) { - return "latest"; + if (requestedApiVersion === "latest") { + return "explicit_latest"; } + if (YYYY_MM_DD_DATE_REGEX.test(requestedApiVersion)) { + return "pinned"; + } + return "unknown"; +} +function getImplicitApiVersionMode(apiVersion?: string): ApiVersionMode { + if (apiVersion === "backwards-compatibility-default") { + return "implicit_legacy"; + } + if (apiVersion === "latest") { + return "implicit_latest"; + } + if (apiVersion === "beta") { + return "beta"; + } return "unknown"; } -export function resolveCanonicalApiVersionLabel( +/** + * We intentionally split version labeling into two dimensions: + * - `api_version`: the effective served surface (`backwards-compatibility-default` + * vs `latest`), which answers "which tenants are still on the legacy/v1 surface?" + * - `api_version_mode`: how the caller selected that surface, which answers + * "which tenants are pinned vs implicitly following a route default?" + * - `api_requested_version`: Analytics Engine-only context with the normalized selector + * the caller actually sent (`latest`, `beta`, a date, or `invalid`) + * - `api_release_date`: the resolved dated release behind that surface, which answers + * "what exact release date would be affected by a deprecation?" + */ +export function resolveApiVersionLabels( request: Request, ctx: ExecutionContext, -): string | undefined { +): ApiVersionLabels { const requestContext = resolveRequestMetricContext(request); if (!VERSIONED_REQUEST_ROUTE_GROUPS.has(requestContext.routeGroup)) { - return undefined; + return {}; } const requestedApiVersion = new URL(request.url).searchParams.get( @@ -117,26 +173,89 @@ export function resolveCanonicalApiVersionLabel( ); const effectiveApiVersion = (ctx as MetricsExecutionContext).props ?.apiVersion; + const effectiveRequestedApiVersion = (ctx as MetricsExecutionContext).props + ?.apiRequestedVersion; + const effectiveApiVersionMode = (ctx as MetricsExecutionContext).props + ?.apiVersionMode; + if (requestedApiVersion) { + const apiVersionSource = + typeof effectiveApiVersion === "string" && effectiveApiVersion.length > 0 + ? effectiveApiVersion + : requestedApiVersion; + try { + return { + apiRequestedVersion: + typeof effectiveRequestedApiVersion === "string" && + effectiveRequestedApiVersion.length > 0 + ? effectiveRequestedApiVersion + : normalizeRequestedApiVersionForAnalytics(requestedApiVersion), + ...resolveApiVersionMetrics(apiVersionSource), + apiVersionMode: + typeof effectiveApiVersionMode === "string" && + effectiveApiVersionMode.length > 0 + ? effectiveApiVersionMode + : resolveRequestedApiVersionMode(requestedApiVersion), + }; + } catch { + return { + apiRequestedVersion: + typeof effectiveRequestedApiVersion === "string" && + effectiveRequestedApiVersion.length > 0 + ? effectiveRequestedApiVersion + : normalizeRequestedApiVersionForAnalytics(requestedApiVersion), + apiReleaseDate: undefined, + apiVersion: "unknown", + apiVersionMode: "unknown", + }; + } + } + if ( typeof effectiveApiVersion === "string" && effectiveApiVersion.length > 0 ) { try { - return getCanonicalResolvedApiVersion(effectiveApiVersion); + const resolved = resolveApiVersionMetrics(effectiveApiVersion); + return { + ...resolved, + apiVersionMode: + typeof effectiveApiVersionMode === "string" && + effectiveApiVersionMode.length > 0 + ? effectiveApiVersionMode + : getImplicitApiVersionMode(resolved.apiVersion), + }; } catch { - return "unknown"; + return { + apiReleaseDate: undefined, + apiVersion: "unknown", + apiVersionMode: "unknown", + }; } } - if (!requestedApiVersion) { - return "default"; + if ( + requestContext.routeGroup === "token_mcp" || + requestContext.routeGroup === "token_sse" + ) { + const resolved = resolveApiVersionMetrics("latest"); + return { + ...resolved, + apiVersionMode: "implicit_latest", + }; } - try { - return getCanonicalResolvedApiVersion(requestedApiVersion); - } catch { - return "unknown"; - } + const resolved = resolveApiVersionMetrics("backwards-compatibility-default"); + return { + ...resolved, + apiVersionMode: "implicit_legacy", + }; +} + +export function resolveCanonicalApiVersionLabel( + request: Request, + ctx: ExecutionContext, +): string | undefined { + return resolveApiVersionLabels(request, ctx).apiVersion; } export function recordStatusMetric( @@ -166,6 +285,15 @@ export function recordBearerAuthRequestMetric( } const requestContext = resolveRequestMetricContext(request); + const requestedApiVersion = new URL(request.url).searchParams.get( + "api-version", + ); + if (requestedApiVersion) { + recorder.setAnalyticsContext({ + apiRequestedVersion: + normalizeRequestedApiVersionForAnalytics(requestedApiVersion), + }); + } recordStatusMetric(recorder, METRIC_NAMES.bearerAuthRequestsTotal, status, { route_group: routeGroupOverride ?? requestContext.routeGroup, transport: routeGroupOverride?.endsWith("_sse") @@ -185,7 +313,13 @@ export function recordHttpRequestMetrics( ): void { const requestContext = resolveRequestMetricContext(request); const outcome = getMetricOutcomeForStatus(response.status); - const apiVersion = resolveCanonicalApiVersionLabel(request, ctx); + const { apiRequestedVersion, apiReleaseDate, apiVersion, apiVersionMode } = + resolveApiVersionLabels(request, ctx); + if (apiRequestedVersion) { + recorder.setAnalyticsContext({ + apiRequestedVersion, + }); + } const baseLabels: MetricLabelInput = { route_group: requestContext.routeGroup, transport: requestContext.transport, @@ -197,6 +331,12 @@ export function recordHttpRequestMetrics( if (apiVersion) { baseLabels.api_version = apiVersion; } + if (apiVersionMode) { + baseLabels.api_version_mode = apiVersionMode; + } + if (apiReleaseDate) { + baseLabels.api_release_date = apiReleaseDate; + } recorder.count(METRIC_NAMES.httpRequestsTotal, 1, { ...baseLabels, @@ -258,21 +398,7 @@ function scheduleRequestMetricsFlush( recorder: MetricsRecorder, ctx: ExecutionContext, ): void { - let flushPromise: Promise; - try { - flushPromise = recorder.flush().catch((error) => { - console.error("[metrics] Failed to execute request metrics flush", error); - }); - } catch (error) { - console.error("[metrics] Failed to execute request metrics flush", error); - return; - } - - try { - ctx.waitUntil(flushPromise); - } catch (error) { - console.error("[metrics] Failed to schedule request metrics flush", error); - } + scheduleMetricsFlush(recorder, ctx.waitUntil.bind(ctx)); } export async function withRequestMetrics( diff --git a/src/metrics/runtime/tool-metrics.ts b/src/metrics/runtime/tool-metrics.ts new file mode 100644 index 0000000..c86b7ee --- /dev/null +++ b/src/metrics/runtime/tool-metrics.ts @@ -0,0 +1,163 @@ +import { ZodError } from "zod"; +import { McpServerError } from "../../utils"; +import { + type ApiVersionMode, + METRIC_NAMES, + type MetricLabelInput, + type MetricOutcome, +} from "./metric-types"; +import type { MetricsRecorder } from "./metrics-recorder"; + +export const UPSTREAM_OPERATION_NAMES = { + getSessionInfo: "get_session_info", + getDataSourceSuggestions: "get_data_source_suggestions", + queryGetDecomposedQuery: "query_get_decomposed_query", + singleAnswer: "single_answer", + exportAnswerReport: "export_answer_report", + getAnswerSession: "get_answer_session", + exportUnsavedAnswerTml: "export_unsaved_answer_tml", + createAgentConversation: "create_agent_conversation", + sendAgentConversationMessageStreaming: + "send_agent_conversation_message_streaming", + importMetadataTml: "import_metadata_tml", + searchMetadata: "search_metadata", +} as const; + +export type UpstreamOperation = + (typeof UPSTREAM_OPERATION_NAMES)[keyof typeof UPSTREAM_OPERATION_NAMES]; + +export type ToolMetricApiSurface = "mcp"; +export type UpstreamStreamMessageType = + | "text" + | "text_chunk" + | "answer" + | "error"; + +function buildToolMetricLabels( + toolName: string, + apiSurface: ToolMetricApiSurface, + outcome: MetricOutcome, + apiVersion?: string, + apiVersionMode?: ApiVersionMode, + apiReleaseDate?: string, +): MetricLabelInput { + const labels: MetricLabelInput = { + tool_name: toolName, + api_surface: apiSurface, + outcome, + }; + + if (apiVersion) { + labels.api_version = apiVersion; + } + if (apiVersionMode) { + labels.api_version_mode = apiVersionMode; + } + if (apiReleaseDate) { + labels.api_release_date = apiReleaseDate; + } + + return labels; +} + +export function getToolMetricOutcomeFromResult(result: unknown): MetricOutcome { + if ( + typeof result === "object" && + result !== null && + "isError" in result && + result.isError === true + ) { + return "error"; + } + + return "success"; +} + +export function getToolMetricOutcomeFromError(error: unknown): MetricOutcome { + if (error instanceof ZodError) { + return "validation_error"; + } + + if (error instanceof McpServerError) { + if (error.statusCode >= 400 && error.statusCode < 500) { + return "client_error"; + } + return "error"; + } + + return "error"; +} + +export function recordToolInvocationMetrics( + recorder: MetricsRecorder, + toolName: string, + apiSurface: ToolMetricApiSurface, + outcome: MetricOutcome, + durationMs: number, + apiVersion?: string, + apiVersionMode?: ApiVersionMode, + apiReleaseDate?: string, +): void { + const labels = buildToolMetricLabels( + toolName, + apiSurface, + outcome, + apiVersion, + apiVersionMode, + apiReleaseDate, + ); + + recorder.count(METRIC_NAMES.toolCallsTotal, 1, labels); + recorder.histogram(METRIC_NAMES.toolDurationMs, durationMs, labels); +} + +export function recordUpstreamCallMetrics( + recorder: MetricsRecorder | undefined, + operation: UpstreamOperation, + outcome: MetricOutcome, + durationMs: number, +): void { + if (!recorder) { + return; + } + + const labels: MetricLabelInput = { + upstream_operation: operation, + outcome, + }; + + recorder.count(METRIC_NAMES.upstreamCallsTotal, 1, labels); + recorder.histogram(METRIC_NAMES.upstreamDurationMs, durationMs, labels); +} + +export function recordUpstreamStreamStartedMetric( + recorder: MetricsRecorder | undefined, + operation: UpstreamOperation, + outcome: MetricOutcome, +): void { + if (!recorder) { + return; + } + + recorder.count(METRIC_NAMES.upstreamStreamsStartedTotal, 1, { + upstream_operation: operation, + outcome, + }); +} + +export function recordUpstreamStreamMessageMetric( + recorder: MetricsRecorder | undefined, + operation: UpstreamOperation, + messageType: UpstreamStreamMessageType, + isThinking: boolean, +): void { + if (!recorder) { + return; + } + + recorder.count(METRIC_NAMES.upstreamStreamMessagesTotal, 1, { + upstream_operation: operation, + message_type: messageType, + is_thinking: isThinking, + }); +} diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index bd6b8af..cae4af2 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -1,21 +1,37 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import { CallToolRequestSchema, - ListToolsRequestSchema, - ToolSchema, ListResourcesRequestSchema, - ReadResourceRequestSchema, + ListToolsRequestSchema, type ListToolsResult, + ReadResourceRequestSchema, + ToolSchema, } from "@modelcontextprotocol/sdk/types.js"; -import type { z } from "zod"; import { type Span, SpanStatusCode } from "@opentelemetry/api"; -import { getActiveSpan, withSpan } from "../metrics/tracing/tracing-utils"; -import { Trackers, type Tracker, TrackEvent } from "../metrics"; -import type { Props } from "../utils"; +import type { z } from "zod"; +import { TrackEvent, type Tracker, Trackers } from "../metrics"; import { MixpanelTracker } from "../metrics/mixpanel/mixpanel"; +import type { ApiVersionMode } from "../metrics/runtime/metric-types"; +import { + type MetricsRecorder, + scheduleMetricsFlush, +} from "../metrics/runtime/metrics-recorder"; +import type { + MetricAnalyticsContext, + MetricEventIdentity, +} from "../metrics/runtime/metrics-sink"; +import { createRequestMetricsRecorder } from "../metrics/runtime/request-metrics"; +import { + type ToolMetricApiSurface, + getToolMetricOutcomeFromError, + getToolMetricOutcomeFromResult, + recordToolInvocationMetrics, +} from "../metrics/runtime/tool-metrics"; +import { getActiveSpan, withSpan } from "../metrics/tracing/tracing-utils"; +import { StorageServiceClient } from "../storage-service/storage-service"; import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client"; import { ThoughtSpotService } from "../thoughtspot/thoughtspot-service"; -import { StorageServiceClient } from "../storage-service/storage-service"; +import type { Props } from "../utils"; const ToolInputSchema = ToolSchema.shape.inputSchema; type ToolInput = z.infer; @@ -41,6 +57,7 @@ export type ToolResponse = SuccessResponse | ErrorResponse; export interface Context { props: Props; env: Env; + ctx?: DurableObjectState; } export abstract class BaseMCPServer extends Server { @@ -181,15 +198,137 @@ export abstract class BaseMCPServer extends Server { ); } - protected getThoughtSpotService() { + protected getThoughtSpotService(recorder?: MetricsRecorder) { return new ThoughtSpotService( getThoughtSpotClient( this.ctx.props.instanceUrl, this.ctx.props.accessToken, ), + { + recorder, + metricsEnv: this.ctx.env as unknown as Record, + waitUntil: this.getMetricsWaitUntil(), + analyticsContext: this.getMetricAnalyticsContext(), + eventIdentity: this.getMetricEventIdentity(), + }, ); } + protected abstract getToolMetricApiSurface(): ToolMetricApiSurface; + + protected getToolMetricApiVersionLabel(): string | undefined { + return undefined; + } + + protected getToolMetricApiVersionModeLabel(): ApiVersionMode | undefined { + return undefined; + } + + protected getToolMetricApiReleaseDateLabel(): string | undefined { + return undefined; + } + + protected getMetricAnalyticsContext(): MetricAnalyticsContext | undefined { + const apiRequestedVersion = this.ctx.props.apiRequestedVersion; + if ( + typeof apiRequestedVersion !== "string" || + apiRequestedVersion.length === 0 + ) { + return undefined; + } + + return { + apiRequestedVersion, + }; + } + + protected getMetricEventIdentity(): MetricEventIdentity | undefined { + if (!this.sessionInfo) { + return undefined; + } + + const tenantId = this.sessionInfo.currentOrgId + ? String(this.sessionInfo.currentOrgId) + : undefined; + const userId = this.sessionInfo.userGUID + ? String(this.sessionInfo.userGUID) + : undefined; + if (!tenantId && !userId) { + return undefined; + } + + return { + tenantId, + userId, + }; + } + + private getMetricsWaitUntil() { + return this.ctx.ctx?.waitUntil?.bind(this.ctx.ctx); + } + + private createToolMetricsRecorder(): MetricsRecorder { + const recorder = createRequestMetricsRecorder( + this.ctx.env as unknown as Record, + ); + recorder.setAnalyticsContext(this.getMetricAnalyticsContext()); + recorder.setEventIdentity(this.getMetricEventIdentity()); + return recorder; + } + + private recordToolMetricsSafe( + recorder: MetricsRecorder, + toolName: string, + outcome: ReturnType, + durationMs: number, + ): void { + try { + recordToolInvocationMetrics( + recorder, + toolName, + this.getToolMetricApiSurface(), + outcome, + durationMs, + this.getToolMetricApiVersionLabel(), + this.getToolMetricApiVersionModeLabel(), + this.getToolMetricApiReleaseDateLabel(), + ); + } catch (error) { + console.error( + `[metrics] Failed to record tool metrics for ${toolName}`, + error, + ); + } + } + + private async withToolMetrics( + request: z.infer, + handler: (recorder: MetricsRecorder) => Promise, + ): Promise { + const recorder = this.createToolMetricsRecorder(); + const startedAt = Date.now(); + let outcome: ReturnType | undefined; + + try { + const result = await handler(recorder); + outcome = getToolMetricOutcomeFromResult(result); + return result; + } catch (error) { + outcome = getToolMetricOutcomeFromError(error); + throw error; + } finally { + if (outcome) { + this.recordToolMetricsSafe( + recorder, + request.params.name, + outcome, + Date.now() - startedAt, + ); + } + scheduleMetricsFlush(recorder, this.getMetricsWaitUntil()); + } + } + protected async initializeService(): Promise { try { this.sessionInfo = await this.getThoughtSpotService().getSessionInfo(); @@ -225,6 +364,7 @@ export abstract class BaseMCPServer extends Server { */ protected abstract callTool( request: z.infer, + recorder: MetricsRecorder, ): Promise; async init() { @@ -265,7 +405,9 @@ export abstract class BaseMCPServer extends Server { async (request: z.infer) => { return withSpan("call-tool", async () => { this.initSpanWithCommonAttributes(); - return this.callTool(request); + return this.withToolMetrics(request, (recorder) => + this.callTool(request, recorder), + ); }); }, ); diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 9684f8b..48dbe8d 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -2,27 +2,34 @@ import type { CallToolRequestSchema, ReadResourceRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; +import { SpanStatusCode, context, trace } from "@opentelemetry/api"; import type { z } from "zod"; -import { McpServerError } from "../utils"; -import type { DataSource } from "../thoughtspot/thoughtspot-service"; import { TrackEvent } from "../metrics"; +import type { ApiVersionMode } from "../metrics/runtime/metric-types"; +import type { MetricsRecorder } from "../metrics/runtime/metrics-recorder"; +import type { ToolMetricApiSurface } from "../metrics/runtime/tool-metrics"; import { WithSpan } from "../metrics/tracing/tracing-utils"; -import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import type { StreamingMessagesStorageWithTtl } from "../streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; +import type { DataSource } from "../thoughtspot/thoughtspot-service"; +import type { Answer, StreamingMessagesState } from "../thoughtspot/types"; +import { McpServerError } from "../utils"; import { BaseMCPServer, type Context } from "./mcp-server-base"; -import { resolveApiVersion, type VersionConfig } from "./version-registry"; import { - GetRelevantQuestionsSchema, - GetAnswerSchema, + CreateAnalysisSessionInputSchema, + CreateDashboardInputSchema, CreateLiveboardSchema, + GetAnswerSchema, GetDataSourceSuggestionsSchema, - ToolName, - CreateAnalysisSessionInputSchema, - SendSessionMessageInputSchema, + GetRelevantQuestionsSchema, GetSessionUpdatesInputSchema, - CreateDashboardInputSchema, + SendSessionMessageInputSchema, + ToolName, } from "./tool-definitions"; -import type { StreamingMessagesStorageWithTtl } from "../streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; -import type { Answer, StreamingMessagesState } from "../thoughtspot/types"; +import { + type VersionConfig, + resolveApiVersion, + resolveApiVersionMetrics, +} from "./version-registry"; export class MCPServer extends BaseMCPServer { constructor( @@ -32,6 +39,64 @@ export class MCPServer extends BaseMCPServer { super(ctx, "ThoughtSpot", "2.0.0"); } + protected getToolMetricApiSurface(): ToolMetricApiSurface { + return "mcp"; + } + + protected getToolMetricApiVersionLabel(): string | undefined { + const apiVersion = this.ctx.props.apiVersion; + if (typeof apiVersion !== "string" || apiVersion.length === 0) { + return "backwards-compatibility-default"; + } + + try { + return resolveApiVersionMetrics(apiVersion).apiVersion; + } catch { + return "unknown"; + } + } + + protected getToolMetricApiVersionModeLabel(): ApiVersionMode | undefined { + const apiVersionMode = this.ctx.props.apiVersionMode; + if (typeof apiVersionMode === "string" && apiVersionMode.length > 0) { + return apiVersionMode; + } + + const apiVersion = this.ctx.props.apiVersion; + if (typeof apiVersion === "string" && apiVersion.length > 0) { + try { + const resolved = resolveApiVersionMetrics(apiVersion); + if (resolved.apiVersion === "backwards-compatibility-default") { + return "implicit_legacy"; + } + if (resolved.apiVersion === "latest") { + return "implicit_latest"; + } + if (resolved.apiVersion === "beta") { + return "beta"; + } + } catch { + return "unknown"; + } + } + + return "implicit_legacy"; + } + + protected getToolMetricApiReleaseDateLabel(): string | undefined { + const apiVersion = this.ctx.props.apiVersion; + if (typeof apiVersion !== "string" || apiVersion.length === 0) { + return resolveApiVersionMetrics("backwards-compatibility-default") + .apiReleaseDate; + } + + try { + return resolveApiVersionMetrics(apiVersion).apiReleaseDate; + } catch { + return undefined; + } + } + @WithSpan("call-list-tools") protected async listTools() { const span = this.initSpanWithCommonAttributes(); @@ -45,7 +110,10 @@ export class MCPServer extends BaseMCPServer { try { versionConfig = resolveApiVersion(this.ctx.props.apiVersion); } catch (error) { - console.error("Error resolving API version, using default:", error); + console.error( + "Error resolving API version, using latest fallback:", + error, + ); span?.recordException(error as Error); versionConfig = resolveApiVersion(); } @@ -114,7 +182,10 @@ export class MCPServer extends BaseMCPServer { }; } - protected async callTool(request: z.infer) { + protected async callTool( + request: z.infer, + recorder: MetricsRecorder, + ) { const { name } = request.params; this.trackers.track(TrackEvent.CallTool, { toolName: name }); @@ -127,19 +198,19 @@ export class MCPServer extends BaseMCPServer { } case ToolName.GetRelevantQuestions: { - return this.callGetRelevantQuestions(request); + return this.callGetRelevantQuestions(request, recorder); } case ToolName.GetAnswer: { - return this.callGetAnswer(request); + return this.callGetAnswer(request, recorder); } case ToolName.CreateLiveboard: { - return this.callCreateLiveboard(request); + return this.callCreateLiveboard(request, recorder); } case ToolName.GetDataSourceSuggestions: { - return this.callGetDataSourceSuggestions(request); + return this.callGetDataSourceSuggestions(request, recorder); } case ToolName.CheckConnectivity: { @@ -156,19 +227,19 @@ export class MCPServer extends BaseMCPServer { } case ToolName.CreateAnalysisSession: { - return this.callCreateAnalysisSession(request); + return this.callCreateAnalysisSession(request, recorder); } case ToolName.SendSessionMessage: { - return this.callSendSessionMessage(request); + return this.callSendSessionMessage(request, recorder); } case ToolName.GetSessionUpdates: { - return this.callGetSessionUpdates(request); + return this.callGetSessionUpdates(request, recorder); } case ToolName.CreateDashboard: { - return this.callCreateDashboard(request); + return this.callCreateDashboard(request, recorder); } default: @@ -179,6 +250,7 @@ export class MCPServer extends BaseMCPServer { @WithSpan("call-get-relevant-questions") async callGetRelevantQuestions( request: z.infer, + recorder: MetricsRecorder, ) { const { query, @@ -190,12 +262,9 @@ export class MCPServer extends BaseMCPServer { sourceIds, ); - const relevantQuestions = - await this.getThoughtSpotService().getRelevantQuestions( - query, - sourceIds!, - additionalContext ?? "", - ); + const relevantQuestions = await this.getThoughtSpotService( + recorder, + ).getRelevantQuestions(query, sourceIds!, additionalContext ?? ""); if (relevantQuestions.error) { console.error( @@ -235,16 +304,17 @@ export class MCPServer extends BaseMCPServer { } @WithSpan("call-get-answer") - async callGetAnswer(request: z.infer) { + async callGetAnswer( + request: z.infer, + recorder: MetricsRecorder, + ) { const { question, datasourceId: sourceId } = GetAnswerSchema.parse( request.params.arguments, ); - const answer = await this.getThoughtSpotService().getAnswerForQuestion( - question, - sourceId, - false, - ); + const answer = await this.getThoughtSpotService( + recorder, + ).getAnswerForQuestion(question, sourceId, false); if (answer.error) { return this.createErrorResponse( @@ -268,7 +338,10 @@ export class MCPServer extends BaseMCPServer { } @WithSpan("call-create-liveboard") - async callCreateLiveboard(request: z.infer) { + async callCreateLiveboard( + request: z.infer, + recorder: MetricsRecorder, + ) { const { name, answers, noteTile } = CreateLiveboardSchema.parse( request.params.arguments, ); @@ -277,12 +350,9 @@ export class MCPServer extends BaseMCPServer { session_identifier: answer.session_identifier, generation_number: answer.generation_number, })); - const liveboard = - await this.getThoughtSpotService().fetchTMLAndCreateLiveboard( - name, - transformedAnswers, - noteTile, - ); + const liveboard = await this.getThoughtSpotService( + recorder, + ).fetchTMLAndCreateLiveboard(name, transformedAnswers, noteTile); if (liveboard.error) { return this.createErrorResponse( @@ -304,6 +374,7 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; @WithSpan("call-create-analysis-session") async callCreateAnalysisSession( request: z.infer, + recorder: MetricsRecorder, ) { const span = trace.getSpan(context.active()); const { data_source_id } = CreateAnalysisSessionInputSchema.parse( @@ -312,7 +383,7 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; span?.setAttribute("data_source_id", data_source_id ?? "(none)"); const response = - await this.getThoughtSpotService().createAgentConversation( + await this.getThoughtSpotService(recorder).createAgentConversation( data_source_id, ); span?.setAttribute("analytical_session_id", response.conversation_id); @@ -327,7 +398,10 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; } @WithSpan("call-send-session-message") - async callSendSessionMessage(request: z.infer) { + async callSendSessionMessage( + request: z.infer, + recorder: MetricsRecorder, + ) { const span = trace.getSpan(context.active()); const { analytical_session_id, message, additional_context } = SendSessionMessageInputSchema.parse(request.params.arguments); @@ -350,7 +424,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; ); } - await this.getThoughtSpotService().sendAgentConversationMessageStreaming( + await this.getThoughtSpotService( + recorder, + ).sendAgentConversationMessageStreaming( analytical_session_id, message, storageService.appendMessages.bind(storageService), @@ -364,7 +440,10 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; } @WithSpan("call-get-session-updates") - async callGetSessionUpdates(request: z.infer) { + async callGetSessionUpdates( + request: z.infer, + _recorder: MetricsRecorder, + ) { const span = trace.getSpan(context.active()); const { analytical_session_id } = GetSessionUpdatesInputSchema.parse( request.params.arguments, @@ -421,7 +500,10 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; } @WithSpan("call-create-dashboard") - async callCreateDashboard(request: z.infer) { + async callCreateDashboard( + request: z.infer, + recorder: MetricsRecorder, + ) { const span = trace.getSpan(context.active()); const { title, answers, note_tile } = CreateDashboardInputSchema.parse( request.params.arguments, @@ -448,12 +530,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; ); } - const liveboard = - await this.getThoughtSpotService().fetchTMLAndCreateLiveboard( - title, - transformedAnswers, - note_tile, - ); + const liveboard = await this.getThoughtSpotService( + recorder, + ).fetchTMLAndCreateLiveboard(title, transformedAnswers, note_tile); if (liveboard.error) { return this.createErrorResponse( @@ -473,12 +552,15 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; @WithSpan("call-get-data-source-suggestions") async callGetDataSourceSuggestions( request: z.infer, + recorder: MetricsRecorder, ) { const { query } = GetDataSourceSuggestionsSchema.parse( request.params.arguments, ); const dataSources = - await this.getThoughtSpotService().getDataSourceSuggestions(query); + await this.getThoughtSpotService(recorder).getDataSourceSuggestions( + query, + ); if (!dataSources || dataSources.length === 0) { return this.createErrorResponse( @@ -506,12 +588,12 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; } | null = null; @WithSpan("get-datasources") - async getDatasources() { + async getDatasources(recorder?: MetricsRecorder) { if (this._sources) { return this._sources; } - const sources = await this.getThoughtSpotService().getDataSources(); + const sources = await this.getThoughtSpotService(recorder).getDataSources(); this._sources = { list: sources, map: new Map(sources.map((s) => [s.id, s])), diff --git a/src/servers/version-registry.ts b/src/servers/version-registry.ts index 5374f98..1ed1c04 100644 --- a/src/servers/version-registry.ts +++ b/src/servers/version-registry.ts @@ -1,6 +1,17 @@ import { toolDefinitionsV1, toolDefinitionsV2 } from "./tool-definitions"; -const YYYY_MM_DD_DATE_REGEX = /^\d{4}-\d{2}-\d{2}$/; +export const YYYY_MM_DD_DATE_REGEX = /^\d{4}-\d{2}-\d{2}$/; + +export type CanonicalApiVersionLabel = + | "backwards-compatibility-default" + | "latest" + | "beta" + | "unknown"; + +export type ResolvedApiVersionMetrics = { + apiReleaseDate?: string; + apiVersion: CanonicalApiVersionLabel; +}; /** * Version configuration interface @@ -27,6 +38,15 @@ function getReleaseDateFromVersion(versions: string[]): Date | null { return new Date(possibleDateVersion); } +export function getReleaseDateIdentifier( + versions: string[], +): string | undefined { + const possibleDateVersion = versions[versions.length - 1]; + return YYYY_MM_DD_DATE_REGEX.test(possibleDateVersion) + ? possibleDateVersion + : undefined; +} + /** * Version registry, with respective tools to expose by API version. The "version" field can * contain multiple identifiers for the same version. Important ordering rules: @@ -96,3 +116,32 @@ export function resolveApiVersion(apiVersion = "latest"): VersionConfig { ); return VERSION_REGISTRY[VERSION_REGISTRY.length - 1]; } + +export function resolveApiVersionMetrics( + apiVersion: string, +): ResolvedApiVersionMetrics { + const versionConfig = resolveApiVersion(apiVersion); + if (versionConfig.version.includes("beta")) { + return { + apiVersion: "beta", + apiReleaseDate: getReleaseDateIdentifier(versionConfig.version), + }; + } + if (versionConfig.version.includes("backwards-compatibility-default")) { + return { + apiVersion: "backwards-compatibility-default", + apiReleaseDate: getReleaseDateIdentifier(versionConfig.version), + }; + } + if (versionConfig.version.includes("latest")) { + return { + apiVersion: "latest", + apiReleaseDate: getReleaseDateIdentifier(versionConfig.version), + }; + } + + return { + apiVersion: "unknown", + apiReleaseDate: getReleaseDateIdentifier(versionConfig.version), + }; +} diff --git a/src/streaming-utils.ts b/src/streaming-utils.ts index ed3eee9..d51fc6b 100644 --- a/src/streaming-utils.ts +++ b/src/streaming-utils.ts @@ -1,6 +1,11 @@ -import type { Message } from "./thoughtspot/types"; -import { withSpan } from "./metrics/tracing/tracing-utils"; import { type Span, SpanStatusCode } from "@opentelemetry/api"; +import type { MetricsRecorder } from "./metrics/runtime/metrics-recorder"; +import { + UPSTREAM_OPERATION_NAMES, + recordUpstreamStreamMessageMetric, +} from "./metrics/runtime/tool-metrics"; +import { withSpan } from "./metrics/tracing/tracing-utils"; +import type { Message } from "./thoughtspot/types"; /* * Handles processing the event stream from a send agent conversation message response. Reads from @@ -16,10 +21,15 @@ export const processSendAgentConversationMessageStreamingResponse = async ( isDone?: boolean, ) => Promise, instanceUrl: string, + recorder?: MetricsRecorder, ) => { return await withSpan( "process-send-agent-conversation-message-streaming-response", async (span: Span) => { + // Include the operation on every stream-message metric so future streamed + // upstream calls do not collapse into the same series. + const upstreamOperation = + UPSTREAM_OPERATION_NAMES.sendAgentConversationMessageStreaming; span.setAttribute("conversation_id", conversationId); let nTextMessagesParsed = 0; let nAnswerMessagesParsed = 0; @@ -71,6 +81,12 @@ export const processSendAgentConversationMessageStreamingResponse = async ( for (const item of data) { if (item.type === "text") { nTextMessagesParsed++; + recordUpstreamStreamMessageMetric( + recorder, + upstreamOperation, + "text", + item.metadata?.type === "thinking", + ); newMessages.push({ is_thinking: item.metadata?.type === "thinking", type: "text", @@ -78,6 +94,12 @@ export const processSendAgentConversationMessageStreamingResponse = async ( }); } else if (item.type === "text-chunk") { nTextMessagesParsed++; + recordUpstreamStreamMessageMetric( + recorder, + upstreamOperation, + "text_chunk", + item.metadata?.type === "thinking", + ); newMessages.push({ is_thinking: item.metadata?.type === "thinking", type: "text_chunk", @@ -85,6 +107,12 @@ export const processSendAgentConversationMessageStreamingResponse = async ( }); } else if (item.type === "answer") { nAnswerMessagesParsed++; + recordUpstreamStreamMessageMetric( + recorder, + upstreamOperation, + "answer", + item.metadata?.type === "thinking", + ); const iframeUrl = `${instanceUrl}/?tsmcp=true#/embed/conv-assist-answer?sessionId=${item.metadata?.session_id}&genNo=${item.metadata?.gen_no}&acSessionId=${item.metadata?.transaction_id}&acGenNo=${item.metadata?.generation_number}`; newMessages.push({ is_thinking: item.metadata?.type === "thinking", @@ -108,7 +136,12 @@ export const processSendAgentConversationMessageStreamingResponse = async ( nMessagesIgnored++; } else if (item.type === "error") { console.error("Error event in event stream: ", item); - nTextMessagesParsed++; + recordUpstreamStreamMessageMetric( + recorder, + upstreamOperation, + "error", + false, + ); spanHasError = true; span.setStatus({ code: SpanStatusCode.ERROR, diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 424331b..04f4622 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -1,23 +1,83 @@ +import { SpanStatusCode, context, trace } from "@opentelemetry/api"; import type { AgentConversation, ThoughtSpotRestApi, } from "@thoughtspot/rest-api-sdk"; -import { SpanStatusCode, trace, context } from "@opentelemetry/api"; -import { getActiveSpan, WithSpan } from "../metrics/tracing/tracing-utils"; +import { + type MetricsRecorder, + scheduleMetricsFlush, +} from "../metrics/runtime/metrics-recorder"; import type { + MetricAnalyticsContext, + MetricEventIdentity, +} from "../metrics/runtime/metrics-sink"; +import { createRequestMetricsRecorder } from "../metrics/runtime/request-metrics"; +import type { MetricsEnvLike } from "../metrics/runtime/runtime-config"; +import { + UPSTREAM_OPERATION_NAMES, + type UpstreamOperation, + recordUpstreamCallMetrics, + recordUpstreamStreamStartedMetric, +} from "../metrics/runtime/tool-metrics"; +import { WithSpan, getActiveSpan } from "../metrics/tracing/tracing-utils"; +import { processSendAgentConversationMessageStreamingResponse } from "../streaming-utils"; +import type { + Answer, DataSource, - SessionInfo, DataSourceSuggestion, Message, - Answer, + SessionInfo, } from "./types"; -import { processSendAgentConversationMessageStreamingResponse } from "../streaming-utils"; + +type ThoughtSpotServiceMetricsOptions = { + recorder?: MetricsRecorder; + metricsEnv?: MetricsEnvLike; + waitUntil?: (promise: Promise) => void; + analyticsContext?: MetricAnalyticsContext; + eventIdentity?: MetricEventIdentity; +}; /** * Main ThoughtSpot service class using decorator pattern for tracing */ export class ThoughtSpotService { - constructor(private client: ThoughtSpotRestApi) {} + constructor( + private client: ThoughtSpotRestApi, + private readonly metrics: ThoughtSpotServiceMetricsOptions = {}, + ) {} + + private async observeUpstreamCall( + operation: UpstreamOperation, + call: () => Promise, + ): Promise { + const startedAt = Date.now(); + let outcome: "success" | "upstream_error" = "success"; + + try { + return await call(); + } catch (error) { + outcome = "upstream_error"; + throw error; + } finally { + recordUpstreamCallMetrics( + this.metrics.recorder, + operation, + outcome, + Date.now() - startedAt, + ); + } + } + + private createStreamMetricsRecorder(): MetricsRecorder { + const recorder = createRequestMetricsRecorder(this.metrics.metricsEnv); + recorder.setAnalyticsContext(this.metrics.analyticsContext); + recorder.setEventIdentity(this.metrics.eventIdentity); + return recorder; + } + + private scheduleBackgroundMetricsFlush(recorder: MetricsRecorder): void { + scheduleMetricsFlush(recorder, this.metrics.waitUntil); + } @WithSpan("discover-data-sources") async discoverDataSources( @@ -49,9 +109,13 @@ export class ThoughtSpotService { span?.setAttribute("query", query); span?.addEvent("query-get-data-source-suggestions"); - const response = await this.client.getDataSourceSuggestions({ - query, - }); + const response = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getDataSourceSuggestions, + () => + this.client.getDataSourceSuggestions({ + query, + }), + ); span?.setStatus({ code: SpanStatusCode.OK, @@ -112,14 +176,18 @@ export class ThoughtSpotService { ); span?.addEvent("get-decomposed-query"); - const resp = await this.client.queryGetDecomposedQuery({ - nlsRequest: { - query: query, - }, - content: [additionalContext], - worksheetIds: sourceIds, - maxDecomposedQueries: 5, - }); + const resp = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.queryGetDecomposedQuery, + () => + this.client.queryGetDecomposedQuery({ + nlsRequest: { + query: query, + }, + content: [additionalContext], + worksheetIds: sourceIds, + maxDecomposedQueries: 5, + }), + ); const questions = resp.decomposedQueryResponse?.decomposedQueries?.map((q) => ({ @@ -183,11 +251,15 @@ export class ThoughtSpotService { "instanceUrl: ", (this.client as any).instanceUrl, ); - const data = await this.client.exportAnswerReport({ - session_identifier, - generation_number, - file_format: "CSV", - }); + const data = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.exportAnswerReport, + () => + this.client.exportAnswerReport({ + session_identifier, + generation_number, + file_format: "CSV", + }), + ); let csvData = await data.text(); // get only the first 100 lines of the csv data @@ -223,10 +295,14 @@ export class ThoughtSpotService { try { span?.setAttribute("session_identifier", session_identifier); console.log("[DEBUG] Getting TML for answer: ", title); - const tml = await (this.client as any).exportUnsavedAnswerTML({ - session_identifier, - generation_number, - }); + const tml = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.exportUnsavedAnswerTml, + () => + (this.client as any).exportUnsavedAnswerTML({ + session_identifier, + generation_number, + }), + ); return tml; } catch (error) { span?.setStatus({ @@ -250,11 +326,13 @@ export class ThoughtSpotService { try { // Use auto mode by default, but support passing an explicit data source context - const response = await ( - this.client as any - ).createAgentConversationWithAutoMode({ - dataSourceId, - }); + const response = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.createAgentConversation, + () => + (this.client as any).createAgentConversationWithAutoMode({ + dataSourceId, + }), + ); span?.setStatus({ code: SpanStatusCode.OK, @@ -298,29 +376,57 @@ export class ThoughtSpotService { ? `${message}\n\nAdditional Context:\n${additionalContext}` : message; - const response = await ( - this.client as any - ).sendAgentConversationMessageStreaming({ - conversation_identifier: conversationId, - message: finalMessage, - }); + // Validate the response body inside the observed call so an unusable streaming + // response is counted as `upstream_error` instead of a false success. + const reader = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.sendAgentConversationMessageStreaming, + async () => { + const response = await ( + this.client as any + ).sendAgentConversationMessageStreaming({ + conversation_identifier: conversationId, + message: finalMessage, + }); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Failed to get reader from response body"); + } + return reader; + }, + ); + + const streamRecorder = this.createStreamMetricsRecorder(); + recordUpstreamStreamStartedMetric( + streamRecorder, + UPSTREAM_OPERATION_NAMES.sendAgentConversationMessageStreaming, + "success", + ); - const reader = response.body?.getReader(); - if (!reader) { - span?.setStatus({ - code: SpanStatusCode.ERROR, - message: "Failed to get reader from response body", + const processStreamPromise = + processSendAgentConversationMessageStreamingResponse( + conversationId, + reader, + appendStoredMessages, + (this.client as any).instanceUrl, + streamRecorder, + ).finally(() => { + this.scheduleBackgroundMetricsFlush(streamRecorder); }); - throw new Error("Failed to get reader from response body"); - } - // We don't await because we want to process the streaming response asynchronously - processSendAgentConversationMessageStreamingResponse( - conversationId, - reader, - appendStoredMessages, - (this.client as any).instanceUrl, - ); + if (this.metrics.waitUntil) { + try { + this.metrics.waitUntil(processStreamPromise); + } catch (error) { + console.error( + "[metrics] Failed to schedule upstream stream processing", + error, + ); + } + } else { + // Tests and non-Worker runtimes may not expose waitUntil. Keep the + // best-effort stream processing running without blocking the caller. + void processStreamPromise; + } span?.setStatus({ code: SpanStatusCode.OK, @@ -363,10 +469,14 @@ export class ThoughtSpotService { ); try { - const answer = await this.client.singleAnswer({ - query: question, - metadata_identifier: sourceId, - }); + const answer = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.singleAnswer, + () => + this.client.singleAnswer({ + query: question, + metadata_identifier: sourceId, + }), + ); const { session_identifier, generation_number } = answer as any; span?.setAttributes({ @@ -376,10 +486,14 @@ export class ThoughtSpotService { const [data, session, tml] = await Promise.all([ this.getAnswerData(question, session_identifier, generation_number), - (this.client as any).getAnswerSession({ - session_identifier, - generation_number, - }), + this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getAnswerSession, + () => + (this.client as any).getAnswerSession({ + session_identifier, + generation_number, + }), + ), shouldGetTML ? this.getAnswerTML(question, session_identifier, generation_number) : Promise.resolve(null), @@ -522,10 +636,14 @@ export class ThoughtSpotService { }, }; - const resp = await this.client.importMetadataTML({ - metadata_tmls: [JSON.stringify(tml)], - import_policy: "ALL_OR_NONE", - }); + const resp = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.importMetadataTml, + () => + this.client.importMetadataTML({ + metadata_tmls: [JSON.stringify(tml)], + import_policy: "ALL_OR_NONE", + }), + ); const liveboardUrl = `${(this.client as any).instanceUrl}/#/pinboard/${resp[0].response.header.id_guid}`; span?.setStatus({ @@ -544,18 +662,22 @@ export class ThoughtSpotService { span?.addEvent("get-data-sources"); - const resp = await this.client.searchMetadata({ - metadata: [ - { - type: "LOGICAL_TABLE", - }, - ], - record_size: 2000, - sort_options: { - field_name: "LAST_ACCESSED", - order: "DESC", - }, - }); + const resp = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.searchMetadata, + () => + this.client.searchMetadata({ + metadata: [ + { + type: "LOGICAL_TABLE", + }, + ], + record_size: 2000, + sort_options: { + field_name: "LAST_ACCESSED", + order: "DESC", + }, + }), + ); const results = resp // Tables can also be used for spotter now @@ -577,7 +699,10 @@ export class ThoughtSpotService { async getSessionInfo(): Promise { const span = getActiveSpan(); - const info = await (this.client as any).getSessionInfo(); + const info = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getSessionInfo, + () => (this.client as any).getSessionInfo(), + ); const devMixpanelToken = info.configInfo.mixpanelConfig.devSdkKey; const prodMixpanelToken = info.configInfo.mixpanelConfig.prodSdkKey; const mixpanelToken = info.configInfo.mixpanelConfig.production @@ -610,18 +735,22 @@ export class ThoughtSpotService { async searchWorksheets(searchTerm: string): Promise { const span = getActiveSpan(); - const resp = await this.client.searchMetadata({ - metadata: [ - { - type: "LOGICAL_TABLE", - }, - ], - record_size: 100, - sort_options: { - field_name: "NAME", - order: "ASC", - }, - }); + const resp = await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.searchMetadata, + () => + this.client.searchMetadata({ + metadata: [ + { + type: "LOGICAL_TABLE", + }, + ], + record_size: 100, + sort_options: { + field_name: "NAME", + order: "ASC", + }, + }), + ); const results = resp .filter((d) => d.metadata_header.type === "WORKSHEET") @@ -645,7 +774,10 @@ export class ThoughtSpotService { @WithSpan("validate-connection") async validateConnection(): Promise { try { - await (this.client as any).getSessionInfo(); + await this.observeUpstreamCall( + UPSTREAM_OPERATION_NAMES.getSessionInfo, + () => (this.client as any).getSessionInfo(), + ); return true; } catch (error) { // The decorator will automatically record the exception diff --git a/src/utils.ts b/src/utils.ts index 0107387..154544d 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1,4 +1,5 @@ import { type Span, SpanStatusCode } from "@opentelemetry/api"; +import type { ApiVersionMode } from "./metrics/runtime/metric-types"; import { getActiveSpan } from "./metrics/tracing/tracing-utils"; export type Props = { @@ -9,7 +10,9 @@ export type Props = { clientName: string; registrationDate: number; }; - apiVersion?: "beta"; + apiVersion?: string; + apiVersionMode?: ApiVersionMode; + apiRequestedVersion?: string; }; export class McpServerError extends Error { diff --git a/test/bearer.spec.ts b/test/bearer.spec.ts index f7ecfc4..a0af829 100644 --- a/test/bearer.spec.ts +++ b/test/bearer.spec.ts @@ -183,6 +183,7 @@ describe("Bearer Handler", () => { instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Custom Test Client", apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", }); // Verify the response @@ -208,6 +209,7 @@ describe("Bearer Handler", () => { instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Bearer Token client", apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", }); // Verify the response @@ -531,8 +533,10 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", }); expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); + expect(mockCtx.props.apiVersionMode).toBe("implicit_legacy"); }); it("should use backwards-compatibility-default apiVersion and ignore query param on /bearer/sse", async () => { @@ -553,8 +557,10 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", }); expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); + expect(mockCtx.props.apiVersionMode).toBe("implicit_legacy"); }); }); @@ -577,7 +583,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", apiVersion: "beta", + apiVersionMode: "beta", }); }); @@ -599,11 +607,13 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", apiVersion: "beta", + apiVersionMode: "beta", }); }); - it("should not inject apiVersion when query param is not present on /token/mcp", async () => { + it("should default unversioned /token/mcp requests to latest", async () => { const appWithBearer = withBearerHandler(app, ThoughtSpotMCP); const request = new Request("https://example.com/token/mcp", { @@ -614,12 +624,13 @@ describe("Bearer Handler", () => { await appWithBearer.fetch(request, mockEnv, mockCtx); - // Verify that props do not have apiVersion + // Verify that props reflect the effective served surface expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiVersion: "latest", + apiVersionMode: "implicit_latest", }); - expect(mockCtx.props.apiVersion).toBeUndefined(); }); it("should inject apiVersion with date format on /token/mcp", async () => { @@ -640,7 +651,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "2025-03-01", apiVersion: "2025-03-01", + apiVersionMode: "pinned", }); }); @@ -662,7 +675,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "2024-12-01", apiVersion: "2024-12-01", + apiVersionMode: "pinned", }); }); @@ -685,7 +700,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", apiVersion: "beta", + apiVersionMode: "beta", }); }); @@ -739,7 +756,8 @@ describe("Bearer Handler", () => { const result = await appWithBearer.fetch(request, mockEnv, mockCtx); expect(result.status).toBe(200); - expect(mockCtx.props.apiVersion).toBeUndefined(); + expect(mockCtx.props.apiVersion).toBe("latest"); + expect(mockCtx.props.apiVersionMode).toBe("implicit_latest"); }); it("should require bearer token on /token/mcp", async () => { diff --git a/test/metrics/runtime/analytics-engine-sink.spec.ts b/test/metrics/runtime/analytics-engine-sink.spec.ts index 7538866..0dcdbb0 100644 --- a/test/metrics/runtime/analytics-engine-sink.spec.ts +++ b/test/metrics/runtime/analytics-engine-sink.spec.ts @@ -1,6 +1,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { ANALYTICS_ENGINE_BLOB_FIELDS, + ANALYTICS_ENGINE_CONTEXT_FIELDS, ANALYTICS_ENGINE_DOUBLE_FIELDS, ANALYTICS_ENGINE_INDEX_FIELDS, ANALYTICS_ENGINE_LABEL_FIELDS, @@ -59,6 +60,8 @@ describe("AnalyticsEngineMetricsSink", () => { "oauth", "mcp", null, + null, + null, "success", null, "create_liveboard", @@ -67,6 +70,7 @@ describe("AnalyticsEngineMetricsSink", () => { null, null, null, + null, "production", "thoughtspot-mcp-server", "thoughtspot", @@ -97,6 +101,9 @@ describe("AnalyticsEngineMetricsSink", () => { it("keeps the Analytics Engine schema aligned with approved labels and resource attributes", () => { expect(ANALYTICS_ENGINE_LABEL_FIELDS).toEqual(APPROVED_METRIC_LABEL_KEYS); + expect(ANALYTICS_ENGINE_CONTEXT_FIELDS).toEqual([ + ["api_requested_version", "apiRequestedVersion"], + ]); expect(ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS).toEqual([ ["deployment_environment", "deployment.environment"], ["service_name", "service.name"], @@ -113,6 +120,7 @@ describe("AnalyticsEngineMetricsSink", () => { expect(ANALYTICS_ENGINE_BLOB_FIELDS).toEqual([ "metric_kind", ...APPROVED_METRIC_LABEL_KEYS, + "api_requested_version", "deployment_environment", "service_name", "service_namespace", @@ -176,6 +184,25 @@ describe("AnalyticsEngineMetricsSink", () => { ); }); + it("maps analytics context into Analytics Engine blobs", async () => { + const dataset = { writeDataPoint: vi.fn() }; + const sink = new AnalyticsEngineMetricsSink(dataset); + + await sink.flush({ + observations: [observation], + resourceAttributes: {}, + analyticsContext: { + apiRequestedVersion: "2026-10-01", + }, + }); + + expect(dataset.writeDataPoint).toHaveBeenCalledWith( + expect.objectContaining({ + blobs: expect.arrayContaining(["2026-10-01"]), + }), + ); + }); + it("continues writing remaining data points when one write fails", async () => { const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); const secondObservation = { diff --git a/test/metrics/runtime/metrics-recorder.spec.ts b/test/metrics/runtime/metrics-recorder.spec.ts index 41c4e4c..5e583fd 100644 --- a/test/metrics/runtime/metrics-recorder.spec.ts +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -57,6 +57,8 @@ describe("RequestMetricsRecorder", () => { expect(flushSpy).toHaveBeenCalledWith({ observations: recorder.snapshot(), resourceAttributes: { "service.name": "thoughtspot-mcp-server" }, + analyticsContext: undefined, + eventIdentity: undefined, }); }); @@ -150,6 +152,8 @@ describe("RequestMetricsRecorder", () => { }), ], resourceAttributes: {}, + analyticsContext: undefined, + eventIdentity: undefined, }); }); @@ -206,4 +210,46 @@ describe("RequestMetricsRecorder", () => { `[metrics] Ignoring metric recorded after flush: ${METRIC_NAMES.httpRequestsTotal}`, ); }); + + it("includes event identity in the flush payload when present", async () => { + const flushSpy = vi.fn().mockResolvedValue(undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.setEventIdentity({ + tenantId: "tenant-123", + }); + recorder.count(METRIC_NAMES.httpRequestsTotal); + await recorder.flush(); + + expect(flushSpy).toHaveBeenCalledWith( + expect.objectContaining({ + eventIdentity: { + tenantId: "tenant-123", + }, + }), + ); + }); + + it("includes analytics context in the flush payload when present", async () => { + const flushSpy = vi.fn().mockResolvedValue(undefined); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.setAnalyticsContext({ + apiRequestedVersion: "2026-10-01", + }); + recorder.count(METRIC_NAMES.httpRequestsTotal); + await recorder.flush(); + + expect(flushSpy).toHaveBeenCalledWith( + expect.objectContaining({ + analyticsContext: { + apiRequestedVersion: "2026-10-01", + }, + }), + ); + }); }); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index 38e83d7..f1935f5 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -7,6 +7,7 @@ import { recordBearerAuthRequestMetric, recordHttpRequestMetrics, recordStatusMetric, + resolveApiVersionLabels, resolveCanonicalApiVersionLabel, setMetricsRecorderOnExecutionContext, withRequestMetrics, @@ -271,7 +272,7 @@ describe("withRequestMetrics", () => { expect(result).toBe("handler-result"); expect(errorSpy).toHaveBeenCalledWith( - "[metrics] Failed to schedule request metrics flush", + "[metrics] Failed to schedule metrics flush", expect.any(Error), ); }); @@ -379,6 +380,7 @@ describe("withRequestMetrics", () => { labels: { api_surface: "mcp", api_version: "beta", + api_version_mode: "beta", auth_mode: "oauth", outcome: "success", route_group: "mcp", @@ -393,6 +395,7 @@ describe("withRequestMetrics", () => { labels: { api_surface: "mcp", api_version: "beta", + api_version_mode: "beta", auth_mode: "oauth", outcome: "success", route_group: "mcp", @@ -402,24 +405,73 @@ describe("withRequestMetrics", () => { ]); }); - it("labels versioned MCP routes as default when no explicit API version is requested", () => { + it("labels unversioned token routes as latest when no explicit API version is requested", () => { const ctx = {} as ExecutionContext; const request = new Request("https://example.com/token/mcp"); - expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("latest"); }); it("uses the effective default surface when bearer routes ignore an api-version query", () => { const ctx = { props: { apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", }, } as unknown as ExecutionContext; const request = new Request( "https://example.com/bearer/mcp?api-version=beta", ); - expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe( + "backwards-compatibility-default", + ); + }); + + it("labels legacy OAuth routes as implicit legacy when no selector is provided", () => { + const request = new Request("https://example.com/mcp"); + + expect(resolveApiVersionLabels(request, {} as ExecutionContext)).toEqual({ + apiReleaseDate: "2025-01-01", + apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", + }); + }); + + it("labels unversioned token routes as following latest", () => { + const request = new Request("https://example.com/token/mcp"); + + expect(resolveApiVersionLabels(request, {} as ExecutionContext)).toEqual({ + apiReleaseDate: "2026-05-01", + apiVersion: "latest", + apiVersionMode: "implicit_latest", + }); + }); + + it("labels date-based token routes as pinned even when they currently resolve to latest", () => { + const request = new Request( + "https://example.com/token/mcp?api-version=2026-05-01", + ); + + expect(resolveApiVersionLabels(request, {} as ExecutionContext)).toEqual({ + apiRequestedVersion: "2026-05-01", + apiReleaseDate: "2026-05-01", + apiVersion: "latest", + apiVersionMode: "pinned", + }); + }); + + it("labels explicit latest selectors separately from implicit latest", () => { + const request = new Request( + "https://example.com/token/mcp?api-version=latest", + ); + + expect(resolveApiVersionLabels(request, {} as ExecutionContext)).toEqual({ + apiRequestedVersion: "latest", + apiReleaseDate: "2026-05-01", + apiVersion: "latest", + apiVersionMode: "explicit_latest", + }); }); it("maps stable date-based versions onto the latest label", () => { @@ -437,13 +489,16 @@ describe("withRequestMetrics", () => { "https://example.com/token/mcp?api-version=2025-12-01", ); - expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("default"); + expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe( + "backwards-compatibility-default", + ); }); it("labels unresolved api-version values as unknown", () => { const ctx = { props: { apiVersion: "garbage", + apiRequestedVersion: "invalid", }, } as unknown as ExecutionContext; const request = new Request( @@ -453,6 +508,26 @@ describe("withRequestMetrics", () => { expect(resolveCanonicalApiVersionLabel(request, ctx)).toBe("unknown"); }); + it("keeps the normalized requested selector even when the served release resolves differently", () => { + const request = new Request( + "https://example.com/bearer/mcp?api-version=beta", + ); + const ctx = { + props: { + apiVersion: "backwards-compatibility-default", + apiRequestedVersion: "beta", + apiVersionMode: "implicit_legacy", + }, + } as unknown as ExecutionContext; + + expect(resolveApiVersionLabels(request, ctx)).toEqual({ + apiRequestedVersion: "beta", + apiReleaseDate: "2025-01-01", + apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", + }); + }); + it("records auth outcome counters from response status", () => { const recorder = createRequestMetricsRecorder(); @@ -489,4 +564,26 @@ describe("withRequestMetrics", () => { }), ]); }); + + it("preserves the requested selector for bearer auth traffic in analytics context", async () => { + const analyticsEngineSink = { flush: vi.fn().mockResolvedValue(undefined) }; + const recorder = createRequestMetricsRecorder( + { METRICS_SINK_MODE: "analytics_engine" }, + { analyticsEngineSink }, + ); + const request = new Request( + "https://example.com/bearer/mcp?api-version=beta", + ); + + recordBearerAuthRequestMetric(recorder, request, 401); + await recorder.flush(); + + expect(analyticsEngineSink.flush).toHaveBeenCalledWith( + expect.objectContaining({ + analyticsContext: { + apiRequestedVersion: "beta", + }, + }), + ); + }); }); diff --git a/test/servers/mcp-server.spec.ts b/test/servers/mcp-server.spec.ts index 39249bd..5fc58f0 100644 --- a/test/servers/mcp-server.spec.ts +++ b/test/servers/mcp-server.spec.ts @@ -1,10 +1,12 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; import { connect } from "mcp-testing-kit"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { MixpanelTracker } from "../../src/metrics/mixpanel/mixpanel"; +import { ANALYTICS_ENGINE_SCHEMA_VERSION } from "../../src/metrics/runtime/analytics-engine-sink"; +import { METRIC_NAMES } from "../../src/metrics/runtime/metric-types"; import { MCPServer } from "../../src/servers/mcp-server"; +import { StreamingMessagesStorageWithTtl } from "../../src/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; import * as thoughtspotClient from "../../src/thoughtspot/thoughtspot-client"; import { ThoughtSpotService } from "../../src/thoughtspot/thoughtspot-service"; -import { MixpanelTracker } from "../../src/metrics/mixpanel/mixpanel"; -import { StreamingMessagesStorageWithTtl } from "../../src/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; // Mock the MixpanelTracker vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ @@ -299,6 +301,84 @@ describe("MCP Server", () => { expect(result.isError).toBeUndefined(); expect((result.content as any[])[0].text).toBe("Pong"); }); + + it("writes tenant-scoped tool metrics without blocking the tool response", async () => { + const analyticsDataset = { + writeDataPoint: vi.fn(), + }; + const waitUntilPromises: Promise[] = []; + const metricsServer = new MCPServer( + { + props: { + ...mockProps, + apiVersion: "2025-03-01", + apiRequestedVersion: "2025-03-01", + apiVersionMode: "pinned", + }, + env: { + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + } as any, + ctx: { + waitUntil(promise: Promise) { + waitUntilPromises.push(promise); + }, + } as any, + }, + new StreamingMessagesStorageWithTtl(null as any, vi.fn(), vi.fn()), + ); + await metricsServer.init(); + + const { callTool } = connect(metricsServer); + const result = await callTool("ping", {}); + + expect(result.isError).toBeUndefined(); + expect(waitUntilPromises).toHaveLength(1); + + await Promise.all(waitUntilPromises); + + const dataPoints = analyticsDataset.writeDataPoint.mock.calls.map( + ([dataPoint]) => dataPoint, + ); + const toolCallCounter = dataPoints.find( + (dataPoint) => dataPoint.indexes?.[2] === METRIC_NAMES.toolCallsTotal, + ); + const toolDuration = dataPoints.find( + (dataPoint) => dataPoint.indexes?.[2] === METRIC_NAMES.toolDurationMs, + ); + + expect(toolCallCounter).toEqual( + expect.objectContaining({ + indexes: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + "test-org", + "test-user-123", + ], + blobs: expect.arrayContaining([ + "ping", + "success", + "mcp", + "backwards-compatibility-default", + "pinned", + "2025-01-01", + "2025-03-01", + ]), + }), + ); + expect(toolDuration).toEqual( + expect.objectContaining({ + indexes: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolDurationMs, + "test-org", + "test-user-123", + ], + }), + ); + }); }); describe("Check Connectivity Tool", () => { diff --git a/test/servers/version-registry.spec.ts b/test/servers/version-registry.spec.ts index fa60659..599b776 100644 --- a/test/servers/version-registry.spec.ts +++ b/test/servers/version-registry.spec.ts @@ -1,7 +1,10 @@ -import { describe, it, expect } from "vitest"; +import { describe, expect, it } from "vitest"; import { - resolveApiVersion, VERSION_REGISTRY, + YYYY_MM_DD_DATE_REGEX, + getReleaseDateIdentifier, + resolveApiVersion, + resolveApiVersionMetrics, } from "../../src/servers/version-registry"; // Helper: Validates if a given API version string is valid @@ -116,6 +119,35 @@ describe("Version Registry", () => { }); }); + describe("metrics helpers", () => { + it("exposes the shared YYYY-MM-DD regex", () => { + expect(YYYY_MM_DD_DATE_REGEX.test("2026-05-01")).toBe(true); + expect(YYYY_MM_DD_DATE_REGEX.test("2026-5-1")).toBe(false); + }); + + it("returns the resolved release date identifier when present", () => { + expect(getReleaseDateIdentifier(["latest", "2026-05-01"])).toBe( + "2026-05-01", + ); + expect(getReleaseDateIdentifier(["beta"])).toBeUndefined(); + }); + + it("maps requested selectors onto canonical metrics labels and release dates", () => { + expect(resolveApiVersionMetrics("latest")).toEqual({ + apiVersion: "latest", + apiReleaseDate: "2026-05-01", + }); + expect(resolveApiVersionMetrics("2025-03-15")).toEqual({ + apiVersion: "backwards-compatibility-default", + apiReleaseDate: "2025-01-01", + }); + expect(resolveApiVersionMetrics("beta")).toEqual({ + apiVersion: "beta", + apiReleaseDate: undefined, + }); + }); + }); + describe("VERSION_REGISTRY", () => { it("should contain beta version", () => { const betaVersion = VERSION_REGISTRY.find((v) => diff --git a/test/streaming-utils.spec.ts b/test/streaming-utils.spec.ts index 92d3dfa..bb0579c 100644 --- a/test/streaming-utils.spec.ts +++ b/test/streaming-utils.spec.ts @@ -1,4 +1,31 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +const tracingState = vi.hoisted(() => ({ + span: undefined as + | { + setAttribute: ReturnType; + setAttributes: ReturnType; + setStatus: ReturnType; + } + | undefined, +})); + +vi.mock("../src/metrics/tracing/tracing-utils", () => ({ + withSpan: async (_name: string, fn: (span: any) => Promise) => { + const span = { + setAttribute: vi.fn(), + setAttributes: vi.fn(), + setStatus: vi.fn(), + }; + tracingState.span = span; + return fn(span); + }, +})); + +import { METRIC_NAMES } from "../src/metrics/runtime/metric-types"; +import { + type MetricsRecorder, + NOOP_METRICS_RECORDER, +} from "../src/metrics/runtime/metrics-recorder"; import { processSendAgentConversationMessageStreamingResponse } from "../src/streaming-utils"; // Helper to build a ReadableStreamDefaultReader from an array of string chunks @@ -35,6 +62,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { beforeEach(() => { vi.clearAllMocks(); + tracingState.span = undefined; consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); }); @@ -85,6 +113,34 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ); }); + it("records upstream operation on stream message metrics", async () => { + const storage = makeMockStorage(); + const recorder: MetricsRecorder = { + ...NOOP_METRICS_RECORDER, + count: vi.fn(), + }; + const line = `data: ${JSON.stringify([{ type: "text", content: "Hello world", metadata: {} }])}\n`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage.appendMessages, + INSTANCE_URL, + recorder, + ); + + expect(recorder.count).toHaveBeenCalledWith( + METRIC_NAMES.upstreamStreamMessagesTotal, + 1, + expect.objectContaining({ + upstream_operation: "send_agent_conversation_message_streaming", + message_type: "text", + is_thinking: false, + }), + ); + }); + it("sets is_thinking=true on a text event when metadata.type is 'thinking'", async () => { const storage = makeMockStorage(); const line = `data: ${JSON.stringify([{ type: "text", content: "Reasoning...", metadata: { type: "thinking" } }])}\n`; @@ -237,6 +293,26 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { ]); }); + it("does not count error events as parsed text messages", async () => { + const storage = makeMockStorage(); + const line = `data: ${JSON.stringify([{ type: "error", display_message: "Something went wrong" }])}\n`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage.appendMessages, + INSTANCE_URL, + ); + + expect(tracingState.span?.setAttributes).toHaveBeenCalledWith({ + total_messages_parsed: 0, + total_text_messages_parsed: 0, + total_answer_messages_parsed: 0, + total_messages_ignored: 0, + }); + }); + it("ignores ack, notification, search_datasets, file, and conv_title events", async () => { const storage = makeMockStorage(); const ignoredTypes = [ diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index b9d3db7..3954204 100644 --- a/test/thoughtspot/thoughtspot-service.spec.ts +++ b/test/thoughtspot/thoughtspot-service.spec.ts @@ -1,13 +1,15 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { METRIC_NAMES } from "../../src/metrics/runtime/metric-types"; +import { createRequestMetricsRecorder } from "../../src/metrics/runtime/request-metrics"; import { - getRelevantQuestions, - getAnswerForQuestion, - fetchTMLAndCreateLiveboard, + ThoughtSpotService, createLiveboard, + fetchTMLAndCreateLiveboard, + getAnswerForQuestion, + getDataSourceSuggestions, getDataSources, + getRelevantQuestions, getSessionInfo, - getDataSourceSuggestions, - ThoughtSpotService, } from "../../src/thoughtspot/thoughtspot-service"; // Mock the ThoughtSpot REST API client @@ -234,14 +236,139 @@ describe("thoughtspot-service", () => { }); describe("sendAgentConversationMessageStreaming", () => { + it("records upstream streaming metrics and flushes stream message metrics in the background", async () => { + vi.useFakeTimers(); + + const analyticsDataset = { + writeDataPoint: vi.fn(), + }; + const waitUntilPromises: Promise[] = []; + const recorder = createRequestMetricsRecorder({ + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + }); + recorder.setAnalyticsContext({ + apiRequestedVersion: "latest", + }); + recorder.setEventIdentity({ + tenantId: "org-123", + userId: "user-123", + }); + + const encoder = new TextEncoder(); + const reader = { + read: vi + .fn() + .mockResolvedValueOnce({ + done: false, + value: encoder.encode( + 'data: [{"type":"text","content":"Hello","metadata":{}}]\n', + ), + }) + .mockResolvedValueOnce({ done: true, value: undefined }), + }; + + mockClient.sendAgentConversationMessageStreaming = vi + .fn() + .mockResolvedValue({ + body: { getReader: vi.fn().mockReturnValue(reader) }, + }); + + const appendMessages = vi.fn().mockResolvedValue(undefined); + + const service = new ThoughtSpotService(mockClient, { + recorder, + metricsEnv: { + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + }, + waitUntil(promise: Promise) { + waitUntilPromises.push(promise); + }, + analyticsContext: { + apiRequestedVersion: "latest", + }, + eventIdentity: { + tenantId: "org-123", + userId: "user-123", + }, + }); + + await service.sendAgentConversationMessageStreaming( + "conv-123", + "Show me revenue", + appendMessages, + ); + await recorder.flush(); + await vi.runAllTimersAsync(); + + for (let index = 0; index < waitUntilPromises.length; index++) { + await waitUntilPromises[index]; + } + + const dataPoints = analyticsDataset.writeDataPoint.mock.calls.map( + ([dataPoint]) => dataPoint, + ); + + expect( + dataPoints.some( + (dataPoint) => + dataPoint.indexes?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("latest") && + dataPoint.indexes?.[3] === "org-123" && + dataPoint.indexes?.[4] === "user-123", + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.indexes?.[2] === + METRIC_NAMES.upstreamStreamsStartedTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ), + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.indexes?.[2] === + METRIC_NAMES.upstreamStreamMessagesTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("text") && + dataPoint.blobs?.includes("false"), + ), + ).toBe(true); + + vi.useRealTimers(); + }); + it("should throw when the response body reader is unavailable", async () => { + const analyticsDataset = { + writeDataPoint: vi.fn(), + }; + const recorder = createRequestMetricsRecorder({ + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + }); mockClient.sendAgentConversationMessageStreaming = vi .fn() .mockResolvedValue({ body: undefined, }); - const service = new ThoughtSpotService(mockClient); + const service = new ThoughtSpotService(mockClient, { + recorder, + metricsEnv: { + METRICS_SINK_MODE: "analytics_engine", + ANALYTICS: analyticsDataset, + }, + }); await expect( service.sendAgentConversationMessageStreaming( @@ -252,6 +379,28 @@ describe("thoughtspot-service", () => { } as any, ), ).rejects.toThrow("Failed to get reader from response body"); + + await recorder.flush(); + + const dataPoints = analyticsDataset.writeDataPoint.mock.calls.map( + ([dataPoint]) => dataPoint, + ); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.indexes?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("upstream_error"), + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.indexes?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal, + ), + ).toBe(false); }); it("should throw when sending the streaming message fails", async () => { From 7f2697e7619e88384ddd58202b1829bb5427d715 Mon Sep 17 00:00:00 2001 From: kanwarbajwa <35466851+kanwarbajwa@users.noreply.github.com> Date: Thu, 7 May 2026 13:02:30 -0700 Subject: [PATCH 20/25] SCAL-310982 Fix Analytics Engine single-index writes (#141) ## Summary Fix the Workers Analytics Engine payload so production writes succeed again. Cloudflare only accepts a single index per data point. The current sink writes five indexes, which causes every Analytics Engine write to fail in production. ## What Changed - switch the Analytics Engine sink to a compact AE-specific row layout with one tenant-scoped index and a bounded blob set - keep is_done, drop is_thinking, and add analytical_session_id as Analytics-Engine-only context for session debugging - propagate analyticalSessionId through the metrics context and attach it to the session creation and send-session-message tool flows - update sink, recorder, request, MCP server, and ThoughtSpot service tests to lock the new schema ## Notes - keep mcp_metrics_v2; production writes were failing, so this becomes the first valid persisted layout rather than a new version bump - local design docs were updated for reference only and are not part of this repo commit --- src/metrics/runtime/analytics-engine-sink.ts | 68 ++++++++++------ src/metrics/runtime/metrics-recorder.ts | 3 + src/metrics/runtime/metrics-sink.ts | 5 +- src/servers/mcp-server-base.ts | 23 +++++- src/servers/mcp-server.ts | 12 ++- .../runtime/analytics-engine-sink.spec.ts | 81 +++++++++---------- test/metrics/runtime/metrics-recorder.spec.ts | 2 + test/metrics/runtime/request-metrics.spec.ts | 3 +- test/servers/mcp-server.spec.ts | 15 ++-- test/thoughtspot/thoughtspot-service.spec.ts | 21 ++--- 10 files changed, 140 insertions(+), 93 deletions(-) diff --git a/src/metrics/runtime/analytics-engine-sink.ts b/src/metrics/runtime/analytics-engine-sink.ts index 78b7b99..44f4c46 100644 --- a/src/metrics/runtime/analytics-engine-sink.ts +++ b/src/metrics/runtime/analytics-engine-sink.ts @@ -1,11 +1,11 @@ import { - APPROVED_METRIC_LABEL_KEYS, METRIC_NAMES, type MetricLabels, type MetricName, } from "./metric-types"; import type { MetricAnalyticsContext, + MetricEventIdentity, MetricObservation, MetricResourceAttributes, MetricsFlushPayload, @@ -23,28 +23,42 @@ export type AnalyticsEngineDatasetLike = { }; export const ANALYTICS_ENGINE_SCHEMA_VERSION = "mcp_metrics_v2"; - -export const ANALYTICS_ENGINE_INDEX_FIELDS = [ - "schema_version", - "event_family", - "metric_name", - "tenant_id", - "user_id", +const ANALYTICS_ENGINE_FALLBACK_INDEX = "shared"; + +// Cloudflare Analytics Engine allows one sampling index and up to 20 blobs. This +// schema is intentionally AE-specific instead of mirroring every approved metric +// label, so we can keep tenant/user + version + tool/upstream context together +// without exceeding those limits. +export const ANALYTICS_ENGINE_INDEX_FIELDS = ["tenant_id"] as const; + +export const ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS = [ + ["tenant_id", "tenantId"], + ["user_id", "userId"], +] as const satisfies readonly (readonly [string, keyof MetricEventIdentity])[]; + +export const ANALYTICS_ENGINE_LABEL_FIELDS = [ + "route_group", + "auth_mode", + "api_version", + "api_version_mode", + "api_release_date", + "outcome", + "status_class", + "tool_name", + "upstream_operation", + "message_type", + "is_done", ] as const; -export const ANALYTICS_ENGINE_LABEL_FIELDS = APPROVED_METRIC_LABEL_KEYS; - export const ANALYTICS_ENGINE_CONTEXT_FIELDS = [ ["api_requested_version", "apiRequestedVersion"], + ["analytical_session_id", "analyticalSessionId"], ] as const satisfies readonly (readonly [ string, keyof MetricAnalyticsContext, ])[]; export const ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS = [ - ["deployment_environment", "deployment.environment"], - ["service_name", "service.name"], - ["service_namespace", "service.namespace"], ["service_version", "service.version"], ] as const satisfies readonly (readonly [ string, @@ -52,7 +66,11 @@ export const ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS = [ ])[]; export const ANALYTICS_ENGINE_BLOB_FIELDS = [ + "schema_version", + "event_family", + "metric_name", "metric_kind", + ...ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS.map(([field]) => field), ...ANALYTICS_ENGINE_LABEL_FIELDS, ...ANALYTICS_ENGINE_CONTEXT_FIELDS.map(([field]) => field), ...ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS.map(([field]) => field), @@ -63,11 +81,6 @@ export const ANALYTICS_ENGINE_DOUBLE_FIELDS = [ "timestamp_ms", ] as const; -type AnalyticsEngineIdentity = { - tenantId?: string; - userId?: string; -}; - type AnalyticsEngineMetricFamily = | "analysis" | "auth" @@ -96,6 +109,13 @@ function getResourceAttribute( return nullableString(resourceAttributes[key]); } +function getEventIdentityField( + identity: MetricEventIdentity, + key: keyof MetricEventIdentity, +): string | null { + return nullableString(identity[key]); +} + function getAnalyticsContextField( analyticsContext: MetricAnalyticsContext, key: keyof MetricAnalyticsContext, @@ -155,18 +175,20 @@ export function toAnalyticsEngineDataPoint( observation: MetricObservation, resourceAttributes: MetricResourceAttributes, analyticsContext: MetricAnalyticsContext = {}, - identity: AnalyticsEngineIdentity = {}, + identity: MetricEventIdentity = {}, ): AnalyticsEngineDataPointLike { return { indexes: [ + nullableString(identity.tenantId) ?? ANALYTICS_ENGINE_FALLBACK_INDEX, + ], + blobs: [ ANALYTICS_ENGINE_SCHEMA_VERSION, getAnalyticsEngineMetricFamily(observation.name), observation.name, - nullableString(identity.tenantId), - nullableString(identity.userId), - ], - blobs: [ observation.kind, + ...ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS.map(([, key]) => + getEventIdentityField(identity, key), + ), ...ANALYTICS_ENGINE_LABEL_FIELDS.map((key) => getLabel(observation.labels, key), ), diff --git a/src/metrics/runtime/metrics-recorder.ts b/src/metrics/runtime/metrics-recorder.ts index 7cb5398..9b15eed 100644 --- a/src/metrics/runtime/metrics-recorder.ts +++ b/src/metrics/runtime/metrics-recorder.ts @@ -78,6 +78,9 @@ export class RequestMetricsRecorder implements MetricsRecorder { if (context.apiRequestedVersion) { nextContext.apiRequestedVersion = context.apiRequestedVersion; } + if (context.analyticalSessionId) { + nextContext.analyticalSessionId = context.analyticalSessionId; + } this.analyticsContext = nextContext; } diff --git a/src/metrics/runtime/metrics-sink.ts b/src/metrics/runtime/metrics-sink.ts index 4d5fe40..4cf833d 100644 --- a/src/metrics/runtime/metrics-sink.ts +++ b/src/metrics/runtime/metrics-sink.ts @@ -28,10 +28,11 @@ export type MetricEventIdentity = { // Sink-specific context for dimensions that should not become generic metric labels. // `api_version`, `api_version_mode`, and `api_release_date` stay in `MetricObservation.labels` // because both Grafana and Analytics Engine should receive them as low-cardinality dimensions. -// `apiRequestedVersion` lives here instead because we only want it in Analytics Engine as -// request-debug context, not as an extra Grafana label. +// `apiRequestedVersion` and `analyticalSessionId` live here instead because they are +// Analytics-Engine-only debug/context fields and should not widen Grafana label cardinality. export type MetricAnalyticsContext = { apiRequestedVersion?: string; + analyticalSessionId?: string; }; export type MetricsFlushPayload = { diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index cae4af2..e246630 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -198,7 +198,10 @@ export abstract class BaseMCPServer extends Server { ); } - protected getThoughtSpotService(recorder?: MetricsRecorder) { + protected getThoughtSpotService( + recorder?: MetricsRecorder, + analyticsContextOverride?: MetricAnalyticsContext, + ) { return new ThoughtSpotService( getThoughtSpotClient( this.ctx.props.instanceUrl, @@ -208,7 +211,9 @@ export abstract class BaseMCPServer extends Server { recorder, metricsEnv: this.ctx.env as unknown as Record, waitUntil: this.getMetricsWaitUntil(), - analyticsContext: this.getMetricAnalyticsContext(), + analyticsContext: this.mergeMetricAnalyticsContext( + analyticsContextOverride, + ), eventIdentity: this.getMetricEventIdentity(), }, ); @@ -242,6 +247,20 @@ export abstract class BaseMCPServer extends Server { }; } + protected mergeMetricAnalyticsContext( + override?: MetricAnalyticsContext, + ): MetricAnalyticsContext | undefined { + const baseContext = this.getMetricAnalyticsContext(); + if (!baseContext && !override) { + return undefined; + } + + return { + ...baseContext, + ...override, + }; + } + protected getMetricEventIdentity(): MetricEventIdentity | undefined { if (!this.sessionInfo) { return undefined; diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 48dbe8d..a7af53a 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -386,6 +386,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; await this.getThoughtSpotService(recorder).createAgentConversation( data_source_id, ); + recorder.setAnalyticsContext({ + analyticalSessionId: response.conversation_id, + }); span?.setAttribute("analytical_session_id", response.conversation_id); // Conversation is initialized in streamingMessageStorage from callSendSessionMessage, @@ -405,6 +408,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; const span = trace.getSpan(context.active()); const { analytical_session_id, message, additional_context } = SendSessionMessageInputSchema.parse(request.params.arguments); + recorder.setAnalyticsContext({ + analyticalSessionId: analytical_session_id, + }); span?.setAttributes({ analytical_session_id, has_additional_context: !!additional_context, @@ -424,9 +430,9 @@ Provide this url to the user as a link to view the liveboard in ThoughtSpot.`; ); } - await this.getThoughtSpotService( - recorder, - ).sendAgentConversationMessageStreaming( + await this.getThoughtSpotService(recorder, { + analyticalSessionId: analytical_session_id, + }).sendAgentConversationMessageStreaming( analytical_session_id, message, storageService.appendMessages.bind(storageService), diff --git a/test/metrics/runtime/analytics-engine-sink.spec.ts b/test/metrics/runtime/analytics-engine-sink.spec.ts index 0dcdbb0..bf1f303 100644 --- a/test/metrics/runtime/analytics-engine-sink.spec.ts +++ b/test/metrics/runtime/analytics-engine-sink.spec.ts @@ -3,6 +3,7 @@ import { ANALYTICS_ENGINE_BLOB_FIELDS, ANALYTICS_ENGINE_CONTEXT_FIELDS, ANALYTICS_ENGINE_DOUBLE_FIELDS, + ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS, ANALYTICS_ENGINE_INDEX_FIELDS, ANALYTICS_ENGINE_LABEL_FIELDS, ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS, @@ -12,10 +13,7 @@ import { getAnalyticsEngineMetricFamily, toAnalyticsEngineDataPoint, } from "../../../src/metrics/runtime/analytics-engine-sink"; -import { - APPROVED_METRIC_LABEL_KEYS, - METRIC_NAMES, -} from "../../../src/metrics/runtime/metric-types"; +import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; import type { MetricObservation } from "../../../src/metrics/runtime/metrics-sink"; describe("AnalyticsEngineMetricsSink", () => { @@ -46,19 +44,16 @@ describe("AnalyticsEngineMetricsSink", () => { "service.version": "0.5.0", }); - expect(dataPoint.indexes).toEqual([ + expect(dataPoint.indexes).toEqual(["shared"]); + expect(dataPoint.blobs).toEqual([ ANALYTICS_ENGINE_SCHEMA_VERSION, "tool", METRIC_NAMES.toolCallsTotal, + "counter", null, null, - ]); - expect(dataPoint.blobs).toEqual([ - "counter", - "mcp", "mcp", "oauth", - "mcp", null, null, null, @@ -70,10 +65,6 @@ describe("AnalyticsEngineMetricsSink", () => { null, null, null, - null, - "production", - "thoughtspot-mcp-server", - "thoughtspot", "0.5.0", ]); expect(dataPoint.doubles).toEqual([2, 1_714_000_000_000]); @@ -99,37 +90,49 @@ describe("AnalyticsEngineMetricsSink", () => { ).toBe("auth"); }); - it("keeps the Analytics Engine schema aligned with approved labels and resource attributes", () => { - expect(ANALYTICS_ENGINE_LABEL_FIELDS).toEqual(APPROVED_METRIC_LABEL_KEYS); + it("keeps the Analytics Engine schema aligned with the compact single-index layout", () => { + expect(ANALYTICS_ENGINE_INDEX_FIELDS).toEqual(["tenant_id"]); + expect(ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS).toEqual([ + ["tenant_id", "tenantId"], + ["user_id", "userId"], + ]); + expect(ANALYTICS_ENGINE_LABEL_FIELDS).toEqual([ + "route_group", + "auth_mode", + "api_version", + "api_version_mode", + "api_release_date", + "outcome", + "status_class", + "tool_name", + "upstream_operation", + "message_type", + "is_done", + ]); expect(ANALYTICS_ENGINE_CONTEXT_FIELDS).toEqual([ ["api_requested_version", "apiRequestedVersion"], + ["analytical_session_id", "analyticalSessionId"], ]); expect(ANALYTICS_ENGINE_RESOURCE_ATTRIBUTE_FIELDS).toEqual([ - ["deployment_environment", "deployment.environment"], - ["service_name", "service.name"], - ["service_namespace", "service.namespace"], ["service_version", "service.version"], ]); - expect(ANALYTICS_ENGINE_INDEX_FIELDS).toEqual([ + expect(ANALYTICS_ENGINE_BLOB_FIELDS).toEqual([ "schema_version", "event_family", "metric_name", + "metric_kind", "tenant_id", "user_id", - ]); - expect(ANALYTICS_ENGINE_BLOB_FIELDS).toEqual([ - "metric_kind", - ...APPROVED_METRIC_LABEL_KEYS, + ...ANALYTICS_ENGINE_LABEL_FIELDS, "api_requested_version", - "deployment_environment", - "service_name", - "service_namespace", + "analytical_session_id", "service_version", ]); expect(ANALYTICS_ENGINE_DOUBLE_FIELDS).toEqual([ "metric_value", "timestamp_ms", ]); + expect(ANALYTICS_ENGINE_BLOB_FIELDS).toHaveLength(20); }); it("writes one data point per observation", async () => { @@ -146,19 +149,13 @@ describe("AnalyticsEngineMetricsSink", () => { expect(dataset.writeDataPoint).toHaveBeenCalledTimes(1); expect(dataset.writeDataPoint).toHaveBeenCalledWith( expect.objectContaining({ - indexes: [ - ANALYTICS_ENGINE_SCHEMA_VERSION, - "tool", - METRIC_NAMES.toolCallsTotal, - null, - null, - ], + indexes: ["shared"], doubles: [2, 1_714_000_000_000], }), ); }); - it("maps request-scoped event identity into Analytics Engine indexes", async () => { + it("maps request-scoped event identity into the Analytics Engine index and blobs", async () => { const dataset = { writeDataPoint: vi.fn() }; const sink = new AnalyticsEngineMetricsSink(dataset); @@ -173,13 +170,8 @@ describe("AnalyticsEngineMetricsSink", () => { expect(dataset.writeDataPoint).toHaveBeenCalledWith( expect.objectContaining({ - indexes: [ - ANALYTICS_ENGINE_SCHEMA_VERSION, - "tool", - METRIC_NAMES.toolCallsTotal, - "tenant-123", - "user-456", - ], + indexes: ["tenant-123"], + blobs: expect.arrayContaining(["tenant-123", "user-456"]), }), ); }); @@ -193,12 +185,13 @@ describe("AnalyticsEngineMetricsSink", () => { resourceAttributes: {}, analyticsContext: { apiRequestedVersion: "2026-10-01", + analyticalSessionId: "conv-123", }, }); expect(dataset.writeDataPoint).toHaveBeenCalledWith( expect.objectContaining({ - blobs: expect.arrayContaining(["2026-10-01"]), + blobs: expect.arrayContaining(["2026-10-01", "conv-123"]), }), ); }); @@ -211,7 +204,7 @@ describe("AnalyticsEngineMetricsSink", () => { } satisfies MetricObservation; const dataset = { writeDataPoint: vi.fn((dataPoint) => { - if (dataPoint?.indexes?.[2] === METRIC_NAMES.toolCallsTotal) { + if (dataPoint?.blobs?.[2] === METRIC_NAMES.toolCallsTotal) { throw new Error("write failed"); } }), diff --git a/test/metrics/runtime/metrics-recorder.spec.ts b/test/metrics/runtime/metrics-recorder.spec.ts index 5e583fd..313da80 100644 --- a/test/metrics/runtime/metrics-recorder.spec.ts +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -240,6 +240,7 @@ describe("RequestMetricsRecorder", () => { recorder.setAnalyticsContext({ apiRequestedVersion: "2026-10-01", + analyticalSessionId: "conv-123", }); recorder.count(METRIC_NAMES.httpRequestsTotal); await recorder.flush(); @@ -248,6 +249,7 @@ describe("RequestMetricsRecorder", () => { expect.objectContaining({ analyticsContext: { apiRequestedVersion: "2026-10-01", + analyticalSessionId: "conv-123", }, }), ); diff --git a/test/metrics/runtime/request-metrics.spec.ts b/test/metrics/runtime/request-metrics.spec.ts index f1935f5..225cb75 100644 --- a/test/metrics/runtime/request-metrics.spec.ts +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -242,7 +242,8 @@ describe("withRequestMetrics", () => { expect(analyticsDataset.writeDataPoint).toHaveBeenCalledTimes(1); expect(analyticsDataset.writeDataPoint).toHaveBeenCalledWith( expect.objectContaining({ - indexes: expect.arrayContaining([METRIC_NAMES.httpRequestsTotal]), + indexes: ["shared"], + blobs: expect.arrayContaining([METRIC_NAMES.httpRequestsTotal]), }), ); }); diff --git a/test/servers/mcp-server.spec.ts b/test/servers/mcp-server.spec.ts index 5fc58f0..0fed040 100644 --- a/test/servers/mcp-server.spec.ts +++ b/test/servers/mcp-server.spec.ts @@ -341,25 +341,23 @@ describe("MCP Server", () => { ([dataPoint]) => dataPoint, ); const toolCallCounter = dataPoints.find( - (dataPoint) => dataPoint.indexes?.[2] === METRIC_NAMES.toolCallsTotal, + (dataPoint) => dataPoint.blobs?.[2] === METRIC_NAMES.toolCallsTotal, ); const toolDuration = dataPoints.find( - (dataPoint) => dataPoint.indexes?.[2] === METRIC_NAMES.toolDurationMs, + (dataPoint) => dataPoint.blobs?.[2] === METRIC_NAMES.toolDurationMs, ); expect(toolCallCounter).toEqual( expect.objectContaining({ - indexes: [ + indexes: ["test-org"], + blobs: expect.arrayContaining([ ANALYTICS_ENGINE_SCHEMA_VERSION, "tool", METRIC_NAMES.toolCallsTotal, "test-org", "test-user-123", - ], - blobs: expect.arrayContaining([ "ping", "success", - "mcp", "backwards-compatibility-default", "pinned", "2025-01-01", @@ -369,13 +367,14 @@ describe("MCP Server", () => { ); expect(toolDuration).toEqual( expect.objectContaining({ - indexes: [ + indexes: ["test-org"], + blobs: expect.arrayContaining([ ANALYTICS_ENGINE_SCHEMA_VERSION, "tool", METRIC_NAMES.toolDurationMs, "test-org", "test-user-123", - ], + ]), }), ); }); diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index 3954204..3143f04 100644 --- a/test/thoughtspot/thoughtspot-service.spec.ts +++ b/test/thoughtspot/thoughtspot-service.spec.ts @@ -249,6 +249,7 @@ describe("thoughtspot-service", () => { }); recorder.setAnalyticsContext({ apiRequestedVersion: "latest", + analyticalSessionId: "conv-123", }); recorder.setEventIdentity({ tenantId: "org-123", @@ -287,6 +288,7 @@ describe("thoughtspot-service", () => { }, analyticsContext: { apiRequestedVersion: "latest", + analyticalSessionId: "conv-123", }, eventIdentity: { tenantId: "org-123", @@ -313,20 +315,20 @@ describe("thoughtspot-service", () => { expect( dataPoints.some( (dataPoint) => - dataPoint.indexes?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamCallsTotal && dataPoint.blobs?.includes( "send_agent_conversation_message_streaming", ) && dataPoint.blobs?.includes("latest") && - dataPoint.indexes?.[3] === "org-123" && - dataPoint.indexes?.[4] === "user-123", + dataPoint.blobs?.includes("conv-123") && + dataPoint.indexes?.[0] === "org-123" && + dataPoint.blobs?.includes("user-123"), ), ).toBe(true); expect( dataPoints.some( (dataPoint) => - dataPoint.indexes?.[2] === - METRIC_NAMES.upstreamStreamsStartedTotal && + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal && dataPoint.blobs?.includes( "send_agent_conversation_message_streaming", ), @@ -335,13 +337,12 @@ describe("thoughtspot-service", () => { expect( dataPoints.some( (dataPoint) => - dataPoint.indexes?.[2] === - METRIC_NAMES.upstreamStreamMessagesTotal && + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamMessagesTotal && dataPoint.blobs?.includes( "send_agent_conversation_message_streaming", ) && dataPoint.blobs?.includes("text") && - dataPoint.blobs?.includes("false"), + dataPoint.blobs?.includes("conv-123"), ), ).toBe(true); @@ -388,7 +389,7 @@ describe("thoughtspot-service", () => { expect( dataPoints.some( (dataPoint) => - dataPoint.indexes?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamCallsTotal && dataPoint.blobs?.includes( "send_agent_conversation_message_streaming", ) && @@ -398,7 +399,7 @@ describe("thoughtspot-service", () => { expect( dataPoints.some( (dataPoint) => - dataPoint.indexes?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal, + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal, ), ).toBe(false); }); From 95f1a0ae57a046f529110aab0b0d9d6c57071ca7 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 7 May 2026 13:48:25 -0700 Subject: [PATCH 21/25] Add more coverage for new tools (#137) - Add coverage for new tools and related functions - Increase min coverage level to 85% for all metrics Co-authored-by: Rifdhan Nazeer Co-authored-by: Mourya Balabhadra <162541770+mouryabalabhadra@users.noreply.github.com> --- test/servers/mcp-server-base.spec.ts | 261 +++++++++---- test/servers/mcp-server.spec.ts | 381 ++++++++++++++++++- test/thoughtspot/thoughtspot-service.spec.ts | 194 ++++++++++ vitest.config.ts | 2 +- 4 files changed, 755 insertions(+), 83 deletions(-) diff --git a/test/servers/mcp-server-base.spec.ts b/test/servers/mcp-server-base.spec.ts index 20b9c9e..8080180 100644 --- a/test/servers/mcp-server-base.spec.ts +++ b/test/servers/mcp-server-base.spec.ts @@ -2,6 +2,8 @@ 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"; +import { TrackEvent, type Tracker } from "../../src/metrics"; +import { StreamingMessagesStorageWithTtl } from "../../src/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; // Mock the MixpanelTracker vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ @@ -10,6 +12,15 @@ vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ })), })); +vi.mock( + "../../src/streaming-message-storage-with-ttl/streaming-message-storage-with-ttl", + () => ({ + StreamingMessagesStorageWithTtl: vi.fn().mockImplementation(() => ({})), + }), +); + +const mockEnv = {} as Env; + // Test subclass to expose protected methods class TestMCPServer extends MCPServer { public testCreateMultiContentSuccessResponse( @@ -47,36 +58,49 @@ class TestMCPServer extends MCPServer { public testIsDatasourceDiscoveryAvailable() { return this.isDatasourceDiscoveryAvailable(); } + + public getTrackers() { + return this.trackers; + } + + public getSessionInfo() { + return this.sessionInfo; + } } describe("MCP Server Base", () => { let server: TestMCPServer; let mockProps: any; + let mockStreamingStorage: any; + + const makeSessionInfo = (overrides: any = {}) => ({ + 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", + enableSpotterDataSourceDiscovery: true, + ...overrides.configInfo, + }, + userName: "test-user", + currentOrgId: "test-org", + privileges: [], + ...overrides, + }); 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", - enableSpotterDataSourceDiscovery: true, - }, - userName: "test-user", - currentOrgId: "test-org", - privileges: [], - }), + getSessionInfo: vi.fn().mockResolvedValue(makeSessionInfo()), searchMetadata: vi.fn().mockResolvedValue([]), instanceUrl: "https://test.thoughtspot.cloud", } as any); @@ -91,7 +115,16 @@ describe("MCP Server Base", () => { }, }; - server = new TestMCPServer({ props: mockProps }); + mockStreamingStorage = new StreamingMessagesStorageWithTtl( + null as any, + vi.fn(), + vi.fn(), + ); + + server = new TestMCPServer( + { props: mockProps, env: mockEnv }, + mockStreamingStorage, + ); }); describe("Response Helper Methods", () => { @@ -111,7 +144,7 @@ describe("MCP Server Base", () => { "Multiple messages", ); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(3); expect(result.content[0].text).toBe("First message"); expect(result.content[1].text).toBe("Second message"); @@ -124,7 +157,7 @@ describe("MCP Server Base", () => { "No messages", ); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(0); }); @@ -136,7 +169,7 @@ describe("MCP Server Base", () => { "Array response", ); - expect(result.isError).toBeUndefined(); + expect((result as any).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" }); @@ -146,7 +179,7 @@ describe("MCP Server Base", () => { it("should create array success response with empty array", () => { const result = server.testCreateArraySuccessResponse([], "Empty array"); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(0); }); @@ -158,7 +191,7 @@ describe("MCP Server Base", () => { "Single item response", ); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(1); expect(result.content[0]).toEqual({ type: "text", text: "Single item" }); }); @@ -185,7 +218,7 @@ describe("MCP Server Base", () => { it("should create success response with message", () => { const result = server.testCreateSuccessResponse("Operation successful"); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(1); expect(result.content[0].text).toBe("Operation successful"); }); @@ -201,7 +234,7 @@ describe("MCP Server Base", () => { "Structured response", ); - expect(result.isError).toBeUndefined(); + expect((result as any).isError).toBeUndefined(); expect(result.content).toHaveLength(1); expect(result.content[0].text).toBe(JSON.stringify(structuredContent)); expect(result.structuredContent).toEqual(structuredContent); @@ -209,6 +242,16 @@ describe("MCP Server Base", () => { }); describe("Datasource Discovery Check", () => { + it("should return false before init is called (sessionInfo not set)", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const result = server.testIsDatasourceDiscoveryAvailable(); + expect(result).toBe(false); + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining("sessionInfo is not initialized"), + ); + warnSpy.mockRestore(); + }); + it("should return true when enableSpotterDataSourceDiscovery is enabled", async () => { await server.init(); const result = server.testIsDatasourceDiscoveryAvailable(); @@ -217,68 +260,61 @@ describe("MCP Server Base", () => { 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", - enableSpotterDataSourceDiscovery: false, - }, - userName: "test-user", - currentOrgId: "test-org", - privileges: [], - }), + getSessionInfo: vi + .fn() + .mockResolvedValue( + makeSessionInfo({ + configInfo: { enableSpotterDataSourceDiscovery: false }, + }), + ), searchMetadata: vi.fn().mockResolvedValue([]), instanceUrl: "https://test.thoughtspot.cloud", } as any); - const testServer = new TestMCPServer({ props: mockProps }); + const testServer = new TestMCPServer( + { props: mockProps, env: mockEnv }, + mockStreamingStorage, + ); await testServer.init(); - const result = testServer.testIsDatasourceDiscoveryAvailable(); - expect(result).toBe(false); + expect(testServer.testIsDatasourceDiscoveryAvailable()).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", - enableSpotterDataSourceDiscovery: undefined, - }, - userName: "test-user", - currentOrgId: "test-org", - privileges: [], - }), + getSessionInfo: vi + .fn() + .mockResolvedValue( + makeSessionInfo({ + configInfo: { enableSpotterDataSourceDiscovery: undefined }, + }), + ), searchMetadata: vi.fn().mockResolvedValue([]), instanceUrl: "https://test.thoughtspot.cloud", } as any); - const testServer = new TestMCPServer({ props: mockProps }); + const testServer = new TestMCPServer( + { props: mockProps, env: mockEnv }, + mockStreamingStorage, + ); await testServer.init(); - const result = testServer.testIsDatasourceDiscoveryAvailable(); - expect(result).toBe(false); + expect(testServer.testIsDatasourceDiscoveryAvailable()).toBe(false); }); }); describe("initializeService", () => { + it("should set sessionInfo and register MixpanelTracker on successful init", async () => { + await server.init(); + + expect(server.getSessionInfo()).toBeDefined(); + expect(server.getSessionInfo().userGUID).toBe("test-user-123"); + expect(MixpanelTracker).toHaveBeenCalledWith( + expect.objectContaining({ userGUID: "test-user-123" }), + mockProps.clientName, + ); + // The tracker was added — the trackers set should have one entry + expect(server.getTrackers().size).toBe(1); + }); + it("should catch and log error if getSessionInfo throws", async () => { vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ getSessionInfo: vi @@ -293,8 +329,8 @@ describe("MCP Server Base", () => { .mockImplementation(() => {}); const testServer = new TestMCPServer( - { props: mockProps } as any, - null as any, + { props: mockProps, env: mockEnv }, + mockStreamingStorage, ); await expect(testServer.init()).resolves.not.toThrow(); @@ -302,24 +338,87 @@ describe("MCP Server Base", () => { "Error initializing session info:", expect.any(Error), ); + // sessionInfo should remain unset, no tracker registered + expect(testServer.getSessionInfo()).toBeUndefined(); + expect(testServer.getTrackers().size).toBe(0); consoleErrorSpy.mockRestore(); }); }); - describe("Server Initialization", () => { - it("should initialize with custom server name and version", () => { - const customServer = new TestMCPServer( - { props: mockProps }, - "CustomServer", - "2.0.0", + describe("init — TrackEvent.Init", () => { + it("should track the Init event after initialization", async () => { + await server.init(); + + // The single tracker added is the MixpanelTracker mock + const mockTrackerInstance = + vi.mocked(MixpanelTracker).mock.results[0].value; + expect(mockTrackerInstance.track).toHaveBeenCalledWith( + TrackEvent.Init, + {}, ); + }); + + it("should track Init even when getSessionInfo fails (trackers may be empty)", async () => { + vi.spyOn(thoughtspotClient, "getThoughtSpotClient").mockReturnValue({ + getSessionInfo: vi.fn().mockRejectedValue(new Error("fail")), + instanceUrl: "https://test.thoughtspot.cloud", + } as any); - expect(customServer).toBeDefined(); + const testServer = new TestMCPServer( + { props: mockProps, env: mockEnv }, + mockStreamingStorage, + ); + // No tracker was registered, but init should not throw + await expect(testServer.init()).resolves.not.toThrow(); }); + }); - it("should initialize with default server name and version", () => { + describe("addTracker", () => { + it("should register a tracker so it receives subsequent track calls", async () => { + await server.init(); + + const customTracker: Tracker = { track: vi.fn() }; + await server.addTracker(customTracker); + + expect(server.getTrackers().has(customTracker)).toBe(true); + }); + + it("should not duplicate a tracker added twice", async () => { + await server.init(); + + const customTracker: Tracker = { track: vi.fn() }; + await server.addTracker(customTracker); + await server.addTracker(customTracker); + + // Trackers extends Set — same reference added twice stays as one entry + const customEntries = [...server.getTrackers()].filter( + (t) => t === customTracker, + ); + expect(customEntries).toHaveLength(1); + }); + }); + + describe("getStorageService", () => { + it("should throw when env.CONVERSATION_STORAGE_OBJECT is not set", () => { + // env is an empty object — accessing CONVERSATION_STORAGE_OBJECT is undefined, + // and StorageServiceClient construction should fail or return an unusable instance. + // We just assert it doesn't throw at construction time (the error surfaces on use). + expect(() => (server as any).getStorageService()).not.toThrow(); + }); + }); + + describe("Server Initialization", () => { + it("should be defined after construction", () => { expect(server).toBeDefined(); }); + + it("should be defined with a fresh env and props", () => { + const freshServer = new TestMCPServer( + { props: mockProps, env: mockEnv }, + mockStreamingStorage, + ); + expect(freshServer).toBeDefined(); + }); }); }); diff --git a/test/servers/mcp-server.spec.ts b/test/servers/mcp-server.spec.ts index 0fed040..55d9d63 100644 --- a/test/servers/mcp-server.spec.ts +++ b/test/servers/mcp-server.spec.ts @@ -1,5 +1,5 @@ import { connect } from "mcp-testing-kit"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { MixpanelTracker } from "../../src/metrics/mixpanel/mixpanel"; import { ANALYTICS_ENGINE_SCHEMA_VERSION } from "../../src/metrics/runtime/analytics-engine-sink"; import { METRIC_NAMES } from "../../src/metrics/runtime/metric-types"; @@ -1323,4 +1323,383 @@ describe("MCP Server", () => { expect(mockClientInstance.searchMetadata).toHaveBeenCalledTimes(1); }); }); + + describe("Send Session Message Tool", () => { + let mockStorageService: { + initializeConversation: ReturnType; + appendMessages: ReturnType; + getNewMessages: ReturnType; + }; + let mockSendAgentConversationMessageStreaming: ReturnType; + + beforeEach(() => { + mockStorageService = { + initializeConversation: vi.fn().mockResolvedValue(undefined), + appendMessages: vi.fn().mockResolvedValue(undefined), + getNewMessages: vi + .fn() + .mockResolvedValue({ messages: [], isDone: true }), + }; + vi.spyOn(server as any, "getStorageService").mockReturnValue( + mockStorageService, + ); + + mockSendAgentConversationMessageStreaming = vi + .spyOn( + ThoughtSpotService.prototype, + "sendAgentConversationMessageStreaming", + ) + .mockResolvedValue(undefined) as unknown as ReturnType; + }); + + it("should send a message and return success", async () => { + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("send_session_message", { + analytical_session_id: "conv-abc-123", + message: "What is the total revenue?", + }); + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).success).toBe(true); + expect(mockStorageService.initializeConversation).toHaveBeenCalledWith( + "conv-abc-123", + ); + expect(mockSendAgentConversationMessageStreaming).toHaveBeenCalledWith( + "conv-abc-123", + "What is the total revenue?", + expect.any(Function), + undefined, + ); + }); + + it("should pass additional_context to the service", async () => { + await server.init(); + const { callTool } = connect(server); + + await callTool("send_session_message", { + analytical_session_id: "conv-abc-123", + message: "Compare revenue by region", + additional_context: "The user is focused on North America", + }); + + expect(mockSendAgentConversationMessageStreaming).toHaveBeenCalledWith( + "conv-abc-123", + "Compare revenue by region", + expect.any(Function), + "The user is focused on North America", + ); + }); + + it("should return error when conversation is already in progress", async () => { + mockStorageService.initializeConversation.mockRejectedValue( + new Error("Conversation already exists and is not marked done"), + ); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("send_session_message", { + analytical_session_id: "conv-abc-123", + message: "Follow-up question", + }); + + expect(result.isError).toBe(true); + expect((result.content as any[])[0].text).toContain( + "ERROR: The analytical session has an ongoing response", + ); + }); + }); + + describe("Get Session Updates Tool", () => { + let mockStorageService: { + initializeConversation: ReturnType; + appendMessages: ReturnType; + getNewMessages: ReturnType; + }; + + beforeEach(() => { + vi.useFakeTimers(); + mockStorageService = { + initializeConversation: vi.fn().mockResolvedValue(undefined), + appendMessages: vi.fn().mockResolvedValue(undefined), + getNewMessages: vi + .fn() + .mockResolvedValue({ messages: [], isDone: true }), + }; + vi.spyOn(server as any, "getStorageService").mockReturnValue( + mockStorageService, + ); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("should return session updates when done", async () => { + const messages = [ + { + type: "text" as const, + is_thinking: false, + text: "The total revenue is $1,000,000", + }, + ]; + mockStorageService.getNewMessages.mockResolvedValue({ + messages, + isDone: true, + }); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("get_session_updates", { + analytical_session_id: "conv-abc-123", + }); + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).is_done).toBe(true); + expect((result.structuredContent as any).session_updates).toEqual( + messages, + ); + }); + + it("should return empty updates when not done and no messages", async () => { + mockStorageService.getNewMessages.mockResolvedValue({ + messages: [], + isDone: false, + }); + + await server.init(); + const { callTool } = connect(server); + + const resultPromise = callTool("get_session_updates", { + analytical_session_id: "conv-abc-123", + }); + await vi.runAllTimersAsync(); + const result = await resultPromise; + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).is_done).toBe(false); + expect((result.structuredContent as any).session_updates).toEqual([]); + }); + + it("should accumulate messages across multiple polls", async () => { + const firstMessages = [ + { type: "text" as const, is_thinking: true, text: "Thinking..." }, + ]; + const secondMessages = [ + { + type: "text" as const, + is_thinking: false, + text: "Here is the answer", + }, + ]; + + mockStorageService.getNewMessages + .mockResolvedValueOnce({ messages: firstMessages, isDone: false }) + .mockResolvedValue({ messages: secondMessages, isDone: true }); + + await server.init(); + const { callTool } = connect(server); + + const resultPromise = callTool("get_session_updates", { + analytical_session_id: "conv-abc-123", + }); + await vi.runAllTimersAsync(); + const result = await resultPromise; + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).is_done).toBe(true); + expect((result.structuredContent as any).session_updates).toEqual([ + ...firstMessages, + ...secondMessages, + ]); + }); + + it("should return answer type updates", async () => { + const answerMessage = { + type: "answer" as const, + is_thinking: false, + answer_id: '{"session_id":"sess-123","gen_no":1}', + answer_title: "Revenue by Region", + answer_query: "revenue by region", + iframe_url: + "https://test.thoughtspot.cloud/#/embed/answer?sessionId=sess-123", + }; + mockStorageService.getNewMessages.mockResolvedValue({ + messages: [answerMessage], + isDone: true, + }); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("get_session_updates", { + analytical_session_id: "conv-abc-123", + }); + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).is_done).toBe(true); + expect( + (result.structuredContent as any).session_updates[0], + ).toMatchObject(answerMessage); + }); + }); + + describe("Create Dashboard Tool", () => { + it("should create dashboard successfully with valid answer_ids", async () => { + const mockFetchTMLAndCreateLiveboard = vi + .spyOn(ThoughtSpotService.prototype, "fetchTMLAndCreateLiveboard") + .mockResolvedValue({ + url: "https://test.thoughtspot.cloud/#/pinboard/dashboard-456", + error: null, + }); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("create_dashboard", { + title: "Revenue Dashboard", + note_tile: "

Revenue Analysis

Generated on May 5, 2026

", + answers: [ + { + answer_id: JSON.stringify({ session_id: "sess-123", gen_no: 1 }), + title: "Total Revenue", + }, + { + answer_id: JSON.stringify({ session_id: "sess-456", gen_no: 2 }), + title: "Revenue by Region", + }, + ], + }); + + expect(result.isError).toBeUndefined(); + expect((result.structuredContent as any).link).toBe( + "https://test.thoughtspot.cloud/#/pinboard/dashboard-456", + ); + expect(mockFetchTMLAndCreateLiveboard).toHaveBeenCalledWith( + "Revenue Dashboard", + [ + { + title: "Total Revenue", + session_identifier: "sess-123", + generation_number: 1, + }, + { + title: "Revenue by Region", + session_identifier: "sess-456", + generation_number: 2, + }, + ], + "

Revenue Analysis

Generated on May 5, 2026

", + ); + }); + + it("should return error when answer_id format is invalid", async () => { + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("create_dashboard", { + title: "Revenue Dashboard", + note_tile: "

Summary

", + answers: [ + { + answer_id: "not-valid-json", + title: "Total Revenue", + }, + ], + }); + + expect(result.isError).toBe(true); + expect((result.content as any[])[0].text).toContain( + "ERROR: Invalid answer_id format", + ); + }); + + it("should return error when answer_id is missing session_id or gen_no", async () => { + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("create_dashboard", { + title: "Revenue Dashboard", + note_tile: "

Summary

", + answers: [ + { + answer_id: JSON.stringify({ foo: "bar" }), + title: "Total Revenue", + }, + ], + }); + + expect(result.isError).toBe(true); + expect((result.content as any[])[0].text).toContain( + "ERROR: Invalid answer_id format", + ); + }); + + it("should return error when liveboard creation fails", async () => { + vi.spyOn( + ThoughtSpotService.prototype, + "fetchTMLAndCreateLiveboard", + ).mockResolvedValue({ + url: undefined, + error: new Error("Failed to import TML"), + }); + + await server.init(); + const { callTool } = connect(server); + + const result = await callTool("create_dashboard", { + title: "Revenue Dashboard", + note_tile: "

Summary

", + answers: [ + { + answer_id: JSON.stringify({ session_id: "sess-123", gen_no: 1 }), + title: "Total Revenue", + }, + ], + }); + + expect(result.isError).toBe(true); + expect((result.content as any[])[0].text).toContain( + "ERROR: Encountered an error while creating the dashboard", + ); + }); + + it("should handle a single answer correctly", async () => { + const mockFetchTMLAndCreateLiveboard = vi + .spyOn(ThoughtSpotService.prototype, "fetchTMLAndCreateLiveboard") + .mockResolvedValue({ + url: "https://test.thoughtspot.cloud/#/pinboard/dashboard-789", + error: null, + }); + + await server.init(); + const { callTool } = connect(server); + + await callTool("create_dashboard", { + title: "Single Answer Dashboard", + note_tile: "

One answer

", + answers: [ + { + answer_id: JSON.stringify({ session_id: "sess-999", gen_no: 3 }), + title: "Key Metric", + }, + ], + }); + + expect(mockFetchTMLAndCreateLiveboard).toHaveBeenCalledWith( + "Single Answer Dashboard", + [ + { + title: "Key Metric", + session_identifier: "sess-999", + generation_number: 3, + }, + ], + "

One answer

", + ); + }); + }); }); diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index 3143f04..e16fea8 100644 --- a/test/thoughtspot/thoughtspot-service.spec.ts +++ b/test/thoughtspot/thoughtspot-service.spec.ts @@ -512,6 +512,200 @@ describe("thoughtspot-service", () => { await vi.runAllTimersAsync(); vi.useRealTimers(); }); + + it("should call appendStoredMessages with parsed text messages and mark done when stream ends", async () => { + const encoder = new TextEncoder(); + const reader = { + read: vi + .fn() + .mockResolvedValueOnce({ + done: false, + value: encoder.encode( + 'data: [{"type":"text","content":"The revenue is $1M"}]\n', + ), + }) + .mockResolvedValueOnce({ done: true, value: undefined }), + }; + + mockClient.sendAgentConversationMessageStreaming = vi + .fn() + .mockResolvedValue({ + body: { getReader: vi.fn().mockReturnValue(reader) }, + }); + + const appendStoredMessages = vi.fn().mockResolvedValue(undefined); + + const service = new ThoughtSpotService(mockClient); + await service.sendAgentConversationMessageStreaming( + "conv-123", + "Show me revenue", + appendStoredMessages, + ); + + // Wait for the fire-and-forget async loop to complete + await vi.waitFor(() => { + expect(appendStoredMessages).toHaveBeenLastCalledWith( + "conv-123", + [], + true, + ); + }); + + expect(appendStoredMessages).toHaveBeenCalledWith("conv-123", [ + { is_thinking: false, type: "text", text: "The revenue is $1M" }, + ]); + // Final call marks the conversation as done + expect(appendStoredMessages).toHaveBeenLastCalledWith( + "conv-123", + [], + true, + ); + }); + + it("should call appendStoredMessages with parsed answer messages", async () => { + const encoder = new TextEncoder(); + const answerEvent = JSON.stringify([ + { + type: "answer", + metadata: { + session_id: "sess-abc", + gen_no: 2, + transaction_id: "txn-xyz", + generation_number: 1, + title: "Revenue by Region", + sage_query: "revenue by region", + }, + }, + ]); + const reader = { + read: vi + .fn() + .mockResolvedValueOnce({ + done: false, + value: encoder.encode(`data: ${answerEvent}\n`), + }) + .mockResolvedValueOnce({ done: true, value: undefined }), + }; + + mockClient.sendAgentConversationMessageStreaming = vi + .fn() + .mockResolvedValue({ + body: { getReader: vi.fn().mockReturnValue(reader) }, + }); + + const appendStoredMessages = vi.fn().mockResolvedValue(undefined); + + const service = new ThoughtSpotService(mockClient); + await service.sendAgentConversationMessageStreaming( + "conv-123", + "Show me revenue by region", + appendStoredMessages, + ); + + await vi.waitFor(() => { + expect(appendStoredMessages).toHaveBeenLastCalledWith( + "conv-123", + [], + true, + ); + }); + + expect(appendStoredMessages).toHaveBeenCalledWith("conv-123", [ + { + is_thinking: false, + type: "answer", + answer_id: JSON.stringify({ session_id: "sess-abc", gen_no: 2 }), + answer_title: "Revenue by Region", + answer_query: "revenue by region", + iframe_url: + "https://test.thoughtspot.com/?tsmcp=true#/embed/conv-assist-answer?sessionId=sess-abc&genNo=2&acSessionId=txn-xyz&acGenNo=1", + }, + ]); + expect(appendStoredMessages).toHaveBeenLastCalledWith( + "conv-123", + [], + true, + ); + }); + + it("should skip heartbeat lines and blank lines during streaming", async () => { + const encoder = new TextEncoder(); + const reader = { + read: vi + .fn() + .mockResolvedValueOnce({ + done: false, + value: encoder.encode( + ': heartbeat\n\ndata: [{"type":"text","content":"Done"}]\n', + ), + }) + .mockResolvedValueOnce({ done: true, value: undefined }), + }; + + mockClient.sendAgentConversationMessageStreaming = vi + .fn() + .mockResolvedValue({ + body: { getReader: vi.fn().mockReturnValue(reader) }, + }); + + const appendStoredMessages = vi.fn().mockResolvedValue(undefined); + + const service = new ThoughtSpotService(mockClient); + await service.sendAgentConversationMessageStreaming( + "conv-123", + "Test", + appendStoredMessages, + ); + + await vi.waitFor(() => { + expect(appendStoredMessages).toHaveBeenLastCalledWith( + "conv-123", + [], + true, + ); + }); + + // Only the actual text message should be stored, not heartbeats/blanks + expect(appendStoredMessages).toHaveBeenCalledWith("conv-123", [ + { is_thinking: false, type: "text", text: "Done" }, + ]); + }); + + it("should correctly identify thinking messages via metadata type", async () => { + const encoder = new TextEncoder(); + const reader = { + read: vi + .fn() + .mockResolvedValueOnce({ + done: false, + value: encoder.encode( + 'data: [{"type":"text","content":"Thinking...","metadata":{"type":"thinking"}}]\n', + ), + }) + .mockResolvedValueOnce({ done: true, value: undefined }), + }; + + mockClient.sendAgentConversationMessageStreaming = vi + .fn() + .mockResolvedValue({ + body: { getReader: vi.fn().mockReturnValue(reader) }, + }); + + const appendStoredMessages = vi.fn().mockResolvedValue(undefined); + + const service = new ThoughtSpotService(mockClient); + await service.sendAgentConversationMessageStreaming( + "conv-123", + "Test", + appendStoredMessages, + ); + + await vi.waitFor(() => { + expect(appendStoredMessages).toHaveBeenCalledWith("conv-123", [ + { is_thinking: true, type: "text", text: "Thinking..." }, + ]); + }); + }); }); describe("getAnswerForQuestion", () => { diff --git a/vitest.config.ts b/vitest.config.ts index e887558..840eb62 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -24,7 +24,7 @@ export default defineWorkersConfig({ thresholds: { lines: 85, functions: 85, - branches: 75, + branches: 85, statements: 85, }, }, From d9aebe61d321ba91ff6682df9b5f9a251ab03876 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 7 May 2026 13:48:36 -0700 Subject: [PATCH 22/25] Switch from legacy KV to SQLite storage (#142) - No implementation change required Co-authored-by: Rifdhan Nazeer --- wrangler.jsonc | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/wrangler.jsonc b/wrangler.jsonc index 9b8d436..20077e0 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -47,6 +47,18 @@ "deleted_classes": [ "ThoughtSpotOpenAIDeepResearchMCP" ] + }, + { + "tag": "v6", + "deleted_classes": [ + "ConversationStorageServer" + ] + }, + { + "tag": "v7", + "new_sqlite_classes": [ + "ConversationStorageServer" + ] } ], "kv_namespaces": [{ From 231d9f2299b1330e6db678673e6198a037a3a102 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 7 May 2026 14:11:57 -0700 Subject: [PATCH 23/25] Rename storage server class to fix deployment issue (#144) Co-authored-by: Rifdhan Nazeer --- src/index.ts | 4 ++-- src/servers/conversation-storage-server.ts | 2 +- test/servers/conversation-storage-server.spec.ts | 8 ++++---- worker-configuration.d.ts | 2 +- wrangler.jsonc | 9 +++------ 5 files changed, 11 insertions(+), 14 deletions(-) diff --git a/src/index.ts b/src/index.ts index a7873e3..dbec08b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -17,10 +17,10 @@ import { withRequestMetrics, } from "./metrics/runtime/request-metrics"; import { PUBLIC_ROUTES } from "./routes"; -import { ConversationStorageServer } from "./servers/conversation-storage-server"; +import { ConversationStorageServerSQLite } from "./servers/conversation-storage-server"; import { MCPServer } from "./servers/mcp-server"; -export { ConversationStorageServer }; +export { ConversationStorageServerSQLite }; // OTEL configuration function const config: ResolveConfigFn = (env: Env, _trigger) => { diff --git a/src/servers/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts index 19ebd75..69e7275 100644 --- a/src/servers/conversation-storage-server.ts +++ b/src/servers/conversation-storage-server.ts @@ -21,7 +21,7 @@ const READ_BOOKMARK_KEY = "read-bookmark"; * POST /storage//append —> appendMessagesAndRestartTtl * GET /storage//messages —> getNewMessagesAndUpdateBookmark */ -export class ConversationStorageServer { +export class ConversationStorageServerSQLite { private conversationId = ""; constructor( diff --git a/test/servers/conversation-storage-server.spec.ts b/test/servers/conversation-storage-server.spec.ts index 3f022f4..effe5c6 100644 --- a/test/servers/conversation-storage-server.spec.ts +++ b/test/servers/conversation-storage-server.spec.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { ConversationStorageServer } from "../../src/servers/conversation-storage-server"; +import { ConversationStorageServerSQLite } from "../../src/servers/conversation-storage-server"; import type { Message, StreamingMessagesState, @@ -69,7 +69,7 @@ function createMockStorage() { function createServer(mock: ReturnType) { const state = { storage: mock.storage } as unknown as DurableObjectState; - return new ConversationStorageServer(state, {} as Env); + return new ConversationStorageServerSQLite(state, {} as Env); } function makeRequest( @@ -121,9 +121,9 @@ const STORAGE_BATCH_SIZE = 127; // Tests // --------------------------------------------------------------------------- -describe("ConversationStorageServer", () => { +describe("ConversationStorageServerSQLite", () => { let mock: ReturnType; - let server: ConversationStorageServer; + let server: ConversationStorageServerSQLite; beforeEach(() => { mock = createMockStorage(); diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 8175aa3..33671c6 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -9,7 +9,7 @@ declare namespace Cloudflare { interface Env { OAUTH_KV: KVNamespace; MCP_OBJECT: DurableObjectNamespace; - CONVERSATION_STORAGE_OBJECT: DurableObjectNamespace; + CONVERSATION_STORAGE_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; ASSETS: Fetcher; } diff --git a/wrangler.jsonc b/wrangler.jsonc index 20077e0..3f853e3 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -18,7 +18,7 @@ "name": "MCP_OBJECT" }, { - "class_name": "ConversationStorageServer", + "class_name": "ConversationStorageServerSQLite", "name": "CONVERSATION_STORAGE_OBJECT" } ] @@ -52,12 +52,9 @@ "tag": "v6", "deleted_classes": [ "ConversationStorageServer" - ] - }, - { - "tag": "v7", + ], "new_sqlite_classes": [ - "ConversationStorageServer" + "ConversationStorageServerSQLite" ] } ], From 31cc0bb2b360ed1a7aaf28ae9f3530f00e7d04c6 Mon Sep 17 00:00:00 2001 From: Rifdhan Nazeer Date: Thu, 7 May 2026 14:17:06 -0700 Subject: [PATCH 24/25] Remove deletion migration --- wrangler.jsonc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index 3f853e3..bf06716 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -52,7 +52,10 @@ "tag": "v6", "deleted_classes": [ "ConversationStorageServer" - ], + ] + }, + { + "tag": "v7", "new_sqlite_classes": [ "ConversationStorageServerSQLite" ] From 8abce48f229794bb30a5e6c9bef81d60083d7a8a Mon Sep 17 00:00:00 2001 From: "cloudflare-workers-and-pages[bot]" <73139402+cloudflare-workers-and-pages[bot]@users.noreply.github.com> Date: Thu, 7 May 2026 22:15:36 +0000 Subject: [PATCH 25/25] Update wrangler config name to test-sb --- wrangler.jsonc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index bf06716..e9922df 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -5,7 +5,7 @@ { "keep_vars": true, "$schema": "node_modules/wrangler/config-schema.json", - "name": "thoughtspot-mcp-server", + "name": "test-sb", "main": "src/index.ts", "compatibility_date": "2025-04-17", "compatibility_flags": [