from __future__ import annotations
import asyncio
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Iterable
from unittest.mock import patch
# The package is a workspace member, so put its source ahead of any editable
# install that may belong to another checkout.
PROJECT_SRC = (Path.cwd() / "projects" / "agent-harness" / "src").resolve()
if not PROJECT_SRC.is_dir():
raise RuntimeError(f"Expected the project source at {PROJECT_SRC}")
sys.path.insert(0, str(PROJECT_SRC))
from agent_harness.client import LLMClient
from agent_harness.config import Config
from agent_harness.events import (
StreamEvent,
StreamEventType,
TextDelta,
TokenUsage,
)Week 2: Streaming Without Silent Loss
A chat completion is not a single string. It is a time-ordered protocol whose partial messages must be assembled without losing content, tool calls, or accounting data. This chapter builds that contract from raw Server-Sent Events (SSE), then exercises the existing agent_harness.client.LLMClient against an offline stand-in for an OpenAI-compatible response. It connects the one-shot loop from Week 1 to the typed tool boundary in Week 3.
Wire contract
Server-Sent Events represent a stream as fields terminated by a blank line. A frame may carry an event name and one or more data: lines. The event name describes the kind of update; the data is the provider payload, often JSON. A client that concatenates only visible text will silently discard tool-call arguments and usage metadata.
The first implementation below is deliberately strict. It handles comments, repeated data fields, and a final frame without a trailing blank line, but it rejects unknown fields and incomplete frames. Strictness is useful at this boundary: a malformed frame should become an observable failure before it can corrupt the conversation history.
SSE parsing
The parser is intentionally independent of OpenAI or Anthropic types. That makes its invariants testable with a string and keeps transport syntax separate from message semantics. JSON decoding belongs to the next layer, because an SSE frame can legally carry non-JSON sentinels such as [DONE].
@dataclass(frozen=True)
class SSEFrame:
event: str | None
data: str
def parse_sse(text: str) -> list[SSEFrame]:
"""Parse the small SSE subset needed by a chat-completion transport."""
frames: list[SSEFrame] = []
event: str | None = None
data_lines: list[str] = []
saw_field = False
def emit_frame() -> None:
nonlocal event, data_lines, saw_field
if not saw_field:
return
if not data_lines:
raise ValueError("SSE frame has no data field")
frames.append(SSEFrame(event=event, data="\n".join(data_lines)))
event = None
data_lines = []
saw_field = False
for raw_line in text.splitlines():
line = raw_line.rstrip("\r")
if line == "":
emit_frame()
continue
if line.startswith(":"):
continue # heartbeat/comment
if ":" not in line:
raise ValueError(f"Malformed SSE line: {line!r}")
field, value = line.split(":", 1)
if field not in {"event", "data"}:
raise ValueError(f"Unsupported SSE field: {field!r}")
if value.startswith(" "):
value = value[1:]
saw_field = True
if field == "event":
event = value
else:
data_lines.append(value)
emit_frame() # handle a stream that ends without a blank line
return framesA valid frame preserves both pieces of protocol information. Repeated data: fields are joined with a newline, as required by SSE, rather than overwritten. The final call to emit_frame is what prevents a server disconnect immediately after a payload from turning into an unreported dropped frame.
raw_sse = (
": keepalive\n"
"event: delta\n"
'data: {"content":"Hel"}\n'
'data: {"content":"lo"}\n\n'
"data: [DONE]"
)
frames = parse_sse(raw_sse)
print([(frame.event, frame.data) for frame in frames])
assert frames == [
SSEFrame("delta", '{"content":"Hel"}\n{"content":"lo"}'),
SSEFrame(None, "[DONE]"),
]The failure cases are part of the API, not incidental parser behavior. An unknown field could signal a provider feature this client does not understand; a frame without data cannot be interpreted. Both should stop the run with a diagnostic instead of producing a plausible but incomplete answer.
malformed_examples = [
'data {"content":"missing colon"}\n\n',
"event: delta\n\n",
"retry: 1000\n\n",
]
for example in malformed_examples:
try:
parse_sse(example)
except ValueError as exc:
print(type(exc).__name__, str(exc))
else:
raise AssertionError("malformed SSE was accepted")
multiline = parse_sse("event: note\ndata: first\ndata: second\n\n")[0]
assert multiline.data == "first\nsecond"Delta accumulation
The raw parser only establishes frames. The existing LLMClient then consumes the SDK’s normalized chunks and emits StreamEvent values. Text deltas can be concatenated in arrival order. Tool-call deltas require a second key: the provider’s integer index. Parallel calls may be interleaved, so grouping by arrival order would merge arguments from different tools.
The next cell uses a fake SDK response, not the network. It targets the library’s _stream_response adapter so the same accumulator used by the agent is exercised while every input remains deterministic and inspectable.
class AsyncItems:
def __init__(self, items: Iterable[Any]) -> None:
self._items = iter(items)
def __aiter__(self) -> "AsyncItems":
return self
async def __anext__(self) -> Any:
try:
return next(self._items)
except StopIteration as exc:
raise StopAsyncIteration from exc
class FakeResponse:
def __init__(self, chunks: Iterable[Any]) -> None:
self._chunks = list(chunks)
def __aiter__(self) -> AsyncItems:
return AsyncItems(self._chunks)
class FakeCompletions:
def __init__(self, chunks: Iterable[Any]) -> None:
self._chunks = list(chunks)
async def create(self, **kwargs: Any) -> FakeResponse:
return FakeResponse(self._chunks)
def function_delta(
index: int,
*,
call_id: str | None = None,
name: str | None = None,
arguments: str = "",
) -> Any:
return SimpleNamespace(
index=index,
id=call_id,
function=SimpleNamespace(name=name, arguments=arguments),
)
def response_chunk(
*,
content: str | None = None,
tool_calls: list[Any] | None = None,
finish_reason: str | None = None,
usage: Any = None,
) -> Any:
return SimpleNamespace(
choices=[
SimpleNamespace(
delta=SimpleNamespace(content=content, tool_calls=tool_calls),
finish_reason=finish_reason,
)
],
usage=usage,
)
chunks = [
response_chunk(content="Reading "),
response_chunk(
tool_calls=[
function_delta(0, call_id="call-read", name="read_file", arguments='{"path":"'),
function_delta(1, call_id="call-shell", name="shell", arguments='{"command":"'),
]
),
response_chunk(
tool_calls=[
function_delta(1, arguments='pwd"}'),
function_delta(0, arguments='notes.py"}'),
]
),
response_chunk(content="the files.", finish_reason="tool_calls"),
SimpleNamespace(
choices=[],
usage=SimpleNamespace(
prompt_tokens=18,
completion_tokens=9,
total_tokens=27,
prompt_tokens_details=SimpleNamespace(cached_tokens=4),
),
),
]
fake_sdk = SimpleNamespace(
chat=SimpleNamespace(completions=FakeCompletions(chunks))
)
client = LLMClient(Config(cwd=Path.cwd()))
stream_events = [
event async for event in client._stream_response(fake_sdk, {"stream": True})
]
text = "".join(
event.text_delta.content
for event in stream_events
if event.type == StreamEventType.TEXT_DELTA and event.text_delta
)
completed_calls = [
event.tool_call
for event in stream_events
if event.type == StreamEventType.TOOL_CALL_COMPLETE
]
summary = {
"event_types": [event.type.value for event in stream_events],
"text": text,
"tool_calls": [
(call.call_id, call.name, call.arguments)
for call in completed_calls
if call is not None
],
"usage": stream_events[-1].usage.__dict__,
}
print(summary)
assert text == "Reading the files."
assert [(call.call_id, call.name, call.arguments) for call in completed_calls if call] == [
("call-read", "read_file", {"path": "notes.py"}),
("call-shell", "shell", {"command": "pwd"}),
]This trace shows why the index is part of the correctness proof. The shell arguments arrived before the final fragment of the read call, yet both completed independently. The library also keeps malformed JSON visible as a raw_arguments value rather than dropping it; the registry in Week 3 will turn that value into a structured validation result.
from agent_harness.events import parse_tool_call_arguments
malformed_arguments = parse_tool_call_arguments('{"path":')
print(malformed_arguments)
assert malformed_arguments == {"raw_arguments": '{"path":'}Token accounting
Usage is not content, but it must follow the same no-silent-loss rule. A provider may report prompt, completion, total, and cached-token counts in a trailing chunk. The ledger below records every usage report exactly once. Its cost formula is an explicit accounting assumption: prompt_tokens includes cached input, so cached tokens are priced at a separate rate and subtracted from ordinary input before billing.
@dataclass
class TokenLedger:
input_tokens: int = 0
output_tokens: int = 0
cached_tokens: int = 0
reported_total_tokens: int = 0
input_rate_per_million: float = 0.50
output_rate_per_million: float = 1.50
cached_rate_per_million: float = 0.05
def record(self, usage: TokenUsage) -> None:
values = (
usage.prompt_tokens,
usage.completion_tokens,
usage.cached_tokens,
usage.total_tokens,
)
if any(value < 0 for value in values):
raise ValueError("token counts cannot be negative")
self.input_tokens += usage.prompt_tokens
self.output_tokens += usage.completion_tokens
self.cached_tokens += usage.cached_tokens
self.reported_total_tokens += usage.total_tokens
@property
def billable_input_tokens(self) -> int:
return max(self.input_tokens - self.cached_tokens, 0)
@property
def cost(self) -> float:
return (
self.billable_input_tokens * self.input_rate_per_million
+ self.cached_tokens * self.cached_rate_per_million
+ self.output_tokens * self.output_rate_per_million
) / 1_000_000
def snapshot(self) -> dict[str, int | float]:
return {
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"cached_tokens": self.cached_tokens,
"reported_total_tokens": self.reported_total_tokens,
"billable_input_tokens": self.billable_input_tokens,
"cost": round(self.cost, 8),
}
ledger = TokenLedger()
usage_reports = [
TokenUsage(prompt_tokens=120, completion_tokens=30, total_tokens=150, cached_tokens=20),
TokenUsage(prompt_tokens=80, completion_tokens=10, total_tokens=90, cached_tokens=0),
]
for report in usage_reports:
ledger.record(report)
print(ledger.snapshot())
assert ledger.input_tokens == sum(report.prompt_tokens for report in usage_reports)
assert ledger.output_tokens == sum(report.completion_tokens for report in usage_reports)
assert ledger.reported_total_tokens == sum(report.total_tokens for report in usage_reports)
assert ledger.cached_tokens == sum(report.cached_tokens for report in usage_reports)The reported total is a reconciliation field, not a value to recompute and silently substitute. If it disagrees with the component counts, retain both values and investigate the provider contract. Likewise, a retry that fails before receiving a usage report contributes no tokens here; a successful retry contributes its one report once. This separation prevents a retry storm from becoming accounting drift.
Retry policy
The current library retries RateLimitError and APIConnectionError with delays 1, 2, 4, ... seconds and returns a typed StreamEventType.ERROR after the final attempt. That is a useful baseline, but synchronized clients can create a retry storm. A production policy adds a cap and jitter, retries only errors known to be transient, and considers request idempotency before repeating a side effect.
The schedule below uses fixed jitter fractions so the experiment is reproducible. In a live client, the fractions would come from a seeded or cryptographically appropriate random source according to the reproducibility requirements of the run.
@dataclass(frozen=True)
class RetryPolicy:
max_retries: int = 3
base_delay: float = 1.0
max_delay: float = 8.0
def delays(self, jitter_fractions: Iterable[float] | None = None) -> list[float]:
fractions = list(jitter_fractions or [0.0] * self.max_retries)
if len(fractions) != self.max_retries:
raise ValueError("one jitter fraction is required per retry")
delays: list[float] = []
for retry_number, fraction in enumerate(fractions):
if fraction <= -1:
raise ValueError("jitter would make the delay non-positive")
exponential = min(self.base_delay * (2**retry_number), self.max_delay)
delays.append(round(exponential * (1 + fraction), 6))
return delays
policy = RetryPolicy()
naive_schedule = policy.delays([0.0, 0.0, 0.0])
jittered_schedule = policy.delays([-0.25, 0.10, -0.15])
print({"naive_seconds": naive_schedule, "jittered_seconds": jittered_schedule})
assert naive_schedule == [1.0, 2.0, 4.0]
assert jittered_schedule == [0.75, 2.2, 3.4]The library’s retry behavior can be tested without sleeping or contacting a provider. The fake response raises a real OpenAI exception twice, then returns the same typed completion event that a successful stream would produce. Patching only the client’s sleep function records the policy while keeping this test fast.
import httpx
from openai import RateLimitError
import agent_harness.client as client_module
def make_rate_limit_error() -> RateLimitError:
request = httpx.Request("POST", "https://offline.invalid/v1/chat/completions")
response = httpx.Response(429, request=request)
return RateLimitError("offline rate limit", response=response, body=None)
class OfflineRetryClient(LLMClient):
def __init__(self, config: Config) -> None:
super().__init__(config)
self.calls = 0
def get_client(self) -> Any:
return object()
async def _stream_response(self, client: Any, kwargs: dict[str, Any]):
self.calls += 1
if self.calls < 3:
raise make_rate_limit_error()
yield StreamEvent(
type=StreamEventType.MESSAGE_COMPLETE,
usage=TokenUsage(prompt_tokens=2, completion_tokens=1, total_tokens=3),
)
retry_client = OfflineRetryClient(Config(cwd=Path.cwd()))
waits: list[float] = []
async def fake_sleep(seconds: float) -> None:
waits.append(seconds)
with patch.object(client_module.asyncio, "sleep", fake_sleep):
retry_events = [
event async for event in retry_client.chat_completion([], stream=True)
]
print({
"attempts": retry_client.calls,
"recorded_waits": waits,
"events": [event.type.value for event in retry_events],
})
assert retry_client.calls == 3
assert waits == [1, 2]
assert [event.type for event in retry_events] == [StreamEventType.MESSAGE_COMPLETE]Timeouts are a separate boundary from retries. A request deadline limits connection and header acquisition; a stream deadline limits how long an established response may remain open. Retrying a timed-out request is safe only when the provider treats the request as idempotent. The helper below demonstrates the control flow with an immediate stream and a never-finishing stream.
@dataclass(frozen=True)
class TimeoutPolicy:
request_seconds: float = 10.0
stream_seconds: float = 60.0
async def collect_stream(stream: Any) -> list[StreamEvent]:
return [event async for event in stream]
async def immediate_stream():
yield StreamEvent(type=StreamEventType.MESSAGE_COMPLETE)
async def never_finishes():
await asyncio.Future()
if False: # keep this function an async generator for the test harness
yield StreamEvent(type=StreamEventType.MESSAGE_COMPLETE)
async def timeout_experiment() -> dict[str, bool | dict[str, float]]:
policy = TimeoutPolicy()
completed = await asyncio.wait_for(
collect_stream(immediate_stream()), timeout=policy.stream_seconds
)
timed_out = False
try:
await asyncio.wait_for(
collect_stream(never_finishes()), timeout=0.001
)
except TimeoutError:
timed_out = True
return {
"policy": policy.__dict__,
"immediate_completed": len(completed) == 1,
"never_finishes_timed_out": timed_out,
}
timeout_result = await timeout_experiment()
print(timeout_result)
assert timeout_result["immediate_completed"] is True
assert timeout_result["never_finishes_timed_out"] is TrueLatency and reliability
Network measurements would depend on provider load, DNS, and the machine running the notebook. To keep the mechanism experiment deterministic, use a fixed transport model instead: first-token latency grows with payload size, while total latency adds one fixed cost per emitted chunk. This is not a benchmark; it is a controlled reminder that reducing total latency can still worsen time to first token if the client buffers too aggressively.
def synthetic_latency(payload_bytes: int, chunks: int) -> dict[str, float | int]:
first_token_ms = 30.0 + 0.02 * payload_bytes
total_ms = first_token_ms + 12.0 * chunks
return {
"payload_bytes": payload_bytes,
"chunks": chunks,
"time_to_first_token_ms": round(first_token_ms, 2),
"total_latency_ms": round(total_ms, 2),
}
profiles = [
synthetic_latency(256, 4),
synthetic_latency(4096, 16),
synthetic_latency(16_384, 64),
]
print(profiles)
assert profiles[0]["time_to_first_token_ms"] == 35.12
assert profiles[-1]["total_latency_ms"] == 1125.68The reliability claim for this week is narrow: a client must not silently drop protocol fields. The strict parser makes malformed transport visible; the indexed accumulator keeps parallel tool calls distinct; the ledger reconciles usage; and the retry test proves that a transient failure is observable rather than duplicated into the conversation. The current library is a solid baseline, with jitter, public timeout controls, and a provider-side billing reconciliation left as explicit hardening work.
Next, Week 3 turns accumulated tool-call arguments into validated calls. Week 4 gives those calls filesystem and process effects, where exact contracts matter even more.