From 2d0ace4ce98f459e392bb23257a9f037aa7d00aa Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Thu, 24 Jul 2025 18:35:38 -0400 Subject: [PATCH 01/21] Add png image formation to MCP --- src/handlers.ts | 33 ++++++++- src/index.ts | 4 +- src/servers/api-server.ts | 8 ++- src/servers/mcp-server-base.ts | 1 + src/servers/mcp-server.ts | 21 +++++- src/stdio.ts | 7 +- src/thoughtspot/thoughtspot-service.ts | 13 +++- src/utils.ts | 58 +++++++++++++++- worker-configuration.d.ts | 96 ++++++++------------------ wrangler.jsonc | 3 +- 10 files changed, 169 insertions(+), 75 deletions(-) diff --git a/src/handlers.ts b/src/handlers.ts index cc8f4d0..331c705 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -1,13 +1,15 @@ import type { AuthRequest, OAuthHelpers } from '@cloudflare/workers-oauth-provider' import { Hono } from 'hono' import type { Props } from './utils'; -import { McpServerError } from './utils'; +import { getFromKV, McpServerError } from './utils'; import { 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 { ThoughtSpotService } from './thoughtspot/thoughtspot-service'; +import { getThoughtSpotClient } from './thoughtspot/thoughtspot-client'; const app = new Hono<{ Bindings: Env & { OAUTH_PROVIDER: OAuthHelpers } }>() @@ -186,6 +188,7 @@ class Handler { const handler = new Handler(); app.get("/", async (c) => { + console.log("[DEBUG] Serving index", c.env); const response = await handler.serveIndex(c.env); return response; }); @@ -263,4 +266,32 @@ app.post("/store-token", async (c) => { } }); +app.get("/data/img", async (c) => { + const token = c.req.query("token") || ''; + if (token === '') { + return c.json({ error: "Token not found" }, 404); + } + console.log("[DEBUG] Token", token); + + const tokenData = await getFromKV(token, c.env); + if (!tokenData) { + return c.json({ error: "Token not found" }, 404); + } + console.log("[DEBUG] Token found", token); + + // Extract values from token data + const sessionId = (tokenData as any).sessionId; + const generationNo = (tokenData as any).GenNo || (tokenData as any).generationNo; // Handle both field names + const instanceURL = (tokenData as any).instanceURL; + const accessToken = (tokenData as any).accessToken; + + console.log("[DEBUG] Session ID", sessionId); + console.log("[DEBUG] Generation No", generationNo); + console.log("[DEBUG] Instance URL", instanceURL); + + const thoughtSpotService = new ThoughtSpotService(getThoughtSpotClient(instanceURL, accessToken)); + const image = await thoughtSpotService.getAnswerImagePNG(sessionId, generationNo); + return c.body(await image.arrayBuffer()); +}); + export default app; \ No newline at end of file diff --git a/src/index.ts b/src/index.ts index 7e3d5cc..03d2b2a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -14,9 +14,9 @@ const config: ResolveConfigFn = (env: Env, _trigger) => { return { exporter: { url: 'https://api.honeycomb.io/v1/traces', - headers: { 'x-honeycomb-team': process.env.HONEYCOMB_API_KEY }, + headers: { 'x-honeycomb-team': env.HONEYCOMB_API_KEY }, }, - service: { name: process.env.HONEYCOMB_DATASET } + service: { name: env.HONEYCOMB_DATASET } }; }; diff --git a/src/servers/api-server.ts b/src/servers/api-server.ts index 563acf1..7a3b1ee 100644 --- a/src/servers/api-server.ts +++ b/src/servers/api-server.ts @@ -1,6 +1,6 @@ import { Hono } from 'hono' import type { Props } from '../utils'; -import { McpServerError } from '../utils'; +import { getFromKV, McpServerError } from '../utils'; import { getDataSources, ThoughtSpotService } from '../thoughtspot/thoughtspot-service'; import { getThoughtSpotClient } from '../thoughtspot/thoughtspot-client'; import { getActiveSpan, WithSpan } from '../metrics/tracing/tracing-utils'; @@ -47,6 +47,12 @@ class ApiHandler { return await service.getDataSources(); } + @WithSpan('api-get-answer-image') + async getAnswerImage(props: Props, sessionId: string, generationNo: number) { + const service = this.getThoughtSpotService(props); + return await service.getAnswerImage(sessionId, generationNo); + } + @WithSpan('api-proxy-post') async proxyPost(props: Props, path: string, body: any) { const span = getActiveSpan(); diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 2105432..78ec3be 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -39,6 +39,7 @@ export type ToolResponse = SuccessResponse | ErrorResponse; export interface Context { props: Props; + env?: Env; // Add access to Cloudflare Workers environment } export abstract class BaseMCPServer extends Server { diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 3a39b30..8259c1e 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -6,7 +6,7 @@ import { import { z } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; -import { McpServerError } from "../utils"; +import { McpServerError, putInKV } from "../utils"; import type { DataSource } from "../thoughtspot/thoughtspot-service"; @@ -224,6 +224,21 @@ export class MCPServer extends BaseMCPServer { const { question, datasourceId: sourceId } = GetAnswerSchema.parse(request.params.arguments); const answer = await this.getThoughtSpotService().getAnswerForQuestion(question, sourceId, false); + let tokenUrl = ""; + + // Generate token and store in KV store + if (!answer.error && this.ctx.env?.OAUTH_KV) { + console.log("[DEBUG] Storing token in KV"); + const token = crypto.randomUUID(); + const tokenData = { + sessionId: answer.session_identifier, + generationNo: answer.generation_number, + instanceURL: this.ctx.props.instanceUrl, + accessToken: this.ctx.props.accessToken + }; + await putInKV(token, tokenData, this.ctx.env); + tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; + } if (answer.error) { return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`); @@ -235,6 +250,10 @@ export class MCPServer extends BaseMCPServer { type: "text" as const, text: `Question: ${question}\nSession Identifier: ${answer.session_identifier}\nGeneration Number: ${answer.generation_number}\n\nUse this information to create a liveboard with the createLiveboard tool, if the user asks.`, }, + ...(tokenUrl ? [{ + type: "text" as const, + text: `URL to download the answer's PNG image visualization: ${tokenUrl}. Use this URL to download the visualization associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data.`, + }] : []), ]; return this.createMultiContentSuccessResponse(content, "Answer found"); } diff --git a/src/stdio.ts b/src/stdio.ts index 3301238..bc4c2ce 100755 --- a/src/stdio.ts +++ b/src/stdio.ts @@ -17,9 +17,14 @@ async function main() { const props: Props = { instanceUrl: validateAndSanitizeUrl(instanceUrl), accessToken, + clientName: { + clientId: "stdio-client", + clientName: "Stdio Client", + registrationDate: Date.now() + } }; - const server = new MCPServer({ props }); + const server = new MCPServer({ props, env: undefined }); await server.init(); const transport = new StdioServerTransport(); diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index 45e5ccf..b088275 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -1,4 +1,4 @@ -import type { ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk"; +import type { HttpFile, ThoughtSpotRestApi } from "@thoughtspot/rest-api-sdk"; import { SpanStatusCode, trace, context } from "@opentelemetry/api"; import { getActiveSpan, WithSpan } from "../metrics/tracing/tracing-utils"; import type { DataSource, SessionInfo } from "./types"; @@ -283,6 +283,17 @@ export class ThoughtSpotService { return liveboardUrl; } + async getAnswerImagePNG(sessionId: string, GenNo: number): Promise { + const span = getActiveSpan(); + span?.addEvent("get-answer-image"); + const data = await this.client.exportAnswerReport({ + session_identifier: sessionId, + generation_number: GenNo, + file_format: "PNG", + }) + return data; + } + /** * Get data sources */ diff --git a/src/utils.ts b/src/utils.ts index 230fd70..7ed6c12 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -92,7 +92,38 @@ export class McpServerError extends Error { export function instrumentedMCPServer(MCPServer: new (ctx: Context) => T, config: ResolveConfigFn) { const Agent = class extends McpAgent { - server = new MCPServer(this); + private _server: T | undefined; + private _env: Env; + + // Lazy getter for server that creates it when first accessed + // + // WHY THIS APPROACH: + // Originally we passed 'this' directly as Context: `server = new MCPServer(this)` + // This worked when Context was just { props: Props }, but broke when we added env. + // + // PROBLEMS WITH ORIGINAL APPROACH: + // 1. McpAgent's 'env' property is protected, but Context expects public + // 2. TypeScript error: "Property 'env' is protected but public in Context" + // + // WHY NOT CONSTRUCTOR CREATION: + // We tried creating server in constructor: `new MCPServer({ props: this.props, env })` + // But this.props is undefined during constructor - it gets set later by McpAgent._init() + // Runtime error: "Cannot read properties of undefined (reading 'instanceUrl')" + // + // SOLUTION - LAZY INITIALIZATION: + // - Store env from constructor (available immediately) + // - Create server only when first accessed (after props are set by McpAgent lifecycle) + // - Combine both props and env into proper Context object + get server(): T { + if (!this._server) { + const context: Context = { + props: this.props, // Available after McpAgent._init() sets it + env: this._env // Stored from constructor + }; + this._server = new MCPServer(context); + } + return this._server; + } // Argument of type 'typeof ThoughtSpotMCPWrapper' is not assignable to parameter of type 'DOClass'. // Cannot assign a 'protected' constructor type to a 'public' constructor type. @@ -100,12 +131,37 @@ export function instrumentedMCPServer(MCPServer: new (c // biome-ignore lint/complexity/noUselessConstructor: required for DOClass public constructor(state: DurableObjectState, env: Env) { super(state, env); + // Store env for later use - props aren't available yet in constructor + // McpAgent lifecycle: constructor → _init(props) → init() + this._env = env; } async init() { + // Access the server property to trigger lazy initialization + // At this point, props have been set by McpAgent._init() await this.server.init(); } } return instrumentDO(Agent, config); +} + +export async function putInKV(key: string, value: any, env: Env) { + if (env?.OAUTH_KV) { + await env.OAUTH_KV.put(key, JSON.stringify(value), { + expirationTtl: 60 * 60 * 3 // 3 hours + }); + } +} + +export async function getFromKV(key: string, env: Env) { + console.log("[DEBUG] Getting from KV", key); + if (env?.OAUTH_KV) { + const value = await env.OAUTH_KV.get(key, { type: "json" }); + if (value) { + console.log("[DEBUG] Value", value); + return value; + } + return null; + } } \ No newline at end of file diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 2307189..e5a5834 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,11 +1,12 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: b5949bb168b75f2fe6e6e5a1db8d9e89) -// Runtime types generated with workerd@1.20250604.0 2025-04-17 nodejs_compat +// Generated by Wrangler by running `wrangler types` (hash: 07dddea59586af0824a29829b938dd68) +// Runtime types generated with workerd@1.20250417.0 2025-04-17 nodejs_compat declare namespace Cloudflare { interface Env { OAUTH_KV: KVNamespace; - HONEYCOMB_DATASET: ""; - HONEYCOMB_API_KEY: ""; + HONEYCOMB_API_KEY: string; + HONEYCOMB_DATASET: string; + HOST_NAME: string; MCP_OBJECT: DurableObjectNamespace; OPENAI_DEEP_RESEARCH_MCP_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; @@ -17,7 +18,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } // Begin runtime types @@ -100,7 +101,7 @@ interface Console { clear(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/count_static) */ count(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countReset_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countreset_static) */ countReset(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/debug_static) */ debug(...data: any[]): void; @@ -112,9 +113,9 @@ interface Console { error(...data: any[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/group_static) */ group(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupCollapsed_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupcollapsed_static) */ groupCollapsed(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupEnd_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupend_static) */ groupEnd(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/info_static) */ info(...data: any[]): void; @@ -124,9 +125,9 @@ interface Console { table(tabularData?: any, properties?: string[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/time_static) */ time(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeEnd_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeend_static) */ timeEnd(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeLog_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timelog_static) */ timeLog(label?: string, ...data: any[]): void; timeStamp(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/trace_static) */ @@ -301,25 +302,25 @@ declare function dispatchEvent(event: WorkerGlobalScopeEventMap[keyof WorkerGlob declare function btoa(data: string): string; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/atob) */ declare function atob(data: string): string; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ declare function setTimeout(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ declare function setTimeout(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearTimeout) */ declare function clearTimeout(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ declare function setInterval(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ declare function setInterval(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearInterval) */ declare function clearInterval(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/queueMicrotask) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/queueMicrotask) */ declare function queueMicrotask(task: Function): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/structuredClone) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/structuredClone) */ declare function structuredClone(value: T, options?: StructuredSerializeOptions): T; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/reportError) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/reportError) */ declare function reportError(error: any): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/fetch) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/fetch) */ declare function fetch(input: RequestInfo | URL, init?: RequestInit): Promise; declare const self: ServiceWorkerGlobalScope; /** @@ -798,7 +799,6 @@ declare class Blob { slice(start?: number, end?: number, type?: string): Blob; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/arrayBuffer) */ arrayBuffer(): Promise; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/text) */ text(): Promise; @@ -1093,15 +1093,10 @@ interface TextEncoderEncodeIntoResult { */ declare class ErrorEvent extends Event { constructor(type: string, init?: ErrorEventErrorEventInit); - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/filename) */ get filename(): string; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/message) */ get message(): string; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/lineno) */ get lineno(): number; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/colno) */ get colno(): number; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/error) */ get error(): any; } interface ErrorEventErrorEventInit { @@ -1275,7 +1270,6 @@ declare abstract class Body { get bodyUsed(): boolean; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/arrayBuffer) */ arrayBuffer(): Promise; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/text) */ text(): Promise; @@ -1387,11 +1381,7 @@ interface Request> e * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/integrity) */ integrity: string; - /** - * Returns a boolean indicating whether or not request can outlive the global in which it was created. - * - * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/keepalive) - */ + /* Returns a boolean indicating whether or not request can outlive the global in which it was created. */ keepalive: boolean; /** * Returns the cache mode associated with request, which is a string indicating how the request will interact with the browser's cache when fetching. @@ -1591,7 +1581,6 @@ interface R2ObjectBody extends R2Object { get body(): ReadableStream; get bodyUsed(): boolean; arrayBuffer(): Promise; - bytes(): Promise; text(): Promise; json(): Promise; blob(): Promise; @@ -3837,18 +3826,8 @@ interface AutoRAGNotFoundError extends Error { } interface AutoRAGUnauthorizedError extends Error { } -type ComparisonFilter = { - key: string; - type: 'eq' | 'ne' | 'gt' | 'gte' | 'lt' | 'lte'; - value: string | number | boolean; -}; -type CompoundFilter = { - type: 'and' | 'or'; - filters: ComparisonFilter[]; -}; type AutoRagSearchRequest = { query: string; - filters?: CompoundFilter | ComparisonFilter; max_num_results?: number; ranking_options?: { ranker?: string; @@ -5198,7 +5177,6 @@ declare module 'cloudflare:workers' { export type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; export type WorkflowDelayDuration = WorkflowSleepDuration; export type WorkflowTimeoutDuration = WorkflowSleepDuration; - export type WorkflowRetentionDuration = WorkflowSleepDuration; export type WorkflowBackoff = 'constant' | 'linear' | 'exponential'; export type WorkflowStepConfig = { retries?: { @@ -5329,7 +5307,6 @@ declare namespace TailStream { readonly type: "onset"; readonly dispatchNamespace?: string; readonly entrypoint?: string; - readonly executionModel: string; readonly scriptName?: string; readonly scriptTags?: string[]; readonly scriptVersion?: ScriptVersion; @@ -5347,8 +5324,8 @@ declare namespace TailStream { } interface SpanOpen { readonly type: "spanOpen"; - readonly name: string; - readonly info?: FetchEventInfo | JsRpcEventInfo | Attributes; + readonly op?: string; + readonly info?: FetchEventInfo | JsRpcEventInfo | Attribute[]; } interface SpanClose { readonly type: "spanClose"; @@ -5372,7 +5349,7 @@ declare namespace TailStream { } interface Return { readonly type: "return"; - readonly info?: FetchResponseInfo; + readonly info?: FetchResponseInfo | Attribute[]; } interface Link { readonly type: "link"; @@ -5382,23 +5359,21 @@ declare namespace TailStream { readonly spanId: string; } interface Attribute { + readonly type: "attribute"; readonly name: string; - readonly value: string | string[] | boolean | boolean[] | number | number[] | bigint | bigint[]; - } - interface Attributes { - readonly type: "attributes"; - readonly info: Attribute[]; + readonly value: string | string[] | boolean | boolean[] | number | number[]; } + type Mark = DiagnosticChannelEvent | Exception | Log | Return | Link | Attribute[]; interface TailEvent { readonly traceId: string; readonly invocationId: string; readonly spanId: string; readonly timestamp: Date; readonly sequence: number; - readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | DiagnosticChannelEvent | Exception | Log | Return | Link | Attributes; + readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | Mark; } type TailEventHandler = (event: TailEvent) => void | Promise; - type TailEventHandlerName = "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attributes"; + type TailEventHandlerName = "onset" | "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attribute"; type TailEventHandlerObject = Record; type TailEventHandlerType = TailEventHandler | TailEventHandlerObject; } @@ -5712,9 +5687,6 @@ declare abstract class Workflow { */ public createBatch(batch: WorkflowInstanceCreateOptions[]): Promise; } -type WorkflowDurationLabel = 'second' | 'minute' | 'hour' | 'day' | 'week' | 'month' | 'year'; -type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; -type WorkflowRetentionDuration = WorkflowSleepDuration; interface WorkflowInstanceCreateOptions { /** * An id for your Workflow instance. Must be unique within the Workflow. @@ -5724,14 +5696,6 @@ interface WorkflowInstanceCreateOptions { * The event payload the Workflow instance is triggered with */ params?: PARAMS; - /** - * The retention policy for Workflow instance. - * Defaults to the maximum retention period available for the owner's account. - */ - retention?: { - successRetention?: WorkflowRetentionDuration; - errorRetention?: WorkflowRetentionDuration; - }; } type InstanceStatus = { status: 'queued' // means that instance is waiting to be started (see concurrency limits) diff --git a/wrangler.jsonc b/wrangler.jsonc index 2bef506..fae8101 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -50,7 +50,8 @@ ], "vars": { "HONEYCOMB_DATASET": "", - "HONEYCOMB_API_KEY": "" + "HONEYCOMB_API_KEY": "", + "HOST_NAME": "https://mcp.thoughtspot.com" }, "assets": { "directory": "./static/", "binding": "ASSETS" }, } From 939d6e11282941f785b8fd8d15be42075ed31f6d Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Thu, 24 Jul 2025 18:45:50 -0400 Subject: [PATCH 02/21] remove value from log line --- src/utils.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/utils.ts b/src/utils.ts index 7ed6c12..23ea5e5 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -159,7 +159,6 @@ export async function getFromKV(key: string, env: Env) { if (env?.OAUTH_KV) { const value = await env.OAUTH_KV.get(key, { type: "json" }); if (value) { - console.log("[DEBUG] Value", value); return value; } return null; From c9e55bb41959c06c0e53ffd6a9577de5f7420949 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Fri, 25 Jul 2025 14:24:46 -0400 Subject: [PATCH 03/21] add png image to chatgpt mcp server --- src/servers/mcp-server.ts | 2 +- src/servers/openai-mcp-server.ts | 27 +++++++++++++++++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 8259c1e..133031b 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -252,7 +252,7 @@ export class MCPServer extends BaseMCPServer { }, ...(tokenUrl ? [{ type: "text" as const, - text: `URL to download the answer's PNG image visualization: ${tokenUrl}. Use this URL to download the visualization associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data.`, + text: `Use this URL : ${tokenUrl} to GET the images associated with the data. It might take time to get the image but use this PNG image of the visualization to do a graphical analysis of the data.`, }] : []), ]; return this.createMultiContentSuccessResponse(content, "Answer found"); diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 0b81235..b2a9b4a 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -8,6 +8,7 @@ import { zodToJsonSchema } from "zod-to-json-schema"; import { BaseMCPServer, type Context } from "./mcp-server-base"; import { z } from "zod"; import { WithSpan } from "../metrics/tracing/tracing-utils"; +import { putInKV } from "../utils"; const ToolInputSchema = ToolSchema.shape.inputSchema; type ToolInput = z.infer; @@ -136,11 +137,33 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`); } + let tokenUrl = ""; + + // Generate token and store in KV store + if (!answer.error && this.ctx.env?.OAUTH_KV) { + console.log("[DEBUG] Storing token in KV"); + const token = crypto.randomUUID(); + const tokenData = { + sessionId: answer.session_identifier, + generationNo: answer.generation_number, + instanceURL: this.ctx.props.instanceUrl, + accessToken: this.ctx.props.accessToken + }; + await putInKV(token, tokenData, this.ctx.env); + tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; + } + const content = [ + { type: "text" as const, text: answer.data }, + ...(tokenUrl ? [{ + type: "text" as const, + text: "Use the URL to GET the images associated with the data. It might take time to get the image but use this PNG image of the visualization to do a graphical analysis of the data.", + }] : []), + ]; const result = { id, title: question, - text: answer.data, - url: `${this.ctx.props.instanceUrl}/#/insights/conv-assist?query=${question.trim()}&worksheet=${datasourceId}&executeSearch=true`, + text: content, + url: tokenUrl, } return this.createStructuredContentSuccessResponse(result, "Answer found"); From fd7da77f040a7ccc7d604b6a16d9513ca48480b2 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 13:54:04 -0400 Subject: [PATCH 04/21] fix tool response text --- src/servers/openai-mcp-server.ts | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index b2a9b4a..eade7ae 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -152,19 +152,26 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { await putInKV(token, tokenData, this.ctx.env); tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; } - const content = [ - { type: "text" as const, text: answer.data }, - ...(tokenUrl ? [{ - type: "text" as const, - text: "Use the URL to GET the images associated with the data. It might take time to get the image but use this PNG image of the visualization to do a graphical analysis of the data.", - }] : []), - ]; + const content = `The data associated with the question is: ${answer.data} Use the URL to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysisto the user.`; const result = { id, title: question, text: content, url: tokenUrl, } + // const content = [ + // { type: "text" as const, text: answer.data }, + // ...(tokenUrl ? [{ + // type: "text" as const, + // text: "Use the URL to GET the images associated with the data. It might take time to get the image but use this PNG image of the visualization to do a graphical analysis of the data.", + // }] : []), + // ]; + // const result = { + // id, + // title: question, + // text: content, + // url: tokenUrl, + // } return this.createStructuredContentSuccessResponse(result, "Answer found"); } From d6040ea6319062beccfc9ff9c4d46831303d2857 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 16:37:56 -0400 Subject: [PATCH 05/21] add url in text to get png-image --- src/servers/openai-mcp-server.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index eade7ae..3c34bce 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -152,12 +152,12 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { await putInKV(token, tokenData, this.ctx.env); tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; } - const content = `The data associated with the question is: ${answer.data} Use the URL to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysisto the user.`; + const content = `The data associated with the question is: ${answer.data} Use this URL : ${tokenUrl} to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysis to the user. Show the graph to the user in the response, not just the URL.`; const result = { id, title: question, text: content, - url: tokenUrl, + url: `${this.ctx.props.instanceUrl}/#/insights/conv-assist?query=${question.trim()}&worksheet=${datasourceId}&executeSearch=true`, } // const content = [ // { type: "text" as const, text: answer.data }, From 04bf35f64b19364a6e9ff23dba4742a8fa311f7d Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 16:50:37 -0400 Subject: [PATCH 06/21] update wrangler.json --- wrangler.jsonc | 5 ----- 1 file changed, 5 deletions(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index fae8101..bf4897f 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -48,10 +48,5 @@ "analytics_engine_datasets": [ { "binding": "ANALYTICS", "dataset": "mcp_events" } ], - "vars": { - "HONEYCOMB_DATASET": "", - "HONEYCOMB_API_KEY": "", - "HOST_NAME": "https://mcp.thoughtspot.com" - }, "assets": { "directory": "./static/", "binding": "ASSETS" }, } From 3d42f4cf3625d4ff8eeca8d596fe5373f92043dc Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 17:36:06 -0400 Subject: [PATCH 07/21] change return text --- src/servers/openai-mcp-server.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 3c34bce..c3d769c 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -152,7 +152,8 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { await putInKV(token, tokenData, this.ctx.env); tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; } - const content = `The data associated with the question is: ${answer.data} Use this URL : ${tokenUrl} to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysis to the user. Show the graph to the user in the response, not just the URL.`; + const content = `The data associated with the question is: ${answer.data} \n Use this URL : \"${tokenUrl}\" to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysis to the user. Show the graph to the user in the response.`; + const result = { id, title: question, From f14bc189add2e99c26ec66f2fd0b5275f5a9d194 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 17:41:15 -0400 Subject: [PATCH 08/21] use a comprehensive prompt to fetch image --- src/servers/openai-mcp-server.ts | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index c3d769c..bd63e22 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -152,7 +152,23 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { await putInKV(token, tokenData, this.ctx.env); tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; } - const content = `The data associated with the question is: ${answer.data} \n Use this URL : \"${tokenUrl}\" to GET the charts associated with the data. Use this PNG image of the visualization to do a graphical analysis of the data and present this graph as part of the analysis to the user. Show the graph to the user in the response.`; + const content = `The data associated with the question is: ${answer.data} + +IMPORTANT: You must fetch and display the visualization chart for this data: +1. Make an HTTP GET request to this URL: ${tokenUrl} +2. Use these headers for the request: + - Accept: image/png, image/*, */* +3. The response will be a PNG image containing a data visualization chart +4. Download the image and display it to the user as part of your response +5. You can use this chart image to provide a comprehensive graphical analysis of the data as well + +URL to fetch the chart image: ${tokenUrl} + +Example request: +GET ${tokenUrl} +Accept: image/png, image/*, */* + +Please fetch this image now and show it to the user along with your analysis of both the data and the visualization.`; const result = { id, From be9def2d6859c966f8b1252209802820f22f67e2 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 18:12:28 -0400 Subject: [PATCH 09/21] add logs and prompt --- src/handlers.ts | 1 + src/servers/openai-mcp-server.ts | 24 ++++++++---------------- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/src/handlers.ts b/src/handlers.ts index 331c705..4cc5485 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -267,6 +267,7 @@ app.post("/store-token", async (c) => { }); app.get("/data/img", async (c) => { + console.log("[DEBUG] Serving data/img", c.env); const token = c.req.query("token") || ''; if (token === '') { return c.json({ error: "Token not found" }, 404); diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index bd63e22..6ae1a94 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -152,24 +152,16 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { await putInKV(token, tokenData, this.ctx.env); tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; } - const content = `The data associated with the question is: ${answer.data} + let content = answer.data; + if (tokenUrl) { + content = `Data: ${answer.data} -IMPORTANT: You must fetch and display the visualization chart for this data: -1. Make an HTTP GET request to this URL: ${tokenUrl} -2. Use these headers for the request: - - Accept: image/png, image/*, */* -3. The response will be a PNG image containing a data visualization chart -4. Download the image and display it to the user as part of your response -5. You can use this chart image to provide a comprehensive graphical analysis of the data as well - -URL to fetch the chart image: ${tokenUrl} - -Example request: -GET ${tokenUrl} -Accept: image/png, image/*, */* - -Please fetch this image now and show it to the user along with your analysis of both the data and the visualization.`; +**CHART VISUALIZATION AVAILABLE:** +There is a data visualization chart for this query available at: ${tokenUrl} +This URL contains a PNG image showing the visual representation of the above data. Please fetch and display this image to provide a complete analysis with both the data and its visual representation.`; + } + console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { id, title: question, From fe53c6eed0362d22d5ebadea1ea2e94a5cda42e0 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Mon, 28 Jul 2025 21:28:24 -0400 Subject: [PATCH 10/21] add question in tool response --- src/servers/openai-mcp-server.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 6ae1a94..016a14c 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -159,7 +159,7 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { **CHART VISUALIZATION AVAILABLE:** There is a data visualization chart for this query available at: ${tokenUrl} -This URL contains a PNG image showing the visual representation of the above data. Please fetch and display this image to provide a complete analysis with both the data and its visual representation.`; +This URL contains a PNG image showing the visual representation of the above data for the question: ${question}. Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question.`; } console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { From 1bc68792065cb4c98a8676cc59150f18653a1669 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 12:08:15 -0400 Subject: [PATCH 11/21] make some prompt changes --- src/servers/openai-mcp-server.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 016a14c..1e8b90f 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -159,7 +159,8 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { **CHART VISUALIZATION AVAILABLE:** There is a data visualization chart for this query available at: ${tokenUrl} -This URL contains a PNG image showing the visual representation of the above data for the question: ${question}. Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question.`; +This URL contains a PNG image showing the visual representation of the above data for the question: ${question}. Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question. +Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response. This image is very important and must be displayed in the response.`; } console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { From f7c92a2698783a11c9141ab485850f60344db8cb Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 13:42:35 -0400 Subject: [PATCH 12/21] make url prominent --- src/servers/openai-mcp-server.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 1e8b90f..8b82f39 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -157,10 +157,10 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { content = `Data: ${answer.data} **CHART VISUALIZATION AVAILABLE:** -There is a data visualization chart for this query available at: ${tokenUrl} +There is a data visualization chart for this query available at: "${tokenUrl}". -This URL contains a PNG image showing the visual representation of the above data for the question: ${question}. Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question. -Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response. This image is very important and must be displayed in the response.`; +This URL contains a PNG image showing the visual representation of the above data for the question: "${question}". Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question. +Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response. This image is very important and must be displayed in the response. If the image cannot be opened, log the error that you are getting.`; } console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { From 9e96814f0e652dd3e6099e1eea1cc27377eb3b88 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 14:57:24 -0400 Subject: [PATCH 13/21] return image blob --- src/servers/openai-mcp-server.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 8b82f39..0230d6d 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -132,7 +132,10 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { 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); + const thoughtSpotService = this.getThoughtSpotService(); + const answer = await thoughtSpotService.getAnswerForQuestion(question, datasourceId, false); + const image = await thoughtSpotService.getAnswerImagePNG(answer.session_identifier, answer.generation_number); + const imageBuffer = await image.arrayBuffer(); if (answer.error) { return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`); } @@ -157,10 +160,9 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { content = `Data: ${answer.data} **CHART VISUALIZATION AVAILABLE:** -There is a data visualization chart for this query available at: "${tokenUrl}". +For question ${question} here is a image blob format of the data visualization chart. Use this image to do a graphical analysis of the data and present this graph as part of the analysis to the user. -This URL contains a PNG image showing the visual representation of the above data for the question: "${question}". Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response for the user to see in the response of this question. -Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response. This image is very important and must be displayed in the response. If the image cannot be opened, log the error that you are getting.`; +${imageBuffer}`; } console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { From 0a850090b78b2d0c785085e8ac825fe0dfca2ca2 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 16:45:27 -0400 Subject: [PATCH 14/21] add tests --- src/handlers.ts | 64 ++-- src/servers/openai-mcp-server.ts | 18 +- src/thoughtspot/thoughtspot-service.ts | 1 + test/handlers.spec.ts | 350 +++++++++++++++++++ test/servers/openai-mcp-server.spec.ts | 241 ++++++++++++- test/thoughtspot/thoughtspot-service.spec.ts | 173 +++++++++ 6 files changed, 801 insertions(+), 46 deletions(-) diff --git a/src/handlers.ts b/src/handlers.ts index 4cc5485..5b7b470 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -183,6 +183,43 @@ class Handler { return { redirectTo }; } + + @WithSpan('get-answer-image') + async getAnswerImage(request: Request, env: Env) { + const span = getActiveSpan(); + const url = new URL(request.url); + const token = url.searchParams.get('token') || ''; + if (token === '') { + return new Response("Token not found", { status: 404 }); + } + console.log("[DEBUG] Token", token); + + const tokenData = await getFromKV(token, env); + if (!tokenData) { + span?.setStatus({ code: SpanStatusCode.ERROR, message: "Token not found" }); + return new Response("Token not found", { status: 404 }); + } + console.log("[DEBUG] Token found", token); + + // Extract values from token data + const sessionId = (tokenData as any).sessionId; + const generationNo = (tokenData as any).GenNo || (tokenData as any).generationNo; // Handle both field names + const instanceURL = (tokenData as any).instanceURL; + const accessToken = (tokenData as any).accessToken; + + console.log("[DEBUG] Session ID", sessionId); + console.log("[DEBUG] Generation No", generationNo); + console.log("[DEBUG] Instance URL", instanceURL); + + const thoughtSpotService = new ThoughtSpotService(getThoughtSpotClient(instanceURL, accessToken)); + const image = await thoughtSpotService.getAnswerImagePNG(sessionId, generationNo); + span?.setStatus({ code: SpanStatusCode.OK, message: "Image fetched successfully" }); + return new Response(await image.arrayBuffer(), { + headers: { + 'Content-Type': 'image/png', + }, + }); + } } const handler = new Handler(); @@ -267,32 +304,7 @@ app.post("/store-token", async (c) => { }); app.get("/data/img", async (c) => { - console.log("[DEBUG] Serving data/img", c.env); - const token = c.req.query("token") || ''; - if (token === '') { - return c.json({ error: "Token not found" }, 404); - } - console.log("[DEBUG] Token", token); - - const tokenData = await getFromKV(token, c.env); - if (!tokenData) { - return c.json({ error: "Token not found" }, 404); - } - console.log("[DEBUG] Token found", token); - - // Extract values from token data - const sessionId = (tokenData as any).sessionId; - const generationNo = (tokenData as any).GenNo || (tokenData as any).generationNo; // Handle both field names - const instanceURL = (tokenData as any).instanceURL; - const accessToken = (tokenData as any).accessToken; - - console.log("[DEBUG] Session ID", sessionId); - console.log("[DEBUG] Generation No", generationNo); - console.log("[DEBUG] Instance URL", instanceURL); - - const thoughtSpotService = new ThoughtSpotService(getThoughtSpotClient(instanceURL, accessToken)); - const image = await thoughtSpotService.getAnswerImagePNG(sessionId, generationNo); - return c.body(await image.arrayBuffer()); + return handler.getAnswerImage(c.req.raw, c.env); }); export default app; \ No newline at end of file diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 0230d6d..1b4d357 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -160,9 +160,10 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { content = `Data: ${answer.data} **CHART VISUALIZATION AVAILABLE:** -For question ${question} here is a image blob format of the data visualization chart. Use this image to do a graphical analysis of the data and present this graph as part of the analysis to the user. +There is a data visualization chart for this query available at: "${tokenUrl}". -${imageBuffer}`; +This URL contains a PNG image showing the visual representation of the above data for the question: "${question}". Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response of this question. +Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response.`; } console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { @@ -171,19 +172,6 @@ ${imageBuffer}`; text: content, url: `${this.ctx.props.instanceUrl}/#/insights/conv-assist?query=${question.trim()}&worksheet=${datasourceId}&executeSearch=true`, } - // const content = [ - // { type: "text" as const, text: answer.data }, - // ...(tokenUrl ? [{ - // type: "text" as const, - // text: "Use the URL to GET the images associated with the data. It might take time to get the image but use this PNG image of the visualization to do a graphical analysis of the data.", - // }] : []), - // ]; - // const result = { - // id, - // title: question, - // text: content, - // url: tokenUrl, - // } return this.createStructuredContentSuccessResponse(result, "Answer found"); } diff --git a/src/thoughtspot/thoughtspot-service.ts b/src/thoughtspot/thoughtspot-service.ts index b088275..adcd0e7 100644 --- a/src/thoughtspot/thoughtspot-service.ts +++ b/src/thoughtspot/thoughtspot-service.ts @@ -283,6 +283,7 @@ export class ThoughtSpotService { return liveboardUrl; } + @WithSpan('get-answer-image') async getAnswerImagePNG(sessionId: string, GenNo: number): Promise { const span = getActiveSpan(); span?.addEvent("get-answer-image"); diff --git a/test/handlers.spec.ts b/test/handlers.spec.ts index e454deb..51f75be 100644 --- a/test/handlers.spec.ts +++ b/test/handlers.spec.ts @@ -11,6 +11,7 @@ import app from "../src/handlers"; // 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 { ThoughtSpotService } from '../src/thoughtspot/thoughtspot-service'; // For correctly-typed Request const IncomingRequest = Request; @@ -799,4 +800,353 @@ describe("Handlers", () => { } }); }); + + describe("GET /data/img (getAnswerImage)", () => { + const mockImageBuffer = new ArrayBuffer(8); + const mockImageData = new Uint8Array(mockImageBuffer); + mockImageData.set([137, 80, 78, 71, 13, 10, 26, 10]); // PNG signature + + // Create a proper HttpFile mock that extends Blob + const createMockHttpFile = () => { + const blob = new Blob([mockImageData], { type: 'image/png' }); + return Object.assign(blob, { + name: 'test-image.png', + arrayBuffer: vi.fn().mockResolvedValue(mockImageBuffer), + }); + }; + + let mockHttpFile: ReturnType; + + beforeEach(() => { + vi.clearAllMocks(); + mockHttpFile = createMockHttpFile(); + }); + + it("should return 404 for missing token parameter", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const result = await runInDurableObject(object, async (instance) => { + const request = new IncomingRequest("https://example.com/data/img"); + return typedWorker.fetch(request, env, mockCtx); + }); + + expect(result.status).toBe(404); + expect(await result.text()).toBe("Token not found"); + }); + + it("should return 404 for empty token parameter", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', ''); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, env, mockCtx); + }); + + expect(result.status).toBe(404); + expect(await result.text()).toBe("Token not found"); + }); + + it("should return 404 when token not found in KV storage", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(null) + } + }; + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'non-existent-token'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(404); + expect(await result.text()).toBe("Token not found"); + expect(mockEnvWithKV.OAUTH_KV.get).toHaveBeenCalledWith('non-existent-token', { type: "json" }); + }); + + it("should successfully return PNG image for valid token", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'test-session-123', + GenNo: 1, + generationNo: 1, // Also test the fallback field name + instanceURL: 'https://test.thoughtspot.cloud', + accessToken: 'test-access-token' + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + // Mock the ThoughtSpotService and client + const mockThoughtSpotService = { + getAnswerImagePNG: vi.fn().mockResolvedValue(mockHttpFile) + }; + + const mockGetThoughtSpotClient = vi.fn().mockReturnValue('mock-client'); + + // We need to mock the imports at the module level + const originalModule = await import('../src/handlers'); + const { ThoughtSpotService } = await import('../src/thoughtspot/thoughtspot-service'); + const { getThoughtSpotClient } = await import('../src/thoughtspot/thoughtspot-client'); + + vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockResolvedValue(mockHttpFile); + vi.doMock('../src/thoughtspot/thoughtspot-client', () => ({ + getThoughtSpotClient: mockGetThoughtSpotClient + })); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'valid-token'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(200); + expect(result.headers.get('Content-Type')).toBe('image/png'); + + const responseBuffer = await result.arrayBuffer(); + expect(responseBuffer).toEqual(mockImageBuffer); + + expect(mockEnvWithKV.OAUTH_KV.get).toHaveBeenCalledWith('valid-token', { type: "json" }); + }); + + it("should handle token data with GenNo field", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'test-session-456', + GenNo: 2, // Using GenNo instead of generationNo + instanceURL: 'https://test.thoughtspot.cloud', + accessToken: 'test-access-token' + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockResolvedValue(mockHttpFile); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'valid-token-genno'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(200); + expect(result.headers.get('Content-Type')).toBe('image/png'); + }); + + it("should handle token data with generationNo field", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'test-session-789', + generationNo: 3, // Using generationNo instead of GenNo + instanceURL: 'https://test.thoughtspot.cloud', + accessToken: 'test-access-token' + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockResolvedValue(mockHttpFile); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'valid-token-generation'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(200); + expect(result.headers.get('Content-Type')).toBe('image/png'); + }); + + it("should handle ThoughtSpotService errors gracefully", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'test-session-error', + GenNo: 1, + instanceURL: 'https://test.thoughtspot.cloud', + accessToken: 'test-access-token' + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + const mockError = new Error('ThoughtSpot API error'); + vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockRejectedValue(mockError); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'error-token'); + const request = new IncomingRequest(url.toString()); + + try { + return await typedWorker.fetch(request, mockEnvWithKV, mockCtx); + } catch (error) { + // If the error is thrown instead of returned as a response, + // create a 500 response to simulate the behavior + return new Response('Internal Server Error', { status: 500 }); + } + }); + + // The exact status might depend on error handling implementation + expect([500, 404]).toContain(result.status); + }); + + it("should properly extract values from complex token data", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'complex-session-id', + GenNo: 5, + generationNo: 10, // Both fields present, GenNo should take precedence + instanceURL: 'https://complex.thoughtspot.cloud', + accessToken: 'complex-access-token', + extraField: 'extra-value', // Extra fields should be ignored + nested: { + field: 'nested-value' + } + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + const getAnswerImagePNGSpy = vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockResolvedValue(mockHttpFile); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'complex-token'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(200); + expect(result.headers.get('Content-Type')).toBe('image/png'); + + // Verify that the service was called with the correct parameters + // GenNo should take precedence over generationNo + expect(getAnswerImagePNGSpy).toHaveBeenCalledWith('complex-session-id', 5); + }); + + it("should handle token data without GenNo or generationNo fields", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockTokenData = { + sessionId: 'session-without-genno', + instanceURL: 'https://test.thoughtspot.cloud', + accessToken: 'test-access-token' + // Missing both GenNo and generationNo + }; + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockResolvedValue(mockTokenData) + } + }; + + const getAnswerImagePNGSpy = vi.spyOn(ThoughtSpotService.prototype, 'getAnswerImagePNG').mockResolvedValue(mockHttpFile); + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'token-no-genno'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithKV, mockCtx); + }); + + expect(result.status).toBe(200); + expect(result.headers.get('Content-Type')).toBe('image/png'); + + // Should be called with undefined for generation number + expect(getAnswerImagePNGSpy).toHaveBeenCalledWith('session-without-genno', undefined); + }); + + it("should handle KV storage errors gracefully", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockEnvWithKV = { + ...env, + OAUTH_KV: { + get: vi.fn().mockRejectedValue(new Error('KV storage error')) + } + }; + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'error-token'); + const request = new IncomingRequest(url.toString()); + + try { + return await typedWorker.fetch(request, mockEnvWithKV, mockCtx); + } catch (error) { + // If the error is thrown instead of returned as a response, + // create a 500 response to simulate the behavior + return new Response('Internal Server Error', { status: 500 }); + } + }); + + // Should handle the error gracefully + expect([500, 404]).toContain(result.status); + }); + + it("should handle missing OAUTH_KV environment variable", async () => { + const id = env.MCP_OBJECT.idFromName("test"); + const object = env.MCP_OBJECT.get(id); + + const mockEnvWithoutKV = { + ...env, + OAUTH_KV: undefined + }; + + const result = await runInDurableObject(object, async (instance) => { + const url = new URL("https://example.com/data/img"); + url.searchParams.append('token', 'some-token'); + const request = new IncomingRequest(url.toString()); + return typedWorker.fetch(request, mockEnvWithoutKV, mockCtx); + }); + + expect(result.status).toBe(404); + expect(await result.text()).toBe("Token not found"); + }); + }); }); \ No newline at end of file diff --git a/test/servers/openai-mcp-server.spec.ts b/test/servers/openai-mcp-server.spec.ts index e56cecc..f42a6d0 100644 --- a/test/servers/openai-mcp-server.spec.ts +++ b/test/servers/openai-mcp-server.spec.ts @@ -4,6 +4,7 @@ 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"; +import * as utils from "../../src/utils"; // Mock the MixpanelTracker vi.mock("../../src/metrics/mixpanel/mixpanel", () => ({ @@ -412,11 +413,19 @@ describe("OpenAI Deep Research MCP Server", () => { // Mock the ThoughtSpot service to return answer const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ data: "The total revenue is $1,000,000", - error: null + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) }); vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); await server.init(); const { callTool } = connect(server); @@ -441,11 +450,19 @@ describe("OpenAI Deep Research MCP Server", () => { // Mock the ThoughtSpot service to return error const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ data: null, - error: { message: "Question not found" } + error: { message: "Question not found" }, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) }); vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); await server.init(); const { callTool } = connect(server); @@ -462,11 +479,19 @@ describe("OpenAI Deep Research MCP Server", () => { // Mock the ThoughtSpot service to return answer const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ data: "The total revenue is $1,000,000", - error: null + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) }); vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); await server.init(); const { callTool } = connect(server); @@ -488,11 +513,19 @@ describe("OpenAI Deep Research MCP Server", () => { // Mock the ThoughtSpot service to return answer const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ data: "The revenue increased by 15%", - error: null + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) }); vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); await server.init(); const { callTool } = connect(server); @@ -509,6 +542,196 @@ describe("OpenAI Deep Research MCP Server", () => { url: "https://test.thoughtspot.cloud/#/insights/conv-assist?query=How much did revenue increase? (in %)&worksheet=ds-123&executeSearch=true" }); }); + + it("should generate token and store in KV when OAUTH_KV is available", async () => { + // Mock the ThoughtSpot service to return answer + const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ + data: "The total revenue is $1,000,000", + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) + }); + + // Mock putInKV + const mockPutInKV = vi.fn().mockResolvedValue(undefined); + vi.spyOn(utils, 'putInKV').mockImplementation(mockPutInKV); + + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') + .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); + + // Mock crypto.randomUUID + const mockToken = "test-token-123"; + vi.spyOn(crypto, 'randomUUID').mockReturnValue(mockToken); + + // Create server with environment that includes OAUTH_KV + const mockEnv = { + OAUTH_KV: {} as any, + HONEYCOMB_API_KEY: "test-key", + HONEYCOMB_DATASET: "test-dataset", + HOST_NAME: "https://test-host.com", + MCP_OBJECT: {} as any, + OPENAI_DEEP_RESEARCH_MCP_OBJECT: {} as any, + ANALYTICS: {} as any, + ASSETS: {} as any + }; + + const serverWithKV = new OpenAIDeepResearchMCPServer({ + props: mockProps, + env: mockEnv + }); + + await serverWithKV.init(); + const { callTool } = connect(serverWithKV); + + const result = await callTool("fetch", { + id: "test-ds: What is the total revenue?" + }); + + expect(result.isError).toBeUndefined(); + + // Verify token generation and storage + expect(crypto.randomUUID).toHaveBeenCalled(); + expect(mockPutInKV).toHaveBeenCalledWith( + mockToken, + { + sessionId: "session-123", + generationNo: 1, + instanceURL: mockProps.instanceUrl, + accessToken: mockProps.accessToken + }, + mockEnv + ); + + // Verify the content includes visualization message + const structuredContent = result.structuredContent as any; + expect(structuredContent.text).toContain("**CHART VISUALIZATION AVAILABLE:**"); + expect(structuredContent.text).toContain(`https://test-host.com/data/img?token=${mockToken}`); + expect(structuredContent.text).toContain("Data: The total revenue is $1,000,000"); + expect(structuredContent.text).toContain("What is the total revenue?"); + }); + + it("should not generate token when OAUTH_KV is not available", async () => { + // Mock the ThoughtSpot service to return answer + const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ + data: "The total revenue is $1,000,000", + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) + }); + + // Mock putInKV + const mockPutInKV = vi.fn().mockResolvedValue(undefined); + vi.spyOn(utils, 'putInKV').mockImplementation(mockPutInKV); + + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') + .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); + + // Mock crypto.randomUUID + vi.spyOn(crypto, 'randomUUID').mockReturnValue("test-token-123"); + + // Create server without OAUTH_KV in environment + const mockEnvWithoutKV = { + OAUTH_KV: undefined as any, + HONEYCOMB_API_KEY: "test-key", + HONEYCOMB_DATASET: "test-dataset", + HOST_NAME: "https://test-host.com", + MCP_OBJECT: {} as any, + OPENAI_DEEP_RESEARCH_MCP_OBJECT: {} as any, + ANALYTICS: {} as any, + ASSETS: {} as any + }; + + const serverWithoutKV = new OpenAIDeepResearchMCPServer({ + props: mockProps, + env: mockEnvWithoutKV + }); + + await serverWithoutKV.init(); + const { callTool } = connect(serverWithoutKV); + + const result = await callTool("fetch", { + id: "test-ds: What is the total revenue?" + }); + + expect(result.isError).toBeUndefined(); + + // Verify token generation and storage were NOT called + expect(crypto.randomUUID).not.toHaveBeenCalled(); + expect(mockPutInKV).not.toHaveBeenCalled(); + + // Verify the content does NOT include visualization message + const structuredContent = result.structuredContent as any; + expect(structuredContent.text).toBe("The total revenue is $1,000,000"); + expect(structuredContent.text).not.toContain("**CHART VISUALIZATION AVAILABLE:**"); + }); + + it("should not generate token when answer has error", async () => { + // Mock the ThoughtSpot service to return error + const mockGetAnswerForQuestion = vi.fn().mockResolvedValue({ + data: null, + error: { message: "Service error" }, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) + }); + + // Mock putInKV + const mockPutInKV = vi.fn().mockResolvedValue(undefined); + vi.spyOn(utils, 'putInKV').mockImplementation(mockPutInKV); + + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') + .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); + + // Mock crypto.randomUUID + vi.spyOn(crypto, 'randomUUID').mockReturnValue("test-token-123"); + + // Create server with environment that includes OAUTH_KV + const mockEnv = { + OAUTH_KV: {} as any, + HONEYCOMB_API_KEY: "test-key", + HONEYCOMB_DATASET: "test-dataset", + HOST_NAME: "https://test-host.com", + MCP_OBJECT: {} as any, + OPENAI_DEEP_RESEARCH_MCP_OBJECT: {} as any, + ANALYTICS: {} as any, + ASSETS: {} as any + }; + + const serverWithKV = new OpenAIDeepResearchMCPServer({ + props: mockProps, + env: mockEnv + }); + + await serverWithKV.init(); + const { callTool } = connect(serverWithKV); + + const result = await callTool("fetch", { + id: "test-ds: What is the total revenue?" + }); + + expect(result.isError).toBe(true); + + // Verify token generation and storage were NOT called because of error + expect(crypto.randomUUID).not.toHaveBeenCalled(); + expect(mockPutInKV).not.toHaveBeenCalled(); + }); }); describe("Error Handling", () => { @@ -516,11 +739,19 @@ describe("OpenAI Deep Research MCP Server", () => { // 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 + error: null, + session_identifier: "session-123", + generation_number: 1 + }); + + const mockGetAnswerImagePNG = vi.fn().mockResolvedValue({ + arrayBuffer: () => Promise.resolve(new ArrayBuffer(8)) }); vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerForQuestion') .mockImplementation(mockGetAnswerForQuestion); + vi.spyOn(thoughtspotService.ThoughtSpotService.prototype, 'getAnswerImagePNG') + .mockImplementation(mockGetAnswerImagePNG); await server.init(); const { callTool } = connect(server); diff --git a/test/thoughtspot/thoughtspot-service.spec.ts b/test/thoughtspot/thoughtspot-service.spec.ts index f5416e8..f3efb59 100644 --- a/test/thoughtspot/thoughtspot-service.spec.ts +++ b/test/thoughtspot/thoughtspot-service.spec.ts @@ -897,4 +897,177 @@ describe('thoughtspot-service', () => { ]); }); }); + + describe('getAnswerImagePNG', () => { + it('should return PNG image data successfully', async () => { + const sessionId = 'session123'; + const genNo = 1; + const mockImageFile = { + blob: vi.fn().mockResolvedValue(new Blob(['mock image data'], { type: 'image/png' })), + arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(8)), + text: vi.fn().mockResolvedValue('mock image text') + }; + + mockClient.exportAnswerReport = vi.fn().mockResolvedValue(mockImageFile); + + const service = new ThoughtSpotService(mockClient); + const result = await service.getAnswerImagePNG(sessionId, genNo); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: sessionId, + generation_number: genNo, + file_format: 'PNG', + }); + + expect(result).toBe(mockImageFile); + }); + + it('should handle API errors when exporting PNG', async () => { + const sessionId = 'session456'; + const genNo = 2; + const error = new Error('PNG Export Error'); + + mockClient.exportAnswerReport = vi.fn().mockRejectedValue(error); + + const service = new ThoughtSpotService(mockClient); + + await expect(service.getAnswerImagePNG(sessionId, genNo)) + .rejects.toThrow('PNG Export Error'); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: sessionId, + generation_number: genNo, + file_format: 'PNG', + }); + }); + + it('should handle different session identifiers and generation numbers', async () => { + const testCases = [ + { sessionId: 'session-abc-123', genNo: 0 }, + { sessionId: 'session-def-456', genNo: 5 }, + { sessionId: 'session-xyz-789', genNo: 100 } + ]; + + const mockImageFile = { + blob: vi.fn().mockResolvedValue(new Blob(['mock image data'], { type: 'image/png' })) + }; + + for (const { sessionId, genNo } of testCases) { + mockClient.exportAnswerReport = vi.fn().mockResolvedValue(mockImageFile); + + const service = new ThoughtSpotService(mockClient); + const result = await service.getAnswerImagePNG(sessionId, genNo); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: sessionId, + generation_number: genNo, + file_format: 'PNG', + }); + + expect(result).toBe(mockImageFile); + } + }); + + it('should handle empty session identifier', async () => { + const sessionId = ''; + const genNo = 1; + const mockImageFile = { + blob: vi.fn().mockResolvedValue(new Blob(['mock image data'], { type: 'image/png' })) + }; + + mockClient.exportAnswerReport = vi.fn().mockResolvedValue(mockImageFile); + + const service = new ThoughtSpotService(mockClient); + const result = await service.getAnswerImagePNG(sessionId, genNo); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: '', + generation_number: genNo, + file_format: 'PNG', + }); + + expect(result).toBe(mockImageFile); + }); + + it('should handle negative generation numbers', async () => { + const sessionId = 'session123'; + const genNo = -1; + const mockImageFile = { + blob: vi.fn().mockResolvedValue(new Blob(['mock image data'], { type: 'image/png' })) + }; + + mockClient.exportAnswerReport = vi.fn().mockResolvedValue(mockImageFile); + + const service = new ThoughtSpotService(mockClient); + const result = await service.getAnswerImagePNG(sessionId, genNo); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: sessionId, + generation_number: genNo, + file_format: 'PNG', + }); + + expect(result).toBe(mockImageFile); + }); + + it('should handle network timeout errors', async () => { + const sessionId = 'session123'; + const genNo = 1; + const timeoutError = new Error('Network timeout'); + timeoutError.name = 'TimeoutError'; + + mockClient.exportAnswerReport = vi.fn().mockRejectedValue(timeoutError); + + const service = new ThoughtSpotService(mockClient); + + await expect(service.getAnswerImagePNG(sessionId, genNo)) + .rejects.toThrow('Network timeout'); + + expect(mockClient.exportAnswerReport).toHaveBeenCalledWith({ + session_identifier: sessionId, + generation_number: genNo, + file_format: 'PNG', + }); + }); + + it('should handle authentication errors', async () => { + const sessionId = 'session123'; + const genNo = 1; + const authError = new Error('Authentication failed'); + authError.name = 'AuthenticationError'; + + mockClient.exportAnswerReport = vi.fn().mockRejectedValue(authError); + + const service = new ThoughtSpotService(mockClient); + + await expect(service.getAnswerImagePNG(sessionId, genNo)) + .rejects.toThrow('Authentication failed'); + }); + + it('should return the exact file object from the API', async () => { + const sessionId = 'session123'; + const genNo = 1; + const mockImageFile = { + name: 'answer-image.png', + size: 1024, + type: 'image/png', + lastModified: Date.now(), + blob: vi.fn().mockResolvedValue(new Blob(['mock image data'], { type: 'image/png' })), + arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(1024)), + text: vi.fn().mockResolvedValue('mock image text'), + stream: vi.fn() + }; + + mockClient.exportAnswerReport = vi.fn().mockResolvedValue(mockImageFile); + + const service = new ThoughtSpotService(mockClient); + const result = await service.getAnswerImagePNG(sessionId, genNo); + + // Verify that the exact object is returned without modification + expect(result).toBe(mockImageFile); + expect(result).toHaveProperty('name', 'answer-image.png'); + expect(result).toHaveProperty('size', 1024); + expect(result).toHaveProperty('type', 'image/png'); + }); + }); }); \ No newline at end of file From 7c10671418943fccba5d7b0a3a606afc7a788ff6 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 17:10:08 -0400 Subject: [PATCH 15/21] restore wrangler --- worker-configuration.d.ts | 96 +++++++++++++++++++++++++++------------ 1 file changed, 66 insertions(+), 30 deletions(-) diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index e5a5834..2307189 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,12 +1,11 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 07dddea59586af0824a29829b938dd68) -// Runtime types generated with workerd@1.20250417.0 2025-04-17 nodejs_compat +// Generated by Wrangler by running `wrangler types` (hash: b5949bb168b75f2fe6e6e5a1db8d9e89) +// Runtime types generated with workerd@1.20250604.0 2025-04-17 nodejs_compat declare namespace Cloudflare { interface Env { OAUTH_KV: KVNamespace; - HONEYCOMB_API_KEY: string; - HONEYCOMB_DATASET: string; - HOST_NAME: string; + HONEYCOMB_DATASET: ""; + HONEYCOMB_API_KEY: ""; MCP_OBJECT: DurableObjectNamespace; OPENAI_DEEP_RESEARCH_MCP_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; @@ -18,7 +17,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } // Begin runtime types @@ -101,7 +100,7 @@ interface Console { clear(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/count_static) */ count(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countreset_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countReset_static) */ countReset(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/debug_static) */ debug(...data: any[]): void; @@ -113,9 +112,9 @@ interface Console { error(...data: any[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/group_static) */ group(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupcollapsed_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupCollapsed_static) */ groupCollapsed(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupend_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupEnd_static) */ groupEnd(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/info_static) */ info(...data: any[]): void; @@ -125,9 +124,9 @@ interface Console { table(tabularData?: any, properties?: string[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/time_static) */ time(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeend_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeEnd_static) */ timeEnd(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timelog_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeLog_static) */ timeLog(label?: string, ...data: any[]): void; timeStamp(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/trace_static) */ @@ -302,25 +301,25 @@ declare function dispatchEvent(event: WorkerGlobalScopeEventMap[keyof WorkerGlob declare function btoa(data: string): string; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/atob) */ declare function atob(data: string): string; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ declare function setTimeout(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ declare function setTimeout(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearTimeout) */ declare function clearTimeout(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ declare function setInterval(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ declare function setInterval(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearInterval) */ declare function clearInterval(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/queueMicrotask) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/queueMicrotask) */ declare function queueMicrotask(task: Function): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/structuredClone) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/structuredClone) */ declare function structuredClone(value: T, options?: StructuredSerializeOptions): T; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/reportError) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/reportError) */ declare function reportError(error: any): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/fetch) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/fetch) */ declare function fetch(input: RequestInfo | URL, init?: RequestInit): Promise; declare const self: ServiceWorkerGlobalScope; /** @@ -799,6 +798,7 @@ declare class Blob { slice(start?: number, end?: number, type?: string): Blob; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/arrayBuffer) */ arrayBuffer(): Promise; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/text) */ text(): Promise; @@ -1093,10 +1093,15 @@ interface TextEncoderEncodeIntoResult { */ declare class ErrorEvent extends Event { constructor(type: string, init?: ErrorEventErrorEventInit); + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/filename) */ get filename(): string; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/message) */ get message(): string; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/lineno) */ get lineno(): number; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/colno) */ get colno(): number; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/error) */ get error(): any; } interface ErrorEventErrorEventInit { @@ -1270,6 +1275,7 @@ declare abstract class Body { get bodyUsed(): boolean; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/arrayBuffer) */ arrayBuffer(): Promise; + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/text) */ text(): Promise; @@ -1381,7 +1387,11 @@ interface Request> e * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/integrity) */ integrity: string; - /* Returns a boolean indicating whether or not request can outlive the global in which it was created. */ + /** + * Returns a boolean indicating whether or not request can outlive the global in which it was created. + * + * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/keepalive) + */ keepalive: boolean; /** * Returns the cache mode associated with request, which is a string indicating how the request will interact with the browser's cache when fetching. @@ -1581,6 +1591,7 @@ interface R2ObjectBody extends R2Object { get body(): ReadableStream; get bodyUsed(): boolean; arrayBuffer(): Promise; + bytes(): Promise; text(): Promise; json(): Promise; blob(): Promise; @@ -3826,8 +3837,18 @@ interface AutoRAGNotFoundError extends Error { } interface AutoRAGUnauthorizedError extends Error { } +type ComparisonFilter = { + key: string; + type: 'eq' | 'ne' | 'gt' | 'gte' | 'lt' | 'lte'; + value: string | number | boolean; +}; +type CompoundFilter = { + type: 'and' | 'or'; + filters: ComparisonFilter[]; +}; type AutoRagSearchRequest = { query: string; + filters?: CompoundFilter | ComparisonFilter; max_num_results?: number; ranking_options?: { ranker?: string; @@ -5177,6 +5198,7 @@ declare module 'cloudflare:workers' { export type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; export type WorkflowDelayDuration = WorkflowSleepDuration; export type WorkflowTimeoutDuration = WorkflowSleepDuration; + export type WorkflowRetentionDuration = WorkflowSleepDuration; export type WorkflowBackoff = 'constant' | 'linear' | 'exponential'; export type WorkflowStepConfig = { retries?: { @@ -5307,6 +5329,7 @@ declare namespace TailStream { readonly type: "onset"; readonly dispatchNamespace?: string; readonly entrypoint?: string; + readonly executionModel: string; readonly scriptName?: string; readonly scriptTags?: string[]; readonly scriptVersion?: ScriptVersion; @@ -5324,8 +5347,8 @@ declare namespace TailStream { } interface SpanOpen { readonly type: "spanOpen"; - readonly op?: string; - readonly info?: FetchEventInfo | JsRpcEventInfo | Attribute[]; + readonly name: string; + readonly info?: FetchEventInfo | JsRpcEventInfo | Attributes; } interface SpanClose { readonly type: "spanClose"; @@ -5349,7 +5372,7 @@ declare namespace TailStream { } interface Return { readonly type: "return"; - readonly info?: FetchResponseInfo | Attribute[]; + readonly info?: FetchResponseInfo; } interface Link { readonly type: "link"; @@ -5359,21 +5382,23 @@ declare namespace TailStream { readonly spanId: string; } interface Attribute { - readonly type: "attribute"; readonly name: string; - readonly value: string | string[] | boolean | boolean[] | number | number[]; + readonly value: string | string[] | boolean | boolean[] | number | number[] | bigint | bigint[]; + } + interface Attributes { + readonly type: "attributes"; + readonly info: Attribute[]; } - type Mark = DiagnosticChannelEvent | Exception | Log | Return | Link | Attribute[]; interface TailEvent { readonly traceId: string; readonly invocationId: string; readonly spanId: string; readonly timestamp: Date; readonly sequence: number; - readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | Mark; + readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | DiagnosticChannelEvent | Exception | Log | Return | Link | Attributes; } type TailEventHandler = (event: TailEvent) => void | Promise; - type TailEventHandlerName = "onset" | "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attribute"; + type TailEventHandlerName = "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attributes"; type TailEventHandlerObject = Record; type TailEventHandlerType = TailEventHandler | TailEventHandlerObject; } @@ -5687,6 +5712,9 @@ declare abstract class Workflow { */ public createBatch(batch: WorkflowInstanceCreateOptions[]): Promise; } +type WorkflowDurationLabel = 'second' | 'minute' | 'hour' | 'day' | 'week' | 'month' | 'year'; +type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; +type WorkflowRetentionDuration = WorkflowSleepDuration; interface WorkflowInstanceCreateOptions { /** * An id for your Workflow instance. Must be unique within the Workflow. @@ -5696,6 +5724,14 @@ interface WorkflowInstanceCreateOptions { * The event payload the Workflow instance is triggered with */ params?: PARAMS; + /** + * The retention policy for Workflow instance. + * Defaults to the maximum retention period available for the owner's account. + */ + retention?: { + successRetention?: WorkflowRetentionDuration; + errorRetention?: WorkflowRetentionDuration; + }; } type InstanceStatus = { status: 'queued' // means that instance is waiting to be started (see concurrency limits) From 5741afc4cadb414f6380ff3c1d38316c7651767b Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Tue, 29 Jul 2025 17:11:03 -0400 Subject: [PATCH 16/21] add HOST_NAME to envirnoment variables --- worker-configuration.d.ts | 96 ++++++++++++--------------------------- wrangler.jsonc | 8 ++++ 2 files changed, 38 insertions(+), 66 deletions(-) diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 2307189..e5a5834 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,11 +1,12 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: b5949bb168b75f2fe6e6e5a1db8d9e89) -// Runtime types generated with workerd@1.20250604.0 2025-04-17 nodejs_compat +// Generated by Wrangler by running `wrangler types` (hash: 07dddea59586af0824a29829b938dd68) +// Runtime types generated with workerd@1.20250417.0 2025-04-17 nodejs_compat declare namespace Cloudflare { interface Env { OAUTH_KV: KVNamespace; - HONEYCOMB_DATASET: ""; - HONEYCOMB_API_KEY: ""; + HONEYCOMB_API_KEY: string; + HONEYCOMB_DATASET: string; + HOST_NAME: string; MCP_OBJECT: DurableObjectNamespace; OPENAI_DEEP_RESEARCH_MCP_OBJECT: DurableObjectNamespace; ANALYTICS: AnalyticsEngineDataset; @@ -17,7 +18,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } // Begin runtime types @@ -100,7 +101,7 @@ interface Console { clear(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/count_static) */ count(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countReset_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/countreset_static) */ countReset(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/debug_static) */ debug(...data: any[]): void; @@ -112,9 +113,9 @@ interface Console { error(...data: any[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/group_static) */ group(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupCollapsed_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupcollapsed_static) */ groupCollapsed(...data: any[]): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupEnd_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/groupend_static) */ groupEnd(): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/info_static) */ info(...data: any[]): void; @@ -124,9 +125,9 @@ interface Console { table(tabularData?: any, properties?: string[]): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/time_static) */ time(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeEnd_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeend_static) */ timeEnd(label?: string): void; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timeLog_static) */ + /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/timelog_static) */ timeLog(label?: string, ...data: any[]): void; timeStamp(label?: string): void; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/console/trace_static) */ @@ -301,25 +302,25 @@ declare function dispatchEvent(event: WorkerGlobalScopeEventMap[keyof WorkerGlob declare function btoa(data: string): string; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/atob) */ declare function atob(data: string): string; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ declare function setTimeout(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setTimeout) */ declare function setTimeout(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearTimeout) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearTimeout) */ declare function clearTimeout(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ declare function setInterval(callback: (...args: any[]) => void, msDelay?: number): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/setInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/setInterval) */ declare function setInterval(callback: (...args: Args) => void, msDelay?: number, ...args: Args): number; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/clearInterval) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/clearInterval) */ declare function clearInterval(timeoutId: number | null): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/queueMicrotask) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/queueMicrotask) */ declare function queueMicrotask(task: Function): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/structuredClone) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/structuredClone) */ declare function structuredClone(value: T, options?: StructuredSerializeOptions): T; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/reportError) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/reportError) */ declare function reportError(error: any): void; -/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Window/fetch) */ +/* [MDN Reference](https://developer.mozilla.org/docs/Web/API/fetch) */ declare function fetch(input: RequestInfo | URL, init?: RequestInit): Promise; declare const self: ServiceWorkerGlobalScope; /** @@ -798,7 +799,6 @@ declare class Blob { slice(start?: number, end?: number, type?: string): Blob; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/arrayBuffer) */ arrayBuffer(): Promise; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Blob/text) */ text(): Promise; @@ -1093,15 +1093,10 @@ interface TextEncoderEncodeIntoResult { */ declare class ErrorEvent extends Event { constructor(type: string, init?: ErrorEventErrorEventInit); - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/filename) */ get filename(): string; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/message) */ get message(): string; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/lineno) */ get lineno(): number; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/colno) */ get colno(): number; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/ErrorEvent/error) */ get error(): any; } interface ErrorEventErrorEventInit { @@ -1275,7 +1270,6 @@ declare abstract class Body { get bodyUsed(): boolean; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/arrayBuffer) */ arrayBuffer(): Promise; - /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/bytes) */ bytes(): Promise; /* [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/text) */ text(): Promise; @@ -1387,11 +1381,7 @@ interface Request> e * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/integrity) */ integrity: string; - /** - * Returns a boolean indicating whether or not request can outlive the global in which it was created. - * - * [MDN Reference](https://developer.mozilla.org/docs/Web/API/Request/keepalive) - */ + /* Returns a boolean indicating whether or not request can outlive the global in which it was created. */ keepalive: boolean; /** * Returns the cache mode associated with request, which is a string indicating how the request will interact with the browser's cache when fetching. @@ -1591,7 +1581,6 @@ interface R2ObjectBody extends R2Object { get body(): ReadableStream; get bodyUsed(): boolean; arrayBuffer(): Promise; - bytes(): Promise; text(): Promise; json(): Promise; blob(): Promise; @@ -3837,18 +3826,8 @@ interface AutoRAGNotFoundError extends Error { } interface AutoRAGUnauthorizedError extends Error { } -type ComparisonFilter = { - key: string; - type: 'eq' | 'ne' | 'gt' | 'gte' | 'lt' | 'lte'; - value: string | number | boolean; -}; -type CompoundFilter = { - type: 'and' | 'or'; - filters: ComparisonFilter[]; -}; type AutoRagSearchRequest = { query: string; - filters?: CompoundFilter | ComparisonFilter; max_num_results?: number; ranking_options?: { ranker?: string; @@ -5198,7 +5177,6 @@ declare module 'cloudflare:workers' { export type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; export type WorkflowDelayDuration = WorkflowSleepDuration; export type WorkflowTimeoutDuration = WorkflowSleepDuration; - export type WorkflowRetentionDuration = WorkflowSleepDuration; export type WorkflowBackoff = 'constant' | 'linear' | 'exponential'; export type WorkflowStepConfig = { retries?: { @@ -5329,7 +5307,6 @@ declare namespace TailStream { readonly type: "onset"; readonly dispatchNamespace?: string; readonly entrypoint?: string; - readonly executionModel: string; readonly scriptName?: string; readonly scriptTags?: string[]; readonly scriptVersion?: ScriptVersion; @@ -5347,8 +5324,8 @@ declare namespace TailStream { } interface SpanOpen { readonly type: "spanOpen"; - readonly name: string; - readonly info?: FetchEventInfo | JsRpcEventInfo | Attributes; + readonly op?: string; + readonly info?: FetchEventInfo | JsRpcEventInfo | Attribute[]; } interface SpanClose { readonly type: "spanClose"; @@ -5372,7 +5349,7 @@ declare namespace TailStream { } interface Return { readonly type: "return"; - readonly info?: FetchResponseInfo; + readonly info?: FetchResponseInfo | Attribute[]; } interface Link { readonly type: "link"; @@ -5382,23 +5359,21 @@ declare namespace TailStream { readonly spanId: string; } interface Attribute { + readonly type: "attribute"; readonly name: string; - readonly value: string | string[] | boolean | boolean[] | number | number[] | bigint | bigint[]; - } - interface Attributes { - readonly type: "attributes"; - readonly info: Attribute[]; + readonly value: string | string[] | boolean | boolean[] | number | number[]; } + type Mark = DiagnosticChannelEvent | Exception | Log | Return | Link | Attribute[]; interface TailEvent { readonly traceId: string; readonly invocationId: string; readonly spanId: string; readonly timestamp: Date; readonly sequence: number; - readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | DiagnosticChannelEvent | Exception | Log | Return | Link | Attributes; + readonly event: Onset | Outcome | Hibernate | SpanOpen | SpanClose | Mark; } type TailEventHandler = (event: TailEvent) => void | Promise; - type TailEventHandlerName = "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attributes"; + type TailEventHandlerName = "onset" | "outcome" | "hibernate" | "spanOpen" | "spanClose" | "diagnosticChannel" | "exception" | "log" | "return" | "link" | "attribute"; type TailEventHandlerObject = Record; type TailEventHandlerType = TailEventHandler | TailEventHandlerObject; } @@ -5712,9 +5687,6 @@ declare abstract class Workflow { */ public createBatch(batch: WorkflowInstanceCreateOptions[]): Promise; } -type WorkflowDurationLabel = 'second' | 'minute' | 'hour' | 'day' | 'week' | 'month' | 'year'; -type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; -type WorkflowRetentionDuration = WorkflowSleepDuration; interface WorkflowInstanceCreateOptions { /** * An id for your Workflow instance. Must be unique within the Workflow. @@ -5724,14 +5696,6 @@ interface WorkflowInstanceCreateOptions { * The event payload the Workflow instance is triggered with */ params?: PARAMS; - /** - * The retention policy for Workflow instance. - * Defaults to the maximum retention period available for the owner's account. - */ - retention?: { - successRetention?: WorkflowRetentionDuration; - errorRetention?: WorkflowRetentionDuration; - }; } type InstanceStatus = { status: 'queued' // means that instance is waiting to be started (see concurrency limits) diff --git a/wrangler.jsonc b/wrangler.jsonc index bf4897f..14631e9 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -48,5 +48,13 @@ "analytics_engine_datasets": [ { "binding": "ANALYTICS", "dataset": "mcp_events" } ], +<<<<<<< HEAD +======= + "vars": { + "HONEYCOMB_DATASET": "", + "HONEYCOMB_API_KEY": "", + "HOST_NAME": "https://agent.thoughtspot.app" + }, +>>>>>>> 188b5de (add HOST_NAME to envirnoment variables) "assets": { "directory": "./static/", "binding": "ASSETS" }, } From ff6b108ebf19685f922409141020cc709f8d1cf2 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Wed, 30 Jul 2025 11:43:48 -0400 Subject: [PATCH 17/21] remove env vars --- wrangler.jsonc | 9 --------- 1 file changed, 9 deletions(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index 14631e9..e06c298 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -3,7 +3,6 @@ * https://developers.cloudflare.com/workers/wrangler/configuration/ */ { - "keep_vars": true, "$schema": "node_modules/wrangler/config-schema.json", "name": "thoughtspot-mcp-server", "main": "src/index.ts", @@ -48,13 +47,5 @@ "analytics_engine_datasets": [ { "binding": "ANALYTICS", "dataset": "mcp_events" } ], -<<<<<<< HEAD -======= - "vars": { - "HONEYCOMB_DATASET": "", - "HONEYCOMB_API_KEY": "", - "HOST_NAME": "https://agent.thoughtspot.app" - }, ->>>>>>> 188b5de (add HOST_NAME to envirnoment variables) "assets": { "directory": "./static/", "binding": "ASSETS" }, } From 60c1e094d9843b1ef33d25db282655750dadcb10 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Wed, 30 Jul 2025 12:14:31 -0400 Subject: [PATCH 18/21] create single implementation for getTokenURL --- src/servers/mcp-server-base.ts | 20 +++++++++++++++++++- src/servers/mcp-server.ts | 18 ++---------------- src/servers/openai-mcp-server.ts | 32 +++++++------------------------- wrangler.jsonc | 5 +++-- 4 files changed, 31 insertions(+), 44 deletions(-) diff --git a/src/servers/mcp-server-base.ts b/src/servers/mcp-server-base.ts index 78ec3be..607488a 100644 --- a/src/servers/mcp-server-base.ts +++ b/src/servers/mcp-server-base.ts @@ -11,7 +11,7 @@ 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 { putInKV, type Props } from "../utils"; import { MixpanelTracker } from "../metrics/mixpanel/mixpanel"; import { getThoughtSpotClient } from "../thoughtspot/thoughtspot-client"; import { ThoughtSpotService } from "../thoughtspot/thoughtspot-service"; @@ -148,6 +148,24 @@ export abstract class BaseMCPServer extends Server { this.addTracker(mixpanel); } + protected async getTokenUrl(sessionId: string, generationNo: number) { + let tokenUrl = ""; + // Generate token and store in KV store + if (this.ctx.env?.OAUTH_KV) { + console.log("[DEBUG] Storing token in KV"); + const token = crypto.randomUUID(); + const tokenData = { + sessionId: sessionId, + generationNo: generationNo, + instanceURL: this.ctx.props.instanceUrl, + accessToken: this.ctx.props.accessToken + }; + await putInKV(token, tokenData, this.ctx.env); + tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?uniqueId=${token}`; + } + return tokenUrl; + } + /** * Abstract method to be implemented by subclasses for listing tools */ diff --git a/src/servers/mcp-server.ts b/src/servers/mcp-server.ts index 133031b..8db23d4 100644 --- a/src/servers/mcp-server.ts +++ b/src/servers/mcp-server.ts @@ -6,7 +6,7 @@ import { import { z } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; -import { McpServerError, putInKV } from "../utils"; +import { McpServerError } from "../utils"; import type { DataSource } from "../thoughtspot/thoughtspot-service"; @@ -224,21 +224,7 @@ export class MCPServer extends BaseMCPServer { const { question, datasourceId: sourceId } = GetAnswerSchema.parse(request.params.arguments); const answer = await this.getThoughtSpotService().getAnswerForQuestion(question, sourceId, false); - let tokenUrl = ""; - - // Generate token and store in KV store - if (!answer.error && this.ctx.env?.OAUTH_KV) { - console.log("[DEBUG] Storing token in KV"); - const token = crypto.randomUUID(); - const tokenData = { - sessionId: answer.session_identifier, - generationNo: answer.generation_number, - instanceURL: this.ctx.props.instanceUrl, - accessToken: this.ctx.props.accessToken - }; - await putInKV(token, tokenData, this.ctx.env); - tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; - } + const tokenUrl = await this.getTokenUrl(answer.session_identifier, answer.generation_number); if (answer.error) { return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`); diff --git a/src/servers/openai-mcp-server.ts b/src/servers/openai-mcp-server.ts index 1b4d357..4ab9b3d 100644 --- a/src/servers/openai-mcp-server.ts +++ b/src/servers/openai-mcp-server.ts @@ -8,7 +8,6 @@ import { zodToJsonSchema } from "zod-to-json-schema"; import { BaseMCPServer, type Context } from "./mcp-server-base"; import { z } from "zod"; import { WithSpan } from "../metrics/tracing/tracing-utils"; -import { putInKV } from "../utils"; const ToolInputSchema = ToolSchema.shape.inputSchema; type ToolInput = z.infer; @@ -134,37 +133,20 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer { const [datasourceId, question = ""] = id.split(":"); const thoughtSpotService = this.getThoughtSpotService(); const answer = await thoughtSpotService.getAnswerForQuestion(question, datasourceId, false); - const image = await thoughtSpotService.getAnswerImagePNG(answer.session_identifier, answer.generation_number); - const imageBuffer = await image.arrayBuffer(); if (answer.error) { return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`); } - let tokenUrl = ""; + const tokenUrl = await this.getTokenUrl(answer.session_identifier, answer.generation_number); - // Generate token and store in KV store - if (!answer.error && this.ctx.env?.OAUTH_KV) { - console.log("[DEBUG] Storing token in KV"); - const token = crypto.randomUUID(); - const tokenData = { - sessionId: answer.session_identifier, - generationNo: answer.generation_number, - instanceURL: this.ctx.props.instanceUrl, - accessToken: this.ctx.props.accessToken - }; - await putInKV(token, tokenData, this.ctx.env); - tokenUrl = `${this.ctx.env?.HOST_NAME}/data/img?token=${token}`; - } - let content = answer.data; - if (tokenUrl) { - content = `Data: ${answer.data} + const content = `Data: ${answer.data} -**CHART VISUALIZATION AVAILABLE:** -There is a data visualization chart for this query available at: "${tokenUrl}". + **CHART VISUALIZATION AVAILABLE:** + There is a data visualization chart for this query available at: "${tokenUrl}". -This URL contains a PNG image showing the visual representation of the above data for the question: "${question}". Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response of this question. -Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response.`; - } + This URL contains a PNG image showing the visual representation of the above data for the question: "${question}". Please fetch and display this image to provide a complete analysis with both the data and its visual representation. The image must be displayed in the response of this question. + Use just this image in the response, do not use any other images or visualizations or charts. It may take time to get the image, so please wait for the image to be fetched and displayed in the response.`; + console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`); const result = { id, diff --git a/wrangler.jsonc b/wrangler.jsonc index e06c298..9735ff5 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -2,7 +2,8 @@ * For more details on how to configure Wrangler, refer to: * https://developers.cloudflare.com/workers/wrangler/configuration/ */ -{ + { + "keep_vars": true, "$schema": "node_modules/wrangler/config-schema.json", "name": "thoughtspot-mcp-server", "main": "src/index.ts", @@ -48,4 +49,4 @@ { "binding": "ANALYTICS", "dataset": "mcp_events" } ], "assets": { "directory": "./static/", "binding": "ASSETS" }, -} +} \ No newline at end of file From 831d19b7ec9ffca430ce8985a12d2136f31db76b Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Wed, 30 Jul 2025 12:16:41 -0400 Subject: [PATCH 19/21] comment migrations --- wrangler.jsonc | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index 9735ff5..945724e 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -23,20 +23,20 @@ } ] }, - "migrations": [ - { - "tag": "v2", - "new_sqlite_classes": [ - "ThoughtSpotMCP" - ] - }, - { - "tag": "v3", - "new_sqlite_classes": [ - "ThoughtSpotOpenAIDeepResearchMCP" - ] - } - ], + // "migrations": [ + // { + // "tag": "v2", + // "new_sqlite_classes": [ + // "ThoughtSpotMCP" + // ] + // }, + // { + // "tag": "v3", + // "new_sqlite_classes": [ + // "ThoughtSpotOpenAIDeepResearchMCP" + // ] + // } + // ], "kv_namespaces": [{ "binding": "OAUTH_KV", "id": "05ca6fed380e4fe48dbfc5c3d03b4070" From 73cc49b7201c8ffc9360a9e72a1f68aeeac2a000 Mon Sep 17 00:00:00 2001 From: Shikhar Bhargava Date: Wed, 30 Jul 2025 12:25:56 -0400 Subject: [PATCH 20/21] use uniqueId as well for fetching image from token url --- src/handlers.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/handlers.ts b/src/handlers.ts index 5b7b470..0113340 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -188,7 +188,7 @@ class Handler { async getAnswerImage(request: Request, env: Env) { const span = getActiveSpan(); const url = new URL(request.url); - const token = url.searchParams.get('token') || ''; + const token = url.searchParams.get('token') || url.searchParams.get('uniqueId') || ''; if (token === '') { return new Response("Token not found", { status: 404 }); } From 8e060c251dbaa53cc9973c42a35a6b533d97b7d8 Mon Sep 17 00:00:00 2001 From: "cloudflare-workers-and-pages[bot]" <73139402+cloudflare-workers-and-pages[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 16:27:55 +0000 Subject: [PATCH 21/21] Update wrangler config name to png-image --- wrangler.jsonc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wrangler.jsonc b/wrangler.jsonc index 945724e..b72eb66 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": "png-image", "main": "src/index.ts", "compatibility_date": "2025-04-17", "compatibility_flags": [