154 lines
5 KiB
Python
154 lines
5 KiB
Python
"""Recording transports for Spotify's httpx2 and Anthropic's httpx."""
|
|
|
|
from collections.abc import AsyncIterator, Callable
|
|
from typing import cast
|
|
|
|
import httpx
|
|
import httpx2
|
|
|
|
from recording.model import RecordedInteraction
|
|
from recording.redaction import is_sensitive_url
|
|
|
|
|
|
class HttpxRecordingTransport(httpx.AsyncBaseTransport):
|
|
"""Wrap an httpx transport and retain completed response chunk sequences."""
|
|
|
|
def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
|
|
"""Bind the real transport and start an empty in-memory cassette."""
|
|
self.transport = transport
|
|
self.interactions: list[RecordedInteraction] = []
|
|
|
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
|
"""Forward one request and wrap its response stream for capture."""
|
|
request.headers["Accept-Encoding"] = "identity"
|
|
request_body = await request.aread()
|
|
response = await self.transport.handle_async_request(request)
|
|
if is_sensitive_url(str(request.url)):
|
|
return response
|
|
stream = _HttpxRecordingStream(
|
|
cast(httpx.AsyncByteStream, response.stream),
|
|
lambda chunks: self._finish(request, request_body, response.status_code, chunks),
|
|
)
|
|
return httpx.Response(
|
|
response.status_code,
|
|
headers=response.headers,
|
|
stream=stream,
|
|
extensions=response.extensions,
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
"""Close the wrapped transport."""
|
|
await self.transport.aclose()
|
|
|
|
def _finish(
|
|
self,
|
|
request: httpx.Request,
|
|
request_body: bytes,
|
|
status: int,
|
|
chunks: tuple[bytes, ...],
|
|
) -> None:
|
|
self.interactions.append(
|
|
RecordedInteraction(request.method, str(request.url), status, request_body, chunks)
|
|
)
|
|
|
|
|
|
class Httpx2RecordingTransport(httpx2.AsyncBaseTransport):
|
|
"""Wrap an httpx2 transport and retain completed response chunk sequences."""
|
|
|
|
def __init__(self, transport: httpx2.AsyncBaseTransport) -> None:
|
|
"""Bind the real transport and start an empty in-memory cassette."""
|
|
self.transport = transport
|
|
self.interactions: list[RecordedInteraction] = []
|
|
|
|
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
|
|
"""Forward one request and wrap its response stream for capture."""
|
|
request.headers["Accept-Encoding"] = "identity"
|
|
request_body = await request.aread()
|
|
response = await self.transport.handle_async_request(request)
|
|
if is_sensitive_url(str(request.url)):
|
|
return response
|
|
stream = _Httpx2RecordingStream(
|
|
cast(httpx2.AsyncByteStream, response.stream),
|
|
lambda chunks: self._finish(request, request_body, response.status_code, chunks),
|
|
)
|
|
return httpx2.Response(
|
|
response.status_code,
|
|
headers=response.headers,
|
|
stream=stream,
|
|
extensions=response.extensions,
|
|
request=request,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
"""Close the wrapped transport."""
|
|
await self.transport.aclose()
|
|
|
|
def _finish(
|
|
self,
|
|
request: httpx2.Request,
|
|
request_body: bytes,
|
|
status: int,
|
|
chunks: tuple[bytes, ...],
|
|
) -> None:
|
|
self.interactions.append(
|
|
RecordedInteraction(request.method, str(request.url), status, request_body, chunks)
|
|
)
|
|
|
|
|
|
class _HttpxRecordingStream(httpx.AsyncByteStream):
|
|
def __init__(
|
|
self,
|
|
stream: httpx.AsyncByteStream,
|
|
finish: Callable[[tuple[bytes, ...]], None],
|
|
) -> None:
|
|
self.stream = stream
|
|
self.finish = finish
|
|
self.chunks: list[bytes] = []
|
|
self.is_finished = False
|
|
|
|
async def __aiter__(self) -> AsyncIterator[bytes]:
|
|
try:
|
|
async for chunk in self.stream:
|
|
self.chunks.append(chunk)
|
|
yield chunk
|
|
finally:
|
|
self._finish()
|
|
|
|
async def aclose(self) -> None:
|
|
await self.stream.aclose()
|
|
self._finish()
|
|
|
|
def _finish(self) -> None:
|
|
if not self.is_finished:
|
|
self.is_finished = True
|
|
self.finish(tuple(self.chunks))
|
|
|
|
|
|
class _Httpx2RecordingStream(httpx2.AsyncByteStream):
|
|
def __init__(
|
|
self,
|
|
stream: httpx2.AsyncByteStream,
|
|
finish: Callable[[tuple[bytes, ...]], None],
|
|
) -> None:
|
|
self.stream = stream
|
|
self.finish = finish
|
|
self.chunks: list[bytes] = []
|
|
self.is_finished = False
|
|
|
|
async def __aiter__(self) -> AsyncIterator[bytes]:
|
|
try:
|
|
async for chunk in self.stream:
|
|
self.chunks.append(chunk)
|
|
yield chunk
|
|
finally:
|
|
self._finish()
|
|
|
|
async def aclose(self) -> None:
|
|
await self.stream.aclose()
|
|
self._finish()
|
|
|
|
def _finish(self) -> None:
|
|
if not self.is_finished:
|
|
self.is_finished = True
|
|
self.finish(tuple(self.chunks))
|