create single implementation for getTokenURL

Author: shikharbsar

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
      */

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

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<typeof ToolInputSchema>;
@@ -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,

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" },
-}
+}