Walkthrough

env = Environment("robolab", capabilities=[Capability.ssh(name="shell", url=...)])

@env.template()
async def fix_bug(difficulty: int):
    answer = yield f"Fix the bug (difficulty {difficulty})"   # 1st yield -> prompt
    yield 1.0 if check(answer) else 0.0                       # 2nd yield -> score

def decorate(func):
    if not inspect.isasyncgenfunction(func):
        raise TypeError(...)                 # must be `async def ... yield`
    task = _TaskFactory(self, id or func.__name__, description, func, ...)
    self.tasks[task_id] = task               # registered on the env
    return task

def __call__(self, *args, **kwargs) -> EvalTask:
    from hud.eval.task import Task
    bound = self.sig.bind(*args, **kwargs)
    return Task(env=self.env.name, id=self.id, args=dict(bound.arguments))

async def serve(env, host="127.0.0.1", port=0):
    await env.start()                              # run @env.initialize hooks
    server = await bind(env, host, port)          # bind the control channel
    print(f"HUD_SERVE_PORT={server.sockets[0].getsockname()[1]}", flush=True)  # announce it
    try:
        async with server:
            await server.serve_forever()
    finally:
        await env.stop()                          # @env.shutdown hooks

async def start(self):
    if self._started:
        return
    self._started = True
    for hook in self._on_start:                     # @env.initialize
        await hook()                                # daemons up, add_capability(...) called here
    self._hooks_done = True

started = env.start()
await asyncio.wait_for(started, ready_timeout)     # @env.initialize, bounded
server = await bind(env, "127.0.0.1", 0)           # in-process control channel
host, port = server.sockets[0].getsockname()[:2]
serve_task = asyncio.create_task(server.serve_forever())
try:
    yield Runtime(f"tcp://{host}:{port}")          # <- this is `addr` back in rollout
finally:
    serve_task.cancel()
    # ...
    await env.stop()                               # @env.shutdown

@asynccontextmanager
async def connect(runtime, *, ready_timeout=240.0):
    parts = urlsplit(runtime.url)
    client = await _connect_ready(parts.hostname, parts.port, ready_timeout=...)
    try:
        yield client
    finally:
        await client.close()

while True:
    try:
        reader, writer = await asyncio.open_connection(host, port)
    except OSError:
        ...; continue                               # not accepting yet, retry
    client = HudClient(reader, writer, endpoint=(host, port))
    try:
        await client.hello()                        # handshake
    except (EOFError, OSError):
        await client.close(); ...                   # accepted but no env behind it, retry
    else:
        return client

async def hello(self, session_id=None):
    params = {} if session_id is None else {"session_id": session_id}
    result = await self._call("hello", params)      # -> server session
    bindings = [Capability.from_manifest(b) for b in result["bindings"]]
    self.manifest = Manifest(session_id=result["session_id"], bindings=bindings, ...)
    for capability in bindings:
        forwarder = await asyncio.start_server(
            partial(forward, capability), "127.0.0.1", 0
        )
        self._forwarders.append(forwarder)
        self._routes[capability.name] = ...          # binding() returns the local route
    return self.manifest

if method == "hello":
    # ... a requested session_id resumes that parked session instead ...
    bindings = [c.to_manifest() for c in env.capabilities]
    await reply_to(msg_id, {"session_id": session_id,
                            "env": {"name": env.name, "version": env.version},
                            "bindings": bindings})

async def __aenter__(self):
    started = await self.client.start_task(self._task_id, self._args)   # tasks.start
    self.prompt = started.get("prompt")
    self.record(Step(source="task", task_call=TaskCall(phase="setup", ...)))
    if self.prompt is not None:
        self.record(Step(source="user", messages=self.prompt_messages))
    return self

async def start_task(self, task_id, args=None):
    return await self._call("tasks.start", {"id": task_id, "args": args or {}})

elif method == "tasks.start":
    prompt = await self.start(session_id, task_id, args)   # _ControlChannel.start
    await reply_to(msg_id, prompt)

async def start(self, session_id, task_id, args):   # _ControlChannel.start
    await self.cancel(session_id)                    # replace this session's own task, if any
    self._runners[session_id] = TaskRunner(self.env.tasks[task_id], args)
    return await self._runners[session_id].start()

async def start(self):
    self._gen = self.task.func(**_coerce_args(self.task.sig, self._args))   # make the generator
    prompt = await self._gen.__anext__()            # run task fn to 1st yield
    frame = prompt if isinstance(prompt, dict) and "prompt" in prompt else {"prompt": prompt}
    return _jsonable(frame)                          # -> back over the wire

run = live
async with file_tracking_observer(client):
    await agent(run)                    # Agent.__call__(run)

class Agent(ABC):
    @abstractmethod
    async def __call__(self, run: Run) -> None:
        """Drive run to completion, filling run.trace (answer is trace.content)."""

async def __aexit__(self, exc_type, exc, tb):
    if exc_type in (asyncio.CancelledError, KeyboardInterrupt):
        await self.client.cancel(); return False
    answer = {"answer": self.trace.content}
    evaluation = await self.client.grade(answer)    # tasks.grade
    self.grade = Grade.from_dict(evaluation)        # score -> run.grade.reward
    self.record(Step(source="task", task_call=TaskCall(phase="evaluate", ...)))
    return False

elif method == "tasks.grade":
    evaluation = await self.grade(session_id, params)   # _ControlChannel.grade
    await reply_to(msg_id, evaluation)

async def grade(self, session_id, payload):          # _ControlChannel.grade
    runner = self._runners.pop(session_id, None)      # this session's own task ...
    if runner is None:
        runner = self._adopt_parked()                 # ... or the lone parked one
    return await runner.grade(payload)

async def grade(self, payload):
    evaluation = await self._gen.asend(_build_answer(self.task.return_type, payload))  # 2nd yield
    # ... normalize dict / EvaluationResult ...
    return {"score": _score_value(evaluation)}