content: ARG TASK_IMAGE
  FROM $TASK_IMAGE
  
  # We install pyhooks as root so that we can run python -m pyhooks.agent_output as root. agent_output.py polls /agent-output
  # for the agent's stdout, stderr, and exit status, then sends any changes to Vivaria in an API request.
  # Pip uses /root/.cache to cache the downloaded packages, so we use --mount=type=cache to share the cache between builds.
  RUN --mount=type=cache,target=/root/.cache \
      python -m pip install "git+https://github.com/METR/pyhooks.git@fc84345493a339c1f066f0c143aa48d86d0898a0"
  
  # Check that root can use pyhooks.
  RUN python -c "import pyhooks"
  
  USER agent
  WORKDIR /home/agent/.agent_code
  
  COPY --chown=agent:agent ./requirements.tx[t] .
  RUN if [ ! -f requirements.txt ]; then touch requirements.txt; fi
  
  # Pip uses /home/agent/.cache to cache the downloaded packages, so we use --mount=type=cache to share the cache between builds.
  RUN --mount=type=cache,target=/home/agent/.cache \
      python -m pip install -r requirements.txt
  
  COPY --chown=agent:agent . .
  
  USER root
  
  # Set PasswordAuthentication to no to avoid confusing users when they try to SSH into a container they don't have access to.
  # If PasswordAuthentication is set to yes, the user will be prompted for a password that they don't know.
  # Set AcceptEnv to * to allow the viv CLI to set environment variables in the container when SSHing in (e.g. agent token,
  # environment variables from secrets.env).
  CMD echo 'PasswordAuthentication no' >> /etc/ssh/sshd_config && \
      echo 'AcceptEnv *' >> /etc/ssh/sshd_config && \
      service ssh restart && \
      su - agent -c "python -m pyhooks.python_server"
  