POC for supporting Spotter 3

package.json

@@ -52,7 +52,7 @@
     "@hono/zod-validator": "^0.7.2",
     "@microlabs/otel-cf-workers": "^1.0.0-rc.52",
     "@modelcontextprotocol/sdk": "^1.20.1",
-    "@thoughtspot/rest-api-sdk": "^2.13.1",
+    "@thoughtspot/rest-api-sdk": "^2.22.0",
     "agents": "^0.2.14",
     "hono": "^4.10.3",
     "rxjs": "^7.8.2",

server.json

@@ -70,6 +70,22 @@
       {
         "name": "getDataSourceSuggestions",
         "description": "Get datasource suggestions for a given query"
+      },
+      {
+        "name": "createAgentConversation",
+        "description": "Create a conversation with a given data source"
+      },
+      {
+        "name": "sendAgentMessage",
+        "description": "Send a message in a conversation"
+      },
+      {
+        "name": "sendAgentMessageAsync",
+        "description": "Send a message in a conversation, and allow the response to be generated asynchronously"
+      },
+      {
+        "name": "getAgentMessageUpdates",
+        "description": "Retrieve the latest messages from an asynchronous conversation"
       }
     ],
     "resources": [

src/cloudflare-utils.ts

@@ -2,10 +2,24 @@ import { McpAgent } from "agents/mcp";
 import { instrumentDO, type ResolveConfigFn } from '@microlabs/otel-cf-workers';
 import type { BaseMCPServer, Context } from "./servers/mcp-server-base";
 import type { Props } from "./utils";
+import { StreamingConversationState } from "./servers/mcp-server";
 
-export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (ctx: Context) => T, config: ResolveConfigFn) {
+export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (
+    ctx: Context,
+    getConversationState: (
+        conversationId: string,
+    ) => Promise<StreamingConversationState | undefined>,
+    updateConversationStateAndResetTtlTimeout: (
+        conversationId: string,
+        newState: StreamingConversationState,
+    ) => Promise<void>,
+) => T, config: ResolveConfigFn) {
     const Agent = class extends McpAgent<Env, any, Props> {
-        server = new MCPServer(this as Context);
+        server = new MCPServer(
+            this as Context,
+            this.getConversationState.bind(this),
+            this.updateConversationStateAndResetTtlTimeout.bind(this),
+        );
 
         // Argument of type 'typeof ThoughtSpotMCPWrapper' is not assignable to parameter of type 'DOClass'.
         // Cannot assign a 'protected' constructor type to a 'public' constructor type.
@@ -17,6 +31,7 @@ export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (c
 
         async init() {
             await this.server.init();
+            this.ctx.storage
         }
 
         public static serve(path: string) {
@@ -36,6 +51,34 @@ export function instrumentedMCPServer<T extends BaseMCPServer>(MCPServer: new (c
             }
             return server;
         }
+
+        private async getConversationState(conversationId: string) {
+            return await this.ctx.storage?.get<StreamingConversationState>(conversationId);
+        }
+
+        private async updateConversationStateAndResetTtlTimeout(
+            conversationId: string,
+            newState: StreamingConversationState,
+        ) {
+            const oldState = await this.getConversationState(conversationId);
+            if (oldState?.ttlTimeoutId) {
+                await this.cancelSchedule(oldState.ttlTimeoutId);
+            }
+
+            const schedule = await this.schedule(30, 'clearConversationState' as any, {
+                conversationId,
+            });
+
+            await this.ctx.storage?.put(conversationId, {
+                ...newState,
+                ttlTimeoutId: schedule.id,
+            });
+        }
+
+        private async clearConversationState(payload: { conversationId: string }) {
+            console.log('>>> clearing conversation state', payload.conversationId);
+            await this.ctx.storage?.delete(payload.conversationId);
+        }
     }
 
     return instrumentDO(Agent, config) as unknown as typeof Agent;