Merge pull request #2111 from mito-ds/dev

Release Dec 17, 2025

Author: ngafar

.github/workflows/test-mitosheet-frontend.yml

@@ -266,4 +266,24 @@ jobs:
         cd tests
         source venv/bin/activate
         python dash-test.py &
-        npm run test:dash
+        npm run test:dash
+
+  test-mitosheet-jest:
+    runs-on: ubuntu-24.04
+    timeout-minutes: 30
+
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-node@v4
+      with:
+        node-version: 22
+        cache: 'npm'
+        cache-dependency-path: mitosheet/package-lock.json
+    - name: Install dependencies
+      working-directory: mitosheet
+      run: |
+        npm install
+    - name: Run Jest tests
+      working-directory: mitosheet
+      run: |
+        npm run test

conductor.json

@@ -0,0 +1,6 @@
+{
+    "scripts": {
+        "setup": "cd mito-ai && (test -d venv || python3 -m venv venv) && source venv/bin/activate && pip install -e \".[test]\" && touch yarn.lock && jlpm install && jlpm build && jupyter labextension develop . --overwrite && jupyter server extension enable --py mito_ai",
+        "run": "cd mito-ai && source venv/bin/activate && jlpm build && jupyter lab --browser=chrome"
+    }
+}

mito-ai/mito_ai/completions/prompt_builders/agent_system_message.py

@@ -3,6 +3,7 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CITATION_RULES,
+    CELL_REFERENCE_RULES,
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
@@ -291,7 +292,7 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 
 Important information:
 1. The message is a short summary of the ALL the work that you've completed on this task. It should not just refer to the final message. It could be something like "I've completed the sales strategy analysis by exploring key relationships in the data and summarizing creating a report with three recommendations to boost sales.""
-2. The message should include citations for any insights that you shared with the user.
+2. The message should include citations for any insights that you shared with the user and cell references for whenever you refer to specific cells that you've updated or created.
 3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Export the cleaned data to CSV", "Analyze revenue per customer", "Convert notebook into an app"].
 4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per age group"].
 5. If you are not sure what the user might want to do next, err on the side of suggesting next steps instead of making an assumption and using more CELL_UPDATES.
@@ -325,19 +326,29 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 
 RULES
 
-- You are working in a Jupyter Lab environment in a .ipynb file. 
+- You are working in a Jupyter Lab environment in a .ipynb file.
 - You can only respond with CELL_UPDATES or FINISHED_TASK responses.
 - In each message you send to the user, you can send one CellModification, one CellAddition, or one FINISHED_TASK response. BUT YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
 - After you send a CELL_UPDATE, the user will send you a message with the updated variables, code, and files in the current directory. You will use this information to decide what to do next, so it is critical that you wait for the user's response after each CELL_UPDATE before deciding your next action.
 - When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
-- When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
 - When writing the message, do not explain to the user how to use the CELL_UPDATE or FINISHED_TASK response, they will already know how to use them. Just provide a summary of your thought process. Do not reference any Cell IDs in the message.
 - When writing the message, do not include leading words like "Explanation:" or "Thought process:". Just provide a summary of your thought process.
 - When writing the message, use tickmarks when referencing specific variable names. For example, write `sales_df` instead of "sales_df" or just sales_df.
 
+====
+
+CODE STYLE
+
+- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+
 ==== 
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Citation Example>
 
 ### User Message 1:
@@ -471,5 +482,4 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 ====
 
 OTHER USEFUL INFORMATION:
-1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
-2. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""
+1. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""

mito-ai/mito_ai/completions/prompt_builders/chat_system_message.py

@@ -3,7 +3,8 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CHAT_CODE_FORMATTING_RULES,
-    CITATION_RULES, 
+    CITATION_RULES,
+    CELL_REFERENCE_RULES,
     ACTIVE_CELL_ID_SECTION_HEADING, 
     CODE_SECTION_HEADING,
     get_database_rules
@@ -28,6 +29,9 @@ def create_chat_system_message_prompt() -> str:
 ====
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Example 1> 
 {ACTIVE_CELL_ID_SECTION_HEADING}
 '7b3a9e2c-5d14-4c83-b2f9-d67891e4a5f2'
@@ -79,6 +83,18 @@ def create_chat_system_message_prompt() -> str:
 ==== 
 {CHAT_CODE_FORMATTING_RULES}
 
+====
+
+CODE STYLE
+
+- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- Write code that preserves the intent of the original code shared with you and the task to complete.
+- Make the solution as simple as possible.
+- Do not add temporary comments like '# Fixed the typo here' or '# Added this line to fix the error'
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+
+====
+
 IMPORTANT RULES:
 - Do not recreate variables that already exist
 - Keep as much of the original code as possible
@@ -87,6 +103,5 @@ def create_chat_system_message_prompt() -> str:
 - Write code that preserves the intent of the original code shared with you and the task to complete.
 - Make the solution as simple as possible.
 - Reuse as much of the existing code as possible.
-- Do not add temporary comments like '# Fixed the typo here' or '# Added this line to fix the error'
 - Whenever writing Python code, it should be a python code block starting with ```python and ending with ```
 """

mito-ai/mito_ai/completions/prompt_builders/prompt_constants.py

@@ -46,6 +46,28 @@
 8. Do not include the citation in the code block as a comment. ONLY include the citation in the message field of your response.
 """
 
+CELL_REFERENCE_RULES = """RULES FOR REFERENCING CELLS
+
+When referring to specific cells in the notebook in your messages, use cell references so the user can easily navigate to the cell you're talking about. The user sees cells numbered as "Cell 1", "Cell 2", etc., but internally cells are identified by their unique IDs.
+
+To reference a cell, use this format inline in your message:
+[MITO_CELL_REF:cell_id]
+
+This will be displayed to the user as a clickable "Cell N" link that navigates to the referenced cell.
+
+Cell Reference Rules:
+
+1. Use cell references when discussing specific cells you've created or modified (e.g., "I've added the data cleaning code in [MITO_CELL_REF:abc123]").
+2. Use cell references when referring to cells the user mentioned or that contain relevant context.
+3. The cell_id must be an actual cell ID from the notebook - do not make up IDs.
+4. Place the reference inline where it makes sense in your message, similar to how you would write "Cell 3" in natural language.
+5. Do not use cell references in code - only in the message field of your responses.
+6. Cell references are different from citations. Use citations for specific line-level insights; use cell references for general cell-level navigation.
+
+Example:
+"I've loaded the sales data in [MITO_CELL_REF:c68fdf19-db8c-46dd-926f-d90ad35bb3bc] and will now calculate the monthly totals."
+"""
+
 def get_active_cell_output_str(has_active_cell_output: bool) -> str:
     """
     Used to tell the AI about the output of the active code cell. 

mito-ai/mito_ai/completions/prompt_builders/utils.py

@@ -39,6 +39,7 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
     selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
     selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
     selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
+    selected_cells = [context["value"] for context in additional_context if context.get("type") == "cell"]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -66,6 +67,12 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
             "The following images have been selected by the user to be used in the task:\n"
             + "\n".join(selected_images)
         )
+        
+    if len(selected_cells) > 0:
+        context_parts.append(
+            "The following cells have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_cells)
+        )
 
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)

mito-ai/mito_ai/utils/open_ai_utils.py

@@ -171,6 +171,9 @@ def get_open_ai_completion_function_params(
         "stream": stream,
         "messages": messages,
     }
+
+    if model == "gpt-5.2":
+        completion_function_params["reasoning_effort"] = "low"
 
     # If a response format is provided, we need to convert it to a json schema.
     # Pydantic models are supported by the OpenAI API, however, we need to be able to 

mito-ai/package.json

@@ -138,7 +138,8 @@
   "jupyterlab": {
     "extension": true,
     "outputDir": "mito_ai/labextension",
-    "schemaDir": "schema"
+    "schemaDir": "schema",
+    "themePath": "style/theme/theme.css"
   },
   "eslintIgnore": [
     "node_modules",

mito-ai/src/Extensions/AiChat/ChatMessage/AssistantCodeBlock.tsx

@@ -13,6 +13,9 @@ import { CodeReviewStatus } from '../ChatTaskpane';
 import CodeIcon from '../../../icons/CodeIcon';
 import CodeBlockToolbar from './CodeBlockToolbar';
 import AgentComponentHeader from '../../../components/AgentComponents/AgentComponentHeader';
+import GoToCellIcon from '../../../icons/GoToCellIcon';
+import { NotebookPanel } from '@jupyterlab/notebook';
+import { scrollToCell, setActiveCellByIDInNotebookPanel } from '../../../utils/notebook';
 
 interface IAssistantCodeBlockProps {
     code: string;
@@ -26,6 +29,8 @@ interface IAssistantCodeBlockProps {
     codeReviewStatus: CodeReviewStatus;
     agentModeEnabled: boolean;
     isErrorFixup?: boolean;
+    cellId?: string;
+    notebookPanel?: NotebookPanel | null;
 }
 
 const AssistantCodeBlock: React.FC<IAssistantCodeBlockProps> = ({
@@ -40,10 +45,35 @@ const AssistantCodeBlock: React.FC<IAssistantCodeBlockProps> = ({
     codeReviewStatus,
     agentModeEnabled,
     isErrorFixup,
+    cellId,
+    notebookPanel,
 }) => {
     const [isCodeExpanded, setIsCodeExpanded] = useState(false);
     const shouldShowToolbar = isLastAiMessage || isCodeComplete;
 
+    const handleGoToCell = (e: React.MouseEvent): void => {
+        e.stopPropagation();
+        if (cellId && notebookPanel) {
+            setActiveCellByIDInNotebookPanel(notebookPanel, cellId);
+            scrollToCell(notebookPanel, cellId, undefined, 'center');
+        }
+    };
+
+    const renderActionButtons = (): React.ReactNode => {
+        if (!cellId || !notebookPanel) {
+            return null;
+        }
+        return (
+            <button
+                className="agent-component-header-action-button"
+                onClick={handleGoToCell}
+                title="Go to cell"
+            >
+                <GoToCellIcon />
+            </button>
+        );
+    };
+
     if (agentModeEnabled) {
 
         // Handle regular code blocks
@@ -58,6 +88,7 @@ const AssistantCodeBlock: React.FC<IAssistantCodeBlockProps> = ({
                     isExpanded={isCodeExpanded}
                     displayBorder={!isErrorFixup}
                     className={isErrorFixup ? 'error-fixup' : undefined}
+                    actionButtons={renderActionButtons()}
                 />
                 {isCodeExpanded && (
                     <PythonCode

mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdown.tsx

@@ -6,7 +6,9 @@
 import React, { useState, useEffect, useRef } from 'react';
 import { ExpandedVariable } from './ChatInput';
 import { getDatabaseConnections, getRules } from '../../../restAPI/RestAPI';
-import { VariableDropdownItem, FileDropdownItem, RuleDropdownItem } from './ChatDropdownItems';
+import { VariableDropdownItem, FileDropdownItem, RuleDropdownItem, CellDropdownItem } from './ChatDropdownItems';
+import { INotebookTracker } from '@jupyterlab/notebook';
+import { getAllCellReferences } from '../../../utils/cellReferences';
 
 interface ChatDropdownProps {
     options: ExpandedVariable[];
@@ -16,6 +18,7 @@ interface ChatDropdownProps {
     isDropdownFromButton?: boolean;
     onFilterChange?: (filterText: string) => void;
     onClose?: () => void;
+    notebookTracker?: INotebookTracker;
 }
 
 interface ChatDropdownVariableOption {
@@ -38,11 +41,19 @@ interface ChatDropdownFileOption {
     file: ExpandedVariable;
 }
 
+interface ChatDropdownCellOption {
+    type: 'cell'
+    cellNumber: number;
+    cellId: string;
+    cellType: string;
+}
+
 export type ChatDropdownOption =
     | ChatDropdownVariableOption
     | ChatDropdownRuleOption
     | ChatDropdownFileOption
-    | ChatDropdownDatabaseOption;
+    | ChatDropdownDatabaseOption
+    | ChatDropdownCellOption;
 
 const priortizeByType = (options: ChatDropdownOption[], maxPerType: number): ChatDropdownOption[] => {
     /* 
@@ -76,6 +87,7 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
     isDropdownFromButton = false,
     onFilterChange,
     onClose,
+    notebookTracker,
 }) => {
     const [selectedIndex, setSelectedIndex] = useState(0);
     const [localFilterText, setLocalFilterText] = useState(filterText);
@@ -107,16 +119,29 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
     // Use local filter text when search input is shown, otherwise use prop
     const effectiveFilterText = isDropdownFromButton ? localFilterText : filterText;
 
+    // Get cell references if notebook tracker is available
+    const cellReferences = notebookTracker?.currentWidget 
+        ? getAllCellReferences(notebookTracker.currentWidget)
+        : [];
+
     // Create a list of all options with the format 
     // ['type': 'variable', "expandedVariable": variable]
     // ['type': 'rule', "rule": rule]
     // ['type': 'file', "file": file]
+    // ['type': 'cell', "cellNumber": number, "cellId": string]
     const allOptions: ChatDropdownOption[] = [
         // Rules first
         ...rules.map((rule): ChatDropdownRuleOption => ({
             type: 'rule',
             rule: rule
         })),
+        // Cells second (when user types @Cell or @cell)
+        ...cellReferences.map((cell): ChatDropdownCellOption => ({
+            type: 'cell',
+            cellNumber: cell.cellNumber,
+            cellId: cell.cellId,
+            cellType: cell.cellType
+        })),
         // Files second
         ...options
             .filter(variable => variable.file_name) // Only files
@@ -167,6 +192,12 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
         } else if (option.type === 'db') {
             return option.variable.variable_name.toLowerCase().includes(effectiveFilterText.toLowerCase()) ||
                 option.variable.value.toLowerCase().includes(effectiveFilterText.toLowerCase());
+        } else if (option.type === 'cell') {
+            // Match "CellN" (no space)
+            const cellText = `cell${option.cellNumber}`;
+            const numberText = String(option.cellNumber);
+            return cellText.includes(effectiveFilterText.toLowerCase()) ||
+                numberText.includes(effectiveFilterText.toLowerCase());
         } else {
             return option.rule.toLowerCase().includes(effectiveFilterText.toLowerCase());
         }
@@ -331,6 +362,20 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
                                 />
                             );
                         }
+                        case 'cell': {
+                            const uniqueKey = `cell-${option.cellNumber}`;
+                            return (
+                                <CellDropdownItem
+                                    key={uniqueKey}
+                                    cellNumber={option.cellNumber}
+                                    cellId={option.cellId}
+                                    cellType={option.cellType}
+                                    index={index}
+                                    selectedIndex={selectedIndex}
+                                    onSelect={() => onSelect(option)}
+                                />
+                            );
+                        }
                         default:
                             return null;
                     }