experiment working, database leaks closed

Author: taylorlynn

app/api/endpoints/analysis_endpoints.py

@@ -1,4 +1,5 @@
 import json
+import asyncio
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
 from typing import List
 from uuid import UUID
@@ -112,24 +113,95 @@ async def stream_claim_analysis_exp(
         claim = await claim_service.get_claim(claim_id=claim_id, user_id=current_user.id)
 
         session = claim_service._claim_repo._session
+        
+        await session.rollback()
 
         async def event_generator():
             try:
                 logger.info(f"Starting analysis stream for claim {claim_id}")
                 yield f"data: {json.dumps({'type': 'status', 'content': 'Initializing analysis...'})}\n\n"
 
-                async for event in analysis_orchestrator.analyze_claim_stream(
+                orchestrator_stream = analysis_orchestrator.analyze_claim_stream(
                     claim=claim, user_id=current_user.id, default=False
-                ):
-                    if isinstance(event, dict):
-                        yield f"data: {json.dumps(event)}\n\n"
+                )
+                # ---------------------------------------------------------
+                # THE HEALTH CHECK LOOP
+                # ---------------------------------------------------------
+                next_event_task = None
+                
+                while True:
+                    # Only create a new task if we don't already have one waiting
+                    if next_event_task is None:
+                        next_event_task = asyncio.create_task(anext(orchestrator_stream))
+
+                    # Wait for the task to finish, but only wait 15 seconds
+                    done, pending = await asyncio.wait(
+                        [next_event_task], 
+                        timeout=15.0, 
+                        return_when=asyncio.FIRST_COMPLETED
+                    )
+                    
+                    if next_event_task in done:
+                        # The LLM yielded a chunk! Let's process it.
+                        try:
+                            event = next_event_task.result()
+                            if isinstance(event, dict):
+                                yield f"data: {json.dumps(event)}\n\n"
+                                
+                            # Reset the task so we grab the next chunk on the next loop
+                            next_event_task = None 
+                            
+                        except StopAsyncIteration:
+                            # The stream finished normally!
+                            break
+                        except Exception as e:
+                            # If the orchestrator crashed, catch it here
+                            raise e
+                            
+                    else:
+                        # The task is in 'pending'. 15 seconds passed, but the LLM is still thinking.
+                        # We yield a heartbeat, but we DO NOT reset next_event_task. 
+                        # It will keep running safely in the background on the next loop!
+                        logger.debug("Stream idle for 15s. Sending health check ping...")
+                        yield ": healthcheck\n\n"
+                        
+                yield "data: [DONE]\n\n"
+
+            except asyncio.CancelledError:
+                logger.warning(f"Client disconnected during stream for claim {claim_id}")
+                raise
+                
+            # async for event in analysis_orchestrator.analyze_claim_stream(
+                #     claim=claim, user_id=current_user.id, default=False
+                # ):
+                #     if isinstance(event, dict):
+                #         yield f"data: {json.dumps(event)}\n\n"
+
+                # yield "data: [DONE]\n\n"
+
+            # except asyncio.CancelledError:
+            #     # THE FIX: The user closed their browser! 
+            #     logger.info(f"Client disconnected during stream for claim {claim_id}")
+            #     await session.rollback()  # Explicitly release the lock!
+            #     raise
 
             except Exception as e:
                 logger.error(f"Error in analysis stream: {str(e)}", exc_info=True)
                 yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
             finally:
-                await session.close()
-                yield "data: [DONE]\n\n"
+                # async def force_cleanup():
+                #     try:
+                #         await session.rollback()
+                #     except Exception as e:
+                #         logger.error(f"Force rollback failed: {e}")
+                #     finally:
+                #         await session.close()
+                
+                # # Fire and forget. FastAPI cannot cancel this!
+                # asyncio.create_task(force_cleanup())
+                if next_event_task and not next_event_task.done():
+                    logger.debug("Cancelling background orchestrator task...")
+                    next_event_task.cancel()
 
         return StreamingResponse(
             event_generator(),
@@ -138,8 +210,8 @@ async def event_generator():
                 "Cache-Control": "no-cache",
                 "Connection": "keep-alive",
                 "X-Accel-Buffering": "no",
-                "Access-Control-Allow-Origin": "*",
-                "Access-Control-Allow-Credentials": "true",
+                # "Access-Control-Allow-Origin": "*",
+                # "Access-Control-Allow-Credentials": "true",
             },
         )
     except Exception as e:

app/main.py

@@ -63,6 +63,7 @@ async def lifespan(app: FastAPI):
         "https://veracity-eval-frontend-git-g2frontend-complex-data-lab.vercel.app",
         "https://veracity-eval-frontend-git-g1frontendmod-complex-data-lab.vercel.app",
         "https://veracity-eval-frontend-git-g1frontendmod2-complex-data-lab.vercel.app",
+        "https://veracity-eval-frontend-git-g0frontendmod-complex-data-lab.vercel.app",
     ],
     allow_credentials=True,
     allow_methods=["*"],

app/models/database/models.py

@@ -145,7 +145,7 @@ class AnalysisModel(Base):
 
     log_probs: Mapped[bytes] = mapped_column(LargeBinary, nullable=True)
 
-    claim: Mapped["ClaimModel"] = relationship(back_populates="analyses", doc="Related claim")
+    claim: Mapped["ClaimModel"] = relationship(back_populates="analyses", doc="Related claim", lazy="selectin")
     searches: Mapped[List["SearchModel"]] = relationship(back_populates="analysis", cascade="all, delete-orphan")
     feedbacks: Mapped[List["FeedbackModel"]] = relationship(
         back_populates="analysis",

app/repositories/implementations/claim_repository.py

@@ -76,6 +76,9 @@ async def update_status(self, claim_id: UUID, status: ClaimStatus) -> Optional[C
 
             claim.status = status
             updated_claim = await self.update(claim)
+
+            await self._session.commit()
+
             return updated_claim
 
         except Exception:

app/repositories/implementations/source_repository.py

@@ -43,10 +43,19 @@ async def create_with_domain(self, source: SourceModel) -> Optional[SourceModel]
         """Create a source with its domain relationship."""
         try:
             self._session.add(source)
-            await self._session.flush()
-            await self._session.refresh(source, ["domain"])
             await self._session.commit()
-            return source
+            
+            # --- THE FIX ---
+            # Re-fetch to guarantee it is fully loaded and not expired by the commit
+            stmt = select(self._model_class).where(self._model_class.id == source.id).options(selectinload(self._model_class.domain))
+            result = await self._session.execute(stmt)
+            loaded_source = result.scalar_one()
+            
+            self._session.expunge(loaded_source)
+            if loaded_source.domain:
+                self._session.expunge(loaded_source.domain)
+                
+            return loaded_source
         except Exception as e:
             await self._session.rollback()
             raise e
@@ -56,7 +65,18 @@ async def update(self, source: SourceModel) -> SourceModel:
         try:
             merged = await self._session.merge(source)
             await self._session.commit()
-            return merged
+            
+            # --- THE FIX ---
+            # Same treatment: refetch eagerly, then detach
+            stmt = select(self._model_class).where(self._model_class.id == merged.id).options(selectinload(self._model_class.domain))
+            result = await self._session.execute(stmt)
+            loaded_source = result.scalar_one()
+            
+            self._session.expunge(loaded_source)
+            if loaded_source.domain:
+                self._session.expunge(loaded_source.domain)
+                
+            return loaded_source
         except Exception as e:
             await self._session.rollback()
             raise e

app/services/analysis_orchestrator.py

@@ -104,13 +104,17 @@ async def _generate_analysis(
             )
             current_analysis = await self._analysis_repo.create(initial_analysis)
 
+            await self._analysis_repo._session.commit()
+
             yield {"type": "status", "content": "Searching for relevant sources..."}
 
             query = self._query_initial(claim_text, language)
             messages = [LLMMessage(role="user", content=query)]
             all_sources = []
             for turns in range(MAX_NUM_TURNS):
 
+                await self._analysis_repo._session.rollback()
+
                 response = await self._llm.generate_response(messages)
 
                 main_agent_message = response.text
@@ -134,12 +138,17 @@ async def _generate_analysis(
                         updated_at=datetime.now(UTC),
                     )
                     current_search = await self._search_repo.create(initial_search)
+
+                    await self._search_repo._session.commit()
+
                     sources = await self._web_search.search_and_create_sources(
                         claim_text=search_request_match.matched_content, search_id=current_search.id, language=language
                     )
 
                     all_sources += sources
 
+                    await self._analysis_repo._session.rollback()
+
                     search_response = self._web_search.format_sources_for_prompt(sources, language)
 
                     if language == "english":
@@ -196,6 +205,8 @@ async def _generate_analysis(
             analysis_text = []
             log_probs = []
 
+            await self._analysis_repo._session.rollback()
+
             async for chunk in self._llm.generate_stream(messages):
                 if not chunk.is_complete:
                     analysis_text.append(chunk.text)
@@ -244,6 +255,8 @@ async def _generate_analysis(
 
                         if not default:
 
+                            await self._analysis_repo._session.rollback()
+
                             con_score = await self._generate_logprob_confidence_score(log_probs=log_probs)
                             logger.info(con_score)
                             current_analysis.confidence_score = float(con_score)
@@ -515,6 +528,7 @@ async def analyze_claim_stream(
             self._analysis_state.current_claim = claim
 
             await self._claim_repo.update_status(claim.id, ClaimStatus.analyzing)
+
             yield {"type": "status", "content": "Starting analysis..."}
 
             # Generate analysis

app/services/implementations/serper_web_search_service.py

@@ -63,7 +63,8 @@ async def search_and_create_sources(
                         except Exception as e:
                             logger.error(f"Error processing search result: {str(e)}", exc_info=True)
                             continue
-
+            await self.source_repository._session.commit()
+            
             return sources
 
         except Exception as e:

infrastructure/terraform/main.tf

@@ -374,6 +374,12 @@ resource "kubernetes_service" "misinformation_mitigation_api" {
   metadata {
     name      = "misinformation-mitigation-api"
     namespace = kubernetes_namespace.misinformation_mitigation.metadata[0].name
+
+    annotations = {
+      "cloud.google.com/backend-config" = jsonencode({
+        "default" = "api-timeout-config"
+      })
+    }
   }
 
   spec {
@@ -391,6 +397,20 @@ resource "kubernetes_service" "misinformation_mitigation_api" {
 
 
 # Load Balancer and DNS
+resource "kubernetes_manifest" "api_backend_config" {
+  manifest = {
+    apiVersion = "cloud.google.com/v1"
+    kind       = "BackendConfig"
+    metadata = {
+      name      = "api-timeout-config"
+      namespace = kubernetes_namespace.misinformation_mitigation.metadata[0].name
+    }
+    spec = {
+      timeoutSec = 600  # 600 seconds = 10 minutes
+    }
+  }
+}
+
 resource "google_compute_global_address" "misinformation_mitigation_api_ip" {
   name = "misinformation-mitigation-api-ip"
 }