Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion src/conductor/client/orkes/orkes_agent_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def __init__(self, configuration: Configuration):
self._schedule_client_instance: Any = None
# Persistent httpx client for async SSE streaming (created lazily).
self._sse_async_client: Any = None
self._sse_async_client_loop: Any = None

# ── internal helpers ─────────────────────────────────────────────────

Expand Down Expand Up @@ -272,10 +273,19 @@ async def stream_sse_async(
await asyncio.sleep(1)

def _get_sse_async_client(self):
import asyncio

import httpx

if self._sse_async_client is None or self._sse_async_client.is_closed:
current_loop = asyncio.get_running_loop()
# httpx transports are bound to the event loop that first uses them.
if (
self._sse_async_client is None
or self._sse_async_client.is_closed
or self._sse_async_client_loop is not current_loop
):
self._sse_async_client = httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=5.0))
self._sse_async_client_loop = current_loop
return self._sse_async_client

@staticmethod
Expand Down Expand Up @@ -356,6 +366,7 @@ async def close_async(self) -> None:
if self._sse_async_client is not None and not self._sse_async_client.is_closed:
await self._sse_async_client.aclose()
self._sse_async_client = None
self._sse_async_client_loop = None
if self._async_api_client is not None:
try:
await self._async_api_client.close()
Expand Down
19 changes: 19 additions & 0 deletions tests/unit/orkes/test_orkes_agent_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
HTTP errors map to the agent SDK exceptions.
"""

import asyncio
import base64
import json
import threading
Expand Down Expand Up @@ -195,3 +196,21 @@ def test_parse_sse_events_and_heartbeats():
assert events[2]["event"] == "done"
assert events[2]["id"] == "42"
assert events[2]["data"]["output"] == "result"


def test_sse_async_client_is_recreated_for_a_new_event_loop():
client = _client("http://localhost:8080/api")

async def get_first_client():
return client._get_sse_async_client()

first_client = asyncio.run(get_first_client())

async def get_second_client():
second_client = client._get_sse_async_client()
await client.close_async()
return second_client

second_client = asyncio.run(get_second_client())

assert second_client is not first_client