fix: synchronize ACP telemetry and refresh remote final state

openhands-sdk/openhands/sdk/agent/acp_agent.py

@@ -65,28 +65,15 @@
     )
 
 
-# Seconds to wait after prompt() for pending session_update notifications
-# to be processed.  This is a best-effort workaround: the ACP protocol does
-# not currently signal when all notifications for a turn have been delivered,
-# so we yield to the event loop and then sleep briefly to allow in-flight
-# handlers to finish.  Override via ACP_NOTIFICATION_DRAIN_DELAY for slow or
-# remote servers.
-# TODO(https://github.com/agentclientprotocol/agent-client-protocol/issues/554):
-#       Replace with protocol-level synchronization once ACP supports a
-#       "turn complete" notification.
-_NOTIFICATION_DRAIN_DELAY: float = float(
-    os.environ.get("ACP_NOTIFICATION_DRAIN_DELAY", "0.1")
+# Maximum seconds to wait for a UsageUpdate notification after prompt()
+# returns. The ACP server writes UsageUpdate to the wire before the
+# PromptResponse, so under normal conditions the notification handler
+# completes almost immediately. This timeout is a safety net for slow
+# or remote servers.
+_USAGE_UPDATE_TIMEOUT: float = float(
+    os.environ.get("ACP_USAGE_UPDATE_TIMEOUT", "2.0")
 )
 
-# Retry configuration for transient ACP connection errors.
-# These errors can occur when the connection drops mid-conversation but the
-# session state is still valid on the server side.
-_ACP_PROMPT_MAX_RETRIES: int = int(os.environ.get("ACP_PROMPT_MAX_RETRIES", "3"))
-_ACP_PROMPT_RETRY_DELAYS: tuple[float, ...] = (5.0, 15.0, 30.0)  # seconds
-
-# Exception types that indicate transient connection issues worth retrying
-_RETRIABLE_CONNECTION_ERRORS = (OSError, ConnectionError, BrokenPipeError, EOFError)
-
 # Limit for asyncio.StreamReader buffers used by the ACP subprocess pipes.
 # The default (64 KiB) is too small for session_update notifications that
 # carry large tool-call outputs (e.g. file contents, test results).  When
@@ -97,17 +84,6 @@
 _STREAM_READER_LIMIT: int = 100 * 1024 * 1024  # 100 MiB
 
 
-async def _drain_notifications() -> None:
-    """Best-effort drain of pending ``session_update`` notifications.
-
-    ACP does not yet signal when all notifications for a turn have been
-    delivered (see TODO above).  We yield to the event loop so already-queued
-    handlers run, then sleep briefly to allow in-flight IO handlers to finish.
-    """
-    await asyncio.sleep(0)
-    await asyncio.sleep(_NOTIFICATION_DRAIN_DELAY)
-
-
 def _make_dummy_llm() -> LLM:
     """Create a dummy LLM that should never be called directly."""
     return LLM(model="acp-managed")
@@ -209,8 +185,12 @@ def __init__(self) -> None:
         self.on_token: Any = None  # ConversationTokenCallbackType | None
         # Telemetry state from UsageUpdate (persists across turns)
         self._last_cost: float = 0.0  # last cumulative cost seen
-        self._context_window: int = 0  # context window size from ACP
-        self._llm_ref: Any = None  # reference to the sentinel LLM
+        self._last_cost_by_session: dict[str, float] = {}
+        self._context_window: int = 0  # last context window seen
+        self._context_window_by_session: dict[str, int] = {}
+        # Per-turn synchronization for UsageUpdate notifications.
+        self._turn_usage_updates: dict[str, Any] = {}
+        self._usage_received: dict[str, asyncio.Event] = {}
         # Fork session state for ask_agent() — guarded by _fork_lock to
         # prevent concurrent ask_agent() calls from colliding.
         self._fork_lock = threading.Lock()
@@ -222,9 +202,27 @@ def reset(self) -> None:
         self.accumulated_thoughts.clear()
         self.accumulated_tool_calls.clear()
         self.on_token = None
+        self._turn_usage_updates.clear()
+        self._usage_received.clear()
         # Note: telemetry state (_last_cost, _context_window, etc.)
         # is intentionally NOT cleared — it accumulates across turns.
 
+    def prepare_usage_sync(self, session_id: str) -> asyncio.Event:
+        """Prepare per-turn UsageUpdate synchronization for a session."""
+        event = asyncio.Event()
+        self._usage_received[session_id] = event
+        self._turn_usage_updates.pop(session_id, None)
+        return event
+
+    def get_turn_usage_update(self, session_id: str) -> Any:
+        """Return the latest UsageUpdate observed for the current turn."""
+        return self._turn_usage_updates.get(session_id)
+
+    def pop_turn_usage_update(self, session_id: str) -> Any:
+        """Consume per-turn UsageUpdate synchronization state for a session."""
+        self._usage_received.pop(session_id, None)
+        return self._turn_usage_updates.pop(session_id, None)
+
     # -- Client protocol methods ------------------------------------------
 
     async def session_update(
@@ -255,14 +253,13 @@ async def session_update(
             if isinstance(update.content, TextContentBlock):
                 self.accumulated_thoughts.append(update.content.text)
         elif isinstance(update, UsageUpdate):
-            # Update context window size
+            # Store the update for step()/ask_agent() to process in one place.
             self._context_window = update.size
-            # Record incremental cost
-            if update.cost is not None and self._llm_ref is not None:
-                delta = update.cost.amount - self._last_cost
-                if delta > 0:
-                    self._llm_ref.metrics.add_cost(delta)
-                self._last_cost = update.cost.amount
+            self._context_window_by_session[session_id] = update.size
+            self._turn_usage_updates[session_id] = update
+            event = self._usage_received.get(session_id)
+            if event is not None:
+                event.set()
         elif isinstance(update, ToolCallStart):
             self.accumulated_tool_calls.append(
                 {
@@ -452,14 +449,24 @@ def _record_usage(
         response: PromptResponse | None,
         session_id: str,
         elapsed: float | None = None,
+        usage_update: UsageUpdate | None = None,
     ) -> None:
-        """Record token usage, latency, and notify stats callback from a PromptResponse.
+        """Record cost, token usage, latency, and notify stats callback once.
 
         Args:
             response: The ACP PromptResponse (may carry a ``usage`` field).
             session_id: Session identifier used as the response_id for metrics.
             elapsed: Wall-clock seconds for this prompt round-trip (optional).
+            usage_update: The synchronized ACP UsageUpdate for this turn, if any.
         """
+        if usage_update is not None and usage_update.cost is not None:
+            last_cost = self._client._last_cost_by_session.get(session_id, 0.0)
+            delta = usage_update.cost.amount - last_cost
+            if delta > 0:
+                self.llm.metrics.add_cost(delta)
+            self._client._last_cost_by_session[session_id] = usage_update.cost.amount
+            self._client._last_cost = usage_update.cost.amount
+
         if response is not None and response.usage is not None:
             usage = response.usage
             self.llm.metrics.add_token_usage(
@@ -468,7 +475,9 @@ def _record_usage(
                 cache_read_tokens=usage.cached_read_tokens or 0,
                 cache_write_tokens=usage.cached_write_tokens or 0,
                 reasoning_tokens=usage.thought_tokens or 0,
-                context_window=self._client._context_window,
+                context_window=self._client._context_window_by_session.get(
+                    session_id, self._client._context_window
+                ),
                 response_id=session_id,
             )
 
@@ -562,7 +571,6 @@ def init_state(
     def _start_acp_server(self, state: ConversationState) -> None:
         """Start the ACP subprocess and initialize the session."""
         client = _OpenHandsACPBridge()
-        client._llm_ref = self.llm
         self._client = client
 
         # Build environment: inherit current env + ACP extras
@@ -706,59 +714,44 @@ def step(
         try:
 
             async def _prompt() -> PromptResponse:
+                usage_sync = self._client.prepare_usage_sync(self._session_id or "")
                 response = await self._conn.prompt(
                     [text_block(user_message)],
                     self._session_id,
                 )
-                await _drain_notifications()
+                if self._client.get_turn_usage_update(self._session_id or "") is None:
+                    try:
+                        await asyncio.wait_for(
+                            usage_sync.wait(), timeout=_USAGE_UPDATE_TIMEOUT
+                        )
+                    except asyncio.TimeoutError:
+                        logger.warning(
+                            "UsageUpdate not received within %.1fs for session %s",
+                            _USAGE_UPDATE_TIMEOUT,
+                            self._session_id,
+                        )
                 return response
 
-            # Send prompt to ACP server with retry logic for connection errors.
-            # Transient connection failures (network blips, server restarts) are
-            # retried to preserve session state and avoid losing progress.
+            # Send prompt to ACP server (with timeout to prevent indefinite hangs)
             logger.info(
                 "Sending ACP prompt (timeout=%.0fs, msg=%d chars)",
                 self.acp_prompt_timeout,
                 len(user_message),
             )
-
-            response: PromptResponse | None = None
-            max_retries = _ACP_PROMPT_MAX_RETRIES
-
-            for attempt in range(max_retries + 1):
-                try:
-                    response = self._executor.run_async(
-                        _prompt, timeout=self.acp_prompt_timeout
-                    )
-                    break  # Success, exit retry loop
-                except TimeoutError:
-                    # Timeout is handled separately below, don't retry
-                    raise
-                except _RETRIABLE_CONNECTION_ERRORS as e:
-                    if attempt < max_retries:
-                        delay = _ACP_PROMPT_RETRY_DELAYS[
-                            min(attempt, len(_ACP_PROMPT_RETRY_DELAYS) - 1)
-                        ]
-                        logger.warning(
-                            "ACP prompt failed with retriable error (attempt %d/%d), "
-                            "retrying in %.0fs: %s",
-                            attempt + 1,
-                            max_retries + 1,
-                            delay,
-                            e,
-                        )
-                        time.sleep(delay)
-                        # Reset accumulators for retry (partial state may be stale)
-                        self._client.reset()
-                        self._client.on_token = on_token
-                    else:
-                        # Max retries exceeded
-                        raise
-
+            response = self._executor.run_async(
+                _prompt, timeout=self.acp_prompt_timeout
+            )
             elapsed = time.monotonic() - t0
             logger.info("ACP prompt returned in %.1fs", elapsed)
 
-            self._record_usage(response, self._session_id or "", elapsed=elapsed)
+            session_id = self._session_id or ""
+            usage_update = self._client.pop_turn_usage_update(session_id)
+            self._record_usage(
+                response,
+                session_id,
+                elapsed=elapsed,
+                usage_update=usage_update,
+            )
 
             # Emit ACPToolCallEvents for each accumulated tool call
             for tc in self._client.accumulated_tool_calls:
@@ -906,15 +899,32 @@ async def _fork_and_prompt() -> str:
             client._fork_accumulated_text.clear()
             try:
                 fork_t0 = time.monotonic()
+                usage_sync = client.prepare_usage_sync(fork_session_id)
                 response = await self._conn.prompt(
                     [text_block(question)],
                     fork_session_id,
                 )
-                await _drain_notifications()
+                if client.get_turn_usage_update(fork_session_id) is None:
+                    try:
+                        await asyncio.wait_for(
+                            usage_sync.wait(), timeout=_USAGE_UPDATE_TIMEOUT
+                        )
+                    except asyncio.TimeoutError:
+                        logger.warning(
+                            "UsageUpdate not received within %.1fs for fork session %s",
+                            _USAGE_UPDATE_TIMEOUT,
+                            fork_session_id,
+                        )
                 fork_elapsed = time.monotonic() - fork_t0
 
                 result = "".join(client._fork_accumulated_text)
-                self._record_usage(response, fork_session_id, elapsed=fork_elapsed)
+                usage_update = client.pop_turn_usage_update(fork_session_id)
+                self._record_usage(
+                    response,
+                    fork_session_id,
+                    elapsed=fork_elapsed,
+                    usage_update=usage_update,
+                )
                 return result
             finally:
                 client._fork_session_id = None