Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
60 changes: 60 additions & 0 deletions src/agent_env/env/gateway/changes.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
"""Who made each change to an env's state (urn:agentenv:changes/v1): changelog rows joined onto the gateway's write windows."""
from __future__ import annotations

import heapq
from typing import Any


class ChangeAttribution:
"""The write windows (changelog ids before and after a call) of the trajectory events logged since start."""

def __init__(self) -> None:
self._calls: dict[str, dict] = {}
self._windows: list[tuple[int, int, int, dict]] = [] # (before, after, order, actor)

def observe(self, event: dict) -> None:
"""Take one trajectory event as `_log_event` writes it."""
kind = event.get("event_type")
if kind == "tool_call":
self._calls[event["event_id"]] = {"kind": "tool_call", "event_id": event["event_id"],
"tool": (event.get("tool_call") or {}).get("function_name"),
**_who(event)}
elif kind in ("tool_call_result", "tool_call_failed"):
actor = self._calls.pop(event.get("tool_call_event_id"), None)
if actor is not None and "changelog_id_before" in event: # a gateway-native tool's result has no window
failed = {"failed": True} if kind == "tool_call_failed" else {}
self._add(event["changelog_id_before"], event.get("changelog_id"), {**actor, **failed})
elif kind == "internal_tool_call":
self._add(event.get("changelog_id_before"), event.get("changelog_id_after"), {
"kind": "trigger_action", "event_id": event["event_id"], "trigger_id": event.get("trigger_id"),
"action_index": event.get("action_index"), "tool": event.get("tool"), "role": event.get("as")})
elif kind == "http_request":
self._add(event.get("changelog_id_before"), event.get("changelog_id_after"), {
"kind": "request", "event_id": event["event_id"],
"request": {k: event.get(k) for k in ("route", "service", "method", "path", "status")},
**_who(event)})

def _add(self, before: Any, after: Any, actor: dict) -> None:
before = 0 if before is None else before # MAX(id) of an empty changelog reads null
if isinstance(before, int) and isinstance(after, int) and after > before:
self._windows.append((before, after, len(self._windows), actor))

def actors(self, row_ids: list[int]) -> list[list[dict]]:
"""For each of ``row_ids`` (ascending), the actors whose ``(before, after]`` window holds it, by finish order."""
windows = sorted(self._windows)
open_windows: list[tuple[int, int, dict]] = [] # a heap of (after, order, actor)
out: list[list[dict]] = []
i = 0
for row_id in row_ids:
while i < len(windows) and windows[i][0] < row_id:
before, after, order, actor = windows[i]
heapq.heappush(open_windows, (after, order, actor))
i += 1
while open_windows and open_windows[0][0] < row_id:
heapq.heappop(open_windows)
out.append([actor for _, _, actor in sorted(open_windows, key=lambda w: w[1])])
return out


def _who(event: dict) -> dict:
return {"role": event.get("role"), **({"caused_by": event["caused_by"]} if event.get("caused_by") else {})}
16 changes: 15 additions & 1 deletion src/agent_env/env/gateway/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,11 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int:
EXT_TRAJECTORY_URI = "urn:agentenv:trajectory/v1"
EXT_STEP_URI = "urn:agentenv:step/v1"
EXT_STATE_URI = "urn:agentenv:state/v1"
EXT_CHANGES_URI = "urn:agentenv:changes/v1"

# Rows per GET /changes page: the default, and the most a client may ask for.
CHANGES_PAGE_DEFAULT = 1000
CHANGES_PAGE_MAX = 10000

# Gateway's compose service name — how backing servers reach it to build env_get_time_url (compose path).
GATEWAY_COMPOSE_HOST = "gateway"
Expand Down Expand Up @@ -198,7 +203,7 @@ def random_mcp_server_name() -> str:
},
{
"uri": EXT_TRAJECTORY_URI,
"description": "Env trajectory (urn:agentenv:trajectory/v1): the env's append-only JSONL record of tool calls, results, and trigger firings, virtual_time-stamped while the clock is armed.",
"description": "Env trajectory (urn:agentenv:trajectory/v1): the env's append-only JSONL record of tool calls, results, trigger firings and the requests that may write through its REST proxy, each write with the changelog ids it spanned, virtual_time-stamped while the clock is armed.",
"params": {"endpoint": "/trajectory", "methods": {
"get": {"method": "GET", "endpoint": "/trajectory"}}},
"required": False,
Expand All @@ -217,4 +222,13 @@ def random_mcp_server_name() -> str:
"get": {"method": "GET", "endpoint": "/state"}}},
"required": False,
},
{
"uri": EXT_CHANGES_URI,
"description": "Changes (urn:agentenv:changes/v1): the env's changelog rows after after_id, oldest first, each with the actors whose write window holds it: the role, the tool call, trigger action or request, and the trigger fires it was made for. None for a write the gateway did not serve; several for one written while calls overlapped.",
"params": {"endpoint": "/changes", "methods": {
"get": {"method": "GET", "endpoint": "/changes", "request": {"type": "object", "properties": {
"after_id": {"type": "integer", "minimum": 0},
"limit": {"type": "integer", "minimum": 1, "maximum": CHANGES_PAGE_MAX}}}}}},
"required": False,
},
]
97 changes: 93 additions & 4 deletions src/agent_env/env/gateway/gateway.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@
AGENT_ENV_ROLE_META_KEY, # noqa: F401 (the server module serves the whole vocabulary, used in this file or not)
AGENT_ENV_SESSION_META_KEY, # noqa: F401
CARD_FETCH_TIMEOUT_S,
CHANGES_PAGE_DEFAULT,
CHANGES_PAGE_MAX,
CHILD_DISCOVERY_TIMEOUT_S,
DATA_PLANE_LOAD_TIMEOUT_MAX_S,
DEFAULT_ROLE,
Expand All @@ -62,6 +64,7 @@
tool_name_clashes_text,
)
from .caller import Caller
from .changes import ChangeAttribution
from .child_transport import child_http_client
from .clock import Clock, ClockError, _iso
from .get_time import GET_TIME_TOOL_NAME, build_get_time_tool
Expand All @@ -70,6 +73,9 @@

logger = logging.getLogger(__name__)

# RFC 9110's safe methods: a request using one asks for no change, so the proxy measures no write window for it.
_SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE"})


class ToolNameClash(ValueError):
"""More than one source serves a tool under one name, so an agent couldn't tell which it calls: the gateway serves
Expand Down Expand Up @@ -139,6 +145,7 @@ def __init__(

self._event_counter = 0
self._event_lock = asyncio.Lock()
self._changes = ChangeAttribution()
Path(GATEWAY_TRAJECTORY_FILE).parent.mkdir(parents=True, exist_ok=True)
self._trajectory_file = open(GATEWAY_TRAJECTORY_FILE, "a")
atexit.register(self._close_trajectory_file)
Expand Down Expand Up @@ -291,6 +298,10 @@ async def clock_state(request: Request) -> Response:
async def env_trajectory(request: Request) -> Response:
return self._serve_trajectory()

@self._mcp.custom_route("/changes", methods=["GET"])
async def env_changes(request: Request) -> Response:
return await self._serve_changes(request)

@self._mcp.custom_route(WELL_KNOWN_PATH, methods=["GET"])
async def get_env_card(request: Request) -> Response:
return await self._serve_env_card()
Expand Down Expand Up @@ -505,8 +516,31 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] |
base_url = url_map.get(service_name)
if base_url is None:
return self._json_response({"error": f"Unknown service: {service_name}"}, 404)
return await forward(request, base_url, role=self._request_role(request), timeout=self.REST_PROXY_TIMEOUT,
mcp_server=url_map is self.rest_proxy_urls and service_name.startswith("mcp-"))
role = self._request_role(request)
mcp_server = url_map is self.rest_proxy_urls and service_name.startswith("mcp-")
if request.method in _SAFE_METHODS:
return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server)
changelog_before = self._query_changelog_id()
attribution = self._attribution(role) # at the start: a fire may be gone by the time the response ends

async def log_write(status: int) -> None:
changelog_after = self._query_changelog_id()
if changelog_after is None or changelog_after <= (changelog_before or 0):
return # only a request that wrote changelog rows is a write window
await self._log_event({
"event_type": "http_request",
"route": "website" if url_map is self.website_urls else "svc",
"service": service_name,
"method": request.method,
"path": "/" + request.path_params.get("path", ""),
"status": status,
**attribution,
"changelog_id_before": changelog_before,
"changelog_id_after": changelog_after,
})

return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server,
on_done=log_write)

@staticmethod
def _request_role(request: Request) -> str:
Expand Down Expand Up @@ -950,14 +984,20 @@ async def _step_call_tool(self, body: dict, role: str) -> Response:
})
try:
async with self._tool_call_guard():
result = await self._call_child_tool(server_url, tool_name, arguments, caller=Caller(role=role))
changelog_before = self._query_changelog_id()
try:
result = await self._call_child_tool(server_url, tool_name, arguments, caller=Caller(role=role))
except Exception as e:
await self._log_failed_call(tool_call_event_id, changelog_before, e, source="step_api")
raise
changelog_id = self._query_changelog_id()
Comment thread
sarah-chitty-1 marked this conversation as resolved.

await self._log_event({
"event_type": "tool_call_result",
"source": "step_api",
"tool_call_event_id": tool_call_event_id,
"tool_call_result": result.model_dump(mode="json"),
"changelog_id_before": changelog_before,
"changelog_id": changelog_id,
})
barrier_tasks, barrier_timeout_s = self._fire_triggers(role, tool_name, arguments, result)
Expand Down Expand Up @@ -1073,13 +1113,19 @@ async def proxy(**kwargs) -> CallToolResult:
})

async with self._tool_call_guard():
result = await self._call_child_tool(server_url, tool_name, kwargs, caller=caller)
changelog_before = self._query_changelog_id()
try:
result = await self._call_child_tool(server_url, tool_name, kwargs, caller=caller)
except Exception as e:
await self._log_failed_call(tool_call_event_id, changelog_before, e)
raise
changelog_id = self._query_changelog_id()

await self._log_event({
"event_type": "tool_call_result",
"tool_call_event_id": tool_call_event_id,
"tool_call_result": result.model_dump(mode="json"),
"changelog_id_before": changelog_before,
"changelog_id": changelog_id,
})
return result
Expand Down Expand Up @@ -1136,6 +1182,16 @@ async def _log_native_tool_result(self, event_id: str | None, result: Any) -> No
except Exception:
logger.exception("Failed to log tool_call_result for a gateway tool")

async def _log_failed_call(self, tool_call_event_id: str, changelog_before: int | None, error: Exception,
**source: str) -> None:
"""Close a failed call's write window, which may hold a write it committed first (fail-open)."""
try:
await self._log_event({"event_type": "tool_call_failed", **source, "tool_call_event_id": tool_call_event_id,
"error": f"{type(error).__name__}: {str(error)[:300]}",
"changelog_id_before": changelog_before, "changelog_id": self._query_changelog_id()})
except Exception:
logger.exception("Failed to log a failed tool call")

def _attribution(self, role: str) -> dict:
"""Who made a tool call, for its trajectory event: the caller's role, and the trigger fires it was made for
when that role is an executor realizing them."""
Expand All @@ -1156,8 +1212,41 @@ async def _log_event(self, event: dict) -> str:
event, default=lambda o: o.model_dump(mode="json") if hasattr(o, "model_dump") else str(o),
) + "\n")
self._trajectory_file.flush()
self._changes.observe(event)
return event_id

async def _serve_changes(self, request: Request) -> Response:
"""The changelog rows after ``after_id`` (at most ``limit``, oldest first), each with its actors (see ``changes``)."""
try:
after_id = int(request.query_params.get("after_id", 0))
limit = int(request.query_params.get("limit", CHANGES_PAGE_DEFAULT))
except ValueError:
return self._json_response({"ok": False, "error": "after_id and limit must be integers"}, 400)
if after_id < 0 or not 1 <= limit <= CHANGES_PAGE_MAX:
return self._json_response(
{"ok": False, "error": f"after_id must be >= 0 and limit between 1 and {CHANGES_PAGE_MAX}"}, 400)
if self._db_conn is None:
return self._json_response({"ok": False, "error": "this env keeps no changelog"}, 404)
try:
rows = await asyncio.to_thread(self._read_changelog, after_id, limit + 1)
except Exception as e:
logger.exception("Reading the changelog failed")
return self._json_response({"ok": False, "error": f"changelog unavailable: {e}"}, 503)
rows, has_more = rows[:limit], len(rows) > limit
actors = self._changes.actors([row["id"] for row in rows])
return self._json_response(
{"changes": [{**row, "actors": who} for row, who in zip(rows, actors)], "has_more": has_more}, 200)

def _read_changelog(self, after_id: int, count: int) -> list[dict]:
with self._db_conn.cursor() as cur:
cur.execute('SELECT id, "timestamp", schema_name, table_name, operation, row_id, summary, changed_fields '
"FROM public._changelog WHERE id > %s ORDER BY id LIMIT %s", (after_id, count))
names = [column[0] for column in cur.description]
rows = [dict(zip(names, values)) for values in cur.fetchall()]
for row in rows:
row["timestamp"] = row["timestamp"].isoformat()
return rows

def _serve_trajectory(self) -> Response:
"""Stream the whole env trajectory as JSONL; only a mid-write tail line can arrive torn."""
try:
Expand Down
26 changes: 19 additions & 7 deletions src/agent_env/env/gateway/rest_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
import json
import logging
import posixpath
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Awaitable, Callable
from urllib.parse import unquote

import httpx
Expand Down Expand Up @@ -52,14 +52,16 @@
ROLE_HEADER_BYTES = AGENT_ENV_ROLE_HEADER.lower().encode()


async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.Timeout, mcp_server: bool) -> Response:
async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.Timeout, mcp_server: bool,
on_done: Callable[[int], Awaitable[None]] | None = None) -> Response:
"""Forward `request` to the service at `base_url` as `role` and stream the response through. `mcp_server` says
the service is a backing MCP server, whose own endpoint takes the role in its JSON-RPC `_meta` as well as in the
header. A 400 names a request-target the proxy cannot route; a 502 names an upstream that never answered."""
header. A 400 names a request-target the proxy cannot route; a 502 names an upstream that never answered.
`on_done` is awaited with the status once the response has ended, after its last byte when streamed."""
try:
target_url = upstream_url(base_url, request)
except ValueError as e: # an escaped slash inside the route prefix; refuse rather than forward a wrong path
return _json_response({"error": str(e)}, 400)
return await _done(_json_response({"error": str(e)}, 400), on_done)
content = await request.body() # /svc and /website request bodies are small (JSON-RPC, forms)
headers = upstream_headers(request, role)
if mcp_server and request.method == "POST" and is_child_mcp_path(request.path_params.get("path", "")):
Expand All @@ -74,8 +76,14 @@ async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.
except Exception as e:
await client.aclose()
logger.error(f"Proxy to {target_url} failed: {e!r}")
return _json_response({"error": f"Proxy request failed: {e!r}"}, 502)
return UpstreamResponse(upstream, client, target_url)
return await _done(_json_response({"error": f"Proxy request failed: {e!r}"}, 502), on_done)
return UpstreamResponse(upstream, client, target_url, on_done)


async def _done(response: Response, on_done: Callable[[int], Awaitable[None]] | None) -> Response:
if on_done is not None:
await on_done(response.status_code)
return response


def upstream_url(base_url: str, request: Request) -> str:
Expand Down Expand Up @@ -161,13 +169,15 @@ class UpstreamResponse(StreamingResponse):

_DROP = HOP_BY_HOP | {b"date", b"server"} # uvicorn adds the gateway's own; forwarding upstream's duplicates them

def __init__(self, upstream: httpx.Response, client: httpx.AsyncClient, target_url: str) -> None:
def __init__(self, upstream: httpx.Response, client: httpx.AsyncClient, target_url: str,
on_done: Callable[[int], Awaitable[None]] | None = None) -> None:
super().__init__(self._body(upstream, target_url), status_code=upstream.status_code)
# Header block forwarded byte for byte: repeated names (Set-Cookie) stay separate, non-latin-1 values
# never hit a str round-trip, and Content-Length/Content-Encoding stay valid for the undecoded body.
self.raw_headers = [(k.lower(), v) for k, v in upstream.headers.raw if k.lower() not in self._DROP]
self._upstream = upstream
self._client = client
self._on_done = on_done

@staticmethod
async def _body(upstream: httpx.Response, target_url: str) -> AsyncIterator[bytes]:
Expand All @@ -186,6 +196,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
finally: # normal end, upstream error, or the downstream client going away
await self._upstream.aclose()
await self._client.aclose()
if self._on_done is not None:
await self._on_done(self.status_code)


async def fetch_child_card(client: httpx.AsyncClient, key: str, base_url: str) -> dict | None:
Expand Down
Loading
Loading