Optimize response stream update cadence
Coalesce expensive full-response parsing and publication behind a 128-character or 50-millisecond cadence while retaining per-chunk security hooks, strict tool validation, immediate canonical completion, and the final pre-Context-Doctor flush. Add deterministic coverage for fast fragments, slow streams, final snapshot ordering, and canonical-root early stopping.
Alessandro committed
Aug 27, 2026 at 04:57 UTC
75bbc1d71adf20fb74d30682a3ce25cf73c2b759
2 files changed
+135
-4
agent.py
+26
-4
@@ -1,4 +1,4 @@
1
-import asyncio, json, random, re, string, threading
1
+import asyncio, json, random, re, string, threading, time
2
3
from collections import OrderedDict
4
from dataclasses import dataclass, field
@@ -42,6 +42,10 @@ from helpers.llm_result import (
42
from helpers.litellm_transport import ResponsesTransport
43
from helpers.responses_tools import build_responses_function_tools, original_tool_name
44
45
+_RESPONSE_STREAM_UPDATE_CHARS = 128
46
+_RESPONSE_STREAM_UPDATE_SECONDS = 0.05
47
+
48
+
49
class AgentContextType(Enum):
50
USER = "user"
51
TASK = "task"
@@ -403,6 +407,9 @@ class Agent:
407
self.loop_data.iteration += 1
408
self.loop_data.params_temporary = {} # clear temporary params
409
last_response_stream_full = ""
410
+ last_response_stream_chars = 0
411
+ last_response_stream_at = time.monotonic()
412
+ response_stream_pending = False
413
414
# call message_loop_start extensions
415
await extension.call_extensions_async(
@@ -440,7 +447,8 @@ class Agent:
447
await self.handle_reasoning_stream(stream_data["full"])
448
449
async def stream_callback(chunk: str, full: str):
443
- nonlocal last_response_stream_full
450
+ nonlocal last_response_stream_full, last_response_stream_chars
451
+ nonlocal last_response_stream_at, response_stream_pending
452
await self.handle_intervention()
453
# output the agent response stream
454
if chunk == full:
@@ -455,6 +463,7 @@ class Agent:
463
pass
464
else:
465
await self.handle_response_stream(full)
466
+ response_stream_pending = False
467
return full.strip()
468
469
await extension.call_extensions_async(
@@ -466,9 +475,19 @@ class Agent:
475
# Stream masked chunk after extensions processed it
476
if stream_data.get("chunk"):
477
printer.stream(stream_data["chunk"])
469
- # Use the potentially modified full text for downstream processing
470
- await self.handle_response_stream(stream_data["full"])
478
last_response_stream_full = stream_data["full"]
479
+ response_stream_pending = True
480
+ now = time.monotonic()
481
+ if (
482
+ len(full) - last_response_stream_chars
483
+ >= _RESPONSE_STREAM_UPDATE_CHARS
484
+ or now - last_response_stream_at
485
+ >= _RESPONSE_STREAM_UPDATE_SECONDS
486
+ ):
487
+ await self.handle_response_stream(last_response_stream_full)
488
+ last_response_stream_chars = len(full)
489
+ last_response_stream_at = time.monotonic()
490
+ response_stream_pending = False
491
492
# call main LLM
493
llm_result = await self.call_chat_model_turn(
@@ -479,6 +498,9 @@ class Agent:
498
agent_response = llm_result.response
499
await self.handle_intervention(agent_response)
500
501
+ if response_stream_pending:
502
+ await self.handle_response_stream(last_response_stream_full)
503
+
504
# Notify extensions to finalize their stream filters
505
await extension.call_extensions_async(
506
"reasoning_stream_end", self, loop_data=self.loop_data
tests/test_response_stream_cadence.py
new
+109
@@ -0,0 +1,109 @@
1
+from types import SimpleNamespace
2
+
3
+import pytest
4
+
5
+import agent as agent_module
6
+from agent import Agent
7
+from helpers import extension
8
+from helpers.llm_result import LLMResult
9
+
10
+
11
+async def _run_monologue(monkeypatch, chunks, clock, *, advance=0.0):
12
+ agent = object.__new__(Agent)
13
+ agent.context = SimpleNamespace(streaming_agent=None, task=None)
14
+ agent.last_user_message = None
15
+ events = []
16
+ provider_callbacks = 0
17
+
18
+ async def call_extensions(name, _agent=None, **kwargs):
19
+ events.append(("extension", name))
20
+ if name == "response_stream_chunk":
21
+ stream_data = kwargs["stream_data"]
22
+ stream_data["full"] = "masked:" + stream_data["full"]
23
+
24
+ async def call_model(*, response_callback, **_kwargs):
25
+ nonlocal provider_callbacks
26
+ full = ""
27
+ for chunk in chunks:
28
+ provider_callbacks += 1
29
+ clock.now += advance
30
+ full += chunk
31
+ stopped = await response_callback(chunk, full)
32
+ if stopped:
33
+ return LLMResult(response=stopped)
34
+ return LLMResult(response=full)
35
+
36
+ async def handle_response_stream(full):
37
+ events.append(("handle", full))
38
+
39
+ async def done(_result):
40
+ return "done"
41
+
42
+ async def no_intervention(*_args):
43
+ return None
44
+
45
+ async def no_prompt(**_kwargs):
46
+ return []
47
+
48
+ monkeypatch.setattr(extension, "call_extensions_async", call_extensions)
49
+ monkeypatch.setattr(agent_module.time, "monotonic", lambda: clock.now)
50
+ agent.prepare_prompt = no_prompt
51
+ agent.handle_intervention = no_intervention
52
+ agent.call_chat_model_turn = call_model
53
+ agent.handle_response_stream = handle_response_stream
54
+ agent.hist_add_ai_response = lambda *_args, **_kwargs: SimpleNamespace(id="")
55
+ agent._remember_llm_result_state = lambda *_args: None
56
+ agent.process_llm_result_tools = done
57
+
58
+ result = await Agent.monologue.__wrapped__(agent)
59
+ return result, events, provider_callbacks
60
+
61
+
62
+@pytest.mark.asyncio
63
+async def test_response_stream_coalesces_fast_fragments_and_flushes_final(monkeypatch):
64
+ clock = SimpleNamespace(now=0.0)
65
+ result, events, provider_callbacks = await _run_monologue(
66
+ monkeypatch, list("x" * 130), clock
67
+ )
68
+
69
+ handled = [value for kind, value in events if kind == "handle"]
70
+ chunk_hooks = [
71
+ event
72
+ for event in events
73
+ if event == ("extension", "response_stream_chunk")
74
+ ]
75
+
76
+ assert result == "done"
77
+ assert provider_callbacks == len(chunk_hooks) == 130
78
+ assert handled == ["masked:" + "x" * size for size in (128, 130)]
79
+ assert events.index(("handle", handled[-1])) < events.index(
80
+ ("extension", "response_stream_end")
81
+ )
82
+ assert events.index(("extension", "response_stream_end")) < events.index(
83
+ ("extension", "message_loop_result")
84
+ )
85
+
86
+
87
+@pytest.mark.asyncio
88
+async def test_response_stream_time_bound_keeps_slow_fragments_live(monkeypatch):
89
+ clock = SimpleNamespace(now=0.0)
90
+ _, events, _ = await _run_monologue(
91
+ monkeypatch, ["a" * 10] * 3, clock, advance=0.03
92
+ )
93
+
94
+ handled = [value for kind, value in events if kind == "handle"]
95
+ assert handled == ["masked:" + "a" * size for size in (20, 30)]
96
+
97
+
98
+@pytest.mark.asyncio
99
+async def test_response_stream_still_stops_on_exact_canonical_root(monkeypatch):
100
+ clock = SimpleNamespace(now=0.0)
101
+ message = '{"tool_name":"response","tool_args":{"text":"ok"}}'
102
+ result, events, provider_callbacks = await _run_monologue(
103
+ monkeypatch, [message[:-1], message[-1], " unreachable"], clock
104
+ )
105
+
106
+ assert result == "done"
107
+ assert provider_callbacks == 2
108
+ assert [value for kind, value in events if kind == "handle"] == [message]
109
+ assert events.count(("extension", "response_stream_chunk")) == 1