zsh fixes (zmodload, k8s, new var type, track zsh prompt)

waveshell/pkg/shellapi/bashapi.go

@@ -113,6 +113,9 @@ func (b bashShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 		if varDecl.IsExport() || varDecl.IsReadOnly() {
 			continue
 		}
+		if varDecl.IsExtVar {
+			continue
+		}
 		rcBuf.WriteString(BashDeclareStmt(varDecl))
 		rcBuf.WriteString("\n")
 	}

waveshell/pkg/shellapi/bashparser.go

@@ -205,7 +205,7 @@ func bashParseDeclareOutput(state *packet.ShellState, declareBytes []byte, pvarB
 			declMap[decl.Name] = decl
 		}
 	}
-	pvarMap := parsePVarOutput(pvarBytes, false)
+	pvarMap := parseExtVarOutput(pvarBytes, "", "")
 	utilfn.CombineMaps(declMap, pvarMap)
 	state.ShellVars = shellenv.SerializeDeclMap(declMap) // this writes out the decls in a canonical order
 	if firstParseErr != nil {

waveshell/pkg/shellapi/shellapi.go

@@ -25,10 +25,13 @@ import (
 	"github.com/wavetermdev/waveterm/waveshell/pkg/base"
 	"github.com/wavetermdev/waveterm/waveshell/pkg/packet"
 	"github.com/wavetermdev/waveterm/waveshell/pkg/shellutil"
+	"github.com/wavetermdev/waveterm/waveshell/pkg/utilfn"
 )
 
 const GetStateTimeout = 5 * time.Second
 const GetGitBranchCmdStr = `printf "GITBRANCH %s\x00" "$(git rev-parse --abbrev-ref HEAD 2>/dev/null)"`
+const GetK8sContextCmdStr = `printf "K8SCONTEXT %s\x00" "$(kubectl config current-context 2>/dev/null)"`
+const GetK8sNamespaceCmdStr = `printf "K8SNAMESPACE %s\x00" "$(kubectl config view --minify --output 'jsonpath={..namespace}' 2>/dev/null)"`
 const RunCommandFmt = `%s`
 const DebugState = false
 
@@ -239,7 +242,7 @@ func RunSimpleCmdInPty(ecmd *exec.Cmd) ([]byte, error) {
 	return outputBuf.Bytes(), nil
 }
 
-func parsePVarOutput(pvarBytes []byte, isZsh bool) map[string]*DeclareDeclType {
+func parseExtVarOutput(pvarBytes []byte, promptOutput string, zmodsOutput string) map[string]*DeclareDeclType {
 	declMap := make(map[string]*DeclareDeclType)
 	pvars := bytes.Split(pvarBytes, []byte{0})
 	for _, pvarBA := range pvars {
@@ -251,11 +254,35 @@ func parsePVarOutput(pvarBytes []byte, isZsh bool) map[string]*DeclareDeclType {
 		if pvarFields[0] == "" {
 			continue
 		}
-		decl := &DeclareDeclType{IsZshDecl: isZsh, Args: "x"}
+		if pvarFields[1] == "" {
+			continue
+		}
+		decl := &DeclareDeclType{IsExtVar: true}
 		decl.Name = "PROMPTVAR_" + pvarFields[0]
 		decl.Value = shellescape.Quote(pvarFields[1])
 		declMap[decl.Name] = decl
 	}
+	if promptOutput != "" {
+		decl := &DeclareDeclType{IsExtVar: true}
+		decl.Name = "PROMPTVAR_PS1"
+		decl.Value = promptOutput
+		declMap[decl.Name] = decl
+	}
+	if zmodsOutput != "" {
+		var zmods []string
+		lines := strings.Split(zmodsOutput, "\n")
+		for _, line := range lines {
+			fields := strings.Fields(line)
+			if len(fields) != 2 || fields[0] != "zmodload" {
+				continue
+			}
+			zmods = append(zmods, fields[1])
+		}
+		decl := &DeclareDeclType{IsExtVar: true}
+		decl.Name = ZModsVarName
+		decl.Value = utilfn.QuickJson(zmods)
+		declMap[decl.Name] = decl
+	}
 	return declMap
 }
 

waveshell/pkg/shellapi/zshapi.go

@@ -30,6 +30,21 @@ const BaseZshOpts = ``
 const ZshShellVersionCmdStr = `echo zsh v$ZSH_VERSION`
 const StateOutputFdNum = 20
 
+const (
+	ZshSection_Version = iota
+	ZshSection_Cwd
+	ZshSection_Env
+	ZshSection_Mods
+	ZshSection_Vars
+	ZshSection_Aliases
+	ZshSection_Fpath
+	ZshSection_Funcs
+	ZshSection_PVars
+	ZshSection_Prompt
+
+	ZshSection_NumFieldsExpected // must be last
+)
+
 // TODO these need updating
 const RunZshSudoCommandFmt = `sudo -n -C %d zsh /dev/fd/%d`
 const RunZshSudoPasswordCommandFmt = `cat /dev/fd/%d | sudo -k -S -C %d zsh -c "echo '[from-mshell]'; exec %d>&-; zsh /dev/fd/%d < /dev/fd/%d"`
@@ -54,6 +69,7 @@ var ZshIgnoreVars = map[string]bool{
 	"SHLVL":                true,
 	"TTY":                  true,
 	"ZDOTDIR":              true,
+	"PPID":                 true,
 	"epochtime":            true,
 	"langinfo":             true,
 	"keymaps":              true,
@@ -77,6 +93,8 @@ var ZshIgnoreVars = map[string]bool{
 	"funcsourcetrace":      true,
 	"funcstack":            true,
 	"functrace":            true,
+	"nameddirs":            true,
+	"userdirs":             true,
 	"parameters":           true,
 	"commands":             true,
 	"functions":            true,
@@ -86,6 +104,25 @@ var ZshIgnoreVars = map[string]bool{
 	"_comps":               true,
 	"_patcomps":            true,
 	"_postpatcomps":        true,
+
+	// zsh/system
+	"errnos":    true,
+	"sysparams": true,
+
+	// zsh/curses
+	"ZCURSES_COLORS":      true,
+	"ZCURSES_COLOR_PAIRS": true,
+	"zcurses_attrs":       true,
+	"zcurses_colors":      true,
+	"zcurses_keycodes":    true,
+	"zcurses_windows":     true,
+
+	// not listed, but we also exclude all ZFTP_* variables
+}
+
+var ZshIgnoreFuncs = map[string]bool{
+	"zftp_chpwd":    true,
+	"zftp_progress": true,
 }
 
 // only options we restore (other than ZshForceOptions)
@@ -131,11 +168,13 @@ var ZshUnsetVars = []string{
 	"ZSH_EXECUTION_STRING",
 }
 
-var ZshLoadMods = []string{
-	"zsh/parameter",
-	"zsh/langinfo",
+var ZshForceLoadMods = map[string]bool{
+	"zsh/parameter": true,
+	"zsh/langinfo":  true,
 }
 
+const ZModsVarName = "WAVESTATE_ZMODS"
+
 // do not use these directly, call GetLocalMajorVersion()
 var localZshMajorVersionOnce = &sync.Once{}
 var localZshMajorVersion = ""
@@ -273,14 +312,29 @@ func (z zshShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 			rcBuf.WriteString(fmt.Sprintf("unsetopt %s\n", optName))
 		}
 	}
-	for _, modName := range ZshLoadMods {
+	for modName := range ZshForceLoadMods {
 		rcBuf.WriteString(fmt.Sprintf("zmodload %s\n", modName))
 	}
+	modDecl := getDeclByName(varDecls, ZModsVarName)
+	if modDecl != nil {
+		modsArr := utilfn.QuickParseJson[[]string](modDecl.Value)
+		for _, modName := range modsArr {
+			if !ZshForceLoadMods[modName] {
+				rcBuf.WriteString(fmt.Sprintf("zmodload %s\n", modName))
+			}
+		}
+	}
 	var postDecls []*shellenv.DeclareDeclType
 	for _, varDecl := range varDecls {
 		if ZshIgnoreVars[varDecl.Name] {
 			continue
 		}
+		if strings.HasPrefix(varDecl.Name, "ZFTP_") {
+			continue
+		}
+		if varDecl.IsExtVar {
+			continue
+		}
 		if ZshUniqueArrayVars[varDecl.Name] && !varDecl.IsUniqueArray() {
 			varDecl.AddFlag("U")
 		}
@@ -332,6 +386,9 @@ func (z zshShellApi) MakeRcFileStr(pk *packet.RunPacketType) string {
 		rcBuf.WriteString("# error decoding zsh functions\n")
 	} else {
 		for fnKey, fnValue := range fnMap {
+			if ZshIgnoreFuncs[fnKey.ParamName] {
+				continue
+			}
 			if fnValue == ZshFnAutoLoad {
 				rcBuf.WriteString(fmt.Sprintf("autoload %s\n", shellescape.Quote(fnKey.ParamName)))
 			} else {
@@ -411,6 +468,8 @@ pwd;
 printf "[%SECTIONSEP%]";
 env -0;
 printf "[%SECTIONSEP%]";
+zmodload -L
+printf "[%SECTIONSEP%]";
 typeset -p +H -m '*';
 printf "[%SECTIONSEP%]";
 for var in "${(@k)aliases}"; do
@@ -448,10 +507,16 @@ for var in "${(@k)dis_functions_source}"; do
 done
 printf "[%SECTIONSEP%]";
 [%GITBRANCH%]
+[%K8SCONTEXT%]
+[%K8SNAMESPACE%]
+printf "[%SECTIONSEP%]";
+print -P "$PS1"
 `
 	cmd = strings.TrimSpace(cmd)
 	cmd = strings.ReplaceAll(cmd, "[%ZSHVERSION%]", ZshShellVersionCmdStr)
 	cmd = strings.ReplaceAll(cmd, "[%GITBRANCH%]", GetGitBranchCmdStr)
+	cmd = strings.ReplaceAll(cmd, "[%K8SCONTEXT%]", GetK8sContextCmdStr)
+	cmd = strings.ReplaceAll(cmd, "[%K8SNAMESPACE%]", GetK8sNamespaceCmdStr)
 	cmd = strings.ReplaceAll(cmd, "[%PARTSEP%]", utilfn.ShellHexEscape(string(sectionSeparator[0:len(sectionSeparator)-1])))
 	cmd = strings.ReplaceAll(cmd, "[%SECTIONSEP%]", utilfn.ShellHexEscape(string(sectionSeparator)))
 	cmd = strings.ReplaceAll(cmd, "[%OUTPUTFD%]", fmt.Sprintf("/dev/fd/%d", fdNum))
@@ -599,6 +664,9 @@ func ParseZshFunctions(fpathArr []string, fnBytes []byte, partSeparator []byte)
 		if fnName == "zshexit" {
 			continue
 		}
+		if ZshIgnoreFuncs[fnName] {
+			continue
+		}
 		if fnType == "functions" || fnType == "dis_functions" {
 			fnBody[ZshParamKey{ParamType: fnType, ParamName: fnName}] = fnValue
 		}
@@ -609,10 +677,13 @@ func ParseZshFunctions(fpathArr []string, fnBytes []byte, partSeparator []byte)
 	// ok, so the trick here is that we want to only include functions that are *not* autoloaded
 	// the ones that are pending autoloading or come from a source file in fpath, can just be set to autoload
 	for fnKey := range fnBody {
+		var inFpath bool
 		source := fnSource[fnKey.ParamName]
-		if isSourceFileInFpath(fpathArr, source) {
-			fnBody[fnKey] = ZshFnAutoLoad
-		} else if strings.TrimSpace(fnBody[fnKey]) == ZshAutoloadFnBody {
+		if source != "" {
+			inFpath = isSourceFileInFpath(fpathArr, source)
+		}
+		isAutoloadFnBody := strings.TrimSpace(fnBody[fnKey]) == ZshAutoloadFnBody
+		if inFpath || isAutoloadFnBody {
 			fnBody[fnKey] = ZshFnAutoLoad
 		}
 	}
@@ -639,11 +710,11 @@ func (z zshShellApi) ParseShellStateOutput(outputBytes []byte) (*packet.ShellSta
 	versionStr := string(outputBytes[0:firstZeroIdx])
 	sectionSeparator := outputBytes[firstZeroIdx+1 : firstDZeroIdx+2]
 	partSeparator := sectionSeparator[0 : len(sectionSeparator)-1]
-	// 8 fields: version [0], cwd [1], env [2], vars [3], aliases [4], fpath [5], functions [6], pvars [7]
-	fields := bytes.Split(outputBytes, sectionSeparator)
-	if len(fields) != 8 {
+	// sections: see ZshSection_* consts
+	sections := bytes.Split(outputBytes, sectionSeparator)
+	if len(sections) != ZshSection_NumFieldsExpected {
 		base.Logf("invalid -- numfields\n")
-		return nil, fmt.Errorf("invalid zsh shell state output, wrong number of fields, fields=%d", len(fields))
+		return nil, fmt.Errorf("invalid zsh shell state output, wrong number of sections, section=%d", len(sections))
 	}
 	rtn := &packet.ShellState{}
 	rtn.Version = strings.TrimSpace(versionStr)
@@ -653,10 +724,10 @@ func (z zshShellApi) ParseShellStateOutput(outputBytes []byte) (*packet.ShellSta
 	if _, _, err := packet.ParseShellStateVersion(rtn.Version); err != nil {
 		return nil, fmt.Errorf("invalid zsh shell state output, invalid version: %v", err)
 	}
-	cwdStr := stripNewLineChars(string(fields[1]))
+	cwdStr := stripNewLineChars(string(sections[ZshSection_Cwd]))
 	rtn.Cwd = cwdStr
-	zshEnv := parseZshEnv(fields[2])
-	zshDecls, err := parseZshDecls(fields[3])
+	zshEnv := parseZshEnv(sections[ZshSection_Env])
+	zshDecls, err := parseZshDecls(sections[ZshSection_Vars])
 	if err != nil {
 		base.Logf("invalid - parsedecls %v\n", err)
 		return nil, err
@@ -666,16 +737,15 @@ func (z zshShellApi) ParseShellStateOutput(outputBytes []byte) (*packet.ShellSta
 			decl.ZshEnvValue = zshEnv[decl.ZshBoundScalar]
 		}
 	}
-	aliasMap := parseZshAliasStateOutput(fields[4], partSeparator)
+	aliasMap := parseZshAliasStateOutput(sections[ZshSection_Aliases], partSeparator)
 	rtn.Aliases = string(EncodeZshMap(aliasMap))
-	fpathStr := stripNewLineChars(string(string(fields[5])))
+	fpathStr := stripNewLineChars(string(string(sections[ZshSection_Fpath])))
 	fpathArr := strings.Split(fpathStr, ":")
-	zshFuncs := ParseZshFunctions(fpathArr, fields[6], partSeparator)
+	zshFuncs := ParseZshFunctions(fpathArr, sections[ZshSection_Funcs], partSeparator)
 	rtn.Funcs = string(EncodeZshMap(zshFuncs))
-	pvarMap := parsePVarOutput(fields[7], true)
+	pvarMap := parseExtVarOutput(sections[ZshSection_PVars], string(sections[ZshSection_Prompt]), string(sections[ZshSection_Mods]))
 	utilfn.CombineMaps(zshDecls, pvarMap)
 	rtn.ShellVars = shellenv.SerializeDeclMap(zshDecls)
-	base.Logf("parse shellstate done\n")
 	return rtn, nil
 }