fix(sdk): resolve async response secrets outside the event loop

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import re
 from collections.abc import Callable
@@ -1022,6 +1023,18 @@ async def _astep(
 
         message: Message = llm_response.message
         response_type = classify_response(message)
+        if response_type is not LLMResponseType.TOOL_CALLS:
+            # Resolve outside the event loop and state lock. A lookup may call
+            # this server, and update_secrets() may register another source while
+            # we await it. Repeat until the registry is stable under the lock.
+            while True:
+                sources = dict(state.secret_registry.secret_sources)
+                async with conversation._released_state_lock_during_io():
+                    message = await asyncio.to_thread(
+                        self._mask_secrets, message, conversation
+                    )
+                if sources == state.secret_registry.secret_sources:
+                    break
 
         match response_type:
             case LLMResponseType.TOOL_CALLS:
@@ -1030,7 +1043,13 @@ async def _astep(
                 )
             case LLMResponseType.CONTENT:
                 self._handle_content_response(
-                    message, llm_response, conversation, state, on_event, stream
+                    message,
+                    llm_response,
+                    conversation,
+                    state,
+                    on_event,
+                    stream,
+                    mask_secrets=False,
                 )
             case LLMResponseType.REASONING_ONLY | LLMResponseType.EMPTY:
                 self._handle_no_content_response(
@@ -1041,6 +1060,7 @@ async def _astep(
                     on_event,
                     stream,
                     response_type=response_type,
+                    mask_secrets=False,
                 )
 
     def _requires_user_confirmation(

openhands-sdk/openhands/sdk/agent/response_dispatch.py

@@ -253,9 +253,18 @@ def _handle_content_response(
         state: ConversationState,
         on_event: ConversationCallbackType,
         stream: StreamContext | None = None,
+        *,
+        mask_secrets: bool = True,
     ) -> None:
         """Handle LLM response with text content — finishes conversation."""
-        self._emit_message_event(message, llm_response, conversation, on_event, stream)
+        self._emit_message_event(
+            message,
+            llm_response,
+            conversation,
+            on_event,
+            stream,
+            mask_secrets=mask_secrets,
+        )
         self._maybe_emit_vllm_tokens(llm_response, on_event)
         logger.debug("LLM produced a message response - awaits user input")
         state.execution_status = ConversationExecutionStatus.FINISHED
@@ -270,6 +279,7 @@ def _handle_no_content_response(
         stream: StreamContext | None = None,
         *,
         response_type: LLMResponseType,
+        mask_secrets: bool = True,
     ) -> None:
         """Handle LLM response with no user-facing content.
 
@@ -279,7 +289,14 @@ 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, stream)
+        self._emit_message_event(
+            message,
+            llm_response,
+            conversation,
+            on_event,
+            stream,
+            mask_secrets=mask_secrets,
+        )
         self._maybe_emit_vllm_tokens(llm_response, on_event)
         self._send_corrective_nudge(on_event)
 
@@ -290,6 +307,8 @@ def _emit_message_event(
         conversation: LocalConversation,
         on_event: ConversationCallbackType,
         stream: StreamContext | None = None,
+        *,
+        mask_secrets: bool = True,
     ) -> MessageEvent:
         """Create and emit a MessageEvent, running critic if configured.
 
@@ -302,7 +321,9 @@ def _emit_message_event(
         msg_event = MessageEvent(
             **minted,
             source="agent",
-            llm_message=self._mask_secrets(message, conversation),
+            llm_message=self._mask_secrets(message, conversation)
+            if mask_secrets
+            else message,
             llm_response_id=llm_response.id,
         )
         if self.critic is not None and self.critic.mode == "finish_and_message":

openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py

@@ -2769,8 +2769,8 @@ def update_secrets(self, secrets: Mapping[str, SecretValue]) -> None:
                      SecretValue = str | Callable[[], str]. Callables are invoked lazily
                      when a command references the secret key.
         """
-        secret_registry = self._state.secret_registry
-        secret_registry.update_secrets(secrets)
+        with self._state:
+            self._state.secret_registry.update_secrets(secrets)
         logger.info(f"Added {len(secrets)} secrets to conversation")
 
     def set_security_analyzer(self, analyzer: SecurityAnalyzerBase | None) -> None:

tests/sdk/agent/test_async_secret_masking.py

@@ -0,0 +1,167 @@
+import asyncio
+
+import pytest
+
+from openhands.sdk import Agent, Conversation
+from openhands.sdk.event import ActionEvent, MessageEvent
+from openhands.sdk.llm import Message, TextContent
+from openhands.sdk.llm.message import MessageToolCall
+from openhands.sdk.secret import LookupSecret
+from openhands.sdk.testing import TestLLM
+
+
+@pytest.mark.asyncio
+async def test_async_response_masks_loopback_lookup_secret_without_blocking(tmp_path):
+    requested = asyncio.Event()
+
+    async def serve_secret(reader, writer):
+        await reader.readuntil(b"\r\n\r\n")
+        requested.set()
+        body = b"loopback-secret-value"
+        writer.write(
+            b"HTTP/1.1 200 OK\r\nContent-Length: "
+            + str(len(body)).encode()
+            + b"\r\nConnection: close\r\n\r\n"
+            + body
+        )
+        await writer.drain()
+        writer.close()
+        await writer.wait_closed()
+
+    server = await asyncio.start_server(serve_secret, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    llm = TestLLM.from_messages(
+        [
+            Message(
+                role="assistant",
+                content=[TextContent(text="value loopback-secret-value")],
+            )
+        ]
+    )
+    conversation = Conversation(
+        agent=Agent(llm=llm, tools=[]),
+        workspace=str(tmp_path),
+        visualizer=None,
+        secrets={"TEST_TOKEN": LookupSecret(url=f"http://127.0.0.1:{port}/secret")},
+    )
+    try:
+        conversation.send_message("hello")
+        await asyncio.wait_for(conversation.arun(), timeout=5)
+        assert requested.is_set()
+        messages = [
+            event
+            for event in conversation.state.events
+            if isinstance(event, MessageEvent) and event.source == "agent"
+        ]
+        assert messages[-1].llm_message.content == [
+            TextContent(text="value <secret-hidden>")
+        ]
+    finally:
+        await asyncio.to_thread(conversation.close)
+        server.close()
+        await server.wait_closed()
+
+
+@pytest.mark.asyncio
+async def test_async_tool_thought_matches_sync_behavior(tmp_path):
+    thoughts = []
+    for use_async in [False, True]:
+        message = Message(
+            role="assistant",
+            content=[TextContent(text="existing tool thought")],
+            tool_calls=[
+                MessageToolCall(
+                    id="finish-call",
+                    origin="completion",
+                    name="finish",
+                    arguments='{"message":"done"}',
+                )
+            ],
+        )
+        conversation = Conversation(
+            agent=Agent(llm=TestLLM.from_messages([message]), tools=[]),
+            workspace=str(tmp_path),
+            visualizer=None,
+            secrets={"TOKEN": "existing tool thought"},
+        )
+        try:
+            conversation.send_message("finish")
+            if use_async:
+                await conversation.arun()
+            else:
+                await asyncio.to_thread(conversation.run)
+            thoughts.append(
+                [
+                    event.thought
+                    for event in conversation.state.events
+                    if isinstance(event, ActionEvent)
+                ]
+            )
+        finally:
+            await asyncio.to_thread(conversation.close)
+    assert thoughts[0]
+    assert thoughts[0] == thoughts[1]
+
+
+@pytest.mark.asyncio
+async def test_async_response_masks_secrets_registered_during_lookup(tmp_path):
+    requested = []
+    values = ["first-secret-value", "second-secret-value", "third-secret-value"]
+
+    async def serve_secret(reader, writer):
+        request = await reader.readuntil(b"\r\n\r\n")
+        index = int(request.split(b" ")[1].removeprefix(b"/"))
+        requested.append(index)
+        if index + 1 < len(values):
+            await asyncio.to_thread(
+                conversation.update_secrets,
+                {
+                    f"TOKEN_{index + 1}": LookupSecret(
+                        url=f"http://127.0.0.1:{port}/{index + 1}"
+                    )
+                },
+            )
+        body = values[index].encode()
+        writer.write(
+            b"HTTP/1.1 200 OK\r\nContent-Length: "
+            + str(len(body)).encode()
+            + b"\r\nConnection: close\r\n\r\n"
+            + body
+        )
+        await writer.drain()
+        writer.close()
+        await writer.wait_closed()
+
+    server = await asyncio.start_server(serve_secret, "127.0.0.1", 0)
+    port = server.sockets[0].getsockname()[1]
+    conversation = Conversation(
+        agent=Agent(
+            llm=TestLLM.from_messages(
+                [
+                    Message(
+                        role="assistant", content=[TextContent(text=" ".join(values))]
+                    )
+                ]
+            ),
+            tools=[],
+        ),
+        workspace=str(tmp_path),
+        visualizer=None,
+        secrets={"TOKEN_0": LookupSecret(url=f"http://127.0.0.1:{port}/0")},
+    )
+    try:
+        conversation.send_message("hello")
+        await asyncio.wait_for(conversation.arun(), timeout=5)
+        assert requested == [0, 1, 2]
+        messages = [
+            event
+            for event in conversation.state.events
+            if isinstance(event, MessageEvent) and event.source == "agent"
+        ]
+        assert messages[-1].llm_message.content == [
+            TextContent(text="<secret-hidden> <secret-hidden> <secret-hidden>")
+        ]
+    finally:
+        await asyncio.to_thread(conversation.close)
+        server.close()
+        await server.wait_closed()