feat: command completions when using bolt

Author: TheNuclearNexus

aegis-server/aegis_server/__main__.py

@@ -35,7 +35,6 @@ def did_change(ls: AegisServer, params: lsp.DidChangeTextDocumentParams):
     def did_open(ls: AegisServer, params: lsp.DidOpenTextDocumentParams):
         asyncio.run(publish_diagnostics(ls, params))
 
-
     @server.thread()
     @server.feature(
         lsp.TEXT_DOCUMENT_COMPLETION,
@@ -82,7 +81,7 @@ def hover(ls: AegisServer, params: lsp.HoverParams):
     @server.feature(lsp.TEXT_DOCUMENT_RENAME)
     def rename(ls: AegisServer, params: lsp.RenameParams):
         return asyncio.run(rename_variable(ls, params))
-    
+
     @server.command("mecha.server.dumpIndices")
     def dump(ls: AegisServer, *args):
         for _, i in ls._instances.values():
@@ -117,6 +116,7 @@ def add_arguments(parser: argparse.ArgumentParser):
         help="Show the AST node for the hovered token",
     )
 
+
 def main():
     print("Starting Aegis Server")
     parser = argparse.ArgumentParser()

aegis-server/aegis_server/server/features/completion.py

@@ -1,7 +1,8 @@
 import builtins
+import logging
 
 from aegis_server.providers.variable import add_raw_definition, add_variable_definition
-from bolt import Runtime, UndefinedIdentifier
+from bolt import Runtime, UnboundLocalIdentifier, UndefinedIdentifier
 from lsprotocol import types as lsp
 from mecha import (
     AstOption,
@@ -133,7 +134,10 @@ def get_diag_completions(
                     pattern if isinstance(pattern, tuple) else [pattern, None]
                 )
                 items += get_token_options(mecha, token_type, value)
-            break
+            
+        if isinstance(diagnostic, UnboundLocalIdentifier):
+            for name in mecha.spec.tree.children.keys(): #type: ignore
+                items.append(lsp.CompletionItem(name, kind = lsp.CompletionItemKind.Keyword))
 
         if isinstance(diagnostic, UndefinedIdentifier):
             for name, variable in diagnostic.lexical_scope.variables.items():
@@ -145,5 +149,4 @@ def get_diag_completions(
             for name in runtime.builtins:
                 add_raw_definition(items, name, getattr(builtins, name))
 
-            break
     return lsp.CompletionList(False, items)

aegis-server/aegis_server/server/features/validate.py

@@ -223,8 +223,9 @@ async def parse_function(
 
     # # Parse the stream
     mecha = ctx.inject(Mecha)
-    compilation_unit = mecha.database[file_instance]
     runtime = ctx.inject(Runtime)
+
+    compilation_unit = mecha.database[file_instance]
     compiled_module = runtime.modules.registry.get(file_instance)
 
     return CompiledDocument(
@@ -355,5 +356,5 @@ async def compile(
                     logging.error("\n".join(traceback.format_tb(cause.__traceback__)))
 
             logging.debug(f"Execution took {time.time() - start}s")
-            await asyncio.sleep(0)
+            
     return indexer.output_ast, diagnostics