fix: add pre-flight module import check to exit early before API calls

[aseembits93]

Summary

• Adds a find_spec import validation check in can_be_optimized() that runs before any expensive API calls (test generation + optimization)
• If the import fails due to a wrong module-root, automatically infers the correct value from the __init__.py chain, updates pyproject.toml, and retries
• If the import fails due to a missing dependency (not a module-root issue), fails early with a clear error message instead of wasting API credits

Test plan

• Set module-root to an incorrect value in pyproject.toml and run codeflash --file <file> --function <func> — verify it auto-corrects and updates pyproject.toml
• Remove a dependency import and verify the pre-flight check catches it before API calls
• Run with a correct module-root and verify no regression — optimization proceeds normally

🤖 Generated with Claude Code

[claude]

Claude finished @aseembits93's task in 5m 18s —— View job

---

PR Review Summary

• Triage PR scope
• Run lint/typecheck — fixed 1 ruff (unsorted imports) + 3 new mypy errors (commit 0404570f)
• Resolve stale review threads — none found
• Code review
• Duplicate detection
• Test coverage
• Merge optimization PRs — skipped (requires elevated permissions not granted)

---

Prek Checks

All ruff and format checks pass after fixing the unsorted-import error in function_optimizer.py.

Code Review

Bugs / Issues found:

1. infer_module_root_from_file return type annotation is misleading (code_utils.py:383): The function is declared -> Path | None but never actually returns None — the fallback always returns file_path.parent. As a result, the if inferred_root is None guard in try_correct_module_root is dead code. The return type should be -> Path. Fix this →

2. Double validate_module_import call when auto-correction succeeds (function_optimizer.py:78,109): try_correct_module_root() internally calls validate_module_import to confirm the new root works (line 78), and then can_be_optimized() calls it again immediately after (line 109) on the same path. When auto-correction succeeds, the import is validated twice unnecessarily. Minor performance issue but worth noting.

3. sys.path mutation in validate_module_import is not thread-safe (code_utils.py:368): In the --worktree parallel optimization mode, concurrent calls could race on sys.path.insert/remove. This is low risk since the pre-flight check runs before parallelism begins, but worth a note if the function is reused elsewhere.

4. Pending reviewer request: pre_flight hook in base class (function_optimizer.py): KRRT7's second review comment ("if pre-flight doesn't exist in the language base add it") hasn't been addressed. The current implementation overrides can_be_optimized() directly. A cleaner design would be to add a pre_flight() -> Result[None, str] | None no-op hook to FunctionOptimizer and call it at the top of can_be_optimized(), then override just that hook in PythonFunctionOptimizer. This would avoid duplicating the super().can_be_optimized() dispatch pattern if other language optimizers need pre-flight checks in the future.

Looks good:

• The find_spec approach (no subprocess, no actual module init) is the right lightweight check
• The sys.path cleanup via finally is correct
• The __init__.py chain walking logic in infer_module_root_from_file is sound
• The pyproject.toml update uses tomlkit (preserves formatting) correctly
• The in-memory config update (self.args.module_root, self.project_root, self.original_module_path) covers all the fields that downstream code uses

Duplicate Detection

No duplicates detected. The validate_module_import helper in code_utils.py is distinct from the simple inline find_spec("pkg") is not None usages scattered elsewhere (comparator, code_extractor) — those are one-line boolean checks while this function adds sys.path management and structured error reporting.

Test Coverage

All 21 tests in tests/code_utils/test_module_import_validation.py pass. The new validate_module_import and infer_module_root_from_file functions are well-covered (lines 357–410 not in the missing-lines list). The try_correct_module_root path is covered by the 5 tests in TestTryCorrectModuleRoot.

---

Last updated: 2026-03-04 | | Branch

[KRRT7]

a full subprocess call is far more expensive than how we check if something is installed in the comparator:
https://github.com/codeflash-ai/codeflash/blob/main/codeflash/verification/comparator.py#L21

you can use that for this

[aseembits93]

this pr is still in progress

[aseembits93]

@claude is the windows unittest failure related to this pull request?

[claude]

Claude finished @aseembits93's task in 48s —— View job

---

Windows unittest failure analysis

• Check CI status for failing workflows
• Analyze failing test details
• Determine if failure is related to this PR

---

No, the Windows unittest failure is not related to this PR.

The failing test is:

FAILED tests/test_languages/test_javascript_test_runner.py::TestBundledJestReporter::test_reporter_produces_valid_junit_xml

The error is a Windows path separator issue in the Jest reporter module path:

Error: Cannot find module 'D:acodeflashcodeflashpackagescodeflashruntimejest-reporter.js'

The path D:acodeflash... is a mangled Windows path — backslashes are being dropped when constructing the path to jest-reporter.js. This is a pre-existing issue in the JavaScript test runner, completely unrelated to this PR's changes.

Evidence that this PR is not the cause:

1. This PR only touches codeflash/code_utils/code_utils.py, codeflash/optimization/function_optimizer.py, and tests/code_utils/test_module_import_validation.py
2. The new test file tests/code_utils/test_module_import_validation.py passed all 21 tests on Windows (.....................)
3. The failing test is in test_javascript_test_runner.py — JavaScript/Jest infrastructure, nothing this PR modified

The Windows Jest reporter path bug is a separate pre-existing issue that would need to be fixed in the JavaScript test runner code.

---

| View failing job

[aseembits93]

@claude fix the failing tests/test_languages/test_javascript_test_runner.py::TestBundledJestReporter::test_reporter_produces_valid_junit_xml

[claude]

Claude finished @aseembits93's task in 1m 12s —— View job

---

Fix: Windows path separator issue in Jest reporter test

• Triage request
• Read failing test file — root cause identified
• Implement the fix
• Commit and push

Root Cause

In test_reporter_produces_valid_junit_xml, Windows Path objects use backslashes which become JavaScript escape sequences when embedded directly in the JS string template:

# Before (broken on Windows):
process.env.JEST_JUNIT_OUTPUT_FILE = '{output_file}';  # D:\a\... → D:a... (backslash escapes mangled)
const Reporter = require('{reporter_path}');

Fix (commit 7014a955)

Used .as_posix() for both paths when embedding in the JS script — Node.js accepts forward slashes on Windows:

# After (works on all platforms):
process.env.JEST_JUNIT_OUTPUT_FILE = '{output_file.as_posix()}';
const Reporter = require('{reporter_path.as_posix()}');

| Branch

[KRRT7]

instead of implementing it in the function optimizer here, can you find a good spot for it in https://github.com/codeflash-ai/codeflash/blob/main/codeflash/languages/python/function_optimizer.py?

[KRRT7]

if pre-flight doesn't exist in the language base add it