Streaming bugs

.changeset/heavy-sloths-start.md

@@ -0,0 +1,5 @@
+---
+"apollo": patch
+---
+
+Fix issue where chat is returned before content in streamed responses

services/job_chat/job_chat.py

@@ -275,7 +275,7 @@ def generate(
                 job_code = None
                 if context and isinstance(context, dict):
                     job_code = context.get("expression")
-                
+
                 with sentry_sdk.start_span(description="parse_and_apply_edits"):
                     text_response, suggested_code, diff = self.parse_and_apply_edits(response=response, content=content, original_code=job_code)
 
@@ -331,21 +331,18 @@ def process_stream_event(
                 accumulated_response += text_chunk
 
                 if suggest_code and not text_started:
-                    # Code edits phase: buffer silently until text_answer starts
-                    delimiter = '"text_answer": "'
+                    # Code edits phase: buffer silently until text_answer starts.
+                    # Tolerant of whitespace variants the model may emit.
+                    match = re.search(r'"text_answer"\s*:\s*"', accumulated_response)
 
-                    if delimiter in accumulated_response:
+                    if match:
                         # Extract code_edits from the JSON before the delimiter
-                        edits_part = accumulated_response.split(delimiter)[0]
+                        edits_part = accumulated_response[:match.start()]
                         # Find the code_edits array value
                         try:
-                            # Build partial JSON to extract code_edits
-                            partial = edits_part.rstrip().rstrip(",")
-                            # Close the partial JSON to make it parseable
-                            if not partial.rstrip().endswith("}"):
-                                partial = partial + "}"
-                            partial_data = json.loads(partial)
-                            code_edits = partial_data.get("code_edits", [])
+                            # Close the partial object and extract code_edits
+                            partial = edits_part.rstrip().rstrip(",") + "}"
+                            code_edits = json.loads(partial).get("code_edits", [])
 
                             if original_code and code_edits:
                                 suggested_code, diff = self.apply_code_edits(
@@ -364,8 +361,7 @@ def process_stream_event(
                             logger.warning(f"Failed to parse code_edits during streaming")
 
                         # Mark where text content starts in the accumulated buffer
-                        text_offset = accumulated_response.find(delimiter) + len(delimiter)
-                        sent_length = text_offset
+                        sent_length = match.end()
                         text_started = True
 
                 if suggest_code and text_started:

services/workflow_chat/workflow_chat.py

@@ -545,18 +545,16 @@ def process_stream_event(self, event, accumulated_response, text_started, sent_l
 
                 if not text_started:
                     # YAML phase: buffer silently until text starts.
-                    # Use '"text": "' as delimiter — safe because inside
-                    # a JSON string quotes are escaped as \", so this can
-                    # only appear as the actual field separator.
-                    delimiter = '"text": "'
-
-                    if delimiter in accumulated_response:
-                        # Extract YAML value (either null or a JSON string)
-                        yaml_part = accumulated_response.split(delimiter)[0]
-                        yaml_raw = yaml_part.strip().rstrip(",").strip()
+                    # Tolerant of whitespace variants the model may emit.
+                    match = re.search(r'"text"\s*:\s*"', accumulated_response)
+
+                    if match:
+                        # Close the partial object and extract the yaml field
+                        yaml_part = accumulated_response[:match.start()]
+                        yaml_raw = yaml_part.rstrip().rstrip(",") + "}"
                         try:
-                            yaml_value = json.loads(yaml_raw)  # None for null, string for "..."
-                        except (json.JSONDecodeError, ValueError):
+                            yaml_value = json.loads(yaml_raw).get("yaml")
+                        except (json.JSONDecodeError, ValueError, AttributeError):
                             yaml_value = None
 
                         if yaml_value:
@@ -569,8 +567,7 @@ def process_stream_event(self, event, accumulated_response, text_started, sent_l
                                 pass  # Invalid YAML, skip changes event
 
                         # Mark where text content starts
-                        text_offset = accumulated_response.find(delimiter) + len(delimiter)
-                        sent_length = text_offset
+                        sent_length = match.end()
                         text_started = True
 
                 if text_started: