fix: keep agent streams alive during long tasks

SovereignAI-Starter/app/main.py

@@ -48,6 +48,7 @@ MAX_PDF_ATTACHMENT_BYTES = 8 * 1024 * 1024
 MAX_ATTACHMENT_TOTAL_BYTES = 16 * 1024 * 1024
 MAX_OCR_CONTEXT_CHARS = 24_000
 MAX_AGENT_TOOL_CALLS = 3
+AGENT_STREAM_HEARTBEAT_SECONDS = 15.0
 MAX_AGENT_KNOWLEDGE_CHUNKS = 4
 READ_ONLY_AGENT_TOOLS = frozenset({"calculator", "search_workspace", "search_knowledge"})
 
@@ -806,7 +807,9 @@ async def index_workspace_knowledge(
                 raise HTTPException(status_code=413, detail="حد الفهرسة 2 ميغابايت لكل طلب.")
             is_pdf = path.suffix.lower() == ".pdf"
             if is_pdf:
-                parsed_pdf = extract_pdf_pages_text(raw, path.name)
+                parsed_pdf = await asyncio.to_thread(
+                    extract_pdf_pages_text, raw, path.name
+                )
                 pdf_truncated = parsed_pdf["truncated"]
                 pdf_has_text_pages = any(page["has_text"] for page in parsed_pdf["pages"])
                 for page in parsed_pdf["pages"]:
@@ -833,14 +836,16 @@ async def index_workspace_knowledge(
             try:
                 selected_scanned_pages = scanned_page_numbers[:MAX_SCANNED_PAGES]
                 ocr_truncated = len(scanned_page_numbers) > len(selected_scanned_pages)
-                rendered = render_scanned_pdf_pages(
+                rendered = await asyncio.to_thread(
+                    render_scanned_pdf_pages,
                     raw,
                     path.name,
                     max_pages=MAX_SCANNED_PAGES,
                     page_numbers=selected_scanned_pages,
                 )
                 for page in rendered["pages"]:
-                    recognized = recognize_image_text(
+                    recognized = await asyncio.to_thread(
+                        recognize_image_text,
                         base64.b64decode(page["data"]),
                         label=f"الصفحة {page['page']} من {path.name}",
                     )
@@ -2362,6 +2367,7 @@ async def run_agent_stream(
         task = asyncio.create_task(
             _execute_agent(request, report_progress=report, user_id=user_id)
         )
+        last_output_at = perf_counter()
         try:
             while True:
                 if task.done() and queue.empty():
@@ -2369,8 +2375,14 @@ async def run_agent_stream(
                 try:
                     message = await asyncio.wait_for(queue.get(), timeout=0.25)
                 except TimeoutError:
+                    if perf_counter() - last_output_at >= AGENT_STREAM_HEARTBEAT_SECONDS:
+                        # Keep otherwise-idle SSE connections alive while the model
+                        # or a local tool is taking several minutes to finish.
+                        yield ": keep-alive\n\n"
+                        last_output_at = perf_counter()
                     continue
                 yield f"event: progress\ndata: {json.dumps({'message': message}, ensure_ascii=False)}\n\n"
+                last_output_at = perf_counter()
             try:
                 result = await task
             except HTTPException as exc: