diff --git a/openhands-agent-server/openhands/agent_server/event_service.py b/openhands-agent-server/openhands/agent_server/event_service.py index f69f4a4d82..85a28e7bdd 100644 --- a/openhands-agent-server/openhands/agent_server/event_service.py +++ b/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) diff --git a/openhands-agent-server/openhands/agent_server/session_protocol.py b/openhands-agent-server/openhands/agent_server/session_protocol.py index b92a3ec162..6f9830beed 100644 --- a/openhands-agent-server/openhands/agent_server/session_protocol.py +++ b/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): diff --git a/openhands-agent-server/openhands/agent_server/session_socket.py b/openhands-agent-server/openhands/agent_server/session_socket.py index 08b76bfc64..4e02c35e6e 100644 --- a/openhands-agent-server/openhands/agent_server/session_socket.py +++ b/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"): diff --git a/openhands-sdk/openhands/sdk/agent/acp_agent.py b/openhands-sdk/openhands/sdk/agent/acp_agent.py index 3f6988b6a3..194dfa4852 100644 --- a/openhands-sdk/openhands/sdk/agent/acp_agent.py +++ b/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, diff --git a/openhands-sdk/openhands/sdk/agent/agent.py b/openhands-sdk/openhands/sdk/agent/agent.py index 828e8f4565..6c3ae89843 100644 --- a/openhands-sdk/openhands/sdk/agent/agent.py +++ b/openhands-sdk/openhands/sdk/agent/agent.py @@ -4,7 +4,7 @@ import re from collections.abc import Callable from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from pydantic import PrivateAttr, ValidationError, model_validator @@ -18,6 +18,7 @@ ResponseDispatchMixin, classify_response, ) +from openhands.sdk.agent.stream_context import StreamContext from openhands.sdk.agent.utils import ( amake_llm_completion, aprepare_llm_messages, @@ -639,6 +640,15 @@ def step( conversation: LocalConversation, on_event: ConversationCallbackType, on_token: ConversationTokenCallbackType | None = None, + ) -> None: + with StreamContext.open(conversation, on_token) as stream: + self._step(conversation, on_event, stream) + + def _step( + self, + conversation: LocalConversation, + on_event: ConversationCallbackType, + stream: StreamContext, ) -> None: state = conversation.state # Check for pending actions (implicit confirmation) @@ -721,7 +731,7 @@ def step( self.llm, _messages, tools=list(self.tools_map.values()), - on_token=on_token, + on_token=stream.token_callback, call_context=call_context, ) except FunctionCallValidationError as e: @@ -807,11 +817,11 @@ def step( match response_type: case LLMResponseType.TOOL_CALLS: self._handle_tool_calls( - message, llm_response, conversation, state, on_event + message, llm_response, conversation, state, on_event, stream ) case LLMResponseType.CONTENT: self._handle_content_response( - message, llm_response, conversation, state, on_event + message, llm_response, conversation, state, on_event, stream ) case LLMResponseType.REASONING_ONLY | LLMResponseType.EMPTY: self._handle_no_content_response( @@ -820,6 +830,7 @@ def step( conversation, state, on_event, + stream, response_type=response_type, ) @@ -839,6 +850,15 @@ async def astep( parallel calls with :func:`asyncio.gather`, keeping the event loop responsive during blocking tool I/O. """ + with StreamContext.open(conversation, on_token) as stream: + await self._astep(conversation, on_event, stream) + + async def _astep( + self, + conversation: LocalConversation, + on_event: ConversationCallbackType, + stream: StreamContext, + ) -> None: state = conversation.state # Check for pending actions (implicit confirmation) pending_actions = ConversationState.get_unmatched_actions(state.active_branch()) @@ -918,7 +938,7 @@ async def astep( self.llm, _messages, tools=list(self.tools_map.values()), - on_token=on_token, + on_token=stream.token_callback, call_context=call_context, ) except FunctionCallValidationError as e: @@ -1006,11 +1026,11 @@ async def astep( match response_type: case LLMResponseType.TOOL_CALLS: await self._ahandle_tool_calls( - message, llm_response, conversation, state, on_event + message, llm_response, conversation, state, on_event, stream ) case LLMResponseType.CONTENT: self._handle_content_response( - message, llm_response, conversation, state, on_event + message, llm_response, conversation, state, on_event, stream ) case LLMResponseType.REASONING_ONLY | LLMResponseType.EMPTY: self._handle_no_content_response( @@ -1019,6 +1039,7 @@ async def astep( conversation, state, on_event, + stream, response_type=response_type, ) @@ -1207,6 +1228,7 @@ def _get_action_event( reasoning_content: str | None = None, thinking_blocks: list[ThinkingBlock | RedactedThinkingBlock] | None = None, responses_reasoning_item: ReasoningItemModel | None = None, + stream: StreamContext | None = None, ) -> ActionEvent | None: """Converts a tool call into an ActionEvent, validating arguments. @@ -1314,8 +1336,15 @@ def _get_action_event( ) return - # Create initial action event + # Create initial action event. Claimed here rather than at the call + # site so an error path, which never builds this event, leaves the slot + # open for the abort in StreamContext.__exit__. + minted: dict[str, Any] = {} + if stream is not None and (item_id := stream.claim()): + minted["id"] = item_id + action_event = ActionEvent( + **minted, action=action, thought=thought or [], reasoning_content=reasoning_content, @@ -1339,6 +1368,8 @@ def _get_action_event( ) on_event(action_event) + if stream is not None and minted: + stream.commit() return action_event def _execute_action_event( diff --git a/openhands-sdk/openhands/sdk/agent/response_dispatch.py b/openhands-sdk/openhands/sdk/agent/response_dispatch.py index aa8abdfad3..7ba1ab2eb5 100644 --- a/openhands-sdk/openhands/sdk/agent/response_dispatch.py +++ b/openhands-sdk/openhands/sdk/agent/response_dispatch.py @@ -9,8 +9,9 @@ from __future__ import annotations from enum import StrEnum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from openhands.sdk.agent.stream_context import StreamContext from openhands.sdk.conversation.state import ConversationExecutionStatus from openhands.sdk.event import MessageEvent from openhands.sdk.llm import LLMResponse, Message, TextContent @@ -106,6 +107,7 @@ def _get_action_event( list[ThinkingBlock | RedactedThinkingBlock] | None ) = None, responses_reasoning_item: ReasoningItemModel | None = None, + stream: StreamContext | None = None, ) -> ActionEvent | None: ... def _execute_actions( @@ -147,6 +149,7 @@ def _handle_tool_calls( conversation: LocalConversation, state: ConversationState, on_event: ConversationCallbackType, + stream: StreamContext | None = None, ) -> None: """Handle LLM response containing tool calls.""" if not all(isinstance(c, TextContent) for c in message.content): @@ -172,6 +175,9 @@ def _handle_tool_calls( responses_reasoning_item=( message.responses_reasoning_item if i == 0 else None ), + # The streamed text is this action's thought, so the first + # action event is what retires the slot. + stream=stream if i == 0 else None, ) if action_event is None: continue @@ -192,6 +198,7 @@ async def _ahandle_tool_calls( conversation: LocalConversation, state: ConversationState, on_event: ConversationCallbackType, + stream: StreamContext | None = None, ) -> None: """Async variant of :meth:`_handle_tool_calls`. @@ -222,6 +229,9 @@ async def _ahandle_tool_calls( responses_reasoning_item=( message.responses_reasoning_item if i == 0 else None ), + # The streamed text is this action's thought, so the first + # action event is what retires the slot. + stream=stream if i == 0 else None, ) if action_event is None: continue @@ -242,9 +252,10 @@ def _handle_content_response( conversation: LocalConversation, state: ConversationState, on_event: ConversationCallbackType, + stream: StreamContext | None = None, ) -> None: """Handle LLM response with text content — finishes conversation.""" - self._emit_message_event(message, llm_response, conversation, on_event) + self._emit_message_event(message, llm_response, conversation, on_event, stream) self._maybe_emit_vllm_tokens(llm_response, on_event) logger.debug("LLM produced a message response - awaits user input") state.execution_status = ConversationExecutionStatus.FINISHED @@ -256,6 +267,7 @@ def _handle_no_content_response( conversation: LocalConversation, state: ConversationState, # noqa: ARG002 on_event: ConversationCallbackType, + stream: StreamContext | None = None, *, response_type: LLMResponseType, ) -> None: @@ -267,7 +279,7 @@ def _handle_no_content_response( """ if response_type is LLMResponseType.EMPTY: logger.warning("LLM produced empty response - continuing agent loop") - self._emit_message_event(message, llm_response, conversation, on_event) + self._emit_message_event(message, llm_response, conversation, on_event, stream) self._maybe_emit_vllm_tokens(llm_response, on_event) self._send_corrective_nudge(on_event) @@ -277,9 +289,18 @@ def _emit_message_event( llm_response: LLMResponse, conversation: LocalConversation, on_event: ConversationCallbackType, + stream: StreamContext | None = None, ) -> MessageEvent: - """Create and emit a MessageEvent, running critic if configured.""" + """Create and emit a MessageEvent, running critic if configured. + + The message carries the id its own stream minted, so a client holding + an open slot retires it on ``frame.event.id == slot.item_id``. + """ + minted: dict[str, Any] = {} + if stream is not None and (item_id := stream.claim()): + minted["id"] = item_id msg_event = MessageEvent( + **minted, source="agent", llm_message=self._mask_secrets(message, conversation), llm_response_id=llm_response.id, @@ -291,6 +312,10 @@ def _emit_message_event( update={"critic_result": critic_result} ) on_event(msg_event) + # Retired only now: on_event can raise while persisting, and a slot + # retired before that leaves the client holding it open forever. + if stream is not None and minted: + stream.commit() return msg_event @staticmethod diff --git a/openhands-sdk/openhands/sdk/agent/stream_context.py b/openhands-sdk/openhands/sdk/agent/stream_context.py new file mode 100644 index 0000000000..67a4c7c218 --- /dev/null +++ b/openhands-sdk/openhands/sdk/agent/stream_context.py @@ -0,0 +1,288 @@ +"""Stream identity for one agent step: mint it, stamp the deltas, close it once. + +The durable event is built with the minted id, so a client retires its open +slot on ``frame.event.id == slot.item_id``. + +Minting up front is safe against the append-only log because it is not a write: +``Event.id`` is already client-minted in-process (``event/base.py``), so this +changes only *when* ``uuid4()`` runs, and an unused id was never on disk. +""" + +from __future__ import annotations + +import asyncio +import itertools +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal + +from openhands.sdk.conversation.secret_registry import StreamOutputMask +from openhands.sdk.llm.streaming import LLMStreamChunk +from openhands.sdk.logger import get_logger + + +if TYPE_CHECKING: + from openhands.sdk.conversation.impl.local_conversation import LocalConversation + + +logger = get_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class StreamStarted: + """A stream is opening. One per attempt, before its first token. + + ``anchor_seq`` is the seq the slot sits after, so a user message landing + mid-stream cannot split it. + """ + + item_id: str + attempt: int + anchor_seq: int | None + + +@dataclass(frozen=True, slots=True) +class StreamDelta: + """One masked increment of a stream. + + ``chunk_id`` is corroboration only: litellm mints a new completion id per + retry attempt, so it cannot identify the durable event. + """ + + item_id: str + attempt: int + order: int + kind: Literal["text", "reasoning"] + content: str + chunk_id: str | None = None + choice_index: int | None = None + + +@dataclass(frozen=True, slots=True) +class StreamAborted: + """A stream ended without producing a durable event.""" + + item_id: str + attempt: int + reason: str + + +StreamProgress = StreamStarted | StreamDelta | StreamAborted +StreamProgressCallbackType = Callable[[StreamProgress], None] + + +class StreamContext: + """One streaming slot: mints its id, stamps its deltas, closes it once.""" + + def __init__( + self, + item_id: str, + anchor_seq: int | None, + on_token: Callable[[Any], None] | None, + on_stream: StreamProgressCallbackType | None, + mask: Callable[[], StreamOutputMask], + ) -> None: + self.item_id = item_id + self._anchor_seq = anchor_seq + self._on_token = on_token + self._on_stream = on_stream + self._mask = mask + # One masker per kind: text and reasoning are separately ordered, and a + # masker holds back a partial secret across the chunks of its own kind. + self._masks: dict[Literal["text", "reasoning"], StreamOutputMask] = {} + self._attempt = 1 + self._order = itertools.count() + self._started = False + self._opened = False + self._reserved = False + self._claimed = False + self._chunk_id: str | None = None + + @classmethod + def open( + cls, + conversation: LocalConversation, + on_token: Callable[[Any], None] | None, + ) -> StreamContext: + """Mint an id and read the anchor. Neither touches the event log.""" + state = conversation.state + length = len(state.events) + return cls( + item_id=str(uuid.uuid4()), + anchor_seq=length - 1 if length else None, + on_token=on_token, + on_stream=conversation.on_stream, + mask=state.secret_registry.compile_stream_mask, + ) + + @property + def token_callback(self) -> Callable[[Any], None] | None: + """The callback to hand the LLM, or ``None`` when nothing consumes it. + + ``llm.completion`` degrades a ``stream=True`` model to a non-streaming + call when ``on_token`` is ``None`` (#4014); an unconditional wrapper + would take that fallback away. + """ + if self._on_token is None and self._on_stream is None: + return None + return self.on_chunk + + def on_chunk(self, chunk: Any) -> None: + """Forward the raw chunk downstream, then emit its stamped deltas. + + The pass-through is what keeps existing ``on_token`` consumers seeing + exactly what they see today. + """ + if self._on_token is not None: + self._on_token(chunk) + if self._on_stream is None: + return + try: + for kind, text, chunk_id, choice_index in _split_chunk(chunk): + self._emit_delta(kind, text, chunk_id, choice_index) + except Exception: + # Progress is a UX affordance; never fail a turn over it. + logger.debug("stream progress emission failed", exc_info=True) + + def new_attempt(self) -> None: + """Start a new attempt for the same item; a higher one supersedes.""" + if not self._opened: + return + self._attempt += 1 + self._order = itertools.count() + self._started = False + self._chunk_id = None + # The new attempt re-streams the item, so the old tail is superseded. + self._masks.clear() + + def claim(self) -> str | None: + """Reserve the minted id for a durable event being built. + + ``None`` once reserved, so a second durable event in the same step + keeps its own id rather than colliding. Reserving does not retire the + slot: an event that is never emitted still owes an abort, so the + caller must :meth:`commit` once emission has succeeded. + """ + if self._reserved or self._on_stream is None: + return None + self._flush() + self._reserved = True + return self.item_id + + def commit(self) -> None: + """Retire the slot: the durable event carrying the id was emitted.""" + if self._reserved: + self._claimed = True + + def close(self, reason: str) -> None: + """Emit ``StreamAborted`` unless a durable event committed the id. + + Keyed on ever-opened, not on the current attempt: a retry that dies + before its first token still owes an answer for the one that streamed. + """ + if not self._opened or self._claimed: + return + self._claimed = True + self._flush() + self._emit(StreamAborted(self.item_id, self._attempt, reason)) + + def __enter__(self) -> StreamContext: + return self + + def __exit__(self, exc_type, exc, tb) -> Literal[False]: + self.close(_abort_reason(exc_type)) + return False + + def _emit_delta( + self, + kind: Literal["text", "reasoning"], + text: str, + chunk_id: str | None, + choice_index: int | None, + ) -> None: + if chunk_id is not None and self._chunk_id is not None: + if chunk_id != self._chunk_id: + # litellm mints a completion id per attempt, so a change of id + # mid-item is a re-stream of the same slot. + self.new_attempt() + self._chunk_id = chunk_id + if not self._started: + self._started = self._opened = True + self._emit(StreamStarted(self.item_id, self._attempt, self._anchor_seq)) + masker = self._masks.get(kind) + if masker is None: + masker = self._masks[kind] = self._mask() + released = masker.feed(text) + if released: + self._emit_content(kind, released, chunk_id, choice_index) + + def _flush(self) -> None: + """Release the tail each masker holds back against a split secret.""" + for kind, masker in self._masks.items(): + held = masker.flush() + if held: + self._emit_content(kind, held, self._chunk_id, None) + self._masks.clear() + + def _emit_content( + self, + kind: Literal["text", "reasoning"], + content: str, + chunk_id: str | None, + choice_index: int | None, + ) -> None: + self._emit( + StreamDelta( + item_id=self.item_id, + attempt=self._attempt, + order=next(self._order), + kind=kind, + content=content, + chunk_id=chunk_id, + choice_index=choice_index, + ) + ) + + def _emit(self, frame: StreamProgress) -> None: + """Never raises: ``close()`` runs from ``__exit__``, where a sink error + would replace the exception the step is unwinding.""" + assert self._on_stream is not None + try: + self._on_stream(frame) + except Exception: + logger.debug("stream progress sink failed", exc_info=True) + + +def _abort_reason(exc_type: type[BaseException] | None) -> str: + if exc_type is None: + return "no_durable_event" + if issubclass(exc_type, (asyncio.CancelledError, KeyboardInterrupt)): + return "cancelled" + return exc_type.__name__ + + +def _split_chunk( + chunk: LLMStreamChunk | str, +) -> list[tuple[Literal["text", "reasoning"], str, str | None, int | None]]: + """Break one token-callback payload into the deltas it carries. + + The ACP bridge passes a bare ``str`` and never enters the LLM layer, which + is why this lives above both rather than in ``llm.py``. + """ + if isinstance(chunk, str): + return [("text", chunk, None, None)] if chunk else [] + + out: list[tuple[Literal["text", "reasoning"], str, str | None, int | None]] = [] + for choice in chunk.choices or (): + delta = choice.delta + if delta is None: + continue + # getattr, not attribute access: litellm *deletes* reasoning_content + # when the provider omits it, declared field or not. + reasoning = getattr(delta, "reasoning_content", None) + if isinstance(reasoning, str) and reasoning: + out.append(("reasoning", reasoning, chunk.id, choice.index)) + if isinstance(delta.content, str) and delta.content: + out.append(("text", delta.content, chunk.id, choice.index)) + return out diff --git a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py index f647e50af9..d705b84a03 100644 --- a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py +++ b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py @@ -10,6 +10,7 @@ from openhands.sdk.agent.acp_agent import ACPAgent from openhands.sdk.agent.base import AgentBase +from openhands.sdk.agent.stream_context import StreamProgressCallbackType from openhands.sdk.context.condenser import CondenserBase, LLMSummarizingCondenser from openhands.sdk.context.memory import load_memory from openhands.sdk.context.prompts.prompt import render_template @@ -183,6 +184,7 @@ class LocalConversation(BaseConversation): _visualizer: ConversationVisualizerBase | None _on_event: ConversationCallbackType _on_token: ConversationTokenCallbackType | None + _on_stream: StreamProgressCallbackType | None max_iteration_per_run: int _stuck_detector: StuckDetector | None llm_registry: LLMRegistry @@ -238,6 +240,7 @@ def __init__( file_store: FileStore | None = None, mcp_tool_provider: MCPToolProvider | None = None, profile_store_dir: str | Path | None = None, + stream_callbacks: list[StreamProgressCallbackType] | None = None, **_: object, ): """Initialize the conversation. @@ -260,6 +263,8 @@ def __init__( suffix their persistent filestore with this ID. callbacks: Optional list of callback functions to handle events token_callbacks: Optional list of callbacks invoked for streaming deltas + stream_callbacks: Optional list of callbacks invoked with the + stream-progress frames minted by ``StreamContext``. hook_config: Optional hook configuration to auto-wire session hooks. If plugins are loaded, their hooks are combined with this config. max_iteration_per_run: Maximum number of iterations per run @@ -468,6 +473,14 @@ def _default_callback(e): if token_callbacks else None ) + self._on_stream = ( + cast( + StreamProgressCallbackType, + BaseConversation.compose_callbacks(stream_callbacks), # type: ignore[arg-type] + ) + if stream_callbacks + else None + ) self.max_iteration_per_run = max_iteration_per_run # Hard cost ceiling (USD) for a run; None disables the budget check. @@ -2538,6 +2551,15 @@ def set_confirmation_policy(self, policy: ConfirmationPolicyBase) -> None: self._state.confirmation_policy = policy logger.info(f"Confirmation policy set to: {policy}") + @property + def on_stream(self) -> StreamProgressCallbackType | None: + """Sink for stream-progress frames, or ``None`` if nothing consumes them. + + Read by the agent rather than passed to ``step``: a new ``step`` + parameter would break every third-party ``AgentBase`` subclass. + """ + return self._on_stream + def set_token_callbacks( self, token_callbacks: list[ConversationTokenCallbackType] | None ) -> None: diff --git a/openhands-sdk/openhands/sdk/conversation/secret_registry.py b/openhands-sdk/openhands/sdk/conversation/secret_registry.py index 1ca3ebdd38..c753a6a489 100644 --- a/openhands-sdk/openhands/sdk/conversation/secret_registry.py +++ b/openhands-sdk/openhands/sdk/conversation/secret_registry.py @@ -1,5 +1,6 @@ """Secrets manager for handling sensitive data in conversations.""" +import re import time from collections.abc import Callable, Collection, Mapping from enum import Enum @@ -59,6 +60,52 @@ def _mask_model[ModelT: BaseModel](model: ModelT, mask: Callable[[str], str]) -> return model.model_copy(update=updates) +class StreamOutputMask: + """Masks text that arrives in pieces. + + A per-chunk masker cannot see a secret split across chunk boundaries, so + this holds back the shortest suffix that could still grow into one: the + longest registered value minus one character. + """ + + def __init__(self, pattern: re.Pattern[str] | None, max_len: int) -> None: + self._pattern = pattern + self._max_len = max_len + self._held = "" + + def feed(self, text: str) -> str: + """Return the masked text that is safe to release now.""" + pattern = self._pattern + if pattern is None: + return text + self._held += text + return self._release(pattern, max(len(self._held) - self._max_len + 1, 0)) + + def flush(self) -> str: + """Return what is still held back; no more input is coming.""" + pattern = self._pattern + if pattern is None: + return "" + return self._release(pattern, len(self._held)) + + def _release(self, pattern: re.Pattern[str], cut: int) -> str: + out: list[str] = [] + pos = 0 + end = cut + for match in pattern.finditer(self._held): + # A match starting before the cut is already maximal: the cut + # leaves max_len-1 characters of lookahead behind it. + if match.start() >= cut: + break + out.append(self._held[pos : match.start()]) + out.append("") + pos = match.end() + end = max(end, match.end()) + out.append(self._held[pos:end]) + self._held = self._held[end:] + return "".join(out) + + class SecretRegistry(OpenHandsModel): """Manages secrets and injects them into bash commands when needed. @@ -219,6 +266,25 @@ def mask_secrets_in_output(self, text: str) -> str: return masked_text + def compile_stream_mask(self) -> StreamOutputMask: + """Return a streaming masker over the values resolved *so far*. + + For hot paths that cannot afford :meth:`mask_secrets_in_output`, which + resolves uncached sources first — ``get_value()`` may block on network + I/O, and on the sync agent path that runs under the state lock. Masks + less in exchange: use it for progress, not for the durable record. + """ + with self._exported_values_lock: + values = {value for value in self._exported_values.values() if value} + if not values: + return StreamOutputMask(None, 0) + # Longest first so an overlapping shorter value cannot mask half of a + # longer one and leave the rest in cleartext. + pattern = re.compile( + "|".join(re.escape(v) for v in sorted(values, key=len, reverse=True)) + ) + return StreamOutputMask(pattern, max(len(v) for v in values)) + def mask_secrets_in_model[ModelT: BaseModel](self, model: ModelT) -> ModelT: """Return ``model`` with secret values masked in every nested string. diff --git a/openhands-sdk/openhands/sdk/utils/models.py b/openhands-sdk/openhands/sdk/utils/models.py index 1c1ead5df0..fbe42e0d7c 100644 --- a/openhands-sdk/openhands/sdk/utils/models.py +++ b/openhands-sdk/openhands/sdk/utils/models.py @@ -174,7 +174,12 @@ def clear_subclass_cache() -> None: Normally not needed — the cache auto-invalidates when new DiscriminatedUnionMixin subclasses are defined. This function exists for edge cases involving non-DiscriminatedUnionMixin hierarchies. + + Entries are dropped, not just superseded: a stale one is a strong + reference to every subclass it captured. """ + _concrete_cache.clear() + _checked_cache.clear() _bump_subclass_generation() diff --git a/tests/agent_server/test_event_streaming.py b/tests/agent_server/test_event_streaming.py index 7a5194d8d7..4a611efbdf 100644 --- a/tests/agent_server/test_event_streaming.py +++ b/tests/agent_server/test_event_streaming.py @@ -16,6 +16,7 @@ from openhands.sdk import Event from openhands.sdk.agent import ACPAgent, Agent from openhands.sdk.agent.acp_agent import ACTIVITY_SIGNAL_INTERVAL +from openhands.sdk.agent.stream_context import StreamProgress, StreamStarted from openhands.sdk.event import StreamingDeltaEvent from openhands.sdk.llm import LLM from openhands.sdk.workspace import LocalWorkspace @@ -366,3 +367,47 @@ async def test_delta_idle_signal_is_throttled(event_service, tmp_path, monkeypat monkeypatch.setattr(server_details_router, "_last_event_time", stale) callback(_make_chunk(content="second")) assert server_details_router._last_event_time == stale + + +async def _start_and_capture_stream_callback(event_service, tmp_path): + """Start the service and return the wired stream-progress callback.""" + (tmp_path / "workspace").mkdir(exist_ok=True) + + with _mock_local_conversation() as MockConv: + mock_conv = MagicMock() + mock_conv.state = MagicMock() + mock_conv.state.execution_status = "idle" + mock_conv._state = MagicMock() + mock_conv._on_event = MagicMock() + MockConv.return_value = mock_conv + + await event_service.start() + return MockConv.call_args.kwargs["stream_callbacks"][0] + + +class _ProgressCollector(Subscriber[StreamProgress]): + def __init__(self): + self.frames: list[StreamProgress] = [] + + async def __call__(self, frame: StreamProgress): + self.frames.append(frame) + + async def close(self): + pass + + +@pytest.mark.asyncio +async def test_stream_progress_reaches_its_own_subscribers(event_service, tmp_path): + """Progress rides a separate fan-out, so no event-bus consumer sees it.""" + callback = await _start_and_capture_stream_callback(event_service, tmp_path) + + progress = _ProgressCollector() + event_service._stream_pub_sub.subscribe(progress) + on_the_event_bus = _CollectorSubscriber() + event_service._pub_sub.subscribe(on_the_event_bus) + + callback(StreamStarted(item_id="item-1", attempt=1, anchor_seq=3)) + await asyncio.sleep(0.05) + + assert progress.frames == [StreamStarted("item-1", 1, 3)] + assert on_the_event_bus.events == [] diff --git a/tests/agent_server/test_session_socket.py b/tests/agent_server/test_session_socket.py index 5958771996..e707defe1c 100644 --- a/tests/agent_server/test_session_socket.py +++ b/tests/agent_server/test_session_socket.py @@ -11,11 +11,17 @@ from openhands.agent_server.session_socket import ( _ConnectionWriter, _inbound_loop, + _ProgressSubscriber, _read_page, _replay, _SessionSubscriber, ) from openhands.sdk import Message, TextContent +from openhands.sdk.agent.stream_context import ( + StreamAborted, + StreamDelta, + StreamStarted, +) from openhands.sdk.event import MessageEvent, StreamingDeltaEvent from openhands.sdk.event.conversation_state import ConversationStateUpdateEvent @@ -323,3 +329,67 @@ async def receive_json(): assert frame["code"] == "JSONDecodeError" # The loop kept going rather than unwinding. assert receives == 2 + + +@pytest.mark.asyncio +async def test_progress_frames_reach_the_wire_with_their_identity(): + """StreamContext's frames map onto the envelope, one for one.""" + ws = _FakeWebSocket() + writer = _ConnectionWriter(ws) # type: ignore[arg-type] + writer.start() + sub = _ProgressSubscriber(writer) + + await sub(StreamStarted(item_id="item-1", attempt=1, anchor_seq=4)) + await sub( + StreamDelta( + item_id="item-1", + attempt=1, + order=0, + kind="reasoning", + content="thinking", + chunk_id="chatcmpl-9", + choice_index=0, + ) + ) + await sub(StreamAborted(item_id="item-1", attempt=1, reason="cancelled")) + await _drain(writer) + await writer.aclose() + + started, delta, aborted = ws.frames() + assert started == { + "type": "item_started", + "item_id": "item-1", + "attempt": 1, + "anchor_seq": 4, + } + assert delta["type"] == "delta" + assert (delta["order"], delta["kind"], delta["content"]) == ( + 0, + "reasoning", + "thinking", + ) + # The provider's own identity, carried through as corroboration. + assert (delta["chunk_id"], delta["choice_index"]) == ("chatcmpl-9", 0) + assert aborted["type"] == "item_aborted" + assert aborted["reason"] == "cancelled" + + +@pytest.mark.asyncio +async def test_a_dropped_progress_frame_does_not_drop_the_connection(): + """Progress is recoverable from the durable event, so it is droppable.""" + ws = _FakeWebSocket() + writer = _ConnectionWriter(ws) # type: ignore[arg-type] + writer.start() + sub = _ProgressSubscriber(writer) + + await sub( + StreamDelta( + item_id="item-1", + attempt=1, + order=0, + kind="text", + content="x" * (MAX_FRAME_BYTES + 1), + ) + ) + await _drain(writer) + assert writer.closed diff --git a/tests/sdk/agent/test_acp_agent.py b/tests/sdk/agent/test_acp_agent.py index 01da4933b9..87181f4bda 100644 --- a/tests/sdk/agent/test_acp_agent.py +++ b/tests/sdk/agent/test_acp_agent.py @@ -60,6 +60,11 @@ ) from openhands.sdk.agent.acp_models import ACPModelInfo from openhands.sdk.agent.base import AgentBase +from openhands.sdk.agent.stream_context import ( + StreamAborted, + StreamDelta, + StreamStarted, +) from openhands.sdk.context import AgentContext from openhands.sdk.conversation.secret_registry import SecretRegistry from openhands.sdk.conversation.state import ( @@ -2234,11 +2239,78 @@ def _fake_run_async(_coro, **_kwargs): agent.step(conversation, on_event=lambda _: None, on_token=on_token) - # Verify on_token was wired during the turn. - assert wired_during_prompt == [on_token] + # The bridge is wired with StreamContext's stamping wrapper, which + # forwards to the caller's callback unchanged. + assert len(wired_during_prompt) == 1 + wired = wired_during_prompt[0] + assert wired is not None and wired is not on_token + wired("chunk") + on_token.assert_called_once_with("chunk") # And unwired afterward so a late token chunk is a no-op. assert mock_client.on_token is None + def test_step_retires_its_stream_on_the_turn_finish_action(self, tmp_path): + """An ACP turn's streamed text lands in the FinishAction, not a message. + + See https://github.com/OpenHands/software-agent-sdk/issues/4682. + """ + agent = _make_agent() + conversation = self._make_conversation_with_message(tmp_path) + frames: list = [] + conversation.on_stream = frames.append + + mock_client = _OpenHandsACPBridge() + agent._client = mock_client + agent._conn = MagicMock() + agent._session_id = "test-session" + + def _fake_run_async(_coro, **_kwargs): + mock_client.on_token("streamed ") + mock_client.on_token("text") + mock_client.accumulated_text.append("streamed text") + + mock_executor = MagicMock() + mock_executor.run_async = _fake_run_async + agent._executor = mock_executor + + events: list = [] + agent.step(conversation, on_event=events.append) + + started = [f for f in frames if isinstance(f, StreamStarted)] + deltas = [f for f in frames if isinstance(f, StreamDelta)] + assert len(started) == 1 + assert [d.content for d in deltas] == ["streamed ", "text"] + + action = next(e for e in events if isinstance(e, ActionEvent)) + assert action.id == started[0].item_id + assert not any(isinstance(f, StreamAborted) for f in frames) + # The context is released with the turn. + assert agent._stream is None + + @pytest.mark.asyncio + async def test_astep_opens_and_retires_the_same_slot(self, tmp_path): + """The async entry point gets the same guarantee as the sync one.""" + agent = _make_agent() + conversation = self._make_conversation_with_message(tmp_path) + frames: list = [] + conversation.on_stream = frames.append + + async def _fake_astep(_self, _conv, _on_event, on_token=None, _prompt=None): + assert on_token is not None + on_token("streamed text") + raise RuntimeError("the prompt died after streaming") + + with patch.object(ACPAgent, "_astep", _fake_astep): + with pytest.raises(RuntimeError): + await agent.astep(conversation, on_event=lambda _: None) + + started = [f for f in frames if isinstance(f, StreamStarted)] + aborted = [f for f in frames if isinstance(f, StreamAborted)] + assert len(started) == len(aborted) == 1 + assert aborted[0].item_id == started[0].item_id + assert aborted[0].reason == "RuntimeError" + assert agent._stream is None + # --------------------------------------------------------------------------- # Async step (astep) — regression coverage for #3348 diff --git a/tests/sdk/agent/test_stream_context.py b/tests/sdk/agent/test_stream_context.py new file mode 100644 index 0000000000..774e814521 --- /dev/null +++ b/tests/sdk/agent/test_stream_context.py @@ -0,0 +1,292 @@ +"""StreamContext mints one identity per step and always closes it. + +See https://github.com/OpenHands/software-agent-sdk/issues/4682. +""" + +import asyncio +import re +import uuid + +import pytest +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + +from openhands.sdk.agent.stream_context import ( + StreamAborted, + StreamContext, + StreamDelta, + StreamStarted, +) +from openhands.sdk.conversation.secret_registry import StreamOutputMask + + +def _no_secrets() -> StreamOutputMask: + return StreamOutputMask(None, 0) + + +def _masking(*values: str) -> StreamOutputMask: + pattern = re.compile( + "|".join(re.escape(v) for v in sorted(values, key=len, reverse=True)) + ) + return StreamOutputMask(pattern, max(len(v) for v in values)) + + +def _chunk( + content: str | None = None, + reasoning_content: str | None = None, + chunk_id: str = "chunk-1", + index: int = 0, +) -> ModelResponseStream: + delta_kwargs: dict = {"role": "assistant"} + if content is not None: + delta_kwargs["content"] = content + delta = Delta(**delta_kwargs) + if reasoning_content is not None: + object.__setattr__(delta, "reasoning_content", reasoning_content) + choice = StreamingChoices(delta=delta, index=index, finish_reason=None) + return ModelResponseStream(id=chunk_id, choices=[choice], model="test-model") + + +def _make( + frames: list, forwarded: list | None = None, mask=_no_secrets +) -> StreamContext: + return StreamContext( + item_id=str(uuid.uuid4()), + anchor_seq=7, + on_token=forwarded.append if forwarded is not None else None, + on_stream=frames.append, + mask=mask, + ) + + +def test_a_step_that_never_streams_opens_nothing(): + frames: list = [] + with _make(frames): + pass + assert frames == [] + + +def test_first_delta_opens_the_slot_with_the_anchor(): + frames: list = [] + with _make(frames) as stream: + stream.on_chunk(_chunk("hello")) + + started, delta, aborted = frames + assert isinstance(started, StreamStarted) + assert started.item_id == stream.item_id + assert started.attempt == 1 + assert started.anchor_seq == 7 + assert isinstance(delta, StreamDelta) + assert (delta.kind, delta.content, delta.order) == ("text", "hello", 0) + assert (delta.chunk_id, delta.choice_index) == ("chunk-1", 0) + # Nothing claimed the id, so the slot is retired by an abort. + assert isinstance(aborted, StreamAborted) + + +def test_a_committed_id_retires_the_slot_without_an_abort(): + frames: list = [] + with _make(frames) as stream: + stream.on_chunk(_chunk("hello")) + claimed = stream.claim() + stream.commit() + + assert claimed == stream.item_id + assert not any(isinstance(f, StreamAborted) for f in frames) + + +def test_a_claim_that_never_commits_still_aborts(): + """on_event can raise while persisting, after the id was handed out.""" + frames: list = [] + with pytest.raises(RuntimeError): + with _make(frames) as stream: + stream.on_chunk(_chunk("hello")) + stream.claim() + raise RuntimeError("persisting the durable event failed") + + aborted = [f for f in frames if isinstance(f, StreamAborted)] + assert [f.reason for f in aborted] == ["RuntimeError"] + + +def test_the_id_is_claimable_once(): + frames: list = [] + stream = _make(frames) + stream.on_chunk(_chunk("hi")) + assert stream.claim() == stream.item_id + assert stream.claim() is None + + +@pytest.mark.parametrize( + "exc, expected", + [ + (RuntimeError("boom"), "RuntimeError"), + (asyncio.CancelledError(), "cancelled"), + (None, "no_durable_event"), + ], + ids=["provider-failure", "cancellation", "returned-without-a-message"], +) +def test_every_opened_slot_is_retired_exactly_once(exc, expected): + frames: list = [] + stream = _make(frames) + try: + with stream: + stream.on_chunk(_chunk("partial")) + if exc is not None: + raise exc + except BaseException as raised: + assert raised is exc + + aborts = [f for f in frames if isinstance(f, StreamAborted)] + assert len(aborts) == 1 + assert aborts[0].reason == expected + assert aborts[0].item_id == stream.item_id + + +def test_a_retry_re_streams_the_same_item_under_a_higher_attempt(): + frames: list = [] + with _make(frames) as stream: + stream.on_chunk(_chunk("half", chunk_id="completion-a")) + # litellm mints a new completion id per attempt. + stream.on_chunk(_chunk("half", chunk_id="completion-b")) + stream.on_chunk(_chunk(" again", chunk_id="completion-b")) + stream.claim() + + starts = [f for f in frames if isinstance(f, StreamStarted)] + deltas = [f for f in frames if isinstance(f, StreamDelta)] + assert [f.attempt for f in starts] == [1, 2] + assert {f.item_id for f in starts} == {stream.item_id} + # order is monotonic within (item_id, attempt), so the retry restarts it. + assert [(d.attempt, d.order) for d in deltas] == [(1, 0), (2, 0), (2, 1)] + + +def test_reasoning_and_text_are_separate_ordered_deltas(): + frames: list = [] + with _make(frames) as stream: + stream.on_chunk(_chunk(content="answer", reasoning_content="thought")) + stream.claim() + + deltas = [f for f in frames if isinstance(f, StreamDelta)] + assert [(d.kind, d.content, d.order) for d in deltas] == [ + ("reasoning", "thought", 0), + ("text", "answer", 1), + ] + + +def test_deltas_are_masked_by_the_snapshot(): + frames: list = [] + with _make(frames, mask=lambda: _masking("hunter2")) as stream: + stream.on_chunk(_chunk("token is hunter2")) + stream.claim() + + delta = next(f for f in frames if isinstance(f, StreamDelta)) + assert delta.content == "token is " + + +def test_the_raw_chunk_still_reaches_the_token_callback(): + """The CLI, the legacy socket and user callbacks must see no change.""" + frames: list = [] + forwarded: list = [] + chunk = _chunk("hello") + with _make(frames, forwarded) as stream: + stream.on_chunk(chunk) + stream.on_chunk("acp bare string") + stream.claim() + + assert forwarded == [chunk, "acp bare string"] + + +def test_no_consumer_means_no_callback_for_the_llm(): + """`llm.completion` degrades a stream=True model when on_token is None.""" + silent = StreamContext( + item_id="item", + anchor_seq=None, + on_token=None, + on_stream=None, + mask=_no_secrets, + ) + assert silent.token_callback is None + + frames: list = [] + assert _make(frames).token_callback is not None + assert _make(frames, forwarded=[]).token_callback is not None + + +def test_without_a_sink_nothing_is_minted_and_chunks_still_flow(): + forwarded: list = [] + stream = StreamContext( + item_id=str(uuid.uuid4()), + anchor_seq=None, + on_token=forwarded.append, + on_stream=None, + mask=_no_secrets, + ) + with stream: + stream.on_chunk(_chunk("hello")) + # No consumer, so the durable event keeps its own id. + assert stream.claim() is None + assert len(forwarded) == 1 + + +def test_a_retry_that_dies_before_its_first_token_still_retires_the_slot(): + """The abort must name the attempt whose start the client observed.""" + frames: list = [] + stream = _make(frames) + try: + with stream: + stream.on_chunk(_chunk("half")) + stream.new_attempt() + raise RuntimeError("the retry never produced a token") + except RuntimeError: + pass + + starts = [f for f in frames if isinstance(f, StreamStarted)] + aborts = [f for f in frames if isinstance(f, StreamAborted)] + assert len(starts) == len(aborts) == 1 + assert (aborts[0].item_id, aborts[0].attempt) == ( + starts[0].item_id, + starts[0].attempt, + ) + + +def test_a_broken_sink_does_not_fail_the_turn(): + def explode(_frame): + raise ValueError("subscriber blew up") + + forwarded: list = [] + stream = StreamContext( + item_id="item", + anchor_seq=None, + on_token=forwarded.append, + on_stream=explode, + mask=_no_secrets, + ) + with stream: + stream.on_chunk(_chunk("hello")) + # Including the abort in __exit__, which would otherwise replace whatever + # exception the step is unwinding. + assert len(forwarded) == 1 + + +def test_a_malformed_chunk_does_not_fail_the_turn(): + class _Exploding: + @property + def choices(self): + raise ValueError("this provider payload is not what we expected") + + frames: list = [] + forwarded: list = [] + stream = _make(frames, forwarded) + stream.on_chunk(_Exploding()) + + assert len(forwarded) == 1 + assert frames == [] + + +def test_a_choice_without_a_delta_is_skipped(): + choice = StreamingChoices(delta=Delta(role="assistant"), index=0) + object.__setattr__(choice, "delta", None) + chunk = ModelResponseStream(id="c", choices=[choice], model="test-model") + + frames: list = [] + stream = _make(frames) + stream.on_chunk(chunk) + + assert frames == [] diff --git a/tests/sdk/agent/test_stream_identity.py b/tests/sdk/agent/test_stream_identity.py new file mode 100644 index 0000000000..cf303922e3 --- /dev/null +++ b/tests/sdk/agent/test_stream_identity.py @@ -0,0 +1,323 @@ +"""The durable event a stream produces carries the id that stream minted. + +Both :class:`Agent` entry points: the sync path streams too, so an async-only +fix would cover half the problem. + +See https://github.com/OpenHands/software-agent-sdk/issues/4682. +""" + +from collections.abc import Sequence +from typing import cast + +import pytest +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices +from pydantic import PrivateAttr + +from openhands.sdk.agent import Agent +from openhands.sdk.agent.stream_context import ( + StreamAborted, + StreamDelta, + StreamStarted, +) +from openhands.sdk.conversation.impl.local_conversation import LocalConversation +from openhands.sdk.event import ActionEvent, MessageEvent +from openhands.sdk.llm import Message, MessageToolCall, TextContent +from openhands.sdk.llm.exceptions import LLMNoResponseError +from openhands.sdk.testing import TestLLM +from openhands.sdk.tool import ( + Action, + Observation, + Tool, + ToolDefinition, + ToolExecutor, + register_tool, +) + + +class _StreamEchoAction(Action): + text: str + + +class _StreamEchoObservation(Observation): + pass + + +class _StreamEchoExecutor(ToolExecutor[_StreamEchoAction, _StreamEchoObservation]): + def __call__( + self, action: _StreamEchoAction, conversation=None + ) -> _StreamEchoObservation: + return _StreamEchoObservation.from_text(action.text) + + +class _StreamEchoTool(ToolDefinition[_StreamEchoAction, _StreamEchoObservation]): + name = "stream_echo" + + @classmethod + def create(cls, conv_state=None) -> Sequence["_StreamEchoTool"]: + return [ + cls( + description="Echo the given text", + action_type=_StreamEchoAction, + observation_type=_StreamEchoObservation, + executor=_StreamEchoExecutor(), + ) + ] + + +register_tool("EchoStreamTool", _StreamEchoTool) + + +class StreamingTestLLM(TestLLM): + """A ``TestLLM`` that emits the scripted text as chunks first.""" + + _chunks: list[str] = PrivateAttr(default_factory=list) + _raises: BaseException | None = PrivateAttr(default=None) + _seen_on_token: list = PrivateAttr(default_factory=list) + + def script(self, chunks: list[str], raises: BaseException | None = None): + self._chunks = chunks + self._raises = raises + return self + + def _stream(self, on_token) -> None: + for text in self._chunks: + if on_token is None: + continue + on_token( + ModelResponseStream( + id="completion-1", + model="test-model", + choices=[ + StreamingChoices( + delta=Delta(role="assistant", content=text), + index=0, + finish_reason=None, + ) + ], + ) + ) + + def completion(self, messages, tools=None, on_token=None, **kwargs): # type: ignore[override] + self._seen_on_token.append(on_token) + self._stream(on_token) + if self._raises is not None: + raise self._raises + return super().completion(messages, tools=tools, **kwargs) + + async def acompletion(self, messages, tools=None, on_token=None, **kwargs): # type: ignore[override] + # TestLLM.acompletion drops on_token, so stream here instead. + self._stream(on_token) + if self._raises is not None: + raise self._raises + return await super().acompletion(messages, tools=tools, **kwargs) + + +def _conversation(tmp_path, llm, frames: list, tools: Sequence[Tool] = ()): + agent = Agent(llm=llm, tools=list(tools)) + convo = LocalConversation( + agent=agent, + workspace=str(tmp_path / "workspace"), + persistence_dir=str(tmp_path / "conversations"), + visualizer=None, + stream_callbacks=[frames.append], + ) + convo._ensure_agent_ready() + return convo + + +def _llm(messages, chunks, raises=None) -> StreamingTestLLM: + llm = cast(StreamingTestLLM, StreamingTestLLM.from_messages(messages)) + llm.stream = True + return llm.script(chunks, raises) + + +def test_the_message_carries_the_id_its_stream_minted(tmp_path): + frames: list = [] + llm = _llm( + [Message(role="assistant", content=[TextContent(text="Hello there")])], + ["Hello ", "there"], + ) + convo = _conversation(tmp_path, llm, frames) + + events: list = [] + convo.agent.step(convo, on_event=events.append) + + started = next(f for f in frames if isinstance(f, StreamStarted)) + message = next(e for e in events if isinstance(e, MessageEvent)) + assert message.id == started.item_id + assert [f.content for f in frames if isinstance(f, StreamDelta)] == [ + "Hello ", + "there", + ] + # Retired by the durable event, so no abort. + assert not any(isinstance(f, StreamAborted) for f in frames) + + +@pytest.mark.asyncio +async def test_astep_mints_the_same_way(tmp_path): + frames: list = [] + llm = _llm( + [Message(role="assistant", content=[TextContent(text="Hello there")])], + ["Hello ", "there"], + ) + convo = _conversation(tmp_path, llm, frames) + + events: list = [] + await convo.agent.astep(convo, on_event=events.append) + + started = next(f for f in frames if isinstance(f, StreamStarted)) + message = next(e for e in events if isinstance(e, MessageEvent)) + assert message.id == started.item_id + + +def test_a_tool_call_turn_retires_the_slot_on_its_first_action(tmp_path): + """The streamed text is that action's thought, so the action closes it.""" + frames: list = [] + llm = _llm( + [ + Message( + role="assistant", + content=[TextContent(text="Listing the directory")], + tool_calls=[ + MessageToolCall( + id="call-1", + name="stream_echo", + arguments='{"text": "hi"}', + origin="completion", + ) + ], + ) + ], + ["Listing ", "the directory"], + ) + convo = _conversation(tmp_path, llm, frames, tools=[Tool(name="EchoStreamTool")]) + + events: list = [] + convo.agent.step(convo, on_event=events.append) + + started = next(f for f in frames if isinstance(f, StreamStarted)) + action = next(e for e in events if isinstance(e, ActionEvent)) + assert action.id == started.item_id + assert not any(isinstance(f, StreamAborted) for f in frames) + + +def test_an_invalid_tool_call_retires_the_stream_with_its_action(tmp_path): + frames: list = [] + llm = _llm( + [ + Message( + role="assistant", + content=[TextContent(text="Trying the requested operation")], + tool_calls=[ + MessageToolCall( + id="call-invalid", + name="missing_tool", + arguments="{}", + origin="completion", + ) + ], + ) + ], + ["Trying the requested operation"], + ) + convo = _conversation(tmp_path, llm, frames) + + events: list = [] + convo.agent.step(convo, on_event=events.append) + + started = next(f for f in frames if isinstance(f, StreamStarted)) + invalid_action = next( + e for e in events if isinstance(e, ActionEvent) and e.action is None + ) + assert invalid_action.id == started.item_id + assert not any(isinstance(f, StreamAborted) for f in frames) + + +def test_a_provider_failure_retires_the_slot_with_an_abort(tmp_path): + frames: list = [] + llm = _llm( + [Message(role="assistant", content=[TextContent(text="unused")])], + ["half a "], + raises=LLMNoResponseError("provider gave up"), + ) + convo = _conversation(tmp_path, llm, frames) + + with pytest.raises(LLMNoResponseError): + convo.agent.step(convo, on_event=lambda _: None) + + started = [f for f in frames if isinstance(f, StreamStarted)] + aborted = [f for f in frames if isinstance(f, StreamAborted)] + assert len(started) == len(aborted) == 1 + assert aborted[0].item_id == started[0].item_id + + +def test_a_persistence_failure_after_streaming_retires_the_slot(tmp_path): + """The id is handed out before the event is emitted; emitting can raise.""" + frames: list = [] + llm = _llm( + [Message(role="assistant", content=[TextContent(text="hello there")])], + ["hello ", "there"], + ) + convo = _conversation(tmp_path, llm, frames) + + def explode(event): + if isinstance(event, MessageEvent) and event.source == "agent": + raise RuntimeError("persisting the durable event failed") + + with pytest.raises(RuntimeError): + convo.agent.step(convo, on_event=explode) + + started = [f for f in frames if isinstance(f, StreamStarted)] + aborted = [f for f in frames if isinstance(f, StreamAborted)] + assert len(started) == len(aborted) == 1 + assert aborted[0].item_id == started[0].item_id + + +def test_no_stream_consumer_leaves_the_llm_free_to_skip_streaming(tmp_path): + """Without a consumer the agent must not force a streaming completion.""" + llm = _llm([Message(role="assistant", content=[TextContent(text="hi")])], ["hi"]) + agent = Agent(llm=llm, tools=[]) + convo = LocalConversation( + agent=agent, + workspace=str(tmp_path / "workspace"), + persistence_dir=str(tmp_path / "conversations"), + visualizer=None, + ) + convo._ensure_agent_ready() + + convo.agent.step(convo, on_event=lambda _: None) + + assert llm._seen_on_token == [None] + + +def test_a_step_that_never_reaches_the_provider_opens_nothing(tmp_path): + """No LLM call, no slot — so there is nothing to abort.""" + frames: list = [] + llm = _llm([Message(role="assistant", content=[TextContent(text="hi")])], []) + convo = _conversation(tmp_path, llm, frames) + + convo.agent.step(convo, on_event=lambda _: None) + + assert frames == [] + + +def test_a_resolved_secret_is_masked_in_the_deltas(tmp_path): + """The LLM stream path masked nothing before this; it does now. + + Only already-resolved values are masked. A secret that can reach the + model's output was exported for a command first, which resolves it. + """ + frames: list = [] + llm = _llm( + [Message(role="assistant", content=[TextContent(text="done")])], + ["the token is ", "hunter2"], + ) + convo = _conversation(tmp_path, llm, frames) + convo.state.secret_registry.update_secrets({"TOKEN": "hunter2"}) + # Resolution happens when a command references the name. + convo.state.secret_registry.get_secrets_as_env_vars("echo $TOKEN") + + convo.agent.step(convo, on_event=lambda _: None) + + streamed = "".join(f.content for f in frames if isinstance(f, StreamDelta)) + assert streamed == "the token is " diff --git a/tests/sdk/conversation/local/test_client_tools_persistence.py b/tests/sdk/conversation/local/test_client_tools_persistence.py index 67571add6f..5ddb7074f1 100644 --- a/tests/sdk/conversation/local/test_client_tools_persistence.py +++ b/tests/sdk/conversation/local/test_client_tools_persistence.py @@ -1,11 +1,13 @@ """Persistence/resume behavior for client-defined tools on LocalConversation.""" +import gc import uuid from pathlib import Path from openhands.sdk import LLM, Agent, Conversation from openhands.sdk.tool import Tool, client_tool as ct, registry as reg from openhands.sdk.tool.client_tool import ClientToolSpec +from openhands.sdk.utils.models import clear_subclass_cache def _make_agent() -> Agent: @@ -23,6 +25,11 @@ def _wipe_client_tool_globals(names: list[str]) -> None: reg._REG.pop(name, None) reg._USABILITY_REG.pop(name, None) reg._MODULE_QUALNAMES.pop(name, None) + # The dynamic ClientAction classes are only weakly reachable once the + # registries drop them, but the subclass cache holds them strongly — and a + # surviving pair with the same name is a duplicate to any later rebuild. + clear_subclass_cache() + gc.collect() def test_persisted_client_tools_resume_without_respecifying(tmp_path: Path) -> None: diff --git a/tests/sdk/conversation/test_secrets_manager.py b/tests/sdk/conversation/test_secrets_manager.py index 0170fdc953..56a6afa254 100644 --- a/tests/sdk/conversation/test_secrets_manager.py +++ b/tests/sdk/conversation/test_secrets_manager.py @@ -38,6 +38,14 @@ def get_value(self): raise OSError("Secret retrieval failed") +class MyCountingLazySource(SecretSource): + attempts: int = 0 + + def get_value(self): + type(self).attempts += 1 + return "lazy-value" + + class MyRecoveringSource(SecretSource): fail: bool = True @@ -451,3 +459,52 @@ def test_mask_secrets_retries_until_source_succeeds(): assert secret_registry.mask_secrets_in_output("leak: recovered-value") == ( "leak: " ) + + +def test_compile_stream_mask_uses_only_already_resolved_values(): + """The stream path cannot afford a resolve: get_value() may block on I/O.""" + MyCountingLazySource.attempts = 0 + secret_registry = SecretRegistry() + secret_registry.update_secrets( + {"RESOLVED": "abc123", "LAZY": MyCountingLazySource()} + ) + secret_registry.get_secrets_as_env_vars("echo $RESOLVED") + + def mask(text: str) -> str: + masker = secret_registry.compile_stream_mask() + return masker.feed(text) + masker.flush() + + assert mask("leak: abc123") == "leak: " + # The unresolved source is not consulted, so its value is not masked here. + assert mask("leak: lazy-value") == "leak: lazy-value" + assert MyCountingLazySource.attempts == 0 + + +def test_compile_stream_mask_prefers_the_longest_overlapping_value(): + secret_registry = SecretRegistry() + secret_registry.update_secrets({"SHORT": "abc", "LONG": "abc123"}) + secret_registry.get_secrets_as_env_vars("echo $SHORT $LONG") + + masker = secret_registry.compile_stream_mask() + + assert masker.feed("leak: abc123") + masker.flush() == "leak: " + + +def test_compile_stream_mask_masks_a_value_split_across_chunks(): + """A per-chunk masker cannot see a secret the provider splits in two.""" + secret_registry = SecretRegistry() + secret_registry.update_secrets({"TOKEN": "hunter2"}) + secret_registry.get_secrets_as_env_vars("echo $TOKEN") + + masker = secret_registry.compile_stream_mask() + released = [masker.feed(piece) for piece in ("leak: ", "hunt", "er2", " done")] + released.append(masker.flush()) + + assert "".join(released) == "leak: done" + + +def test_compile_stream_mask_releases_text_when_nothing_is_registered(): + masker = SecretRegistry().compile_stream_mask() + + assert masker.feed("nothing held back") == "nothing held back" + assert masker.flush() == ""