ROB-1751 Strip embedded tool-call markers from the answer text only

src/robusta/core/playbooks/internal/ai_integration.py

@@ -352,9 +352,8 @@ def holmes_chat(event: ExecutionBaseEvent, params: HolmesChatParams):
                     if tool.tool_name not in GRAPH_TOOLS:
                         continue
 
-                    # removes all embedded tool calls in this case as its not supported. 
-                    holmes_result.analysis = re.sub(
-                        r"<<.*?>>", "", holmes_result.analysis
+                    holmes_result.analysis = strip_embedded_tool_calls(
+                        holmes_result.analysis
                     ).strip()
                     json_content = json.loads(tool.result["data"])
                     contents, name = get_png_from_graph_tool(json_content)
@@ -432,6 +431,32 @@ def holmes_oauth(event: ExecutionBaseEvent, params: HolmesOAuthParams):
         handle_holmes_error(e)
 
 
+# Embedded tool calls (<< {"type": "promql", ...} >>) are not supported when
+# graphs are rendered as images. Only the marker shape is matched: a bare `<<`
+# (a shell/Terraform heredoc, a C++ operator) must survive untouched, and a
+# match must not run from one marker across another (ROB-1751).
+EMBEDDED_TOOL_CALL_RE = re.compile(r"<<\s*\{(?:(?!<<|>>).)*?\}\s*>>", re.DOTALL)
+ANSWER_TEXT_FIELDS = ("analysis", "content")
+
+
+def strip_embedded_tool_calls(text: str) -> str:
+    return EMBEDDED_TOOL_CALL_RE.sub("", text)
+
+
+def strip_embedded_tool_calls_from_answer(
+    event_type: str, data_line: str, raw_event: str
+) -> str:
+    """Strip markers from the answer text fields only. Running the regex over
+    the raw SSE line also eats conversation_history and breaks the JSON."""
+    data = parse_sse_data(data_line)
+    if not isinstance(data, dict):
+        return raw_event
+    for field in ANSWER_TEXT_FIELDS:
+        if isinstance(data.get(field), str):
+            data[field] = strip_embedded_tool_calls(data[field])
+    return create_sse_message(event_type, data)
+
+
 def stream_and_render_graphs(url, holmes_req, event):
     with requests.post(
         url,
@@ -451,9 +476,10 @@ def stream_and_render_graphs(url, holmes_req, event):
                     continue
 
                 event_type = parse_sse_event_type(event_lines[0])
-                # removes all embedded tool calls in this case as its not supported. 
                 if event_type == StreamEvents.ANSWER_END.value:
-                    stream_event = re.sub(r"<<.*?>>", "", stream_event)
+                    stream_event = strip_embedded_tool_calls_from_answer(
+                        event_type, event_lines[1], stream_event
+                    )
 
                 if event_type != StreamEvents.TOOL_RESULT.value:
                     event.ws(data=stream_event)

tests/test_ai_integration.py

@@ -259,3 +259,30 @@ def test_get_png_from_graph_tool_output_types(output_type):
 
     assert name == "node_memory_used"
     assert_valid_png(contents)
+
+
+@patch("robusta.core.playbooks.internal.ai_integration.requests.post")
+def test_holmes_chat_streaming_keeps_unmatched_heredoc_in_answer(
+    mock_post, mock_event, holmes_chat_params
+):
+    """ROB-1751: an unmatched `<<` in the answer must not eat the rest of the
+    event (up to the next `>>` in conversation_history) and corrupt the JSON."""
+    history = [
+        {"role": "system", "content": 'ids look like << {"type": "issue"} >>'},
+        {"role": "assistant", "content": "answer"},
+    ]
+    answer = 'set message = <<EOT\nhello\nEOT or << {\nx\n} then << {"type": "promql", "tool_call_id": "x"} >> done'
+    sse_events = [
+        create_sse_message(
+            StreamEvents.ANSWER_END.value,
+            {"analysis": answer, "conversation_history": history},
+        ),
+    ]
+    mock_post.return_value = MockResponse(sse_events)
+
+    holmes_chat(mock_event, holmes_chat_params)
+
+    assert mock_event.ws.call_count == 1
+    _, data = parse_sse_message(mock_event.ws.call_args_list[0][1]["data"])
+    assert data["analysis"] == "set message = <<EOT\nhello\nEOT or << {\nx\n} then  done"
+    assert data["conversation_history"] == history