diff --git a/JSONRPC-DataChannel-UniSocket b/JSONRPC-DataChannel-UniSocket
index d1e8947..e6afb4c 160000
--- a/JSONRPC-DataChannel-UniSocket
+++ b/JSONRPC-DataChannel-UniSocket
@@ -1 +1 @@
-Subproject commit d1e894785e5b9f00d9403ab3e6c4c7a02c0b3020
+Subproject commit e6afb4cf5b3d5cfabff49fea9d7d0a6c05c4c42b
diff --git a/LanguageClient b/LanguageClient
index 0543134..7383de1 160000
--- a/LanguageClient
+++ b/LanguageClient
@@ -1 +1 @@
-Subproject commit 05431343c5c0867398baadeb427610d3ab8fbf58
+Subproject commit 7383de1e6176519f64c69d13ff7a4e7892101eb1
diff --git a/LanguageServerProtocol b/LanguageServerProtocol
index 6372a00..84add60 160000
--- a/LanguageServerProtocol
+++ b/LanguageServerProtocol
@@ -1 +1 @@
-Subproject commit 6372a003b1884ff95773e63fe8663ee5cf7e99c7
+Subproject commit 84add603d1cce79bff14f4e0bb16dce8e51b624f
diff --git a/Package.resolved b/Package.resolved
index 6a3d360..a90fe8f 100644
--- a/Package.resolved
+++ b/Package.resolved
@@ -36,6 +36,15 @@
         "version" : "0.1.1"
       }
     },
+    {
+      "identity" : "jsonrpc",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ChimeHQ/JSONRPC",
+      "state" : {
+        "revision" : "c6ec759d41a76ac88fe7327c41a77d9033943374",
+        "version" : "0.9.0"
+      }
+    },
     {
       "identity" : "processenv",
       "kind" : "remoteSourceControl",
@@ -86,8 +95,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/apple/swift-collections.git",
       "state" : {
-        "revision" : "937e904258d22af6e447a0b72c0bc67583ef64a2",
-        "version" : "1.0.4"
+        "revision" : "a902f1823a7ff3c9ab2fba0f992396b948eda307",
+        "version" : "1.0.5"
       }
     },
     {
diff --git a/Package.swift b/Package.swift
index 7a4c7af..95bf9d7 100644
--- a/Package.swift
+++ b/Package.swift
@@ -5,7 +5,8 @@ import PackageDescription
 import Foundation
 
 let commonCompileSetting: SwiftSetting =
-	.unsafeFlags(["-strict-concurrency=complete", "-warn-concurrency"])
+	.unsafeFlags([])
+	// .unsafeFlags(["-strict-concurrency=complete", "-warn-concurrency"])
 
 
 
diff --git a/Sources/hylo-lsp-client/hylo-lsp-client-cli.swift b/Sources/hylo-lsp-client/hylo-lsp-client-cli.swift
index 86a8ede..ecd6d23 100644
--- a/Sources/hylo-lsp-client/hylo-lsp-client-cli.swift
+++ b/Sources/hylo-lsp-client/hylo-lsp-client-cli.swift
@@ -58,7 +58,7 @@ struct Options: ParsableArguments {
 
       #if os(Windows)
       // let search1 = try Regex(#"(.+)(?::(\d+)(?:\.(\d+))?)"#)
-      logger.warning("Document path parsing not currently supported on Windows, assuming normal filepath")
+      print("Document path parsing not currently supported on Windows, assuming normal filepath")
       let path = docLocation
       let url = resolveDocumentUrl(path)
       let uri = url.absoluteString
@@ -210,6 +210,28 @@ func resolveDocumentUri(_ uri: String) -> DocumentUri {
   return resolveDocumentUrl(uri).absoluteString
 }
 
+func printDiagnostic(_ d: LanguageServerProtocol.Diagnostic, in filepath: String) {
+  print("\(cliLink(uri: filepath, range: d.range)) \(d.severity ?? .information): \(d.message)")
+  for ri in d.relatedInformation ?? [] {
+    print("  \(cliLink(uri: ri.location.uri, range: ri.location.range)) \(ri.message)")
+  }
+}
+
+func withDiagnosticsCheck<T>(_ fn: () async throws -> T) async throws -> T {
+  do {
+    return try await fn()
+  }
+  catch let d as DiagnosticSet {
+
+    for d in d.elements {
+      let _d = LanguageServerProtocol.Diagnostic(d)
+      printDiagnostic(_d, in: d.site.file.url.path)
+    }
+
+    throw d
+  }
+}
+
 protocol DocumentCommand : AsyncParsableCommand {
   func process(doc: DocumentLocation, using server: Server) async throws
 }
@@ -242,7 +264,6 @@ extension HyloLspCommand {
 
       let response = try await server.documentSymbol(params: params)
 
-
       switch response {
         case nil:
           print("No symbols")
@@ -348,10 +369,13 @@ extension HyloLspCommand {
     func process(doc: DocumentLocation, using server: Server) async throws {
       let params = DocumentDiagnosticParams(textDocument: TextDocumentIdentifier(uri: doc.uri))
       let report = try await server.diagnostics(params: params)
-      for i in report.items ?? [] {
-        print("\(cliLink(uri: doc.filepath, range: i.range)) \(i.severity ?? .information): \(i.message)")
-        for ri in i.relatedInformation ?? [] {
-          print("  \(cliLink(uri: ri.location.uri, range: ri.location.range)) \(ri.message)")
+      for d in report.items ?? [] {
+        printDiagnostic(d, in: doc.filepath)
+      }
+
+      for (f, r) in report.relatedDocuments ?? [:] {
+        for d in r.items ?? [] {
+          printDiagnostic(d, in: f)
         }
       }
     }
@@ -442,7 +466,7 @@ extension HyloLspCommand {
       try socket.bind()
       try socket.listen()
       print("Created socket pipe: \(pipe)")
-      let client = try socket.accept()
+      _ = try socket.accept()
       print("LSP attached")
       // client.timeout = (connect: 5, read: nil, write: 5)
       // let clientChannel = DataChannel(socket: client)
diff --git a/Sources/hylo-lsp-server/hylo-lsp-server-cli.swift b/Sources/hylo-lsp-server/hylo-lsp-server-cli.swift
index d91bde1..b93d505 100644
--- a/Sources/hylo-lsp-server/hylo-lsp-server-cli.swift
+++ b/Sources/hylo-lsp-server/hylo-lsp-server-cli.swift
@@ -74,7 +74,7 @@ struct HyloLspCommand: AsyncParsableCommand {
     }
 
     func logHandlerFactory(_ label: String, fileLogger: FileLogger) -> LogHandler {
-      if ServerState.disableLogging {
+      if HyloServer.disableLogging {
         return NullLogHandler(label: label)
       }
 
diff --git a/Sources/hylo-lsp/AST+StdLibrary.swift b/Sources/hylo-lsp/AST+StdLibrary.swift
index c68aed4..e6d66b2 100644
--- a/Sources/hylo-lsp/AST+StdLibrary.swift
+++ b/Sources/hylo-lsp/AST+StdLibrary.swift
@@ -3,14 +3,20 @@ import FrontEnd
 import Foundation
 
 extension AST {
-  internal init(libraryRoot: URL) throws {
+  internal init(sourceFiles: [SourceFile]) throws {
     self.init()
     var diagnostics = DiagnosticSet()
     coreLibrary = try makeModule(
       "Hylo",
-      sourceCode: sourceFiles(in: [libraryRoot]),
+      sourceCode: sourceFiles,
       builtinModuleAccess: true,
       diagnostics: &diagnostics)
+
     assert(isCoreModuleLoaded)
   }
+
+  internal init(libraryRoot: URL) throws {
+    let sourceFiles = try sourceFiles(in: [libraryRoot])
+    try self.init(sourceFiles: sourceFiles)
+  }
 }
diff --git a/Sources/hylo-lsp/Document.swift b/Sources/hylo-lsp/Document.swift
index 66e15f3..0b500d8 100644
--- a/Sources/hylo-lsp/Document.swift
+++ b/Sources/hylo-lsp/Document.swift
@@ -3,91 +3,116 @@ import Core
 import FrontEnd
 import Foundation
 
-public struct DocumentProfiling {
-  public let stdlibParsing: TimeInterval
-  public let ASTParsing: TimeInterval
-  public let typeChecking: TimeInterval
-}
 
-public struct AnalyzedDocument {
+public struct Document {
   public let uri: DocumentUri
-  public let program: TypedProgram
-  public let ast: AST
-  public let profiling: DocumentProfiling
+  public let version: Int?
+  public let text: String
 
-  public init(uri: DocumentUri, ast: AST, program: TypedProgram, profiling: DocumentProfiling) {
+  public init(uri: DocumentUri, version: Int?, text: String) {
     self.uri = uri
-    self.ast = ast
-    self.program = program
-    self.profiling = profiling
+    self.version = version
+    self.text = text
   }
-}
 
-// public struct CachedDocumentResult: Codable {
-//   public var uri: DocumentUri
-//   public var symbols: DocumentSymbolResponse?
-//   public var semanticTokens: SemanticTokensResponse?
-// }
+  public init(textDocument: TextDocumentItem) {
+    uri = textDocument.uri
+    version = textDocument.version
+    text = textDocument.text
+  }
+}
 
+struct InvalidDocumentChangeRange : Error {
+  public let range: LSPRange
+}
 
-public actor DocumentContext {
-  public var uri: DocumentUri { request.uri }
-  public let request: DocumentBuildRequest
-  private var analyzedDocument: Result<AnalyzedDocument, Error>?
+extension Document {
 
-  public init(_ request: DocumentBuildRequest) {
-    self.request = request
-    Task {
-      await self.monitorTasks()
+  public func withAppliedChanges(_ changes: [TextDocumentContentChangeEvent], nextVersion: Int?) throws -> Document {
+    var text = self.text
+    for c in changes {
+      try Document.applyChange(c, on: &text)
     }
-  }
 
-  public func pollAnalyzedDocument() -> Result<AnalyzedDocument, Error>? {
-    return analyzedDocument
+    return Document(uri: uri, version: nextVersion, text: text)
   }
 
-  public func getAnalyzedDocument() async -> Result<AnalyzedDocument, Error> {
-    do {
-      let doc = try await request.buildTask.value
-      return .success(doc)
-    }
-    catch {
-      return .failure(error)
+  private static func findPosition(_ position: Position, in text: String, startingFrom: String.Index) -> String.Index? {
+
+    var it = text[startingFrom...]
+    for _ in 0..<position.line {
+      guard let i = it.firstIndex(of: "\n") else {
+        return nil
+      }
+
+      it = it[it.index(after: i)...]
     }
+
+    return text.index(it.startIndex, offsetBy: position.character)
   }
 
-  public func getAST() async -> Result<AST, Error> {
-    do {
-      let ast = try await request.astTask.value
-      return .success(ast)
+  private static func findRange(_ range: LSPRange, in text: String) -> Range<String.Index>? {
+    guard let startIndex = findPosition(range.start, in: text, startingFrom: text.startIndex) else {
+      return nil
     }
-    catch {
-      return .failure(error)
+
+    guard let endIndex = findPosition(range.end, in: text, startingFrom: startIndex) else {
+      return nil
     }
+
+    return startIndex..<endIndex
   }
 
-  private func monitorTasks() {
+  private static func applyChange(_ change: TextDocumentContentChangeEvent, on text: inout String) throws {
+    if let range = change.range {
+      guard let range = findRange(range, in: text) else {
+        throw InvalidDocumentChangeRange(range: range)
+      }
 
-    Task {
-      self.analyzedDocument = await getAnalyzedDocument()
+      text.replaceSubrange(range, with: change.text)
+    }
+    else {
+      text = change.text
     }
   }
 }
 
-public struct DocumentBuildRequest {
+public struct DocumentProfiling {
+  public let stdlibParsing: TimeInterval
+  public let ASTParsing: TimeInterval
+  public let typeChecking: TimeInterval
+}
+
+public struct AnalyzedDocument {
   public let uri: DocumentUri
-  public let astTask: Task<AST, Error>
-  public let buildTask: Task<AnalyzedDocument, Error>
+  public let program: TypedProgram
+  public let ast: AST
+  public let profiling: DocumentProfiling
 
-  public init(uri: DocumentUri, astTask: Task<AST, Error>, buildTask: Task<AnalyzedDocument, Error>) {
+  public init(uri: DocumentUri, ast: AST, program: TypedProgram, profiling: DocumentProfiling) {
     self.uri = uri
-    self.astTask = astTask
-    self.buildTask = buildTask
+    self.ast = ast
+    self.program = program
+    self.profiling = profiling
   }
 }
 
+extension DocumentProvider {
+  // This should really be a struct since we are building for Hylo
+  class DocumentContext {
+    public var doc: Document
+    public var uri: DocumentUri { doc.uri }
+    var astTask: Task<AST, Error>?
+    var buildTask: Task<AnalyzedDocument, Error>?
+
+    public init(_ doc: Document) {
+      self.doc = doc
+    }
+  }
+}
+
+
 public enum DocumentError : Error {
   case diagnostics(DiagnosticSet)
   case other(Error)
 }
-
diff --git a/Sources/hylo-lsp/HyloServerState.swift b/Sources/hylo-lsp/DocumentProvider.swift
similarity index 69%
rename from Sources/hylo-lsp/HyloServerState.swift
rename to Sources/hylo-lsp/DocumentProvider.swift
index 8d2fc29..fa93a3e 100644
--- a/Sources/hylo-lsp/HyloServerState.swift
+++ b/Sources/hylo-lsp/DocumentProvider.swift
@@ -15,7 +15,12 @@ extension TextDocumentIdentifier : TextDocumentProtocol {}
 extension TextDocumentItem : TextDocumentProtocol {}
 extension VersionedTextDocumentIdentifier : TextDocumentProtocol {}
 
-public actor ServerState {
+public enum GetDocumentContextError : Error {
+  case invalidUri(DocumentUri)
+  case documentNotOpened(DocumentUri)
+}
+
+public actor DocumentProvider {
   private var documents: [DocumentUri:DocumentContext]
   public let logger: Logger
   let lsp: JSONRPCServer
@@ -24,7 +29,6 @@ public actor ServerState {
   var stdlibCache: [URL:AST]
 
   public let defaultStdlibFilepath: URL
-  public static let disableLogging = if let disableLogging = ProcessInfo.processInfo.environment["HYLO_LSP_DISABLE_LOGGING"] { !disableLogging.isEmpty } else { false }
 
   public init(lsp: JSONRPCServer, logger: Logger) {
     self.logger = logger
@@ -32,7 +36,7 @@ public actor ServerState {
     stdlibCache = [:]
     self.lsp = lsp
     self.workspaceFolders = []
-    defaultStdlibFilepath = ServerState.loadDefaultStdlibFilepath(logger: logger)
+    defaultStdlibFilepath = DocumentProvider.loadDefaultStdlibFilepath(logger: logger)
   }
 
 
@@ -113,16 +117,6 @@ public actor ServerState {
     workspaceFolders.append(contentsOf: added)
   }
 
-  // private static func loadStdlibProgram() throws -> TypedProgram {
-  //   let ast = try AST(libraryRoot: defaultStdlibFilepath)
-
-  //   var diagnostics = DiagnosticSet()
-  //   return try TypedProgram(
-  //   annotating: ScopedProgram(ast), inParallel: true,
-  //   reportingDiagnosticsTo: &diagnostics,
-  //   tracingInferenceIf: nil)
-  // }
-
   private static func loadDefaultStdlibFilepath(logger: Logger) -> URL {
     if let path = ProcessInfo.processInfo.environment["HYLO_STDLIB_PATH"] {
       logger.info("Hylo stdlib filepath from HYLO_STDLIB_PATH: \(path)")
@@ -200,26 +194,6 @@ public actor ServerState {
     return closest
   }
 
-
-  private func requestDocument(_ uri: DocumentUri) -> DocumentContext {
-    let (stdlibPath, isStdlibDocument) = getStdlibPath(uri)
-
-    let input = URL.init(string: uri)!
-    let inputs: [URL] = if !isStdlibDocument { [input] } else { [] }
-
-    let astTask = Task {
-      return try buildAst(uri: uri, stdlibPath: stdlibPath, inputs: inputs)
-    }
-
-    let buildTask = Task {
-      let ast = try await astTask.value
-      return try buildProgram(uri: uri, ast: ast)
-    }
-
-    let request = DocumentBuildRequest(uri: uri, astTask: astTask, buildTask: buildTask)
-    return DocumentContext(request)
-  }
-
   func uriAsFilepath(_ uri: DocumentUri) -> String? {
     guard let url = URL.init(string: uri) else {
       return nil
@@ -228,21 +202,44 @@ public actor ServerState {
     return url.path
   }
 
-  private func buildAst(uri: DocumentUri, stdlibPath: URL, inputs: [URL]) throws -> AST {
-    var diagnostics = DiagnosticSet()
-    logger.debug("Build ast for document: \(uri), with stdlibPath: \(stdlibPath), inputs: \(inputs)")
+  private func buildStdlibAST(_ stdlibPath: URL) throws -> AST {
+    let sourceFiles = try sourceFiles(in: [stdlibPath]).map { file in
 
-    var ast: AST
-    if let stdlib = stdlibCache[stdlibPath] {
-      ast = stdlib
+      // We need to replace stdlib files from in memory buffers if they have unsaved changes
+      if let context = documents[file.url.absoluteString] {
+        logger.debug("Replace content for stdlib source file: \(file.url)")
+        return SourceFile(filePath: file.url, withContent: context.doc.text)
+      }
+      else {
+        return file
+      }
+    }
+
+    return try AST(sourceFiles: sourceFiles)
+  }
+
+  // We cache stdlib AST, and since AST is struct the cache values are implicitly immutable (thanks MVS!)
+  private func getStdlibAST(_ stdlibPath: URL) throws -> AST {
+    if let ast = stdlibCache[stdlibPath] {
+      return ast
     }
     else {
-      ast = try AST(libraryRoot: stdlibPath)
+      let ast = try buildStdlibAST(stdlibPath)
       stdlibCache[stdlibPath] = ast
+      return ast
     }
+  }
 
-    if !inputs.isEmpty {
-      _ = try ast.makeModule(HyloNotificationHandler.productName, sourceCode: sourceFiles(in: inputs), builtinModuleAccess: false, diagnostics: &diagnostics)
+  private func buildAST(uri: DocumentUri, stdlibPath: URL, sourceFiles: [SourceFile]) throws -> AST {
+    var diagnostics = DiagnosticSet()
+    logger.debug("Build ast for document: \(uri), with stdlibPath: \(stdlibPath)")
+
+    var ast = try getStdlibAST(stdlibPath)
+
+    if !sourceFiles.isEmpty {
+        let productName = "lsp-build"
+        // let sourceFiles = try sourceFiles(in: inputs)
+      _ = try ast.makeModule(productName, sourceCode: sourceFiles, builtinModuleAccess: false, diagnostics: &diagnostics)
     }
 
     return ast
@@ -294,30 +291,92 @@ public actor ServerState {
     return nil
   }
 
-  // public func preloadDocument(_ textDocument: TextDocumentProtocol) -> DocumentBuildRequest {
-  //   let uri = DocumentProvider.resolveDocumentUri(textDocument.uri)
-  //   return preloadDocument(uri)
-  // }
+  public func updateDocument(_ params: TextDocumentDidChangeParams) {
+    let uri = params.textDocument.uri
+    guard let context = documents[uri] else {
+      logger.error("Could not find opened document: \(uri)")
+      return
+    }
+
+    do {
+      let updatedDoc = try context.doc.withAppliedChanges(params.contentChanges, nextVersion: params.textDocument.version)
+      context.doc = updatedDoc
+      context.astTask = nil
+      context.buildTask = nil
+      logger.debug("Updated changed document: \(uri), version: \(updatedDoc.version ?? -1)")
+
+      // NOTE: We also need to invalidate cached stdlib AST if the edited document is part of the stdlib
+      let (stdlibPath, isStdlibDocument) = getStdlibPath(uri)
+      if isStdlibDocument {
+        stdlibCache[stdlibPath] = nil
+      }
+    }
+    catch {
+      logger.error("Failed to apply document changes")
+    }
+  }
+
+  public func registerDocument(_ params: TextDocumentDidOpenParams) {
+    let doc = Document(textDocument: params.textDocument)
+    let context = DocumentContext(doc)
+    // requestDocument(doc)
+    logger.debug("Register opened document: \(doc.uri)")
+    documents[doc.uri] = context
+  }
+
+  public func unregisterDocument(_ params: TextDocumentDidCloseParams) {
+    let uri = params.textDocument.uri
+    documents[uri] = nil
+  }
+
+
+  func implicitlyRegisterDocument(_ uri: DocumentUri)-> DocumentContext? {
+    guard let url = URL.init(string: uri) else {
+      return nil
+    }
+
+    guard let text = try? String(contentsOf: url) else {
+      return nil
+    }
+
+    let doc = Document(uri: uri, version: 0, text: text)
+    return DocumentContext(doc)
+  }
+
 
-  private func preloadDocument(_ uri: DocumentUri) -> DocumentContext {
-    let document = requestDocument(uri)
-    logger.debug("Register opened document: \(uri)")
-    documents[uri] = document
-    return document
+  func getDocumentContext(_ textDocument: TextDocumentProtocol) -> Result<DocumentContext, GetDocumentContextError> {
+    getDocumentContext(textDocument.uri)
   }
 
-  public func getDocumentContext(_ textDocument: TextDocumentProtocol) -> Result<DocumentContext, InvalidUri> {
-    guard let uri = ServerState.validateDocumentUri(textDocument.uri) else {
-      return .failure(InvalidUri(textDocument.uri))
+  func getDocumentContext(_ uri: DocumentUri) -> Result<DocumentContext, GetDocumentContextError> {
+    guard let uri = DocumentProvider.validateDocumentUri(uri) else {
+      return .failure(.invalidUri(uri))
+    }
+
+    guard let context = documents[uri] else {
+      // NOTE: We can not assume document is opened, VSCode apparently does not guarantee ordering
+      // Specifically `textDocument/diagnostic` -> `textDocument/didOpen` has been observed
+
+      // return .failure(.documentNotOpened(uri))
+
+      logger.warning("Implicitly registering unopened document: \(uri)")
+      if let context = implicitlyRegisterDocument(uri) {
+        return .success(context)
+      }
+      else {
+        return .failure(.invalidUri(uri))
+      }
     }
 
-    // Check for cached document
-    if let request = documents[uri] {
-      logger.info("Found cached document: \(uri)")
-      return .success(request)
-    } else {
-      let request = preloadDocument(uri)
-      return .success(request)
+    return .success(context)
+  }
+
+  public func getAST(_ textDocument: TextDocumentProtocol) async -> Result<AST, DocumentError> {
+    switch getDocumentContext(textDocument) {
+      case let .failure(error):
+        return .failure(.other(error))
+      case let .success(context):
+        return await getAST(context)
     }
   }
 
@@ -326,14 +385,60 @@ public actor ServerState {
       case let .failure(error):
         return .failure(.other(error))
       case let .success(context):
-        return await resolveDocumentRequest(context.request)
+        return await getAnalyzedDocument(context)
+    }
+  }
+
+  private func createASTTask(_ context: DocumentContext) -> Task<AST, Error> {
+    if context.astTask == nil {
+      let uri = context.uri
+      let (stdlibPath, isStdlibDocument) = getStdlibPath(uri)
+
+      let sourceFiles: [SourceFile]
+
+      if isStdlibDocument {
+        sourceFiles = []
+      }
+      else {
+        let url = URL.init(string: uri)!
+        sourceFiles = [SourceFile(filePath: url, withContent: context.doc.text)]
+      }
+
+      context.astTask = Task {
+        return try buildAST(uri: uri, stdlibPath: stdlibPath, sourceFiles: sourceFiles)
+      }
+    }
+
+    return context.astTask!
+  }
+
+  private func getAST(_ context: DocumentContext) async -> Result<AST, DocumentError> {
+    do {
+      let astTask = createASTTask(context)
+      let ast = try await astTask.value
+      return .success(ast)
+    }
+    catch let d as DiagnosticSet {
+      // NOTE: We want to be able to use partial AST, need to update Hylo call to not throw
+      return .failure(.diagnostics(d))
+    }
+    catch {
+      return .failure(.other(error))
     }
   }
 
-  public func resolveDocumentRequest(_ request: DocumentBuildRequest) async -> Result<AnalyzedDocument, DocumentError> {
+  private func getAnalyzedDocument(_ context: DocumentContext) async -> Result<AnalyzedDocument, DocumentError> {
+    if context.buildTask == nil {
+      context.buildTask = Task {
+        let astTask = createASTTask(context)
+        let ast = try await astTask.value
+        return try buildProgram(uri: context.uri, ast: ast)
+      }
+    }
+
     do {
-      let document = try await request.buildTask.value
-      return .success(document)
+      let doc = try await context.buildTask!.value
+      return .success(doc)
     }
     catch let d as DiagnosticSet {
       return .failure(.diagnostics(d))
@@ -343,6 +448,8 @@ public actor ServerState {
     }
   }
 
+
+#if false
   // NOTE: We currently write cached results inside the workspace
   // These should perhaps be stored outside workspace, but then it is more important
   // to implement some kind of garbage collection for out-dated workspace cache entries
@@ -350,7 +457,6 @@ public actor ServerState {
     NSString.path(withComponents: [uriAsFilepath(wsFile.workspace)!, ".hylo-lsp", "cache", wsFile.relativePath + ".json"])
   }
 
-#if false
   private func loadCachedDocumentResult(_ uri: DocumentUri) -> CachedDocumentResult? {
     do {
       guard let filepath = uriAsFilepath(uri) else {
diff --git a/Sources/hylo-lsp/HyloServer.swift b/Sources/hylo-lsp/HyloServer.swift
index 100caab..470a683 100644
--- a/Sources/hylo-lsp/HyloServer.swift
+++ b/Sources/hylo-lsp/HyloServer.swift
@@ -17,10 +17,8 @@ enum BuildError : Error {
 public struct HyloNotificationHandler : NotificationHandler {
   public let lsp: JSONRPCServer
   public let logger: Logger
-  var state: ServerState
-  // var ast: AST { state.ast }
-  static let productName = "lsp-build"
-
+  var documentProvider: DocumentProvider
+  // var ast: AST { documentProvider.ast }
 
   private func withErrorLogging(_ fn: () throws -> Void) {
     do {
@@ -51,16 +49,15 @@ public struct HyloNotificationHandler : NotificationHandler {
 
 
   public func textDocumentDidOpen(_ params: TextDocumentDidOpenParams) async {
-    // _ = await state.documentProvider.preloadDocument(params.textDocument)
+    await documentProvider.registerDocument(params)
   }
 
   public func textDocumentDidChange(_ params: TextDocumentDidChangeParams) async {
-    // _ = await state.documentProvider.preloadDocument(params.textDocument)
-    // TODO: Handle changes from input (not stored on disk)
+    await documentProvider.updateDocument(params)
   }
 
   public func textDocumentDidClose(_ params: TextDocumentDidCloseParams) async {
-
+    await documentProvider.unregisterDocument(params)
   }
 
   public func textDocumentWillSave(_ params: TextDocumentWillSaveParams) async {
@@ -71,7 +68,8 @@ public struct HyloNotificationHandler : NotificationHandler {
   }
 
   public func protocolCancelRequest(_ params: CancelParams) async {
-
+    // NOTE: For cancel to work we must pass JSONRPC request ids to handlers
+    logger.debug("Cancel request: \(params.id)")
   }
 
   public func protocolSetTrace(_ params: SetTraceParams) async {
@@ -87,7 +85,7 @@ public struct HyloNotificationHandler : NotificationHandler {
   }
 
   public func workspaceDidChangeWorkspaceFolders(_ params: DidChangeWorkspaceFoldersParams) async {
-    await state.workspaceDidChangeWorkspaceFolders(params)
+    await documentProvider.workspaceDidChangeWorkspaceFolders(params)
   }
 
   public func workspaceDidChangeConfiguration(_ params: DidChangeConfigurationParams)  async {
@@ -113,21 +111,21 @@ public struct HyloRequestHandler : RequestHandler {
   public let lsp: JSONRPCServer
   public let logger: Logger
 
-  var state: ServerState
-  // var ast: AST { state.ast }
-  // var program: TypedProgram? { state.program }
+  var documentProvider: DocumentProvider
+  // var ast: AST { documentProvider.ast }
+  // var program: TypedProgram? { documentProvider.program }
   // var initTask: Task<TypedProgram, Error>
 
 
-  public init(lsp: JSONRPCServer, logger: Logger, state: ServerState) {
+  public init(lsp: JSONRPCServer, logger: Logger, documentProvider: DocumentProvider) {
     self.lsp = lsp
     self.logger = logger
-    self.state = state
+    self.documentProvider = documentProvider
   }
 
 
   public func initialize(_ params: InitializeParams) async -> Result<InitializationResponse, AnyJSONRPCResponseError> {
-    return await state.initialize(params)
+    return await documentProvider.initialize(params)
   }
 
   public func shutdown() async {
@@ -161,7 +159,7 @@ public struct HyloRequestHandler : RequestHandler {
 
   public func definition(_ params: TextDocumentPositionParams, _ doc: AnalyzedDocument) async -> Result<DefinitionResponse, AnyJSONRPCResponseError> {
 
-    guard let url = ServerState.validateDocumentUrl(params.textDocument.uri) else {
+    guard let url = DocumentProvider.validateDocumentUrl(params.textDocument.uri) else {
       return .failure(JSONRPCResponseError(code: ErrorCodes.InvalidParams, message: "Invalid document uri: \(params.textDocument.uri)"))
     }
 
@@ -325,18 +323,14 @@ public struct HyloRequestHandler : RequestHandler {
 
   public func documentSymbol(_ params: DocumentSymbolParams) async -> Result<DocumentSymbolResponse, AnyJSONRPCResponseError> {
 
-    await withDocumentContext(params.textDocument) { context in
-      let astResult = await context.getAST()
-      return await withDocument(astResult) { ast in
-        await documentSymbol(params, ast: ast)
-      }
+    await withDocumentAST(params.textDocument) { ast in
+      await documentSymbol(params, ast: ast)
     }
   }
 
   public func diagnostics(_ params: DocumentDiagnosticParams) async -> Result<DocumentDiagnosticReport, AnyJSONRPCResponseError> {
 
-
-    let docResult = await state.getAnalyzedDocument(params.textDocument)
+    let docResult = await documentProvider.getAnalyzedDocument(params.textDocument)
 
     switch docResult {
     case .success:
@@ -344,14 +338,26 @@ public struct HyloRequestHandler : RequestHandler {
     case let .failure(error):
       switch error {
       case let .diagnostics(d):
-      let dList = d.elements.map { LanguageServerProtocol.Diagnostic($0) }
-      return .success(RelatedDocumentDiagnosticReport(kind: .full, items: dList))
+      return .success(buildDiagnosticReport(uri: params.textDocument.uri, diagnostics: d))
       case .other:
         return .failure(JSONRPCResponseError(code: ErrorCodes.InternalError, message: "Unknown build error: \(error)"))
       }
     }
   }
 
+  func buildDiagnosticReport(uri: DocumentUri, diagnostics: DiagnosticSet) -> RelatedDocumentDiagnosticReport {
+    let matching = diagnostics.elements.filter { $0.site.file.url.absoluteString == uri }
+    let nonMatching = diagnostics.elements.filter { $0.site.file.url.absoluteString != uri }
+
+    let items = matching.map { LanguageServerProtocol.Diagnostic($0) }
+    let related = nonMatching.reduce(into: [String: LanguageServerProtocol.DocumentDiagnosticReport]()) {
+      let d = LanguageServerProtocol.Diagnostic($1)
+      $0[$1.site.file.url.absoluteString] = DocumentDiagnosticReport(kind: .full, items: [d])
+    }
+
+    return RelatedDocumentDiagnosticReport(kind: .full, items: items, relatedDocuments: related)
+  }
+
   func trySendDiagnostics(_ diagnostics: DiagnosticSet, in uri: DocumentUri) async {
     do {
       logger.debug("[\(uri)] send diagnostics")
@@ -395,7 +401,7 @@ public struct HyloRequestHandler : RequestHandler {
 
   func withAnalyzedDocument<ResponseT>(_ textDocument: TextDocumentIdentifier, fn: (AnalyzedDocument) async -> Result<ResponseT?, AnyJSONRPCResponseError>) async -> Result<ResponseT?, AnyJSONRPCResponseError> {
 
-    let docResult = await state.getAnalyzedDocument(textDocument)
+    let docResult = await documentProvider.getAnalyzedDocument(textDocument)
 
     switch docResult {
     case let .success(doc):
@@ -411,14 +417,18 @@ public struct HyloRequestHandler : RequestHandler {
     }
   }
 
-  func withDocumentContext<ResponseT>(_ textDocument: TextDocumentIdentifier, fn: (DocumentContext) async -> Result<ResponseT?, AnyJSONRPCResponseError>) async -> Result<ResponseT?, AnyJSONRPCResponseError> {
-    let result = await state.getDocumentContext(textDocument)
+  func withDocumentAST<ResponseT>(_ textDocument: TextDocumentIdentifier, fn: (AST) async -> Result<ResponseT?, AnyJSONRPCResponseError>) async -> Result<ResponseT?, AnyJSONRPCResponseError> {
+    let result = await documentProvider.getAST(textDocument)
 
     switch result {
       case let .failure(error):
-        return .failure(JSONRPCResponseError(code: ErrorCodes.InvalidParams, message: error.localizedDescription))
-      case let .success(context):
-        return await fn(context)
+        let errorMsg = switch error {
+        case .diagnostics: "Failed to build AST"
+        case let .other(e): e.localizedDescription
+        }
+        return .failure(JSONRPCResponseError(code: ErrorCodes.InvalidParams, message: errorMsg))
+      case let .success(ast):
+        return await fn(ast)
     }
   }
 
@@ -427,11 +437,8 @@ public struct HyloRequestHandler : RequestHandler {
   // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens
   public func semanticTokensFull(_ params: SemanticTokensParams) async -> Result<SemanticTokensResponse, AnyJSONRPCResponseError> {
 
-    await withDocumentContext(params.textDocument) { context in
-      let astResult = await context.getAST()
-      return await withDocument(astResult) { ast in
-        await semanticTokensFull(params, ast: ast)
-      }
+    await withDocumentAST(params.textDocument) { ast in
+      await semanticTokensFull(params, ast: ast)
     }
   }
 
@@ -445,26 +452,19 @@ public struct HyloRequestHandler : RequestHandler {
 
 public actor HyloServer {
   let lsp: JSONRPCServer
-  // private var ast: AST
   private let logger: Logger
-  private var state: ServerState
+  private var documentProvider: DocumentProvider
   private let requestHandler: HyloRequestHandler
   private let notificationHandler: HyloNotificationHandler
 
+  public static let disableLogging = if let disableLogging = ProcessInfo.processInfo.environment["HYLO_LSP_DISABLE_LOGGING"] { !disableLogging.isEmpty } else { false }
+
   public init(_ dataChannel: DataChannel, logger: Logger) {
     self.logger = logger
     lsp = JSONRPCServer(dataChannel)
-    self.state = ServerState(lsp: lsp, logger: logger)
-    requestHandler = HyloRequestHandler(lsp: lsp, logger: logger, state: state)
-    notificationHandler = HyloNotificationHandler(lsp: lsp, logger: logger, state: state)
-
-    // Task {
-		// 	await monitorRequests()
-		// }
-
-    // Task {
-		// 	await monitorNotifications()
-		// }
+    self.documentProvider = DocumentProvider(lsp: lsp, logger: logger)
+    requestHandler = HyloRequestHandler(lsp: lsp, logger: logger, documentProvider: documentProvider)
+    notificationHandler = HyloNotificationHandler(lsp: lsp, logger: logger, documentProvider: documentProvider)
   }
 
   public func run() async {
diff --git a/Sources/hylo-lsp/Range.swift b/Sources/hylo-lsp/Range.swift
index 0808171..6e8e0fa 100644
--- a/Sources/hylo-lsp/Range.swift
+++ b/Sources/hylo-lsp/Range.swift
@@ -9,9 +9,16 @@ public extension LanguageServerProtocol.Location {
 
 public extension LanguageServerProtocol.LSPRange {
   init(_ range: SourceRange) {
-    var (first, last) = (range.first(), range.last()!)
-    let incLast = range.file.text.index(after: last.index)
-    last = SourcePosition(incLast, in: last.file)
+    let first = range.first()
+    let last: SourcePosition
+
+    if let l = range.last() {
+      let incLast = range.file.text.index(after: l.index)
+      last = SourcePosition(incLast, in: l.file)
+    }
+    else {
+      last = SourcePosition(range.file.text.endIndex, in: range.file)
+    }
 
     self.init(start: Position(first), end: Position(last))
   }
diff --git a/Tests/hylo-lspTests/HyloServerTests.swift b/Tests/hylo-lspTests/HyloServerTests.swift
index 866c193..918df6d 100644
--- a/Tests/hylo-lspTests/HyloServerTests.swift
+++ b/Tests/hylo-lspTests/HyloServerTests.swift
@@ -13,50 +13,89 @@ func XCTUnwrapAsync<T>(_ expression: @autoclosure () async throws -> T?, _ messa
 }
 
 
-final class val_lspTests: XCTestCase {
-
-    func testFindDocumentRelativeWorkspacePath() async throws {
-        logger = Logger(label: loggerLabel) { label in
-          StreamLogHandler.standardOutput(label: label)
-        }
-
-        logger.logLevel = .debug
-
-        let dataChannel = DataChannel.stdioPipe()
-        let lsp = JSONRPCServer(dataChannel)
-        let state = ServerState(lsp: lsp)
-
-        let caps = ClientCapabilities(workspace: nil, textDocument: nil, window: nil, general: nil, experimental: nil)
-
-        let initParam = InitializeParams(
-          processId: 1,
-          locale: nil,
-          rootPath: nil,
-          rootUri: "/foo/a",
-          initializationOptions: nil,
-          capabilities: caps,
-          trace: nil,
-          workspaceFolders: [
-            WorkspaceFolder(uri: "/foo/b", name: "b"),
-            WorkspaceFolder(uri: "/foo/b/c", name: "b/c"),
-          ]
-        )
-
-        _ = await state.initialize(initParam)
-
-        var ws1 = await state.getWorkspaceFile("/foo/a/x.hylo")
-        var ws = try XCTUnwrap(ws1)
-        XCTAssert(ws.relativePath == "x.hylo")
-        XCTAssert(ws.workspace == "/foo/a")
-
-        ws1 = await state.getWorkspaceFile("/foo/b/x.hylo")
-        ws = try XCTUnwrap(ws1)
-        XCTAssert(ws.relativePath == "x.hylo")
-        XCTAssert(ws.workspace == "/foo/b")
-
-        ws1 = await state.getWorkspaceFile("/foo/b/c/x.hylo")
-        ws = try XCTUnwrap(ws1)
-        XCTAssert(ws.relativePath == "x.hylo")
-        XCTAssert(ws.workspace == "/foo/b/c")
+final class hyloLspTests: XCTestCase {
+  func createLogger() -> Logger {
+    var logger = Logger(label: loggerLabel) { label in
+      StreamLogHandler.standardOutput(label: label)
     }
+
+    logger.logLevel = .debug
+    return logger
+  }
+
+  func testApplyDocumentChanges() async throws {
+    let uri = "file:///factorial.hylo"
+    let beforeEdit = """
+      fun factorial(_ n: Int) -> Int {
+        if n < 2 { 1 } else { n * factorial(n - 1) }
+      }
+
+      public fun main() {
+        let _ = factorial(6)
+      }
+    """
+
+    let afterEdit = """
+      fun foo(_ n: Int) -> Int {
+        if n < 2 { 1 } else { n * factorial(n - 1) }
+      }
+      public fun main() {
+        let _ = foo(123)
+      }
+    """
+
+    let textDocument = TextDocumentItem(uri: uri, languageId: "hylo", version: 0, text: beforeEdit)
+
+    let doc = Document(textDocument: textDocument)
+
+    let changes = [
+      TextDocumentContentChangeEvent(range: LSPRange(startPair: (0, 4), endPair: (0, 13)), rangeLength: nil, text: "foo"),
+      TextDocumentContentChangeEvent(range: LSPRange(startPair: (3, 0), endPair: (3, 1)), rangeLength: nil, text: ""),
+      TextDocumentContentChangeEvent(range: LSPRange(startPair: (4, 10), endPair: (4, 19)), rangeLength: nil, text: "foo"),
+      TextDocumentContentChangeEvent(range: LSPRange(startPair: (4, 14), endPair: (4, 15)), rangeLength: nil, text: "123"),
+    ]
+
+    let updatedDoc = try doc.withAppliedChanges(changes, nextVersion: 2)
+    XCTAssertNotEqual(updatedDoc.text, afterEdit)
+  }
+
+  func testFindDocumentRelativeWorkspacePath() async throws {
+    let logger = createLogger()
+    let dataChannel = DataChannel.stdioPipe()
+    let lsp = JSONRPCServer(dataChannel)
+    let documentProvider = DocumentProvider(lsp: lsp, logger: logger)
+
+    let caps = ClientCapabilities(workspace: nil, textDocument: nil, window: nil, general: nil, experimental: nil)
+
+    let initParam = InitializeParams(
+      processId: 1,
+      locale: nil,
+      rootPath: nil,
+      rootUri: "/foo/a",
+      initializationOptions: nil,
+      capabilities: caps,
+      trace: nil,
+      workspaceFolders: [
+        WorkspaceFolder(uri: "/foo/b", name: "b"),
+        WorkspaceFolder(uri: "/foo/b/c", name: "b/c"),
+      ]
+    )
+
+    _ = await documentProvider.initialize(initParam)
+
+    var ws1 = await documentProvider.getWorkspaceFile("/foo/a/x.hylo")
+    var ws = try XCTUnwrap(ws1)
+    XCTAssert(ws.relativePath == "x.hylo")
+    XCTAssert(ws.workspace == "/foo/a")
+
+    ws1 = await documentProvider.getWorkspaceFile("/foo/b/x.hylo")
+    ws = try XCTUnwrap(ws1)
+    XCTAssert(ws.relativePath == "x.hylo")
+    XCTAssert(ws.workspace == "/foo/b")
+
+    ws1 = await documentProvider.getWorkspaceFile("/foo/b/c/x.hylo")
+    ws = try XCTUnwrap(ws1)
+    XCTAssert(ws.relativePath == "x.hylo")
+    XCTAssert(ws.workspace == "/foo/b/c")
+  }
 }
diff --git a/hylo b/hylo
index 2c393d6..fe7c1f4 160000
--- a/hylo
+++ b/hylo
@@ -1 +1 @@
-Subproject commit 2c393d6c959933bb7c0d51ba9bd50546e1d60248
+Subproject commit fe7c1f4d3c495600a5c50abcb4f1dd8eb9b028b4
diff --git a/hylo-vscode-extension b/hylo-vscode-extension
index 1af1fb5..3542951 160000
--- a/hylo-vscode-extension
+++ b/hylo-vscode-extension
@@ -1 +1 @@
-Subproject commit 1af1fb50027652220b77959792a48b8bde330d0e
+Subproject commit 35429510205a2c6b081910a493e8dabba6ba8fda