Update main.py

main.py

@@ -1,6 +1,7 @@
 import os
 import math
 import logging
+from contextlib import asynccontextmanager  # <--- NEW IMPORT
 from pyrogram import Client
 from pyrogram.session import Session, Auth
 from pyrogram import raw
@@ -10,17 +11,12 @@ from fastapi.responses import StreamingResponse
 import uvicorn
 
 # --- CONFIGURATION ---
-# These are loaded from the Secrets you set in Step 2
 API_ID = int(os.environ.get("API_ID"))
 API_HASH = os.environ.get("API_HASH")
 BOT_TOKEN = os.environ.get("BOT_TOKEN")
 STORAGE_CHANNEL = int(os.environ.get("STORAGE_CHANNEL"))
 
-app = FastAPI()
-
-# --- BOT CLIENT (No Login Required) ---
-# We use a Bot Token, so no phone number login is needed.
-# 'no_updates=True' makes it lighter (it won't read chat messages, only files).
+# --- BOT CLIENT ---
 client = Client(
     "worker_bot",
     api_id=API_ID,
@@ -30,13 +26,29 @@ client = Client(
     no_updates=True 
 )
 
-@app.on_event("startup")
-async def start_client():
+# --- 🚀 NEW LIFESPAN LOGIC ---
+# This replaces @app.on_event("startup")
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     print("⏳ Connecting to Telegram...")
-    await client.start()
-    print("✅ Worker Started & Connected!")
+    try:
+        await client.start()
+        print("✅ Worker Started & Connected!")
+    except Exception as e:
+        print(f"❌ Failed to connect: {e}")
+    
+    yield  # The application runs here
+    
+    print("🛑 Stopping Worker...")
+    try:
+        await client.stop()
+    except:
+        pass
+
+# Initialize FastAPI with the lifespan handler
+app = FastAPI(lifespan=lifespan)
 
-# --- STREAMING LOGIC ---
+# --- STREAMING LOGIC (Unchanged) ---
 class ByteStreamer:
     def __init__(self, client: Client):
         self.client = client
@@ -44,7 +56,6 @@ class ByteStreamer:
     async def yield_file(self, file_id, offset, first_part_cut, last_part_cut, part_count, chunk_size):
         client = self.client
         
-        # 1. Create a direct Media Session (High Performance)
         ms = client.media_sessions.get(file_id.dc_id)
         if not ms:
             if file_id.dc_id != await client.storage.dc_id():
@@ -55,7 +66,6 @@ class ByteStreamer:
                 ms = client.session
             client.media_sessions[file_id.dc_id] = ms
         
-        # 2. Define File Location
         loc = raw.types.InputDocumentFileLocation(
             id=file_id.media_id,
             access_hash=file_id.access_hash,
@@ -63,7 +73,6 @@ class ByteStreamer:
             thumb_size=file_id.thumbnail_size
         )
         
-        # 3. The Loop (Fetch Chunks)
         curr = 1
         while curr <= part_count:
             try:
@@ -75,7 +84,6 @@ class ByteStreamer:
                     chunk = r.bytes
                     if not chunk: break
                     
-                    # Cut the first/last chunks if needed for seeking
                     if part_count == 1:
                         yield chunk[first_part_cut:last_part_cut]
                     elif curr == 1:
@@ -95,8 +103,6 @@ class ByteStreamer:
 @app.get("/stream/{message_id}/{filename}")
 async def stream_handler(req: Request, message_id: int, filename: str):
     try:
-        # 1. Fetch File Info directly from Channel (No Database needed)
-        # We trust the Message ID passed in the URL
         msg = await client.get_messages(STORAGE_CHANNEL, message_id)
         media = msg.document or msg.video or msg.audio
         
@@ -106,8 +112,6 @@ async def stream_handler(req: Request, message_id: int, filename: str):
         file_id = FileId.decode(media.file_id)
         file_size = media.file_size
         
-        # 2. Handle Video Seeking (Range Headers)
-        # This allows the user to jump forward/backward in the video
         range_header = req.headers.get("Range", 0)
         from_bytes, until_bytes = 0, file_size - 1
         if range_header:
@@ -116,15 +120,13 @@ async def stream_handler(req: Request, message_id: int, filename: str):
             if s[1]: until_bytes = int(s[1])
             
         req_len = until_bytes - from_bytes + 1
-        chunk_size = 1048576 # 1 MB Chunks
+        chunk_size = 1048576 
         
-        # Math for calculating where to start downloading
         offset = (from_bytes // chunk_size) * chunk_size
         first_part_cut = from_bytes - offset
         last_part_cut = (until_bytes % chunk_size) + 1
         part_count = math.ceil(req_len / chunk_size)
         
-        # 3. Start Streaming
         streamer = ByteStreamer(client)
         body = streamer.yield_file(file_id, offset, first_part_cut, last_part_cut, part_count, chunk_size)
         
@@ -143,6 +145,4 @@ async def stream_handler(req: Request, message_id: int, filename: str):
         raise HTTPException(404, "File Not Found")
 
 if __name__ == "__main__":
-    # Hugging Face Spaces expect the app to run on port 7860
     uvicorn.run(app, host="0.0.0.0", port=7860)
-