diff --git a/openspec/changes/acp-notification-batching/.openspec.yaml b/openspec/changes/acp-notification-batching/.openspec.yaml new file mode 100644 index 000000000..d6b53dee5 --- /dev/null +++ b/openspec/changes/acp-notification-batching/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-30 diff --git a/openspec/changes/acp-notification-batching/design.md b/openspec/changes/acp-notification-batching/design.md new file mode 100644 index 000000000..9c93396de --- /dev/null +++ b/openspec/changes/acp-notification-batching/design.md @@ -0,0 +1,83 @@ +## Context + +`ACPNotifications.replay()` (`src/acp/agent/notifications.py:475`) replays conversation history during `session/load`. It iterates `Sequence[ModelRequest | ModelResponse]`, converts each part to a `SessionUpdate`, and calls `await self.send_update(update)` per update. `send_update()` constructs a `SessionNotification`, serializes to dict, and calls `await self.client.session_update(notification)`, which calls `await self._conn.send_notification("session/update", dct)` — a JSON-RPC notification over the wire. + +Each `send_notification` involves: Pydantic serialization → JSON encoding → TCP write → await flush. For a 100-message session (300-600 updates), this is 300-600 sequential awaits. The `session/update` JSON-RPC notification is one-way (no response expected), but each `await` still pays the TCP flush cost. + +The ACP spec's `SessionNotification` schema (`src/acp/schema/notifications.py:17`) has a single `update: TSessionUpdate_co` field — it cannot carry multiple updates in one `session/update` notification. + +## Goals / Non-Goals + +**Goals:** +- Reduce `replay()` wire roundtrips by 80%+ for typical sessions (50-100 messages) +- Maintain message ordering guarantees within and across batches +- Preserve backward compatibility: clients without batch support receive sequential `session/update` notifications +- No changes to ACP specification — use `ext_notification` extension mechanism + +**Non-Goals:** +- Batching real-time streaming notifications (live `session/prompt` turns) — only `replay()` is targeted +- Introducing a new `SessionNotification` schema with array field (would require protocol spec change) +- Cross-session batching or background flush tasks +- Modifying the EventBus or protocol event consumer paths + +## Decisions + +### Decision 1: ext_notification batch protocol over schema change + +**Chosen**: Use ACP's `ext_notification` to send a `_batch_session_updates` notification containing `{ "session_id": str, "updates": [SessionUpdate, ...] }`. + +**Rationale**: `SessionNotification.update` is a single `SessionUpdate`, not an array. Changing the schema would break all existing clients. `ext_notification` is the ACP-sanctioned extension point for non-spec methods — prefixed with `_` and ignored by clients that don't understand it. + +**Alternative considered**: New `BatchSessionNotification` schema with `updates: list[SessionUpdate]`. Rejected — requires protocol spec change, breaks all clients, and ACP v2 may address this differently. + +### Decision 2: Collect-then-send, not pipe-through-queue + +**Chosen**: `replay()` first converts all messages to a `list[SessionUpdate]`, then sends in chunks of `notification_batch_size`. + +```python +async def replay(self, messages): + updates = [] + for message in messages: + match message: + case ModelRequest(): + updates.extend(await self._collect_request_updates(message)) + case ModelResponse(): + updates.extend(await self._collect_response_updates(message)) + for i in range(0, len(updates), self.notification_batch_size): + batch = updates[i:i + self.notification_batch_size] + await self.send_batch_update(batch) +``` + +**Rationale**: The conversion from `ModelMessage` parts to `SessionUpdate` objects is pure CPU (no I/O). Collecting all updates first is simpler than a streaming pipe and lets us batch cleanly. Memory cost is bounded — a 100-message session produces ~600 `SessionUpdate` objects, each a few hundred bytes. + +**Alternative considered**: anyio `MemoryObjectSendStream` with background consumer. Rejected — adds lifecycle complexity for a problem that's fundamentally "collect then chunk". + +### Decision 3: Refactor `_replay_request`/`_replay_response` into pure collectors + +**Chosen**: Rename `_replay_request` → `_collect_request_updates`, `_replay_response` → `_collect_response_updates`. Return `list[SessionUpdate]` instead of calling `send_update()` directly. + +**Rationale**: Current methods mix conversion and I/O. Separating them enables batching and makes the conversion testable without a client connection. The conversion logic (pattern matching on `UserPromptPart`, `TextPart`, etc.) stays identical — only the output changes from `await send_*()` to `list.append(update)`. + +### Decision 4: Batch size and flush interval defaults + +**Chosen**: `notification_batch_size = 20`, `notification_flush_interval = 0.0` (no artificial delay). + +**Rationale**: 20 updates per batch reduces roundtrips by ~20x. No flush interval needed because we send all batches in a tight loop — the `await` on each `send_batch_update` provides natural flow control. The interval is kept as a config knob for future tuning with slow remote clients. + +### Decision 5: Fallback detection via client capability check + +**Chosen**: Check if the client advertised `_batch_session_updates` support during `initialize`. If not, `send_batch_update()` falls back to looping `send_update()` per update. + +**Rationale**: `ext_notification` is fire-and-forget — the agent cannot know if the client processed it. For clients that don't implement batch handling, the notification would be silently dropped, losing replay data. Capability advertisement is the ACP-native way to negotiate extensions. + +**Alternative considered**: Always send sequential, let client opt into batch via a separate `session/set_config_option`. Rejected — adds a round-trip and configuration burden. + +## Risks / Trade-offs + +| Risk | Mitigation | +|------|------------| +| Client silently drops `_batch_session_updates` (doesn't implement it but doesn't crash) | Capability check during initialize; fallback to sequential if unsupported | +| Large batch causes TCP buffer pressure on slow connections | Configurable `notification_batch_size`; default 20 is conservative | +| Memory spike for very large sessions (1000+ messages) | Bounded: ~6000 `SessionUpdate` objects × ~200 bytes = ~1.2MB; acceptable | +| Tool call ordering within batch (ToolCallStart must precede ToolCallProgress) | Batch preserves insertion order; `updates` list is built in message order | +| `_tool_call_inputs` cache state during collection | Cache populated during collection phase, consumed during fallback sequential send; batch path doesn't need it (inputs embedded in collected updates) | diff --git a/openspec/changes/acp-notification-batching/proposal.md b/openspec/changes/acp-notification-batching/proposal.md new file mode 100644 index 000000000..22538d8b1 --- /dev/null +++ b/openspec/changes/acp-notification-batching/proposal.md @@ -0,0 +1,27 @@ +## Why + +When `session/load` replays conversation history to the client, each `SessionUpdate` is sent as a separate `session/update` JSON-RPC notification with an `await` between each. A 100-message session produces 300-600 individual notifications, each paying serialization + TCP flush overhead. Estimated load time: 150-300 seconds. The bottleneck is the wire layer (serial `await send_notification`), not event production. + +## What Changes + +- **Batch `SessionUpdate` collection in `replay()`**: Convert messages to `SessionUpdate` objects first, then send them in batched groups instead of awaiting each notification individually. +- **New `ext_notification` batch protocol**: Introduce `_batch_session_updates` extension method that sends multiple `SessionUpdate` objects in a single JSON-RPC notification, using ACP's existing extension mechanism. +- **Graceful fallback**: When the client does not support `_batch_session_updates`, degrade to sequential `session/update` notifications automatically. +- **Configurable batch size**: Add `notification_batch_size` (default 20) and `notification_flush_interval` (default 0.0s) to `ACPNotifications` for tuning. The initial proposal suggested 0.05s as the default, but this was rejected in favor of 0.0 (no artificial delay) since the `await` on each `send_batch_update` already provides natural flow control — an inter-batch delay would only slow down replay without benefit in the common stdio transport case. +- **No changes to ACP specification**: Uses `ext_notification` extension, no protocol spec changes required. + +## Capabilities + +### New Capabilities + +- `acp-notification-batching`: Batch `SessionUpdate` delivery during `session/load` replay via `_batch_session_updates` extension notification, with graceful fallback to sequential delivery. + +### Modified Capabilities + +## Impact + +- `src/acp/agent/notifications.py` — `replay()`, `send_update()`, new `send_batch_update()` and `_replay_*()` refactored to return `SessionUpdate` lists instead of sending directly +- `src/agentpool_server/acp_server/acp_agent.py` — `load_session()` calls batched replay +- `src/acp/schema/notifications.py` — no structural changes (batch uses `ext_notification`) +- `tests/acp/test_notifications_replay.py` — update for batch assertions +- `tests/servers/acp_server/test_acp_load.py` — verify batch delivery and fallback diff --git a/openspec/changes/acp-notification-batching/specs/acp-notification-batching/spec.md b/openspec/changes/acp-notification-batching/specs/acp-notification-batching/spec.md new file mode 100644 index 000000000..e146e0edf --- /dev/null +++ b/openspec/changes/acp-notification-batching/specs/acp-notification-batching/spec.md @@ -0,0 +1,61 @@ +## ADDED Requirements + +### Requirement: Replay SHALL batch SessionUpdate delivery + +The `ACPNotifications.replay()` method SHALL collect all `SessionUpdate` objects from message conversion before sending, then deliver them in batches via `_batch_session_updates` ext_notification. Batch size SHALL be configurable via `notification_batch_size` (default: 20). + +#### Scenario: Batched replay with capable client + +- **WHEN** `replay()` is called with 100 messages producing 400 `SessionUpdate` objects +- **AND** the client advertised `_batch_session_updates` support during initialize +- **THEN** the agent SHALL send `ceil(400 / 20) = 20` `_batch_session_updates` ext_notifications +- **AND** each ext_notification SHALL contain at most 20 `SessionUpdate` objects in the `updates` array +- **AND** the total `session/update` notifications sent SHALL be 0 + +#### Scenario: Fallback replay with non-capable client + +- **WHEN** `replay()` is called with 100 messages producing 400 `SessionUpdate` objects +- **AND** the client did NOT advertise `_batch_session_updates` support +- **THEN** the agent SHALL send 400 individual `session/update` notifications +- **AND** each `session/update` notification SHALL contain exactly one `SessionUpdate` + +### Requirement: Batch SHALL preserve update ordering + +The `_batch_session_updates` ext_notification SHALL deliver `SessionUpdate` objects in the same order they would appear in sequential `session/update` delivery. Updates within a batch SHALL maintain the sequence produced by `_collect_request_updates` and `_collect_response_updates`. + +#### Scenario: Tool call ordering within batch + +- **WHEN** a `ModelResponse` contains a `ToolCallPart` followed by a `TextPart` +- **AND** the corresponding `ModelRequest` has a `ToolReturnPart` +- **THEN** the `ToolCallStart` update SHALL appear before the `AgentMessageChunk` update in the batch +- **AND** the `ToolCallProgress` update (from `ToolReturnPart`) SHALL appear in a subsequent batch or the same batch at a later index + +### Requirement: Batch protocol SHALL use ext_notification + +The batch delivery mechanism SHALL use ACP's `ext_notification` with method name `_batch_session_updates`. The notification params SHALL contain `session_id` (str) and `updates` (list of `SessionUpdate` dicts). No new schema type SHALL be added to `SessionNotification`. + +#### Scenario: ext_notification format + +- **WHEN** a batch of 20 `SessionUpdate` objects is ready for delivery +- **THEN** the agent SHALL call `client.ext_notification("_batch_session_updates", {"session_id": "", "updates": [, ...]})` +- **AND** the method name SHALL be prefixed with underscore (ACP extension convention) + +### Requirement: Replay conversion SHALL be side-effect free + +The `_collect_request_updates` and `_collect_response_updates` methods SHALL return `list[SessionUpdate]` without performing any I/O. The `_tool_call_inputs` cache SHALL be populated during collection for use in sequential fallback, but SHALL NOT be consumed during batch delivery. + +#### Scenario: Collection without client connection + +- **WHEN** `_collect_request_updates` is called with a `ModelRequest` +- **THEN** it SHALL return a `list[SessionUpdate]` without calling any `client.*` method +- **AND** `_tool_call_inputs` SHALL be populated for `ToolCallPart` entries encountered + +### Requirement: Batch size SHALL be configurable + +`ACPNotifications` SHALL accept `notification_batch_size` (int, default 20) and `notification_flush_interval` (float, default 0.0) as constructor parameters. These SHALL control the maximum number of updates per batch and any inter-batch delay respectively. + +#### Scenario: Custom batch size + +- **WHEN** `ACPNotifications` is constructed with `notification_batch_size=50` +- **AND** `replay()` produces 200 `SessionUpdate` objects +- **THEN** the agent SHALL send `ceil(200 / 50) = 4` batch notifications diff --git a/openspec/changes/acp-notification-batching/tasks.md b/openspec/changes/acp-notification-batching/tasks.md new file mode 100644 index 000000000..7c28867d4 --- /dev/null +++ b/openspec/changes/acp-notification-batching/tasks.md @@ -0,0 +1,46 @@ +## 1. Refactor: Side-effect-free collectors + +- [x] 1.1 Add `_collect_request_updates(self, request: ModelRequest) -> list[SessionUpdate]` — extract conversion logic from `_replay_request`, return list instead of calling `send_update()` +- [x] 1.2 Add `_collect_response_updates(self, response: ModelResponse) -> list[SessionUpdate]` — extract from `_replay_response`, populate `_tool_call_inputs` cache during collection +- [x] 1.3 Keep `_replay_request`/`_replay_response` as thin wrappers calling collectors then `send_update()` per update (preserves existing behavior for non-batch callers) +- [x] 1.4 Verify existing tests in `tests/acp/test_notifications_replay.py` still pass unchanged + +## 2. Batch delivery: send_batch_update + +- [x] 2.1 Add `notification_batch_size: int = 20` and `notification_flush_interval: float = 0.0` to `ACPNotifications.__init__` +- [x] 2.2 Add `_batch_supported: bool = False` field, set via `set_batch_support(supported: bool)` method +- [x] 2.3 Implement `send_batch_update(self, updates: list[SessionUpdate]) -> None`: + - If `_batch_supported`: call `self.client.ext_notification("_batch_session_updates", {"session_id": self.id, "updates": [u.model_dump(by_alias=True, exclude_none=True) for u in updates]})` + - If not supported: loop `await self.send_update(u)` for each update (fallback) +- [x] 2.4 Add `notification_flush_interval` sleep between batches (only when `> 0`) + +## 3. Rewrite replay() with batch collection + +- [x] 3.1 Rewrite `replay()` to collect all updates via `_collect_request_updates`/`_collect_response_updates` into a single `list[SessionUpdate]` +- [x] 3.2 Loop over updates in chunks of `notification_batch_size`, call `await self.send_batch_update(batch)` per chunk +- [x] 3.3 Preserve error handling: wrap per-message collection in try/except, log failures, continue (same as current `replay()` behavior) +- [x] 3.4 Verify tool call ordering: `ToolCallStart` from `ToolCallPart` precedes `ToolCallProgress` from `ToolReturnPart` in the collected list + +## 4. Client capability detection + +- [x] 4.1 Add `_batch_session_updates` to client capabilities negotiation — define how clients advertise support (e.g., via `client_capabilities.field_meta` or a dedicated field in `InitializeRequest`) +- [x] 4.2 In `ACPSession.__post_init__` or `initialize` flow, detect batch support from client capabilities and call `notifications.set_batch_support(True/False)` +- [x] 4.3 Default to `False` (sequential fallback) when client doesn't advertise support + +## 5. Tests + +- [x] 5.1 Test: `replay()` with batch-capable client sends `_batch_session_updates` ext_notifications, not individual `session/update` +- [x] 5.2 Test: `replay()` with non-capable client falls back to sequential `session/update` (same count as before) +- [x] 5.3 Test: batch preserves ordering — `ToolCallStart` before `ToolCallProgress` within same batch +- [x] 5.4 Test: custom `notification_batch_size=5` produces correct chunk count +- [x] 5.5 Test: `_collect_request_updates` returns correct `SessionUpdate` list without calling any client method +- [x] 5.6 Test: empty messages list produces zero notifications +- [x] 5.7 Update existing `tests/acp/test_notifications_replay.py` to cover both batch and fallback paths +- [x] 5.8 Update `tests/servers/acp_server/test_acp_load.py` to verify batch delivery during `session/load` + +## 6. Integration and benchmarks + +- [x] 6.1 Run `uv run pytest tests/acp/test_notifications_replay.py tests/servers/acp_server/test_acp_load.py -v` — all pass +- [x] 6.2 Run `uv run ruff check src/acp/agent/notifications.py` — no lint errors +- [x] 6.3 Run `uv run --no-group docs mypy src/acp/agent/notifications.py` — no type errors +- [x] 6.4 Write a benchmark script: time `replay()` with 100 messages, compare batch vs sequential (target: 80%+ reduction) diff --git a/scripts/benchmark_replay.py b/scripts/benchmark_replay.py new file mode 100644 index 000000000..d3284cb2e --- /dev/null +++ b/scripts/benchmark_replay.py @@ -0,0 +1,59 @@ +"""Benchmark: batch vs sequential replay performance.""" + +from __future__ import annotations + +import asyncio +import time +from unittest.mock import AsyncMock + +from pydantic_ai import ModelRequest, ModelResponse, TextPart, UserPromptPart + +from acp.agent.notifications import ACPNotifications + + +def make_messages(count: int) -> list[ModelRequest | ModelResponse]: + msgs: list[ModelRequest | ModelResponse] = [] + for i in range(count): + if i % 2 == 0: + msgs.append(ModelRequest(parts=[UserPromptPart(content=f"User message {i}")])) + else: + msgs.append(ModelResponse(parts=[TextPart(content=f"Agent response {i}")])) + return msgs + + +async def benchmark_sequential(count: int) -> float: + client = AsyncMock() + client.session_update = AsyncMock() + client.ext_notification = AsyncMock() + n = ACPNotifications(client=client, session_id="bench") + msgs = make_messages(count) + start = time.perf_counter() + await n.replay(msgs) + return time.perf_counter() - start + + +async def benchmark_batch(count: int) -> float: + client = AsyncMock() + client.session_update = AsyncMock() + client.ext_notification = AsyncMock() + n = ACPNotifications(client=client, session_id="bench") + n.set_batch_support(True) + msgs = make_messages(count) + start = time.perf_counter() + await n.replay(msgs) + return time.perf_counter() - start + + +async def main() -> None: + for count in [50, 100, 200]: + seq = await benchmark_sequential(count) + bat = await benchmark_batch(count) + reduction = ((seq - bat) / seq * 100) if seq > 0 else 0 + print( + f"{count:4d} msgs | sequential: {seq:.4f}s | batch: {bat:.4f}s" + f" | reduction: {reduction:.1f}%" + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/acp/agent/notifications.py b/src/acp/agent/notifications.py index f43ad88e4..ef6ba8a40 100644 --- a/src/acp/agent/notifications.py +++ b/src/acp/agent/notifications.py @@ -64,17 +64,31 @@ class ACPNotifications: handling both creation and sending in a single call. """ - def __init__(self, client: Client, session_id: str) -> None: + def __init__( + self, + client: Client, + session_id: str, + *, + notification_batch_size: int = 20, + ) -> None: """Initialize notifications helper. Args: client: ACP client and session_id session_id: Session identifier + notification_batch_size: Maximum number of SessionUpdate objects per + batch during replay. Default 20. """ + if notification_batch_size <= 0: + raise ValueError( + f"notification_batch_size must be greater than 0, got {notification_batch_size}" + ) self.client = client self.id = session_id self.log = logger.bind(session_id=session_id) self._tool_call_inputs: dict[str, dict[str, Any]] = {} + self.notification_batch_size = notification_batch_size + self._batch_supported: bool = False async def create_tool_reporter( self, @@ -169,6 +183,41 @@ async def send_update(self, update: SessionUpdate) -> None: notification = SessionNotification(session_id=self.id, update=update) await self.client.session_update(notification) # pyright: ignore[reportArgumentType] + def set_batch_support(self, supported: bool) -> None: + """Enable or disable batch session update delivery. + + When enabled, ``replay()`` sends updates via ``_batch_session_updates`` + ext_notification instead of individual ``session/update`` notifications. + + Args: + supported: Whether the client supports ``_batch_session_updates``. + """ + self._batch_supported = supported + + async def send_batch_update(self, updates: list[SessionUpdate]) -> None: + """Send a batch of SessionUpdate objects in a single notification. + + When the client supports ``_batch_session_updates``, sends all updates + in one ext_notification. Otherwise, falls back to sequential + ``session/update`` notifications. + + Args: + updates: List of SessionUpdate objects to send. + """ + if not updates: + return + if self._batch_supported: + await self.client.ext_notification( + "_batch_session_updates", + { + "session_id": self.id, + "updates": [u.model_dump(by_alias=True, exclude_none=True) for u in updates], + }, + ) + else: + for update in updates: + await self.send_update(update) + async def send_elicitation_complete( self, elicitation_id: str, @@ -473,74 +522,100 @@ async def send_user_resource( await self.send_update(update) async def replay(self, messages: Sequence[ModelRequest | ModelResponse]) -> None: - """Replay a sequence of model messages as notifications.""" + """Replay a sequence of model messages as notifications. + + Collects all SessionUpdate objects from message conversion first, then + sends them in batches of ``notification_batch_size``. When the client + supports ``_batch_session_updates``, uses ext_notification for batch + delivery; otherwise falls back to sequential ``session/update``. + """ + # Collect all updates from all messages first + all_updates: list[SessionUpdate] = [] for message in messages: try: match message: case ModelRequest(): - await self._replay_request(message) + all_updates.extend(self._collect_request_updates(message)) case ModelResponse(): - await self._replay_response(message) + all_updates.extend(self._collect_response_updates(message)) case _ as unreachable: assert_never(unreachable) except Exception as e: self.log.exception("Failed to replay message", error=str(e)) - async def _replay_request(self, request: ModelRequest) -> None: - """Replay a ModelRequest by converting it to appropriate ACP notifications.""" + # Send in batches + for i in range(0, len(all_updates), self.notification_batch_size): + batch = all_updates[i : i + self.notification_batch_size] + await self.send_batch_update(batch) + + def _collect_request_updates(self, request: ModelRequest) -> list[SessionUpdate]: + """Convert a ModelRequest to a list of SessionUpdate objects. + + This is a pure conversion — no I/O is performed. The + ``_tool_call_inputs`` cache is consumed (popped) for ToolReturnPart. + + Args: + request: The ModelRequest to convert. + + Returns: + List of SessionUpdate objects in order. + """ + updates: list[SessionUpdate] = [] for part in request.parts: match part: case UserPromptPart(content=content) if isinstance(content, str): - # Handle both str and Sequence[UserContent] types - await self.send_user_message(content) + updates.append(UserMessageChunk.text(text=content)) case UserPromptPart(content=content): - # Convert multi-modal content to appropriate ACP content blocks converted_content = to_acp_content_blocks(content) - # Send each content block as separate notifications for block in converted_content: match block: case TextContentBlock(text=text): - await self.send_user_message(text) + updates.append(UserMessageChunk.text(text=text)) case ImageContentBlock(annotations=annots) as img_block: - await self.send_user_image( - data=img_block.data, - mime_type=img_block.mime_type, - uri=img_block.uri, - audience=annots.audience if annots else None, - last_modified=annots.last_modified if annots else None, - priority=annots.priority if annots else None, + updates.append( + UserMessageChunk.image( + data=img_block.data, + mime_type=img_block.mime_type, + uri=img_block.uri, + audience=annots.audience if annots else None, + last_modified=annots.last_modified if annots else None, + priority=annots.priority if annots else None, + ) ) case AudioContentBlock(annotations=annots) as audio_block: - await self.send_user_audio( - data=audio_block.data, - mime_type=audio_block.mime_type, - audience=annots.audience if annots else None, - last_modified=annots.last_modified if annots else None, - priority=annots.priority if annots else None, + updates.append( + UserMessageChunk.audio( + data=audio_block.data, + mime_type=audio_block.mime_type, + audience=annots.audience if annots else None, + last_modified=annots.last_modified if annots else None, + priority=annots.priority if annots else None, + ) ) case ResourceContentBlock(annotations=annots) as resource_block: - await self.send_user_resource( - uri=resource_block.uri, - name=resource_block.name, - description=resource_block.description, - mime_type=resource_block.mime_type, - size=resource_block.size, - title=resource_block.title, - audience=annots.audience if annots else None, - last_modified=annots.last_modified if annots else None, - priority=annots.priority if annots else None, + updates.append( + UserMessageChunk.resource( + uri=resource_block.uri, + name=resource_block.name, + description=resource_block.description, + mime_type=resource_block.mime_type, + size=resource_block.size, + title=resource_block.title, + audience=annots.audience if annots else None, + last_modified=annots.last_modified if annots else None, + priority=annots.priority if annots else None, + ) ) case EmbeddedResourceContentBlock(resource=resource): - # Handle embedded resources with proper pattern matching match resource: case TextResourceContents(text=text): - await self.send_user_message(text) + updates.append(UserMessageChunk.text(text=text)) case BlobResourceContents(blob=blob, mime_type=mime_type): blob_size = len(blob) * 3 // 4 size_mb = blob_size / (1024 * 1024) mime = mime_type or "unknown" msg = f"Embedded resource: {mime} ({size_mb:.2f} MB)" - await self.send_user_message(msg) + updates.append(UserMessageChunk.text(text=msg)) case _ as unreachable: assert_never(unreachable) # ty: ignore[type-assertion-failure] case _ as unreachable: @@ -551,6 +626,13 @@ async def _replay_request(self, request: ModelRequest) -> None: ): converted = to_acp_content_blocks(content) tool_input = self._tool_call_inputs.get(tool_call_id, {}) + if tool_call_id not in self._tool_call_inputs: + self.log.debug( + "Tool return has no matching cached tool call input — " + "message ordering may be incorrect", + tool_call_id=tool_call_id, + tool_name=tool_name, + ) acp_content = [ContentToolCallContent(content=block) for block in converted] locations = [ ToolCallLocation(path=value) @@ -558,47 +640,94 @@ async def _replay_request(self, request: ModelRequest) -> None: if key in {"path", "file_path", "filepath"} and isinstance(value, str) ] title = generate_tool_title(tool_name, tool_input) - await self.tool_call_progress( - tool_call_id=tool_call_id, - title=title, - status="completed", - locations=locations or None, - content=acp_content or None, - raw_output=converted, + updates.append( + ToolCallProgress( + tool_call_id=tool_call_id, + status="completed", + title=title, + locations=locations or None, + content=acp_content or None, + raw_output=converted, + ) ) self._tool_call_inputs.pop(tool_call_id, None) case _: typ = type(part).__name__ self.log.debug("Unhandled request part type", part_type=typ) + return updates - async def _replay_response(self, response: ModelResponse) -> None: - """Replay a ModelResponse by converting it to appropriate ACP notifications.""" + def _collect_response_updates(self, response: ModelResponse) -> list[SessionUpdate]: + """Convert a ModelResponse to a list of SessionUpdate objects. + + This is a pure conversion — no I/O is performed. The + ``_tool_call_inputs`` cache is populated for ToolCallPart entries + for later use by ToolReturnPart conversion. + + Args: + response: The ModelResponse to convert. + + Returns: + List of SessionUpdate objects in order. + """ from pydantic_ai import TextPart, ThinkingPart, ToolCallPart + updates: list[SessionUpdate] = [] for part in response.parts: match part: case TextPart(content=content): - await self.send_agent_text(content) + updates.append(AgentMessageChunk.text(text=content)) case ThinkingPart(content=content): - await self.send_agent_thought(content) + updates.append(AgentThoughtChunk.text(text=content)) case ToolCallPart(tool_call_id=tool_call_id, tool_name=tool_name): - # Store tool call inputs for later use with ToolReturnPart tool_input = safe_args_as_dict(part) self._tool_call_inputs[tool_call_id] = tool_input - # Send tool_call_start so UI can track the tool call title = generate_tool_title(tool_name, tool_input) - await self.tool_call_start( - tool_call_id=tool_call_id, - title=title, - kind=infer_tool_kind(tool_name), - raw_input=tool_input, + updates.append( + ToolCallStart( + tool_call_id=tool_call_id, + status="pending", + title=title, + kind=infer_tool_kind(tool_name), + locations=None, + content=[], + raw_input=tool_input, + ) ) case _: typ = type(part).__name__ self.log.debug("Unhandled response part type", part_type=typ) + return updates + + async def _replay_request(self, request: ModelRequest) -> None: + """Replay a ModelRequest by sending collected updates sequentially. + + Thin wrapper around ``_collect_request_updates`` that sends each + update via ``send_update()``. Preserves legacy behavior for + non-batch callers. + + Args: + request: The ModelRequest to replay. + """ + updates = self._collect_request_updates(request) + for update in updates: + await self.send_update(update) + + async def _replay_response(self, response: ModelResponse) -> None: + """Replay a ModelResponse by sending collected updates sequentially. + + Thin wrapper around ``_collect_response_updates`` that sends each + update via ``send_update()``. Preserves legacy behavior for + non-batch callers. + + Args: + response: The ModelResponse to replay. + """ + updates = self._collect_response_updates(response) + for update in updates: + await self.send_update(update) async def send_agent_image( self, diff --git a/src/agentpool_server/acp_server/session.py b/src/agentpool_server/acp_server/session.py index a43cc1955..8ea42cb64 100644 --- a/src/agentpool_server/acp_server/session.py +++ b/src/agentpool_server/acp_server/session.py @@ -231,6 +231,10 @@ def __post_init__(self) -> None: # CRITICAL: Initialize requests and acp_env BEFORE agent mutation self.notifications = ACPNotifications(client=self.client, session_id=self.session_id) + # Detect batch session update support from client capabilities + if self.client_capabilities and self.client_capabilities.field_meta: + batch_meta = self.client_capabilities.field_meta.get("_batch_session_updates") + self.notifications.set_batch_support(bool(batch_meta)) self.requests = ACPRequests(client=self.client, session_id=self.session_id) self.input_provider = ACPInputProvider(self) self.acp_env = ACPExecutionEnvironment(fs=self.fs, requests=self.requests, cwd=self.cwd) diff --git a/tests/acp/test_notifications_replay.py b/tests/acp/test_notifications_replay.py index 47de383b0..b6b950b0a 100644 --- a/tests/acp/test_notifications_replay.py +++ b/tests/acp/test_notifications_replay.py @@ -35,6 +35,7 @@ def mock_client(): """Create a mock ACP client that captures sent notifications.""" client = AsyncMock() client.session_update = AsyncMock() + client.ext_notification = AsyncMock() return client @@ -44,6 +45,14 @@ def notifications(mock_client): return ACPNotifications(client=mock_client, session_id="test-session") +@pytest.fixture +def batch_notifications(mock_client): + """Create an ACPNotifications with batch support enabled.""" + n = ACPNotifications(client=mock_client, session_id="test-session") + n.set_batch_support(True) + return n + + @pytest.mark.unit async def test_replay_model_request_with_user_prompt_part(notifications, mock_client): """Test replay of ModelRequest with simple UserPromptPart (string content).""" @@ -349,3 +358,115 @@ async def test_replay_multiple_request_parts(notifications, mock_client): assert isinstance(updates[1], UserMessageChunk) assert isinstance(updates[1].content, TextContentBlock) assert updates[1].content.text == "Second message" + + +@pytest.mark.unit +async def test_replay_batch_mode_sends_ext_notification(batch_notifications, mock_client): + """Batch-capable client receives _batch_session_updates, not session/update.""" + messages = [ModelRequest(parts=[UserPromptPart(content="Hello")])] + + await batch_notifications.replay(messages) + + mock_client.session_update.assert_not_awaited() + mock_client.ext_notification.assert_awaited_once() + method, params = mock_client.ext_notification.call_args[0] + assert method == "_batch_session_updates" + assert params["session_id"] == "test-session" + assert len(params["updates"]) == 1 + + +@pytest.mark.unit +async def test_replay_fallback_mode_sends_session_update(notifications, mock_client): + """Non-capable client falls back to sequential session/update.""" + messages = [ModelRequest(parts=[UserPromptPart(content="Hello")])] + + await notifications.replay(messages) + + mock_client.ext_notification.assert_not_awaited() + assert mock_client.session_update.call_count == 1 + + +@pytest.mark.unit +async def test_replay_batch_preserves_tool_call_ordering(batch_notifications, mock_client): + """ToolCallStart must appear before ToolCallProgress in the batch.""" + messages = [ + ModelResponse( + parts=[ + TextPart(content="Let me help"), + ToolCallPart( + tool_call_id="tc-1", + tool_name="read_file", + args={"path": "/tmp/test.txt"}, + ), + ] + ), + ModelRequest( + parts=[ + ToolReturnPart( + tool_call_id="tc-1", + tool_name="read_file", + content="file contents", + ) + ] + ), + ] + + await batch_notifications.replay(messages) + + mock_client.ext_notification.assert_awaited_once() + _, params = mock_client.ext_notification.call_args[0] + updates = params["updates"] + assert len(updates) == 3 + assert updates[0]["sessionUpdate"] == "agent_message_chunk" + assert updates[1]["sessionUpdate"] == "tool_call" + assert updates[1]["toolCallId"] == "tc-1" + assert updates[2]["sessionUpdate"] == "tool_call_update" + assert updates[2]["toolCallId"] == "tc-1" + + +@pytest.mark.unit +async def test_replay_batch_custom_size(mock_client): + """Custom notification_batch_size produces correct chunk count.""" + n = ACPNotifications(client=mock_client, session_id="test", notification_batch_size=5) + n.set_batch_support(True) + messages = [ModelRequest(parts=[UserPromptPart(content=f"msg {i}")]) for i in range(12)] + + await n.replay(messages) + + assert mock_client.ext_notification.call_count == 3 + batch_sizes = [ + len(call[0][1]["updates"]) for call in mock_client.ext_notification.call_args_list + ] + assert batch_sizes == [5, 5, 2] + + +@pytest.mark.unit +async def test_collect_request_updates_is_pure(notifications, mock_client): + """_collect_request_updates returns list without calling client methods.""" + request = ModelRequest(parts=[UserPromptPart(content="Hello")]) + + updates = notifications._collect_request_updates(request) + + assert len(updates) == 1 + assert isinstance(updates[0], UserMessageChunk) + mock_client.session_update.assert_not_awaited() + mock_client.ext_notification.assert_not_awaited() + + +@pytest.mark.unit +async def test_replay_empty_messages_batch_mode(batch_notifications, mock_client): + """Empty messages list produces zero notifications in batch mode.""" + await batch_notifications.replay([]) + + mock_client.ext_notification.assert_not_awaited() + mock_client.session_update.assert_not_awaited() + + +@pytest.mark.unit +async def test_init_rejects_non_positive_batch_size(mock_client): + """notification_batch_size must be greater than 0.""" + with pytest.raises(ValueError, match="notification_batch_size must be greater than 0"): + ACPNotifications(mock_client, "session-1", notification_batch_size=0) + + with pytest.raises(ValueError, match="notification_batch_size must be greater than 0"): + ACPNotifications(mock_client, "session-1", notification_batch_size=-1) diff --git a/tests/agents/native_agent/test_get_agentlet_capabilities.py b/tests/agents/native_agent/test_get_agentlet_capabilities.py index d80360fed..61af9dcb1 100644 --- a/tests/agents/native_agent/test_get_agentlet_capabilities.py +++ b/tests/agents/native_agent/test_get_agentlet_capabilities.py @@ -627,7 +627,7 @@ async def test_capability_config_build_called(mock_agent: Agent[Any]) -> None: @pytest.mark.anyio -async def test_from_config_capabilities_not_duplicated() -> None: +async def test_from_config_capabilities_not_duplicated(monkeypatch: pytest.MonkeyPatch) -> None: """Capabilities built in from_config() must not be re-built in get_agentlet(). from_config() pre-builds capabilities from config.capabilities and stores @@ -639,6 +639,7 @@ async def test_from_config_capabilities_not_duplicated() -> None: This test calls from_config() with a config containing a capability, then calls get_agentlet() and verifies the capability appears exactly once. """ + monkeypatch.setenv("OPENAI_API_KEY", "test-key-for-ci") from llmling_models_config import TestModelConfig from pydantic_ai.capabilities import Instrumentation