fix(backends,markdown): stream turn events live, no blank lines for tool blocks
This commit is contained in:
@@ -296,15 +296,21 @@ class ClaudeSdkBackend:
|
||||
live = await self._acquire(key, session_id=session_id, history=prior, spec=spec)
|
||||
message_id = f"msg_{uuid.uuid4().hex}"
|
||||
yield build_message_start(message_id=message_id, model=self._agent.model)
|
||||
turn = _Turn()
|
||||
async with live.lock:
|
||||
live.running_turn = turn_id or message_id
|
||||
live.last_used = time.monotonic()
|
||||
try:
|
||||
turn = await self._run_turn(live, prompt, observer=observer)
|
||||
# Events go out as the CLI produces them: the frontends
|
||||
# stream text and thinking live, the turn is not buffered.
|
||||
async for event in self._run_turn(live, prompt, turn, observer):
|
||||
yield event
|
||||
except Exception:
|
||||
live.running_turn = None
|
||||
await self._pool.close(key)
|
||||
if not (live.resumed and live.turns == 0):
|
||||
# A dead resume can be reseeded from history, but only
|
||||
# while nothing of this turn has reached the caller yet.
|
||||
if not (live.resumed and live.turns == 0) or turn.events:
|
||||
raise
|
||||
_log.exception(
|
||||
"resume of %s failed, reseeding from history", live.session_id
|
||||
@@ -312,10 +318,10 @@ class ClaudeSdkBackend:
|
||||
live = await self._acquire(
|
||||
key, session_id=None, history=prior, spec=spec
|
||||
)
|
||||
turn = _Turn()
|
||||
async with live.lock:
|
||||
live.running_turn = turn_id or message_id
|
||||
turn = await self._run_turn(live, prompt, observer=observer)
|
||||
for event in turn.events:
|
||||
async for event in self._run_turn(live, prompt, turn, observer):
|
||||
yield event
|
||||
live.turns += 1
|
||||
live.last_used = time.monotonic()
|
||||
@@ -373,11 +379,16 @@ class ClaudeSdkBackend:
|
||||
self,
|
||||
live: Session,
|
||||
prompt: str,
|
||||
*,
|
||||
turn: _Turn,
|
||||
observer: Callable[[Any], None] | None = None,
|
||||
) -> _Turn:
|
||||
) -> AsyncIterator[MessageStreamEvent]:
|
||||
"""Run one prompt, yielding wire events as they arrive.
|
||||
|
||||
``turn`` is filled in place (result, synthesized history, count of
|
||||
events already yielded) so the caller can finish bookkeeping - and
|
||||
decide whether a retry is still possible - after a failure.
|
||||
"""
|
||||
streaming = self._agent.options.include_partial_messages
|
||||
turn = _Turn()
|
||||
raw: list[Any] = []
|
||||
next_index = 0
|
||||
offset = 0
|
||||
@@ -404,12 +415,16 @@ class ClaudeSdkBackend:
|
||||
if isinstance(index, int):
|
||||
next_index = max(next_index, offset + index + 1)
|
||||
if streaming:
|
||||
turn.events.extend(_emit_stream_event(event, offset + index))
|
||||
for out in _emit_stream_event(event, offset + index):
|
||||
turn.events += 1
|
||||
yield out
|
||||
elif isinstance(message, AssistantMessage):
|
||||
raw.append(message)
|
||||
if not streaming:
|
||||
for block in message.content:
|
||||
turn.events.extend(_emit_block(block, next_index))
|
||||
for out in _emit_block(block, next_index):
|
||||
turn.events += 1
|
||||
yield out
|
||||
next_index += 1
|
||||
elif isinstance(message, UserMessage):
|
||||
raw.append(message)
|
||||
@@ -423,11 +438,10 @@ class ClaudeSdkBackend:
|
||||
"turn: agent=%s session=%s events=%d synthesized=%d stop=%s",
|
||||
self._agent.name,
|
||||
live.session_id,
|
||||
len(turn.events),
|
||||
turn.events,
|
||||
len(turn.synthesized),
|
||||
turn.stop_reason,
|
||||
)
|
||||
return turn
|
||||
|
||||
async def _acquire(
|
||||
self,
|
||||
@@ -606,7 +620,8 @@ class _SessionSpec:
|
||||
|
||||
@dataclass
|
||||
class _Turn:
|
||||
events: list[Any] = field(default_factory=list)
|
||||
events: int = 0
|
||||
"""Wire events already yielded to the caller."""
|
||||
synthesized: list[dict[str, Any]] = field(default_factory=list)
|
||||
result: ResultMessage | None = None
|
||||
stop_reason: StopReason = "end_turn"
|
||||
|
||||
@@ -33,7 +33,9 @@ __all__ = [
|
||||
# Empty ``### User:`` block appended after each assistant reply so the
|
||||
# human has an obvious place to type the next turn. Parser drops empty
|
||||
# user blocks, so this doesn't re-trigger dispatch on its own.
|
||||
USER_SCAFFOLD = "### User:\n"
|
||||
# Blank line after the header, like every rendered turn - the file stays
|
||||
# symmetric whether the human or the gateway wrote the marker.
|
||||
USER_SCAFFOLD = "### User:\n\n"
|
||||
|
||||
|
||||
# Default 4-backtick fence so tool results that contain literal ```` ``` ````
|
||||
@@ -73,7 +75,12 @@ def render_assistant_message(message: Message) -> str:
|
||||
"""
|
||||
parts: list[str] = ["### Assistant:", ""]
|
||||
for block in message.content:
|
||||
parts.extend(_render_block(block))
|
||||
lines = list(_render_block(block))
|
||||
if not lines:
|
||||
# Tool calls render to nothing - no separator for them either,
|
||||
# or every tool leaves a blank line behind.
|
||||
continue
|
||||
parts.extend(lines)
|
||||
parts.append("")
|
||||
return "\n".join(parts).rstrip() + "\n"
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -419,3 +420,41 @@ async def test_close_disconnects(cwd: Path) -> None:
|
||||
)
|
||||
assert FakeClient.instances[0].connected is False
|
||||
assert backend.sessions == {}
|
||||
|
||||
|
||||
class GatedClient(FakeClient):
|
||||
"""Streams one delta, then waits for ``gate`` before finishing the turn."""
|
||||
|
||||
gate: asyncio.Event
|
||||
|
||||
async def receive_response(self):
|
||||
yield StreamEvent(
|
||||
uuid="u", session_id="s", event={"type": "message_start", "message": {}}
|
||||
)
|
||||
for e in _stream(0, "first"):
|
||||
yield e
|
||||
await GatedClient.gate.wait()
|
||||
yield AssistantMessage(content=[TextBlock(text="first")], model="m")
|
||||
yield _result(self.session_id)
|
||||
|
||||
|
||||
async def test_deltas_reach_the_caller_before_the_turn_ends(cwd: Path) -> None:
|
||||
"""The turn is not buffered: a delta is observable while the CLI still runs."""
|
||||
GatedClient.gate = asyncio.Event()
|
||||
backend = _backend(cwd, InMemorySessionStore())
|
||||
backend._factory = GatedClient
|
||||
events = backend.complete(
|
||||
agent=backend.agent,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
conversation_id="conv-live",
|
||||
)
|
||||
seen: list[Any] = []
|
||||
async for ev in events:
|
||||
seen.append(ev)
|
||||
if isinstance(ev, RawContentBlockDeltaEvent):
|
||||
break
|
||||
assert seen[-1].delta.text == "first"
|
||||
assert not GatedClient.gate.is_set()
|
||||
GatedClient.gate.set()
|
||||
rest = [e async for e in events]
|
||||
assert isinstance(rest[-1], RawMessageStopEvent)
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
from anthropic.types import Message, TextBlock, ThinkingBlock, ToolUseBlock, Usage
|
||||
|
||||
from beaver_gateway.frontends.markdown import renderer
|
||||
|
||||
|
||||
def _message(*blocks: object) -> Message:
|
||||
return Message(
|
||||
id="m",
|
||||
type="message",
|
||||
role="assistant",
|
||||
model="x",
|
||||
content=list(blocks), # type: ignore[arg-type]
|
||||
stop_reason="end_turn",
|
||||
stop_sequence=None,
|
||||
usage=Usage(input_tokens=1, output_tokens=1),
|
||||
)
|
||||
|
||||
|
||||
def test_tool_calls_leave_no_blank_lines() -> None:
|
||||
message = _message(
|
||||
TextBlock(type="text", text="first"),
|
||||
ToolUseBlock(type="tool_use", id="t1", name="Bash", input={"command": "ls"}),
|
||||
ToolUseBlock(type="tool_use", id="t2", name="Read", input={"file_path": "x"}),
|
||||
ToolUseBlock(type="tool_use", id="t3", name="Grep", input={"pattern": "y"}),
|
||||
TextBlock(type="text", text="second"),
|
||||
)
|
||||
assert renderer.render_assistant_message(message) == (
|
||||
"### Assistant:\n\nfirst\n\nsecond\n"
|
||||
)
|
||||
|
||||
|
||||
def test_thinking_renders_as_collapsed_callout() -> None:
|
||||
message = _message(
|
||||
ThinkingBlock(type="thinking", thinking="a\nb", signature="s"),
|
||||
TextBlock(type="text", text="answer"),
|
||||
)
|
||||
assert renderer.render_assistant_message(message) == (
|
||||
"### Assistant:\n\n> [!thinking]-\n> a\n> b\n\nanswer\n"
|
||||
)
|
||||
|
||||
|
||||
def test_user_scaffold_matches_rendered_user_turn_shape() -> None:
|
||||
body = renderer.append_to_body("### Assistant:\n\nhi\n", renderer.USER_SCAFFOLD)
|
||||
assert body.endswith("\n\n---\n\n### User:\n\n")
|
||||
assert renderer.render_user_text("q").startswith("### User:\n\n")
|
||||
Reference in New Issue
Block a user