feat: LSP root and document URI integration from backend

frontend/src/core/codemirror/language/languages/python.ts

@@ -36,7 +36,7 @@ import { FederatedLanguageServerClient } from "../../lsp/federated-lsp";
 import { NotebookLanguageServerClient } from "../../lsp/notebook-lsp";
 import { createTransport } from "../../lsp/transports";
 import { CellDocumentUri, type ILanguageServerClient } from "../../lsp/types";
-import { getLSPDocumentRootUri } from "../../lsp/utils";
+import { getLspRootUri, getLspWorkspaceFolders } from "../../lsp/utils";
 import {
   clickablePlaceholderExtension,
   smartPlaceholderExtension,
@@ -54,8 +54,8 @@ const pylspClient = once((lspConfig: LSPConfig) => {
 
   const lspClientOpts = {
     transport,
-    rootUri: getLSPDocumentRootUri(),
-    workspaceFolders: [],
+    rootUri: getLspRootUri(),
+    workspaceFolders: getLspWorkspaceFolders(),
   };
   const config = lspConfig?.pylsp;
 
@@ -161,8 +161,8 @@ const tyLspClient = once((_: LSPConfig) => {
 
   const lspClientOpts = {
     transport,
-    rootUri: getLSPDocumentRootUri(),
-    workspaceFolders: [],
+    rootUri: getLspRootUri(),
+    workspaceFolders: getLspWorkspaceFolders(),
   };
 
   // We wrap the client in a NotebookLanguageServerClient to add some
@@ -192,8 +192,8 @@ const pyreflyClient = once(
 
     const lspClientOpts = {
       transport,
-      rootUri: getLSPDocumentRootUri(),
-      workspaceFolders: [],
+      rootUri: getLspRootUri(),
+      workspaceFolders: getLspWorkspaceFolders(),
     };
 
     // We wrap the client in a NotebookLanguageServerClient to add some
@@ -230,8 +230,8 @@ const pyrightClient = once((_: LSPConfig) => {
 
   const lspClientOpts = {
     transport,
-    rootUri: getLSPDocumentRootUri(),
-    workspaceFolders: [],
+    rootUri: getLspRootUri(),
+    workspaceFolders: getLspWorkspaceFolders(),
   };
 
   // We wrap the client in a NotebookLanguageServerClient to add some

frontend/src/core/codemirror/lsp/__tests__/notebook-lsp.test.ts

@@ -12,6 +12,7 @@ import { cellId } from "@/__tests__/branded";
 import type { CellId } from "@/core/cells/ids";
 import { store } from "@/core/state/jotai";
 import { topologicalCodesAtom } from "../../copilot/getCodes";
+import { lspWorkspaceAtom } from "@/core/saving/file-state";
 import { languageAdapterState } from "../../language/extension";
 import { PythonLanguageAdapter } from "../../language/languages/python";
 import { languageMetadataField } from "../../language/metadata";
@@ -285,6 +286,12 @@ describe("NotebookLanguageServerClient", () => {
           },
         };
       }
+      if (atom === lspWorkspaceAtom) {
+        return {
+          rootUri: "file:///project",
+          documentUri: "file:///project/__marimo_notebook__.py",
+        };
+      }
       return undefined;
     });
 
@@ -421,7 +428,7 @@ describe("NotebookLanguageServerClient", () => {
       expect(result).toEqual(mockCompletionResponse);
       expect(mockClient.textDocumentCompletion).toHaveBeenCalledWith(
         expect.objectContaining({
-          textDocument: { uri: "file:///__marimo_notebook__.py" },
+          textDocument: { uri: "file:///project/__marimo_notebook__.py" },
         }),
       );
     });

frontend/src/core/codemirror/lsp/federated-lsp.ts

@@ -3,7 +3,7 @@
 import type * as LSP from "vscode-languageserver-protocol";
 import { Objects } from "@/utils/objects";
 import type { ILanguageServerClient } from "./types";
-import { getLSPDocument } from "./utils";
+import { getLspDocumentUri } from "./utils";
 
 function removeFalseyValues<T extends object>(obj: T): T {
   return Objects.filter(obj, (value) => value !== false && value !== null) as T;
@@ -20,7 +20,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
 
   constructor(clients: ILanguageServerClient[]) {
     this.clients = clients;
-    this.documentUri = getLSPDocument();
+    this.documentUri = getLspDocumentUri();
   }
 
   onNotification(

frontend/src/core/codemirror/lsp/notebook-lsp.ts

@@ -22,7 +22,7 @@ import {
   type ILanguageServerClient,
   isClientWithNotify,
 } from "./types";
-import { getLSPDocument } from "./utils";
+import { getLspDocumentUri } from "./utils";
 
 /**
  * Check if a variable name is private (starts with underscore but not dunder).
@@ -189,7 +189,7 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       EditorView | null | undefined
     > = defaultGetNotebookEditors,
   ) {
-    this.documentUri = getLSPDocument();
+    this.documentUri = getLspDocumentUri();
     this.getNotebookEditors = getNotebookEditors;
     this.initialSettings = initialSettings;
     this.client = client;

frontend/src/core/codemirror/lsp/utils.ts

@@ -1,11 +1,26 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { getFilenameFromDOM } from "@/core/dom/htmlUtils";
-import { Paths } from "@/utils/paths";
+import { lspWorkspaceAtom } from "@/core/saving/file-state";
+import { store } from "@/core/state/jotai";
 
-export function getLSPDocument() {
-  return `file://${getFilenameFromDOM() ?? "/__marimo_notebook__.py"}`;
+export function getLspRootUri() {
+  const lspWorkspace = store.get(lspWorkspaceAtom);
+  // The backend provides rootUri for active notebook sessions.
+  // For non-notebook pages (home, gallery), lspWorkspace is null,
+  // so return a valid file URI fallback.
+  return lspWorkspace?.rootUri ?? "file:///";
 }
 
-export function getLSPDocumentRootUri() {
-  return `file://${Paths.dirname(getFilenameFromDOM() ?? "/")}`;
+export function getLspWorkspaceFolders() {
+  const lspWorkspace = store.get(lspWorkspaceAtom);
+  const rootUri = lspWorkspace?.rootUri;
+  // Return workspace folders only if rootUri is set; empty array otherwise.
+  return rootUri ? [{ uri: rootUri, name: "marimo" }] : [];
+}
+
+export function getLspDocumentUri() {
+  const lspWorkspace = store.get(lspWorkspaceAtom);
+  // The backend provides documentUri for active notebook sessions.
+  // For non-notebook pages (home, gallery), lspWorkspace is null,
+  // so return a valid file URI fallback.
+  return lspWorkspace?.documentUri ?? "file:///__marimo_notebook__.py";
 }