diff --git a/README.md b/README.md index 0682c87..ab2f46e 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" ] } } @@ -374,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):** @@ -391,28 +389,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. 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/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/bearer.ts b/src/bearer.ts index 5a18de1..9b9478f 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -1,6 +1,26 @@ 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"; + +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 @@ -8,69 +28,128 @@ import { validateAndSanitizeUrl } from "./oauth-manager/oauth-utils"; * @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( +async function handleTokenAuth( req: Request, env: Env, ctx: ExecutionContext, MCPServer: typeof ThoughtSpotMCP, - supportApiVersion: boolean, -): Response { - const authHeader = req.headers.get("authorization"); - if (!authHeader) { - return new Response("Bearer token is required", { status: 400 }); - } + apiVersionOverride?: string, + 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, - }; + // Build props object + const props: any = { + accessToken: accessToken, + instanceUrl: validateAndSanitizeUrl(tsHost), + clientName, + }; + const requestedApiVersion = url.searchParams.get("api-version"); - // Add api-version support only for /token endpoints (supports "beta" or "YYYY-MM-DD" format) - if (supportApiVersion) { - const apiVersion = 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 ?? + 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; + (ctx as any).props = props; - // Route to appropriate handler - const pathname = url.pathname; - if (pathname.endsWith("/mcp")) { - return MCPServer.serve("/mcp").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 }); + } - if (pathname.endsWith("/sse")) { - return MCPServer.serveSSE("/sse").fetch(req, env, ctx); + recordBearerAuthRequestMetric( + recorder, + req, + response.status, + authMetricRouteGroup, + ); + return response; + } catch (error) { + recordBearerAuthRequestMetric(recorder, req, 500, authMetricRouteGroup); + throw error; } - - return new Response("Not found", { status: 404 }); } export function withBearerHandler( @@ -79,14 +158,21 @@ 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) => { - return handleTokenAuth(req, env, ctx, MCPServer, false); + app.mount(PUBLIC_ROUTE_PREFIXES.bearer, (req, env, ctx) => { + return handleTokenAuth( + req, + env, + ctx, + MCPServer, + "backwards-compatibility-default", + "bearer", + ); }); // NEW: /token endpoints - supports api-version query params // Recommended for all new implementations - app.mount("/token", (req, env, ctx) => { - return handleTokenAuth(req, env, ctx, MCPServer, true); + app.mount(PUBLIC_ROUTE_PREFIXES.token, (req, env, ctx) => { + return handleTokenAuth(req, env, ctx, MCPServer, undefined, "token"); }); return app; 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/handlers.ts b/src/handlers.ts index 95d847d..0ade3fa 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -2,23 +2,59 @@ 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 { 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, 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 } }>(); +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) { @@ -235,104 +271,174 @@ 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, 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("/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); + 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; } }); -app.get("/callback", async (c) => { +app.get(PUBLIC_ROUTES.callback, async (c) => { try { const htmlContent = await handler.handleCallback( c.req.raw, 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("/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), { + 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; } }); -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); - export default app; diff --git a/src/index.ts b/src/index.ts index 1510805..dbec08b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,17 +1,26 @@ -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 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"; +import { ConversationStorageServerSQLite } from "./servers/conversation-storage-server"; import { MCPServer } from "./servers/mcp-server"; -import { apiServer } from "./servers/api-server"; -import { withBearerHandler } from "./bearer"; -import { OpenAIDeepResearchMCPServer } from "./servers/openai-mcp-server"; + +export { ConversationStorageServerSQLite }; // OTEL configuration function const config: ResolveConfigFn = (env: Env, _trigger) => { @@ -27,11 +36,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, @@ -46,17 +50,30 @@ function createMCPRouter( ctx: ExecutionContext, ): Promise { const url = new URL(request.url); - const apiVersion = url.searchParams.get("api-version"); + const requestedApiVersion = url.searchParams.get("api-version"); + let apiVersion = requestedApiVersion; + let apiVersionMode: ApiVersionMode; - // 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"; + apiVersionMode = "implicit_legacy"; + } else { + apiVersionMode = resolveRequestedApiVersionMode(apiVersion); } + // Inject apiVersion into props + const originalProps = (ctx as any).props || {}; + (ctx as any).props = { + ...originalProps, + apiVersion, + apiRequestedVersion: requestedApiVersion + ? normalizeRequestedApiVersionForAnalytics(requestedApiVersion) + : undefined, + apiVersionMode, + }; + // Route to the appropriate serve method return serverClass[serveMethod](path, options).fetch(request, env, ctx); }, @@ -66,20 +83,21 @@ 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, }, 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 @@ -120,6 +138,38 @@ 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 (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/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/analytics-engine-sink.ts b/src/metrics/runtime/analytics-engine-sink.ts new file mode 100644 index 0000000..44f4c46 --- /dev/null +++ b/src/metrics/runtime/analytics-engine-sink.ts @@ -0,0 +1,249 @@ +import { + METRIC_NAMES, + type MetricLabels, + type MetricName, +} from "./metric-types"; +import type { + MetricAnalyticsContext, + MetricEventIdentity, + 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_v2"; +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_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 = [ + ["service_version", "service.version"], +] as const satisfies readonly (readonly [ + string, + keyof MetricResourceAttributes, +])[]; + +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), +] as const; + +export const ANALYTICS_ENGINE_DOUBLE_FIELDS = [ + "metric_value", + "timestamp_ms", +] as const; + +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]); +} + +function getEventIdentityField( + identity: MetricEventIdentity, + key: keyof MetricEventIdentity, +): string | null { + return nullableString(identity[key]); +} + +function getAnalyticsContextField( + analyticsContext: MetricAnalyticsContext, + key: keyof MetricAnalyticsContext, +): string | null { + return nullableString(analyticsContext[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, + analyticsContext: MetricAnalyticsContext = {}, + identity: MetricEventIdentity = {}, +): AnalyticsEngineDataPointLike { + return { + indexes: [ + nullableString(identity.tenantId) ?? ANALYTICS_ENGINE_FALLBACK_INDEX, + ], + blobs: [ + ANALYTICS_ENGINE_SCHEMA_VERSION, + getAnalyticsEngineMetricFamily(observation.name), + observation.name, + observation.kind, + ...ANALYTICS_ENGINE_IDENTITY_BLOB_FIELDS.map(([, key]) => + getEventIdentityField(identity, key), + ), + ...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), + ), + ], + 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.analyticsContext, + 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/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/grafana-otlp-sink.ts b/src/metrics/runtime/grafana-otlp-sink.ts new file mode 100644 index 0000000..14fe1bc --- /dev/null +++ b/src/metrics/runtime/grafana-otlp-sink.ts @@ -0,0 +1,417 @@ +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) + // 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 { + 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/metric-context.ts b/src/metrics/runtime/metric-context.ts new file mode 100644 index 0000000..b3fb4d0 --- /dev/null +++ b/src/metrics/runtime/metric-context.ts @@ -0,0 +1,243 @@ +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.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", + }, +} as const satisfies Record; + +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 ( + 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 + ) { + 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) { + 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 + ) { + return "none"; + } + return "unknown"; +} + +export function resolvePathMetricContext( + pathname: string, +): RequestMetricContext { + const explicitContext = getExplicitRouteContext(pathname); + if (explicitContext) { + return explicitContext; + } + + 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..2fb4110 --- /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", + "api_version_mode", + "api_release_date", + "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_apps_challenge" + | "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" | "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) { + 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..9b15eed --- /dev/null +++ b/src/metrics/runtime/metrics-recorder.ts @@ -0,0 +1,195 @@ +import { + type MetricKind, + type MetricLabelInput, + type MetricName, + getMetricKind, + normalizeMetricLabels, +} from "./metric-types"; +import type { + MetricAnalyticsContext, + MetricEventIdentity, + 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; + setAnalyticsContext(context?: MetricAnalyticsContext): void; + setEventIdentity(identity?: MetricEventIdentity): void; + 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 {}, + setAnalyticsContext(_context): void {}, + setEventIdentity(_identity): void {}, + flush(): Promise { + return NOOP_FLUSH_PROMISE; + }, + snapshot(): readonly MetricObservation[] { + return NOOP_METRIC_OBSERVATIONS; + }, +}; + +export class RequestMetricsRecorder implements MetricsRecorder { + private readonly observations: MetricObservation[] = []; + private analyticsContext?: MetricAnalyticsContext; + private eventIdentity?: MetricEventIdentity; + 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); + } + + setAnalyticsContext(context?: MetricAnalyticsContext): void { + if (!context) { + return; + } + + const nextContext: MetricAnalyticsContext = { + ...this.analyticsContext, + }; + if (context.apiRequestedVersion) { + nextContext.apiRequestedVersion = context.apiRequestedVersion; + } + if (context.analyticalSessionId) { + nextContext.analyticalSessionId = context.analyticalSessionId; + } + 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]; + } + + flush(): Promise { + if (!this.flushPromise) { + this.flushPromise = this.flushInternal(); + } + + 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 }, + 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 new file mode 100644 index 0000000..4cf833d --- /dev/null +++ b/src/metrics/runtime/metrics-sink.ts @@ -0,0 +1,47 @@ +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 MetricEventIdentity = { + tenantId?: string; + 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` 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 = { + observations: readonly MetricObservation[]; + resourceAttributes: MetricResourceAttributes; + eventIdentity?: MetricEventIdentity; + analyticsContext?: MetricAnalyticsContext; +}; + +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..b50e254 --- /dev/null +++ b/src/metrics/runtime/request-metrics.ts @@ -0,0 +1,419 @@ +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, + type MetricOutcome, +} from "./metric-types"; +import { + type MetricsRecorder, + NOOP_METRICS_RECORDER, + RequestMetricsRecorder, + scheduleMetricsFlush, +} from "./metrics-recorder"; +import type { MetricsSink } from "./metrics-sink"; +import { + type ConfiguredMetricsSinks, + type MetricsEnvLike, + createConfiguredMetricsSink, + resolveMetricsRuntimeConfig, +} from "./runtime-config"; + +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; + apiRequestedVersion?: unknown; + apiVersionMode?: unknown; + }; +}; + +export function setMetricsRecorderOnExecutionContext( + ctx: ExecutionContext, + recorder: MetricsRecorder, +): MetricsRecorder { + (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL] = recorder; + 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 { + return (ctx as MetricsExecutionContext)[METRICS_RECORDER_SYMBOL]; +} + +export function clearMetricsRecorderFromExecutionContext( + ctx: ExecutionContext, +): void { + 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"; +} + +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 (requestedApiVersion === "backwards-compatibility-default") { + return "explicit_legacy"; + } + 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"; +} + +/** + * 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, +): ApiVersionLabels { + const requestContext = resolveRequestMetricContext(request); + if (!VERSIONED_REQUEST_ROUTE_GROUPS.has(requestContext.routeGroup)) { + return {}; + } + + const requestedApiVersion = new URL(request.url).searchParams.get( + "api-version", + ); + 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 { + const resolved = resolveApiVersionMetrics(effectiveApiVersion); + return { + ...resolved, + apiVersionMode: + typeof effectiveApiVersionMode === "string" && + effectiveApiVersionMode.length > 0 + ? effectiveApiVersionMode + : getImplicitApiVersionMode(resolved.apiVersion), + }; + } catch { + return { + apiReleaseDate: undefined, + apiVersion: "unknown", + apiVersionMode: "unknown", + }; + } + } + + if ( + requestContext.routeGroup === "token_mcp" || + requestContext.routeGroup === "token_sse" + ) { + const resolved = resolveApiVersionMetrics("latest"); + return { + ...resolved, + apiVersionMode: "implicit_latest", + }; + } + + 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( + 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); + 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") + ? "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 { apiRequestedVersion, apiReleaseDate, apiVersion, apiVersionMode } = + resolveApiVersionLabels(request, ctx); + if (apiRequestedVersion) { + recorder.setAnalyticsContext({ + apiRequestedVersion, + }); + } + 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; + } + if (apiVersionMode) { + baseLabels.api_version_mode = apiVersionMode; + } + if (apiReleaseDate) { + baseLabels.api_release_date = apiReleaseDate; + } + + 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 { + if (!env || typeof env !== "object") { + 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; + } + + const sink = createGrafanaOtlpMetricsSink(env); + if (sink) { + GRAFANA_SINK_CACHE.set(env, sink); + } + return sink; +} + +export function createRequestMetricsRecorder( + env?: MetricsEnvLike, + sinks: ConfiguredMetricsSinks = {}, +): MetricsRecorder { + try { + const config = resolveMetricsRuntimeConfig(env); + const sink = createConfiguredMetricsSink( + config, + createDefaultConfiguredMetricsSinks(env, sinks), + ); + + 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 { + scheduleMetricsFlush(recorder, ctx.waitUntil.bind(ctx)); +} + +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 { + scheduleRequestMetricsFlush(recorder, 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..ec4b932 --- /dev/null +++ b/src/metrics/runtime/runtime-config.ts @@ -0,0 +1,149 @@ +import { CompositeMetricsSink } from "./composite-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"; +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/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/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/src/routes.ts b/src/routes.ts new file mode 100644 index 0000000..7952a2f --- /dev/null +++ b/src/routes.ts @@ -0,0 +1,38 @@ +export const PUBLIC_ROUTES = { + root: "/", + hello: "/hello", + authorize: "/authorize", + callback: "/callback", + storeToken: "/store-token", + oauthToken: "/token", + register: "/register", + mcp: "/mcp", + sse: "/sse", + bearerMcp: "/bearer/mcp", + bearerSse: "/bearer/sse", + tokenMcp: "/token/mcp", + tokenSse: "/token/sse", + openaiAppsChallenge: "/.well-known/openai-apps-challenge", +} as const; + +export const PUBLIC_ROUTE_PREFIXES = { + bearer: "/bearer", + 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.bearerMcp, + PUBLIC_ROUTES.bearerSse, + PUBLIC_ROUTES.tokenMcp, + PUBLIC_ROUTES.tokenSse, + PUBLIC_ROUTES.openaiAppsChallenge, +] 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/conversation-storage-server.ts b/src/servers/conversation-storage-server.ts new file mode 100644 index 0000000..69e7275 --- /dev/null +++ b/src/servers/conversation-storage-server.ts @@ -0,0 +1,246 @@ +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"; +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. + * + * 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 ConversationStorageServerSQLite { + private conversationId = ""; + + 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", "", ""] + this.conversationId = parts[2]; + 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 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. We never delete + * messages in the conversation, instead the next messages begin at the existing bookmark. + */ + private async initializeConversation(): Promise { + 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.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. 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[], + isDone = false, + ): Promise { + const existingIsDone = await this.state.storage.get(IS_DONE_KEY); + if (!isBoolean(existingIsDone)) { + throw new Error(`Conversation ${this.conversationId} not found`); + } + if (existingIsDone) { + throw new Error( + `Cannot append messages to conversation ${this.conversationId} marked done`, + ); + } + + let idx = (await this.state.storage.get(WRITE_BOOKMARK_KEY)) ?? 0; + const entriesToStore = {} as Record; + for (const message of newMessages) { + entriesToStore[`${MESSAGE_KEY_PREFIX}${idx}`] = message; + idx++; + } + entriesToStore[WRITE_BOOKMARK_KEY] = idx; + + if (isDone) { + 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 + * use WRITE_BOOKMARK to know up to which index new messages have been written. + */ + private async getNewMessagesAndUpdateBookmark(): Promise { + const [isDone, readBookmark, writeBookmark] = + await this.getIsDoneAndReadWriteBookmarks(); + if (!isBoolean(isDone)) { + throw new Error(`Conversation ${this.conversationId} not found`); + } + + const keys = []; + for (let i = readBookmark; i < writeBookmark; i++) { + keys.push(MESSAGE_KEY_PREFIX + i); + } + + const newMessages: Message[] = []; + const messagesMap = await this.getInBatches(keys); + for (let i = readBookmark; i < writeBookmark; i++) { + const message = messagesMap.get(MESSAGE_KEY_PREFIX + i); + if (!message) { + console.warn( + `Expected message at index ${i} for conversation ${this.conversationId} not found`, + { readBookmark, writeBookmark }, + ); + continue; + } + newMessages.push(message); + } + + await this.state.storage.put(READ_BOOKMARK_KEY, writeBookmark); + + return { + messages: newMessages, + isDone, + }; + } + + /* + * 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 { + 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, readBookmark, writeBookmark] = + await this.getIsDoneAndReadWriteBookmarks(); + if (!isBoolean(isDone) || !isDone) { + console.warn( + `Conversation ${this.conversationId} expired without being marked done`, + { + isDone, + readBookmark, + writeBookmark, + }, + ); + } + if (writeBookmark !== readBookmark) { + console.warn( + `Conversation ${this.conversationId} expired with unread messages`, + { + isDone, + readBookmark, + writeBookmark, + }, + ); + } + + // 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/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index c8c0ddf..e246630 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -1,20 +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 Span, SpanStatusCode } from "@opentelemetry/api"; import type { z } from "zod"; -import { context, 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 { 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 type { Props } from "../utils"; const ToolInputSchema = ToolSchema.shape.inputSchema; type ToolInput = z.infer; @@ -39,6 +56,8 @@ export type ToolResponse = SuccessResponse | ErrorResponse; export interface Context { props: Props; + env: Env; + ctx?: DurableObjectState; } export abstract class BaseMCPServer extends Server { @@ -172,22 +191,174 @@ export abstract class BaseMCPServer extends Server { }; } - protected getThoughtSpotService() { + protected getStorageService(): StorageServiceClient { + return new StorageServiceClient( + this.ctx.env + .CONVERSATION_STORAGE_OBJECT as unknown as DurableObjectNamespace, + ); + } + + protected getThoughtSpotService( + recorder?: MetricsRecorder, + analyticsContextOverride?: MetricAnalyticsContext, + ) { 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.mergeMetricAnalyticsContext( + analyticsContextOverride, + ), + eventIdentity: this.getMetricEventIdentity(), + }, ); } - protected async initializeService(): Promise { - this.sessionInfo = await this.getThoughtSpotService().getSessionInfo(); - const mixpanel = new MixpanelTracker( - this.sessionInfo, - this.ctx.props.clientName, + 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 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; + } + + 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, ); - this.addTracker(mixpanel); + 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(); + const mixpanel = new MixpanelTracker( + this.sessionInfo, + this.ctx.props.clientName, + ); + this.addTracker(mixpanel); + } catch (error) { + console.error("Error initializing session info:", error); + } } /** @@ -212,6 +383,7 @@ export abstract class BaseMCPServer extends Server { */ protected abstract callTool( request: z.infer, + recorder: MetricsRecorder, ): Promise; async init() { @@ -252,7 +424,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 5b5a57b..a7af53a 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 } 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,9 +39,89 @@ 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(); + 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 latest fallback:", + 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]; @@ -95,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 }); @@ -108,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: { @@ -137,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: @@ -160,6 +250,7 @@ export class MCPServer extends BaseMCPServer { @WithSpan("call-get-relevant-questions") async callGetRelevantQuestions( request: z.infer, + recorder: MetricsRecorder, ) { const { query, @@ -171,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( @@ -216,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( @@ -249,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, ); @@ -258,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( @@ -285,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( @@ -293,9 +383,12 @@ 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, ); + recorder.setAnalyticsContext({ + analyticalSessionId: response.conversation_id, + }); span?.setAttribute("analytical_session_id", response.conversation_id); // Conversation is initialized in streamingMessageStorage from callSendSessionMessage, @@ -308,30 +401,41 @@ 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); + recorder.setAnalyticsContext({ + analyticalSessionId: analytical_session_id, + }); span?.setAttributes({ analytical_session_id, 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}`, ); } - await this.getThoughtSpotService().sendAgentConversationMessageStreaming( + await this.getThoughtSpotService(recorder, { + analyticalSessionId: analytical_session_id, + }).sendAgentConversationMessageStreaming( analytical_session_id, message, - this.streamingMessageStorage, + storageService.appendMessages.bind(storageService), additional_context, ); @@ -342,7 +446,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, @@ -356,6 +463,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 +471,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; @@ -399,7 +506,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, @@ -426,12 +536,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( @@ -451,12 +558,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( @@ -484,12 +594,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/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/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/servers/version-registry.ts b/src/servers/version-registry.ts index 0be8760..1ed1c04 100644 --- a/src/servers/version-registry.ts +++ b/src/servers/version-registry.ts @@ -1,5 +1,18 @@ import { toolDefinitionsV1, toolDefinitionsV2 } from "./tool-definitions"; +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 */ @@ -13,26 +26,32 @@ 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); +} + +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 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 +60,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 +83,65 @@ 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]; +} + +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/storage-service/storage-service.ts b/src/storage-service/storage-service.ts new file mode 100644 index 0000000..805698d --- /dev/null +++ b/src/storage-service/storage-service.ts @@ -0,0 +1,95 @@ +import type { Message, StreamingMessagesState } from "../thoughtspot/types"; + +/** + * Client for the ConversationStorageServer Durable Object. + * + * 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 namespace: DurableObjectNamespace) {} + + private headers(): HeadersInit { + return { + "Content-Type": "application/json", + Accept: "application/json", + }; + } + + 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 `https://internal/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 this.stubFor(conversationId).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 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(); + 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 this.stubFor(conversationId).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/src/streaming-utils.ts b/src/streaming-utils.ts index 8efb208..d51fc6b 100644 --- a/src/streaming-utils.ts +++ b/src/streaming-utils.ts @@ -1,7 +1,11 @@ -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"; +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 @@ -11,12 +15,21 @@ 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, + 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; @@ -33,11 +46,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; } @@ -72,20 +81,41 @@ 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", text: item.content, }); } 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", text: item.content, }); } 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", type: "answer", answer_id: JSON.stringify({ session_id: item.metadata?.session_id, @@ -106,13 +136,19 @@ 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, message: item, }); newMessages.push({ + is_thinking: false, type: "text", text: item.display_message || "Something went wrong", }); @@ -125,10 +161,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..04f4622 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -1,24 +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 type { StreamingMessagesStorageWithTtl } from "../streaming-message-storage-with-ttl/streaming-message-storage-with-ttl"; -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( @@ -50,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, @@ -113,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) => ({ @@ -184,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 @@ -224,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({ @@ -251,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, @@ -281,7 +358,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()); @@ -295,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, - streamingMessageStorage, - (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, @@ -360,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({ @@ -373,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), @@ -519,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({ @@ -541,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 @@ -574,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 @@ -607,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") @@ -642,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/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/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 b64276a..a0af829 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; @@ -177,6 +182,8 @@ describe("Bearer Handler", () => { accessToken: "my-access-token", instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Custom Test Client", + apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", }); // Verify the response @@ -201,6 +208,8 @@ describe("Bearer Handler", () => { accessToken: "my-access-token", instanceUrl: "https://my-instance.thoughtspot.cloud", clientName: "Bearer Token client", + apiVersion: "backwards-compatibility-default", + apiVersionMode: "implicit_legacy", }); // Verify the response @@ -210,6 +219,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); @@ -479,8 +514,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 +529,17 @@ 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", + apiRequestedVersion: "beta", }); - expect(mockCtx.props.apiVersion).toBeUndefined(); + expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); + expect(mockCtx.props.apiVersionMode).toBe("implicit_legacy"); }); - 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 +553,14 @@ 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", + apiRequestedVersion: "beta", }); - expect(mockCtx.props.apiVersion).toBeUndefined(); + expect(mockCtx.props.apiVersion).toBe("backwards-compatibility-default"); + expect(mockCtx.props.apiVersionMode).toBe("implicit_legacy"); }); }); @@ -544,7 +583,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", apiVersion: "beta", + apiVersionMode: "beta", }); }); @@ -566,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", { @@ -581,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 () => { @@ -607,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", }); }); @@ -629,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", }); }); @@ -652,7 +700,9 @@ describe("Bearer Handler", () => { expect(mockCtx.props).toMatchObject({ accessToken: "test-token", instanceUrl: "https://test.thoughtspot.cloud", + apiRequestedVersion: "beta", apiVersion: "beta", + apiVersionMode: "beta", }); }); @@ -706,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/handlers.spec.ts b/test/handlers.spec.ts index da9b144..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; @@ -105,13 +110,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", @@ -139,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/mixpanel/integration.spec.ts b/test/metrics/mixpanel/integration.spec.ts index d7503ff..d3810f8 100644 --- a/test/metrics/mixpanel/integration.spec.ts +++ b/test/metrics/mixpanel/integration.spec.ts @@ -1,10 +1,7 @@ -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"; -// 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/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/analytics-engine-sink.spec.ts b/test/metrics/runtime/analytics-engine-sink.spec.ts new file mode 100644 index 0000000..bf1f303 --- /dev/null +++ b/test/metrics/runtime/analytics-engine-sink.spec.ts @@ -0,0 +1,236 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +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, + ANALYTICS_ENGINE_SCHEMA_VERSION, + AnalyticsEngineMetricsSink, + createAnalyticsEngineMetricsSink, + getAnalyticsEngineMetricFamily, + toAnalyticsEngineDataPoint, +} from "../../../src/metrics/runtime/analytics-engine-sink"; +import { 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(["shared"]); + expect(dataPoint.blobs).toEqual([ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + "counter", + null, + null, + "mcp", + "oauth", + null, + null, + null, + "success", + null, + "create_liveboard", + null, + null, + null, + null, + null, + "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 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([ + ["service_version", "service.version"], + ]); + expect(ANALYTICS_ENGINE_BLOB_FIELDS).toEqual([ + "schema_version", + "event_family", + "metric_name", + "metric_kind", + "tenant_id", + "user_id", + ...ANALYTICS_ENGINE_LABEL_FIELDS, + "api_requested_version", + "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 () => { + 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: ["shared"], + doubles: [2, 1_714_000_000_000], + }), + ); + }); + + it("maps request-scoped event identity into the Analytics Engine index and blobs", 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: ["tenant-123"], + blobs: expect.arrayContaining(["tenant-123", "user-456"]), + }), + ); + }); + + 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", + analyticalSessionId: "conv-123", + }, + }); + + expect(dataset.writeDataPoint).toHaveBeenCalledWith( + expect.objectContaining({ + blobs: expect.arrayContaining(["2026-10-01", "conv-123"]), + }), + ); + }); + + 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?.blobs?.[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/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/metric-context.spec.ts b/test/metrics/runtime/metric-context.spec.ts new file mode 100644 index 0000000..d7e69b9 --- /dev/null +++ b/test/metrics/runtime/metric-context.spec.ts @@ -0,0 +1,87 @@ +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, +} 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("/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("/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.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..313da80 --- /dev/null +++ b/test/metrics/runtime/metrics-recorder.spec.ts @@ -0,0 +1,257 @@ +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" }, + analyticsContext: undefined, + eventIdentity: undefined, + }); + }); + + it("returns the same in-flight flush promise when called repeatedly", async () => { + let resolveFlush!: () => void; + const flushSpy = vi.fn().mockResolvedValue(undefined); + flushSpy.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFlush = resolve; + }), + ); + const recorder = new RequestMetricsRecorder({ + sink: { flush: flushSpy }, + }); + + recorder.count(METRIC_NAMES.httpRequestsTotal); + const firstFlushPromise = recorder.flush(); + const secondFlushPromise = recorder.flush(); + + expect(secondFlushPromise).toBe(firstFlushPromise); + expect(flushSpy).toHaveBeenCalledTimes(1); + + resolveFlush(); + await firstFlushPromise; + 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: {}, + analyticsContext: undefined, + eventIdentity: undefined, + }); + }); + + 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}`, + ); + }); + + 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", + analyticalSessionId: "conv-123", + }); + recorder.count(METRIC_NAMES.httpRequestsTotal); + await recorder.flush(); + + expect(flushSpy).toHaveBeenCalledWith( + 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 new file mode 100644 index 0000000..225cb75 --- /dev/null +++ b/test/metrics/runtime/request-metrics.spec.ts @@ -0,0 +1,590 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { METRIC_NAMES } from "../../../src/metrics/runtime/metric-types"; +import { + clearMetricsRecorderFromExecutionContext, + createRequestMetricsRecorder, + getMetricsRecorderFromExecutionContext, + recordBearerAuthRequestMetric, + recordHttpRequestMetrics, + recordStatusMetric, + resolveApiVersionLabels, + resolveCanonicalApiVersionLabel, + setMetricsRecorderOnExecutionContext, + withRequestMetrics, +} 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) }; + 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("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) }; + 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("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( + { + 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("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("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("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: ["shared"], + blobs: expect.arrayContaining([METRIC_NAMES.httpRequestsTotal]), + }), + ); + }); + + 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 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("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(); + + expect(setMetricsRecorderOnExecutionContext(ctx, recorder)).toBe(recorder); + expect(getMetricsRecorderFromExecutionContext(ctx)).toBe(recorder); + + clearMetricsRecorderFromExecutionContext(ctx); + + 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", + api_version_mode: "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", + api_version_mode: "beta", + auth_mode: "oauth", + outcome: "success", + route_group: "mcp", + transport: "mcp", + }, + }), + ]); + }); + + 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("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( + "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", () => { + 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( + "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( + "https://example.com/token/mcp?api-version=garbage", + ); + + 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(); + + 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", + }, + }), + ]); + }); + + 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/metrics/runtime/runtime-config.spec.ts b/test/metrics/runtime/runtime-config.spec.ts new file mode 100644 index 0000000..308968b --- /dev/null +++ b/test/metrics/runtime/runtime-config.spec.ts @@ -0,0 +1,131 @@ +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..5ad8e00 --- /dev/null +++ b/test/metrics/runtime/sinks.spec.ts @@ -0,0 +1,44 @@ +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(); + }); +}); 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"; 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/conversation-storage-server.spec.ts b/test/servers/conversation-storage-server.spec.ts new file mode 100644 index 0000000..effe5c6 --- /dev/null +++ b/test/servers/conversation-storage-server.spec.ts @@ -0,0 +1,536 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { ConversationStorageServerSQLite } 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 ( + 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); + } + }), + setAlarm: vi.fn(async (scheduledTime: number): Promise => { + alarm = scheduledTime; + }), + deleteAlarm: vi.fn(async (): Promise => { + alarm = null; + }), + deleteAll: vi.fn(async (): Promise => { + store.clear(); + }), + }, + }; +} + +function createServer(mock: ReturnType) { + const state = { storage: mock.storage } as unknown as DurableObjectState; + return new ConversationStorageServerSQLite(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", + 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, +}; + +// 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 +// --------------------------------------------------------------------------- + +describe("ConversationStorageServerSQLite", () => { + let mock: ReturnType; + let server: ConversationStorageServerSQLite; + + 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")); + + 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")); + + // 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 () => { + 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); + + expect(mock.store.get("is-done")).toBe(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] }), + ); + + expect(mock.store.get("message-0")).toEqual(textMessage); + expect(mock.store.get("write-bookmark")).toBe(1); + }); + + it("accumulates messages across multiple calls", async () => { + await server.fetch( + makeRequest("POST", "append", { messages: [textMessage] }), + ); + await server.fetch( + makeRequest("POST", "append", { + messages: [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 () => { + await server.fetch( + makeRequest("POST", "append", { + messages: [textMessage], + isDone: true, + }), + ); + + expect(mock.store.get("is-done")).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); + }); + + // ------------------------------------------------------------------- + // 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]); + } + }); + }); + + // ------------------------------------------------------------------------- + // 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); + }); + + // ------------------------------------------------------------------- + // 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); + }); + }); + + // ------------------------------------------------------------------------- + // 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("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 () => { + 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/test/servers/mcp-server-base.spec.ts b/test/servers/mcp-server-base.spec.ts index d7c5fea..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,80 +260,165 @@ 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("Server Initialization", () => { - it("should initialize with custom server name and version", () => { - const customServer = new TestMCPServer( - { props: mockProps }, - "CustomServer", - "2.0.0", + 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 + .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, env: mockEnv }, + mockStreamingStorage, + ); + await expect(testServer.init()).resolves.not.toThrow(); + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Error initializing session info:", + expect.any(Error), ); + // sessionInfo should remain unset, no tracker registered + expect(testServer.getSessionInfo()).toBeUndefined(); + expect(testServer.getTrackers().size).toBe(0); - expect(customServer).toBeDefined(); + consoleErrorSpy.mockRestore(); }); + }); + + 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); + + 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(); + }); + }); + + 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); - it("should initialize with default server name and version", () => { + 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 f1e8ac5..55d9d63 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 { 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"; 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", () => ({ @@ -178,14 +180,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 +197,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 +253,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(); }); }); @@ -309,6 +301,83 @@ 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.blobs?.[2] === METRIC_NAMES.toolCallsTotal, + ); + const toolDuration = dataPoints.find( + (dataPoint) => dataPoint.blobs?.[2] === METRIC_NAMES.toolDurationMs, + ); + + expect(toolCallCounter).toEqual( + expect.objectContaining({ + indexes: ["test-org"], + blobs: expect.arrayContaining([ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolCallsTotal, + "test-org", + "test-user-123", + "ping", + "success", + "backwards-compatibility-default", + "pinned", + "2025-01-01", + "2025-03-01", + ]), + }), + ); + expect(toolDuration).toEqual( + expect.objectContaining({ + indexes: ["test-org"], + blobs: expect.arrayContaining([ + ANALYTICS_ENGINE_SCHEMA_VERSION, + "tool", + METRIC_NAMES.toolDurationMs, + "test-org", + "test-user-123", + ]), + }), + ); + }); }); describe("Check Connectivity Tool", () => { @@ -1254,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/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/test/servers/version-registry.spec.ts b/test/servers/version-registry.spec.ts index f370b07..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 @@ -16,14 +19,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 +33,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 +54,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 +66,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 +95,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,12 +112,42 @@ 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); }); }); + 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) => @@ -170,9 +206,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 +220,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 +241,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 +273,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 +290,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"); }); }); }); diff --git a/test/storage-service/storage-service.spec.ts b/test/storage-service/storage-service.spec.ts new file mode 100644 index 0000000..67194c4 --- /dev/null +++ b/test/storage-service/storage-service.spec.ts @@ -0,0 +1,248 @@ +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 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", +}; + +// 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 lastRequest(): Request { + if (!lastStubRequest) throw new Error("No stub request recorded"); + return lastStubRequest; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("StorageServiceClient", () => { + let client: StorageServiceClient; + let namespaceMock: DurableObjectNamespace; + + beforeEach(() => { + vi.restoreAllMocks(); + namespaceMock = makeNamespaceMock(); + client = new StorageServiceClient(namespaceMock); + }); + + // ------------------------------------------------------------------------- + // initializeConversation + // ------------------------------------------------------------------------- + + describe("initializeConversation", () => { + it("sends POST to /storage//initialize", async () => { + await client.initializeConversation(CONVERSATION_ID); + + 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 () => { + await client.initializeConversation("conv with spaces/and-slash"); + + 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 () => { + await expect( + client.initializeConversation(CONVERSATION_ID), + ).resolves.toBeUndefined(); + }); + + it("throws when the server returns a non-ok status", async () => { + namespaceMock = makeNamespaceMock("Something went wrong", 500); + client = new StorageServiceClient(namespaceMock); + + await expect( + client.initializeConversation(CONVERSATION_ID), + ).rejects.toThrow("Failed to initialize conversation (500)"); + }); + + it("includes the error body in the thrown error message", async () => { + namespaceMock = makeNamespaceMock( + "Conversation already exists and is not marked done", + 400, + ); + client = new StorageServiceClient(namespaceMock); + + 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 () => { + await client.appendMessages(CONVERSATION_ID, [textMessage]); + + 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 () => { + await client.appendMessages(CONVERSATION_ID, [textMessage, chunkMessage]); + + 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 () => { + await client.appendMessages(CONVERSATION_ID, [answerMessage], true); + + const body = (await lastRequest().json()) as StreamingMessagesState; + expect(body.isDone).toBe(true); + }); + + it("sends Content-Type: application/json", async () => { + await client.appendMessages(CONVERSATION_ID, []); + + expect(lastRequest().headers.get("Content-Type")).toBe( + "application/json", + ); + }); + + it("resolves without error on a 200 response", async () => { + await expect( + client.appendMessages(CONVERSATION_ID, [textMessage]), + ).resolves.toBeUndefined(); + }); + + it("throws when the server returns a non-ok status", async () => { + namespaceMock = makeNamespaceMock("Conversation not found", 500); + client = new StorageServiceClient(namespaceMock); + + 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 () => { + namespaceMock = makeNamespaceMock( + "Cannot append messages to a conversation marked done", + 400, + ); + client = new StorageServiceClient(namespaceMock); + + 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 () => { + namespaceMock = makeNamespaceMock({ messages: [textMessage], isDone: false }); + client = new StorageServiceClient(namespaceMock); + + await client.getNewMessages(CONVERSATION_ID); + + const req = lastRequest(); + expect(req.url).toBe( + `https://internal/storage/${CONVERSATION_ID}/messages`, + ); + expect(req.method).toBe("GET"); + }); + + it("returns the parsed StreamingMessagesState", async () => { + const state: StreamingMessagesState = { + messages: [textMessage, answerMessage], + isDone: true, + }; + namespaceMock = makeNamespaceMock(state); + client = new StorageServiceClient(namespaceMock); + + const result = await client.getNewMessages(CONVERSATION_ID); + + expect(result).toEqual(state); + }); + + it("returns an empty messages array when there are no new messages", async () => { + namespaceMock = makeNamespaceMock({ messages: [], isDone: false }); + client = new StorageServiceClient(namespaceMock); + + 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 () => { + namespaceMock = makeNamespaceMock("Conversation not found", 404); + client = new StorageServiceClient(namespaceMock); + + await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( + "Failed to get messages (404)", + ); + }); + + it("includes the error body in the thrown error message", async () => { + namespaceMock = makeNamespaceMock("Internal error", 500); + client = new StorageServiceClient(namespaceMock); + + await expect(client.getNewMessages(CONVERSATION_ID)).rejects.toThrow( + "Internal error", + ); + }); + }); +}); 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..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 @@ -19,8 +46,10 @@ function makeReader(chunks: string[]): ReadableStreamDefaultReader { // Mock storage function makeMockStorage() { + const fn = vi.fn(async () => {}); return { - appendMessagesAndRestartTtl: vi.fn(async () => {}), + appendMessages: fn, + appendMessagesAndRestartTtl: fn, }; } @@ -33,6 +62,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { beforeEach(() => { vi.clearAllMocks(); + tracingState.span = undefined; consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); }); @@ -48,7 +78,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -62,18 +92,18 @@ 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( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); 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,20 +113,82 @@ 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`; + const reader = makeReader([line]); + + await processSendAgentConversationMessageStreamingResponse( + CONV_ID, + reader, + storage.appendMessages, + 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( + CONV_ID, + reader, + storage.appendMessages, + INSTANCE_URL, + ); + + expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ + { 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, + storage.appendMessages, INSTANCE_URL, ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text_chunk", text: "chunk content" }, + { is_thinking: true, type: "text_chunk", text: "thinking chunk" }, ]); }); @@ -116,13 +208,14 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); 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 +225,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.appendMessages, + 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`; @@ -140,12 +267,12 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "Something went wrong" }, + { is_thinking: false, type: "text", text: "Something went wrong" }, ]); }); @@ -157,15 +284,35 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "Something went wrong" }, + { is_thinking: false, type: "text", text: "Something went wrong" }, ]); }); + 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 = [ @@ -182,7 +329,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -198,18 +345,18 @@ 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( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); expect(storage.appendMessagesAndRestartTtl).toHaveBeenCalledWith(CONV_ID, [ - { type: "text", text: "hi" }, + { is_thinking: false, type: "text", text: "hi" }, ]); }); @@ -221,7 +368,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -241,7 +388,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -259,7 +406,7 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); @@ -273,7 +420,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]); @@ -281,37 +428,37 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); 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( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); 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 +471,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]); @@ -333,13 +480,13 @@ describe("processSendAgentConversationMessageStreamingResponse", () => { await processSendAgentConversationMessageStreamingResponse( CONV_ID, reader, - storage as any, + storage.appendMessages, INSTANCE_URL, ); 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..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,13 +162,11 @@ 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 - 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); diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index b9d3db7..e16fea8 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,140 @@ 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", + analyticalSessionId: "conv-123", + }); + 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", + analyticalSessionId: "conv-123", + }, + 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.blobs?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("latest") && + dataPoint.blobs?.includes("conv-123") && + dataPoint.indexes?.[0] === "org-123" && + dataPoint.blobs?.includes("user-123"), + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ), + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamMessagesTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("text") && + dataPoint.blobs?.includes("conv-123"), + ), + ).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 +380,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.blobs?.[2] === METRIC_NAMES.upstreamCallsTotal && + dataPoint.blobs?.includes( + "send_agent_conversation_message_streaming", + ) && + dataPoint.blobs?.includes("upstream_error"), + ), + ).toBe(true); + expect( + dataPoints.some( + (dataPoint) => + dataPoint.blobs?.[2] === METRIC_NAMES.upstreamStreamsStartedTotal, + ), + ).toBe(false); }); it("should throw when sending the streaming message fails", async () => { @@ -362,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, }, }, diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 4e97ea3..33671c6 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -4,12 +4,12 @@ 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 bf4897f..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": [ @@ -18,8 +18,8 @@ "name": "MCP_OBJECT" }, { - "class_name": "ThoughtSpotOpenAIDeepResearchMCP", - "name": "OPENAI_DEEP_RESEARCH_MCP_OBJECT" + "class_name": "ConversationStorageServerSQLite", + "name": "CONVERSATION_STORAGE_OBJECT" } ] }, @@ -35,6 +35,30 @@ "new_sqlite_classes": [ "ThoughtSpotOpenAIDeepResearchMCP" ] + }, + { + "tag": "v4", + "new_classes": [ + "ConversationStorageServer" + ] + }, + { + "tag": "v5", + "deleted_classes": [ + "ThoughtSpotOpenAIDeepResearchMCP" + ] + }, + { + "tag": "v6", + "deleted_classes": [ + "ConversationStorageServer" + ] + }, + { + "tag": "v7", + "new_sqlite_classes": [ + "ConversationStorageServerSQLite" + ] } ], "kv_namespaces": [{