Implement agent-driven profiling

camelid · camelid · commit b58f1c2f8f82 · 2025-11-06T00:18:41.000-05:00
diff --git a/accelerant/agent.py b/accelerant/agent.py
@@ -34,7 +34,7 @@ def run_agent(
     ag_tools: list[Tool] = [
         tools.edit_code,
         tools.check_codebase_for_errors,
-        tools.get_profiler_data,
+        tools.run_perf_profiler,
         tools.get_info,
         tools.get_references,
         tools.get_surrounding_code,
@@ -47,22 +47,17 @@ def run_agent(
         tools=ag_tools,
     )
 
-    with project.new_fs_sandbox() as fs:
-        ag_context = AgentContext(
-            project=project,
-            active_fs=fs,
-            initial_perf_data_path=ag_input["perf_data_path"],
-        )
-        prompt = user_prompt(
-            lang=project.lang(), hotspot_lines=ag_input["hotspot_lines"] or []
-        )
-        result = Runner.run_sync(
-            agent,
-            prompt,
-            context=ag_context,
-            max_turns=100,
-        ).final_output
-        assert result is not None
-        final_message = str(result)
-        fs.persist_all()
-        return AgentResult(final_message=final_message)
+    ag_context = AgentContext(project=project)
+    prompt = user_prompt(
+        lang=project.lang(), hotspot_lines=ag_input["hotspot_lines"] or []
+    )
+    result = Runner.run_sync(
+        agent,
+        prompt,
+        context=ag_context,
+        max_turns=100,
+    ).final_output
+    assert result is not None
+    final_message = str(result)
+    project.fs_sandbox().persist_all()
+    return AgentResult(final_message=final_message)
diff --git a/accelerant/fs_sandbox.py b/accelerant/fs_sandbox.py
@@ -1,15 +1,24 @@
+from dataclasses import dataclass
+import hashlib
 from pathlib import Path
 from typing import Literal
 
 
+@dataclass(frozen=True)
+class FsVersion:
+    hash: str
+
+
 class FsSandbox:
     base_dir: Path
     old_versions: dict[Path, str]
+    cur_hashes: dict[Path, str]
     status: Literal["fresh"] | Literal["entered"] | Literal["done"] = "fresh"
 
     def __init__(self, base_dir: Path) -> None:
         self.base_dir = base_dir
         self.old_versions = {}
+        self.cur_hashes = {}
 
     def __enter__(self) -> "FsSandbox":
         self.status = "entered"
@@ -38,6 +47,10 @@ def write_file(self, relpath: Path, new_text: str) -> None:
                 self.old_versions[relpath] = f.read()
         with open(abspath, "w") as f:
             f.write(new_text)
+        self.cur_hashes[relpath] = hashlib.sha256(new_text.encode()).hexdigest()
+        if self.old_versions[relpath] == new_text:
+            del self.old_versions[relpath]
+            del self.cur_hashes[relpath]
 
     def persist(self, relpath: Path) -> None:
         assert self.status == "entered"
@@ -47,3 +60,12 @@ def persist(self, relpath: Path) -> None:
     def persist_all(self) -> None:
         assert self.status == "entered"
         self.old_versions = {}
+
+    def version(self) -> FsVersion:
+        hasher = hashlib.sha256()
+        for relpath in sorted(self.cur_hashes.keys()):
+            hasher.update(relpath.as_posix().encode())
+            hasher.update(b"\0")
+            hasher.update(self.cur_hashes[relpath].encode())
+            hasher.update(b"\0")
+        return FsVersion(hash=hasher.hexdigest()[:8])
diff --git a/accelerant/project.py b/accelerant/project.py
@@ -1,38 +1,117 @@
+import os
+import shutil
+import subprocess
+import time
 from multilspy.lsp_protocol_handler import lsp_types
 
 
 from pathlib import Path
 from typing import List, Optional
 
-from accelerant.fs_sandbox import FsSandbox
+from accelerant.fs_sandbox import FsSandbox, FsVersion
 from accelerant.lsp import LSP
 from accelerant.perf import PerfData
 
 
 class Project:
     _root: Path
+    # FIXME: this should probably not be here to allow for multiple targets
+    _target_binary: Path
     _lang: str
+    _fs: FsSandbox
     _lsp: Optional[LSP]
-    _perf_data: dict[Path, PerfData]
+    _perf_per_version: dict[FsVersion, Path]
+    _perf_data_map: dict[Path, PerfData]
 
-    def __init__(self, root, lang):
+    def __init__(self, root: Path, target_binary: Path, lang: str) -> None:
         self._root = root
+        self._target_binary = target_binary
         self._lang = lang
+        self._fs = FsSandbox(root)
         self._lsp = None
-        self._perf_data = {}
+        self._perf_per_version = {}
+        self._perf_data_map = {}
+
+    def target_binary(self) -> Path:
+        return self._target_binary
 
     def lsp(self) -> LSP:
         if self._lsp is None:
             self._lsp = LSP(self._root, self._lang)
         return self._lsp
 
-    def perf_data(self, perf_data_path: Path) -> PerfData:
-        if perf_data_path not in self._perf_data:
-            self._perf_data[perf_data_path] = PerfData(perf_data_path, self._root)
-        return self._perf_data[perf_data_path]
+    def perf_data(self, version: Optional[FsVersion] = None) -> Optional[PerfData]:
+        if version is None:
+            version = self.fs_sandbox().version()
+        if version not in self._perf_per_version:
+            return None
+        perf_data_path = self._perf_per_version[version]
+        if perf_data_path not in self._perf_data_map:
+            self._perf_data_map[perf_data_path] = PerfData(perf_data_path, self._root)
+        return self._perf_data_map[perf_data_path]
+
+    def add_perf_data(self, version: FsVersion, perf_data_path: Path) -> None:
+        self._perf_per_version[version] = perf_data_path
+
+    def build_for_profiling(self) -> None:
+        if self._lang != "rust":
+            raise NotImplementedError(
+                f"Build for profiling not implemented for language: {self._lang}"
+            )
+
+        cargo_path = shutil.which("cargo")
+        assert cargo_path is not None, "cargo not found in PATH"
+
+        path_env_var = os.environ.get("PATH")
+        assert path_env_var is not None, "PATH environment variable is not set"
+
+        subprocess.run(
+            [
+                cargo_path,
+                "build",
+                "--config",
+                "profile.release.debug=true",
+                "--release",
+                "--all-targets",
+            ],
+            check=True,
+            cwd=str(self._root),
+            env={"PATH": path_env_var},
+        )
+
+    def run_profiler(
+        self,
+    ) -> None:
+        if self._lang != "rust":
+            raise NotImplementedError(
+                f"Profiler run not implemented for language: {self._lang}"
+            )
+
+        perf_data_path = self._root / f"perf{time.time_ns()}.data"
+
+        path_env_var = os.environ.get("PATH")
+        assert path_env_var is not None, "PATH environment variable is not set"
+
+        subprocess.run(
+            [
+                "perf",
+                "record",
+                "-F99",
+                "--call-graph",
+                "dwarf",
+                "-o",
+                str(perf_data_path),
+                str(self._target_binary),
+            ],
+            check=True,
+            cwd=str(self._root),
+            env={"PATH": path_env_var},
+        )
+        version = self.fs_sandbox().version()
+        self.add_perf_data(version, perf_data_path)
 
-    def new_fs_sandbox(self) -> FsSandbox:
-        return FsSandbox(self._root)
+    def fs_sandbox(self) -> FsSandbox:
+        return self._fs
 
     def get_line(self, filename: str, line: int) -> str:
         assert line >= 0
diff --git a/accelerant/prompts.py b/accelerant/prompts.py
@@ -11,7 +11,8 @@
         "Use these tools effectively to gather the necessary information before making optimization suggestions.\n"
         "Always provide clear, concise, and actionable suggestions that can be directly implemented in the codebase.\n"
         "Take full control and apply edits to the code without needing approval from the user.\n"
-        "Check the codebase for errors after making edits to ensure correctness."
+        "Check the codebase for errors after making edits to ensure correctness.\n"
+        "Rerun the performance profiler after making edits to measure improvements."
     ),
 }
 
diff --git a/accelerant/tools.py b/accelerant/tools.py
@@ -1,14 +1,13 @@
 from dataclasses import dataclass
 from itertools import islice
-from pathlib import Path
+import shutil
 import subprocess
 from typing import Any, Optional
 from agents import RunContextWrapper, function_tool
 from llm_utils import number_group_of_lines
 from perfparser import LineLoc
 
 from accelerant.chat_interface import CodeSuggestion
-from accelerant.fs_sandbox import FsSandbox
 from accelerant.lsp import TOP_LEVEL_SYMBOL_KINDS, uri_to_relpath
 from accelerant.patch import apply_simultaneous_suggestions
 from accelerant.util import find_symbol, truncate_for_llm
@@ -18,8 +17,6 @@
 @dataclass
 class AgentContext:
     project: Project
-    active_fs: FsSandbox
-    initial_perf_data_path: Optional[Path]
 
 
 @function_tool
@@ -32,7 +29,9 @@ def edit_code(
     Args:
         suggs: A list of code suggestions that should be applied.
     """
-    apply_simultaneous_suggestions(ctx.context.project, ctx.context.active_fs, suggs)
+    apply_simultaneous_suggestions(
+        ctx.context.project, ctx.context.project.fs_sandbox(), suggs
+    )
 
 
 @function_tool
@@ -43,26 +42,34 @@ def check_codebase_for_errors(
     assert ctx.context.project._lang == "rust", (
         "Only Rust is supported for code checking"
     )
+
+    cargo_path = shutil.which("cargo")
+    assert cargo_path is not None, "cargo not found in PATH"
     try:
         subprocess.run(
-            ["cargo", "check", "--all"], check=True, cwd=str(ctx.context.project._root)
+            [cargo_path, "check", "--all-targets"],
+            check=True,
+            cwd=str(ctx.context.project._root),
         )
     except subprocess.CalledProcessError as e:
         return f"ERROR: Codebase has errors:\n\n{e}"
     return "OK: Codebase has no errors!"
 
 
 @function_tool
-def get_profiler_data(
+def run_perf_profiler(
     ctx: RunContextWrapper[AgentContext],
 ) -> list[dict[str, Any]]:
-    """Get a summary of the objective performance data gathered by a profiler."""
+    """Run a performance profiler on the target binary and return the top hotspots."""
     try:
-        perf_data_path = ctx.context.initial_perf_data_path
-        if perf_data_path is None:
-            raise ValueError("No initial performance data path provided")
         project = ctx.context.project
-        perf_data = project.perf_data(perf_data_path)
+        version = project.fs_sandbox().version()
+        perf_data = project.perf_data(version)
+        if perf_data is None:
+            project.build_for_profiling()
+            project.run_profiler()
+            perf_data = project.perf_data(version)
+        assert perf_data is not None, "perf data should be available after profiling"
         perf_tabulated = perf_data.tabulate()
         NUM_HOTSPOTS = 5
 
@@ -76,18 +83,20 @@ def get_parent_region(loc: LineLoc) -> Optional[str]:
                 return None
             return parent_sym["name"]
 
-        hotspots = islice(
-            map(
-                lambda x: {
-                    "parent_region": get_parent_region(x[0]) or "<unknown>",
-                    "loc": x[0],
-                    "pct_time": x[1] * 100,
-                },
-                filter(lambda x: x[0].line > 0, perf_tabulated),
-            ),
-            NUM_HOTSPOTS,
+        hotspots = list(
+            islice(
+                map(
+                    lambda x: {
+                        "parent_region": get_parent_region(x[0]) or "<unknown>",
+                        "loc": x[0],
+                        "pct_time": x[1] * 100,
+                    },
+                    filter(lambda x: x[0].line > 0, perf_tabulated),
+                ),
+                NUM_HOTSPOTS,
+            )
         )
-        return list(hotspots)
+        return hotspots
     except Exception as e:
         print("ERROR", e)
         raise e
@@ -240,8 +249,8 @@ def get_surrounding_code(
             filename, line - 1, TOP_LEVEL_SYMBOL_KINDS
         ),
     )
-    # FIXME: avoid crashing
-    assert parent_sym is not None
+    if parent_sym is None:
+        raise ValueError(f"no surrounding top-level symbol found at {filename}:{line}")
     sline = parent_sym["range"]["start"]["line"] + 1
     lines = project.get_range(filename, parent_sym["range"])
     return {
diff --git a/accelerant_server.py b/accelerant_server.py

Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,8 @@`
`11`	`11`	`"Use these tools effectively to gather the necessary information before making optimization suggestions.\n"`
`12`	`12`	`"Always provide clear, concise, and actionable suggestions that can be directly implemented in the codebase.\n"`
`13`	`13`	`"Take full control and apply edits to the code without needing approval from the user.\n"`
`14`		`- "Check the codebase for errors after making edits to ensure correctness."`
	`14`	`+ "Check the codebase for errors after making edits to ensure correctness.\n"`
	`15`	`+ "Rerun the performance profiler after making edits to measure improvements."`
`15`	`16`	`),`
`16`	`17`	`}`
`17`	`18`