Update bot/youtube/resume.py

bot/youtube/resume.py

@@ -1,8 +1,10 @@
 # PATH: bot/youtube/resume.py
+from __future__ import annotations
+
+import asyncio
 import os
 import re
-import asyncio
-from typing import Callable, Awaitable
+from typing import Awaitable, Callable, Optional, Tuple
 
 import httpx
 
@@ -13,7 +15,7 @@ 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:
+def _parse_range_last_byte(headers: dict) -> Optional[int]:
     """
     Range header examples:
       - "bytes=0-1048575"
@@ -27,24 +29,50 @@ def _parse_range_last_byte(headers: dict) -> int | None:
     return int(m.group(2))
 
 
-async def start_resumable_session(access_token: str, metadata: dict) -> str:
+async def start_resumable_session(
+    *,
+    access_token: str,
+    metadata: dict,
+    total_bytes: int = 0,
+) -> dict:
     """
-    Returns upload_url (Location header)
+    Starts a YouTube resumable upload session.
+
+    Returns:
+      {"ok": True, "session_url": "..."}
+      {"ok": False, "err": "...", "detail": "..."}
     """
+    if not access_token:
+        return {"ok": False, "err": "missing_access_token"}
+
     headers = {
         "Authorization": f"Bearer {access_token}",
         "Content-Type": "application/json; charset=UTF-8",
         "X-Upload-Content-Type": "video/*",
     }
+    # Helps YouTube know size (recommended)
+    if isinstance(total_bytes, int) and total_bytes > 0:
+        headers["X-Upload-Content-Length"] = str(total_bytes)
+
     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:
-            raise RuntimeError(f"resumable_init_failed:{r.status_code}:{r.text[:200]}")
-        upload_url = r.headers.get("Location") or r.headers.get("location")
-        if not upload_url:
-            raise RuntimeError("resumable_init_no_location")
-        return upload_url
+
+    timeout = httpx.Timeout(connect=30.0, read=60.0, write=60.0, pool=30.0)
+
+    try:
+        async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as c:
+            r = await c.post(YOUTUBE_RESUMABLE_INIT, params=params, headers=headers, json=metadata)
+            if r.status_code >= 400:
+                return {
+                    "ok": False,
+                    "err": f"resumable_init_failed:{r.status_code}",
+                    "detail": (r.text or "")[:600],
+                }
+            upload_url = r.headers.get("Location") or r.headers.get("location")
+            if not upload_url:
+                return {"ok": False, "err": "resumable_init_no_location", "detail": "Missing Location header"}
+            return {"ok": True, "session_url": upload_url}
+    except Exception as e:
+        return {"ok": False, "err": f"resumable_init_exception:{type(e).__name__}", "detail": str(e)[:600]}
 
 
 async def _query_upload_status(
@@ -52,7 +80,7 @@ async def _query_upload_status(
     upload_url: str,
     headers_base: dict,
     total: int,
-) -> tuple[int, dict | None]:
+) -> Tuple[int, Optional[dict]]:
     """
     Asks server what it has already received.
     Returns (last_received_byte, final_json_if_completed)
@@ -74,22 +102,41 @@ async def _query_upload_status(
         return total - 1, j
 
     if r.status_code >= 400:
-        raise RuntimeError(f"resumable_status_failed:{r.status_code}:{r.text[:200]}")
+        raise RuntimeError(f"resumable_status_failed:{r.status_code}:{(r.text or '')[:200]}")
 
     return -1, None
 
 
 async def upload_resumable(
-    upload_url: str,
+    *,
+    session_url: str,
     file_path: str,
     access_token: str,
-    progress_cb: ProgressCB | None = None,
+    total_bytes: int = 0,
+    progress_cb: Optional[ProgressCB] = None,
 ) -> dict:
     """
-    Upload file in chunks to upload_url.
-    Returns YouTube API JSON (contains id).
+    Upload file in chunks to session_url.
+
+    Returns:
+      {"ok": True, "response": {...}}
+      {"ok": False, "err": "...", "detail": "..."}
     """
-    total = os.path.getsize(file_path)
+    if not session_url:
+        return {"ok": False, "err": "missing_session_url"}
+    if not access_token:
+        return {"ok": False, "err": "missing_access_token"}
+    if not file_path or not os.path.exists(file_path):
+        return {"ok": False, "err": "file_not_found"}
+
+    try:
+        total = int(total_bytes) if total_bytes else int(os.path.getsize(file_path))
+    except Exception:
+        total = 0
+
+    if total <= 0:
+        return {"ok": False, "err": "bad_total_bytes"}
+
     sent = 0
 
     headers_base = {
@@ -102,98 +149,99 @@ async def upload_resumable(
 
     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:
-                # 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 = start + len(chunk) - 1
-
-                headers = dict(headers_base)
-                headers["Content-Length"] = str(len(chunk))
-                headers["Content-Range"] = f"bytes {start}-{end}/{total}"
-
-                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)
+    try:
+        async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as c:
+            with open(file_path, "rb") as f:
+                while sent < total:
+                    f.seek(sent)
+
+                    to_read = min(YOUTUBE_CHUNK_SIZE, total - sent)
+                    chunk = f.read(to_read)
+                    if not chunk:
+                        break
+
+                    start = sent
+                    end = start + len(chunk) - 1
+
+                    headers = dict(headers_base)
+                    headers["Content-Length"] = str(len(chunk))
+                    headers["Content-Range"] = f"bytes {start}-{end}/{total}"
+
+                    attempt = 0
+                    while True:
+                        try:
+                            r = await c.put(session_url, headers=headers, content=chunk)
+                        except (httpx.TimeoutException, httpx.NetworkError):
+                            # Ask server status, then resume
+                            last, done_json = await _query_upload_status(c, session_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:
+                                return {"ok": True, "response": done_json}
+                            sent = max(sent, last + 1)
                             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]}")
+                                await progress_cb(sent, total)
+                            break  # next outer loop
+
+                        if r.status_code == 308:
+                            last = _parse_range_last_byte(r.headers)
+                            if last is None:
+                                last, done_json = await _query_upload_status(c, session_url, headers_base, total)
+                                if done_json is not None:
+                                    if progress_cb:
+                                        await progress_cb(total, total)
+                                    return {"ok": True, "response": done_json}
+                            sent = max(sent, (last or -1) + 1)
+                            if progress_cb:
+                                await progress_cb(sent, total)
+                            break
 
-                    # Unexpected status: treat like error
-                    raise RuntimeError(f"resumable_put_unexpected:{r.status_code}:{r.text[:200]}")
+                        if r.status_code in (200, 201):
+                            j = r.json() if r.text else {}
+                            if progress_cb:
+                                await progress_cb(total, total)
+                            return {"ok": True, "response": j}
 
-            # 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 r.status_code in retryable and attempt < max_retries:
+                            attempt += 1
+                            last, done_json = await _query_upload_status(c, session_url, headers_base, total)
+                            if done_json is not None:
+                                if progress_cb:
+                                    await progress_cb(total, total)
+                                return {"ok": True, "response": done_json}
+                            sent = max(sent, last + 1)
+                            if progress_cb:
+                                await progress_cb(sent, total)
+                            await asyncio.sleep(min(2 ** attempt, 20))
+                            break
+
+                        if r.status_code >= 400:
+                            return {
+                                "ok": False,
+                                "err": f"resumable_put_failed:{r.status_code}",
+                                "detail": (r.text or "")[:600],
+                            }
+
+                        return {
+                            "ok": False,
+                            "err": f"resumable_put_unexpected:{r.status_code}",
+                            "detail": (r.text or "")[:600],
+                        }
+
+            # Final status check
+            last, done_json = await _query_upload_status(c, session_url, headers_base, total)
             if done_json is not None:
                 if progress_cb:
                     await progress_cb(total, total)
-                return done_json
+                return {"ok": True, "response": 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)
+                last2, done_json2 = await _query_upload_status(c, session_url, headers_base, total)
                 if done_json2 is not None:
                     if progress_cb:
                         await progress_cb(total, total)
-                    return done_json2
+                    return {"ok": True, "response": done_json2}
 
-    raise RuntimeError("resumable_upload_incomplete")
+        return {"ok": False, "err": "resumable_upload_incomplete"}
+    except Exception as e:
+        return {"ok": False, "err": f"resumable_exception:{type(e).__name__}", "detail": str(e)[:600]}