"""Code World - run coding agents on git repositories."""
from __future__ import annotations
import os
import shutil
import subprocess
from pathlib import Path
from typing import Annotated, ClassVar
from plato.agents import run_agent
from plato.worlds import (
Agent,
AgentConfig,
BaseWorld,
Observation,
RunConfig,
Secret,
StepResult,
register_world,
)
class CodeWorldConfig(RunConfig):
"""Configuration for CodeWorld.
All fields are typed - no .get() methods needed.
"""
# World-specific fields
repository_url: str
prompt: str
checkout: str = "main"
workspace_dir: str = "/workspace" # Can be overridden for local testing
# Agents (typed)
coder: Annotated[AgentConfig, Agent(description="Coding agent to run tasks on the repository")]
# Secrets (typed, optional)
git_token: Annotated[str | None, Secret(description="GitHub/GitLab token for HTTPS auth")] = None
git_ssh_key: Annotated[str | None, Secret(description="SSH private key for SSH auth")] = None
@register_world("code")
class CodeWorld(BaseWorld[CodeWorldConfig]):
"""World for running coding agents on git repositories.
Simple workflow:
1. reset(): Clone the repository
2. step(): Run the agent docker container (single step, then done)
"""
name: ClassVar[str] = "code"
description: ClassVar[str] = "Run coding agents on git repositories"
async def reset(self) -> Observation:
"""Clone the repository to workspace."""
repo_url = self.config.repository_url
checkout = self.config.checkout
workspace = Path(self.config.workspace_dir)
self.logger.info(f"Cloning repository: {repo_url}")
# Ensure workspace exists and is empty
if workspace.exists():
shutil.rmtree(workspace)
workspace.mkdir(parents=True)
# Setup git credentials
env = os.environ.copy()
git_token = self.config.git_token
git_ssh_key = self.config.git_ssh_key
if git_ssh_key:
ssh_key_path = workspace.parent / ".ssh_key"
ssh_key_path.write_text(git_ssh_key)
ssh_key_path.chmod(0o600)
env["GIT_SSH_COMMAND"] = f"ssh -i {ssh_key_path} -o StrictHostKeyChecking=no"
elif git_token and repo_url.startswith("https://"):
# Insert token into URL for HTTPS auth
if "gitlab.com" in repo_url:
repo_url = repo_url.replace("https://", f"https://oauth2:{git_token}@")
else:
# Works for GitHub, Gitea, and most git hosts
repo_url = repo_url.replace("https://", f"https://{git_token}@")
# Clone
clone_cmd = ["git", "clone", "--depth", "1", "-b", checkout, repo_url, str(workspace)]
result = subprocess.run(clone_cmd, env=env, capture_output=True, text=True)
if result.returncode != 0:
raise RuntimeError(f"Failed to clone: {result.stderr}")
self.logger.info(f"Cloned to {workspace}")
return Observation(data={"workspace": str(workspace), "checkout": checkout})
async def step(self) -> StepResult:
"""Run the agent container - single step then done."""
agent = self.config.coder
prompt = self.config.prompt
workspace = Path(self.config.workspace_dir)
self.logger.info(f"Running agent: {agent.image}")
self.logger.info(f"Agent config: {agent.config}")
# Create logs directory
logs_path = workspace / ".agent_logs"
logs_path.mkdir(parents=True, exist_ok=True)
# Run agent (trajectory and logs uploaded automatically via logging singleton)
await run_agent(
image=agent.image,
config=agent.config,
secrets=self.config.all_secrets,
instruction=prompt,
workspace=str(workspace),
logs_dir=str(logs_path),
pull=False,
)
return StepResult(
observation=Observation(data={"completed": True}),
done=True,
)
