Update bot/youtube/resume.py

bot/youtube/resume.py

@@ -1,7 +1,11 @@
 # PATH: bot/youtube/resume.py
 import os
+import re
+import asyncio
 from typing import Callable, Awaitable
+
 import httpx
+
 from bot.core.settings import YOUTUBE_CHUNK_SIZE
 
 ProgressCB = Callable[[int, int], Awaitable[None]]
@@ -9,6 +13,20 @@ ProgressCB = Callable[[int, int], Awaitable[None]]
 YOUTUBE_RESUMABLE_INIT = "https://www.googleapis.com/upload/youtube/v3/videos"
 
 
+def _parse_range_last_byte(headers: dict) -> int | None:
+    """
+    Range header examples:
+      - "bytes=0-1048575"
+      - "0-1048575"
+    Return last received byte index, or None if missing/unparseable.
+    """
+    v = headers.get("Range") or headers.get("range") or ""
+    m = re.search(r"(\d+)\s*-\s*(\d+)", v)
+    if not m:
+        return None
+    return int(m.group(2))
+
+
 async def start_resumable_session(access_token: str, metadata: dict) -> str:
     """
     Returns upload_url (Location header)
@@ -19,7 +37,6 @@ async def start_resumable_session(access_token: str, metadata: dict) -> str:
         "X-Upload-Content-Type": "video/*",
     }
     params = {"uploadType": "resumable", "part": "snippet,status"}
-
     async with httpx.AsyncClient(timeout=60, follow_redirects=True) as c:
         r = await c.post(YOUTUBE_RESUMABLE_INIT, params=params, headers=headers, json=metadata)
         if r.status_code >= 400:
@@ -30,12 +47,43 @@ async def start_resumable_session(access_token: str, metadata: dict) -> str:
         return upload_url
 
 
+async def _query_upload_status(
+    c: httpx.AsyncClient,
+    upload_url: str,
+    headers_base: dict,
+    total: int,
+) -> tuple[int, dict | None]:
+    """
+    Asks server what it has already received.
+    Returns (last_received_byte, final_json_if_completed)
+    """
+    headers = dict(headers_base)
+    headers["Content-Length"] = "0"
+    headers["Content-Range"] = f"bytes */{total}"
+
+    r = await c.put(upload_url, headers=headers)
+
+    if r.status_code == 308:
+        last = _parse_range_last_byte(r.headers)
+        if last is None:
+            return -1, None
+        return last, None
+
+    if r.status_code in (200, 201):
+        j = r.json() if r.text else {}
+        return total - 1, j
+
+    if r.status_code >= 400:
+        raise RuntimeError(f"resumable_status_failed:{r.status_code}:{r.text[:200]}")
+
+    return -1, None
+
+
 async def upload_resumable(
     upload_url: str,
     file_path: str,
     access_token: str,
-    *,
-    progress_cb: ProgressCB | None = None,   # ✅ keyword-only (prevents “multiple values”)
+    progress_cb: ProgressCB | None = None,
 ) -> dict:
     """
     Upload file in chunks to upload_url.
@@ -46,38 +94,106 @@ async def upload_resumable(
 
     headers_base = {
         "Authorization": f"Bearer {access_token}",
+        "Content-Type": "video/*",
     }
 
-    async with httpx.AsyncClient(timeout=120, follow_redirects=True) as c:
+    retryable = {429, 500, 502, 503, 504}
+    max_retries = 8
+
+    timeout = httpx.Timeout(connect=30.0, read=300.0, write=300.0, pool=30.0)
+
+    async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as c:
         with open(file_path, "rb") as f:
             while sent < total:
-                chunk = f.read(YOUTUBE_CHUNK_SIZE)
+                # Always seek to where server says it is
+                f.seek(sent)
+
+                to_read = min(YOUTUBE_CHUNK_SIZE, total - sent)
+                chunk = f.read(to_read)
                 if not chunk:
                     break
 
                 start = sent
-                end = sent + len(chunk) - 1
+                end = start + len(chunk) - 1
 
                 headers = dict(headers_base)
                 headers["Content-Length"] = str(len(chunk))
-                headers["Content-Type"] = "video/*"
                 headers["Content-Range"] = f"bytes {start}-{end}/{total}"
 
-                r = await c.put(upload_url, headers=headers, content=chunk)
-
-                # 308 Resume Incomplete is normal for chunked upload
-                if r.status_code == 308:
-                    sent = end + 1
-                    if progress_cb:
-                        await progress_cb(sent, total)
-                    continue
-
-                if r.status_code >= 400:
-                    raise RuntimeError(f"resumable_put_failed:{r.status_code}:{r.text[:200]}")
-
-                j = r.json()
+                attempt = 0
+                while True:
+                    try:
+                        r = await c.put(upload_url, headers=headers, content=chunk)
+                    except (httpx.TimeoutException, httpx.NetworkError):
+                        # Network issue: ask server what it has, then resume
+                        last, done_json = await _query_upload_status(c, upload_url, headers_base, total)
+                        if done_json is not None:
+                            if progress_cb:
+                                await progress_cb(total, total)
+                            return done_json
+                        sent = max(sent, last + 1)
+                        if progress_cb:
+                            await progress_cb(sent, total)
+                        break  # go to next outer loop iteration
+
+                    # Normal: 308 means "keep going" but must read Range
+                    if r.status_code == 308:
+                        last = _parse_range_last_byte(r.headers)
+                        if last is None:
+                            # If Range missing, do an explicit status query
+                            last, done_json = await _query_upload_status(c, upload_url, headers_base, total)
+                            if done_json is not None:
+                                if progress_cb:
+                                    await progress_cb(total, total)
+                                return done_json
+                        sent = max(sent, last + 1)
+                        if progress_cb:
+                            await progress_cb(sent, total)
+                        break  # next outer loop chunk (from sent)
+
+                    # Success: final response JSON
+                    if r.status_code in (200, 201):
+                        j = r.json() if r.text else {}
+                        if progress_cb:
+                            await progress_cb(total, total)
+                        return j
+
+                    # Retryable server responses
+                    if r.status_code in retryable and attempt < max_retries:
+                        attempt += 1
+                        # Before retrying, query status so we don't duplicate wrong chunk
+                        last, done_json = await _query_upload_status(c, upload_url, headers_base, total)
+                        if done_json is not None:
+                            if progress_cb:
+                                await progress_cb(total, total)
+                            return done_json
+                        sent = max(sent, last + 1)
+                        if progress_cb:
+                            await progress_cb(sent, total)
+                        await asyncio.sleep(min(2 ** attempt, 20))
+                        # break to outer loop (rebuild chunk from new sent)
+                        break
+
+                    # Non-retryable error
+                    if r.status_code >= 400:
+                        raise RuntimeError(f"resumable_put_failed:{r.status_code}:{r.text[:200]}")
+
+                    # Unexpected status: treat like error
+                    raise RuntimeError(f"resumable_put_unexpected:{r.status_code}:{r.text[:200]}")
+
+            # If loop ended, do a final status check instead of blindly erroring
+            last, done_json = await _query_upload_status(c, upload_url, headers_base, total)
+            if done_json is not None:
                 if progress_cb:
                     await progress_cb(total, total)
-                return j
+                return done_json
+
+            if last >= total - 1:
+                # Server claims full data but didn't return JSON (rare), try one more status query
+                last2, done_json2 = await _query_upload_status(c, upload_url, headers_base, total)
+                if done_json2 is not None:
+                    if progress_cb:
+                        await progress_cb(total, total)
+                    return done_json2
 
     raise RuntimeError("resumable_upload_incomplete")