datarobot_genai.dragent.frontends.stream_cancellation

stream_cancellation

Stop the agent run when a streaming client goes away.

NAT's generate_streaming_response runs the workflow in a separate producer task and hands its output to the HTTP response through a queue. When the client disconnects, Starlette cancels the response generator, but that generator's finally only closes the queue. The producer task is never cancelled. It keeps calling the LLM and tools until its next q.put raises QueueClosed, which can be one slow LLM call (tens of seconds) or more later.

While it lingers, the abandoned run overlaps the next request in the same worker. That overlap is what breaks:

• trace isolation: the abandoned run's spans finish late, outside its request, and land in the next request's export window (often with a parentless guard span);
• NAT's per-run span stack ("Current span ID stack is not equal to the previous stack");
• any framework state that is process-global rather than per-request (e.g. CrewAI's event bus scope).

:func:patch_nat_streaming_disconnect_cancellation swaps in a copy of NAT's function that cancels the producer when the consumer exits early, and waits briefly (shielded from the caller's cancellation) for it to unwind before the runner context closes.

The copy also keeps the NAT 1.7 wire format for our AG-UI events. NAT 1.8+ wraps every /generate/stream chunk in a {"value": "<json string>"} envelope, which AG-UI consumers of that route (our e2e helpers, the agent template's server) do not unwrap.

generate_streaming_response async

generate_streaming_response(payload: Any, *, session: Any, streaming: bool, step_adaptor: Any = None, result_type: type | None = None, output_type: type | None = None, wrap_output_in_payload: bool = False) -> AsyncGenerator[Any]

Copy of NAT 1.9's generate_streaming_response that cancels its producer on exit.

Two differences from NAT. When the consumer stops early (client disconnect, cancellation, aclose()), the producer task is cancelled and awaited instead of being left to run; the wait is shielded because the caller's cancel scope is already cancelled and would otherwise interrupt it at once. And DRAgentEventResponse chunks are never wrapped in the {"value": ...} envelope, so /generate/stream keeps its 1.7 shape.

Source code in datarobot_genai/dragent/frontends/stream_cancellation.py

async def generate_streaming_response(
    payload: Any,
    *,
    session: Any,
    streaming: bool,
    step_adaptor: Any = None,
    result_type: type | None = None,
    output_type: type | None = None,
    wrap_output_in_payload: bool = False,
) -> AsyncGenerator[Any]:
    """Copy of NAT 1.9's ``generate_streaming_response`` that cancels its producer on exit.

    Two differences from NAT. When the consumer stops early (client disconnect,
    cancellation, ``aclose()``), the producer task is cancelled and awaited instead of being
    left to run; the wait is shielded because the caller's cancel scope is already cancelled
    and would otherwise interrupt it at once. And ``DRAgentEventResponse`` chunks are never
    wrapped in the ``{"value": ...}`` envelope, so ``/generate/stream`` keeps its 1.7 shape.
    """

    def _wrap(item: Any) -> Any:
        if wrap_output_in_payload and not isinstance(item, DRAgentEventResponse):
            return ResponsePayloadOutput(payload=item)
        return item

    if step_adaptor is None:
        step_adaptor = StepAdaptor(StepAdaptorConfig())

    async with session.run(payload) as runner:
        q: AsyncIOProducerConsumerQueue[Any] = AsyncIOProducerConsumerQueue()

        intermediate_complete = await pull_intermediate(q, step_adaptor)

        async def pull_result() -> None:
            try:
                if session.workflow.has_streaming_output and streaming:
                    async for chunk in runner.result_stream(to_type=output_type):
                        await q.put(_wrap(chunk))
                else:
                    result = await runner.result(to_type=result_type)
                    await q.put(_wrap(runner.convert(result, output_type)))

                await intermediate_complete.wait()
            finally:
                await q.close()

        task: asyncio.Task[None] | None = None
        try:
            task = asyncio.create_task(pull_result())

            async for item in q:
                if isinstance(item, ResponseSerializable):
                    yield item
                else:
                    yield ResponsePayloadOutput(payload=item)

            # Re-raise any exception from the producer so callers can handle it
            await task
        finally:
            with anyio.CancelScope(shield=True):
                await _stop_producer(task)
                await q.close()

patch_nat_streaming_disconnect_cancellation

patch_nat_streaming_disconnect_cancellation() -> None

Make NAT's streaming responses cancel the workflow when the client goes away.

Idempotent. Skips (with a warning) when NAT's function no longer has the signature the copy above was written against.

Source code in datarobot_genai/dragent/frontends/stream_cancellation.py

def patch_nat_streaming_disconnect_cancellation() -> None:
    """Make NAT's streaming responses cancel the workflow when the client goes away.

    Idempotent. Skips (with a warning) when NAT's function no longer has the signature the
    copy above was written against.
    """
    if _PATCH_STATE["patched"]:
        return

    original = response_helpers.generate_streaming_response
    parameters = tuple(inspect.signature(original).parameters)
    if parameters != _EXPECTED_PARAMETERS:
        logger.warning(
            "NAT generate_streaming_response signature changed to %s; client disconnects "
            "will not cancel the agent run. Update %s for this NAT version.",
            parameters,
            __name__,
        )
        return

    for module_name in _PATCHED_MODULES:
        try:
            module = importlib.import_module(module_name)
        except Exception:
            logger.debug("Could not import %s for stream cancellation patch", module_name)
            continue
        if getattr(module, "generate_streaming_response", None) is original:
            module.generate_streaming_response = generate_streaming_response  # type: ignore[attr-defined]

    _PATCH_STATE["patched"] = True
    logger.debug("Patched NAT generate_streaming_response to cancel on client disconnect")