Added in_sage command and removed texSagetexArg from mathzone

autoload/vimtex/syntax.vim

@@ -24,12 +24,17 @@ function! vimtex#syntax#in_comment(...) abort " {{{1
   return call('vimtex#syntax#in', ['texComment'] + a:000)
 endfunction
 
+" }}}1
+function! vimtex#syntax#in_sage(...) abort " {{{1
+  return call('vimtex#syntax#in', ['texSagetexArg'] + a:000) || call('vimtex#syntax#in', ['texSagetexZone'] + a:000)
+endfunction
+
 " }}}1
 function! vimtex#syntax#in_mathzone(...) abort " {{{1
   " The following checks if we are inside a texMathZone environment. The
-  " arguments to \label{...}, the texRefArg group, and \text{...} like
-  " commands, the texMathTextArg group, are actively ignored as these should
-  " not be considered to be math environments.
+  " arguments to \label{...}, the texRefArg group, the texSagetexArg group,
+  " and \text{...} like commands, the texMathTextArg group, are actively
+  " ignored as these should not be considered to be math environments.
   let l:groups = reverse(call('vimtex#syntax#stack', a:000))
   let l:group = matchstr(l:groups, '\v^tex%(Math%(Zone|Text|Tag)|RefArg|SagetexArg)')
   return l:group =~# '^texMathZone'

autoload/vimtex/syntax/p/sagetex.vim

@@ -7,8 +7,6 @@
 function! vimtex#syntax#p#sagetex#load(cfg) abort " {{{1
   call vimtex#syntax#nested#include('python')
 
-  syntax match texCmdSagetex /\\sagestr\>/
-        \ nextgroup=texSagetexArg skipwhite skipnl
   syntax match texCmdSagetex /\\sageplot\>/
         \ nextgroup=texSagetexOpt,texSagetexArg skipwhite skipnl
 
@@ -36,6 +34,8 @@ function! vimtex#syntax#p#sagetex#load(cfg) abort " {{{1
         \ 'sage',
         \ 'sagestr'
         \]
+    execute 'syntax match texCmdSagetex /\\' . l:env .
+          \ '\>/ nextgroup=texSagetexArg skipwhite skipnl'
     call vimtex#syntax#core#new_cmd({
           \ 'name': l:env,
           \ 'mathmode': 1,