Compare commits

...

2 Commits

Author SHA1 Message Date
harry-yao_data 02a168b3ba perf(claude-native): mirror a run of transcript items in one request
The delta forwarder now batches, but transcript items still cost a
request each — and they are posted sequentially inside a single poll
iteration. A poll that finds thirty items therefore spends thirty round
trips before the loop comes back around, which from another region both
delays the tool cards and starves the streamed-text tail queued behind
them.

Send the leading run of plain items through the batch route. The loop
itself is unchanged: it consumes the outcome for an item the batch
attempted and posts anything the batch did not cover, so ordering,
cursor advancement, retries and dead-lettering behave exactly as they
did per item. A rejection is rebuilt into the exception a single post
would have raised, so it flows through one classification path rather
than acquiring a parallel set of rules.

The run is capped at eight because a batch shares one fate in the rare
case where its response is lost: any prefix may be committed, external
items are not deduped, so none can be re-posted. That is what a single
ambiguous post already does for one item — and those drops are now
dead-lettered, which they were not before, so the enlarged case is
recoverable instead of silent. A request that provably never left is
unambiguous and falls back to per-item posts, losing nothing.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: harry-yao_data <harry.yao@databricks.com>
2026-08-22 02:52:48 +00:00
harry-yao_data ae13397e5f perf(sessions): batch native-forwarder event posts
A native harness's forwarder mirrors the session one HTTP POST at a
time — a request per transcript item, a request per streamed text
chunk. On loopback that is invisible. From another region it caps the
forwarder at roughly one event per round trip, so a turn that the pane
finished in seconds keeps trickling into the web UI for tens of
seconds.

Add `POST /v1/sessions/{id}/events/batch`, which runs a run of events
through the same handler as the single-post route, in order, and
reports each one's outcome so a caller can advance a durable cursor
over the delivered prefix instead of re-sending events that already
landed. Point the claude-native delta forwarder at it: a poll's chunks
now cost one round trip instead of one each, with a per-event fallback
for deployments older than the route.

At `--network-delay-ms 100`, mirroring one poll of streamed text drops
from ~2.52s to ~0.14s (24 requests to 1); the two new
`forward_native_deltas_*` benchmark journeys are the matched pair that
measures it and guards the regression.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: harry-yao_data <harry.yao@databricks.com>
2026-08-22 02:52:37 +00:00
10 changed files with 2016 additions and 89 deletions
+9
View File
@@ -62,6 +62,15 @@ see *Network* below).
| `add_comment` | `POST /v1/sessions/{id}/comments` — create a review comment | write path |
| `list_projects` | `GET /v1/sessions/projects` — sidebar project list (dual-read union) | project count |
| `list_project_sessions` | `GET /v1/sessions?project=` — a project folder's sessions (dual-read filter) | sessions/project |
| `forward_native_deltas_per_event` | `POST /v1/sessions/{id}/events` × 24 — a native forwarder mirroring one poll's streamed text one request at a time | round-trip count |
| `forward_native_deltas_batched` | `POST /v1/sessions/{id}/events/batch` — the same poll's text in one request | round-trip count |
The two `forward_native_deltas_*` journeys are a matched pair: same work, 24
round trips versus 1. Run them with `--network-delay-ms` to price the wire,
which is what a user in another region actually pays — at `--network-delay-ms
100` the per-event journey takes ~2.5s per poll against ~0.14s batched, and a
turn is many polls. They are also the regression guard: if the forwarder ever
fans back out into per-event posts, the batched journey's `HTTP/op` climbs off 1.
Read journeys target a **pre-seeded** session when the DB has a corpus; against
an empty DB they self-seed a small fallback session over HTTP (the
+88
View File
@@ -51,6 +51,7 @@ import subprocess
import sys
import tempfile
import time
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Literal, cast
@@ -789,6 +790,68 @@ async def _measure_policy_evaluate(env: BenchEnvironment, ctx: JourneyContext) -
resp.raise_for_status()
# ── native-forwarder event mirroring (per-event vs batched) ──────────────────
# Chunks in one simulated poll's worth of streamed assistant text. A native
# harness produces this many in well under a second, and the forwarder used to
# spend one round trip on each — which is why a far-from-server client's live
# view fell behind. Pair these two journeys with ``--network-delay-ms`` to see
# the round-trip cost directly: the per-event journey pays it 24 times, the
# batched one once.
_FORWARD_DELTA_COUNT = 24
# One op is 24 requests for the per-event journey, so the default 100 iterations
# would let this pair dominate the CI leg (and add tail noise to the journeys
# after it). The pair exists to price round trips, which a handful of samples
# establishes.
_FORWARD_DELTA_MAX_ITERATIONS = 20
def _forward_delta_events(message_id: str) -> list[dict[str, object]]:
"""
Build one poll's worth of streamed-text events.
:param message_id: Assistant message the chunks belong to, so the
server scopes them to one in-flight buffer.
:returns: ``external_output_text_delta`` event bodies, in order.
"""
return [
{
"type": "external_output_text_delta",
"data": {
"delta": f"chunk-{index} ",
"message_id": message_id,
"index": index,
"final": index == _FORWARD_DELTA_COUNT - 1,
},
}
for index in range(_FORWARD_DELTA_COUNT)
]
async def _measure_forward_deltas_per_event(env: BenchEnvironment, ctx: JourneyContext) -> None:
"""Mirror a poll's chunks one request at a time (the pre-batching path)."""
session_id = cast(str, ctx) # _setup_target_session
assert env.client is not None
for event in _forward_delta_events(f"bench-{uuid.uuid4().hex}"):
resp = await env.client.post(f"/v1/sessions/{session_id}/events", json=event)
resp.raise_for_status()
async def _measure_forward_deltas_batched(env: BenchEnvironment, ctx: JourneyContext) -> None:
"""Mirror the same chunks in one request."""
session_id = cast(str, ctx) # _setup_target_session
assert env.client is not None
resp = await env.client.post(
f"/v1/sessions/{session_id}/events/batch",
json={
"events": _forward_delta_events(f"bench-{uuid.uuid4().hex}"),
"on_error": "continue",
},
)
resp.raise_for_status()
# ── CLI startup (omnigent polly against the local bench server) ──────────────
# Signal that the REPL is ready — the last spinner message before the prompt.
@@ -1015,6 +1078,31 @@ ALL_JOURNEYS: dict[str, Journey] = {
description="POST /v1/sessions/{id}/policies/evaluate — PreToolUse hook "
"(single tree scan, preloaded conversation row, caches warm).",
),
Journey(
name="forward_native_deltas_per_event",
kind="latency",
measure=_measure_forward_deltas_per_event,
setup=_setup_target_session,
concurrency_safe=True,
max_iterations=_FORWARD_DELTA_MAX_ITERATIONS,
description=(
f"POST /v1/sessions/{{id}}/events × {_FORWARD_DELTA_COUNT} — a native "
"forwarder mirroring one poll's streamed text one request at a time."
),
),
Journey(
name="forward_native_deltas_batched",
kind="latency",
measure=_measure_forward_deltas_batched,
setup=_setup_target_session,
concurrency_safe=True,
max_iterations=_FORWARD_DELTA_MAX_ITERATIONS,
description=(
"POST /v1/sessions/{id}/events/batch — the same poll's streamed text "
"in one request. Compare against forward_native_deltas_per_event "
"under --network-delay-ms."
),
),
# Runner (full-turn) journeys — with_runner=True, openai-agents, mock LLM.
Journey(
name="session_cold_start",
+194
View File
@@ -23,6 +23,7 @@ import contextlib
import json
import logging
import time
import weakref
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from pathlib import Path
@@ -218,6 +219,199 @@ async def post_external_session_status(
resp.raise_for_status()
# Batch event ingestion (``POST /v1/sessions/{id}/events/batch``). A forwarder
# far from the server is otherwise capped at one event per round trip, so a
# turn's live text trickles in for tens of seconds after the harness produced
# it. Batching collapses a poll's worth of events into one request.
#
# Kept at half the server's ``MAX_SESSION_EVENTS_PER_BATCH`` (256) so a client
# never trips the server's envelope validation, and so one request stays small
# enough to retry cheaply. Longer runs are chunked.
MAX_EVENTS_PER_BATCH = 128
# Clients whose server has no batch route. Latched on the first route-miss 404,
# so talking to an older deployment costs one wasted request per client, not one
# per event. Keyed by client (weakly) rather than by base URL: a forwarder holds
# one client for its whole life, and nothing leaks into an unrelated client that
# happens to share a server.
_EVENTS_BATCH_UNSUPPORTED: weakref.WeakSet[httpx.AsyncClient] = weakref.WeakSet()
def events_batch_supported(client: httpx.AsyncClient) -> bool:
"""
Report whether this server is still believed to serve the batch route.
:param client: Omnigent HTTP client.
:returns: ``False`` once a route-miss 404 latched the fallback.
"""
return client not in _EVENTS_BATCH_UNSUPPORTED
def reset_events_batch_support(client: httpx.AsyncClient | None = None) -> None:
"""
Clear the batch-unsupported latch (tests; and after a server upgrade).
:param client: Client to clear, or ``None`` to clear every latch.
:returns: None.
"""
if client is None:
_EVENTS_BATCH_UNSUPPORTED.clear()
return
_EVENTS_BATCH_UNSUPPORTED.discard(client)
def _is_route_miss(response: httpx.Response) -> bool:
"""
Distinguish "this server has no such route" from "no such session".
Omnigent's error handler shapes application 404s as
``{"error": {...}}``; Starlette's route miss answers
``{"detail": "Not Found"}``. Only the latter means the deployment
predates the batch endpoint.
:param response: The 404 response to classify.
:returns: ``True`` when the 404 came from routing, not the handler.
"""
try:
body = response.json()
except ValueError:
return True
return not (isinstance(body, dict) and "error" in body)
@dataclass(frozen=True)
class BatchedEventOutcome:
"""
Per-event outcome of one batched session-event POST.
:param index: Position of the event in the submitted run.
:param delivered: ``True`` when the server accepted the event (the
status the equivalent single post would have returned was 2xx).
:param status: That status, or ``None`` when the event was never
attempted because an earlier failure stopped the batch.
:param error: Failure reason reported by the server, when present.
"""
index: int
delivered: bool
status: int | None = None
error: str | None = None
def _parse_batch_results(
payload: object,
*,
total: int,
offset: int,
) -> list[BatchedEventOutcome]:
"""
Convert one batch response body into per-event outcomes.
Events the server never attempted (it stopped at a failure) are
reported as not-delivered with a ``None`` status so callers keep them
for the next attempt instead of treating silence as success.
:param payload: Decoded JSON body of the batch response.
:param total: Number of events submitted in this request.
:param offset: Index of this request's first event within the caller's
full run, so outcomes are numbered in the caller's terms.
:returns: One outcome per submitted event, in order.
"""
raw_results = payload.get("results") if isinstance(payload, dict) else None
by_index: dict[int, BatchedEventOutcome] = {}
if isinstance(raw_results, list):
for position, entry in enumerate(raw_results):
if not isinstance(entry, dict):
continue
raw_index = entry.get("index")
index = raw_index if isinstance(raw_index, int) else position
raw_status = entry.get("status")
status = raw_status if isinstance(raw_status, int) else None
raw_error = entry.get("error")
by_index[index] = BatchedEventOutcome(
index=offset + index,
delivered=status is not None and 200 <= status < 400,
status=status,
error=raw_error if isinstance(raw_error, str) else None,
)
return [
by_index.get(index, BatchedEventOutcome(index=offset + index, delivered=False))
for index in range(total)
]
async def post_session_events_batch(
client: httpx.AsyncClient,
*,
session_id: str,
events: list[dict[str, object]],
on_error: str = "stop",
) -> list[BatchedEventOutcome] | None:
"""
POST a run of session events in as few round trips as possible.
Each event is the same ``{"type": ..., "data": ...}`` body a single
``POST /v1/sessions/{id}/events`` would carry, and the server
dispatches them in order through that same handler — so this changes
only how many times the wire is crossed. That is the whole point: on
a client far from the server, per-event posting caps the forwarder at
a few events per second and its live view falls behind the harness.
Runs longer than :data:`MAX_EVENTS_PER_BATCH` are chunked. With
``on_error="stop"`` a chunk that ends in a failure stops the run
there, so a caller advancing a cursor over the delivered prefix
behaves exactly as it did when posting one event at a time.
:param client: Omnigent HTTP client.
:param session_id: Omnigent session/conversation id.
:param events: Event bodies to dispatch, in order. Empty is a no-op.
:param on_error: ``"stop"`` to leave the rest of a chunk unattempted
after a failure (the default, matching a sequential caller), or
``"continue"`` to attempt every event — for best-effort streams
where dropping one chunk beats stalling the tail.
:returns: One outcome per event, in order; or ``None`` when this
server has no batch route and the caller should post singly.
:raises httpx.HTTPError: On a transport failure or an
envelope-level rejection. Delivery of the events in the failed
chunk is then unknown, exactly as for a single post whose
response was lost.
"""
if not events:
return []
if not events_batch_supported(client):
return None
outcomes: list[BatchedEventOutcome] = []
for offset in range(0, len(events), MAX_EVENTS_PER_BATCH):
chunk = events[offset : offset + MAX_EVENTS_PER_BATCH]
response = await client.post(
f"/v1/sessions/{session_id}/events/batch",
json={"events": chunk, "on_error": on_error},
)
if response.status_code == 404 and _is_route_miss(response):
_EVENTS_BATCH_UNSUPPORTED.add(client)
_logger.debug(
"Server has no session-events batch route; falling back to per-event posts for %s",
client.base_url,
)
return None
response.raise_for_status()
note_native_post_success()
try:
payload = response.json()
except ValueError as exc:
raise httpx.HTTPError(f"malformed session-events batch response: {exc}") from exc
chunk_outcomes = _parse_batch_results(payload, total=len(chunk), offset=offset)
outcomes.extend(chunk_outcomes)
if on_error == "stop" and not all(outcome.delivered for outcome in chunk_outcomes):
# The server stopped at a failure; everything after it in this
# chunk was never attempted, and later chunks must not jump the
# queue ahead of it.
for index in range(offset + len(chunk), len(events)):
outcomes.append(BatchedEventOutcome(index=index, delivered=False))
break
return outcomes
async def post_session_event_with_retry(
*,
client: httpx.AsyncClient,
+360 -30
View File
@@ -20,6 +20,7 @@ from omnigent._native_post_delivery import (
append_dead_letter,
post_external_session_status,
post_may_have_been_delivered,
post_session_events_batch,
)
from omnigent.claude_native_bridge import (
BRIDGE_ID_LABEL_KEY,
@@ -3274,6 +3275,189 @@ async def _handle_compact_summary_item(
return True
# Transcript items per batch. A batch's items share one fate only in the rare
# case where its response is lost: they may or may not have been committed, so
# none can be safely re-posted and all are dead-lettered together — which is
# what a single ambiguous post does today, just for one item. A small cap keeps
# that blast radius near today's while still collapsing a burst's round trips,
# which is what matters: a poll that dumps thirty items used to spend thirty
# round trips inside one iteration, starving the delta stream behind it.
_ITEM_BATCH_MAX = 8
def _item_event(item: ClaudeTranscriptItem) -> dict[str, object]:
"""
Build the ``external_conversation_item`` event body for one item.
:param item: Transcript-derived conversation item.
:returns: The event body a single or batched post carries.
"""
return {
"type": "external_conversation_item",
"data": {
"item_type": item.item_type,
"item_data": item.data,
"response_id": item.response_id,
},
}
@dataclass(frozen=True)
class _AttemptedItem:
"""
Outcome of an item a batch already attempted.
:param error: ``None`` when the server accepted it. Otherwise the
exception the equivalent single post would have raised, so the
caller's existing failure handling classifies it identically —
a synthesized status error for a rejection, or the batch's real
transport error when the response was lost.
"""
error: httpx.HTTPError | None
def _synthesized_status_error(
status: int,
*,
session_id: str,
message: str | None,
code: str | None,
) -> httpx.HTTPStatusError:
"""
Rebuild the exception a single post would have raised for one status.
Lets a batched rejection flow through exactly the same permanent /
transient / ambiguous classification as a per-item post, instead of a
parallel set of rules that could drift.
:param status: HTTP status the server reported for the event.
:param session_id: Session the post was for (URL reconstruction only).
:param message: Server-reported failure message, if any.
:param code: Server-reported Omnigent error code, if any.
:returns: An ``httpx.HTTPStatusError`` carrying that status.
"""
request = httpx.Request("POST", f"/v1/sessions/{session_id}/events")
response = httpx.Response(
status,
request=request,
json={"error": {"code": code or "", "message": message or ""}},
)
return httpx.HTTPStatusError(
f"batched session event rejected with {status}",
request=request,
response=response,
)
def _batchable_item_run(
items: list[ClaudeTranscriptItem],
*,
seen: set[str],
skip_user_messages: bool,
retry_tracker: _PostRetryTracker,
limit: int = _ITEM_BATCH_MAX,
) -> list[ClaudeTranscriptItem]:
"""
Take the leading run of items that can go out together.
Ordering is load-bearing, so the run stops at the first item the caller
handles specially: a compaction boundary (its own persist path, which
must land before later items) or an item still inside its retry
backoff. Items already forwarded, and user messages this session skips,
post nothing at all — they neither join the run nor end it.
:param items: Items read from the transcript this poll, in order.
:param seen: Source ids already forwarded.
:param skip_user_messages: Whether user messages are mirrored.
:param retry_tracker: Per-item retry/backoff state.
:param limit: Most items to take.
:returns: The leading batchable items, in order (possibly empty).
"""
run: list[ClaudeTranscriptItem] = []
for item in items:
if item.source_id in seen:
continue
if item.is_compact_summary:
break
if skip_user_messages and item.item_type == "message" and item.data.get("role") == "user":
continue
if retry_tracker.retry_delay_s(f"item:{item.source_id}") is not None:
break
run.append(item)
if len(run) >= limit:
break
return run
async def _attempt_item_batch(
client: httpx.AsyncClient,
*,
session_id: str,
items: list[ClaudeTranscriptItem],
seen: set[str],
skip_user_messages: bool,
retry_tracker: _PostRetryTracker,
) -> dict[str, _AttemptedItem] | None:
"""
Forward the leading run of items in one request, if that helps.
:param client: Omnigent HTTP client.
:param session_id: Omnigent session/conversation id.
:param items: Items read from the transcript this poll, in order.
:param seen: Source ids already forwarded.
:param skip_user_messages: Whether user messages are mirrored.
:param retry_tracker: Per-item retry/backoff state.
:returns: Outcomes by source id for the items the batch attempted, or
``None`` when nothing was attempted — a run too short to be worth a
batch, a server without the batch route, or a request that provably
never left, all of which leave the caller to post item by item.
"""
run = _batchable_item_run(
items,
seen=seen,
skip_user_messages=skip_user_messages,
retry_tracker=retry_tracker,
)
if len(run) < 2:
return None
try:
outcomes = await post_session_events_batch(
client,
session_id=session_id,
events=[_item_event(item) for item in run],
on_error="stop",
)
except httpx.HTTPError as exc:
if not post_may_have_been_delivered(exc):
# Proven undelivered (the request never left): nothing landed, so
# the caller may post these normally and retry as it always has.
return None
# The response was lost. Any prefix of the run may be committed and
# external items are not deduped, so re-posting would duplicate
# bubbles; hand every item back as ambiguous and let the caller's
# existing ambiguous path dead-letter them.
return {item.source_id: _AttemptedItem(error=exc) for item in run}
if outcomes is None:
return None
attempted: dict[str, _AttemptedItem] = {}
for item, outcome in zip(run, outcomes, strict=False):
if outcome.delivered:
attempted[item.source_id] = _AttemptedItem(error=None)
elif outcome.status is not None:
attempted[item.source_id] = _AttemptedItem(
error=_synthesized_status_error(
outcome.status,
session_id=session_id,
message=outcome.error,
code=None,
)
)
# An item the batch stopped short of is absent, so the caller posts it
# itself — the same order, one request later.
return attempted
async def _forward_available_items(
*,
client: httpx.AsyncClient,
@@ -3338,6 +3522,18 @@ async def _forward_available_items(
# their own ``response_id`` (see :func:`_post_external_conversation_item`),
# so the transcript's job here is items, not status.
updated = state
# Forward the leading run of plain items in one request. Everything below is
# unchanged: the loop consumes the outcome for an item the batch attempted
# and posts the rest itself, so ordering, cursor advancement, retries and
# dead-lettering all behave as they did per item.
batch = await _attempt_item_batch(
client,
session_id=session_id,
items=items,
seen=seen,
skip_user_messages=skip_user_messages,
retry_tracker=retry_tracker,
)
for item in items:
if item.source_id in seen:
continue
@@ -3390,13 +3586,23 @@ async def _forward_available_items(
retry_key = f"item:{item.source_id}"
if retry_tracker.retry_delay_s(retry_key) is not None:
return updated
try:
await _post_external_conversation_item(
client,
session_id=session_id,
item=item,
)
except httpx.HTTPError as exc:
# The batch above may already have attempted this item. Consume its
# outcome instead of re-posting; anything it did not cover (an older
# server, or an item past the batch's cap) falls back to a single post.
attempted = batch.pop(item.source_id, None) if batch is not None else None
exc: httpx.HTTPError | None = None
if attempted is not None:
exc = attempted.error
else:
try:
await _post_external_conversation_item(
client,
session_id=session_id,
item=item,
)
except httpx.HTTPError as post_exc:
exc = post_exc
if exc is not None:
decision = retry_tracker.record_failure(retry_key, exc)
if decision.exhausted:
_logger.error(
@@ -3462,7 +3668,27 @@ async def _forward_available_items(
item.source_id,
item.item_type,
_http_status_for_log(exc),
exc_info=True,
# The failure may have come from the batch above rather than
# a call in this frame, so name it explicitly.
exc_info=exc,
)
# Record the drop so it is recoverable. The replay classifier
# never auto-replays an ambiguous record (re-posting could
# duplicate a bubble), but without a record the item is simply
# gone — and a batch resolves this ambiguity for its whole run
# at once, so the gap would be a poll's worth of transcript.
append_dead_letter(
bridge_dir,
session_id=session_id,
event_type="external_conversation_item",
payload={
"item_type": item.item_type,
"item_data": item.data,
"response_id": item.response_id,
},
reason="ambiguous POST failure (may already be committed)",
delivered_ambiguous=True,
http_status=_http_status_for_log(exc),
)
retry_tracker.clear(retry_key)
seen.add(item.source_id)
@@ -3491,7 +3717,7 @@ async def _forward_available_items(
decision.permanent,
decision.delay_s,
_http_status_for_log(exc),
exc_info=True,
exc_info=exc,
)
return updated
retry_tracker.clear(retry_key)
@@ -4021,19 +4247,129 @@ async def _post_external_output_text_delta(
"""
resp = await client.post(
f"/v1/sessions/{session_id}/events",
json={
"type": "external_output_text_delta",
"data": {
"delta": delta.delta,
"message_id": delta.message_id,
"index": delta.index,
"final": delta.final,
},
},
json=_output_text_delta_event(delta),
)
resp.raise_for_status()
def _output_text_delta_event(delta: ClaudeMessageDelta) -> dict[str, object]:
"""
Build the ``external_output_text_delta`` event body for one chunk.
:param delta: Parsed streamed chunk.
:returns: The event body a single or batched post carries.
"""
return {
"type": "external_output_text_delta",
"data": {
"delta": delta.delta,
"message_id": delta.message_id,
"index": delta.index,
"final": delta.final,
},
}
def _log_dropped_delta(
*,
session_id: str,
bridge_dir: Path,
delta: ClaudeMessageDelta,
http_status: int | None = None,
error_type: str | None = None,
) -> None:
"""
Note one streamed chunk that never reached the server.
Records the status and the exception's class, never a rendered
exception or server message: an httpx error's text carries the request
it was made with, and this log is not the place to spill anything that
travelled in a header.
:param session_id: Omnigent session/conversation id.
:param bridge_dir: Native Claude bridge directory.
:param delta: The chunk that was dropped.
:param http_status: Status the server reported, when it responded.
:param error_type: Exception class name, for a transport failure.
:returns: None.
"""
_logger.debug(
"Dropping Claude streamed delta after HTTP failure; session=%s "
"bridge_dir=%s message_id=%s index=%s http_status=%s error_type=%s",
session_id,
bridge_dir,
delta.message_id,
delta.index,
http_status,
error_type,
)
async def _forward_delta_run(
client: httpx.AsyncClient,
*,
session_id: str,
bridge_dir: Path,
deltas: list[ClaudeMessageDelta],
) -> None:
"""
Publish a poll's worth of streamed chunks in as few round trips as possible.
One request for the whole run when the server serves the batch route,
else one per chunk as before. Deltas are best-effort live preview, so
a failure is logged and dropped (the authoritative final text still
arrives via ``external_conversation_item``) — never retried, so a
transient blip can't wedge the tail. That is also why the batch runs
with ``on_error="continue"``: one rejected chunk must not shadow the
rest of the run.
:param client: Omnigent HTTP client.
:param session_id: Omnigent session/conversation id.
:param bridge_dir: Native Claude bridge directory (logging only).
:param deltas: The new chunks to publish, in order.
:returns: None.
"""
try:
outcomes = await post_session_events_batch(
client,
session_id=session_id,
events=[_output_text_delta_event(delta) for delta in deltas],
on_error="continue",
)
except httpx.HTTPError as exc:
for delta in deltas:
_log_dropped_delta(
session_id=session_id,
bridge_dir=bridge_dir,
delta=delta,
http_status=_http_status_for_log(exc),
error_type=type(exc).__name__,
)
return
if outcomes is not None:
for delta, outcome in zip(deltas, outcomes, strict=False):
if not outcome.delivered:
_log_dropped_delta(
session_id=session_id,
bridge_dir=bridge_dir,
delta=delta,
http_status=outcome.status,
)
return
# Server predates the batch route: post one at a time.
for delta in deltas:
try:
await _post_external_output_text_delta(client, session_id=session_id, delta=delta)
except httpx.HTTPError as exc:
_log_dropped_delta(
session_id=session_id,
bridge_dir=bridge_dir,
delta=delta,
http_status=_http_status_for_log(exc),
error_type=type(exc).__name__,
)
async def _forward_available_deltas(
*,
client: httpx.AsyncClient,
@@ -4076,6 +4412,7 @@ async def _forward_available_deltas(
)
if result.byte_offset == state.byte_offset and not result.deltas:
return state
fresh: list[ClaudeMessageDelta] = []
for delta in result.deltas:
key = (delta.message_id, delta.index)
if key in seen_keys:
@@ -4086,18 +4423,11 @@ async def _forward_available_deltas(
# limit.
while len(seen_keys) > _MAX_SEEN_DELTA_KEYS:
del seen_keys[next(iter(seen_keys))]
try:
await _post_external_output_text_delta(client, session_id=session_id, delta=delta)
except httpx.HTTPError as exc:
_logger.debug(
"Dropping Claude streamed delta after HTTP failure; session=%s "
"bridge_dir=%s message_id=%s index=%s http_status=%s",
session_id,
bridge_dir,
delta.message_id,
delta.index,
_http_status_for_log(exc),
)
fresh.append(delta)
if fresh:
await _forward_delta_run(
client, session_id=session_id, bridge_dir=bridge_dir, deltas=fresh
)
updated = DeltaForwardState(byte_offset=result.byte_offset)
await _write_delta_forward_state_async(bridge_dir, updated)
return updated
+48
View File
@@ -1026,6 +1026,54 @@ mirror that status into the parent stream as `session.child_session.updated`
when the child is registered for fan-out. New user-facing event types should
default to the queue.
### Post Events (batch)
```
POST /v1/sessions/{session_id}/events/batch
Content-Type: application/json
{
"events": [
{"type": "external_output_text_delta", "data": {"delta": "Hel"}},
{"type": "external_output_text_delta", "data": {"delta": "lo"}}
],
"on_error": "stop"
}
```
Dispatches a run of events through the same handler as `POST
.../events`, in request order — one round trip instead of
`len(events)`. Internal (hidden from the OpenAPI reference): the native
harness forwarders use it to mirror a poll's transcript items and text
deltas at once. Without it a client far from the server can push only
about one event per round trip, so a long turn's live view falls tens of
seconds behind the harness.
`events` carries 1256 `SessionEventInput` bodies. `on_error` selects
what happens after a failure: `"stop"` (default) leaves the rest of the
batch unattempted, matching a sequential caller that retries from its
first failure; `"continue"` attempts every event, for best-effort
streams (live deltas) where dropping one chunk beats stalling the tail.
Always `200` with per-event outcomes; a rejected event does not fail the
request, so a caller can advance a durable cursor over the delivered
prefix instead of re-sending — and duplicating — events that landed:
```
{
"results": [
{"index": 0, "status": 202, "body": {"queued": false}},
{"index": 1, "status": 400, "error": "...", "code": "invalid_input"}
],
"stopped_at": 1
}
```
`results` is shorter than `events` when `stopped_at` is set — entries
past it were never attempted. Clients tell "this deployment has no batch
route" from "no such session" by the 404 body: a route miss has no
`error` envelope, an application 404 does.
### Resolve Elicitation (URL-based)
```
@@ -14,7 +14,9 @@ from fastapi import (
Request,
)
from fastapi.responses import StreamingResponse
from sqlalchemy.exc import StatementError
from omnigent.db.db_models import InvalidUuidError
from omnigent.entities import (
ErrorData,
NewConversationItem,
@@ -200,6 +202,7 @@ from omnigent.server.schemas import (
ElicitationRequestParams,
ErrorDetail,
McpServerStartup,
SessionEventBatchInput,
SessionEventInput,
)
from omnigent.session_lifecycle import (
@@ -1738,6 +1741,120 @@ def register_events_routes(
response["pending_id"] = dispatch.pending_id
return response
# ── POST /sessions/{session_id}/events/batch ─────────────
@router.post(
"/sessions/{session_id}/events/batch",
# Internal event ingestion — hidden from the public API reference.
include_in_schema=False,
# response_model=None: the body is a small per-event outcome
# envelope, not a domain model.
response_model=None,
)
async def post_events_batch(
request: Request,
session_id: str,
body: SessionEventBatchInput,
) -> dict[str, Any]:
"""
Dispatch a run of session events in one round trip.
Every event runs through the same handler as
``POST /sessions/{session_id}/events``, in request order, so a
batch is observationally equivalent to the sequence of single
posts it replaces. What it saves is round trips: a native
harness's forwarder posts one transcript item or text delta per
request, so a client an ocean away from the server can only push
a handful of events per second and its live view falls far behind
a long turn.
Per-event failures land in ``results`` instead of failing the
request, so a caller can advance its cursor over the successful
prefix and retry from the first failure rather than re-sending
(and duplicating) events that already landed. ``on_error="stop"``
leaves everything after the failure unattempted, matching a
sequential caller; ``"continue"`` attempts all of them, for
best-effort delta streams. Envelope-level problems (unparseable
body, oversized batch) are rejected by validation as usual.
:param session_id: Session/conversation identifier.
:param body: Events to dispatch plus the on-error policy.
:returns: ``{"results": [...], "stopped_at": int | None}`` — see
:class:`SessionEventBatchResponse`.
"""
results: list[dict[str, Any]] = []
stopped_at: int | None = None
for index, event in enumerate(body.events):
try:
outcome = await post_event(request, session_id, event)
except OmnigentError as exc:
# A wrong-replica landing is a fact about the *request*, not the
# event: on a multi-replica deployment the caller is expected to
# re-address and re-send. Reporting it per-event would hand the
# caller an opaque 400 it reads as a permanent rejection — the
# forwarder would dead-letter and drop a live item instead of
# retrying it at the right replica. Every event in a batch
# shares the session's replica binding, so this fires before
# anything is dispatched; propagate then, exactly as a single
# post does. If it somehow lands mid-batch (a tunnel that moved
# under us), keep the delivered prefix accounted for and report
# it with its code so the caller can still re-address.
if exc.code == ErrorCode.WRONG_REPLICA and not results:
raise
results.append(
{
"index": index,
"status": exc.http_status,
"error": str(exc),
"code": exc.code,
}
)
except StatementError as exc:
# A malformed id can address no row. The app's handler maps
# this to 404 for a single post; mirror it so a batched event
# reports the same status instead of a bogus 500.
if not isinstance(exc.orig, InvalidUuidError):
raise
results.append(
{
"index": index,
"status": 404,
"error": "Not found.",
"code": ErrorCode.NOT_FOUND,
}
)
except Exception:
# A single post would have failed only its own request;
# propagating here would fail the whole batch and make the
# caller re-send events that already landed. Report it
# per-event and keep the log loud so the bug stays visible.
#
# The body carries the same generic message the app's 500
# handler returns, never the exception's text: an internal
# failure's detail (paths, SQL, stack context) must not reach
# a client just because it happened inside a batch.
_logger.exception(
"Batched session event failed; session=%s index=%s type=%s",
session_id,
index,
event.type,
)
results.append(
{
"index": index,
"status": 500,
"error": "An internal error occurred.",
"code": ErrorCode.INTERNAL_ERROR,
}
)
else:
results.append({"index": index, "status": 202, "body": dict(outcome)})
continue
if body.on_error == "stop":
stopped_at = index
break
return {"results": results, "stopped_at": stopped_at}
# ── GET /sessions/{session_id}/stream ────────────────────────
# Live-tail only. Clients reconnect via GET /v1/sessions/{id}
+68
View File
@@ -1197,6 +1197,74 @@ class SessionEventInput(BaseModel):
created_by: str | None = None
# Cap on how many events one batch may carry. Bounds the worst-case work a
# single request can queue behind the event loop (each event runs the full
# ``POST /events`` dispatch) while staying far above what a forwarder poll
# produces in practice.
MAX_SESSION_EVENTS_PER_BATCH = 256
class SessionEventBatchInput(BaseModel):
"""
Body of ``POST /v1/sessions/{id}/events/batch``.
Carries a run of events that would otherwise be posted one at a
time. Events are dispatched **in order** through the same code path
as ``POST /v1/sessions/{id}/events``, so a batch is observationally
equivalent to the sequence of single posts it replaces — it just
costs one round trip instead of ``len(events)``. This matters most
for clients far from the server, where a native harness's transcript
and delta forwarders are otherwise rate-limited to one event per RTT.
:param events: The events to dispatch, in order. At least one, at
most :data:`MAX_SESSION_EVENTS_PER_BATCH`.
:param on_error: What to do when an event fails. ``"stop"`` (the
default) leaves the remaining events unattempted, mirroring a
sequential caller that stops at its first failure and retries
from there. ``"continue"`` attempts every event regardless, for
best-effort streams (live text deltas) where a dropped chunk is
preferable to stalling the tail.
"""
events: list[SessionEventInput] = Field(min_length=1, max_length=MAX_SESSION_EVENTS_PER_BATCH)
on_error: Literal["stop", "continue"] = "stop"
class SessionEventBatchResult(BaseModel):
"""
Outcome of one event inside a ``POST /events/batch`` request.
:param index: Position of the event in the request's ``events``.
:param status: HTTP status the equivalent single post would have
returned — ``202`` on success, otherwise the error's status.
:param body: The single-post response body, on success.
:param error: Human-readable failure reason, on failure.
:param code: Machine-readable Omnigent error code, on failure.
"""
index: int
status: int
body: dict[str, Any] | None = None
error: str | None = None
code: str | None = None
class SessionEventBatchResponse(BaseModel):
"""
Response of ``POST /v1/sessions/{id}/events/batch``.
:param results: One entry per *attempted* event, in request order.
Shorter than the request's ``events`` when ``on_error="stop"``
tripped: entries after ``stopped_at`` were never attempted and
the caller should re-send them.
:param stopped_at: Index of the failure that ended the batch early,
or ``None`` when every event was attempted.
"""
results: list[SessionEventBatchResult]
stopped_at: int | None = None
class SessionGitOptions(BaseModel):
"""
Git worktree options for ``POST /v1/sessions``.
@@ -23,6 +23,7 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
from omnigent.errors import ErrorCode, OmnigentError
from omnigent.llms.context_window import ModelPricing
from omnigent.runtime.tool_output import MAX_TOOL_OUTPUT_BYTES
from omnigent.server.background_session_titles import BackgroundTitleRequest
@@ -30,6 +31,7 @@ from omnigent.server.routes._sessions.helpers import (
_NativeTerminalEnsureOutcome,
_RunnerForwardResult,
)
from omnigent.server.schemas import MAX_SESSION_EVENTS_PER_BATCH
from omnigent.spec.types import SkillSpec
from omnigent.stores.conversation_store.sqlalchemy_store import (
SqlAlchemyConversationStore,
@@ -3967,6 +3969,295 @@ async def test_post_external_output_text_delta_rejects_malformed_delta(
assert published == []
def _text_delta_event(text: str) -> dict[str, Any]:
"""
Build one ``external_output_text_delta`` event body.
:param text: The chunk text, e.g. ``"hel"``.
:returns: An event body for ``/events`` or ``/events/batch``.
"""
return {"type": "external_output_text_delta", "data": {"delta": text}}
async def test_post_events_batch_dispatches_every_event_in_order(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
A batch is the sequence of single posts it replaces, in one round trip.
Native forwarders post one transcript item or text chunk per request,
which caps a far-away client at a few events per second and leaves its
live view far behind the harness. Batching only changes the number of
round trips, so ordering and per-event acknowledgement must match the
single-post path exactly. Fails if events are reordered, dropped, or
reported without their individual status.
"""
published: list[tuple[str, dict[str, Any]]] = []
monkeypatch.setattr(
"omnigent.server.routes.sessions.session_stream.publish",
lambda sid, ev: published.append((sid, ev)),
)
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={"events": [_text_delta_event(text) for text in ("a", "b", "c")]},
)
assert resp.status_code == 200, resp.text
body = resp.json()
assert body["stopped_at"] is None
assert [entry["status"] for entry in body["results"]] == [202, 202, 202]
assert [entry["index"] for entry in body["results"]] == [0, 1, 2]
assert [entry["body"] for entry in body["results"]] == [{"queued": False}] * 3
assert [event["delta"] for _sid, event in published] == ["a", "b", "c"]
async def test_post_events_batch_persists_items_like_single_posts(
client: httpx.AsyncClient,
) -> None:
"""
Batched ``external_conversation_item`` events land in history, in order.
This is the equivalence that lets a forwarder mirror a whole poll's
transcript in one request. Fails if the batch path skips persistence
or scrambles item order (resume would replay a garbled transcript).
"""
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={
"events": [
{
"type": "external_conversation_item",
"data": {
"item_type": "message",
"item_data": {
"role": "user",
"content": [{"type": "input_text", "text": text}],
"is_meta": True,
},
"response_id": "turn_1",
"source_id": f"src-{index}",
},
}
for index, text in enumerate(("first", "second", "third"))
]
},
)
assert resp.status_code == 200, resp.text
assert [entry["status"] for entry in resp.json()["results"]] == [202, 202, 202], resp.text
items = (await client.get(f"/v1/sessions/{session['id']}/items")).json()["data"]
texts = [item["content"][0]["text"] for item in items if item["type"] == "message"]
assert texts == ["first", "second", "third"]
async def test_post_events_batch_stops_at_the_first_failure(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
By default a failed event leaves the rest of the batch unattempted.
A sequential caller stops at its first failure and retries from
there; the batch must behave the same so a forwarder can advance its
cursor over the delivered prefix without re-sending and therefore
duplicating events that already landed. Fails if later events are
dispatched anyway (the caller then can't tell what to retry) or if
one bad event fails the whole request.
"""
published: list[tuple[str, dict[str, Any]]] = []
monkeypatch.setattr(
"omnigent.server.routes.sessions.session_stream.publish",
lambda sid, ev: published.append((sid, ev)),
)
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={
"events": [
_text_delta_event("ok"),
{"type": "external_output_text_delta", "data": {"delta": {"bad": 1}}},
_text_delta_event("never"),
]
},
)
assert resp.status_code == 200, resp.text
body = resp.json()
assert body["stopped_at"] == 1
assert [entry["status"] for entry in body["results"]] == [202, 400]
assert "requires string data.delta" in body["results"][1]["error"]
# The third event was never attempted.
assert [event["delta"] for _sid, event in published] == ["ok"]
async def test_post_events_batch_continue_attempts_every_event(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
``on_error="continue"`` keeps going past a rejected event.
Live text deltas are a preview the final message supersedes, so
dropping one chunk beats stalling the tail behind it. Fails if a
mid-run rejection swallows the chunks after it.
"""
published: list[tuple[str, dict[str, Any]]] = []
monkeypatch.setattr(
"omnigent.server.routes.sessions.session_stream.publish",
lambda sid, ev: published.append((sid, ev)),
)
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={
"events": [
_text_delta_event("a"),
{"type": "external_output_text_delta", "data": {"delta": None}},
_text_delta_event("c"),
],
"on_error": "continue",
},
)
assert resp.status_code == 200, resp.text
body = resp.json()
assert body["stopped_at"] is None
assert [entry["status"] for entry in body["results"]] == [202, 400, 202]
assert [event["delta"] for _sid, event in published] == ["a", "c"]
async def test_post_events_batch_rejects_an_oversized_or_empty_batch(
client: httpx.AsyncClient,
) -> None:
"""
The envelope caps how much work one request can queue.
Each event runs the full single-post dispatch, so an unbounded batch
is an unbounded amount of work behind one request. Fails if the cap
or the non-empty requirement stops being enforced.
"""
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
too_many = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={"events": [_text_delta_event("x")] * (MAX_SESSION_EVENTS_PER_BATCH + 1)},
)
assert too_many.status_code == 422, too_many.text
empty = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={"events": []},
)
assert empty.status_code == 422, empty.text
async def test_post_events_batch_reports_a_missing_session_per_event(
client: httpx.AsyncClient,
) -> None:
"""
An application 404 lands in ``results``, not as a route miss.
Clients probe this endpoint to learn whether the deployment has it,
and tell "no such route" from "no such session" by the error
envelope. Fails if a missing session answers with a bare route-miss
shape, which would make a client permanently downgrade to per-event
posts against a server that supports batching.
"""
resp = await client.post(
f"/v1/sessions/{'0' * 32}/events/batch",
json={"events": [_text_delta_event("a")]},
)
assert resp.status_code == 200, resp.text
body = resp.json()
assert body["stopped_at"] == 0
assert body["results"][0]["status"] == 404
assert body["results"][0]["code"]
async def test_post_events_batch_propagates_a_wrong_replica_landing(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
A wrong-replica miss fails the request, not one event inside it.
On a multi-replica deployment a session-scoped request can land on a
replica that doesn't hold the session's tunnel; the contract is a 400
carrying ``wrong_replica`` so the caller re-addresses and re-sends.
Reporting it per-event instead hands the caller an opaque 400 it reads
as a permanent rejection the forwarder would dead-letter and drop a
live transcript item rather than retry it at the right replica. Fails
if the code is swallowed into ``results``.
"""
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
def _wrong_replica(*_args: Any, **_kwargs: Any) -> None:
"""Raise the routing miss the real handler raises."""
raise OmnigentError(
"session runner is on another replica; retry",
code=ErrorCode.WRONG_REPLICA,
)
monkeypatch.setattr(
"omnigent.server.routes.sessions.session_stream.publish",
_wrong_replica,
)
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={"events": [_text_delta_event("a"), _text_delta_event("b")]},
)
assert resp.status_code == 400, resp.text
assert resp.json()["error"]["code"] == ErrorCode.WRONG_REPLICA
async def test_post_events_batch_never_leaks_an_internal_error_detail(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
An unexpected server-side failure reports a generic message, not its text.
The batch route has to catch per-event failures rather than let them fail
the whole request otherwise a caller re-sends events that already
landed but that catch must not become a channel for internal detail
(paths, SQL, stack context) that the app's own 500 handler withholds.
Fails if the exception's text reaches the response body.
"""
agent = await create_test_agent(client)
session = await _create_session(client, agent["id"])
_LEAK_MARKER = "internal-detail-that-must-not-escape"
def _boom(*_args: Any, **_kwargs: Any) -> None:
"""Raise an error whose text must not reach the client."""
raise RuntimeError(f"{_LEAK_MARKER} at /srv/omnigent/state row 42")
monkeypatch.setattr(
"omnigent.server.routes.sessions.session_stream.publish",
_boom,
)
resp = await client.post(
f"/v1/sessions/{session['id']}/events/batch",
json={"events": [_text_delta_event("a")]},
)
assert resp.status_code == 200, resp.text
entry = resp.json()["results"][0]
assert entry["status"] == 500
assert entry["error"] == "An internal error occurred."
assert _LEAK_MARKER not in resp.text
assert "RuntimeError" not in resp.text
async def test_post_external_tool_output_delta_publishes_transient_delta(
client: httpx.AsyncClient,
monkeypatch: pytest.MonkeyPatch,
+625 -59
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import contextlib
import itertools
import json
import logging
import os
@@ -13,6 +14,7 @@ from collections.abc import Callable, Generator, Iterator
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
@@ -24,6 +26,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
import omnigent.claude_native_forwarder as forwarder
from omnigent._native_post_delivery import reset_events_batch_support
from omnigent.claude_native_bridge import (
BRIDGE_ID_LABEL_KEY,
ClaudeMessageDelta,
@@ -84,6 +87,11 @@ def _handler_factory(
:returns: A concrete :class:`BaseHTTPRequestHandler` subclass.
"""
# Monotonic per-HTTP-request counter stamped on every recorded entry. A
# batch envelope expands into several entries that share one index, which is
# how a test can still distinguish "two events" from "two requests".
request_counter = itertools.count()
class _Handler(BaseHTTPRequestHandler):
"""Request handler for the test Omnigent endpoint."""
@@ -101,16 +109,50 @@ def _handler_factory(
"""
Record a JSON POST body and return HTTP 202.
A batch envelope is recorded as its individual events, one queue
entry each, so tests that read the mirrored event sequence do not
have to know how many requests it arrived in.
:returns: None.
"""
length = int(self.headers.get("Content-Length", "0"))
raw = self.rfile.read(length)
body = json.loads(raw.decode("utf-8"))
request_index = next(request_counter)
if self.path.endswith("/events/batch"):
events = body["events"]
single_path = self.path[: -len("/batch")]
for event in events:
requests.put(
{
"method": "POST",
"path": single_path,
"body": event,
"authorization": self.headers.get("Authorization"),
"request_index": request_index,
}
)
payload = json.dumps(
{
"results": [
{"index": index, "status": 202} for index in range(len(events))
],
"stopped_at": None,
}
).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
return
requests.put(
{
"method": "POST",
"path": self.path,
"body": json.loads(raw.decode("utf-8")),
"body": body,
"authorization": self.headers.get("Authorization"),
"request_index": request_index,
}
)
self.send_response(202)
@@ -131,6 +173,7 @@ def _handler_factory(
"path": self.path,
"body": json.loads(raw.decode("utf-8")),
"authorization": self.headers.get("Authorization"),
"request_index": next(request_counter),
}
)
self.send_response(200)
@@ -1945,11 +1988,17 @@ async def test_forwarder_uses_auth_to_refresh_token_per_request(tmp_path: Path)
)
)
try:
# Two external_conversation_item POSTs (one per assistant item).
# The PATCH that mirrors the Claude session id is filtered out
# by ``_get_recorded_request``'s default ``method="POST"``.
first = await _get_recorded_request(server)
second = await _get_recorded_request(server)
# Two *requests*, not two events: the forwarder mirrors Claude's session
# id with a one-shot PATCH and ships the run of transcript items in one
# batched POST. Both go through the same client, so they are what proves
# the header is minted per request rather than snapshotted once. (Reading
# two events out of one batch could not: one request carries one header.)
collected: dict[str, dict[str, Any]] = {}
while len(collected) < 2:
recorded = await asyncio.to_thread(server.requests.get, True, 5.0)
collected.setdefault(str(recorded["method"]), recorded)
first = collected["PATCH"]
second = collected["POST"]
finally:
task.cancel()
with pytest.raises(asyncio.CancelledError):
@@ -1965,22 +2014,22 @@ async def test_forwarder_uses_auth_to_refresh_token_per_request(tmp_path: Path)
assert first["authorization"] is not None and first["authorization"].startswith(
"Bearer token-"
), (
f"First POST must carry a bearer minted by the counting auth, "
f"First request must carry a bearer minted by the counting auth, "
f"got {first['authorization']!r}. ``None`` means auth was not "
f"threaded into httpx.AsyncClient."
)
assert second["authorization"] is not None and second["authorization"].startswith(
"Bearer token-"
), (
f"Second POST must carry a bearer minted by the counting auth, "
f"Second request must carry a bearer minted by the counting auth, "
f"got {second['authorization']!r}."
)
# The load-bearing assertion: the two POSTs carry DIFFERENT
# The load-bearing assertion: the two requests carry DIFFERENT
# bearers. If they were equal, httpx would be reusing a
# construction-time header snapshot instead of consulting the
# auth flow per request — that is exactly the production bug.
assert first["authorization"] != second["authorization"], (
f"Two consecutive POSTs share the same Authorization "
f"Two consecutive requests share the same Authorization "
f"({first['authorization']!r}). The AsyncClient is reusing a "
f"snapshot of the original header instead of consulting auth "
f"on each request — this is the production token-refresh bug."
@@ -2941,6 +2990,320 @@ async def test_forwarder_survives_unhandled_loop_exceptions(
}
def _write_assistant_transcript(path: Path, count: int) -> None:
"""
Write *count* assistant items to a transcript file.
:param path: Transcript path to create.
:param count: How many assistant items to write.
:returns: None.
"""
path.write_text(
"".join(
json.dumps(
{
"type": "assistant",
"uuid": f"item-{index}",
"message": {"role": "assistant", "content": f"line {index}"},
}
)
+ "\n"
for index in range(count)
),
encoding="utf-8",
)
def _item_forward_state(transcript_path: Path) -> forwarder.TranscriptForwardState:
"""
Build a fresh transcript cursor for *transcript_path*.
:param transcript_path: Transcript the cursor points at.
:returns: A state positioned at the start of the file.
"""
return forwarder.TranscriptForwardState(
transcript_path=transcript_path,
line_cursor=0,
byte_offset=0,
cursor_fingerprint=forwarder._jsonl_cursor_fingerprint(transcript_path, 0),
)
async def _forward_items_once(
client: httpx.AsyncClient,
*,
bridge_dir: Path,
state: forwarder.TranscriptForwardState,
retry_tracker: forwarder._PostRetryTracker | None = None,
) -> forwarder.TranscriptForwardState:
"""
Run one item-forward pass with the usual test arguments.
:param client: Omnigent HTTP client.
:param bridge_dir: Bridge directory.
:param state: Transcript cursor to start from.
:param retry_tracker: Retry tracker, defaulting to a no-delay one.
:returns: The updated cursor.
"""
return await forwarder._forward_available_items(
client=client,
session_id="conv_abc",
bridge_dir=bridge_dir,
agent_name="claude-native-ui",
state=state,
retry_tracker=retry_tracker
or forwarder._PostRetryTracker(base_delay_s=0.0, max_delay_s=0.0),
dedupe=forwarder._ForwardDedupeState(),
)
@pytest.mark.asyncio
async def test_a_run_of_items_is_mirrored_in_one_request(tmp_path: Path) -> None:
"""
Items available in one poll go out together, capped, in order.
A burst used to cost a round trip per item *inside a single poll
iteration*, which from another region both delayed the cards and starved
the streamed-text tail queued behind it. Fails if the run fans back out
into per-item requests, if the cap stops bounding one request's blast
radius, or if the remainder past the cap is dropped instead of following.
"""
bridge_dir = tmp_path / "bridge"
transcript_path = tmp_path / "session.jsonl"
total = forwarder._ITEM_BATCH_MAX + 2
_write_assistant_transcript(transcript_path, total)
paths: list[str] = []
events: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Accept everything, recording the request path and each event."""
paths.append(request.url.path)
batched = _record_forwarder_post(events, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await _forward_items_once(
client, bridge_dir=bridge_dir, state=_item_forward_state(transcript_path)
)
# One batch for the capped run, then a single post for each leftover.
assert paths == [
"/v1/sessions/conv_abc/events/batch",
*["/v1/sessions/conv_abc/events"] * 2,
]
# Every item still mirrored, exactly once, in transcript order.
assert [event["data"]["item_data"]["content"][0]["text"] for event in events] == [
f"line {index}" for index in range(total)
]
@pytest.mark.asyncio
async def test_an_item_rejected_inside_a_batch_retries_like_a_single_post(
tmp_path: Path,
) -> None:
"""
A batched rejection is classified exactly as a per-item one.
The retry, dead-letter, and cursor rules for a rejected item are
load-bearing and must not acquire a second, divergent implementation for
the batched path. Fails if a batched 4xx escapes the permanent-failure
budget (the item would be retried forever, or dropped without a
dead-letter record).
"""
bridge_dir = tmp_path / "bridge"
transcript_path = tmp_path / "session.jsonl"
_write_assistant_transcript(transcript_path, 3)
retry_tracker = forwarder._PostRetryTracker(
max_permanent_attempts=2, base_delay_s=0.0, max_delay_s=0.0
)
events: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Reject the first item of every batch; accept everything else."""
payload = json.loads(request.content.decode("utf-8"))
if not request.url.path.endswith("/events/batch"):
events.append(payload)
return httpx.Response(202, json={})
events.extend(payload["events"])
results: list[dict[str, Any]] = [
{"index": 0, "status": 422, "error": "bad item", "code": "invalid_input"}
]
return httpx.Response(200, json={"results": results, "stopped_at": 0})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
state = _item_forward_state(transcript_path)
for _ in range(2):
state = await _forward_items_once(
client, bridge_dir=bridge_dir, state=state, retry_tracker=retry_tracker
)
# Two attempts (the permanent budget), then the forwarder-failed status —
# the same shape the per-item path produces for a poison item.
assert [event["type"] for event in events].count("external_conversation_item") >= 2
assert "external_session_status" in [event["type"] for event in events]
dead_letters = (bridge_dir / "dead_letter.jsonl").read_text(encoding="utf-8").splitlines()
assert [json.loads(line)["event_type"] for line in dead_letters] == [
"external_conversation_item"
]
@pytest.mark.asyncio
async def test_a_lost_batch_response_is_dead_lettered_not_reposted(tmp_path: Path) -> None:
"""
When a batch's response is lost, its items are never posted again.
External items are not deduped server-side, so any prefix of that batch
may already be committed and a re-post would render duplicate bubbles.
The per-item path already resolves this ambiguity by dropping the item
and dead-lettering it for recovery; a batch must resolve it the same way
for its whole run. Fails if the items are re-sent (duplicates) or
dropped without a dead-letter record (unrecoverable).
"""
bridge_dir = tmp_path / "bridge"
transcript_path = tmp_path / "session.jsonl"
_write_assistant_transcript(transcript_path, 3)
attempts: list[str] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Lose the batch response; accept the failure status that follows."""
attempts.append(request.url.path)
if request.url.path.endswith("/events/batch"):
# A response the client never sees: request sent, reply lost.
raise httpx.ReadTimeout("response lost", request=request)
return httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
state = _item_forward_state(transcript_path)
state = await _forward_items_once(client, bridge_dir=bridge_dir, state=state)
# A second pass must not re-attempt the ambiguous items.
await _forward_items_once(client, bridge_dir=bridge_dir, state=state)
assert attempts.count("/v1/sessions/conv_abc/events/batch") == 1
assert "/v1/sessions/conv_abc/events" not in attempts[1:2]
dead_letters = (bridge_dir / "dead_letter.jsonl").read_text(encoding="utf-8").splitlines()
records = [json.loads(line) for line in dead_letters]
assert len(records) == 3
assert all(record["delivered_ambiguous"] is True for record in records)
@pytest.mark.asyncio
async def test_a_batch_that_never_left_falls_back_to_single_posts(tmp_path: Path) -> None:
"""
A connect failure loses nothing: the run is retried item by item.
A request that provably never reached the server committed nothing, so
unlike the lost-response case there is no ambiguity to resolve and the
items must still be delivered. Fails if a refused connection
dead-letters a poll's transcript that the server never saw.
"""
bridge_dir = tmp_path / "bridge"
transcript_path = tmp_path / "session.jsonl"
_write_assistant_transcript(transcript_path, 3)
events: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Refuse the batch outright; accept single posts."""
if request.url.path.endswith("/events/batch"):
raise httpx.ConnectError("refused", request=request)
events.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
await _forward_items_once(
client, bridge_dir=bridge_dir, state=_item_forward_state(transcript_path)
)
assert [event["data"]["item_data"]["content"][0]["text"] for event in events] == [
"line 0",
"line 1",
"line 2",
]
assert not (bridge_dir / "dead_letter.jsonl").exists()
@pytest.mark.asyncio
async def test_items_post_one_by_one_when_the_server_lacks_the_batch_route(
tmp_path: Path,
) -> None:
"""
An older deployment still receives every item, one request each.
A user's CLI can be newer than the server it talks to; the mirror must
keep working there, just without the round-trip saving. Fails if a
missing batch route drops the run.
"""
bridge_dir = tmp_path / "bridge"
transcript_path = tmp_path / "session.jsonl"
_write_assistant_transcript(transcript_path, 3)
events: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Answer the batch route with Starlette's route-miss shape."""
if request.url.path.endswith("/events/batch"):
return httpx.Response(404, json={"detail": "Not Found"})
events.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
reset_events_batch_support(client)
await _forward_items_once(
client, bridge_dir=bridge_dir, state=_item_forward_state(transcript_path)
)
assert [event["data"]["item_data"]["content"][0]["text"] for event in events] == [
"line 0",
"line 1",
"line 2",
]
def test_a_compaction_boundary_ends_a_batchable_run() -> None:
"""
The run stops before a compaction boundary, never spanning it.
The boundary has its own durable persist path and must land in order
batching past it would reorder the transcript, and a resume would reload
history the compaction was supposed to drop. Fails if a summary record
is swept into the batch or merely skipped over.
"""
def _item(source_id: str, *, compact: bool = False) -> Any:
"""Build a minimal transcript item stand-in."""
return SimpleNamespace(
source_id=source_id,
item_type="message",
data={"role": "assistant"},
is_compact_summary=compact,
)
items = [_item("a"), _item("b"), _item("summary", compact=True), _item("c")]
run = forwarder._batchable_item_run(
items,
seen=set(),
skip_user_messages=False,
retry_tracker=forwarder._PostRetryTracker(base_delay_s=0.0, max_delay_s=0.0),
)
assert [item.source_id for item in run] == ["a", "b"]
# Already-forwarded items and skipped user messages neither join the run
# nor end it — they post nothing at all.
already = [_item("a"), _item("b")]
assert [
item.source_id
for item in forwarder._batchable_item_run(
already,
seen={"a"},
skip_user_messages=False,
retry_tracker=forwarder._PostRetryTracker(base_delay_s=0.0, max_delay_s=0.0),
)
] == ["b"]
@pytest.mark.asyncio
async def test_forwarder_drops_poison_item_after_bounded_permanent_retries(
tmp_path: Path,
@@ -3523,8 +3886,14 @@ async def test_forward_model_from_status_posts_the_status_model_verbatim(
requests: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
requests.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
"""
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
dedupe = forwarder._ForwardDedupeState()
transport = httpx.MockTransport(_handle_request)
@@ -3556,8 +3925,14 @@ async def test_model_reports_keep_generation_and_context_marker(tmp_path: Path)
requests: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
requests.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
"""
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
dedupe = forwarder._ForwardDedupeState()
transport = httpx.MockTransport(_handle_request)
@@ -3626,13 +4001,13 @@ async def test_forwarder_reports_the_launch_model_then_a_switch(tmp_path: Path)
def _handle_request(request: httpx.Request) -> httpx.Response:
"""
Accept every forwarder POST and record its payload.
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: 202 for every event.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
requests.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -3705,13 +4080,13 @@ async def test_forwarder_mirrors_tui_rename_on_first_observation(tmp_path: Path)
def _handle_request(request: httpx.Request) -> httpx.Response:
"""
Accept every forwarder POST and record its payload.
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: 202 for every event.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
requests.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -3958,9 +4333,14 @@ async def test_forwarder_mirrors_in_pane_permission_mode_switch(
posts: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Accept every POST and record its payload."""
posts.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
"""
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
batched = _record_forwarder_post(posts, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -4023,9 +4403,14 @@ async def test_forwarder_posts_manual_launch_mode_so_picker_renders(
posts: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
"""Accept the POST and record its payload."""
posts.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
"""
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
batched = _record_forwarder_post(posts, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -5880,6 +6265,38 @@ class _CapturedDeltaPost:
body: dict[str, Any]
def _record_forwarder_post(
sink: list[dict[str, Any]], request: httpx.Request
) -> httpx.Response | None:
"""
Record a forwarder POST as individual events, unwrapping a batch.
The forwarder ships runs of events through ``/events/batch``, so a stub
that wants to assert on the *event* sequence has to look inside the
envelope. Recording both framings the same way keeps these tests about
what the forwarder mirrors rather than how many requests it took.
:param sink: List appended to with each event body, in order.
:param request: The outbound request to record.
:returns: The response for a batch envelope (every event accepted), or
``None`` for a single post so the caller can pick its own status.
"""
payload = json.loads(request.content.decode("utf-8"))
assert isinstance(payload, dict)
if not request.url.path.endswith("/events/batch"):
sink.append(payload)
return None
events = payload["events"]
sink.extend(events)
return httpx.Response(
200,
json={
"results": [{"index": index, "status": 202} for index in range(len(events))],
"stopped_at": None,
},
)
def _write_deltas_file(bridge_dir: Path, records: list[dict[str, Any]]) -> None:
"""
Append delta records to ``message_deltas.jsonl`` as the hook would.
@@ -5898,34 +6315,60 @@ def _write_deltas_file(bridge_dir: Path, records: list[dict[str, Any]]) -> None:
def _delta_capture_client(
captured: list[_CapturedDeltaPost],
status_code: int = 202,
*,
batch_supported: bool = True,
batch_event_statuses: list[int] | None = None,
) -> httpx.AsyncClient:
"""
Build an AsyncClient whose ``/events`` POSTs are captured.
Build an AsyncClient whose ``/events`` and ``/events/batch`` POSTs are captured.
:param captured: List appended to with each observed POST body.
:param status_code: HTTP status the stub returns, e.g. ``202`` for
success or ``500`` to exercise the best-effort drop path.
:param batch_supported: ``False`` answers the batch route with the
route-miss 404 an older deployment returns, so the caller's
per-event fallback runs.
:param batch_event_statuses: Per-event statuses for the batch reply,
to exercise a batch that partially fails. Defaults to all-202.
:returns: An ``httpx.AsyncClient`` bound to the capturing transport.
"""
def handler(request: httpx.Request) -> httpx.Response:
captured.append(
_CapturedDeltaPost(url_path=request.url.path, body=json.loads(request.content))
body = json.loads(request.content)
captured.append(_CapturedDeltaPost(url_path=request.url.path, body=body))
if not request.url.path.endswith("/events/batch"):
return httpx.Response(status_code, json={"queued": False})
if not batch_supported:
# Starlette's route miss — no ``error`` envelope, which is how
# the client tells "no such route" from "no such session".
return httpx.Response(404, json={"detail": "Not Found"})
if status_code >= 400:
return httpx.Response(status_code, json={"detail": "boom"})
statuses = batch_event_statuses or [202] * len(body["events"])
return httpx.Response(
200,
json={
"results": [
{"index": index, "status": status} for index, status in enumerate(statuses)
],
"stopped_at": None,
},
)
return httpx.Response(status_code, json={"queued": False})
return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url="http://ap")
async def test_forward_available_deltas_posts_each_and_advances_offset(tmp_path: Path) -> None:
"""
Each appended chunk is POSTed as an ``external_output_text_delta``.
A poll's chunks go out as one batch of ``external_output_text_delta`` events.
Proves the forwarder turns deltas-file lines into the exact event
shape the Omnigent route expects (delta + message_id + index + final) and
advances+persists the byte offset so the next poll resumes after
them. Fails if a field is dropped (UI can't scope/order the buffer)
or the offset doesn't persist (chunks re-POST on restart).
shape the Omnigent route expects (delta + message_id + index + final),
ships the whole run in a single round trip, and advances+persists the
byte offset so the next poll resumes after them. Fails if a field is
dropped (UI can't scope/order the buffer), if the run costs a request
per chunk again (a far-from-server client then streams at one chunk
per RTT), or if the offset doesn't persist (chunks re-POST on restart).
"""
bridge_dir = prepare_bridge_dir("conv_x", bridge_id="b1", workspace=tmp_path)
_write_deltas_file(
@@ -5946,12 +6389,10 @@ async def test_forward_available_deltas_posts_each_and_advances_offset(tmp_path:
seen_keys=seen,
)
assert [c.url_path for c in captured] == [
"/v1/sessions/conv_x/events",
"/v1/sessions/conv_x/events",
]
# One request, not one per chunk.
assert [c.url_path for c in captured] == ["/v1/sessions/conv_x/events/batch"]
# Full event shape proves every field survived hook → file → POST.
assert [c.body for c in captured] == [
assert captured[0].body["events"] == [
{
"type": "external_output_text_delta",
"data": {"delta": "Hello ", "message_id": "m1", "index": 0, "final": False},
@@ -5961,6 +6402,8 @@ async def test_forward_available_deltas_posts_each_and_advances_offset(tmp_path:
"data": {"delta": "world", "message_id": "m1", "index": 1, "final": True},
},
]
# Best-effort preview: a rejected chunk must not shadow the rest of the run.
assert captured[0].body["on_error"] == "continue"
# Offset advanced to EOF and was persisted, so a reload resumes past
# the two chunks instead of re-POSTing them.
assert new_state.byte_offset == os.path.getsize(bridge_dir / "message_deltas.jsonl")
@@ -5996,11 +6439,12 @@ async def test_forward_available_deltas_dedupes_by_message_id_and_index(tmp_path
seen_keys=seen,
)
# The duplicate (m1, 0) is collapsed: only the first (m1,0) and the
# distinct (m1,1) are POSTed — 2 requests, not 3.
assert [(c.body["data"]["message_id"], c.body["data"]["index"]) for c in captured] == [
("m1", 0),
("m1", 1),
]
# distinct (m1,1) are sent — 2 events in the batch, not 3.
assert len(captured) == 1
assert [
(event["data"]["message_id"], event["data"]["index"])
for event in captured[0].body["events"]
] == [("m1", 0), ("m1", 1)]
async def test_forward_available_deltas_drops_on_http_error(tmp_path: Path) -> None:
@@ -6030,6 +6474,126 @@ async def test_forward_available_deltas_drops_on_http_error(tmp_path: Path) -> N
# The POST was attempted (and 500'd) but no exception escaped, and
# the offset moved past the chunk so it won't be retried endlessly.
assert len(captured) == 1
assert captured[0].url_path == "/v1/sessions/conv_x/events/batch"
assert new_state.byte_offset == os.path.getsize(bridge_dir / "message_deltas.jsonl")
async def test_forward_available_deltas_batches_a_long_run_in_one_round_trip(
tmp_path: Path,
) -> None:
"""
A whole turn's worth of chunks costs one request, not one per chunk.
This is the property that makes streaming usable from far away: an
assistant message is hundreds of chunks, and a per-chunk POST caps the
live preview at one chunk per round trip, so the web view falls tens of
seconds behind the pane. Fails the moment the run fans back out into
per-chunk requests.
"""
bridge_dir = prepare_bridge_dir("conv_x", bridge_id="b1", workspace=tmp_path)
_write_deltas_file(
bridge_dir,
[
{"message_id": "m1", "index": index, "final": index == 79, "delta": f"c{index}"}
for index in range(80)
],
)
captured: list[_CapturedDeltaPost] = []
seen: dict[tuple[str, int], None] = {}
async with _delta_capture_client(captured) as client:
await forwarder._forward_available_deltas(
client=client,
session_id="conv_x",
bridge_dir=bridge_dir,
state=forwarder.DeltaForwardState(),
seen_keys=seen,
)
assert len(captured) == 1
assert len(captured[0].body["events"]) == 80
async def test_forward_available_deltas_falls_back_when_server_lacks_batch_route(
tmp_path: Path,
) -> None:
"""
An older server's route-miss 404 falls back to one POST per chunk.
A user's CLI can be newer than the deployment it talks to, and the
live preview must keep working there just without the round-trip
saving. Fails if a missing batch route silently drops the run (no
preview at all) or raises out of the poll loop.
"""
bridge_dir = prepare_bridge_dir("conv_x", bridge_id="b1", workspace=tmp_path)
_write_deltas_file(
bridge_dir,
[
{"message_id": "m1", "index": 0, "final": False, "delta": "Hello "},
{"message_id": "m1", "index": 1, "final": True, "delta": "world"},
],
)
captured: list[_CapturedDeltaPost] = []
seen: dict[tuple[str, int], None] = {}
async with _delta_capture_client(captured, batch_supported=False) as client:
reset_events_batch_support(client)
await forwarder._forward_available_deltas(
client=client,
session_id="conv_x",
bridge_dir=bridge_dir,
state=forwarder.DeltaForwardState(),
seen_keys=seen,
)
# One probe, then both chunks the old way.
assert [c.url_path for c in captured] == [
"/v1/sessions/conv_x/events/batch",
"/v1/sessions/conv_x/events",
"/v1/sessions/conv_x/events",
]
# The next poll must not re-probe: the miss is latched per client.
_write_deltas_file(
bridge_dir, [{"message_id": "m2", "index": 0, "final": True, "delta": "again"}]
)
captured.clear()
await forwarder._forward_available_deltas(
client=client,
session_id="conv_x",
bridge_dir=bridge_dir,
state=forwarder.DeltaForwardState(),
seen_keys=seen,
)
# The one new chunk went straight to the per-event route — no second probe.
assert [c.url_path for c in captured] == ["/v1/sessions/conv_x/events"]
async def test_forward_available_deltas_keeps_going_when_one_batched_chunk_is_rejected(
tmp_path: Path,
) -> None:
"""
One rejected chunk inside a batch doesn't stop the rest of the run.
Deltas are a preview the final message supersedes, so the tail matters
more than any single chunk. Fails if a mid-run rejection stalls the
cursor (chunks re-POST forever) or raises into the poll loop.
"""
bridge_dir = prepare_bridge_dir("conv_x", bridge_id="b1", workspace=tmp_path)
_write_deltas_file(
bridge_dir,
[
{"message_id": "m1", "index": 0, "final": False, "delta": "a"},
{"message_id": "m1", "index": 1, "final": False, "delta": "b"},
{"message_id": "m1", "index": 2, "final": True, "delta": "c"},
],
)
captured: list[_CapturedDeltaPost] = []
seen: dict[tuple[str, int], None] = {}
async with _delta_capture_client(captured, batch_event_statuses=[202, 400, 202]) as client:
new_state = await forwarder._forward_available_deltas(
client=client,
session_id="conv_x",
bridge_dir=bridge_dir,
state=forwarder.DeltaForwardState(),
seen_keys=seen,
)
assert len(captured) == 1
assert new_state.byte_offset == os.path.getsize(bridge_dir / "message_deltas.jsonl")
@@ -6163,15 +6727,13 @@ async def test_scheduled_wake_forwards_marker_under_a_new_turn_id(tmp_path: Path
def _handle_request(request: httpx.Request) -> httpx.Response:
"""
Accept every forwarder POST, recording its payload.
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
payload = json.loads(request.content.decode("utf-8"))
assert isinstance(payload, dict)
requests.append(payload)
return httpx.Response(202, json={})
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -8166,13 +8728,13 @@ async def test_short_turn_poll_posts_items_without_a_status_edge(tmp_path: Path)
def _handle_request(request: httpx.Request) -> httpx.Response:
"""
Record every forwarder POST body.
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
posted.append(json.loads(request.content.decode("utf-8")))
return httpx.Response(202, json={})
batched = _record_forwarder_post(posted, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
@@ -8259,10 +8821,14 @@ async def test_forwarder_does_not_leave_running_open_for_slash_command_only_turn
requests: list[dict[str, Any]] = []
def _handle_request(request: httpx.Request) -> httpx.Response:
payload = json.loads(request.content.decode("utf-8"))
assert isinstance(payload, dict)
requests.append(payload)
return httpx.Response(202, json={})
"""
Accept every forwarder POST, recording each event it carries.
:param request: Outbound HTTP request from the forwarder.
:returns: HTTP 202 for the mock Omnigent endpoint.
"""
batched = _record_forwarder_post(requests, request)
return batched if batched is not None else httpx.Response(202, json={})
transport = httpx.MockTransport(_handle_request)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
+216
View File
@@ -12,12 +12,16 @@ from omnigent._native_post_delivery import (
_DEAD_LETTER_BACKUP_FILE,
_DEAD_LETTER_FILE,
_DEAD_LETTER_MAX_BYTES,
MAX_EVENTS_PER_BATCH,
RepostResult,
_dead_letter_record_replayable,
append_dead_letter,
events_batch_supported,
post_external_session_status,
post_may_have_been_delivered,
post_session_events_batch,
replay_dead_letters,
reset_events_batch_support,
)
# Status codes a forwarder treats as transient/retryable, used by the replay
@@ -676,3 +680,215 @@ async def test_post_external_session_status_attaches_failure_reason() -> None:
"output": "transcript item item-1 rejected",
}
assert captured[1]["data"] == {"status": "idle"}
def _batch_client(
requests: list[dict[str, object]],
*,
responder: object = None,
) -> httpx.AsyncClient:
"""
Build an AsyncClient recording every batch request body.
:param requests: List appended to with each parsed request body.
:param responder: Optional ``(index, body) -> httpx.Response`` override,
for exercising route misses and partial failures. Defaults to
accepting every event.
:returns: An ``httpx.AsyncClient`` bound to the recording transport.
"""
def handler(request: httpx.Request) -> httpx.Response:
body = json.loads(request.content)
requests.append({"path": request.url.path, "body": body})
if responder is not None:
return responder(len(requests) - 1, body) # type: ignore[operator]
return httpx.Response(
200,
json={
"results": [
{"index": index, "status": 202} for index in range(len(body["events"]))
],
"stopped_at": None,
},
)
return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url="http://ap")
@pytest.mark.asyncio
async def test_batch_post_sends_one_request_and_reports_each_event() -> None:
"""
A run of events costs one request, with one outcome per event.
This is the round-trip saving the whole change exists for: a client
far from the server can otherwise only push one event per RTT. Fails
if the helper fans back out into per-event posts or loses the
per-event acknowledgement callers key their cursor off.
"""
requests: list[dict[str, object]] = []
events = [{"type": "external_output_text_delta", "data": {"delta": str(i)}} for i in range(5)]
async with _batch_client(requests) as client:
reset_events_batch_support(client)
outcomes = await post_session_events_batch(client, session_id="conv_x", events=events)
assert len(requests) == 1
assert requests[0]["path"] == "/v1/sessions/conv_x/events/batch"
assert outcomes is not None
assert [(o.index, o.delivered) for o in outcomes] == [(i, True) for i in range(5)]
@pytest.mark.asyncio
async def test_batch_post_chunks_runs_over_the_cap() -> None:
"""
A run longer than the cap is split, not truncated or rejected.
A resumed session can replay a large backlog in one poll; the helper
must ship all of it. Fails if events past the cap are silently
dropped (the live view would be missing a chunk of the turn).
"""
requests: list[dict[str, object]] = []
total = MAX_EVENTS_PER_BATCH + 3
events = [
{"type": "external_output_text_delta", "data": {"delta": str(i)}} for i in range(total)
]
async with _batch_client(requests) as client:
reset_events_batch_support(client)
outcomes = await post_session_events_batch(client, session_id="conv_x", events=events)
assert [len(r["body"]["events"]) for r in requests] == [MAX_EVENTS_PER_BATCH, 3] # type: ignore[index]
assert outcomes is not None
assert [o.index for o in outcomes] == list(range(total))
assert all(o.delivered for o in outcomes)
@pytest.mark.asyncio
async def test_batch_post_stops_after_a_failed_chunk() -> None:
"""
With ``on_error="stop"`` a failure ends the run there.
A caller advancing a durable cursor must not have later events land
ahead of a failed one — on retry it re-sends from the failure, and a
delivered successor would then be duplicated. Fails if the helper
keeps sending chunks past a failure, or reports unattempted events as
delivered.
"""
requests: list[dict[str, object]] = []
total = MAX_EVENTS_PER_BATCH + 2
events = [{"type": "external_conversation_item", "data": {"i": i}} for i in range(total)]
def responder(index: int, body: dict[str, object]) -> httpx.Response:
"""Fail the second event of the first chunk."""
del index
results = [
{"index": position, "status": 202}
for position in range(len(body["events"])) # type: ignore[arg-type]
]
results[1] = {"index": 1, "status": 400, "error": "bad item"}
return httpx.Response(200, json={"results": results[:2], "stopped_at": 1})
async with _batch_client(requests, responder=responder) as client:
reset_events_batch_support(client)
outcomes = await post_session_events_batch(client, session_id="conv_x", events=events)
# Only the first chunk was sent.
assert len(requests) == 1
assert outcomes is not None
assert outcomes[0].delivered is True
assert outcomes[1].delivered is False
assert outcomes[1].error == "bad item"
# Everything after the failure is reported unattempted, never delivered.
assert not any(o.delivered for o in outcomes[1:])
assert [o.index for o in outcomes] == list(range(total))
@pytest.mark.asyncio
async def test_batch_post_returns_none_on_route_miss_and_latches_it() -> None:
"""
An older server's route miss tells the caller to post singly, once.
A user's CLI can be newer than the deployment it talks to. Fails if
the miss surfaces as an error (the forwarder would drop the run) or
if every later run re-probes (a wasted round trip per poll, on the
very link this change is trying to spare).
"""
requests: list[dict[str, object]] = []
events = [{"type": "external_output_text_delta", "data": {"delta": "a"}}]
def responder(index: int, body: dict[str, object]) -> httpx.Response:
"""Answer with Starlette's route-miss shape."""
del index, body
return httpx.Response(404, json={"detail": "Not Found"})
async with _batch_client(requests, responder=responder) as client:
reset_events_batch_support(client)
assert await post_session_events_batch(client, session_id="conv_x", events=events) is None
assert events_batch_supported(client) is False
assert await post_session_events_batch(client, session_id="conv_x", events=events) is None
assert len(requests) == 1
@pytest.mark.asyncio
async def test_batch_post_keeps_batching_when_the_session_is_missing() -> None:
"""
An application 404 is not a route miss and must not latch the fallback.
``POST /events/batch`` answers 404 both when the deployment lacks the
route and when the session is gone; only the first should downgrade
the client. Fails if a deleted session permanently costs a
far-away client its batching.
"""
requests: list[dict[str, object]] = []
events = [{"type": "external_output_text_delta", "data": {"delta": "a"}}]
def responder(index: int, body: dict[str, object]) -> httpx.Response:
"""Answer with Omnigent's application-error envelope."""
del index, body
return httpx.Response(
404, json={"error": {"code": "not_found", "message": "Session not found."}}
)
async with _batch_client(requests, responder=responder) as client:
reset_events_batch_support(client)
with pytest.raises(httpx.HTTPStatusError):
await post_session_events_batch(client, session_id="conv_x", events=events)
assert events_batch_supported(client) is True
@pytest.mark.asyncio
async def test_batch_post_treats_a_malformed_response_as_an_http_error() -> None:
"""
A non-JSON batch reply is an error, not silent success.
Callers advance cursors on delivery; an unparseable reply proves
nothing about delivery, so it must reach the caller's failure path.
Fails if the helper swallows it and reports the run as delivered.
"""
requests: list[dict[str, object]] = []
events = [{"type": "external_output_text_delta", "data": {"delta": "a"}}]
def responder(index: int, body: dict[str, object]) -> httpx.Response:
"""Answer 200 with a body that is not JSON."""
del index, body
return httpx.Response(200, content=b"<html>gateway</html>")
async with _batch_client(requests, responder=responder) as client:
reset_events_batch_support(client)
with pytest.raises(httpx.HTTPError):
await post_session_events_batch(client, session_id="conv_x", events=events)
@pytest.mark.asyncio
async def test_batch_post_of_nothing_is_a_noop() -> None:
"""
An empty run makes no request.
The forwarder calls this every poll; an idle poll must not cross the
wire at all. Fails if an empty run posts an empty envelope (which the
server rejects) on every poll.
"""
requests: list[dict[str, object]] = []
async with _batch_client(requests) as client:
reset_events_batch_support(client)
assert await post_session_events_batch(client, session_id="conv_x", events=[]) == []
assert requests == []