Log warm-up cleanup failures and close the tests' scan clients

A cleanup step that fails in warm_up_llm's finally now logs at debug
instead of being silently suppressed. Each step still runs on its own,
and none of them replaces an error raised by the preflight.

The regression tests now close their scan-side client inside the scan's
own loop. Also drop a noqa on reset_subscription_client that ruff
reports as unused (RUF100).
This commit is contained in:
drakeo338 2026-09-25 00:45:13 +00:00
parent a9dbe6a32d
commit 70793a9a2e
3 changed files with 24 additions and 15 deletions

View file

@ -394,7 +394,7 @@ def get_subscription_client() -> AsyncOpenAI:
async def reset_subscription_client() -> None:
"""Close the cached client so the next ``get_subscription_client`` builds a
fresh one. Its pooled connections are bound to the event loop that opened them."""
global _subscription_client # noqa: PLW0603
global _subscription_client
client, _subscription_client = _subscription_client, None
if client is not None:
await client.close()

View file

@ -259,21 +259,27 @@ async def warm_up_llm(show_model_warning: bool = True) -> None:
# The scan runs under a different event loop, and pooled connections
# opened here would still be bound to this one. Close the shared
# clients now so the scan starts from fresh ones.
with contextlib.suppress(Exception):
try:
from agents.models import openai_provider
http_client = openai_provider._http_client # pyright: ignore[reportPrivateUsage]
openai_provider._http_client = None # pyright: ignore[reportPrivateUsage]
if http_client is not None:
await http_client.aclose()
with contextlib.suppress(Exception):
except Exception:
logger.debug("Closing the SDK HTTP client after warm-up failed", exc_info=True)
try:
await codex.reset_subscription_client()
except Exception:
logger.debug("Closing the subscription client after warm-up failed", exc_info=True)
# LiteLLM keys its client cache by event loop; this only releases the
# warm-up's sockets.
with contextlib.suppress(Exception):
try:
import litellm
await litellm.close_litellm_async_clients() # type: ignore[attr-defined,no-untyped-call]
except Exception:
logger.debug("Closing LiteLLM clients after warm-up failed", exc_info=True)
def display_completion_message(args: argparse.Namespace, results_path: Path) -> None:

View file

@ -98,17 +98,20 @@ class _KeepAliveHandler(BaseHTTPRequestHandler):
async def _stream_turn(client: AsyncOpenAI) -> str:
"""One streamed scan turn. The SDK's model retry disables the OpenAI client's
own retries, so a stale connection is not quietly replaced by a fresh one."""
stream = await client.with_options(max_retries=0).chat.completions.create(
model="gw-model",
messages=[{"role": "user", "content": "go"}],
stream=True,
)
text = ""
async for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
text += chunk.choices[0].delta.content
await stream.close()
return text
try:
stream = await client.with_options(max_retries=0).chat.completions.create(
model="gw-model",
messages=[{"role": "user", "content": "go"}],
stream=True,
)
text = ""
async for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
text += chunk.choices[0].delta.content
await stream.close()
return text
finally:
await client.close()
@pytest.fixture