fix(googledrive): reduce upload_from_url chunk buffering (CUS-199)

@zen-agentchecks n/achecks…zen/cherrypick-cus199-upload-chunks-pu14zr → master2 files · +246 −34updated 4w ago

▸Description· 2 comments · 1 noise

Description

Cherry-picks commit 2e484c9 from closed PR #23568 so the CUS-199 follow-up optimization for GOOGLEDRIVE_UPLOAD_FROM_URL can land on master.

This is an incremental optimization on top of the merged base PR #23536 (which switched the action from full-file multipart buffering to a Drive resumable upload). The follow-up tightens the common known-size (Content-Length) path:

Uploads each known-size chunk to Drive as soon as it is read instead of reading one chunk ahead.
Keeps the unknown-size path's one-chunk lookahead so the final Content-Range can still name the resolved total.
Adds tests pinning bounded chunk sizes, Content-Range values, and the no-lookahead behavior for Content-Length sources.

Net effect: lower peak memory (one chunk instead of two) and faster time-to-first-byte for the common Content-Length path.

How did I test this PR

uv run pytest tests/test_apps/test_googledrive/test_upload_from_url.py -q ✅ — 11 passed in 1.34s
uv run --with ruff ruff --config config/ruff.toml check apps/googledrive/actions/upload_from_url.py tests/test_apps/test_googledrive/test_upload_from_url.py ✅ — all checks passed
uv run --with ruff ruff --config config/ruff.toml format --check apps/googledrive/actions/upload_from_url.py tests/test_apps/test_googledrive/test_upload_from_url.py ✅ — 2 files already formatted
make fmt + MYPY_DIFF_ONLY=1 nox -s chk ✅ — all green
Codex review ✅ — 1 P2 finding (boundary-aligned EOF gap on the no-lookahead path) documented in a follow-up comment; not fixed here so the cherry-pick stays faithful to #23568.

E2E: GOOGLEDRIVE_UPLOAD_FROM_URL requires real Drive OAuth + an upload target, so it can only be exercised in prod. The underlying base PR #23536 has been running in prod since merge; CI tests + lint-and-type-checks (both green) cover the regression surface this change introduces.

Pre-existing CI failure: secrets-detection.yml fails on every Mercury PR (including the original #23568 and recent cortex/zen PRs) due to a workflow-file/secret-inheritance issue on master — not caused by this PR. ci_status.py's rollup doesn't see it; visible via gh api repos/ComposioHQ/mercury/actions/runs?workflow_id=225585440.

Cherry-pick of: 2e484c9a961c16af728142259677161b278a041e Original (closed) PR: https://github.com/ComposioHQ/mercury/pull/23568

Triggered by: samvit@composio.dev | Source: mcp Session: https://zen-api-production-4c98.up.railway.app/dashboard/#/chat/zen-3b5e181f25e4

@zen-agent4w ago

Codex review — 1 P2 finding (documented, not fixed in this PR)

[P2] Boundary-aligned EOF check in known-size path — apps/googledrive/actions/upload_from_url.py:421 When the advertised Content-Length is an exact multiple of _UPLOAD_CHUNK_SIZE, the loop returns immediately after PUTting the boundary-aligned chunk without attempting to read past it. If the source over-reads (delivers more bytes than advertised), the end + 1 > total_size guard never runs and Drive is finalized with what may be a silently truncated upload (relative to the source's actual content).

Why not fixed here: This PR is a faithful cherry-pick of #23568, whose explicit goal is to eliminate lookahead on the known-size path. Adding even a one-byte EOF peek would re-introduce a (minimal) lookahead, diverging from the original PR. The condition only fires for inconsistent sources whose advertised length lands on the 256 KiB chunk boundary, so the impact is bounded.

Follow-up: Worth a small follow-up PR after this merges — either a one-byte EOF peek after the final chunk, or a chunked transfer-encoding sentinel check before trusting Content-Length as authoritative.

@zen-agent4w ago

Post-PR lifecycle complete

Step	Result
`/build-checks`	✅ `make fmt` + `MYPY_DIFF_ONLY=1 nox -s chk` both pass
`/codex-review-loop`	✅ 1 iteration, 1 P2 finding (documented above); not auto-fixed to preserve cherry-pick fidelity to #23568
`/gh-fix-ci-loop`	✅ All 14 real CI checks green (CodeQL `Analyze (python)` was the last to land)
Unit tests	✅ `tests/test_apps/test_googledrive/test_upload_from_url.py` — 11/11 pass

▸ 1 bot/status comment hidden

@linear-code[bot]4w ago

Problem

GOOGLEDRIVE_UPLOAD_FROM_URL in composiox-googledrive v20260506_0 OOMs Lambda for files >300MB downloaded from a URL. Reported by Vikram Balaji (AWS) on 2026-05-13 and reproduced with a 1GB test file, hitting 1024/1024 MB max memory and Runtime.OutOfMemory at 159s. Causing an availability drop on AWS's side.

Reporter quote (2026-05-13, Vikram Balaji)

There is a bug on the GOOGLEDRIVE_UPLOAD_FROM_URL ... that is leading to an availability drop on our end. The file_content = download_response.content defeats the stream=True and reads the entire response body into RAM leading to an Out of Memory error for files copied from URL's into GDrive >300Mb.

Additionally the multipart_body also does a header.encode() + file_content + trailer.encode() - further bloating the memory available to the Lambda.

Repro

2026-05-13 07:26 UTC
Source: https://nbg1-speed.hetzner.com/1GB.bin
Max Memory Used: 1024/1024 MB
Runtime.OutOfMemory at 159s

Slack thread

https://composioworkspace.slack.com/archives/C09FSDNS1MF/p1778659668177109

Root cause (verified in repo)

Full-file buffering, twice, in upload_from_url.py:

Lines 141-149 request the download with stream=True (correct).
Line 226 immediately reads the entire body via — fully negates .

loading diff…

Description

Cherry-picks commit 2e484c9 from closed PR #23568 so the CUS-199 follow-up optimization for GOOGLEDRIVE_UPLOAD_FROM_URL can land on master.

Uploads each known-size chunk to Drive as soon as it is read instead of reading one chunk ahead.

Keeps the unknown-size path's one-chunk lookahead so the final Content-Range can still name the resolved total.

Adds tests pinning bounded chunk sizes, Content-Range values, and the no-lookahead behavior for Content-Length sources.

Net effect: lower peak memory (one chunk instead of two) and faster time-to-first-byte for the common Content-Length path.

How did I test this PR

uv run pytest tests/test_apps/test_googledrive/test_upload_from_url.py -q ✅ — 11 passed in 1.34s

uv run --with ruff ruff --config config/ruff.toml check apps/googledrive/actions/upload_from_url.py tests/test_apps/test_googledrive/test_upload_from_url.py

✅ — all checks passed

uv run --with ruff ruff --config config/ruff.toml format --check apps/googledrive/actions/upload_from_url.py tests/test_apps/test_googledrive/test_upload_from_url.py

✅ — 2 files already formatted

make fmt + MYPY_DIFF_ONLY=1 nox -s chk ✅ — all green

Codex review ✅ — 1 P2 finding (boundary-aligned EOF gap on the no-lookahead path) documented in a follow-up comment; not fixed here so the cherry-pick stays faithful to #23568.

Cherry-pick of: 2e484c9a961c16af728142259677161b278a041e Original (closed) PR: https://github.com/ComposioHQ/mercury/pull/23568

Step

Result

/build-checks

✅ make fmt + MYPY_DIFF_ONLY=1 nox -s chk both pass

/codex-review-loop

✅ 1 iteration, 1 P2 finding (documented above); not auto-fixed to preserve cherry-pick fidelity to #23568

/gh-fix-ci-loop

✅ All 14 real CI checks green (CodeQL Analyze (python) was the last to land)

Unit tests

✅ tests/test_apps/test_googledrive/test_upload_from_url.py — 11/11 pass

Reporter quote (2026-05-13, Vikram Balaji)

There is a bug on the GOOGLEDRIVE_UPLOAD_FROM_URL ... that is leading to an availability drop on our end. The file_content = download_response.content defeats the stream=True and reads the entire response body into RAM leading to an Out of Memory error for files copied from URL's into GDrive >300Mb.

Additionally the multipart_body also does a header.encode() + file_content + trailer.encode() - further bloating the memory available to the Lambda.

fix(googledrive): reduce upload_from_url chunk buffering (CUS-199)

Description

How did I test this PR

Codex review — 1 P2 finding (documented, not fixed in this PR)

Post-PR lifecycle complete

Problem

Reporter quote (2026-05-13, Vikram Balaji)

Repro

Slack thread

Root cause (verified in repo)

Description

How did I test this PR

Codex review — 1 P2 finding (documented, not fixed in this PR)

Post-PR lifecycle complete

Problem

Reporter quote (2026-05-13, Vikram Balaji)

Repro

Slack thread

Root cause (verified in repo)

Recommended fix

Additional notes