feat(sdk): StreamContext mints the stream identity and closes every stream

openhands-agent-server/openhands/agent_server/event_service.py

@@ -31,6 +31,7 @@
     CODEX_AUTH_SECRET_NAME,
     is_valid_codex_auth,
 )
+from openhands.sdk.agent.stream_context import StreamProgress
 from openhands.sdk.conversation.base import BaseConversation
 from openhands.sdk.conversation.events_list_base import EventsListBase
 from openhands.sdk.conversation.exceptions import ConversationRunError
@@ -135,6 +136,11 @@ class EventService:
     _pub_sub: PubSub[Event] = field(
         default_factory=lambda: PubSub[Event](max_subscribers=50), init=False
     )
+    # Its own fan-out, not the event bus: frames are not events, and only the
+    # session socket consumes them.
+    _stream_pub_sub: PubSub[StreamProgress] = field(
+        default_factory=lambda: PubSub[StreamProgress](max_subscribers=50), init=False
+    )
     _run_task: asyncio.Task | None = field(default=None, init=False)
     # Set when a send_message(run=True) is rejected because a run is still
     # wrapping up; consumed by _run_and_publish to re-run the stranded message.
@@ -848,6 +854,19 @@ async def subscribe_to_events(self, subscriber: Subscriber[Event]) -> UUID:
     async def unsubscribe_from_events(self, subscriber_id: UUID) -> bool:
         return self._pub_sub.unsubscribe(subscriber_id)
 
+    async def subscribe_to_stream_progress(
+        self, subscriber: Subscriber[StreamProgress]
+    ) -> UUID:
+        """Register for stream-progress frames.
+
+        No initial push, unlike :meth:`subscribe_to_events`: a client that
+        connects mid-stream gets the real text with the durable event.
+        """
+        return self._stream_pub_sub.subscribe(subscriber)
+
+    async def unsubscribe_from_stream_progress(self, subscriber_id: UUID) -> bool:
+        return self._stream_pub_sub.unsubscribe(subscriber_id)
+
     def _emit_event_from_thread(self, event: Event) -> None:
         """Helper to safely emit events from non-async contexts (e.g., callbacks).
 
@@ -1075,6 +1094,16 @@ def _publish_stream_delta(
             with suppress(RuntimeError):  # main loop already closed during teardown
                 asyncio.run_coroutine_threadsafe(self._pub_sub(event), self._main_loop)
 
+        def _publish_stream_progress(frame: StreamProgress) -> None:
+            # Same cross-thread hop as _publish_stream_delta: called from the
+            # run thread, or the ACP portal thread.
+            if not self._main_loop or not self._main_loop.is_running():
+                return
+            with suppress(RuntimeError):  # main loop already closed during teardown
+                asyncio.run_coroutine_threadsafe(
+                    self._stream_pub_sub(frame), self._main_loop
+                )
+
         def _token_streaming_callback(chunk: LLMStreamChunk | str) -> None:
             if isinstance(chunk, str):
                 _publish_stream_delta(content=chunk)
@@ -1099,6 +1128,7 @@ def _token_streaming_callback(chunk: LLMStreamChunk | str) -> None:
             conversation_id=self.stored.id,
             callbacks=[self._callback_wrapper],
             token_callbacks=([_token_streaming_callback] if streaming_enabled else []),
+            stream_callbacks=[_publish_stream_progress],
             max_iteration_per_run=self.stored.max_iterations,
             stuck_detection=self.stored.stuck_detection,
             visualizer=None,
@@ -1774,6 +1804,7 @@ async def close(self):
             self._run_task = None
 
         await self._pub_sub.close()
+        await self._stream_pub_sub.close()
         if self._conversation:
             loop = asyncio.get_running_loop()
             await loop.run_in_executor(None, self._conversation.close)

openhands-agent-server/openhands/agent_server/session_protocol.py

@@ -118,6 +118,16 @@ class DeltaFrame(SessionFrameBase):
     order: int
     kind: Literal["text", "reasoning"] = "text"
     content: str
+    chunk_id: str | None = Field(
+        default=None,
+        description=(
+            "The provider's completion id for this chunk. Corroboration only: "
+            "litellm mints a new one per retry attempt."
+        ),
+    )
+    choice_index: int | None = Field(
+        default=None, description="The provider's choice index for this chunk."
+    )
 
 
 class ItemAbortedFrame(SessionFrameBase):

openhands-agent-server/openhands/agent_server/session_socket.py

@@ -5,11 +5,6 @@
 and live traffic cannot interleave; and a slow consumer cannot wedge the
 publisher, because admission is byte-bounded and only this connection's writer
 task awaits the socket.
-
-``ItemStarted`` / ``Delta`` / ``ItemAborted`` are carried but never produced
-yet — that needs ``StreamContext`` (#4682). Until then this is a durable-only
-channel and ``StreamingDeltaEvent`` is dropped rather than forwarded, since
-putting it back on the wire would restore the coupling this endpoint removes.
 """
 
 import asyncio
@@ -30,8 +25,11 @@
 from openhands.agent_server.session_protocol import (
     MAX_FRAME_BYTES,
     MAX_PENDING_BYTES,
+    DeltaFrame,
     DurableFrame,
     ErrorFrame,
+    ItemAbortedFrame,
+    ItemStartedFrame,
     SessionFrameBase,
     SyncFrame,
     TransientFrame,
@@ -44,6 +42,12 @@
     _safe_close_websocket,
 )
 from openhands.sdk import Event, Message
+from openhands.sdk.agent.stream_context import (
+    StreamAborted,
+    StreamDelta,
+    StreamProgress,
+    StreamStarted,
+)
 from openhands.sdk.conversation.event_store import EventLog
 from openhands.sdk.event import StreamingDeltaEvent
 from openhands.sdk.event.conversation_state import ConversationStateUpdateEvent
@@ -181,8 +185,8 @@ class _SessionSubscriber(Subscriber[Event]):
 
     async def __call__(self, event: Event) -> None:
         if isinstance(event, StreamingDeltaEvent):
-            # Deltas never ride the durable channel; progress frames will come
-            # from StreamContext instead.
+            # Deltas never ride the durable channel; progress frames come from
+            # StreamContext, over their own fan-out.
             return
         if self._buffer is not None:
             self._buffer.append(event)
@@ -229,6 +233,47 @@ def go_live(self, through_seq: int | None) -> None:
             self._emit(event)
 
 
+def _to_wire(frame: StreamProgress) -> SessionFrameBase:
+    """Map one SDK progress frame onto its envelope."""
+    match frame:
+        case StreamStarted():
+            return ItemStartedFrame(
+                item_id=frame.item_id,
+                attempt=frame.attempt,
+                anchor_seq=frame.anchor_seq,
+            )
+        case StreamDelta():
+            return DeltaFrame(
+                item_id=frame.item_id,
+                attempt=frame.attempt,
+                order=frame.order,
+                kind=frame.kind,
+                content=frame.content,
+                chunk_id=frame.chunk_id,
+                choice_index=frame.choice_index,
+            )
+        case StreamAborted():
+            return ItemAbortedFrame(
+                item_id=frame.item_id,
+                attempt=frame.attempt,
+                reason=frame.reason,
+            )
+
+
+class _ProgressSubscriber(Subscriber[StreamProgress]):
+    """Forwards stream progress straight to the writer.
+
+    No buffering and no replay boundary: a dropped frame costs a repaint, and
+    the real text arrives with the durable event.
+    """
+
+    def __init__(self, writer: _ConnectionWriter) -> None:
+        self._writer = writer
+
+    async def __call__(self, frame: StreamProgress) -> None:
+        self._writer.send(_to_wire(frame))
+
+
 def _read_page(events: EventLog, start: int, stop: int) -> list[tuple[int, Event]]:
     """Read one page by index, skipping unreadable events.
 
@@ -340,7 +385,17 @@ async def session_socket(
         await _safe_close_websocket(websocket, code=1011, reason="subscribe_failed")
         return
 
+    progress_id: UUID | None = None
     try:
+        try:
+            progress_id = await event_service.subscribe_to_stream_progress(
+                _ProgressSubscriber(writer)
+            )
+        except MaxSubscribersError:
+            # Durable delivery is what the client cannot recover on its own;
+            # losing progress only costs it the live typing effect.
+            logger.warning("session_socket_progress_limit: %s", conversation_id)
+
         length = len(events)
         through_seq = length - 1 if length else None
 
@@ -361,6 +416,8 @@ async def session_socket(
 
         await _inbound_loop(conversation_id, websocket, event_service, writer)
     finally:
+        if progress_id is not None:
+            await event_service.unsubscribe_from_stream_progress(progress_id)
         await event_service.unsubscribe_from_events(subscriber_id)
         await writer.aclose()
         if writer.drop_reason in ("slow_consumer", "frame_too_large"):

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

@@ -82,6 +82,7 @@
 from openhands.sdk.agent.acp_models import ACPModelInfo
 from openhands.sdk.agent.acp_tracing import ACPTurnTrace
 from openhands.sdk.agent.base import AgentBase
+from openhands.sdk.agent.stream_context import StreamContext
 from openhands.sdk.context import AgentContext
 from openhands.sdk.conversation.state import ConversationExecutionStatus
 from openhands.sdk.credential import (
@@ -1926,6 +1927,10 @@ def model_post_init(self, __context: object) -> None:
     _suffix_install_state: str = PrivateAttr(default="unused")
     _installed_suffix: str | None = PrivateAttr(default=None)
     _restart_session_on_next_turn: bool = PrivateAttr(default=False)
+    # Stream identity for the turn in flight; see stream_context.py. Held on
+    # the agent rather than threaded through the finalizers because the ACP
+    # turn already resolves through four of them.
+    _stream: StreamContext | None = PrivateAttr(default=None)
     _resumed_existing_session: bool = PrivateAttr(default=False)
     _file_credential_lifecycles: dict[str, ACPFileCredentialLifecycle] = PrivateAttr(
         default_factory=dict
@@ -3688,7 +3693,13 @@ def _finalize_successful_turn(
         # completed turn for eval/remote consumers, matching #2190.
         finish_action = FinishAction(message=response_text)
         tc_id = str(uuid.uuid4())
+        # An ACP turn's streamed text lands here, not in a MessageEvent, so
+        # this is the event that retires the stream's slot.
+        minted: dict[str, Any] = {}
+        if self._stream is not None and (item_id := self._stream.claim()):
+            minted["id"] = item_id
         action_event = ActionEvent(
+            **minted,
             source="agent",
             thought=[],
             reasoning_content=thought_text or None,
@@ -3704,6 +3715,8 @@ def _finalize_successful_turn(
             llm_response_id=str(uuid.uuid4()),
         )
         on_event(action_event)
+        if self._stream is not None and minted:
+            self._stream.commit()
         on_event(
             ObservationEvent(
                 observation=FinishObservation.from_text(text=response_text),
@@ -3865,6 +3878,19 @@ def step(
         (``LocalConversation.arun``) goes through :meth:`astep`, which
         avoids the cross-thread state-lock deadlock described in #3348.
         """
+        with StreamContext.open(conversation, on_token) as stream:
+            self._stream = stream
+            try:
+                self._step(conversation, on_event, stream.token_callback)
+            finally:
+                self._stream = None
+
+    def _step(
+        self,
+        conversation: LocalConversation,
+        on_event: ConversationCallbackType,
+        on_token: ConversationTokenCallbackType | None = None,
+    ) -> None:
         state = conversation.state
 
         if self._restart_session_on_next_turn:
@@ -3933,6 +3959,8 @@ async def _prompt() -> PromptResponse | None:
                         )
                         time.sleep(delay)
                         self._cancel_inflight_tool_calls()
+                        if self._stream is not None:
+                            self._stream.new_attempt()
                         self._reset_client_for_turn(
                             on_token,
                             on_event,
@@ -3963,6 +3991,8 @@ async def _prompt() -> PromptResponse | None:
                         )
                         time.sleep(delay)
                         self._cancel_inflight_tool_calls()
+                        if self._stream is not None:
+                            self._stream.new_attempt()
                         self._reset_client_for_turn(
                             on_token,
                             on_event,
@@ -4023,6 +4053,22 @@ async def astep(
         supplied by ``LocalConversation.arun`` is responsible for taking
         the state lock around each individual event.
         """
+        with StreamContext.open(conversation, on_token) as stream:
+            self._stream = stream
+            try:
+                await self._astep(
+                    conversation, on_event, stream.token_callback, prompt_message
+                )
+            finally:
+                self._stream = None
+
+    async def _astep(
+        self,
+        conversation: LocalConversation,
+        on_event: ConversationCallbackType,
+        on_token: ConversationTokenCallbackType | None = None,
+        prompt_message: MessageEvent | None = None,
+    ) -> None:
         state = conversation.state
 
         if self._restart_session_on_next_turn:
@@ -4099,6 +4145,8 @@ async def astep(
                         )
                         await asyncio.sleep(delay)
                         self._cancel_inflight_tool_calls()
+                        if self._stream is not None:
+                            self._stream.new_attempt()
                         self._reset_client_for_turn(
                             on_token,
                             on_event,
@@ -4126,6 +4174,8 @@ async def astep(
                         )
                         await asyncio.sleep(delay)
                         self._cancel_inflight_tool_calls()
+                        if self._stream is not None:
+                            self._stream.new_attempt()
                         self._reset_client_for_turn(
                             on_token,
                             on_event,