Update name in Wrangler configuration file to match deployed Worker

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 } }>()
 
@@ -181,11 +183,49 @@ 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') || url.searchParams.get('uniqueId') || '';
+        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();
 
 app.get("/", async (c) => {
+    console.log("[DEBUG] Serving index", c.env);
     const response = await handler.serveIndex(c.env);
     return response;
 });
@@ -263,4 +303,8 @@ app.post("/store-token", async (c) => {
     }
 });
 
+app.get("/data/img", async (c) => {
+    return handler.getAnswerImage(c.req.raw, c.env);
+});
+
 export default app;

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 }
     };
 };
 

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();

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";
@@ -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 {
@@ -147,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
      */

src/servers/mcp-server.ts

@@ -224,6 +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);
+        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}`);
@@ -235,6 +236,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: `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");
     }

src/servers/openai-mcp-server.ts

@@ -131,15 +131,27 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer {
         const { id } = fetchInputSchema.parse(request.params.arguments);
         // id is of the form "<datasource-id>:<question>"
         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);
         if (answer.error) {
             return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`);
         }
 
+        const tokenUrl = await this.getTokenUrl(answer.session_identifier, answer.generation_number);
+
+        const content = `Data: ${answer.data}
+
+        **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.`;
+
+        console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`);
         const result = {
             id,
             title: question,
-            text: answer.data,
+            text: content,
             url: `${this.ctx.props.instanceUrl}/#/insights/conv-assist?query=${question.trim()}&worksheet=${datasourceId}&executeSearch=true`,
         }
 

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();

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,18 @@ export class ThoughtSpotService {
         return liveboardUrl;
     }
 
+    @WithSpan('get-answer-image')
+    async getAnswerImagePNG(sessionId: string, GenNo: number): Promise<HttpFile> {
+        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
      */

src/utils.ts

@@ -92,20 +92,75 @@ export class McpServerError extends Error {
 
 export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (ctx: Context) => T, config: ResolveConfigFn) {
     const Agent = class extends McpAgent<Env, any, Props> {
-        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.
         // Created to satisfy the DOClass type.
         // 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) {
+            return value;
+        }
+        return null;
+    }
 }