Update bot/youtube/resume.py

bot/youtube/resume.py

@@ -1,79 +1,118 @@
 # PATH: bot/youtube/resume.py
 import os
-from typing import Callable, Awaitable
 import httpx
-from bot.core.settings import YOUTUBE_CHUNK_SIZE
-
-ProgressCB = Callable[[int, int], Awaitable[None]]
+import json
+import asyncio
+from typing import Callable, Awaitable
 
-YOUTUBE_RESUMABLE_INIT = "https://www.googleapis.com/upload/youtube/v3/videos"
+# Standard chunk size (multiple of 256KB). 10MB is a good balance for RAM/Speed.
+CHUNK_SIZE = 10 * 1024 * 1024 
 
 async def start_resumable_session(access_token: str, metadata: dict) -> str:
     """
-    Returns upload_url (Location header)
+    Initializes the upload and gets the unique Upload URL.
     """
+    url = "https://www.googleapis.com/upload/youtube/v3/videos?uploadType=resumable&part=snippet,status"
     headers = {
         "Authorization": f"Bearer {access_token}",
         "Content-Type": "application/json; charset=UTF-8",
-        "X-Upload-Content-Type": "video/*",
+        "X-Upload-Content-Type": "video/mp4",
+        "X-Upload-Content-Length": "0" # Initial draft is 0
     }
-    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")
+    
+    async with httpx.AsyncClient() as client:
+        r = await client.post(url, headers=headers, json=metadata, timeout=30.0)
+        
+        if r.status_code != 200:
+            raise RuntimeError(f"Initiate failed: {r.status_code} - {r.text}")
+            
+        upload_url = r.headers.get("Location")
         if not upload_url:
-            raise RuntimeError("resumable_init_no_location")
+            raise RuntimeError("No upload URL returned from YouTube.")
+            
         return upload_url
 
 async def upload_resumable(
-    upload_url: str,
-    file_path: str,
-    access_token: str,
-    progress_cb: ProgressCB | None = None,
+    upload_url: str, 
+    file_path: str, 
+    access_token: str, 
+    progress_cb: Callable[[int, int], Awaitable[None]] = None
 ) -> dict:
     """
-    Upload file in chunks to upload_url.
-    Returns YouTube API JSON (contains id).
+    Uploads the file chunks. Handles 308 Resume Incomplete robustly.
     """
-    total = os.path.getsize(file_path)
-    sent = 0
-
-    headers_base = {
-        "Authorization": f"Bearer {access_token}",
-    }
-
-    async with httpx.AsyncClient(timeout=120, follow_redirects=True) as c:
+    file_size = os.path.getsize(file_path)
+    
+    # Create a persistent client
+    async with httpx.AsyncClient(timeout=120.0) as client:
         with open(file_path, "rb") as f:
-            while sent < total:
-                chunk = f.read(YOUTUBE_CHUNK_SIZE)
-                if not chunk:
-                    break
-                start = sent
-                end = sent + len(chunk) - 1
+            while True:
+                # 1. Query status (Ask YouTube: "How much have you got?")
+                # We do this by sending an empty PUT with Content-Range: bytes */total
+                r_status = await client.put(
+                    upload_url, 
+                    headers={"Content-Range": f"bytes */{file_size}"},
+                    content=None
+                )
+                
+                # 2. Determine Offset
+                if r_status.status_code == 308:
+                    range_header = r_status.headers.get("Range")
+                    if range_header:
+                        # Header format: "bytes=0-999" -> next byte is 1000
+                        next_byte = int(range_header.split("-")[1]) + 1
+                    else:
+                        # No bytes uploaded yet
+                        next_byte = 0
+                elif r_status.status_code in [200, 201]:
+                    # Already done!
+                    if progress_cb: await progress_cb(file_size, file_size)
+                    return r_status.json()
+                else:
+                    # Some other error (4xx/5xx)
+                    # Try to start from 0 if status failed, or raise error
+                    next_byte = 0
 
-                headers = dict(headers_base)
-                headers["Content-Length"] = str(len(chunk))
-                headers["Content-Type"] = "video/*"
-                headers["Content-Range"] = f"bytes {start}-{end}/{total}"
+                if next_byte >= file_size:
+                    break # Should be handled by 200/201 check, but safety first
 
-                r = await c.put(upload_url, headers=headers, content=chunk)
+                # 3. Read Chunk
+                f.seek(next_byte)
+                chunk = f.read(CHUNK_SIZE)
+                chunk_len = len(chunk)
+                
+                if chunk_len == 0:
+                    break
 
-                # 308 Resume Incomplete is normal for chunked upload
-                if r.status_code == 308:
-                    sent = end + 1
-                    if progress_cb:
-                        await progress_cb(sent, total)
+                # 4. Upload Chunk
+                content_range = f"bytes {next_byte}-{next_byte + chunk_len - 1}/{file_size}"
+                headers = {
+                    "Content-Length": str(chunk_len),
+                    "Content-Range": content_range
+                }
+                
+                try:
+                    r_upload = await client.put(upload_url, headers=headers, content=chunk)
+                    
+                    if r_upload.status_code in [200, 201]:
+                        # Upload Complete!
+                        if progress_cb: await progress_cb(file_size, file_size)
+                        return r_upload.json()
+                    
+                    elif r_upload.status_code == 308:
+                        # Chunk received, loop again to send next
+                        if progress_cb: await progress_cb(next_byte + chunk_len, file_size)
+                        continue
+                        
+                    else:
+                        # Retryable error? Wait and loop (the status check at top will fix offset)
+                        print(f"Chunk error {r_upload.status_code}, retrying...")
+                        await asyncio.sleep(5)
+                        continue
+                        
+                except Exception as e:
+                    print(f"Network error: {e}, retrying...")
+                    await asyncio.sleep(5)
                     continue
 
-                if r.status_code >= 400:
-                    raise RuntimeError(f"resumable_put_failed:{r.status_code}:{r.text[:200]}")
-
-                # success: final response JSON
-                j = r.json()
-                if progress_cb:
-                    await progress_cb(total, total)
-                return j
-
-    raise RuntimeError("resumable_upload_incomplete")
+    raise RuntimeError("resumable_upload_incomplete: Loop finished without 200 OK")