From 2326bfcea0fb7fd95ef2c88a4405fee454716ac8 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Mon, 2 Nov 2026 11:47:22 -0800 Subject: [PATCH] feat(prime-traces): raise the client line limit to 1 GiB The SDK rejects any line over MAX_LINE_BYTES before making a request, so episodes above 64 MiB never reach the service. Match the service's new limits: 1 GiB per line and 1100 MiB per request. Co-Authored-By: Claude Opus 5.5 --- packages/prime-traces/src/prime_traces/batching.py | 14 ++++++-------- .../src/prime_traces/core/async_client.py | 2 +- .../prime-traces/src/prime_traces/core/client.py | 4 ++-- packages/prime-traces/tests/test_batching.py | 11 +++++++++++ packages/prime/src/prime_cli/commands/traces.py | 2 +- 5 files changed, 21 insertions(+), 12 deletions(-) diff --git a/packages/prime-traces/src/prime_traces/batching.py b/packages/prime-traces/src/prime_traces/batching.py index 44f60589f..13d716192 100644 --- a/packages/prime-traces/src/prime_traces/batching.py +++ b/packages/prime-traces/src/prime_traces/batching.py @@ -38,20 +38,18 @@ MIB = 1024 * 1024 #: Server-side cap on uncompressed bytes per request (-> upload_too_large). -MAX_BATCH_BYTES = 256 * MIB +MAX_BATCH_BYTES = 1100 * MIB #: Server-side cap on one JSONL line (-> trace_too_large). -MAX_LINE_BYTES = 64 * MIB +MAX_LINE_BYTES = 1024 * MIB #: Server-side cap on staged rows per request (-> too_many_traces_in_upload). #: Exact for bare-trace uploads, where one line is one row. Episode lines #: stage one episode row plus each nested trace, which opaque line bytes #: cannot count — there the byte caps are the practical bound. MAX_BATCH_LINES = 1_000_000 -#: Default chunk-close threshold. Held under Cloud Run's 32 MiB request cap so -#: a batch fits the deployed transport even with compression disabled; the -#: 256 MiB service contract above is not reachable through that cap today. A -#: batch may exceed this only when a single line does — such a batch (or any -#: caller-raised target) still risks transport-level rejection before the -#: service sees it. +#: Default chunk-close threshold. Far below the request cap so that a failed +#: request costs little to resend and a batch stays cheap to hold in memory. A +#: batch may exceed this only when a single line does, and such a batch is held +#: whole in memory for hashing and compression. DEFAULT_TARGET_BATCH_BYTES = 30 * MIB diff --git a/packages/prime-traces/src/prime_traces/core/async_client.py b/packages/prime-traces/src/prime_traces/core/async_client.py index 16ce9a91b..1e308bd18 100644 --- a/packages/prime-traces/src/prime_traces/core/async_client.py +++ b/packages/prime-traces/src/prime_traces/core/async_client.py @@ -172,7 +172,7 @@ async def delete(self, endpoint: str, params: Optional[Dict[str, Any]] = None) - async def stream_bytes( self, endpoint: str, params: Optional[Dict[str, Any]] = None ) -> AsyncIterator[bytes]: - """Stream a response body in chunks (a raw trace can be 64 MiB).""" + """Stream a response body in chunks (a raw trace can be 1 GiB).""" self._check_auth() for attempt in range(IDEMPOTENT_RETRY_ATTEMPTS): last = attempt == IDEMPOTENT_RETRY_ATTEMPTS - 1 diff --git a/packages/prime-traces/src/prime_traces/core/client.py b/packages/prime-traces/src/prime_traces/core/client.py index d76ec80ab..af31ebb92 100644 --- a/packages/prime-traces/src/prime_traces/core/client.py +++ b/packages/prime-traces/src/prime_traces/core/client.py @@ -233,7 +233,7 @@ def raise_for_response(response: httpx.Response) -> None: raise APIError(f"HTTP {status}: {message}", status_code=status, code=code) -#: Generous read/write budget: one request body can be 256 MiB. +#: Generous read/write budget: one request body can be 1100 MiB. DEFAULT_TIMEOUT = httpx.Timeout(120.0, connect=10.0) @@ -473,7 +473,7 @@ def delete(self, endpoint: str, params: Optional[Dict[str, Any]] = None) -> None def stream_bytes( self, endpoint: str, params: Optional[Dict[str, Any]] = None ) -> Iterator[bytes]: - """Stream a response body in chunks (a raw trace can be 64 MiB). + """Stream a response body in chunks (a raw trace can be 1 GiB). Transient failures are retried only until the first body byte has been yielded: a mid-stream retry would silently restart the body under the diff --git a/packages/prime-traces/tests/test_batching.py b/packages/prime-traces/tests/test_batching.py index 6640f4d44..3b522bb2b 100644 --- a/packages/prime-traces/tests/test_batching.py +++ b/packages/prime-traces/tests/test_batching.py @@ -6,6 +6,7 @@ from prime_traces import ( MAX_BATCH_BYTES, + MAX_LINE_BYTES, TraceTooLargeError, aiter_batches, iter_batches, @@ -96,6 +97,16 @@ def test_oversized_line_rejected_locally(self): assert exc_info.value.line_number == 2 assert exc_info.value.limit == 32 + def test_default_caps_match_the_service(self): + assert MAX_LINE_BYTES == 1024 * 1024 * 1024 + assert MAX_BATCH_BYTES == 1100 * 1024 * 1024 + + def test_line_over_the_previous_64_mib_cap_is_accepted_by_default(self): + content = b"x" * (64 * 1024 * 1024 + 1) + [batch] = iter_batches([content + b"\n"]) + assert batch.size == len(content) + 1 + assert batch.num_lines == 1 + def test_line_cap_measured_without_terminator(self): # Content is exactly at the cap; the trailing newline must not tip it. content = b'{"p":"' + b"x" * 24 + b'"}' diff --git a/packages/prime/src/prime_cli/commands/traces.py b/packages/prime/src/prime_cli/commands/traces.py index 82c89cacc..69e1ed6a1 100644 --- a/packages/prime/src/prime_cli/commands/traces.py +++ b/packages/prime/src/prime_cli/commands/traces.py @@ -794,7 +794,7 @@ def get_trace( else: console.print(f"[green]Wrote {written} bytes to {escape(str(dest))}[/green]") return - # Raw documents can be tens of MiB. Preserve their exact bytes and + # Raw documents can reach 1 GiB. Preserve their exact bytes and # do not append a newline so redirected output is a faithful copy. stdout = click.get_binary_stream("stdout") stdout.write(client.get_raw(trace_id))