diff --git a/pylsp_ruff/plugin.py b/pylsp_ruff/plugin.py
index a36babf..590ff49 100644
--- a/pylsp_ruff/plugin.py
+++ b/pylsp_ruff/plugin.py
@@ -516,6 +516,7 @@ def build_arguments(
     args.append("--exit-zero")
     # Use the json formatting for easier evaluation
     args.append("--output-format=json")
+    args.append("--extension=ipynb:python")
     if fix:
         args.append("--fix")
     else:
diff --git a/pyproject.toml b/pyproject.toml
index cfdf720..867c621 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,7 +13,7 @@ readme = "README.md"
 requires-python = ">=3.7"
 license = {text = "MIT"}
 dependencies = [
-  "ruff>=0.1.0, <0.2.0",
+  "ruff>=0.1.5, <0.2.0",
   "python-lsp-server",
   "lsprotocol>=2022.0.0a1",
   "tomli>=1.1.0; python_version < '3.11'",
diff --git a/tests/test_ruff_lint.py b/tests/test_ruff_lint.py
index 41026fb..c56a1ac 100644
--- a/tests/test_ruff_lint.py
+++ b/tests/test_ruff_lint.py
@@ -181,6 +181,7 @@ def f():
         "--quiet",
         "--exit-zero",
         "--output-format=json",
+        "--extension=ipynb:python",
         "--no-fix",
         "--force-exclude",
         f"--stdin-filename={os.path.join(workspace.root_path, '__init__.py')}",
@@ -242,3 +243,23 @@ def f():
         assert diag["code"] != "F401"
 
     os.unlink(os.path.join(workspace.root_path, "pyproject.toml"))
+
+
+def test_notebook_input(workspace):
+    doc_str = r"""
+print('hi')
+import os
+def f():
+    a = 2
+"""
+    # attribute the python code to a notebook file name per jupyterlab-lsp
+    doc_uri = uris.from_fs_path(os.path.join(workspace.root_path, "Untitled.ipynb"))
+    workspace.put_document(doc_uri, doc_str)
+    doc = workspace.get_document(doc_uri)
+
+    diags = ruff_lint.pylsp_lint(workspace, doc)
+    diag_codes = [diag["code"] for diag in diags]
+    assert "E999" not in diag_codes
+    assert "E402" in diag_codes
+    assert "F401" in diag_codes
+    assert "F841" in diag_codes