Png image

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/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;
 }
 
 export abstract class BaseMCPServer extends Server {
@@ -147,6 +148,45 @@ export abstract class BaseMCPServer extends Server {
         this.addTracker(mixpanel);
     }
 
+    protected async createImageUrl(sessionId: string, generationNo: number) {
+        const hostName = this.ctx.props.hostName;
+        console.log(`[DEBUG] hostName: ${hostName}`);
+
+        // If hostName is undefined return empty string
+        if (!hostName) {
+            return "";
+        }
+
+        // Generate token and store in KV store
+        const uniqueId = crypto.randomUUID();
+        const sessionData = {
+            sessionId: sessionId,
+            generationNo: generationNo,
+            instanceURL: this.ctx.props.instanceUrl,
+            accessToken: this.ctx.props.accessToken
+        };
+        try {
+            await putInKV(uniqueId, sessionData, this.ctx.env);
+        } catch (error) {
+            console.error(`[ERROR] Error storing session data in KV: ${error}`);
+            return "";
+        }
+
+        return `${hostName}/data/img?uniqueId=${uniqueId}`;
+    }
+
+    protected async getAnswerContent(answer: any, question: string) {
+        const tokenUrl = await this.createImageUrl(answer.session_identifier, answer.generation_number);
+        console.log(`[DEBUG] question: ${question} tokenUrl: ${tokenUrl}`);
+        return `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. `;
+    }
+
     /**
      * Abstract method to be implemented by subclasses for listing tools
      */

src/servers/mcp-server.ts

@@ -230,7 +230,7 @@ export class MCPServer extends BaseMCPServer {
         }
 
         const content = [
-            { type: "text" as const, text: answer.data },
+            { type: "text" as const, text: await this.getAnswerContent(answer, question) },
             {
                 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.`,

src/servers/openai-mcp-server.ts

@@ -135,11 +135,11 @@ export class OpenAIDeepResearchMCPServer extends BaseMCPServer {
         if (answer.error) {
             return this.createErrorResponse(answer.error.message, `Error getting answer ${answer.error.message}`);
         }
-
+        const content = await this.getAnswerContent(answer, question);
         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

@@ -16,10 +16,10 @@ async function main() {
 
     const props: Props = {
         instanceUrl: validateAndSanitizeUrl(instanceUrl),
-        accessToken,
+        accessToken
     };
 
-    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, generationNo: number): Promise<HttpFile> {
+        const span = getActiveSpan();
+        span?.addEvent("get-answer-image");
+        const data = await this.client.exportAnswerReport({
+            session_identifier: sessionId,
+            generation_number: generationNo,
+            file_format: "PNG",
+        })
+        return data;
+    }
+
     /**
      * Get data sources
      */

src/utils.ts

@@ -13,6 +13,7 @@ export type Props = {
         clientName: string;
         registrationDate: number;
     };
+    hostName: string;
 };
 
 export class McpServerError extends Error {
@@ -92,20 +93,46 @@ 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);
+        // added ! to satisfy the type checker
+        public server: T;
+        public env: Env;
 
-        // 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);
+            this.env = env;
         }
 
         async init() {
+            // Create the server directly here after props have been set by McpAgent._init()
+            const context: Context = {
+                props: this.props, // Available after McpAgent._init() sets it
+                env: this.env     // Stored from constructor
+            };
+            this.server = new MCPServer(context);
             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
+        });
+    } else {
+        throw new McpServerError("OAUTH_KV is not available", 500);
+    }
+}
+
+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;
+    }
 }

static/oauth-callback.js

@@ -120,7 +120,7 @@
                             const errorText = await storeResponse.text();
                             throw new Error(`Failed to store token (Status: ${storeResponse.status}): ${errorText}`);
                         }
-                        window.location.href = responseData.redirectTo;
+                        window.location.href = responseData.data.redirectTo;
                     } catch (err) {
                         document.getElementById('status').textContent = err.message;
                         document.getElementById('status').style.color = '#dc3545';
@@ -149,14 +149,15 @@
             })
         });
         const responseData = await storeResponse.json();
+        console.log('Response data:', responseData);
 
         if (!storeResponse.ok) {
             const errorText = await storeResponse.text();
             throw new Error(`Failed to store token (Status: ${storeResponse.status}): ${errorText}`);
         }
 
-        console.log('Redirecting to:', responseData.redirectTo);
-        window.location.href = responseData.redirectTo;
+        console.log('Redirecting to:', responseData.data.redirectTo);
+        window.location.href = responseData.data.redirectTo;
 
     } catch (error) {
         console.error('Error:', error);