From 4b9cac1400c8292d6b0f3310e5355222961f79f7 Mon Sep 17 00:00:00 2001 From: "qiang.zeng" <285093074@qq.com> Date: Tue, 16 Jun 2026 09:23:04 +0800 Subject: [PATCH 1/5] fix(acp): increase default read timeout for MCP transport and session handling --- src/agentpool_server/acp_server/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/agentpool_server/acp_server/session.py b/src/agentpool_server/acp_server/session.py index 6836b24d3..ad00b5410 100644 --- a/src/agentpool_server/acp_server/session.py +++ b/src/agentpool_server/acp_server/session.py @@ -431,7 +431,7 @@ async def _init_server(server: Any) -> None: ) transport = AcpMcpTransport( - conn, timeout=getattr(server, "timeout", None) or 300.0 + conn, timeout=getattr(server, "timeout", None) or 600.0 ) cfg = convert_acp_mcp_server_to_config(server) provider = MCPResourceProvider( From 3991f2a122073a4d4073aaa6d0b7acfd6da095f0 Mon Sep 17 00:00:00 2001 From: "qiang.zeng" <285093074@qq.com> Date: Tue, 16 Jun 2026 09:39:29 +0800 Subject: [PATCH 2/5] fix(acp): increase timeout for mcp/connect and mcp/message to 300 seconds --- src/agentpool_server/acp_server/acp_agent.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/agentpool_server/acp_server/acp_agent.py b/src/agentpool_server/acp_server/acp_agent.py index 52107b608..7aaa0859c 100644 --- a/src/agentpool_server/acp_server/acp_agent.py +++ b/src/agentpool_server/acp_server/acp_agent.py @@ -760,13 +760,13 @@ async def connect_acp_mcp_server(self, server: AcpMcpServer) -> str: Raises: ValueError: If the client does not return a connectionId. - TimeoutError: If the client does not respond to mcp/connect within 10s. + TimeoutError: If the client does not respond to mcp/connect within 300s. """ params = { "server": server.model_dump(by_alias=True, exclude_none=True), "acpId": server.id, } - with anyio.fail_after(10): + with anyio.fail_after(300): response = await self.client.send_request("mcp/connect", params) connection_id = str(response.get("connectionId", "")) if not connection_id: @@ -776,7 +776,7 @@ async def connect_acp_mcp_server(self, server: AcpMcpServer) -> str: async def send_to_client(message: dict[str, Any]) -> Any: # message is already wrapped as {"connectionId": conn_id, "message": mcp_msg} # by AcpMcpConnection.send_to_client. Pass through directly. - with anyio.fail_after(30): + with anyio.fail_after(300): return await self.client.send_request("mcp/message", message) await self._mcp_manager.create_connection( From 591e71019f9a9e9cf1efa71664aebf35ccc97dd8 Mon Sep 17 00:00:00 2001 From: "qiang.zeng" <285093074@qq.com> Date: Wed, 17 Jun 2026 12:03:42 +0800 Subject: [PATCH 3/5] feat: Add WebSocket heartbeat configuration support - Add heartbeat parameters (ping_interval, pong_timeout, max_missed_pongs) to transport classes - Implement _websocket_heartbeat function for proactive connection monitoring - Update CLI with --ws-ping-interval, --ws-pong-timeout, --ws-max-missed-pongs flags - Add ACPPoolServerConfig heartbeat fields for YAML configuration - Integrate heartbeat logic into _serve_websocket and _serve_streamable_http - Support disabling heartbeat via ping_interval=None Defaults: ping_interval=60s, pong_timeout=30s, max_missed_pongs=3 --- src/acp/transports.py | 185 ++++++++++++++++++++++++++-- src/agentpool_cli/serve_acp.py | 39 +++++- src/agentpool_config/pool_server.py | 24 ++++ 3 files changed, 237 insertions(+), 11 deletions(-) diff --git a/src/acp/transports.py b/src/acp/transports.py index cdb292533..33073867f 100644 --- a/src/acp/transports.py +++ b/src/acp/transports.py @@ -13,7 +13,7 @@ import logging import os import subprocess -from typing import TYPE_CHECKING, Any, Literal, assert_never +from typing import TYPE_CHECKING, Any, Literal, Protocol, assert_never import uuid import anyio @@ -21,7 +21,7 @@ if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Mapping + from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from pathlib import Path from anyio.abc import Process @@ -32,6 +32,10 @@ logger = logging.getLogger(__name__) +DEFAULT_WEBSOCKET_PING_INTERVAL = 60.0 +DEFAULT_WEBSOCKET_PONG_TIMEOUT = 30.0 +DEFAULT_WEBSOCKET_MAX_MISSED_PONGS = 3 + # ============================================================================= # Transport Configuration Classes @@ -57,10 +61,49 @@ class WebSocketTransport: Attributes: host: Host to bind the WebSocket server to. port: Port for the WebSocket server. + ping_interval: Seconds between server-initiated WebSocket ping frames. + Set to None to disable AgentPool's heartbeat. + pong_timeout: Seconds to wait for a pong response to each ping. + max_missed_pongs: Consecutive missed pongs before closing the connection. """ host: str = "localhost" port: int = 8765 + ping_interval: float | None = DEFAULT_WEBSOCKET_PING_INTERVAL + pong_timeout: float = DEFAULT_WEBSOCKET_PONG_TIMEOUT + max_missed_pongs: int = DEFAULT_WEBSOCKET_MAX_MISSED_PONGS + + def __post_init__(self) -> None: + if self.ping_interval is not None and self.ping_interval <= 0: + msg = "ping_interval must be positive or None" + raise ValueError(msg) + _validate_websocket_heartbeat(self.pong_timeout, self.max_missed_pongs) + + +class _HeartbeatWebSocket(Protocol): + async def ping(self) -> Awaitable[float]: ... + + async def close(self, code: int = 1000, reason: str = "") -> None: ... + + +def _validate_websocket_heartbeat(pong_timeout: float, max_missed_pongs: int) -> None: + if pong_timeout <= 0: + msg = "pong_timeout must be positive" + raise ValueError(msg) + if max_missed_pongs <= 0: + msg = "max_missed_pongs must be positive" + raise ValueError(msg) + + +def _effective_websocket_pong_timeout( + ping_interval: float | None, + pong_timeout: float, + max_missed_pongs: int, +) -> float | None: + """Return a single keepalive timeout matching the multi-miss tolerance window.""" + if ping_interval is None: + return None + return pong_timeout * max_missed_pongs + ping_interval * (max_missed_pongs - 1) @dataclass @@ -89,10 +132,23 @@ class ACPWebSocketTransport: Attributes: host: Host to bind the WebSocket server to. port: Port for the WebSocket server. + ping_interval: Seconds between server-initiated WebSocket ping frames. + Set to None to disable server-side heartbeat. + pong_timeout: Seconds to wait for each expected pong. + max_missed_pongs: Consecutive missed pongs to tolerate before disconnecting. """ host: str = "localhost" port: int = 8080 + ping_interval: float | None = DEFAULT_WEBSOCKET_PING_INTERVAL + pong_timeout: float = DEFAULT_WEBSOCKET_PONG_TIMEOUT + max_missed_pongs: int = DEFAULT_WEBSOCKET_MAX_MISSED_PONGS + + def __post_init__(self) -> None: + if self.ping_interval is not None and self.ping_interval <= 0: + msg = "ping_interval must be positive or None" + raise ValueError(msg) + _validate_websocket_heartbeat(self.pong_timeout, self.max_missed_pongs) # Type alias for all supported transports @@ -162,10 +218,42 @@ async def serve( match transport: case StdioTransport(): await _serve_stdio(agent, shutdown_event, debug_file, **kwargs) - case WebSocketTransport(host=host, port=port): - await _serve_websocket(agent, host, port, shutdown_event, debug_file, **kwargs) - case ACPWebSocketTransport(host=host, port=port): - await _serve_streamable_http(agent, host, port, shutdown_event, debug_file, **kwargs) + case WebSocketTransport( + host=host, + port=port, + ping_interval=ping_interval, + pong_timeout=pong_timeout, + max_missed_pongs=max_missed_pongs, + ): + await _serve_websocket( + agent, + host, + port, + shutdown_event, + debug_file, + ping_interval=ping_interval, + pong_timeout=pong_timeout, + max_missed_pongs=max_missed_pongs, + **kwargs, + ) + case ACPWebSocketTransport( + host=host, + port=port, + ping_interval=ping_interval, + pong_timeout=pong_timeout, + max_missed_pongs=max_missed_pongs, + ): + await _serve_streamable_http( + agent, + host, + port, + shutdown_event, + debug_file, + ping_interval=ping_interval, + pong_timeout=pong_timeout, + max_missed_pongs=max_missed_pongs, + **kwargs, + ) case StreamTransport(reader=reader, writer=writer): await _serve_streams(agent, reader, writer, shutdown_event, debug_file, **kwargs) case _ as unreachable: @@ -227,6 +315,10 @@ async def _serve_websocket( port: int, shutdown_event: asyncio.Event | None, debug_file: str | None, + *, + ping_interval: float | None = DEFAULT_WEBSOCKET_PING_INTERVAL, + pong_timeout: float = DEFAULT_WEBSOCKET_PONG_TIMEOUT, + max_missed_pongs: int = DEFAULT_WEBSOCKET_MAX_MISSED_PONGS, **kwargs: Any, ) -> None: """Run agent as WebSocket server.""" @@ -251,6 +343,17 @@ async def handle_client(websocket: ServerConnection) -> None: ) connections.append(conn) + heartbeat_task: asyncio.Task[None] | None = None + if ping_interval is not None: + heartbeat_task = asyncio.create_task( + _websocket_heartbeat( + websocket, + ping_interval=ping_interval, + pong_timeout=pong_timeout, + max_missed_pongs=max_missed_pongs, + ) + ) + try: # Wait for shutdown or for the receive loop to end (client disconnect) _recv_conn = getattr(conn, "_conn", None) @@ -258,6 +361,8 @@ async def handle_client(websocket: ServerConnection) -> None: waitables: list[asyncio.Future[Any]] = [asyncio.create_task(shutdown.wait())] if isinstance(recv_task, asyncio.Task): waitables.append(recv_task) + if heartbeat_task is not None: + waitables.append(heartbeat_task) done, _ = await asyncio.wait(waitables, return_when=asyncio.FIRST_COMPLETED) @@ -270,11 +375,15 @@ async def handle_client(websocket: ServerConnection) -> None: except websockets.exceptions.ConnectionClosed: logger.info("WebSocket client disconnected") finally: + if heartbeat_task is not None and not heartbeat_task.done(): + heartbeat_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat_task connections.remove(conn) await conn.close() logger.info("Starting WebSocket server on ws://%s:%d", host, port) - async with websockets.serve(handle_client, host, port): + async with websockets.serve(handle_client, host, port, ping_interval=None): logger.info("WebSocket server running on ws://%s:%d", host, port) await shutdown.wait() @@ -283,12 +392,61 @@ async def handle_client(websocket: ServerConnection) -> None: await conn.close() +async def _websocket_heartbeat( + websocket: _HeartbeatWebSocket, + *, + ping_interval: float, + pong_timeout: float, + max_missed_pongs: int, +) -> None: + """Close a WebSocket only after several consecutive missed pong responses.""" + import websockets + + missed_pongs = 0 + while True: + await asyncio.sleep(ping_interval) + try: + pong_waiter: Awaitable[float] = await websocket.ping() + await asyncio.wait_for(pong_waiter, timeout=pong_timeout) + if missed_pongs: + logger.info( + "WebSocket heartbeat recovered after %d missed pong(s)", + missed_pongs, + ) + missed_pongs = 0 + except TimeoutError: + missed_pongs += 1 + logger.warning( + "WebSocket heartbeat missed pong %d/%d", + missed_pongs, + max_missed_pongs, + ) + if missed_pongs >= max_missed_pongs: + logger.warning( + "Closing WebSocket after %d consecutive missed pong(s)", + missed_pongs, + ) + await websocket.close(code=1011, reason="pong timeout") + return + except websockets.exceptions.ConnectionClosed: + return + except Exception: + logger.exception("WebSocket heartbeat failed") + with contextlib.suppress(Exception): + await websocket.close(code=1011, reason="heartbeat failed") + return + + async def _serve_streamable_http( agent: Agent | Callable[[AgentSideConnection], Agent], host: str, port: int, shutdown_event: asyncio.Event | None, debug_file: str | None, + *, + ping_interval: float | None = DEFAULT_WEBSOCKET_PING_INTERVAL, + pong_timeout: float = DEFAULT_WEBSOCKET_PONG_TIMEOUT, + max_missed_pongs: int = DEFAULT_WEBSOCKET_MAX_MISSED_PONGS, **kwargs: Any, ) -> None: """Run agent as a streamable HTTP WebSocket server (Starlette-based).""" @@ -339,7 +497,18 @@ async def handle_acp(websocket: Any) -> None: await conn.close() app = Starlette(routes=[WebSocketRoute("/acp", handle_acp)]) - config = uvicorn.Config(app, host=host, port=port, log_level="warning") + config = uvicorn.Config( + app, + host=host, + port=port, + log_level="warning", + ws_ping_interval=ping_interval, + ws_ping_timeout=_effective_websocket_pong_timeout( + ping_interval, + pong_timeout, + max_missed_pongs, + ), + ) server = uvicorn.Server(config) async def shutdown_watcher() -> None: diff --git a/src/agentpool_cli/serve_acp.py b/src/agentpool_cli/serve_acp.py index 7872e7dae..533952257 100644 --- a/src/agentpool_cli/serve_acp.py +++ b/src/agentpool_cli/serve_acp.py @@ -82,7 +82,7 @@ def acp_command( # noqa: PLR0915 t.Option( "--skills/--no-skills", help="Load client-side skills from .claude/skills directory. " - "Defaults to the manifest's skills.include_default setting.", + "Defaults to the manifest's skills.include_default setting.", ), ] = None, transport: Annotated[ @@ -123,6 +123,27 @@ def acp_command( # noqa: PLR0915 help="WebSocket port (only used with --transport websocket, deprecated)", ), ] = 8765, + ws_ping_interval: Annotated[ + float, + t.Option( + "--ws-ping-interval", + help="Seconds between WebSocket ping frames for WebSocket transports", + ), + ] = 60.0, + ws_pong_timeout: Annotated[ + float, + t.Option( + "--ws-pong-timeout", + help="Seconds to wait for each WebSocket pong", + ), + ] = 30.0, + ws_max_missed_pongs: Annotated[ + int, + t.Option( + "--ws-max-missed-pongs", + help="Consecutive missed WebSocket pongs before disconnecting", + ), + ] = 3, mcp_config: Annotated[ str | None, t.Option( @@ -177,14 +198,26 @@ def acp_command( # noqa: PLR0915 # Build transport config if transport == "streamable-http": - transport_config: Transport = ACPWebSocketTransport(host=host, port=port) + transport_config: Transport = ACPWebSocketTransport( + host=host, + port=port, + ping_interval=ws_ping_interval, + pong_timeout=ws_pong_timeout, + max_missed_pongs=ws_max_missed_pongs, + ) elif transport == "websocket": warnings.warn( "--transport websocket is deprecated; use --transport streamable-http instead", DeprecationWarning, stacklevel=2, ) - transport_config = WebSocketTransport(host=ws_host, port=ws_port) + transport_config = WebSocketTransport( + host=ws_host, + port=ws_port, + ping_interval=ws_ping_interval, + pong_timeout=ws_pong_timeout, + max_missed_pongs=ws_max_missed_pongs, + ) elif transport == "stdio": transport_config = StdioTransport() diff --git a/src/agentpool_config/pool_server.py b/src/agentpool_config/pool_server.py index a0623ca2c..b7ac02d62 100644 --- a/src/agentpool_config/pool_server.py +++ b/src/agentpool_config/pool_server.py @@ -202,6 +202,30 @@ class ACPPoolServerConfig(BasePoolServerConfig): ) """Port to listen on (streamable-http only).""" + websocket_ping_interval: float | None = Field( + default=60.0, + gt=0, + title="WebSocket ping interval", + ) + """Seconds between server-initiated WebSocket ping frames. + + Set to null to disable server-side WebSocket heartbeat. + """ + + websocket_pong_timeout: float = Field( + default=30.0, + gt=0, + title="WebSocket pong timeout", + ) + """Seconds to wait for each expected WebSocket pong.""" + + websocket_max_missed_pongs: int = Field( + default=3, + gt=0, + title="WebSocket max missed pongs", + ) + """Consecutive missed WebSocket pongs to tolerate before disconnecting.""" + model_config = ConfigDict(frozen=True) From 8e00feb94310cfdd8523062a5e2878a15540107c Mon Sep 17 00:00:00 2001 From: "qiang.zeng" <285093074@qq.com> Date: Wed, 17 Jun 2026 15:10:38 +0800 Subject: [PATCH 4/5] feat: Enhance WebSocket heartbeat logging with detailed info and debug messages --- src/acp/transports.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/acp/transports.py b/src/acp/transports.py index 33073867f..34376ac80 100644 --- a/src/acp/transports.py +++ b/src/acp/transports.py @@ -345,6 +345,12 @@ async def handle_client(websocket: ServerConnection) -> None: heartbeat_task: asyncio.Task[None] | None = None if ping_interval is not None: + logger.info( + "Starting WebSocket heartbeat with interval=%s, timeout=%s, max_missed=%s", + ping_interval, + pong_timeout, + max_missed_pongs, + ) heartbeat_task = asyncio.create_task( _websocket_heartbeat( websocket, @@ -403,11 +409,16 @@ async def _websocket_heartbeat( import websockets missed_pongs = 0 + ping_count = 0 + logger.info("WebSocket heartbeat started") while True: await asyncio.sleep(ping_interval) + ping_count += 1 + logger.debug("Sending WebSocket ping #%d", ping_count) try: pong_waiter: Awaitable[float] = await websocket.ping() await asyncio.wait_for(pong_waiter, timeout=pong_timeout) + logger.debug("Pong #%d received in time", ping_count) if missed_pongs: logger.info( "WebSocket heartbeat recovered after %d missed pong(s)", From 8d22ff5ce891a7ecd3f6cf0b2219e49e9745a7ce Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Mon, 22 Jun 2026 10:56:17 +0800 Subject: [PATCH 5/5] =?UTF-8?q?fix:=20address=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20wrap=20finally=20cleanup=20in=20try-except,=20reduc?= =?UTF-8?q?e=20MCP=20transport=20fallback=20timeout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- src/acp/transports.py | 16 +++++++++++++--- src/agentpool_server/acp_server/session.py | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/acp/transports.py b/src/acp/transports.py index 34376ac80..c150c25f6 100644 --- a/src/acp/transports.py +++ b/src/acp/transports.py @@ -383,10 +383,20 @@ async def handle_client(websocket: ServerConnection) -> None: finally: if heartbeat_task is not None and not heartbeat_task.done(): heartbeat_task.cancel() - with contextlib.suppress(asyncio.CancelledError): + try: await heartbeat_task - connections.remove(conn) - await conn.close() + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Unexpected error during heartbeat task cleanup") + try: + connections.remove(conn) + except ValueError: + pass + try: + await conn.close() + except Exception: + logger.exception("Unexpected error closing WebSocket connection") logger.info("Starting WebSocket server on ws://%s:%d", host, port) async with websockets.serve(handle_client, host, port, ping_interval=None): diff --git a/src/agentpool_server/acp_server/session.py b/src/agentpool_server/acp_server/session.py index ad00b5410..42fe5d686 100644 --- a/src/agentpool_server/acp_server/session.py +++ b/src/agentpool_server/acp_server/session.py @@ -431,7 +431,7 @@ async def _init_server(server: Any) -> None: ) transport = AcpMcpTransport( - conn, timeout=getattr(server, "timeout", None) or 600.0 + conn, timeout=getattr(server, "timeout", None) or 120.0 ) cfg = convert_acp_mcp_server_to_config(server) provider = MCPResourceProvider(