Code mirror plugin lsp interface (#1444)

* better named dirs

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* move some stuff around

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* more logging

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* less logging

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* add fs in

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* file reader

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* workspace

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* start of workspace folders

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* start of workspace folders

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* cleanup workspace folders

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* cleanup logs

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

---------

Signed-off-by: Jess Frazelle <github@jessfraz.com>

src/components/TextEditor.tsx

@@ -3,9 +3,9 @@ import ReactCodeMirror, {
   ViewUpdate,
   keymap,
 } from '@uiw/react-codemirror'
-import { FromServer, IntoServer } from 'editor/lsp/codec'
-import Server from '../editor/lsp/server'
-import Client from '../editor/lsp/client'
+import { FromServer, IntoServer } from 'editor/plugins/lsp/codec'
+import Server from '../editor/plugins/lsp/server'
+import Client from '../editor/plugins/lsp/client'
 import { TEST } from 'env'
 import { useCommandsContext } from 'hooks/useCommandsContext'
 import { useGlobalStateContext } from 'hooks/useGlobalStateContext'
@@ -15,8 +15,8 @@ import { useMemo, useRef } from 'react'
 import { linter, lintGutter } from '@codemirror/lint'
 import { useStore } from 'useStore'
 import { processCodeMirrorRanges } from 'lib/selections'
-import { LanguageServerClient } from 'editor/lsp'
-import kclLanguage from 'editor/lsp/language'
+import { LanguageServerClient } from 'editor/plugins/lsp'
+import kclLanguage from 'editor/plugins/lsp/kcl/language'
 import { EditorView, lineHighlightField } from 'editor/highlightextension'
 import { roundOff } from 'lib/utils'
 import { kclErrToDiagnostic } from 'lang/errors'
@@ -27,7 +27,9 @@ import { engineCommandManager } from '../lang/std/engineConnection'
 import { kclManager, useKclContext } from 'lang/KclSingleton'
 import { ModelingMachineEvent } from 'machines/modelingMachine'
 import { sceneInfra } from 'clientSideScene/sceneInfra'
-import { copilotBundle } from 'editor/copilot'
+import { copilotPlugin } from 'editor/plugins/lsp/copilot'
+import { isTauri } from 'lib/isTauri'
+import type * as LSP from 'vscode-languageserver-protocol'
 
 export const editorShortcutMeta = {
   formatCode: {
@@ -40,6 +42,15 @@ export const editorShortcutMeta = {
   },
 }
 
+function getWorkspaceFolders(): LSP.WorkspaceFolder[] {
+  // We only use workspace folders in Tauri since that is where we use more than
+  // one file.
+  if (isTauri()) {
+    return [{ uri: 'file://', name: 'ProjectRoot' }]
+  }
+  return []
+}
+
 export const TextEditor = ({
   theme,
 }: {
@@ -91,7 +102,7 @@ export const TextEditor = ({
       })
     }
 
-    const lspClient = new LanguageServerClient({ client })
+    const lspClient = new LanguageServerClient({ client, name: 'kcl' })
     return { lspClient }
   }, [setIsKclLspServerReady])
 
@@ -107,7 +118,7 @@ export const TextEditor = ({
       const lsp = kclLanguage({
         // When we have more than one file, we'll need to change this.
         documentUri: `file:///we-just-have-one-file-for-now.kcl`,
-        workspaceFolders: null,
+        workspaceFolders: getWorkspaceFolders(),
         client: kclLspClient,
       })
 
@@ -128,7 +139,7 @@ export const TextEditor = ({
       })
     }
 
-    const lspClient = new LanguageServerClient({ client })
+    const lspClient = new LanguageServerClient({ client, name: 'copilot' })
     return { lspClient }
   }, [setIsCopilotLspServerReady])
 
@@ -141,10 +152,10 @@ export const TextEditor = ({
     let plugin = null
     if (isCopilotLspServerReady && !TEST) {
       // Set up the lsp plugin.
-      const lsp = copilotBundle({
+      const lsp = copilotPlugin({
         // When we have more than one file, we'll need to change this.
         documentUri: `file:///we-just-have-one-file-for-now.kcl`,
-        workspaceFolders: null,
+        workspaceFolders: getWorkspaceFolders(),
         client: copilotLspClient,
         allowHTMLContent: true,
       })

src/editor/plugins/lsp/client.ts

@@ -65,6 +65,7 @@ export default class Client extends jsrpc.JSONRPCServerAndClient {
   afterInitializedHooks: (() => Promise<void>)[] = []
   #fromServer: FromServer
   private serverCapabilities: LSP.ServerCapabilities<any> = {}
+  private notifyFn: ((message: LSP.NotificationMessage) => void) | null = null
 
   constructor(fromServer: FromServer, intoServer: IntoServer) {
     super(
@@ -167,9 +168,15 @@ export default class Client extends jsrpc.JSONRPCServerAndClient {
     return this.serverCapabilities
   }
 
+  setNotifyFn(fn: (message: LSP.NotificationMessage) => void): void {
+    this.notifyFn = fn
+  }
+
   async processNotifications(): Promise<void> {
     for await (const notification of this.#fromServer.notifications) {
-      await this.receiveAndSend(notification)
+      if (this.notifyFn) {
+        this.notifyFn(notification)
+      }
     }
   }
 

src/editor/plugins/lsp/copilot/index.ts

@@ -11,30 +11,20 @@ import {
   Annotation,
   EditorState,
   Extension,
-  Facet,
   Prec,
   StateEffect,
   StateField,
   Transaction,
 } from '@codemirror/state'
 import { completionStatus } from '@codemirror/autocomplete'
-import { docPathFacet, offsetToPos, posToOffset } from 'editor/lsp/util'
-import { LanguageServerPlugin } from 'editor/lsp/plugin'
-import { LanguageServerOptions } from 'editor/lsp/plugin'
-import { LanguageServerClient } from 'editor/lsp'
-
-// Create Facet for the current docPath
-export const docPath = Facet.define<string, string>({
-  combine(value: readonly string[]) {
-    return value[value.length - 1]
-  },
-})
-
-export const relDocPath = Facet.define<string, string>({
-  combine(value: readonly string[]) {
-    return value[value.length - 1]
-  },
-})
+import { offsetToPos, posToOffset } from 'editor/plugins/lsp/util'
+import { LanguageServerOptions, LanguageServerClient } from 'editor/plugins/lsp'
+import {
+  LanguageServerPlugin,
+  documentUri,
+  languageId,
+  workspaceFolders,
+} from 'editor/plugins/lsp/plugin'
 
 const ghostMark = Decoration.mark({ class: 'cm-ghostText' })
 
@@ -361,9 +351,9 @@ const completionRequester = (client: LanguageServerClient) => {
       const pos = state.selection.main.head
       const source = state.doc.toString()
 
-      const path = state.facet(docPath)
-      const relativePath = state.facet(relDocPath)
-      const languageId = 'kcl'
+      const dUri = state.facet(documentUri)
+      const path = dUri.split('/').pop()!
+      const relativePath = dUri.replace('file://', '')
 
       // Set a new timeout to request completion
       timeout = setTimeout(async () => {
@@ -378,9 +368,9 @@ const completionRequester = (client: LanguageServerClient) => {
                 indentSize: 1,
                 insertSpaces: true,
                 path,
-                uri: `file://${path}`,
+                uri: dUri,
                 relativePath,
-                languageId,
+                languageId: state.facet(languageId),
                 position: offsetToPos(state.doc, pos),
               },
             })
@@ -483,21 +473,24 @@ const completionRequester = (client: LanguageServerClient) => {
   })
 }
 
-export function copilotServer(options: LanguageServerOptions) {
-  let plugin: LanguageServerPlugin
-  return ViewPlugin.define(
-    (view) =>
-      (plugin = new LanguageServerPlugin(view, options.allowHTMLContent))
-  )
-}
+export const copilotPlugin = (options: LanguageServerOptions): Extension => {
+  let plugin: LanguageServerPlugin | null = null
 
-export const copilotBundle = (options: LanguageServerOptions): Extension => [
-  docPath.of(options.documentUri.split('/').pop()!),
-  docPathFacet.of(options.documentUri.split('/').pop()!),
-  relDocPath.of(options.documentUri.replace('file://', '')),
+  return [
+    documentUri.of(options.documentUri),
+    languageId.of('kcl'),
+    workspaceFolders.of(options.workspaceFolders),
+    ViewPlugin.define(
+      (view) =>
+        (plugin = new LanguageServerPlugin(
+          options.client,
+          view,
+          options.allowHTMLContent
+        ))
+    ),
     completionDecoration,
     Prec.highest(completionPlugin(options.client)),
     Prec.highest(viewCompletionPlugin(options.client)),
     completionRequester(options.client),
-  copilotServer(options),
   ]
+}

src/editor/plugins/lsp/index.ts

@@ -1,7 +1,7 @@
 import type * as LSP from 'vscode-languageserver-protocol'
 import Client from './client'
-import { LanguageServerPlugin } from './plugin'
-import { SemanticToken, deserializeTokens } from './semantic_tokens'
+import { SemanticToken, deserializeTokens } from './kcl/semantic_tokens'
+import { LanguageServerPlugin } from 'editor/plugins/lsp/plugin'
 
 export interface CopilotGetCompletionsParams {
   doc: {
@@ -90,26 +90,22 @@ interface LSPNotifyMap {
   'textDocument/didOpen': LSP.DidOpenTextDocumentParams
 }
 
-// Server to client
-interface LSPEventMap {
-  'textDocument/publishDiagnostics': LSP.PublishDiagnosticsParams
-}
-
-export type Notification = {
-  [key in keyof LSPEventMap]: {
-    jsonrpc: '2.0'
-    id?: null | undefined
-    method: key
-    params: LSPEventMap[key]
-  }
-}[keyof LSPEventMap]
-
 export interface LanguageServerClientOptions {
   client: Client
+  name: string
+}
+
+export interface LanguageServerOptions {
+  // We assume this is the main project directory, we are currently working in.
+  workspaceFolders: LSP.WorkspaceFolder[]
+  documentUri: string
+  allowHTMLContent: boolean
+  client: LanguageServerClient
 }
 
 export class LanguageServerClient {
   private client: Client
+  private name: string
 
   public ready: boolean
 
@@ -124,6 +120,7 @@ export class LanguageServerClient {
   constructor(options: LanguageServerClientOptions) {
     this.plugins = []
     this.client = options.client
+    this.name = options.name
 
     this.ready = false
 
@@ -133,11 +130,16 @@ export class LanguageServerClient {
 
   async initialize() {
     // Start the client in the background.
+    this.client.setNotifyFn(this.processNotifications.bind(this))
     this.client.start()
 
     this.ready = true
   }
 
+  getName(): string {
+    return this.name
+  }
+
   getServerCapabilities(): LSP.ServerCapabilities<any> {
     return this.client.getServerCapabilities()
   }
@@ -156,6 +158,11 @@ export class LanguageServerClient {
   }
 
   async updateSemanticTokens(uri: string) {
+    const serverCapabilities = this.getServerCapabilities()
+    if (!serverCapabilities.semanticTokensProvider) {
+      return
+    }
+
     // Make sure we can only run, if we aren't already running.
     if (!this.isUpdatingSemanticTokens) {
       this.isUpdatingSemanticTokens = true
@@ -180,10 +187,18 @@ export class LanguageServerClient {
   }
 
   async textDocumentHover(params: LSP.HoverParams) {
+    const serverCapabilities = this.getServerCapabilities()
+    if (!serverCapabilities.hoverProvider) {
+      return
+    }
     return await this.request('textDocument/hover', params)
   }
 
   async textDocumentCompletion(params: LSP.CompletionParams) {
+    const serverCapabilities = this.getServerCapabilities()
+    if (!serverCapabilities.completionProvider) {
+      return
+    }
     return await this.request('textDocument/completion', params)
   }
 
@@ -234,11 +249,12 @@ export class LanguageServerClient {
   async acceptCompletion(params: CopilotAcceptCompletionParams) {
     return await this.request('notifyAccepted', params)
   }
+
   async rejectCompletions(params: CopilotRejectCompletionParams) {
     return await this.request('notifyRejected', params)
   }
 
-  private processNotification(notification: Notification) {
+  private processNotifications(notification: LSP.NotificationMessage) {
     for (const plugin of this.plugins) plugin.processNotification(notification)
   }
 }

src/editor/plugins/lsp/kcl/index.ts

@@ -0,0 +1,75 @@
+import { autocompletion } from '@codemirror/autocomplete'
+import { Extension } from '@codemirror/state'
+import { ViewPlugin, hoverTooltip, tooltips } from '@codemirror/view'
+import { CompletionTriggerKind } from 'vscode-languageserver-protocol'
+import { offsetToPos } from 'editor/plugins/lsp/util'
+import { LanguageServerOptions } from 'editor/plugins/lsp'
+import {
+  LanguageServerPlugin,
+  documentUri,
+  languageId,
+  workspaceFolders,
+} from 'editor/plugins/lsp/plugin'
+
+export function kclPlugin(options: LanguageServerOptions): Extension {
+  let plugin: LanguageServerPlugin | null = null
+
+  return [
+    documentUri.of(options.documentUri),
+    languageId.of('kcl'),
+    workspaceFolders.of(options.workspaceFolders),
+    ViewPlugin.define(
+      (view) =>
+        (plugin = new LanguageServerPlugin(
+          options.client,
+          view,
+          options.allowHTMLContent
+        ))
+    ),
+    hoverTooltip(
+      (view, pos) =>
+        plugin?.requestHoverTooltip(view, offsetToPos(view.state.doc, pos)) ??
+        null
+    ),
+    tooltips({
+      position: 'absolute',
+    }),
+    autocompletion({
+      override: [
+        async (context) => {
+          if (plugin == null) return null
+
+          const { state, pos, explicit } = context
+          const line = state.doc.lineAt(pos)
+          let trigKind: CompletionTriggerKind = CompletionTriggerKind.Invoked
+          let trigChar: string | undefined
+          if (
+            !explicit &&
+            plugin.client
+              .getServerCapabilities()
+              .completionProvider?.triggerCharacters?.includes(
+                line.text[pos - line.from - 1]
+              )
+          ) {
+            trigKind = CompletionTriggerKind.TriggerCharacter
+            trigChar = line.text[pos - line.from - 1]
+          }
+          if (
+            trigKind === CompletionTriggerKind.Invoked &&
+            !context.matchBefore(/\w+$/)
+          ) {
+            return null
+          }
+          return await plugin.requestCompletion(
+            context,
+            offsetToPos(state.doc, pos),
+            {
+              triggerKind: trigKind,
+              triggerCharacter: trigChar,
+            }
+          )
+        },
+      ],
+    }),
+  ]
+}

src/editor/plugins/lsp/kcl/language.ts

@@ -5,8 +5,8 @@ import {
   defineLanguageFacet,
   LanguageSupport,
 } from '@codemirror/language'
-import { LanguageServerClient } from '.'
-import { kclPlugin } from './plugin'
+import { LanguageServerClient } from 'editor/plugins/lsp'
+import { kclPlugin } from '.'
 import type * as LSP from 'vscode-languageserver-protocol'
 import { parser as jsParser } from '@lezer/javascript'
 import { EditorState } from '@uiw/react-codemirror'
@@ -14,7 +14,7 @@ import { EditorState } from '@uiw/react-codemirror'
 const data = defineLanguageFacet({})
 
 export interface LanguageOptions {
-  workspaceFolders: LSP.WorkspaceFolder[] | null
+  workspaceFolders: LSP.WorkspaceFolder[]
   documentUri: string
   client: LanguageServerClient
 }

src/editor/plugins/lsp/kcl/parser.ts

@@ -9,8 +9,8 @@ import {
   NodeType,
   NodeSet,
 } from '@lezer/common'
-import { LanguageServerClient } from '.'
-import { posToOffset } from 'editor/lsp/util'
+import { LanguageServerClient } from 'editor/plugins/lsp'
+import { posToOffset } from 'editor/plugins/lsp/util'
 import { SemanticToken } from './semantic_tokens'
 import { DocInput } from '@codemirror/language'
 import { tags, styleTags } from '@lezer/highlight'

src/editor/plugins/lsp/plugin.ts

@@ -1,13 +1,7 @@
-import { autocompletion, completeFromList } from '@codemirror/autocomplete'
+import { completeFromList } from '@codemirror/autocomplete'
 import { setDiagnostics } from '@codemirror/lint'
 import { Facet } from '@codemirror/state'
-import {
-  EditorView,
-  ViewPlugin,
-  Tooltip,
-  hoverTooltip,
-  tooltips,
-} from '@codemirror/view'
+import { EditorView, Tooltip } from '@codemirror/view'
 import {
   DiagnosticSeverity,
   CompletionItemKind,
@@ -23,9 +17,17 @@ import type {
 import type { PublishDiagnosticsParams } from 'vscode-languageserver-protocol'
 import type { ViewUpdate, PluginValue } from '@codemirror/view'
 import type * as LSP from 'vscode-languageserver-protocol'
-import { LanguageServerClient, Notification } from '.'
+import { LanguageServerClient } from 'editor/plugins/lsp'
 import { Marked } from '@ts-stack/markdown'
-import { offsetToPos, posToOffset } from 'editor/lsp/util'
+import { posToOffset } from 'editor/plugins/lsp/util'
+
+const useLast = (values: readonly any[]) => values.reduce((_, v) => v, '')
+export const documentUri = Facet.define<string, string>({ combine: useLast })
+export const languageId = Facet.define<string, string>({ combine: useLast })
+export const workspaceFolders = Facet.define<
+  LSP.WorkspaceFolder[],
+  LSP.WorkspaceFolder[]
+>({ combine: useLast })
 
 const changesDelay = 500
 
@@ -33,31 +35,22 @@ const CompletionItemKindMap = Object.fromEntries(
   Object.entries(CompletionItemKind).map(([key, value]) => [value, key])
 ) as Record<CompletionItemKind, string>
 
-const useLast = (values: readonly any[]) => values.reduce((_, v) => v, '')
-const documentUri = Facet.define<string, string>({ combine: useLast })
-const languageId = Facet.define<string, string>({ combine: useLast })
-const client = Facet.define<LanguageServerClient, LanguageServerClient>({
-  combine: useLast,
-})
-
-export interface LanguageServerOptions {
-  workspaceFolders: LSP.WorkspaceFolder[] | null
-  documentUri: string
-  allowHTMLContent: boolean
-  client: LanguageServerClient
-}
-
 export class LanguageServerPlugin implements PluginValue {
   public client: LanguageServerClient
-
   private documentUri: string
   private languageId: string
+  private workspaceFolders: LSP.WorkspaceFolder[]
   private documentVersion: number
 
-  constructor(private view: EditorView, private allowHTMLContent: boolean) {
-    this.client = this.view.state.facet(client)
+  constructor(
+    client: LanguageServerClient,
+    private view: EditorView,
+    private allowHTMLContent: boolean
+  ) {
+    this.client = client
     this.documentUri = this.view.state.facet(documentUri)
     this.languageId = this.view.state.facet(languageId)
+    this.workspaceFolders = this.view.state.facet(workspaceFolders)
     this.documentVersion = 0
 
     this.client.attachPlugin(this)
@@ -238,11 +231,28 @@ export class LanguageServerPlugin implements PluginValue {
     return completeFromList(options)(context)
   }
 
-  processNotification(notification: Notification) {
+  processNotification(notification: LSP.NotificationMessage) {
     try {
       switch (notification.method) {
         case 'textDocument/publishDiagnostics':
-          this.processDiagnostics(notification.params)
+          this.processDiagnostics(
+            notification.params as PublishDiagnosticsParams
+          )
+          break
+        case 'window/logMessage':
+          console.log(
+            '[lsp] [window/logMessage]',
+            this.client.getName(),
+            notification.params
+          )
+          break
+        case 'window/showMessage':
+          console.log(
+            '[lsp] [window/showMessage]',
+            this.client.getName(),
+            notification.params
+          )
+          break
       }
     } catch (error) {
       console.error(error)
@@ -284,65 +294,6 @@ export class LanguageServerPlugin implements PluginValue {
   }
 }
 
-export function kclPlugin(options: LanguageServerOptions) {
-  let plugin: LanguageServerPlugin | null = null
-
-  return [
-    client.of(options.client),
-    documentUri.of(options.documentUri),
-    languageId.of('kcl'),
-    ViewPlugin.define(
-      (view) =>
-        (plugin = new LanguageServerPlugin(view, options.allowHTMLContent))
-    ),
-    hoverTooltip(
-      (view, pos) =>
-        plugin?.requestHoverTooltip(view, offsetToPos(view.state.doc, pos)) ??
-        null
-    ),
-    tooltips({
-      position: 'absolute',
-    }),
-    autocompletion({
-      override: [
-        async (context) => {
-          if (plugin == null) return null
-
-          const { state, pos, explicit } = context
-          const line = state.doc.lineAt(pos)
-          let trigKind: CompletionTriggerKind = CompletionTriggerKind.Invoked
-          let trigChar: string | undefined
-          if (
-            !explicit &&
-            plugin.client
-              .getServerCapabilities()
-              .completionProvider?.triggerCharacters?.includes(
-                line.text[pos - line.from - 1]
-              )
-          ) {
-            trigKind = CompletionTriggerKind.TriggerCharacter
-            trigChar = line.text[pos - line.from - 1]
-          }
-          if (
-            trigKind === CompletionTriggerKind.Invoked &&
-            !context.matchBefore(/\w+$/)
-          ) {
-            return null
-          }
-          return await plugin.requestCompletion(
-            context,
-            offsetToPos(state.doc, pos),
-            {
-              triggerKind: trigKind,
-              triggerCharacter: trigChar,
-            }
-          )
-        },
-      ],
-    }),
-  ]
-}
-
 function formatContents(
   contents: LSP.MarkupContent | LSP.MarkedString | LSP.MarkedString[]
 ): string {

src/editor/plugins/lsp/server-capability-registration.ts

@@ -34,6 +34,8 @@ const ServerCapabilitiesProviders: IMethodServerCapabilityProviderDictionary = {
   'textDocument/foldingRange': 'foldingRangeProvider',
   'textDocument/declaration': 'declarationProvider',
   'textDocument/executeCommand': 'executeCommandProvider',
+  'textDocument/semanticTokens/full': 'semanticTokensProvider',
+  'textDocument/publishDiagnostics': 'diagnosticsProvider',
 }
 
 function registerServerCapability(

src/editor/plugins/lsp/server.ts

@@ -3,8 +3,9 @@ import init, {
   InitOutput,
   kcl_lsp_run,
   ServerConfig,
-} from '../../wasm-lib/pkg/wasm_lib'
+} from 'wasm-lib/pkg/wasm_lib'
 import { FromServer, IntoServer } from './codec'
+import { fileSystemManager } from 'lang/std/fileSystemManager'
 
 export default class Server {
   readonly initOutput: InitOutput
@@ -31,7 +32,11 @@ export default class Server {
   }
 
   async start(type_: 'kcl' | 'copilot', token?: string): Promise<void> {
-    const config = new ServerConfig(this.#intoServer, this.#fromServer)
+    const config = new ServerConfig(
+      this.#intoServer,
+      this.#fromServer,
+      fileSystemManager
+    )
     if (type_ === 'copilot') {
       if (!token) {
         throw new Error('auth token is required for copilot')

src/editor/plugins/lsp/util.ts

@@ -1,4 +1,4 @@
-import { Facet, Text } from '@codemirror/state'
+import { Text } from '@codemirror/state'
 
 export function posToOffset(
   doc: Text,
@@ -17,7 +17,3 @@ export function offsetToPos(doc: Text, offset: number) {
     character: offset - line.from,
   }
 }
-
-export const docPathFacet = Facet.define<string, string>({
-  combine: (values) => values[values.length - 1],
-})

src/lang/std/fileSystemManager.ts

@@ -1,4 +1,8 @@
-import { readBinaryFile, exists as tauriExists } from '@tauri-apps/api/fs'
+import {
+  readDir,
+  readBinaryFile,
+  exists as tauriExists,
+} from '@tauri-apps/api/fs'
 import { isTauri } from 'lib/isTauri'
 import { join } from '@tauri-apps/api/path'
 
@@ -53,6 +57,30 @@ class FileSystemManager {
         return tauriExists(file)
       })
   }
+
+  getAllFiles(path: string): Promise<string[] | void> {
+    // Using local file system only works from Tauri.
+    if (!isTauri()) {
+      throw new Error(
+        'This function can only be called from a Tauri application'
+      )
+    }
+
+    return join(this.dir, path)
+      .catch((error) => {
+        throw new Error(`Error joining dir: ${error}`)
+      })
+      .then((p) => {
+        readDir(p, { recursive: true })
+          .catch((error) => {
+            throw new Error(`Error reading dir: ${error}`)
+          })
+
+          .then((files) => {
+            return files.map((file) => file.path)
+          })
+      })
+  }
 }
 
 export const fileSystemManager = new FileSystemManager()

src/machines/authMachine.ts

@@ -37,7 +37,10 @@ export type Events =
     }
 
 export const TOKEN_PERSIST_KEY = 'TOKEN_PERSIST_KEY'
-const persistedToken = localStorage?.getItem(TOKEN_PERSIST_KEY) || ''
+const persistedToken =
+  localStorage?.getItem(TOKEN_PERSIST_KEY) ||
+  getCookie('__Secure-next-auth.session-token') ||
+  ''
 
 export const authMachine = createMachine<UserContext, Events>(
   {
@@ -135,3 +138,23 @@ async function getUser(context: UserContext) {
 
   return user
 }
+
+function getCookie(cname: string): string {
+  if (isTauri()) {
+    return ''
+  }
+
+  let name = cname + '='
+  let decodedCookie = decodeURIComponent(document.cookie)
+  let ca = decodedCookie.split(';')
+  for (let i = 0; i < ca.length; i++) {
+    let c = ca[i]
+    while (c.charAt(0) === ' ') {
+      c = c.substring(1)
+    }
+    if (c.indexOf(name) === 0) {
+      return c.substring(name.length, c.length)
+    }
+  }
+  return ''
+}

src/wasm-lib/kcl/src/fs/local.rs

@@ -53,4 +53,38 @@ impl FileSystem for FileManager {
             }
         })
     }
+
+    async fn get_all_files<P: AsRef<std::path::Path>>(
+        &self,
+        path: P,
+        source_range: crate::executor::SourceRange,
+    ) -> Result<Vec<std::path::PathBuf>, crate::errors::KclError> {
+        let mut files = vec![];
+        let mut stack = vec![path.as_ref().to_path_buf()];
+
+        while let Some(path) = stack.pop() {
+            if !path.is_dir() {
+                continue;
+            }
+
+            let mut read_dir = tokio::fs::read_dir(&path).await.map_err(|e| {
+                KclError::Engine(KclErrorDetails {
+                    message: format!("Failed to read directory `{}`: {}", path.display(), e),
+                    source_ranges: vec![source_range],
+                })
+            })?;
+
+            while let Ok(Some(entry)) = read_dir.next_entry().await {
+                let path = entry.path();
+                if path.is_dir() {
+                    // Iterate over the directory.
+                    stack.push(path);
+                } else {
+                    files.push(path);
+                }
+            }
+        }
+
+        Ok(files)
+    }
 }

src/wasm-lib/kcl/src/fs/mod.rs

@@ -28,4 +28,11 @@ pub trait FileSystem: Clone {
         path: P,
         source_range: crate::executor::SourceRange,
     ) -> Result<bool, crate::errors::KclError>;
+
+    /// Get all the files in a directory recursively.
+    async fn get_all_files<P: AsRef<std::path::Path>>(
+        &self,
+        path: P,
+        source_range: crate::executor::SourceRange,
+    ) -> Result<Vec<std::path::PathBuf>, crate::errors::KclError>;
 }

src/wasm-lib/kcl/src/fs/wasm.rs

@@ -18,6 +18,9 @@ extern "C" {
 
     #[wasm_bindgen(method, js_name = exists, catch)]
     fn exists(this: &FileSystemManager, path: String) -> Result<js_sys::Promise, js_sys::Error>;
+
+    #[wasm_bindgen(method, js_name = getAllFiles, catch)]
+    fn get_all_files(this: &FileSystemManager, path: String) -> Result<js_sys::Promise, js_sys::Error>;
 }
 
 #[derive(Debug, Clone)]
@@ -31,6 +34,9 @@ impl FileManager {
     }
 }
 
+unsafe impl Send for FileManager {}
+unsafe impl Sync for FileManager {}
+
 #[async_trait::async_trait(?Send)]
 impl FileSystem for FileManager {
     async fn read<P: AsRef<std::path::Path>>(
@@ -112,4 +118,53 @@ impl FileSystem for FileManager {
 
         Ok(it_exists)
     }
+
+    async fn get_all_files<P: AsRef<std::path::Path>>(
+        &self,
+        path: P,
+        source_range: crate::executor::SourceRange,
+    ) -> Result<Vec<std::path::PathBuf>, crate::errors::KclError> {
+        let promise = self
+            .manager
+            .get_all_files(
+                path.as_ref()
+                    .to_str()
+                    .ok_or_else(|| {
+                        KclError::Engine(KclErrorDetails {
+                            message: "Failed to convert path to string".to_string(),
+                            source_ranges: vec![source_range],
+                        })
+                    })?
+                    .to_string(),
+            )
+            .map_err(|e| {
+                KclError::Engine(KclErrorDetails {
+                    message: e.to_string().into(),
+                    source_ranges: vec![source_range],
+                })
+            })?;
+
+        let value = wasm_bindgen_futures::JsFuture::from(promise).await.map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to wait for promise from javascript: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        let s = value.as_string().ok_or_else(|| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to get string from response from javascript: `{:?}`", value),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        let files: Vec<String> = serde_json::from_str(&s).map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to parse json from javascript: `{}` `{:?}`", s, e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        Ok(files.into_iter().map(|s| std::path::PathBuf::from(s)).collect())
+    }
 }

src/wasm-lib/kcl/src/lib.rs

@@ -10,7 +10,7 @@ pub mod engine;
 pub mod errors;
 pub mod executor;
 pub mod fs;
+pub mod lsp;
 pub mod parser;
-pub mod server;
 pub mod std;
 pub mod token;

src/wasm-lib/kcl/src/lsp/backend.rs

@@ -13,6 +13,8 @@ use tower_lsp::lsp_types::{
 pub trait Backend {
     fn client(&self) -> tower_lsp::Client;
 
+    fn fs(&self) -> crate::fs::FileManager;
+
     /// Get the current code map.
     fn current_code_map(&self) -> DashMap<String, String>;
 

src/wasm-lib/kcl/src/lsp/copilot/cache.rs

@@ -6,7 +6,7 @@ use std::{
     sync::{Mutex, RwLock},
 };
 
-use crate::server::copilot::types::CopilotCompletionResponse;
+use crate::lsp::copilot::types::CopilotCompletionResponse;
 
 // if file changes, keep the cache.
 // if line number is different for an existing file, clean.

src/wasm-lib/kcl/src/lsp/copilot/mod.rs

@@ -23,7 +23,7 @@ use tower_lsp::{
     LanguageServer,
 };
 
-use crate::server::{
+use crate::lsp::{
     backend::Backend as _,
     copilot::types::{CopilotCompletionResponse, CopilotEditorInfo, CopilotLspCompletionParams, DocParams},
 };
@@ -42,6 +42,8 @@ impl Success {
 pub struct Backend {
     /// The client is used to send notifications and requests to the client.
     pub client: tower_lsp::Client,
+    /// The file system client to use.
+    pub fs: crate::fs::FileManager,
     /// Current code.
     pub current_code_map: DashMap<String, String>,
     /// The token is used to authenticate requests to the API server.
@@ -54,11 +56,15 @@ pub struct Backend {
 
 // Implement the shared backend trait for the language server.
 #[async_trait::async_trait]
-impl crate::server::backend::Backend for Backend {
+impl crate::lsp::backend::Backend for Backend {
     fn client(&self) -> tower_lsp::Client {
         self.client.clone()
     }
 
+    fn fs(&self) -> crate::fs::FileManager {
+        self.fs.clone()
+    }
+
     fn current_code_map(&self) -> DashMap<String, String> {
         self.current_code_map.clone()
     }
@@ -125,15 +131,15 @@ impl Backend {
         let pos = params.doc.position;
         let uri = params.doc.uri.to_string();
         let rope = ropey::Rope::from_str(&params.doc.source);
-        let offset = crate::server::util::position_to_offset(pos, &rope).unwrap_or_default();
+        let offset = crate::lsp::util::position_to_offset(pos, &rope).unwrap_or_default();
 
         Ok(DocParams {
             uri: uri.to_string(),
             pos,
             language: params.doc.language_id.to_string(),
-            prefix: crate::server::util::get_text_before(offset, &rope).unwrap_or_default(),
-            suffix: crate::server::util::get_text_after(offset, &rope).unwrap_or_default(),
-            line_before: crate::server::util::get_line_before(pos, &rope).unwrap_or_default(),
+            prefix: crate::lsp::util::get_text_before(offset, &rope).unwrap_or_default(),
+            suffix: crate::lsp::util::get_text_after(offset, &rope).unwrap_or_default(),
+            line_before: crate::lsp::util::get_line_before(pos, &rope).unwrap_or_default(),
             rope,
         })
     }

src/wasm-lib/kcl/src/lsp/kcl/mod.rs

@@ -29,7 +29,7 @@ use tower_lsp::{
     Client, LanguageServer,
 };
 
-use crate::{ast::types::VariableKind, executor::SourceRange, parser::PIPE_OPERATOR, server::backend::Backend as _};
+use crate::{ast::types::VariableKind, executor::SourceRange, lsp::backend::Backend as _, parser::PIPE_OPERATOR};
 
 /// A subcommand for running the server.
 #[derive(Clone, Debug)]
@@ -48,6 +48,8 @@ pub struct Server {
 pub struct Backend {
     /// The client for the backend.
     pub client: Client,
+    /// The file system client to use.
+    pub fs: crate::fs::FileManager,
     /// The stdlib completions for the language.
     pub stdlib_completions: HashMap<String, CompletionItem>,
     /// The stdlib signatures for the language.
@@ -70,11 +72,15 @@ pub struct Backend {
 
 // Implement the shared backend trait for the language server.
 #[async_trait::async_trait]
-impl crate::server::backend::Backend for Backend {
+impl crate::lsp::backend::Backend for Backend {
     fn client(&self) -> Client {
         self.client.clone()
     }
 
+    fn fs(&self) -> crate::fs::FileManager {
+        self.fs.clone()
+    }
+
     fn current_code_map(&self) -> DashMap<String, String> {
         self.current_code_map.clone()
     }

src/wasm-lib/kcl/src/lsp/mod.rs

@@ -2,5 +2,5 @@
 
 mod backend;
 pub mod copilot;
-pub mod lsp;
+pub mod kcl;
 mod util;

src/wasm-lib/src/lib.rs

@@ -129,16 +129,22 @@ pub fn recast_wasm(json_str: &str) -> Result<JsValue, JsError> {
 pub struct ServerConfig {
     into_server: js_sys::AsyncIterator,
     from_server: web_sys::WritableStream,
+    fs: kcl_lib::fs::wasm::FileSystemManager,
 }
 
 #[cfg(target_arch = "wasm32")]
 #[wasm_bindgen]
 impl ServerConfig {
     #[wasm_bindgen(constructor)]
-    pub fn new(into_server: js_sys::AsyncIterator, from_server: web_sys::WritableStream) -> Self {
+    pub fn new(
+        into_server: js_sys::AsyncIterator,
+        from_server: web_sys::WritableStream,
+        fs: kcl_lib::fs::wasm::FileSystemManager,
+    ) -> Self {
         Self {
             into_server,
             from_server,
+            fs,
         }
     }
 }
@@ -156,17 +162,19 @@ pub async fn kcl_lsp_run(config: ServerConfig) -> Result<(), JsValue> {
     let ServerConfig {
         into_server,
         from_server,
+        fs,
     } = config;
 
     let stdlib = kcl_lib::std::StdLib::new();
-    let stdlib_completions = kcl_lib::server::lsp::get_completions_from_stdlib(&stdlib).map_err(|e| e.to_string())?;
-    let stdlib_signatures = kcl_lib::server::lsp::get_signatures_from_stdlib(&stdlib).map_err(|e| e.to_string())?;
+    let stdlib_completions = kcl_lib::lsp::kcl::get_completions_from_stdlib(&stdlib).map_err(|e| e.to_string())?;
+    let stdlib_signatures = kcl_lib::lsp::kcl::get_signatures_from_stdlib(&stdlib).map_err(|e| e.to_string())?;
     // We can unwrap here because we know the tokeniser is valid, since
     // we have a test for it.
     let token_types = kcl_lib::token::TokenType::all_semantic_token_types().unwrap();
 
-    let (service, socket) = LspService::new(|client| kcl_lib::server::lsp::Backend {
+    let (service, socket) = LspService::new(|client| kcl_lib::lsp::kcl::Backend {
         client,
+        fs: kcl_lib::fs::FileManager::new(fs),
         stdlib_completions,
         stdlib_signatures,
         token_types,
@@ -211,24 +219,24 @@ pub async fn copilot_lsp_run(config: ServerConfig, token: String) -> Result<(),
     let ServerConfig {
         into_server,
         from_server,
+        fs,
     } = config;
 
-    let (service, socket) = LspService::build(|client| kcl_lib::server::copilot::Backend {
+    let (service, socket) = LspService::build(|client| kcl_lib::lsp::copilot::Backend {
         client,
+        fs: kcl_lib::fs::FileManager::new(fs),
         current_code_map: Default::default(),
-        editor_info: Arc::new(RwLock::new(
-            kcl_lib::server::copilot::types::CopilotEditorInfo::default(),
-        )),
-        cache: kcl_lib::server::copilot::cache::CopilotCache::new(),
+        editor_info: Arc::new(RwLock::new(kcl_lib::lsp::copilot::types::CopilotEditorInfo::default())),
+        cache: kcl_lib::lsp::copilot::cache::CopilotCache::new(),
         token,
     })
-    .custom_method("setEditorInfo", kcl_lib::server::copilot::Backend::set_editor_info)
+    .custom_method("setEditorInfo", kcl_lib::lsp::copilot::Backend::set_editor_info)
     .custom_method(
         "getCompletions",
-        kcl_lib::server::copilot::Backend::get_completions_cycling,
+        kcl_lib::lsp::copilot::Backend::get_completions_cycling,
     )
-    .custom_method("notifyAccepted", kcl_lib::server::copilot::Backend::accept_completions)
-    .custom_method("notifyRejected", kcl_lib::server::copilot::Backend::reject_completions)
+    .custom_method("notifyAccepted", kcl_lib::lsp::copilot::Backend::accept_completions)
+    .custom_method("notifyRejected", kcl_lib::lsp::copilot::Backend::reject_completions)
     .finish();
 
     let input = wasm_bindgen_futures::stream::JsStream::from(into_server);