feat: add --daemon to codegen run

Author: caroljung-cg

pyproject.toml

@@ -74,6 +74,7 @@ dependencies = [
   "httpx>=0.28.1",
   "docker>=6.1.3",
   "urllib3>=2.0.0",
+  "colorlog>=6.9.0",
 ]
 
 license = { text = "Apache-2.0" }

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

@@ -14,16 +14,22 @@
 @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("--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(
     session: CodegenSession,
     label: str,
     web: bool = False,
+    daemon: bool = False,
     diff_preview: int | None = None,
     arguments: str | None = None,
 ):
     """Run a codegen function by its label."""
+    if web and daemon:
+        msg = "Cannot enable run on both the web and daemon"
+        raise ValueError(msg)
+
     # Ensure venv is initialized
     venv = VenvManager(session.codegen_dir)
     if not venv.is_initialized():
@@ -54,6 +60,10 @@ def run_command(
         from codegen.cli.commands.run.run_cloud import run_cloud
 
         run_cloud(session, codemod, diff_preview=diff_preview)
+    elif daemon:
+        from codegen.cli.commands.run.run_daemon import run_daemon
+
+        run_daemon(session, codemod, diff_preview=diff_preview)
     else:
         from codegen.cli.commands.run.run_local import run_local
 

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

@@ -0,0 +1,87 @@
+import rich
+import rich_click as click
+from rich.panel import Panel
+
+from codegen.cli.auth.session import CodegenSession
+from codegen.cli.commands.start.docker_container import DockerContainer
+from codegen.cli.errors import ServerError
+from codegen.cli.rich.codeblocks import format_command
+from codegen.cli.rich.spinners import create_spinner
+from codegen.runner.clients.docker_client import DockerClient
+from codegen.runner.enums.warmup_state import WarmupState
+
+
+def run_daemon(session: CodegenSession, function, diff_preview: int | None = None):
+    """Run a function on the cloud service.
+
+    Args:
+        session: The current codegen session
+        function: The function to run
+        diff_preview: Number of lines of diff to preview (None for all)
+    """
+    with create_spinner(f"Running {function.name}...") as status:
+        try:
+            client = _get_docker_client(session)
+            run_output = client.run_function(function, commit=not diff_preview)
+            rich.print(f"✅ Ran {function.name} successfully")
+
+            if run_output.logs:
+                rich.print("")
+                panel = Panel(run_output.logs, title="[bold]Logs[/bold]", border_style="blue", padding=(1, 2), expand=False)
+                rich.print(panel)
+
+            if run_output.error:
+                rich.print("")
+                panel = Panel(run_output.error, title="[bold]Error[/bold]", border_style="red", padding=(1, 2), expand=False)
+                rich.print(panel)
+
+            if run_output.observation:
+                # Only show diff preview if requested
+                if diff_preview:
+                    rich.print("")  # Add some spacing
+
+                    # Split and limit diff to requested number of lines
+                    diff_lines = run_output.observation.splitlines()
+                    truncated = len(diff_lines) > diff_preview
+                    limited_diff = "\n".join(diff_lines[:diff_preview])
+
+                    if truncated:
+                        limited_diff += "\n\n...\n\n[yellow]diff truncated to {diff_preview} lines, view the full change set on your local file system after using run with `--apply-local`[/yellow]"
+
+                    panel = Panel(limited_diff, title="[bold]Diff Preview[/bold]", border_style="blue", padding=(1, 2), expand=False)
+                    rich.print(panel)
+            else:
+                rich.print("")
+                rich.print("[yellow] No changes were produced by this codemod[/yellow]")
+
+            if diff_preview:
+                rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")
+                rich.print("[yellow]→ Don't forget to commit your changes:[/yellow]")
+                rich.print(format_command("git add ."))
+                rich.print(format_command("git commit -m 'Applied codemod changes'"))
+
+        except ServerError as e:
+            status.stop()
+            raise click.ClickException(str(e))
+
+
+def _get_docker_client(session: CodegenSession) -> DockerClient:
+    repo_name = session.config.repository.name
+    if (container := DockerContainer.get(repo_name)) is None:
+        msg = f"Codegen runner does not exist for {repo_name}. Please run 'codegen start' from {session.config.repository.path}."
+        raise click.ClickException(msg)
+
+    if not container.is_running():
+        msg = f"Codegen runner for {repo_name} is not running. Please run 'codegen start' from {session.config.repository.path}."
+        raise click.ClickException(msg)
+
+    client = DockerClient(container)
+    if not client.is_running():
+        msg = "Codebase server is not running. Please stop the container and restart."
+        raise click.ClickException(msg)
+
+    if client.server_info().warmup_state != WarmupState.COMPLETED:
+        msg = "Runner has not finished parsing the codebase. Please wait a moment and try again."
+        raise click.ClickException(msg)
+
+    return client

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

@@ -1,31 +1,74 @@
+from functools import cached_property
+
 import docker
+from docker import DockerClient
+from docker.errors import APIError, NotFound
+from docker.models.containers import Container
 
 
 class DockerContainer:
-    _client: docker.DockerClient
-    host: str | None
-    port: int | None
-    name: str
+    _client: DockerClient
+    _container: Container | None
 
-    def __init__(self, client: docker.DockerClient, name: str, port: int | None = None, host: str | None = None):
+    def __init__(self, client: DockerClient, container: Container) -> None:
         self._client = client
-        self.host = host
-        self.port = port
-        self.name = name
+        self._container = container
+
+    @classmethod
+    def get(cls, name: str) -> "DockerContainer | None":
+        try:
+            client = docker.from_env()
+            container = client.containers.get(name)
+            return cls(client=client, container=container)
+        except NotFound:
+            return None
+
+    @cached_property
+    def name(self) -> str:
+        return self._container.name
+
+    @cached_property
+    def host(self) -> str | None:
+        if not self.is_running():
+            return None
+
+        host_config = next(iter(self._container.ports.values()))[0]
+        return host_config["HostIp"]
+
+    @cached_property
+    def port(self) -> int | None:
+        if not self.is_running():
+            return None
+
+        host_config = next(iter(self._container.ports.values()))[0]
+        return host_config["HostPort"]
 
     def is_running(self) -> bool:
         try:
-            container = self._client.containers.get(self.name)
-            return container.status == "running"
-        except docker.errors.NotFound:
+            return self._container.status == "running"
+        except NotFound:
             return False
 
     def start(self) -> bool:
         try:
-            container = self._client.containers.get(self.name)
-            container.start()
+            self._container.start()
+            return True
+        except (NotFound, APIError):
+            return False
+
+    def stop(self) -> bool:
+        try:
+            self._container.stop()
+            return True
+        except (NotFound, APIError):
+            return False
+
+    def remove(self) -> bool:
+        try:
+            self.stop()
+            self._container.remove()
             return True
-        except (docker.errors.NotFound, docker.errors.APIError):
+        except (NotFound, APIError):
             return False
 
     def __str__(self) -> str:

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

@@ -1,4 +1,5 @@
 import docker
+from docker.errors import NotFound
 
 from codegen.cli.commands.start.docker_container import DockerContainer
 
@@ -15,19 +16,14 @@ def __init__(self, containers: list[DockerContainer]):
     def load(cls) -> "DockerFleet":
         try:
             client = docker.from_env()
-            containers = client.containers.list(all=True, filters={"ancestor": CODEGEN_RUNNER_IMAGE})
-            codegen_containers = []
-            for container in containers:
+            filters = {"ancestor": CODEGEN_RUNNER_IMAGE}
+            containers = []
+            for container in client.containers.list(all=True, filters=filters):
                 if container.attrs["Config"]["Image"] == CODEGEN_RUNNER_IMAGE:
-                    if container.status == "running":
-                        host_config = next(iter(container.ports.values()))[0]
-                        codegen_container = DockerContainer(client=client, host=host_config["HostIp"], port=host_config["HostPort"], name=container.name)
-                    else:
-                        codegen_container = DockerContainer(client=client, name=container.name)
-                    codegen_containers.append(codegen_container)
-
-            return cls(containers=codegen_containers)
-        except docker.errors.NotFound:
+                    containers.append(DockerContainer(client=client, container=container))
+
+            return cls(containers=containers)
+        except NotFound:
             return cls(containers=[])
 
     @property

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

@@ -8,7 +8,7 @@
 from rich.panel import Panel
 
 from codegen.cli.commands.start.docker_container import DockerContainer
-from codegen.cli.commands.start.docker_fleet import CODEGEN_RUNNER_IMAGE, DockerFleet
+from codegen.cli.commands.start.docker_fleet import CODEGEN_RUNNER_IMAGE
 from codegen.configs.models.secrets import SecretsConfig
 from codegen.git.repo_operator.local_git_repo import LocalGitRepo
 from codegen.git.schemas.repo_config import RepoConfig
@@ -20,13 +20,19 @@
 @click.command(name="start")
 @click.option("--platform", "-t", type=click.Choice(["linux/amd64", "linux/arm64", "linux/amd64,linux/arm64"]), default="linux/amd64,linux/arm64", help="Target platform(s) for the Docker image")
 @click.option("--port", "-p", type=int, default=None, help="Port to run the server on")
-def start_command(port: int | None, platform: str):
+@click.option("--detached", "-d", is_flag=True, help="Run the server in detached mode")
+@click.option("--skip-build", is_flag=True, help="Skip building the Docker image")
+@click.option("--force", "-f", is_flag=True, help="Force start the server even if it is already running")
+def start_command(port: int | None, platform: str, detached: bool = False, skip_build: bool = False, force: bool = False) -> None:
     """Starts a local codegen server"""
     repo_path = Path.cwd().resolve()
     repo_config = RepoConfig.from_repo_path(str(repo_path))
-    fleet = DockerFleet.load()
-    if (container := fleet.get(repo_config.name)) is not None:
-        return _handle_existing_container(repo_config, container)
+    if (container := DockerContainer.get(repo_config.name)) is not None:
+        if force:
+            rich.print(f"[yellow]Removing existing runner {repo_config.name} to force restart[/yellow]")
+            container.remove()
+        else:
+            return _handle_existing_container(repo_config, container, force)
 
     codegen_version = version("codegen")
     rich.print(f"[bold green]Codegen version:[/bold green] {codegen_version}")
@@ -35,10 +41,11 @@ def start_command(port: int | None, platform: str):
         port = get_free_port()
 
     try:
-        rich.print("[bold blue]Building Docker image...[/bold blue]")
-        _build_docker_image(codegen_root, platform)
+        if not skip_build:
+            rich.print("[bold blue]Building Docker image...[/bold blue]")
+            _build_docker_image(codegen_root, platform)
         rich.print("[bold blue]Starting Docker container...[/bold blue]")
-        _run_docker_container(repo_config, port)
+        _run_docker_container(repo_config, port, detached)
         rich.print(Panel(f"[green]Server started successfully![/green]\nAccess the server at: [bold]http://{_default_host}:{port}[/bold]", box=ROUNDED, title="Codegen Server"))
         # TODO: memory snapshot here
     except subprocess.CalledProcessError as e:
@@ -49,7 +56,7 @@ def start_command(port: int | None, platform: str):
         raise click.Abort()
 
 
-def _handle_existing_container(repo_config: RepoConfig, container: DockerContainer) -> None:
+def _handle_existing_container(repo_config: RepoConfig, container: DockerContainer, force: bool) -> None:
     if container.is_running():
         rich.print(
             Panel(
@@ -86,19 +93,26 @@ def _build_docker_image(codegen_root: Path, platform: str) -> None:
     subprocess.run(build_cmd, check=True)
 
 
-def _run_docker_container(repo_config: RepoConfig, port: int) -> None:
+def _run_docker_container(repo_config: RepoConfig, port: int, detached: bool) -> None:
     container_repo_path = f"/app/git/{repo_config.name}"
     name_args = ["--name", f"{repo_config.name}"]
     envvars = {
         "REPOSITORY_LANGUAGE": repo_config.language.value,
         "REPOSITORY_OWNER": LocalGitRepo(repo_config.repo_path).owner,
         "REPOSITORY_PATH": container_repo_path,
         "GITHUB_TOKEN": SecretsConfig().github_token,
+        "PYTHONUNBUFFERED": "1",  # Ensure Python output is unbuffered
     }
     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}"]
-    entry_point = f"uv run --frozen uvicorn codegen.runner.sandbox.server:app --host {_default_host} --port {port}"
-    run_cmd = ["docker", "run", "-d", "-p", f"{port}:{port}", *name_args, *mount_args, *envvars_args, CODEGEN_RUNNER_IMAGE, entry_point]
+    entry_point = f"uv run --frozen uvicorn codegen.runner.servers.local_daemon:app --host {_default_host} --port {port}"
+    port_args = ["-p", f"{port}:{port}"]
+    detached_args = ["-d"] if detached else []
+    run_cmd = ["docker", "run", *detached_args, *port_args, *name_args, *mount_args, *envvars_args, CODEGEN_RUNNER_IMAGE, entry_point]
 
     rich.print(f"run_cmd: {str.join(' ', run_cmd)}")
     subprocess.run(run_cmd, check=True)
+
+    if detached:
+        rich.print("[yellow]Container started in detached mode. To view logs, run:[/yellow]")
+        rich.print(f"[bold]docker logs -f {repo_config.name}[/bold]")

src/codegen/git/repo_operator/repo_operator.py

@@ -113,16 +113,20 @@ def viz_file_path(self) -> str:
         return os.path.join(self.viz_path, "graph.json")
 
     def _set_bot_email(self, git_cli: GitCLI) -> None:
+        logging.info(f"****** Setting bot email to {CODEGEN_BOT_EMAIL} ******")
         with git_cli.config_writer("repository") as writer:
             if not writer.has_section("user"):
                 writer.add_section("user")
             writer.set("user", "email", CODEGEN_BOT_EMAIL)
+            logging.info(f"****** [DONE] Setting bot email to {CODEGEN_BOT_EMAIL} ******")
 
     def _set_bot_username(self, git_cli: GitCLI) -> None:
+        logging.info(f"****** Setting bot USERNAME to {CODEGEN_BOT_NAME} ******")
         with git_cli.config_writer("repository") as writer:
             if not writer.has_section("user"):
                 writer.add_section("user")
             writer.set("user", "name", CODEGEN_BOT_NAME)
+            logging.info(f"****** [DONE] Setting bot USERNAME to {CODEGEN_BOT_NAME} ******")
 
     def _unset_bot_email(self, git_cli: GitCLI) -> None:
         with git_cli.config_writer("repository") as writer:
@@ -136,6 +140,7 @@ def _unset_bot_username(self, git_cli: GitCLI) -> None:
 
     @cached_property
     def git_cli(self) -> GitCLI:
+        logger.info("**** Initializing git_cli!!!*****")
         git_cli = GitCLI(self.repo_path)
         username = None
         user_level = None