Skip to content

agentplane: replace test Client SSE with typed History reads #16223

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
agentydragon wants to merge 1 commit into
base: devel
Choose a base branch
from
+897 −423
Open

agentplane: replace test Client SSE with typed History reads #16223

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion agentplane/app/testing/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ py_library(
"//agentplane/app:app_service_pb2",
"//agentplane/environments/v1:environment_service_pb2",
"//agentplane/grpc:calls",
"//agentplane/proto:event_log_pb2",
"//agentplane/history_service/api/history:history_service_pb2",
"//agentplane/proto:identity_pb2",
"//agentplane/proto:session_pb2",
"//agentplane/runner_gateway/dispatcher:testing_client",
Expand All @@ -222,15 +222,32 @@ py_library(
],
)

py_library(
name = "history_reads",
srcs = ["history_reads.py"],
visibility = ["//agentplane/system_tests:__subpackages__"],
deps = [
":client",
"//agentplane/history_service/api/history:history_service_pb2",
"//agentplane/history_service/testing:unstored_ranges",
"//agentplane/proto:event_log_pb2",
"@pypi//httpx",
"@pypi//protobuf",
"@pypi//tenacity",
],
)

py_test(
name = "test_client",
srcs = ["test_client.py"],
deps = [
":client",
":history_reads",
":models",
"//:conftest",
"//agentplane/env/runtime:harness",
"//agentplane/environments/v1:environment_service_pb2",
"//agentplane/history_service/api/history:history_service_pb2",
"//agentplane/proto:event_log_pb2",
"//agentplane/proto:identity_pb2",
"//agentplane/proto:session_pb2",
Expand All @@ -239,6 +256,7 @@ py_test(
"@pypi//protobuf",
"@pypi//pytest",
"@pypi//pytest_bazel",
"@pypi//tenacity",
],
)

Expand Down
117 changes: 25 additions & 92 deletions agentplane/app/testing/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

import asyncio
import json
import logging
import time
from collections.abc import AsyncGenerator
from http import HTTPStatus
Expand All @@ -19,29 +18,26 @@

import grpc
import httpx
from google.protobuf.json_format import MessageToDict, ParseDict, ParseError
from google.protobuf.json_format import MessageToDict, ParseDict
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from tenacity import AsyncRetrying, retry_if_exception, stop_after_delay, wait_fixed

from agentplane.app import app_service_pb2
from agentplane.app.testing.models import ThreadMetadata, ThreadView, thread_view_from_proto_json
from agentplane.environments.v1 import environment_service_pb2
from agentplane.grpc.calls import ServiceError
from agentplane.proto import event_log_pb2, identity_pb2, session_pb2
from agentplane.history_service.api.history import history_service_pb2
from agentplane.proto import identity_pb2, session_pb2
from agentplane.runner_gateway.dispatcher.api.session import session_dispatcher_pb2
from agentplane.runner_gateway.dispatcher.testing_client import SessionDispatcherClient
from agentplane.sandboxes.runtimeconfig.model_catalog import ModelCatalog

# The generated protocol stubs' own stub chain, which the mypy aspect resolves for direct deps only.

logger = logging.getLogger(__name__)

REQUEST_SECONDS = 30.0


# How long a Thread's stream may stay down, reconnecting, before the caller is told. Testing's app
# Deployment recreates its only replica, which took about three minutes to answer again (migrate
# initContainer, then startup) when its pod was deleted on 2026-10-13.
# How long idempotent History reads may retry while the app has no ready replica.
RECONNECT_SECONDS = 600.0
RECONNECT_PAUSE_SECONDS = 2.0
GRPC_UNAVAILABLE = 14
Expand All @@ -63,14 +59,6 @@ def _thread_not_found(error: BaseException) -> bool:
return isinstance(error, httpx.HTTPStatusError) and error.response.status_code == HTTPStatus.NOT_FOUND


class SessionStreamError(Exception):
"""The app ended a session's event stream with an error frame, or it stayed down too long."""


class _StreamDroppedError(Exception):
"""The stream's replica went away; the stream is to be reopened after its last entry."""


class Client:
"""Async context manager over one deployment of the app."""

Expand Down Expand Up @@ -345,6 +333,11 @@ async def wait_for_thread_projection(self, thread_id: UUID, *, timeout_s: float)
reraise=True,
)(self.thread, thread_id)

async def thread_metadata(self, thread_id: UUID) -> ThreadMetadata:
"""The app-owned Thread row, independent of History's projection."""
data = await self._json("GET", f"/upstream/app/v1/threads/{thread_id}")
return ThreadMetadata.model_validate(data)

async def threads(self, *, sandbox_uid: str) -> AsyncGenerator[ThreadView]:
"""Every retained Thread of this Sandbox incarnation, including archived Threads."""
params = {"include_archived": "true", "limit": "200"}
Expand Down Expand Up @@ -372,80 +365,20 @@ async def thread(self, thread_id: UUID) -> ThreadView:
summary = await self._json("GET", f"/upstream/history/v1/sessions/{thread_id}/summary")
return thread_view_from_proto_json(metadata, summary)

async def events(
self, thread_id: UUID, *, after: int, read_seconds: float, awaiting_command_id: str | None = None
) -> AsyncGenerator[event_log_pb2.EventEntry]:
"""The Thread's events from `after`, ending when the app sends the `end` frame, and not before
the admission of `awaiting_command_id` (a command submitted after `after`).

A stream that drops without one -- an app rollout or crash, a closed connection, a `502`, `503` or `504`, or the proxy's plain-text `403`, while no replica is up -- is reopened after the last entry yielded, so a
caller sees every entry once and in order. An outage that outlasts `RECONNECT_SECONDS`
without delivering an entry raises. `read_seconds` bounds how long a single frame may take
to arrive, so a wedged session fails a caller rather than hanging it.

Server-sent events, parsed as the format actually specifies: a frame is terminated by a
blank line, `data:` may repeat and is joined with newlines, and a comment line (`:`) is a
keepalive. A frame carrying no data resets the parser like any other, so a name never
leaks into the frame after it.
"""
timeout = httpx.Timeout(REQUEST_SECONDS, read=read_seconds)
url = f"/threads/{thread_id}/events/stream"
cursor = after
outage_began: float | None = None
while True:
try:
params: dict[str, str | int] = {"after": cursor}
if awaiting_command_id is not None:
params["awaiting_command_id"] = awaiting_command_id
async with self._http.stream(
"GET", url, params=params, headers={"Accept": "text/event-stream"}, timeout=timeout
) as response:
if _no_app_answered(response):
raise _StreamDroppedError(f"{response.status_code=}")
response.raise_for_status()
kind, data = "message", []
async for line in response.aiter_lines():
if line.startswith(":"):
continue
if line.startswith("event:"):
kind = line.removeprefix("event:").strip()
elif line.startswith("data:"):
data.append(line.removeprefix("data:").strip())
elif not line:
frame, payload = kind, "\n".join(data)
kind, data = "message", []
if frame == "end":
return
if frame == "error":
raise SessionStreamError(f"Thread {thread_id}: {payload}")
if frame == "event" and payload:
entry = _event(payload)
cursor = entry.cursor
outage_began = None
if (
entry.event.HasField("command_admitted")
and entry.event.command_admitted.command.command_id == awaiting_command_id
):
# A reopened stream starts past the admission, so it must not await it.
awaiting_command_id = None
yield entry
reason = "the stream closed without an end frame"
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadError, httpx.RemoteProtocolError) as error:
reason = f"{error!r}"
except _StreamDroppedError as error:
reason = str(error)
outage_began = outage_began if outage_began is not None else time.monotonic()
if time.monotonic() - outage_began > RECONNECT_SECONDS:
raise SessionStreamError(
f"Thread {thread_id}: no event for {RECONNECT_SECONDS}s after {cursor=}: {reason}"
)
logger.warning("Thread %s stream dropped (%s); resuming after cursor %d", thread_id, reason, cursor)
await asyncio.sleep(RECONNECT_PAUSE_SECONDS)

async def read_session_events(
self, thread_id: UUID, *, after_cursor: int = 0, limit: int = 128
) -> history_service_pb2.ReadSessionEventsResponse:
"""Read one typed History page through the authenticated App facade."""
if after_cursor < 0 or not 1 <= limit <= 1000:
raise ValueError("invalid session history page")
data = await self._json(
"GET",
f"/upstream/history/v1/sessions/{quote(str(thread_id), safe='')}/events",
params={"after_cursor": after_cursor, "limit": limit},
)
return ParseDict(data, history_service_pb2.ReadSessionEventsResponse())

def _event(payload: str) -> event_log_pb2.EventEntry:
"""One `event` frame's data as an EventEntry; a payload that is not one says so."""
try:
return ParseDict(json.loads(payload), event_log_pb2.EventEntry())
except (ValueError, ParseError) as error:
raise SessionStreamError(f"not an EventEntry: {error}: {payload[:400]!r}") from error
async def session_summary(self, thread_id: UUID) -> history_service_pb2.SessionSummary:
"""Read History's raw commit and feed state, without substituting the folded Thread view."""
data = await self._json("GET", f"/upstream/history/v1/sessions/{quote(str(thread_id), safe='')}/summary")
return ParseDict(data, history_service_pb2.SessionSummary())
166 changes: 166 additions & 0 deletions agentplane/app/testing/history_reads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
"""Paged, correction-aware reads for App-authenticated Session history consumers."""

from __future__ import annotations

import asyncio
from collections.abc import Sequence
from http import HTTPStatus
from uuid import UUID

import httpx
from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_delay, wait_fixed

from agentplane.app.testing.client import RECONNECT_PAUSE_SECONDS, RECONNECT_SECONDS, Client, _no_app_answered
from agentplane.history_service.api.history import history_service_pb2
from agentplane.history_service.testing.unstored_ranges import unstored_ranges
from agentplane.proto import event_log_pb2

# The generated protobuf stubs need the protobuf runtime as a direct mypy dependency.
# gazelle:include_dep @pypi//protobuf

HISTORY_POLL_SECONDS = 0.2
_TRANSPORT_FAILURES = (
httpx.ConnectError,
httpx.ConnectTimeout,
httpx.ReadError,
httpx.ReadTimeout,
httpx.RemoteProtocolError,
)


class _RetryableHistoryReadError(Exception):
def __init__(self, cause: BaseException) -> None:
super().__init__(str(cause))
self.cause = cause


def covered_cursor(after_cursor: int, limit: int, page: history_service_pb2.ReadSessionEventsResponse) -> int:
"""Return the page's covered cursor, validating every entry or reported cursor gap.

A full page covers only through its final entry. Its `last_cursor` is a global watermark and
may include later entries that the page did not return.
"""
if after_cursor < 0 or not 1 <= limit <= 1000:
raise ValueError("invalid session history cursor or page limit")
entries = list(page.entries)
if len(entries) > limit or any(
entry.cursor <= after_cursor or (index > 0 and entry.cursor <= entries[index - 1].cursor)
for index, entry in enumerate(entries)
):
raise ValueError("History Service returned unordered or out-of-window entries")
if page.last_cursor < after_cursor or (entries and entries[-1].cursor > page.last_cursor):
raise ValueError("History Service returned an invalid global cursor")

end = entries[-1].cursor if len(entries) == limit else page.last_cursor
if end <= after_cursor:
return after_cursor

intervals = sorted(
[*((entry.cursor, entry.cursor) for entry in entries), *unstored_ranges(page)], key=lambda interval: interval[0]
)
cursor = after_cursor
for first, last in intervals:
if last <= cursor:
continue
if first > cursor + 1:
raise ValueError(f"History page did not account for cursor {cursor + 1}")
cursor = min(end, max(cursor, last))
if cursor >= end:
return cursor
if cursor < end:
raise ValueError(f"History page did not account for cursor {cursor + 1}")
return cursor


class SessionEventReader:
"""Read typed History pages while retaining corrections reported below the read cursor.

Callers evaluate their predicates after each page and keep their existing deadline around the
loop. Archive readers stop at their previously captured committed watermark.
"""

def __init__(self, client: Client, thread_id: UUID, *, after_cursor: int = 0, limit: int = 128) -> None:
if after_cursor < 0 or not 1 <= limit <= 1000:
raise ValueError("invalid session history cursor or page limit")
self.client = client
self.thread_id = thread_id
self.cursor = after_cursor
self.limit = limit
self.pending_attachment = False
self.more_available = False
self._entries: dict[int, event_log_pb2.EventEntry] = {}
self._has_read = False

@property
def entries(self) -> Sequence[event_log_pb2.EventEntry]:
return tuple(self._entries[cursor] for cursor in sorted(self._entries))

async def read_page(self) -> list[event_log_pb2.EventEntry]:
"""Read one page, retrying only idempotent GETs during an App outage."""
self.pending_attachment = False
page = await self._read_page()
if page is None:
self.more_available = False
return []

for first, last in unstored_ranges(page):
for cursor in tuple(self._entries):
if first <= cursor <= last:
del self._entries[cursor]
newly_read = list(page.entries)
for entry in newly_read:
self._entries[entry.cursor] = entry

next_cursor = covered_cursor(self.cursor, self.limit, page)
self.cursor = next_cursor
self.more_available = len(page.entries) == self.limit and self.cursor < page.last_cursor
return newly_read

async def wait_for_next_page(self) -> None:
"""Poll after a quiet page, but continue paging immediately when its watermark is ahead."""
if not self.more_available:
await asyncio.sleep(HISTORY_POLL_SECONDS)

async def _read_page(self) -> history_service_pb2.ReadSessionEventsResponse | None:
try:
return await AsyncRetrying(
stop=stop_after_delay(RECONNECT_SECONDS),
wait=wait_fixed(RECONNECT_PAUSE_SECONDS),
retry=retry_if_exception_type(_RetryableHistoryReadError),
reraise=True,
)(self._read_page_once)
except _RetryableHistoryReadError as error:
raise TimeoutError(
f"History read for Thread {self.thread_id} stayed unavailable after cursor {self.cursor}: {error.cause}"
) from error.cause

async def _read_page_once(self) -> history_service_pb2.ReadSessionEventsResponse | None:
try:
page = await self.client.read_session_events(self.thread_id, after_cursor=self.cursor, limit=self.limit)
except httpx.HTTPStatusError as error:
if error.response.status_code == HTTPStatus.NOT_FOUND:
if self._has_read:
raise
try:
await self.client.thread_metadata(self.thread_id)
except httpx.HTTPStatusError as metadata_error:
if metadata_error.response.status_code == HTTPStatus.NOT_FOUND:
raise error
if not _no_app_answered(metadata_error.response):
raise
raise _RetryableHistoryReadError(metadata_error) from metadata_error
except _TRANSPORT_FAILURES as metadata_error:
raise _RetryableHistoryReadError(metadata_error) from metadata_error
# The app owns this Thread, but History has not attached its Session yet.
self.pending_attachment = True
return None
if not _no_app_answered(error.response):
raise
raise _RetryableHistoryReadError(error) from error
except _TRANSPORT_FAILURES as error:
raise _RetryableHistoryReadError(error) from error

# Every healthy History response, including an empty page, ends this retry window. The
# caller's fixed operation deadline remains in force when no semantic progress is made.
self._has_read = True
return page
Loading
Loading