feat(skills/telegram): opt-in live streaming edit-in-place for TG bot (FPLAN-0297, DPLAN-0229)
Repurpose the heartbeat into a ~2s transcript-tail loop that edits the "Processing" message in place (block-level: thinking/tool/text), plain text, coalesced, no-op-skipped, 429 retry_after aware, with 4096 rollover. Opt-in per-bot "stream" flag, default OFF; batch path byte-for-byte unchanged. @hooks reviewed: no change needed (already edits processing_message_id for the single-chunk final). Race hardened: re-check delivered before each edit. 37/37 streaming + 653/653 TG tests green; live-proven on the devpulse bot. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q3mZT61WsKVN3srCwVDBiW
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
337f31ddab
commit
a5ede6fbf4
@@ -96,6 +96,21 @@
|
||||
"standard": "encapsulation",
|
||||
"reason": "Test file — imports handler directly for unit testing. Same pattern as test_multi_bot.py and test_attach_only.py."
|
||||
},
|
||||
{
|
||||
"file": "lib/telegram/tests/test_streaming.py",
|
||||
"standard": "architecture",
|
||||
"reason": "Test file — lives in tests/ by convention. Test files are exempt from layer architecture standard."
|
||||
},
|
||||
{
|
||||
"file": "lib/telegram/tests/test_streaming.py",
|
||||
"standard": "encapsulation",
|
||||
"reason": "Test file — imports handler directly for unit testing. Same pattern as all other TG test files."
|
||||
},
|
||||
{
|
||||
"file": "lib/telegram/tests/test_streaming.py",
|
||||
"standard": "documentation",
|
||||
"reason": "Test file — pytest test methods are self-documenting via class/method names."
|
||||
},
|
||||
{
|
||||
"file": "lib/telegram/tests/test_heartbeat_delivered.py",
|
||||
"standard": "architecture",
|
||||
|
||||
@@ -57,7 +57,7 @@ import uuid
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.error import URLError
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
# Logging
|
||||
@@ -133,6 +133,7 @@ RATE_LIMIT_WINDOW = 60
|
||||
POLL_TIMEOUT = 30
|
||||
SEND_KEYS_DELAY = 0.5
|
||||
HEARTBEAT_INTERVAL = 30 # seconds
|
||||
STREAM_INTERVAL = 2 # seconds between streaming edits
|
||||
CLAUDE_BIN = str(Path.home() / ".local" / "bin" / "claude")
|
||||
MIRROR_SESSION_TYPE = "interactive-mirror"
|
||||
TEMP_DIR = Path("/tmp/telegram_uploads")
|
||||
@@ -164,6 +165,7 @@ class BaseBot:
|
||||
branch_name: Optional[str] = None,
|
||||
shared_session: Optional[str] = None,
|
||||
attach_only: bool = False,
|
||||
stream: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize BaseBot.
|
||||
@@ -200,6 +202,7 @@ class BaseBot:
|
||||
self._shared_session_name = shared_session
|
||||
self._using_shared_session = False
|
||||
self._attach_only = attach_only
|
||||
self._stream = stream
|
||||
self._mirror_mapping_written = False
|
||||
self._last_transcript_path: str | None = None
|
||||
self._config_chat_id: int | None = None
|
||||
@@ -1439,6 +1442,8 @@ class BaseBot:
|
||||
"transcript_path": str(self._active_transcript_path) if self._active_transcript_path else None,
|
||||
"session_id": self._active_session_id,
|
||||
}
|
||||
if self._stream:
|
||||
pending_data["streaming"] = True
|
||||
|
||||
try:
|
||||
self.pending_file.write_text(
|
||||
@@ -1856,8 +1861,10 @@ class BaseBot:
|
||||
|
||||
def _start_heartbeat(self, chat_id: int, processing_msg_id: int) -> None:
|
||||
"""
|
||||
Start a background thread that updates the "Processing..." message
|
||||
with elapsed time.
|
||||
Start a background thread that updates the "Processing..." message.
|
||||
|
||||
In batch mode (default): edits with elapsed time every 30s.
|
||||
In stream mode: tails transcript and edits with live content every 2s.
|
||||
|
||||
Args:
|
||||
chat_id: Chat ID where the processing message was sent
|
||||
@@ -1868,6 +1875,13 @@ class BaseBot:
|
||||
|
||||
def _heartbeat_loop():
|
||||
start = time.time()
|
||||
|
||||
# Streaming mode (FPLAN-0297): live transcript tail
|
||||
if self._stream and self._active_transcript_path:
|
||||
self._streaming_loop(chat_id, processing_msg_id, start)
|
||||
return
|
||||
|
||||
# Batch mode (default): elapsed-time updates
|
||||
while not self._heartbeat_stop.is_set():
|
||||
self._heartbeat_stop.wait(HEARTBEAT_INTERVAL)
|
||||
if self._heartbeat_stop.is_set():
|
||||
@@ -1921,6 +1935,183 @@ class BaseBot:
|
||||
minutes, secs = divmod(total, 60)
|
||||
return f"{minutes}m {secs}s"
|
||||
|
||||
# =============================================
|
||||
# STREAMING EDIT-IN-PLACE (FPLAN-0297)
|
||||
# =============================================
|
||||
|
||||
def _streaming_loop(self, chat_id: int, msg_id: int, start_time: float) -> None:
|
||||
"""Stream transcript content into the processing message via edit-in-place."""
|
||||
path = self._active_transcript_path
|
||||
try:
|
||||
byte_offset = path.stat().st_size if path else 0
|
||||
except OSError as exc:
|
||||
logger.info("Cannot stat transcript for streaming: %s", exc)
|
||||
byte_offset = 0
|
||||
|
||||
buffer = ""
|
||||
last_sent = ""
|
||||
current_msg_id = msg_id
|
||||
retry_after_until = 0.0
|
||||
|
||||
while not self._heartbeat_stop.is_set():
|
||||
self._heartbeat_stop.wait(STREAM_INTERVAL)
|
||||
if self._heartbeat_stop.is_set():
|
||||
break
|
||||
if self._is_pending_delivered():
|
||||
break
|
||||
if not self._tmux_session_exists():
|
||||
break
|
||||
|
||||
new_entries, byte_offset = self._tail_transcript_bytes(path, byte_offset)
|
||||
new_text = self._format_stream_entries(new_entries)
|
||||
if new_text:
|
||||
buffer += new_text
|
||||
|
||||
if self._is_pending_delivered():
|
||||
break
|
||||
|
||||
if not buffer:
|
||||
elapsed = time.time() - start_time
|
||||
placeholder = f"Processing... ({self._format_elapsed(elapsed)})"
|
||||
now = time.time()
|
||||
if placeholder != last_sent and now >= retry_after_until:
|
||||
ok, retry = self._stream_edit(chat_id, current_msg_id, placeholder)
|
||||
retry_after_until = now + retry if retry > 0 else retry_after_until
|
||||
last_sent = placeholder if ok else last_sent
|
||||
continue
|
||||
|
||||
if buffer == last_sent:
|
||||
continue
|
||||
|
||||
now = time.time()
|
||||
if now < retry_after_until:
|
||||
continue
|
||||
|
||||
if len(buffer) > TELEGRAM_CHAR_LIMIT:
|
||||
break_at = buffer.rfind("\n", 0, TELEGRAM_CHAR_LIMIT)
|
||||
if break_at < TELEGRAM_CHAR_LIMIT // 2:
|
||||
break_at = TELEGRAM_CHAR_LIMIT
|
||||
final_chunk = buffer[:break_at].rstrip()
|
||||
self._stream_edit(chat_id, current_msg_id, final_chunk)
|
||||
buffer = buffer[break_at:].lstrip()
|
||||
initial = buffer[:TELEGRAM_CHAR_LIMIT] or "..."
|
||||
result = self.send_message(chat_id, initial)
|
||||
if result:
|
||||
current_msg_id = result.get("message_id", current_msg_id)
|
||||
last_sent = initial
|
||||
continue
|
||||
|
||||
ok, retry = self._stream_edit(chat_id, current_msg_id, buffer)
|
||||
if retry > 0:
|
||||
retry_after_until = now + retry
|
||||
if ok:
|
||||
last_sent = buffer
|
||||
|
||||
def _tail_transcript_bytes(self, path: Path | None, byte_offset: int) -> tuple[list[dict], int]:
|
||||
"""Incrementally tail a JSONL transcript from byte_offset."""
|
||||
if not path:
|
||||
return [], byte_offset
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
f.seek(0, 2)
|
||||
file_size = f.tell()
|
||||
if file_size <= byte_offset:
|
||||
return [], byte_offset
|
||||
f.seek(byte_offset)
|
||||
new_bytes = f.read()
|
||||
except OSError as exc:
|
||||
logger.info("Cannot read transcript for streaming tail: %s", exc)
|
||||
return [], byte_offset
|
||||
|
||||
last_newline = new_bytes.rfind(b"\n")
|
||||
if last_newline == -1:
|
||||
return [], byte_offset
|
||||
|
||||
complete = new_bytes[: last_newline + 1]
|
||||
entries = []
|
||||
for line_bytes in complete.split(b"\n"):
|
||||
line_str = line_bytes.decode("utf-8", errors="replace").strip()
|
||||
if not line_str:
|
||||
continue
|
||||
try:
|
||||
entries.append(json.loads(line_str))
|
||||
except json.JSONDecodeError:
|
||||
logger.info("Malformed JSONL line in streaming tail")
|
||||
|
||||
return entries, byte_offset + len(complete)
|
||||
|
||||
@staticmethod
|
||||
def _format_content_block(block: dict) -> str | None:
|
||||
"""Map a single transcript content block to plain-text display."""
|
||||
if not isinstance(block, dict):
|
||||
return None
|
||||
btype = block.get("type", "")
|
||||
if btype == "text":
|
||||
text = block.get("text", "")
|
||||
return text if text.strip() else None
|
||||
if btype == "thinking":
|
||||
return "Thinking..."
|
||||
if btype == "tool_use":
|
||||
return f"Running {block.get('name', 'tool')}..."
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _format_stream_entries(entries: list[dict]) -> str:
|
||||
"""Block-map transcript entries to plain-text streaming content."""
|
||||
parts: list[str] = []
|
||||
for entry in entries:
|
||||
if entry.get("isSidechain"):
|
||||
continue
|
||||
msg = entry.get("message", {})
|
||||
if msg.get("role") != "assistant":
|
||||
continue
|
||||
content = msg.get("content", [])
|
||||
if isinstance(content, str):
|
||||
if content.strip():
|
||||
parts.append(content)
|
||||
continue
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for block in content:
|
||||
mapped = BaseBot._format_content_block(block)
|
||||
if mapped is not None:
|
||||
parts.append(mapped)
|
||||
if not parts:
|
||||
return ""
|
||||
return "\n".join(parts) + "\n"
|
||||
|
||||
def _stream_edit(self, chat_id: int, message_id: int, text: str) -> tuple[bool, float]:
|
||||
"""Edit a message with 429/not-modified handling for streaming.
|
||||
|
||||
Returns (success, retry_after_seconds). retry_after > 0 means rate-limited.
|
||||
"""
|
||||
url = f"https://api.telegram.org/bot{self.bot_token}/editMessageText"
|
||||
payload = {"chat_id": chat_id, "message_id": message_id, "text": text}
|
||||
|
||||
try:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = Request(url, data=data, headers={"Content-Type": "application/json"})
|
||||
with urlopen(req, timeout=15) as resp:
|
||||
result = json.loads(resp.read().decode("utf-8"))
|
||||
return result.get("ok", False), 0.0
|
||||
except HTTPError as e:
|
||||
try:
|
||||
body = json.loads(e.read().decode("utf-8"))
|
||||
except Exception as parse_err:
|
||||
logger.info("Cannot parse stream edit error body: %s", parse_err)
|
||||
body = {}
|
||||
if e.code == 429:
|
||||
retry = body.get("parameters", {}).get("retry_after", 30)
|
||||
logger.info("Stream edit 429 — backing off %ds", retry)
|
||||
return False, float(retry)
|
||||
if e.code == 400 and "not modified" in body.get("description", ""):
|
||||
return True, 0.0
|
||||
logger.warning("Stream edit HTTP %d: %s", e.code, body.get("description", ""))
|
||||
return False, 0.0
|
||||
except Exception as e:
|
||||
logger.warning("Stream edit error: %s", e)
|
||||
return False, 0.0
|
||||
|
||||
# =============================================
|
||||
# OVERRIDABLE HOOKS
|
||||
# =============================================
|
||||
@@ -2300,6 +2491,7 @@ if __name__ == "__main__":
|
||||
branch_name=config.get("branch_name"),
|
||||
shared_session=config.get("shared_session"),
|
||||
attach_only=config.get("attach_only", False),
|
||||
stream=config.get("stream", False),
|
||||
)
|
||||
|
||||
if config.get("chat_id"):
|
||||
|
||||
@@ -0,0 +1,595 @@
|
||||
# =================== AIPass ====================
|
||||
# Name: test_streaming.py
|
||||
# Description: Tests for streaming edit-in-place (FPLAN-0297)
|
||||
# Version: 1.0.0
|
||||
# Created: 2026-07-01
|
||||
# Modified: 2026-07-01
|
||||
# =============================================
|
||||
|
||||
"""
|
||||
Tests for streaming edit-in-place (FPLAN-0297).
|
||||
|
||||
Covers:
|
||||
- _format_content_block: single block → plain text
|
||||
- _format_stream_entries: batch block-mapping
|
||||
- _tail_transcript_bytes: incremental byte-offset JSONL tail
|
||||
- _stream_edit: 429/not-modified handling
|
||||
- _streaming_loop: full loop behaviour (throttle, rollover, fallback)
|
||||
- write_pending_file: streaming flag in pending file
|
||||
- Batch mode unchanged when stream=False
|
||||
"""
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from apps.handlers.base_bot import BaseBot # type: ignore[import-not-found]
|
||||
|
||||
|
||||
# =============================================
|
||||
# Fixtures
|
||||
# =============================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _patch_deps(tmp_path):
|
||||
"""Patch signal and atexit for safe BaseBot construction."""
|
||||
patches = [
|
||||
patch("apps.handlers.base_bot.PENDING_DIR", tmp_path),
|
||||
patch("apps.handlers.base_bot.signal.signal"),
|
||||
patch("apps.handlers.base_bot.atexit.register"),
|
||||
]
|
||||
for p in patches:
|
||||
p.start()
|
||||
yield
|
||||
for p in patches:
|
||||
p.stop()
|
||||
|
||||
|
||||
def _make_bot(tmp_path, _patch_deps, stream=False):
|
||||
"""Create a BaseBot with test defaults."""
|
||||
workdir = tmp_path / "workdir"
|
||||
workdir.mkdir(exist_ok=True)
|
||||
with patch("apps.handlers.base_bot.PENDING_DIR", tmp_path):
|
||||
bot = BaseBot(
|
||||
bot_id="stream_test",
|
||||
bot_token="123:FAKETOKEN",
|
||||
work_dir=workdir,
|
||||
bot_name="Stream Test Bot",
|
||||
allowed_user_ids=[111],
|
||||
branch_name="testbranch",
|
||||
stream=stream,
|
||||
)
|
||||
bot.send_message = MagicMock(return_value={"ok": True, "message_id": 1})
|
||||
bot.edit_message = MagicMock(return_value=True)
|
||||
return bot
|
||||
|
||||
|
||||
def _write_transcript(path, entries):
|
||||
"""Write JSONL entries to a file. Returns the file path."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
content = "\n".join(json.dumps(e) for e in entries) + "\n"
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def _assistant_entry(content_blocks, sidechain=False):
|
||||
"""Build a transcript entry with assistant role."""
|
||||
entry: dict = {"message": {"role": "assistant", "content": content_blocks}}
|
||||
if sidechain:
|
||||
entry["isSidechain"] = True # type: ignore[assignment]
|
||||
return entry
|
||||
|
||||
|
||||
def _user_entry(text="hello"):
|
||||
"""Build a transcript entry with user role."""
|
||||
return {"message": {"role": "user", "content": text}}
|
||||
|
||||
|
||||
# =============================================
|
||||
# 1. _format_content_block
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestFormatContentBlock:
|
||||
"""Unit tests for _format_content_block."""
|
||||
|
||||
def test_text_block(self):
|
||||
result = BaseBot._format_content_block({"type": "text", "text": "Hello world"})
|
||||
assert result == "Hello world"
|
||||
|
||||
def test_text_block_empty(self):
|
||||
result = BaseBot._format_content_block({"type": "text", "text": " "})
|
||||
assert result is None
|
||||
|
||||
def test_thinking_block(self):
|
||||
result = BaseBot._format_content_block({"type": "thinking", "thinking": "Let me think..."})
|
||||
assert result == "Thinking..."
|
||||
|
||||
def test_tool_use_block(self):
|
||||
result = BaseBot._format_content_block({"type": "tool_use", "name": "Bash", "input": {}})
|
||||
assert result == "Running Bash..."
|
||||
|
||||
def test_tool_use_no_name(self):
|
||||
result = BaseBot._format_content_block({"type": "tool_use"})
|
||||
assert result == "Running tool..."
|
||||
|
||||
def test_unknown_type(self):
|
||||
result = BaseBot._format_content_block({"type": "server_tool_use"})
|
||||
assert result is None
|
||||
|
||||
def test_not_dict(self):
|
||||
result = BaseBot._format_content_block("not a dict")
|
||||
assert result is None
|
||||
|
||||
|
||||
# =============================================
|
||||
# 2. _format_stream_entries
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestFormatStreamEntries:
|
||||
"""Unit tests for _format_stream_entries."""
|
||||
|
||||
def test_text_entry(self):
|
||||
entries = [_assistant_entry([{"type": "text", "text": "Answer here"}])]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == "Answer here\n"
|
||||
|
||||
def test_thinking_then_text(self):
|
||||
entries = [
|
||||
_assistant_entry(
|
||||
[
|
||||
{"type": "thinking", "thinking": "hmm"},
|
||||
{"type": "text", "text": "Done"},
|
||||
]
|
||||
)
|
||||
]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == "Thinking...\nDone\n"
|
||||
|
||||
def test_tool_use_entry(self):
|
||||
entries = [_assistant_entry([{"type": "tool_use", "name": "Read"}])]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == "Running Read...\n"
|
||||
|
||||
def test_skips_user_role(self):
|
||||
entries = [_user_entry("hello")]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == ""
|
||||
|
||||
def test_skips_sidechain(self):
|
||||
entries = [_assistant_entry([{"type": "text", "text": "side"}], sidechain=True)]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == ""
|
||||
|
||||
def test_multiple_entries(self):
|
||||
entries = [
|
||||
_assistant_entry([{"type": "tool_use", "name": "Bash"}]),
|
||||
_assistant_entry([{"type": "text", "text": "Result"}]),
|
||||
]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == "Running Bash...\nResult\n"
|
||||
|
||||
def test_string_content(self):
|
||||
entries = [{"message": {"role": "assistant", "content": "plain string"}}]
|
||||
result = BaseBot._format_stream_entries(entries)
|
||||
assert result == "plain string\n"
|
||||
|
||||
def test_empty_entries(self):
|
||||
result = BaseBot._format_stream_entries([])
|
||||
assert result == ""
|
||||
|
||||
|
||||
# =============================================
|
||||
# 3. _tail_transcript_bytes
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestTailTranscriptBytes:
|
||||
"""Unit tests for _tail_transcript_bytes."""
|
||||
|
||||
def test_new_lines_parsed(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
transcript = tmp_path / "test.jsonl"
|
||||
entries = [_assistant_entry([{"type": "text", "text": "Hello"}])]
|
||||
_write_transcript(transcript, entries)
|
||||
|
||||
result_entries, new_offset = bot._tail_transcript_bytes(transcript, 0)
|
||||
assert len(result_entries) == 1
|
||||
assert result_entries[0]["message"]["content"][0]["text"] == "Hello"
|
||||
assert new_offset > 0
|
||||
|
||||
def test_no_new_content(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
transcript = tmp_path / "test.jsonl"
|
||||
_write_transcript(transcript, [_assistant_entry([{"type": "text", "text": "Hi"}])])
|
||||
size = transcript.stat().st_size
|
||||
|
||||
result_entries, new_offset = bot._tail_transcript_bytes(transcript, size)
|
||||
assert result_entries == []
|
||||
assert new_offset == size
|
||||
|
||||
def test_partial_line_not_consumed(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
transcript = tmp_path / "test.jsonl"
|
||||
complete = json.dumps(_assistant_entry([{"type": "text", "text": "A"}]))
|
||||
partial = '{"incomplete": true'
|
||||
transcript.write_text(complete + "\n" + partial, encoding="utf-8")
|
||||
|
||||
result_entries, new_offset = bot._tail_transcript_bytes(transcript, 0)
|
||||
assert len(result_entries) == 1
|
||||
assert new_offset == len(complete.encode("utf-8")) + 1 # includes the \n
|
||||
|
||||
def test_none_path(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
result_entries, offset = bot._tail_transcript_bytes(None, 0)
|
||||
assert result_entries == []
|
||||
assert offset == 0
|
||||
|
||||
def test_missing_file(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
missing = tmp_path / "nonexistent.jsonl"
|
||||
result_entries, offset = bot._tail_transcript_bytes(missing, 0)
|
||||
assert result_entries == []
|
||||
assert offset == 0
|
||||
|
||||
def test_incremental_reads(self, tmp_path, _patch_deps):
|
||||
"""Two reads: first gets initial content, second gets appended content."""
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
transcript = tmp_path / "test.jsonl"
|
||||
entry1 = _assistant_entry([{"type": "text", "text": "First"}])
|
||||
_write_transcript(transcript, [entry1])
|
||||
|
||||
entries1, offset1 = bot._tail_transcript_bytes(transcript, 0)
|
||||
assert len(entries1) == 1
|
||||
|
||||
entry2 = _assistant_entry([{"type": "text", "text": "Second"}])
|
||||
with open(transcript, "a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(entry2) + "\n")
|
||||
|
||||
entries2, offset2 = bot._tail_transcript_bytes(transcript, offset1)
|
||||
assert len(entries2) == 1
|
||||
assert entries2[0]["message"]["content"][0]["text"] == "Second"
|
||||
assert offset2 > offset1
|
||||
|
||||
|
||||
# =============================================
|
||||
# 4. _stream_edit
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestStreamEdit:
|
||||
"""Unit tests for _stream_edit (429/not-modified handling)."""
|
||||
|
||||
def test_success(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
response = json.dumps({"ok": True}).encode("utf-8")
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.read.return_value = response
|
||||
mock_resp.__enter__ = MagicMock(return_value=mock_resp)
|
||||
mock_resp.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
with patch("apps.handlers.base_bot.urlopen", return_value=mock_resp):
|
||||
ok, retry = bot._stream_edit(123, 456, "hello")
|
||||
assert ok is True
|
||||
assert retry == 0.0
|
||||
|
||||
def test_429_returns_retry_after(self, tmp_path, _patch_deps):
|
||||
from urllib.error import HTTPError
|
||||
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
|
||||
body = json.dumps(
|
||||
{
|
||||
"ok": False,
|
||||
"error_code": 429,
|
||||
"description": "Too Many Requests: retry after 15",
|
||||
"parameters": {"retry_after": 15},
|
||||
}
|
||||
).encode("utf-8")
|
||||
err = HTTPError("url", 429, "Too Many Requests", None, None) # type: ignore[arg-type]
|
||||
err.read = MagicMock(return_value=body)
|
||||
|
||||
with patch("apps.handlers.base_bot.urlopen", side_effect=err):
|
||||
ok, retry = bot._stream_edit(123, 456, "hello")
|
||||
assert ok is False
|
||||
assert retry == 15.0
|
||||
|
||||
def test_400_not_modified_treated_as_ok(self, tmp_path, _patch_deps):
|
||||
from urllib.error import HTTPError
|
||||
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
|
||||
body = json.dumps(
|
||||
{
|
||||
"ok": False,
|
||||
"error_code": 400,
|
||||
"description": "Bad Request: message is not modified",
|
||||
}
|
||||
).encode("utf-8")
|
||||
err = HTTPError("url", 400, "Bad Request", None, None) # type: ignore[arg-type]
|
||||
err.read = MagicMock(return_value=body)
|
||||
|
||||
with patch("apps.handlers.base_bot.urlopen", side_effect=err):
|
||||
ok, retry = bot._stream_edit(123, 456, "same text")
|
||||
assert ok is True
|
||||
assert retry == 0.0
|
||||
|
||||
def test_other_error(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
with patch("apps.handlers.base_bot.urlopen", side_effect=ConnectionError("fail")):
|
||||
ok, retry = bot._stream_edit(123, 456, "hello")
|
||||
assert ok is False
|
||||
assert retry == 0.0
|
||||
|
||||
|
||||
# =============================================
|
||||
# 5. Pending file streaming flag
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestPendingFileStreaming:
|
||||
"""Tests for streaming flag in write_pending_file."""
|
||||
|
||||
def test_no_streaming_key_when_off(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=False)
|
||||
with patch("apps.handlers.base_bot.PENDING_DIR", tmp_path):
|
||||
bot.pending_file = tmp_path / "bot-stream_test.json"
|
||||
bot.write_pending_file(123, 1, 2)
|
||||
data = json.loads(bot.pending_file.read_text(encoding="utf-8"))
|
||||
assert "streaming" not in data
|
||||
|
||||
def test_streaming_true_when_on(self, tmp_path, _patch_deps):
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
with patch("apps.handlers.base_bot.PENDING_DIR", tmp_path):
|
||||
bot.pending_file = tmp_path / "bot-stream_test.json"
|
||||
bot.write_pending_file(123, 1, 2)
|
||||
data = json.loads(bot.pending_file.read_text(encoding="utf-8"))
|
||||
assert data["streaming"] is True
|
||||
|
||||
|
||||
# =============================================
|
||||
# 6. Batch mode unchanged
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestBatchModeUnchanged:
|
||||
"""Prove batch mode (stream=False) is byte-for-byte unchanged."""
|
||||
|
||||
def test_heartbeat_uses_edit_message_not_stream(self, tmp_path, _patch_deps):
|
||||
"""When stream=False, heartbeat calls edit_message (batch), not _stream_edit."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=False)
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def delivered_after_one():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 2
|
||||
|
||||
bot._is_pending_delivered = delivered_after_one
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
bot._start_heartbeat(123, 456)
|
||||
if bot._heartbeat_thread:
|
||||
bot._heartbeat_thread.join(timeout=35)
|
||||
bot._stream_edit.assert_not_called()
|
||||
bot.edit_message.assert_called()
|
||||
|
||||
def test_stream_false_no_transcript_tail(self, tmp_path, _patch_deps):
|
||||
"""When stream=False, _tail_transcript_bytes is never called."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=False)
|
||||
bot._tail_transcript_bytes = MagicMock()
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def delivered_after_first():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 1
|
||||
|
||||
bot._is_pending_delivered = delivered_after_first
|
||||
bot._start_heartbeat(123, 456)
|
||||
if bot._heartbeat_thread:
|
||||
bot._heartbeat_thread.join(timeout=5)
|
||||
bot._tail_transcript_bytes.assert_not_called()
|
||||
|
||||
|
||||
# =============================================
|
||||
# 7. Streaming loop integration
|
||||
# =============================================
|
||||
|
||||
|
||||
class TestStreamingLoop:
|
||||
"""Integration tests for _streaming_loop."""
|
||||
|
||||
def test_edits_with_transcript_content(self, tmp_path, _patch_deps):
|
||||
"""Streaming loop edits the processing message with transcript text."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
transcript = tmp_path / "transcript.jsonl"
|
||||
bot._active_transcript_path = transcript
|
||||
|
||||
entry = _assistant_entry([{"type": "text", "text": "Live answer"}])
|
||||
_write_transcript(transcript, [entry])
|
||||
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def stop_after_three():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 3
|
||||
|
||||
bot._is_pending_delivered = stop_after_three
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
|
||||
# Write transcript AFTER byte_offset init (simulating turn content arriving)
|
||||
original_stat = Path.stat
|
||||
|
||||
def fake_stat(self_path):
|
||||
if self_path == transcript:
|
||||
# Return 0 size initially so the loop reads from start
|
||||
return type("FakeStat", (), {"st_size": 0})()
|
||||
return original_stat(self_path)
|
||||
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
with patch.object(Path, "stat", fake_stat):
|
||||
bot._streaming_loop(123, 456, time.time())
|
||||
|
||||
assert bot._stream_edit.call_count >= 1
|
||||
edit_text = bot._stream_edit.call_args[0][2]
|
||||
assert "Live answer" in edit_text
|
||||
|
||||
def test_skips_unchanged_buffer(self, tmp_path, _patch_deps):
|
||||
"""Streaming loop does not edit when buffer hasn't changed."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
transcript = tmp_path / "transcript.jsonl"
|
||||
bot._active_transcript_path = transcript
|
||||
_write_transcript(transcript, [])
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def stop_after_three():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 3
|
||||
|
||||
bot._is_pending_delivered = stop_after_three
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
|
||||
bot._streaming_loop(123, 456, time.time())
|
||||
|
||||
# Should show "Processing..." but not re-edit identical text
|
||||
first_call_text = bot._stream_edit.call_args_list[0][0][2] if bot._stream_edit.call_count > 0 else ""
|
||||
for call_args in bot._stream_edit.call_args_list[1:]:
|
||||
assert call_args[0][2] != first_call_text or "Processing..." not in first_call_text
|
||||
|
||||
def test_4096_rollover(self, tmp_path, _patch_deps):
|
||||
"""Buffer exceeding 4096 chars triggers message rollover."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
transcript = tmp_path / "transcript.jsonl"
|
||||
bot._active_transcript_path = transcript
|
||||
|
||||
big_text = "A" * 5000
|
||||
entry = _assistant_entry([{"type": "text", "text": big_text}])
|
||||
_write_transcript(transcript, [entry])
|
||||
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
bot.send_message = MagicMock(return_value={"message_id": 999})
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def stop_after_three():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 3
|
||||
|
||||
bot._is_pending_delivered = stop_after_three
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
|
||||
with patch.object(Path, "stat", lambda s: type("S", (), {"st_size": 0})()):
|
||||
bot._streaming_loop(123, 456, time.time())
|
||||
|
||||
# Should have called send_message for the rollover
|
||||
bot.send_message.assert_called()
|
||||
|
||||
def test_429_pauses_edits(self, tmp_path, _patch_deps):
|
||||
"""After a 429, streaming loop skips edits until retry_after elapses."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
transcript = tmp_path / "transcript.jsonl"
|
||||
bot._active_transcript_path = transcript
|
||||
|
||||
entry = _assistant_entry([{"type": "text", "text": "data"}])
|
||||
_write_transcript(transcript, [entry])
|
||||
|
||||
edit_calls = [0]
|
||||
|
||||
def rate_limited_edit(_chat, _msg, _text):
|
||||
edit_calls[0] += 1
|
||||
if edit_calls[0] == 1:
|
||||
return False, 9999.0 # huge backoff
|
||||
return True, 0.0
|
||||
|
||||
bot._stream_edit = MagicMock(side_effect=rate_limited_edit)
|
||||
|
||||
tick = [0]
|
||||
|
||||
def stop_after_three():
|
||||
tick[0] += 1
|
||||
return tick[0] >= 3
|
||||
|
||||
bot._is_pending_delivered = stop_after_three
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
|
||||
with patch.object(Path, "stat", lambda s: type("S", (), {"st_size": 0})()):
|
||||
bot._streaming_loop(123, 456, time.time())
|
||||
|
||||
# Only 1 edit attempt — subsequent ticks skipped due to 429 backoff
|
||||
assert bot._stream_edit.call_count == 1
|
||||
|
||||
def test_fallback_to_batch_when_no_transcript(self, tmp_path, _patch_deps):
|
||||
"""When stream=True but no transcript path, falls back to batch heartbeat."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
bot._active_transcript_path = None # no transcript
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def delivered_after_one():
|
||||
call_count[0] += 1
|
||||
return call_count[0] >= 2
|
||||
|
||||
bot._is_pending_delivered = delivered_after_one
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
bot._start_heartbeat(123, 456)
|
||||
if bot._heartbeat_thread:
|
||||
bot._heartbeat_thread.join(timeout=35)
|
||||
|
||||
# Batch mode: edit_message called, _stream_edit not called
|
||||
bot.edit_message.assert_called()
|
||||
bot._stream_edit.assert_not_called()
|
||||
|
||||
def test_mid_loop_delivery_breaks_without_edit(self, tmp_path, _patch_deps):
|
||||
"""If _is_pending_delivered flips True mid-loop, loop breaks with no _stream_edit."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
transcript = tmp_path / "transcript.jsonl"
|
||||
bot._active_transcript_path = transcript
|
||||
|
||||
entry = _assistant_entry([{"type": "text", "text": "Should not be sent"}])
|
||||
_write_transcript(transcript, [entry])
|
||||
|
||||
delivered = [False]
|
||||
|
||||
def flip_on_second_check():
|
||||
if delivered[0]:
|
||||
return True
|
||||
delivered[0] = True
|
||||
return False
|
||||
|
||||
bot._is_pending_delivered = flip_on_second_check
|
||||
bot._tmux_session_exists = MagicMock(return_value=True)
|
||||
bot._stream_edit = MagicMock(return_value=(True, 0.0))
|
||||
bot._heartbeat_stop = threading.Event()
|
||||
|
||||
with patch.object(Path, "stat", lambda s: type("S", (), {"st_size": 0})()):
|
||||
bot._streaming_loop(123, 456, time.time())
|
||||
|
||||
bot._stream_edit.assert_not_called()
|
||||
|
||||
def test_stream_flag_wired_from_init(self, tmp_path, _patch_deps):
|
||||
"""stream=True in constructor sets self._stream."""
|
||||
bot = _make_bot(tmp_path, _patch_deps, stream=True)
|
||||
assert bot._stream is True
|
||||
|
||||
def test_stream_default_false(self, tmp_path, _patch_deps):
|
||||
"""stream defaults to False."""
|
||||
bot = _make_bot(tmp_path, _patch_deps)
|
||||
assert bot._stream is False
|
||||
Reference in New Issue
Block a user