Support concurrent Python SDK turns across threads

This commit is contained in:
Shaqayeq
2026-03-19 16:16:02 -07:00
parent 6b8175c734
commit 083243dca1
8 changed files with 584 additions and 189 deletions

View File

@@ -41,8 +41,6 @@ class AsyncAppServerClient:
def __init__(self, config: AppServerConfig | None = None) -> None:
self._sync = AppServerClient(config=config)
# Single stdio transport cannot be read safely from multiple threads.
self._transport_lock = asyncio.Lock()
async def __aenter__(self) -> "AsyncAppServerClient":
await self.start()
@@ -58,8 +56,7 @@ class AsyncAppServerClient:
*args: ParamsT.args,
**kwargs: ParamsT.kwargs,
) -> ReturnT:
async with self._transport_lock:
return await asyncio.to_thread(fn, *args, **kwargs)
return await asyncio.to_thread(fn, *args, **kwargs)
@staticmethod
def _next_from_iterator(
@@ -79,11 +76,11 @@ class AsyncAppServerClient:
async def initialize(self) -> InitializeResponse:
return await self._call_sync(self._sync.initialize)
def acquire_turn_consumer(self, turn_id: str) -> None:
self._sync.acquire_turn_consumer(turn_id)
def acquire_turn_consumer(self, thread_id: str, turn_id: str) -> None:
self._sync.acquire_turn_consumer(thread_id, turn_id)
def release_turn_consumer(self, turn_id: str) -> None:
self._sync.release_turn_consumer(turn_id)
def release_turn_consumer(self, thread_id: str, turn_id: str) -> None:
self._sync.release_turn_consumer(thread_id, turn_id)
async def request(
self,
@@ -184,6 +181,9 @@ class AsyncAppServerClient:
async def next_notification(self) -> Notification:
return await self._call_sync(self._sync.next_notification)
async def next_turn_notification(self, thread_id: str, turn_id: str) -> Notification:
return await self._call_sync(self._sync.next_turn_notification, thread_id, turn_id)
async def wait_for_turn_completed(self, turn_id: str) -> TurnCompletedNotification:
return await self._call_sync(self._sync.wait_for_turn_completed, turn_id)
@@ -196,13 +196,12 @@ class AsyncAppServerClient:
text: str,
params: V2TurnStartParams | JsonObject | None = None,
) -> AsyncIterator[AgentMessageDeltaNotification]:
async with self._transport_lock:
iterator = self._sync.stream_text(thread_id, text, params)
while True:
has_value, chunk = await asyncio.to_thread(
self._next_from_iterator,
iterator,
)
if not has_value:
break
yield chunk
iterator = self._sync.stream_text(thread_id, text, params)
while True:
has_value, chunk = await asyncio.to_thread(
self._next_from_iterator,
iterator,
)
if not has_value:
break
yield chunk