feat: add maxTimeMS support to the database tools COMPASS-10235

packages/compass-assistant/src/compass-assistant-provider.spec.tsx

@@ -368,6 +368,7 @@ describe('CompassAssistantProvider', function () {
       enableToolCalling = true,
       enableGenAIToolCallingAtlasProject = true,
       enableGenAIToolCalling = true,
+      maxTimeMS,
       query,
       pipeline,
       connections,
@@ -382,6 +383,7 @@ describe('CompassAssistantProvider', function () {
       enableToolCalling?: boolean;
       enableGenAIToolCallingAtlasProject?: boolean;
       enableGenAIToolCalling?: boolean;
+      maxTimeMS?: number;
       query?: string;
       pipeline?: string;
       connections?: ActiveConnectionInfo[];
@@ -412,6 +414,7 @@ describe('CompassAssistantProvider', function () {
             enableToolCalling,
             enableGenAIToolCallingAtlasProject,
             enableGenAIToolCalling,
+            maxTimeMS,
           },
         }
       );
@@ -704,6 +707,7 @@ describe('CompassAssistantProvider', function () {
       expect(mockToolsController.setContext.callCount).to.equal(1);
       expect(mockToolsController.setContext.firstCall.args[0]).to.deep.equal({
         enableTelemetry: true,
+        maxTimeMS: undefined,
         connections: [],
         query,
         pipeline,
@@ -770,6 +774,7 @@ describe('CompassAssistantProvider', function () {
       expect(mockToolsController.setContext.callCount).to.equal(1);
       expect(mockToolsController.setContext.firstCall.args[0]).to.deep.equal({
         enableTelemetry: true,
+        maxTimeMS: undefined,
         connections: [],
         query,
         pipeline,
@@ -844,6 +849,7 @@ describe('CompassAssistantProvider', function () {
       expect(mockToolsController.setContext.callCount).to.equal(1);
       expect(mockToolsController.setContext.firstCall.args[0]).to.deep.equal({
         enableTelemetry: true,
+        maxTimeMS: undefined,
         connections: [],
         query,
         pipeline,
@@ -929,6 +935,7 @@ describe('CompassAssistantProvider', function () {
       expect(mockToolsController.setContext.callCount).to.equal(1);
       expect(mockToolsController.setContext.firstCall.args[0]).to.deep.equal({
         enableTelemetry: true,
+        maxTimeMS: undefined,
         connections: [
           {
             connectionId: 'connection-1',
@@ -944,6 +951,56 @@ describe('CompassAssistantProvider', function () {
       });
     });
 
+    it('forwards maxTimeMS preference to toolsController.setContext', async function () {
+      const mockChat = new Chat<AssistantMessage>({
+        messages: [
+          {
+            id: 'assistant',
+            role: 'assistant',
+            parts: [{ type: 'text', text: 'Hello user!' }],
+          },
+        ],
+      });
+
+      const sendMessageSpy = sinon.spy(mockChat, 'sendMessage');
+
+      const mockToolsController = {
+        setActiveTools: sinon.stub(),
+        getActiveTools: sinon.stub(),
+        setContext: sinon.stub(),
+        startServer: sinon.stub().resolves(),
+        stopServer: sinon.stub().resolves(),
+        setConnectionIdForToolCall: sinon.stub(),
+      };
+
+      await renderOpenAssistantDrawer({
+        chat: mockChat,
+        toolsController: mockToolsController,
+        enableToolCalling: true,
+        enableGenAIToolCalling: true,
+        maxTimeMS: 5000,
+      });
+
+      const input = screen.getByPlaceholderText('Ask a question');
+      const sendButton = screen.getByLabelText('Send message');
+
+      userEvent.type(input, 'Hello assistant');
+      userEvent.click(sendButton);
+
+      await waitFor(() => {
+        expect(sendMessageSpy.calledOnce).to.be.true;
+      });
+
+      expect(mockToolsController.setContext.callCount).to.equal(1);
+      expect(mockToolsController.setContext.firstCall.args[0]).to.deep.equal({
+        enableTelemetry: true,
+        maxTimeMS: 5000,
+        connections: [],
+        query: undefined,
+        pipeline: undefined,
+      });
+    });
+
     describe('error handling with default chat', function () {
       it('fires a telemetry event and displays error banner when error occurs', async function () {
         const track = sinon.stub();

packages/compass-assistant/src/compass-assistant-provider.tsx

@@ -425,6 +425,7 @@ export function ensureOptInAndSendThunk(
       : null;
     setToolsContext(toolsController, {
       enableTelemetry: prefs.trackUsageStatistics,
+      maxTimeMS: prefs.maxTimeMS,
       activeConnection,
       connections: activeConnections,
       query,
@@ -836,6 +837,7 @@ export function setToolsContext(
   toolsController: ToolsController,
   {
     enableTelemetry,
+    maxTimeMS,
     activeConnection,
     connections,
     query,
@@ -845,6 +847,7 @@ export function setToolsContext(
     activeTab,
   }: {
     enableTelemetry: boolean;
+    maxTimeMS?: number;
     activeConnection: ActiveConnectionInfo | null;
     connections: ActiveConnectionInfo[];
     query?: string | null;
@@ -870,6 +873,7 @@ export function setToolsContext(
     toolsController.setActiveTools(toolGroups);
     toolsController.setContext({
       enableTelemetry,
+      maxTimeMS,
       connections: connections.map((connection) => {
         if (!connection.connectOptions) {
           throw new Error(
@@ -889,6 +893,7 @@ export function setToolsContext(
     toolsController.setActiveTools(new Set([]));
     toolsController.setContext({
       enableTelemetry,
+      maxTimeMS,
       connections: [],
       query: undefined,
       pipeline: undefined,

packages/compass-generative-ai/src/tools-controller.spec.ts

@@ -409,6 +409,26 @@ describe('ToolsController', function () {
         'disabled'
       );
     });
+
+    it('syncs maxTimeMS setting with runner userConfig', function () {
+      toolsController.setContext({
+        enableTelemetry: false,
+        maxTimeMS: 5000,
+        connections: [],
+      });
+      expect((toolsController as any).runner.userConfig.maxTimeMS).to.equal(
+        5000
+      );
+
+      toolsController.setContext({
+        enableTelemetry: false,
+        maxTimeMS: undefined,
+        connections: [],
+      });
+      expect((toolsController as any).runner.userConfig.maxTimeMS).to.equal(
+        undefined
+      );
+    });
   });
 
   describe('setConnectionIdForToolCall', function () {

packages/compass-generative-ai/src/tools-controller.ts

@@ -25,6 +25,7 @@ export type ToolGroup = 'querybar' | 'aggregation-builder' | 'db-read';
 
 type CompassContext = {
   enableTelemetry: boolean;
+  maxTimeMS?: number;
   query?: string;
   pipeline?: string;
 };
@@ -75,6 +76,7 @@ type ToolsControllerConfig = {
   logger: Logger;
   getTelemetryAnonymousId: () => string;
   enableTelemetry: boolean;
+  maxTimeMS?: number;
 };
 
 export class ToolsController {
@@ -90,15 +92,17 @@ export class ToolsController {
     logger,
     getTelemetryAnonymousId,
     enableTelemetry,
+    maxTimeMS,
   }: ToolsControllerConfig) {
     this.logger = logger;
     const mcpConfig = UserConfigSchema.parse({
       disabledTools: ['connect'],
       loggers: ['mcp'],
       readOnly: true,
-      // NOTE: the preference could change at runtime. As a best-effort way of
-      // keeping it in sync we'll change it every time we set the tools' context
+      // NOTE: the preferences could change at runtime. As a best-effort way of
+      // keeping them in sync we'll change them every time we set the tools' context
       telemetry: enableTelemetry ? 'enabled' : 'disabled',
+      maxTimeMS,
     });
 
     this.runner = new InMemoryRunner({
@@ -271,10 +275,11 @@ export class ToolsController {
   }
 
   setContext(context: ToolsContext): void {
-    // make sure this property is always in sync with the tools' config
+    // make sure these properties are always in sync with the tools' config
     this.runner.userConfig.telemetry = context.enableTelemetry
       ? 'enabled'
       : 'disabled';
+    this.runner.userConfig.maxTimeMS = context.maxTimeMS;
     this.context = context;
   }