fix agent stream audit lifecycle

SovereignAI-Starter/app/main.py

@@ -186,12 +186,8 @@ async def audit_agent_routes(request: Request, call_next: Any):
     event_id = str(uuid4())
     started = perf_counter()
     status_code = 500
-    try:
-        response = await call_next(request)
-        status_code = response.status_code
-        response.headers["X-Agent-Audit-ID"] = event_id
-        return response
-    finally:
+
+    def record() -> None:
         try:
             database.record_agent_audit_event(
                 event_id=event_id,
@@ -204,6 +200,43 @@ async def audit_agent_routes(request: Request, call_next: Any):
         except Exception:
             logger.exception("Unable to write agent audit metadata")
 
+    try:
+        response = await call_next(request)
+        status_code = response.status_code
+        response.headers["X-Agent-Audit-ID"] = event_id
+
+        body_iterator = getattr(response, "body_iterator", None)
+        if body_iterator is None:
+            record()
+            return response
+
+        async def audited_body():
+            nonlocal status_code
+            try:
+                async for chunk in body_iterator:
+                    yield chunk
+            except (asyncio.CancelledError, GeneratorExit):
+                # 499 is a diagnostic marker for a client-closed stream; it is
+                # not an HTTP status sent to the client after headers were sent.
+                status_code = 499
+                raise
+            except Exception:
+                status_code = 500
+                raise
+            finally:
+                record()
+
+        response.body_iterator = audited_body()
+        return response
+    except asyncio.CancelledError:
+        status_code = 499
+        record()
+        raise
+    except Exception:
+        status_code = 500
+        record()
+        raise
+
 
 class Message(BaseModel):
     role: str = Field(description="system أو user أو assistant")