Project-2 baseline

worker-toolkit-potion-polyglot/scripts/dnsjail.py

@@ -4,14 +4,19 @@ Opt-in with RACCOON_DNS_JAIL=1. Runs from the agent's own turn rather than from
 overlay — the allowlist comes from the proxy URL this process already holds (plus any hosts
 RACCOON_DNS_JAIL_ALLOW adds), so nothing has to be injected into the container, and the jail works on every harbor backend. Deliberately
 after agent-setup: a harness that downloads its CLI there still reaches the network to do it.
+
+Covers the agent's turn and nothing else -- `jailed` lifts it again before harbor's verifier
+phase, which shares the container and runs test suites we do not control.
 """
 
+import functools
 import logging
 import os
 import shlex
 from typing import Any
 
 JAIL = "/usr/local/bin/raccoon-dns-jail"
+STATE = "/tmp/.dnsjail"   # the container script's own state dir
 _NO_SCRIPT = "raccoon-dns-jail: not in this image"
 
 _URL_VARS = (
@@ -96,3 +101,44 @@ async def apply_dns_jail(agent: Any, environment: Any) -> None:
         _log.warning(
             "DNS jail: this task's image ships no resolver — the trial keeps normal network access"
         )
+
+
+async def lift_dns_jail(agent: Any, environment: Any) -> None:
+    """Restore the container's own resolver once the agent's turn is over.
+
+    Harbor grades a timed-out or crashed agent turn too, so a jail left standing would
+    reach the verifier and change what the task's own test suite can do.
+    """
+    if not dns_jail_enabled():
+        return
+    try:
+        # Keyed on the file the apply wrote, not on the baked script: a task image frozen
+        # before this feature has nothing to invoke, and must still end up unjailed.
+        await agent.exec_as_root(
+            environment,
+            command=(
+                f"if [ -s {STATE}/resolv.orig ]; then "
+                f"cat {STATE}/resolv.orig > /etc/resolv.conf; fi"
+            ),
+        )
+    except Exception as exc:
+        # Raising here would replace whatever ended the agent's turn with this.
+        _log.warning("DNS jail: could not lift before the verifier (%s)", exc)
+
+
+def jailed(run: Any) -> Any:
+    """Wrap an agent `run()` so the jail covers exactly the agent's turn.
+
+    A decorator rather than two calls in the body: the pair is what matters, and an
+    `apply` whose `lift` was forgotten looks like a working trial.
+    """
+
+    @functools.wraps(run)
+    async def wrapper(self, instruction, environment, context):
+        await apply_dns_jail(self, environment)
+        try:
+            return await run(self, instruction, environment, context)
+        finally:
+            await lift_dns_jail(self, environment)
+
+    return wrapper