fix(ci): stabilize exec streaming behavior

Author: usnavy13

src/api/exec.py

@@ -15,7 +15,13 @@
 from fastapi.responses import StreamingResponse
 
 from ..models import ExecRequest, ExecResponse
-from ..models.errors import ErrorResponse, ValidationError, ServiceUnavailableError
+from ..models.errors import (
+    CodeInterpreterException,
+    ErrorResponse,
+    ErrorType,
+    ValidationError,
+    ServiceUnavailableError,
+)
 from ..services.orchestrator import ExecutionOrchestrator
 from ..dependencies.services import (
     SessionServiceDep,
@@ -106,6 +112,14 @@ async def _execute() -> ExecResponse:
         )
         return response
     except asyncio.TimeoutError:
+        if execution_task.done():
+            response = await execution_task
+            logger.info(
+                "Code execution completed",
+                request_id=request_id,
+                session_id=response.session_id,
+            )
+            return response
         # Fall through to streamed keepalives for genuinely long-running work.
         pass
     except (ValidationError, ServiceUnavailableError):
@@ -119,6 +133,12 @@ async def _stream_response():
         whitespace is ignored by JSON parsers, so this is transparent
         to clients.
         """
+        # The endpoint already spent one interval deciding whether to switch to
+        # streaming. Emit a first keepalive immediately so long-running
+        # requests stay under client-side socket timeout thresholds.
+        if not execution_task.done():
+            yield b" "
+
         # Send keepalive spaces while execution is running
         while not execution_task.done():
             try:
@@ -127,17 +147,24 @@ async def _stream_response():
                 )
             except asyncio.TimeoutError:
                 # Execution still running — send keepalive space
-                yield b" "
+                if not execution_task.done():
+                    yield b" "
+            except Exception:
+                # Task raised an exception — it will be handled below.
+                break
 
         # Ensure the task is complete
         try:
             response = await execution_task
         except Exception as err:
             # Once the streaming response has started, surface failures as a JSON
             # error payload instead of raising after headers have been sent.
+            error_type = ErrorType.INTERNAL_SERVER
+            if isinstance(err, CodeInterpreterException):
+                error_type = err.error_type
             error_resp = ErrorResponse(
                 error=str(err),
-                error_type="execution",
+                error_type=error_type,
             )
             yield error_resp.model_dump_json().encode()
             return

tests/integration/test_auth_integration.py

@@ -217,6 +217,15 @@ def test_file_upload_flow_without_auth(self, client, mock_services):
 class TestAuthenticationEdgeCases:
     """Test edge cases in authentication."""
 
+    @staticmethod
+    def _protected_exec_request(client, headers):
+        """Hit a real authenticated endpoint with a minimal valid payload."""
+        return client.post(
+            "/exec",
+            headers=headers,
+            json={"code": "print('auth edge')", "lang": "py"},
+        )
+
     def test_auth_with_special_characters_in_key(self, client, mock_services):
         """Test authentication with special characters in API key."""
         special_key = "test-key-with-special-chars!@#$%^&*()"
@@ -225,10 +234,10 @@ def test_auth_with_special_characters_in_key(self, client, mock_services):
             mock_settings.api_key = special_key
             headers = {"x-api-key": special_key}
 
-            response = client.get("/sessions", headers=headers)
+            response = self._protected_exec_request(client, headers)
 
             # Should handle special characters correctly
-            # If 401, it means auth failed, but we want to ensure no 500 error
+            # If 401, auth rejected the key. If 200, auth accepted it.
             assert response.status_code in [200, 401]
 
     def test_auth_with_very_long_key(self, client, mock_services):
@@ -239,7 +248,7 @@ def test_auth_with_very_long_key(self, client, mock_services):
             mock_settings.api_key = long_key
             headers = {"x-api-key": long_key}
 
-            response = client.get("/sessions", headers=headers)
+            response = self._protected_exec_request(client, headers)
 
             # Should handle long keys (within reason)
             assert response.status_code in [200, 401]

tests/integration/test_exec_api.py

@@ -1,5 +1,6 @@
 """Integration tests for the /exec endpoint."""
 
+import asyncio
 import pytest
 from fastapi.testclient import TestClient
 from unittest.mock import AsyncMock, patch, MagicMock
@@ -8,7 +9,13 @@
 from datetime import datetime, timezone, timedelta
 
 from src.main import app
-from src.models import CodeExecution, ExecutionStatus, ExecutionOutput, OutputType
+from src.models import (
+    CodeExecution,
+    ExecutionStatus,
+    ExecutionOutput,
+    OutputType,
+    ServiceUnavailableError,
+)
 
 
 @pytest.fixture
@@ -381,6 +388,33 @@ def test_exec_service_error(self, client, auth_headers, mock_execution_service):
         assert response.status_code == 503
         assert "error" in response.json()
 
+    def test_exec_delayed_service_error_after_stream_start(
+        self, client, auth_headers
+    ):
+        """Delayed failures should return a JSON error payload, not crash the stream."""
+
+        async def _delayed_failure(*args, **kwargs):
+            await asyncio.sleep(3.2)
+            raise ServiceUnavailableError(
+                service="Code Execution",
+                message="Delayed backend failure",
+            )
+
+        with patch(
+            "src.services.orchestrator.ExecutionOrchestrator.execute",
+            side_effect=_delayed_failure,
+        ):
+            response = client.post(
+                "/exec",
+                json={"code": "print('Hello')", "lang": "py"},
+                headers=auth_headers,
+            )
+
+        assert response.status_code == 200
+        response_data = json.loads(response.text.lstrip())
+        assert response_data["error_type"] == "service_unavailable"
+        assert "Delayed backend failure" in response_data["error"]
+
     def test_exec_response_format_compatibility(
         self, client, auth_headers, mock_execution_service
     ):