@mongoosejs/studio 0.3.11 → 0.4.0

package/backend/actions/ChatThread/createChatMessage.js

@@ -3,7 +3,9 @@
 const Archetype = require('archetype');
 const assert = require('assert');
 const authorize = require('../../authorize');
+const agentSystemPrompt = require('../../chatAgent/agentSystemPrompt');
 const callLLM = require('../../integrations/callLLM');
+const getAgentTools = require('../../chatAgent/getAgentTools');
 const getModelDescriptions = require('../../helpers/getModelDescriptions');
 const mongoose = require('mongoose');
 
@@ -81,12 +83,14 @@ module.exports = ({ db, studioConnection, options }) => async function createCha
 
   const modelDescriptions = getModelDescriptions(db);
   const system = [
-    systemPrompt,
+    chatThread.agentMode ? agentSystemPrompt : systemPrompt,
     currentDateTime ? `Current date: ${currentDateTime}` : null,
     modelDescriptions,
     options?.context
   ].filter(Boolean).join('\n\n');
 
+  const llmOptions = chatThread.agentMode ? { ...options, tools: getAgentTools(db) } : options;
+
   // Create the chat message and get LLM response in parallel
   const chatMessages = await Promise.all([
     ChatMessage.create({
@@ -96,12 +100,20 @@ module.exports = ({ db, studioConnection, options }) => async function createCha
       script,
       executionResult: null
     }),
-    callLLM(llmMessages, system, options).then(res => {
+    callLLM(llmMessages, system, llmOptions).then(res => {
       const content = res.text;
+      const toolCalls = (res.steps || []).flatMap(step =>
+        (step.toolCalls || []).map(tc => ({
+          toolName: tc.toolName,
+          input: tc.toolName === 'typeCheck' ? {} : tc.args,
+          status: 'done'
+        }))
+      );
       return ChatMessage.create({
         chatThreadId,
         role: 'assistant',
-        content
+        content,
+        toolCalls: toolCalls.length ? toolCalls : undefined
       });
     })
   ]);

package/backend/actions/ChatThread/createChatThread.js

@@ -22,11 +22,15 @@ const CreateChatThreadParams = new Archetype({
   dashboardId: {
     $type: mongoose.Types.ObjectId,
     $required: false
+  },
+  agentMode: {
+    $type: 'boolean',
+    $required: false
   }
 }).compile('CreateChatThreadParams');
 
 module.exports = ({ studioConnection }) => async function createChatThread(params) {
-  const { initiatedById, roles, $workspaceId, initialMessage, dashboardId } = new CreateChatThreadParams(params);
+  const { initiatedById, roles, $workspaceId, initialMessage, dashboardId, agentMode } = new CreateChatThreadParams(params);
   const ChatThread = studioConnection.model('__Studio_ChatThread');
   const ChatMessage = studioConnection.model('__Studio_ChatMessage');
 
@@ -42,6 +46,9 @@ module.exports = ({ studioConnection }) => async function createChatThread(param
   if (dashboardId) {
     doc.dashboardId = dashboardId;
   }
+  if (agentMode != null) {
+    doc.agentMode = agentMode;
+  }
   const chatThread = await ChatThread.create(doc);
 
   if (initialMessage != null && initialMessage.trim() !== '') {

package/backend/actions/ChatThread/index.js

@@ -6,3 +6,4 @@ exports.getChatThread = require('./getChatThread');
 exports.listChatThreads = require('./listChatThreads');
 exports.shareChatThread = require('./shareChatThread');
 exports.streamChatMessage = require('./streamChatMessage');
+exports.toggleAgentMode = require('./toggleAgentMode');

package/backend/actions/ChatThread/streamChatMessage.js

@@ -4,6 +4,8 @@ const Archetype = require('archetype');
 const assert = require('assert');
 const authorize = require('../../authorize');
 const callLLM = require('../../integrations/callLLM');
+const agentSystemPrompt = require('../../chatAgent/agentSystemPrompt');
+const getAgentTools = require('../../chatAgent/getAgentTools');
 const streamLLM = require('../../integrations/streamLLM');
 const getModelDescriptions = require('../../helpers/getModelDescriptions');
 const mongoose = require('mongoose');
@@ -27,7 +29,7 @@ const CreateChatMessageParams = new Archetype({
   }
 }).compile('CreateChatMessageParams');
 
-module.exports = ({ db, studioConnection, options }) => async function* createChatMessage(params) {
+module.exports = ({ db, studioConnection, options }) => async function* streamChatMessage(params) {
   const { chatThreadId, initiatedById, content, currentDateTime, script, roles } = new CreateChatMessageParams(params);
   const ChatThread = studioConnection.model('__Studio_ChatThread');
   const ChatMessage = studioConnection.model('__Studio_ChatMessage');
@@ -73,16 +75,20 @@ module.exports = ({ db, studioConnection, options }) => async function* createCh
       }],
       'You are a helpful assistant that summarizes chat threads into titles.',
       options
-    ).then(res => {
-      const title = res.text;
-      chatThread.title = title;
-      return chatThread.save();
-    });
+    ).then(
+      res => {
+        const title = res.text;
+        chatThread.title = title;
+        return chatThread.save();
+      },
+      // Avoid unhandled promise rejection
+      () => {}
+    );
   }
 
   const modelDescriptions = getModelDescriptions(db);
   const system = [
-    systemPrompt,
+    chatThread.agentMode ? agentSystemPrompt : systemPrompt,
     currentDateTime ? `Current date: ${currentDateTime}` : null,
     modelDescriptions,
     options?.context
@@ -105,16 +111,43 @@ module.exports = ({ db, studioConnection, options }) => async function* createCh
     script: null,
     executionResult: null
   });
-  const textStream = streamLLM(llmMessages, system, options);
+  const llmOptions = chatThread.agentMode ? { ...options, tools: getAgentTools(db) } : options;
+  let textStream;
+  try {
+    textStream = streamLLM(llmMessages, system, llmOptions);
+  } catch (err) {
+    yield { message: err?.message || 'Failed to stream response' };
+    return {};
+  }
+  const toolCalls = [];
   let count = 0;
-  for await (const textPart of textStream) {
-    assistantChatMessage.content += textPart;
-    // Only save every 10th chunk for performance
-    if (count++ % 10 === 0) {
-      await assistantChatMessage.save();
+  try {
+    for await (const event of textStream) {
+      if (typeof event === 'string') {
+        assistantChatMessage.content += event;
+        // Only save every 10th chunk for performance
+        if (count++ % 10 === 0) {
+          await assistantChatMessage.save();
+        }
+        yield { textPart: event };
+      } else if (event.toolCall) {
+        toolCalls.push({ toolName: event.toolCall.toolName, input: event.toolCall.input, status: 'running' });
+        yield { toolCall: event.toolCall };
+      } else if (event.toolResult) {
+        const tc = toolCalls.find(t => t.toolName === event.toolResult.toolName && t.status === 'running');
+        if (tc) {
+          tc.status = 'done';
+        }
+        yield { toolResult: event.toolResult };
+      } else {
+        yield { textPart: event };
+      }
     }
-    yield { textPart };
+  } catch (err) {
+    yield { message: err?.message || 'Failed to stream chat response' };
+    return {};
   }
+  assistantChatMessage.toolCalls = toolCalls;
 
   await assistantChatMessage.save();
   yield { chatMessage: assistantChatMessage };
@@ -139,6 +172,8 @@ const systemPrompt = `
 
   Assume the user has pre-defined schemas and models. Do not define any new schemas or models for the user.
 
+  Avoid using the aggregation framework unless explicitly asked by the user.
+
   Use async/await where possible. Assume top-level await is allowed.
 
   Write at most one script, unless the user explicitly asks for multiple scripts.

package/backend/actions/ChatThread/toggleAgentMode.js

@@ -0,0 +1,37 @@
+'use strict';
+
+const Archetype = require('archetype');
+const authorize = require('../../authorize');
+const mongoose = require('mongoose');
+
+const ToggleAgentModeParams = new Archetype({
+  chatThreadId: {
+    $type: mongoose.Types.ObjectId
+  },
+  agentMode: {
+    $type: 'boolean'
+  },
+  initiatedById: {
+    $type: mongoose.Types.ObjectId
+  },
+  roles: {
+    $type: ['string']
+  }
+}).compile('ToggleAgentModeParams');
+
+module.exports = ({ studioConnection }) => async function toggleAgentMode(params) {
+  const { chatThreadId, agentMode, initiatedById, roles } = new ToggleAgentModeParams(params);
+  const ChatThread = studioConnection.model('__Studio_ChatThread');
+
+  await authorize('ChatThread.toggleAgentMode', roles);
+
+  const chatThread = await ChatThread.findById(chatThreadId).orFail();
+  if (initiatedById != null && chatThread.userId?.toString() !== initiatedById.toString()) {
+    throw new Error('Not authorized');
+  }
+
+  chatThread.agentMode = agentMode;
+  await chatThread.save();
+
+  return { chatThread };
+};

package/backend/actions/Model/streamChatMessage.js

@@ -3,6 +3,7 @@
 const Archetype = require('archetype');
 const assert = require('assert');
 const authorize = require('../../authorize');
+const getAgentTools = require('../../chatAgent/getAgentTools');
 const streamLLM = require('../../integrations/streamLLM');
 const getModelDescriptions = require('../../helpers/getModelDescriptions');
 
@@ -18,6 +19,10 @@ const StreamChatMessageParams = new Archetype({
   documentData: {
     $type: 'string'
   },
+  agentMode: {
+    $type: 'boolean',
+    $default: true
+  },
   currentDateTime: {
     $type: 'string',
     $transform: v => v == null ? null : decodeURIComponent(v),
@@ -29,7 +34,7 @@ const StreamChatMessageParams = new Archetype({
 }).compile('StreamChatMessageParams');
 
 module.exports = ({ db, options }) => async function* streamChatMessage(params) {
-  const { model, content, documentData, currentDateTime, roles } = new StreamChatMessageParams(params);
+  const { model, content, documentData, currentDateTime, agentMode, roles } = new StreamChatMessageParams(params);
 
   await authorize('Model.streamChatMessage', roles);
 
@@ -51,10 +56,16 @@ module.exports = ({ db, options }) => async function* streamChatMessage(params)
   ].filter(Boolean).join('\n\n');
 
   const llmMessages = [{ role: 'user', content: [{ type: 'text', text: content }] }];
-  const textStream = streamLLM(llmMessages, system, options);
+  const llmOptions = agentMode ? { ...options, tools: getAgentTools(db) } : options;
+  const textStream = streamLLM(llmMessages, system, llmOptions);
 
-  for await (const textPart of textStream) {
-    yield { textPart };
+  try {
+    for await (const textPart of textStream) {
+      yield { textPart };
+    }
+  } catch (err) {
+    yield { message: err?.message || 'Failed to stream chat response' };
+    return {};
   }
 
   return {};

package/backend/authorize.js

@@ -7,6 +7,7 @@ const actionsToRequiredRoles = {
   'ChatThread.getChatThread': ['owner', 'admin', 'member'],
   'ChatThread.listChatThreads': ['owner', 'admin', 'member'],
   'ChatThread.shareChatThread': ['owner', 'admin', 'member'],
+  'ChatThread.toggleAgentMode': ['owner', 'admin', 'member'],
   'Dashboard.createDashboard': ['owner', 'admin', 'member'],
   'Dashboard.deleteDashboard': ['owner', 'admin', 'member'],
   'Dashboard.getDashboard': ['owner', 'admin', 'member', 'readonly', 'dashboards'],

package/backend/chatAgent/agentSystemPrompt.js

@@ -0,0 +1,45 @@
+'use strict';
+
+module.exports = `
+You are a data querying assistant with access to tools that can query MongoDB directly. Use these tools to investigate the user's data before writing a final script. The final script must be written in JavaScript.
+
+Your tools are for EXPLORATION ONLY — use them to understand the data before writing a script. NEVER use tool results as the data source for the final script. The final script must be fully self-contained: it must query MongoDB itself and not depend on any data you retrieved via tools.
+
+Follow this process for each query:
+
+1. **Identify models**: Based on the user's question and the model descriptions below, identify which models are relevant.
+2. **Check document counts**: Use estimatedDocumentCount on the relevant models to understand data volume. This helps you decide whether a query is safe to run or needs limits/filters.
+3. **Test assumptions**: Use find or findOne to sample real documents and verify your assumptions about the shape of the data, field values, and relationships. Do not guess field values when you can look them up.
+4. **Type-check**: Before presenting the script, run it through the typeCheck tool to catch errors. The sandbox globals are: db (a mongoose.Connection — use db.model('ModelName') to get models), mongoose, ObjectId, console, and MongooseStudioChartColors. Fix any errors and re-check until clean.
+5. **Write the final script**: Once you understand the data, produce a self-contained script that queries MongoDB directly. The script must NOT use hard-coded data from your tool calls — it must run its own queries.
+
+Keep scripts concise. Avoid unnecessary comments, error handling, and temporary variables.
+
+Do not write any imports or require() statements, that will cause the script to break.
+
+If the user approves the script, the script will run in the Node.js server in a sandboxed vm.createContext() call with the following globals: db (the Mongoose connection), mongoose, ObjectId (mongoose.Types.ObjectId), console, and MongooseStudioChartColors (an array of 8 hex color strings for chart dataset colors). The script return value will then send the response via JSON to the client. Be aware that the result of the query will be serialized to JSON before being displayed to the user. MAKE SURE TO RETURN A VALUE FROM THE SCRIPT.
+
+Optimize scripts for readability first, followed by reliability, followed by performance. Avoid using the aggregation framework unless explicitly requested by the user. Use indexed fields in queries where possible.
+
+Avoid using the aggregation framework unless explicitly asked by the user.
+
+Assume the user has pre-defined schemas and models. Do not define any new schemas or models for the user.
+
+Use async/await where possible. Assume top-level await is allowed.
+
+Write at most one script, unless the user explicitly asks for multiple scripts.
+
+Format output as Markdown, including code fences for any scripts the user requested.
+
+Add a brief text description of what the script does.
+
+If the user's query is best answered with a chart, return a Chart.js 4 configuration as \`return { $chart: chartJSConfig };\`. Disable ChartJS animation by default unless user asks for it. Set responsive: true, maintainAspectRatio: false options unless the user explicitly asks. Use MongooseStudioChartColors for dataset backgroundColor and borderColor by default. For line/bar charts, use MongooseStudioChartColors[i] as borderColor and MongooseStudioChartColors[i] + '33' as backgroundColor for each dataset. For pie/doughnut charts, use MongooseStudioChartColors.slice(0, data.length) as backgroundColor. Only use custom colors if the user explicitly requests specific colors.
+
+If the user's query is best answered by a map, return an object { $featureCollection } which contains a GeoJSON FeatureCollection
+
+If the user's query is best answered by a table, return an object { $table: { columns: string[], rows: any[][] } }
+
+-----------
+
+Here is a description of the user's models. Assume these are the only models available in the system unless explicitly instructed otherwise by the user.
+`.trim();

package/backend/chatAgent/agentToolMetadata.js

@@ -0,0 +1,20 @@
+'use strict';
+
+module.exports = [
+  {
+    name: 'estimatedDocumentCount',
+    description: 'Count the estimated number of documents in a collection. Times out after 5 seconds.'
+  },
+  {
+    name: 'find',
+    description: 'Run a find() query on a Mongoose model and return an array of documents. Times out after 5 seconds.'
+  },
+  {
+    name: 'findOne',
+    description: 'Run a findOne() query on a Mongoose model and return a single document or null. Times out after 5 seconds.'
+  },
+  {
+    name: 'typeCheck',
+    description: 'Type-checks the script and validates syntax using TypeScript and node --check'
+  }
+];

package/backend/chatAgent/getAgentTools.js

@@ -0,0 +1,140 @@
+'use strict';
+
+const vm = require('vm');
+const ts = require('typescript');
+const { tool } = require('ai');
+const { jsonSchema } = require('@ai-sdk/provider-utils');
+const agentToolMetadata = require('./agentToolMetadata');
+
+module.exports = function getAgentTools(db) {
+  const modelNames = Object.keys(db.models).filter(name => !name.startsWith('__Studio'));
+  const toolDescriptions = Object.fromEntries(agentToolMetadata.map(tool => [tool.name, tool.description]));
+
+  return {
+    estimatedDocumentCount: tool({
+      description: toolDescriptions.estimatedDocumentCount + ' Takes a model name.',
+      inputSchema: jsonSchema({
+        type: 'object',
+        properties: {
+          modelName: { type: 'string', description: 'The Mongoose model name' }
+        },
+        required: ['modelName']
+      }),
+      execute: async({ modelName }) => {
+        console.log(`estimatedDocumentCount: modelName=${modelName}`);
+        const Model = db.models[modelName];
+        if (Model == null) {
+          return { error: `Model ${modelName} not found. Available models: ${modelNames.join(', ')}` };
+        }
+        const count = await Model.estimatedDocumentCount().maxTimeMS(5000);
+        return { modelName, count };
+      }
+    }),
+    find: tool({
+      description: toolDescriptions.find,
+      inputSchema: jsonSchema({
+        type: 'object',
+        properties: {
+          modelName: { type: 'string', description: 'The Mongoose model name' },
+          filter: { type: 'object', description: 'MongoDB query filter', default: {} },
+          limit: { type: 'number', description: 'Maximum number of documents to return', default: 10 }
+        },
+        required: ['modelName']
+      }),
+      execute: async({ modelName, filter = {}, limit = 10 }) => {
+        console.log(`find: modelName=${modelName}, filter=${JSON.stringify(filter)}, limit=${limit}`);
+        const Model = db.models[modelName];
+        if (Model == null) {
+          return { error: `Model ${modelName} not found. Available models: ${modelNames.join(', ')}` };
+        }
+        const docs = await Model.find(filter).limit(limit).maxTimeMS(5000);
+        return { modelName, count: docs.length, docs };
+      }
+    }),
+    findOne: tool({
+      description: toolDescriptions.findOne,
+      inputSchema: jsonSchema({
+        type: 'object',
+        properties: {
+          modelName: { type: 'string', description: 'The Mongoose model name' },
+          filter: { type: 'object', description: 'MongoDB query filter', default: {} }
+        },
+        required: ['modelName']
+      }),
+      execute: async({ modelName, filter = {} }) => {
+        console.log(`findOne: modelName=${modelName}, filter=${JSON.stringify(filter)}`);
+        const Model = db.models[modelName];
+        if (Model == null) {
+          return { error: `Model ${modelName} not found. Available models: ${modelNames.join(', ')}` };
+        }
+        const doc = await Model.findOne(filter).maxTimeMS(5000);
+        return { modelName, doc };
+      }
+    }),
+    typeCheck: tool({
+      description: toolDescriptions.typeCheck + ' The script will run in a sandbox with globals: db (mongoose.Connection), mongoose, ObjectId (mongoose.Types.ObjectId), console, and MongooseStudioChartColors (string[]). Pass the raw script body (no imports, no wrapping function). Returns any TypeScript errors or JavaScript syntax errors found. Remember that you should write JavaScript, NOT TypeScript. This tool is just to check for obvious errors.',
+      inputSchema: jsonSchema({
+        type: 'object',
+        properties: {
+          script: { type: 'string', description: 'The script to type-check' }
+        },
+        required: ['script']
+      }),
+      execute: async({ script }) => {
+        const wrapped = wrapScriptForTypeCheck(script);
+        const fileName = '__check.ts';
+        const compilerOptions = {
+          noEmit: true,
+          strict: false,
+          target: ts.ScriptTarget.ESNext,
+          module: ts.ModuleKind.ESNext,
+          moduleResolution: ts.ModuleResolutionKind.Node10,
+          esModuleInterop: true,
+          baseUrl: require('path').resolve(__dirname, '..', '..'),
+          paths: { mongoose: ['./node_modules/mongoose'] }
+        };
+        const host = ts.createCompilerHost(compilerOptions);
+        const originalGetSourceFile = host.getSourceFile;
+        host.getSourceFile = (name, languageVersion) => {
+          if (name === fileName) {
+            return ts.createSourceFile(name, wrapped, languageVersion);
+          }
+          return originalGetSourceFile.call(host, name, languageVersion);
+        };
+        const program = ts.createProgram([fileName], compilerOptions, host);
+        const diagnostics = ts.getPreEmitDiagnostics(program)
+          .filter(d => d.file?.fileName === fileName);
+
+        const errors = diagnostics.map(d => ts.flattenDiagnosticMessageText(d.messageText, '\n'));
+
+        // Also check that the script is valid JavaScript (not TypeScript) via vm.Script
+        try {
+          new vm.Script('(async () => {\n' + script + '\n})()');
+        } catch (e) {
+          errors.push('JavaScript syntax error: ' + e.message);
+        }
+
+        if (errors.length === 0) {
+          return { ok: true };
+        }
+
+        console.log('Errors', errors);
+
+        return { ok: false, errors };
+      }
+    })
+  };
+};
+
+const wrapScriptForTypeCheck = (script) => `
+import mongoose from 'mongoose';
+declare const db: Omit<mongoose.Connection, 'model'> & {
+  model(name: string): mongoose.Model<any>;
+};
+declare const ObjectId: typeof mongoose.Types.ObjectId;
+declare const console: Console;
+declare const MongooseStudioChartColors: string[];
+async function __script() {
+${script}
+}
+`;

package/backend/db/chatMessageSchema.js

@@ -20,6 +20,11 @@ const chatMessageSchema = new mongoose.Schema({
     type: String,
     required: false
   },
+  toolCalls: [{
+    toolName: String,
+    input: mongoose.Schema.Types.Mixed,
+    status: String
+  }],
   executionResult: {
     output: mongoose.Schema.Types.Mixed,
     logs: String,

package/backend/db/chatThreadSchema.js

@@ -18,6 +18,10 @@ const chatThreadSchema = new mongoose.Schema({
     type: mongoose.ObjectId,
     ref: 'Workspace'
   },
+  agentMode: {
+    type: Boolean,
+    default: false
+  },
   sharingOptions: {
     sharedWithWorkspace: {
       type: Boolean,

package/backend/integrations/callLLM.js

@@ -3,7 +3,7 @@
 const { createAnthropic } = require('@ai-sdk/anthropic');
 const { createGoogleGenerativeAI } = require('@ai-sdk/google');
 const { createOpenAI } = require('@ai-sdk/openai');
-const { generateText } = require('ai');
+const { generateText, stepCountIs } = require('ai');
 const { defaultMothershipURL } = require('../../constants');
 
 module.exports = async function callLLM(messages, system, options) {
@@ -46,7 +46,9 @@ module.exports = async function callLLM(messages, system, options) {
     return generateText({
       model: provider(model),
       system,
-      messages
+      messages,
+      tools: options?.tools,
+      stopWhen: options?.tools ? stepCountIs(10) : undefined
     });
   }
 

package/backend/integrations/streamLLM.js

@@ -3,7 +3,7 @@
 const { createAnthropic } = require('@ai-sdk/anthropic');
 const { createGoogleGenerativeAI } = require('@ai-sdk/google');
 const { createOpenAI } = require('@ai-sdk/openai');
-const { streamText } = require('ai');
+const { streamText, stepCountIs } = require('ai');
 const { defaultMothershipURL } = require('../../constants');
 
 module.exports = async function* streamLLM(messages, system, options) {
@@ -44,16 +44,24 @@ module.exports = async function* streamLLM(messages, system, options) {
 
   if (provider) {
     let error = null;
-    const { textStream } = streamText({
+    const { fullStream } = streamText({
       model: provider(model),
       system,
       messages,
+      tools: options?.tools,
+      stopWhen: options?.tools ? stepCountIs(10) : undefined,
       onError(err) {
         error = err.error;
       }
     });
-    for await (const chunk of textStream) {
-      yield chunk;
+    for await (const chunk of fullStream) {
+      if (chunk.type === 'text-delta') {
+        yield chunk.text;
+      } else if (chunk.type === 'tool-call') {
+        yield { toolCall: { toolName: chunk.toolName, input: chunk.input } };
+      } else if (chunk.type === 'tool-result') {
+        yield { toolResult: { toolName: chunk.toolName, output: chunk.output } };
+      }
     }
     if (error) {
       throw error;

package/eslint.config.js

@@ -15,7 +15,8 @@ module.exports = defineConfig([
         fetch: true,
         setTimeout: true,
         process: true,
-        console: true
+        console: true,
+        __dirname: true
       }
     },
     extends: [config]