[FIX] Correctly resolve test files paths when adding runtime comments

[mohammedahmed18]

PR Type

Bug fix, Enhancement

this will fix the bug where runtime comments not shown in worktree mode

---

Description

• Correct test path resolution for comments

• Pass tests rootdir into runtime mapping

• Simplify optimizer file/path initialization

---

Diagram Walkthrough

flowchart LR
  runtimes["Invocation runtimes"] -- with tests root --> unique["unique_inv_id(...)"]
  unique -- ids mapped --> annotate["add_runtime_comments_to_generated_tests"]
  lspinit["LSP init"] -- set file/function early --> optimizer["Optimizer args"]

Loading

File Walkthrough

	Relevant files
Bug fix	edit_generated_tests.py Resolve test paths using provided tests rootdir                    codeflash/code_utils/edit_generated_tests.py Add Optional Path import and parameter. Compute absolute test paths under provided rootdir. Update unique_inv_id signature to accept rootdir. Thread rootdir through add_runtime_comments_to_generated_tests. +7/-5
Enhancement	beta.py Simplify optimizer path setup in LSP init                                codeflash/lsp/beta.py Set optimizer args.file directly to current file. Remove relative-to-project_root computation. Reorder to call worktree_mode after args set. +3/-6      function_optimizer.py Provide tests rootdir to comment annotator                              codeflash/optimization/function_optimizer.py Pass tests_project_rootdir into add_runtime_comments_to_generated_tests. +1/-1

---

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

Codeflash Bot seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

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

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Path Handling The new resolution uses tests_project_rootdir combined with a dot-path and then strips the suffix again with with_suffix(""). This double with_suffix call is unusual and may yield unexpected results if module paths or symlinks are involved. Validate that abs_path consistently matches how test paths are later looked up and keyed, especially across platforms and when tests_project_rootdir is None. def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path) -> dict[str, int]: unique_inv_ids: dict[str, int] = {} for inv_id, runtimes in inv_id_runtimes.items(): test_qualified_name = ( inv_id.test_class_name + "." + inv_id.test_function_name # type: ignore[operator] if inv_id.test_class_name else inv_id.test_function_name ) abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") abs_path = str(abs_path.resolve().with_suffix("")) if "__unit_test_" not in abs_path: Behavior Change Optimizer args initialization was simplified to set args.file directly to the file_path and call worktree_mode earlier. This removes previous logic that relativized to project_root/current_worktree. Confirm this still points to the correct worktree file and doesn’t break scenarios where optimizer expects paths relative to project_root. def initialize_function_optimization( server: CodeflashLanguageServer, params: FunctionOptimizationInitParams ) -> dict[str, str]: file_path = Path(uris.to_fs_path(params.textDocument.uri)) server.show_message_log(f"Initializing optimization for function: {params.functionName} in {file_path}", "Info") if server.optimizer is None: _initialize_optimizer_if_api_key_is_valid(server) server.optimizer.args.file = file_path server.optimizer.args.function = params.functionName server.optimizer.args.previous_checkpoint_functions = False server.optimizer.worktree_mode() server.show_message_log( f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info" ) Optional Rootdir Default When tests_project_rootdir is None, Path() defaults to the CWD. Ensure this default matches intended behavior in all call sites to avoid mismatches between runtime mapping keys and generated test paths when the working directory differs from the tests project root. def add_runtime_comments_to_generated_tests( generated_tests: GeneratedTestsList, original_runtimes: dict[InvocationId, list[int]], optimized_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Optional[Path] = None, ) -> GeneratedTestsList: """Add runtime performance comments to function calls in generated tests.""" original_runtimes_dict = unique_inv_id(original_runtimes, tests_project_rootdir or Path()) optimized_runtimes_dict = unique_inv_id(optimized_runtimes, tests_project_rootdir or Path()) # Process each generated test

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Set file after worktree switch Setting the file before entering worktree mode may desync the path if worktree_mode remaps paths. Move the file assignment after worktree_mode and ensure it is translated to the current worktree root, preserving the original behavior. This prevents optimizer from targeting the wrong file. codeflash/lsp/beta.py [273-279] -server.optimizer.args.file = file_path server.optimizer.args.function = params.functionName server.optimizer.args.previous_checkpoint_functions = False server.optimizer.worktree_mode() +original_args, _ = server.optimizer.original_args_and_test_cfg +original_relative_file_path = file_path.relative_to(original_args.project_root) +server.optimizer.args.file = server.optimizer.current_worktree / original_relative_file_path + Suggestion importance[1-10]: 8 __ Why: The PR changed behavior by setting args.file directly and calling worktree_mode afterward; restoring mapping via current_worktree maintains previous semantics and prevents targeting the wrong file.	Medium
	Avoid fragile path resolution Avoid calling resolve on paths derived from possibly non-existent files; it can raise on some environments and also collapses symlinks, breaking matching with stored runtimes. Normalize the path with as_posix() and absolute() relative to the provided root, without stripping the suffix twice. This ensures consistent keys and prevents runtime errors when files are absent. codeflash/code_utils/edit_generated_tests.py [152-166] def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path) -> dict[str, int]: unique_inv_ids: dict[str, int] = {} for inv_id, runtimes in inv_id_runtimes.items(): test_qualified_name = ( - inv_id.test_class_name + "." + inv_id.test_function_name # type: ignore[operator] + inv_id.test_class_name + "." + inv_id.test_function_name if inv_id.test_class_name else inv_id.test_function_name ) - abs_path = tests_project_rootdir / Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") - abs_path = str(abs_path.resolve().with_suffix("")) - if "__unit_test_" not in abs_path: + rel_path = Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py") + abs_path = (tests_project_rootdir / rel_path).absolute() + abs_path_str = abs_path.as_posix() + if "__unit_test_" not in abs_path_str: continue - key = test_qualified_name + "#" + abs_path # type: ignore[operator] + key = f"{test_qualified_name}#{abs_path_str}" parts = inv_id.iteration_id.split("_").__len__() # type: ignore[union-attr] + ... Suggestion importance[1-10]: 7 __ Why: Replacing resolve().with_suffix("") with absolute/as_posix avoids symlink collapse and potential resolution issues while preserving consistent key formation; it's a reasonable robustness improvement aligned with the diffed code.	Medium
General	Enforce explicit project root Using an empty Path as fallback can generate absolute paths relative to CWD, causing key mismatches. Explicitly require and validate tests_project_rootdir, or use a stable default tied to generated_tests for consistent path computation. This avoids silently mismatched runtime annotations. codeflash/code_utils/edit_generated_tests.py [181-182] -original_runtimes_dict = unique_inv_id(original_runtimes, tests_project_rootdir or Path()) -optimized_runtimes_dict = unique_inv_id(optimized_runtimes, tests_project_rootdir or Path()) +if tests_project_rootdir is None: + raise ValueError("tests_project_rootdir must be provided to resolve test file paths reliably.") +original_runtimes_dict = unique_inv_id(original_runtimes, tests_project_rootdir) +optimized_runtimes_dict = unique_inv_id(optimized_runtimes, tests_project_rootdir) Suggestion importance[1-10]: 6 __ Why: Requiring a non-None tests_project_rootdir can prevent subtle key mismatches, but it tightens API and may be too strict; still a valid, moderate-impact correctness safeguard.	Low

[aseembits93]

might want to test with different test dirs, some which are a subdirectory of the module root some which are not. I remember i was using the test_dir early on but it didn't work for all scenarioes @mohammedahmed18

[mohammedahmed18]

@aseembits93
tested with a tests dir inside the module root, and the runtime diff is shown correctly

[aseembits93]

outside the module root? @mohammedahmed18

[KRRT7]

@mohammedahmed18 branch conflicts after pygls migration

[aseembits93]

async e2e test doesn't show it yet, let me investigate

[aseembits93]

codeflash --file code_to_optimize/code_directories/async_e2e/main.py --verbose --module-root /Users/aseemsaxena/Downloads/codeflash_dev/codeflash --tests-root /Users/aseemsaxena/Downloads/codeflash_dev/codeflash/code_to_optimize/code_directories/async_e2e/tests/ --async does not resolve to the correct path @mohammedahmed18

[aseembits93]

@mohammedahmed18

[aseembits93]

@KRRT7 discover_tests() is somehow mutating the tests_project_rootdir in async mode

[aseembits93]

ok the mutation is expected, it's somehow working as expected downstream for everything except async mode. somehow the async mode invocation ids have a different test_module_path (should be code_to_optimize.code_directories.async_e2e.tests.test_retry_with_backoff__perf_test_1 instead of tests.test_retry_with_backoff__perf_test_1)

[aseembits93]

works for non-async. async issue is unrelated.

[mohammedahmed18]

@aseembits93 found the issue, the tests_project_rootdir is being overridden during discovering tests

    if pytest_rootdir is not None:
        cfg.tests_project_rootdir = Path(pytest_rootdir)

EDIT: this is weird it really does work with the non-async code even with tests_project_rootdir being overriden