fix: inter-chunk stream timeout and download part guard (RFC)

[planetf1]

Closes #650. Addresses all three concerns in #1235.

Fixes #650
Fixes #1235

What changed

Streaming hang fix (helpers/async_helpers.py)

A stalled backend iterator previously blocked the process forever. send_to_queue had no per-chunk timeout, so its iteration loop waited indefinitely and no sentinel ever reached the queue. Every downstream consumer blocked with it.

This wraps each anext() call in asyncio.timeout(). A stalled stream now surfaces as a TimeoutError through the queue. On timeout, the underlying iterator's close method is called via duck-typing (getattr(ait, "aclose", None) or getattr(ait, "close", None)) to release HTTP connections cleanly — this covers both async def-style AsyncGenerator and class-based iterators such as OpenAI's AsyncStream and HF's AsyncTextIteratorStreamer.

The timeout context manager is captured (async with asyncio.timeout(...) as cm) and cm.expired() is checked before rewriting the error — so a TimeoutError raised by the backend iterator itself (e.g. an httpx read timeout) is forwarded verbatim rather than being misreported as a misconfigured STREAM_TIMEOUT.

DEFAULT_CHUNK_TIMEOUT = 120.0 is a public constant exported from helpers/. All eight backend call sites import and use it — no hardcoded values. The new ModelOption.STREAM_TIMEOUT key lets callers override per-call or at the backend/session level. chunk_timeout is keyword-only in the send_to_queue signature.

The 120 s default covers time-to-first-token as well as inter-chunk gaps. Slow local inference (large models on CPU, long prompts) can legitimately exceed this before the first token. Use a higher value or None for those deployments. Note: this timeout only applies to streaming responses; non-streaming coroutines bypass the per-chunk loop entirely and are unaffected.

Retriever download guard (formatters/granite/retrievers/util.py)

The parquet download loop now stops at _MAX_PARTS + 1 attempts. The extra iteration allows a corpus of exactly _MAX_PARTS parts to receive the terminating 404 rather than raising spuriously. The largest corpus currently has 20 parts; 50 gives headroom. _MAX_PARTS is a module-level constant (patchable in tests).

Documentation (docs/)

ModelOption reference table updated: STREAM_TIMEOUT description notes it covers TTFT and that non-streaming calls are unaffected. MAX_NEW_TOKENS row flags that backend defaults vary widely (vLLM defaults to 16). Both streaming how-to pages show three examples: tight timeout, slow-model value (300 s), and None to disable.

Tests

Four new cases in test/helpers/test_async_helpers.py: timeout fires and puts TimeoutError with no trailing sentinel; None disables the timeout and clean completion still gets a sentinel; DEFAULT_CHUNK_TIMEOUT is 120.0. One new case in test/formatters/granite/test_retrievers_util.py: _MAX_PARTS exhaustion raises RuntimeError.

Follow-up

#1242 tracks the cancel hook not firing for direct avalue()/astream() callers on HF STREAM_TIMEOUT — the worker thread keeps generating until natural completion on those paths. Marked with TODO(#1242) at the two _cancel_hook arming sites. The stream_with_chunking path already calls cancel_generation() and is not affected.

🤖 Generated with Claude Code

[ajbozarth]

A quick review from Claude to start:

Nice scoped fix — the streaming hang really did need bounding, and the API surface (single DEFAULT_CHUNK_TIMEOUT constant, single STREAM_TIMEOUT ModelOption, None to disable) is restrained. Tests pin the right invariants, especially that a timeout uses the exception itself as the stream terminator with no trailing sentinel.

Four action items inline. The (model_options or {}).get(...) consistency item and the aclose() reach-through are worth addressing in this PR; the cancel-hook miss can be a follow-up.

[ajbozarth]

LGTM on a code read through and Claude says so too

[planetf1]

@jakelorocco can you take a look - I need extra approval on this one?

[planetf1]

@AngeloDanducci @psschwei — this PR has been open for 10 days and has one approval. It needs a sign-off from the mellea-intrinsics team (CODEOWNERS for mellea/formatters/granite). Would appreciate a look when you can.

[AngeloDanducci]

I'm not in the mellea-intrinsics approver list - do you want additional reviews other than that? I saw one approve and considering my review is non blocking I did not follow up.

[frreiss]

The core issue being addressed, #650, appears to involve the system operating correctly as designed.

My understanding of the scenario is:

• A test case issues a combination of requests that triggers a bug in ollama.
• The bug causes the ollama server to hang.
• The frozen server causes the test case to freeze waiting for a response.
• The test harness correctly identifies that the test has frozen and correctly fails the session.

The proposed solution for this issue involves adding a timeout to some backend operations but not to others. In the event that an inference server becomes unresponsive, the new timeout logic may or may not trigger. If the timeout logic does trigger, it will cause the entire application to fail with a stack trace involving multiple threads and code from obscure asyncio classes. This failure mode does not seem like an improvement to me.

Knowledgeable users could make their code catch TimeoutError to prevent the entire application from crashing due to a timeout. However, a hypothetical user who adds a try-except block to catch TimeoutError could just as easily add their own call to asyncio.timeout() at a location that is more application-appropriate than deep inside certain backend streaming operations.

Claude's flagging of the in mellea/formatters/granite/retrievers/util.py is a false positive. That file-downloading code is specfically related to the MTRAG benchmark data set. This code would only perform unbounded amounts of downloads if the MTRAG benchmark were to become infinitely large. That said, Claude's fix doesn't appear to do any harm.

[planetf1]

Thanks for the detailed review — let me take the points in turn.

On #650 being "correct as designed": You're right that the test harness caught the hang — but the fix is aimed at production use, not tests. A user streaming from a stalled backend has nothing watching them; today they hang indefinitely with no error and no way out. It doesn't have to be an ollama bug — any unresponsive server produces the same result.

On "some backends but not others": The timeout lives in send_to_queue rather than in each backend — every streaming path goes through that function. The backend changes just wire the argument through. All five streaming backends are covered (huggingface, openai, ollama, litellm, watsonx); bedrock and dummy don't implement streaming. The one genuine gap is the cancel-hook cleanup on timeout, which is already split out as #1242.

On the asyncio stack trace: The asyncio TimeoutError is caught inside send_to_queue and discarded — a fresh one is constructed before it ever leaves the function, so there's no chained __context__ from the internal cancellation machinery. Here's the actual uncaught traceback:

Traceback (most recent call last):
  File "myapp.py", line 23, in <module>
    asyncio.run(main())
  File ".../asyncio/runners.py", line 194, in run
    return runner.run(main)
  File ".../asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
  File ".../asyncio/base_events.py", line 686, in run_until_complete
    return future.result()
  File "myapp.py", line 21, in main
    raise item2
TimeoutError: Stream timed out after 60s without a chunk (covers time-to-first-token and inter-chunk gaps). Set ModelOption.STREAM_TIMEOUT to a larger value or None to disable.

The asyncio.run frames at the top appear in any unhandled async exception — they're not specific to this timeout. The raise is at the caller's own site and the message is the actionable one. No CancelledError chain, no task cancellation internals. If you've seen a different trace in practice I'd want to fix it.

On "users could add their own asyncio.timeout()": Fair point, and STREAM_TIMEOUT=None is there for anyone who wants the old behaviour. But there's a practical difference: a user-level asyncio.timeout() budgets the entire response, which would cut off legitimate long outputs. The per-chunk timeout only triggers when the backend goes completely silent — a slow model that's actively generating is unaffected however long it runs. That's hard to replicate cleanly from outside the library.

On the retriever guard: You're right — it's your repo so you know the data. The while True looked unbounded in isolation. Happy to revert it; since you say it does no harm I'll leave the call to you.

[jakelorocco]

I agree that a timeout is needed / nice. I think the documentation / comments in relation to this code are misleading though. All generation uses the send_to_queue function; it doesn't matter if STREAM==True.

I think the default value of 60s is likely too low. I think it should either default to None or be some very high value. A large prompt to a sluggish / overloaded inference provider can easily take more time. In the non-streaming case, the client will have to wait for the whole response back while fitting within this timeout window.

To educate our users, I would expect most examples to have some sort of try-catch block for these TimeoutErrors as well. For that reason, I think it makes more sense for end users to opt-in to a reasonable timeout.

[planetf1]

The send_to_queue observation is correct — it's called for all backends regardless of STREAM. However, the timeout itself only fires inside the isinstance(aresponse, AsyncIterator) branch, so non-streaming paths (where the coroutine resolves to a plain response object) skip the per-chunk loop entirely. The STREAM_TIMEOUT name and docs are scoped accordingly, but you're right that neither explicitly says "this is a no-op for non-streaming callers" — worth making that clearer in the docstring.

The non-streaming point in your second concern doesn't apply for the same reason. The streaming concern is fair though: the same window covers TTFT as well as inter-chunk gaps, and 60 s can be tight for a large model loading locally before the first token arrives. Bumping the default to 120 s seems like a reasonable middle ground — still bounded (an unbounded default feels risky), but less likely to surprise users on slower hardware. On the TTFT vs inter-chunk distinction: they do have different profiles and a split timeout would give finer control — worth a follow-up issue if that's of interest?

On error handling in examples — agreed that showing what to do when a timeout fires would be useful. We could update the existing streaming examples to include a try-catch, though a dedicated error-handling example might cover it better. Worth a follow-up issue?