[wip] Claude code container

examples/code-runner/claude_code/claude-code.Dockerfile

@@ -0,0 +1,15 @@
+# Use Node.js Alpine as base image for smaller size
+FROM node:20-alpine
+
+# Install bash and other required tools
+RUN apk add --no-cache bash git curl python3 py3-pip
+
+# Install Claude Code globally
+RUN npm install -g @anthropic-ai/claude-code
+
+
+# Ensure claude-code is in PATH
+ENV PATH="/usr/local/bin:$PATH"
+
+# Default command
+CMD ["bash"]

examples/code-runner/claude_code/claude_code_agent.py

@@ -0,0 +1,130 @@
+import pathlib
+import tempfile
+from pathlib import Path
+
+import flyte
+from flyte.extras import ContainerTask
+from flyte.io import Dir
+
+# Create a custom image with Claude Code pre-installed
+claude_code_image = flyte.Image.from_dockerfile(
+    name="flyte",
+    file=pathlib.Path(__file__).parent / "claude-code.Dockerfile",
+    registry="ghcr.io/flyteorg",
+    platform=("linux/amd64", "linux/arm64"),
+)
+
+claude_code_env = flyte.TaskEnvironment(
+    name="claude_code_env",
+    image=claude_code_image,
+)
+
+# Claude Code container task
+claude_code_task = ContainerTask(
+    name="run_claude_code",
+    image=claude_code_image,
+    input_data_dir="/var/inputs",
+    output_data_dir="/var/outputs",
+    inputs={"prompt": str, "working_dir": Dir, "max_turns": int, "budget_limit": float},
+    outputs={"generated_code": Dir, "exit_code": int, "log": str},
+    command=[
+        "/bin/bash",
+        "-c",
+        (
+            "set -o pipefail && "
+            "cd /var/inputs && "
+            "mkdir -p /var/outputs/generated_code && "
+            "cp -r . /var/outputs/generated_code/ && "
+            "cd /var/outputs/generated_code && "
+            'echo "{{.inputs.prompt}}" | claude '
+            "--print "
+            "--max-turns {{.inputs.max_turns}} "
+            "--dangerously-skip-permissions "
+            "--output-format text "
+            "> /var/outputs/log 2>&1; "
+            "echo $? > /var/outputs/exit_code"
+        ),
+    ],
+    secrets=[flyte.Secret(key="anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
+    resources=flyte.Resources(cpu=2, memory="4Gi"),
+)
+
+claude_code_env.add_task(claude_code_task)
+
+env = flyte.TaskEnvironment(
+    name="claude_code_agent",
+    resources=flyte.Resources(cpu=1, memory="2Gi"),
+    depends_on=[claude_code_env],
+)
+
+
+@env.task
+async def create_working_directory() -> Dir:
+    """
+    Create a temporary working directory for Claude Code to operate in.
+    """
+    temp_dir = Path(tempfile.mkdtemp(prefix="claude_code_workspace_"))
+
+    # Create a basic project structure
+    (temp_dir / "README.md").write_text(
+        "# Claude Code Generated Project\n\nThis project was generated by Claude Code.\n"
+    )
+
+    # Create basic directories
+    (temp_dir / "src").mkdir(exist_ok=True)
+    (temp_dir / "tests").mkdir(exist_ok=True)
+
+    return await Dir.from_local(str(temp_dir))
+
+
+@env.task
+async def main(
+        prompt: str, max_turns: int = 10, budget_limit: float = 5.0, timeout_minutes: int = 30
+) -> tuple[Dir, str, str, bool]:
+    """
+    Generate code using Claude Code based on the user's prompt.
+
+    Args:
+        prompt: The user's prompt/request for Claude Code
+        max_turns: Maximum number of agentic turns to allow
+        budget_limit: Maximum budget in dollars to spend on API calls
+        timeout_minutes: Maximum time to allow Claude Code to run
+
+    Returns:
+        tuple of (generated_code_dir, execution_log, exit_code, success)
+    """
+    print("---GENERATING CODE WITH CLAUDE CODE---")
+    print(f"Prompt: {prompt}")
+    print(f"Max turns: {max_turns}")
+    print(f"Budget limit: ${budget_limit}")
+
+    working_dir = await create_working_directory()
+
+    try:
+        # Run Claude Code
+        generated_code, exit_code, log = await claude_code_task(
+            prompt=prompt, working_dir=working_dir, max_turns=max_turns, budget_limit=budget_limit
+        )
+
+        success = exit_code.strip() == "0"
+
+        if success:
+            print("---CLAUDE CODE: SUCCESS---")
+        else:
+            raise RuntimeError(f"Claude Code failed with exit code {exit_code.strip()}")
+
+        return generated_code, log, exit_code, success
+
+    except Exception as e:
+        print(f"---CLAUDE CODE: ERROR - {e!s}---")
+        # Return the working directory as fallback
+        return working_dir, f"Error: {e!s}", "1", False
+
+
+if __name__ == "__main__":
+    prompt: str = "Create a simple Python web API using FastAPI that has endpoints for creating, "
+    "reading, updating and deleting users. Include proper error handling and validation."
+    parent = Path(__file__).parent
+    flyte.init_from_config(str(parent / "../../../config.yaml"), root_dir=parent)
+    run = flyte.with_runcontext(copy_style="all").run(main, prompt=prompt)
+    print(run.url)