barebones LSP Server implementation

[KRRT7]

PR Type

Enhancement

---

Description

• Integrate LSP server for function optimization

• Define LSP methods for optimization workflow

• Modularize & refactor FunctionOptimizer methods

• Introduce benchmark and test discovery utilities

---

Changes walkthrough 📝

	Relevant files
Enhancement	beta.py Implement LSP optimization endpoints                                          codeflash/lsp/beta.py Add LSP feature handlers for optimization Implement getOptimizableFunctions and initialization Support test discovery, preparation, and execution Define performFunctionOptimization endpoint +231/-0  server.py Add custom LSP server protocol and initialization                codeflash/lsp/server.py Define custom LanguageServerProtocol with initialize Locate and load pyproject.toml for optimizer Implement CodeflashLanguageServer initializer +56/-0    function_optimizer.py Modularize FunctionOptimizer methods                                          codeflash/optimization/function_optimizer.py Extract can_be_optimized helper method Add modular test generation and baseline setup Refactor optimize_function to use new steps +296/-160 optimizer.py Enhance Optimizer with LSP integration                                      codeflash/optimization/optimizer.py Introduce run_benchmarks and discover_tests Add prepare_module_for_optimization helper Update create_function_optimizer signature Maintain current function context for LSP +186/-153
Dependencies	pyproject.toml Update dependencies and lint ignore                                            pyproject.toml Add direct dependency on pygls Update lint ignore rules +3/-1
Additional files	__init__.py [link]

---

Need help?

Type /help how to ... in the comments thread for any questions about PR-Agent usage.

Check out the documentation for more information.

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 39% (0.39x) speedup for function_kind in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 1.77 milliseconds → 1.27 milliseconds (best of 158 runs)

I created a new dependent PR with the suggested changes. Please review:

• ⚡️ Speed up function function_kind by 39% in PR #309 (updated-vsc-extension) #330

If you approve, it will be merged into this PR (branch updated-vsc-extension).

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 89ceca8)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Missing None-handling LSP handlers like prepareOptimization and performFunctionOptimization unpack the result of prepare_module_for_optimization without checking for None, which can lead to runtime errors if parsing fails. @server.feature("prepareOptimization") def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: current_function = server.optimizer.current_function_being_optimized module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) validated_original_code, original_module_ast = module_prep_result function_optimizer = server.optimizer.create_function_optimizer( current_function, function_to_optimize_source_code=validated_original_code[current_function.file_path].source_code, original_module_ast=original_module_ast, original_module_path=current_function.file_path, ) Unimplemented method The newly declared submit_test_generation_tasks method is referenced in test generation but its implementation is not shown, likely causing failures at runtime. def submit_test_generation_tasks( self, Stale state The attribute generate_and_instrument_tests_results is set but never cleared or used elsewhere, which may retain outdated data between optimization runs. self.generate_and_instrument_tests_results: ( tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] | None ) = None

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 89ceca8
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Define missing context before returning The variables code_context and original_helper_code are not defined, causing a NameError. You should call prepare_module_for_optimization to obtain the validated code and AST, then construct a CodeOptimizationContext and extract the helper code mapping before returning Success. codeflash/optimization/function_optimizer.py [150-157] def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]: + prep = self.prepare_module_for_optimization(self.function_to_optimize.file_path) + if prep is None: + return Failure("Module preparation failed") + validated_code, module_ast = prep + code_context = CodeOptimizationContext(validated_code, module_ast) + original_helper_code = {path: vc.source_code for path, vc in validated_code.items()} should_run_experiment = self.experiment_id is not None logger.debug(f"Function Trace ID: {self.function_trace_id}") ph("cli-optimize-function-start", {"function_trace_id": self.function_trace_id}) self.cleanup_leftover_test_return_values() - if random.random() > REPEAT_OPTIMIZATION_PROBABILITY and was_function_previously_optimized( self.function_to_optimize, code_context ): return Failure("Function optimization previously attempted, skipping.") - return Success((should_run_experiment, code_context, original_helper_code)) Suggestion importance[1-10]: 9 __ Why: The new can_be_optimized function references undefined code_context and original_helper_code, causing a NameError; fixing this is critical for the optimization pipeline.	High
	Handle failed module preparation prepare_module_for_optimization can return None on syntax errors, leading to a TypeError when unpacking. Add a check for None and return an LSP error response if module preparation failed. codeflash/lsp/beta.py [76-77] -@server.feature("prepareOptimization") -def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: - current_function = server.optimizer.current_function_being_optimized +module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) +if module_prep_result is None: + return {"functionName": params.functionName, "status": "error", "message": "Module preparation failed"} +validated_original_code, original_module_ast = module_prep_result - module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) - validated_original_code, original_module_ast = module_prep_result - Suggestion importance[1-10]: 7 __ Why: Adding a check for None prevents a TypeError when unpacking a failed prepare_module_for_optimization, improving robustness.	Medium
	Guard uninitialized current function If current_function_being_optimized is None, accessing file_path will raise an AttributeError. Add a guard to return an error when no function is initialized before discovery. codeflash/lsp/beta.py [61-65] -@server.feature("discoverFunctionTests") -def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: - current_function = server.optimizer.current_function_being_optimized +if not server.optimizer.current_function_being_optimized: + return {"functionName": params.functionName, "status": "error", "message": "No function initialized"} +current_function = server.optimizer.current_function_being_optimized +optimizable_funcs = {current_function.file_path: [current_function]} - optimizable_funcs = {current_function.file_path: [current_function]} - Suggestion importance[1-10]: 6 __ Why: Without the guard, current_function may be None and cause an AttributeError; adding this check prevents a potential runtime crash.	Low

---

Previous suggestions

Suggestions up to commit 3006d62

Category	Suggestion	Impact
Possible issue	Select function by name Instead of arbitrarily popping an item, filter the optimizable_funcs mapping by the requested function name. This ensures you select the correct function rather than an arbitrary one. Also handle the case where the function isn't found. codeflash/lsp/beta.py [56] -fto = optimizable_funcs.popitem()[1][0] +fto = next( + ( + func + for funcs in optimizable_funcs.values() + for func in funcs + if func.qualified_name == params.functionName + ), + None, +) +if not fto: + return {"functionName": params.functionName, "status": "not found", "info": None} Suggestion importance[1-10]: 8 __ Why: The code currently uses popitem arbitrarily which can select the wrong function; filtering by params.functionName ensures the correct function is chosen.	Medium
	Prevent divide-by-zero error Guard against division by zero when computing speedup to avoid runtime errors. Check that best_optimization.runtime is non-zero before dividing. codeflash/lsp/beta.py [223] -"extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster" +if best_optimization.runtime: + speedup = original_code_baseline.runtime / best_optimization.runtime + extra_msg = f"Speedup: {speedup:.2f}x faster" +else: + extra_msg = "Speedup: N/A" +... +"extra": extra_msg Suggestion importance[1-10]: 7 __ Why: Adding a guard against best_optimization.runtime being zero prevents potential runtime exceptions when calculating the speedup.	Medium
General	Return actual discovered tests Rather than returning the count as a string, return the actual number or the list of discovered tests to the client. This provides richer context and keeps types correct. codeflash/lsp/beta.py [69] -return {"functionName": params.functionName, "status": "success", "generated_tests": str(num_discovered_tests)} +discovered_tests = function_to_tests.get(current_function.file_path, []) +return { + "functionName": params.functionName, + "status": "success", + "discovered_tests": discovered_tests +} Suggestion importance[1-10]: 6 __ Why: Returning the actual list of discovered tests (or proper numeric type) provides more useful context and corrects the type mismatch.	Low
	Consistent response schema Align the error response field names with the success case by using "info" instead of "args". This keeps your response schema consistent. codeflash/lsp/beta.py [54-55] if not optimizable_funcs: - return {"functionName": params.functionName, "status": "not found", "args": None} + return {"functionName": params.functionName, "status": "not found", "info": None} Suggestion importance[1-10]: 5 __ Why: Changing the error response key from args to info matches the success case and keeps the response payload consistent.	Low

[github-actions]

Persistent review updated to latest commit 89ceca8

[Saga4]

LGTM