feat(acp): batch SessionUpdate delivery during replay

openspec/changes/acp-notification-batching/.openspec.yaml

@@ -0,0 +1,2 @@
+schema: spec-driven
+created: 2026-06-30

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) |

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

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": "<id>", "updates": [<update1>, ...]})`
+- **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

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)

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())