Update app/routers/screenshot/screenshot.py

app/routers/screenshot/screenshot.py

@@ -2,13 +2,14 @@ from __future__ import annotations
 
 import logging
 
-from core.service.playwright import AsyncPlaywrightContext, TimeoutError
+from core.service.playwright import PlaywrightInstance, TimeoutError
 from core.service.playwright.models import PageModel, ScreenshotModel
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import Response
 
 router = APIRouter()
-playwright_context = AsyncPlaywrightContext()
+
+playwright_instance = PlaywrightInstance()
 
 
 @router.post("/screenshot", tags=["V1"])
@@ -17,14 +18,19 @@ async def screenshot(
     page_model: PageModel | None = None,
 ) -> Response:
     page_model = page_model if page_model else PageModel()
-    async with playwright_context as sc:
-        try:
-            response = await sc.screenshot(screenshot_model=screenshot_model, page_model=page_model)
-            return Response(content=response, media_type="image/jpeg")
-        except TimeoutError as e:
-            status_code = 504 if isinstance(e, TimeoutError) else 400
-            detail = f"An error occurred while generating the screenshot: {e}"
-            raise HTTPException(status_code=status_code, detail=detail) from e
-        except Exception as e:
-            logging.exception("Unhandled error in screenshot generation")
-            raise HTTPException(status_code=500, detail="An unexpected error occurred") from e
+
+    await playwright_instance
+
+    try:
+        response = await playwright_instance.screenshot(
+            screenshot_model=screenshot_model, page_model=page_model,
+        )
+        return Response(content=response, media_type="image/png")
+    except TimeoutError as e:
+        status_code = 504 if isinstance(e, TimeoutError) else 400
+        detail = f"An error occurred while generating the screenshot: {e}"
+        raise HTTPException(status_code=status_code, detail=detail) from e
+    except Exception as e:
+        logging.exception("Unhandled error in screenshot generation")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred") from e
+