feat: sync changes to local state before run (#696)

Author: caroljung-cg

src/codegen/cli/commands/config/main.py

@@ -104,7 +104,7 @@ def set_command(key: str, value: str):
         return
 
     cur_value = config.get(key)
-    if cur_value is None or cur_value.lower() != value.lower():
+    if cur_value is None or str(cur_value).lower() != value.lower():
         try:
             config.set(key, value)
         except Exception as e:

src/codegen/cli/commands/run/main.py

@@ -14,7 +14,7 @@
 @requires_init
 @click.argument("label", required=True)
 @click.option("--web", is_flag=True, help="Run the function on the web service instead of locally")
-@click.option("--daemon", is_flag=True, help="Run the function against a running daemon")
+@click.option("--daemon", "-d", is_flag=True, help="Run the function against a running daemon")
 @click.option("--diff-preview", type=int, help="Show a preview of the first N lines of the diff")
 @click.option("--arguments", type=str, help="Arguments as a json string to pass as the function's 'arguments' parameter")
 def run_command(

src/codegen/cli/commands/start/main.py

@@ -132,6 +132,7 @@ def _run_docker_container(repo_config: RepoConfig, port: int, detached: bool) ->
         "REPOSITORY_PATH": container_repo_path,
         "GITHUB_TOKEN": SecretsConfig().github_token,
         "PYTHONUNBUFFERED": "1",  # Ensure Python output is unbuffered
+        "CODEBASE_SYNC_ENABLED": "True",
     }
     envvars_args = [arg for k, v in envvars.items() for arg in ("--env", f"{k}={v}")]
     mount_args = ["-v", f"{repo_config.repo_path}:{container_repo_path}"]

src/codegen/git/repo_operator/repo_operator.py

@@ -458,12 +458,18 @@ def get_diffs(self, ref: str | GitCommit, reverse: bool = True) -> list[Diff]:
         return [diff for diff in self.git_cli.index.diff(ref, R=reverse)]
 
     @stopwatch
-    def stage_and_commit_all_changes(self, message: str, verify: bool = False) -> bool:
+    def stage_and_commit_all_changes(self, message: str, verify: bool = False, exclude_paths: list[str] | None = None) -> bool:
         """TODO: rename to stage_and_commit_changes
         Stage all changes and commit them with the given message.
         Returns True if a commit was made and False otherwise.
         """
         self.git_cli.git.add(A=True)
+        # Unstage the excluded paths
+        for path in exclude_paths or []:
+            try:
+                self.git_cli.git.reset("HEAD", "--", path)
+            except GitCommandError as e:
+                logger.warning(f"Failed to exclude path {path}: {e}")
         return self.commit_changes(message, verify)
 
     def commit_changes(self, message: str, verify: bool = False) -> bool:

src/codegen/runner/sandbox/runner.py

@@ -1,7 +1,6 @@
 import sys
 
-from git import Commit as GitCommit
-
+from codegen.configs.models.codebase import CodebaseConfig
 from codegen.git.repo_operator.repo_operator import RepoOperator
 from codegen.git.schemas.enums import SetupOption
 from codegen.git.schemas.repo_config import RepoConfig
@@ -21,7 +20,6 @@ class SandboxRunner:
 
     # =====[ __init__ instance attributes ]=====
     repo: RepoConfig
-    commit: GitCommit
     op: RepoOperator | None
 
     # =====[ computed instance attributes ]=====
@@ -31,20 +29,19 @@ class SandboxRunner:
     def __init__(self, repo_config: RepoConfig, op: RepoOperator | None = None) -> None:
         self.repo = repo_config
         self.op = op or RepoOperator(repo_config=self.repo, setup_option=SetupOption.PULL_OR_CLONE, bot_commit=True)
-        self.commit = self.op.git_cli.head.commit
 
-    async def warmup(self) -> None:
+    async def warmup(self, codebase_config: CodebaseConfig | None = None) -> None:
         """Warms up this runner by cloning the repo and parsing the graph."""
         logger.info(f"===== Warming runner for {self.repo.full_name or self.repo.name} =====")
         sys.setrecursionlimit(10000)  # for graph parsing
 
-        self.codebase = await self._build_graph()
+        self.codebase = await self._build_graph(codebase_config)
         self.executor = SandboxExecutor(self.codebase)
 
-    async def _build_graph(self) -> Codebase:
+    async def _build_graph(self, codebase_config: CodebaseConfig | None = None) -> Codebase:
         logger.info("> Building graph...")
         projects = [ProjectConfig(programming_language=self.repo.language, repo_operator=self.op, base_path=self.repo.base_path, subdirectories=self.repo.subdirectories)]
-        return Codebase(projects=projects)
+        return Codebase(projects=projects, config=codebase_config)
 
     async def get_diff(self, request: GetDiffRequest) -> GetDiffResponse:
         custom_scope = {"context": request.codemod.codemod_context} if request.codemod.codemod_context else {}

src/codegen/runner/sandbox/server.py

@@ -45,7 +45,7 @@ async def lifespan(server: FastAPI):
         runner = SandboxRunner(repo_config=repo_config)
         server_info.warmup_state = WarmupState.PENDING
         await runner.warmup()
-        server_info.synced_commit = runner.commit.hexsha
+        server_info.synced_commit = runner.op.git_cli.head.commit.hexsha
         server_info.warmup_state = WarmupState.COMPLETED
     except Exception:
         logger.exception("Failed to build graph during warmup")

src/codegen/runner/servers/local_daemon.py

@@ -1,9 +1,9 @@
 import logging
-import os
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
+from codegen.configs.models.codebase import DefaultCodebaseConfig
 from codegen.git.configs.constants import CODEGEN_BOT_EMAIL, CODEGEN_BOT_NAME
 from codegen.git.repo_operator.repo_operator import RepoOperator
 from codegen.git.schemas.enums import SetupOption
@@ -47,11 +47,12 @@ async def lifespan(server: FastAPI):
         runner.op.git_cli.git.config("user.email", CODEGEN_BOT_EMAIL)
         runner.op.git_cli.git.config("user.name", CODEGEN_BOT_NAME)
 
-        # Parse the codebase
+        # Parse the codebase with sync enabled
         logger.info(f"Starting up fastapi server for repo_name={repo_config.name}")
         server_info.warmup_state = WarmupState.PENDING
-        await runner.warmup()
-        server_info.synced_commit = runner.commit.hexsha
+        codebase_config = DefaultCodebaseConfig.model_copy(update={"sync_enabled": True})
+        await runner.warmup(codebase_config=codebase_config)
+        server_info.synced_commit = runner.op.head_commit.hexsha
         server_info.warmup_state = WarmupState.COMPLETED
 
     except Exception:
@@ -73,27 +74,24 @@ def health() -> ServerInfo:
 
 @app.post(RUN_FUNCTION_ENDPOINT)
 async def run(request: RunFunctionRequest) -> CodemodRunResult:
-    # TODO: Sync graph to whatever changes are in the repo currently
-
-    # Run the request
+    _save_uncommitted_changes_and_sync()
     diff_req = GetDiffRequest(codemod=Codemod(user_code=request.codemod_source))
     diff_response = await runner.get_diff(request=diff_req)
     if request.commit:
-        if _should_skip_commit(request.function_name):
-            logger.info(f"Skipping commit because only changes to {request.function_name} were made")
-        elif commit_sha := runner.codebase.git_commit(f"[Codegen] {request.function_name}"):
+        if commit_sha := runner.codebase.git_commit(f"[Codegen] {request.function_name}", exclude_paths=[".codegen/*"]):
             logger.info(f"Committed changes to {commit_sha.hexsha}")
     return diff_response.result
 
 
-def _should_skip_commit(function_name: str) -> bool:
-    changed_files = runner.op.get_modified_files(runner.commit)
-    if len(changed_files) != 1:
-        return False
+def _save_uncommitted_changes_and_sync() -> None:
+    if commit := runner.codebase.git_commit("[Codegen] Save uncommitted changes", exclude_paths=[".codegen/*"]):
+        logger.info(f"Saved uncommitted changes to {commit.hexsha}")
 
-    file_path = changed_files[0]
-    if not file_path.startswith(".codegen/codemods/"):
-        return False
+    cur_commit = runner.op.head_commit
+    if cur_commit != runner.codebase.ctx.synced_commit:
+        logger.info(f"Syncing codebase to head commit: {cur_commit.hexsha}")
+        runner.codebase.sync_to_commit(target_commit=cur_commit)
+    else:
+        logger.info("Codebase is already synced to head commit")
 
-    changed_file_name = os.path.splitext(os.path.basename(file_path))[0]
-    return changed_file_name == function_name.replace("-", "_")
+    server_info.synced_commit = cur_commit.hexsha

src/codegen/sdk/core/codebase.py

@@ -734,7 +734,7 @@ def get_relative_path(self, from_file: str, to_file: str) -> str:
     # State/Git
     ####################################################################################################################
 
-    def git_commit(self, message: str, *, verify: bool = False) -> GitCommit | None:
+    def git_commit(self, message: str, *, verify: bool = False, exclude_paths: list[str] | None = None) -> GitCommit | None:
         """Stages + commits all changes to the codebase and git.
 
         Args:
@@ -745,7 +745,7 @@ def git_commit(self, message: str, *, verify: bool = False) -> GitCommit | None:
             GitCommit | None: The commit object if changes were committed, None otherwise.
         """
         self.ctx.commit_transactions(sync_graph=False)
-        if self._op.stage_and_commit_all_changes(message, verify):
+        if self._op.stage_and_commit_all_changes(message, verify, exclude_paths):
             logger.info(f"Commited repository to {self._op.head_commit} on {self._op.get_active_branch_or_commit()}")
             return self._op.head_commit
         else:

src/codegen/shared/logging/get_logger.py

@@ -35,7 +35,7 @@ def get_logger(name: str, level: int = logging.INFO) -> logging.Logger:
     handler.setFormatter(formatter)
     logger.addHandler(handler)
     # Ensure the logger propagates to the root logger
-    logger.propagate = False
+    logger.propagate = True
     # Set the level on the logger itself
     logger.setLevel(level)
     return logger

tests/unit/codegen/runner/sandbox/test_runner.py

@@ -40,7 +40,7 @@ async def test_sandbox_runner_warmup_starts_with_default_branch(mock_executor, r
     # assert len(runner.codebase._op.git_cli.branches) == 1  TODO: fix GHA creating master and main branch
     assert not runner.codebase._op.git_cli.head.is_detached
     assert runner.codebase._op.git_cli.active_branch.name == runner.codebase.default_branch
-    assert runner.codebase._op.git_cli.head.commit == runner.commit
+    assert runner.codebase._op.git_cli.head.commit == runner.op.head_commit
 
 
 @pytest.mark.asyncio