wehub-resource-sync
265 lines
8.1 KiB
Python
265 lines
8.1 KiB
Python
"""Provider streams raise canonical failures after closing committed blocks."""
|
|
|
|
from collections import deque
|
|
from collections.abc import AsyncIterator
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from free_claude_code.config.nim import NimSettings
|
|
from free_claude_code.core.anthropic.stream_contracts import parse_sse_text
|
|
from free_claude_code.core.async_iterators import AsyncCloseable
|
|
from free_claude_code.core.failures import ExecutionFailure, FailureKind
|
|
from free_claude_code.providers.base import ProviderConfig
|
|
from free_claude_code.providers.http import close_provider_stream
|
|
from free_claude_code.providers.nvidia_nim import NvidiaNimProvider
|
|
from tests.providers.request_factory import make_messages_request
|
|
from tests.providers.support import passthrough_rate_limiter
|
|
|
|
|
|
class _FailingStream:
|
|
def __init__(
|
|
self,
|
|
chunks: list[object],
|
|
error: Exception | None,
|
|
*,
|
|
close_error: Exception | None = None,
|
|
) -> None:
|
|
self._chunks = chunks
|
|
self._error = error
|
|
self._close_error = close_error
|
|
self.close_calls = 0
|
|
|
|
def __aiter__(self) -> AsyncIterator[object]:
|
|
return self._iterate()
|
|
|
|
async def _iterate(self) -> AsyncIterator[object]:
|
|
for chunk in self._chunks:
|
|
yield chunk
|
|
if self._error is not None:
|
|
raise self._error
|
|
|
|
async def aclose(self) -> None:
|
|
self.close_calls += 1
|
|
if self._close_error is not None:
|
|
raise self._close_error
|
|
|
|
|
|
def _chunk(*, content: str | None = None, finish_reason: str | None = None) -> object:
|
|
delta = MagicMock(content=content, tool_calls=None, reasoning_content=None)
|
|
choice = MagicMock(delta=delta, finish_reason=finish_reason)
|
|
return MagicMock(choices=[choice], usage=None)
|
|
|
|
|
|
def _provider() -> NvidiaNimProvider:
|
|
return NvidiaNimProvider(
|
|
ProviderConfig(
|
|
api_key="test_key",
|
|
base_url="https://test.api.nvidia.com/v1",
|
|
rate_limit=10,
|
|
rate_window=60,
|
|
),
|
|
nim_settings=NimSettings(),
|
|
rate_limiter=passthrough_rate_limiter(),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_committed_provider_failure_closes_block_then_raises_canonical_value() -> (
|
|
None
|
|
):
|
|
provider = _provider()
|
|
request = make_messages_request(
|
|
"test-model",
|
|
messages=[],
|
|
max_tokens=32,
|
|
)
|
|
# Crossing the recovery buffer's byte threshold makes the content block
|
|
# downstream-visible before the failure, so its close prelude must escape.
|
|
stream = _FailingStream(
|
|
[_chunk(content="x" * 65_536)],
|
|
RuntimeError("connection lost after commit"),
|
|
)
|
|
emitted: deque[str] = deque()
|
|
|
|
with (
|
|
patch.object(
|
|
provider._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=stream,
|
|
),
|
|
pytest.raises(ExecutionFailure) as exc_info,
|
|
):
|
|
async for event in provider.stream_response(
|
|
request,
|
|
request_id="req_committed_failure",
|
|
):
|
|
emitted.append(event)
|
|
|
|
events = parse_sse_text("".join(emitted))
|
|
assert [event.event for event in events][-1] == "content_block_stop"
|
|
assert not any(event.event in {"error", "message_stop"} for event in events)
|
|
assert exc_info.value.kind is FailureKind.UPSTREAM
|
|
assert exc_info.value.status_code == 502
|
|
assert exc_info.value.retryable is False
|
|
assert "connection lost after commit" in exc_info.value.message
|
|
assert "Request ID: req_committed_failure" in exc_info.value.message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_openai_stream_close_failure_cannot_mask_execution_failure() -> None:
|
|
provider = _provider()
|
|
request = make_messages_request(
|
|
"test-model",
|
|
messages=[],
|
|
max_tokens=32,
|
|
)
|
|
stream = _FailingStream(
|
|
[],
|
|
RuntimeError("original provider failure"),
|
|
close_error=RuntimeError("cleanup api_key=SECRET"),
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
provider._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=stream,
|
|
),
|
|
patch("free_claude_code.providers.http.trace_event") as trace_event,
|
|
pytest.raises(ExecutionFailure) as exc_info,
|
|
):
|
|
[
|
|
event
|
|
async for event in provider.stream_response(
|
|
request,
|
|
request_id="req_close_failure",
|
|
)
|
|
]
|
|
|
|
assert stream.close_calls == 1
|
|
assert exc_info.value.status_code == 502
|
|
assert "original provider failure" in exc_info.value.message
|
|
assert "cleanup" not in exc_info.value.message
|
|
assert "SECRET" not in exc_info.value.message
|
|
trace_event.assert_called_once_with(
|
|
stage="provider",
|
|
event="provider.stream.close_failed",
|
|
source="provider",
|
|
provider="NIM",
|
|
request_id="req_close_failure",
|
|
close_exc_type="RuntimeError",
|
|
preserved_exc_type="ExecutionFailure",
|
|
)
|
|
assert "SECRET" not in repr(trace_event.call_args)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_close_failure_without_active_error_is_observability_only() -> (
|
|
None
|
|
):
|
|
stream = _FailingStream(
|
|
[],
|
|
RuntimeError("unused"),
|
|
close_error=RuntimeError("normal close failed"),
|
|
)
|
|
|
|
with patch("free_claude_code.providers.http.trace_event") as trace_event:
|
|
await close_provider_stream(
|
|
stream,
|
|
active_error=None,
|
|
provider_name="TEST",
|
|
request_id="req_normal_close",
|
|
)
|
|
|
|
assert stream.close_calls == 1
|
|
trace_event.assert_called_once_with(
|
|
stage="provider",
|
|
event="provider.stream.close_failed",
|
|
source="provider",
|
|
provider="TEST",
|
|
request_id="req_normal_close",
|
|
close_exc_type="RuntimeError",
|
|
preserved_exc_type=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_stream_close_failure_preserves_success_lifecycle() -> None:
|
|
provider = _provider()
|
|
request = make_messages_request(
|
|
"test-model",
|
|
messages=[],
|
|
max_tokens=32,
|
|
)
|
|
stream = _FailingStream(
|
|
[_chunk(content="complete", finish_reason="stop")],
|
|
None,
|
|
close_error=RuntimeError("cleanup api_key=SECRET"),
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
provider._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=stream,
|
|
),
|
|
patch("free_claude_code.providers.http.trace_event") as trace_event,
|
|
):
|
|
emitted = [
|
|
event
|
|
async for event in provider.stream_response(
|
|
request,
|
|
request_id="req_successful_close_failure",
|
|
)
|
|
]
|
|
|
|
events = parse_sse_text("".join(emitted))
|
|
assert events[-1].event == "message_stop"
|
|
assert sum(event.event == "message_stop" for event in events) == 1
|
|
assert not any(event.event == "error" for event in events)
|
|
assert stream.close_calls == 1
|
|
trace_event.assert_called_once_with(
|
|
stage="provider",
|
|
event="provider.stream.close_failed",
|
|
source="provider",
|
|
provider="NIM",
|
|
request_id="req_successful_close_failure",
|
|
close_exc_type="RuntimeError",
|
|
preserved_exc_type=None,
|
|
)
|
|
assert "SECRET" not in repr(trace_event.call_args)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_closing_public_openai_stream_closes_raw_stream_once() -> None:
|
|
provider = _provider()
|
|
request = make_messages_request(
|
|
"test-model",
|
|
messages=[],
|
|
max_tokens=32,
|
|
)
|
|
raw_stream = _FailingStream(
|
|
[
|
|
_chunk(content="x" * 65_536),
|
|
_chunk(content="done", finish_reason="stop"),
|
|
],
|
|
None,
|
|
)
|
|
|
|
with patch.object(
|
|
provider._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=raw_stream,
|
|
):
|
|
stream = provider.stream_response(request, request_id="req_early_close")
|
|
await anext(stream)
|
|
assert isinstance(stream, AsyncCloseable)
|
|
await stream.aclose()
|
|
|
|
assert raw_stream.close_calls == 1
|