fix: continue recovered streams in a new segment

maintainer edit: streamed timeout recovery was returning the retried response internally while the channel still treated the final outbound as already streamed. End the current stream segment before retry/fallback recovery so subsequent deltas are delivered in a new segment.
This commit is contained in:
chengyongru
2026-06-10 18:10:44 +08:00
committed by Xubin Ren
parent 2c5a4e0703
commit bc4bb508a1
6 changed files with 162 additions and 13 deletions
+4
View File
@@ -754,11 +754,15 @@ class AgentRunner:
context.streamed_reasoning = True
await hook.emit_reasoning(delta)
async def _stream_recover() -> None:
await hook.on_stream_end(context, resuming=True)
coro = self.provider.chat_stream_with_retry(
**kwargs,
on_content_delta=_stream,
on_thinking_delta=_thinking,
on_tool_call_delta=_tool_call_delta if live_file_edits is not None else None,
on_stream_recover=_stream_recover,
)
elif wants_progress_streaming:
stream_buf = ""
+27 -9
View File
@@ -631,6 +631,7 @@ class LLMProvider(ABC):
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_stream_recover: Callable[[], Awaitable[None]] | None = None,
retry_mode: str = "standard",
on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
) -> LLMResponse:
@@ -651,6 +652,12 @@ class LLMProvider(ABC):
if on_content_delta:
await on_content_delta(text)
async def _recover_stream() -> None:
nonlocal has_streamed_content
if on_stream_recover:
await on_stream_recover()
has_streamed_content = False
kw: dict[str, Any] = dict(
messages=messages, tools=tools, model=model,
max_tokens=max_tokens, temperature=temperature,
@@ -659,6 +666,8 @@ class LLMProvider(ABC):
on_thinking_delta=on_thinking_delta,
on_tool_call_delta=on_tool_call_delta,
)
if on_stream_recover and getattr(self, "supports_stream_recover_callback", False):
kw["on_stream_recover"] = _recover_stream
return await self._run_with_retry(
self._safe_chat_stream,
kw,
@@ -666,6 +675,7 @@ class LLMProvider(ABC):
retry_mode=retry_mode,
on_retry_wait=on_retry_wait,
should_retry_guard=lambda: not has_streamed_content,
on_stream_recover=_recover_stream if on_stream_recover else None,
)
async def chat_with_retry(
@@ -813,6 +823,7 @@ class LLMProvider(ABC):
retry_mode: str,
on_retry_wait: Callable[[str], Awaitable[None]] | None,
should_retry_guard: Callable[[], bool] | None = None,
on_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse:
attempt = 0
delays = list(self._CHAT_RETRY_DELAYS)
@@ -829,15 +840,22 @@ class LLMProvider(ABC):
if should_retry_guard is not None and not should_retry_guard():
is_timeout = (response.error_kind or "").lower() == "timeout"
if is_timeout:
logger.warning(
"LLM stream stalled after content was emitted; "
"suppressing delta callbacks and retrying"
)
kw.setdefault("on_content_delta", None)
kw["on_content_delta"] = None
kw["on_thinking_delta"] = None
kw["on_tool_call_delta"] = None
should_retry_guard = None
if on_stream_recover:
logger.warning(
"LLM stream stalled after content was emitted; "
"starting a new stream segment and retrying"
)
await on_stream_recover()
else:
logger.warning(
"LLM stream stalled after content was emitted; "
"suppressing delta callbacks and retrying"
)
kw.setdefault("on_content_delta", None)
kw["on_content_delta"] = None
kw["on_thinking_delta"] = None
kw["on_tool_call_delta"] = None
should_retry_guard = None
else:
logger.warning(
"LLM stream failed after content was emitted; skipping retry"
+26 -3
View File
@@ -71,6 +71,8 @@ class FallbackProvider(LLMProvider):
wasting requests on a known-bad endpoint.
"""
supports_stream_recover_callback = True
def __init__(
self,
primary: LLMProvider,
@@ -116,6 +118,7 @@ class FallbackProvider(LLMProvider):
)
async def chat_stream(self, **kwargs: Any) -> LLMResponse:
on_stream_recover = kwargs.pop("on_stream_recover", None)
if not self._has_fallbacks:
return await self._primary.chat_stream(**kwargs)
@@ -130,7 +133,10 @@ class FallbackProvider(LLMProvider):
kwargs["on_content_delta"] = _tracking_delta
return await self._try_with_fallback(
lambda p, kw: p.chat_stream(**kw), kwargs, has_streamed=has_streamed
lambda p, kw: p.chat_stream(**kw),
kwargs,
has_streamed=has_streamed,
on_stream_recover=on_stream_recover,
)
async def _try_with_fallback(
@@ -138,6 +144,7 @@ class FallbackProvider(LLMProvider):
call: Callable[[LLMProvider, dict[str, Any]], Awaitable[LLMResponse]],
kwargs: dict[str, Any],
has_streamed: list[bool] | None,
on_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse:
primary_model = kwargs.get("model") or self._primary.get_default_model()
@@ -157,7 +164,10 @@ class FallbackProvider(LLMProvider):
primary_model,
)
has_streamed[0] = False
kwargs["on_content_delta"] = None
if on_stream_recover:
await on_stream_recover()
else:
kwargs["on_content_delta"] = None
else:
logger.warning(
"Primary model error but content already streamed; skipping failover"
@@ -187,7 +197,20 @@ class FallbackProvider(LLMProvider):
for idx, fallback in enumerate(self._fallback_presets):
fallback_model = fallback.model
if has_streamed is not None and has_streamed[0]:
break
is_timeout = (
last_response is not None
and (last_response.error_kind or "").lower() == "timeout"
)
if is_timeout and on_stream_recover:
logger.warning(
"Fallback model '{}' stream stalled after content was emitted; "
"starting a new stream segment and trying next fallback",
self._fallback_presets[idx - 1].model if idx > 0 else primary_model,
)
has_streamed[0] = False
await on_stream_recover()
else:
break
if idx == 0 and primary_skipped:
logger.info(
"Primary model '{}' circuit open, trying fallback '{}'",