secondshift 0.1.0b12__py3-none-any.whl

secondshift/runner/agent_image/subagent_hook.py

@@ -0,0 +1,71 @@
+"""Claude Code calls this before it starts a sub-agent, and it may refuse.
+
+Claude Code runs it as a PreToolUse hook for its Agent tool, with the run's
+limit and a counter file as arguments. It allows only the server-owned worker,
+only from the run's own agent, in the foreground, up to the limit. Any error
+exits with 2, which makes Claude Code refuse the call. Other exit codes would
+let it through. The agent reads the event stream too, and stops the CLI if a
+sub-agent runs that this should have refused.
+
+Claude Code runs this file by path, so it uses only the standard library.
+"""
+
+import fcntl
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+# The one sub-agent type a run may start. Its tools are the broker's only.
+WORKER = "secondshift-worker"
+# Claude Code's name for the tool in --tools, and in tool calls.
+AGENT_TOOLS = ("Task", "Agent")
+
+
+def refusal(event: dict[str, Any], limit: int, counter: Path) -> str:
+    """Why this start is refused, or an empty string to allow it."""
+    if event.get("agent_id"):
+        return "A sub-agent cannot start sub-agents. Do this yourself."
+    task = event.get("tool_input") or {}
+    if task.get("subagent_type") != WORKER:
+        return f"Start sub-agents with subagent_type {WORKER}."
+    # Claude Code runs a sub-agent in the background unless the call says not
+    # to, so only an explicit false keeps it in the foreground.
+    if task.get("run_in_background") is not False:
+        return (
+            "Start sub-agents with run_in_background set to false, and wait "
+            "for their report."
+        )
+    with counter.open("a+") as stream:
+        fcntl.flock(stream, fcntl.LOCK_EX)
+        stream.seek(0)
+        started = int(stream.read() or 0)
+        if started >= limit:
+            return (
+                f"This run has started the most sub-agents it may, {limit}. "
+                "Do the rest yourself."
+            )
+        stream.seek(0)
+        stream.truncate()
+        stream.write(str(started + 1))
+    return ""
+
+
+def main() -> None:
+    try:
+        limit, counter = int(sys.argv[1]), Path(sys.argv[2])
+        reason = refusal(json.load(sys.stdin), limit, counter)
+    except Exception:
+        print("The sub-agent check failed.", file=sys.stderr)
+        sys.exit(2)
+    if reason:
+        decision = {
+            "hookEventName": "PreToolUse",
+            "permissionDecision": "deny",
+            "permissionDecisionReason": reason,
+        }
+        print(json.dumps({"hookSpecificOutput": decision}))
+
+
+if __name__ == "__main__":
+    main()

secondshift/runner/agent_protocol.py

@@ -0,0 +1,41 @@
+"""The control plane sends only fixed agent choices, never Docker options."""
+
+from typing import Literal
+
+from ninja import Schema
+from pydantic import ConfigDict, Field, SecretStr
+
+from secondshift.runner.agent_image.outcomes import AgentOutcome
+
+
+class AgentInput(Schema):
+    model_config = ConfigDict(extra="forbid")
+    backend: Literal["claude_code", "codex"]
+    credential_key: str = Field(pattern=r"^[0-9a-f]{32}$")
+    auth: SecretStr = Field(max_length=64000, exclude=True)
+    broker_token: SecretStr = Field(max_length=200, exclude=True)
+    prompt: str = Field(max_length=250000)
+    max_turns: int = Field(ge=1, le=10000)
+    token_budget: int = Field(ge=1)
+    timeout_seconds: int = Field(ge=1, le=86400)
+    # The most sub-agents the agent may start. Only Claude Code can so far.
+    subagents: int = Field(default=0, ge=0, le=100)
+
+
+class AgentCheckInput(Schema):
+    """A login check. The runner owns the prompt and every limit."""
+
+    model_config = ConfigDict(extra="forbid")
+    backend: Literal["claude_code", "codex"]
+    credential_key: str = Field(pattern=r"^[0-9a-f]{32}$")
+    auth: SecretStr = Field(max_length=64000, exclude=True)
+
+
+class AgentOutput(Schema):
+    report: str = ""
+    succeeded: bool = False
+    reason: str = AgentOutcome.AGENT_FAILED
+    usage: dict[str, int] = Field(default_factory=dict)
+    # Each sub-agent's model requests and tool calls, in the order they started.
+    subagents: list[dict[str, int]] = Field(default_factory=list, max_length=100)
+    versions: dict[str, str] = Field(default_factory=dict)

secondshift/runner/agents.py

@@ -0,0 +1,254 @@
+"""Server-owned agent containers. Repository command sessions stay offline."""
+
+import fcntl
+import json
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+from django.conf import settings
+
+from secondshift.config.constants import (
+    HOST_ALIAS,
+    AgentBackend,
+)
+from secondshift.runner.agent_credentials import decrypt_agent_auth
+from secondshift.runner.agent_image.outcomes import (
+    AUTH_FILES,
+    AgentConfig,
+    AgentOutcome,
+)
+from secondshift.runner.agent_protocol import AgentCheckInput, AgentInput, AgentOutput
+from secondshift.runner.dao import is_session_active
+from secondshift.runner.images import resolve_image_id
+from secondshift.runner.protocol import CommandRequest, RunnerProtocolError
+from secondshift.runner.services import DockerRunnerService, agent_name, sandbox_flags
+
+# A login check asks for one short answer, with no tools.
+CHECK_PROMPT = "Reply with the single word OK."
+CHECK_TOKEN_BUDGET = 50_000
+CHECK_TIMEOUT_SECONDS = 20
+
+
+def require_agent_configuration() -> str:
+    """The exact image to run. Only a local runner may resolve a tag."""
+    if not settings.AGENT_IMAGE or not settings.AGENT_BROKER_URL:
+        raise RunnerProtocolError(
+            "CLI agent execution is not configured on the runner."
+        )
+    image = settings.AGENT_IMAGE
+    if not re.fullmatch(r"(?:[^@\s]+@)?sha256:[0-9a-f]{64}", image):
+        if not settings.AGENT_IMAGE_TAGS_ALLOWED:
+            raise RunnerProtocolError("The agent image must be pinned by digest.")
+        try:
+            image = resolve_image_id(image)
+        except (subprocess.SubprocessError, OSError, ValueError) as error:
+            raise RunnerProtocolError(
+                "The agent image isn't built. Restart SecondShift to build it."
+            ) from error
+    DockerRunnerService().verify_runtime()
+    return image
+
+
+def agent_command(session_id: str, auth_path: Path, image: str) -> list[str]:
+    command = [
+        "docker",
+        "run",
+        "--rm",
+        "--interactive",
+        "--name",
+        agent_name(session_id),
+        *sandbox_flags(),
+    ]
+    if settings.AGENT_RESOLV_CONF:
+        resolver = Path(settings.AGENT_RESOLV_CONF)
+        if (
+            not resolver.is_absolute()
+            or not resolver.is_file()
+            or resolver.is_symlink()
+        ):
+            raise RunnerProtocolError(
+                "The agent resolver must be an existing absolute file."
+            )
+        command += [f"--mount=type=bind,src={resolver},dst=/etc/resolv.conf,readonly"]
+    uid, gid = settings.RUNNER_SANDBOX_UID, settings.RUNNER_SANDBOX_GID
+    return [
+        *command,
+        f"--network={settings.AGENT_NETWORK}",
+        # Docker Desktop defines this alias. Native Linux Docker needs it added.
+        f"--add-host={HOST_ALIAS}:host-gateway",
+        f"--tmpfs=/tmp:rw,nosuid,size=256m,uid={uid},gid={gid}",
+        f"--mount=type=bind,src={auth_path},dst=/auth",
+        "--workdir=/tmp",
+        "--env=HOME=/tmp",
+        "--env=CODEX_HOME=/auth",
+        "--env=CLAUDE_CONFIG_DIR=/tmp/claude",
+        "--env=DISABLE_AUTOUPDATER=1",
+        f"--label=secondshift.agent-credential={auth_path.name}",
+        "--log-driver=none",
+        "--entrypoint=python3",
+        image,
+        "-m",
+        "secondshift.runner.agent_image.agent",
+    ]
+
+
+def provision_auth(request: AgentInput | AgentCheckInput) -> Path:
+    root = settings.AGENT_AUTH_ROOT.resolve()
+    root.mkdir(parents=True, exist_ok=True, mode=0o700)
+    path = root / request.credential_key
+    # Keys are random, server-issued, immutable per reconnect. No paths from repos.
+    if path.exists():
+        if path.is_symlink() or not path.is_dir():
+            raise RunnerProtocolError("Invalid agent credential directory.")
+        return path
+    auth = decrypt_agent_auth(request.auth.get_secret_value())
+    path.mkdir(mode=0o700)
+    target = path / AUTH_FILES[request.backend]
+    target.write_text(auth)
+    target.chmod(0o600)
+    uid, gid = settings.RUNNER_SANDBOX_UID, settings.RUNNER_SANDBOX_GID
+    if os.geteuid() == 0:
+        os.chown(target, uid, gid)
+        os.chown(path, uid, gid)
+    return path
+
+
+@contextmanager
+def credential_lock(key: str) -> Iterator[None]:
+    root = settings.AGENT_AUTH_ROOT.resolve()
+    root.mkdir(parents=True, exist_ok=True, mode=0o700)
+    with (root / f"{key}.lock").open("a") as stream:
+        try:
+            fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            raise RunnerProtocolError(
+                "This agent login is already running another job."
+            ) from None
+        try:
+            yield
+        finally:
+            fcntl.flock(stream, fcntl.LOCK_UN)
+
+
+def execute_agent(session_id: str, request: AgentInput) -> AgentOutput:
+    with credential_lock(request.credential_key):
+        return execute_locked_agent(session_id, request)
+
+
+def execute_locked_agent(session_id: str, request: AgentInput) -> AgentOutput:
+    image = require_agent_configuration()
+    if not is_session_active(session_id):
+        raise RunnerProtocolError("Unknown or inactive runner session.")
+    # Reserve once, before starting the process: an ambiguous outcome is not retried.
+    reservation = CommandRequest(
+        call_id="agent-execution", argv=(request.backend, request.credential_key)
+    )
+    service = DockerRunnerService()
+    if service.reserve_command(session_id, reservation) is not None:
+        raise RunnerProtocolError("This agent execution cannot be repeated.")
+    auth_path = provision_auth(request)
+    configuration: AgentConfig = {
+        "backend": request.backend,
+        "broker_url": settings.AGENT_BROKER_URL,
+        "broker_token": request.broker_token.get_secret_value(),
+        "prompt": request.prompt,
+        "max_turns": request.max_turns,
+        "token_budget": request.token_budget,
+        "subagents": request.subagents
+        if AgentBackend(request.backend).starts_subagents
+        else 0,
+    }
+    return run_agent_container(
+        agent_name(session_id),
+        agent_command(session_id, auth_path, image),
+        configuration,
+        request.timeout_seconds,
+    )
+
+
+def check_agent(request: AgentCheckInput) -> AgentOutput:
+    """Run the login's CLI once with a fixed prompt, no tools and no repository."""
+    with credential_lock(request.credential_key):
+        image = require_agent_configuration()
+        auth_path = provision_auth(request)
+        name = f"check-{request.credential_key}"
+        configuration: AgentConfig = {
+            "backend": request.backend,
+            "prompt": CHECK_PROMPT,
+            "max_turns": 1,
+            "token_budget": CHECK_TOKEN_BUDGET,
+        }
+        return run_agent_container(
+            agent_name(name),
+            agent_command(name, auth_path, image),
+            configuration,
+            CHECK_TIMEOUT_SECONDS,
+        )
+
+
+def run_agent_container(
+    name: str, command: list[str], configuration: AgentConfig, timeout: int
+) -> AgentOutput:
+    # Auth and the broker token never appear in argv, Docker env, logs, or DBOS.
+    with tempfile.TemporaryFile() as output:
+        try:
+            completed = subprocess.run(
+                command,
+                input=json.dumps(configuration).encode(),
+                stdout=output,
+                stderr=subprocess.DEVNULL,
+                timeout=timeout,
+                check=False,
+            )
+            output.seek(0)
+            data = output.read(settings.RUNNER_COMMAND_OUTPUT_MAX_BYTES + 1)
+            if completed.returncode:
+                return AgentOutput(reason=AgentOutcome.AGENT_STOPPED)
+            if len(data) > settings.RUNNER_COMMAND_OUTPUT_MAX_BYTES:
+                return AgentOutput(reason=AgentOutcome.AGENT_FAILED)
+            try:
+                return AgentOutput.model_validate_json(data)
+            except ValueError:
+                return AgentOutput(reason=AgentOutcome.INVALID_AGENT_OUTPUT)
+        except subprocess.TimeoutExpired:
+            return AgentOutput(reason=AgentOutcome.TIMED_OUT)
+        finally:
+            DockerRunnerService().remove_container(name)
+
+
+def remove_agent_credentials(key: str) -> None:
+    """Stop consumers before removing a disconnected login's CLI cache."""
+    if not re.fullmatch(r"[0-9a-f]{32}", key):
+        raise RunnerProtocolError("Invalid credential key.")
+    service = DockerRunnerService()
+    listed = service.run_command(
+        [
+            "docker",
+            "ps",
+            "--all",
+            "--quiet",
+            "--filter",
+            f"label=secondshift.agent-credential={key}",
+        ],
+        timeout=15,
+    )
+    if listed.returncode:
+        raise RunnerProtocolError("The runner could not revoke agent credentials.")
+    for container in listed.stdout.splitlines():
+        if not re.fullmatch(r"[0-9a-f]{12,64}", container):
+            raise RunnerProtocolError(
+                "Docker returned an invalid container identifier."
+            )
+        service.remove_container(container)
+    with credential_lock(key):
+        path = settings.AGENT_AUTH_ROOT.resolve() / key
+        if path.is_symlink():
+            raise RunnerProtocolError("Invalid agent credential directory.")
+        if path.exists():
+            shutil.rmtree(path)

secondshift/runner/api.py

@@ -0,0 +1,123 @@
+import hmac
+
+from django.conf import settings
+from django.http import FileResponse, HttpRequest
+from ninja import NinjaAPI, Status
+from ninja.security import HttpBearer
+
+from secondshift.runner.agent_credentials import provisioning_public_key
+from secondshift.runner.agent_protocol import AgentCheckInput, AgentInput, AgentOutput
+from secondshift.runner.agents import (
+    check_agent,
+    execute_agent,
+    remove_agent_credentials,
+)
+from secondshift.runner.dao import has_sessions
+from secondshift.runner.protocol import (
+    CommandRequest,
+    RunnerProtocolError,
+    RunnerSessionSpec,
+)
+from secondshift.runner.schemas import (
+    CommandInput,
+    CommandOutput,
+    SessionInput,
+    SessionOutput,
+)
+from secondshift.runner.services import get_runner_service
+from secondshift.runner.sources import has_source, receive_source
+
+
+class RunnerBearer(HttpBearer):
+    def authenticate(self, request: HttpRequest, token: str) -> str | None:
+        if hmac.compare_digest(token, settings.RUNNER_SERVICE_TOKEN):
+            return token
+        return None
+
+
+api = NinjaAPI(
+    title="SecondShift Runner API",
+    version="1.0",
+    auth=RunnerBearer(),
+    urls_namespace="runner",
+    docs_url=None,
+    openapi_url=None,
+)
+
+
+@api.exception_handler(RunnerProtocolError)
+def protocol_error(request: HttpRequest, error: RunnerProtocolError):
+    return api.create_response(request, {"error": str(error)}, status=409)
+
+
+@api.exception_handler(ValueError)
+def invalid_protocol_input(request: HttpRequest, error: ValueError):
+    return api.create_response(request, {"error": str(error)}, status=422)
+
+
+@api.get("/-/health/", auth=None)
+def health(request: HttpRequest) -> dict[str, bool]:
+    has_sessions()
+    return {"ok": True}
+
+
+@api.post("/v1/sessions", response={201: SessionOutput})
+def start_session(request: HttpRequest, payload: SessionInput):
+    spec = RunnerSessionSpec(**payload.model_dump())
+    session_id = get_runner_service().start(spec)
+    return Status(201, {"session_id": session_id})
+
+
+@api.post("/v1/sessions/{session_id}/commands", response=CommandOutput)
+def execute_command(request: HttpRequest, session_id: str, payload: CommandInput):
+    command = CommandRequest(**{**payload.model_dump(), "argv": tuple(payload.argv)})
+    return get_runner_service().execute(session_id, command)
+
+
+@api.post("/v1/sessions/{session_id}/collection")
+def collect_session(request: HttpRequest, session_id: str) -> FileResponse:
+    """Stop the sandbox and return its changes and saved files as one archive."""
+    archive = get_runner_service().collect(session_id)
+    return FileResponse(archive.open("rb"), content_type="application/gzip")
+
+
+@api.delete("/v1/sessions/{session_id}", response={204: None})
+def stop_session(request: HttpRequest, session_id: str):
+    get_runner_service().stop(session_id)
+    return Status(204, None)
+
+
+@api.api_operation(["HEAD"], "/v1/sources/{sha256}", response={204: None, 404: None})
+def source_exists(request: HttpRequest, sha256: str) -> Status[None]:
+    return Status(204 if has_source(sha256) else 404, None)
+
+
+@api.put("/v1/sources/{sha256}", response={204: None})
+def upload_source(request: HttpRequest, sha256: str) -> Status[None]:
+    # Read the body as a stream, so a snapshot never sits in memory whole.
+    length = request.headers.get("Content-Length")
+    receive_source(sha256, request, int(length) if length else None)
+    return Status(204, None)
+
+
+@api.post("/v1/sessions/{session_id}/agent", response=AgentOutput)
+def run_agent(
+    request: HttpRequest, session_id: str, payload: AgentInput
+) -> AgentOutput:
+    return execute_agent(session_id, payload)
+
+
+@api.post("/v1/agent-checks", response=AgentOutput)
+def run_agent_check(request: HttpRequest, payload: AgentCheckInput) -> AgentOutput:
+    return check_agent(payload)
+
+
+@api.delete("/v1/agent-credentials/{key}", response={204: None})
+def revoke_agent_credentials(request: HttpRequest, key: str) -> Status[None]:
+    remove_agent_credentials(key)
+    return Status(204, None)
+
+
+@api.get("/v1/agent-provisioning-key", response=dict[str, str])
+def agent_provisioning_key(request: HttpRequest) -> dict[str, str]:
+    return {"public_key": provisioning_public_key()}

secondshift/runner/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class RunnerConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "secondshift.runner"

secondshift/runner/client.py

@@ -0,0 +1,168 @@
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
+
+import httpx
+from django.conf import settings
+
+from secondshift.runner.agent_credentials import encrypt_agent_auth
+from secondshift.runner.agent_protocol import AgentCheckInput, AgentInput, AgentOutput
+from secondshift.runner.protocol import (
+    CommandRequest,
+    CommandResult,
+    RunnerProtocolError,
+    RunnerSessionSpec,
+)
+
+
+class HttpRunnerClient:
+    """Trusted control-plane client for a separately deployed runner."""
+
+    def __init__(
+        self,
+        base_url: str,
+        auth_token: str,
+        timeout: int = 30,
+        transport: httpx.BaseTransport | None = None,
+    ) -> None:
+        if not base_url.strip() or not auth_token:
+            raise ValueError("Runner URL and authentication token are required.")
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.session = httpx.Client(
+            base_url=self.base_url,
+            trust_env=False,
+            transport=transport,
+            headers={
+                "Authorization": f"Bearer {auth_token}",
+                "Accept": "application/json",
+            },
+        )
+
+    def send(
+        self, method: str, path: str, timeout: float | None = None, **kwargs: Any
+    ) -> httpx.Response:
+        """One call to the runner. Transport errors become RunnerProtocolError."""
+        try:
+            return self.session.request(
+                method, path, timeout=timeout or self.timeout, **kwargs
+            )
+        except httpx.TimeoutException as error:
+            raise RunnerProtocolError("The runner did not answer in time.") from error
+        except httpx.HTTPError as error:
+            raise RunnerProtocolError(
+                "The runner service could not be reached."
+            ) from error
+
+    def upload_source(self, sha256: str, path: Path) -> None:
+        """Send a run's snapshot as a raw stream, unless the runner has it."""
+        if self.send("HEAD", f"/v1/sources/{sha256}").status_code == 204:
+            return
+        with path.open("rb") as stream:
+            response = self.send(
+                "PUT",
+                f"/v1/sources/{sha256}",
+                timeout=180,
+                content=stream,
+                # Django reads a body only up to its Content-Length.
+                headers={
+                    "Content-Type": "application/gzip",
+                    "Content-Length": str(path.stat().st_size),
+                },
+            )
+        if response.status_code != 204:
+            self.response_payload(response)
+
+    def start(self, spec: RunnerSessionSpec) -> str:
+        response = self.send("POST", "/v1/sessions", json=asdict(spec))
+        payload = self.response_payload(response)
+        return str(payload["session_id"])
+
+    def execute(self, session_id: str, request: CommandRequest) -> CommandResult:
+        response = self.send(
+            "POST",
+            f"/v1/sessions/{session_id}/commands",
+            timeout=max(self.timeout, request.timeout_seconds + 5),
+            json=asdict(request),
+        )
+        return CommandResult(**self.response_payload(response))
+
+    def collect(self, session_id: str, destination: Path) -> None:
+        """Stop the sandbox and save its archive of changes and files."""
+        try:
+            with self.session.stream(
+                "POST", f"/v1/sessions/{session_id}/collection", timeout=360
+            ) as response:
+                if response.status_code != 200:
+                    response.read()
+                    self.response_payload(response)
+                    raise RunnerProtocolError("Runner returned an invalid response.")
+                with destination.open("wb") as output:
+                    for chunk in response.iter_bytes():
+                        output.write(chunk)
+        except httpx.TimeoutException as error:
+            raise RunnerProtocolError("The runner did not answer in time.") from error
+        except httpx.HTTPError as error:
+            raise RunnerProtocolError(
+                "The runner service could not be reached."
+            ) from error
+
+    def encrypted_auth(self, auth: str) -> str:
+        """An agent login, encrypted so that only the runner can read it."""
+        key = self.response_payload(self.send("GET", "/v1/agent-provisioning-key"))
+        return encrypt_agent_auth(str(key["public_key"]), auth)
+
+    def run_agent(self, session_id: str, request: AgentInput) -> AgentOutput:
+        response = self.send(
+            "POST",
+            f"/v1/sessions/{session_id}/agent",
+            timeout=request.timeout_seconds + 30,
+            json=request.model_dump(mode="json")
+            | {
+                "auth": request.auth.get_secret_value(),
+                "broker_token": request.broker_token.get_secret_value(),
+            },
+        )
+        return AgentOutput.model_validate(self.response_payload(response))
+
+    def check_agent(self, request: AgentCheckInput) -> AgentOutput:
+        response = self.send(
+            "POST",
+            "/v1/agent-checks",
+            # The runner stops the CLI after its own limit. This covers startup.
+            timeout=28,
+            json=request.model_dump(mode="json")
+            | {"auth": request.auth.get_secret_value()},
+        )
+        return AgentOutput.model_validate(self.response_payload(response))
+
+    def stop(self, session_id: str) -> None:
+        response = self.send("DELETE", f"/v1/sessions/{session_id}")
+        if response.status_code != 204:
+            self.response_payload(response)
+
+    @staticmethod
+    def response_payload(response: httpx.Response) -> dict[str, object]:
+        try:
+            payload = response.json()
+        except ValueError as error:
+            raise RunnerProtocolError(
+                f"Runner answered HTTP {response.status_code}."
+            ) from error
+        if response.is_error:
+            message = payload.get("error") if isinstance(payload, dict) else None
+            raise RunnerProtocolError(
+                str(message)
+                if message
+                else f"Runner answered HTTP {response.status_code}."
+            )
+        if not isinstance(payload, dict):
+            raise RunnerProtocolError("Runner returned an invalid response.")
+        return payload
+
+
+def get_runner_client() -> HttpRunnerClient:
+    return HttpRunnerClient(
+        base_url=settings.RUNNER_SERVICE_URL,
+        auth_token=settings.RUNNER_SERVICE_TOKEN,
+    )