notebook-intelligence 2.4.2__tar.gz → 2.5.0__tar.gz

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/.gitignore

@@ -83,6 +83,14 @@ target/
 # pyenv
 .python-version
 
+# Virtual environments
+venv/
+env/
+ENV/
+.venv/
+.env/
+venv*/
+.vscode/
 # celery beat schedule file
 celerybeat-schedule
 

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: notebook_intelligence
-Version: 2.4.2
+Version: 2.5.0
 Summary: AI coding assistant for JupyterLab
 Project-URL: Homepage, https://github.com/notebook-intelligence/notebook-intelligence
 Project-URL: Bug Tracker, https://github.com/notebook-intelligence/notebook-intelligence/issues
@@ -801,18 +801,18 @@ jupyter lab --NotebookIntelligence.notebook_execute_tool=env_enabled
 
 ### Configuration files
 
-NBI saves configuration at `~/.jupyter/nbi-config.json`. It also supports environment wide base configuration at `<env-prefix>/share/jupyter/nbi-config.json`. Organizations can ship default configuration at this environment wide config path. User's changes will be stored as overrides at `~/.jupyter/nbi-config.json`.
+NBI saves configuration at `~/.jupyter/nbi/config.json`. It also supports environment wide base configuration at `<env-prefix>/share/jupyter/nbi/config.json`. Organizations can ship default configuration at this environment wide config path. User's changes will be stored as overrides at `~/.jupyter/nbi/config.json`.
 
 These config files are used for saving LLM provider, model and MCP configuration. Note that API keys you enter for your custom LLM providers will also be stored in these config files.
 
 > [!IMPORTANT]
-> Note that updating nbi-config.json manually requires restarting JupyterLab to take effect.
+> Note that updating config.json manually requires restarting JupyterLab to take effect.
 
 ### Model Context Protocol ([MCP](https://modelcontextprotocol.io)) Support
 
 NBI seamlessly integrates with MCP servers. It supports servers with both Standard Input/Output (stdio) and Server-Sent Events (SSE) transports. The MCP support is limited to server tools at the moment.
 
-You can easily add MCP servers to NBI by editing the configuration file [nbi-config.json](#configuration-files). Simply add a key "mcp" and "mcpServers" under it as shown below.
+You can easily add MCP servers to NBI by editing the configuration file [~/.jupyter/nbi/mcp.json](#configuration-files). Environment wide base configuration is also support using the file at `<env-prefix>/share/jupyter/nbi/mcp.json`.
 
 > [!NOTE]
 > Using MCP servers requires an LLM model with tool calling capabilities. All of the GitHub Copilot models provided in NBI support this feature. If you are using other providers make sure you choose a tool calling capable model.
@@ -820,47 +820,24 @@ You can easily add MCP servers to NBI by editing the configuration file [nbi-con
 > [!CAUTION]
 > Note that most MCP servers are run on the same computer as your JupyterLab installation and they can make irreversible changes to your computer and/or access private data. Make sure that you only install MCP servers from trusted sources.
 
+### MCP Config file example
+
 ```json
 {
-    "chat_model": {
-        ...
-    },
-    ...<other configuration>,
-
-    "mcp": {
-        "mcpServers": {
-            "filesystem": {
-                "command": "npx",
-                "args": [
-                    "-y",
-                    "@modelcontextprotocol/server-filesystem",
-                    "/Users/mbektas/mcp-test"
-                ]
-            },
-        }
+  "mcpServers": {
+    "filesystem": {
+      "command": "npx",
+      "args": [
+        "-y",
+        "@modelcontextprotocol/server-filesystem",
+        "/Users/mbektas/mcp-test"
+      ]
     }
+  }
 }
 ```
 
-This will automatically create a new chat participant in NBI and you can access it by starting your prompts with `@mcp`. Use `@mcp /info` prompt to get information on the tools provided by the MCP servers you configured. This chat participant will have access all the tools provided by the servers you configure.
-
-<img src="media/mcp-prompt.png" alt="Settings dialog" width=600 />
-
-By default, each tool call to MCP servers will require approval. If you would like to auto approve tools, you can do so by using the `"alwaysAllow": []` configuration key in the nbi-config.json. Simply list the names of tools.
-
-```json
-"mcpServers": {
-    "filesystem": {
-        "command": "npx",
-        "args": [
-            "-y",
-            "@modelcontextprotocol/server-filesystem",
-            "/Users/mbektas/mcp-test"
-        ],
-        "alwaysAllow": ["list_allowed_directories", "list_directory"]
-    },
-}
-```
+You can use Agent mode to access tools provided by MCP servers you configured.
 
 For servers with stdio transport, you can also set additional environment variables by using the `env` key. Environment variables are specified as key value pairs.
 
@@ -876,12 +853,12 @@ For servers with stdio transport, you can also set additional environment variab
 }
 ```
 
-Below is an example of a server configuration with SSE transport. For SSE transport servers, you can also specify headers to be sent as part of the requests.
+Below is an example of a server configuration with Streamable HTTP transport. For Streamable HTTP transport servers, you can also specify headers to be sent as part of the requests.
 
 ```json
 "mcpServers": {
     "remoterservername": {
-        "url": "http://127.0.0.1:8080/sse",
+        "url": "http://127.0.0.1:8080/mcp",
         "headers": {
             "Authorization": "Bearer mysecrettoken"
         }
@@ -905,77 +882,6 @@ If you have multiple servers configured but you would like to disable some for a
 }
 ```
 
-#### Grouping MCP servers
-
-When you integrate multiple MCP servers to NBI, all of their tools will be available under the same chat participant `@mcp`. However, this may not be ideal in many situations. You may want to group certain servers and their tools based on their functionality. NBI lets you do that easily by configuring MCP chat participants. You can list the servers for each custom participant. If there are any unassigned MCP servers, then they will be used the default `@mcp` chat participant.
-
-Below is an example of creating a custom MCP participant. This configuration results in two chat participants `@mcp-fs` with `filesytem` MC server tools and `@mcp` with `servername1` and `servername1` MCP server tools.
-
-```json
-{
-    "chat_model": {
-        ...
-    },
-    ...<other configuration>,
-
-    "mcp": {
-        "mcpServers": {
-            "filesystem": {
-                "command": "npx",
-                "args": [
-                    "-y",
-                    "@modelcontextprotocol/server-filesystem",
-                    "/Users/mbektas/mcp-test"
-                ]
-            },
-            "servername1": {
-                "command": "",
-                "args": [],
-            },
-            "servername2": {
-                "command": "",
-                "args": [],
-                "disabled": true
-            }
-        },
-        "participants": {
-            "fs": {
-                "name": "MCP - File system",
-                "servers": ["filesystem"]
-            }
-        }
-    }
-}
-```
-
-#### Using NBI tools within MCP chat participants
-
-NBI allows you to access built-in tools from an MCP participant. You can do that by adding the list of built in NBI tools to your MCP participant configuration. The built-in tools available to MCP are `create_new_notebook`, `add_markdown_cell_to_notebook`, `add_code_cell_to_notebook`. Below is an example that integrates all these tools to MCP participant `@mcp-fs`.
-
-```json
-"participants": {
-    "fs": {
-        "name": "MCP - File system",
-        "servers": ["filesystem"],
-        "nbiTools": [
-            "create_new_notebook",
-            "add_markdown_cell_to_notebook",
-            "add_code_cell_to_notebook"
-        ]
-    }
-}
-```
-
-This chat participant will allow you to run example prompts like below.
-
-```
-@mcp-fs list the directories I have access to.
-```
-
-```
-@mcp-fs add a code cell which demonstrates ipywidgets Button to this notebook.
-```
-
 ### Developer documentation
 
 For building locally and contributing see the [developer documentatation](CONTRIBUTING.md).

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/README.md

@@ -94,18 +94,18 @@ jupyter lab --NotebookIntelligence.notebook_execute_tool=env_enabled
 
 ### Configuration files
 
-NBI saves configuration at `~/.jupyter/nbi-config.json`. It also supports environment wide base configuration at `<env-prefix>/share/jupyter/nbi-config.json`. Organizations can ship default configuration at this environment wide config path. User's changes will be stored as overrides at `~/.jupyter/nbi-config.json`.
+NBI saves configuration at `~/.jupyter/nbi/config.json`. It also supports environment wide base configuration at `<env-prefix>/share/jupyter/nbi/config.json`. Organizations can ship default configuration at this environment wide config path. User's changes will be stored as overrides at `~/.jupyter/nbi/config.json`.
 
 These config files are used for saving LLM provider, model and MCP configuration. Note that API keys you enter for your custom LLM providers will also be stored in these config files.
 
 > [!IMPORTANT]
-> Note that updating nbi-config.json manually requires restarting JupyterLab to take effect.
+> Note that updating config.json manually requires restarting JupyterLab to take effect.
 
 ### Model Context Protocol ([MCP](https://modelcontextprotocol.io)) Support
 
 NBI seamlessly integrates with MCP servers. It supports servers with both Standard Input/Output (stdio) and Server-Sent Events (SSE) transports. The MCP support is limited to server tools at the moment.
 
-You can easily add MCP servers to NBI by editing the configuration file [nbi-config.json](#configuration-files). Simply add a key "mcp" and "mcpServers" under it as shown below.
+You can easily add MCP servers to NBI by editing the configuration file [~/.jupyter/nbi/mcp.json](#configuration-files). Environment wide base configuration is also support using the file at `<env-prefix>/share/jupyter/nbi/mcp.json`.
 
 > [!NOTE]
 > Using MCP servers requires an LLM model with tool calling capabilities. All of the GitHub Copilot models provided in NBI support this feature. If you are using other providers make sure you choose a tool calling capable model.
@@ -113,47 +113,24 @@ You can easily add MCP servers to NBI by editing the configuration file [nbi-con
 > [!CAUTION]
 > Note that most MCP servers are run on the same computer as your JupyterLab installation and they can make irreversible changes to your computer and/or access private data. Make sure that you only install MCP servers from trusted sources.
 
+### MCP Config file example
+
 ```json
 {
-    "chat_model": {
-        ...
-    },
-    ...<other configuration>,
-
-    "mcp": {
-        "mcpServers": {
-            "filesystem": {
-                "command": "npx",
-                "args": [
-                    "-y",
-                    "@modelcontextprotocol/server-filesystem",
-                    "/Users/mbektas/mcp-test"
-                ]
-            },
-        }
+  "mcpServers": {
+    "filesystem": {
+      "command": "npx",
+      "args": [
+        "-y",
+        "@modelcontextprotocol/server-filesystem",
+        "/Users/mbektas/mcp-test"
+      ]
     }
+  }
 }
 ```
 
-This will automatically create a new chat participant in NBI and you can access it by starting your prompts with `@mcp`. Use `@mcp /info` prompt to get information on the tools provided by the MCP servers you configured. This chat participant will have access all the tools provided by the servers you configure.
-
-<img src="media/mcp-prompt.png" alt="Settings dialog" width=600 />
-
-By default, each tool call to MCP servers will require approval. If you would like to auto approve tools, you can do so by using the `"alwaysAllow": []` configuration key in the nbi-config.json. Simply list the names of tools.
-
-```json
-"mcpServers": {
-    "filesystem": {
-        "command": "npx",
-        "args": [
-            "-y",
-            "@modelcontextprotocol/server-filesystem",
-            "/Users/mbektas/mcp-test"
-        ],
-        "alwaysAllow": ["list_allowed_directories", "list_directory"]
-    },
-}
-```
+You can use Agent mode to access tools provided by MCP servers you configured.
 
 For servers with stdio transport, you can also set additional environment variables by using the `env` key. Environment variables are specified as key value pairs.
 
@@ -169,12 +146,12 @@ For servers with stdio transport, you can also set additional environment variab
 }
 ```
 
-Below is an example of a server configuration with SSE transport. For SSE transport servers, you can also specify headers to be sent as part of the requests.
+Below is an example of a server configuration with Streamable HTTP transport. For Streamable HTTP transport servers, you can also specify headers to be sent as part of the requests.
 
 ```json
 "mcpServers": {
     "remoterservername": {
-        "url": "http://127.0.0.1:8080/sse",
+        "url": "http://127.0.0.1:8080/mcp",
         "headers": {
             "Authorization": "Bearer mysecrettoken"
         }
@@ -198,77 +175,6 @@ If you have multiple servers configured but you would like to disable some for a
 }
 ```
 
-#### Grouping MCP servers
-
-When you integrate multiple MCP servers to NBI, all of their tools will be available under the same chat participant `@mcp`. However, this may not be ideal in many situations. You may want to group certain servers and their tools based on their functionality. NBI lets you do that easily by configuring MCP chat participants. You can list the servers for each custom participant. If there are any unassigned MCP servers, then they will be used the default `@mcp` chat participant.
-
-Below is an example of creating a custom MCP participant. This configuration results in two chat participants `@mcp-fs` with `filesytem` MC server tools and `@mcp` with `servername1` and `servername1` MCP server tools.
-
-```json
-{
-    "chat_model": {
-        ...
-    },
-    ...<other configuration>,
-
-    "mcp": {
-        "mcpServers": {
-            "filesystem": {
-                "command": "npx",
-                "args": [
-                    "-y",
-                    "@modelcontextprotocol/server-filesystem",
-                    "/Users/mbektas/mcp-test"
-                ]
-            },
-            "servername1": {
-                "command": "",
-                "args": [],
-            },
-            "servername2": {
-                "command": "",
-                "args": [],
-                "disabled": true
-            }
-        },
-        "participants": {
-            "fs": {
-                "name": "MCP - File system",
-                "servers": ["filesystem"]
-            }
-        }
-    }
-}
-```
-
-#### Using NBI tools within MCP chat participants
-
-NBI allows you to access built-in tools from an MCP participant. You can do that by adding the list of built in NBI tools to your MCP participant configuration. The built-in tools available to MCP are `create_new_notebook`, `add_markdown_cell_to_notebook`, `add_code_cell_to_notebook`. Below is an example that integrates all these tools to MCP participant `@mcp-fs`.
-
-```json
-"participants": {
-    "fs": {
-        "name": "MCP - File system",
-        "servers": ["filesystem"],
-        "nbiTools": [
-            "create_new_notebook",
-            "add_markdown_cell_to_notebook",
-            "add_code_cell_to_notebook"
-        ]
-    }
-}
-```
-
-This chat participant will allow you to run example prompts like below.
-
-```
-@mcp-fs list the directories I have access to.
-```
-
-```
-@mcp-fs add a code cell which demonstrates ipywidgets Button to this notebook.
-```
-
 ### Developer documentation
 
 For building locally and contributing see the [developer documentatation](CONTRIBUTING.md).

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '2.4.2'
+__version__ = VERSION = '2.5.0'

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/ai_service_manager.py

@@ -16,6 +16,7 @@ from notebook_intelligence.llm_providers.litellm_compatible_llm_provider import
 from notebook_intelligence.llm_providers.ollama_llm_provider import OllamaLLMProvider
 from notebook_intelligence.llm_providers.openai_compatible_llm_provider import OpenAICompatibleLLMProvider
 from notebook_intelligence.mcp_manager import MCPManager
+from notebook_intelligence.util import ThreadSafeWebSocketConnector
 
 log = logging.getLogger(__name__)
 
@@ -40,6 +41,7 @@ class AIServiceManager(Host):
         self._litellm_compatible_llm_provider = LiteLLMCompatibleLLMProvider()
         self._ollama_llm_provider = OllamaLLMProvider()
         self._extensions = []
+        self._websocket_connector: ThreadSafeWebSocketConnector = None
         self.initialize()
 
     @property
@@ -50,6 +52,15 @@ class AIServiceManager(Host):
     def ollama_llm_provider(self) -> OllamaLLMProvider:
         return self._ollama_llm_provider
 
+    @property
+    def websocket_connector(self) -> ThreadSafeWebSocketConnector:
+        return self._websocket_connector
+    
+    @websocket_connector.setter
+    def websocket_connector(self, _websocket_connector: ThreadSafeWebSocketConnector):
+        self._websocket_connector = _websocket_connector
+        self._mcp_manager.websocket_connector = _websocket_connector
+
     def initialize(self):
         self.chat_participants = {}
         self.register_llm_provider(GitHubCopilotLLMProvider())
@@ -370,3 +381,15 @@ class AIServiceManager(Host):
             if extension.id == extension_id:
                 return extension
         return None
+
+    def handle_stop_request(self):
+        self._mcp_manager.handle_stop_request()
+
+    def update_mcp_server_connections(self, disabled_mcp_servers: list[str]):
+        self._mcp_manager.update_mcp_server_connections(disabled_mcp_servers)
+
+    def connect_mcp_server(self, server_name: str):
+        self._mcp_manager.connect_mcp_server(server_name)
+    
+    def disconnect_mcp_server(self, server_name: str):
+        self._mcp_manager.disconnect_mcp_server(server_name)

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/api.py

@@ -1,7 +1,6 @@
 # Copyright (c) Mehmet Bektas <mbektasgh@outlook.com>
 
 import asyncio
-import json
 from typing import Any, Callable, Dict, Union
 from dataclasses import asdict, dataclass
 from enum import Enum
@@ -29,6 +28,7 @@ class BackendMessageType(str, Enum):
     StreamEnd = 'stream-end'
     RunUICommand = 'run-ui-command'
     GitHubCopilotLoginStatusChange = 'github-copilot-login-status-change'
+    MCPServerStatusChange = 'mcp-server-status-change'
 
 class ResponseStreamDataType(str, Enum):
     LLMRaw = 'llm-raw'
@@ -49,6 +49,15 @@ class BuiltinToolset(str, Enum):
     NotebookExecute = 'nbi-notebook-execute'
     PythonFileEdit = 'nbi-python-file-edit'
 
+class MCPServerStatus(str, Enum):
+    NotConnected = 'not-connected'
+    Connecting = 'connecting'
+    Disconnecting = 'disconnecting'
+    FailedToConnect = 'failed-to-connect'
+    Connected = 'connected'
+    UpdatingToolList = 'updating-tool-list'
+    UpdatedToolList = 'updated-tool-list'
+
 class Signal:
     def __init__(self):
         self._listeners = []
@@ -376,14 +385,18 @@ class MCPServer:
     @property
     def name(self) -> str:
         return NotImplemented
+
+    @property
+    def status(self) -> MCPServerStatus:
+        return NotImplemented
     
-    async def connect(self):
+    def connect(self):
         return NotImplemented
 
-    async def disconnect(self):
+    def disconnect(self):
         return NotImplemented
 
-    async def update_tool_list(self):
+    def update_tool_list(self):
         return NotImplemented
     
     def get_tools(self) -> list[Tool]:
@@ -392,7 +405,7 @@ class MCPServer:
     def get_tool(self, tool_name: str) -> Tool:
         return NotImplemented
 
-    async def call_tool(self, tool_name: str, tool_args: dict):
+    def call_tool(self, tool_name: str, tool_args: dict):
         return NotImplemented
 
 def auto_approve(tool: SimpleTool):

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/config.py

@@ -120,3 +120,22 @@ class NBIConfig:
     def using_github_copilot_service(self) -> bool:
         return self.chat_model.get("provider") == 'github-copilot' or \
             self.inline_completion_model.get("provider") == 'github-copilot'
+
+    @property
+    def mcp_server_settings(self):
+        return self.get('mcp_server_settings', {})
+
+    @property
+    def rules_enabled(self) -> bool:
+        """Check if the ruleset system is enabled."""
+        return self.get('rules_enabled', True)
+
+    @property
+    def rules_directory(self) -> str:
+        """Get the rules directory path."""
+        return os.path.join(self.nbi_user_dir, 'rules')
+
+    @property
+    def active_rules(self) -> dict:
+        """Get dictionary of active rule states (filename -> bool)."""
+        return self.get('active_rules', {})

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/extension.py

@@ -27,6 +27,7 @@ from notebook_intelligence.util import ThreadSafeWebSocketConnector
 ai_service_manager: AIServiceManager = None
 log = logging.getLogger(__name__)
 tiktoken_encoding = tiktoken.encoding_for_model('gpt-4o')
+thread_safe_websocket_connector: ThreadSafeWebSocketConnector = None
 
 class GetCapabilitiesHandler(APIHandler):
     notebook_execute_tool = 'enabled'
@@ -39,8 +40,7 @@ class GetCapabilitiesHandler(APIHandler):
         notebook_execute_tool_enabled = self.notebook_execute_tool == 'enabled' or (self.notebook_execute_tool == 'env_enabled' and os.getenv('NBI_NOTEBOOK_EXECUTE_TOOL', 'disabled') == 'enabled')
         allowed_builtin_toolsets = [{"id": toolset.id, "name": toolset.name} for toolset in built_in_toolsets.values() if toolset.id != BuiltinToolset.NotebookExecute or notebook_execute_tool_enabled]
         mcp_servers = ai_service_manager.get_mcp_servers()
-        mcp_server_tools = [{"id": mcp_server.name, "tools": [{"name": tool.name, "description": tool.description} for tool in mcp_server.get_tools()]} for mcp_server in mcp_servers]
-        mcp_server_tools = [tool for tool in mcp_server_tools if len(tool["tools"]) > 0]
+        mcp_server_tools = [{"id": mcp_server.name, "status": mcp_server.status, "tools": [{"name": tool.name, "description": tool.description} for tool in mcp_server.get_tools()]} for mcp_server in mcp_servers]
         # sort by server id
         mcp_server_tools.sort(key=lambda server: server["id"])
 
@@ -88,6 +88,7 @@ class GetCapabilitiesHandler(APIHandler):
                 "mcpServers": mcp_server_tools,
                 "extensions": extensions
             },
+            "mcp_server_settings": nbi_config.mcp_server_settings,
             "default_chat_mode": nbi_config.default_chat_mode
         }
         for participant_id in ai_service_manager.chat_participants:
@@ -105,7 +106,8 @@ class ConfigHandler(APIHandler):
     @tornado.web.authenticated
     def post(self):
         data = json.loads(self.request.body)
-        valid_keys = set(["default_chat_mode", "chat_model", "inline_completion_model", "store_github_access_token"])
+        valid_keys = set(["default_chat_mode", "chat_model", "inline_completion_model", "store_github_access_token", "mcp_server_settings"])
+        has_model_change = "chat_model" in data or "inline_completion_model" in data
         for key in data:
             if key in valid_keys:
                 ai_service_manager.nbi_config.set(key, data[key])
@@ -114,8 +116,17 @@ class ConfigHandler(APIHandler):
                         github_copilot.store_github_access_token()
                     else:
                         github_copilot.delete_stored_github_access_token()
+                elif key == "mcp_server_settings":
+                    disabled_mcp_servers = []
+                    for server_id in data[key]:
+                        server_settings = data[key][server_id]
+                        if server_settings.get("disabled") == True:
+                            disabled_mcp_servers.append(server_id)
+                    ai_service_manager.update_mcp_server_connections(disabled_mcp_servers)
+
         ai_service_manager.nbi_config.save()
-        ai_service_manager.update_models_from_config()
+        if has_model_change:
+            ai_service_manager.update_models_from_config()
         self.finish(json.dumps({}))
 
 class UpdateProviderModelsHandler(APIHandler):
@@ -126,18 +137,10 @@ class UpdateProviderModelsHandler(APIHandler):
             ai_service_manager.ollama_llm_provider.update_chat_model_list()
         self.finish(json.dumps({}))
 
-class ReloadMCPServersHandler(APIHandler):
-    @tornado.web.authenticated
-    def post(self):
-        ai_service_manager.nbi_config.load()
-        ai_service_manager.update_mcp_servers()
-        self.finish(json.dumps({
-            "mcpServers": [{"id": server.name} for server in ai_service_manager.get_mcp_servers()]
-        }))
-
 class MCPConfigFileHandler(APIHandler):
     @tornado.web.authenticated
     def get(self):
+        ai_service_manager.nbi_config.load()
         mcp_config = ai_service_manager.nbi_config.mcp.copy()
         if "mcpServers" not in mcp_config:
             mcp_config["mcpServers"] = {}
@@ -156,6 +159,15 @@ class MCPConfigFileHandler(APIHandler):
             self.finish(json.dumps({"status": "error", "message": str(e)}))
             return
 
+class ReloadMCPServersHandler(APIHandler):
+    @tornado.web.authenticated
+    def post(self):
+        ai_service_manager.nbi_config.load()
+        ai_service_manager.update_mcp_servers()
+        self.finish(json.dumps({
+            "mcpServers": [{"id": server.name} for server in ai_service_manager.get_mcp_servers()]
+        }))
+
 class EmitTelemetryEventHandler(APIHandler):
     @tornado.web.authenticated
     def post(self):
@@ -453,7 +465,9 @@ class WebsocketCopilotHandler(websocket.WebSocketHandler):
         # TODO: cleanup
         self._messageCallbackHandlers: dict[str, MessageCallbackHandlers] = {}
         self.chat_history = ChatHistory()
-        github_copilot.websocket_connector = ThreadSafeWebSocketConnector(self)
+        ws_connector = ThreadSafeWebSocketConnector(self)
+        ai_service_manager.websocket_connector = ws_connector
+        github_copilot.websocket_connector = ws_connector
 
     def open(self):
         pass
@@ -631,6 +645,7 @@ class NotebookIntelligence(ExtensionApp):
     async def stop_extension(self):
         log.info(f"Stopping {self.name} extension...")
         github_copilot.handle_stop_request()
+        ai_service_manager.handle_stop_request()
 
     def _setup_handlers(self, web_app):
         host_pattern = ".*$"
@@ -639,8 +654,8 @@ class NotebookIntelligence(ExtensionApp):
         route_pattern_capabilities = url_path_join(base_url, "notebook-intelligence", "capabilities")
         route_pattern_config = url_path_join(base_url, "notebook-intelligence", "config")
         route_pattern_update_provider_models = url_path_join(base_url, "notebook-intelligence", "update-provider-models")
-        route_pattern_reload_mcp_servers = url_path_join(base_url, "notebook-intelligence", "reload-mcp-servers")
         route_pattern_mcp_config_file = url_path_join(base_url, "notebook-intelligence", "mcp-config-file")
+        route_pattern_reload_mcp_servers = url_path_join(base_url, "notebook-intelligence", "reload-mcp-servers")
         route_pattern_emit_telemetry_event = url_path_join(base_url, "notebook-intelligence", "emit-telemetry-event")
         route_pattern_github_login_status = url_path_join(base_url, "notebook-intelligence", "gh-login-status")
         route_pattern_github_login = url_path_join(base_url, "notebook-intelligence", "gh-login")
@@ -651,8 +666,8 @@ class NotebookIntelligence(ExtensionApp):
             (route_pattern_capabilities, GetCapabilitiesHandler),
             (route_pattern_config, ConfigHandler),
             (route_pattern_update_provider_models, UpdateProviderModelsHandler),
-            (route_pattern_reload_mcp_servers, ReloadMCPServersHandler),
             (route_pattern_mcp_config_file, MCPConfigFileHandler),
+            (route_pattern_reload_mcp_servers, ReloadMCPServersHandler),
             (route_pattern_emit_telemetry_event, EmitTelemetryEventHandler),
             (route_pattern_github_login_status, GetGitHubLoginStatusHandler),
             (route_pattern_github_login, PostGitHubLoginHandler),

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/labextension/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@notebook-intelligence/notebook-intelligence",
-  "version": "2.4.2",
+  "version": "2.5.0",
   "description": "AI coding assistant for JupyterLab",
   "keywords": [
     "AI",
@@ -130,7 +130,7 @@
       }
     },
     "_build": {
-      "load": "static/remoteEntry.d9d709f6862dfa2e5a69.js",
+      "load": "static/remoteEntry.84c6d4df1ff766bc53bf.js",
       "extension": "./extension",
       "style": "./style"
     }

{notebook_intelligence-2.4.2 → notebook_intelligence-2.5.0}/notebook_intelligence/labextension/schemas/@notebook-intelligence/notebook-intelligence/package.json.orig

@@ -1,6 +1,6 @@
 {
     "name": "@notebook-intelligence/notebook-intelligence",
-    "version": "2.4.2",
+    "version": "2.5.0",
     "description": "AI coding assistant for JupyterLab",
     "keywords": [
         "AI",