Skip to content

x402 Client reuse #162

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adambalogh merged 1 commit into from
Feb 18, 2026
Merged

x402 Client reuse #162

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 10 additions & 0 deletions src/opengradient/client/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,3 +122,13 @@ def __init__(

if twins_api_key is not None:
self.twins = Twins(api_key=twins_api_key)

def close(self) -> None:
"""Close underlying SDK resources."""
self.llm.close()

def __enter__(self):
return self

def __exit__(self, exc_type, exc, tb):
self.close()
Comment on lines +130 to +134
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Client class's context manager implementation (enter and exit) is missing documentation. The class docstring should be updated to mention that the Client supports the context manager protocol and that users should either use it with a with statement or explicitly call close() when done. The existing usage examples in the docstring don't demonstrate proper resource cleanup.

Copilot uses AI. Check for mistakes.
Comment on lines +126 to +134
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new close() method and context manager functionality lack test coverage. The existing tests in tests/client_test.py don't test the resource cleanup, thread lifecycle, or context manager behavior. Given the complexity of the threading and async event loop management, comprehensive tests should be added to verify: (1) proper cleanup when close() is called, (2) context manager behavior, (3) thread joining, (4) event loop shutdown, and (5) behavior when methods are called after close().

Copilot uses AI. Check for mistakes.
321 changes: 166 additions & 155 deletions src/opengradient/client/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

import asyncio
import json
import threading
from queue import Queue
Comment on lines +5 to +6
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The imports threading and Queue are now only used within the _tee_llm_chat_stream_sync method (line 380-381), but are imported at the module level. Since these imports were moved from inside the method to the top of the file, consider whether they should remain at module level or be moved back inside the method where they're used. Module-level imports are generally preferred for performance and clarity, but in this case, the imports are only used in one specific method and were previously scoped there.

Copilot uses AI. Check for mistakes.
from typing import AsyncGenerator, Dict, List, Optional, Union

import httpx
Expand Down Expand Up @@ -62,6 +64,58 @@ def __init__(self, wallet_account: LocalAccount, og_llm_server_url: str, og_llm_
self._og_llm_server_url = og_llm_server_url
self._og_llm_streaming_server_url = og_llm_streaming_server_url

signer = EthAccountSignerv2(self._wallet_account)
self._x402_client = x402Clientv2()
register_exact_evm_clientv2(self._x402_client, signer, networks=[BASE_TESTNET_NETWORK])
register_upto_evm_clientv2(self._x402_client, signer, networks=[BASE_TESTNET_NETWORK])

self._request_client_ctx = None
self._request_client = None
self._stream_client_ctx = None
self._stream_client = None
self._closed = False

self._loop = asyncio.new_event_loop()
self._loop_thread = threading.Thread(target=self._run_event_loop, daemon=True)
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The event loop thread is created as a daemon thread (line 79), which means it will be abruptly terminated when the main program exits. This could leave HTTP connections open or in-progress requests incomplete. Since a close() method is now provided, consider making this a non-daemon thread to ensure proper cleanup. Users who don't call close() would then need to ensure they do so, which is enforced by the context manager pattern added in client.py.

Suggested change
self._loop_thread = threading.Thread(target=self._run_event_loop, daemon=True)
self._loop_thread = threading.Thread(target=self._run_event_loop)

Copilot uses AI. Check for mistakes.
self._loop_thread.start()
self._run_coroutine(self._initialize_http_clients())
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If an exception occurs during _run_coroutine(self._initialize_http_clients()) on line 81, the event loop thread will be left running because close() won't be called. Consider wrapping the initialization in a try-except block to ensure cleanup on failure, or implementing a del method as a safety net to stop the event loop thread if the object is garbage collected without being properly closed.

Copilot uses AI. Check for mistakes.
Comment on lines +78 to +81
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a potentially breaking API change. Previously, users could create an LLM client and let it be garbage collected without explicitly managing resources. Now, with the persistent event loop thread and HTTP clients, users should either use the context manager pattern or explicitly call close(). However, the daemon thread setting (line 79) means the old behavior still works (thread exits when program exits), but it's not clean. Consider updating the class docstring to document the new resource management requirements and provide usage examples with the context manager.

Copilot uses AI. Check for mistakes.

def _run_event_loop(self):
asyncio.set_event_loop(self._loop)
self._loop.run_forever()

def _run_coroutine(self, coroutine):
if self._closed:
raise OpenGradientError("LLM client is closed.")
future = asyncio.run_coroutine_threadsafe(coroutine, self._loop)
return future.result()
Comment on lines +87 to +91
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The _run_coroutine method calls future.result() without a timeout on line 91. If the coroutine hangs or takes a very long time, this will block indefinitely. Consider adding a timeout parameter to future.result() to prevent indefinite blocking, especially since network operations are involved in the HTTP client initialization and requests.

Suggested change
def _run_coroutine(self, coroutine):
if self._closed:
raise OpenGradientError("LLM client is closed.")
future = asyncio.run_coroutine_threadsafe(coroutine, self._loop)
return future.result()
def _run_coroutine(self, coroutine, timeout: Optional[float] = None):
if self._closed:
raise OpenGradientError("LLM client is closed.")
future = asyncio.run_coroutine_threadsafe(coroutine, self._loop)
try:
# Default to the configured overall HTTP timeout if none is provided.
effective_timeout = TIMEOUT.timeout if timeout is None else timeout
return future.result(timeout=effective_timeout)
except asyncio.TimeoutError as exc:
future.cancel()
raise OpenGradientError(
f"Timed out after {effective_timeout} seconds while waiting for asynchronous operation to complete."
) from exc

Copilot uses AI. Check for mistakes.

async def _initialize_http_clients(self) -> None:
if self._request_client is None:
self._request_client_ctx = x402HttpxClientv2(self._x402_client)
self._request_client = await self._request_client_ctx.__aenter__()
if self._stream_client is None:
self._stream_client_ctx = x402HttpxClientv2(self._x402_client)
self._stream_client = await self._stream_client_ctx.__aenter__()

async def _close_http_clients(self) -> None:
if self._request_client_ctx is not None:
await self._request_client_ctx.__aexit__(None, None, None)
self._request_client_ctx = None
self._request_client = None
if self._stream_client_ctx is not None:
await self._stream_client_ctx.__aexit__(None, None, None)
self._stream_client_ctx = None
self._stream_client = None

def close(self) -> None:
if self._closed:
return
self._run_coroutine(self._close_http_clients())
self._closed = True
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop_thread.join(timeout=5)
Comment on lines +111 to +117
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's no synchronization between close() and ongoing requests. If close() is called while a request is in progress via _run_coroutine(), the HTTP clients will be closed (line 114) while the request coroutine is still using them. This could lead to errors in the ongoing request. Consider either: (1) using a lock to prevent close() from proceeding while requests are in progress, (2) keeping track of active requests and waiting for them to complete in close(), or (3) documenting that users must ensure no requests are in progress before calling close().

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The thread join on line 117 has a 5-second timeout, but there's no handling for the case where the thread doesn't join within that time. If the event loop thread is still running after 5 seconds (e.g., due to a stuck coroutine), the daemon thread will be left running. Consider either removing the timeout to wait indefinitely, or adding logging/error handling for timeout cases. Additionally, consider forcefully canceling all pending tasks in the loop before stopping it to ensure a clean shutdown.

Suggested change
self._loop_thread.join(timeout=5)
self._loop_thread.join()

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The _loop variable holds a reference to an event loop that is never explicitly closed. While asyncio.new_event_loop() creates a loop, the loop should be closed using loop.close() after it has been stopped to properly release resources. Consider adding self._loop.close() after self._loop_thread.join(timeout=5) in the close() method to ensure the event loop is properly disposed of.

Suggested change
self._loop_thread.join(timeout=5)
self._loop_thread.join(timeout=5)
if not self._loop.is_running() and not self._loop.is_closed():
self._loop.close()

Copilot uses AI. Check for mistakes.
Comment on lines +111 to +117
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's a potential race condition in the close() method. If multiple threads call close() concurrently, both could pass the check on line 112 before either sets _closed = True on line 115. This could lead to attempting to close the HTTP clients multiple times or stopping the event loop multiple times. Consider using a lock to protect the close() method or checking _closed again after closing the HTTP clients but before stopping the loop.

Copilot uses AI. Check for mistakes.
Comment on lines +111 to +117
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The close() method is a new public API method but lacks documentation. All other public methods in this class have comprehensive docstrings (e.g., ensure_opg_approval, completion, chat). The close() method should include a docstring explaining its purpose, when it should be called, and what resources it releases. This is especially important since this is a breaking change affecting resource management.

Copilot uses AI. Check for mistakes.
Comment on lines +115 to +117
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The order of operations in close() could be problematic. The method sets _closed = True (line 115) before stopping the event loop (line 116). If any coroutine running in the event loop tries to spawn a new coroutine via _run_coroutine after _closed is set but before the loop stops, it will raise "LLM client is closed" error. Consider setting _closed = True after successfully stopping the event loop and joining the thread.

Suggested change
self._closed = True
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop_thread.join(timeout=5)
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop_thread.join(timeout=5)
self._closed = True

Copilot uses AI. Check for mistakes.

def ensure_opg_approval(self, opg_amount: float) -> Permit2ApprovalResult:
"""Ensure the Permit2 allowance for OPG is at least ``opg_amount``.

Expand Down Expand Up @@ -139,45 +193,40 @@ def _tee_llm_completion(
"""

async def make_request_v2():
x402_client = x402Clientv2()
register_exact_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])
register_upto_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])

# Security Fix: verify=True enabled
async with x402HttpxClientv2(x402_client) as client:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {X402_PLACEHOLDER_API_KEY}",
"X-SETTLEMENT-TYPE": x402_settlement_mode,
}

payload = {
"model": model,
"prompt": prompt,
"max_tokens": max_tokens,
"temperature": temperature,
}

if stop_sequence:
payload["stop"] = stop_sequence

try:
response = await client.post(self._og_llm_server_url + "/v1/completions", json=payload, headers=headers, timeout=60)

# Read the response content
content = await response.aread()
result = json.loads(content.decode())

return TextGenerationOutput(
transaction_hash="external",
completion_output=result.get("completion"),
)

except Exception as e:
raise OpenGradientError(f"TEE LLM completion request failed: {str(e)}")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {X402_PLACEHOLDER_API_KEY}",
"X-SETTLEMENT-TYPE": x402_settlement_mode,
}

payload = {
"model": model,
"prompt": prompt,
"max_tokens": max_tokens,
"temperature": temperature,
}

if stop_sequence:
payload["stop"] = stop_sequence

try:
response = await self._request_client.post(
self._og_llm_server_url + "/v1/completions", json=payload, headers=headers, timeout=60
)
Comment on lines +213 to +215
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using a shared HTTP client across multiple requests could cause issues if a request fails and leaves the client in an inconsistent state. Consider whether connection pooling errors or HTTP connection issues could affect subsequent requests. The httpx library's connection pool should handle this gracefully, but it would be helpful to verify this behavior with the x402HttpxClientv2 wrapper. Additionally, ensure that the client can handle concurrent requests if multiple threads call completion() or chat() simultaneously.

Suggested change
response = await self._request_client.post(
self._og_llm_server_url + "/v1/completions", json=payload, headers=headers, timeout=60
)
async with x402HttpxClientv2(timeout=TIMEOUT, limits=LIMITS) as client:
response = await client.post(
self._og_llm_server_url + "/v1/completions",
json=payload,
headers=headers,
timeout=60,
)

Copilot uses AI. Check for mistakes.

content = await response.aread()
result = json.loads(content.decode())

return TextGenerationOutput(
transaction_hash="external",
completion_output=result.get("completion"),
)

except Exception as e:
raise OpenGradientError(f"TEE LLM completion request failed: {str(e)}")

try:
return asyncio.run(make_request_v2())
return self._run_coroutine(make_request_v2())
except OpenGradientError:
raise
except Exception as e:
Expand Down Expand Up @@ -262,55 +311,50 @@ def _tee_llm_chat(
"""

async def make_request_v2():
x402_client = x402Clientv2()
register_exact_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])
register_upto_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])

# Security Fix: verify=True enabled
async with x402HttpxClientv2(x402_client) as client:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {X402_PLACEHOLDER_API_KEY}",
"X-SETTLEMENT-TYPE": x402_settlement_mode,
}

payload = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}

if stop_sequence:
payload["stop"] = stop_sequence

if tools:
payload["tools"] = tools
payload["tool_choice"] = tool_choice or "auto"

try:
# Non-streaming with x402
endpoint = "/v1/chat/completions"
response = await client.post(self._og_llm_server_url + endpoint, json=payload, headers=headers, timeout=60)

content = await response.aread()
result = json.loads(content.decode())

choices = result.get("choices")
if not choices:
raise OpenGradientError(f"Invalid response: 'choices' missing or empty in {result}")

return TextGenerationOutput(
transaction_hash="external",
finish_reason=choices[0].get("finish_reason"),
chat_output=choices[0].get("message"),
)

except Exception as e:
raise OpenGradientError(f"TEE LLM chat request failed: {str(e)}")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {X402_PLACEHOLDER_API_KEY}",
"X-SETTLEMENT-TYPE": x402_settlement_mode,
}

payload = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}

if stop_sequence:
payload["stop"] = stop_sequence

if tools:
payload["tools"] = tools
payload["tool_choice"] = tool_choice or "auto"

try:
endpoint = "/v1/chat/completions"
response = await self._request_client.post(
self._og_llm_server_url + endpoint, json=payload, headers=headers, timeout=60
)

content = await response.aread()
result = json.loads(content.decode())

choices = result.get("choices")
if not choices:
raise OpenGradientError(f"Invalid response: 'choices' missing or empty in {result}")

return TextGenerationOutput(
transaction_hash="external",
finish_reason=choices[0].get("finish_reason"),
chat_output=choices[0].get("message"),
)

except Exception as e:
raise OpenGradientError(f"TEE LLM chat request failed: {str(e)}")

try:
return asyncio.run(make_request_v2())
return self._run_coroutine(make_request_v2())
except OpenGradientError:
raise
except Exception as e:
Expand All @@ -333,72 +377,44 @@ def _tee_llm_chat_stream_sync(
Yields StreamChunk objects as they arrive from the background thread.
NO buffering, NO conversion, just direct pass-through.
"""
import threading
from queue import Queue

queue = Queue()
exception_holder = []
sentinel = object()

def _run_async():
"""Run async streaming in background thread"""
loop = None
async def _stream():
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

async def _stream():
try:
async for chunk in self._tee_llm_chat_stream_async(
model=model,
messages=messages,
max_tokens=max_tokens,
stop_sequence=stop_sequence,
temperature=temperature,
tools=tools,
tool_choice=tool_choice,
x402_settlement_mode=x402_settlement_mode,
):
queue.put(chunk) # Put chunk immediately
except Exception as e:
exception_holder.append(e)
finally:
queue.put(None) # Signal completion

loop.run_until_complete(_stream())
async for chunk in self._tee_llm_chat_stream_async(
model=model,
messages=messages,
max_tokens=max_tokens,
stop_sequence=stop_sequence,
temperature=temperature,
tools=tools,
tool_choice=tool_choice,
x402_settlement_mode=x402_settlement_mode,
):
queue.put(chunk)
except Exception as e:
exception_holder.append(e)
queue.put(None)
queue.put(e)
finally:
if loop:
try:
pending = asyncio.all_tasks(loop)
for task in pending:
task.cancel()
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
# Properly close async generators to avoid RuntimeWarning
loop.run_until_complete(loop.shutdown_asyncgens())
finally:
loop.close()

# Start background thread
thread = threading.Thread(target=_run_async, daemon=True)
thread.start()

# Yield chunks DIRECTLY as they arrive - NO buffering
queue.put(sentinel)

future = asyncio.run_coroutine_threadsafe(_stream(), self._loop)
Comment on lines 380 to +401
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The _tee_llm_chat_stream_sync method submits a coroutine to the shared event loop for streaming. If multiple threads call this method concurrently, they'll all be using the same event loop thread. While asyncio event loops are designed to handle multiple coroutines, the queue-based communication pattern could lead to mixed chunks from different streams if not handled carefully. Verify that each streaming call gets its own queue instance (which it does on line 380) and that there's no cross-contamination between concurrent streaming requests.

Copilot uses AI. Check for mistakes.

try:
while True:
chunk = queue.get() # Blocks until chunk available
if chunk is None:
chunk = queue.get()
if chunk is sentinel:
break
yield chunk # Yield immediately!

thread.join(timeout=5)

if exception_holder:
raise exception_holder[0]
if isinstance(chunk, Exception):
raise chunk
yield chunk
except Exception:
thread.join(timeout=1)
if not future.done():
future.cancel()
raise
finally:
if not future.done():
future.cancel()
Comment on lines +415 to +417
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the finally block, the future is cancelled on line 417 even if it has already completed successfully. Cancelling a completed future is harmless but unnecessary. Consider only cancelling if the future is not done, or restructure the error handling to avoid double cancellation (both in the except block on line 413 and in the finally block on line 417).

Copilot uses AI. Check for mistakes.

async def _tee_llm_chat_stream_async(
self,
Expand Down Expand Up @@ -474,18 +490,13 @@ async def _parse_sse_response(response) -> AsyncGenerator[StreamChunk, None]:
except json.JSONDecodeError:
continue

x402_client = x402Clientv2()
register_exact_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])
register_upto_evm_clientv2(x402_client, EthAccountSignerv2(self._wallet_account), networks=[BASE_TESTNET_NETWORK])

async with x402HttpxClientv2(x402_client) as client:
endpoint = "/v1/chat/completions"
async with client.stream(
"POST",
self._og_llm_streaming_server_url + endpoint,
json=payload,
headers=headers,
timeout=60,
) as response:
async for parsed_chunk in _parse_sse_response(response):
yield parsed_chunk
endpoint = "/v1/chat/completions"
async with self._stream_client.stream(
"POST",
self._og_llm_streaming_server_url + endpoint,
json=payload,
headers=headers,
timeout=60,
) as response:
async for parsed_chunk in _parse_sse_response(response):
yield parsed_chunk
Loading