Skip to content

Commit bc692f1

Browse files
committed
Fix some performance issues.
1 parent 78988c8 commit bc692f1

7 files changed

Lines changed: 138 additions & 14 deletions

File tree

‎python_agent_harness/core/agent.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -471,8 +471,9 @@ def run(self) -> str | None:
471471
# after successful responses, so without this the
472472
# interrupted turn would never reach the session
473473
# file — the next turn's save would overwrite it
474-
# without ever containing it.
475-
session.auto_save(salvaged, self.system)
474+
# without ever containing it. flush bypasses the
475+
# debounce so the partial turn is not lost.
476+
session.flush_auto_save(salvaged)
476477
# The title is generated on the first save, not
477478
# only on clean completion — an interrupted session
478479
# still gets a meaningful name (one-shot; no-op when
@@ -484,6 +485,9 @@ def run(self) -> str | None:
484485
# no-op when the title already exists or generation
485486
# is in flight)
486487
session.generate_session_title()
488+
# The per-round auto_save is debounced; flush at run
489+
# end so the final state always reaches the file.
490+
session.flush_auto_save(salvaged)
487491
return self.result
488492

489493
def _next_state(self) -> str:

‎python_agent_harness/core/models.py‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -69,20 +69,26 @@ class ImagePart:
6969
path: str | None = None
7070
url: str | None = None
7171

72+
def __post_init__(self) -> None:
73+
# to_api() runs on every round (token estimation + API payload),
74+
# so the base64 data URL is encoded once, not per call.
75+
self._data_url: str | None = None
76+
if self.data is not None:
77+
self._data_url = (
78+
f"data:{self.media_type};base64,"
79+
f"{base64.b64encode(self.data).decode('ascii')}"
80+
)
81+
7282
@classmethod
7383
def from_url(cls, url: str) -> ImagePart:
7484
return cls(url=url)
7585

7686
def to_api(self) -> dict[str, Any]:
7787
if self.url:
7888
return {"type": "image_url", "image_url": {"url": self.url}}
79-
if self.data is None:
89+
if self._data_url is None:
8090
raise ValueError("ImagePart requires either data or url")
81-
b64 = base64.b64encode(self.data).decode("ascii")
82-
return {
83-
"type": "image_url",
84-
"image_url": {"url": f"data:{self.media_type};base64,{b64}"},
85-
}
91+
return {"type": "image_url", "image_url": {"url": self._data_url}}
8692

8793

8894
# Type alias for the individual parts a Message.content list may contain.
@@ -155,6 +161,11 @@ class Message:
155161
injected: bool = False # harness-injected (nudge/plan/build-switch), not user input
156162

157163
def to_api(self) -> dict[str, Any]:
164+
# Not cached: the per-message dict is cheap to rebuild and
165+
# caching it (returning the same object each call) would make a
166+
# mutation by any caller corrupt later rounds and token
167+
# estimation. The one expensive step — base64-encoding an image
168+
# — is cached on ``ImagePart`` instead (see ImagePart.__post_init__).
158169
d: dict[str, Any] = {"role": self.role}
159170
if self.content is not None:
160171
content = self._api_content()

‎python_agent_harness/core/token_estimator.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,10 @@ def estimate_tokens(text: str) -> int:
4545
"""Estimate tokens in TEXT: Latin ~4 chars/token, CJK ~2 chars/token."""
4646
if not text:
4747
return 0
48-
cjk = len(_CJK_RE.findall(text))
48+
# subn counts CJK chars in one C-speed pass without allocating the
49+
# per-match list findall would build (this runs on the full payload
50+
# every round).
51+
cjk = _CJK_RE.subn("", text)[1]
4952
latin = len(text) - cjk
5053
return round(latin / 4.0 + cjk / 2.0)
5154

‎python_agent_harness/session/config.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,14 @@
5353
SESSION_DIR = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local" / "share"))
5454
SESSION_SUBDIR = "python-agent-harness/sessions"
5555
AUTO_SAVE_SESSION = True
56+
# auto_save is called after every round; the full-conversation rewrite is
57+
# debounced so a long run does not rewrite the whole file each round.
58+
# A save is skipped while fewer than AUTO_SAVE_MIN_INTERVAL seconds have
59+
# passed since the last one AND the message count has not grown by
60+
# AUTO_SAVE_MIN_GROWTH_MSGS. flush_auto_save() bypasses the debounce
61+
# (run end, cancel, compaction).
62+
AUTO_SAVE_MIN_INTERVAL = 2.0
63+
AUTO_SAVE_MIN_GROWTH_MSGS = 8
5664

5765
# ---- LLM interaction logs ---------------------------------------------------
5866
LLM_LOG_ENABLED = False

‎python_agent_harness/session/session.py‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,10 @@ def __init__(
195195
self.todos: list[dict] = []
196196
self.pending_user_prompts: list[str] = []
197197
self._save_error: str | None = None
198+
# auto_save debounce state (see config.AUTO_SAVE_*): the last
199+
# successful save's wall time and conversation length
200+
self._last_save_at: float | None = None
201+
self._last_save_len: int = 0
198202
self.last_messages: list = []
199203
self.cancel_event = threading.Event()
200204
# Named LLM profiles for runtime switching via /model command
@@ -680,6 +684,39 @@ def remember_user_text(self, messages: list) -> None:
680684
def auto_save(self, messages: list, system: str | None) -> None:
681685
if not config.AUTO_SAVE_SESSION:
682686
return
687+
# Debounce the full-conversation rewrite: auto_save runs after
688+
# every round, but the file only needs refreshing when enough
689+
# time has passed AND the conversation grew. The growth check
690+
# is on the message count (cheap) before the expensive
691+
# _conversation_text serialization.
692+
now = time.monotonic()
693+
# getattr: test doubles built via __new__ skip __init__ and never
694+
# set the debounce state; treat a missing state as "never saved".
695+
last = getattr(self, "_last_save_at", None)
696+
if last is not None:
697+
elapsed = now - last
698+
grew = (
699+
len(messages) - getattr(self, "_last_save_len", 0)
700+
>= config.AUTO_SAVE_MIN_GROWTH_MSGS
701+
)
702+
# skip while inside the min interval and the conversation has
703+
# not grown enough; a big round (grew) or the interval
704+
# elapsing forces a save
705+
if elapsed < config.AUTO_SAVE_MIN_INTERVAL and not grew:
706+
return
707+
self._do_save(messages)
708+
709+
def flush_auto_save(self, messages: list) -> None:
710+
"""Save immediately, bypassing the auto_save debounce.
711+
712+
Called at run end / cancel / compaction, where the conversation
713+
must reach the session file regardless of the debounce window.
714+
"""
715+
if not config.AUTO_SAVE_SESSION:
716+
return
717+
self._do_save(messages)
718+
719+
def _do_save(self, messages: list) -> None:
683720
text = self._conversation_text(messages)
684721
# retry once: transient failures (NFS hiccup, brief lock) clear
685722
# on the second attempt; permanent ones (disk full, read-only)
@@ -689,6 +726,8 @@ def auto_save(self, messages: list, system: str | None) -> None:
689726
try:
690727
self.store.save(text)
691728
self._save_error = None
729+
self._last_save_at = time.monotonic()
730+
self._last_save_len = len(messages)
692731
return
693732
except (OSError, UnicodeError) as e:
694733
# UnicodeError too: a save failure must never escape into
@@ -983,7 +1022,7 @@ def compact_conversation(self) -> tuple[bool, str]:
9831022
# preserved verbatim after the frame, so the model keeps
9841023
# the actual requests.
9851024
self.last_messages = compacted_messages(summary, user_prompt_texts(messages))
986-
self.auto_save(self.last_messages, self.system_prompt)
1025+
self.flush_auto_save(self.last_messages)
9871026
self.notify("compact")
9881027
return True, "Session compacted successfully."
9891028
except Exception as e: # noqa: BLE001 - compaction failure is non-fatal
@@ -1019,7 +1058,7 @@ def summarize_conversation(self) -> str:
10191058
if not summary:
10201059
return "Summary failed: empty response."
10211060
self.last_messages.append(Msg(role="assistant", content=summary))
1022-
self.auto_save(self.last_messages, self.system_prompt)
1061+
self.flush_auto_save(self.last_messages)
10231062
return "Summary appended."
10241063

10251064

‎python_agent_harness/tui/render.py‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,17 @@ def sanitize_terminal(text: str) -> str:
7575
return _TERMINAL_CONTROL.sub("", api_safe(text))
7676

7777

78+
# Characters sanitize_terminal would change: any C0/C1 control except
79+
# tab/newline, or a lone surrogate. A cheap pre-check lets callers skip
80+
# the full regex pass on the common (clean) case.
81+
_NEEDS_SANITIZE = re.compile(r"[\x00-\x08\x0b-\x1f\x7f-\x9f\ud800-\udfff]")
82+
83+
84+
def needs_sanitize(text: str) -> bool:
85+
"""Whether sanitize_terminal would change TEXT (cheap pre-check)."""
86+
return _NEEDS_SANITIZE.search(text) is not None
87+
88+
7889
def safe_markup(value: object) -> str:
7990
"""VALUE made safe to embed in a Rich markup string.
8091
@@ -355,12 +366,24 @@ def _stream_row(self) -> Text | None:
355366
"""Live stream row (cheap Text, tail-capped)."""
356367
with self.lock:
357368
stream = self.stream_text
358-
stream = _strip_final_check(sanitize_terminal(stream))
359-
if not stream:
360-
return None
361369
cap = self._visible_row_cap()
362370
width = getattr(self.console, "width", None) or 80
363371
lines = max(3, cap - 3)
372+
# Sanitize BEFORE stripping: a control sequence inside the
373+
# [FINAL CHECK] header ("[FINAL\x1b[0m CHECK]") would otherwise
374+
# defeat the header match and leak the block. The full-buffer
375+
# sanitize is skipped when the text is already clean (the common
376+
# case), so the expensive regex pass only runs when needed.
377+
if needs_sanitize(stream):
378+
stream = sanitize_terminal(stream)
379+
# Strip on the FULL buffer: the header can sit arbitrarily far
380+
# before the tail, so windowing first would leak the block.
381+
stream = _strip_final_check(stream)
382+
window = max(lines * max(1, width) * 2, 8192)
383+
if len(stream) > window:
384+
stream = stream[-window:]
385+
if not stream:
386+
return None
364387
preview = _tail_lines(stream, lines)
365388
preview = _tail_chars(preview, lines * max(1, width))
366389
row = Text(f"assistant: {preview}", style=ASSISTANT_STYLE)

‎tests/tui/test_tui_render.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,42 @@ def test_stream_pure_final_check_hidden(self):
428428
row = tui._stream_row()
429429
self.assertIsNone(row)
430430

431+
def test_stream_final_check_longer_than_window_hidden(self):
432+
"""A [FINAL CHECK] block whose header sits further back than the
433+
sanitize window must still be stripped: the strip runs on the
434+
full buffer, only the sanitize is windowed. Windowing first
435+
leaked the block's content (the header fell outside the window)."""
436+
tui, _ = make_tui()
437+
block = (
438+
"[FINAL CHECK]\n"
439+
+ ("- Goal: " + "x" * 100 + "\n") * 200
440+
+ "- Status: SUCCESS\n- Evidence: y"
441+
)
442+
tui.stream_text = "answer text\n\n" + block
443+
row = tui._stream_row()
444+
self.assertIsNotNone(row)
445+
self.assertIn("answer text", row.plain)
446+
self.assertNotIn("FINAL CHECK", row.plain)
447+
self.assertNotIn("Status: SUCCESS", row.plain)
448+
self.assertNotIn("Evidence:", row.plain)
449+
450+
def test_stream_final_check_with_escape_in_header_hidden(self):
451+
"""A control sequence inside the [FINAL CHECK] header must not
452+
defeat the strip: sanitize runs BEFORE the strip, so the header
453+
matches. Stripping first leaked the block (the escape broke the
454+
header regex)."""
455+
tui, _ = make_tui()
456+
tui.stream_text = (
457+
"answer text\n\n[FINAL\x1b[0m CHECK]\n"
458+
"- Goal: x\n- Status: SUCCESS\n- Evidence: y"
459+
)
460+
row = tui._stream_row()
461+
self.assertIsNotNone(row)
462+
self.assertIn("answer text", row.plain)
463+
self.assertNotIn("FINAL CHECK", row.plain)
464+
self.assertNotIn("Status: SUCCESS", row.plain)
465+
self.assertNotIn("Evidence:", row.plain)
466+
431467
def test_reasoning_streams_normally(self):
432468
"""While streaming, reasoning content shows up live like any
433469
other text — the collapse only happens in the final history."""

0 commit comments

Comments
 (0)