secondshift 0.1.0b12__py3-none-any.whl

secondshift/runner/protocol.py

@@ -0,0 +1,97 @@
+import hashlib
+import json
+import re
+from dataclasses import asdict, dataclass
+from pathlib import PurePosixPath
+
+
+class RunnerProtocolError(RuntimeError):
+    pass
+
+
+REPOSITORY_NAME = re.compile(r"^[A-Za-z0-9._-]+(?:/[A-Za-z0-9._-]+)+$")
+COMMIT_SHA = re.compile(r"^[0-9a-f]{40,64}$")
+SESSION_KEY = re.compile(r"^[0-9a-f]{32}$")
+SOURCE_SHA256 = re.compile(r"^[0-9a-f]{64}$")
+
+
+@dataclass(frozen=True)
+class RunnerSessionSpec:
+    run_id: int
+    session_key: str
+    repository: str
+    commit_sha: str
+    image: str
+    # The uploaded snapshot of the commit's files the session unpacks.
+    source_sha256: str
+    # Record the unpacked files, so the agent's changes can be collected.
+    # The control plane sets it from the run's plan, never from repository input.
+    collect_changes: bool = False
+    # An uploaded archive of earlier runs' saved files, mounted read-only.
+    # Empty when the run has none.
+    history_sha256: str = ""
+    # An uploaded archive of a skill from another repository, mounted
+    # read-only. Empty when the skill is in the run's repository or there is
+    # none. The control plane builds it, so the skill's repository never
+    # chooses where it goes or how it is mounted.
+    skill_sha256: str = ""
+
+    def __post_init__(self) -> None:
+        if self.run_id < 1:
+            raise ValueError("run_id must be positive")
+        if not SESSION_KEY.fullmatch(self.session_key):
+            raise ValueError("session_key must be 32 lowercase hexadecimal characters")
+        if not REPOSITORY_NAME.fullmatch(self.repository):
+            raise ValueError("repository must be a normalized owner/name path")
+        if any(part in {".", ".."} for part in self.repository.split("/")):
+            raise ValueError("repository must be a normalized owner/name path")
+        if not COMMIT_SHA.fullmatch(self.commit_sha):
+            raise ValueError("commit_sha must be a full lowercase hexadecimal SHA")
+        if not self.image.strip():
+            raise ValueError("a server-owned image is required")
+        if not SOURCE_SHA256.fullmatch(self.source_sha256):
+            raise ValueError("source_sha256 must be a lowercase SHA-256")
+        if self.history_sha256 and not SOURCE_SHA256.fullmatch(self.history_sha256):
+            raise ValueError("history_sha256 must be empty or a lowercase SHA-256")
+        if self.skill_sha256 and not SOURCE_SHA256.fullmatch(self.skill_sha256):
+            raise ValueError("skill_sha256 must be empty or a lowercase SHA-256")
+
+    @property
+    def session_id(self) -> str:
+        # The key is random per run, so a reset or restored database that
+        # reuses a run ID never matches an older session.
+        return f"run-{self.run_id}-{self.session_key}"
+
+
+@dataclass(frozen=True)
+class CommandRequest:
+    call_id: str
+    argv: tuple[str, ...]
+    cwd: str = "."
+    timeout_seconds: int = 120
+
+    def __post_init__(self) -> None:
+        if not self.call_id.strip():
+            raise ValueError("call_id is required")
+        if not self.argv or any(
+            not argument or "\0" in argument for argument in self.argv
+        ):
+            raise ValueError("argv must contain non-empty arguments")
+        path = PurePosixPath(self.cwd)
+        if path.is_absolute() or ".." in path.parts:
+            raise ValueError("cwd must stay inside the workspace")
+        if not 1 <= self.timeout_seconds <= 900:
+            raise ValueError("timeout_seconds must be between 1 and 900")
+
+    @property
+    def fingerprint(self) -> str:
+        encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(encoded.encode()).hexdigest()
+
+
+@dataclass(frozen=True)
+class CommandResult:
+    call_id: str
+    exit_code: int | None
+    output: str
+    timed_out: bool = False

secondshift/runner/schemas.py

@@ -0,0 +1,36 @@
+from ninja import Schema
+from pydantic import ConfigDict, Field
+
+
+class StrictSchema(Schema):
+    model_config = ConfigDict(extra="forbid")
+
+
+class SessionInput(StrictSchema):
+    run_id: int = Field(gt=0)
+    session_key: str = Field(min_length=32, max_length=32)
+    repository: str = Field(min_length=3, max_length=300)
+    commit_sha: str = Field(min_length=40, max_length=64)
+    image: str = Field(min_length=1, max_length=300)
+    source_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
+    collect_changes: bool = False
+    history_sha256: str = Field(default="", pattern=r"^(?:[0-9a-f]{64})?$")
+    skill_sha256: str = Field(default="", pattern=r"^(?:[0-9a-f]{64})?$")
+
+
+class SessionOutput(Schema):
+    session_id: str
+
+
+class CommandInput(StrictSchema):
+    call_id: str = Field(min_length=1, max_length=200)
+    argv: list[str] = Field(min_length=1, max_length=256)
+    cwd: str = Field(default=".", max_length=1000)
+    timeout_seconds: int = Field(default=120, ge=1, le=900)
+
+
+class CommandOutput(Schema):
+    call_id: str
+    exit_code: int | None
+    output: str
+    timed_out: bool

secondshift/runner/services.py

@@ -0,0 +1,465 @@
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+from django.conf import settings
+from django.db import IntegrityError, transaction
+
+from secondshift.config.constants import (
+    SANDBOX_HISTORY_DIR,
+    SANDBOX_OUTPUTS_DIR,
+    SANDBOX_SKILL_DIR,
+    SandboxRuntime,
+)
+from secondshift.runner.collection import OUTPUTS_DIR, collect, record_baseline
+from secondshift.runner.dao import (
+    complete_command,
+    create_command,
+    create_session,
+    get_session,
+    lock_session,
+    session_command,
+    set_session_status,
+)
+from secondshift.runner.models import RunnerCommand, RunnerSession
+from secondshift.runner.protocol import (
+    CommandRequest,
+    CommandResult,
+    RunnerProtocolError,
+    RunnerSessionSpec,
+)
+from secondshift.runner.sources import sweep_sources, unpack_source
+
+TRUNCATION_MARKER = b"\n[... command output truncated by SecondShift ...]\n"
+# Beside the workspace: the files earlier runs of the job saved.
+HISTORY_DIR = "history"
+# Beside the workspace: a skill from another repository.
+SKILL_DIR = "skill"
+
+
+def sandbox_name(session_id: str) -> str:
+    return f"secondshift-{session_id}"
+
+
+def agent_name(session_id: str) -> str:
+    return f"secondshift-agent-{session_id}"
+
+
+def sandbox_flags() -> list[str]:
+    """The runtime, hardening, limits and user every sandbox container gets.
+
+    Callers add only their own mounts, network and temporary files.
+    """
+    flags = []
+    if settings.RUNNER_RUNTIME != SandboxRuntime.DOCKER:
+        flags += ["--runtime", settings.RUNNER_RUNTIME]
+    return [
+        *flags,
+        "--read-only",
+        "--cap-drop=ALL",
+        "--security-opt=no-new-privileges:true",
+        f"--pids-limit={settings.RUNNER_PIDS_LIMIT}",
+        f"--memory={settings.RUNNER_MEMORY}",
+        # Docker otherwise lets the sandbox swap as much again as its memory.
+        f"--memory-swap={settings.RUNNER_MEMORY}",
+        f"--cpus={settings.RUNNER_CPUS}",
+        f"--user={settings.RUNNER_SANDBOX_UID}:{settings.RUNNER_SANDBOX_GID}",
+    ]
+
+
+class RunnerConfigurationError(RuntimeError):
+    pass
+
+
+class DockerRunnerService:
+    """Container operations and ORM-backed idempotency for the runner service."""
+
+    def verify_configuration(self) -> None:
+        settings.RUNNER_WORK_ROOT.mkdir(parents=True, exist_ok=True, mode=0o700)
+        sweep_sources()
+        self.verify_runtime()
+        self.ensure_image()
+
+    def ensure_image(self) -> None:
+        """Pull the sandbox image at startup so the first run doesn't wait on it."""
+        image = settings.RUNNER_IMAGE
+        if not self.run_command(
+            ["docker", "image", "inspect", image], timeout=15
+        ).returncode:
+            return
+        pulled = self.run_command(["docker", "pull", image], timeout=600)
+        if pulled.returncode:
+            detail = (pulled.stderr or pulled.stdout).strip().splitlines()
+            reason = detail[-1] if detail else "no output"
+            raise RunnerConfigurationError(
+                f"Couldn't pull the sandbox image {image}: {reason}"
+            )
+
+    def verify_runtime(self) -> None:
+        completed = self.run_command(
+            ["docker", "info", "--format", "{{json .Runtimes}}"], timeout=15
+        )
+        if completed.returncode:
+            raise RunnerConfigurationError("Docker is unavailable to the runner.")
+        if settings.RUNNER_RUNTIME == SandboxRuntime.DOCKER:
+            return
+        try:
+            runtimes = json.loads(completed.stdout)
+        except (TypeError, ValueError) as error:
+            raise RunnerConfigurationError(
+                "Docker returned an invalid runtime list."
+            ) from error
+        if settings.RUNNER_RUNTIME not in runtimes:
+            raise RunnerConfigurationError(
+                f"Docker runtime {settings.RUNNER_RUNTIME} is not registered."
+            )
+
+    def start(self, spec: RunnerSessionSpec) -> str:
+        if spec.image != settings.RUNNER_IMAGE:
+            raise RunnerProtocolError(
+                "The requested image is not enabled by this runner."
+            )
+        session_id = spec.session_id
+        container_name = sandbox_name(session_id)
+        session_root = (settings.RUNNER_WORK_ROOT / "sessions" / session_id).resolve()
+        self.require_owned_path(session_root)
+        workspace = session_root / "workspace"
+
+        with transaction.atomic():
+            session = lock_session(session_id)
+            if session is None:
+                try:
+                    session = create_session(
+                        session_id=session_id,
+                        run_id=spec.run_id,
+                        repository=spec.repository,
+                        commit_sha=spec.commit_sha,
+                        image=spec.image,
+                        source_sha256=spec.source_sha256,
+                        collect_changes=spec.collect_changes,
+                        history_sha256=spec.history_sha256,
+                        skill_sha256=spec.skill_sha256,
+                        container_name=container_name,
+                        workspace_path=str(workspace),
+                    )
+                except IntegrityError as error:
+                    raise RunnerProtocolError(
+                        "The run already has a runner session."
+                    ) from error
+            else:
+                self.require_matching_spec(session, spec)
+                if session.status == RunnerSession.Status.ACTIVE:
+                    if not self.container_exists(container_name):
+                        raise RunnerProtocolError(
+                            "The recorded runner session has lost its container."
+                        )
+                    return session_id
+                raise RunnerProtocolError(
+                    "The previous session outcome is not reusable; it will not be repeated."
+                )
+
+        self.remove_container(container_name)
+        self.remove_workspace(session_root)
+        workspace.mkdir(parents=True, mode=0o700)
+        outputs = session_root / OUTPUTS_DIR
+        outputs.mkdir(mode=0o700)
+        (session_root / HISTORY_DIR).mkdir(mode=0o700)
+        (session_root / SKILL_DIR).mkdir(mode=0o700)
+        try:
+            self.prepare_sandbox(spec, container_name, workspace, outputs)
+        except Exception:
+            self.remove_container(container_name)
+            self.remove_workspace(session_root)
+            set_session_status(session_id, RunnerSession.Status.ERROR)
+            raise
+        set_session_status(session_id, RunnerSession.Status.ACTIVE)
+        return session_id
+
+    def prepare_sandbox(
+        self,
+        spec: RunnerSessionSpec,
+        container_name: str,
+        workspace: Path,
+        outputs: Path,
+    ) -> None:
+        """Unpack the commit, any earlier files and any skill from another
+        repository, record the commit if asked, and start the container."""
+        unpack_source(spec.source_sha256, workspace)
+        history = workspace.parent / HISTORY_DIR
+        if spec.history_sha256:
+            unpack_source(spec.history_sha256, history)
+        skill = workspace.parent / SKILL_DIR
+        if spec.skill_sha256:
+            unpack_source(spec.skill_sha256, skill)
+        if spec.collect_changes:
+            record_baseline(workspace.parent, workspace)
+        self.prepare_workspace_permissions(workspace)
+        self.prepare_workspace_permissions(outputs)
+        self.prepare_workspace_permissions(history)
+        self.prepare_workspace_permissions(skill)
+        completed = self.run_command(
+            self.create_command(spec, container_name, workspace)
+        )
+        if completed.returncode:
+            raise RunnerProtocolError("The runner could not create the sandbox.")
+        completed = self.run_command(["docker", "start", container_name])
+        if completed.returncode:
+            raise RunnerProtocolError("The runner could not start the sandbox.")
+
+    def execute(self, session_id: str, request: CommandRequest) -> CommandResult:
+        cached = self.reserve_command(session_id, request)
+        if cached is not None:
+            return cached
+        session = get_session(session_id)
+        timeout = min(request.timeout_seconds, settings.RUNNER_COMMAND_TIMEOUT_SECONDS)
+        workdir = "/workspace"
+        if request.cwd != ".":
+            workdir += f"/{request.cwd}"
+        command = [
+            "docker",
+            "exec",
+            "--workdir",
+            workdir,
+            session.container_name,
+            *request.argv,
+        ]
+        output_descriptor, output_name = tempfile.mkstemp(
+            dir=settings.RUNNER_WORK_ROOT, prefix="command-output-"
+        )
+        output_path = Path(output_name)
+        timed_out = False
+        exit_code: int | None = None
+        try:
+            with os.fdopen(output_descriptor, "wb") as output:
+                try:
+                    completed = subprocess.run(
+                        command,
+                        stdout=output,
+                        stderr=subprocess.STDOUT,
+                        timeout=timeout,
+                        check=False,
+                    )
+                    exit_code = completed.returncode
+                except subprocess.TimeoutExpired:
+                    timed_out = True
+            result = CommandResult(
+                call_id=request.call_id,
+                exit_code=exit_code,
+                output=self.read_output(output_path),
+                timed_out=timed_out,
+            )
+            complete_command(
+                session,
+                request.call_id,
+                result.exit_code,
+                result.output,
+                result.timed_out,
+            )
+            if timed_out:
+                self.stop(session_id)
+            elif exit_code:
+                # gVisor ends the whole sandbox when it runs out of memory, so
+                # later commands could never succeed. Docker sets OOMKilled at
+                # once but reports the container running a little longer.
+                if self.container_oom_killed(session.container_name):
+                    self.stop(session_id)
+                    raise RunnerProtocolError("The sandbox ran out of memory.")
+                if not self.container_running(session.container_name):
+                    self.stop(session_id)
+                    raise RunnerProtocolError("The sandbox stopped unexpectedly.")
+            return result
+        finally:
+            output_path.unlink(missing_ok=True)
+
+    def collect(self, session_id: str) -> Path:
+        """Stop the sandbox and archive its changes and saved files.
+
+        The sandbox stops first, so nothing changes the files while the runner
+        reads them. Asking again returns the same archive.
+        """
+        try:
+            session = get_session(session_id)
+        except RunnerSession.DoesNotExist as error:
+            raise RunnerProtocolError("Unknown runner session.") from error
+        if session.status != RunnerSession.Status.ACTIVE:
+            raise RunnerProtocolError("The session has ended, so its files are gone.")
+        workspace = Path(session.workspace_path)
+        session_root = workspace.parent.resolve()
+        self.require_owned_path(session_root)
+        self.remove_container(agent_name(session_id))
+        self.run_command(["docker", "kill", session.container_name])
+        if self.container_running(session.container_name):
+            raise RunnerProtocolError("The runner could not stop the sandbox.")
+        return collect(session_root, workspace)
+
+    def stop(self, session_id: str) -> None:
+        with transaction.atomic():
+            session = lock_session(session_id)
+            if session is None:
+                return
+            session.status = RunnerSession.Status.STOPPED
+            session.save(update_fields=("status", "updated_at"))
+            container_name = session.container_name
+            workspace = Path(session.workspace_path)
+        self.remove_container(agent_name(session_id))
+        self.remove_container(container_name)
+        self.remove_workspace(workspace.parent)
+
+    def reserve_command(
+        self, session_id: str, request: CommandRequest
+    ) -> CommandResult | None:
+        with transaction.atomic():
+            session = lock_session(session_id, active_only=True)
+            if session is None:
+                raise RunnerProtocolError("Unknown or inactive runner session.")
+            command = session_command(session, request.call_id)
+            if command is None:
+                create_command(session, request.call_id, request.fingerprint)
+                return None
+            if command.fingerprint != request.fingerprint:
+                raise RunnerProtocolError(
+                    "A tool call ID was reused with different arguments."
+                )
+            if command.status != RunnerCommand.Status.COMPLETED:
+                raise RunnerProtocolError(
+                    "The previous execution outcome is unknown; it will not be repeated."
+                )
+            return CommandResult(
+                call_id=command.call_id,
+                exit_code=command.exit_code,
+                output=command.output,
+                timed_out=command.timed_out,
+            )
+
+    @staticmethod
+    def require_matching_spec(session: RunnerSession, spec: RunnerSessionSpec) -> None:
+        recorded = (
+            session.run_id,
+            session.repository,
+            session.commit_sha,
+            session.image,
+            session.source_sha256,
+            session.collect_changes,
+            session.history_sha256,
+            session.skill_sha256,
+        )
+        requested = (
+            spec.run_id,
+            spec.repository,
+            spec.commit_sha,
+            spec.image,
+            spec.source_sha256,
+            spec.collect_changes,
+            spec.history_sha256,
+            spec.skill_sha256,
+        )
+        if recorded != requested:
+            raise RunnerProtocolError("The run already has a different session spec.")
+
+    @staticmethod
+    def create_command(
+        spec: RunnerSessionSpec, container_name: str, workspace: Path
+    ) -> list[str]:
+        return [
+            "docker",
+            "create",
+            "--name",
+            container_name,
+            *sandbox_flags(),
+            "--network=none",
+            "--tmpfs=/tmp:rw,noexec,nosuid,size=256m",
+            # Bind mounts are writable by default. Docker rejects a bare "rw" field.
+            f"--mount=type=bind,src={workspace},dst=/workspace",
+            # Files the agent hands back. The runner collects them at the end.
+            f"--mount=type=bind,src={workspace.parent / OUTPUTS_DIR},"
+            f"dst={SANDBOX_OUTPUTS_DIR}",
+            # Files earlier runs of the job saved. The sandbox can only read them.
+            f"--mount=type=bind,src={workspace.parent / HISTORY_DIR},"
+            f"dst={SANDBOX_HISTORY_DIR},readonly",
+            # A skill from another repository. The sandbox can only read it,
+            # and it is never part of the repository's changes.
+            f"--mount=type=bind,src={workspace.parent / SKILL_DIR},"
+            f"dst={SANDBOX_SKILL_DIR},readonly",
+            "--workdir=/workspace",
+            "--env=HOME=/tmp",
+            # Skills that run Python must not leave bytecode in the repository.
+            "--env=PYTHONDONTWRITEBYTECODE=1",
+            spec.image,
+            "sleep",
+            "infinity",
+        ]
+
+    @staticmethod
+    def prepare_workspace_permissions(workspace: Path) -> None:
+        uid = settings.RUNNER_SANDBOX_UID
+        gid = settings.RUNNER_SANDBOX_GID
+        for root, directories, files in os.walk(workspace):
+            for name in [*directories, *files]:
+                os.chown(Path(root) / name, uid, gid, follow_symlinks=False)
+        os.chown(workspace, uid, gid, follow_symlinks=False)
+
+    @staticmethod
+    def read_output(path: Path) -> str:
+        limit = settings.RUNNER_COMMAND_OUTPUT_MAX_BYTES
+        size = path.stat().st_size
+        with path.open("rb") as stream:
+            if size <= limit:
+                content = stream.read()
+            else:
+                remaining = max(limit - len(TRUNCATION_MARKER), 0)
+                head_size = remaining // 2
+                tail_size = remaining - head_size
+                head = stream.read(head_size)
+                stream.seek(size - tail_size)
+                content = head + TRUNCATION_MARKER + stream.read(tail_size)
+        return content[:limit].decode(errors="replace")
+
+    def container_exists(self, name: str) -> bool:
+        return self.run_command(["docker", "inspect", name], timeout=10).returncode == 0
+
+    def container_state(self, name: str, field: str) -> bool:
+        completed = self.run_command(
+            ["docker", "inspect", "--format", f"{{{{.State.{field}}}}}", name],
+            timeout=10,
+        )
+        return completed.returncode == 0 and completed.stdout.strip() == "true"
+
+    def container_running(self, name: str) -> bool:
+        return self.container_state(name, "Running")
+
+    def container_oom_killed(self, name: str) -> bool:
+        return self.container_state(name, "OOMKilled")
+
+    def remove_container(self, name: str) -> None:
+        self.run_command(["docker", "rm", "--force", name], timeout=30)
+
+    def remove_workspace(self, path: Path) -> None:
+        path = path.resolve()
+        self.require_owned_path(path)
+        if path.exists():
+            shutil.rmtree(path)
+
+    @staticmethod
+    def require_owned_path(path: Path) -> None:
+        sessions_root = (settings.RUNNER_WORK_ROOT / "sessions").resolve()
+        if path == sessions_root or not path.is_relative_to(sessions_root):
+            raise RunnerProtocolError("Runner workspace escaped its owned root.")
+
+    @staticmethod
+    def run_command(
+        command: list[str], timeout: int = 30
+    ) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            command,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            check=False,
+        )
+
+
+def get_runner_service() -> DockerRunnerService:
+    return DockerRunnerService()