21 · 异步能力检查与协议边界

def _accepted_kwargs(function, kwargs: dict[str, Any]) -> dict[str, Any]:
    signature = inspect.signature(function)
    if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()):
        return kwargs
    return {key: value for key, value in kwargs.items() if key in signature.parameters}


async def _apply_to_sample(args, sample: Sample, paths: list[str], **kwargs) -> Sample:
    for path in paths:
        hook = load_function(path)
        result = hook(args, sample, **_accepted_kwargs(hook, kwargs))
        if inspect.isawaitable(result):
            result = await result
        if result is not None:
            if not isinstance(result, Sample):
                raise TypeError(
                    f"Rollout sample hook {path!r} returned {type(result).__name__}, expected Sample or None."
                )
            sample = result
    return sample

def trace_function(
    name: str,
    *,
    target: str | None = None,
    target_getter: Callable[..., Sample | TraceHandle | list[Sample | TraceHandle] | None] | None = None,
    attrs_getter: Callable[..., dict[str, Any] | None] | None = None,
    record_error: bool = True,
):
    def decorator(fn):
        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args, **kwargs):
                resolved_target = _resolve_trace_function_target(
                    fn,
                    args,
                    kwargs,
                    target=target,
                    target_getter=target_getter,
                )
                if resolved_target is None:
                    return await fn(*args, **kwargs)
                attrs = _resolve_trace_function_attrs(fn, args, kwargs, attrs_getter=attrs_getter)
                with trace_span(resolved_target, name, attrs=attrs, record_error=record_error):
                    return await fn(*args, **kwargs)

            return async_wrapper

        @functools.wraps(fn)
        def sync_wrapper(*args, **kwargs):
            resolved_target = _resolve_trace_function_target(
                fn,
                args,
                kwargs,
                target=target,
                target_getter=target_getter,
            )
            if resolved_target is None:
                return fn(*args, **kwargs)
            attrs = _resolve_trace_function_attrs(fn, args, kwargs, attrs_getter=attrs_getter)
            with trace_span(resolved_target, name, attrs=attrs, record_error=record_error):
                return fn(*args, **kwargs)

        return sync_wrapper

    return decorator

async def _run_request_abortable_generate(
    state: GenerateState,
    sample: Sample,
    generate_call: Awaitable[Sample | list[Sample]],
) -> Sample | list[Sample]:
    task = asyncio.current_task()
    assert task is not None
    state.cancellable_tasks.add(task)
    try:
        return await generate_call
    except asyncio.CancelledError:
        if task in state.cancellable_tasks:
            raise
        sample.status = Sample.Status.ABORTED
        return sample
    finally:
        state.cancellable_tasks.discard(task)

@asynccontextmanager
async def boot_agent_sandbox(image: str, instance_id: str) -> AsyncIterator[E2BSandbox]:
    """Boot a fresh E2B sandbox and install the selected harness toolchain.

    Create the sandbox from the dataset image, install Node 22 + the harness CLI
    from host tarballs, retry transient boot/install failures, and close the
    sandbox when the caller leaves the context.
    """
    sb = None

    if sb is None:
        assert last_err is not None
        raise last_err
    try:
        yield sb
    finally:
        await sb.__aexit__(None, None, None)