diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
index 71dc5e1e2b..b5623f3ff8 100755
--- a/sqlmesh/lsp/main.py
+++ b/sqlmesh/lsp/main.py
@@ -1074,8 +1074,8 @@ def _context_get_or_load(self, document_uri: t.Optional[URI] = None) -> LSPConte
                     loaded_sqlmesh_message(self.server)
             else:
                 self._ensure_context_for_document(document_uri)
-        if isinstance(state, ContextLoaded):
-            return state.lsp_context
+        if isinstance(self.context_state, ContextLoaded):
+            return self.context_state.lsp_context
         raise RuntimeError("Context failed to load")
 
     def _ensure_context_for_document(