patch_with_retry drops the original exception when it retries an async-generator (streaming) call

Version

1.9.0

Which installation method(s) does this occur on?

PyPi

Describe the bug.

patch_with_retry wraps async-generator methods in _agen_with_retry. When such a call fails and the wrapper retries it, the retry re-invokes the wrapped method with the same arguments. If any argument is a one-shot async iterator, which is exactly how LangChain chains streaming steps, the replay consumes an already exhausted iterator: the upstream call is never made again, the generator yields nothing and hits the return inside the loop, and the caught exception is discarded. last_exception is never re-raised, because return leaves the loop before the trailing if last_exception: raise last_exception.

The consumer sees a stream that ended cleanly with zero chunks. The provider error is gone.

Output, identical on 1.8.0 and 1.9.0:

retries=1 is the only configuration that behaves correctly, and only because the loop's last attempt raises before the replay can happen.

The same thing through the real stack

With a ChatOpenAI pointed at a local HTTP server that answers 429, consumed through ToolCallAgentGraph:

Note the HTTP column: for a streaming call the wrapper never retries anything. num_retries buys no resilience here, it only decides whether the error survives.

Impact

ToolCallAgentGraph._invoke_llm raises RuntimeError('No response received from agent') when the stream produced zero chunks, so every retryable provider failure on a streaming call arrives as that one message. A caller cannot tell a rate limit from a genuinely empty completion, and cannot back off, fall back to another deployment, or tell the user what happened.

In our product this reached end users as "the response failed with an internal error" while the deployment was simply out of TPM quota (429, exceeded rate limit), and our own rate-limit handling never fired.

Expected

A failed streaming call surfaces its exception. Either re-raise last_exception when a retry attempt completes without producing anything, or do not retry async-generator methods at all (raise immediately, as retries=1 effectively does today), or refuse to replay arguments that cannot be replayed.

Environment

• nvidia-nat 1.9.0 (latest) and 1.8.0, both installed from PyPI. The reproduction above prints the same output on both.

• nvidia-nat-core, nvidia-nat-langchain

• langchain-core 1.6.2, langchain-openai 1.2.1

• Python 3.13.13 and 3.12.10, Windows

Workaround

num_retries: 1 on every streaming LLM config. Note that the documented switch, do_auto_retry: false, does not help today because of #2212.

Possibly related

• #2177 (_in_retry_context races across concurrent requests), same function, different failure.

• #2212 (do_auto_retry: false ignored by the adapters), with #2215 open as its fix.

Minimum reproducible example

import asyncio
from nat.utils.exception_handlers.automatic_retries import patch_with_retry
class Boom(Exception):
    def __init__(self) -> None:
        super().__init__("Error code: 429 - rate limit exceeded")
        self.status_code = 429
class Service:
    """A streaming call whose input is a one-shot async iterator.
    That is how LangChain chains streaming steps: every step receives the
    previous step's async iterator, and it can only be consumed once.
    """
    def __init__(self) -> None:
        self.attempts = 0
    async def astream(self, source):
        async for item in source:
            self.attempts += 1
            raise Boom()
            yield item  # unreachable; makes this an async generator function
        # A second attempt lands here: the iterator is spent, so this produces
        # nothing and raises nothing.
async def one_shot_input():
    yield "prompt"
async def main() -> None:
    for retries in (1, 2, 5):
        service = Service()
        patched = patch_with_retry(
            service, retries=retries, retry_codes=[429], retry_on_messages=None
        )
        chunks = []
        try:
            async for chunk in patched.astream(one_shot_input()):
                chunks.append(chunk)
            print(f"retries={retries}: NO exception, {len(chunks)} chunk(s), "
                  f"upstream attempts={service.attempts}   <-- 429 lost")
        except Boom:
            print(f"retries={retries}: Boom raised, upstream attempts={service.attempts}")
asyncio.run(main())

Relevant log output

[Paste the error here, it will be hidden by default]

Other/Misc.

No response

Code of Conduct

• I agree to follow the NeMo Agent Toolkit Code of Conduct
• I have searched the open bugs and have found no duplicates for this bug report

Start by tracing patch_with_retry into _agen_with_retry and reproduce a failed async-generator call with a one-shot async iterator. Verify that the original provider exception is surfaced instead of the stream ending with zero chunks, and add or update coverage for the retry path.