akseljoonas HF Staff commited on
Commit
babd3b5
·
1 Parent(s): 8c7924f

Fix streaming: line-buffered output instead of per-token Live re-render

Browse files
Files changed (1) hide show
  1. agent/main.py +16 -31
agent/main.py CHANGED
@@ -212,50 +212,35 @@ class _ThinkingShimmer:
212
 
213
 
214
  class _StreamBuffer:
215
- """Buffers streaming chunks and renders markdown line-by-line via rich Live."""
216
 
217
  def __init__(self, console):
218
  self._console = console
219
  self._buffer = ""
220
- self._live = None
221
- self._lines_printed = 0
222
-
223
- def _start_live(self):
224
- if self._live is None:
225
- from rich.live import Live
226
- self._live = Live(
227
- "",
228
- console=self._console,
229
- refresh_per_second=8,
230
- vertical_overflow="visible",
231
- )
232
- self._live.start()
233
 
234
  def add_chunk(self, text: str):
235
  self._buffer += text
236
- self._start_live()
237
- self._update()
238
-
239
- def _update(self):
240
- from rich.markdown import Markdown
241
- if self._live:
242
- self._live.update(Markdown(self._buffer))
243
 
244
  def finish(self):
245
- """Finalize: stop live display (final frame is already rendered)."""
246
- if self._live:
247
- self._live.stop()
248
- self._live = None
 
249
  self._buffer = ""
250
- self._lines_printed = 0
251
 
252
  def discard(self):
253
- """Discard without final render (e.g. for tool-only turns)."""
254
- if self._live:
255
- self._live.stop()
256
- self._live = None
257
  self._buffer = ""
258
- self._lines_printed = 0
259
 
260
 
261
  async def event_listener(
 
212
 
213
 
214
  class _StreamBuffer:
215
+ """Buffers streaming tokens, flushes complete lines. No flicker."""
216
 
217
  def __init__(self, console):
218
  self._console = console
219
  self._buffer = ""
220
+ self._started = False
 
 
 
 
 
 
 
 
 
 
 
 
221
 
222
  def add_chunk(self, text: str):
223
  self._buffer += text
224
+ # Flush every complete line
225
+ while "\n" in self._buffer:
226
+ line, self._buffer = self._buffer.split("\n", 1)
227
+ if not self._started:
228
+ self._console.print() # blank line before first output
229
+ self._started = True
230
+ self._console.print(line)
231
 
232
  def finish(self):
233
+ """Print any remaining partial line, then reset."""
234
+ if self._buffer.strip():
235
+ if not self._started:
236
+ self._console.print()
237
+ self._console.print(self._buffer)
238
  self._buffer = ""
239
+ self._started = False
240
 
241
  def discard(self):
 
 
 
 
242
  self._buffer = ""
243
+ self._started = False
244
 
245
 
246
  async def event_listener(