codeflash-admin/codeflash
1
0
Fork 0
mirror of https://github.com/codeflash-ai/codeflash.git synced 2026-05-04 18:25:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge main, resolve prek conflict (keep shared workflow)

Browse source
This commit is contained in:
Kevin Turcios 2026-04-23 05:27:38 -05:00
commit c613fda5de
42 changed files with 5495 additions and 1735 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
.claude/hooks/bash-guard.sh Executable file
View file

@ -0,0 +1,45 @@
#!/usr/bin/env bash
# PreToolUse hook: Block Bash calls that should use dedicated tools.
# Exit 0 = allow, Exit 2 = block (message on stderr).
INPUT=$(cat 2>/dev/null || true)
COMMAND=$(echo "$INPUT" | jq -r '.tool_input.command // empty' 2>/dev/null || true)
[ -z "$COMMAND" ] && exit 0
# Strip leading env vars (FOO=bar cmd ...) and whitespace to get the actual command
STRIPPED=$(echo "$COMMAND" | sed 's/^[[:space:]]*\([A-Za-z_][A-Za-z0-9_]*=[^[:space:]]*[[:space:]]*\)*//')
FIRST_CMD=$(echo "$STRIPPED" | awk '{print $1}')
case "$FIRST_CMD" in
grep|egrep|fgrep|rg)
echo "BLOCKED: Use the Grep tool instead of \`$FIRST_CMD\`. It provides better output and permissions handling." >&2
exit 2
;;
find)
echo "BLOCKED: Use the Glob tool instead of \`find\`. Glob is faster and returns results sorted by modification time." >&2
exit 2
;;
cat|head|tail)
echo "BLOCKED: Use the Read tool instead of \`$FIRST_CMD\`. Read provides line numbers and supports images/PDFs." >&2
exit 2
;;
awk)
echo "BLOCKED: Use the Grep tool or Read tool instead of \`awk\`." >&2
exit 2
;;
sed)
if echo "$COMMAND" | grep -qE '(^|[[:space:]])sed[[:space:]]+-i'; then
echo "BLOCKED: Use the Edit tool instead of \`sed -i\`. Edit tracks changes properly." >&2
exit 2
fi
;;
esac
# echo with file redirection (echo "..." > file)
if echo "$STRIPPED" | grep -qE '^echo\b.*[[:space:]]>'; then
echo "BLOCKED: Use the Write tool instead of \`echo >\`. Write provides proper file creation." >&2
exit 2
fi
exit 0

47
.claude/hooks/post-compact.sh Executable file
View file

@ -0,0 +1,47 @@
#!/usr/bin/env bash
# PreCompact hook: Inject state preservation guidance before context compaction.
cd "$CLAUDE_PROJECT_DIR" 2>/dev/null || exit 0
STATE=""
BRANCH=$(git branch --show-current 2>/dev/null)
[ -n "$BRANCH" ] && STATE="${STATE}Branch: ${BRANCH}\n"
DIRTY=$(git status --porcelain 2>/dev/null)
if [ -n "$DIRTY" ]; then
COUNT=$(echo "$DIRTY" | wc -l | tr -d ' ')
STATE="${STATE}Uncommitted files (${COUNT}):\n${DIRTY}\n"
fi
UPSTREAM=$(git rev-parse --abbrev-ref '@{upstream}' 2>/dev/null)
if [ -n "$UPSTREAM" ]; then
AHEAD=$(git rev-list --count "${UPSTREAM}..HEAD" 2>/dev/null)
[ "$AHEAD" -gt 0 ] 2>/dev/null && STATE="${STATE}Unpushed commits: ${AHEAD}\n"
fi
RECENT=$(git log --oneline -5 2>/dev/null)
[ -n "$RECENT" ] && STATE="${STATE}Recent commits:\n${RECENT}\n"
LATEST_HANDOFF=$(ls -t "$CLAUDE_PROJECT_DIR/.claude/handoffs/"*.md 2>/dev/null | head -1)
if [ -n "$LATEST_HANDOFF" ] && [ -f "$LATEST_HANDOFF" ]; then
HANDOFF_CONTENT=$(head -40 "$LATEST_HANDOFF" 2>/dev/null)
[ -n "$HANDOFF_CONTENT" ] && STATE="${STATE}\nHandoff context:\n${HANDOFF_CONTENT}\n"
fi
STATE="${STATE}\nProject conventions to preserve:\n"
STATE="${STATE}- Python 3.9+, uv for all tooling, ruff + mypy via prek\n"
STATE="${STATE}- Verification: uv run prek (single command for lint/format/types)\n"
STATE="${STATE}- Pre-push: uv run prek run --from-ref origin/<base>\n"
STATE="${STATE}- Conventional commits: fix:, feat:, refactor:, test:, chore:\n"
STATE="${STATE}- Result type: Success(value) / Failure(error), check with is_successful()\n"
STATE="${STATE}- Language singleton: set_current_language() / current_language()\n"
STATE="${STATE}- libcst for code transforms, ast for read-only analysis\n"
[ -z "$STATE" ] && exit 0
EXPANDED=$(printf '%b' "$STATE")
jq -n --arg msg "PRESERVE the following session state through compaction:
$EXPANDED" '{"systemMessage": $msg}'
exit 0

2
.claude/hooks/post-edit-lint.sh
View file

@ -1,5 +1,4 @@
#!/usr/bin/env bash
# Everyone is on macOS so this should be fine, we don't account for Windows
set -euo pipefail
input=$(cat)
@ -10,6 +9,5 @@ if [[ -z "$file_path" || ! -f "$file_path" ]]; then
fi
if [[ "$file_path" == *.py ]]; then
# First run auto-fixes formatting; second run catches real lint errors
uv run prek --files "$file_path" 2>/dev/null || uv run prek --files "$file_path"
fi

25
.claude/hooks/require-read.sh Executable file
View file

@ -0,0 +1,25 @@
#!/usr/bin/env bash
# PreToolUse hook: Block Write/Edit on existing files that haven't been Read first.
# Exit 0 = allow, Exit 2 = block (message on stderr).
INPUT=$(cat 2>/dev/null || true)
FILE_PATH=$(echo "$INPUT" | jq -r '.tool_input.file_path // empty' 2>/dev/null || true)
[ -z "$FILE_PATH" ] && exit 0
# New files don't need prior reads
[ ! -f "$FILE_PATH" ] && exit 0
TRACKER="$CLAUDE_PROJECT_DIR/.claude/.read-tracker"
if [ ! -f "$TRACKER" ]; then
echo "BLOCKED: Read \`$(basename "$FILE_PATH")\` first before modifying it." >&2
exit 2
fi
if grep -qxF "$FILE_PATH" "$TRACKER"; then
exit 0
fi
echo "BLOCKED: Read \`$(basename "$FILE_PATH")\` first before modifying it." >&2
exit 2

50
.claude/hooks/status-line.sh Executable file
View file

@ -0,0 +1,50 @@
#!/usr/bin/env bash
# Status line: derive context from git state.
input=$(cat)
project_dir=$(echo "$input" | jq -r '.workspace.project_dir')
user=$(whoami)
branch=$(git -C "$project_dir" branch --show-current 2>/dev/null)
changed=$(git -C "$project_dir" diff --name-only HEAD 2>/dev/null)
[ -z "$changed" ] && changed=$(git -C "$project_dir" diff --name-only 2>/dev/null)
[ -z "$changed" ] && changed=$(git -C "$project_dir" diff --name-only --cached 2>/dev/null)
if [ -n "$changed" ]; then
area=$(echo "$changed" | sed 's|/.*||' | sort | uniq -c | sort -rn | head -1 | awk '{print $2}')
else
area=""
fi
context=""
case "$area" in
codeflash)
subsystem=$(echo "$changed" | grep '^codeflash/' | sed 's|^codeflash/||; s|/.*||' | sort | uniq -c | sort -rn | head -1 | awk '{print $2}')
[ -n "$subsystem" ] && context="editing $subsystem" ;;
tests)
target=$(echo "$changed" | grep '^tests/' | sed 's|^tests/||; s|/.*||' | sort -u | head -1)
[ -n "$target" ] && context="testing $target" ;;
.claude)
context="configuring claude" ;;
esac
if [ -z "$context" ] && [ -n "$branch" ]; then
case "$branch" in
feat/*|cf-*) context="building: ${branch#feat/}" ;;
fix/*) context="fixing: ${branch#fix/}" ;;
refactor/*) context="refactoring: ${branch#refactor/}" ;;
test/*) context="testing: ${branch#test/}" ;;
chore/*) context="chore: ${branch#chore/}" ;;
esac
fi
dirty=""
if [ -n "$(git -C "$project_dir" status --porcelain 2>/dev/null)" ]; then
dirty=" *"
fi
status="$user | codeflash"
[ -n "$context" ] && status="$status | $context"
[ -n "$branch" ] && status="$status | $branch$dirty"
echo "$status"

11
.claude/hooks/track-read.sh Executable file
View file

@ -0,0 +1,11 @@
#!/usr/bin/env bash
# PostToolUse hook: Track Read calls for the require-read guard.
INPUT=$(cat 2>/dev/null || true)
FILE_PATH=$(echo "$INPUT" | jq -r '.tool_input.file_path // empty' 2>/dev/null || true)
[ -z "$FILE_PATH" ] && exit 0
TRACKER="$CLAUDE_PROJECT_DIR/.claude/.read-tracker"
grep -qxF "$FILE_PATH" "$TRACKER" 2>/dev/null || echo "$FILE_PATH" >> "$TRACKER"
exit 0

13
.claude/rules/code-style.md
View file

@ -4,10 +4,11 @@
- **Python**: 3.9+ syntax
- **Package management**: Always use `uv`, never `pip`
- **Tooling**: Ruff for linting/formatting, mypy strict mode, prek for pre-commit checks
- **Comments**: Minimal - only explain "why", not "what"
- **Docstrings**: Do not add docstrings to new or changed code unless the user explicitly asks for them — not even one-liners. The codebase intentionally keeps functions self-documenting through clear naming and type annotations
- **Types**: Match the type annotation style of surrounding code — the codebase uses annotations, so add them in new code
- **Naming**: NEVER use leading underscores (`_function_name`) - Python has no true private functions, use public names
- **Comments**: Minimal — only explain "why", not "what"
- **Docstrings**: Do not add docstrings unless the user explicitly asks
- **Types**: Match the type annotation style of surrounding code
- **Naming**: No leading underscores (`_function_name`) — Python has no true private functions
- **Paths**: Always use absolute paths
- **Encoding**: Always pass `encoding="utf-8"` to `open()`, `read_text()`, `write_text()`, etc. in new or changed code — Windows defaults to `cp1252` which breaks on non-ASCII content. Don't flag pre-existing code that lacks it unless you're already modifying that line.
- **Verification**: Use `uv run prek` to verify code — it handles ruff, ty, mypy in one pass. Don't run `ruff`, `mypy`, or `python -c "import ..."` separately; `prek` is the single verification command
- **Encoding**: Always pass `encoding="utf-8"` to `open()`, `read_text()`, `write_text()` in new or changed code
- **Verification**: Use `uv run prek` — it handles ruff, ty, mypy in one pass. Don't run them separately
- **Code transforms**: Use `libcst` for code modification/transformation. `ast` is acceptable for read-only analysis

19
.claude/rules/debugging.md Normal file
View file

@ -0,0 +1,19 @@
# Debugging
## Root cause first
When encountering a bug, investigate the root cause. Don't patch symptoms. If you're about to add a try/except, a fallback default, or a defensive check — ask whether the real fix is upstream.
## Isolated testing
Prefer running individual test functions over full suites. Only run the full suite when explicitly asked or before pushing.
- Single function: `uv run pytest tests/test_foo.py::TestBar::test_baz -v`
- Single module: `uv run pytest tests/test_foo.py -v`
- Full suite: only when asked, or before `git push`
When debugging a specific endpoint or integration, test it directly instead of running the entire pipeline end-to-end.
## Subprocess failures
When a subprocess fails, always log stdout and stderr. "Exit code 1" with no output is useless.

42
.claude/rules/git.md
View file

@ -1,19 +1,35 @@
# Git Commits & Pull Requests
# Git
## Commits
- Never commit, amend, or push without explicit permission
- Don't commit intermediate states — wait until the full implementation is complete, reviewed, and explicitly approved before committing. If the user corrects direction mid-implementation, incorporate the correction before any commit
- Always create a new branch from `main` before starting any new work — never commit directly to `main` or reuse an existing feature branch for unrelated changes
- Use conventional commit format: `fix:`, `feat:`, `refactor:`, `docs:`, `test:`, `chore:`
- Keep commits atomic - one logical change per commit
- Commit message body should be concise (1-2 sentences max)
- Merge for simple syncs, rebase when branches have diverged significantly
- When committing to an external/third-party repo, follow that repo's own conventions for versioning, changelog, and CI
- Pre-commit: Run `uv run prek` before committing — fix any issues before creating the commit
- Pre-push: Run `uv run prek run --from-ref origin/<base>` to check all changed files against the PR base — this matches CI behavior and catches issues that per-commit prek misses. To detect the base branch: `gh pr view --json baseRefName -q .baseRefName 2>/dev/null || echo main`
- Don't commit intermediate states — wait until the full implementation is complete and approved
- Always create a new branch from `main` — never commit directly to `main`
- Conventional format: `fix:`, `feat:`, `refactor:`, `docs:`, `test:`, `chore:`
- First line: imperative verb + what changed, under 72 characters
- Body for *why*, not *what* — the diff shows what changed
- One purpose per commit: a bug fix, a new function, a refactor — not all three
- A commit that adds a function also adds its tests and exports — that's one logical change
## Sizing
- Too small: renaming a variable in one commit, updating its references in another
- Right size: adding a function with its tests, `__init__` export, and usage update
- Too large: implementing an entire subsystem in one commit
## Pre-commit / Pre-push
- Pre-commit: Run `uv run prek` before committing
- Pre-push: Run `uv run prek run --from-ref origin/<base>` to check all changed files against the PR base
## Pull Requests
- PR titles should use conventional format
- Keep the PR body short and straight to the point
- PR titles use conventional format
- Keep the PR body short and to the point
- If related to a Linear issue, include `CF-#` in the body
- Branch naming: `cf-#-title` (lowercase, hyphenated), no other prefixes/suffixes
- Branch naming: `cf-#-title` (lowercase, hyphenated)
## Branch Hygiene
- Delete feature branches locally after merging (`git branch -d <branch>`)
- Use `/clean_gone` to prune local branches whose remote tracking branch has been deleted

5
.claude/rules/github.md Normal file
View file

@ -0,0 +1,5 @@
# GitHub Interactions
ALWAYS use MCP GitHub tools (`mcp__github__*`) for GitHub operations. Check for a matching MCP tool first — only fall back to `gh` via Bash when no MCP tool exists for the operation.
This also applies to other MCP-connected services (Linear, Granola). MCP first, CLI second.

8
.claude/rules/language-patterns.md
View file

@ -6,8 +6,8 @@ paths:
# Language Support Patterns
- Current language is a module-level singleton in `languages/current.py` — use `set_current_language()` / `current_language()`, never pass language as a parameter through call chains
- Use `get_language_support(identifier)` from `languages/registry.py` to get a `LanguageSupport` instance — never import language classes directly
- New language support classes must use the `@register_language` decorator to register with the extension and language registries
- `languages/__init__.py` uses `__getattr__` for lazy imports to avoid circular dependencies — follow this pattern when adding new exports
- Prefer `LanguageSupport` protocol dispatch over `is_python()`/`is_javascript()` guards — remaining guards are being migrated to protocol methods
- Use `get_language_support(identifier)` from `languages/registry.py` — never import language classes directly
- New language support classes must use the `@register_language` decorator
- `languages/__init__.py` uses `__getattr__` for lazy imports to avoid circular dependencies
- Prefer `LanguageSupport` protocol dispatch over `is_python()`/`is_javascript()` guards
- `is_javascript()` returns `True` for both JavaScript and TypeScript (still used in ~15 call sites pending migration)

27
.claude/rules/sessions.md Normal file
View file

@ -0,0 +1,27 @@
# Session Discipline
## Scope
One task per session. Don't mix implementation with communication drafting, transcript search, or strategic planning.
## Duration
Cap sessions at 2-3 hours. Use `/handoff` at natural breakpoints rather than letting auto-compaction degrade context.
- After 1 compaction: consider wrapping up the current task and handing off
- After 3 compactions: stop, and tell the user to start a fresh session
- Never continue past 5 compactions — context is too degraded
## Context preservation
When compacting, preserve: modified files list, current branch, test commands used, key decisions made. Use subagents for exploration to keep main context clean.
## No polling
Never poll background tasks. No `wc -l`, no `tail -f`, no `sleep` loops. Use `run_in_background` and wait for the completion notification.
## File read budget
If you've read the same file 3+ times in a session, either:
- The session is too long and compaction destroyed your context — write a handoff
- You're not retaining key information — write it down in your response before it compacts away

8
.claude/rules/source-code.md
View file

@ -1,8 +0,0 @@
---
paths:
- "codeflash/**/*.py"
---
# Source Code Rules
- Use `libcst` for code modification/transformation to preserve formatting. `ast` is acceptable for read-only analysis and parsing.

19
.claude/rules/testing.md
View file

@ -4,13 +4,14 @@ paths:
- "codeflash/**/*test*.py"
---
# Testing Conventions
# Testing
- Code context extraction and replacement tests must always assert for full string equality, no substring matching.
- Use pytest's `tmp_path` fixture for temp directories — do not use `tempfile.mkdtemp()`, `tempfile.TemporaryDirectory()`, or `NamedTemporaryFile`. Some existing tests still use `tempfile` but new tests must use `tmp_path`.
- Always call `.resolve()` on Path objects before passing them to functions under test — this ensures absolute paths and resolves symlinks. Example: `source_file = (tmp_path / "example.py").resolve()`
- Use `.as_posix()` when converting resolved paths to strings (normalizes to forward slashes).
- Any new feature or bug fix that can be tested automatically must have test cases.
- If changes affect existing test expectations, update the tests accordingly. Tests must always pass after changes.
- The pytest plugin patches `time`, `random`, `uuid`, and `datetime` for deterministic test execution — never assume real randomness or real time in verification tests.
- `conftest.py` uses an autouse fixture that calls `reset_current_language()` — tests always start with Python as the default language.
- Full string equality for context extraction/replacement tests — no substring matching
- Use pytest's `tmp_path` fixture — not `tempfile.mkdtemp()` or `NamedTemporaryFile`
- Always call `.resolve()` on Path objects before passing to functions under test
- Use `.as_posix()` when converting resolved paths to strings
- New features and bug fixes must have test cases
- The pytest plugin patches `time`, `random`, `uuid`, `datetime` for deterministic execution
- `conftest.py` autouse fixture calls `reset_current_language()` — tests start with Python as default
- Prefer running individual tests over full suites: `uv run pytest tests/test_foo.py::TestBar::test_baz -v`
- Only run the full suite when explicitly asked or before pushing

12
.claude/rules/workflow.md
View file

@ -1,13 +1,17 @@
# Workflow
## Code Changes
- Before making any changes, outline your approach in 3-5 numbered steps. Include which repo/branch you'll work in, what commands you'll run, and what success looks like. Wait for approval before starting
Before making any changes, outline your approach in 3-5 numbered steps. Include which branch you'll work on, what commands you'll run, and what success looks like. Wait for approval before starting.
## Response Style
- When listing items (PRs, functions, optimization targets), always provide the complete list ordered by priority on the first attempt. Do not give partial lists
When listing items (PRs, functions, optimization targets), provide the complete list ordered by priority on the first attempt. No partial lists.
## Commands
- When running long-running commands (benchmarks, profiling, optimizers like codeflash), always run them in the foreground. Do not use background processes
Long-running commands (benchmarks, profiling, optimizers) always run in the foreground. Do not use background processes.
## Debugging
- When claiming something is a pre-existing issue (e.g., test failures on main), verify by checking out main and running the tests before making that claim
When claiming something is a pre-existing issue (e.g., test failures on main), verify by checking out main and running the tests before making that claim.

75
.claude/settings.json
View file

@ -1,16 +1,89 @@
{
"attribution": {
"commit": "",
"pr": ""
},
"includeCoAuthoredBy": false,
"permissions": {
"allow": [
"Bash(git status*)",
"Bash(git diff*)",
"Bash(git log*)",
"Bash(git branch*)",
"Bash(git show*)",
"Bash(git fetch*)",
"Bash(git checkout*)",
"Bash(uv run*)",
"Bash(uv sync*)",
"Bash(uv pip*)",
"Bash(prek*)",
"Bash(make*)",
"Bash(gh *)"
]
},
"hooks": {
"PreToolUse": [
{
"matcher": "Bash",
"hooks": [
{
"type": "command",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/bash-guard.sh",
"timeout": 5
}
]
},
{
"matcher": "Write",
"hooks": [
{
"type": "command",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/require-read.sh",
"timeout": 5
}
]
}
],
"PostToolUse": [
{
"matcher": "Read",
"hooks": [
{
"type": "command",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/track-read.sh",
"timeout": 5
}
]
},
{
"matcher": "Edit|Write",
"hooks": [
{
"type": "command",
"command": ".claude/hooks/post-edit-lint.sh",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/post-edit-lint.sh",
"timeout": 30
}
]
}
],
"PreCompact": [
{
"hooks": [
{
"type": "command",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/post-compact.sh",
"timeout": 10
}
]
}
]
},
"statusLine": {
"type": "command",
"command": "$CLAUDE_PROJECT_DIR/.claude/hooks/status-line.sh"
},
"enableAllProjectMcpServers": true,
"env": {
"ENABLE_LSP_TOOL": "1"
}
}

4
.claude/skills/fix-mypy.md
View file

@ -7,6 +7,6 @@ uv run mypy --non-interactive --config-file pyproject.toml <changed_files>
```
- Fix type annotation issues: missing return types, incorrect types, Optional/None unions, import errors for type hints
- Do NOT add `# type: ignore` comments — always fix the root cause
- Do NOT add `# type: ignore` comments -- always fix the root cause
- Do NOT fix type errors that require logic changes, complex generic type rework, or anything that could change runtime behavior
- Files in `mypy_allowlist.txt` are checked in CI — ensure they remain error-free
- Files in `mypy_allowlist.txt` are checked in CI -- ensure they remain error-free

2
.claude/skills/fix-prek.md
View file

@ -5,5 +5,5 @@ When prek (pre-commit) checks fail:
1. Run `uv run prek run` to see failures (local, checks staged files)
2. In CI, the equivalent is `uv run prek run --from-ref origin/main`
3. prek runs ruff format, ruff check, and mypy on changed files
4. Fix issues in order: formatting → lint → type errors
4. Fix issues in order: formatting -> lint -> type errors
5. Re-run `uv run prek run` to verify all checks pass

28
.github/workflows/ci.yaml vendored
View file

@ -36,6 +36,7 @@ jobs:
env:
PR_BODY: ${{ github.event.pull_request.body }}
PR_AUTHOR: ${{ github.event.pull_request.user.login }}
AUTHOR_ASSOCIATION: ${{ github.event.pull_request.author_association }}
run: |
# Skip for bots (dependabot, renovate, github-actions)
if [[ "$PR_AUTHOR" == *"[bot]"* || "$PR_AUTHOR" == "dependabot" ]]; then
@ -43,6 +44,12 @@ jobs:
exit 0
fi
# Skip for org members
if [[ "$AUTHOR_ASSOCIATION" == "MEMBER" ]]; then
echo "Org member ($PR_AUTHOR) — skipping linked issue check."
exit 0
fi
if [ -z "$PR_BODY" ]; then
echo "::error::PR body is empty. Every PR must link an issue or discussion."
echo "Use 'Closes #<number>', 'Fixes #<number>', 'Relates to #<number>', or include a discussion URL."
@ -124,7 +131,7 @@ jobs:
token: ${{ secrets.GITHUB_TOKEN }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: ${{ matrix.python-version }}
enable-cache: true
@ -157,7 +164,7 @@ jobs:
token: ${{ secrets.GITHUB_TOKEN }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: "3.13"
enable-cache: true
@ -173,7 +180,7 @@ jobs:
- name: Upload coverage report
if: always()
uses: actions/upload-artifact@v4
uses: actions/upload-artifact@v7
with:
name: coverage-report
path: coverage.xml
@ -193,7 +200,7 @@ jobs:
token: ${{ secrets.GITHUB_TOKEN }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
enable-cache: true
@ -230,6 +237,7 @@ jobs:
if: >-
fromJSON(needs.determine-changes.outputs.flags).e2e == 'true'
&& github.event_name != 'push'
&& github.actor != 'dependabot[bot]'
strategy:
fail-fast: false
matrix:
@ -289,7 +297,7 @@ jobs:
pr_state: ${{ github.event.pull_request.state }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: 3.11.6
enable-cache: true
@ -333,6 +341,7 @@ jobs:
if: >-
fromJSON(needs.determine-changes.outputs.flags).e2e_js == 'true'
&& github.event_name != 'push'
&& github.actor != 'dependabot[bot]'
strategy:
fail-fast: false
matrix:
@ -345,10 +354,12 @@ jobs:
script: end_to_end_test_js_esm_async.py
js_project_dir: code_to_optimize/js/code_to_optimize_js_esm
expected_improvement: 10
allow_failure: true
- name: js-ts-class
script: end_to_end_test_js_ts_class.py
js_project_dir: code_to_optimize/js/code_to_optimize_ts
expected_improvement: 30
continue-on-error: ${{ matrix.allow_failure || false }}
environment: ${{ ((github.event_name == 'workflow_dispatch' && github.actor != 'misrasaurabh1' && github.actor != 'KRRT7') || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
runs-on: ubuntu-latest
env:
@ -397,7 +408,7 @@ jobs:
npm install
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: 3.11.6
enable-cache: true
@ -414,6 +425,7 @@ jobs:
if: >-
fromJSON(needs.determine-changes.outputs.flags).e2e_java == 'true'
&& github.event_name != 'push'
&& github.actor != 'dependabot[bot]'
strategy:
fail-fast: false
matrix:
@ -466,7 +478,7 @@ jobs:
cache: maven
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: 3.11.6
enable-cache: true
@ -476,7 +488,7 @@ jobs:
- name: Cache codeflash-runtime JAR
id: runtime-jar-cache
uses: actions/cache@v4
uses: actions/cache@v5
with:
path: ~/.m2/repository/io/codeflash
key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}

4
.github/workflows/claude.yml vendored
View file

@ -57,7 +57,7 @@ jobs:
ref: ${{ github.event.pull_request.head.ref || github.ref }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
enable-cache: true
@ -319,7 +319,7 @@ jobs:
ref: ${{ steps.pr-ref.outputs.ref }}
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
enable-cache: true

2
.github/workflows/codeflash-optimize.yaml vendored
View file

@ -31,7 +31,7 @@ jobs:
fetch-depth: 0
- name: 🐍 Set up Python 3.11 for CLI
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: 3.11.6
enable-cache: true

4
.github/workflows/java-e2e.yaml vendored
View file

@ -41,7 +41,7 @@ jobs:
cache: maven
- name: Install uv
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
python-version: 3.11.6
enable-cache: true
@ -51,7 +51,7 @@ jobs:
- name: Cache codeflash-runtime JAR
id: runtime-jar-cache
uses: actions/cache@v4
uses: actions/cache@v5
with:
path: ~/.m2/repository/io/codeflash
key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}

2
.github/workflows/label-workflow-changes.yml vendored
View file

@ -13,7 +13,7 @@ jobs:
pull-requests: write
steps:
- name: Label PR with workflow changes
uses: actions/github-script@v7
uses: actions/github-script@v9
with:
script: |
const labelName = 'workflow-modified';

4
.github/workflows/publish.yml vendored
View file

@ -110,7 +110,7 @@ jobs:
- name: Install uv
if: steps.should_run.outputs.run == 'true' && steps.check_tag.outputs.exists == 'false'
uses: astral-sh/setup-uv@v8.0.0
uses: astral-sh/setup-uv@v8.1.0
with:
enable-cache: true
@ -124,7 +124,7 @@ jobs:
- name: Create GitHub Release
if: steps.should_run.outputs.run == 'true' && steps.check_tag.outputs.exists == 'false'
uses: softprops/action-gh-release@v2
uses: softprops/action-gh-release@v3
with:
tag_name: ${{ steps.extract_version.outputs.tag }}
name: ${{ matrix.release_name_prefix }} ${{ steps.extract_version.outputs.tag }}

5
.gitignore vendored
View file

@ -266,11 +266,12 @@ WARP.MD
.tessl/
tessl.json
# Claude Code - track shared rules, ignore local config
# Claude Code - track shared config, ignore local state
.claude/*
!.claude/rules/
!.claude/settings.json
!.claude/hooks/
!.claude/skills/
!.claude/settings.json
**/node_modules/**
**/dist-nuitka/**
**/.npmrc

13
.pre-commit-config.yaml
View file

@ -1,8 +1,15 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.8
hooks:
# Run the linter.
- id: ruff-check
# Run the formatter.
- id: ruff-format
- repo: local
hooks:
- id: mypy
name: mypy
entry: uv run mypy --non-interactive --config-file pyproject.toml
language: system
types: [python]
require_serial: true

28
CLAUDE.md
View file

@ -7,19 +7,29 @@ CodeFlash is an AI-powered code optimizer that automatically improves performanc
## Optimization Pipeline
```
Discovery → Ranking → Context Extraction → Test Gen + Optimization → Baseline → Candidate Evaluation → PR
Discovery -> Ranking -> Context Extraction -> Test Gen + Optimization -> Baseline -> Candidate Evaluation -> PR
```
See `.claude/rules/architecture.md` for directory mapping and entry points.
# Instructions
- **Bug fix workflow** — follow these steps in order, do not skip ahead:
1. Read the relevant code to understand the bug
2. Write a test that reproduces the bug (run it to confirm it fails)
3. Spawn subagents (using the Agent tool) to attempt the fix — each subagent should apply a fix and run the test to prove it passes
4. Review the subagent results, pick the best fix, and apply it
5. Never jump straight to writing a fix yourself — always go through steps 1-4
- Everything that can be tested should have tests.
## Setup
```bash
uv sync # Install all dependencies
uv run prek install # Install git pre-commit hooks (ruff + mypy)
```
## Bug Fix Workflow
Follow these steps in order, do not skip ahead:
1. Read the relevant code to understand the bug
2. Write a test that reproduces the bug (run it to confirm it fails)
3. Spawn subagents (using the Agent tool) to attempt the fix — each subagent should apply a fix and run the test to prove it passes
4. Review the subagent results, pick the best fix, and apply it
5. Never jump straight to writing a fix yourself — always go through steps 1-4
Everything that can be tested should have tests.
<!-- Section below is auto-generated by `tessl install` - do not edit manually -->

244
codeflash/languages/javascript/replay_test.py Normal file
View file

@ -0,0 +1,244 @@
from __future__ import annotations
import json
import re
import sqlite3
import textwrap
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional
if TYPE_CHECKING:
from collections.abc import Generator
@dataclass
class JavaScriptFunctionModule:
function_name: str
file_name: Path
module_name: str
class_name: Optional[str] = None
line_no: Optional[int] = None
def get_next_arg_and_return(
trace_file: str, function_name: str, file_name: str, class_name: Optional[str] = None, num_to_get: int = 25
) -> Generator[Any]:
db = sqlite3.connect(trace_file)
cur = db.cursor()
try:
cur.execute("SELECT name FROM sqlite_master WHERE type='table'")
tables = {row[0] for row in cur.fetchall()}
if "function_calls" in tables:
if class_name:
cursor = cur.execute(
"SELECT args FROM function_calls WHERE function = ? AND filename = ? AND classname = ? AND type = 'call' ORDER BY time_ns ASC LIMIT ?",
(function_name, file_name, class_name, num_to_get),
)
else:
cursor = cur.execute(
"SELECT args FROM function_calls WHERE function = ? AND filename = ? AND type = 'call' ORDER BY time_ns ASC LIMIT ?",
(function_name, file_name, num_to_get),
)
while (val := cursor.fetchone()) is not None:
args_data = val[0]
if isinstance(args_data, bytes):
yield args_data
else:
yield args_data
elif "traces" in tables:
if class_name:
cursor = cur.execute(
"SELECT args FROM traces WHERE function = ? AND file = ? ORDER BY id ASC LIMIT ?",
(function_name, file_name, num_to_get),
)
else:
cursor = cur.execute(
"SELECT args FROM traces WHERE function = ? AND file = ? ORDER BY id ASC LIMIT ?",
(function_name, file_name, num_to_get),
)
while (val := cursor.fetchone()) is not None:
yield val[0]
finally:
db.close()
def get_function_alias(module: str, function_name: str, class_name: Optional[str] = None) -> str:
module_alias = re.sub(r"[^a-zA-Z0-9]", "_", module).strip("_")
if class_name:
return f"{module_alias}_{class_name}_{function_name}"
return f"{module_alias}_{function_name}"
def create_javascript_replay_test(
trace_file: str,
functions: list[JavaScriptFunctionModule],
max_run_count: int = 100,
framework: str = "jest",
project_root: Optional[Path] = None,
) -> str:
is_vitest = framework.lower() == "vitest"
imports = []
if is_vitest:
imports.append("import { describe, test } from 'vitest';")
imports.append("const { getNextArg } = require('codeflash/replay');")
imports.append("")
for func in functions:
if func.function_name in ("__init__", "constructor"):
continue
alias = get_function_alias(func.module_name, func.function_name, func.class_name)
if func.class_name:
imports.append(f"const {{ {func.class_name}: {alias}_class }} = require('./{func.module_name}');")
else:
imports.append(f"const {{ {func.function_name}: {alias} }} = require('./{func.module_name}');")
imports.append("")
functions_to_test = [f.function_name for f in functions if f.function_name not in ("__init__", "constructor")]
metadata = f"""const traceFilePath = '{trace_file}';
const functions = {json.dumps(functions_to_test)};
"""
test_cases = []
for func in functions:
if func.function_name in ("__init__", "constructor"):
continue
alias = get_function_alias(func.module_name, func.function_name, func.class_name)
test_name = f"{func.class_name}.{func.function_name}" if func.class_name else func.function_name
if func.class_name:
class_arg = f"'{func.class_name}'"
test_body = textwrap.dedent(f"""
describe('Replay: {test_name}', () => {{
const traces = getNextArg(traceFilePath, '{func.function_name}', '{func.file_name.as_posix()}', {max_run_count}, {class_arg});
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {{
const instance = new {alias}_class();
instance.{func.function_name}(...args);
}});
}});
""")
else:
test_body = textwrap.dedent(f"""
describe('Replay: {test_name}', () => {{
const traces = getNextArg(traceFilePath, '{func.function_name}', '{func.file_name.as_posix()}', {max_run_count});
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {{
{alias}(...args);
}});
}});
""")
test_cases.append(test_body)
return "\n".join(
[
"// Auto-generated replay test by Codeflash",
"// Do not edit this file directly",
"",
*imports,
metadata,
*test_cases,
]
)
def get_traced_functions_from_db(trace_file: Path) -> list[JavaScriptFunctionModule]:
if not trace_file.exists():
return []
try:
conn = sqlite3.connect(trace_file)
cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
tables = {row[0] for row in cursor.fetchall()}
functions = []
if "function_calls" in tables:
cursor.execute(
"SELECT DISTINCT function, filename, classname, line_number FROM function_calls WHERE type = 'call'"
)
for row in cursor.fetchall():
func_name = row[0]
file_name = row[1]
class_name = row[2]
line_number = row[3]
module_path = file_name.replace("\\", "/").replace(".js", "").replace(".ts", "")
module_path = module_path.removeprefix("./")
functions.append(
JavaScriptFunctionModule(
function_name=func_name,
file_name=Path(file_name),
module_name=module_path,
class_name=class_name,
line_no=line_number,
)
)
elif "traces" in tables:
cursor.execute("SELECT DISTINCT function, file FROM traces")
for row in cursor.fetchall():
func_name = row[0]
file_name = row[1]
module_path = file_name.replace("\\", "/").replace(".js", "").replace(".ts", "")
module_path = module_path.removeprefix("./")
functions.append(
JavaScriptFunctionModule(
function_name=func_name, file_name=Path(file_name), module_name=module_path
)
)
conn.close()
return functions
except Exception:
return []
def create_replay_test_file(
trace_file: Path,
output_path: Path,
framework: str = "jest",
max_run_count: int = 100,
project_root: Optional[Path] = None,
) -> Optional[Path]:
functions = get_traced_functions_from_db(trace_file)
if not functions:
return None
content = create_javascript_replay_test(
trace_file=str(trace_file),
functions=functions,
max_run_count=max_run_count,
framework=framework,
project_root=project_root,
)
try:
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(content, encoding="utf-8")
return output_path
except Exception:
return None

19
codeflash/languages/javascript/support.py
View file

@ -1702,8 +1702,9 @@ class JavaScriptSupport:
) -> str:
"""Add behavior instrumentation to capture inputs/outputs.
For JavaScript, this wraps functions to capture their arguments
and return values.
For JavaScript, instrumentation is handled at runtime by the Babel tracer plugin
(babel-tracer-plugin.js) via trace-runner.js. This method returns the source
unchanged since no source-level transformation is needed.
Args:
source: Source code to instrument.
@ -1711,22 +1712,12 @@ class JavaScriptSupport:
output_file: Optional output file for traces.
Returns:
Instrumented source code.
Source code unchanged (Babel handles instrumentation at runtime).
"""
if not functions:
# JavaScript tracing is done at runtime via Babel plugin, not source transformation
return source
from codeflash.languages.javascript.tracer import JavaScriptTracer
# Use first function's file path if output_file not specified
if output_file is None:
file_path = functions[0].file_path
output_file = file_path.parent / ".codeflash" / "traces.db"
tracer = JavaScriptTracer(output_file)
return tracer.instrument_source(source, functions[0].file_path, list(functions))
def instrument_for_benchmarking(self, test_source: str, target_function: FunctionToOptimize) -> str:
"""Add timing instrumentation to test code.

543
codeflash/languages/javascript/tracer.py
View file

@ -1,35 +1,56 @@
"""Function tracing instrumentation for JavaScript.
This module provides functionality to wrap JavaScript functions to capture their
inputs, outputs, and execution behavior. This is used for generating replay tests
and verifying optimization correctness.
This module provides functionality to parse JavaScript function traces and generate
replay tests. Tracing is performed via Babel AST transformation using the
babel-tracer-plugin.js and trace-runner.js in the npm package.
The tracer uses Babel plugin for AST transformation which:
- Works with both CommonJS and ESM
- Handles async functions, arrow functions, methods correctly
- Preserves source maps and formatting
Database Schema (matches Python tracer):
- function_calls: Main trace data (type, function, classname, filename, line_number, time_ns, args)
- metadata: Key-value metadata about the trace session
"""
from __future__ import annotations
import json
import logging
import re
import sqlite3
from typing import TYPE_CHECKING, Any
import textwrap
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional
if TYPE_CHECKING:
from pathlib import Path
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
logger = logging.getLogger(__name__)
class JavaScriptTracer:
"""Instruments JavaScript code to capture function inputs and outputs.
@dataclass
class JavaScriptFunctionInfo:
function_name: str
file_name: str
module_path: str
class_name: Optional[str] = None
line_number: Optional[int] = None
Similar to Python's tracing system, this wraps functions to record:
- Input arguments
- Return values
- Exceptions thrown
- Execution time
class JavaScriptTracer:
"""Parses JavaScript function traces and generates replay tests.
Tracing is performed via Babel AST transformation (trace-runner.js).
This class handles:
- Parsing trace results from SQLite database
- Extracting traced function information
- Generating replay test files for Jest/Vitest
"""
SCHEMA_VERSION = "1.0.0"
def __init__(self, output_db: Path) -> None:
"""Initialize the tracer.
@ -38,322 +59,15 @@ class JavaScriptTracer:
"""
self.output_db = output_db
self.tracer_var = "__codeflash_tracer__"
def instrument_source(self, source: str, file_path: Path, functions: list[FunctionToOptimize]) -> str:
"""Instrument JavaScript source code with function tracing.
Wraps specified functions to capture their inputs and outputs.
Args:
source: Original JavaScript source code.
file_path: Path to the source file.
functions: List of functions to instrument.
Returns:
Instrumented source code with tracing.
"""
if not functions:
return source
# Add tracer initialization at the top
tracer_init = self._generate_tracer_init()
# Add instrumentation to each function
lines = source.splitlines(keepends=True)
# Process functions in reverse order to preserve line numbers
for func in sorted(functions, key=lambda f: f.starting_line, reverse=True):
instrumented = self._instrument_function(func, lines, file_path)
start_idx = func.starting_line - 1
end_idx = func.ending_line
lines = lines[:start_idx] + instrumented + lines[end_idx:]
instrumented_source = "".join(lines)
# Add tracer save at the end
tracer_save = self._generate_tracer_save()
return tracer_init + "\n" + instrumented_source + "\n" + tracer_save
def _generate_tracer_init(self) -> str:
"""Generate JavaScript code for tracer initialization."""
return f"""
// Codeflash function tracer initialization
const {self.tracer_var} = {{
traces: [],
callId: 0,
serialize: function(value) {{
try {{
// Handle special cases
if (value === undefined) return {{ __type__: 'undefined' }};
if (value === null) return null;
if (typeof value === 'function') return {{ __type__: 'function', name: value.name }};
if (typeof value === 'symbol') return {{ __type__: 'symbol', value: value.toString() }};
if (value instanceof Error) return {{
__type__: 'error',
name: value.name,
message: value.message,
stack: value.stack
}};
if (typeof value === 'bigint') return {{ __type__: 'bigint', value: value.toString() }};
if (value instanceof Date) return {{ __type__: 'date', value: value.toISOString() }};
if (value instanceof RegExp) return {{ __type__: 'regexp', value: value.toString() }};
if (value instanceof Map) return {{
__type__: 'map',
value: Array.from(value.entries()).map(([k, v]) => [this.serialize(k), this.serialize(v)])
}};
if (value instanceof Set) return {{
__type__: 'set',
value: Array.from(value).map(v => this.serialize(v))
}};
// Handle circular references with a simple check
return JSON.parse(JSON.stringify(value));
}} catch (e) {{
return {{ __type__: 'unserializable', error: e.message }};
}}
}},
wrap: function(originalFunc, funcName, filePath) {{
const self = this;
if (originalFunc.constructor.name === 'AsyncFunction') {{
return async function(...args) {{
const callId = self.callId++;
const start = process.hrtime.bigint();
let result, error;
try {{
result = await originalFunc.apply(this, args);
}} catch (e) {{
error = e;
}}
const end = process.hrtime.bigint();
self.traces.push({{
call_id: callId,
function: funcName,
file: filePath,
args: args.map(a => self.serialize(a)),
result: error ? null : self.serialize(result),
error: error ? self.serialize(error) : null,
runtime_ns: (end - start).toString(),
timestamp: Date.now()
}});
if (error) throw error;
return result;
}};
}}
return function(...args) {{
const callId = self.callId++;
const start = process.hrtime.bigint();
let result, error;
try {{
result = originalFunc.apply(this, args);
}} catch (e) {{
error = e;
}}
const end = process.hrtime.bigint();
self.traces.push({{
call_id: callId,
function: funcName,
file: filePath,
args: args.map(a => self.serialize(a)),
result: error ? null : self.serialize(result),
error: error ? self.serialize(error) : null,
runtime_ns: (end - start).toString(),
timestamp: Date.now()
}});
if (error) throw error;
return result;
}};
}},
saveToDb: function() {{
const sqlite3 = require('sqlite3').verbose();
const fs = require('fs');
const path = require('path');
const dbPath = '{self.output_db.as_posix()}';
const dbDir = path.dirname(dbPath);
if (!fs.existsSync(dbDir)) {{
fs.mkdirSync(dbDir, {{ recursive: true }});
}}
const db = new sqlite3.Database(dbPath);
db.serialize(() => {{
// Create table
db.run(`
CREATE TABLE IF NOT EXISTS traces (
id INTEGER PRIMARY KEY AUTOINCREMENT,
call_id INTEGER,
function TEXT,
file TEXT,
args TEXT,
result TEXT,
error TEXT,
runtime_ns TEXT,
timestamp INTEGER
)
`);
// Insert traces
const stmt = db.prepare(`
INSERT INTO traces (call_id, function, file, args, result, error, runtime_ns, timestamp)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`);
for (const trace of this.traces) {{
stmt.run(
trace.call_id,
trace.function,
trace.file,
JSON.stringify(trace.args),
JSON.stringify(trace.result),
JSON.stringify(trace.error),
trace.runtime_ns,
trace.timestamp
);
}}
stmt.finalize();
}});
db.close();
}},
saveToJson: function() {{
const fs = require('fs');
const path = require('path');
const jsonPath = '{self.output_db.with_suffix(".json").as_posix()}';
const jsonDir = path.dirname(jsonPath);
if (!fs.existsSync(jsonDir)) {{
fs.mkdirSync(jsonDir, {{ recursive: true }});
}}
fs.writeFileSync(jsonPath, JSON.stringify(this.traces, null, 2));
}}
}};
"""
def _generate_tracer_save(self) -> str:
"""Generate JavaScript code to save tracer results."""
return f"""
// Save tracer results on process exit
process.on('exit', () => {{
try {{
{self.tracer_var}.saveToJson();
// Try SQLite, but don't fail if sqlite3 is not installed
try {{
{self.tracer_var}.saveToDb();
}} catch (e) {{
// SQLite not available, JSON is sufficient
}}
}} catch (e) {{
console.error('Failed to save traces:', e);
}}
}});
"""
def _instrument_function(self, func: FunctionToOptimize, lines: list[str], file_path: Path) -> list[str]:
"""Instrument a single function with tracing.
Args:
func: Function to instrument.
lines: Source lines.
file_path: Path to source file.
Returns:
Instrumented function lines.
"""
func_lines = lines[func.starting_line - 1 : func.ending_line]
func_text = "".join(func_lines)
# Detect function pattern
func_name = func.function_name
is_arrow = "=>" in func_text.split("\n")[0]
is_method = func.is_method
is_async = func.is_async
# Generate wrapper code based on function type
if is_arrow:
# For arrow functions: const foo = (a, b) => { ... }
# Replace with: const foo = __codeflash_tracer__.wrap((a, b) => { ... }, 'foo', 'file.js')
return self._wrap_arrow_function(func_lines, func_name, file_path)
if is_method:
# For methods: methodName(a, b) { ... }
# Wrap the method body
return self._wrap_method(func_lines, func_name, file_path, is_async)
# For regular functions: function foo(a, b) { ... }
# Wrap the entire function
return self._wrap_regular_function(func_lines, func_name, file_path, is_async)
def _wrap_arrow_function(self, func_lines: list[str], func_name: str, file_path: Path) -> list[str]:
"""Wrap an arrow function with tracing."""
# Find the assignment line
first_line = func_lines[0]
indent = len(first_line) - len(first_line.lstrip())
indent_str = " " * indent
# Insert wrapper call
func_text = "".join(func_lines).rstrip()
# Find the '=' and wrap everything after it
if "=" in func_text:
parts = func_text.split("=", 1)
wrapped = f"{parts[0]}= {self.tracer_var}.wrap({parts[1]}, '{func_name}', '{file_path.as_posix()}');\n"
return [wrapped]
return func_lines
def _wrap_method(self, func_lines: list[str], func_name: str, file_path: Path, is_async: bool) -> list[str]:
"""Wrap a class method with tracing."""
# For methods, we wrap by reassigning them after definition
# This is complex, so for now we'll return unwrapped
# TODO: Implement method wrapping
logger.warning("Method wrapping not fully implemented for %s", func_name)
return func_lines
def _wrap_regular_function(
self, func_lines: list[str], func_name: str, file_path: Path, is_async: bool
) -> list[str]:
"""Wrap a regular function declaration with tracing."""
# Replace: function foo(a, b) { ... }
# With: const __original_foo = function foo(a, b) { ... }; const foo = __codeflash_tracer__.wrap(__original_foo, 'foo', 'file.js');
func_text = "".join(func_lines).rstrip()
first_line = func_lines[0]
indent = len(first_line) - len(first_line.lstrip())
indent_str = " " * indent
wrapped = (
f"{indent_str}const __original_{func_name}__ = {func_text};\n"
f"{indent_str}const {func_name} = {self.tracer_var}.wrap(__original_{func_name}__, '{func_name}', '{file_path.as_posix()}');\n"
)
return [wrapped]
@staticmethod
def parse_results(trace_file: Path) -> list[dict[str, Any]]:
"""Parse tracing results from output file.
Supports both the new function_calls schema and legacy traces schema.
Args:
trace_file: Path to traces JSON file.
trace_file: Path to traces file (SQLite or JSON).
Returns:
List of trace records.
@ -363,22 +77,45 @@ process.on('exit', () => {{
if json_file.exists():
try:
with json_file.open("r") as f:
return json.load(f)
with json_file.open("r", encoding="utf-8") as f:
data: list[dict[str, Any]] = json.load(f)
return data
except Exception as e:
logger.exception("Failed to parse trace JSON: %s", e)
return []
# Try SQLite database
if not trace_file.exists():
return []
try:
conn = sqlite3.connect(trace_file)
cursor = conn.cursor()
cursor.execute("SELECT * FROM traces ORDER BY id")
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
tables = {row[0] for row in cursor.fetchall()}
traces = []
if "function_calls" in tables:
cursor.execute(
"SELECT type, function, classname, filename, line_number, "
"last_frame_address, time_ns, args FROM function_calls ORDER BY time_ns"
)
for row in cursor.fetchall():
traces.append(
{
"type": row[0],
"function": row[1],
"classname": row[2],
"filename": row[3],
"line_number": row[4],
"last_frame_address": row[5],
"time_ns": row[6],
"args": json.loads(row[7]) if row[7] else [],
}
)
elif "traces" in tables:
cursor.execute("SELECT * FROM traces ORDER BY id")
for row in cursor.fetchall():
traces.append(
{
@ -386,11 +123,11 @@ process.on('exit', () => {{
"call_id": row[1],
"function": row[2],
"file": row[3],
"args": json.loads(row[4]),
"result": json.loads(row[5]),
"error": json.loads(row[6]) if row[6] != "null" else None,
"runtime_ns": int(row[7]),
"timestamp": row[8],
"args": json.loads(row[4]) if row[4] else [],
"result": json.loads(row[5]) if row[5] else None,
"error": json.loads(row[6]) if row[6] and row[6] != "null" else None,
"runtime_ns": int(row[7]) if row[7] else 0,
"timestamp": row[8] if len(row) > 8 else None,
}
)
@ -400,3 +137,145 @@ process.on('exit', () => {{
except Exception as e:
logger.exception("Failed to parse trace database: %s", e)
return []
@staticmethod
def get_traced_functions(trace_file: Path) -> list[JavaScriptFunctionInfo]:
if not trace_file.exists():
return []
try:
conn = sqlite3.connect(trace_file)
cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
tables = {row[0] for row in cursor.fetchall()}
functions = []
if "function_calls" in tables:
cursor.execute(
"SELECT DISTINCT function, filename, classname, line_number FROM function_calls WHERE type = 'call'"
)
for row in cursor.fetchall():
func_name = row[0]
file_name = row[1]
class_name = row[2]
line_number = row[3]
module_path = file_name.replace("\\", "/").replace(".js", "").replace(".ts", "")
module_path = module_path.removeprefix("./")
functions.append(
JavaScriptFunctionInfo(
function_name=func_name,
file_name=file_name,
module_path=module_path,
class_name=class_name,
line_number=line_number,
)
)
conn.close()
return functions
except Exception as e:
logger.exception("Failed to get traced functions: %s", e)
return []
def create_replay_test(
self,
trace_file: Path,
output_path: Path,
framework: str = "jest",
max_run_count: int = 100,
project_root: Optional[Path] = None,
) -> Optional[str]:
functions = self.get_traced_functions(trace_file)
if not functions:
logger.warning("No traced functions found in %s", trace_file)
return None
is_vitest = framework.lower() == "vitest"
imports = []
if is_vitest:
imports.append("import { describe, test } from 'vitest';")
imports.append("const { getNextArg } = require('codeflash/replay');")
imports.append("")
for func in functions:
alias = self._get_function_alias(func.module_path, func.function_name, func.class_name)
if func.class_name:
imports.append(f"const {{ {func.class_name}: {alias}_class }} = require('./{func.module_path}');")
else:
imports.append(f"const {{ {func.function_name}: {alias} }} = require('./{func.module_path}');")
imports.append("")
trace_path = trace_file.as_posix()
metadata = [
f"const traceFilePath = '{trace_path}';",
f"const functions = {json.dumps([f.function_name for f in functions])};",
"",
]
test_cases = []
for func in functions:
alias = self._get_function_alias(func.module_path, func.function_name, func.class_name)
test_name = f"{func.class_name}.{func.function_name}" if func.class_name else func.function_name
class_arg = f"'{func.class_name}'" if func.class_name else "null"
if func.class_name:
test_cases.append(
textwrap.dedent(f"""
describe('Replay: {test_name}', () => {{
const traces = getNextArg(traceFilePath, '{func.function_name}', '{func.file_name}', {max_run_count}, {class_arg});
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {{
const instance = new {alias}_class();
instance.{func.function_name}(...args);
}});
}});
""")
)
else:
test_cases.append(
textwrap.dedent(f"""
describe('Replay: {test_name}', () => {{
const traces = getNextArg(traceFilePath, '{func.function_name}', '{func.file_name}', {max_run_count});
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {{
{alias}(...args);
}});
}});
""")
)
content = "\n".join(
[
"// Auto-generated replay test by Codeflash",
"// Do not edit this file directly",
"",
*imports,
*metadata,
*test_cases,
]
)
try:
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(content, encoding="utf-8")
logger.info("Generated replay test: %s", output_path)
return str(output_path)
except Exception as e:
logger.exception("Failed to write replay test: %s", e)
return None
@staticmethod
def _get_function_alias(module_path: str, function_name: str, class_name: Optional[str] = None) -> str:
module_alias = re.sub(r"[^a-zA-Z0-9]", "_", module_path).strip("_")
if class_name:
return f"{module_alias}_{class_name}_{function_name}"
return f"{module_alias}_{function_name}"

231
codeflash/languages/javascript/tracer_runner.py Normal file
View file

@ -0,0 +1,231 @@
from __future__ import annotations
import json
import logging
import os
import shutil
import subprocess
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional
if TYPE_CHECKING:
from argparse import Namespace
logger = logging.getLogger(__name__)
def find_node_executable() -> Optional[Path]:
node_path = shutil.which("node")
if node_path:
return Path(node_path)
npx_path = shutil.which("npx")
if npx_path:
return Path(npx_path)
return None
def find_trace_runner() -> Optional[Path]:
cwd = Path.cwd()
local_path = cwd / "node_modules" / "codeflash" / "runtime" / "trace-runner.js"
if local_path.exists():
return local_path
try:
result = subprocess.run(["npm", "root", "-g"], capture_output=True, text=True, check=True)
global_modules = Path(result.stdout.strip())
global_path = global_modules / "codeflash" / "runtime" / "trace-runner.js"
if global_path.exists():
return global_path
except Exception:
pass
bundled_path = Path(__file__).parent.parent.parent.parent / "packages" / "codeflash" / "runtime" / "trace-runner.js"
if bundled_path.exists():
return bundled_path
return None
def run_javascript_tracer(args: Namespace, config: dict[str, Any], project_root: Path) -> dict[str, Any]:
result: dict[str, Any] = {"success": False, "trace_file": None, "replay_test_file": None, "error": None}
node_path = find_node_executable()
if not node_path:
result["error"] = "Node.js not found. Please install Node.js to use JavaScript tracing."
logger.error(result["error"])
return result
trace_runner_path = find_trace_runner()
if not trace_runner_path:
result["error"] = "trace-runner.js not found. Please install the codeflash npm package."
logger.error(result["error"])
return result
outfile = getattr(args, "outfile", None) or "codeflash.trace.sqlite"
trace_file = Path(outfile).resolve()
env = os.environ.copy()
env["CODEFLASH_TRACE_DB"] = str(trace_file)
env["CODEFLASH_PROJECT_ROOT"] = str(project_root)
max_count = getattr(args, "max_function_count", 256)
env["CODEFLASH_MAX_FUNCTION_COUNT"] = str(max_count)
timeout = getattr(args, "tracer_timeout", None)
if timeout:
env["CODEFLASH_TRACER_TIMEOUT"] = str(timeout)
only_functions = getattr(args, "only_functions", None)
if only_functions:
env["CODEFLASH_FUNCTIONS"] = json.dumps(only_functions)
cmd = [str(node_path), str(trace_runner_path)]
cmd.extend(["--trace-db", str(trace_file)])
cmd.extend(["--project-root", str(project_root)])
if max_count:
cmd.extend(["--max-function-count", str(max_count)])
if timeout:
cmd.extend(["--timeout", str(timeout)])
if only_functions:
cmd.extend(["--functions", json.dumps(only_functions)])
is_module = getattr(args, "module", False)
script_args = []
if hasattr(args, "script_args"):
script_args = args.script_args
elif hasattr(args, "unknown_args"):
script_args = args.unknown_args
if is_module and script_args and script_args[0] == "jest":
cmd.append("--jest")
cmd.append("--")
cmd.extend(script_args[1:])
elif is_module and script_args and script_args[0] == "vitest":
cmd.append("--vitest")
cmd.append("--")
cmd.extend(script_args[1:])
elif script_args:
cmd.extend(script_args)
logger.info("Running JavaScript tracer: %s", " ".join(cmd))
try:
process = subprocess.run(cmd, cwd=project_root, env=env, capture_output=False, check=False)
if process.returncode != 0:
result["error"] = f"Tracing failed with exit code {process.returncode}"
logger.error(result["error"])
return result
except Exception as e:
result["error"] = f"Failed to run tracer: {e}"
logger.exception(result["error"])
return result
if not trace_file.exists():
result["error"] = f"Trace file not created: {trace_file}"
logger.error(result["error"])
return result
result["success"] = True
result["trace_file"] = str(trace_file)
trace_only = getattr(args, "trace_only", False)
if not trace_only:
replay_test_path = generate_replay_test(trace_file=trace_file, project_root=project_root, config=config)
if replay_test_path:
result["replay_test_file"] = str(replay_test_path)
logger.info("Generated replay test: %s", replay_test_path)
return result
def generate_replay_test(
trace_file: Path, project_root: Path, config: dict[str, Any], output_path: Optional[Path] = None
) -> Optional[Path]:
from codeflash.languages.javascript.replay_test import create_replay_test_file
framework = detect_test_framework(project_root, config)
if output_path is None:
tests_root = config.get("tests_root", "tests")
tests_dir = project_root / tests_root
output_path = tests_dir / "codeflash_replay.test.js"
return create_replay_test_file(
trace_file=trace_file,
output_path=output_path,
framework=framework,
max_run_count=100,
project_root=project_root,
)
def detect_test_framework(project_root: Path, config: dict[str, Any]) -> str:
if "test_framework" in config:
framework: str = config["test_framework"]
return framework
vitest_configs = ["vitest.config.js", "vitest.config.ts", "vitest.config.mjs"]
for conf in vitest_configs:
if (project_root / conf).exists():
return "vitest"
jest_configs = ["jest.config.js", "jest.config.ts", "jest.config.mjs", "jest.config.json"]
for conf in jest_configs:
if (project_root / conf).exists():
return "jest"
package_json = project_root / "package.json"
if package_json.exists():
try:
with package_json.open(encoding="utf-8") as f:
pkg = json.load(f)
test_script = pkg.get("scripts", {}).get("test", "")
if "vitest" in test_script:
return "vitest"
if "jest" in test_script:
return "jest"
deps = {**pkg.get("dependencies", {}), **pkg.get("devDependencies", {})}
if "vitest" in deps:
return "vitest"
if "jest" in deps:
return "jest"
except Exception:
pass
return "jest"
def check_javascript_tracer_available() -> bool:
if not find_node_executable():
return False
if not find_trace_runner():
return False
return True
def get_tracer_requirements_message() -> str:
missing = []
if not find_node_executable():
missing.append("Node.js (v18+)")
if not find_trace_runner():
missing.append("codeflash npm package (npm install codeflash)")
if not missing:
return "All requirements met for JavaScript tracing."
return "Missing requirements for JavaScript tracing:\n- " + "\n- ".join(missing)

93
codeflash/tracer.py
View file

@ -149,6 +149,11 @@ def main(args: Namespace | None = None) -> ArgumentParser:
parser.add_argument(
"--limit", type=int, default=None, help="Limit the number of test files to process (for -m pytest mode)"
)
parser.add_argument(
"--language",
help="Language to trace (python, javascript, typescript). Auto-detected if not specified.",
default=None,
)
if args is not None:
parsed_args = args
@ -182,6 +187,13 @@ def main(args: Namespace | None = None) -> ArgumentParser:
outfile = parsed_args.outfile
config, found_config_path = parse_config_file(parsed_args.codeflash_config)
project_root = project_root_from_module_root(Path(config["module_root"]), found_config_path)
language = getattr(parsed_args, "language", None) or config.get("language", "python")
if language in ("javascript", "typescript"):
return run_javascript_tracer_main(
parsed_args=parsed_args, config=config, project_root=project_root, unknown_args=unknown_args
)
if len(unknown_args) > 0:
args_dict = {
"functions": parsed_args.only_functions,
@ -332,6 +344,87 @@ def main(args: Namespace | None = None) -> ArgumentParser:
return parser
def run_javascript_tracer_main(
parsed_args: Namespace, config: dict, project_root: Path, unknown_args: list[str]
) -> ArgumentParser:
from codeflash.languages.javascript.tracer_runner import (
check_javascript_tracer_available,
detect_test_framework,
get_tracer_requirements_message,
run_javascript_tracer,
)
if not check_javascript_tracer_available():
console.print("[bold red]Error:[/] JavaScript tracer requirements not met.")
console.print(get_tracer_requirements_message())
sys.exit(1)
trace_only = getattr(parsed_args, "trace_only", False)
only_functions = getattr(parsed_args, "only_functions", None)
max_function_count = getattr(parsed_args, "max_function_count", 256)
timeout = getattr(parsed_args, "tracer_timeout", None)
framework = detect_test_framework(project_root, config)
logger.info("JavaScript tracer: framework=%s, project_root=%s", framework, project_root)
trace_db_path = get_run_tmp_file(Path("js_trace.sqlite"))
script_or_test_args = unknown_args if unknown_args else []
replay_test_path = run_javascript_tracer(
script_args=script_or_test_args,
trace_db_path=trace_db_path,
project_root=project_root,
functions=only_functions,
max_function_count=max_function_count,
timeout=int(timeout) if timeout else 0,
framework=framework,
)
if replay_test_path and not trace_only:
from codeflash.cli_cmds.cli import parse_args as cli_parse_args
from codeflash.cli_cmds.cli import process_pyproject_config
from codeflash.cli_cmds.console import paneled_text
from codeflash.cli_cmds.console_constants import CODEFLASH_LOGO
from codeflash.languages import Language, set_current_language
from codeflash.optimization import optimizer
from codeflash.telemetry import posthog_cf
from codeflash.telemetry.sentry import init_sentry
language = getattr(parsed_args, "language", None) or config.get("language", "javascript")
if language == "typescript":
set_current_language(Language.TYPESCRIPT)
else:
set_current_language(Language.JAVASCRIPT)
sys.argv = ["codeflash", "--replay-test", str(replay_test_path)]
args = cli_parse_args()
paneled_text(
CODEFLASH_LOGO,
panel_args={"title": "https://codeflash.ai", "expand": False},
text_args={"style": "bold gold3"},
)
args = process_pyproject_config(args)
args.previous_checkpoint_functions = None
init_sentry(enabled=not args.disable_telemetry, exclude_errors=True)
posthog_cf.initialize_posthog(enabled=not args.disable_telemetry)
args.effort = EffortLevel.HIGH.value
optimizer.run_with_args(args)
# Clean up
trace_db_path.unlink(missing_ok=True)
if replay_test_path:
Path(replay_test_path).unlink(missing_ok=True)
elif replay_test_path:
console.print(f"[bold green]Trace complete.[/] Replay test: {replay_test_path}")
else:
console.print("[bold yellow]Warning:[/] No functions were traced.")
return ArgumentParser()
def _run_java_tracer(existing_args: Namespace | None = None) -> ArgumentParser:
"""Run the Java two-stage tracer (JFR + argument capture) and optionally optimize."""
from codeflash.cli_cmds.cli import parse_args, process_pyproject_config

21
packages/codeflash/package.json
View file

@ -6,7 +6,8 @@
"types": "runtime/index.d.ts",
"bin": {
"codeflash": "./bin/codeflash.js",
"codeflash-setup": "./bin/codeflash-setup.js"
"codeflash-setup": "./bin/codeflash-setup.js",
"codeflash-trace": "./runtime/trace-runner.js"
},
"publishConfig": {
"access": "public"
@ -36,6 +37,18 @@
"./jest-reporter": {
"require": "./runtime/jest-reporter.js",
"import": "./runtime/jest-reporter.js"
},
"./tracer": {
"require": "./runtime/tracer.js",
"import": "./runtime/tracer.js"
},
"./replay": {
"require": "./runtime/replay.js",
"import": "./runtime/replay.js"
},
"./babel-tracer-plugin": {
"require": "./runtime/babel-tracer-plugin.js",
"import": "./runtime/babel-tracer-plugin.js"
}
},
"scripts": {
@ -92,5 +105,11 @@
"dependencies": {
"better-sqlite3": "^12.0.0",
"@msgpack/msgpack": "^3.0.0"
},
"optionalDependencies": {
"@babel/core": "^7.24.0",
"@babel/register": "^7.24.0",
"@babel/preset-env": "^7.24.0",
"@babel/preset-typescript": "^7.24.0"
}
}

434
packages/codeflash/runtime/babel-tracer-plugin.js Normal file
View file

@ -0,0 +1,434 @@
/**
* Codeflash Babel Tracer Plugin
*
* A Babel plugin that instruments JavaScript/TypeScript functions for tracing.
* This plugin wraps functions with tracing calls to capture:
* - Function arguments
* - Return values
* - Execution time
*
* The plugin transforms:
* function foo(a, b) { return a + b; }
*
* Into:
* const __codeflash_tracer__ = require('codeflash/tracer');
* function foo(a, b) {
* return __codeflash_tracer__.wrap(function foo(a, b) { return a + b; }, 'foo', '/path/file.js', 1)
* .apply(this, arguments);
* }
*
* Supported function types:
* - FunctionDeclaration: function foo() {}
* - FunctionExpression: const foo = function() {}
* - ArrowFunctionExpression: const foo = () => {}
* - ClassMethod: class Foo { bar() {} }
* - ObjectMethod: const obj = { foo() {} }
*
* Configuration (via plugin options or environment variables):
* - functions: Array of function names to trace (traces all if not set)
* - files: Array of file patterns to trace (traces all if not set)
* - exclude: Array of patterns to exclude from tracing
*
* Usage with @babel/register:
* require('@babel/register')({
* plugins: [['codeflash/babel-tracer-plugin', { functions: ['myFunc'] }]],
* });
*
* Environment Variables:
* CODEFLASH_FUNCTIONS - JSON array of functions to trace
* CODEFLASH_TRACE_FILES - JSON array of file patterns to trace
* CODEFLASH_TRACE_EXCLUDE - JSON array of patterns to exclude
*/
'use strict';
const path = require('path');
// Parse environment variables for configuration
function getEnvConfig() {
const config = {
functions: null,
files: null,
exclude: null,
};
try {
if (process.env.CODEFLASH_FUNCTIONS) {
config.functions = JSON.parse(process.env.CODEFLASH_FUNCTIONS);
}
} catch (e) {
console.error('[codeflash-babel] Failed to parse CODEFLASH_FUNCTIONS:', e.message);
}
try {
if (process.env.CODEFLASH_TRACE_FILES) {
config.files = JSON.parse(process.env.CODEFLASH_TRACE_FILES);
}
} catch (e) {
console.error('[codeflash-babel] Failed to parse CODEFLASH_TRACE_FILES:', e.message);
}
try {
if (process.env.CODEFLASH_TRACE_EXCLUDE) {
config.exclude = JSON.parse(process.env.CODEFLASH_TRACE_EXCLUDE);
}
} catch (e) {
console.error('[codeflash-babel] Failed to parse CODEFLASH_TRACE_EXCLUDE:', e.message);
}
return config;
}
/**
* Check if a function should be traced based on configuration.
*
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {string|null} className - Class name (for methods)
* @param {Object} config - Plugin configuration
* @returns {boolean} - True if function should be traced
*/
function shouldTraceFunction(funcName, fileName, className, config) {
// Check exclude patterns first
if (config.exclude && config.exclude.length > 0) {
for (const pattern of config.exclude) {
if (typeof pattern === 'string') {
if (funcName === pattern || fileName.includes(pattern)) {
return false;
}
} else if (pattern instanceof RegExp) {
if (pattern.test(funcName) || pattern.test(fileName)) {
return false;
}
}
}
}
// Check file patterns
if (config.files && config.files.length > 0) {
const matchesFile = config.files.some(pattern => {
if (typeof pattern === 'string') {
return fileName.includes(pattern);
}
if (pattern instanceof RegExp) {
return pattern.test(fileName);
}
return false;
});
if (!matchesFile) return false;
}
// Check function names
if (config.functions && config.functions.length > 0) {
const matchesName = config.functions.some(f => {
if (typeof f === 'string') {
return f === funcName || f === `${className}.${funcName}`;
}
// Support object format: { function: 'name', file: 'path', class: 'className' }
if (typeof f === 'object' && f !== null) {
if (f.function && f.function !== funcName) return false;
if (f.file && !fileName.includes(f.file)) return false;
if (f.class && f.class !== className) return false;
return true;
}
return false;
});
if (!matchesName) return false;
}
return true;
}
/**
* Check if a path should be excluded from tracing (node_modules, etc.)
*
* @param {string} fileName - File path
* @returns {boolean} - True if file should be excluded
*/
function isExcludedPath(fileName) {
// Always exclude node_modules
if (fileName.includes('node_modules')) return true;
// Exclude common test runner internals
if (fileName.includes('jest-runner') || fileName.includes('jest-jasmine')) return true;
if (fileName.includes('@vitest')) return true;
// Exclude this plugin itself
if (fileName.includes('codeflash/runtime')) return true;
if (fileName.includes('babel-tracer-plugin')) return true;
return false;
}
/**
* Create the Babel plugin.
*
* @param {Object} babel - Babel object with types (t)
* @returns {Object} - Babel plugin configuration
*/
module.exports = function codeflashTracerPlugin(babel) {
const { types: t } = babel;
// Merge environment config with plugin options
const envConfig = getEnvConfig();
return {
name: 'codeflash-tracer',
visitor: {
Program: {
enter(programPath, state) {
// Merge options from plugin config and environment
state.codeflashConfig = {
...envConfig,
...(state.opts || {}),
};
// Track whether we've added the tracer import
state.tracerImportAdded = false;
// Get file info
state.fileName = state.filename || state.file.opts.filename || 'unknown';
// Check if entire file should be excluded
if (isExcludedPath(state.fileName)) {
state.skipFile = true;
return;
}
state.skipFile = false;
},
exit(programPath, state) {
// Add tracer import if we instrumented any functions
if (state.tracerImportAdded) {
const tracerRequire = t.variableDeclaration('const', [
t.variableDeclarator(
t.identifier('__codeflash_tracer__'),
t.callExpression(
t.identifier('require'),
[t.stringLiteral('codeflash/tracer')]
)
),
]);
// Add at the beginning of the program
programPath.unshiftContainer('body', tracerRequire);
}
},
},
// Handle: function foo() {}
FunctionDeclaration(path, state) {
if (state.skipFile) return;
if (!path.node.id) return; // Skip anonymous functions
const funcName = path.node.id.name;
const lineNumber = path.node.loc ? path.node.loc.start.line : 0;
if (!shouldTraceFunction(funcName, state.fileName, null, state.codeflashConfig)) {
return;
}
// Transform the function body to wrap with tracing
wrapFunctionBody(t, path, funcName, state.fileName, lineNumber, null);
state.tracerImportAdded = true;
},
// Handle: const foo = function() {} or const foo = () => {}
VariableDeclarator(path, state) {
if (state.skipFile) return;
if (!t.isIdentifier(path.node.id)) return;
if (!path.node.init) return;
const init = path.node.init;
if (!t.isFunctionExpression(init) && !t.isArrowFunctionExpression(init)) {
return;
}
const funcName = path.node.id.name;
const lineNumber = path.node.loc ? path.node.loc.start.line : 0;
if (!shouldTraceFunction(funcName, state.fileName, null, state.codeflashConfig)) {
return;
}
// Wrap the function expression with tracer.wrap()
path.node.init = createWrapperCall(t, init, funcName, state.fileName, lineNumber, null);
state.tracerImportAdded = true;
},
// Handle: class Foo { bar() {} }
ClassMethod(path, state) {
if (state.skipFile) return;
if (path.node.kind === 'constructor') return; // Skip constructors for now
const funcName = path.node.key.name || (path.node.key.value && String(path.node.key.value));
if (!funcName) return;
// Get class name from parent
const classPath = path.findParent(p => t.isClassDeclaration(p) || t.isClassExpression(p));
const className = classPath && classPath.node.id ? classPath.node.id.name : null;
const lineNumber = path.node.loc ? path.node.loc.start.line : 0;
if (!shouldTraceFunction(funcName, state.fileName, className, state.codeflashConfig)) {
return;
}
// Wrap the method body
wrapMethodBody(t, path, funcName, state.fileName, lineNumber, className);
state.tracerImportAdded = true;
},
// Handle: const obj = { foo() {} }
ObjectMethod(path, state) {
if (state.skipFile) return;
const funcName = path.node.key.name || (path.node.key.value && String(path.node.key.value));
if (!funcName) return;
const lineNumber = path.node.loc ? path.node.loc.start.line : 0;
if (!shouldTraceFunction(funcName, state.fileName, null, state.codeflashConfig)) {
return;
}
// Wrap the method body
wrapMethodBody(t, path, funcName, state.fileName, lineNumber, null);
state.tracerImportAdded = true;
},
},
};
};
/**
* Create a __codeflash_tracer__.wrap() call expression.
*
* @param {Object} t - Babel types
* @param {Object} funcNode - The function AST node
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name
* @returns {Object} - Call expression AST node
*/
function createWrapperCall(t, funcNode, funcName, fileName, lineNumber, className) {
const args = [
funcNode,
t.stringLiteral(funcName),
t.stringLiteral(fileName),
t.numericLiteral(lineNumber),
];
if (className) {
args.push(t.stringLiteral(className));
} else {
args.push(t.nullLiteral());
}
return t.callExpression(
t.memberExpression(
t.identifier('__codeflash_tracer__'),
t.identifier('wrap')
),
args
);
}
/**
* Wrap a function declaration's body with tracing.
* Transforms:
* function foo(a, b) { return a + b; }
* Into:
* function foo(a, b) {
* const __original__ = function(a, b) { return a + b; };
* return __codeflash_tracer__.wrap(__original__, 'foo', 'file.js', 1, null).apply(this, arguments);
* }
*
* @param {Object} t - Babel types
* @param {Object} path - Babel path
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name
*/
function wrapFunctionBody(t, path, funcName, fileName, lineNumber, className) {
const node = path.node;
const isAsync = node.async;
const isGenerator = node.generator;
// Create a copy of the original function as an expression
const originalFunc = t.functionExpression(
null, // anonymous
node.params,
node.body,
isGenerator,
isAsync
);
// Create the wrapper call
const wrapperCall = createWrapperCall(t, originalFunc, funcName, fileName, lineNumber, className);
// Create: return __codeflash_tracer__.wrap(...).apply(this, arguments)
const applyCall = t.callExpression(
t.memberExpression(wrapperCall, t.identifier('apply')),
[t.thisExpression(), t.identifier('arguments')]
);
const returnStatement = t.returnStatement(applyCall);
// Replace the function body
node.body = t.blockStatement([returnStatement]);
}
/**
* Wrap a method's body with tracing.
* Similar to wrapFunctionBody but preserves method semantics.
*
* @param {Object} t - Babel types
* @param {Object} path - Babel path
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name
*/
function wrapMethodBody(t, path, funcName, fileName, lineNumber, className) {
const node = path.node;
const isAsync = node.async;
const isGenerator = node.generator;
// Create a copy of the original function as an expression
const originalFunc = t.functionExpression(
null, // anonymous
node.params,
node.body,
isGenerator,
isAsync
);
// Create the wrapper call
const wrapperCall = createWrapperCall(t, originalFunc, funcName, fileName, lineNumber, className);
// Create: return __codeflash_tracer__.wrap(...).apply(this, arguments)
const applyCall = t.callExpression(
t.memberExpression(wrapperCall, t.identifier('apply')),
[t.thisExpression(), t.identifier('arguments')]
);
let returnStatement;
if (isAsync) {
// For async methods, we need to await the result
returnStatement = t.returnStatement(t.awaitExpression(applyCall));
} else {
returnStatement = t.returnStatement(applyCall);
}
// Replace the function body
node.body = t.blockStatement([returnStatement]);
}
// Export helper functions for testing
module.exports.shouldTraceFunction = shouldTraceFunction;
module.exports.isExcludedPath = isExcludedPath;
module.exports.getEnvConfig = getEnvConfig;

38
packages/codeflash/runtime/index.js
View file

@ -8,6 +8,8 @@
* - capturePerf: Capture performance metrics (timing only)
* - serialize/deserialize: Value serialization for storage
* - comparator: Deep equality comparison
* - tracer: Function tracing for replay test generation
* - replay: Replay test utilities
*
* Usage (CommonJS):
* const { capture, capturePerf } = require('codeflash');
@ -30,6 +32,22 @@ const comparator = require('./comparator');
// Result comparison (used by CLI)
const compareResults = require('./compare-results');
// Function tracing (for replay test generation)
let tracer = null;
try {
tracer = require('./tracer');
} catch (e) {
// Tracer may not be available if better-sqlite3 is not installed
}
// Replay test utilities
let replay = null;
try {
replay = require('./replay');
} catch (e) {
// Replay may not be available
}
// Re-export all public APIs
module.exports = {
// === Main Instrumentation API ===
@ -88,4 +106,24 @@ module.exports = {
// === Feature Detection ===
hasV8: serializer.hasV8,
hasMsgpack: serializer.hasMsgpack,
// === Function Tracing (for replay test generation) ===
tracer: tracer ? {
init: tracer.init,
wrap: tracer.wrap,
createWrapper: tracer.createWrapper,
disable: tracer.disable,
enable: tracer.enable,
getStats: tracer.getStats,
} : null,
// === Replay Test Utilities ===
replay: replay ? {
getNextArg: replay.getNextArg,
getTracesWithMetadata: replay.getTracesWithMetadata,
getTracedFunctions: replay.getTracedFunctions,
getTraceMetadata: replay.getTraceMetadata,
generateReplayTest: replay.generateReplayTest,
createReplayTestFromTrace: replay.createReplayTestFromTrace,
} : null,
};

454
packages/codeflash/runtime/replay.js Normal file
View file

@ -0,0 +1,454 @@
/**
* Codeflash Replay Test Utilities
*
* This module provides utilities for generating and running replay tests
* from traced function calls. Replay tests allow verifying that optimized
* code produces the same results as the original code.
*
* Usage:
* const { getNextArg, createReplayTest } = require('codeflash/replay');
*
* // In a test file:
* describe('Replay tests', () => {
* test.each(getNextArg(traceFile, 'myFunction', '/path/file.js', 25))
* ('myFunction replay %#', (args) => {
* myFunction(...args);
* });
* });
*
* The module supports both Jest and Vitest test frameworks.
*/
'use strict';
const path = require('path');
const fs = require('fs');
// Load the codeflash serializer for argument deserialization
const serializer = require('./serializer');
// ============================================================================
// DATABASE ACCESS
// ============================================================================
/**
* Open a SQLite database connection.
*
* @param {string} dbPath - Path to the SQLite database
* @returns {Object|null} - Database connection or null if failed
*/
function openDatabase(dbPath) {
try {
const Database = require('better-sqlite3');
return new Database(dbPath, { readonly: true });
} catch (e) {
console.error('[codeflash-replay] Failed to open database:', e.message);
return null;
}
}
/**
* Get traced function calls from the database.
*
* @param {string} traceFile - Path to the trace SQLite database
* @param {string} functionName - Name of the function
* @param {string} fileName - Path to the source file
* @param {string|null} className - Class name (for methods)
* @param {number} limit - Maximum number of traces to retrieve
* @returns {Array} - Array of traced arguments
*/
function getNextArg(traceFile, functionName, fileName, limit = 25, className = null) {
const db = openDatabase(traceFile);
if (!db) {
return [];
}
try {
let stmt;
let rows;
if (className) {
stmt = db.prepare(`
SELECT args FROM function_calls
WHERE function = ? AND filename = ? AND classname = ? AND type = 'call'
ORDER BY time_ns ASC
LIMIT ?
`);
rows = stmt.all(functionName, fileName, className, limit);
} else {
stmt = db.prepare(`
SELECT args FROM function_calls
WHERE function = ? AND filename = ? AND type = 'call'
ORDER BY time_ns ASC
LIMIT ?
`);
rows = stmt.all(functionName, fileName, limit);
}
db.close();
// Deserialize arguments
return rows.map((row, index) => {
try {
const args = serializer.deserialize(row.args);
return args;
} catch (e) {
console.warn(`[codeflash-replay] Failed to deserialize args at index ${index}:`, e.message);
return [];
}
});
} catch (e) {
console.error('[codeflash-replay] Database query failed:', e.message);
db.close();
return [];
}
}
/**
* Get traced function calls with full metadata.
*
* @param {string} traceFile - Path to the trace SQLite database
* @param {string} functionName - Name of the function
* @param {string} fileName - Path to the source file
* @param {string|null} className - Class name (for methods)
* @param {number} limit - Maximum number of traces to retrieve
* @returns {Array} - Array of trace objects with args and metadata
*/
function getTracesWithMetadata(traceFile, functionName, fileName, limit = 25, className = null) {
const db = openDatabase(traceFile);
if (!db) {
return [];
}
try {
let stmt;
let rows;
if (className) {
stmt = db.prepare(`
SELECT type, function, classname, filename, line_number, time_ns, args
FROM function_calls
WHERE function = ? AND filename = ? AND classname = ? AND type = 'call'
ORDER BY time_ns ASC
LIMIT ?
`);
rows = stmt.all(functionName, fileName, className, limit);
} else {
stmt = db.prepare(`
SELECT type, function, classname, filename, line_number, time_ns, args
FROM function_calls
WHERE function = ? AND filename = ? AND type = 'call'
ORDER BY time_ns ASC
LIMIT ?
`);
rows = stmt.all(functionName, fileName, limit);
}
db.close();
// Deserialize arguments and return with metadata
return rows.map((row, index) => {
let args;
try {
args = serializer.deserialize(row.args);
} catch (e) {
console.warn(`[codeflash-replay] Failed to deserialize args at index ${index}:`, e.message);
args = [];
}
return {
args,
function: row.function,
className: row.classname,
fileName: row.filename,
lineNumber: row.line_number,
timeNs: row.time_ns,
};
});
} catch (e) {
console.error('[codeflash-replay] Database query failed:', e.message);
db.close();
return [];
}
}
/**
* Get all traced functions from the database.
*
* @param {string} traceFile - Path to the trace SQLite database
* @returns {Array} - Array of { function, fileName, className, count } objects
*/
function getTracedFunctions(traceFile) {
const db = openDatabase(traceFile);
if (!db) {
return [];
}
try {
const stmt = db.prepare(`
SELECT function, filename, classname, COUNT(*) as count
FROM function_calls
WHERE type = 'call'
GROUP BY function, filename, classname
ORDER BY count DESC
`);
const rows = stmt.all();
db.close();
return rows.map(row => ({
function: row.function,
fileName: row.filename,
className: row.classname,
count: row.count,
}));
} catch (e) {
console.error('[codeflash-replay] Failed to get traced functions:', e.message);
db.close();
return [];
}
}
/**
* Get metadata from the trace database.
*
* @param {string} traceFile - Path to the trace SQLite database
* @returns {Object} - Metadata key-value pairs
*/
function getTraceMetadata(traceFile) {
const db = openDatabase(traceFile);
if (!db) {
return {};
}
try {
const stmt = db.prepare('SELECT key, value FROM metadata');
const rows = stmt.all();
db.close();
const metadata = {};
for (const row of rows) {
metadata[row.key] = row.value;
}
return metadata;
} catch (e) {
console.error('[codeflash-replay] Failed to get metadata:', e.message);
db.close();
return {};
}
}
// ============================================================================
// TEST GENERATION
// ============================================================================
/**
* Generate a Jest/Vitest replay test file.
*
* @param {string} traceFile - Path to the trace SQLite database
* @param {Array} functions - Array of { function, fileName, className, modulePath } to test
* @param {Object} options - Generation options
* @returns {string} - Generated test file content
*/
function generateReplayTest(traceFile, functions, options = {}) {
const {
framework = 'jest', // 'jest' or 'vitest'
maxRunCount = 100,
outputPath = null,
} = options;
const isVitest = framework === 'vitest';
// Build imports section
const imports = [];
if (isVitest) {
imports.push("import { describe, test } from 'vitest';");
}
imports.push("const { getNextArg } = require('codeflash/replay');");
imports.push('');
// Build function imports
for (const func of functions) {
const alias = getFunctionAlias(func.modulePath, func.function, func.className);
if (func.className) {
// Import class for method testing
imports.push(`const { ${func.className}: ${alias}_class } = require('${func.modulePath}');`);
} else {
// Import function directly
imports.push(`const { ${func.function}: ${alias} } = require('${func.modulePath}');`);
}
}
imports.push('');
// Metadata
const metadata = [
`const traceFilePath = '${traceFile}';`,
`const functions = ${JSON.stringify(functions.map(f => f.function))};`,
'',
];
// Build test cases
const testCases = [];
for (const func of functions) {
const alias = getFunctionAlias(func.modulePath, func.function, func.className);
const testName = func.className
? `${func.className}.${func.function}`
: func.function;
if (func.className) {
// Method test
testCases.push(`
describe('Replay: ${testName}', () => {
const traces = getNextArg(traceFilePath, '${func.function}', '${func.fileName}', ${maxRunCount}, '${func.className}');
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {
// For instance methods, first arg is 'this' context
const [thisArg, ...methodArgs] = args;
const instance = thisArg || new ${alias}_class();
instance.${func.function}(...methodArgs);
});
});
`);
} else {
// Function test
testCases.push(`
describe('Replay: ${testName}', () => {
const traces = getNextArg(traceFilePath, '${func.function}', '${func.fileName}', ${maxRunCount});
test.each(traces.map((args, i) => [i, args]))('call %i', (index, args) => {
${alias}(...args);
});
});
`);
}
}
// Combine all parts
const content = [
'// Auto-generated replay test by Codeflash',
'// Do not edit this file directly',
'',
...imports,
...metadata,
...testCases,
].join('\n');
// Write to file if outputPath provided
if (outputPath) {
const dir = path.dirname(outputPath);
if (!fs.existsSync(dir)) {
fs.mkdirSync(dir, { recursive: true });
}
fs.writeFileSync(outputPath, content);
console.log(`[codeflash-replay] Generated test file: ${outputPath}`);
}
return content;
}
/**
* Create a function alias for imports to avoid naming conflicts.
*
* @param {string} modulePath - Module path
* @param {string} functionName - Function name
* @param {string|null} className - Class name
* @returns {string} - Alias name
*/
function getFunctionAlias(modulePath, functionName, className = null) {
// Normalize module path to valid identifier
const moduleAlias = modulePath
.replace(/[^a-zA-Z0-9]/g, '_')
.replace(/^_+|_+$/g, '');
if (className) {
return `${moduleAlias}_${className}_${functionName}`;
}
return `${moduleAlias}_${functionName}`;
}
/**
* Create replay tests from a trace file.
* This is the main entry point for Python integration.
*
* @param {string} traceFile - Path to the trace SQLite database
* @param {string} outputPath - Path to write the test file
* @param {Object} options - Generation options
* @returns {Object} - { success, outputPath, functions }
*/
function createReplayTestFromTrace(traceFile, outputPath, options = {}) {
const {
framework = 'jest',
maxRunCount = 100,
projectRoot = process.cwd(),
} = options;
// Get all traced functions
const tracedFunctions = getTracedFunctions(traceFile);
if (tracedFunctions.length === 0) {
console.warn('[codeflash-replay] No traced functions found in database');
return { success: false, outputPath: null, functions: [] };
}
// Convert to the format expected by generateReplayTest
const functions = tracedFunctions.map(tf => {
// Calculate module path from file name
let modulePath = tf.fileName;
// Make relative to project root
if (path.isAbsolute(modulePath)) {
modulePath = path.relative(projectRoot, modulePath);
}
// Convert to module path (remove .js extension, use forward slashes)
modulePath = './' + modulePath
.replace(/\\/g, '/')
.replace(/\.js$/, '')
.replace(/\.ts$/, '');
return {
function: tf.function,
fileName: tf.fileName,
className: tf.className,
modulePath,
};
});
// Generate the test file
const testContent = generateReplayTest(traceFile, functions, {
framework,
maxRunCount,
outputPath,
});
return {
success: true,
outputPath,
functions: functions.map(f => f.function),
content: testContent,
};
}
// ============================================================================
// EXPORTS
// ============================================================================
module.exports = {
// Core API
getNextArg,
getTracesWithMetadata,
getTracedFunctions,
getTraceMetadata,
// Test generation
generateReplayTest,
createReplayTestFromTrace,
getFunctionAlias,
// Database utilities
openDatabase,
};

381
packages/codeflash/runtime/trace-runner.js Normal file
View file

@ -0,0 +1,381 @@
#!/usr/bin/env node
/**
* Codeflash Trace Runner
*
* Entry point script that runs JavaScript/TypeScript code with function tracing enabled.
* This script:
* 1. Registers Babel with the tracer plugin for AST transformation
* 2. Sets up environment variables for tracing configuration
* 3. Runs the user's script, tests, or module
*
* Usage:
* # Run a script with tracing
* node trace-runner.js script.js
*
* # Run tests with tracing (Jest)
* node trace-runner.js --jest -- --testPathPattern=mytest
*
* # Run tests with tracing (Vitest)
* node trace-runner.js --vitest -- --run
*
* # Run with specific functions to trace
* node trace-runner.js --functions='["myFunc","otherFunc"]' script.js
*
* Environment Variables (also settable via command line):
* CODEFLASH_TRACE_DB - Path to SQLite database for storing traces
* CODEFLASH_PROJECT_ROOT - Project root for relative path calculation
* CODEFLASH_FUNCTIONS - JSON array of functions to trace
* CODEFLASH_MAX_FUNCTION_COUNT - Maximum traces per function (default: 256)
* CODEFLASH_TRACER_TIMEOUT - Timeout in seconds for tracing
*
* For ESM (ECMAScript modules), use the loader flag:
* node --loader ./esm-loader.mjs trace-runner.js script.mjs
*/
'use strict';
const path = require('path');
const fs = require('fs');
// ============================================================================
// ARGUMENT PARSING
// ============================================================================
function parseArgs(args) {
const config = {
traceDb: process.env.CODEFLASH_TRACE_DB || path.join(process.cwd(), 'codeflash.trace.sqlite'),
projectRoot: process.env.CODEFLASH_PROJECT_ROOT || process.cwd(),
functions: process.env.CODEFLASH_FUNCTIONS || null,
maxFunctionCount: process.env.CODEFLASH_MAX_FUNCTION_COUNT || '256',
tracerTimeout: process.env.CODEFLASH_TRACER_TIMEOUT || null,
traceFiles: process.env.CODEFLASH_TRACE_FILES || null,
traceExclude: process.env.CODEFLASH_TRACE_EXCLUDE || null,
jest: false,
vitest: false,
module: false,
script: null,
scriptArgs: [],
};
let i = 0;
while (i < args.length) {
const arg = args[i];
if (arg === '--trace-db') {
config.traceDb = args[++i];
} else if (arg.startsWith('--trace-db=')) {
config.traceDb = arg.split('=')[1];
} else if (arg === '--project-root') {
config.projectRoot = args[++i];
} else if (arg.startsWith('--project-root=')) {
config.projectRoot = arg.split('=')[1];
} else if (arg === '--functions') {
config.functions = args[++i];
} else if (arg.startsWith('--functions=')) {
config.functions = arg.split('=')[1];
} else if (arg === '--max-function-count') {
config.maxFunctionCount = args[++i];
} else if (arg.startsWith('--max-function-count=')) {
config.maxFunctionCount = arg.split('=')[1];
} else if (arg === '--timeout') {
config.tracerTimeout = args[++i];
} else if (arg.startsWith('--timeout=')) {
config.tracerTimeout = arg.split('=')[1];
} else if (arg === '--trace-files') {
config.traceFiles = args[++i];
} else if (arg.startsWith('--trace-files=')) {
config.traceFiles = arg.split('=')[1];
} else if (arg === '--trace-exclude') {
config.traceExclude = args[++i];
} else if (arg.startsWith('--trace-exclude=')) {
config.traceExclude = arg.split('=')[1];
} else if (arg === '--jest') {
config.jest = true;
} else if (arg === '--vitest') {
config.vitest = true;
} else if (arg === '-m' || arg === '--module') {
config.module = true;
} else if (arg === '--') {
// Everything after -- is passed to the script/test runner
config.scriptArgs = args.slice(i + 1);
break;
} else if (arg === '--help' || arg === '-h') {
printHelp();
process.exit(0);
} else if (!arg.startsWith('-')) {
// First non-flag argument is the script
config.script = arg;
config.scriptArgs = args.slice(i + 1);
break;
}
i++;
}
return config;
}
function printHelp() {
console.log(`
Codeflash Trace Runner - JavaScript Function Tracing
Usage:
trace-runner [options] <script> [script-args...]
trace-runner [options] --jest -- [jest-args...]
trace-runner [options] --vitest -- [vitest-args...]
Options:
--trace-db <path> Path to SQLite database for traces (default: ./codeflash.trace.sqlite)
--project-root <path> Project root directory (default: cwd)
--functions <json> JSON array of functions to trace (traces all if not set)
--max-function-count <n> Maximum traces per function (default: 256)
--timeout <seconds> Timeout for tracing
--trace-files <json> JSON array of file patterns to trace
--trace-exclude <json> JSON array of patterns to exclude from tracing
--jest Run with Jest test framework
--vitest Run with Vitest test framework
-m, --module Run a module (like python -m)
-h, --help Show this help message
Examples:
# Trace a script
trace-runner --functions='["processData"]' ./src/main.js
# Trace Jest tests
trace-runner --jest --functions='["myFunc"]' -- --testPathPattern=mytest
# Trace Vitest tests
trace-runner --vitest -- --run
Environment Variables:
CODEFLASH_TRACE_DB Path to SQLite database
CODEFLASH_PROJECT_ROOT Project root directory
CODEFLASH_FUNCTIONS JSON array of functions to trace
CODEFLASH_MAX_FUNCTION_COUNT Maximum traces per function
CODEFLASH_TRACER_TIMEOUT Timeout in seconds
`);
}
// ============================================================================
// BABEL REGISTRATION
// ============================================================================
function registerBabel(config) {
// Set environment variables before loading Babel
process.env.CODEFLASH_TRACE_DB = config.traceDb;
process.env.CODEFLASH_PROJECT_ROOT = config.projectRoot;
process.env.CODEFLASH_MAX_FUNCTION_COUNT = config.maxFunctionCount;
if (config.functions) {
process.env.CODEFLASH_FUNCTIONS = config.functions;
}
if (config.tracerTimeout) {
process.env.CODEFLASH_TRACER_TIMEOUT = config.tracerTimeout;
}
if (config.traceFiles) {
process.env.CODEFLASH_TRACE_FILES = config.traceFiles;
}
if (config.traceExclude) {
process.env.CODEFLASH_TRACE_EXCLUDE = config.traceExclude;
}
// Try to find @babel/register
let babelRegister;
try {
babelRegister = require('@babel/register');
} catch (e) {
console.error('[codeflash] Error: @babel/register is required for tracing.');
console.error('Install it with: npm install --save-dev @babel/register @babel/core');
process.exit(1);
}
// Get the path to our Babel plugin
const pluginPath = path.join(__dirname, 'babel-tracer-plugin.js');
// Configure Babel
const babelConfig = {
// Use our tracer plugin
plugins: [pluginPath],
// Compile only project files, not node_modules
ignore: [/node_modules/],
// Only compile files in project root
only: [config.projectRoot],
// Don't look for .babelrc files - we provide all config
babelrc: false,
configFile: false,
// Support TypeScript and modern JS
presets: [],
// Enable source maps for better error messages
sourceMaps: 'inline',
// Cache for faster repeated runs
cache: true,
// File extensions to process
extensions: ['.js', '.jsx', '.ts', '.tsx', '.mjs', '.cjs'],
};
// Try to add TypeScript support if available
try {
require.resolve('@babel/preset-typescript');
babelConfig.presets.push('@babel/preset-typescript');
} catch (e) {
// TypeScript preset not available, skip
}
// Try to add modern JS support
try {
require.resolve('@babel/preset-env');
babelConfig.presets.push(['@babel/preset-env', { targets: { node: 'current' } }]);
} catch (e) {
// preset-env not available, skip
}
// Register Babel
babelRegister(babelConfig);
console.log(`[codeflash] Tracing enabled. Output: ${config.traceDb}`);
if (config.functions) {
console.log(`[codeflash] Tracing functions: ${config.functions}`);
}
}
// ============================================================================
// SCRIPT EXECUTION
// ============================================================================
function runScript(config) {
if (!config.script) {
console.error('[codeflash] Error: No script specified');
printHelp();
process.exit(1);
}
// Resolve script path
const scriptPath = path.resolve(config.projectRoot, config.script);
if (!fs.existsSync(scriptPath)) {
console.error(`[codeflash] Error: Script not found: ${scriptPath}`);
process.exit(1);
}
// Update process.argv for the script
process.argv = [process.argv[0], scriptPath, ...config.scriptArgs];
// Run the script
require(scriptPath);
}
function runJest(config) {
// Find Jest
let jestPath;
try {
jestPath = require.resolve('jest');
} catch (e) {
console.error('[codeflash] Error: Jest not found. Install it with: npm install --save-dev jest');
process.exit(1);
}
// Get Jest CLI path
const jestCli = path.join(path.dirname(jestPath), 'cli');
// Update process.argv for Jest
process.argv = [process.argv[0], 'jest', ...config.scriptArgs];
// Run Jest
const jest = require(jestCli);
jest.run();
}
function runVitest(config) {
// Vitest needs special handling as it's ESM-first
// We'll spawn it as a subprocess with our loader
const { spawn } = require('child_process');
const args = [
'--experimental-vm-modules',
require.resolve('vitest/vitest.mjs'),
'run',
...config.scriptArgs,
];
const env = {
...process.env,
CODEFLASH_TRACE_DB: config.traceDb,
CODEFLASH_PROJECT_ROOT: config.projectRoot,
CODEFLASH_MAX_FUNCTION_COUNT: config.maxFunctionCount,
};
if (config.functions) {
env.CODEFLASH_FUNCTIONS = config.functions;
}
if (config.tracerTimeout) {
env.CODEFLASH_TRACER_TIMEOUT = config.tracerTimeout;
}
console.log('[codeflash] Running Vitest with tracing...');
console.log('[codeflash] Note: ESM tracing requires additional setup. See documentation.');
const child = spawn(process.execPath, args, {
env,
stdio: 'inherit',
cwd: config.projectRoot,
});
child.on('exit', (code) => {
process.exit(code || 0);
});
}
function runModule(config) {
if (!config.script) {
console.error('[codeflash] Error: No module specified');
printHelp();
process.exit(1);
}
// For module mode, we resolve the module from the project root
const modulePath = require.resolve(config.script, { paths: [config.projectRoot] });
// Update process.argv
process.argv = [process.argv[0], modulePath, ...config.scriptArgs];
// Run the module
require(modulePath);
}
// ============================================================================
// MAIN
// ============================================================================
function main() {
// Parse command line arguments (skip node and script name)
const args = process.argv.slice(2);
const config = parseArgs(args);
// Register Babel with tracer plugin
registerBabel(config);
// Initialize the tracer
const tracer = require('./tracer');
tracer.init(config.traceDb, config.projectRoot);
// Run based on mode
if (config.jest) {
runJest(config);
} else if (config.vitest) {
runVitest(config);
} else if (config.module) {
runModule(config);
} else {
runScript(config);
}
}
main();

558
packages/codeflash/runtime/tracer.js Normal file
View file

@ -0,0 +1,558 @@
/**
* Codeflash JavaScript Function Tracer
*
* This module provides function tracing instrumentation that captures:
* - Function inputs (arguments)
* - Return values
* - Exceptions thrown
* - Execution time (nanosecond precision)
*
* Traces are stored in SQLite database for later replay test generation.
* This mirrors the Python tracer functionality in codeflash/tracing/.
*
* Database Schema (matches Python tracer):
* - function_calls: Main trace data (type, function, classname, filename, line_number, time_ns, args)
* - metadata: Key-value metadata about the trace session
* - pstats: Profiling statistics (optional)
*
* Usage:
* const tracer = require('codeflash/tracer');
* tracer.init('/path/to/output.sqlite', ['/path/to/project']);
*
* // Wrap a function for tracing
* const tracedFunc = tracer.wrap(originalFunc, 'funcName', '/path/to/file.js', 10);
*
* // Or use the decorator pattern
* tracer.trace('funcName', '/path/to/file.js', 10, () => {
* // function body
* });
*
* Environment Variables:
* CODEFLASH_TRACE_DB - Path to SQLite database for storing traces
* CODEFLASH_PROJECT_ROOT - Project root for relative path calculation
* CODEFLASH_FUNCTIONS - JSON array of functions to trace (optional, traces all if not set)
* CODEFLASH_MAX_FUNCTION_COUNT - Maximum traces per function (default: 256)
* CODEFLASH_TRACER_TIMEOUT - Timeout in seconds for tracing (optional)
*/
'use strict';
const path = require('path');
const fs = require('fs');
// Load the codeflash serializer for robust value serialization
const serializer = require('./serializer');
// ============================================================================
// CONFIGURATION
// ============================================================================
// Configuration from environment
const TRACE_DB = process.env.CODEFLASH_TRACE_DB;
const PROJECT_ROOT = process.env.CODEFLASH_PROJECT_ROOT || process.cwd();
const MAX_FUNCTION_COUNT = parseInt(process.env.CODEFLASH_MAX_FUNCTION_COUNT || '256', 10);
const TRACER_TIMEOUT = process.env.CODEFLASH_TRACER_TIMEOUT
? parseFloat(process.env.CODEFLASH_TRACER_TIMEOUT) * 1000
: null;
// Parse functions to trace from environment
let FUNCTIONS_TO_TRACE = null;
try {
if (process.env.CODEFLASH_FUNCTIONS) {
FUNCTIONS_TO_TRACE = JSON.parse(process.env.CODEFLASH_FUNCTIONS);
}
} catch (e) {
console.error('[codeflash-tracer] Failed to parse CODEFLASH_FUNCTIONS:', e.message);
}
// ============================================================================
// STATE
// ============================================================================
// SQLite database (lazy initialized)
let db = null;
let dbInitialized = false;
// Track function call counts for MAX_FUNCTION_COUNT limit
const functionCallCounts = new Map();
// Track start time for timeout
let tracingStartTime = null;
// Track if tracing is enabled
let tracingEnabled = true;
// Address counter for unique call identification
let lastFrameAddress = 0;
// Prepared statements (cached for performance)
let insertCallStmt = null;
let insertMetadataStmt = null;
// ============================================================================
// DATABASE INITIALIZATION
// ============================================================================
/**
* Initialize the SQLite database for storing traces.
*
* @param {string} dbPath - Path to the SQLite database file
* @returns {boolean} - True if initialization succeeded
*/
function initDatabase(dbPath) {
if (dbInitialized) return true;
if (!dbPath) {
console.error('[codeflash-tracer] No database path provided');
return false;
}
try {
const Database = require('better-sqlite3');
// Ensure directory exists
const dbDir = path.dirname(dbPath);
if (!fs.existsSync(dbDir)) {
fs.mkdirSync(dbDir, { recursive: true });
}
db = new Database(dbPath);
// Create tables matching Python tracer schema
db.exec(`
CREATE TABLE IF NOT EXISTS function_calls (
type TEXT,
function TEXT,
classname TEXT,
filename TEXT,
line_number INTEGER,
last_frame_address INTEGER,
time_ns INTEGER,
args BLOB
);
CREATE TABLE IF NOT EXISTS metadata (
key TEXT PRIMARY KEY,
value TEXT
);
CREATE TABLE IF NOT EXISTS pstats (
filename TEXT,
line_number INTEGER,
function TEXT,
class_name TEXT,
call_count_nonrecursive INTEGER,
num_callers INTEGER,
total_time_ns INTEGER,
cumulative_time_ns INTEGER,
callers BLOB
);
CREATE INDEX IF NOT EXISTS idx_function_calls_function
ON function_calls(function, filename);
CREATE INDEX IF NOT EXISTS idx_function_calls_time
ON function_calls(time_ns);
`);
// Prepare statements for performance
insertCallStmt = db.prepare(`
INSERT INTO function_calls (type, function, classname, filename, line_number, last_frame_address, time_ns, args)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`);
insertMetadataStmt = db.prepare(`
INSERT OR REPLACE INTO metadata (key, value) VALUES (?, ?)
`);
// Record metadata
insertMetadataStmt.run('tracer_version', '1.0.0');
insertMetadataStmt.run('language', 'javascript');
insertMetadataStmt.run('project_root', PROJECT_ROOT);
insertMetadataStmt.run('node_version', process.version);
insertMetadataStmt.run('start_time', new Date().toISOString());
dbInitialized = true;
tracingStartTime = Date.now();
return true;
} catch (e) {
console.error('[codeflash-tracer] Failed to initialize database:', e.message);
return false;
}
}
/**
* Close the database and finalize traces.
*/
function closeDatabase() {
if (db) {
try {
// Record end time
if (insertMetadataStmt) {
insertMetadataStmt.run('end_time', new Date().toISOString());
insertMetadataStmt.run('total_traces', getTotalTraceCount());
}
db.close();
} catch (e) {
console.error('[codeflash-tracer] Error closing database:', e.message);
}
db = null;
dbInitialized = false;
insertCallStmt = null;
insertMetadataStmt = null;
}
}
// ============================================================================
// TIMING UTILITIES
// ============================================================================
/**
* Get high-resolution time in nanoseconds.
*
* @returns {bigint} - Time in nanoseconds
*/
function getTimeNs() {
return process.hrtime.bigint();
}
/**
* Calculate duration in nanoseconds.
*
* @param {bigint} start - Start time in nanoseconds
* @param {bigint} end - End time in nanoseconds
* @returns {number} - Duration in nanoseconds (as Number for SQLite compatibility)
*/
function getDurationNs(start, end) {
return Number(end - start);
}
// ============================================================================
// TRACING UTILITIES
// ============================================================================
/**
* Check if tracing is still enabled (not timed out or disabled).
*
* @returns {boolean} - True if tracing should continue
*/
function isTracingEnabled() {
if (!tracingEnabled) return false;
if (TRACER_TIMEOUT && tracingStartTime) {
const elapsed = Date.now() - tracingStartTime;
if (elapsed >= TRACER_TIMEOUT) {
console.log('[codeflash-tracer] Tracing timeout reached, stopping tracer');
tracingEnabled = false;
return false;
}
}
return true;
}
/**
* Check if a function should be traced based on configuration.
*
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {string|null} className - Class name (for methods)
* @returns {boolean} - True if function should be traced
*/
function shouldTraceFunction(funcName, fileName, className = null) {
if (!isTracingEnabled()) return false;
// Check if we've exceeded the max call count for this function
const key = `${fileName}:${className || ''}:${funcName}`;
const count = functionCallCounts.get(key) || 0;
if (count >= MAX_FUNCTION_COUNT) {
return false;
}
// Check if function is in the filter list
if (FUNCTIONS_TO_TRACE && FUNCTIONS_TO_TRACE.length > 0) {
// Check by function name only, or by full qualified name
const matchesName = FUNCTIONS_TO_TRACE.some(f => {
if (typeof f === 'string') {
return f === funcName || f === `${className}.${funcName}`;
}
// Support object format: { function: 'name', file: 'path', class: 'className' }
if (typeof f === 'object') {
if (f.function && f.function !== funcName) return false;
if (f.file && !fileName.includes(f.file)) return false;
if (f.class && f.class !== className) return false;
return true;
}
return false;
});
if (!matchesName) return false;
}
return true;
}
/**
* Increment the call count for a function.
*
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {string|null} className - Class name (for methods)
*/
function incrementCallCount(funcName, fileName, className = null) {
const key = `${fileName}:${className || ''}:${funcName}`;
const count = functionCallCounts.get(key) || 0;
functionCallCounts.set(key, count + 1);
}
/**
* Get total trace count across all functions.
*
* @returns {number} - Total number of traces
*/
function getTotalTraceCount() {
let total = 0;
for (const count of functionCallCounts.values()) {
total += count;
}
return total;
}
/**
* Safely serialize arguments for storage.
*
* @param {Array} args - Arguments to serialize
* @returns {Buffer} - Serialized arguments
*/
function serializeArgs(args) {
try {
return serializer.serialize(args);
} catch (e) {
console.warn('[codeflash-tracer] Serialization failed:', e.message);
return Buffer.from(JSON.stringify({ __error__: 'SerializationError', message: e.message }));
}
}
// ============================================================================
// CORE TRACING API
// ============================================================================
/**
* Record a function call trace.
*
* @param {string} type - Event type ('call' or 'return')
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name (for methods)
* @param {bigint} timeNs - Timestamp in nanoseconds
* @param {any} argsOrResult - Arguments (for 'call') or return value (for 'return')
*/
function recordTrace(type, funcName, fileName, lineNumber, className, timeNs, argsOrResult) {
if (!dbInitialized) {
if (!initDatabase(TRACE_DB)) {
return;
}
}
try {
const serializedData = serializeArgs(argsOrResult);
const frameAddress = lastFrameAddress++;
insertCallStmt.run(
type,
funcName,
className,
fileName,
lineNumber,
frameAddress,
Number(timeNs),
serializedData
);
} catch (e) {
console.error('[codeflash-tracer] Failed to record trace:', e.message);
}
}
/**
* Wrap a function with tracing instrumentation.
*
* @param {Function} fn - The function to wrap
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name (for methods)
* @returns {Function} - Wrapped function
*/
function wrap(fn, funcName, fileName, lineNumber, className = null) {
// Don't wrap if function shouldn't be traced
if (typeof fn !== 'function') {
return fn;
}
// Check if it's an async function
const isAsync = fn.constructor.name === 'AsyncFunction';
if (isAsync) {
return async function codeflashTracedAsync(...args) {
if (!shouldTraceFunction(funcName, fileName, className)) {
return fn.apply(this, args);
}
incrementCallCount(funcName, fileName, className);
const startTime = getTimeNs();
// Record call
recordTrace('call', funcName, fileName, lineNumber, className, startTime, args);
try {
const result = await fn.apply(this, args);
return result;
} catch (error) {
throw error;
}
};
}
return function codeflashTraced(...args) {
if (!shouldTraceFunction(funcName, fileName, className)) {
return fn.apply(this, args);
}
incrementCallCount(funcName, fileName, className);
const startTime = getTimeNs();
// Record call
recordTrace('call', funcName, fileName, lineNumber, className, startTime, args);
try {
const result = fn.apply(this, args);
// Handle promise returns from non-async functions
if (result instanceof Promise) {
return result.then(
(resolved) => resolved,
(error) => { throw error; }
);
}
return result;
} catch (error) {
throw error;
}
};
}
/**
* Create a wrapper function factory for use with Babel transformation.
*
* @param {string} funcName - Function name
* @param {string} fileName - File path
* @param {number} lineNumber - Line number
* @param {string|null} className - Class name (for methods)
* @returns {Function} - A function that takes the original function and returns a wrapped version
*/
function createWrapper(funcName, fileName, lineNumber, className = null) {
return function(fn) {
return wrap(fn, funcName, fileName, lineNumber, className);
};
}
/**
* Initialize the tracer.
*
* @param {string} dbPath - Path to SQLite database
* @param {string} projectRoot - Project root path
*/
function init(dbPath, projectRoot) {
if (projectRoot) {
process.env.CODEFLASH_PROJECT_ROOT = projectRoot;
}
initDatabase(dbPath || TRACE_DB);
}
/**
* Disable tracing.
*/
function disable() {
tracingEnabled = false;
}
/**
* Enable tracing.
*/
function enable() {
tracingEnabled = true;
}
/**
* Get tracing statistics.
*
* @returns {Object} - Statistics object
*/
function getStats() {
return {
totalTraces: getTotalTraceCount(),
functionCounts: Object.fromEntries(functionCallCounts),
tracingEnabled,
dbInitialized,
};
}
// ============================================================================
// PROCESS EXIT HANDLER
// ============================================================================
// Ensure database is closed on process exit
process.on('exit', () => {
closeDatabase();
});
process.on('SIGINT', () => {
closeDatabase();
process.exit(0);
});
process.on('SIGTERM', () => {
closeDatabase();
process.exit(0);
});
// Also handle uncaught exceptions
process.on('uncaughtException', (err) => {
console.error('[codeflash-tracer] Uncaught exception:', err);
closeDatabase();
process.exit(1);
});
// ============================================================================
// EXPORTS
// ============================================================================
module.exports = {
// Core API
init,
wrap,
createWrapper,
recordTrace,
// Control
disable,
enable,
isTracingEnabled,
// Database
initDatabase,
closeDatabase,
// Utilities
shouldTraceFunction,
getStats,
serializeArgs,
getTimeNs,
getDurationNs,
// Configuration
TRACE_DB,
PROJECT_ROOT,
MAX_FUNCTION_COUNT,
FUNCTIONS_TO_TRACE,
};

114
pyproject.toml
View file

@ -16,45 +16,45 @@ keywords = [
"LLM",
]
dependencies = [
"unidiff>=0.7.4",
"pytest>=7.0.0",
"gitpython>=3.1.31",
"libcst>=1.0.1",
"jedi>=0.19.1",
"unidiff>=0.7.5",
"pytest>=8.4.2",
"gitpython>=3.1.47",
"libcst>=1.8.6",
"jedi>=0.19.2",
# Tree-sitter for multi-language support
"tree-sitter>=0.23.0",
"tree-sitter-javascript>=0.23.0",
"tree-sitter-typescript>=0.23.0",
"tree-sitter-java>=0.23.0",
"tree-sitter-groovy>=0.1.0",
"tree-sitter-kotlin>=1.0.0",
"pytest-timeout>=2.1.0",
"tomlkit>=0.11.7",
"attrs>=23.1.0",
"requests>=2.28.0",
"junitparser>=3.1.0",
"pydantic>=1.10.1",
"humanize>=4.0.0",
"posthog>=3.0.0",
"click>=8.1.0",
"inquirer>=3.0.0",
"sentry-sdk>=1.40.6,<3.0.0",
"tree-sitter>=0.23.2",
"tree-sitter-javascript>=0.23.1",
"tree-sitter-typescript>=0.23.2",
"tree-sitter-java>=0.23.5",
"tree-sitter-groovy>=0.1.2",
"tree-sitter-kotlin>=1.1.0",
"pytest-timeout>=2.4.0",
"tomlkit>=0.14.0",
"attrs>=26.1.0",
"requests>=2.32.5",
"junitparser>=4.0.2",
"pydantic>=2.13.3",
"humanize>=4.13.0",
"posthog>=6.9.3",
"click>=8.1.8",
"inquirer>=3.4.0",
"sentry-sdk>=2.58.0,<3.0.0",
"parameterized>=0.9.0",
"isort>=5.11.0",
"dill>=0.3.8",
"rich>=13.8.1",
"lxml>=5.3.0",
"crosshair-tool>=0.0.78; python_version < '3.15'",
"coverage>=7.6.4",
"line_profiler>=4.2.0",
"platformdirs>=4.3.7",
"pygls>=2.0.0,<3.0.0",
"isort>=6.1.0",
"dill>=0.4.1",
"rich>=15.0.0",
"lxml>=6.1.0",
"crosshair-tool>=0.0.103; python_version < '3.15'",
"coverage>=7.10.7",
"line_profiler>=5.0.2",
"platformdirs>=4.4.0",
"pygls>=2.1.1,<3.0.0",
"codeflash-benchmark",
"filelock>=3.20.3; python_version >= '3.10'",
"filelock>=3.19.1; python_version >= '3.10'",
"filelock<3.20.3; python_version < '3.10'",
"pytest-asyncio>=0.18.0",
"memray>=1.12; sys_platform != 'win32'",
"pytest-memray>=1.7; sys_platform != 'win32'",
"pytest-asyncio>=1.2.0",
"memray>=1.19.3; sys_platform != 'win32'",
"pytest-memray>=1.8.0; sys_platform != 'win32'",
]
[project.urls]
@ -68,40 +68,40 @@ codeflash = "codeflash.main:main"
[dependency-groups]
dev = [
"ipython>=8.12.0",
"mypy>=1.13",
"ruff>=0.7.0",
"mypy>=1.19.1",
"ruff>=0.15.11",
"lxml-stubs>=0.5.1",
"pandas-stubs>=2.2.2.240807, <2.2.3.241009",
"types-Pygments>=2.18.0.20240506",
"types-colorama>=0.4.15.20240311",
"types-decorator>=5.1.8.20240310",
"types-jsonschema>=4.23.0.20240813",
"types-requests>=2.32.0.20241016",
"types-six>=1.16.21.20241009",
"types-cffi>=1.16.0.20240331",
"types-openpyxl>=3.1.5.20241020",
"types-regex>=2024.9.11.20240912",
"types-python-dateutil>=2.9.0.20241003",
"types-gevent>=24.11.0.20241230,<25",
"types-greenlet>=3.1.0.20241221,<4",
"types-pexpect>=4.9.0.20241208,<5",
"types-unidiff>=0.7.0.20240505,<0.8",
"prek>=0.2.25",
"ty>=0.0.14",
"uv>=0.9.29",
"types-Pygments>=2.19.0.20251121",
"types-colorama>=0.4.15.20250801",
"types-decorator>=5.2.0.20251101",
"types-jsonschema>=4.26.0.20260202",
"types-requests>=2.32.4.20260107",
"types-six>=1.17.0.20251009",
"types-cffi>=1.17.0.20250915",
"types-openpyxl>=3.1.5.20250919",
"types-regex>=2026.1.15.20260116",
"types-python-dateutil>=2.9.0.20260124",
"types-gevent>=25.9.0.20251228",
"types-greenlet>=3.3.0.20251206",
"types-pexpect>=4.9.0.20260127",
"types-unidiff>=0.7.0.20240505",
"prek>=0.3.10",
"ty>=0.0.32",
"uv>=0.11.7",
"pytest-cov>=7.1.0",
]
tests = [
"black>=25.9.0",
"black>=25.11.0",
"jax>=0.4.30",
"numpy>=2.0.2",
"pandas>=2.3.3",
"pyarrow>=15.0.0",
"pyarrow>=21.0.0",
"pyrsistent>=0.20.0",
"scipy>=1.13.1",
"torch>=2.8.0",
"xarray>=2024.7.0",
"eval_type_backport",
"eval_type_backport>=0.3.1",
"numba>=0.60.0",
"tensorflow>=2.20.0; python_version >= '3.10'",
]

50
tests/test_languages/test_javascript_instrumentation.py
View file

@ -94,48 +94,6 @@ class TestJavaScriptTracer:
tracer = JavaScriptTracer(output_db)
assert tracer.output_db == output_db
assert tracer.tracer_var == "__codeflash_tracer__"
def test_tracer_generates_init_code(self):
"""Test tracer generates initialization code."""
output_db = Path("/tmp/test_traces.db")
tracer = JavaScriptTracer(output_db)
init_code = tracer._generate_tracer_init()
assert tracer.tracer_var in init_code
assert "serialize" in init_code
assert "wrap" in init_code
assert output_db.as_posix() in init_code
def test_tracer_instruments_simple_function(self):
"""Test tracer can instrument a simple function."""
source = """
function multiply(x, y) {
return x * y;
}
"""
with tempfile.NamedTemporaryFile(suffix=".js", mode="w", delete=False) as f:
f.write(source)
f.flush()
file_path = Path(f.name)
func_info = FunctionInfo(
function_name="multiply", file_path=file_path, starting_line=2, ending_line=4, language="javascript"
)
output_db = Path("/tmp/test_traces.db")
tracer = JavaScriptTracer(output_db)
instrumented = tracer.instrument_source(source, file_path, [func_info])
# Check that tracer initialization is added
assert tracer.tracer_var in instrumented
assert "wrap" in instrumented
# Clean up
file_path.unlink()
def test_tracer_parse_results_empty(self):
"""Test parsing results when file doesn't exist."""
@ -149,7 +107,10 @@ class TestJavaScriptSupportInstrumentation:
"""Integration tests for JavaScript support instrumentation methods."""
def test_javascript_support_instrument_for_behavior(self):
"""Test JavaScriptSupport.instrument_for_behavior method."""
"""Test JavaScriptSupport.instrument_for_behavior returns source unchanged.
JavaScript tracing is now handled at runtime via Babel plugin, not source transformation.
"""
from codeflash.languages import get_language_support
js_support = get_language_support(Language.JAVASCRIPT)
@ -172,8 +133,7 @@ function greet(name) {
output_file = file_path.parent / ".codeflash" / "traces.db"
instrumented = js_support.instrument_for_behavior(source, [func_info], output_file=output_file)
assert "__codeflash_tracer__" in instrumented
assert "wrap" in instrumented
assert instrumented == source
# Clean up
file_path.unlink()

545
tests/test_languages/test_javascript_tracer.py Normal file
View file

@ -0,0 +1,545 @@
"""Tests for JavaScript function tracing.
Tests the JavaScript tracer implementation including:
- Unit tests for Python-side trace parsing and replay test generation
- End-to-end tests for the full tracing pipeline via trace-runner.js
"""
from __future__ import annotations
import json
import os
import shutil
import sqlite3
import subprocess
from pathlib import Path
import pytest
from codeflash.languages.javascript.replay_test import (
JavaScriptFunctionModule,
create_javascript_replay_test,
get_function_alias,
get_traced_functions_from_db,
)
from codeflash.languages.javascript.tracer import JavaScriptTracer
def node_available() -> bool:
return shutil.which("node") is not None
def skip_if_node_not_available() -> None:
if not node_available():
pytest.skip("Node.js not available")
class TestJavaScriptTracerParsing:
@pytest.fixture
def trace_db_with_function_calls(self, tmp_path: Path) -> Path:
db_path = (tmp_path / "trace.sqlite").resolve()
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE function_calls (
type TEXT,
function TEXT,
classname TEXT,
filename TEXT,
line_number INTEGER,
last_frame_address INTEGER,
time_ns INTEGER,
args BLOB
)
""")
cursor.execute("""
CREATE TABLE metadata (
key TEXT PRIMARY KEY,
value TEXT
)
""")
cursor.execute("INSERT INTO metadata (key, value) VALUES ('language', 'javascript')")
test_args = json.dumps([1, 2, 3])
cursor.execute(
"""
INSERT INTO function_calls (type, function, classname, filename, line_number, last_frame_address, time_ns, args)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
("call", "add", None, "/project/src/math.js", 10, 1, 1000000, test_args),
)
cursor.execute(
"""
INSERT INTO function_calls (type, function, classname, filename, line_number, last_frame_address, time_ns, args)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
("call", "multiply", "Calculator", "/project/src/calc.js", 25, 2, 2000000, json.dumps([5, 10])),
)
conn.commit()
conn.close()
return db_path
@pytest.fixture
def trace_db_legacy_schema(self, tmp_path: Path) -> Path:
db_path = (tmp_path / "legacy_trace.sqlite").resolve()
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE traces (
id INTEGER PRIMARY KEY,
call_id INTEGER,
function TEXT,
file TEXT,
args TEXT,
result TEXT,
error TEXT,
runtime_ns TEXT,
timestamp INTEGER
)
""")
cursor.execute(
"""
INSERT INTO traces (call_id, function, file, args, result, error, runtime_ns, timestamp)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(1, "legacyFunc", "/old/path.js", json.dumps(["arg1"]), json.dumps("result"), "null", "5000", 1234567890),
)
conn.commit()
conn.close()
return db_path
def test_parse_results_function_calls_schema(self, trace_db_with_function_calls: Path) -> None:
traces = JavaScriptTracer.parse_results(trace_db_with_function_calls)
assert len(traces) == 2
add_trace = next(t for t in traces if t["function"] == "add")
assert add_trace["type"] == "call"
assert add_trace["filename"] == "/project/src/math.js"
assert add_trace["line_number"] == 10
assert add_trace["args"] == [1, 2, 3]
assert add_trace["classname"] is None
multiply_trace = next(t for t in traces if t["function"] == "multiply")
assert multiply_trace["classname"] == "Calculator"
assert multiply_trace["args"] == [5, 10]
def test_parse_results_legacy_schema(self, trace_db_legacy_schema: Path) -> None:
traces = JavaScriptTracer.parse_results(trace_db_legacy_schema)
assert len(traces) == 1
trace = traces[0]
assert trace["function"] == "legacyFunc"
assert trace["file"] == "/old/path.js"
assert trace["args"] == ["arg1"]
assert trace["result"] == "result"
assert trace["runtime_ns"] == 5000
def test_parse_results_nonexistent_file(self, tmp_path: Path) -> None:
traces = JavaScriptTracer.parse_results((tmp_path / "nonexistent.sqlite").resolve())
assert traces == []
def test_parse_results_json_file(self, tmp_path: Path) -> None:
json_path = (tmp_path / "trace.json").resolve()
trace_data = [{"function": "jsonFunc", "args": [1, 2], "time_ns": 1000}]
json_path.write_text(json.dumps(trace_data), encoding="utf-8")
sqlite_path = (tmp_path / "trace.sqlite").resolve()
traces = JavaScriptTracer.parse_results(sqlite_path)
assert len(traces) == 1
assert traces[0]["function"] == "jsonFunc"
def test_get_traced_functions(self, trace_db_with_function_calls: Path) -> None:
functions = JavaScriptTracer.get_traced_functions(trace_db_with_function_calls)
assert len(functions) == 2
func_names = {f.function_name for f in functions}
assert func_names == {"add", "multiply"}
add_func = next(f for f in functions if f.function_name == "add")
assert add_func.file_name == "/project/src/math.js"
assert add_func.class_name is None
assert add_func.line_number == 10
assert "math" in add_func.module_path
multiply_func = next(f for f in functions if f.function_name == "multiply")
assert multiply_func.class_name == "Calculator"
class TestJavaScriptReplayTestGeneration:
def test_get_function_alias_simple(self) -> None:
alias = get_function_alias("src/utils", "processData")
assert alias == "src_utils_processData"
def test_get_function_alias_with_class(self) -> None:
alias = get_function_alias("src/calculator", "add", "Calculator")
assert alias == "src_calculator_Calculator_add"
def test_get_function_alias_special_chars(self) -> None:
alias = get_function_alias("@scope/package/lib", "func")
assert "_" in alias
assert "func" in alias
def test_create_javascript_replay_test_jest(self, tmp_path: Path) -> None:
functions = [
JavaScriptFunctionModule(
function_name="add", file_name=(tmp_path / "math.js").resolve(), module_name="src/math"
),
JavaScriptFunctionModule(
function_name="multiply",
file_name=(tmp_path / "calc.js").resolve(),
module_name="src/calc",
class_name="Calculator",
),
]
content = create_javascript_replay_test(
trace_file=str((tmp_path / "trace.sqlite").resolve()),
functions=functions,
max_run_count=50,
framework="jest",
)
assert "// Auto-generated replay test by Codeflash" in content
assert "require('codeflash/replay')" in content
assert "describe('Replay: add'" in content
assert "describe('Replay: Calculator.multiply'" in content
assert "test.each" in content
assert "50" in content
def test_create_javascript_replay_test_vitest(self, tmp_path: Path) -> None:
functions = [
JavaScriptFunctionModule(
function_name="process", file_name=(tmp_path / "data.js").resolve(), module_name="src/data"
)
]
content = create_javascript_replay_test(
trace_file=str((tmp_path / "trace.sqlite").resolve()), functions=functions, framework="vitest"
)
assert "import { describe, test } from 'vitest'" in content
assert "describe('Replay: process'" in content
def test_create_javascript_replay_test_skips_constructors(self, tmp_path: Path) -> None:
functions = [
JavaScriptFunctionModule(
function_name="constructor",
file_name=(tmp_path / "class.js").resolve(),
module_name="src/class",
class_name="MyClass",
),
JavaScriptFunctionModule(
function_name="__init__", file_name=(tmp_path / "class.js").resolve(), module_name="src/class"
),
JavaScriptFunctionModule(
function_name="doWork", file_name=(tmp_path / "class.js").resolve(), module_name="src/class"
),
]
content = create_javascript_replay_test(
trace_file=str((tmp_path / "trace.sqlite").resolve()), functions=functions
)
assert "constructor" not in content.lower() or "Replay: constructor" not in content
assert "__init__" not in content or "Replay: __init__" not in content
assert "describe('Replay: doWork'" in content
class TestJavaScriptTracerCreateReplayTest:
@pytest.fixture
def trace_db(self, tmp_path: Path) -> Path:
db_path = (tmp_path / "trace.sqlite").resolve()
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE function_calls (
type TEXT,
function TEXT,
classname TEXT,
filename TEXT,
line_number INTEGER,
last_frame_address INTEGER,
time_ns INTEGER,
args BLOB
)
""")
cursor.execute(
"""
INSERT INTO function_calls (type, function, classname, filename, line_number, last_frame_address, time_ns, args)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
("call", "fibonacci", None, "./src/math.js", 5, 1, 1000, json.dumps([10])),
)
conn.commit()
conn.close()
return db_path
def test_create_replay_test_generates_file(self, trace_db: Path, tmp_path: Path) -> None:
tracer = JavaScriptTracer(trace_db)
output_path = (tmp_path / "tests" / "replay.test.js").resolve()
result = tracer.create_replay_test(trace_db, output_path)
assert result is not None
assert output_path.exists()
content = output_path.read_text(encoding="utf-8")
assert "fibonacci" in content
assert "describe" in content
def test_create_replay_test_empty_db(self, tmp_path: Path) -> None:
db_path = (tmp_path / "empty.sqlite").resolve()
conn = sqlite3.connect(db_path)
conn.execute("""
CREATE TABLE function_calls (
type TEXT, function TEXT, classname TEXT, filename TEXT,
line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB
)
""")
conn.commit()
conn.close()
tracer = JavaScriptTracer(db_path)
result = tracer.create_replay_test(db_path, (tmp_path / "test.js").resolve())
assert result is None
@pytest.mark.skipif(os.name == "nt" or not node_available(), reason="Skipped on Windows or if Node.js not available")
class TestJavaScriptTracerE2E:
@pytest.fixture
def js_project(self, tmp_path: Path) -> Path:
project_dir = (tmp_path / "js_project").resolve()
project_dir.mkdir()
src_dir = project_dir / "src"
src_dir.mkdir()
math_js = src_dir / "math.js"
math_js.write_text(
"""
function add(a, b) {
return a + b;
}
function multiply(a, b) {
return a * b;
}
module.exports = { add, multiply };
""",
encoding="utf-8",
)
main_js = project_dir / "main.js"
main_js.write_text(
"""
const { add, multiply } = require('./src/math.js');
console.log('Running calculations...');
console.log('add(2, 3) =', add(2, 3));
console.log('add(10, 20) =', add(10, 20));
console.log('multiply(4, 5) =', multiply(4, 5));
""",
encoding="utf-8",
)
package_json = project_dir / "package.json"
package_json.write_text(
json.dumps({"name": "test-project", "version": "1.0.0", "main": "main.js"}), encoding="utf-8"
)
return project_dir
@pytest.fixture
def trace_runner_path(self) -> Path:
runner_path = Path(__file__).parent.parent.parent / "packages" / "codeflash" / "runtime" / "trace-runner.js"
if not runner_path.exists():
pytest.skip("trace-runner.js not found")
return runner_path
def test_trace_runner_help(self, trace_runner_path: Path) -> None:
result = subprocess.run(
["node", str(trace_runner_path), "--help"], capture_output=True, text=True, timeout=30, check=False
)
assert "Usage:" in result.stdout or result.returncode == 0
def test_trace_javascript_file(self, js_project: Path, trace_runner_path: Path, tmp_path: Path) -> None:
trace_db = (tmp_path / "trace.sqlite").resolve()
package_json = js_project / "package.json"
package_json.write_text(
json.dumps(
{
"name": "test-project",
"version": "1.0.0",
"main": "main.js",
"dependencies": {
"@babel/core": "^7.24.0",
"@babel/register": "^7.24.0",
"better-sqlite3": "^11.0.0",
},
}
),
encoding="utf-8",
)
npm_install = subprocess.run(
["npm", "install", "--silent"], capture_output=True, text=True, timeout=120, cwd=js_project, check=False
)
if npm_install.returncode != 0:
pytest.skip(f"npm install failed: {npm_install.stderr}")
codeflash_runtime = trace_runner_path.parent
node_modules = js_project / "node_modules"
codeflash_pkg = node_modules / "codeflash"
if codeflash_pkg.exists():
shutil.rmtree(codeflash_pkg)
codeflash_pkg.mkdir()
runtime_dst = codeflash_pkg / "runtime"
shutil.copytree(codeflash_runtime, runtime_dst)
(codeflash_pkg / "package.json").write_text(
json.dumps(
{
"name": "codeflash",
"version": "1.0.0",
"main": "runtime/index.js",
"exports": {
".": {"require": "./runtime/index.js"},
"./tracer": {"require": "./runtime/tracer.js"},
"./replay": {"require": "./runtime/replay.js"},
"./babel-tracer-plugin": {"require": "./runtime/babel-tracer-plugin.js"},
},
}
),
encoding="utf-8",
)
env = os.environ.copy()
env["NODE_PATH"] = str(node_modules.resolve())
result = subprocess.run(
[
"node",
str(trace_runner_path),
"--trace-db",
str(trace_db),
"--project-root",
str(js_project),
"--functions",
json.dumps(["add", "multiply"]),
str(js_project / "main.js"),
],
capture_output=True,
text=True,
timeout=60,
cwd=js_project,
env=env,
check=False,
)
if result.returncode != 0:
print(f"STDOUT: {result.stdout}")
print(f"STDERR: {result.stderr}")
assert "add(2, 3) =" in result.stdout, f"Expected output not found. stderr: {result.stderr}"
assert trace_db.exists(), "Trace database was not created"
conn = sqlite3.connect(trace_db)
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM function_calls WHERE type = 'call'")
trace_count = cursor.fetchone()[0]
conn.close()
assert trace_count >= 2, f"Expected at least 2 traced calls, got {trace_count}"
def test_tracer_runner_python_integration(self, js_project: Path, tmp_path: Path) -> None:
from codeflash.languages.javascript.tracer_runner import (
check_javascript_tracer_available,
detect_test_framework,
)
assert check_javascript_tracer_available() is True
framework = detect_test_framework(js_project, {})
assert framework in ("jest", "vitest")
def test_detect_jest_framework(self, tmp_path: Path) -> None:
from codeflash.languages.javascript.tracer_runner import detect_test_framework
(tmp_path / "jest.config.js").write_text("module.exports = {};", encoding="utf-8")
framework = detect_test_framework(tmp_path, {})
assert framework == "jest"
def test_detect_vitest_framework(self, tmp_path: Path) -> None:
from codeflash.languages.javascript.tracer_runner import detect_test_framework
(tmp_path / "vitest.config.js").write_text("export default {};", encoding="utf-8")
framework = detect_test_framework(tmp_path, {})
assert framework == "vitest"
def test_detect_framework_from_package_json(self, tmp_path: Path) -> None:
from codeflash.languages.javascript.tracer_runner import detect_test_framework
(tmp_path / "package.json").write_text(
json.dumps({"scripts": {"test": "vitest run"}, "devDependencies": {"vitest": "^1.0.0"}}),
encoding="utf-8",
)
framework = detect_test_framework(tmp_path, {})
assert framework == "vitest"
class TestGetTracedFunctionsFromDb:
def test_get_traced_functions_from_db(self, tmp_path: Path) -> None:
db_path = (tmp_path / "trace.sqlite").resolve()
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE function_calls (
type TEXT, function TEXT, classname TEXT, filename TEXT,
line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB
)
""")
cursor.execute(
"INSERT INTO function_calls VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
("call", "testFunc", None, "./src/test.js", 1, 1, 1000, "[]"),
)
conn.commit()
conn.close()
functions = get_traced_functions_from_db(db_path)
assert len(functions) == 1
assert functions[0].function_name == "testFunc"
assert functions[0].module_name == "src/test"
def test_get_traced_functions_nonexistent_file(self, tmp_path: Path) -> None:
functions = get_traced_functions_from_db((tmp_path / "nonexistent.sqlite").resolve())
assert functions == []

2975
uv.lock
View file

File diff suppressed because it is too large Load diff