fix(agent): count only substantive auxiliary progress
Salvages the stream-progress fix from #80122 on current main while incorporating review feedback for empty provider deltas and tool-call scaffolding.
This commit is contained in:
StanleyStetson
+105
-15
@@ -437,7 +437,7 @@ class _AuxiliaryCancellationDecision:
|
||||
# deadline punishes SLOW summary models exactly as hard as HUNG ones: a
|
||||
# reasoning model happily streaming a large summary is killed mid-generation.
|
||||
# This thread-local hook lets the host observe liveness instead: the wire
|
||||
# consumers below tick it on every streamed token/SSE event, and the host
|
||||
# consumers below tick it only for non-empty streamed payloads, and the host
|
||||
# extends its deadline while tokens are moving (see gateway/run.py session
|
||||
# hygiene + CompressionCommitFence.touch_progress). Thread-local matches the
|
||||
# call topology — the aux call and its stream consumption run synchronously
|
||||
@@ -468,14 +468,26 @@ def _notify_aux_dispatch() -> None:
|
||||
logger.debug("aux dispatch hook failed", exc_info=True)
|
||||
|
||||
|
||||
def _notify_aux_provider_response() -> None:
|
||||
"""Record a provider response/chunk, then preserve the liveness signal."""
|
||||
def _notify_aux_timing_response() -> None:
|
||||
"""Record a provider response/chunk WITHOUT claiming forward progress.
|
||||
|
||||
Same timing slot as :func:`_notify_aux_provider_response`, minus the
|
||||
forward-progress chain: used for content-free frames (keepalives,
|
||||
lifecycle events, typed-but-empty deltas) that must still count toward
|
||||
``time_to_first_progress_ms`` telemetry but must not reset a compression
|
||||
inactivity fence.
|
||||
"""
|
||||
hook = getattr(_aux_provider_response, "hook", None)
|
||||
if hook is not None:
|
||||
try:
|
||||
hook()
|
||||
except Exception:
|
||||
logger.debug("aux provider response hook failed", exc_info=True)
|
||||
|
||||
|
||||
def _notify_aux_provider_response() -> None:
|
||||
"""Record a provider response/chunk, then preserve the liveness signal."""
|
||||
_notify_aux_timing_response()
|
||||
_notify_aux_progress()
|
||||
|
||||
|
||||
@@ -483,6 +495,55 @@ def _aux_progress_active() -> bool:
|
||||
return getattr(_aux_progress, "hook", None) is not None
|
||||
|
||||
|
||||
def _event_field(event: Any, name: str) -> Any:
|
||||
if isinstance(event, dict):
|
||||
return event.get(name)
|
||||
return getattr(event, name, None)
|
||||
|
||||
|
||||
def _anthropic_event_has_content(event: Any) -> bool:
|
||||
"""Whether an Anthropic stream event carries a non-empty payload."""
|
||||
event_type = _event_field(event, "type")
|
||||
if event_type == "content_block_delta":
|
||||
delta = _event_field(event, "delta")
|
||||
return any(
|
||||
bool(_event_field(delta, field))
|
||||
for field in ("text", "thinking", "partial_json")
|
||||
)
|
||||
if event_type == "content_block_start":
|
||||
block = _event_field(event, "content_block")
|
||||
return _event_field(block, "type") == "tool_use" and any(
|
||||
bool(_event_field(block, field)) for field in ("id", "name")
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
_CODEX_PROGRESS_DELTA_TYPES = frozenset(
|
||||
{
|
||||
"response.output_text.delta",
|
||||
"response.reasoning_summary_text.delta",
|
||||
"response.text.delta",
|
||||
"response.audio.delta",
|
||||
"response.function_call_arguments.delta",
|
||||
"response.reasoning_text.delta",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _codex_event_has_content(event: Any) -> bool:
|
||||
"""Whether a Codex Responses event carries a non-empty payload."""
|
||||
event_type = _event_field(event, "type")
|
||||
if event_type in _CODEX_PROGRESS_DELTA_TYPES:
|
||||
return bool(_event_field(event, "delta"))
|
||||
if event_type == "response.output_item.added":
|
||||
item = _event_field(event, "item")
|
||||
return "function_call" in str(_event_field(item, "type") or "") and any(
|
||||
bool(_event_field(item, field))
|
||||
for field in ("id", "call_id", "name", "arguments")
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _aux_thread_local_hook(local: threading.local, hook):
|
||||
"""Install one thread-local hook callback and restore its prior value.
|
||||
@@ -1906,10 +1967,15 @@ class _CodexCompletionsAdapter:
|
||||
def _on_each_event(_event: Any) -> None:
|
||||
# Re-check timeout/cancellation per event, matching the
|
||||
# cadence the old in-line ``_check_cancelled()`` used.
|
||||
# Each SSE event is also forward progress for hosts watching
|
||||
# a progress hook (gateway session hygiene): a reasoning
|
||||
# model streaming a long summary must not look hung.
|
||||
_notify_aux_provider_response()
|
||||
# Provider response timing (TTFP telemetry) records every
|
||||
# frame; forward progress for hosts watching liveness (the
|
||||
# compression commit fence) counts only substantive
|
||||
# payloads — lifecycle and keepalive events must not reset
|
||||
# the compression idle clock.
|
||||
if _codex_event_has_content(_event):
|
||||
_notify_aux_provider_response()
|
||||
else:
|
||||
_notify_aux_timing_response()
|
||||
_check_cancelled()
|
||||
|
||||
event_stream = self._client.responses.create(**stream_kwargs)
|
||||
@@ -2268,8 +2334,15 @@ class _AnthropicCompletionsAdapter:
|
||||
# slow-but-generating summary model. No-op when no hook is
|
||||
# installed (None keeps the fast get_final_message path).
|
||||
on_stream_event=(
|
||||
(lambda _event: _notify_aux_provider_response())
|
||||
if _aux_progress_active() else None
|
||||
(
|
||||
lambda event: (
|
||||
_notify_aux_provider_response()
|
||||
if _anthropic_event_has_content(event)
|
||||
else _notify_aux_timing_response()
|
||||
)
|
||||
)
|
||||
if _aux_progress_active()
|
||||
else None
|
||||
),
|
||||
)
|
||||
_transport = get_transport("anthropic_messages")
|
||||
@@ -9353,9 +9426,9 @@ def _create_with_progress(
|
||||
neither trigger applies (every existing caller/task) or when the client's
|
||||
wire adapter streams internally. With a hook + a chunk-capable client,
|
||||
the request is sent with ``stream=True`` and aggregated, ticking the hook
|
||||
per chunk — so the configured ``timeout`` acts per stream read (idle)
|
||||
rather than as a total budget, and outer liveness watchdogs see tokens
|
||||
moving. ``force_stream=True`` (stream-only providers such as Tencent
|
||||
only for substantive chunks. The configured ``timeout`` acts per stream
|
||||
read (idle) rather than as a total budget, and outer liveness watchdogs see
|
||||
tokens moving. ``force_stream=True`` (stream-only providers such as Tencent
|
||||
Copilot — credit @kudi88, PR #60686) takes the same streamed path even
|
||||
without a hook. Providers that reject the streamed request fall back to
|
||||
the plain non-streaming call — except under ``force_stream``, where a
|
||||
@@ -9403,7 +9476,9 @@ def _create_with_progress(
|
||||
return response
|
||||
|
||||
# Some shims (MoA virtual provider under quiet mode, defensive adapters)
|
||||
# return a complete response even when stream=True was requested.
|
||||
# return a complete response even when stream=True was requested. A
|
||||
# complete response object carries the full summary payload, so it counts
|
||||
# as provider response progress (TTFP) and forward progress alike.
|
||||
if hasattr(chunks, "choices"):
|
||||
_notify_aux_provider_response()
|
||||
return chunks
|
||||
@@ -9420,7 +9495,8 @@ def _aggregate_chat_stream(
|
||||
) -> Any:
|
||||
"""Consume a chat.completions chunk stream into a complete response.
|
||||
|
||||
Ticks the thread-local aux progress hook on every chunk. Raises
|
||||
Ticks the thread-local aux progress hook only for non-empty content,
|
||||
reasoning, or tool-call fragments. Raises
|
||||
TimeoutError when *total_ceiling* seconds elapse before the stream
|
||||
finishes — phrased with "timed out" so existing timeout classification
|
||||
(``_is_timeout_error``) treats it exactly like a request timeout.
|
||||
@@ -9461,7 +9537,11 @@ class _ChatStreamAccumulator:
|
||||
self.resp_model = model or ""
|
||||
|
||||
def feed(self, chunk: Any) -> None:
|
||||
_notify_aux_provider_response()
|
||||
# Every provider frame records transport-level timing (TTFP
|
||||
# telemetry, first-frame-wins); only a substantive payload below
|
||||
# ticks the forward-progress hook that keeps compression alive.
|
||||
_notify_aux_timing_response()
|
||||
made_progress = False
|
||||
if (
|
||||
self._total_ceiling is not None
|
||||
and (time.monotonic() - self._started) >= self._total_ceiling
|
||||
@@ -9486,25 +9566,35 @@ class _ChatStreamAccumulator:
|
||||
piece = getattr(delta, "content", None)
|
||||
if piece:
|
||||
self.content_parts.append(piece)
|
||||
made_progress = True
|
||||
reasoning_piece = (
|
||||
getattr(delta, "reasoning", None)
|
||||
or getattr(delta, "reasoning_content", None)
|
||||
)
|
||||
if reasoning_piece and isinstance(reasoning_piece, str):
|
||||
self.reasoning_parts.append(reasoning_piece)
|
||||
made_progress = True
|
||||
for tc in (getattr(delta, "tool_calls", None) or []):
|
||||
idx = getattr(tc, "index", 0) or 0
|
||||
acc = self.tool_calls_acc.setdefault(
|
||||
idx, {"id": "", "name": "", "arguments": []}
|
||||
)
|
||||
tool_fragment = False
|
||||
if getattr(tc, "id", None):
|
||||
acc["id"] = tc.id
|
||||
tool_fragment = True
|
||||
fn = getattr(tc, "function", None)
|
||||
if fn is not None:
|
||||
if getattr(fn, "name", None):
|
||||
acc["name"] = fn.name
|
||||
tool_fragment = True
|
||||
if getattr(fn, "arguments", None):
|
||||
acc["arguments"].append(fn.arguments)
|
||||
tool_fragment = True
|
||||
made_progress = made_progress or tool_fragment
|
||||
|
||||
if made_progress:
|
||||
_notify_aux_progress()
|
||||
|
||||
def finish(self) -> Any:
|
||||
tool_calls = None
|
||||
|
||||
@@ -2,9 +2,9 @@
|
||||
|
||||
Slow summary models must not be punished like hung ones (#see PR): when a
|
||||
forward-progress hook is installed (context compression), the primary
|
||||
auxiliary call streams and ticks the hook per chunk, so outer watchdogs
|
||||
(gateway session hygiene) can extend their deadline on liveness. Without a
|
||||
hook, behavior is byte-for-byte the old non-streaming call.
|
||||
auxiliary call streams and ticks the hook only for substantive payloads, so
|
||||
outer watchdogs (gateway session hygiene) can extend their deadline on
|
||||
liveness. Without a hook, behavior is byte-for-byte the old non-streaming call.
|
||||
"""
|
||||
|
||||
import threading
|
||||
@@ -15,10 +15,15 @@ from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
|
||||
from agent.auxiliary_client import (
|
||||
_AnthropicCompletionsAdapter,
|
||||
_ChatStreamAccumulator,
|
||||
_CodexCompletionsAdapter,
|
||||
_acreate_with_stream,
|
||||
_aggregate_chat_stream,
|
||||
_aggregate_chat_stream_async,
|
||||
_anthropic_event_has_content,
|
||||
_aux_stream_total_ceiling,
|
||||
_codex_event_has_content,
|
||||
_create_with_progress,
|
||||
_notify_aux_progress,
|
||||
_provider_requires_stream,
|
||||
@@ -112,8 +117,15 @@ class TestAuxProgressHook:
|
||||
|
||||
class TestCreateWithProgress:
|
||||
|
||||
def test_hook_upgrades_to_streaming_and_ticks_per_chunk(self):
|
||||
def test_hook_upgrades_to_streaming_and_ticks_only_for_payload(self):
|
||||
empty_tool_call = SimpleNamespace(
|
||||
index=0,
|
||||
id=None,
|
||||
function=SimpleNamespace(name=None, arguments=""),
|
||||
)
|
||||
chunks = [
|
||||
SimpleNamespace(id=None, model=None, choices=[], usage=None),
|
||||
_chunk(content="", reasoning="", tool_calls=[empty_tool_call]),
|
||||
_chunk(reasoning="thinking..."),
|
||||
_chunk(content="Hello "),
|
||||
_chunk(content="world", finish_reason="stop",
|
||||
@@ -131,8 +143,26 @@ class TestCreateWithProgress:
|
||||
assert result.choices[0].message.reasoning == "thinking..."
|
||||
assert result.choices[0].finish_reason == "stop"
|
||||
assert result.usage.total_tokens == 7
|
||||
# 1 dispatch tick + 1 per chunk
|
||||
assert len(ticks) >= len(chunks) + 1
|
||||
assert ticks == [1, 1, 1]
|
||||
|
||||
def test_completed_response_without_stream_payload_does_not_tick(self):
|
||||
calls = []
|
||||
|
||||
def _create(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return _COMPLETE
|
||||
|
||||
client = SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=_create))
|
||||
)
|
||||
ticks = []
|
||||
|
||||
with aux_progress_hook(lambda: ticks.append(1)):
|
||||
result = _create_with_progress(client, {"model": "m1", "messages": []})
|
||||
|
||||
assert calls[0]["stream"] is True
|
||||
assert result is _COMPLETE
|
||||
assert ticks == []
|
||||
|
||||
def test_streaming_rejected_falls_back_to_plain_call(self):
|
||||
client = _FakeClient(
|
||||
@@ -195,6 +225,205 @@ class TestAggregateChatStream:
|
||||
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Content-bearing progress classification
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestContentBearingProgress:
|
||||
@pytest.mark.parametrize(
|
||||
"event",
|
||||
[
|
||||
SimpleNamespace(type="response.created"),
|
||||
SimpleNamespace(type="response.output_item.added", item={"type": "message"}),
|
||||
SimpleNamespace(type="response.content_part.added", part={"type": "output_text"}),
|
||||
SimpleNamespace(type="response.output_text.delta", delta=""),
|
||||
SimpleNamespace(type="response.reasoning_summary_text.delta", delta=""),
|
||||
SimpleNamespace(type="response.function_call_arguments.delta", delta=""),
|
||||
SimpleNamespace(
|
||||
type="response.output_item.added",
|
||||
item=SimpleNamespace(type="function_call"),
|
||||
),
|
||||
{"type": "response.output_text.delta", "delta": ""},
|
||||
],
|
||||
)
|
||||
def test_codex_empty_and_structural_events_are_not_progress(self, event):
|
||||
assert _codex_event_has_content(event) is False
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"event",
|
||||
[
|
||||
SimpleNamespace(type="response.output_text.delta", delta="token"),
|
||||
SimpleNamespace(type="response.reasoning_summary_text.delta", delta="thought"),
|
||||
SimpleNamespace(type="response.function_call_arguments.delta", delta='{"x"'),
|
||||
SimpleNamespace(
|
||||
type="response.output_item.added",
|
||||
item=SimpleNamespace(
|
||||
type="function_call",
|
||||
id="item_1",
|
||||
call_id="call_1",
|
||||
name="lookup",
|
||||
),
|
||||
),
|
||||
{"type": "response.output_text.delta", "delta": "token"},
|
||||
],
|
||||
)
|
||||
def test_codex_nonempty_deltas_are_progress(self, event):
|
||||
assert _codex_event_has_content(event) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("delta", "expected"),
|
||||
[
|
||||
(SimpleNamespace(type="text_delta", text="", thinking=None), False),
|
||||
(SimpleNamespace(type="text_delta", text="token", thinking=None), True),
|
||||
(SimpleNamespace(type="thinking_delta", text=None, thinking="thought"), True),
|
||||
(SimpleNamespace(type="input_json_delta", partial_json=""), False),
|
||||
(SimpleNamespace(type="input_json_delta", partial_json='{"x"'), True),
|
||||
],
|
||||
)
|
||||
def test_anthropic_requires_nonempty_delta_payload(self, delta, expected):
|
||||
event = SimpleNamespace(type="content_block_delta", delta=delta)
|
||||
assert _anthropic_event_has_content(event) is expected
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("block", "expected"),
|
||||
[
|
||||
(SimpleNamespace(type="text"), False),
|
||||
(SimpleNamespace(type="tool_use", id=None, name=None), False),
|
||||
(SimpleNamespace(type="tool_use", id="toolu_1", name="lookup"), True),
|
||||
],
|
||||
)
|
||||
def test_anthropic_tool_start_requires_identity(self, block, expected):
|
||||
event = SimpleNamespace(type="content_block_start", content_block=block)
|
||||
assert _anthropic_event_has_content(event) is expected
|
||||
|
||||
def test_chat_stream_empty_tool_scaffolding_is_not_progress(self):
|
||||
ticks = []
|
||||
empty_tool_call = SimpleNamespace(
|
||||
index=0,
|
||||
id=None,
|
||||
function=SimpleNamespace(name=None, arguments=""),
|
||||
)
|
||||
accumulator = _ChatStreamAccumulator()
|
||||
|
||||
with aux_progress_hook(lambda: ticks.append(1)):
|
||||
accumulator.feed(_chunk(tool_calls=[empty_tool_call]))
|
||||
|
||||
assert ticks == []
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_call",
|
||||
[
|
||||
SimpleNamespace(index=0, id="call_1", function=None),
|
||||
SimpleNamespace(
|
||||
index=0,
|
||||
id=None,
|
||||
function=SimpleNamespace(name="lookup", arguments=""),
|
||||
),
|
||||
SimpleNamespace(
|
||||
index=0,
|
||||
id=None,
|
||||
function=SimpleNamespace(name=None, arguments='{"q"'),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_chat_stream_substantive_tool_fragments_are_progress(self, tool_call):
|
||||
ticks = []
|
||||
accumulator = _ChatStreamAccumulator()
|
||||
|
||||
with aux_progress_hook(lambda: ticks.append(1)):
|
||||
accumulator.feed(_chunk(tool_calls=[tool_call]))
|
||||
|
||||
assert ticks == [1]
|
||||
|
||||
def test_codex_adapter_updates_fence_only_for_substantive_events(self):
|
||||
events = [
|
||||
SimpleNamespace(type="response.created"),
|
||||
SimpleNamespace(type="response.output_text.delta", delta=""),
|
||||
SimpleNamespace(type="response.output_text.delta", delta="token"),
|
||||
SimpleNamespace(
|
||||
type="response.output_item.added",
|
||||
item=SimpleNamespace(
|
||||
type="function_call", id="item_1", call_id="call_1", name="lookup"
|
||||
),
|
||||
),
|
||||
]
|
||||
real_client = SimpleNamespace(
|
||||
base_url="https://chatgpt.com/backend-api/codex",
|
||||
responses=SimpleNamespace(create=lambda **_kwargs: iter(events)),
|
||||
)
|
||||
adapter = _CodexCompletionsAdapter(real_client, "gpt-5.6-sol")
|
||||
fence = CompressionCommitFence()
|
||||
touches = []
|
||||
|
||||
def _touch():
|
||||
touches.append(1)
|
||||
fence.touch_progress()
|
||||
|
||||
def _consume(stream, *, model, on_event):
|
||||
del model
|
||||
for event in stream:
|
||||
on_event(event)
|
||||
return SimpleNamespace(output=[], usage=None)
|
||||
|
||||
with (
|
||||
patch("agent.codex_runtime._consume_codex_event_stream", _consume),
|
||||
aux_progress_hook(_touch),
|
||||
):
|
||||
adapter.create(messages=[{"role": "user", "content": "summarize"}])
|
||||
|
||||
assert touches == [1, 1]
|
||||
|
||||
def test_anthropic_adapter_updates_fence_only_for_substantive_events(self):
|
||||
events = [
|
||||
SimpleNamespace(type="ping"),
|
||||
SimpleNamespace(
|
||||
type="content_block_delta",
|
||||
delta=SimpleNamespace(type="text_delta", text=""),
|
||||
),
|
||||
SimpleNamespace(
|
||||
type="content_block_delta",
|
||||
delta=SimpleNamespace(type="input_json_delta", partial_json='{"q"'),
|
||||
),
|
||||
SimpleNamespace(
|
||||
type="content_block_start",
|
||||
content_block=SimpleNamespace(
|
||||
type="tool_use", id="toolu_1", name="lookup"
|
||||
),
|
||||
),
|
||||
]
|
||||
adapter = _AnthropicCompletionsAdapter(
|
||||
MagicMock(), "claude-sonnet-4-6", is_oauth=False
|
||||
)
|
||||
fence = CompressionCommitFence()
|
||||
touches = []
|
||||
|
||||
def _touch():
|
||||
touches.append(1)
|
||||
fence.touch_progress()
|
||||
|
||||
def _create_message(*_args, **kwargs):
|
||||
for event in events:
|
||||
kwargs["on_stream_event"](event)
|
||||
raise RuntimeError("stop after callback verification")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"agent.anthropic_adapter.build_anthropic_kwargs",
|
||||
return_value={"model": "claude-sonnet-4-6", "messages": []},
|
||||
),
|
||||
patch(
|
||||
"agent.anthropic_adapter.create_anthropic_message",
|
||||
side_effect=_create_message,
|
||||
),
|
||||
aux_progress_hook(_touch),
|
||||
pytest.raises(RuntimeError, match="stop after callback verification"),
|
||||
):
|
||||
adapter.create(messages=[{"role": "user", "content": "summarize"}])
|
||||
|
||||
assert touches == [1, 1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Ceiling arithmetic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user