fix: resolve LLM returning wrong file paths in multi-file optimization

[KRRT7]

Summary

• Ensure the target file is always first in the read-writable markdown sent to the LLM
• Add basename fallback in get_optimized_code_for_module when exact path lookup fails
• Add 11 unit tests covering all lookup/fallback paths

Fixes a warning seen in CI where the LLM returns code blocks tagged with context file paths (e.g. console.py, aiservice.py) instead of the target file (version.py).

Companion backend PR: codeflash-ai/codeflash-internal#2469

Test plan

• Unit tests for get_optimized_code_for_module — exact match, None fallback, basename fallback, ambiguous, empty, CI reproduction

[claude]

Claude finished @KRRT7's task in 4m 12s —— View job

---

PR Review Summary

Tasks

• Gather context and read changed files
• Run linting and type checks
• Review code for bugs/security/correctness
• Check for duplicate logic
• Analyze test coverage
• Handle optimization PRs
• Post summary

---

Prek Checks

All checks passed — no linting or formatting issues.

ruff check...  Passed
ruff format... Passed

mypy: no issues found in changed source files.

---

Code Review

Overall: The approach is sound — putting the target file first in the read-writable context is a clean preventive fix, and the basename fallback in get_optimized_code_for_module is a good recovery heuristic. No critical bugs found.

Minor observations:

1. code_context_extractor.py:127 — uncaught ValueError on path outside project root

   target_relative = function_to_optimize.file_path.resolve().relative_to(project_root_path.resolve())

   If file_path resolves to something outside project_root_path (e.g. a symlink pointing outside the project), this raises ValueError and crashes get_code_optimization_context. The rest of the function continues to run with unmodified ordering — a try/except ValueError or Path.is_relative_to() guard would make this more robust. This is low-risk given the existing invariant, but worth noting.

2. code_replacer.py:31 — "None" string key is fragile
   The "None" key comes from str(None) in file_to_path(). This works because str(None) == "None" in Python, but it's an implicit contract between two files. A named constant or sentinel would make the coupling explicit and prevent silent breakage if the serialization changes.

3. CodeStringsMarkdown._cache invalidation risk
   CodeStringsMarkdown has a private _cache populated lazily by flat, file_to_path. The PR mutates code_strings in place after construction:

   final_read_writable_code.code_strings = target_blocks + other_blocks

   This is safe here because the cache is empty at that point. However, this is a fragile pattern — if the ordering code runs after any property access, the cache silently serves stale data. Not a bug today, but worth watching if this pattern is reused elsewhere.

---

Duplicate Detection

No duplicates detected. get_optimized_code_for_module is the single canonical implementation imported into languages/python/static_analysis/code_replacer.py — no parallel reimplementation exists.

---

Test Coverage

All 11 new unit tests pass. Coverage for the changed function path:

File	Coverage
codeflash/languages/code_replacer.py	34% (full file; new logic well-covered by unit tests)
codeflash/languages/python/context/code_context_extractor.py	70%

The 17 test failures in test_code_context_extractor.py are all pre-existing API key failures (CODEFLASH_API_KEY not available in this environment), not related to the PR changes. The new unit tests in test_get_optimized_code_for_module.py all pass.

---

Optimization PRs

No open optimization PRs from codeflash-ai[bot].

---

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 33% (0.33x) speedup for get_optimized_code_for_module in codeflash/languages/code_replacer.py

⏱️ Runtime : 9.16 milliseconds → 6.86 milliseconds (best of 23 runs)

A dependent PR with the suggested changes has been created. Please review:

• ⚡️ Speed up function get_optimized_code_for_module by 33% in PR #1772 (fix-markdown-code-path-lookup) #1773

If you approve, it will be merged into this PR (branch fix-markdown-code-path-lookup).